diff --git a/.github/actions/check-rebuild/action.yml b/.github/actions/check-rebuild/action.yml index c9fc5ddd23f..2d3b6e2df41 100644 --- a/.github/actions/check-rebuild/action.yml +++ b/.github/actions/check-rebuild/action.yml @@ -2,8 +2,12 @@ name: "Build and verify rebuild is a no-op" description: "Run a cargo command, then re-run it to verify all artifacts are cached" inputs: command: - description: "The cargo build command to run and verify (--message-format json is appended on the verification run)" + description: "The cargo build command to run and verify" required: true + message-format-flag: + description: "The flag used to request Cargo compiler-artifact JSON during the verification run" + required: false + default: "--message-format" runs: using: "composite" steps: @@ -14,7 +18,7 @@ runs: shell: bash run: | stale=$(${{ inputs.command }} \ - --message-format json 2>/dev/null \ + ${{ inputs.message-format-flag }} json 2>/dev/null \ | jq -r 'select(.reason == "compiler-artifact" and .fresh == false) | .target.name') if [ -n "$stale" ]; then echo "ERROR: Rebuild recompiled crates that should have been cached:" diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index a6fe8732a0f..ddefd50673a 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,55 +1,23 @@ - - -## Summary +## Rationale for this change -Closes: #000 +- Closes: #000 +## What changes are included in this PR? -## Testing +## What APIs are changed? Are there any user-facing changes? diff --git a/.github/workflows/bench-dispatch.yml b/.github/workflows/bench-dispatch.yml index 48bf6fbe4af..d1116d964b5 100644 --- a/.github/workflows/bench-dispatch.yml +++ b/.github/workflows/bench-dispatch.yml @@ -47,3 +47,21 @@ jobs: needs: remove-sql-label uses: ./.github/workflows/sql-pr.yml secrets: inherit + + remove-sql-full-label: + runs-on: ubuntu-latest + timeout-minutes: 10 + if: github.event.label.name == 'action/benchmark-sql-full' + steps: + - uses: actions-ecosystem/action-remove-labels@2ce5d41b4b6aa8503e285553f75ed56e0a40bae0 # v1 + if: github.event.pull_request.head.repo.full_name == 'vortex-data/vortex' + with: + labels: action/benchmark-sql-full + fail_on_error: true + + sql-full-bench: + needs: remove-sql-full-label + uses: ./.github/workflows/sql-pr.yml + secrets: inherit + with: + benchmark_profile: "full" diff --git a/.github/workflows/bench-pr.yml b/.github/workflows/bench-pr.yml index cfab4f7fd3b..8223ea80a56 100644 --- a/.github/workflows/bench-pr.yml +++ b/.github/workflows/bench-pr.yml @@ -38,7 +38,7 @@ jobs: if: github.event.pull_request.head.repo.fork == false with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: ref: ${{ github.event.pull_request.head.sha }} - name: Setup benchmark environment @@ -72,7 +72,18 @@ jobs: profiling_frequency: 199 extra_args: "--off-cpu-threshold=0.03" # Personally tuned by @brancz + - name: Run ${{ matrix.benchmark.name }} benchmark (per-combination) + if: matrix.benchmark.id == 'random-access-bench' + shell: bash + env: + RUST_BACKTRACE: full + VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1" + FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" + run: | + python3 scripts/random-access-split.py + - name: Run ${{ matrix.benchmark.name }} benchmark + if: matrix.benchmark.id != 'random-access-bench' shell: bash env: RUST_BACKTRACE: full diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml index 839db948add..fc59a8156ea 100644 --- a/.github/workflows/bench.yml +++ b/.github/workflows/bench.yml @@ -16,7 +16,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 10 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: Setup AWS CLI uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6 with: @@ -54,7 +54,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: Setup benchmark environment run: sudo bash scripts/setup-benchmark.sh - uses: ./.github/actions/setup-rust @@ -85,7 +85,18 @@ jobs: profiling_frequency: 199 extra_args: "--off-cpu-threshold=0.03" # Personally tuned by @brancz + - name: Run ${{ matrix.benchmark.name }} benchmark (per-combination) + if: matrix.benchmark.id == 'random-access-bench' + shell: bash + env: + RUST_BACKTRACE: full + VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1" + FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" + run: | + python3 scripts/random-access-split.py --v3 + - name: Run ${{ matrix.benchmark.name }} benchmark + if: matrix.benchmark.id != 'random-access-bench' shell: bash env: RUST_BACKTRACE: full @@ -117,6 +128,53 @@ jobs: --benchmark-id "${{ matrix.benchmark.id }}" \ --repo-url "${{ github.server_url }}/${{ github.repository }}" + # v4 (Postgres) dual-write -- BEST-EFFORT during the migration soak. The + # proven v3 step above is hard-required; a v4 failure must NOT fail the job + # (v4 is promoted to required at cutover, PR-5.1). Gated on the ingest-role + # ARN var (the assume-role input that MUST exist for OIDC to succeed) so it + # no-ops until v4 infra is wired, and every step is + # `continue-on-error` so an OIDC / uv / connect hiccup never breaks the v3 + # pipeline. post-ingest.py mints the RDS IAM token internally (boto3) from + # the assumed GitHubBenchmarkIngestRole; sslmode=verify-full validates the cert. + # + # ORDER MATTERS: "Install uv" runs BEFORE "Configure AWS credentials". + # configure-aws-credentials persists the assumed ingest-role (rds-db:connect + # only) as the job's ambient AWS creds; the uv setup compiles via sccache + # (S3-backed), so running it after the role switch fails with S3 AccessDenied. + # Installing uv first keeps sccache on the original S3-capable creds; the role + # is assumed immediately before the ingest, which needs only rds-db:connect. + - name: Install uv for v4 ingest + if: vars.GH_BENCH_INGEST_ROLE_ARN != '' + continue-on-error: true + uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 + - name: Configure AWS credentials for v4 ingest (OIDC) + if: vars.GH_BENCH_INGEST_ROLE_ARN != '' + continue-on-error: true + uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6 + with: + role-to-assume: ${{ vars.GH_BENCH_INGEST_ROLE_ARN }} + aws-region: ${{ vars.RDS_BENCH_REGION }} + - name: Ingest results to v4 Postgres (best-effort) + if: vars.GH_BENCH_INGEST_ROLE_ARN != '' + continue-on-error: true + shell: bash + env: + RDS_BENCH_INSTANCE_ENDPOINT: ${{ vars.RDS_BENCH_INSTANCE_ENDPOINT }} + RDS_BENCH_DB_NAME: ${{ vars.RDS_BENCH_DB_NAME }} + AWS_REGION: ${{ vars.RDS_BENCH_REGION }} + BENCH_SITE_BASE_URL: ${{ vars.BENCH_SITE_BASE_URL }} + BENCH_REVALIDATE_TOKEN: ${{ secrets.BENCH_REVALIDATE_TOKEN }} + run: | + set -Eeuo pipefail + curl -fsSL https://truststore.pki.rds.amazonaws.com/global/global-bundle.pem \ + -o "${RUNNER_TEMP}/rds-global-bundle.pem" + DSN="postgresql://bench_ingest@${RDS_BENCH_INSTANCE_ENDPOINT}:5432/${RDS_BENCH_DB_NAME}?sslmode=verify-full&sslrootcert=${RUNNER_TEMP}/rds-global-bundle.pem" + uv run --no-project --with 'psycopg[binary]' --with boto3 --with xxhash \ + scripts/post-ingest.py results.v3.jsonl \ + --postgres "${DSN}" \ + --commit-sha "${{ github.sha }}" \ + --region "${AWS_REGION}" + - name: Alert incident.io if: failure() uses: ./.github/actions/alert-incident-io diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5798d1d56b0..44797183cee 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -25,18 +25,37 @@ env: NIGHTLY_TOOLCHAIN: nightly-2026-02-05 jobs: + duckdb-mirror: + name: "Mirror DuckDB to R2" + if: github.event_name == 'pull_request' + uses: ./.github/workflows/duckdb-r2.yml + secrets: inherit + + duckdb-ready: + name: "DuckDB libraries available in R2" + needs: duckdb-mirror + if: ${{ !cancelled() }} + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - name: Verify DuckDB mirror + if: ${{ needs.duckdb-mirror.result == 'failure' }} + run: | + echo "DuckDB mirror failed; downstream builds would 404" + exit 1 + lint-toml: runs-on: ubuntu-latest timeout-minutes: 10 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: spiraldb/actions/.github/actions/lint-toml@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 validate-workflow-yaml: runs-on: ubuntu-latest timeout-minutes: 10 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: Validate YAML file run: | # Lint the workflows and yamllint's configuration file. @@ -57,7 +76,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Install uv if: github.repository != 'vortex-data/vortex' @@ -93,7 +112,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Pytest - Vortex @@ -119,8 +138,42 @@ jobs: uv run --all-packages make html working-directory: docs/ + python-cuda-test: + name: "Python CUDA (test)" + if: github.repository == 'vortex-data/vortex' + runs-on: >- + ${{ format('runs-on={0}/runner=gpu/tag=python-cuda-test', github.run_id) }} + timeout-minutes: 30 + env: + RUST_LOG: "info,maturin=off,uv=debug" + MATURIN_PEP517_ARGS: "--profile ci" + steps: + - uses: runs-on/action@v2 + with: + sccache: s3 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + - uses: ./.github/actions/setup-rust + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + components: cargo + - name: Pin rustup proxy to repository toolchain + run: | + TOOLCHAIN="$(grep '^channel' rust-toolchain.toml | cut -d '"' -f 2)" + echo "RUSTUP_TOOLCHAIN=$TOOLCHAIN" >> "$GITHUB_ENV" + - name: Install uv + uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 + with: + sync: false + + - name: Pytest - PyVortex CUDA bridge + run: | + uv run --extra cuda \ + pytest --benchmark-disable ../vortex-python-cuda/test/test_native_bridge.py + working-directory: vortex-python/ + rust-docs: name: "Rust (docs)" + needs: duckdb-ready timeout-minutes: 30 runs-on: >- ${{ github.repository == 'vortex-data/vortex' @@ -131,7 +184,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Docs env: @@ -140,9 +193,9 @@ jobs: run: | # Build docs for the whole workspace (aside from vortex-python that is having hard time with the python-docs syntax), # including all private docs. - cargo doc --profile ci --no-deps --document-private-items --workspace --exclude vortex-python + cargo doc --profile ci --no-deps --document-private-items --workspace --exclude vortex-python --exclude vortex-python-cuda # nextest doesn't support doc tests, so we run it here - cargo test --profile ci --doc --workspace --all-features --exclude vortex-cxx --exclude vortex-jni --exclude vortex-ffi --exclude xtask --no-fail-fast + cargo test --profile ci --doc --workspace --all-features --exclude vortex-cxx --exclude vortex-jni --exclude vortex-ffi --exclude xtask --exclude vortex-python-cuda --no-fail-fast build-rust: name: "Rust build (${{matrix.config.name}})" @@ -173,13 +226,13 @@ jobs: target: wasm32-unknown-unknown env: rustflags: "RUSTFLAGS='-A warnings --cfg getrandom_backend=\"unsupported\"'" - args: "--target wasm32-unknown-unknown --exclude vortex --exclude vortex-cuda --exclude vortex-cub --exclude vortex-nvcomp --exclude vortex-datafusion --exclude vortex-duckdb --exclude vortex-tui --exclude vortex-zstd --exclude vortex-test-e2e-cuda --exclude vortex-sqllogictest --exclude vortex-parquet-variant" + args: "--target wasm32-unknown-unknown --exclude vortex --exclude vortex-cuda --exclude vortex-cub --exclude vortex-nvcomp --exclude vortex-datafusion --exclude vortex-duckdb --exclude vortex-tui --exclude vortex-zstd --exclude vortex-test-e2e-cuda --exclude vortex-python-cuda --exclude vortex-sqllogictest --exclude vortex-parquet-variant" steps: - uses: runs-on/action@v2 if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Install wasm32 target if: ${{ matrix.config.target == 'wasm32-unknown-unknown' }} @@ -204,12 +257,13 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - run: cargo minimal-versions check --direct --workspace --ignore-private rust-lint: name: "Rust (lint)" + needs: duckdb-ready timeout-minutes: 30 runs-on: >- ${{ github.repository == 'vortex-data/vortex' @@ -220,7 +274,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Install nightly for fmt run: rustup toolchain install $NIGHTLY_TOOLCHAIN --component rustfmt @@ -276,7 +330,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 10 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: C/C++ Lint - clang-format run: | git ls-files vortex-cuda vortex-cxx vortex-duckdb vortex-ffi \ @@ -298,7 +352,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Install cargo-hack if: github.repository != 'vortex-data/vortex' @@ -311,6 +365,7 @@ jobs: rust-test-other: name: "Rust tests (${{ matrix.os }})" + needs: duckdb-ready timeout-minutes: 30 strategy: fail-fast: false @@ -330,7 +385,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: Setup (Windows) if: matrix.os == 'windows-x64' run: | @@ -343,7 +398,7 @@ jobs: --exclude vortex-bench ` --exclude vortex-python --exclude vortex-duckdb ` --exclude vortex-fuzz --exclude vortex-cuda --exclude vortex-cuda-ffi ` - --exclude vortex-nvcomp --exclude vortex-cub --exclude vortex-test-e2e-cuda ` + --exclude vortex-nvcomp --exclude vortex-cub --exclude vortex-test-e2e-cuda --exclude vortex-python-cuda ` --exclude duckdb-bench ` --exclude lance-bench --exclude datafusion-bench --exclude random-access-bench ` --exclude compress-bench --exclude xtask --exclude vortex-datafusion ` @@ -377,7 +432,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - run: ./gradlew javadoc working-directory: ./java @@ -396,7 +451,7 @@ jobs: # Prevent sudden announcement of a new advisory from failing ci: continue-on-error: ${{ matrix.checks == 'advisories' }} steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: EmbarkStudios/cargo-deny-action@bb137d7af7e4fb67e5f82a49c4fce4fad40782fe # v2 with: command: check ${{ matrix.checks }} @@ -413,7 +468,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Build and run C++ unit tests run: | @@ -432,6 +487,7 @@ jobs: sqllogic-test: name: "SQL logic tests" + needs: duckdb-ready runs-on: >- ${{ github.repository == 'vortex-data/vortex' && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=sql-logic-test', github.run_id) @@ -442,22 +498,23 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Run sqllogictest tests run: | ./vortex-sqllogictest/slt/tpch/generate_data.sh - cargo test --profile ci -p vortex-sqllogictest --test sqllogictests + cargo nextest run --cargo-profile ci -p vortex-sqllogictest - uses: ./.github/actions/check-rebuild with: - command: "cargo test --profile ci -p vortex-sqllogictest --test sqllogictests --no-run" + command: "cargo nextest run --cargo-profile ci -p vortex-sqllogictest --no-run" + message-format-flag: "--cargo-message-format" wasm-integration: name: "WASM integration smoke test" runs-on: ubuntu-latest timeout-minutes: 30 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} @@ -484,7 +541,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Install nightly for cbindgen macro expansion run: rustup toolchain install $NIGHTLY_TOOLCHAIN @@ -523,7 +580,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: "regenerate FFI header file" run: | diff --git a/.github/workflows/claude-review.yml b/.github/workflows/claude-review.yml index c95031a19e6..fbcfb3f42ae 100644 --- a/.github/workflows/claude-review.yml +++ b/.github/workflows/claude-review.yml @@ -232,7 +232,7 @@ jobs: actions: read steps: - name: Checkout same-repo PR contents - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: ref: ${{ needs.gate.outputs.checkout_ref }} fetch-depth: 1 diff --git a/.github/workflows/claude-write.yml b/.github/workflows/claude-write.yml index 5d6224fc18b..acf5add785c 100644 --- a/.github/workflows/claude-write.yml +++ b/.github/workflows/claude-write.yml @@ -167,7 +167,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: ref: ${{ needs.gate.outputs.checkout_ref }} fetch-depth: 0 diff --git a/.github/workflows/close-fixed-fuzzer-issues.yml b/.github/workflows/close-fixed-fuzzer-issues.yml index 9e7deda1f97..ee203553ac7 100644 --- a/.github/workflows/close-fixed-fuzzer-issues.yml +++ b/.github/workflows/close-fixed-fuzzer-issues.yml @@ -36,7 +36,7 @@ jobs: with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index c819d5178d6..2cdcbaa9dbf 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -27,7 +27,7 @@ jobs: strategy: matrix: include: - - { shard: 1, name: "Core foundation", packages: "vortex-buffer vortex-error vortex-mask" } + - { shard: 1, name: "Core foundation", packages: "vortex-buffer vortex-error vortex-mask vortex-compute" } - { shard: 2, name: "Arrays", packages: "vortex-array", features: "--features _test-harness" } - { shard: 3, name: "Main library", packages: "vortex" } - { shard: 4, name: "Encodings 1", packages: "vortex-alp vortex-bytebool vortex-datetime-parts" } @@ -46,7 +46,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: Setup benchmark environment run: sudo bash scripts/setup-benchmark.sh - uses: ./.github/actions/setup-prebuild @@ -69,8 +69,48 @@ jobs: token: ${{ secrets.CODSPEED_TOKEN }} mode: "simulation" + # Getting a GPU box is slow, in the future we can build on a box without one and only run + # on GPU machines. + bench-codspeed-cuda-build: + if: github.repository == 'vortex-data/vortex' + name: "Build Codspeed CUDA benchmarks" + timeout-minutes: 30 + runs-on: >- + runs-on=${{ github.run_id }}/family=g5/cpu=8/image=ubuntu24-gpu-x64/tag=bench-codspeed-cuda-build + steps: + - uses: runs-on/action@v2 + with: + sccache: s3 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + - uses: ./.github/actions/setup-rust + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + - name: Install Codspeed + uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 + with: + tool: cargo-codspeed + - name: Build benchmarks + run: | + cargo codspeed build -m walltime \ + --bench bitpacked_cuda \ + --bench dynamic_dispatch_cuda \ + --bench alp_cuda \ + --bench date_time_parts_cuda \ + --bench dict_cuda \ + --bench fsst_cuda \ + --bench runend_cuda \ + --profile bench + - name: Upload benchmark executables + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7 + with: + name: codspeed-cuda-benchmarks + path: target/codspeed/ + retention-days: 1 + if-no-files-found: error + bench-codspeed-cuda: if: github.repository == 'vortex-data/vortex' + needs: [bench-codspeed-cuda-build] strategy: matrix: include: @@ -85,7 +125,7 @@ jobs: - uses: runs-on/action@v2 with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} @@ -101,8 +141,13 @@ jobs: uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 with: tool: cargo-codspeed - - name: Build benchmarks - run: cargo codspeed build -m walltime $(printf -- '--bench %s ' ${{ matrix.benches }}) --profile bench + - name: Download benchmark executables + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8 + with: + name: codspeed-cuda-benchmarks + path: target/codspeed + - name: Restore executable permissions + run: find target/codspeed -type f -exec chmod +x {} + - name: Run benchmarks uses: CodSpeedHQ/action@d872884a306dd4853acf0f584f4b706cf0cc72a2 env: diff --git a/.github/workflows/commit-metadata.yml b/.github/workflows/commit-metadata.yml new file mode 100644 index 00000000000..272bbaad9c9 --- /dev/null +++ b/.github/workflows/commit-metadata.yml @@ -0,0 +1,82 @@ +# Uploads commit metadata (an empty ingest envelope -- no benchmark records) on +# every push to develop, to both the v3 server and the v4 Postgres, so the +# `commits` dim stays populated even when no benchmark ran. Needed by both +# backends, hence not named for either. + +name: Commit metadata + +on: + push: + branches: [develop] + workflow_dispatch: { } + +permissions: + id-token: write # enables AWS-GitHub OIDC for the best-effort v4 ingest step + contents: read + +jobs: + commit-metadata: + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + with: + fetch-depth: 2 + + - name: Ingest commit metadata to v3 server + if: vars.V3_INGEST_URL != '' + shell: bash + env: + INGEST_BEARER_TOKEN: ${{ secrets.INGEST_BEARER_TOKEN }} + run: | + echo -n > empty.jsonl + python3 scripts/post-ingest.py empty.jsonl \ + --server "${{ vars.V3_INGEST_URL }}" \ + --commit-sha "${{ github.sha }}" \ + --benchmark-id "commit-metadata" \ + --repo-url "${{ github.server_url }}/${{ github.repository }}" + + # v4 (Postgres) dual-write -- BEST-EFFORT (see bench.yml rationale). Empty records: + # post-ingest.py --postgres upserts the commit row only. v3 above stays required; + # a v4 failure never fails the job (promoted to required at cutover, PR-5.1). + # Gated on the ingest-role ARN var (the assume-role input that MUST exist) so + # it no-ops until v4 infra is wired. + # + # ORDER MATTERS: "Install uv" runs BEFORE "Configure AWS credentials". + # configure-aws-credentials persists the assumed ingest-role (rds-db:connect + # only) as the job's ambient AWS creds; the uv setup compiles via sccache + # (S3-backed), so running it after the role switch fails with S3 AccessDenied. + # Installing uv first keeps sccache on the original S3-capable creds; the role + # is assumed immediately before the ingest, which needs only rds-db:connect. + - name: Install uv for v4 ingest + if: vars.GH_BENCH_INGEST_ROLE_ARN != '' + continue-on-error: true + uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 + - name: Configure AWS credentials for v4 ingest (OIDC) + if: vars.GH_BENCH_INGEST_ROLE_ARN != '' + continue-on-error: true + uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6 + with: + role-to-assume: ${{ vars.GH_BENCH_INGEST_ROLE_ARN }} + aws-region: ${{ vars.RDS_BENCH_REGION }} + - name: Ingest commit metadata to v4 Postgres (best-effort) + if: vars.GH_BENCH_INGEST_ROLE_ARN != '' + continue-on-error: true + shell: bash + env: + RDS_BENCH_INSTANCE_ENDPOINT: ${{ vars.RDS_BENCH_INSTANCE_ENDPOINT }} + RDS_BENCH_DB_NAME: ${{ vars.RDS_BENCH_DB_NAME }} + AWS_REGION: ${{ vars.RDS_BENCH_REGION }} + BENCH_SITE_BASE_URL: ${{ vars.BENCH_SITE_BASE_URL }} + BENCH_REVALIDATE_TOKEN: ${{ secrets.BENCH_REVALIDATE_TOKEN }} + run: | + set -Eeuo pipefail + echo -n > empty.jsonl + curl -fsSL https://truststore.pki.rds.amazonaws.com/global/global-bundle.pem \ + -o "${RUNNER_TEMP}/rds-global-bundle.pem" + DSN="postgresql://bench_ingest@${RDS_BENCH_INSTANCE_ENDPOINT}:5432/${RDS_BENCH_DB_NAME}?sslmode=verify-full&sslrootcert=${RUNNER_TEMP}/rds-global-bundle.pem" + uv run --no-project --with 'psycopg[binary]' --with boto3 --with xxhash \ + scripts/post-ingest.py empty.jsonl \ + --postgres "${DSN}" \ + --commit-sha "${{ github.sha }}" \ + --region "${AWS_REGION}" diff --git a/.github/workflows/compat-gen-upload.yml b/.github/workflows/compat-gen-upload.yml index d3aafe0dfbb..8dd4cd1999e 100644 --- a/.github/workflows/compat-gen-upload.yml +++ b/.github/workflows/compat-gen-upload.yml @@ -38,7 +38,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: fetch-depth: 0 - uses: ./.github/actions/setup-prebuild @@ -88,7 +88,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: fetch-depth: 0 - uses: ./.github/actions/setup-prebuild diff --git a/.github/workflows/compat-validation.yml b/.github/workflows/compat-validation.yml index e722d6f9f34..8732bdfbf6d 100644 --- a/.github/workflows/compat-validation.yml +++ b/.github/workflows/compat-validation.yml @@ -33,7 +33,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Run compat tests run: | diff --git a/.github/workflows/cuda.yaml b/.github/workflows/cuda.yaml index 00faca7fd43..2bd9afbbdf1 100644 --- a/.github/workflows/cuda.yaml +++ b/.github/workflows/cuda.yaml @@ -21,8 +21,37 @@ env: RUST_BACKTRACE: 1 jobs: + changes: + name: "Detect CUDA changes" + runs-on: ubuntu-latest + timeout-minutes: 10 + permissions: + pull-requests: read + outputs: + run-cuda-san: ${{ github.event_name != 'pull_request' || steps.filter.outputs.cuda == 'true' }} + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + - uses: dorny/paths-filter@fbd0ab8f3e69293af611ebaee6363fc25e6d187d # v4 + id: filter + if: github.event_name == 'pull_request' + with: + filters: | + cuda: + - "vortex-cuda/**" + - "vortex-python/python/vortex/__init__.py" + - "vortex-python/pyproject.toml" + - "vortex-python/test/test_cuda.py" + - "vortex-python-cuda/**" + - "vortex-test/**" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/**" + cuda-build-lint: - if: github.repository == 'vortex-data/vortex' + needs: [changes] + if: >- + always() && github.repository == 'vortex-data/vortex' && + needs.changes.outputs.run-cuda-san == 'true' name: "CUDA build & lint" timeout-minutes: 30 runs-on: runs-on=${{ github.run_id }}/runner=gpu/tag=cuda-build @@ -30,7 +59,7 @@ jobs: - uses: runs-on/action@v2 with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} @@ -39,7 +68,7 @@ jobs: command: >- cargo build --profile ci --locked --all-features --all-targets -p vortex-cuda -p vortex-cuda-ffi -p vortex-cub -p vortex-nvcomp - -p gpu-scan-cli -p vortex-test-e2e-cuda + -p gpu-scan-cli -p vortex-test-e2e-cuda -p vortex-python-cuda - name: Clippy CUDA crates run: | cargo clippy --profile ci --locked --all-features --all-targets \ @@ -49,10 +78,14 @@ jobs: -p vortex-nvcomp \ -p gpu-scan-cli \ -p vortex-test-e2e-cuda \ + -p vortex-python-cuda \ -- -D warnings cuda-test: - if: github.repository == 'vortex-data/vortex' + needs: [changes] + if: >- + always() && github.repository == 'vortex-data/vortex' && + needs.changes.outputs.run-cuda-san == 'true' name: "CUDA tests" timeout-minutes: 30 runs-on: runs-on=${{ github.run_id }}/runner=gpu/tag=cuda-tests @@ -65,12 +98,16 @@ jobs: nvidia-smi nvidia-smi -L nvidia-smi -q -d Memory - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} + - name: Install uv + uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 + with: + sync: false - name: Install nextest - uses: taiki-e/install-action@0631aa6515c7d545823c67cfae7ef4fc7f490154 # v2 + uses: taiki-e/install-action@9e1e5806d4a4822de933115878265be9aaa786d9 # v2 with: tool: nextest - name: Rust Tests @@ -86,13 +123,25 @@ jobs: -p vortex-cub \ -p vortex-nvcomp \ -p vortex-test-e2e-cuda \ + -p vortex-python-cuda \ --all-features \ --no-fail-fast \ --target x86_64-unknown-linux-gnu \ --verbose + - name: Python CUDA extension checks + env: + MATURIN_PEP517_ARGS: "--profile ci" + run: | + # --all-packages installs the shared dev tooling (basedpyright) from the + # root `dev` group; --extra cuda adds the vortex-data-cuda extension. + uv run --all-packages --extra cuda basedpyright vortex-python vortex-python-cuda + uv run --all-packages --extra cuda pytest --benchmark-disable vortex-python/test/test_cuda.py vortex-python-cuda/test cuda-test-sanitizer: - if: github.repository == 'vortex-data/vortex' + needs: [changes] + if: >- + always() && github.repository == 'vortex-data/vortex' && + needs.changes.outputs.run-cuda-san == 'true' name: "CUDA tests (${{ matrix.sanitizer }})" timeout-minutes: 30 runs-on: runs-on=${{ github.run_id }}/runner=gpu/tag=cuda-test-sanitizer @@ -118,7 +167,7 @@ jobs: nvidia-smi nvidia-smi -L nvidia-smi -q -d Memory - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} @@ -130,7 +179,10 @@ jobs: run: cargo test --profile ci --locked -p vortex-cuda --all-features --target x86_64-unknown-linux-gnu cuda-test-cudf: - if: github.repository == 'vortex-data/vortex' + needs: [changes] + if: >- + always() && github.repository == 'vortex-data/vortex' && + needs.changes.outputs.run-cuda-san == 'true' name: "CUDA tests (cudf)" timeout-minutes: 30 runs-on: runs-on=${{ github.run_id }}/runner=gpu/tag=cuda-test-cudf @@ -143,7 +195,7 @@ jobs: nvidia-smi nvidia-smi -L nvidia-smi -q -d Memory - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index eadc0839c11..1a0d376f319 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -19,13 +19,13 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 30 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} enable-sccache: "false" - name: Set up JDK 17 - uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654 # v5 + uses: actions/setup-java@ad2b38190b15e4d6bdf0c97fb4fca8412226d287 # v5 with: java-version: "17" distribution: "temurin" diff --git a/.github/workflows/duckdb-r2.yml b/.github/workflows/duckdb-r2.yml new file mode 100644 index 00000000000..6670085ac91 --- /dev/null +++ b/.github/workflows/duckdb-r2.yml @@ -0,0 +1,83 @@ +name: DuckDB R2 mirror + +# If we have an archive for a commit or a tag in R2, return it. +# Otherwise, if it's a tag, download release archive from Duckdb Github +# page and upload it to R2. +# If it's a commit, build Duckdb from source and upload it to R2. +on: + workflow_call: { } + +concurrency: + group: duckdb-r2-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: false + +permissions: + contents: read + +env: + PUBLIC_BASE_URL: "https://ci-builds.vortex.dev" + R2_BUCKET: "duckdb-builds" + R2_ENDPOINT_URL: "https://52bdeab5651e1584747feefd051fd566.r2.cloudflarestorage.com" + +jobs: + check: + name: "Resolve DuckDB version and check R2" + runs-on: ubuntu-latest + timeout-minutes: 10 + outputs: + version: ${{ steps.resolve.outputs.version }} + ref_dir: ${{ steps.resolve.outputs.ref_dir }} + release: ${{ steps.resolve.outputs.release }} + matrix: ${{ steps.resolve.outputs.matrix }} + any_missing: ${{ steps.resolve.outputs.any_missing }} + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - name: Resolve version and check R2 + id: resolve + run: sh scripts/duckdb-r2-resolve.sh + + mirror: + name: "Mirror DuckDB ${{ matrix.archive }} to R2" + needs: check + if: >- + needs.check.outputs.any_missing == 'true' && + github.repository == 'vortex-data/vortex' && + github.event.pull_request.head.repo.full_name == github.repository + environment: duckdb-build + timeout-minutes: 120 + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.check.outputs.matrix) }} + runs-on: ${{ matrix.runner }} + steps: + - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + + - name: Install build dependencies (Linux) + if: needs.check.outputs.release != 'true' && runner.os == 'Linux' + run: | + sudo apt-get update + sudo apt-get install -y ninja-build libcurl4-openssl-dev zip unzip + + # MacOS already has ninja and p7zip + + - name: Prepare ${{ matrix.archive }} + env: + ARCHIVE: ${{ matrix.archive }} + REF_DIR: ${{ needs.check.outputs.ref_dir }} + RELEASE: ${{ needs.check.outputs.release }} + PLATFORM_OS: ${{ matrix.os }} + run: sh scripts/duckdb-r2-prepare.sh + + - name: Upload to R2 + env: + AWS_ACCESS_KEY_ID: ${{ secrets.DUCKDB_R2_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.DUCKDB_R2_SECRET_ACCESS_KEY }} + AWS_REGION: "us-east-1" + AWS_ENDPOINT_URL: ${{ env.R2_ENDPOINT_URL }} + run: | + set -Eeuo pipefail + python3 scripts/s3-upload.py \ + --bucket "$R2_BUCKET" \ + --key "${{ needs.check.outputs.ref_dir }}/${{ matrix.archive }}" \ + --body "${{ matrix.archive }}" \ + --checksum-algorithm CRC32 diff --git a/.github/workflows/fuzz-coverage.yml b/.github/workflows/fuzz-coverage.yml index a62577313cc..fd6ca8aed1b 100644 --- a/.github/workflows/fuzz-coverage.yml +++ b/.github/workflows/fuzz-coverage.yml @@ -26,7 +26,7 @@ jobs: with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: diff --git a/.github/workflows/fuzz.yml b/.github/workflows/fuzz.yml index 8d6186a9e8f..f421045100e 100644 --- a/.github/workflows/fuzz.yml +++ b/.github/workflows/fuzz.yml @@ -200,6 +200,10 @@ jobs: fuzz_target: compress_gpu runner: gpu extra_features: "cuda" + # ASan mprotects the shadow gap PROT_NONE by default, which collides with the large + # virtual-address range the CUDA driver reserves and makes device init fail with an + # OOM-style error. protect_shadow_gap=0 frees that VA range for CUDA. + extra_env: "ASAN_OPTIONS=protect_shadow_gap=0" jobs: 1 secrets: R2_FUZZ_ACCESS_KEY_ID: ${{ secrets.R2_FUZZ_ACCESS_KEY_ID }} diff --git a/.github/workflows/fuzzer-fix-automation.yml b/.github/workflows/fuzzer-fix-automation.yml index 01df7162e78..7b8e7012511 100644 --- a/.github/workflows/fuzzer-fix-automation.yml +++ b/.github/workflows/fuzzer-fix-automation.yml @@ -63,7 +63,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: # Don't leave the built-in GITHUB_TOKEN in git config; the fix is # committed and pushed with the App token so its PR triggers CI. @@ -273,7 +273,7 @@ jobs: CRASH_FILE: ${{ steps.extract.outputs.crash_file }} CRASH_FILE_PATH: ${{ steps.download.outputs.crash_file_path }} ARTIFACT_URL: ${{ steps.extract.outputs.artifact_url }} - uses: anthropics/claude-code-action@fbda2eb1bdc90d319b8d853f5deb53bca199a7c1 # v1 + uses: anthropics/claude-code-action@2fee15510437d71399d9139ed60433470484a8fb # v1 with: claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }} # Use the App token (not GITHUB_TOKEN) so the committed fix branch and diff --git a/.github/workflows/minimize_fuzz_corpus_workflow.yml b/.github/workflows/minimize_fuzz_corpus_workflow.yml index f6c1e8d8df1..00f9ba73b84 100644 --- a/.github/workflows/minimize_fuzz_corpus_workflow.yml +++ b/.github/workflows/minimize_fuzz_corpus_workflow.yml @@ -51,7 +51,7 @@ jobs: with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: diff --git a/.github/workflows/musl.yml b/.github/workflows/musl.yml index cf8ff08eeb6..6ab9ef566f4 100644 --- a/.github/workflows/musl.yml +++ b/.github/workflows/musl.yml @@ -52,7 +52,7 @@ jobs: pkgconf protobuf protobuf-dev openssl-dev zstd-dev git curl ca-certificates \ python3 python3-dev tar tzdata bash - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: Install nextest shell: bash @@ -66,7 +66,7 @@ jobs: run: | cargo nextest run --cargo-profile ci --locked --workspace --no-fail-fast \ --exclude vortex-cuda --exclude vortex-cub --exclude vortex-nvcomp \ - --exclude gpu-scan-cli --exclude vortex-test-e2e-cuda \ + --exclude gpu-scan-cli --exclude vortex-test-e2e-cuda --exclude vortex-python-cuda \ --exclude vortex-duckdb --exclude duckdb-bench --exclude vortex-sqllogictest \ --exclude vortex-bench --exclude lance-bench --exclude datafusion-bench --exclude vortex-datafusion \ --exclude compress-bench --exclude random-access-bench --exclude vortex-bench-server diff --git a/.github/workflows/package.yml b/.github/workflows/package.yml index cc4e68d0eff..3de9ac5c81d 100644 --- a/.github/workflows/package.yml +++ b/.github/workflows/package.yml @@ -31,7 +31,7 @@ jobs: - { os: ubuntu, runs-on: "ubuntu-latest", target: aarch64-unknown-linux-gnu } - { os: ubuntu, runs-on: "ubuntu-latest", target: x86_64-unknown-linux-gnu } steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: fetch-depth: 0 @@ -50,6 +50,34 @@ jobs: run: | cargo set-version --workspace ${{ inputs.version }} + # cargo set-version bumps Cargo.toml crate versions but not the ==0.1.0 pins + # between these two wheels in their pyproject.toml. The base and CUDA extension + # wheels must require each other at the exact released version (they exchange + # Vortex IPC buffers), so rewrite those pins here. + - name: Set Python extension version pins + env: + VERSION: ${{ inputs.version }} + run: | + python3 - <<'PY' + import os + import re + from pathlib import Path + + version = os.environ["VERSION"] + replacements = { + Path("vortex-python/pyproject.toml"): "vortex-data-cuda", + Path("vortex-python-cuda/pyproject.toml"): "vortex-data", + } + for path, package in replacements.items(): + text = path.read_text() + text, count = re.subn( + f"{re.escape(package)}==[^\"\\]]+", f"{package}=={version}", text + ) + if count != 1: + raise RuntimeError(f"Expected one {package} version pin in {path}, found {count}") + path.write_text(text) + PY + - name: Set up Python uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6 # Latest macOS doesn't allow maturin to install stuff into the global Python interpreter @@ -86,10 +114,10 @@ jobs: runs-on: "macos-latest" timeout-minutes: 30 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: fetch-depth: 0 - - uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654 # v5 + - uses: actions/setup-java@ad2b38190b15e4d6bdf0c97fb4fca8412226d287 # v5 with: distribution: "corretto" java-version: "17" @@ -121,10 +149,10 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: fetch-depth: 0 - - uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654 # v5 + - uses: actions/setup-java@ad2b38190b15e4d6bdf0c97fb4fca8412226d287 # v5 with: distribution: "corretto" java-version: "17" diff --git a/.github/workflows/publish-benchmarks-website.yml b/.github/workflows/publish-benchmarks-website.yml index be1def626a0..5555dcf841a 100644 --- a/.github/workflows/publish-benchmarks-website.yml +++ b/.github/workflows/publish-benchmarks-website.yml @@ -14,7 +14,7 @@ jobs: contents: read packages: write steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: Log in to GHCR uses: docker/login-action@650006c6eb7dba73a995cc03b0b2d7f5ca915bee # v4 diff --git a/.github/workflows/publish-dry-runs.yml b/.github/workflows/publish-dry-runs.yml index 7b6a750260b..49f539a19e4 100644 --- a/.github/workflows/publish-dry-runs.yml +++ b/.github/workflows/publish-dry-runs.yml @@ -27,7 +27,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 30 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: Rust Dependency Cache uses: Swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2 with: @@ -79,10 +79,10 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: fetch-depth: 0 - - uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654 # v5 + - uses: actions/setup-java@ad2b38190b15e4d6bdf0c97fb4fca8412226d287 # v5 with: distribution: "corretto" java-version: "17" @@ -112,7 +112,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Install cargo-edit uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 6398c54c20e..aba77190690 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -23,7 +23,7 @@ jobs: timeout-minutes: 30 needs: [package] steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: fetch-depth: 0 - uses: ./.github/actions/setup-rust @@ -88,12 +88,12 @@ jobs: run: working-directory: ./java steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 - - uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654 # v5 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + - uses: actions/setup-java@ad2b38190b15e4d6bdf0c97fb4fca8412226d287 # v5 with: distribution: "corretto" java-version: "17" - - uses: gradle/actions/setup-gradle@50e97c2cd7a37755bbfafc9c5b7cafaece252f6e # v6 + - uses: gradle/actions/setup-gradle@3f131e8634966bd73d06cc69884922b02e6faf92 # v6 - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8 with: pattern: libvortex_jni_*.zip diff --git a/.github/workflows/release-binaries.yml b/.github/workflows/release-binaries.yml index 64ecafd2e77..dcedec293c3 100644 --- a/.github/workflows/release-binaries.yml +++ b/.github/workflows/release-binaries.yml @@ -44,7 +44,7 @@ jobs: with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: fetch-depth: 0 diff --git a/.github/workflows/report-fuzz-crash.yml b/.github/workflows/report-fuzz-crash.yml index ea68e9a832a..ae4b894e3cd 100644 --- a/.github/workflows/report-fuzz-crash.yml +++ b/.github/workflows/report-fuzz-crash.yml @@ -49,7 +49,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: Download fuzzer logs uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8 @@ -118,7 +118,7 @@ jobs: steps.dedup.outputs.duplicate != 'true' || steps.dedup.outputs.confidence != 'exact' continue-on-error: true - uses: anthropics/claude-code-action@fbda2eb1bdc90d319b8d853f5deb53bca199a7c1 # v1 + uses: anthropics/claude-code-action@2fee15510437d71399d9139ed60433470484a8fb # v1 with: claude_code_oauth_token: ${{ secrets.claude_code_oauth_token }} github_token: ${{ secrets.gh_token }} diff --git a/.github/workflows/reuse.yml b/.github/workflows/reuse.yml index 3028dbbd8f1..f5515f96cfa 100644 --- a/.github/workflows/reuse.yml +++ b/.github/workflows/reuse.yml @@ -17,6 +17,6 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 10 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: REUSE Compliance Check uses: fsfe/reuse-action@676e2d560c9a403aa252096d99fcab3e1132b0f5 # v6 diff --git a/.github/workflows/run-fuzzer.yml b/.github/workflows/run-fuzzer.yml index a3a41a1a53e..e563dbefbee 100644 --- a/.github/workflows/run-fuzzer.yml +++ b/.github/workflows/run-fuzzer.yml @@ -76,7 +76,7 @@ jobs: with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: @@ -116,6 +116,7 @@ jobs: - name: Run fuzzing target id: fuzz run: | + CORPUS_DIR="fuzz/corpus/${FUZZ_NAME}" FEATURES_FLAG="" if [ -n "${{ inputs.extra_features }}" ]; then FEATURES_FLAG="--features ${{ inputs.extra_features }}" @@ -127,7 +128,7 @@ jobs: ${{ inputs.extra_env }} RUST_BACKTRACE=1 \ cargo +$NIGHTLY_TOOLCHAIN fuzz run --release --debug-assertions \ $FEATURES_FLAG \ - ${{ inputs.fuzz_target }} -- \ + ${{ inputs.fuzz_target }} "$CORPUS_DIR" -- \ $FORK_FLAG -max_total_time=${{ inputs.max_time }} -rss_limit_mb=0 \ 2>&1 | tee fuzz_output.log continue-on-error: true diff --git a/.github/workflows/rust-instrumented.yml b/.github/workflows/rust-instrumented.yml index d58463d53d2..bf45fc7be13 100644 --- a/.github/workflows/rust-instrumented.yml +++ b/.github/workflows/rust-instrumented.yml @@ -22,8 +22,28 @@ env: NIGHTLY_TOOLCHAIN: nightly-2026-02-05 jobs: + duckdb-mirror: + name: "Mirror DuckDB to R2" + if: github.event_name == 'pull_request' + uses: ./.github/workflows/duckdb-r2.yml + secrets: inherit + + duckdb-ready: + name: "DuckDB libraries available in R2" + needs: duckdb-mirror + if: ${{ !cancelled() }} + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - name: Verify DuckDB mirror + if: ${{ needs.duckdb-mirror.result == 'failure' }} + run: | + echo "DuckDB mirror failed" + exit 1 + rust-coverage: name: "Rust tests (coverage) (${{ matrix.suite }})" + needs: duckdb-ready timeout-minutes: 30 permissions: id-token: write @@ -45,7 +65,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Ensure llvm-tools are installed run: | @@ -68,7 +88,7 @@ jobs: - name: Rust Tests if: ${{ matrix.suite == 'tests' }} run: | - cargo nextest run --locked --workspace --all-features --no-fail-fast + cargo nextest run --locked --workspace --all-features --no-fail-fast --exclude vortex-python-cuda - name: Generate coverage report run: | RUST_SYSROOT="$(rustc --print sysroot)" @@ -128,7 +148,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Install Rust nightly toolchain run: | @@ -195,7 +215,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Install rustfilt run: | @@ -221,7 +241,8 @@ jobs: - name: Run tests run: | set -o pipefail - ./vortex-ffi/build/test/vortex_ffi_test 2>&1 | rustfilt +# re-enable once we can fix this hang +# ./vortex-ffi/build/test/vortex_ffi_test 2>&1 | rustfilt - name: Run examples run: | set -o pipefail @@ -252,7 +273,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - name: Install nightly with miri run: rustup toolchain install $NIGHTLY_TOOLCHAIN --component rust-src,rustfmt,clippy,miri diff --git a/.github/workflows/sql-benchmarks.yml b/.github/workflows/sql-benchmarks.yml index b46fe455a34..b8c2d08f64f 100644 --- a/.github/workflows/sql-benchmarks.yml +++ b/.github/workflows/sql-benchmarks.yml @@ -13,7 +13,7 @@ on: benchmark_matrix: required: false type: string - description: "JSON string containing the matrix configuration" + description: "JSON string containing the full matrix configuration" default: | [ { @@ -39,6 +39,29 @@ on: {"engine": "duckdb", "format": "duckdb"} ] }, + { + "id": "clickbench-sorted-nvme", + "subcommand": "clickbench-sorted", + "name": "Clickbench Sorted on NVME", + "data_formats": ["parquet", "vortex", "vortex-compact", "duckdb"], + "pr_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"}, + {"engine": "duckdb", "format": "duckdb"} + ], + "develop_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "datafusion", "format": "vortex-compact"}, + {"engine": "datafusion", "format": "lance"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"}, + {"engine": "duckdb", "format": "vortex-compact"}, + {"engine": "duckdb", "format": "duckdb"} + ] + }, { "id": "tpch-nvme", "subcommand": "tpch", @@ -277,6 +300,212 @@ on: "iterations": "10" } ] + base_benchmark_matrix: + required: false + type: string + description: "JSON string containing the base matrix configuration" + default: | + [ + { + "id": "clickbench-nvme", + "subcommand": "clickbench", + "name": "Clickbench on NVME", + "data_formats": ["parquet", "vortex"], + "pr_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "develop_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "datafusion", "format": "lance"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ] + }, + { + "id": "clickbench-sorted-nvme", + "subcommand": "clickbench-sorted", + "name": "Clickbench Sorted on NVME", + "data_formats": ["parquet", "vortex"], + "pr_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "develop_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "datafusion", "format": "lance"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ] + }, + { + "id": "tpch-nvme", + "subcommand": "tpch", + "name": "TPC-H SF=1 on NVME", + "data_formats": ["parquet", "vortex"], + "pr_targets": [ + {"engine": "datafusion", "format": "arrow"}, + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "develop_targets": [ + {"engine": "datafusion", "format": "arrow"}, + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "datafusion", "format": "lance"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "scale_factor": "1.0", + "iterations": "10" + }, + { + "id": "tpch-s3", + "subcommand": "tpch", + "name": "TPC-H SF=1 on S3", + "local_dir": "vortex-bench/data/tpch/1.0", + "remote_storage": "s3://vortex-ci-benchmark-datasets/${{github.ref_name}}/${{github.run_id}}/tpch/1.0/", + "data_formats": ["parquet", "vortex"], + "pr_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "develop_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "scale_factor": "1.0", + "iterations": "10" + }, + { + "id": "tpch-nvme-10", + "subcommand": "tpch", + "name": "TPC-H SF=10 on NVME", + "data_formats": ["parquet", "vortex"], + "pr_targets": [ + {"engine": "datafusion", "format": "arrow"}, + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "develop_targets": [ + {"engine": "datafusion", "format": "arrow"}, + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "datafusion", "format": "lance"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "scale_factor": "10.0", + "iterations": "10" + }, + { + "id": "tpcds-nvme", + "subcommand": "tpcds", + "name": "TPC-DS SF=1 on NVME", + "data_formats": ["parquet", "vortex"], + "pr_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "develop_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "scale_factor": "1.0" + }, + { + "id": "statpopgen", + "subcommand": "statpopgen", + "name": "Statistical and Population Genetics", + "local_dir": "vortex-bench/data/statpopgen", + "data_formats": ["parquet", "vortex"], + "pr_targets": [ + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "develop_targets": [ + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "scale_factor": "100" + }, + { + "id": "fineweb", + "subcommand": "fineweb", + "name": "FineWeb NVMe", + "data_formats": ["parquet", "vortex"], + "pr_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "develop_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "scale_factor": "100" + }, + { + "id": "fineweb-s3", + "subcommand": "fineweb", + "name": "FineWeb S3", + "local_dir": "vortex-bench/data/fineweb", + "remote_storage": "s3://vortex-ci-benchmark-datasets/${{github.ref_name}}/${{github.run_id}}/fineweb/", + "data_formats": ["parquet", "vortex"], + "pr_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "develop_targets": [ + {"engine": "datafusion", "format": "parquet"}, + {"engine": "datafusion", "format": "vortex"}, + {"engine": "duckdb", "format": "parquet"}, + {"engine": "duckdb", "format": "vortex"} + ], + "scale_factor": "100" + }, + { + "id": "polarsignals", + "subcommand": "polarsignals", + "name": "PolarSignals Profiling", + "data_formats": ["vortex"], + "pr_targets": [ + {"engine": "datafusion", "format": "vortex"} + ], + "develop_targets": [ + {"engine": "datafusion", "format": "vortex"} + ], + "scale_factor": "1" + } + ] + benchmark_profile: + required: false + type: string + description: "Benchmark profile to run: full or base" + default: "full" jobs: bench: @@ -289,7 +518,7 @@ jobs: strategy: fail-fast: false matrix: - include: ${{ fromJSON(inputs.benchmark_matrix) }} + include: ${{ fromJSON(inputs.benchmark_profile == 'base' && inputs.base_benchmark_matrix || inputs.benchmark_matrix) }} runs-on: >- ${{ github.repository == 'vortex-data/vortex' @@ -300,12 +529,12 @@ jobs: if: inputs.mode != 'pr' || github.event.pull_request.head.repo.fork == false with: sccache: s3 - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 if: inputs.mode == 'pr' with: ref: ${{ github.event.pull_request.head.sha }} - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 if: inputs.mode != 'pr' - name: Setup benchmark environment run: sudo bash scripts/setup-benchmark.sh @@ -321,7 +550,7 @@ jobs: run: | wget -qO- https://github.com/duckdb/duckdb/releases/download/v1.5.3/duckdb_cli-linux-amd64.zip | funzip > duckdb chmod +x duckdb - echo "$PWD" >> $GITHUB_PATH + echo "$PWD" >> "$GITHUB_PATH" - uses: ./.github/actions/system-info @@ -345,11 +574,11 @@ jobs: env: RUSTFLAGS: "-C target-cpu=native" run: | - packages="--bin data-gen --bin datafusion-bench --bin duckdb-bench" + packages=(--bin data-gen --bin datafusion-bench --bin duckdb-bench) if [ "${{ inputs.mode }}" != "pr" ]; then - packages="$packages --bin lance-bench" + packages+=(--bin lance-bench) fi - cargo build $packages --profile release_debug --features unstable_encodings + cargo build "${packages[@]}" --profile release_debug --features unstable_encodings - name: Generate data shell: bash @@ -366,6 +595,8 @@ jobs: with: role-to-assume: arn:aws:iam::245040174862:role/GitHubBenchmarkRole aws-region: us-east-1 + # Expand session window, otherwise some large scale runs time out. + role-duration-seconds: 7200 - name: Upload data if: matrix.remote_storage != null && (inputs.mode != 'pr' || github.event.pull_request.head.repo.fork == false) @@ -446,11 +677,11 @@ jobs: python3 scripts/s3-download.py s3://vortex-ci-benchmark-results/data.json.gz data.json.gz --no-sign-request gzip -d -c data.json.gz > base.json - echo '# Benchmarks: ${{ matrix.name }}' > comment.md + echo "# Benchmarks: ${{ matrix.name }}" > comment.md echo '' >> comment.md uv run --no-project scripts/compare-benchmark-jsons.py base.json results.json "${{ matrix.name }}" \ >> comment.md - cat comment.md >> $GITHUB_STEP_SUMMARY + cat comment.md >> "$GITHUB_STEP_SUMMARY" - name: Comment PR if: inputs.mode == 'pr' && github.event.pull_request.head.repo.fork == false @@ -458,8 +689,7 @@ jobs: with: file-path: comment.md # There is exactly one comment per comment-tag. If a comment with this tag already exists, - # this action will *update* the comment instead of posting a new comment. Therefore, each - # unique benchmark configuration must have a unique comment-tag. + # this action will *update* the comment instead of posting a new comment. comment-tag: bench-pr-comment-${{ matrix.id }} - name: Comment PR on failure @@ -469,7 +699,7 @@ jobs: message: | # 🚨🚨🚨❌❌❌ SQL BENCHMARK FAILED ❌❌❌🚨🚨🚨 - Benchmark `${{ matrix.name }}` failed! Check the [workflow run](${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}) for details. + Benchmark `${{ matrix.name }}` (${{ inputs.benchmark_profile }}) failed! Check the [workflow run](${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}) for details. comment-tag: bench-pr-comment-${{ matrix.id }} - name: Upload Benchmark Results @@ -490,6 +720,52 @@ jobs: --benchmark-id "${{ matrix.id }}" \ --repo-url "${{ github.server_url }}/${{ github.repository }}" + # v4 (Postgres) dual-write -- BEST-EFFORT during the migration soak (see bench.yml + # for the full rationale). The v3 step above is hard-required; a v4 failure must NOT + # fail the job (v4 is promoted to required at cutover, PR-5.1). Gated on + # inputs.mode == 'develop' (matching v3) + the ingest-role ARN var (the assume-role + # input that MUST exist for OIDC to succeed; it no-ops until v4 infra is wired), and + # every step is continue-on-error. post-ingest.py mints the RDS IAM token (boto3) from + # the assumed GitHubBenchmarkIngestRole; sslmode=verify-full validates the cert. + # + # ORDER MATTERS: "Install uv" runs BEFORE "Configure AWS credentials". + # configure-aws-credentials persists the assumed ingest-role (rds-db:connect + # only) as the job's ambient AWS creds; the uv setup compiles via sccache + # (S3-backed), so running it after the role switch fails with S3 AccessDenied. + # Installing uv first keeps sccache on the original S3-capable creds; the role + # is assumed immediately before the ingest, which needs only rds-db:connect. + - name: Install uv for v4 ingest + if: inputs.mode == 'develop' && vars.GH_BENCH_INGEST_ROLE_ARN != '' + continue-on-error: true + uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 + - name: Configure AWS credentials for v4 ingest (OIDC) + if: inputs.mode == 'develop' && vars.GH_BENCH_INGEST_ROLE_ARN != '' + continue-on-error: true + uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6 + with: + role-to-assume: ${{ vars.GH_BENCH_INGEST_ROLE_ARN }} + aws-region: ${{ vars.RDS_BENCH_REGION }} + - name: Ingest results to v4 Postgres (best-effort) + if: inputs.mode == 'develop' && vars.GH_BENCH_INGEST_ROLE_ARN != '' + continue-on-error: true + shell: bash + env: + RDS_BENCH_INSTANCE_ENDPOINT: ${{ vars.RDS_BENCH_INSTANCE_ENDPOINT }} + RDS_BENCH_DB_NAME: ${{ vars.RDS_BENCH_DB_NAME }} + AWS_REGION: ${{ vars.RDS_BENCH_REGION }} + BENCH_SITE_BASE_URL: ${{ vars.BENCH_SITE_BASE_URL }} + BENCH_REVALIDATE_TOKEN: ${{ secrets.BENCH_REVALIDATE_TOKEN }} + run: | + set -Eeuo pipefail + curl -fsSL https://truststore.pki.rds.amazonaws.com/global/global-bundle.pem \ + -o "${RUNNER_TEMP}/rds-global-bundle.pem" + DSN="postgresql://bench_ingest@${RDS_BENCH_INSTANCE_ENDPOINT}:5432/${RDS_BENCH_DB_NAME}?sslmode=verify-full&sslrootcert=${RUNNER_TEMP}/rds-global-bundle.pem" + uv run --no-project --with 'psycopg[binary]' --with boto3 --with xxhash \ + scripts/post-ingest.py results.v3.jsonl \ + --postgres "${DSN}" \ + --commit-sha "${{ github.sha }}" \ + --region "${AWS_REGION}" + - name: Alert incident.io if: failure() && inputs.mode == 'develop' uses: ./.github/actions/alert-incident-io diff --git a/.github/workflows/sql-pr.yml b/.github/workflows/sql-pr.yml index 420425a4321..45b0ed1d675 100644 --- a/.github/workflows/sql-pr.yml +++ b/.github/workflows/sql-pr.yml @@ -1,15 +1,32 @@ # Runs SQL benchmarks once for a pull request. -# Called from bench-dispatch.yml when the `action/benchmark-sql` label is added. +# Called from bench-dispatch.yml when SQL benchmark labels are added. name: PR SQL Benchmarks concurrency: - group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + group: >- + ${{ github.workflow }}-${{ inputs.benchmark_profile || 'base' }}-${{ + github.head_ref || github.run_id + }} cancel-in-progress: true on: - workflow_call: { } - workflow_dispatch: { } + workflow_call: + inputs: + benchmark_profile: + required: false + type: string + default: "base" + workflow_dispatch: + inputs: + benchmark_profile: + description: "SQL benchmark profile to run" + required: false + type: choice + default: "base" + options: + - "base" + - "full" permissions: contents: read @@ -22,3 +39,4 @@ jobs: secrets: inherit with: mode: "pr" + benchmark_profile: ${{ inputs.benchmark_profile || 'base' }} diff --git a/.github/workflows/typos.yml b/.github/workflows/typos.yml index e05120f8843..8f288be51a4 100644 --- a/.github/workflows/typos.yml +++ b/.github/workflows/typos.yml @@ -19,6 +19,6 @@ jobs: timeout-minutes: 10 steps: - name: Checkout Actions Repository - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: Spell Check Repo uses: crate-ci/typos@37bb98842b0d8c4ffebdb75301a13db0267cef89 # v1.47.2 diff --git a/.github/workflows/v3-commit-metadata.yml b/.github/workflows/v3-commit-metadata.yml deleted file mode 100644 index d95da69207b..00000000000 --- a/.github/workflows/v3-commit-metadata.yml +++ /dev/null @@ -1,34 +0,0 @@ -# Posts a v3 ingest envelope with no records on every push to develop, so the -# `commits` dim stays populated even when no benchmark ran. - -name: v3 commit metadata - -on: - push: - branches: [develop] - workflow_dispatch: { } - -permissions: - contents: read - -jobs: - commit-metadata: - runs-on: ubuntu-latest - timeout-minutes: 10 - steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 - with: - fetch-depth: 2 - - - name: Ingest commit metadata to v3 server - if: vars.V3_INGEST_URL != '' - shell: bash - env: - INGEST_BEARER_TOKEN: ${{ secrets.INGEST_BEARER_TOKEN }} - run: | - echo -n > empty.jsonl - python3 scripts/post-ingest.py empty.jsonl \ - --server "${{ vars.V3_INGEST_URL }}" \ - --commit-sha "${{ github.sha }}" \ - --benchmark-id "commit-metadata" \ - --repo-url "${{ github.server_url }}/${{ github.repository }}" diff --git a/.github/workflows/wasm-fuzz.yml b/.github/workflows/wasm-fuzz.yml index a2cec6c6c14..d61b3b60598 100644 --- a/.github/workflows/wasm-fuzz.yml +++ b/.github/workflows/wasm-fuzz.yml @@ -25,7 +25,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 240 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: diff --git a/.github/workflows/web.yml b/.github/workflows/web.yml index cee88d0e43b..d03a2b45d14 100644 --- a/.github/workflows/web.yml +++ b/.github/workflows/web.yml @@ -28,7 +28,7 @@ jobs: outputs: web: ${{ steps.filter.outputs.web }} steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: dorny/paths-filter@fbd0ab8f3e69293af611ebaee6363fc25e6d187d # v4 id: filter with: @@ -44,7 +44,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 30 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} @@ -71,7 +71,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 30 steps: - - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index aef2e0c33ac..6d32374293a 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -12,7 +12,7 @@ We ask that you read the guidelines below in order to make the process as stream > > For the full AI policy — including disclosure requirements, review standards for > AI-generated PRs, and rules for autonomous agents — see the -> [contributing guide](https://vortex.dev/project/contributing.html#ai-assistance). +> [contributing guide](https://docs.vortex.dev/project/contributing#ai-assistance). ## Code Contributions diff --git a/Cargo.lock b/Cargo.lock index e2d1292066a..cb9e255846b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -61,9 +61,9 @@ checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" [[package]] name = "alloc-stdlib" -version = "0.2.2" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94fb8275041c72129eb51b7d0322c29b8387a0386127718b096429201a5d6ece" +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195" dependencies = [ "alloc-no-stdlib", ] @@ -141,9 +141,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.102" +version = "1.0.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" [[package]] name = "approx" @@ -156,9 +156,9 @@ dependencies = [ [[package]] name = "ar_archive_writer" -version = "0.5.1" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7eb93bbb63b9c227414f6eb3a0adfddca591a8ce1e9b60661bb08969b87e340b" +checksum = "4087686b4b0a3427190bae57a1d9a478dbb2d40c5dc1bd6e2b6d797913bdd348" dependencies = [ "object", ] @@ -171,9 +171,9 @@ checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" [[package]] name = "arc-swap" -version = "1.9.1" +version = "1.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" +checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" dependencies = [ "rustversion", ] @@ -192,9 +192,27 @@ checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" [[package]] name = "arrayvec" -version = "0.7.6" +version = "0.7.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" +checksum = "f02882884d3e1bc524fb12c79f107f6ad0e1cfd498c536ffb494301740995dfe" + +[[package]] +name = "arrow" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fb98341a7e051bb79731ecb33ec00cbd6e0e315a542d6732b46d462c9215ea2" +dependencies = [ + "arrow-arith 56.2.1", + "arrow-array 56.2.1", + "arrow-buffer 56.2.1", + "arrow-cast 56.2.1", + "arrow-data 56.2.1", + "arrow-ord 56.2.1", + "arrow-row 56.2.1", + "arrow-schema 56.2.1", + "arrow-select 56.2.1", + "arrow-string 56.2.1", +] [[package]] name = "arrow" @@ -202,19 +220,33 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "378530e55cd479eda3c14eb345310799717e6f76d0c332041e8487022166b471" dependencies = [ - "arrow-arith", - "arrow-array", - "arrow-buffer", - "arrow-cast", + "arrow-arith 58.3.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-cast 58.3.0", "arrow-csv", - "arrow-data", - "arrow-ipc", + "arrow-data 58.3.0", + "arrow-ipc 58.3.0", "arrow-json", - "arrow-ord", - "arrow-row", - "arrow-schema", - "arrow-select", - "arrow-string", + "arrow-ord 58.3.0", + "arrow-row 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", + "arrow-string 58.3.0", +] + +[[package]] +name = "arrow-arith" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce4751cbc4bcccfeeea79df9571ff1dc066d61e44723c7604d11c7937f5b560" +dependencies = [ + "arrow-array 56.2.1", + "arrow-buffer 56.2.1", + "arrow-data 56.2.1", + "arrow-schema 56.2.1", + "chrono", + "num", ] [[package]] @@ -223,14 +255,30 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a0ab212d2c1886e802f51c5212d78ebbcbb0bec980fff9dadc1eb8d45cd0b738" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", "chrono", "num-traits", ] +[[package]] +name = "arrow-array" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b02ccba2e977a3aabb4384036109ca32f552399a2bc0588f925f91ed073ce70c" +dependencies = [ + "ahash 0.8.12", + "arrow-buffer 56.2.1", + "arrow-data 56.2.1", + "arrow-schema 56.2.1", + "chrono", + "half", + "hashbrown 0.16.1", + "num", +] + [[package]] name = "arrow-array" version = "58.3.0" @@ -238,9 +286,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cfd33d3e92f207444098c75b42de99d329562be0cf686b307b097cc52b4e999e" dependencies = [ "ahash 0.8.12", - "arrow-buffer", - "arrow-data", - "arrow-schema", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", "chrono", "chrono-tz", "half", @@ -256,9 +304,9 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "049230728cd6e093088c8d231b4beede184e35cad7777c1505c0d5a8571f4376" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-schema 58.3.0", "bytes", "bzip2", "crc", @@ -274,6 +322,17 @@ dependencies = [ "zstd", ] +[[package]] +name = "arrow-buffer" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a90f8bece6a9ee316a699fbbfde368a206676a1206ce89b50f07937648e76c3c" +dependencies = [ + "bytes", + "half", + "num", +] + [[package]] name = "arrow-buffer" version = "58.3.0" @@ -286,18 +345,39 @@ dependencies = [ "num-traits", ] +[[package]] +name = "arrow-cast" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61ffe645cfb4e80b1ca37a3a106ce7b4af66ccdd60c655a57e6b9aab096164a7" +dependencies = [ + "arrow-array 56.2.1", + "arrow-buffer 56.2.1", + "arrow-data 56.2.1", + "arrow-schema 56.2.1", + "arrow-select 56.2.1", + "atoi", + "base64", + "chrono", + "comfy-table", + "half", + "lexical-core", + "num", + "ryu", +] + [[package]] name = "arrow-cast" version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4c5aefb56a2c02e9e2b30746241058b85f8983f0fcff2ba0c6d09006e1cded7f" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-ord", - "arrow-schema", - "arrow-select", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-ord 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "atoi", "base64", "chrono", @@ -314,41 +394,67 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e94e8cf7e517657a52b91ea1263acf38c4ca62a84655d72458a3359b12ab97de" dependencies = [ - "arrow-array", - "arrow-cast", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-cast 58.3.0", + "arrow-schema 58.3.0", "chrono", "csv", "csv-core", "regex", ] +[[package]] +name = "arrow-data" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78468c813909465dd0f858950c8a0614eb63608134acf95c602ec21381258b28" +dependencies = [ + "arrow-buffer 56.2.1", + "arrow-schema 56.2.1", + "half", + "num", +] + [[package]] name = "arrow-data" version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c88210023a2bfee1896af366309a3028fc3bcbd6515fa29a7990ee1baa08ee0" dependencies = [ - "arrow-buffer", - "arrow-schema", + "arrow-buffer 58.3.0", + "arrow-schema 58.3.0", "half", "num-integer", "num-traits", ] +[[package]] +name = "arrow-ipc" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31f88b0fbb33af28089ccd3e4dcd0ff09de46842168d00220b920f7231feddf5" +dependencies = [ + "arrow-array 56.2.1", + "arrow-buffer 56.2.1", + "arrow-data 56.2.1", + "arrow-schema 56.2.1", + "arrow-select 56.2.1", + "flatbuffers", +] + [[package]] name = "arrow-ipc" version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "238438f0834483703d88896db6fe5a7138b2230debc31b34c0336c2996e3c64f" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "arrow-select", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "flatbuffers", - "lz4_flex", + "lz4_flex 0.13.1", "zstd", ] @@ -358,12 +464,12 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "205ca2119e6d679d5c133c6f30e68f027738d95ed948cf77677ea69c7800036b" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-ord", - "arrow-schema", - "arrow-select", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-cast 58.3.0", + "arrow-ord 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "chrono", "half", "indexmap 2.14.0", @@ -377,17 +483,43 @@ dependencies = [ "simdutf8", ] +[[package]] +name = "arrow-ord" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aed58a38c3db0a2cf75ef70e3cb6bc4bd0da0a3d390de37c36139b31fae826e8" +dependencies = [ + "arrow-array 56.2.1", + "arrow-buffer 56.2.1", + "arrow-data 56.2.1", + "arrow-schema 56.2.1", + "arrow-select 56.2.1", +] + [[package]] name = "arrow-ord" version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1bffd8fd2579286a5d63bac898159873e5094a79009940bcb42bbfce4f19f1d0" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "arrow-select", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", +] + +[[package]] +name = "arrow-row" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "079ced0517daf4f09b070d09ff641cee7cc331aa216bebcb25d1a6474ad53086" +dependencies = [ + "arrow-array 56.2.1", + "arrow-buffer 56.2.1", + "arrow-data 56.2.1", + "arrow-schema 56.2.1", + "half", ] [[package]] @@ -396,24 +528,44 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bab5994731204603c73ba69267616c50f80780774c6bb0476f1f830625115e0c" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", "half", ] +[[package]] +name = "arrow-schema" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a0d5eb3fe25337ff83e8333a08379bdd1540b0961b1c888f6e505d971c198e1" + [[package]] name = "arrow-schema" version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f633dbfdf39c039ada1bf9e34c694816eb71fbb7dc78f613993b7245e078a1ed" dependencies = [ - "bitflags", + "bitflags 2.13.0", "serde_core", "serde_json", ] +[[package]] +name = "arrow-select" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2368a78bd32902dba39d52519d70f63799c8b5dc8a9477129a30c2fd3dc70c19" +dependencies = [ + "ahash 0.8.12", + "arrow-array 56.2.1", + "arrow-buffer 56.2.1", + "arrow-data 56.2.1", + "arrow-schema 56.2.1", + "num", +] + [[package]] name = "arrow-select" version = "58.3.0" @@ -421,24 +573,41 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8cd065c54172ac787cf3f2f8d4107e0d3fdc26edba76fdf4f4cc170258942222" dependencies = [ "ahash 0.8.12", - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", "num-traits", ] +[[package]] +name = "arrow-string" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dece58a130b9187756ded8bc071bd8ee9dd7a146566af244b297c7e632fd1ef7" +dependencies = [ + "arrow-array 56.2.1", + "arrow-buffer 56.2.1", + "arrow-data 56.2.1", + "arrow-schema 56.2.1", + "arrow-select 56.2.1", + "memchr", + "num", + "regex", + "regex-syntax", +] + [[package]] name = "arrow-string" version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "29dd7cda3ab9692f43a2e4acc444d760cc17b12bb6d8232ddf64e9bab7c06b42" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "arrow-select", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "memchr", "num-traits", "regex", @@ -560,7 +729,7 @@ checksum = "3b43422f69d8ff38f95f1b2bb76517c91589a924d1559a0e935d7c8ce0274c11" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -600,7 +769,7 @@ checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -617,7 +786,7 @@ checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -684,32 +853,63 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "beamterm-core" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0375f7f3df163a52a06807f2150bac9a5da1933864ed40d52dda7dcd7cfa193b" +dependencies = [ + "beamterm-data", + "beamterm-unicode", + "bitflags 2.13.0", + "compact_str", + "glow", + "lru 0.16.4", + "rustc-hash", + "thiserror 2.0.18", + "unicode-width 0.2.2", +] + [[package]] name = "beamterm-data" -version = "0.10.0" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bf74a7c20e83a40e542c032c918f1b5653ad4a137b49b7063e61030983ce243" +checksum = "f21ff1fc630079352bae9b024f85519bf1f641cf7f326623f4c0b59f7ea834fd" dependencies = [ "compact_str", - "miniz_oxide", + "miniz_oxide 0.9.1", + "thiserror 2.0.18", ] [[package]] name = "beamterm-renderer" -version = "0.10.0" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c27cc93582c7ce0f7350adcd2e0ffe21237687ba6bfb67cdea15e6e37001d19f" +checksum = "d8db76d04d575a0534ad5dba63acc05ccbe69597f977f11c59a612cbea199112" dependencies = [ + "beamterm-core", "beamterm-data", + "bitflags 2.13.0", "compact_str", "console_error_panic_hook", + "glow", "js-sys", "thiserror 2.0.18", + "unicode-width 0.2.2", "wasm-bindgen", "wasm-bindgen-futures", "web-sys", ] +[[package]] +name = "beamterm-unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f44155933cdd0625a699c554c504aab30a571e4df68f3532fb95bf9cf53b935" +dependencies = [ + "unicode-width 0.2.2", +] + [[package]] name = "better_io" version = "0.2.0" @@ -735,7 +935,7 @@ version = "0.72.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "993776b509cfb49c750f11b8f07a46fa23e0a1386ffc01fb1e7d343efc387895" dependencies = [ - "bitflags", + "bitflags 2.13.0", "cexpr", "clang-sys", "itertools 0.13.0", @@ -746,7 +946,7 @@ dependencies = [ "regex", "rustc-hash", "shlex 1.3.0", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -758,6 +958,12 @@ dependencies = [ "serde", ] +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.13.0" @@ -775,9 +981,9 @@ dependencies = [ [[package]] name = "bitvec" -version = "1.0.1" +version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bc2832c24239b0141d5674bb9174f9d68a8b5b3f2753311927c172ca46f7e9c" +checksum = "ddcec3d12c579d40898fe0a9a358a803c23e9c52ca3c425707f81c9436211837" dependencies = [ "funty", "radium", @@ -819,9 +1025,9 @@ dependencies = [ [[package]] name = "block-buffer" -version = "0.12.0" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cdd35008169921d80bc60d3d0ab416eecb028c4cd653352907921d95084790be" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" dependencies = [ "hybrid-array", "zeroize", @@ -842,9 +1048,9 @@ dependencies = [ [[package]] name = "borsh" -version = "1.6.1" +version = "1.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfd1e3f8955a5d7de9fab72fc8373fade9fb8a703968cb200ae3dc6cf08e185a" +checksum = "2f3f6da4992df95bbcd9af42a6c7dcb994498fc9048230405f3b36ff7cd3f145" dependencies = [ "borsh-derive", "bytes", @@ -853,22 +1059,22 @@ dependencies = [ [[package]] name = "borsh-derive" -version = "1.6.1" +version = "1.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfcfdc083699101d5a7965e49925975f2f55060f94f9a05e7187be95d530ca59" +checksum = "3ae8fb4fb5740e4b2c4884ff95f5f32f5e8479db1e8fd8eb49ddbe09eb09bb7c" dependencies = [ "once_cell", "proc-macro-crate", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] name = "brotli" -version = "8.0.3" +version = "8.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8119e4516436f5708bbc474a9d395bf12f1b5395e93a92a56e647ac3388c8610" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" dependencies = [ "alloc-no-stdlib", "alloc-stdlib", @@ -877,9 +1083,9 @@ dependencies = [ [[package]] name = "brotli-decompressor" -version = "5.0.1" +version = "5.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5962523e1b92ce1b5e793d9169b9943eece10d39f62550bc04bb605d75b94924" +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583" dependencies = [ "alloc-no-stdlib", "alloc-stdlib", @@ -896,13 +1102,13 @@ dependencies = [ [[package]] name = "bstr" -version = "1.12.1" +version = "1.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63044e1ae8e69f3b5a92c736ca6269b8d12fa7efe39bf34ddb06d102cf0e2cab" +checksum = "5cee35f73844aa3014bb606320a6c1f010249dbdf43342fe54b5a4f6a8ed4b79" dependencies = [ "memchr", "regex-automata", - "serde", + "serde_core", ] [[package]] @@ -925,9 +1131,9 @@ checksum = "7575182f7272186991736b70173b0ea045398f984bf5ebbb3804736ce1330c9d" [[package]] name = "byte-unit" -version = "5.2.0" +version = "5.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8c6d47a4e2961fb8721bcfc54feae6455f2f64e7054f9bc67e875f0e77f4c58d" +checksum = "4a813de7f2bbedb7dce265b64f1cf5908ebe4d56281ece8d847e98113788b9b0" dependencies = [ "rust_decimal", "schemars 1.2.1", @@ -977,9 +1183,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.11.1" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" [[package]] name = "bzip2" @@ -992,9 +1198,9 @@ dependencies = [ [[package]] name = "camino" -version = "1.2.2" +version = "1.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e629a66d692cb9ff1a1c664e41771b3dcaf961985a9774c0eb0bd1b51cf60a48" +checksum = "5f2d30e4173c4026932d51d31d6b0613b1fd3014bf3f9f8943d4ba139c437ba0" dependencies = [ "serde_core", ] @@ -1040,9 +1246,9 @@ dependencies = [ [[package]] name = "cbindgen" -version = "0.29.3" +version = "0.29.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c95537b45400390270fae69ac098d057c8f5399001cde9d04f700c105ddfff2d" +checksum = "2ecb53484c9c167ba674026b656d8a27d7657a58e6066aa902bfb1a4aa00ae20" dependencies = [ "clap", "heck", @@ -1052,16 +1258,16 @@ dependencies = [ "quote", "serde", "serde_json", - "syn 2.0.117", + "syn 2.0.118", "tempfile", "toml", ] [[package]] name = "cc" -version = "1.2.63" +version = "1.2.65" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "556e016178bb5662a08681bbe0f00f8e17631781a4dfc8c45e466e4b185ec27f" +checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" dependencies = [ "find-msvc-tools", "jobserver", @@ -1092,9 +1298,9 @@ checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" [[package]] name = "chacha20" -version = "0.10.0" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" dependencies = [ "cfg-if", "cpufeatures 0.3.0", @@ -1205,7 +1411,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -1242,14 +1448,14 @@ dependencies = [ [[package]] name = "codspeed" -version = "4.7.0" +version = "5.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57af92d1db7f6871b7e82c79cd87f2501db66f36b0eab924be6ea83dd6b2f3f3" +checksum = "7083f253260bcb4aaa3b4aa4c52973703dabc1a85c2f193997e2689aafa8a919" dependencies = [ "anyhow", "cc", "colored", - "getrandom 0.2.17", + "getrandom 0.4.3", "glob", "libc", "nix", @@ -1260,9 +1466,9 @@ dependencies = [ [[package]] name = "codspeed-criterion-compat-walltime" -version = "4.7.0" +version = "5.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc8605e40bab5114dcb0f76268e18880082b5798dec10757b5b58d2c3bbc7a1c" +checksum = "5c38205d56e2cb4fe04b708de7f9653a3f1b89edbe3a20b28f21e9e525e9e061" dependencies = [ "anes", "cast", @@ -1287,9 +1493,9 @@ dependencies = [ [[package]] name = "codspeed-divan-compat" -version = "4.7.0" +version = "5.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4ea79fd0b1f2128cfac6308369013dba92df47baf4a4f66b57d8158224a361d" +checksum = "bc1065d507e1cbab731a7976db4cef7e47e49b87b4dbc0a925df07d343558420" dependencies = [ "clap", "codspeed", @@ -1300,23 +1506,23 @@ dependencies = [ [[package]] name = "codspeed-divan-compat-macros" -version = "4.7.0" +version = "5.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f70e4ddd6beedefeb48f59d5f85fc21365a66e7976408c3d39f6cbbc4f03e08c" +checksum = "2fd05482a95823ffe421e8a9ba24fa22a6a30d594e2c60455cbb43a41bf2d8fa" dependencies = [ "divan-macros", "itertools 0.14.0", "proc-macro-crate", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] name = "codspeed-divan-compat-walltime" -version = "4.7.0" +version = "5.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "490c04f6076be6eacfafb496b8b237f3efbbed93838f2689115cc6f35fcf81c9" +checksum = "d2f8eae75b8fa85357020a404899c4280d590c9fd1c47640b3ce53106c62d2ee" dependencies = [ "cfg-if", "clap", @@ -1335,12 +1541,11 @@ checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" [[package]] name = "colored" -version = "2.2.0" +version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" +checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "lazy_static", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -1355,11 +1560,12 @@ dependencies = [ [[package]] name = "comfy-table" -version = "7.1.4" +version = "7.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a65ebfec4fb190b6f90e944a817d60499ee0744e582530e2c9900a22e591d9a" +checksum = "e0d05af1e006a2407bedef5af410552494ce5be9090444dbbcb57258c1af3d56" dependencies = [ - "unicode-segmentation", + "strum 0.26.3", + "strum_macros 0.26.4", "unicode-width 0.2.2", ] @@ -1382,8 +1588,8 @@ name = "compress-bench" version = "0.1.0" dependencies = [ "anyhow", - "arrow-array", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", "async-trait", "bytes", "clap", @@ -1391,7 +1597,7 @@ dependencies = [ "indicatif", "itertools 0.14.0", "lance-bench", - "parquet", + "parquet 58.3.0", "regex", "tokio", "tracing", @@ -1692,7 +1898,7 @@ version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b" dependencies = [ - "bitflags", + "bitflags 2.13.0", "crossterm_winapi", "derive_more", "document-features", @@ -1771,9 +1977,9 @@ dependencies = [ [[package]] name = "cudarc" -version = "0.19.7" +version = "0.19.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1cea5f10a99e025c1b44ae2354c2d8326b25ddbd0baf76bde8e55cfd4018a2cc" +checksum = "42310153e06cf4cd532901f7096beb27504d681736a29ee90728ae4e2d93b2a8" dependencies = [ "half", "libloading 0.9.0", @@ -1818,7 +2024,7 @@ dependencies = [ "proc-macro2", "quote", "scratch", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -1832,7 +2038,7 @@ dependencies = [ "indexmap 2.14.0", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -1850,7 +2056,7 @@ dependencies = [ "indexmap 2.14.0", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -1873,7 +2079,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -1884,7 +2090,7 @@ checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" dependencies = [ "darling_core", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -1907,8 +2113,8 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "93db0e623840612f7f2cd757f7e8a8922064192363732c88692e0870016e141b" dependencies = [ - "arrow", - "arrow-schema", + "arrow 58.3.0", + "arrow-schema 58.3.0", "async-trait", "bytes", "chrono", @@ -1939,7 +2145,7 @@ dependencies = [ "futures", "itertools 0.14.0", "log", - "object_store 0.13.2", + "object_store", "parking_lot", "rand 0.9.4", "regex", @@ -1956,8 +2162,8 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "997a31e15872606a49478e670c58302094c97cb96abb0a7d60720f8e92170040" dependencies = [ - "arrow", - "arrow-schema", + "arrow 58.3.0", + "arrow-schema 58.3.0", "async-trait", "bzip2", "chrono", @@ -1993,9 +2199,9 @@ dependencies = [ "itertools 0.14.0", "liblzma", "log", - "object_store 0.13.2", + "object_store", "parking_lot", - "parquet", + "parquet 58.3.0", "sqlparser 0.62.0", "tempfile", "tokio", @@ -2016,7 +2222,7 @@ dependencies = [ "datafusion-physical-plan 54.0.0", "futures", "itertools 0.14.0", - "object_store 0.13.2", + "object_store", "opentelemetry", "opentelemetry-otlp", "opentelemetry_sdk", @@ -2036,7 +2242,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37cefde60b26a7f4ff61e9d2ff2833322f91df2b568d7238afe67bde5bdffb66" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "dashmap", "datafusion-common 53.1.0", @@ -2050,7 +2256,7 @@ dependencies = [ "futures", "itertools 0.14.0", "log", - "object_store 0.13.2", + "object_store", "parking_lot", "tokio", ] @@ -2061,7 +2267,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f7dd61161508f8f5fa1107774ea687bd753c22d83a32eebf963549f89de14139" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "dashmap", "datafusion-common 54.0.0", @@ -2075,7 +2281,7 @@ dependencies = [ "futures", "itertools 0.14.0", "log", - "object_store 0.13.2", + "object_store", "parking_lot", "tokio", ] @@ -2086,7 +2292,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "17e112307715d6a7a331111a4c2330ff54bc237183511c319e3708a4cff431fb" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "datafusion-catalog 53.1.0", "datafusion-common 53.1.0", @@ -2100,7 +2306,7 @@ dependencies = [ "futures", "itertools 0.14.0", "log", - "object_store 0.13.2", + "object_store", ] [[package]] @@ -2109,7 +2315,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "897c70f871277f9ce99aa38347be0d679bbe3e617156c4d2a8378cec8a2a0891" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "datafusion-catalog 54.0.0", "datafusion-common 54.0.0", @@ -2123,7 +2329,7 @@ dependencies = [ "futures", "itertools 0.14.0", "log", - "object_store 0.13.2", + "object_store", ] [[package]] @@ -2133,8 +2339,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d72a11ca44a95e1081870d3abb80c717496e8a7acb467a1d3e932bb636af5cc2" dependencies = [ "ahash 0.8.12", - "arrow", - "arrow-ipc", + "arrow 58.3.0", + "arrow-ipc 58.3.0", "chrono", "half", "hashbrown 0.16.1", @@ -2142,7 +2348,7 @@ dependencies = [ "itertools 0.14.0", "libc", "log", - "object_store 0.13.2", + "object_store", "paste", "sqlparser 0.61.0", "tokio", @@ -2155,9 +2361,9 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "121c9ded5d87d9172319e006f2afdb9928d72dbacd6a90a458d8acb1e3b43a65" dependencies = [ - "arrow", - "arrow-ipc", - "arrow-schema", + "arrow 58.3.0", + "arrow-ipc 58.3.0", + "arrow-schema 58.3.0", "chrono", "foldhash 0.2.0", "half", @@ -2166,8 +2372,8 @@ dependencies = [ "itertools 0.14.0", "libc", "log", - "object_store 0.13.2", - "parquet", + "object_store", + "parquet 58.3.0", "recursive", "sqlparser 0.62.0", "tokio", @@ -2203,7 +2409,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e9fb386e1691355355a96419978a0022b7947b44d4a24a6ea99f00b6b485cbb6" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "bytes", "chrono", @@ -2220,7 +2426,7 @@ dependencies = [ "glob", "itertools 0.14.0", "log", - "object_store 0.13.2", + "object_store", "rand 0.9.4", "tokio", "url", @@ -2232,7 +2438,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ffd7d295b2ec7c00d8a56562f41ed41062cf0af75549ed891c12a0a09eddfefe" dependencies = [ - "arrow", + "arrow 58.3.0", "async-compression", "async-trait", "bytes", @@ -2253,7 +2459,7 @@ dependencies = [ "itertools 0.14.0", "liblzma", "log", - "object_store 0.13.2", + "object_store", "parking_lot", "rand 0.9.4", "tokio", @@ -2268,8 +2474,8 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ffa6c52cfed0734c5f93754d1c0175f558175248bf686c944fb05c373e5fc096" dependencies = [ - "arrow", - "arrow-ipc", + "arrow 58.3.0", + "arrow-ipc 58.3.0", "async-trait", "bytes", "datafusion-common 53.1.0", @@ -2282,7 +2488,7 @@ dependencies = [ "datafusion-session 53.1.0", "futures", "itertools 0.14.0", - "object_store 0.13.2", + "object_store", "tokio", ] @@ -2292,8 +2498,8 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "552b0b3f342f7ec41b3fbd70f6339dc82a30cfd0349e7f280e7852528085349f" dependencies = [ - "arrow", - "arrow-ipc", + "arrow 58.3.0", + "arrow-ipc 58.3.0", "async-trait", "bytes", "datafusion-common 54.0.0", @@ -2306,7 +2512,7 @@ dependencies = [ "datafusion-session 54.0.0", "futures", "itertools 0.14.0", - "object_store 0.13.2", + "object_store", "tokio", ] @@ -2316,7 +2522,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fb517d08967d536284ce70afb5fe8583133779249f2d7b90587d339741a7f195" dependencies = [ - "arrow", + "arrow 58.3.0", "arrow-avro", "async-trait", "bytes", @@ -2326,7 +2532,7 @@ dependencies = [ "datafusion-physical-plan 54.0.0", "datafusion-session 54.0.0", "futures", - "object_store 0.13.2", + "object_store", ] [[package]] @@ -2335,7 +2541,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "503f29e0582c1fc189578d665ff57d9300da1f80c282777d7eb67bb79fb8cdca" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "bytes", "datafusion-common 53.1.0", @@ -2347,7 +2553,7 @@ dependencies = [ "datafusion-physical-plan 53.1.0", "datafusion-session 53.1.0", "futures", - "object_store 0.13.2", + "object_store", "regex", "tokio", ] @@ -2358,7 +2564,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68850aa426b897e879c8b87e512ea8124f1d0a2869a4e51808ddaaddf1bc0ada" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "bytes", "datafusion-common 54.0.0", @@ -2370,7 +2576,7 @@ dependencies = [ "datafusion-physical-plan 54.0.0", "datafusion-session 54.0.0", "futures", - "object_store 0.13.2", + "object_store", "regex", "tokio", ] @@ -2381,7 +2587,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e33804749abc8d0c8cb7473228483cb8070e524c6f6086ee1b85a64debe2b3d2" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "bytes", "datafusion-common 53.1.0", @@ -2393,7 +2599,7 @@ dependencies = [ "datafusion-physical-plan 53.1.0", "datafusion-session 53.1.0", "futures", - "object_store 0.13.2", + "object_store", "serde_json", "tokio", "tokio-stream", @@ -2405,7 +2611,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "402f93242ae08ef99139ee2c528a49d087efe88d5c7b2c3ff5480855a40ce54f" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "bytes", "datafusion-common 54.0.0", @@ -2417,7 +2623,7 @@ dependencies = [ "datafusion-physical-plan 54.0.0", "datafusion-session 54.0.0", "futures", - "object_store 0.13.2", + "object_store", "tokio", "tokio-stream", ] @@ -2428,7 +2634,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ffd2499c1bee0eeccf6a57156105700eeeb17bc701899ac719183c4e74231450" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "bytes", "datafusion-common 54.0.0", @@ -2447,9 +2653,9 @@ dependencies = [ "futures", "itertools 0.14.0", "log", - "object_store 0.13.2", + "object_store", "parking_lot", - "parquet", + "parquet 58.3.0", "tokio", ] @@ -2471,8 +2677,8 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c03c7fbdaefcca4ef6ffe425a5fc2325763bfb426599bb0bf4536466efabe709" dependencies = [ - "arrow", - "arrow-buffer", + "arrow 58.3.0", + "arrow-buffer 58.3.0", "async-trait", "chrono", "dashmap", @@ -2481,7 +2687,7 @@ dependencies = [ "datafusion-physical-expr-common 53.1.0", "futures", "log", - "object_store 0.13.2", + "object_store", "parking_lot", "rand 0.9.4", "tempfile", @@ -2494,8 +2700,8 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37a8643ab852eb68864e1b72ae789e8066282dce48eea6347ffb0aee33d1ccc0" dependencies = [ - "arrow", - "arrow-buffer", + "arrow 58.3.0", + "arrow-buffer 58.3.0", "async-trait", "dashmap", "datafusion-common 54.0.0", @@ -2503,7 +2709,7 @@ dependencies = [ "datafusion-physical-expr-common 54.0.0", "futures", "log", - "object_store 0.13.2", + "object_store", "parking_lot", "rand 0.9.4", "tempfile", @@ -2516,7 +2722,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "574b9b6977fedbd2a611cbff12e5caf90f31640ad9dc5870f152836d94bad0dd" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "chrono", "datafusion-common 53.1.0", @@ -2538,8 +2744,8 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6932f4d71eed9c8d9341476a2b845aadfabde5495d08dbcd8fc23881f49fa7a0" dependencies = [ - "arrow", - "arrow-schema", + "arrow 58.3.0", + "arrow-schema 58.3.0", "async-trait", "chrono", "datafusion-common 54.0.0", @@ -2561,7 +2767,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7d7c3adf3db8bf61e92eb90cb659c8e8b734593a8f7c8e12a843c7ddba24b87e" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 53.1.0", "indexmap 2.14.0", "itertools 0.14.0", @@ -2574,7 +2780,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0225491839a31b1f7d2cb8092c2d50792e2fe1c1724e4e6d08e011f5feaf4ed2" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 54.0.0", "indexmap 2.14.0", "itertools 0.14.0", @@ -2586,8 +2792,8 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f28aa4e10384e782774b10e72aca4d93ef7b31aa653095d9d4536b0a3dbc51b6" dependencies = [ - "arrow", - "arrow-buffer", + "arrow 58.3.0", + "arrow-buffer 58.3.0", "base64", "blake2", "blake3", @@ -2618,8 +2824,8 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "14872c47bfc3d21e53ec82f57074e6987a15941c1e2f43cde4ac6ae2746634e3" dependencies = [ - "arrow", - "arrow-buffer", + "arrow 58.3.0", + "arrow-buffer 58.3.0", "base64", "blake2", "blake3", @@ -2651,7 +2857,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "00aa6217e56098ba84e0a338176fe52f0a84cca398021512c6c8c5eff806d0ad" dependencies = [ "ahash 0.8.12", - "arrow", + "arrow 58.3.0", "datafusion-common 53.1.0", "datafusion-doc 53.1.0", "datafusion-execution 53.1.0", @@ -2672,7 +2878,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "75a2ca14e1b609be21e657e2d3130b2f446456b08393b377bb721a33952d2e09" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 54.0.0", "datafusion-doc 54.0.0", "datafusion-execution 54.0.0", @@ -2694,7 +2900,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b511250349407db7c43832ab2de63f5557b19a20dfd236b39ca2c04468b50d47" dependencies = [ "ahash 0.8.12", - "arrow", + "arrow 58.3.0", "datafusion-common 53.1.0", "datafusion-expr-common 53.1.0", "datafusion-physical-expr-common 53.1.0", @@ -2706,7 +2912,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ece74ba09092d2ef9c9b54a38445450aea292a1f8b04faf531936b723a24b3c" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 54.0.0", "datafusion-expr-common 54.0.0", "datafusion-physical-expr-common 54.0.0", @@ -2718,8 +2924,8 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ef13a858e20d50f0a9bb5e96e7ac82b4e7597f247515bccca4fdd2992df0212a" dependencies = [ - "arrow", - "arrow-ord", + "arrow 58.3.0", + "arrow-ord 58.3.0", "datafusion-common 53.1.0", "datafusion-doc 53.1.0", "datafusion-execution 53.1.0", @@ -2743,8 +2949,8 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f3e3f9ee8ca59bf70518802107de6f1b88a9509efdc629fadc5de9d6b2d5ef5" dependencies = [ - "arrow", - "arrow-ord", + "arrow 58.3.0", + "arrow-ord 58.3.0", "datafusion-common 54.0.0", "datafusion-doc 54.0.0", "datafusion-execution 54.0.0", @@ -2768,7 +2974,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b40d3f5bbb3905f9ccb1ce9485a9595c77b69758a7c24d3ba79e334ff51e7e" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "datafusion-catalog 53.1.0", "datafusion-common 53.1.0", @@ -2784,7 +2990,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "89161dffc22cf2b50f9f4b1bee83b5221d3b4ed7c2e37fd7aa2b22a5297b3a26" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "datafusion-catalog 54.0.0", "datafusion-common 54.0.0", @@ -2800,7 +3006,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4e88ec9d57c9b685d02f58bfee7be62d72610430ddcedb82a08e5d9925dbfb6" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 53.1.0", "datafusion-doc 53.1.0", "datafusion-expr 53.1.0", @@ -2818,7 +3024,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d7339345b226b3874037708bf5023ba1c2de705128f8457a095aae5ae9cb9c78" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 54.0.0", "datafusion-doc 54.0.0", "datafusion-expr 54.0.0", @@ -2857,7 +3063,7 @@ checksum = "2e367e6a71051d0ebdd29b2f85d12059b38b1d1f172c6906e80016da662226bd" dependencies = [ "datafusion-doc 53.1.0", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -2868,7 +3074,7 @@ checksum = "587164e03ad68732aa9e7bfe5686e3f25970d4c64fd4bd80790749840892dae5" dependencies = [ "datafusion-doc 54.0.0", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -2877,7 +3083,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e929015451a67f77d9d8b727b2bf3a40c4445fdef6cdc53281d7d97c76888ace" dependencies = [ - "arrow", + "arrow 58.3.0", "chrono", "datafusion-common 53.1.0", "datafusion-expr 53.1.0", @@ -2896,7 +3102,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77f20e8cf9e8654d92f4c16b24c487353ee5bf153ffc12d5772cd399ab8cd281" dependencies = [ - "arrow", + "arrow 58.3.0", "chrono", "datafusion-common 54.0.0", "datafusion-expr 54.0.0", @@ -2917,7 +3123,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4b1e68aba7a4b350401cfdf25a3d6f989ad898a7410164afe9ca52080244cb59" dependencies = [ "ahash 0.8.12", - "arrow", + "arrow 58.3.0", "datafusion-common 53.1.0", "datafusion-expr 53.1.0", "datafusion-expr-common 53.1.0", @@ -2939,7 +3145,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f015a4a82f6f7ff7e1d8d4bf3870a936752fa38b17705dfcc14adef95aa8922c" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 54.0.0", "datafusion-expr 54.0.0", "datafusion-expr-common 54.0.0", @@ -2961,7 +3167,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ea22315f33cf2e0adc104e8ec42e285f6ed93998d565c65e82fec6a9ee9f9db4" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 53.1.0", "datafusion-expr 53.1.0", "datafusion-functions 53.1.0", @@ -2976,7 +3182,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "51e6ffff8acdfe54e0ea15ccf38115c4a9184433b0439f42907637928d00a235" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 54.0.0", "datafusion-expr 54.0.0", "datafusion-functions 54.0.0", @@ -2992,7 +3198,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b04b45ea8ad3ac2d78f2ea2a76053e06591c9629c7a603eda16c10649ecf4362" dependencies = [ "ahash 0.8.12", - "arrow", + "arrow 58.3.0", "chrono", "datafusion-common 53.1.0", "datafusion-expr-common 53.1.0", @@ -3008,7 +3214,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7967a3e171c6a4bf09474b3f7a14f1a3db13ed1714ba12156f33fcce2bba54e8" dependencies = [ - "arrow", + "arrow 58.3.0", "chrono", "datafusion-common 54.0.0", "datafusion-expr-common 54.0.0", @@ -3025,7 +3231,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cb13397809a425918f608dfe8653f332015a3e330004ab191b4404187238b95" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 53.1.0", "datafusion-execution 53.1.0", "datafusion-expr 53.1.0", @@ -3043,7 +3249,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "59ff803e2a96054cb6d83f35f9e60fd4f42eac515e1932bd1b2dbc91d5fcbf36" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 54.0.0", "datafusion-execution 54.0.0", "datafusion-expr 54.0.0", @@ -3063,9 +3269,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5edc023675791af9d5fb4cc4c24abf5f7bd3bd4dcf9e5bd90ea1eff6976dcc79" dependencies = [ "ahash 0.8.12", - "arrow", - "arrow-ord", - "arrow-schema", + "arrow 58.3.0", + "arrow-ord 58.3.0", + "arrow-schema 58.3.0", "async-trait", "datafusion-common 53.1.0", "datafusion-common-runtime 53.1.0", @@ -3094,11 +3300,11 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "776ee54d47d15bdb126452f9ca17b03761e3b004682914beaedd3f86eb507fbc" dependencies = [ - "arrow", - "arrow-data", - "arrow-ipc", - "arrow-ord", - "arrow-schema", + "arrow 58.3.0", + "arrow-data 58.3.0", + "arrow-ipc 58.3.0", + "arrow-ord 58.3.0", + "arrow-schema 58.3.0", "async-trait", "datafusion-common 54.0.0", "datafusion-common-runtime 54.0.0", @@ -3127,7 +3333,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac8c76860e355616555081cab5968cec1af7a80701ff374510860bcd567e365a" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 53.1.0", "datafusion-datasource 53.1.0", "datafusion-expr-common 53.1.0", @@ -3144,7 +3350,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d5fb9e5774660aa69c3ba93c610f175f75b65cb8c3776edb3626de8f3a4f4ee3" dependencies = [ - "arrow", + "arrow 58.3.0", "datafusion-common 54.0.0", "datafusion-datasource 54.0.0", "datafusion-expr-common 54.0.0", @@ -3188,7 +3394,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "390bb0bf37cb2b95ffd65eacd66f60df50793d1f94097799e416f39477a51957" dependencies = [ - "arrow", + "arrow 58.3.0", "bigdecimal", "chrono", "crc32fast", @@ -3218,7 +3424,7 @@ version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fa0d133ddf8b9b3b872acac900157f783e7b879fe9a6bccf389abebbfac45ec1" dependencies = [ - "arrow", + "arrow 58.3.0", "bigdecimal", "chrono", "datafusion-common 53.1.0", @@ -3236,7 +3442,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6094ad36a3ed6d7ac87b20b479b2d0b118250f66cf997603829fdc65b44a7099" dependencies = [ - "arrow", + "arrow 58.3.0", "bigdecimal", "chrono", "datafusion-common 54.0.0", @@ -3255,7 +3461,7 @@ version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e0c08025966108056d3547d879c4d39e246277494f59ca12920f78187d95eea1" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "bigdecimal", "clap", @@ -3266,7 +3472,7 @@ dependencies = [ "indicatif", "itertools 0.14.0", "log", - "object_store 0.13.2", + "object_store", "sqllogictest", "sqlparser 0.62.0", "tempfile", @@ -3300,13 +3506,44 @@ version = "0.1.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6b926516df9c60bfa16e107b21086399f8285a44ca9711344b9e553c5146e2" +[[package]] +name = "defmt" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6e524506490a1953d237cb87b1cfc1e46f88c18f10a22dfe0f507dc6bfc7f7f" +dependencies = [ + "bitflags 1.3.2", + "defmt-macros", +] + +[[package]] +name = "defmt-macros" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0a27770e9c8f719a79d8b638281f4d828f77d8fd61e0bd94451b9b85e576a0b" +dependencies = [ + "defmt-parser", + "proc-macro-error2", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "defmt-parser" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" +dependencies = [ + "thiserror 2.0.18", +] + [[package]] name = "deranged" version = "0.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" dependencies = [ - "powerfmt", "serde_core", ] @@ -3329,7 +3566,7 @@ dependencies = [ "proc-macro2", "quote", "rustc_version", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -3349,7 +3586,7 @@ version = "0.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" dependencies = [ - "block-buffer 0.12.0", + "block-buffer 0.12.1", "const-oid", "crypto-common 0.2.2", "ctutils", @@ -3385,7 +3622,7 @@ checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -3396,7 +3633,7 @@ checksum = "8dc51d98e636f5e3b0759a39257458b22619cac7e96d932da6eeb052891bb67c" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -3441,6 +3678,16 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" +[[package]] +name = "earcutr" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79127ed59a85d7687c409e9978547cffb7dc79675355ed22da6b66fd5f6ead01" +dependencies = [ + "itertools 0.11.0", + "num-traits", +] + [[package]] name = "educe" version = "0.6.0" @@ -3450,7 +3697,7 @@ dependencies = [ "enum-ordinalize", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -3491,34 +3738,34 @@ checksum = "685adfa4d6f3d765a26bc5dbc936577de9abf756c1feeb3089b01dd395034842" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] name = "enum-ordinalize" -version = "4.3.2" +version = "4.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a1091a7bb1f8f2c4b28f1fe2cef4980ca2d410a3d727d67ecc3178c9b0800f0" +checksum = "07f808d588c10e464ea6f7d3eaed500049eff30aaac103460f61828c2d65b3eb" dependencies = [ "enum-ordinalize-derive", ] [[package]] name = "enum-ordinalize-derive" -version = "4.3.2" +version = "4.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ca9601fb2d62598ee17836250842873a413586e5d7ed88b356e38ddbb0ec631" +checksum = "42e528e2d34ba8a67a1a650b86beae8ef69fc5fdb638016f386b973226590432" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] name = "env_filter" -version = "1.0.1" +version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32e90c2accc4b07a8456ea0debdc2e7587bdd890680d71173a15d4ae604f6eef" +checksum = "900d271a03799a1ee8d1ca9b19893b48ca674a9284fefcfb85f05e74ed314217" dependencies = [ "log", "regex", @@ -3526,9 +3773,9 @@ dependencies = [ [[package]] name = "env_logger" -version = "0.11.10" +version = "0.11.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0621c04f2196ac3f488dd583365b9c09be011a4ab8b9f37248ffcc8f6198b56a" +checksum = "de671bd27a75a797dc9ae289ba1e77276e75e2026408aab65185384e2d5cd3f6" dependencies = [ "anstream", "anstyle", @@ -3629,11 +3876,10 @@ checksum = "dd2e7510819d6fbf51a5545c8f922716ecfb14df168a3242f7d33e0239efe6a1" [[package]] name = "fastlanes" -version = "0.5.1" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "20c597e23b8ec8506f589d18bc701ca83a3def6086748f628ad23092e1dfe577" +checksum = "b9272a674b446f53253f66582596614e41552d14968a3d1aaa590576a91c1188" dependencies = [ - "arrayref", "const_for", "core_detect", "num-traits", @@ -3688,7 +3934,7 @@ version = "25.12.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" dependencies = [ - "bitflags", + "bitflags 2.13.0", "rustc_version", ] @@ -3699,10 +3945,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ "crc32fast", - "miniz_oxide", + "miniz_oxide 0.8.9", "zlib-rs", ] +[[package]] +name = "float_next_after" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8bf7cc16383c4b8d58b9905a8509f02926ce3058053c056376248d958c9df1e8" + [[package]] name = "fnv" version = "1.0.7" @@ -3747,11 +3999,11 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "6.0.1" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "83cf860f6a6bf0a6a60fdfe5a36c75121fad5ea4332d1d12deee3e65b6047727" +checksum = "bcd0ce0249ac12fd44fcde62d435c36d881952c2f0df4d1de24b45e1dbba5ddb" dependencies = [ - "arrow-array", + "arrow-array 58.3.0", "rand 0.9.4", ] @@ -3848,7 +4100,7 @@ checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -3920,6 +4172,24 @@ dependencies = [ "version_check", ] +[[package]] +name = "geo" +version = "0.31.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fc1a1678e54befc9b4bcab6cd43b8e7f834ae8ea121118b0fd8c42747675b4a" +dependencies = [ + "earcutr", + "float_next_after", + "geo-types", + "geographiclib-rs", + "i_overlay", + "log", + "num-traits", + "robust", + "rstar", + "spade", +] + [[package]] name = "geo-traits" version = "0.3.0" @@ -3937,6 +4207,8 @@ checksum = "94776032c45f950d30a13af6113c2ad5625316c9abfbccee4dd5a6695f8fe0f5" dependencies = [ "approx", "num-traits", + "rayon", + "rstar", "serde", ] @@ -3956,14 +4228,14 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dafe7b7de3fab1a8b7099fd6a6434ca955fa65065f9c19f0f8a133693f3c2b0e" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-schema 58.3.0", "geo-traits", "geoarrow-schema", "num-traits", "wkb", - "wkt", + "wkt 0.14.0", ] [[package]] @@ -3972,13 +4244,49 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4d4a7edb2a1d87024a93805332a9c8184a0354836271d42c0d18cf628a5e3cd0" dependencies = [ - "arrow-schema", + "arrow-schema 58.3.0", "geo-traits", "serde", "serde_json", "thiserror 1.0.69", ] +[[package]] +name = "geographiclib-rs" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5a7f08910fd98737a6eda7568e7c5e645093e073328eeef49758cfe8b0489c7" +dependencies = [ + "libm", +] + +[[package]] +name = "geojson" +version = "0.24.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e26f3c45b36fccc9cf2805e61d4da6bc4bbd5a3a9589b01afa3a40eff703bd79" +dependencies = [ + "log", + "serde", + "serde_json", + "thiserror 2.0.18", +] + +[[package]] +name = "geozero" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5f28f34864745eb2f123c990c6ffd92c1584bd39439b3f27ff2a0f4ea5b309b" +dependencies = [ + "geo-types", + "geojson", + "log", + "scroll", + "serde_json", + "thiserror 1.0.69", + "wkt 0.11.1", +] + [[package]] name = "get_dir" version = "0.5.0" @@ -4014,17 +4322,15 @@ dependencies = [ [[package]] name = "getrandom" -version = "0.4.2" +version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", "js-sys", "libc", "r-efi 6.0.0", "rand_core 0.10.1", - "wasip2", - "wasip3", "wasm-bindgen", ] @@ -4034,6 +4340,18 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" +[[package]] +name = "glow" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29038e1c483364cc6bb3cf78feee1816002e127c331a1eec55a4d202b9e1adb5" +dependencies = [ + "js-sys", + "slotmap", + "wasm-bindgen", + "web-sys", +] + [[package]] name = "goldenfile" version = "1.11.0" @@ -4070,9 +4388,9 @@ checksum = "b40ca9252762c466af32d0b1002e91e4e1bc5398f77455e55474deb466355ff5" [[package]] name = "h2" -version = "0.4.14" +version = "0.4.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "171fefbc92fe4a4de27e0698d6a5b392d6a0e333506bc49133760b3bcf948733" +checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" dependencies = [ "atomic-waker", "bytes", @@ -4101,6 +4419,15 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "hash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" +dependencies = [ + "byteorder", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -4151,6 +4478,16 @@ dependencies = [ "foldhash 0.2.0", ] +[[package]] +name = "heapless" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" +dependencies = [ + "hash32", + "stable_deref_trait", +] + [[package]] name = "heck" version = "0.5.0" @@ -4180,9 +4517,9 @@ dependencies = [ [[package]] name = "http" -version = "1.4.1" +version = "1.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8be7462df143984c4598a256ef469b251d7d7f9e271135073e78fc535414f3d0" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" dependencies = [ "bytes", "itoa", @@ -4234,9 +4571,9 @@ checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" [[package]] name = "hybrid-array" -version = "0.4.12" +version = "0.4.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9155a582abd142abc056962c29e3ce5ff2ad5469f4246b537ed42c5deba857da" +checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" dependencies = [ "typenum", ] @@ -4312,6 +4649,49 @@ dependencies = [ "serde", ] +[[package]] +name = "i_float" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "010025c2c532c8d82e42d0b8bb5184afa449fa6f06c709ea9adcb16c49ae405b" +dependencies = [ + "libm", +] + +[[package]] +name = "i_key_sort" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9190f86706ca38ac8add223b2aed8b1330002b5cdbbce28fb58b10914d38fc27" + +[[package]] +name = "i_overlay" +version = "4.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413183068e6e0289e18d7d0a1f661b81546e6918d5453a44570b9ab30cbed1b3" +dependencies = [ + "i_float", + "i_key_sort", + "i_shape", + "i_tree", + "rayon", +] + +[[package]] +name = "i_shape" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ea154b742f7d43dae2897fcd5ead86bc7b5eefcedd305a7ebf9f69d44d61082" +dependencies = [ + "i_float", +] + +[[package]] +name = "i_tree" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35e6d558e6d4c7b82bc51d9c771e7a927862a161a7d87bf2b0541450e0e20915" + [[package]] name = "iana-time-zone" version = "0.1.65" @@ -4418,12 +4798,6 @@ dependencies = [ "zerovec", ] -[[package]] -name = "id-arena" -version = "2.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" - [[package]] name = "ident_case" version = "1.0.1" @@ -4468,7 +4842,7 @@ checksum = "a0eb5a3343abf848c0984fe4604b2b105da9539376e24fc0a3b0007411ae4fd9" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -4528,9 +4902,9 @@ dependencies = [ [[package]] name = "insta" -version = "1.47.2" +version = "1.48.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b4a6248eb93a4401ed2f37dfe8ea592d3cf05b7cf4f8efa867b6895af7e094e" +checksum = "86f0f8fee8c926415c58d6ae43a08523a26faccb2323f5e6b644fe7dd4ef6b82" dependencies = [ "console 0.16.3", "once_cell", @@ -4548,7 +4922,7 @@ dependencies = [ "indoc", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -4568,11 +4942,11 @@ dependencies = [ [[package]] name = "io-uring" -version = "0.7.12" +version = "0.7.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d09b98f7eace8982db770e4408e7470b028ce513ac28fecdc6bf4c30fe92b62" +checksum = "9080b15e63775b9a2ac7dca720f7050a8b955e092ea0f6020a4a80f69998cdc0" dependencies = [ - "bitflags", + "bitflags 2.13.0", "cfg-if", "libc", ] @@ -4609,6 +4983,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.12.1" @@ -4653,10 +5036,11 @@ dependencies = [ [[package]] name = "jiff" -version = "0.2.28" +version = "0.2.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4603d3033e49e2b0e31229fcab20a5d40089c607d975cd9c80551dc69eed9102" +checksum = "34f877a98676d2fb664698d74cc6a51ce6c484ce8c770f05d0108ec9090aeb46" dependencies = [ + "defmt", "jiff-static", "jiff-tzdb-platform", "log", @@ -4668,13 +5052,13 @@ dependencies = [ [[package]] name = "jiff-static" -version = "0.2.28" +version = "0.2.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "782d32378dddf207193ac91cefb848ad41abb58195c95168e1291227a0832b47" +checksum = "0666b5ab5ecaca213fc2a85b8c0083d9004e84ee2d5f9a7e0017aaf50986f25f" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -4721,7 +5105,7 @@ dependencies = [ "quote", "rustc_version", "simd_cesu8", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -4740,7 +5124,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" dependencies = [ "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -4755,13 +5139,12 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.99" +version = "0.3.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "142bc4740e452c1e57ade0cbc129f139c9093e354346f0872ef985f4f5cf5f11" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" dependencies = [ "cfg-if", "futures-util", - "once_cell", "wasm-bindgen", ] @@ -4823,26 +5206,29 @@ checksum = "a4933f3f57a8e9d9da04db23fb153356ecaf00cbd14aee46279c33dc80925c37" [[package]] name = "lance" -version = "6.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d34e854994e84d043897f5ec9fb609221e9e69e3fd52996cd715d979fcd349f6" -dependencies = [ - "arrow", - "arrow-arith", - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-ipc", - "arrow-ord", - "arrow-row", - "arrow-schema", - "arrow-select", +version = "7.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3944aca86f4c78f4da04af1c2bf33e664a2826b7af72972ad200d6b9de59019f" +dependencies = [ + "arc-swap", + "arrow 58.3.0", + "arrow-arith 58.3.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-cast 58.3.0", + "arrow-ipc 58.3.0", + "arrow-ord 58.3.0", + "arrow-row 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "async-recursion", "async-trait", "async_cell", + "bitpacking", "byteorder", "bytes", "chrono", + "crossbeam-queue", "crossbeam-skiplist", "dashmap", "datafusion 53.1.0", @@ -4869,13 +5255,14 @@ dependencies = [ "lance-tokenizer", "log", "moka", - "object_store 0.12.5", + "object_store", "permutation", "pin-project", "prost 0.14.4", "prost-build", "prost-types", "rand 0.9.4", + "rayon", "roaring", "semver", "serde", @@ -4891,18 +5278,17 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "6.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7827fe404358c27d120ee8ea8ef7b9415c2911d54072bec83dd689d750ae65da" -dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-data", - "arrow-ipc", - "arrow-ord", - "arrow-schema", - "arrow-select", +version = "7.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "253f4a0a70580c985b91e65e9ca6cad644825a4078de28d8efbacf3ffbd7ecdc" +dependencies = [ + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-ipc 58.3.0", + "arrow-ord 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "bytes", "futures", "getrandom 0.2.17", @@ -4917,13 +5303,13 @@ name = "lance-bench" version = "0.1.0" dependencies = [ "anyhow", - "arrow-cast", + "arrow-cast 58.3.0", "async-trait", "clap", "futures", "lance", "lance-encoding", - "parquet", + "parquet 58.3.0", "tempfile", "tokio", "tracing", @@ -4932,9 +5318,9 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "6.0.1" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2cd0b31570d50fe13c7e4e36b03e1f1c99c3d8e5a34845b24b0665b51b40570d" +checksum = "80c4d12521b1945041dd515a56aa0854973138e7ac12111c92572e33e4ecb593" dependencies = [ "arrayref", "paste", @@ -4943,17 +5329,16 @@ dependencies = [ [[package]] name = "lance-core" -version = "6.0.1" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b128c213c676cb8e03c62a68670642770825171e64097cc2da97cbb19fe35d29" +checksum = "13f84020da5a484e2f07dd1796e09785ed7cd889857ebc4cb77e32ef214ee594" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-schema 58.3.0", "async-trait", "byteorder", "bytes", - "chrono", "datafusion-common 53.1.0", "datafusion-sql 53.1.0", "deepsize", @@ -4962,10 +5347,9 @@ dependencies = [ "lance-arrow", "libc", "log", - "mock_instant", "moka", "num_cpus", - "object_store 0.12.5", + "object_store", "pin-project", "prost 0.14.4", "rand 0.9.4", @@ -4982,17 +5366,17 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "6.0.1" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e03b2de71cbcd09b10bf1a17c83cacbc0176ecd97203fb72b9e59d9b8f9a3743" +checksum = "7460597a66534a75987993d4dac5bc330586d99c5b79ae73367dbcbd4e29e576" dependencies = [ - "arrow", - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-ord", - "arrow-schema", - "arrow-select", + "arrow 58.3.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-cast 58.3.0", + "arrow-ord 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "async-trait", "chrono", "datafusion 53.1.0", @@ -5008,21 +5392,20 @@ dependencies = [ "pin-project", "prost 0.14.4", "prost-build", - "snafu", "tokio", "tracing", ] [[package]] name = "lance-datagen" -version = "6.0.1" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2fe7c7ea7fd397e495a1646fec360e46ee0cbd75718f1c0e887aad657c5f2944" +checksum = "046f5506ed2271cd941a050de7bf535dd3aedc291aadec836a63fa56c5926e3b" dependencies = [ - "arrow", - "arrow-array", - "arrow-cast", - "arrow-schema", + "arrow 58.3.0", + "arrow-array 58.3.0", + "arrow-cast 58.3.0", + "arrow-schema 58.3.0", "chrono", "futures", "half", @@ -5035,17 +5418,17 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "6.0.1" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe3f8070835b407d8db9ea8728386bc3207ba23c66a9c22d344e231ef12b77ca" +checksum = "7af54edf43dcf9d6a56cc636eb35d457e68373c6448dca3f0891b3325b4a24e6" dependencies = [ - "arrow-arith", - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-data", - "arrow-schema", - "arrow-select", + "arrow-arith 58.3.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-cast 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "bytemuck", "byteorder", "bytes", @@ -5062,9 +5445,7 @@ dependencies = [ "num-traits", "prost 0.14.4", "prost-build", - "prost-types", "rand 0.9.4", - "snafu", "strum 0.26.3", "tokio", "tracing", @@ -5074,16 +5455,16 @@ dependencies = [ [[package]] name = "lance-file" -version = "6.0.1" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6dfcf654549330df3aef708cd7c12e170feecddd34d6c19dd005b4153213268" +checksum = "0772ae2d6207995dc1eb28aff9507f78e90b3362b58f311da001e9dc25f3d736" dependencies = [ - "arrow-arith", - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "arrow-select", + "arrow-arith 58.3.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "async-recursion", "async-trait", "byteorder", @@ -5097,27 +5478,27 @@ dependencies = [ "lance-io", "log", "num-traits", - "object_store 0.12.5", + "object_store", "prost 0.14.4", "prost-build", "prost-types", - "snafu", "tokio", "tracing", ] [[package]] name = "lance-index" -version = "6.0.1" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fb8ad0bd10efa2608634a2518b7dd501231e76c56a65fbd6519e23914cc425a" +checksum = "e71fbfb51096a903cb524fe0da716f5f15fbc4a6b6f84cd6dec21abf319c5e84" dependencies = [ - "arrow", - "arrow-arith", - "arrow-array", - "arrow-ord", - "arrow-schema", - "arrow-select", + "arc-swap", + "arrow 58.3.0", + "arrow-arith 58.3.0", + "arrow-array 58.3.0", + "arrow-ord 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "async-channel", "async-recursion", "async-trait", @@ -5130,7 +5511,6 @@ dependencies = [ "datafusion-common 53.1.0", "datafusion-expr 53.1.0", "datafusion-physical-expr 53.1.0", - "datafusion-sql 53.1.0", "deepsize", "dirs", "fst", @@ -5152,7 +5532,7 @@ dependencies = [ "log", "ndarray", "num-traits", - "object_store 0.12.5", + "object_store", "prost 0.14.4", "prost-build", "prost-types", @@ -5164,7 +5544,6 @@ dependencies = [ "serde", "serde_json", "smallvec", - "snafu", "tempfile", "tokio", "tracing", @@ -5174,18 +5553,18 @@ dependencies = [ [[package]] name = "lance-io" -version = "6.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef5314703fa8c8baed04193cc669da80ab42521c6319d3cc921a4a997690dcc0" -dependencies = [ - "arrow", - "arrow-arith", - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-data", - "arrow-schema", - "arrow-select", +version = "7.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bab8c98ef1b870b20541d27f3ca4efdf7c9f5c25214233be07d231ba88900219" +dependencies = [ + "arrow 58.3.0", + "arrow-arith 58.3.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-cast 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "async-recursion", "async-trait", "byteorder", @@ -5198,16 +5577,14 @@ dependencies = [ "lance-arrow", "lance-core", "lance-namespace", - "libc", "log", "moka", - "object_store 0.12.5", + "object_store", "path_abs", "pin-project", "prost 0.14.4", "rand 0.9.4", "serde", - "snafu", "tempfile", "tokio", "tracing", @@ -5216,13 +5593,13 @@ dependencies = [ [[package]] name = "lance-linalg" -version = "6.0.1" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51aa9b73279f505b2bec0f194c7a2390ca74ad3260131e631a7bef8d97d54b2e" +checksum = "6b4c51cad0ac780b02dc4da48528479e7693c03e8d05390510bbc69ca2a9a1f1" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-schema 58.3.0", "cc", "deepsize", "half", @@ -5234,16 +5611,15 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "6.0.1" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39cd01581f55ce45c49cbe494ee86c7ba7ca4ca3654690fd820941cd9105a46e" +checksum = "014e8332ca0615506342e0d3af608639864b68396973be14239f09c9f21f1fc2" dependencies = [ - "arrow", + "arrow 58.3.0", "async-trait", "bytes", "lance-core", "lance-namespace-reqwest-client", - "serde", "snafu", ] @@ -5263,15 +5639,15 @@ dependencies = [ [[package]] name = "lance-table" -version = "6.0.1" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5db70650465a1af174b7dfe6948ec91a3d466ada12e11274eb66e51132173aa0" +checksum = "b16f1355904aea4ebb04ffc70c58c97901e10bde44452b4b021de4a1f329250d" dependencies = [ - "arrow", - "arrow-array", - "arrow-buffer", - "arrow-ipc", - "arrow-schema", + "arrow 58.3.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-ipc 58.3.0", + "arrow-schema 58.3.0", "async-trait", "byteorder", "bytes", @@ -5283,7 +5659,7 @@ dependencies = [ "lance-file", "lance-io", "log", - "object_store 0.12.5", + "object_store", "prost 0.14.4", "prost-build", "prost-types", @@ -5302,9 +5678,9 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "6.0.1" +version = "7.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb08ef9382c9d58036c323db2c19cc097e02d1d0d87714fc7176b5d3b36a31aa" +checksum = "b39b7f5ed9d0c0b716bf599b559d888267ed1dfe4c4e29d3648b51d2a28940cf" dependencies = [ "rust-stemmers", "serde", @@ -5327,12 +5703,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "leb128fmt" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" - [[package]] name = "lending-iterator" version = "0.1.7" @@ -5459,18 +5829,18 @@ dependencies = [ [[package]] name = "liblzma" -version = "0.4.6" +version = "0.4.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6033b77c21d1f56deeae8014eb9fbe7bdf1765185a6c508b5ca82eeaed7f899" +checksum = "45aec2360b3933207e27908049d8e4df4e476b58180afb1e56b2a4fb72efe4ba" dependencies = [ "liblzma-sys", ] [[package]] name = "liblzma-sys" -version = "0.4.6" +version = "0.4.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a60851d15cd8c5346eca4ab8babff585be2ae4bc8097c067291d3ffe2add3b6" +checksum = "a046c7f353ba30f810545151e04f63545833803f5b86ee3ddf1517247fe560a5" dependencies = [ "cc", "libc", @@ -5519,7 +5889,7 @@ version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f50e8f47623268b5407192d26876c4d7f89d686ca130fdc53bced4814cd29f8" dependencies = [ - "bitflags", + "bitflags 2.13.0", ] [[package]] @@ -5560,9 +5930,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.32" +version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "953f07c43838f8e6f9758cab68bf5bed85465e7587ebe0b823f1bcd81978ad3a" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" [[package]] name = "loom" @@ -5577,6 +5947,15 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "lru" +version = "0.16.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f66e8d5d03f609abc3a39e6f08e4164ebf1447a732906d39eb9b99b7919ef39" +dependencies = [ + "hashbrown 0.16.1", +] + [[package]] name = "lru" version = "0.18.0" @@ -5611,6 +5990,15 @@ dependencies = [ "libc", ] +[[package]] +name = "lz4_flex" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "373f5eceeeab7925e0c1098212f2fbc4d416adec9d35051a6ab251e824c1854a" +dependencies = [ + "twox-hash", +] + [[package]] name = "lz4_flex" version = "0.13.1" @@ -5689,15 +6077,15 @@ dependencies = [ [[package]] name = "memchr" -version = "2.8.1" +version = "2.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b947ae49db0d222b1dbc6b113ce7248a3fc3a6ca21b696717bfc000ba4484d8" +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" [[package]] name = "memmap2" -version = "0.9.10" +version = "0.9.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "714098028fe011992e1c3962653c96b2d578c4b4bce9036e15ff220319b1e0e3" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" dependencies = [ "libc", ] @@ -5743,6 +6131,15 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "miniz_oxide" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" +dependencies = [ + "adler2", +] + [[package]] name = "mio" version = "1.2.1" @@ -5755,12 +6152,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "mock_instant" -version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dce6dd36094cac388f119d2e9dc82dc730ef91c32a6222170d630e5414b956e6" - [[package]] name = "moka" version = "0.12.15" @@ -5814,7 +6205,7 @@ version = "0.31.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf20d2fde8ff38632c426f1165ed7436270b44f199fc55284c38276f9db47c3d" dependencies = [ - "bitflags", + "bitflags 2.13.0", "cfg-if", "cfg_aliases", "libc", @@ -5950,6 +6341,20 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" +dependencies = [ + "num-bigint", + "num-complex", + "num-integer", + "num-iter", + "num-rational", + "num-traits", +] + [[package]] name = "num-bigint" version = "0.4.6" @@ -5984,6 +6389,28 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-iter" +version = "0.1.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1429034a0490724d0075ebb2bc9e875d6503c3cf69e235a8941aa757d83ef5bf" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -6023,7 +6450,7 @@ dependencies = [ "proc-macro-crate", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -6041,7 +6468,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" dependencies = [ - "bitflags", + "bitflags 2.13.0", ] [[package]] @@ -6063,30 +6490,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "object_store" -version = "0.12.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbfbfff40aeccab00ec8a910b57ca8ecf4319b335c542f2edcd19dd25a1e2a00" -dependencies = [ - "async-trait", - "bytes", - "chrono", - "futures", - "http", - "humantime", - "itertools 0.14.0", - "parking_lot", - "percent-encoding", - "thiserror 2.0.18", - "tokio", - "tracing", - "url", - "walkdir", - "wasm-bindgen-futures", - "web-time", -] - [[package]] name = "object_store" version = "0.13.2" @@ -6291,7 +6694,7 @@ dependencies = [ "by_address", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -6330,7 +6733,7 @@ dependencies = [ "proc-macro-crate", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -6362,6 +6765,41 @@ dependencies = [ "windows-link", ] +[[package]] +name = "parquet" +version = "56.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3abbfef8a25900f4925c86e4cb881ea24672ca3c31ee4fb50a8083c4c56d313" +dependencies = [ + "ahash 0.8.12", + "arrow-array 56.2.1", + "arrow-buffer 56.2.1", + "arrow-cast 56.2.1", + "arrow-data 56.2.1", + "arrow-ipc 56.2.1", + "arrow-schema 56.2.1", + "arrow-select 56.2.1", + "base64", + "brotli", + "bytes", + "chrono", + "flate2", + "futures", + "half", + "hashbrown 0.16.1", + "lz4_flex 0.11.6", + "num", + "num-bigint", + "paste", + "seq-macro", + "simdutf8", + "snap", + "thrift", + "tokio", + "twox-hash", + "zstd", +] + [[package]] name = "parquet" version = "58.3.0" @@ -6369,12 +6807,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5dafa7d01085b62a47dd0c1829550a0a36710ea9c4fe358a05a85477cec8a908" dependencies = [ "ahash 0.8.12", - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-ipc", - "arrow-schema", - "arrow-select", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-ipc 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "base64", "brotli", "bytes", @@ -6383,11 +6821,11 @@ dependencies = [ "futures", "half", "hashbrown 0.17.1", - "lz4_flex", + "lz4_flex 0.13.1", "num-bigint", "num-integer", "num-traits", - "object_store 0.13.2", + "object_store", "paste", "seq-macro", "simdutf8", @@ -6404,8 +6842,8 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "74c8db065291f088a2aad8ab831853eae1871c0d311c8d0b83bbc3b7e735d0fc" dependencies = [ - "arrow", - "arrow-schema", + "arrow 58.3.0", + "arrow-schema 58.3.0", "chrono", "half", "indexmap 2.14.0", @@ -6420,8 +6858,8 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a530e8d5b5e14efcb39c9a6ec55432ad11f6afb7dc4455a79be0dc615fe3cc31" dependencies = [ - "arrow", - "arrow-schema", + "arrow 58.3.0", + "arrow-schema 58.3.0", "chrono", "half", "indexmap 2.14.0", @@ -6437,7 +6875,7 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "00ed89908289f67caa2ca078f9ff9aacd6229a313ec92b12bf4f48f613dc2b97" dependencies = [ - "arrow-schema", + "arrow-schema 58.3.0", "base64", "chrono", "parquet-variant", @@ -6544,7 +6982,7 @@ checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -6681,7 +7119,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -6723,7 +7161,7 @@ dependencies = [ "proc-macro-error-attr2", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -6770,7 +7208,7 @@ dependencies = [ "prost 0.14.4", "prost-types", "regex", - "syn 2.0.117", + "syn 2.0.118", "tempfile", ] @@ -6784,7 +7222,7 @@ dependencies = [ "itertools 0.12.1", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -6797,7 +7235,7 @@ dependencies = [ "itertools 0.14.0", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -6841,9 +7279,9 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.28.3" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" dependencies = [ "chrono", "indexmap 2.14.0", @@ -6857,9 +7295,9 @@ dependencies = [ [[package]] name = "pyo3-async-runtimes" -version = "0.28.0" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e7364a95bf00e8377bbf9b0f09d7ff9715a29d8fcf93b47d1a967363b973178" +checksum = "b3ef68daa7316a3fac65e5e18b2203f010346de1c1c53456811a2624673ab046" dependencies = [ "futures-channel", "futures-util", @@ -6871,18 +7309,18 @@ dependencies = [ [[package]] name = "pyo3-build-config" -version = "0.28.3" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" dependencies = [ "target-lexicon", ] [[package]] name = "pyo3-bytes" -version = "0.6.0" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d98190ac733bf1c9a3bdbd2856d85abee243f16b7ef1155feefcd613c8704a1" +checksum = "a3fbca288ba42ffd1b9199fa39af7423bfd9e4418e8ac9ce23ad07fe8956504f" dependencies = [ "bytes", "pyo3", @@ -6890,9 +7328,9 @@ dependencies = [ [[package]] name = "pyo3-ffi" -version = "0.28.3" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" dependencies = [ "libc", "pyo3-build-config", @@ -6900,9 +7338,9 @@ dependencies = [ [[package]] name = "pyo3-log" -version = "0.13.3" +version = "0.13.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26c2ec80932c5c3b2d4fbc578c9b56b2d4502098587edb8bef5b6bfcad43682e" +checksum = "f64083bd3a16a353d9d62335808e8e13d0552d2a2b83fdb084496192dcfa9fcd" dependencies = [ "arc-swap", "log", @@ -6911,34 +7349,33 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.28.3" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] name = "pyo3-macros-backend" -version = "0.28.3" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" dependencies = [ "heck", "proc-macro2", - "pyo3-build-config", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] name = "pyo3-object_store" -version = "0.9.0" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3056dc8a77db5d44d32e7f740db0b01e5b65a43d181487a6917576959c6eb9a" +checksum = "13aaa9b876f9b02773cc18fd03dd359bddf3ca21978068024511f166c9c0c6a5" dependencies = [ "async-trait", "bytes", @@ -6947,7 +7384,7 @@ dependencies = [ "http", "humantime", "itertools 0.14.0", - "object_store 0.13.2", + "object_store", "percent-encoding", "pyo3", "pyo3-async-runtimes", @@ -6969,9 +7406,9 @@ dependencies = [ [[package]] name = "quinn" -version = "0.11.9" +version = "0.11.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" dependencies = [ "bytes", "cfg_aliases", @@ -6989,9 +7426,9 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.14" +version = "0.11.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" +checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" dependencies = [ "aws-lc-rs", "bytes", @@ -7025,9 +7462,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.45" +version = "1.0.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" dependencies = [ "proc-macro2", ] @@ -7078,7 +7515,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d2e8e8bcc7961af1fdac401278c6a831614941f6164ee3bf4ce61b7edb162207" dependencies = [ "chacha20", - "getrandom 0.4.2", + "getrandom 0.4.3", "rand_core 0.10.1", ] @@ -7200,9 +7637,9 @@ checksum = "973443cf09a9c8656b574a866ab68dfa19f0867d0340648c7d2f6a71b8a8ea68" [[package]] name = "ratatui" -version = "0.30.1" +version = "0.30.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1695748e3a735b34968c887ceea5a380b43545903868ae8f5b666593100f6b68" +checksum = "3274ba0a2c5e1bcad2a2005d20f4dc59dad26b2eb0940fb094500dba4099d57d" dependencies = [ "instability", "ratatui-core", @@ -7213,18 +7650,17 @@ dependencies = [ [[package]] name = "ratatui-core" -version = "0.1.1" +version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42d3603f354bba8c595fa47860e60142d7372b7210c27044c6a7d0e1a4336b44" +checksum = "cbb175c433c8e28a809d1f5773a2ae96e68c0ce40db865cbab1020bf33ae479c" dependencies = [ - "bitflags", + "bitflags 2.13.0", "compact_str", "critical-section", "hashbrown 0.17.1", - "indoc", "itertools 0.14.0", "kasuari", - "lru", + "lru 0.18.0", "palette", "serde", "strum 0.28.0", @@ -7236,9 +7672,9 @@ dependencies = [ [[package]] name = "ratatui-crossterm" -version = "0.1.1" +version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b2867bedcbd6a690ca4f8672a687b730ec07660c79844517b084311b529980c" +checksum = "567584a3b0e6a8203c23de40b4861497266725eb5363dbfd18a1edd603cca9f0" dependencies = [ "cfg-if", "crossterm", @@ -7248,11 +7684,11 @@ dependencies = [ [[package]] name = "ratatui-widgets" -version = "0.3.1" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ef4f17dd7ac3abf5adc2b920a03c61eee4bfe6a88fa5191936895525371d79c" +checksum = "66e3d19bcc9130ca376277d93b60767ff121ace3be06f5f95f81dd68956407d1" dependencies = [ - "bitflags", + "bitflags 2.13.0", "hashbrown 0.17.1", "indoc", "instability", @@ -7268,9 +7704,9 @@ dependencies = [ [[package]] name = "ratzilla" -version = "0.3.0" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfc729362a3a8714b512f330c8dfaa03979813e9491e21b8e49b81e1e5d78014" +checksum = "d53309a50b2cfb9828d1ea653b77fc6d2fe08137b05e65aa070ab87630e35b49" dependencies = [ "beamterm-renderer", "bitvec", @@ -7325,7 +7761,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76009fbe0614077fc1a2ce255e3a1881a2e3a3527097d5dc6d8212c585e7e38b" dependencies = [ "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -7334,7 +7770,7 @@ version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags", + "bitflags 2.13.0", ] [[package]] @@ -7365,14 +7801,14 @@ checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] name = "regex" -version = "1.12.3" +version = "1.12.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +checksum = "f1292b7759ae1cb9ec195452d1390a074f0cd8541ab7a5a8c31cd6db45d4a6ba" dependencies = [ "aho-corasick", "memchr", @@ -7399,9 +7835,9 @@ checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" [[package]] name = "regex-syntax" -version = "0.8.10" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "relative-path" @@ -7558,6 +7994,23 @@ dependencies = [ "byteorder", ] +[[package]] +name = "robust" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e27ee8bb91ca0adcf0ecb116293afa12d393f9c2b9b9cd54d33e8078fe19839" + +[[package]] +name = "rstar" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "421400d13ccfd26dfa5858199c30a5d76f9c54e0dba7575273025b43c5175dbb" +dependencies = [ + "heapless", + "num-traits", + "smallvec", +] + [[package]] name = "rstest" version = "0.26.1" @@ -7583,7 +8036,7 @@ dependencies = [ "regex", "relative-path", "rustc_version", - "syn 2.0.117", + "syn 2.0.118", "unicode-ident", ] @@ -7595,7 +8048,7 @@ checksum = "b3a8fb4672e840a587a66fc577a5491375df51ddb88f2a2c2a792598c326fe14" dependencies = [ "quote", "rand 0.8.6", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -7610,9 +8063,9 @@ dependencies = [ [[package]] name = "rust_decimal" -version = "1.42.0" +version = "1.42.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c5108e3d4d903e21aac27f12ba5377b6b34f9f44b325e4894c7924169d06995" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" dependencies = [ "arrayvec", "borsh", @@ -7637,6 +8090,12 @@ version = "2.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3e75f6a532d0fd9f7f13144f392b6ad56a32696bfcd9c78f797f16bbb6f072d6" +[[package]] +name = "rustc-std-workspace-core" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa9c45b374136f52f2d6311062c7146bff20fec063c3f5d46a410bd937746955" + [[package]] name = "rustc_version" version = "0.4.1" @@ -7652,7 +8111,7 @@ version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ - "bitflags", + "bitflags 2.13.0", "errno", "libc", "linux-raw-sys", @@ -7661,9 +8120,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.40" +version = "0.23.41" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" +checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" dependencies = [ "aws-lc-rs", "once_cell", @@ -7807,6 +8266,12 @@ version = "1.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d68f2ec51b097e4c1a75b681a8bec621909b5e91f15bb7b840c4f2f7b01148b2" +[[package]] +name = "scroll" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04c565b551bafbef4157586fa379538366e4385d42082f255bfd96e4fe8519da" + [[package]] name = "seahash" version = "4.1.0" @@ -7819,7 +8284,7 @@ version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ - "bitflags", + "bitflags 2.13.0", "core-foundation 0.10.1", "core-foundation-sys", "libc", @@ -7879,7 +8344,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -7903,7 +8368,7 @@ checksum = "175ee3e80ae9982737ca543e96133087cbd9a485eecc3bc4de9c1a37b47ea59c" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -7965,7 +8430,7 @@ dependencies = [ "darling", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -7990,7 +8455,7 @@ checksum = "94e153fc76e1c6a068703d6d29c508a0b15c061c4b7e43da59cc097bc342673c" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -8158,9 +8623,9 @@ dependencies = [ [[package]] name = "smallvec" -version = "1.15.1" +version = "1.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" [[package]] name = "smol" @@ -8197,7 +8662,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -8216,6 +8681,42 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "spade" +version = "2.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9699399fd9349b00b184f5635b074f9ec93afffef30c853f8c875b32c0f8c7fa" +dependencies = [ + "hashbrown 0.16.1", + "num-traits", + "robust", + "smallvec", +] + +[[package]] +name = "spatialbench" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07f3f4b67ccf571f183d3695aa6b9d6f996864c31782a480e97a23ed0f2f6f18" +dependencies = [ + "geo", + "once_cell", + "rand 0.8.6", + "serde", +] + +[[package]] +name = "spatialbench-arrow" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad89c32ed9e258bcc89713c296c7437963ce31f511eb8a408d2046e853294206" +dependencies = [ + "arrow 56.2.1", + "geo", + "geozero", + "spatialbench", +] + [[package]] name = "sqllogictest" version = "0.29.1" @@ -8270,7 +8771,7 @@ checksum = "a6dd45d8fc1c79299bfbb7190e42ccbbdf6a5f52e4a6ad98d92357ea965bd289" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -8354,7 +8855,7 @@ dependencies = [ "proc-macro2", "quote", "rustversion", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -8366,7 +8867,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -8398,9 +8899,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.117" +version = "2.0.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" dependencies = [ "proc-macro2", "quote", @@ -8424,7 +8925,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -8447,7 +8948,7 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b" dependencies = [ - "bitflags", + "bitflags 2.13.0", "core-foundation 0.9.4", "system-configuration-sys", ] @@ -8535,7 +9036,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.2", + "getrandom 0.4.3", "once_cell", "rustix", "windows-sys 0.61.2", @@ -8568,9 +9069,9 @@ checksum = "d4d1330fe7f7f872cd05165130b10602d667b205fd85be09be2814b115d4ced9" [[package]] name = "test-with" -version = "0.16.1" +version = "0.16.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c292571fb159e65c78c7c88a96e5ddfc326d6819e49b3a040d7b0aa3497925b" +checksum = "964c2c219d0e7658fc3cd7354826e9e8d997c1e7d592cc32f349fc1125b64a84" dependencies = [ "byte-unit", "chrono", @@ -8585,20 +9086,19 @@ dependencies = [ [[package]] name = "test-with-derive" -version = "0.16.1" +version = "0.16.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "01b95b557b54c0d50b04688ab86aaeaf16f201a0341936f3bfa19497dee081bd" +checksum = "52884ff2163f79cd0a7af841fd2e9aca103f2ff1207092d317bf00565abdc59c" dependencies = [ "byte-unit", "chrono", "num_cpus", "ping", - "proc-macro-error2", "proc-macro2", "quote", "regex", "reqwest 0.13.4", - "syn 2.0.117", + "syn 2.0.118", "sysinfo", "uzers", "which", @@ -8639,7 +9139,7 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -8650,7 +9150,7 @@ checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -8694,12 +9194,11 @@ dependencies = [ [[package]] name = "time" -version = "0.3.47" +version = "0.3.51" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "743bd48c283afc0388f9b8827b976905fb217ad9e647fae3a379a9283c4def2c" +checksum = "85c17d80feb7334b40c484e45ed1a5273dfd8bfda537c3be2e74a06a6686f327" dependencies = [ "deranged", - "itoa", "js-sys", "libc", "num-conv", @@ -8712,15 +9211,15 @@ dependencies = [ [[package]] name = "time-core" -version = "0.1.8" +version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7694e1cfe791f8d31026952abf09c69ca6f6fa4e1a1229e18988f06a04a12dca" +checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" [[package]] name = "time-macros" -version = "0.2.27" +version = "0.2.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e70e4c5a0e0a8a4823ad65dfe1a6930e4f4d756dcd9dd7939022b5e8c501215" +checksum = "dcef1a61bdb119096e153208ec5cbec23944ce8bca13be5c7f60c634f7403935" dependencies = [ "num-conv", "time-core", @@ -8795,7 +9294,7 @@ checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -8915,7 +9414,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ "async-compression", - "bitflags", + "bitflags 2.13.0", "bytes", "futures-core", "futures-util", @@ -8953,7 +9452,7 @@ name = "tpchgen-arrow" version = "2.0.2" source = "git+https://github.com/clflushopt/tpchgen-rs.git?rev=438e9c2dbc25b2fff82c0efc08b3f13b5707874f#438e9c2dbc25b2fff82c0efc08b3f13b5707874f" dependencies = [ - "arrow", + "arrow 58.3.0", "tpchgen", ] @@ -8977,7 +9476,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -9137,6 +9636,9 @@ name = "unicode-width" version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" +dependencies = [ + "rustc-std-workspace-core", +] [[package]] name = "unicode-xid" @@ -9194,11 +9696,11 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.23.2" +version = "1.23.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d258b83ceec21034727ecee8c382cfa6c3e133699b0742c64571814fb420c9f7" +checksum = "bf80a72845275afea99e7f2b434723d3bc7e38470fcd1c7ed39a599c73319a53" dependencies = [ - "getrandom 0.4.2", + "getrandom 0.4.3", "js-sys", "serde_core", "wasm-bindgen", @@ -9206,9 +9708,9 @@ dependencies = [ [[package]] name = "uzers" -version = "0.11.3" +version = "0.12.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76d283dc7e8c901e79e32d077866eaf599156cbf427fffa8289aecc52c5c3f63" +checksum = "0b8275fb1afee25b4111d2dc8b5c505dbbc4afd0b990cb96deb2d88bff8be18d" dependencies = [ "libc", "log", @@ -9231,12 +9733,12 @@ name = "vortex" version = "0.1.0" dependencies = [ "anyhow", - "arrow-array", + "arrow-array 58.3.0", "codspeed-divan-compat", "fastlanes", "futures", "mimalloc", - "parquet", + "parquet 58.3.0", "paste", "rand 0.10.1", "rand_distr 0.6.0", @@ -9302,15 +9804,15 @@ dependencies = [ "arbitrary", "arc-swap", "arcref", - "arrow-arith", - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-data", - "arrow-ord", - "arrow-schema", - "arrow-select", - "arrow-string", + "arrow-arith 58.3.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-cast 58.3.0", + "arrow-data 58.3.0", + "arrow-ord 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", + "arrow-string 58.3.0", "async-lock", "bytes", "cfg-if", @@ -9352,6 +9854,7 @@ dependencies = [ "vortex-array", "vortex-array-macros", "vortex-buffer", + "vortex-compute", "vortex-error", "vortex-flatbuffers", "vortex-mask", @@ -9366,7 +9869,7 @@ version = "0.1.0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -9374,9 +9877,9 @@ name = "vortex-bench" version = "0.1.0" dependencies = [ "anyhow", - "arrow-array", - "arrow-schema", - "arrow-select", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "async-trait", "bzip2", "clap", @@ -9391,12 +9894,15 @@ dependencies = [ "noodles-bgzf", "noodles-vcf", "parking_lot", - "parquet", + "parquet 56.2.1", + "parquet 58.3.0", "rand 0.10.1", "regex", "reqwest 0.13.4", "serde", "serde_json", + "spatialbench", + "spatialbench-arrow", "sysinfo", "tabled", "target-lexicon", @@ -9411,6 +9917,7 @@ dependencies = [ "url", "uuid", "vortex", + "vortex-geo", "vortex-tensor", ] @@ -9449,13 +9956,14 @@ dependencies = [ name = "vortex-buffer" version = "0.1.0" dependencies = [ - "arrow-buffer", + "arrow-buffer 58.3.0", "bitvec", "bytes", "codspeed-divan-compat", "itertools 0.14.0", "memmap2", "num-traits", + "rand 0.10.1", "rstest", "serde", "simdutf8", @@ -9480,13 +9988,13 @@ dependencies = [ name = "vortex-compat" version = "0.1.0" dependencies = [ - "arrow-array", - "arrow-select", + "arrow-array 58.3.0", + "arrow-select 58.3.0", "base16ct", "bytes", "clap", "futures", - "parquet", + "parquet 58.3.0", "reqwest 0.13.4", "serde", "serde_json", @@ -9499,6 +10007,7 @@ dependencies = [ "vortex-array", "vortex-buffer", "vortex-error", + "vortex-fsst", "vortex-session", ] @@ -9523,6 +10032,21 @@ dependencies = [ "vortex-utils", ] +[[package]] +name = "vortex-compute" +version = "0.1.0" +dependencies = [ + "arrow-arith 58.3.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-cast 58.3.0", + "arrow-schema 58.3.0", + "codspeed-divan-compat", + "num-traits", + "rand 0.10.1", + "vortex-buffer", +] + [[package]] name = "vortex-cub" version = "0.1.0" @@ -9539,7 +10063,7 @@ name = "vortex-cuda" version = "0.1.0" dependencies = [ "arc-swap", - "arrow-schema", + "arrow-schema 58.3.0", "async-trait", "bindgen", "bytes", @@ -9550,7 +10074,7 @@ dependencies = [ "itertools 0.14.0", "kanal", "num-traits", - "object_store 0.13.2", + "object_store", "parking_lot", "prost 0.14.4", "rstest", @@ -9570,7 +10094,7 @@ dependencies = [ name = "vortex-cuda-ffi" version = "0.1.0" dependencies = [ - "arrow-schema", + "arrow-schema 58.3.0", "futures", "vortex", "vortex-array", @@ -9584,7 +10108,7 @@ name = "vortex-cuda-macros" version = "0.1.0" dependencies = [ "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -9592,8 +10116,8 @@ name = "vortex-cxx" version = "0.1.0" dependencies = [ "anyhow", - "arrow-array", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", "async-fs", "cxx", "futures", @@ -9607,7 +10131,8 @@ name = "vortex-datafusion" version = "0.1.0" dependencies = [ "anyhow", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", "async-trait", "datafusion 54.0.0", "datafusion-catalog 54.0.0", @@ -9617,6 +10142,7 @@ dependencies = [ "datafusion-execution 54.0.0", "datafusion-expr 54.0.0", "datafusion-functions 54.0.0", + "datafusion-functions-nested 54.0.0", "datafusion-physical-expr 54.0.0", "datafusion-physical-expr-adapter 54.0.0", "datafusion-physical-expr-common 54.0.0", @@ -9625,7 +10151,7 @@ dependencies = [ "futures", "insta", "itertools 0.14.0", - "object_store 0.13.2", + "object_store", "rstest", "tempfile", "tokio", @@ -9683,7 +10209,7 @@ dependencies = [ "jiff", "kanal", "num-traits", - "object_store 0.13.2", + "object_store", "parking_lot", "paste", "reqwest 0.13.4", @@ -9707,10 +10233,10 @@ dependencies = [ name = "vortex-error" version = "0.1.0" dependencies = [ - "arrow-schema", + "arrow-schema 58.3.0", "flatbuffers", "jiff", - "object_store 0.13.2", + "object_store", "prost 0.14.4", "serial_test", "temp-env", @@ -9742,14 +10268,15 @@ dependencies = [ name = "vortex-ffi" version = "0.1.0" dependencies = [ - "arrow-array", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", "async-fs", + "bytes", "cbindgen", "futures", "itertools 0.14.0", "mimalloc", - "object_store 0.13.2", + "object_store", "paste", "prost 0.14.4", "rand 0.10.1", @@ -9772,7 +10299,7 @@ dependencies = [ "itertools 0.14.0", "kanal", "moka", - "object_store 0.13.2", + "object_store", "oneshot", "parking_lot", "pin-project-lite", @@ -9822,6 +10349,7 @@ version = "0.1.0" dependencies = [ "codspeed-divan-compat", "fsst-rs", + "num-traits", "prost 0.14.4", "rand 0.10.1", "rstest", @@ -9863,12 +10391,14 @@ dependencies = [ name = "vortex-geo" version = "0.1.0" dependencies = [ - "arrow-array", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", + "geo", "geo-traits", "geo-types", "geoarrow", "prost 0.14.4", + "rstest", "vortex-array", "vortex-error", "vortex-session", @@ -9890,7 +10420,7 @@ dependencies = [ "itertools 0.14.0", "kanal", "log", - "object_store 0.13.2", + "object_store", "oneshot", "parking_lot", "pin-project-lite", @@ -9929,12 +10459,12 @@ dependencies = [ name = "vortex-jni" version = "0.1.0" dependencies = [ - "arrow-array", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", "async-fs", "futures", "jni", - "object_store 0.13.2", + "object_store", "parking_lot", "roaring", "thiserror 2.0.18", @@ -9949,10 +10479,12 @@ dependencies = [ name = "vortex-json" version = "0.1.0" dependencies = [ - "arrow-array", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", + "prost 0.14.4", "vortex-array", "vortex-error", + "vortex-proto", "vortex-session", ] @@ -9961,8 +10493,8 @@ name = "vortex-layout" version = "0.1.0" dependencies = [ "arcref", - "arrow-array", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", "async-stream", "async-trait", "bit-vec", @@ -10052,9 +10584,9 @@ dependencies = [ name = "vortex-parquet-variant" version = "0.1.0" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-schema 58.3.0", "chrono", "parquet-variant", "parquet-variant-compute", @@ -10066,6 +10598,7 @@ dependencies = [ "vortex-error", "vortex-file", "vortex-io", + "vortex-json", "vortex-layout", "vortex-mask", "vortex-proto", @@ -10098,14 +10631,14 @@ dependencies = [ name = "vortex-python" version = "0.1.0" dependencies = [ - "arrow-array", - "arrow-data", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", "async-fs", "bytes", "itertools 0.14.0", "log", - "object_store 0.13.2", + "object_store", "parking_lot", "pyo3", "pyo3-bytes", @@ -10115,16 +10648,33 @@ dependencies = [ "url", "vortex", "vortex-array", + "vortex-python-abi", "vortex-tui", ] +[[package]] +name = "vortex-python-abi" +version = "0.1.0" + +[[package]] +name = "vortex-python-cuda" +version = "0.1.0" +dependencies = [ + "arrow-schema 58.3.0", + "bytes", + "pyo3", + "vortex", + "vortex-cuda", + "vortex-python-abi", +] + [[package]] name = "vortex-row" version = "0.1.0" dependencies = [ - "arrow-array", - "arrow-row", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-row 58.3.0", + "arrow-schema 58.3.0", "bytes", "codspeed-divan-compat", "mimalloc", @@ -10143,8 +10693,8 @@ name = "vortex-runend" version = "0.1.0" dependencies = [ "arbitrary", - "arrow-array", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", "codspeed-divan-compat", "itertools 0.14.0", "num-traits", @@ -10194,7 +10744,7 @@ dependencies = [ name = "vortex-session" version = "0.1.0" dependencies = [ - "dashmap", + "arc-swap", "lasso", "parking_lot", "vortex-error", @@ -10224,11 +10774,11 @@ dependencies = [ "anyhow", "async-trait", "bigdecimal", - "clap", "datafusion 54.0.0", + "datafusion-functions-nested 54.0.0", "datafusion-sqllogictest", - "futures", "indicatif", + "regex", "rstest", "sqllogictest", "thiserror 2.0.18", @@ -10242,8 +10792,8 @@ dependencies = [ name = "vortex-tensor" version = "0.1.0" dependencies = [ - "arrow-array", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", "codspeed-divan-compat", "half", "itertools 0.14.0", @@ -10266,8 +10816,8 @@ dependencies = [ name = "vortex-test-e2e-cuda" version = "0.1.0" dependencies = [ - "arrow-array", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", "futures", "vortex", "vortex-cuda", @@ -10278,8 +10828,8 @@ name = "vortex-tui" version = "0.1.0" dependencies = [ "anyhow", - "arrow-array", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-schema 58.3.0", "clap", "console_error_panic_hook", "crossterm", @@ -10292,7 +10842,7 @@ dependencies = [ "indicatif", "itertools 0.14.0", "js-sys", - "parquet", + "parquet 58.3.0", "ratatui", "ratzilla", "serde", @@ -10319,9 +10869,9 @@ dependencies = [ name = "vortex-web-wasm" version = "0.1.0" dependencies = [ - "arrow-array", - "arrow-ipc", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-ipc 58.3.0", + "arrow-schema 58.3.0", "console_error_panic_hook", "futures", "js-sys", @@ -10389,27 +10939,18 @@ checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" [[package]] name = "wasip2" -version = "1.0.3+wasi-0.2.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" -dependencies = [ - "wit-bindgen 0.57.1", -] - -[[package]] -name = "wasip3" -version = "0.4.0+wasi-0.3.0-rc-2026-01-06" +version = "1.0.4+wasi-0.2.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" dependencies = [ - "wit-bindgen 0.51.0", + "wit-bindgen", ] [[package]] name = "wasm-bindgen" -version = "0.2.122" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ed04576f974d2b2fba0f38c51dbc5518011e38c36bf1143164be765528fd409" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" dependencies = [ "cfg-if", "once_cell", @@ -10420,9 +10961,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.72" +version = "0.4.76" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9473dbd2991ae90b6291c3c32c30c6187ac49aa32f9905d1cce280ec1e110b0f" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" dependencies = [ "js-sys", "wasm-bindgen", @@ -10430,9 +10971,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.122" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "916151b09da36bd82f6615cbf3a419e2f0ba23a03c6160e8e92eb6bd4aa1dec6" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -10440,48 +10981,26 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.122" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "299047362ccbfce148b67ab7e73349f77748e00c8296f9542adfad2ad82c5c5e" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" -version = "0.2.122" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a929b2c61f11ba3e9bc35b50c1f25cb38e0e892c0c231ae2b8cf78d5dad4437" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" dependencies = [ "unicode-ident", ] -[[package]] -name = "wasm-encoder" -version = "0.244.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" -dependencies = [ - "leb128fmt", - "wasmparser", -] - -[[package]] -name = "wasm-metadata" -version = "0.244.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" -dependencies = [ - "anyhow", - "indexmap 2.14.0", - "wasm-encoder", - "wasmparser", -] - [[package]] name = "wasm-streams" version = "0.4.2" @@ -10508,23 +11027,11 @@ dependencies = [ "web-sys", ] -[[package]] -name = "wasmparser" -version = "0.244.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" -dependencies = [ - "bitflags", - "hashbrown 0.15.5", - "indexmap 2.14.0", - "semver", -] - [[package]] name = "web-sys" -version = "0.3.99" +version = "0.3.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6d621441cfc37b84979402712047321980c178f299193a3589d05b99e8763436" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" dependencies = [ "js-sys", "wasm-bindgen", @@ -10542,18 +11049,18 @@ dependencies = [ [[package]] name = "webpki-root-certs" -version = "1.0.7" +version = "1.0.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f31141ce3fc3e300ae89b78c0dd67f9708061d1d2eda54b8209346fd6be9a92c" +checksum = "0d46a5a140e6f7afeccd8eae97eff335163939eac8b929834875168b29b3d267" dependencies = [ "rustls-pki-types", ] [[package]] name = "which" -version = "8.0.3" +version = "8.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c789537cf2f7f55be8e6192f92e464174ee55f91af622777f7f1ceb0dbccd03e" +checksum = "48d7cd18d4acb58fb3cdfe9ea54e6cd96a4e7d4cc45c56338b236e82dad47248" dependencies = [ "libc", ] @@ -10642,7 +11149,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -10653,7 +11160,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -10890,100 +11397,12 @@ dependencies = [ "memchr", ] -[[package]] -name = "wit-bindgen" -version = "0.51.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" -dependencies = [ - "wit-bindgen-rust-macro", -] - [[package]] name = "wit-bindgen" version = "0.57.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" -[[package]] -name = "wit-bindgen-core" -version = "0.51.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" -dependencies = [ - "anyhow", - "heck", - "wit-parser", -] - -[[package]] -name = "wit-bindgen-rust" -version = "0.51.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" -dependencies = [ - "anyhow", - "heck", - "indexmap 2.14.0", - "prettyplease", - "syn 2.0.117", - "wasm-metadata", - "wit-bindgen-core", - "wit-component", -] - -[[package]] -name = "wit-bindgen-rust-macro" -version = "0.51.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" -dependencies = [ - "anyhow", - "prettyplease", - "proc-macro2", - "quote", - "syn 2.0.117", - "wit-bindgen-core", - "wit-bindgen-rust", -] - -[[package]] -name = "wit-component" -version = "0.244.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" -dependencies = [ - "anyhow", - "bitflags", - "indexmap 2.14.0", - "log", - "serde", - "serde_derive", - "serde_json", - "wasm-encoder", - "wasm-metadata", - "wasmparser", - "wit-parser", -] - -[[package]] -name = "wit-parser" -version = "0.244.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" -dependencies = [ - "anyhow", - "id-arena", - "indexmap 2.14.0", - "log", - "semver", - "serde", - "serde_derive", - "serde_json", - "unicode-xid", - "wasmparser", -] - [[package]] name = "wkb" version = "0.9.2" @@ -10996,6 +11415,18 @@ dependencies = [ "thiserror 1.0.69", ] +[[package]] +name = "wkt" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54f7f1ff4ea4c18936d6cd26a6fd24f0003af37e951a8e0e8b9e9a2d0bd0a46d" +dependencies = [ + "geo-types", + "log", + "num-traits", + "thiserror 1.0.69", +] + [[package]] name = "wkt" version = "0.14.0" @@ -11091,28 +11522,28 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", "synstructure", ] [[package]] name = "zerocopy" -version = "0.8.50" +version = "0.8.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b065d4f0e55f82fae73202e189638116a87c55ab6b8e6c2721e13dd9d854ad1" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.50" +version = "0.8.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b631b19d36a892ab55420c92dbc83ccd79274f25be714855d3074aa71cab639" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -11132,15 +11563,15 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", "synstructure", ] [[package]] name = "zeroize" -version = "1.8.2" +version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" [[package]] name = "zerotrie" @@ -11172,7 +11603,7 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.118", ] [[package]] @@ -11196,7 +11627,7 @@ dependencies = [ "crc32fast", "deflate64", "flate2", - "getrandom 0.4.2", + "getrandom 0.4.3", "hmac", "indexmap 2.14.0", "lzma-rust2", @@ -11213,9 +11644,9 @@ dependencies = [ [[package]] name = "zlib-rs" -version = "0.6.3" +version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3be3d40e40a133f9c916ee3f9f4fa2d9d63435b5fbe1bfc6d9dae0aa0ada1513" +checksum = "977347db8caa080403f6b6b7c1cda9479a8e869316f7e13a59b19076a40f94e3" [[package]] name = "zmij" diff --git a/Cargo.toml b/Cargo.toml index a587ef1db16..bf5057432a6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ members = [ "vortex-array-macros", "vortex-error", "vortex-buffer", + "vortex-compute", "vortex-mask", "vortex-utils", "vortex-session", @@ -35,7 +36,9 @@ members = [ "vortex-ffi", "fuzz", "vortex-jni", + "vortex-python-abi", "vortex-python", + "vortex-python-cuda", "vortex-tui", "vortex-web/crate", "vortex-sqllogictest", @@ -124,7 +127,7 @@ cc = "1.2" cfg-if = "1.0.1" chrono = "0.4.44" clap = "4.5" -criterion = "0.8" +criterion = { package = "codspeed-criterion-compat-walltime", version = "5.0.0" } crossterm = "0.29" cudarc = { version = "0.19.0", features = [ # The NSight Compute version available on lambda.ai hosts does not inject a symbol for @@ -142,6 +145,7 @@ datafusion-datasource = { version = "54", default-features = false } datafusion-execution = { version = "54" } datafusion-expr = { version = "54" } datafusion-functions = { version = "54" } +datafusion-functions-nested = { version = "54" } datafusion-physical-expr = { version = "54" } datafusion-physical-expr-adapter = { version = "54" } datafusion-physical-expr-common = { version = "54" } @@ -149,7 +153,7 @@ datafusion-physical-plan = { version = "54" } datafusion-pruning = { version = "54" } datafusion-sqllogictest = { version = "54" } dirs = "6.0.0" -divan = { package = "codspeed-divan-compat", version = "4.0.4" } +divan = { package = "codspeed-divan-compat", version = "5.0.0" } enum-iterator = "2.0.0" env_logger = "0.11" fastlanes = "0.5.1" @@ -157,6 +161,7 @@ flatbuffers = "25.2.10" fsst-rs = "0.5.11" futures = { version = "0.3.31", default-features = false } fuzzy-matcher = "0.3" +geo = "0.31.0" geo-traits = "0.3.0" geo-types = "0.7.19" geoarrow = "0.8.0" @@ -174,7 +179,7 @@ jetscii = "0.5.3" jiff = "0.2.0" jni = { version = "0.22.0" } kanal = "0.1.1" -lasso = { version = "0.7", features = ["multi-threaded"] } +lasso = { version = "0.7", features = ["multi-threaded", "inline-more"] } lending-iterator = "0.1.7" libfuzzer-sys = "0.4" libloading = "0.8" @@ -209,10 +214,10 @@ proc-macro2 = "1.0.95" prost = "0.14" prost-build = "0.14" prost-types = "0.14" -pyo3 = { version = "0.28.0" } -pyo3-bytes = "0.6" +pyo3 = { version = "0.29.0" } +pyo3-bytes = "0.7" pyo3-log = "0.13.0" -pyo3-object_store = "0.9.0" +pyo3-object_store = "0.11.0" quote = "1.0.44" rand = "0.10.1" rand_distr = "0.6" @@ -239,6 +244,14 @@ similar = "3.0.0" sketches-ddsketch = "0.4.0" smallvec = "1.15.1" smol = "2.0.2" +spatialbench = "0.2" +spatialbench-arrow = "0.2" +# spatialbench still pins arrow 56, two majors behind the workspace arrow. Until upstream +# catches up, write its generated batches with a matching parquet instead of converting +# arrow versions at the boundary. +spatialbench-parquet = { package = "parquet", version = "56", features = [ + "async", +] } static_assertions = "1.1" strum = "0.28" syn = { version = "2.0.117", features = ["full"] } @@ -282,6 +295,7 @@ vortex-btrblocks = { version = "0.1.0", path = "./vortex-btrblocks", default-fea vortex-buffer = { version = "0.1.0", path = "./vortex-buffer", default-features = false } vortex-bytebool = { version = "0.1.0", path = "./encodings/bytebool", default-features = false } vortex-compressor = { version = "0.1.0", path = "./vortex-compressor", default-features = false } +vortex-compute = { version = "0.1.0", path = "./vortex-compute", default-features = false } vortex-datafusion = { version = "0.1.0", path = "./vortex-datafusion", default-features = false } vortex-datetime-parts = { version = "0.1.0", path = "./encodings/datetime-parts", default-features = false } vortex-decimal-byte-parts = { version = "0.1.0", path = "encodings/decimal-byte-parts", default-features = false } diff --git a/bench-orchestrator/bench_orchestrator/config.py b/bench-orchestrator/bench_orchestrator/config.py index fb90ac8bd6a..c597e84c6be 100644 --- a/bench-orchestrator/bench_orchestrator/config.py +++ b/bench-orchestrator/bench_orchestrator/config.py @@ -46,6 +46,7 @@ class Benchmark(Enum): TPCH = "tpch" TPCDS = "tpcds" CLICKBENCH = "clickbench" + CLICKBENCH_SORTED = "clickbench-sorted" FINEWEB = "fineweb" GHARCHIVE = "gh-archive" POLARSIGNALS = "polarsignals" diff --git a/benchmarks-website/package-lock.json b/benchmarks-website/package-lock.json index 5e1682ef02a..85c5abe0a1f 100644 --- a/benchmarks-website/package-lock.json +++ b/benchmarks-website/package-lock.json @@ -29,21 +29,21 @@ } }, "node_modules/@emnapi/core": { - "version": "1.10.0", - "resolved": "https://registry.npmjs.org/@emnapi/core/-/core-1.10.0.tgz", - "integrity": "sha512-yq6OkJ4p82CAfPl0u9mQebQHKPJkY7WrIuk205cTYnYe+k2Z8YBh11FrbRG/H6ihirqcacOgl2BIO8oyMQLeXw==", + "version": "1.11.1", + "resolved": "https://registry.npmjs.org/@emnapi/core/-/core-1.11.1.tgz", + "integrity": "sha512-RSvbQmHzdKzNsLYa/wHrbc3KN4sYLKAdPZxqiM2HATqv/SBk2/ENSHpvXGaLOMcsAyz0poEGqkmmKYG3OWiJEQ==", "dev": true, "license": "MIT", "optional": true, "dependencies": { - "@emnapi/wasi-threads": "1.2.1", + "@emnapi/wasi-threads": "1.2.2", "tslib": "^2.4.0" } }, "node_modules/@emnapi/runtime": { - "version": "1.10.0", - "resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.10.0.tgz", - "integrity": "sha512-ewvYlk86xUoGI0zQRNq/mC+16R1QeDlKQy21Ki3oSYXNgLb45GV1P6A0M+/s6nyCuNDqe5VpaY84BzXGwVbwFA==", + "version": "1.11.1", + "resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.11.1.tgz", + "integrity": "sha512-vgj7R3y3Wgx24IQaGPA/R6YFXLHVMOZ0uVEyIQPaWs+rd1AzfEMXlAC22FYwO1XkKR6NPsq7mUandH8oIRdZFw==", "dev": true, "license": "MIT", "optional": true, @@ -52,9 +52,9 @@ } }, "node_modules/@emnapi/wasi-threads": { - "version": "1.2.1", - "resolved": "https://registry.npmjs.org/@emnapi/wasi-threads/-/wasi-threads-1.2.1.tgz", - "integrity": "sha512-uTII7OYF+/Mes/MrcIOYp5yOtSMLBWSIoLPpcgwipoiKbli6k322tcoFsxoIIxPDqW01SQGAgko4EzZi2BNv2w==", + "version": "1.2.2", + "resolved": "https://registry.npmjs.org/@emnapi/wasi-threads/-/wasi-threads-1.2.2.tgz", + "integrity": "sha512-c95qOXkHdydNKhscBTebqEC1CVAZpyqOfVfBzQ1qgzyl3gfeldUjIggDbIZgDKsHLgnsM+igH7TJ/eAasaVuMA==", "dev": true, "license": "MIT", "optional": true, @@ -69,14 +69,14 @@ "license": "MIT" }, "node_modules/@napi-rs/wasm-runtime": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@napi-rs/wasm-runtime/-/wasm-runtime-1.1.4.tgz", - "integrity": "sha512-3NQNNgA1YSlJb/kMH1ildASP9HW7/7kYnRI2szWJaofaS1hWmbGI4H+d3+22aGzXXN9IJ+n+GiFVcGipJP18ow==", + "version": "1.1.6", + "resolved": "https://registry.npmjs.org/@napi-rs/wasm-runtime/-/wasm-runtime-1.1.6.tgz", + "integrity": "sha512-ZLv/JdUfkvOy9eCnnBaGfiO+XimbjebAeO+MRQqD/B+FR1tnRN0tpKSJHRbE8sFfS6aqsXZ67TQjfwfsxULVbg==", "dev": true, "license": "MIT", "optional": true, "dependencies": { - "@tybys/wasm-util": "^0.10.1" + "@tybys/wasm-util": "^0.10.3" }, "funding": { "type": "github", @@ -88,9 +88,9 @@ } }, "node_modules/@oxc-project/types": { - "version": "0.133.0", - "resolved": "https://registry.npmjs.org/@oxc-project/types/-/types-0.133.0.tgz", - "integrity": "sha512-KzkdCd6Uxqnf6l3HOw1xfatAlUURA0g14cvBYFyJ5SaNOQbOUvBr9PKArcPcrNIeRsBdgcUzOGrhKveVpvOIGA==", + "version": "0.137.0", + "resolved": "https://registry.npmjs.org/@oxc-project/types/-/types-0.137.0.tgz", + "integrity": "sha512-WT+Gb24i8hmvo85AIv2oEYouEXkRlKAlT9WaCa3TfLgNCN+GhrJOGZuIlMouAh38Qe4QOx26eUOVsq70qXrywA==", "dev": true, "license": "MIT", "funding": { @@ -98,9 +98,9 @@ } }, "node_modules/@rolldown/binding-android-arm64": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-android-arm64/-/binding-android-arm64-1.0.3.tgz", - "integrity": "sha512-454rs7jHngixp/NMxd5srYD57OnzSlZ/eFTETjORQHLwJG1lRtmNOJcBerZlfu4GjKqeq8aCCIQrMdHyhI51Hw==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-android-arm64/-/binding-android-arm64-1.1.3.tgz", + "integrity": "sha512-DT6Z3PhvioeHMvxo+xHc3KtqggrI7CCTXCmC2h/5zUlp5jVitv7XEy+9q5/7v8IolhlioawpMo8Kg0EEBy7J0g==", "cpu": [ "arm64" ], @@ -115,9 +115,9 @@ } }, "node_modules/@rolldown/binding-darwin-arm64": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-arm64/-/binding-darwin-arm64-1.0.3.tgz", - "integrity": "sha512-PcAhP+ynjURNyy8SKGl5DQP94aGuB/7JrXJb/t7P+hanXvQVMWzUvRRhBAcg/lNRadBhoUPqSoP4xw5tR/KBEA==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-arm64/-/binding-darwin-arm64-1.1.3.tgz", + "integrity": "sha512-0NwgwsjM7LrsuVnXMK3koTpagBNOhloc/BNjKqZjv4V5zI5r13qx69uVhRx+o5Z0yy4Hzq+lpy7TAgUG/ocvrw==", "cpu": [ "arm64" ], @@ -132,9 +132,9 @@ } }, "node_modules/@rolldown/binding-darwin-x64": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-x64/-/binding-darwin-x64-1.0.3.tgz", - "integrity": "sha512-9YpfeUvSE2RS7wysJ81uOZkXJz7f7Q55H2Gvp3VEw/EsahqDtrphrZ0EwDLK5vvKOzaCrBsjF8JmnMLcUt78Gg==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-x64/-/binding-darwin-x64-1.1.3.tgz", + "integrity": "sha512-YtiBp4disu6V560loT6PjMdiRaWmVvDNrUunAalbiFx2ggeJwxdAsgZMcoGP17uyAsTwAj5V1niksxlHnVQ1Sw==", "cpu": [ "x64" ], @@ -149,9 +149,9 @@ } }, "node_modules/@rolldown/binding-freebsd-x64": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-freebsd-x64/-/binding-freebsd-x64-1.0.3.tgz", - "integrity": "sha512-yB1IlAsSNHncV6SCTL27/MVGR5htvQsoGxIv5KMGXALp+Ll1wYsn+x98M9MW7qa+NdSbvrrY7ANI4wLJ0n1e6g==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-freebsd-x64/-/binding-freebsd-x64-1.1.3.tgz", + "integrity": "sha512-yD3EkEdXk2LypPxnf/kSZHirarsI8gcPzc62SukhR9VJTyvV+F9Q/GxWNuCojc7sXyuVC4DxRGhdDK4X8VSsbw==", "cpu": [ "x64" ], @@ -166,9 +166,9 @@ } }, "node_modules/@rolldown/binding-linux-arm-gnueabihf": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-1.0.3.tgz", - "integrity": "sha512-Yi30IVAAfLUCy2MseFjbB1jAMDl1VMCAas5StnYp8da9+CKvMd2H2cbEjWcw5NPaPqzvYkVIaF1nNUG+b7u/sw==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-1.1.3.tgz", + "integrity": "sha512-c+8vieQbsD7HNAHKIA34w0GJ9FedFFuJGD+7E6vz7Q3uqAIugL5p45fhlsj4UaAsHpcmlqugBWMhA0/j7o0sIg==", "cpu": [ "arm" ], @@ -183,9 +183,9 @@ } }, "node_modules/@rolldown/binding-linux-arm64-gnu": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-1.0.3.tgz", - "integrity": "sha512-jsO7R8To+AdlYgUmN5sHSCZbfhtMBkO0WUx8iORQnPcMMdgr7qM2DQmMwgabs3GhNztdmoKkMKQFHD6DTMCIQw==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-1.1.3.tgz", + "integrity": "sha512-50jD0uUwLvur7Zz9LHz17kaAdTPjn5wN93hEgjvmYFRZwiR7ZJYovTd5ipyWJDAnXKvZ+wgc+/Ika6dwSF5OcA==", "cpu": [ "arm64" ], @@ -200,9 +200,9 @@ } }, "node_modules/@rolldown/binding-linux-arm64-musl": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-musl/-/binding-linux-arm64-musl-1.0.3.tgz", - "integrity": "sha512-VWkUHwWriDciit80wleYwKILoR/KMvxh/IdwS/paX+ZgpuRpCrKLUdadJbc0NpBEiyhpYawsJ73j9aCvOH+f7Q==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-musl/-/binding-linux-arm64-musl-1.1.3.tgz", + "integrity": "sha512-BO9+oPL8K9poZJBfYPsXNtYjPE5uM3qeehT3aFcW4LITOl+iSqhp0abzjR2nWBUNjIZeKXjAEWBZ64WjNoHd6w==", "cpu": [ "arm64" ], @@ -217,9 +217,9 @@ } }, "node_modules/@rolldown/binding-linux-ppc64-gnu": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-1.0.3.tgz", - "integrity": "sha512-5f1laC0SlIR0yDbFCd8acUhvJIag6N3zC5P7oUPN6wX0aOma+uKJ0wBDH5aq7I1PVI2ttTlhJwzwRIBnLiSGEg==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-1.1.3.tgz", + "integrity": "sha512-f3VpLB1vQ0Eo6ecr/6cekLnvYMFF4YBFoVGkfkvPLq1bAkbAwHYQPZKoAmG6OJyTcxxoC+AvezGx/S1obNC0Mw==", "cpu": [ "ppc64" ], @@ -234,9 +234,9 @@ } }, "node_modules/@rolldown/binding-linux-s390x-gnu": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-1.0.3.tgz", - "integrity": "sha512-Iq4ko0r4XsgbrF/LunNgHtAGLRRVE2kXonAXQ/MV0mC6jQpMOhW1SvtZja2EhC/kd05++bP78dsqBeIQyYJ6Yg==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-1.1.3.tgz", + "integrity": "sha512-AmurZ26Pqx/RI9N1gzEOCklkKXl927yjfXWUUS0O7Puh8ARM/Ob8qfrD3qnWksScdw6cSrW5PSHE9DyLu7+PtA==", "cpu": [ "s390x" ], @@ -251,9 +251,9 @@ } }, "node_modules/@rolldown/binding-linux-x64-gnu": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-gnu/-/binding-linux-x64-gnu-1.0.3.tgz", - "integrity": "sha512-B8m6tD5+/N5FeNQFbKlLA/2yVq9ycQP1SeedyEYYKWBNR3ZQbkvIUcNnDNM03lO1l5F2roiiFJGgvoLLyZXtSg==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-gnu/-/binding-linux-x64-gnu-1.1.3.tgz", + "integrity": "sha512-JJpqs8bRGITDOdbkNKnlojzBabbOHrqjSvDr0IVsZObE1lBcPjxItUEY9eWIDbxaJ3cGrXPWGfGkIxFijg/URg==", "cpu": [ "x64" ], @@ -268,9 +268,9 @@ } }, "node_modules/@rolldown/binding-linux-x64-musl": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-musl/-/binding-linux-x64-musl-1.0.3.tgz", - "integrity": "sha512-pSdpdUJHkuCxun9LE7jvgUB9qsRgaiyNNCX7m/AvHTcq67AiT/Yhoxvw5zPfhrM8k/BfP8ce/hMOpthKDpEUow==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-musl/-/binding-linux-x64-musl-1.1.3.tgz", + "integrity": "sha512-rSJcdjPxzA/by/6/rYs+v+bXU7UjvnbUWz8MJb6kh6+knqB1dCrtHg0uu7C/4haqJvqdkYHQ5IGn+tCH9GLW/g==", "cpu": [ "x64" ], @@ -285,9 +285,9 @@ } }, "node_modules/@rolldown/binding-openharmony-arm64": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-openharmony-arm64/-/binding-openharmony-arm64-1.0.3.tgz", - "integrity": "sha512-OXXS3RKJgX2uLwM+gYyuH5omcH8fL1LJs96pZGgtetVCahON57+d4SJHzTgZiOjxgGkSnpXpOsWuPDGAKAigEg==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-openharmony-arm64/-/binding-openharmony-arm64-1.1.3.tgz", + "integrity": "sha512-hQ3/PYkDJICgevvyNcVrihVeqq7k1Pp3VZ9lY+dauAYUJKO+auqApvANhvR1An9BhmqYKvW2Mu1F9u4DXSMLxQ==", "cpu": [ "arm64" ], @@ -302,9 +302,9 @@ } }, "node_modules/@rolldown/binding-wasm32-wasi": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-wasm32-wasi/-/binding-wasm32-wasi-1.0.3.tgz", - "integrity": "sha512-JTtb8BWFynicNSoPrehsCzBtOKjZ6jhMiPFEmOiuXg1Fl8dn2KHQob+GuPSGR0dryQa1PQJbzjF3dqO/whhjLg==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-wasm32-wasi/-/binding-wasm32-wasi-1.1.3.tgz", + "integrity": "sha512-Elcv/BtML9lXrV6JuKITc/grN2kYV9gjsQpW8Jfw4ioK0TOkjBjye0nnyqQNy9STNaI20lXNaQBRrD5gSgR0Yg==", "cpu": [ "wasm32" ], @@ -312,18 +312,18 @@ "license": "MIT", "optional": true, "dependencies": { - "@emnapi/core": "1.10.0", - "@emnapi/runtime": "1.10.0", - "@napi-rs/wasm-runtime": "^1.1.4" + "@emnapi/core": "1.11.1", + "@emnapi/runtime": "1.11.1", + "@napi-rs/wasm-runtime": "^1.1.6" }, "engines": { "node": "^20.19.0 || >=22.12.0" } }, "node_modules/@rolldown/binding-win32-arm64-msvc": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-1.0.3.tgz", - "integrity": "sha512-gEdFFEN70A/jxb2svrWsN3aDL7OUtmvlOy+6fa2jxG8K0wQ1ZbdeLGnidov6Yu5/733dI5ySfzFlQ/cb0bSz1g==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-1.1.3.tgz", + "integrity": "sha512-2DrEfhluH9yhiaFApmsjsjwrSYbNcY1oFTzYSP1a535jDbV98zCFanA/96TBUd0iDFcxGmw9QRExwGCXz3U+/g==", "cpu": [ "arm64" ], @@ -338,9 +338,9 @@ } }, "node_modules/@rolldown/binding-win32-x64-msvc": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-x64-msvc/-/binding-win32-x64-msvc-1.0.3.tgz", - "integrity": "sha512-eXB7CHuaQdqmJcc3koCNtNPmT/bj2gc999kUFgBxG8Ac0NdgXc4rkCHhqrgrhN3zddvvvrgzj1e90SuSfmyIXA==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-x64-msvc/-/binding-win32-x64-msvc-1.1.3.tgz", + "integrity": "sha512-OL4OMk7UPXOeVGGd3qo5zJyPIljf4AFgk5QAkPPS+OoLuOOozhuaQGC18MxVTnw/06q93gShAJzlwnSCY9YtqA==", "cpu": [ "x64" ], @@ -362,9 +362,9 @@ "license": "MIT" }, "node_modules/@tybys/wasm-util": { - "version": "0.10.2", - "resolved": "https://registry.npmjs.org/@tybys/wasm-util/-/wasm-util-0.10.2.tgz", - "integrity": "sha512-RoBvJ2X0wuKlWFIjrwffGw1IqZHKQqzIchKaadZZfnNpsAYp2mM0h36JtPCjNDAHGgYez/15uMBpfGwchhiMgg==", + "version": "0.10.3", + "resolved": "https://registry.npmjs.org/@tybys/wasm-util/-/wasm-util-0.10.3.tgz", + "integrity": "sha512-F3fo1MYrRJYL3zER0OUOmkutjr1Vp23m7OsSgp7nq4SP6OqX6C/56XFIPAl5bt3zaBRjmW7SGz3u/6LwFpYcOg==", "dev": true, "license": "MIT", "optional": true, @@ -399,13 +399,13 @@ } }, "node_modules/@vitejs/plugin-react": { - "version": "6.0.2", - "resolved": "https://registry.npmjs.org/@vitejs/plugin-react/-/plugin-react-6.0.2.tgz", - "integrity": "sha512-DlSMqo4WhThw4vB8Mpn0Woe9J+Jfq1geJ61AKW0QEgLzGMNwtIMdxbDUzLxcun8W7NbJO0e2Jg/Nxm3cCSVzzg==", + "version": "6.0.3", + "resolved": "https://registry.npmjs.org/@vitejs/plugin-react/-/plugin-react-6.0.3.tgz", + "integrity": "sha512-vmFvco5/QuC2f9Oj+wTk0+9XeDFkHxSamwZKYc7MxYwKICfvUvlMhqKI0VuICPltGqh1neqBKDvO4kes1ya8vg==", "dev": true, "license": "MIT", "dependencies": { - "@rolldown/pluginutils": "^1.0.0" + "@rolldown/pluginutils": "^1.0.1" }, "engines": { "node": "^20.19.0 || >=22.12.0" @@ -895,18 +895,18 @@ } }, "node_modules/lucide-react": { - "version": "1.17.0", - "resolved": "https://registry.npmjs.org/lucide-react/-/lucide-react-1.17.0.tgz", - "integrity": "sha512-9FA9evdox/JQL5PT57fdA1x/yg8T7knJ98+zjTL3UfKza6pflQUUh3XtaQIHKvnsJw1lmsEyHVlt5jchYxOQ5w==", + "version": "1.22.0", + "resolved": "https://registry.npmjs.org/lucide-react/-/lucide-react-1.22.0.tgz", + "integrity": "sha512-c9o3l0PiNcgOQDW4F31BEYHudE7kgxVt3o30qMl36ZPwTxXlGB4QnLilhERvVM4uh/pl5MDyY1/gzZSYcHDtBg==", "license": "ISC", "peerDependencies": { "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "node_modules/nanoid": { - "version": "3.3.12", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.12.tgz", - "integrity": "sha512-ZB9RH/39qpq5Vu6Y+NmUaFhQR6pp+M2Xt76XBnEwDaGcVAqhlvxrl3B2bKS5D3NH3QR76v3aSrKaF/Kiy7lEtQ==", + "version": "3.3.15", + "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.15.tgz", + "integrity": "sha512-y7Wygv/7mEOvxTuEQDB8StXdMRBWf1kR/tlhAzBRUFkB2jfcLOAxO/SHmOO2zgz1pVgK29/kyupn059/bCHdjA==", "dev": true, "funding": [ { @@ -943,9 +943,9 @@ } }, "node_modules/postcss": { - "version": "8.5.15", - "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.15.tgz", - "integrity": "sha512-FfR8sjd4em2T6fb3I2MwAJU7HWVMr9zba+enmQeeWFfCbm+UOC/0X4DS8XtpUTMwWMGbjKYP7xjfNekzyGmB3A==", + "version": "8.5.16", + "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.16.tgz", + "integrity": "sha512-vuwillviilfKZsg0VGj5R/YwwcHx4SLsIOI/7K6mQkWx+l5cUHTjj5g0AasTBcyXsbfTgrwsUNmVUb5xVwyPwg==", "dev": true, "funding": [ { @@ -1003,13 +1003,13 @@ } }, "node_modules/rolldown": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/rolldown/-/rolldown-1.0.3.tgz", - "integrity": "sha512-i00lAJ2ks1BYr7rjNjKC7BcqAS7nVfiT3QX1SI5aY+AFHblCmaUf9OE9dbdzDvW6dJxbi2ZCZiy9v3CcwOiX3g==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/rolldown/-/rolldown-1.1.3.tgz", + "integrity": "sha512-1F1eEtUBtFvcGm1HQ9TiUIUHPQG7mSAODrhIzjxoUEFuo8OcbrGLiVLkevNgj84TE4lnHvnumwFjhJO5Eu135g==", "dev": true, "license": "MIT", "dependencies": { - "@oxc-project/types": "=0.133.0", + "@oxc-project/types": "=0.137.0", "@rolldown/pluginutils": "^1.0.0" }, "bin": { @@ -1019,21 +1019,21 @@ "node": "^20.19.0 || >=22.12.0" }, "optionalDependencies": { - "@rolldown/binding-android-arm64": "1.0.3", - "@rolldown/binding-darwin-arm64": "1.0.3", - "@rolldown/binding-darwin-x64": "1.0.3", - "@rolldown/binding-freebsd-x64": "1.0.3", - "@rolldown/binding-linux-arm-gnueabihf": "1.0.3", - "@rolldown/binding-linux-arm64-gnu": "1.0.3", - "@rolldown/binding-linux-arm64-musl": "1.0.3", - "@rolldown/binding-linux-ppc64-gnu": "1.0.3", - "@rolldown/binding-linux-s390x-gnu": "1.0.3", - "@rolldown/binding-linux-x64-gnu": "1.0.3", - "@rolldown/binding-linux-x64-musl": "1.0.3", - "@rolldown/binding-openharmony-arm64": "1.0.3", - "@rolldown/binding-wasm32-wasi": "1.0.3", - "@rolldown/binding-win32-arm64-msvc": "1.0.3", - "@rolldown/binding-win32-x64-msvc": "1.0.3" + "@rolldown/binding-android-arm64": "1.1.3", + "@rolldown/binding-darwin-arm64": "1.1.3", + "@rolldown/binding-darwin-x64": "1.1.3", + "@rolldown/binding-freebsd-x64": "1.1.3", + "@rolldown/binding-linux-arm-gnueabihf": "1.1.3", + "@rolldown/binding-linux-arm64-gnu": "1.1.3", + "@rolldown/binding-linux-arm64-musl": "1.1.3", + "@rolldown/binding-linux-ppc64-gnu": "1.1.3", + "@rolldown/binding-linux-s390x-gnu": "1.1.3", + "@rolldown/binding-linux-x64-gnu": "1.1.3", + "@rolldown/binding-linux-x64-musl": "1.1.3", + "@rolldown/binding-openharmony-arm64": "1.1.3", + "@rolldown/binding-wasm32-wasi": "1.1.3", + "@rolldown/binding-win32-arm64-msvc": "1.1.3", + "@rolldown/binding-win32-x64-msvc": "1.1.3" } }, "node_modules/rxjs": { @@ -1157,16 +1157,16 @@ "license": "0BSD" }, "node_modules/vite": { - "version": "8.0.16", - "resolved": "https://registry.npmjs.org/vite/-/vite-8.0.16.tgz", - "integrity": "sha512-h9bXPmJichP5fLmVQo3PyaGSDE2n3aPuomeAlVRm0JLmt4rY6zmPKd59HYI4LNW8oTK7tlTsuC7l/m7awx9Jcw==", + "version": "8.1.0", + "resolved": "https://registry.npmjs.org/vite/-/vite-8.1.0.tgz", + "integrity": "sha512-BuJcQK/56NQTWDGn4ABea3q4SSBdNPWwNZKTkkUpcMPnLoquSYH8llRtSUIgoL1KSCpHt5eghLShn50mH36y7Q==", "dev": true, "license": "MIT", "dependencies": { "lightningcss": "^1.32.0", "picomatch": "^4.0.4", "postcss": "^8.5.15", - "rolldown": "1.0.3", + "rolldown": "~1.1.2", "tinyglobby": "^0.2.17" }, "bin": { @@ -1183,7 +1183,7 @@ }, "peerDependencies": { "@types/node": "^20.19.0 || >=22.12.0", - "@vitejs/devtools": "^0.1.18", + "@vitejs/devtools": "^0.3.0", "esbuild": "^0.27.0 || ^0.28.0", "jiti": ">=1.21.0", "less": "^4.0.0", diff --git a/benchmarks-website/src/config.js b/benchmarks-website/src/config.js index c6ce9060428..dfe5555019d 100644 --- a/benchmarks-website/src/config.js +++ b/benchmarks-website/src/config.js @@ -14,6 +14,15 @@ export const QUERY_SUITES = [ tags: ["Queries (NVMe)"], hiddenDatasets: ["datafusion:lance"], }, + { + prefix: "clickbench-sorted", + displayName: "Clickbench Sorted", + queryPrefix: "CLICKBENCH SORTED", + description: + "ClickBench queries over data globally sorted by event date and event time", + tags: ["Queries (NVMe)"], + hiddenDatasets: ["datafusion:lance"], + }, { prefix: "statpopgen", displayName: "Statistical and Population Genetics", diff --git a/benchmarks-website/src/utils.js b/benchmarks-website/src/utils.js index 2d6575fe814..140ebebeb6c 100644 --- a/benchmarks-website/src/utils.js +++ b/benchmarks-website/src/utils.js @@ -97,6 +97,7 @@ export function getBenchmarkDescription(categoryName) { 'Compression': 'Measures encoding and decoding throughput (MB/s) for Vortex and Parquet files', 'Compression Size': 'Compares compressed file sizes across different encoding strategies', 'Clickbench': "ClickHouse's analytical benchmark suite on web analytics data", + 'Clickbench Sorted': 'ClickBench queries over data globally sorted by event date and event time', 'Statistical and Population Genetics': 'Statistical and population genetics queries on gnomAD dataset', }; return descriptions[categoryName] || ''; diff --git a/benchmarks/datafusion-bench/src/lib.rs b/benchmarks/datafusion-bench/src/lib.rs index c8353eb1f85..1100d38d24e 100644 --- a/benchmarks/datafusion-bench/src/lib.rs +++ b/benchmarks/datafusion-bench/src/lib.rs @@ -45,10 +45,7 @@ pub fn get_session_context() -> SessionContext { .build_arc() .expect("could not build runtime environment"); - let factory = VortexFormatFactory::new().with_options(VortexTableOptions { - projection_pushdown: true, - ..Default::default() - }); + let factory = VortexFormatFactory::new().with_options(vortex_table_options()); let mut session_state_builder = SessionStateBuilder::new() .with_config(SessionConfig::from_env().expect("shouldn't fail")) @@ -114,11 +111,21 @@ pub fn format_to_df_format(format: Format) -> Arc { Format::Csv => Arc::new(CsvFormat::default()) as _, Format::Arrow => Arc::new(ArrowFormat), Format::Parquet => Arc::new(ParquetFormat::new()), - Format::OnDiskVortex | Format::VortexCompact => { - Arc::new(VortexFormat::new(SESSION.clone())) - } + Format::OnDiskVortex | Format::VortexCompact => Arc::new(VortexFormat::new_with_options( + SESSION.clone(), + vortex_table_options(), + )), Format::OnDiskDuckDB | Format::Lance => { unimplemented!("Format {format} cannot be turned into a DataFusion `FileFormat`") } } } + +fn vortex_table_options() -> VortexTableOptions { + let mut opts = VortexTableOptions::default(); + + opts.predicate_pushdown = true; + opts.predicate_pushdown = true; + + opts +} diff --git a/benchmarks/duckdb-bench/src/lib.rs b/benchmarks/duckdb-bench/src/lib.rs index fed9f82b004..4ec1efd0993 100644 --- a/benchmarks/duckdb-bench/src/lib.rs +++ b/benchmarks/duckdb-bench/src/lib.rs @@ -16,7 +16,6 @@ use vortex_bench::Format; use vortex_bench::IdempotentPath; use vortex_bench::generate_duckdb_registration_sql; use vortex_bench::runner::BenchmarkQueryResult; -use vortex_duckdb::duckdb::Config; use vortex_duckdb::duckdb::Connection; use vortex_duckdb::duckdb::Database; use vortex_duckdb::duckdb::QueryResult; @@ -27,6 +26,8 @@ pub struct DuckClient { connection: Option, pub db_path: PathBuf, pub threads: Option, + /// `INSTALL spatial; LOAD spatial;` for SpatialBench. + init_sql: Vec, } impl DuckClient { @@ -68,37 +69,38 @@ impl DuckClient { connection: Some(connection), db_path, threads, + init_sql: Vec::new(), }) } + /// Run `statements` now and after every subsequent [`DuckClient::reopen`]. + pub fn set_init_sql(&mut self, statements: Vec) -> Result<()> { + for stmt in &statements { + self.connection().query(stmt)?; + } + self.init_sql = statements; + Ok(()) + } + pub fn open_and_setup_database( path: Option, threads: Option, ) -> Result<(Database, Connection)> { - let mut config = Config::new().vortex_expect("failed to create duckdb config"); - - // Set DuckDB thread count if specified - if let Some(thread_count) = threads { - config.set("threads", &format!("{}", thread_count))?; - } - let db = match path { - Some(path) => Database::open_with_config(path, config), - None => Database::open_in_memory_with_config(config), + Some(path) => Database::open(path), + None => Database::open_in_memory(), }?; let connection = db.connect()?; vortex_duckdb::initialize(&db)?; - // Enable Parquet metadata cache for all benchmark runs. - // + if let Some(thread_count) = threads { + connection.query(&format!("SET threads = {thread_count}"))?; + } + // `parquet_metadata_cache` is an extension-specific option that's // only available after the Parquet extension is loaded. The Parquet // extension is loaded after the connection is established. - // - // Passing the option to `open_with_config` before leads to - // "Invalid Input Error: The following options were not recognized: - // parquet_metadata_cache" when running DuckDB in debug mode. connection.query("SET parquet_metadata_cache = true")?; Ok((db, connection)) @@ -118,6 +120,14 @@ impl DuckClient { self.db = Some(db); self.connection = Some(connection); + // Replay init SQL (e.g. LOAD spatial). + for stmt in &self.init_sql { + self.connection + .as_ref() + .vortex_expect("connection just opened") + .query(stmt)?; + } + Ok(()) } @@ -133,6 +143,7 @@ impl DuckClient { connection: Some(connection), db_path, threads: None, + init_sql: Vec::new(), }) } diff --git a/benchmarks/duckdb-bench/src/main.rs b/benchmarks/duckdb-bench/src/main.rs index 8ba4937f566..cf4fa071067 100644 --- a/benchmarks/duckdb-bench/src/main.rs +++ b/benchmarks/duckdb-bench/src/main.rs @@ -171,12 +171,13 @@ fn main() -> anyhow::Result<()> { &filtered_queries, mode, |format| { - let ctx = DuckClient::new( + let mut ctx = DuckClient::new( &*benchmark, format, args.delete_duckdb_database, args.threads, )?; + ctx.set_init_sql(benchmark.engine_init_sql(Engine::DuckDB))?; ctx.register_tables(&*benchmark, format)?; // Duckdb doesn't support octet_length for strings but we need this diff --git a/benchmarks/lance-bench/Cargo.toml b/benchmarks/lance-bench/Cargo.toml index a32b27d97cd..7cab6af2119 100644 --- a/benchmarks/lance-bench/Cargo.toml +++ b/benchmarks/lance-bench/Cargo.toml @@ -15,8 +15,8 @@ version.workspace = true publish = false [dependencies] -lance = { version = "6", default-features = false } -lance-encoding = { version = "6" } +lance = { version = "7", default-features = false } +lance-encoding = { version = "7" } anyhow = { workspace = true } arrow-cast = { version = "58" } diff --git a/benchmarks/random-access-bench/src/main.rs b/benchmarks/random-access-bench/src/main.rs index d64482c01a3..6722265fde6 100644 --- a/benchmarks/random-access-bench/src/main.rs +++ b/benchmarks/random-access-bench/src/main.rs @@ -43,7 +43,7 @@ mod render; // --------------------------------------------------------------------------- /// Access pattern for random access benchmarks. -#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, ValueEnum)] pub enum AccessPattern { /// Multiple clusters of sequential indices scattered across the dataset, /// simulating workloads with spatial locality (e.g. scanning nearby records). @@ -62,8 +62,6 @@ impl AccessPattern { } } -const ACCESS_PATTERNS: [AccessPattern; 2] = [AccessPattern::Correlated, AccessPattern::Uniform]; - /// Number of clusters for the correlated pattern. const NUM_CLUSTERS: usize = 5; @@ -190,6 +188,14 @@ struct Args { default_values_t = vec![DatasetArg::Taxi, DatasetArg::FeatureVectors, DatasetArg::NestedLists, DatasetArg::NestedStructs] )] datasets: Vec, + /// Which access patterns to benchmark. + #[arg( + long, + value_delimiter = ',', + value_enum, + default_values_t = vec![AccessPattern::Correlated, AccessPattern::Uniform] + )] + patterns: Vec, /// Whether to reopen the file on each iteration, use a cached handle, or run both. #[arg(long, value_enum, default_value_t = OpenMode::Both)] open_mode: OpenMode, @@ -201,22 +207,22 @@ async fn main() -> Result<()> { setup_logging_and_tracing(args.verbose, args.tracing)?; - let datasets: Vec> = args - .datasets - .into_iter() - .map(|d| d.into_dataset()) - .collect(); - - run_random_access( - &datasets, - args.formats, - args.time_limit, - args.open_mode, - args.display_format, - args.output_path, - args.gh_json_v3, - ) - .await + let config = RunConfig { + datasets: args + .datasets + .into_iter() + .map(|d| d.into_dataset()) + .collect(), + formats: args.formats, + patterns: args.patterns, + time_limit: args.time_limit, + open_mode: args.open_mode, + display_format: args.display_format, + output_path: args.output_path, + gh_json_v3: args.gh_json_v3, + }; + + run_random_access(config).await } // --------------------------------------------------------------------------- @@ -379,15 +385,30 @@ const BENCHMARK_ID: &str = "random-access"; /// Fixed indices used by the original taxi benchmark (preserved for historical continuity). const FIXED_TAXI_INDICES: [u64; 6] = [10, 11, 12, 13, 100_000, 3_000_000]; -async fn run_random_access( - datasets: &[Box], +/// Resolved configuration for a single random-access benchmark invocation. +struct RunConfig { + datasets: Vec>, formats: Vec, + patterns: Vec, time_limit: u64, open_mode: OpenMode, display_format: DisplayFormat, output_path: Option, gh_json_v3: Option, -) -> Result<()> { +} + +async fn run_random_access(config: RunConfig) -> Result<()> { + let RunConfig { + datasets, + formats, + patterns, + time_limit, + open_mode, + display_format, + output_path, + gh_json_v3, + } = config; + let reopen_variants: &[bool] = match open_mode { OpenMode::Cached => &[false], OpenMode::Reopen => &[true], @@ -398,7 +419,7 @@ async fn run_random_access( .iter() .map(|d| { let legacy_extra = if d.name() == "taxi" { formats.len() } else { 0 }; - (formats.len() * ACCESS_PATTERNS.len() + legacy_extra) * reopen_variants.len() + (formats.len() * patterns.len() + legacy_extra) * reopen_variants.len() }) .sum(); let progress = ProgressBar::new(total_steps as u64); @@ -408,7 +429,7 @@ async fn run_random_access( // Iteration order matters for the table renderer: row order is set by the // first time each `(dataset, pattern)` pair is observed. - for dataset in datasets { + for dataset in &datasets { for format in &formats { if dataset.name() == "taxi" { let name = measurement_name(dataset.name(), None, *format); @@ -436,7 +457,7 @@ async fn run_random_access( } } - for pattern in &ACCESS_PATTERNS { + for pattern in &patterns { let indices = generate_indices(dataset.as_ref(), *pattern); let name = measurement_name(dataset.name(), Some(*pattern), *format); for &reopen in reopen_variants { diff --git a/clippy.toml b/clippy.toml index e761d3c17f7..9f6d5bd21b8 100644 --- a/clippy.toml +++ b/clippy.toml @@ -15,4 +15,6 @@ disallowed-methods = [ { path = "itertools::Itertools::counts", reason = "It uses the default hasher which is slow for primitives. Just inline the loop for better performance.", allow-invalid = true }, { path = "std::result::Result::and", reason = "This method is a footgun, especially when working with `Result`." }, { path = "std::thread::available_parallelism", reason = "This function might do an unbounded amount of work, use `vortex_utils::parallelism::get_available_parallelism instead" }, + { path = "vortex_session::registry::Id::new", reason = "Interning a static id on every call grabs the interner lock (#8380). Use a `CachedId` static for static ids; for a dynamic string, annotate the call with `#[expect(clippy::disallowed_methods)]`.", allow-invalid = true }, + { path = "vortex_session::registry::Id::new_static", reason = "Interning a static id on every call grabs the interner lock; use a `CachedId` static instead (#8380).", allow-invalid = true }, ] diff --git a/docs/api/python/runtime.rst b/docs/api/python/runtime.rst index 1c404c74d32..b5307a4041b 100644 --- a/docs/api/python/runtime.rst +++ b/docs/api/python/runtime.rst @@ -9,6 +9,7 @@ otherwise to the number of available CPU cores minus one. Use .. autosummary:: :nosignatures: + ~vortex.cuda_extension_installed ~vortex.set_worker_threads ~vortex.worker_threads @@ -16,6 +17,8 @@ otherwise to the number of available CPU cores minus one. Use
+.. autofunction:: vortex.cuda_extension_installed + .. autofunction:: vortex.set_worker_threads .. autofunction:: vortex.worker_threads diff --git a/docs/developer-guide/benchmarking.md b/docs/developer-guide/benchmarking.md index 0034965a056..dfa0fa1cd45 100644 --- a/docs/developer-guide/benchmarking.md +++ b/docs/developer-guide/benchmarking.md @@ -212,8 +212,10 @@ Benchmarks run automatically on all commits to `develop` and can be run on-deman `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 a parametric matrix - of suites, engines, formats, and storage backends (NVMe, S3). +- **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). 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/index.md b/docs/developer-guide/embedding/index.md index f09d1841762..e204939431b 100644 --- a/docs/developer-guide/embedding/index.md +++ b/docs/developer-guide/embedding/index.md @@ -2,8 +2,8 @@ :::{warning} This section is under construction. For guidance on embedding Vortex, please join the -[Vortex Slack channel](https://join.slack.com/t/vortex-array/shared_invite/zt-2ycp2w24h-sRdrGbMGPmQwCuPQT40Jig) -or start a [GitHub Discussion](https://github.com/spiraldb/vortex/discussions). +[Vortex Slack channel](https://vortex.dev/slack) +or start a [GitHub Discussion](https://github.com/vortex-data/vortex/discussions). ::: Vortex can be embedded into applications and services via its C FFI, C++ wrapper, or the Scan API. diff --git a/docs/developer-guide/extending/index.md b/docs/developer-guide/extending/index.md index 78b7a1829e3..ff004a69bc5 100644 --- a/docs/developer-guide/extending/index.md +++ b/docs/developer-guide/extending/index.md @@ -3,7 +3,7 @@ :::{warning} This section is under construction. For guidance on extending Vortex, please join the [Vortex Slack channel](https://vortex.dev/slack) -or start a [GitHub Discussion](https://github.com/spiraldb/vortex/discussions). +or start a [GitHub Discussion](https://github.com/vortex-data/vortex/discussions). ::: Vortex is designed to be extended with custom types, encodings, layouts, and compute functions. diff --git a/docs/developer-guide/index.md b/docs/developer-guide/index.md index 37ad276044d..e6eec4bfb70 100644 --- a/docs/developer-guide/index.md +++ b/docs/developer-guide/index.md @@ -24,6 +24,7 @@ internals/session internals/async-runtime internals/vtables internals/execution +internals/stats-pruning internals/io internals/serialization internals/cuda @@ -38,4 +39,4 @@ caption: Integrations integrations/datafusion integrations/duckdb integrations/spark -``` \ No newline at end of file +``` diff --git a/docs/developer-guide/internals/execution.md b/docs/developer-guide/internals/execution.md index be92e29ceef..15ba16a8135 100644 --- a/docs/developer-guide/internals/execution.md +++ b/docs/developer-guide/internals/execution.md @@ -115,8 +115,9 @@ Each child is given the opportunity to execute its parent in a fused manner via `ExecuteParentKernel`. Unlike reduce rules, parent kernels may read buffers and perform real computation. -An encoding declares its parent kernels in a `ParentKernelSet`, specifying which parent types -each kernel handles via a `Matcher`: +An encoding declares parent kernels by implementing `ExecuteParentKernel` and registering them +with the session-scoped kernel registry, specifying which parent types each kernel handles via a +`Matcher`: ```rust pub trait ExecuteParentKernel { @@ -124,12 +125,18 @@ pub trait ExecuteParentKernel { fn execute_parent( &self, - array: &V::Array, // the child + array: ArrayView<'_, V>, // the child parent: ::Match<'_>, // the matched parent child_idx: usize, ctx: &mut ExecutionCtx, ) -> VortexResult>; } + +pub fn initialize(session: &VortexSession) { + session + .kernels() + .register_execute_parent_kernel(parent_id, Child, Kernel); +} ``` Examples: @@ -195,17 +202,23 @@ execute_until(root): │ restore builder, │ │ loop ▼ │ ┌────────────────────────────────────────────┐ - │ │ Step 2a: current_array.execute_parent( │ - │ │ stack.top.parent_array ) │ - │ │ child looks UP at the suspended parent │ + │ │ Step 2a: execute_parent_for_child( │ + │ │ stack.top.parent_array, │ + │ │ current_array ) │ + │ │ lookup session kernels by key: │ + │ │ (parent.encoding_id(), │ + │ │ child.encoding_id()) │ │ ├────────────┬───────────────────────────────┘ │ │ Some │ None │ │ │ │ │ ▼ │ │ ┌─────────────────────────────────────────┐ - │ │ │ Step 2b: each child.execute_parent( │ - │ │ │ current_array ) │ - │ │ │ children look UP at current_array │ + │ │ │ Step 2b: for each child: │ + │ │ │ execute_parent_for_child( │ + │ │ │ current_array, child ) │ + │ │ │ lookup session kernels by key: │ + │ │ │ (parent.encoding_id(), │ + │ │ │ child.encoding_id()) │ │ │ ├──────────┬──────────────────────────────┘ │ │ │ Some │ None │ │ │ │ @@ -236,6 +249,9 @@ Step 2a and Step 2b are skipped while `current_builder` is active. `AppendChild` consumes `current_array`: some slots already live in the builder, so a parent rewrite would observe inconsistent state and could discard accumulated builder data. +Both Step 2a and Step 2b route through `execute_parent_for_child`, which looks up +`(parent.encoding_id(), child.encoding_id())` in the session kernel snapshot. + ## Incremental Execution Execution is incremental: each call to `execute` moves the array one step closer to canonical @@ -256,12 +272,12 @@ codes through the slice, missing the Dict-RLE optimization entirely. Incremental avoids this: 1. First iteration: the slice `execute` returns `ExecuteSlot` for its `RunEndArray` child. - Once that child is in focus, Step 2a gives it a chance to rewrite the suspended slice - parent before the child is forced toward canonical form. + Once that child is in focus, Step 2a looks up a registered parent kernel for the + `(slice, runend)` pair before the child is forced toward canonical form. 2. Second iteration: the `RunEndArray` codes child now matches the Dict-RLE pattern. Its - `execute_parent` provides a fused kernel that expands runs while performing dictionary - lookups in a single pass, returning the canonical array directly. + registered execute-parent kernel expands runs while performing dictionary lookups in a single + pass, returning the canonical array directly. ## Walkthrough: Executing a RunEnd-Encoded Array diff --git a/docs/developer-guide/internals/session.md b/docs/developer-guide/internals/session.md index 5aff27845af..71158eaae20 100644 --- a/docs/developer-guide/internals/session.md +++ b/docs/developer-guide/internals/session.md @@ -12,14 +12,15 @@ or other shared resources. The session is built on two primitives from the `vortex-session` crate: -- **`VortexSession`** -- a cloneable, thread-safe map from Rust `TypeId` to a boxed value. Any - type that is `Send + Sync + Debug + 'static` can be stored as a session variable. +- **`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. -Because `VortexSession` is backed by an `Arc`, cloning is cheap and all clones share -the same state. This makes it safe to hand the session to multiple threads, tasks, or I/O -operations without coordination. +Because `VortexSession` is 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. ## Component Registries @@ -95,8 +96,8 @@ all built-in components and encodings: let session = VortexSession::default(); ``` -For tests or specialized use-cases, sessions can be assembled from individual components using -the `.with::()` builder: +For tests or specialized use-cases, sessions can be assembled from individual components by +starting from an empty session and chaining the `.with::()` helpers: ```rust let session = VortexSession::empty() diff --git a/docs/developer-guide/internals/stats-pruning.md b/docs/developer-guide/internals/stats-pruning.md new file mode 100644 index 00000000000..4acaf9efc2a --- /dev/null +++ b/docs/developer-guide/internals/stats-pruning.md @@ -0,0 +1,39 @@ +# Stats Pruning + +Vortex uses statistics to prove when a filter cannot match a row group, zone, or +file. The proof expression returns `true` when the input can be skipped. It +returns `false` or `null` when pruning is not proven. + +Both `false` and `null` are non-pruning outcomes, but they mean different +things. `false` means the available stats disproved the skip proof. `null` means +the proof was unknown, usually because a required stat was missing or inexact. + +The pruning pipeline has two phases: + +1. `Expression::falsify(scope, session)` asks the session's + `StatsRewriteRule`s to rewrite a filter into an abstract proof expression. + Rules describe semantics in terms of `vortex.stat(input, aggregate_fn)` + placeholders. These placeholders name the statistic needed by the proof, but + not where that statistic is stored. +2. `bind_stats` lowers those abstract stat placeholders with a `StatBinder`. + The binder maps stats to the representation used by the caller, such as + zone-map table fields, file-level stat literals, or typed null literals for + missing stats. + +Missing stats lower to typed null literals. This preserves the three-valued +logic used by pruning: only a non-null `true` value proves that the scope can be +skipped. A missing stat therefore cannot accidentally prune data. + +## Binding Targets + +Zone maps bind stats to fields in their per-zone stats table. The lowered +expression is evaluated against that table and produces a mask where `true` +means the zone can be skipped. + +File-level stats bind stats to literal values from the file footer. The lowered +expression is reduced and evaluated once for the full file. If it evaluates to +`true`, the file stats reader can return an all-false pruning mask without +reading child layouts. + +For the layout model around these pruning points, see +[Layouts](../../concepts/layouts.md) and [Scanning](../../concepts/scanning.md). diff --git a/encodings/alp/benches/alp_compress.rs b/encodings/alp/benches/alp_compress.rs index 70a566ab0f8..6e769ffc22e 100644 --- a/encodings/alp/benches/alp_compress.rs +++ b/encodings/alp/benches/alp_compress.rs @@ -19,7 +19,6 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::NativePType; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; @@ -51,8 +50,11 @@ const BENCH_ARGS: &[(usize, f64, f64)] = &[ (10_000, 0.1, 1.0), ]; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_alp::initialize(&session); + session +}); #[divan::bench(types = [f32, f64], args = BENCH_ARGS)] fn compress_alp(bencher: Bencher, args: (usize, f64, f64)) { @@ -144,8 +146,8 @@ fn compress_rd(bencher: Bencher, args: (usize, f64) let encoder = RDEncoder::new(primitive.as_slice::()); bencher - .with_inputs(|| (&primitive, &encoder, SESSION.create_execution_ctx())) - .bench_refs(|(primitive, encoder, ctx)| encoder.encode(primitive.as_view(), ctx)) + .with_inputs(|| (&primitive, &encoder)) + .bench_refs(|(primitive, encoder)| encoder.encode(primitive.as_view())) } #[divan::bench(types = [f32, f64], args = RD_BENCH_ARGS)] @@ -153,7 +155,7 @@ fn decompress_rd(bencher: Bencher, args: (usize, f6 let (n, fraction_patch) = args; let primitive = make_rd_array::(n, fraction_patch); let encoder = RDEncoder::new(primitive.as_slice::()); - let encoded = encoder.encode(primitive.as_view(), &mut SESSION.create_execution_ctx()); + let encoded = encoder.encode(primitive.as_view()); bencher .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array.rs index 3f7c2ca3e6d..5b0984dce7a 100644 --- a/encodings/alp/src/alp/array.rs +++ b/encodings/alp/src/alp/array.rs @@ -48,7 +48,6 @@ use vortex_session::registry::CachedId; use crate::ALPFloat; use crate::alp::Exponents; use crate::alp::decompress::execute_decompress; -use crate::alp::rules::PARENT_KERNELS; use crate::alp::rules::RULES; /// A [`ALP`]-encoded Vortex array. @@ -103,6 +102,14 @@ impl VTable for ALP { 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, @@ -188,15 +195,6 @@ impl VTable for ALP { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[array_slots(ALP)] @@ -481,11 +479,9 @@ mod tests { use rstest::rstest; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; - use vortex_array::session::ArraySession; use vortex_error::VortexExpect; use vortex_session::VortexSession; @@ -493,8 +489,11 @@ mod tests { use crate::alp_encode; use crate::decompress_into_array; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[rstest] #[case(0)] @@ -519,10 +518,9 @@ mod tests { .unwrap() }; // Compare against the traditional array-based decompress path - let expected = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(result_canonical.into_array(), expected); + assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx); } #[rstest] @@ -537,26 +535,19 @@ mod tests { #[case(2049)] fn test_execute_f64(#[case] size: usize) { let values = PrimitiveArray::from_iter((0..size).map(|i| i as f64)); - let encoded = alp_encode( - values.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(values.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); - let result_canonical = { - let mut ctx = SESSION.create_execution_ctx(); - encoded - .clone() - .into_array() - .execute::(&mut ctx) - .unwrap() - }; + let mut ctx = SESSION.create_execution_ctx(); + let result_canonical = encoded + .clone() + .into_array() + .execute::(&mut ctx) + .unwrap(); // Compare against the traditional array-based decompress path - let expected = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(result_canonical.into_array(), expected); + assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx); } #[rstest] @@ -576,27 +567,20 @@ mod tests { .collect(); let array = PrimitiveArray::from_iter(values); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().unwrap().array_len() > 0); - let result_canonical = { - let mut ctx = SESSION.create_execution_ctx(); - encoded - .clone() - .into_array() - .execute::(&mut ctx) - .unwrap() - }; + let mut ctx = SESSION.create_execution_ctx(); + let result_canonical = encoded + .clone() + .into_array() + .execute::(&mut ctx) + .unwrap(); // Compare against the traditional array-based decompress path - let expected = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(result_canonical.into_array(), expected); + assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx); } #[rstest] @@ -615,26 +599,19 @@ mod tests { .collect(); let array = PrimitiveArray::from_option_iter(values); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); - let result_canonical = { - let mut ctx = SESSION.create_execution_ctx(); - encoded - .clone() - .into_array() - .execute::(&mut ctx) - .unwrap() - }; + let mut ctx = SESSION.create_execution_ctx(); + let result_canonical = encoded + .clone() + .into_array() + .execute::(&mut ctx) + .unwrap(); // Compare against the traditional array-based decompress path - let expected = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(result_canonical.into_array(), expected); + assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx); } #[rstest] @@ -655,27 +632,20 @@ mod tests { .collect(); let array = PrimitiveArray::from_option_iter(values); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().unwrap().array_len() > 0); - let result_canonical = { - let mut ctx = SESSION.create_execution_ctx(); - encoded - .clone() - .into_array() - .execute::(&mut ctx) - .unwrap() - }; + let mut ctx = SESSION.create_execution_ctx(); + let result_canonical = encoded + .clone() + .into_array() + .execute::(&mut ctx) + .unwrap(); // Compare against the traditional array-based decompress path - let expected = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(result_canonical.into_array(), expected); + assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx); } #[rstest] @@ -732,7 +702,7 @@ mod tests { #[case(1000, 200)] #[case(2048, 512)] fn test_sliced_to_primitive(#[case] size: usize, #[case] slice_start: usize) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Vec> = (0..size) .map(|i| { if i % 5 == 0 { @@ -796,7 +766,7 @@ mod tests { let normally_encoded = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); assert!( @@ -830,7 +800,7 @@ mod tests { // The legacy decompress_into_array path should work correctly. let result_legacy = decompress_into_array( alp_without_chunk_offsets.clone(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let legacy_slice = result_legacy.as_slice::(); diff --git a/encodings/alp/src/alp/compress.rs b/encodings/alp/src/alp/compress.rs index dff1c3dda2d..819fc4b3dc4 100644 --- a/encodings/alp/src/alp/compress.rs +++ b/encodings/alp/src/alp/compress.rs @@ -133,59 +133,64 @@ where mod tests { use core::f32; use core::f64; + use std::sync::LazyLock; use f64::consts::E; use f64::consts::PI; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::dtype::NativePType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; + use vortex_session::VortexSession; use super::*; use crate::alp::array::ALPArrayExt; use crate::alp::array::ALPArraySlotsExt; use crate::decompress_into_array; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_compress() { let array = PrimitiveArray::new(buffer![1.234f32; 1025], Validity::NonNullable); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_none()); let expected_encoded = PrimitiveArray::from_iter(vec![1234i32; 1025]); - assert_arrays_eq!(encoded.encoded(), expected_encoded); + assert_arrays_eq!( + encoded.encoded(), + expected_encoded, + &mut SESSION.create_execution_ctx() + ); assert_eq!(encoded.exponents(), Exponents { e: 9, f: 6 }); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } #[test] fn test_nullable_compress() { let array = PrimitiveArray::from_option_iter([None, Some(1.234f32), None]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_none()); let expected_encoded = PrimitiveArray::from_option_iter([None, Some(1234i32), None]); - assert_arrays_eq!(encoded.encoded(), expected_encoded); + assert_arrays_eq!( + encoded.encoded(), + expected_encoded, + &mut SESSION.create_execution_ctx() + ); assert_eq!(encoded.exponents(), Exponents { e: 9, f: 6 }); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); let expected = PrimitiveArray::from_option_iter(vec![None, Some(1.234f32), None]); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -193,21 +198,24 @@ mod tests { fn test_patched_compress() { let values = buffer![1.234f64, 2.718, PI, 4.0]; let array = PrimitiveArray::new(values.clone(), Validity::NonNullable); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_some()); let expected_encoded = PrimitiveArray::from_iter(vec![1234i64, 2718, 1234, 4000]); - assert_arrays_eq!(encoded.encoded(), expected_encoded); + assert_arrays_eq!( + encoded.encoded(), + expected_encoded, + &mut SESSION.create_execution_ctx() + ); assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 }); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); let expected_decoded = PrimitiveArray::new(values, Validity::NonNullable); - assert_arrays_eq!(decoded, expected_decoded); + assert_arrays_eq!( + decoded, + expected_decoded, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -215,21 +223,20 @@ mod tests { fn test_compress_ignores_invalid_exceptional_values() { let values = buffer![1.234f64, 2.718, PI, 4.0]; let array = PrimitiveArray::new(values, Validity::from_iter([true, true, false, true])); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_none()); let expected_encoded = PrimitiveArray::from_option_iter(buffer![Some(1234i64), Some(2718), None, Some(4000)]); - assert_arrays_eq!(encoded.encoded(), expected_encoded); + assert_arrays_eq!( + encoded.encoded(), + expected_encoded, + &mut SESSION.create_execution_ctx() + ); assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 }); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } #[test] @@ -242,20 +249,15 @@ mod tests { Some(4.0), None, ]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_some()); assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 }); - assert_arrays_eq!(encoded, array); + assert_arrays_eq!(encoded, array, &mut SESSION.create_execution_ctx()); - let _decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let _decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); } #[test] @@ -264,15 +266,15 @@ mod tests { let alp_arr = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); - assert_arrays_eq!(alp_arr, original); + assert_arrays_eq!(alp_arr, original, &mut SESSION.create_execution_ctx()); } #[test] fn roundtrips_all_null() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let original = PrimitiveArray::new(buffer![195.26274f64, PI, -48.815685], Validity::AllInvalid); let alp_arr = alp_encode(original.as_view(), None, &mut ctx).unwrap(); @@ -287,12 +289,12 @@ mod tests { decompressed.as_slice::() ); - assert_arrays_eq!(decompressed, original); + assert_arrays_eq!(decompressed, original, &mut ctx); } #[test] fn non_finite_numbers() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let original = PrimitiveArray::new( buffer![0.0f32, -0.0, f32::NAN, f32::NEG_INFINITY, f32::INFINITY], Validity::NonNullable, @@ -315,7 +317,7 @@ mod tests { #[test] fn test_chunk_offsets() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut values = vec![1.0f64; 3072]; values[1023] = PI; @@ -333,7 +335,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 3]); - assert_arrays_eq!(chunk_offsets, expected_offsets); + assert_arrays_eq!(chunk_offsets, expected_offsets, &mut ctx); let patch_indices = patches .indices() @@ -341,7 +343,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_indices = PrimitiveArray::from_iter(vec![1023u64, 1024, 1025]); - assert_arrays_eq!(patch_indices, expected_indices); + assert_arrays_eq!(patch_indices, expected_indices, &mut ctx); let patch_values = patches .values() @@ -349,12 +351,12 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_values = PrimitiveArray::from_iter(vec![PI, E, PI]); - assert_arrays_eq!(patch_values, expected_values); + assert_arrays_eq!(patch_values, expected_values, &mut ctx); } #[test] fn test_chunk_offsets_no_patches_in_middle() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut values = vec![1.0f64; 3072]; values[0] = PI; values[2048] = E; @@ -370,7 +372,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 1]); - assert_arrays_eq!(chunk_offsets, expected_offsets); + assert_arrays_eq!(chunk_offsets, expected_offsets, &mut ctx); let patch_indices = patches .indices() @@ -378,7 +380,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_indices = PrimitiveArray::from_iter(vec![0u64, 2048]); - assert_arrays_eq!(patch_indices, expected_indices); + assert_arrays_eq!(patch_indices, expected_indices, &mut ctx); let patch_values = patches .values() @@ -386,12 +388,12 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_values = PrimitiveArray::from_iter(vec![PI, E]); - assert_arrays_eq!(patch_values, expected_values); + assert_arrays_eq!(patch_values, expected_values, &mut ctx); } #[test] fn test_chunk_offsets_trailing_empty_chunks() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut values = vec![1.0f64; 3072]; values[0] = PI; @@ -406,7 +408,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 1]); - assert_arrays_eq!(chunk_offsets, expected_offsets); + assert_arrays_eq!(chunk_offsets, expected_offsets, &mut ctx); let patch_indices = patches .indices() @@ -414,7 +416,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_indices = PrimitiveArray::from_iter(vec![0u64]); - assert_arrays_eq!(patch_indices, expected_indices); + assert_arrays_eq!(patch_indices, expected_indices, &mut ctx); let patch_values = patches .values() @@ -422,12 +424,12 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_values = PrimitiveArray::from_iter(vec![PI]); - assert_arrays_eq!(patch_values, expected_values); + assert_arrays_eq!(patch_values, expected_values, &mut ctx); } #[test] fn test_chunk_offsets_single_chunk() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut values = vec![1.0f64; 512]; values[0] = PI; values[100] = E; @@ -443,7 +445,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_offsets = PrimitiveArray::from_iter(vec![0u64]); - assert_arrays_eq!(chunk_offsets, expected_offsets); + assert_arrays_eq!(chunk_offsets, expected_offsets, &mut ctx); let patch_indices = patches .indices() @@ -451,7 +453,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_indices = PrimitiveArray::from_iter(vec![0u64, 100]); - assert_arrays_eq!(patch_indices, expected_indices); + assert_arrays_eq!(patch_indices, expected_indices, &mut ctx); let patch_values = patches .values() @@ -459,7 +461,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_values = PrimitiveArray::from_iter(vec![PI, E]); - assert_arrays_eq!(patch_values, expected_values); + assert_arrays_eq!(patch_values, expected_values, &mut ctx); } #[test] @@ -470,14 +472,18 @@ mod tests { let encoded = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let sliced_alp = encoded.slice(512..1024).unwrap(); let expected_slice = original.slice(512..1024).unwrap(); - assert_arrays_eq!(sliced_alp, expected_slice); + assert_arrays_eq!( + sliced_alp, + expected_slice, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -487,14 +493,18 @@ mod tests { let encoded = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let sliced_alp = encoded.slice(512..1024).unwrap(); let expected_slice = original.slice(512..1024).unwrap(); - assert_arrays_eq!(sliced_alp, expected_slice); + assert_arrays_eq!( + sliced_alp, + expected_slice, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -508,14 +518,18 @@ mod tests { let encoded = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let sliced_alp = encoded.slice(512..1024).unwrap(); let expected_slice = original.slice(512..1024).unwrap(); - assert_arrays_eq!(sliced_alp, expected_slice); + assert_arrays_eq!( + sliced_alp, + expected_slice, + &mut SESSION.create_execution_ctx() + ); assert!(encoded.patches().is_some()); } @@ -533,20 +547,24 @@ mod tests { let encoded = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let sliced_alp = encoded.slice(1023..1025).unwrap(); let expected_slice = original.slice(1023..1025).unwrap(); - assert_arrays_eq!(sliced_alp, expected_slice); + assert_arrays_eq!( + sliced_alp, + expected_slice, + &mut SESSION.create_execution_ctx() + ); assert!(encoded.patches().is_some()); } #[test] fn test_slice_half_chunk_nullable_roundtrip() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = (0..1024) .map(|i| if i % 3 == 0 { None } else { Some(2.5f32) }) .collect::>(); @@ -558,41 +576,31 @@ mod tests { let decoded = sliced_alp.execute::(&mut ctx).unwrap(); let expected_slice = original.slice(512..1024).unwrap(); - assert_arrays_eq!(decoded, expected_slice); + assert_arrays_eq!(decoded, expected_slice, &mut ctx); } #[test] fn test_large_f32_array_uniform_values() { let size = 10_000; let array = PrimitiveArray::new(buffer![42.125f32; size], Validity::NonNullable); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_none()); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } #[test] fn test_large_f64_array_uniform_values() { let size = 50_000; let array = PrimitiveArray::new(buffer![123.456789f64; size], Validity::NonNullable); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_none()); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } #[test] @@ -605,17 +613,12 @@ mod tests { values[4500] = f32::INFINITY; let array = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_some()); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } #[test] @@ -632,16 +635,11 @@ mod tests { values[7000] = 999.999999999; let array = PrimitiveArray::new(Buffer::from(values.clone()), Validity::NonNullable); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_some()); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); for idx in 0..size { let decoded_val = decoded.as_slice::()[idx]; @@ -667,16 +665,11 @@ mod tests { .collect(); let array = PrimitiveArray::from_option_iter(values); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } #[test] @@ -693,16 +686,11 @@ mod tests { let validity = Validity::from_iter((0..size).map(|i| !matches!(i, 500 | 2500))); let array = PrimitiveArray::new(Buffer::from(values), validity); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } /// Regression test for patch_chunk index-out-of-bounds when slicing a multi-chunk @@ -732,7 +720,7 @@ mod tests { let encoded = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); assert!(encoded.patches().is_some()); @@ -742,11 +730,11 @@ mod tests { // but patches_indices only includes patches up to element 2500. let sliced_alp = encoded.slice(0..2500).unwrap(); let expected = original.slice(0..2500).unwrap(); - assert_arrays_eq!(sliced_alp, expected); + assert_arrays_eq!(sliced_alp, expected, &mut SESSION.create_execution_ctx()); // Also test slicing that starts mid-chunk (both start and end mid-chunk). let sliced_alp = encoded.slice(500..2500).unwrap(); let expected = original.slice(500..2500).unwrap(); - assert_arrays_eq!(sliced_alp, expected); + assert_arrays_eq!(sliced_alp, expected, &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/alp/src/alp/compute/between.rs b/encodings/alp/src/alp/compute/between.rs index c1390288e84..dbb6fec9a22 100644 --- a/encodings/alp/src/alp/compute/between.rs +++ b/encodings/alp/src/alp/compute/between.rs @@ -135,8 +135,9 @@ fn encode_upper_bound( #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; @@ -145,12 +146,19 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::between::BetweenOptions; use vortex_array::scalar_fn::fns::between::StrictComparison; + use vortex_session::VortexSession; use crate::ALPArray; use crate::alp::array::ALPArrayExt; use crate::alp::compute::between::between_impl; use crate::alp_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn assert_between( arr: &ALPArray, lower: f32, @@ -163,7 +171,8 @@ mod tests { assert_arrays_eq!( res, ConstantArray::new(Scalar::bool(expected, res.dtype().nullability()), arr.len()) - .into_array() + .into_array(), + &mut SESSION.create_execution_ctx() ); } @@ -171,12 +180,8 @@ mod tests { fn comparison_range() { let value = 0.0605_f32; let array = PrimitiveArray::from_iter([value; 1]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_none()); assert_between( @@ -237,7 +242,7 @@ mod tests { #[test] fn non_finite_bounds_use_total_order() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([1.234f32; 10]); let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap(); assert!(encoded.patches().is_none()); diff --git a/encodings/alp/src/alp/compute/cast.rs b/encodings/alp/src/alp/compute/cast.rs index 71b51ecf9af..d77a7257544 100644 --- a/encodings/alp/src/alp/compute/cast.rs +++ b/encodings/alp/src/alp/compute/cast.rs @@ -45,9 +45,10 @@ impl CastReduce for ALP { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -59,13 +60,20 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::alp::array::ALPArrayExt; use crate::alp_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn issue_5766_test_cast_alp_with_patches_to_nullable() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = buffer![1.234f32, f32::NAN, 2.345, f32::INFINITY, 3.456].into_array(); let values_primitive = values.clone().execute::(&mut ctx)?; let alp = alp_encode(values_primitive.as_view(), None, &mut ctx)?; @@ -81,14 +89,14 @@ mod tests { let expected = values.cast(nullable_dtype)?; let casted_prim = casted.execute::(&mut ctx)?; - assert_arrays_eq!(casted_prim, expected); + assert_arrays_eq!(casted_prim, expected, &mut ctx); Ok(()) } #[test] fn test_cast_alp_f32_to_f64() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = buffer![1.5f32, 2.5, 3.5, 4.5].into_array(); let values_primitive = values.execute::(&mut ctx)?; let alp = alp_encode(values_primitive.as_view(), None, &mut ctx)?; @@ -112,7 +120,7 @@ mod tests { #[test] fn test_cast_alp_to_int() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = buffer![1.0f32, 2.0, 3.0, 4.0].into_array(); let values_primitive = values.execute::(&mut ctx)?; let alp = alp_encode(values_primitive.as_view(), None, &mut ctx)?; @@ -126,7 +134,11 @@ mod tests { ); let decoded = casted.execute::(&mut ctx)?; - assert_arrays_eq!(decoded, PrimitiveArray::from_iter([1i32, 2, 3, 4])); + assert_arrays_eq!( + decoded, + PrimitiveArray::from_iter([1i32, 2, 3, 4]), + &mut ctx + ); Ok(()) } @@ -138,7 +150,7 @@ mod tests { #[case(buffer![42.42f64].into_array())] #[case(buffer![0.0f32, -1.5, 2.5, -3.5, 4.5].into_array())] fn test_cast_alp_conformance(#[case] array: vortex_array::ArrayRef) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array_primitive = array.execute::(&mut ctx)?; let alp = alp_encode(array_primitive.as_view(), None, &mut ctx).vortex_expect("cannot fail"); diff --git a/encodings/alp/src/alp/compute/compare.rs b/encodings/alp/src/alp/compute/compare.rs index 239e0558692..83e38e19bda 100644 --- a/encodings/alp/src/alp/compute/compare.rs +++ b/encodings/alp/src/alp/compute/compare.rs @@ -151,10 +151,10 @@ where #[cfg(test)] mod tests { use std::f32; + use std::sync::LazyLock; use rstest::rstest; use vortex_array::ArrayRef; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; @@ -167,10 +167,17 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_session::VortexSession; use super::*; use crate::alp_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn test_alp_compare>( alp: ArrayView, value: F, @@ -185,7 +192,7 @@ mod tests { #[test] fn basic_comparison_test() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([1.234f32; 1025]); let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap(); assert!(encoded.patches().is_none()); @@ -200,18 +207,18 @@ mod tests { .unwrap() .unwrap(); let expected = BoolArray::from_iter([false; 1025]); - assert_arrays_eq!(r, expected); + assert_arrays_eq!(r, expected, &mut ctx); let r = alp_scalar_compare(encoded.as_view(), 1.234f32, CompareOperator::Eq) .unwrap() .unwrap(); let expected = BoolArray::from_iter([true; 1025]); - assert_arrays_eq!(r, expected); + assert_arrays_eq!(r, expected, &mut ctx); } #[test] fn comparison_with_unencodable_value() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([1.234f32; 1025]); let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap(); assert!(encoded.patches().is_none()); @@ -226,18 +233,18 @@ mod tests { .unwrap() .unwrap(); let expected = BoolArray::from_iter([false; 1025]); - assert_arrays_eq!(r_eq, expected); + assert_arrays_eq!(r_eq, expected, &mut ctx); let r_neq = alp_scalar_compare(encoded.as_view(), 1.234444f32, CompareOperator::NotEq) .unwrap() .unwrap(); let expected = BoolArray::from_iter([true; 1025]); - assert_arrays_eq!(r_neq, expected); + assert_arrays_eq!(r_neq, expected, &mut ctx); } #[test] fn comparison_range() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([0.0605_f32; 10]); let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap(); assert!(encoded.patches().is_none()); @@ -253,33 +260,33 @@ mod tests { .unwrap() .unwrap(); let expected = BoolArray::from_iter([false; 10]); - assert_arrays_eq!(r_gte, expected); + assert_arrays_eq!(r_gte, expected, &mut ctx); // (0.0605_f32 > 0.06051_f32); let r_gt = alp_scalar_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Gt) .unwrap() .unwrap(); let expected = BoolArray::from_iter([false; 10]); - assert_arrays_eq!(r_gt, expected); + assert_arrays_eq!(r_gt, expected, &mut ctx); // 0.0605_f32 <= 0.06051_f32; let r_lte = alp_scalar_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Lte) .unwrap() .unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_lte, expected); + assert_arrays_eq!(r_lte, expected, &mut ctx); - //0.0605_f32 < 0.06051_f32; + // 0.0605_f32 < 0.06051_f32; let r_lt = alp_scalar_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Lt) .unwrap() .unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_lt, expected); + assert_arrays_eq!(r_lt, expected, &mut ctx); } #[test] fn comparison_zeroes() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([0.0_f32; 10]); let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap(); assert!(encoded.patches().is_none()); @@ -293,43 +300,39 @@ mod tests { let r_gte = test_alp_compare(encoded.as_view(), -0.00000001_f32, CompareOperator::Gte).unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_gte, expected); + assert_arrays_eq!(r_gte, expected, &mut ctx); let r_gte = test_alp_compare(encoded.as_view(), -0.0_f32, CompareOperator::Gte).unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_gte, expected); + assert_arrays_eq!(r_gte, expected, &mut ctx); let r_gt = test_alp_compare(encoded.as_view(), -0.0000000001f32, CompareOperator::Gt).unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_gt, expected); + assert_arrays_eq!(r_gt, expected, &mut ctx); let r_gte = test_alp_compare(encoded.as_view(), -0.0_f32, CompareOperator::Gt).unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_gte, expected); + assert_arrays_eq!(r_gte, expected, &mut ctx); let r_lte = test_alp_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Lte).unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_lte, expected); + assert_arrays_eq!(r_lte, expected, &mut ctx); let r_lt = test_alp_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Lt).unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_lt, expected); + assert_arrays_eq!(r_lt, expected, &mut ctx); let r_lt = test_alp_compare(encoded.as_view(), -0.00001_f32, CompareOperator::Lt).unwrap(); let expected = BoolArray::from_iter([false; 10]); - assert_arrays_eq!(r_lt, expected); + assert_arrays_eq!(r_lt, expected, &mut ctx); } #[test] fn compare_with_patches() { let array = PrimitiveArray::from_iter([1.234f32, 1.5, 19.0, f32::consts::E, 1_000_000.9]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_some()); // Not supported! @@ -343,12 +346,8 @@ mod tests { #[test] fn compare_to_null() { let array = PrimitiveArray::from_iter([1.234f32; 10]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); let other = ConstantArray::new( Scalar::null(DType::Primitive(PType::F32, Nullability::Nullable)), @@ -361,7 +360,7 @@ mod tests { .unwrap(); // Comparing to null yields null results let expected = BoolArray::from_iter([None::; 10]); - assert_arrays_eq!(r, expected); + assert_arrays_eq!(r, expected, &mut SESSION.create_execution_ctx()); } #[rstest] @@ -371,16 +370,12 @@ mod tests { #[case(f32::NEG_INFINITY, true)] fn compare_to_non_finite_gt(#[case] value: f32, #[case] result: bool) { let array = PrimitiveArray::from_iter([1.234f32; 10]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); let r = test_alp_compare(encoded.as_view(), value, CompareOperator::Gt).unwrap(); let expected = BoolArray::from_iter([result; 10]); - assert_arrays_eq!(r, expected); + assert_arrays_eq!(r, expected, &mut SESSION.create_execution_ctx()); } #[rstest] @@ -390,15 +385,11 @@ mod tests { #[case(f32::NEG_INFINITY, false)] fn compare_to_non_finite_lt(#[case] value: f32, #[case] result: bool) { let array = PrimitiveArray::from_iter([1.234f32; 10]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); let r = test_alp_compare(encoded.as_view(), value, CompareOperator::Lt).unwrap(); let expected = BoolArray::from_iter([result; 10]); - assert_arrays_eq!(r, expected); + assert_arrays_eq!(r, expected, &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/alp/src/alp/compute/filter.rs b/encodings/alp/src/alp/compute/filter.rs index 97b5b4af49e..f38a87f19b0 100644 --- a/encodings/alp/src/alp/compute/filter.rs +++ b/encodings/alp/src/alp/compute/filter.rs @@ -44,8 +44,8 @@ mod test { use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::filter::test_filter_conformance; use vortex_buffer::buffer; @@ -62,7 +62,7 @@ mod test { 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0 ].into_array())] fn test_filter_alp_conformance(#[case] array: ArrayRef) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array_primitive = array.execute::(&mut ctx).unwrap(); let alp = alp_encode(array_primitive.as_view(), None, &mut ctx).unwrap(); test_filter_conformance(&alp.into_array()); diff --git a/encodings/alp/src/alp/compute/mask.rs b/encodings/alp/src/alp/compute/mask.rs index 31f131f6e0c..2a528987143 100644 --- a/encodings/alp/src/alp/compute/mask.rs +++ b/encodings/alp/src/alp/compute/mask.rs @@ -36,18 +36,13 @@ impl MaskKernel for ALP { ) -> VortexResult> { let vortex_mask = Validity::Array(mask.not()?).execute_mask(array.len(), ctx)?; let masked_encoded = array.encoded().clone().mask(mask.clone())?; - let masked_dtype = array - .dtype() - .with_nullability(masked_encoded.dtype().nullability()); let masked_patches = array .patches() .map(|p| p.mask(&vortex_mask, ctx)) .transpose()? - .flatten() - .map(|patches| patches.cast_values(&masked_dtype)) - .transpose()?; + .flatten(); Ok(Some( - ALP::new(masked_encoded, array.exponents(), masked_patches).into_array(), + ALP::try_new(masked_encoded, array.exponents(), masked_patches)?.into_array(), )) } } @@ -56,8 +51,8 @@ impl MaskKernel for ALP { mod test { use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::mask::test_mask_conformance; @@ -78,7 +73,7 @@ mod test { 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0 ].into_array())] fn test_mask_alp_conformance(#[case] array: vortex_array::ArrayRef) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array_primitive = array.execute::(&mut ctx).unwrap(); let alp = alp_encode(array_primitive.as_view(), None, &mut ctx).unwrap(); test_mask_conformance(&alp.into_array()); @@ -87,7 +82,7 @@ mod test { #[test] fn test_mask_alp_with_patches() { use std::f64::consts::PI; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // PI doesn't encode cleanly with ALP, so it creates patches. let values: Vec = (0..100) .map(|i| if i % 4 == 3 { PI } else { 1.0 }) @@ -100,7 +95,7 @@ mod test { #[test] fn test_mask_alp_with_patches_casts_surviving_patch_values_to_nullable() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::from_iter([1.234f32, f32::NAN, 2.345, f32::INFINITY, 3.456]); let alp = alp_encode(values.as_view(), None, &mut ctx).unwrap(); assert!(alp.patches().is_some(), "expected patches"); diff --git a/encodings/alp/src/alp/compute/mod.rs b/encodings/alp/src/alp/compute/mod.rs index e5e5b657740..83a6c667893 100644 --- a/encodings/alp/src/alp/compute/mod.rs +++ b/encodings/alp/src/alp/compute/mod.rs @@ -12,41 +12,50 @@ mod take; #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; 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_session::VortexSession; use crate::ALPArray; use crate::alp_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + crate::initialize(&session); + session + }); + #[rstest] // Basic float arrays - #[case::f32_array(alp_encode(PrimitiveArray::from_iter([1.23f32, 4.56, 7.89, 10.11, 12.13]).as_view(), None, &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] - #[case::f64_array(alp_encode(PrimitiveArray::from_iter([100.1f64, 200.2, 300.3, 400.4, 500.5]).as_view(), None, &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] + #[case::f32_array(alp_encode(PrimitiveArray::from_iter([1.23f32, 4.56, 7.89, 10.11, 12.13]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] + #[case::f64_array(alp_encode(PrimitiveArray::from_iter([100.1f64, 200.2, 300.3, 400.4, 500.5]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] // Nullable arrays - #[case::nullable_f32(alp_encode(PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]).as_view(), None, &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] - #[case::nullable_f64(alp_encode(PrimitiveArray::from_option_iter([Some(1.1f64), None, Some(2.2), Some(3.3), None]).as_view(), None, &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] + #[case::nullable_f32(alp_encode(PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] + #[case::nullable_f64(alp_encode(PrimitiveArray::from_option_iter([Some(1.1f64), None, Some(2.2), Some(3.3), None]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] // Edge cases - #[case::single_element(alp_encode(PrimitiveArray::from_iter([42.42f64]).as_view(), None, &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] + #[case::single_element(alp_encode(PrimitiveArray::from_iter([42.42f64]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] // Large arrays - #[case::large_f32(alp_encode(PrimitiveArray::from_iter((0..1000).map(|i| i as f32 * 0.1)).as_view(), None, &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] + #[case::large_f32(alp_encode(PrimitiveArray::from_iter((0..1000).map(|i| i as f32 * 0.1)).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] // Arrays with patterns - #[case::repeating_pattern(alp_encode(PrimitiveArray::from_iter([1.1f32, 2.2, 3.3, 1.1, 2.2, 3.3, 1.1, 2.2, 3.3]).as_view(), None, &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] - #[case::close_values(alp_encode(PrimitiveArray::from_iter([100.001f64, 100.002, 100.003, 100.004, 100.005]).as_view(), None, &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] + #[case::repeating_pattern(alp_encode(PrimitiveArray::from_iter([1.1f32, 2.2, 3.3, 1.1, 2.2, 3.3, 1.1, 2.2, 3.3]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] + #[case::close_values(alp_encode(PrimitiveArray::from_iter([100.001f64, 100.002, 100.003, 100.004, 100.005]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] fn test_alp_consistency(#[case] array: ALPArray) { - test_array_consistency(&array.into_array()); + test_array_consistency(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[rstest] - #[case::f32_basic(alp_encode(PrimitiveArray::from_iter([1.23f32, 4.56, 7.89, 10.11, 12.13]).as_view(), None, &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] - #[case::f64_basic(alp_encode(PrimitiveArray::from_iter([100.1f64, 200.2, 300.3, 400.4, 500.5]).as_view(), None, &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] - #[case::f32_large(alp_encode(PrimitiveArray::from_iter((0..100).map(|i| i as f32 * 1.5)).as_view(), None, &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] - #[case::f64_large(alp_encode(PrimitiveArray::from_iter((0..100).map(|i| i as f64 * 2.5)).as_view(), None, &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] + #[case::f32_basic(alp_encode(PrimitiveArray::from_iter([1.23f32, 4.56, 7.89, 10.11, 12.13]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] + #[case::f64_basic(alp_encode(PrimitiveArray::from_iter([100.1f64, 200.2, 300.3, 400.4, 500.5]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] + #[case::f32_large(alp_encode(PrimitiveArray::from_iter((0..100).map(|i| i as f32 * 1.5)).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] + #[case::f64_large(alp_encode(PrimitiveArray::from_iter((0..100).map(|i| i as f64 * 2.5)).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] fn test_alp_binary_numeric(#[case] array: ALPArray) { - test_binary_numeric_array(array.into_array()); + test_binary_numeric_array(&array.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/alp/src/alp/compute/take.rs b/encodings/alp/src/alp/compute/take.rs index 3b2b35ceaa9..2e071eaa2aa 100644 --- a/encodings/alp/src/alp/compute/take.rs +++ b/encodings/alp/src/alp/compute/take.rs @@ -23,17 +23,9 @@ impl TakeExecute for ALP { .patches() .map(|p| p.take(indices, ctx)) .transpose()? - .flatten() - .map(|patches| { - patches.cast_values( - &array - .dtype() - .with_nullability(taken_encoded.dtype().nullability()), - ) - }) - .transpose()?; + .flatten(); Ok(Some( - ALP::new(taken_encoded, array.exponents(), taken_patches).into_array(), + ALP::try_new(taken_encoded, array.exponents(), taken_patches)?.into_array(), )) } } @@ -42,8 +34,8 @@ impl TakeExecute for ALP { mod test { use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::take::test_take_conformance; use vortex_buffer::buffer; @@ -56,7 +48,7 @@ mod test { #[case(PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]).into_array())] #[case(buffer![42.42f64].into_array())] fn test_take_alp_conformance(#[case] array: vortex_array::ArrayRef) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array_primitive = array.execute::(&mut ctx).unwrap(); let alp = alp_encode(array_primitive.as_view(), None, &mut ctx).unwrap(); test_take_conformance(&alp.into_array()); diff --git a/encodings/alp/src/alp/mod.rs b/encodings/alp/src/alp/mod.rs index 88375fe2ae6..9992a9ddbe8 100644 --- a/encodings/alp/src/alp/mod.rs +++ b/encodings/alp/src/alp/mod.rs @@ -30,8 +30,49 @@ 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() { @@ -61,9 +102,14 @@ 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, @@ -83,6 +129,13 @@ mod private { 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; @@ -115,20 +168,17 @@ pub trait ALPFloat: private::Sealed + Float + Display + NativePType { .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 (_, encoded, _, exc_patches, _) = Self::encode( - sample.as_deref().unwrap_or(values), - Some(Exponents { e, f }), - ); - - let size = Self::estimate_encoded_size(&encoded, &exc_patches); + let exp = Exponents { e, f }; + let size = Self::estimate_encoded_size_for_exponents(sample, exp); if size < best_nbytes { best_nbytes = size; - best_exp = Exponents { e, f }; + best_exp = exp; } else if size == best_nbytes && e - f < best_exp.e - best_exp.f { - best_exp = Exponents { e, f }; + best_exp = exp; } } } @@ -136,6 +186,51 @@ pub trait ALPFloat: private::Sealed + Float + Display + NativePType { 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 diff --git a/encodings/alp/src/alp/plugin.rs b/encodings/alp/src/alp/plugin.rs index 25266bef821..c14133109d1 100644 --- a/encodings/alp/src/alp/plugin.rs +++ b/encodings/alp/src/alp/plugin.rs @@ -98,7 +98,6 @@ mod tests { use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::patched::PatchedArraySlotsExt; use vortex_array::buffer::BufferHandle; - use vortex_array::session::ArraySession; use vortex_array::session::ArraySessionExt; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -111,7 +110,7 @@ mod tests { use crate::alp_encode; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); session.arrays().register(ALPPatchedPlugin); session }); diff --git a/encodings/alp/src/alp/rules.rs b/encodings/alp/src/alp/rules.rs index fe13682ccd8..3b888ae8269 100644 --- a/encodings/alp/src/alp/rules.rs +++ b/encodings/alp/src/alp/rules.rs @@ -1,26 +1,35 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; +use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::optimizer::rules::ParentRuleSet; +use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::between::BetweenReduceAdaptor; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; +use vortex_array::scalar_fn::fns::mask::Mask; use vortex_array::scalar_fn::fns::mask::MaskExecuteAdaptor; use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; +use vortex_session::VortexSession; use crate::ALP; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(ALP)), - ParentKernelSet::lift(&FilterExecuteAdaptor(ALP)), - ParentKernelSet::lift(&MaskExecuteAdaptor(ALP)), - ParentKernelSet::lift(&SliceExecuteAdaptor(ALP)), - ParentKernelSet::lift(&TakeExecuteAdaptor(ALP)), -]); +pub(super) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), ALP, CompareExecuteAdaptor(ALP)); + kernels.register_execute_parent_kernel(Filter.id(), ALP, FilterExecuteAdaptor(ALP)); + kernels.register_execute_parent_kernel(Mask.id(), ALP, MaskExecuteAdaptor(ALP)); + kernels.register_execute_parent_kernel(Slice.id(), ALP, SliceExecuteAdaptor(ALP)); + kernels.register_execute_parent_kernel(Dict.id(), ALP, TakeExecuteAdaptor(ALP)); +} pub(super) const RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&BetweenReduceAdaptor(ALP)), diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index eaf98851665..28ca30f909c 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -17,14 +17,11 @@ 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::LEGACY_SESSION; use vortex_array::TypedArrayRef; -use vortex_array::VortexSessionExecute; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; @@ -53,7 +50,6 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::alp_rd::kernel::PARENT_KERNELS; use crate::alp_rd::rules::RULES; use crate::alp_rd_decode; @@ -130,6 +126,14 @@ impl VTable for ALPRD { 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, @@ -210,13 +214,6 @@ impl VTable for ALPRD { ) }) .transpose()?; - // NOTE: `VTable::deserialize` has a fixed trait signature without `ExecutionCtx`, so we - // cannot plumb a ctx in here. We construct a legacy ctx locally at this trait boundary. - let left_parts_patches = ALPRDData::canonicalize_patches( - &left_parts, - left_parts_patches, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; let slots = ALPRDData::make_slots(&left_parts, &right_parts, left_parts_patches.as_ref()); let data = ALPRDData::new( left_parts_dictionary, @@ -307,15 +304,6 @@ impl VTable for ALPRD { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } /// The left (most significant) parts of the real-double encoded values. @@ -378,11 +366,8 @@ impl ALPRD { right_parts: ArrayRef, right_bit_width: u8, left_parts_patches: Option, - ctx: &mut ExecutionCtx, ) -> VortexResult { let len = left_parts.len(); - let left_parts_patches = - ALPRDData::canonicalize_patches(&left_parts, left_parts_patches, ctx)?; let slots = ALPRDData::make_slots(&left_parts, &right_parts, left_parts_patches.as_ref()); let data = ALPRDData::new(left_parts_dictionary, right_bit_width, left_parts_patches); Array::try_from_parts(ArrayParts::new(ALPRD, dtype, len, data).with_slots(slots)) @@ -410,28 +395,6 @@ impl ALPRD { } impl ALPRDData { - fn canonicalize_patches( - left_parts: &ArrayRef, - left_parts_patches: Option, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - left_parts_patches - .map(|patches| { - if !patches.values().all_valid(ctx)? { - vortex_bail!("patches must be all valid: {}", patches.values()); - } - // TODO(ngates): assert the DType, don't cast it. - // TODO(joe): assert the DType, don't cast it in the next PR. - let mut patches = patches.cast_values(&left_parts.dtype().as_nonnullable())?; - // Force execution of the lazy cast so patch values are materialized - // before serialization. - let canonical = patches.values().clone().execute::(ctx)?; - *patches.values_mut() = canonical.into_array(); - Ok(patches) - }) - .transpose() - } - /// Build a new `ALPRDArray` from components. pub fn new( left_parts_dictionary: Buffer, @@ -558,18 +521,16 @@ fn validate_parts( "patches array_len {} != outer len {len}", patches.array_len(), ); + // Left-parts exceptions are always all-valid and are stored as the non-nullable left-parts + // dtype. Requiring that exact dtype (rather than ignoring nullability) means each + // construction path must produce correct patches, removing the need to normalize them. + // Non-nullable also implies all-valid, so no separate validity check is required. + let expected = left_parts.dtype().as_nonnullable(); vortex_ensure!( - patches.dtype().eq_ignore_nullability(left_parts.dtype()), - "patches dtype {} does not match left_parts dtype {}", + patches.dtype() == &expected, + "patches dtype {} must be the non-nullable left_parts dtype {}", patches.dtype(), - left_parts.dtype(), - ); - vortex_ensure!( - patches - .values() - .all_valid(&mut LEGACY_SESSION.create_execution_ctx())?, - "patches must be all valid: {}", - patches.values() + expected, ); } @@ -631,20 +592,28 @@ impl ValidityChild for ALPRD { #[cfg(test)] mod test { + use std::sync::LazyLock; + use prost::Message; use rstest::rstest; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::PType; use vortex_array::patches::PatchesMetadata; use vortex_array::test_harness::check_metadata; + use vortex_session::VortexSession; use super::ALPRDMetadata; use crate::ALPRDFloat; use crate::alp_rd; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[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)] @@ -652,7 +621,7 @@ mod test { #[case] reals: Vec, #[case] seed: T, ) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_eq!(reals.len(), 1024, "test expects 1024-length fixture"); // Null out some of the values. let mut reals: Vec> = reals.into_iter().map(Some).collect(); @@ -666,7 +635,7 @@ mod test { // Pick a seed that we know will trigger lots of patches. let encoder: alp_rd::RDEncoder = alp_rd::RDEncoder::new(&[seed.powi(-2)]); - let rd_array = encoder.encode(real_array.as_view(), &mut ctx); + let rd_array = encoder.encode(real_array.as_view()); let decoded = rd_array .as_array() @@ -674,7 +643,7 @@ mod test { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(decoded, PrimitiveArray::from_option_iter(reals)); + assert_arrays_eq!(decoded, PrimitiveArray::from_option_iter(reals), &mut ctx); } #[cfg_attr(miri, ignore)] diff --git a/encodings/alp/src/alp_rd/compute/cast.rs b/encodings/alp/src/alp_rd/compute/cast.rs index c110475b880..68ed39dd390 100644 --- a/encodings/alp/src/alp_rd/compute/cast.rs +++ b/encodings/alp/src/alp_rd/compute/cast.rs @@ -50,8 +50,8 @@ impl CastReduce for ALPRD { mod tests { use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::compute::conformance::cast::test_cast_conformance; @@ -63,11 +63,11 @@ mod tests { #[test] fn test_cast_alprd_f32_to_f64() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = vec![1.0f32, 1.1, 1.2, 1.3, 1.4]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - let alprd = encoder.encode(arr.as_view(), &mut ctx); + let alprd = encoder.encode(arr.as_view()); let casted = alprd .into_array() @@ -87,12 +87,12 @@ mod tests { #[test] fn test_cast_alprd_nullable() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([Some(10.0f64), None, Some(10.1), Some(10.2), None]); let values = vec![10.0f64, 10.1, 10.2]; let encoder = RDEncoder::new(&values); - let alprd = encoder.encode(arr.as_view(), &mut ctx); + let alprd = encoder.encode(arr.as_view()); // Cast to NonNullable should fail since we have nulls. The failure surfaces during // execution since the reduce path defers when the validity stat is not cached. @@ -122,31 +122,31 @@ mod tests { let values = vec![1.23f32, 4.56, 7.89, 10.11, 12.13]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] #[case::f64({ let values = vec![100.1f64, 200.2, 300.3, 400.4, 500.5]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] #[case::single({ let values = vec![42.42f64]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] #[case::negative({ let values = vec![0.0f32, -1.5, 2.5, -3.5, 4.5]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] #[case::nullable({ let arr = PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]); let values = vec![1.1f32, 2.2, 3.3]; let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] fn test_cast_alprd_conformance(#[case] alprd: crate::alp_rd::ALPRDArray) { test_cast_conformance(&alprd.into_array()); diff --git a/encodings/alp/src/alp_rd/compute/filter.rs b/encodings/alp/src/alp_rd/compute/filter.rs index 5deaa7c466b..7aa8c247f2a 100644 --- a/encodings/alp/src/alp_rd/compute/filter.rs +++ b/encodings/alp/src/alp_rd/compute/filter.rs @@ -32,7 +32,6 @@ impl FilterKernel for ALPRD { array.right_parts().filter(mask.clone())?, array.right_bit_width(), left_parts_exceptions, - ctx, )? .into_array(), )) @@ -41,28 +40,37 @@ impl FilterKernel for ALPRD { #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::filter::test_filter_conformance; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::ALPRDArrayExt; use crate::ALPRDFloat; use crate::RDEncoder; + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + crate::initialize(&session); + session + }); + #[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) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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(), &mut ctx); + let encoded = RDEncoder::new(&[a, b]).encode(array.as_view()); // Make sure that we're testing the exception pathway. assert!(encoded.left_parts_patches().is_some()); @@ -71,20 +79,16 @@ mod test { let filtered = encoded .filter(Mask::from_iter([true, false, true])) .unwrap(); - assert_arrays_eq!(filtered, PrimitiveArray::from_iter([a, outlier])); + assert_arrays_eq!(filtered, PrimitiveArray::from_iter([a, outlier]), &mut ctx); } #[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) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); test_filter_conformance( &RDEncoder::new(&[a, b]) - .encode( - PrimitiveArray::from_iter([a, b, outlier, b, outlier]).as_view(), - &mut ctx, - ) + .encode(PrimitiveArray::from_iter([a, b, outlier, b, outlier]).as_view()) .into_array(), ); } @@ -93,13 +97,11 @@ mod test { #[case(0.1f32, 3e25f32)] #[case(0.5f64, 1e100f64)] fn test_filter_with_nulls(#[case] a: T, #[case] outlier: T) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); test_filter_conformance( &RDEncoder::new(&[a]) .encode( PrimitiveArray::from_option_iter([Some(a), None, Some(outlier), Some(a), None]) .as_view(), - &mut ctx, ) .into_array(), ); diff --git a/encodings/alp/src/alp_rd/compute/mask.rs b/encodings/alp/src/alp_rd/compute/mask.rs index 47a2f2a6b67..2773387c738 100644 --- a/encodings/alp/src/alp_rd/compute/mask.rs +++ b/encodings/alp/src/alp_rd/compute/mask.rs @@ -4,8 +4,6 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; -use vortex_array::VortexSessionExecute; use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::mask::Mask as MaskExpr; @@ -22,8 +20,6 @@ impl MaskReduce for ALPRD { EmptyOptions, [array.left_parts().clone(), mask.clone()], )?; - // NOTE: `MaskReduce::mask` has a fixed trait signature without `ExecutionCtx`, so we - // construct a legacy ctx locally at this trait boundary. Ok(Some( ALPRD::try_new( array.dtype().as_nullable(), @@ -32,7 +28,6 @@ impl MaskReduce for ALPRD { array.right_parts().clone(), array.right_bit_width(), array.left_parts_patches(), - &mut LEGACY_SESSION.create_execution_ctx(), )? .into_array(), )) @@ -43,8 +38,6 @@ impl MaskReduce for ALPRD { mod tests { use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; - use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::mask::test_mask_conformance; @@ -55,13 +48,9 @@ mod tests { #[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) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); test_mask_conformance( &RDEncoder::new(&[a, b]) - .encode( - PrimitiveArray::from_iter([a, b, outlier, b, outlier]).as_view(), - &mut ctx, - ) + .encode(PrimitiveArray::from_iter([a, b, outlier, b, outlier]).as_view()) .into_array(), ); } @@ -70,13 +59,11 @@ mod tests { #[case(0.1f32, 3e25f32)] #[case(0.5f64, 1e100f64)] fn test_mask_with_nulls(#[case] a: T, #[case] outlier: T) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); test_mask_conformance( &RDEncoder::new(&[a]) .encode( PrimitiveArray::from_option_iter([Some(a), None, Some(outlier), Some(a), None]) .as_view(), - &mut ctx, ) .into_array(), ); diff --git a/encodings/alp/src/alp_rd/compute/mod.rs b/encodings/alp/src/alp_rd/compute/mod.rs index 7d03eddaf40..696e622b873 100644 --- a/encodings/alp/src/alp_rd/compute/mod.rs +++ b/encodings/alp/src/alp_rd/compute/mod.rs @@ -8,67 +8,77 @@ mod take; #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; 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_session::VortexSession; use crate::ALPRDArray; use crate::RDEncoder; + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + crate::initialize(&session); + session + }); + #[rstest] // Basic float arrays with RD (reference delta) encoding #[case::f32_array({ let values = vec![1.0f32, 1.1, 1.2, 1.3, 1.4]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] #[case::f64_array({ let values = vec![100.0f64, 100.01, 100.02, 100.03, 100.04]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] // Nullable arrays #[case::nullable_f32({ let values = vec![1.0f32, 1.2, 1.3]; let arr = PrimitiveArray::from_option_iter([Some(1.0f32), None, Some(1.2), Some(1.3), None]); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] #[case::nullable_f64({ let values = vec![10.0f64, 10.2, 10.3]; let arr = PrimitiveArray::from_option_iter([Some(10.0f64), None, Some(10.2), Some(10.3), None]); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] // Edge cases #[case::single_element({ let values = vec![42.42f64]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] // Arrays with small deltas (good for RD encoding) #[case::small_deltas({ let values = vec![1000.0f32, 1000.001, 1000.002, 1000.003, 1000.004]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] // Large arrays #[case::large_f32({ let values: Vec = (0..1000).map(|i| 100.0 + i as f32 * 0.01).collect(); let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] fn test_alp_rd_consistency(#[case] array: ALPRDArray) { - test_array_consistency(&array.into_array()); + let ctx = &mut SESSION.create_execution_ctx(); + test_array_consistency(&array.into_array(), ctx); } #[rstest] @@ -76,27 +86,27 @@ mod tests { let values = vec![1.0f32, 1.1, 1.2, 1.3, 1.4]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] #[case::f64_basic({ let values = vec![100.0f64, 100.01, 100.02, 100.03, 100.04]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] #[case::f32_large({ let values: Vec = (0..100).map(|i| 50.0 + i as f32 * 0.1).collect(); let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] #[case::f64_large({ let values: Vec = (0..100).map(|i| 1000.0 + i as f64 * 0.01).collect(); let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut LEGACY_SESSION.create_execution_ctx()) + encoder.encode(arr.as_view()) })] fn test_alp_rd_binary_numeric(#[case] array: ALPRDArray) { - test_binary_numeric_array(array.into_array()); + test_binary_numeric_array(&array.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/alp/src/alp_rd/compute/take.rs b/encodings/alp/src/alp_rd/compute/take.rs index 0d980e0abfc..8d12bfbb316 100644 --- a/encodings/alp/src/alp_rd/compute/take.rs +++ b/encodings/alp/src/alp_rd/compute/take.rs @@ -20,19 +20,21 @@ impl TakeExecute for ALPRD { ctx: &mut ExecutionCtx, ) -> VortexResult> { let taken_left_parts = array.left_parts().take(indices.clone())?; + // With nullable take indices, `Patches::take` widens the (still all-valid) exception + // values to nullable, but ALPRD stores exceptions as the non-nullable left-parts dtype. + // Narrow them back so the patches satisfy the invariant `validate_parts` enforces. This is + // a no-op when the take indices are non-nullable and the values already match. + let exceptions_dtype = taken_left_parts.dtype().as_nonnullable(); let left_parts_exceptions = array .left_parts_patches() - .map(|patches| patches.take(indices, ctx)) - .transpose()? - .flatten() - .map(|p| { - let values_dtype = p - .values() - .dtype() - .with_nullability(taken_left_parts.dtype().nullability()); - p.cast_values(&values_dtype) + .map(|patches| { + patches + .take(indices, ctx)? + .map(|taken| taken.map_values(|values| values.cast(exceptions_dtype.clone()))) + .transpose() }) - .transpose()?; + .transpose()? + .flatten(); let right_parts = array .right_parts() .take(indices.clone())? @@ -48,7 +50,6 @@ impl TakeExecute for ALPRD { right_parts, array.right_bit_width(), left_parts_exceptions, - ctx, )? .into_array(), )) @@ -57,18 +58,27 @@ impl TakeExecute for ALPRD { #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::take::test_take_conformance; + use vortex_session::VortexSession; use crate::ALPRDArrayExt; use crate::ALPRDFloat; use crate::RDEncoder; + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + crate::initialize(&session); + session + }); + #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] @@ -76,9 +86,9 @@ mod test { use vortex_array::IntoArray as _; use vortex_buffer::buffer; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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(), &mut ctx); + let encoded = RDEncoder::new(&[a, b]).encode(array.as_view()); assert!(encoded.left_parts_patches().is_some()); assert!( @@ -95,16 +105,16 @@ mod test { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(taken, PrimitiveArray::from_iter([a, outlier])); + assert_arrays_eq!(taken, PrimitiveArray::from_iter([a, outlier]), &mut ctx); } #[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) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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(), &mut ctx); + let encoded = RDEncoder::new(&[a, b]).encode(array.as_view()); assert!(encoded.left_parts_patches().is_some()); assert!( @@ -123,7 +133,8 @@ mod test { assert_arrays_eq!( taken, - PrimitiveArray::from_option_iter([Some(a), Some(outlier), None]) + PrimitiveArray::from_option_iter([Some(a), Some(outlier), None]), + &mut ctx ); } @@ -131,13 +142,9 @@ mod test { #[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) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); test_take_conformance( &RDEncoder::new(&[a, b]) - .encode( - PrimitiveArray::from_iter([a, b, outlier, b, outlier]).as_view(), - &mut ctx, - ) + .encode(PrimitiveArray::from_iter([a, b, outlier, b, outlier]).as_view()) .into_array(), ); } @@ -146,13 +153,11 @@ mod test { #[case(0.1f32, 3e25f32)] #[case(0.5f64, 1e100f64)] fn test_take_with_nulls_conformance(#[case] a: T, #[case] outlier: T) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); test_take_conformance( &RDEncoder::new(&[a]) .encode( PrimitiveArray::from_option_iter([Some(a), None, Some(outlier), Some(a), None]) .as_view(), - &mut ctx, ) .into_array(), ); diff --git a/encodings/alp/src/alp_rd/kernel.rs b/encodings/alp/src/alp_rd/kernel.rs index 66721cbb1f9..73670163954 100644 --- a/encodings/alp/src/alp_rd/kernel.rs +++ b/encodings/alp/src/alp_rd/kernel.rs @@ -1,15 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; +use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_session::VortexSession; use crate::alp_rd::ALPRD; -pub(crate) static PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&SliceExecuteAdaptor(ALPRD)), - ParentKernelSet::lift(&FilterExecuteAdaptor(ALPRD)), - ParentKernelSet::lift(&TakeExecuteAdaptor(ALPRD)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Slice.id(), ALPRD, SliceExecuteAdaptor(ALPRD)); + kernels.register_execute_parent_kernel(Filter.id(), ALPRD, FilterExecuteAdaptor(ALPRD)); + kernels.register_execute_parent_kernel(Dict.id(), ALPRD, TakeExecuteAdaptor(ALPRD)); +} diff --git a/encodings/alp/src/alp_rd/mod.rs b/encodings/alp/src/alp_rd/mod.rs index 5f2e703cc0a..07afe6b8adf 100644 --- a/encodings/alp/src/alp_rd/mod.rs +++ b/encodings/alp/src/alp_rd/mod.rs @@ -36,6 +36,7 @@ use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; +use vortex_session::VortexSession; use vortex_utils::aliases::hash_map::HashMap; use crate::match_each_alp_float_ptype; @@ -55,6 +56,10 @@ const CUT_LIMIT: usize = 16; const MAX_DICT_SIZE: u8 = 8; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + mod private { pub trait Sealed {} @@ -182,15 +187,11 @@ impl RDEncoder { /// /// Each value will be split into a left and right component, which are compressed individually. // TODO(joe): make fallible - pub fn encode(&self, array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx) -> ALPRDArray { - match_each_alp_float_ptype!(array.ptype(), |P| { self.encode_generic::

(array, ctx) }) + 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>, - ctx: &mut ExecutionCtx, - ) -> ALPRDArray + fn encode_generic(&self, array: ArrayView<'_, Primitive>) -> ALPRDArray where T: ALPRDFloat + NativePType, T::UINT: NativePType, @@ -289,7 +290,6 @@ impl RDEncoder { packed_right, self.right_bit_width, exceptions, - ctx, ) .vortex_expect("ALPRDArray construction in encode") } diff --git a/encodings/alp/src/alp_rd/ops.rs b/encodings/alp/src/alp_rd/ops.rs index 0351d446424..0fc0dc40fa1 100644 --- a/encodings/alp/src/alp_rd/ops.rs +++ b/encodings/alp/src/alp_rd/ops.rs @@ -65,27 +65,35 @@ impl OperationsVTable for ALPRD { #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::rstest; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::scalar::Scalar; + use vortex_session::VortexSession; use crate::ALPRDArrayExt; use crate::ALPRDFloat; use crate::RDEncoder; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[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) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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(), &mut ctx); + let encoded = RDEncoder::new(&[a, b]).encode(array.as_view()); assert!(encoded.left_parts_patches().is_some()); - assert_arrays_eq!(encoded, array); + assert_arrays_eq!(encoded, array, &mut ctx); } #[rstest] @@ -96,25 +104,26 @@ mod test { #[case] b: T, #[case] outlier: T, ) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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(), &mut ctx); + let encoded = RDEncoder::new(&[a, b]).encode(array.as_view()); assert!(encoded.left_parts_patches().is_some()); - assert_arrays_eq!(encoded, array); + assert_arrays_eq!(encoded, array, &mut ctx); } #[test] fn nullable_scalar_at() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let a = 0.1f64; let b = 0.2f64; let outlier = 3e100f64; let array = PrimitiveArray::from_option_iter([Some(a), Some(b), Some(outlier)]); - let encoded = RDEncoder::new(&[a, b]).encode(array.as_view(), &mut ctx); + let encoded = RDEncoder::new(&[a, b]).encode(array.as_view()); assert!(encoded.left_parts_patches().is_some()); assert_arrays_eq!( encoded, - PrimitiveArray::from_option_iter([Some(a), Some(b), Some(outlier)]) + PrimitiveArray::from_option_iter([Some(a), Some(b), Some(outlier)]), + &mut ctx ); } } diff --git a/encodings/alp/src/lib.rs b/encodings/alp/src/lib.rs index 7da676d6c6c..1286f3e6e2b 100644 --- a/encodings/alp/src/lib.rs +++ b/encodings/alp/src/lib.rs @@ -39,6 +39,8 @@ pub fn initialize(session: &VortexSession) { session.arrays().register(ALP); } session.arrays().register(ALPRD); + alp::initialize(session); + alp_rd::initialize(session); // Register the ALP-specific NaN count aggregate kernel. session.aggregate_fns().register_aggregate_kernel( diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs index 722f48f7376..fe00bd5fe4d 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -39,8 +39,6 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::kernel::PARENT_KERNELS; - /// A [`ByteBool`]-encoded Vortex array. pub type ByteBoolArray = Array; @@ -96,6 +94,23 @@ impl VTable for ByteBool { } } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + let data = ByteBoolData::new(buffers[0].clone()); + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + fn serialize( _array: ArrayView<'_, Self>, _session: &VortexSession, @@ -157,15 +172,6 @@ impl VTable for ByteBool { BoolArray::new(boolean_buffer, validity).into_array(), )) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } /// The validity bitmap indicating which elements are non-null. @@ -315,21 +321,26 @@ impl OperationsVTable for ByteBool { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::ArrayContext; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::serde::SerializeOptions; use vortex_array::serde::SerializedArray; - use vortex_array::session::ArraySession; use vortex_array::session::ArraySessionExt; use vortex_buffer::ByteBufferMut; - use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use super::*; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_validity_construction() { let v = vec![true, false]; @@ -338,7 +349,7 @@ mod tests { let arr = ByteBool::from_vec(v, Validity::AllValid); assert_eq!(v_len, arr.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); for idx in 0..arr.len() { assert!(arr.is_valid(idx, &mut ctx).unwrap()); } @@ -367,7 +378,7 @@ mod tests { let array = ByteBool::from_option_vec(vec![Some(true), None, Some(false), None]); let dtype = array.dtype().clone(); let len = array.len(); - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); session.arrays().register(ByteBool); let ctx = ArrayContext::empty(); @@ -387,6 +398,6 @@ mod tests { .decode(&dtype, len, &ReadContext::new(ctx.to_ids()), &session) .unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs index e81b59dc9a4..112abc921a4 100644 --- a/encodings/bytebool/src/compute.rs +++ b/encodings/bytebool/src/compute.rs @@ -6,17 +6,24 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::Constant; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::TakeExecute; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::match_each_integer_ptype; +use vortex_array::scalar_fn::fns::binary::BooleanKernel; +use vortex_array::scalar_fn::fns::binary::kleene_boolean_buffer_scalar; +use vortex_array::scalar_fn::fns::binary::kleene_boolean_buffers; use vortex_array::scalar_fn::fns::cast::CastKernel; use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_array::scalar_fn::fns::mask::MaskReduce; +use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; +use vortex_error::vortex_err; use super::ByteBool; @@ -106,9 +113,62 @@ impl TakeExecute for ByteBool { } } +impl BooleanKernel for ByteBool { + fn boolean( + lhs: ArrayView<'_, Self>, + rhs: &ArrayRef, + operator: Operator, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); + let lhs_values = truthy_bit_buffer(lhs); + + if let Some(rhs) = rhs.as_opt::() { + let rhs = rhs + .scalar() + .as_bool_opt() + .ok_or_else(|| vortex_err!("expected boolean scalar"))?; + return kleene_boolean_buffer_scalar( + lhs_values, + lhs.validity()?, + &rhs, + operator, + nullability, + ctx, + ) + .map(Some); + } + + let Some(rhs) = rhs.as_opt::() else { + return Ok(None); + }; + + kleene_boolean_buffers( + lhs_values, + lhs.validity()?, + truthy_bit_buffer(rhs), + rhs.validity()?, + operator, + nullability, + ctx, + ) + .map(Some) + } +} + +fn truthy_bit_buffer(array: ArrayView<'_, ByteBool>) -> BitBuffer { + let bytes = array.truthy_bytes(); + BitBuffer::collect_bool(bytes.len(), |idx| bytes[idx] != 0) +} + #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::BoolArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; use vortex_array::compute::conformance::cast::test_cast_conformance; @@ -119,10 +179,18 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_error::vortex_err; + use vortex_session::VortexSession; use super::*; use crate::ByteBoolArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + crate::initialize(&session); + session + }); + fn bb(v: Vec) -> ByteBoolArray { ByteBool::from_vec(v, Validity::AllValid) } @@ -139,7 +207,11 @@ mod tests { let sliced_arr = vortex_arr.slice(1..4).unwrap(); let expected = bb_opt(vec![Some(true), None, Some(false)]); - assert_arrays_eq!(sliced_arr, expected.into_array()); + assert_arrays_eq!( + sliced_arr, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -153,7 +225,11 @@ mod tests { .unwrap(); let expected = bb(vec![true; 5]); - assert_arrays_eq!(arr, expected.into_array()); + assert_arrays_eq!( + arr, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -167,7 +243,11 @@ mod tests { .unwrap(); let expected = bb(vec![false; 5]); - assert_arrays_eq!(arr, expected.into_array()); + assert_arrays_eq!( + arr, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -181,7 +261,38 @@ mod tests { .unwrap(); let expected = bb_opt(vec![Some(true), Some(true), Some(true), Some(false), None]); - assert_arrays_eq!(arr, expected.into_array()); + assert_arrays_eq!( + arr, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); + } + + #[test] + fn test_boolean_kernel_kleene() -> VortexResult<()> { + let lhs = bb_opt(vec![Some(false), Some(true), None, Some(false), None]); + let rhs = bb_opt(vec![None, None, Some(true), Some(false), None]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let and_result = + ::boolean(lhs.as_view(), &rhs, Operator::And, &mut ctx)? + .ok_or_else(|| vortex_err!("ByteBool should handle ByteBool boolean AND"))?; + assert_arrays_eq!( + and_result, + BoolArray::from_iter([Some(false), None, None, Some(false), None]), + &mut ctx + ); + + let or_result = + ::boolean(lhs.as_view(), &rhs, Operator::Or, &mut ctx)? + .ok_or_else(|| vortex_err!("ByteBool should handle ByteBool boolean OR"))?; + assert_arrays_eq!( + or_result, + BoolArray::from_iter([None, Some(true), Some(true), Some(false), None]), + &mut ctx + ); + + Ok(()) } #[test] @@ -240,6 +351,7 @@ mod tests { #[case::single_null(bb_opt(vec![None]))] #[case::mixed_with_nulls(bb_opt(vec![Some(true), None, Some(false), None, Some(true)]))] fn test_bytebool_consistency(#[case] array: ByteBoolArray) { - test_array_consistency(&array.into_array()); + let ctx = &mut array_session().create_execution_ctx(); + test_array_consistency(&array.into_array(), ctx); } } diff --git a/encodings/bytebool/src/kernel.rs b/encodings/bytebool/src/kernel.rs index ddf9679b3d5..678da8642a7 100644 --- a/encodings/bytebool/src/kernel.rs +++ b/encodings/bytebool/src/kernel.rs @@ -1,13 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::binary::BooleanExecuteAdaptor; +use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; +use vortex_session::VortexSession; use crate::ByteBool; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(ByteBool)), - ParentKernelSet::lift(&TakeExecuteAdaptor(ByteBool)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), ByteBool, BooleanExecuteAdaptor(ByteBool)); + kernels.register_execute_parent_kernel(Cast.id(), ByteBool, CastExecuteAdaptor(ByteBool)); + kernels.register_execute_parent_kernel(Dict.id(), ByteBool, TakeExecuteAdaptor(ByteBool)); +} diff --git a/encodings/bytebool/src/lib.rs b/encodings/bytebool/src/lib.rs index 3258341cdb9..a803dacd776 100644 --- a/encodings/bytebool/src/lib.rs +++ b/encodings/bytebool/src/lib.rs @@ -16,7 +16,7 @@ //! to a canonical [`BoolArray`][vortex_array::arrays::BoolArray]: //! //! ``` -//! # use vortex_array::{IntoArray, LEGACY_SESSION, VortexSessionExecute}; +//! # use vortex_array::{IntoArray, VortexSessionExecute, array_session}; //! # use vortex_array::arrays::BoolArray; //! # use vortex_array::arrays::bool::BoolArrayExt; //! # use vortex_array::buffer::BufferHandle; @@ -25,7 +25,7 @@ //! # use vortex_bytebool::ByteBool; //! # use vortex_error::VortexResult; //! # fn main() -> VortexResult<()> { -//! # let mut ctx = LEGACY_SESSION.create_execution_ctx(); +//! # let mut ctx = array_session().create_execution_ctx(); //! let handle = BufferHandle::new_host(ByteBuffer::from(vec![0u8, 1, 42, 0])); //! let array = ByteBool::new(handle, Validity::NonNullable); //! @@ -41,9 +41,17 @@ //! [spec]: https://arrow.apache.org/docs/format/CanonicalExtensions.html#bit-boolean pub use array::*; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; mod array; mod compute; mod kernel; mod rules; mod slice; + +/// Initialize bytebool encoding in the given session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(ByteBool); + kernel::initialize(session); +} diff --git a/encodings/datetime-parts/src/array.rs b/encodings/datetime-parts/src/array.rs index ddc5ee52981..42569da9728 100644 --- a/encodings/datetime-parts/src/array.rs +++ b/encodings/datetime-parts/src/array.rs @@ -42,7 +42,6 @@ use vortex_session::registry::CachedId; use crate::TemporalParts; use crate::canonical::decode_to_temporal; -use crate::compute::kernel::PARENT_KERNELS; use crate::compute::rules::PARENT_RULES; use crate::split_temporal; @@ -123,6 +122,14 @@ impl VTable for DateTimeParts { vortex_panic!("DateTimePartsArray buffer_name index {idx} out of bounds") } + 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, @@ -200,15 +207,6 @@ impl VTable for DateTimeParts { ) -> VortexResult> { PARENT_RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[array_slots(DateTimeParts)] diff --git a/encodings/datetime-parts/src/canonical.rs b/encodings/datetime-parts/src/canonical.rs index 3a46eb6f328..b08b76b148c 100644 --- a/encodings/datetime-parts/src/canonical.rs +++ b/encodings/datetime-parts/src/canonical.rs @@ -99,9 +99,11 @@ pub fn decode_to_temporal( #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::rstest; - use vortex_array::ExecutionCtx; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::assert_arrays_eq; @@ -116,6 +118,12 @@ mod test { use crate::array::DateTimePartsParts; use crate::canonical::decode_to_temporal; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[rstest] #[case(Validity::NonNullable)] #[case(Validity::AllValid)] @@ -133,7 +141,7 @@ mod test { ], validity.clone(), ); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = SESSION.create_execution_ctx(); let date_times = DateTimeParts::try_from_temporal( TemporalArray::new_timestamp( milliseconds.clone().into_array(), @@ -161,7 +169,7 @@ mod test { .clone() .execute::(&mut ctx)?; - assert_arrays_eq!(primitive_values, milliseconds); + assert_arrays_eq!(primitive_values, milliseconds, &mut ctx); Ok(()) } } diff --git a/encodings/datetime-parts/src/compress.rs b/encodings/datetime-parts/src/compress.rs index 6caf3c9f643..a3ab3e6f23b 100644 --- a/encodings/datetime-parts/src/compress.rs +++ b/encodings/datetime-parts/src/compress.rs @@ -62,8 +62,8 @@ pub fn split_temporal(array: TemporalArray, ctx: &mut ExecutionCtx) -> VortexRes mod tests { use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::extension::datetime::TimeUnit; @@ -80,7 +80,7 @@ mod tests { #[case(Validity::AllInvalid)] #[case(Validity::from_iter([true, false, true]))] fn test_split_temporal(#[case] validity: Validity) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let milliseconds = PrimitiveArray::new( buffer![ 86_400i64, // element with only day component diff --git a/encodings/datetime-parts/src/compute/cast.rs b/encodings/datetime-parts/src/compute/cast.rs index ccc3cc0cd91..90aba343ae3 100644 --- a/encodings/datetime-parts/src/compute/cast.rs +++ b/encodings/datetime-parts/src/compute/cast.rs @@ -37,8 +37,8 @@ mod tests { use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::builtins::ArrayBuiltins; @@ -66,7 +66,7 @@ mod tests { TimeUnit::Milliseconds, Some("UTC".into()), ), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() .into_array() @@ -95,7 +95,7 @@ mod tests { #[case(Validity::AllInvalid)] #[case(Validity::from_iter([true, false, true]))] fn test_bad_cast_fails(#[case] validity: Validity) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = date_time_array(validity); // Cast to incompatible type - force evaluation via execute:: let result = array @@ -121,7 +121,7 @@ mod tests { ].into_array(), TimeUnit::Milliseconds, Some("UTC".into()) - ), &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] + ), &mut array_session().create_execution_ctx()).unwrap())] #[case(DateTimeParts::try_from_temporal(TemporalArray::new_timestamp( PrimitiveArray::from_option_iter([ Some(0i64), @@ -132,12 +132,12 @@ mod tests { ]).into_array(), TimeUnit::Milliseconds, Some("UTC".into()) - ), &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] + ), &mut array_session().create_execution_ctx()).unwrap())] #[case(DateTimeParts::try_from_temporal(TemporalArray::new_timestamp( buffer![86_400_000_000_000i64].into_array(), // 1 day in ns TimeUnit::Nanoseconds, Some("UTC".into()) - ), &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] + ), &mut array_session().create_execution_ctx()).unwrap())] fn test_cast_datetime_parts_conformance(#[case] array: DateTimePartsArray) { use vortex_array::compute::conformance::cast::test_cast_conformance; test_cast_conformance(&array.into_array()); diff --git a/encodings/datetime-parts/src/compute/compare.rs b/encodings/datetime-parts/src/compute/compare.rs index 9f854b0e0f9..0791861df79 100644 --- a/encodings/datetime-parts/src/compute/compare.rs +++ b/encodings/datetime-parts/src/compute/compare.rs @@ -204,9 +204,9 @@ mod test { use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::sum::sum; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::dtype::IntegerPType; @@ -228,7 +228,7 @@ mod test { TimeUnit::Seconds, Some("UTC".into()), ), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .expect("Failed to construct DateTimePartsArray from TemporalArray") } @@ -248,7 +248,7 @@ mod test { #[case(Validity::AllValid, Validity::NonNullable)] #[case(Validity::AllValid, Validity::AllValid)] fn compare_date_time_parts_eq(#[case] lhs_validity: Validity, #[case] rhs_validity: Validity) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = dtp_array_from_timestamp(86400i64, lhs_validity); // January 2, 1970, 00:00:00 UTC let rhs = dtp_array_from_timestamp(86400i64, rhs_validity.clone()); // January 2, 1970, 00:00:00 UTC let comparison = lhs @@ -272,7 +272,7 @@ mod test { #[case(Validity::AllValid, Validity::NonNullable)] #[case(Validity::AllValid, Validity::AllValid)] fn compare_date_time_parts_ne(#[case] lhs_validity: Validity, #[case] rhs_validity: Validity) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = dtp_array_from_timestamp(86400i64, lhs_validity); // January 2, 1970, 00:00:00 UTC let rhs = dtp_array_from_timestamp(86401i64, rhs_validity.clone()); // January 2, 1970, 00:00:01 UTC let comparison = lhs @@ -296,7 +296,7 @@ mod test { #[case(Validity::AllValid, Validity::NonNullable)] #[case(Validity::AllValid, Validity::AllValid)] fn compare_date_time_parts_lt(#[case] lhs_validity: Validity, #[case] rhs_validity: Validity) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = dtp_array_from_timestamp(0i64, lhs_validity); // January 1, 1970, 01:00:00 UTC let rhs = dtp_array_from_timestamp(86400i64, rhs_validity); // January 2, 1970, 00:00:00 UTC @@ -313,7 +313,7 @@ mod test { #[case(Validity::AllValid, Validity::NonNullable)] #[case(Validity::AllValid, Validity::AllValid)] fn compare_date_time_parts_gt(#[case] lhs_validity: Validity, #[case] rhs_validity: Validity) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = dtp_array_from_timestamp(86400i64, lhs_validity); // January 2, 1970, 02:00:00 UTC let rhs = dtp_array_from_timestamp(0i64, rhs_validity); // January 1, 1970, 01:00:00 UTC @@ -333,7 +333,7 @@ mod test { #[case] lhs_validity: Validity, #[case] rhs_validity: Validity, ) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let temporal_array = TemporalArray::new_timestamp( PrimitiveArray::new(buffer![0i64], lhs_validity.clone()).into_array(), TimeUnit::Seconds, diff --git a/encodings/datetime-parts/src/compute/filter.rs b/encodings/datetime-parts/src/compute/filter.rs index 149574f53d4..696660510d5 100644 --- a/encodings/datetime-parts/src/compute/filter.rs +++ b/encodings/datetime-parts/src/compute/filter.rs @@ -27,8 +27,8 @@ impl FilterReduce for DateTimeParts { #[cfg(test)] mod test { use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::compute::conformance::filter::test_filter_conformance; @@ -39,7 +39,7 @@ mod test { #[test] fn test_filter_datetime_parts() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Create temporal arrays and convert to DateTimePartsArray let timestamps = buffer![ 0i64, diff --git a/encodings/datetime-parts/src/compute/kernel.rs b/encodings/datetime-parts/src/compute/kernel.rs index e07f1d00f38..4bd0cf2e1b3 100644 --- a/encodings/datetime-parts/src/compute/kernel.rs +++ b/encodings/datetime-parts/src/compute/kernel.rs @@ -1,13 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::VortexSession; use crate::DateTimeParts; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(DateTimeParts)), - ParentKernelSet::lift(&TakeExecuteAdaptor(DateTimeParts)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Binary.id(), + DateTimeParts, + CompareExecuteAdaptor(DateTimeParts), + ); + kernels.register_execute_parent_kernel( + Dict.id(), + DateTimeParts, + TakeExecuteAdaptor(DateTimeParts), + ); +} diff --git a/encodings/datetime-parts/src/compute/mod.rs b/encodings/datetime-parts/src/compute/mod.rs index 5d25757c20d..4196b6b9832 100644 --- a/encodings/datetime-parts/src/compute/mod.rs +++ b/encodings/datetime-parts/src/compute/mod.rs @@ -15,8 +15,8 @@ mod take; mod tests { use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::compute::conformance::consistency::test_array_consistency; @@ -27,7 +27,7 @@ mod tests { use crate::DateTimePartsArray; fn dtp_from_temporal(temporal: TemporalArray) -> DateTimePartsArray { - DateTimeParts::try_from_temporal(temporal, &mut LEGACY_SESSION.create_execution_ctx()) + DateTimeParts::try_from_temporal(temporal, &mut array_session().create_execution_ctx()) .unwrap() } @@ -79,6 +79,7 @@ mod tests { )))] fn test_datetime_parts_consistency(#[case] array: DateTimePartsArray) { - test_array_consistency(&array.into_array()); + let ctx = &mut array_session().create_execution_ctx(); + test_array_consistency(&array.into_array(), ctx); } } diff --git a/encodings/datetime-parts/src/compute/rules.rs b/encodings/datetime-parts/src/compute/rules.rs index bba19ceea47..84f58310ac1 100644 --- a/encodings/datetime-parts/src/compute/rules.rs +++ b/encodings/datetime-parts/src/compute/rules.rs @@ -180,9 +180,9 @@ fn is_constant_zero(array: &ArrayRef) -> bool { mod tests { use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::sum::sum; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; @@ -230,7 +230,7 @@ mod tests { time_unit, None, ); - DateTimeParts::try_from_temporal(temporal, &mut LEGACY_SESSION.create_execution_ctx()) + DateTimeParts::try_from_temporal(temporal, &mut array_session().create_execution_ctx()) .vortex_expect("TemporalArray must produce valid DateTimeParts") } @@ -296,7 +296,7 @@ mod tests { ); // Verify correctness: days [0, 1, 2] <= 1 should give [true, true, false] - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(true_count(&optimized, &mut ctx), 2); } @@ -325,7 +325,7 @@ mod tests { let optimized = between.optimize().unwrap(); // Verify correctness: days [0, 1, 2, 3, 4] between 1 and 3 should give [false, true, true, true, false] - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(true_count(&optimized, &mut ctx), 3); } @@ -348,7 +348,7 @@ mod tests { // (optimization doesn't apply, so we keep the original structure) // Just verify it still computes correctly // days [0, 1, 2] at midnight <= day 1 at noon: [true, true, false] - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(true_count(&optimized, &mut ctx), 2); } @@ -366,7 +366,7 @@ mod tests { TimeUnit::Seconds, None, ); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtp = DateTimeParts::try_from_temporal(temporal, &mut ctx).unwrap(); let len = dtp.len(); diff --git a/encodings/datetime-parts/src/compute/take.rs b/encodings/datetime-parts/src/compute/take.rs index 8b09e2d78ef..9f5a50d3a9b 100644 --- a/encodings/datetime-parts/src/compute/take.rs +++ b/encodings/datetime-parts/src/compute/take.rs @@ -96,8 +96,8 @@ impl TakeExecute for DateTimeParts { mod tests { use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::compute::conformance::take::test_take_conformance; @@ -118,7 +118,7 @@ mod tests { ].into_array(), TimeUnit::Milliseconds, Some("UTC".into()) - ), &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] + ), &mut array_session().create_execution_ctx()).unwrap())] #[case(DateTimeParts::try_from_temporal(TemporalArray::new_timestamp( PrimitiveArray::from_option_iter([ Some(0i64), @@ -129,12 +129,12 @@ mod tests { ]).into_array(), TimeUnit::Milliseconds, Some("UTC".into()) - ), &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] + ), &mut array_session().create_execution_ctx()).unwrap())] #[case(DateTimeParts::try_from_temporal(TemporalArray::new_timestamp( buffer![86_400_000i64].into_array(), TimeUnit::Milliseconds, Some("UTC".into()) - ), &mut LEGACY_SESSION.create_execution_ctx()).unwrap())] + ), &mut array_session().create_execution_ctx()).unwrap())] fn test_take_datetime_parts_conformance(#[case] array: DateTimePartsArray) { test_take_conformance(&array.into_array()); } diff --git a/encodings/datetime-parts/src/lib.rs b/encodings/datetime-parts/src/lib.rs index e061d69bccd..cdc65a44a71 100644 --- a/encodings/datetime-parts/src/lib.rs +++ b/encodings/datetime-parts/src/lib.rs @@ -21,6 +21,7 @@ use vortex_session::VortexSession; /// Initialize datetime-parts encoding in the given session. pub fn initialize(session: &VortexSession) { session.arrays().register(DateTimeParts); + compute::kernel::initialize(session); session.aggregate_fns().register_aggregate_kernel( DateTimeParts.id(), 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 882d07bbaef..376ad12f564 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 @@ -40,8 +40,8 @@ mod tests { use rstest::rstest; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; @@ -56,7 +56,7 @@ mod tests { #[test] fn test_cast_decimal_byte_parts_nullability() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(10, 2); let array = DecimalByteParts::try_new(buffer![100i32, 200, 300, 400].into_array(), decimal_dtype) @@ -79,7 +79,7 @@ mod tests { #[test] fn test_cast_decimal_byte_parts_nullable_to_non_nullable() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(10, 2); let array = DecimalByteParts::try_new( PrimitiveArray::from_option_iter([Some(100i32), None, Some(300)]).into_array(), 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 e6cc8d4d72d..04f321cf88f 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 @@ -144,7 +144,10 @@ where #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; @@ -159,9 +162,16 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::DecimalByteParts; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn compare_decimal_const() { let decimal_dtype = DecimalDType::new(8, 2); @@ -180,7 +190,7 @@ mod tests { let res = lhs.binary(rhs.into_array(), Operator::Eq).unwrap(); let expected = BoolArray::from_iter([Some(false), Some(false), Some(true)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -208,7 +218,7 @@ mod tests { let res = lhs.into_array().binary(rhs, Operator::Lte)?; let expected = BoolArray::from_iter([None, Some(true), Some(true), Some(true)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -235,15 +245,15 @@ mod tests { let res = lhs.binary(rhs.clone().into_array(), Operator::Eq).unwrap(); let expected = BoolArray::from_iter([Some(false), Some(false), Some(false)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); let res = lhs.binary(rhs.clone().into_array(), Operator::Gt).unwrap(); let expected = BoolArray::from_iter([Some(true), Some(true), Some(true)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); let res = lhs.binary(rhs.into_array(), Operator::Lt).unwrap(); let expected = BoolArray::from_iter([Some(false), Some(false), Some(false)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); // This cannot be converted to a i32. let rhs = ConstantArray::new( @@ -253,14 +263,14 @@ mod tests { let res = lhs.binary(rhs.clone().into_array(), Operator::Eq).unwrap(); let expected = BoolArray::from_iter([Some(false), Some(false), Some(false)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); let res = lhs.binary(rhs.clone().into_array(), Operator::Gt).unwrap(); let expected = BoolArray::from_iter([Some(false), Some(false), Some(false)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); let res = lhs.binary(rhs.into_array(), Operator::Lt).unwrap(); let expected = BoolArray::from_iter([Some(true), Some(true), Some(true)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs index 9e16b4f81dc..5e8d28e3526 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs @@ -1,13 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::VortexSession; use crate::DecimalByteParts; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(DecimalByteParts)), - ParentKernelSet::lift(&TakeExecuteAdaptor(DecimalByteParts)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Binary.id(), + DecimalByteParts, + CompareExecuteAdaptor(DecimalByteParts), + ); + kernels.register_execute_parent_kernel( + Dict.id(), + DecimalByteParts, + TakeExecuteAdaptor(DecimalByteParts), + ); +} 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 35da5be40b6..0a5477e93cf 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 @@ -13,6 +13,8 @@ mod take; mod tests { use rstest::rstest; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::dtype::DecimalDType; @@ -69,6 +71,7 @@ mod tests { ).unwrap())] fn test_decimal_byte_parts_consistency(#[case] array: DecimalBytePartsArray) { - test_array_consistency(&array.into_array()); + let ctx = &mut array_session().create_execution_ctx(); + test_array_consistency(&array.into_array(), ctx); } } 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 42135a9ecaa..864cd604885 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -46,7 +46,6 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::decimal_byte_parts::compute::kernel::PARENT_KERNELS; use crate::decimal_byte_parts::rules::PARENT_RULES; /// A [`DecimalByteParts`]-encoded Vortex array. @@ -109,6 +108,14 @@ impl VTable for DecimalByteParts { vortex_panic!("DecimalBytePartsArray buffer_name index {idx} out of bounds") } + 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, @@ -165,15 +172,6 @@ impl VTable for DecimalByteParts { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { to_canonical_decimal(&array, ctx).map(ExecutionResult::done) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } /// The most significant parts of the decimal values. @@ -330,8 +328,8 @@ impl ValidityChild for DecimalByteParts { #[cfg(test)] mod tests { use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DType; @@ -363,7 +361,7 @@ mod tests { assert_eq!( Scalar::null(dtype.clone()), array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() ); assert_eq!( @@ -373,13 +371,13 @@ mod tests { ) .unwrap(), array - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() ); assert_eq!( Scalar::try_new(dtype, Some(ScalarValue::Decimal(DecimalValue::I64(400)))).unwrap(), array - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap() ); } diff --git a/encodings/decimal-byte-parts/src/lib.rs b/encodings/decimal-byte-parts/src/lib.rs index 86566027afa..36a53c3a614 100644 --- a/encodings/decimal-byte-parts/src/lib.rs +++ b/encodings/decimal-byte-parts/src/lib.rs @@ -23,6 +23,7 @@ use vortex_session::VortexSession; /// Initialize decimal-byte-parts encoding in the given session. pub fn initialize(session: &VortexSession) { session.arrays().register(DecimalByteParts); + compute::kernel::initialize(session); session.aggregate_fns().register_aggregate_kernel( DecimalByteParts.id(), diff --git a/encodings/experimental/onpair/benches/decode.rs b/encodings/experimental/onpair/benches/decode.rs index ff8db8c22c2..ff4158cda32 100644 --- a/encodings/experimental/onpair/benches/decode.rs +++ b/encodings/experimental/onpair/benches/decode.rs @@ -33,6 +33,7 @@ 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::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; @@ -41,7 +42,6 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; -use vortex_array::session::ArraySession; use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_mask::Mask; @@ -80,8 +80,11 @@ impl DecodeInputs { use vortex_onpair::onpair_compress; use vortex_session::VortexSession; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + vortex_onpair::initialize(&session); + session +}); #[derive(Copy, Clone, Debug)] enum Shape { @@ -151,13 +154,13 @@ fn corpus(n: usize, shape: Shape) -> Vec { out } -fn compress(n: usize, shape: Shape) -> OnPairArray { +fn compress(n: usize, shape: Shape, ctx: &mut ExecutionCtx) -> OnPairArray { let strings = corpus(n, shape); let varbin = VarBinArray::from_iter( strings.iter().map(|s| Some(s.as_bytes())), DType::Utf8(Nullability::NonNullable), ); - onpair_compress(&varbin, varbin.len(), varbin.dtype(), DEFAULT_DICT12_CONFIG) + onpair_compress(varbin.as_array(), DEFAULT_DICT12_CONFIG, ctx) .unwrap_or_else(|e| panic!("onpair_compress failed: {e}")) } @@ -170,13 +173,12 @@ fn widen(arr: &ArrayRef, ctx: &mut ExecutionCtx) -> Buffer { .into_buffer::() } -fn materialise(arr: &OnPairArray) -> (DecodeInputs, usize) { - let mut ctx = SESSION.create_execution_ctx(); +fn materialise(arr: &OnPairArray, ctx: &mut ExecutionCtx) -> (DecodeInputs, usize) { let view = arr.as_view(); let inputs = DecodeInputs { dict_bytes: view.dict_bytes().clone(), - dict_offsets: widen::(view.dict_offsets(), &mut ctx), - codes: widen::(view.codes(), &mut ctx), + dict_offsets: widen::(view.dict_offsets(), ctx), + codes: widen::(view.codes(), ctx), bits: view.bits(), }; let total = inputs.decompressed_len(); @@ -195,9 +197,10 @@ const CASES: &[(Shape, usize)] = &[ /// Hits `onpair::decompress_into` directly. #[divan::bench(args = CASES)] fn decompress_into_bench(bencher: Bencher, case: (Shape, usize)) { + let mut ctx = SESSION.create_execution_ctx(); let (shape, n) = case; - let arr = compress(n, shape); - let (inputs, total) = materialise(&arr); + 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()); @@ -210,8 +213,9 @@ fn decompress_into_bench(bencher: Bencher, case: (Shape, usize)) { /// building the view buffer + `BinaryView` list, etc. #[divan::bench(args = CASES)] fn canonicalize_to_varbinview(bencher: Bencher, case: (Shape, usize)) { + let mut ctx = SESSION.create_execution_ctx(); let (shape, n) = case; - let arr = compress(n, shape); + let arr = compress(n, shape, &mut ctx); bencher .with_inputs(|| (arr.clone().into_array(), SESSION.create_execution_ctx())) .bench_local_values(|(arr, mut ctx)| { @@ -230,8 +234,9 @@ const COMPUTE_CASES: &[(Shape, usize)] = &[(Shape::UrlLog, 100_000), (Shape::Url /// rows; the cost is dominated by the `codes` segment copy + offsets. #[divan::bench(args = COMPUTE_CASES)] fn filter_share_dict(bencher: Bencher, case: (Shape, usize)) { + let mut ctx = SESSION.create_execution_ctx(); let (shape, n) = case; - let arr = compress(n, shape); + let arr = compress(n, shape, &mut ctx); let mask = Mask::from_iter((0..n).map(|i| i % 7 == 0)); bencher .with_inputs(|| SESSION.create_execution_ctx()) diff --git a/encodings/experimental/onpair/src/array.rs b/encodings/experimental/onpair/src/array.rs index f9e18c36ac4..58fd0332ac3 100644 --- a/encodings/experimental/onpair/src/array.rs +++ b/encodings/experimental/onpair/src/array.rs @@ -43,7 +43,6 @@ use vortex_session::registry::CachedId; use crate::canonical::canonicalize_onpair; use crate::canonical::onpair_decode_views; -use crate::kernel::PARENT_KERNELS; use crate::rules::RULES; /// An [`OnPair`]-encoded Vortex array. @@ -368,6 +367,24 @@ impl VTable for OnPair { } } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + let mut data = array.data().clone(); + data.dict_bytes = buffers[0].clone(); + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, @@ -499,15 +516,6 @@ impl VTable for OnPair { Ok(()) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/encodings/experimental/onpair/src/compress.rs b/encodings/experimental/onpair/src/compress.rs index 226922cf593..adf9e73461c 100644 --- a/encodings/experimental/onpair/src/compress.rs +++ b/encodings/experimental/onpair/src/compress.rs @@ -8,20 +8,20 @@ use onpair::Offset; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; -use vortex_array::VortexSessionExecute; -use vortex_array::accessor::ArrayAccessor; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbinview::BinaryView; use vortex_array::buffer::BufferHandle; -use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; 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; @@ -29,49 +29,66 @@ use crate::OnPairArray; /// Default OnPair training configuration: 12-bit codes ("dict-12"). pub const DEFAULT_DICT12_CONFIG: Config = onpair::DEFAULT_CONFIG; -/// Compress an iterable of optional byte strings via the OnPair encoder. -pub fn onpair_compress_iter<'a, I>( - iter: I, - len: usize, - dtype: DType, - config: Config, -) -> VortexResult -where - I: Iterator>, -{ - onpair_compress_iter_with_offsets::(iter, len, dtype, config) -} - -fn onpair_compress_iter_with_offsets<'a, O, I>( - iter: I, - len: usize, - dtype: DType, +fn onpair_compress_varbinview( + array: VarBinViewArray, config: Config, + ctx: &mut ExecutionCtx, ) -> VortexResult where O: Offset, - I: Iterator>, { + 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); - let mut validity_bits: Vec = Vec::with_capacity(len); - offsets.push(::from_usize(0)); + offsets.push(O::from_usize(0)); + let views = array.views(); + let buffers = array + .data_buffers() + .as_ref() + .iter() + .map(|b| b.as_host()) + .collect::>(); - for item in iter { - match item { - Some(bytes) => { + match mask.bit_buffer() { + AllOr::All => { + for view in views { + let bytes = view_bytes(view, &buffers); flat.extend_from_slice(bytes); - offsets.push(::from_usize(flat.len())); - uncompressed_lengths.push( - i32::try_from(bytes.len()).vortex_expect("string length must fit in i32"), - ); - validity_bits.push(true); + offsets.push(O::from_usize(flat.len())); + uncompressed_lengths + .push(i32::try_from(view.len()).vortex_expect("must fit in i32")); } - None => { - offsets.push(::from_usize(flat.len())); - uncompressed_lengths.push(0); - validity_bits.push(false); + } + 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); + } } } } @@ -80,44 +97,53 @@ where .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)?; + 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 codes_offsets = Buffer::from(codes_offsets).into_array(); let uncompressed_lengths = uncompressed_lengths.into_array(); - let validity = match dtype.nullability() { - Nullability::NonNullable => Validity::NonNullable, - Nullability::Nullable => Validity::from_iter(validity_bits), - }; OnPair::try_new( - dtype, + array.dtype().clone(), dict_bytes, dict_offsets, codes, codes_offsets, uncompressed_lengths, - validity, + 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. - let mut padded = Vec::with_capacity(dict_bytes.len() + onpair::MAX_TOKEN_SIZE); - padded.extend_from_slice(&dict_bytes); - padded.resize(dict_bytes.len() + onpair::MAX_TOKEN_SIZE, 0); + // // 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`. - BufferHandle::new_host(ByteBuffer::from(padded).aligned(vortex_buffer::Alignment::new(8))) + 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 @@ -128,9 +154,9 @@ fn build_codes_offsets( codes: &[u16], dict_offsets: &[u32], row_byte_offsets: &[O], -) -> VortexResult> { +) -> VortexResult> { let nrows = row_byte_offsets.len() - 1; - let mut codes_offsets = Vec::with_capacity(nrows + 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; @@ -149,38 +175,16 @@ fn build_codes_offsets( .map_err(|_| vortex_err!("OnPair: code boundary {code_idx} does not fit u32"))?, ); } - Ok(codes_offsets) -} - -/// Compress a byte-string accessor (typically a `VarBinArray` or -/// `VarBinViewArray`). -pub fn onpair_compress>( - array: A, - len: usize, - dtype: &DType, - config: Config, -) -> VortexResult { - array.with_iterator(|iter| onpair_compress_iter(iter, len, dtype.clone(), config)) + Ok(codes_offsets.freeze()) } /// Compress any [`ArrayRef`] whose canonical form is a string array, by first /// canonicalising to `VarBinViewArray`. -pub fn onpair_compress_array( +pub fn onpair_compress( array: &ArrayRef, config: Config, ctx: &mut ExecutionCtx, ) -> VortexResult { let view = array.clone().execute::(ctx)?; - let len = view.len(); - let dtype = view.dtype().clone(); - onpair_compress(&view, len, &dtype, config) -} - -/// Convenience: build a default `ExecutionCtx` from `LEGACY_SESSION`. -pub fn onpair_compress_array_default( - array: &ArrayRef, - config: Config, -) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - onpair_compress_array(array, config, &mut ctx) + onpair_compress_varbinview::(view, config, ctx) } diff --git a/encodings/experimental/onpair/src/compute/byte_length.rs b/encodings/experimental/onpair/src/compute/byte_length.rs new file mode 100644 index 00000000000..933e8bccd85 --- /dev/null +++ b/encodings/experimental/onpair/src/compute/byte_length.rs @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::IntoArray; +use vortex_array::ValidityVTable; +use vortex_array::arrays::ConstantArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::byte_length::ByteLengthKernel; +use vortex_array::validity::Validity; + +use crate::OnPair; +use crate::OnPairArraySlotsExt; + +// TODO(myrrc): this and FSST comparison should be in the same trait. +// https://github.com/vortex-data/vortex/tree/myrrc/onpair-compare-trait +impl ByteLengthKernel for OnPair { + fn byte_length( + array: vortex_array::ArrayView<'_, Self>, + _ctx: &mut vortex_array::ExecutionCtx, + ) -> vortex_error::VortexResult> { + let nullable = array.dtype().nullability(); + let dtype = DType::Primitive(PType::U64, nullable); + // Uncompressed lengths are non-nullable and may be less than u64 each + let lengths = array.uncompressed_lengths().cast(dtype.clone())?; + Ok(Some(match OnPair::validity(array)? { + Validity::NonNullable | Validity::AllValid => lengths, + Validity::Array(v) => lengths.mask(v)?, + Validity::AllInvalid => { + ConstantArray::new(Scalar::null(dtype), lengths.len()).into_array() + } + })) + } +} diff --git a/encodings/experimental/onpair/src/compute/compare.rs b/encodings/experimental/onpair/src/compute/compare.rs new file mode 100644 index 00000000000..939a1a2fb42 --- /dev/null +++ b/encodings/experimental/onpair/src/compute/compare.rs @@ -0,0 +1,145 @@ +// 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/compute/mod.rs b/encodings/experimental/onpair/src/compute/mod.rs index 4cb15868625..4ad5f48f578 100644 --- a/encodings/experimental/onpair/src/compute/mod.rs +++ b/encodings/experimental/onpair/src/compute/mod.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod byte_length; mod cast; +mod compare; mod filter; mod slice; diff --git a/encodings/experimental/onpair/src/kernel.rs b/encodings/experimental/onpair/src/kernel.rs index fdd521e887e..8863d750a72 100644 --- a/encodings/experimental/onpair/src/kernel.rs +++ b/encodings/experimental/onpair/src/kernel.rs @@ -1,11 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Filter; use vortex_array::arrays::filter::FilterExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +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_session::VortexSession; use crate::OnPair; // TODO: implement ListExecute & TakeExecute for OnPair -pub(super) const PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&FilterExecuteAdaptor(OnPair))]); +pub(super) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + 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( + ByteLength.id(), + OnPair, + ByteLengthExecuteAdaptor(OnPair), + ); +} diff --git a/encodings/experimental/onpair/src/lib.rs b/encodings/experimental/onpair/src/lib.rs index 94c18b6dec8..11b22f63bc1 100644 --- a/encodings/experimental/onpair/src/lib.rs +++ b/encodings/experimental/onpair/src/lib.rs @@ -27,3 +27,11 @@ pub use onpair::Bits; pub use onpair::Config; pub use onpair::Error as OnPairError; pub use onpair::Threshold; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; + +/// Initialize OnPair encoding in the given session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(OnPair); + kernel::initialize(session); +} diff --git a/encodings/experimental/onpair/src/tests.rs b/encodings/experimental/onpair/src/tests.rs index dd6fe4b0116..4f2124025f2 100644 --- a/encodings/experimental/onpair/src/tests.rs +++ b/encodings/experimental/onpair/src/tests.rs @@ -6,7 +6,6 @@ use std::sync::LazyLock; use prost::Message; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; @@ -15,7 +14,6 @@ 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::session::ArraySession; use vortex_array::test_harness::check_metadata; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; @@ -28,8 +26,7 @@ use crate::OnPairMetadata; use crate::compress::DEFAULT_DICT12_CONFIG; use crate::compress::onpair_compress; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn sample_input() -> VarBinArray { VarBinArray::from_iter( @@ -64,74 +61,65 @@ fn test_onpair_metadata_golden() { #[cfg_attr(miri, ignore)] #[test] -fn test_onpair_roundtrip() { +fn test_onpair_roundtrip() -> vortex_error::VortexResult<()> { let input = sample_input(); - let len = input.len(); - let dtype = input.dtype().clone(); - let compressed = onpair_compress(&input, len, &dtype, DEFAULT_DICT12_CONFIG).expect("compress"); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = onpair_compress(&input.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; assert!(compressed.clone().into_array().is::()); - let mut ctx = SESSION.create_execution_ctx(); let decoded = compressed .into_array() - .execute::(&mut ctx) - .expect("canonicalize"); + .execute::(&mut ctx)?; - decoded - .with_iterator(|iter| { - let got: Vec>> = iter.map(|b| b.map(|s| s.to_vec())).collect(); - assert_eq!(got.len(), 5); - assert_eq!( - got[0].as_deref(), - Some(b"https://www.example.com/page".as_ref()) - ); - assert_eq!( - got[3].as_deref(), - Some(b"ftp://files.example.com/x".as_ref()) - ); - Ok::<_, vortex_error::VortexError>(()) - }) - .unwrap(); + let mask = decoded.validity()?.execute_mask(decoded.len(), &mut ctx)?; + let got: Vec>> = (0..decoded.len()) + .map(|i| mask.value(i).then(|| decoded.bytes_at(i).to_vec())) + .collect(); + assert_eq!(got.len(), 5); + assert_eq!( + got[0].as_deref(), + Some(b"https://www.example.com/page".as_ref()) + ); + assert_eq!( + got[3].as_deref(), + Some(b"ftp://files.example.com/x".as_ref()) + ); + Ok(()) } #[cfg_attr(miri, ignore)] #[test] -fn test_onpair_nullable_canonicalize() { +fn test_onpair_nullable_canonicalize() -> vortex_error::VortexResult<()> { let input = VarBinArray::from_iter( [Some("a"), None, Some("bbb"), None, Some("ccccc")], DType::Utf8(Nullability::Nullable), ); - let len = input.len(); - let dtype = input.dtype().clone(); - let arr = onpair_compress(&input, len, &dtype, DEFAULT_DICT12_CONFIG).unwrap(); let mut ctx = SESSION.create_execution_ctx(); - let canonical = arr - .into_array() - .execute::(&mut ctx) - .unwrap(); - canonical - .with_iterator(|iter| { - let got: Vec>> = iter.map(|b| b.map(|s| s.to_vec())).collect(); - assert_eq!(got[1], None); - assert_eq!(got[3], None); - assert_eq!(got[4].as_deref(), Some(b"ccccc".as_ref())); - Ok::<_, vortex_error::VortexError>(()) - }) - .unwrap(); + let arr = onpair_compress(&input.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let canonical = arr.into_array().execute::(&mut ctx)?; + let mask = canonical + .validity()? + .execute_mask(canonical.len(), &mut ctx)?; + let got: Vec>> = (0..canonical.len()) + .map(|i| mask.value(i).then(|| canonical.bytes_at(i).to_vec())) + .collect(); + assert_eq!(got[1], None); + assert_eq!(got[3], None); + assert_eq!(got[4].as_deref(), Some(b"ccccc".as_ref())); + Ok(()) } #[cfg_attr(miri, ignore)] #[test] -fn test_onpair_scalar_at() { +fn test_onpair_scalar_at() -> vortex_error::VortexResult<()> { let input = sample_input(); - let len = input.len(); - let dtype = input.dtype().clone(); - let arr = onpair_compress(&input, len, &dtype, DEFAULT_DICT12_CONFIG).unwrap(); let mut ctx = SESSION.create_execution_ctx(); - let s = arr.into_array().execute_scalar(2, &mut ctx).unwrap(); + let arr = onpair_compress(&input.into_array(), DEFAULT_DICT12_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"); + Ok(()) } /// `scalar_at` must decode only the requested row's code window — fetching @@ -151,10 +139,9 @@ fn test_onpair_scalar_at_window() -> vortex_error::VortexResult<()> { strings.iter().map(|s| Some(s.as_bytes())), DType::Utf8(Nullability::NonNullable), ); - let arr = - onpair_compress(&varbin, varbin.len(), varbin.dtype(), DEFAULT_DICT12_CONFIG)?.into_array(); - let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?.into_array(); + for &i in &[0usize, 1, 999, 1000, n - 1] { let got = arr.execute_scalar(i, &mut ctx)?; assert_eq!( @@ -194,51 +181,43 @@ fn test_onpair_scalar_at_window() -> vortex_error::VortexResult<()> { #[case::n_7(7)] #[case::n_8(8)] #[case::n_9(9)] -fn test_onpair_unroll_tail_boundaries(#[case] n: usize) { +fn test_onpair_unroll_tail_boundaries(#[case] n: usize) -> vortex_error::VortexResult<()> { let words: &[&str] = &["a", "bb", "ccc", "https://www.example.com/x"]; let strings: Vec<&str> = (0..n).map(|i| words[i % words.len()]).collect(); let input = VarBinArray::from_iter( strings.iter().map(|s| Some(*s)), DType::Utf8(Nullability::NonNullable), ); - let len = input.len(); - let dtype = input.dtype().clone(); - let arr = onpair_compress(&input, len, &dtype, DEFAULT_DICT12_CONFIG).unwrap(); let mut ctx = SESSION.create_execution_ctx(); - let canonical = arr - .into_array() - .execute::(&mut ctx) - .unwrap(); - canonical - .with_iterator(|iter| { - let got: Vec>> = iter.map(|b| b.map(|s| s.to_vec())).collect(); - assert_eq!(got.len(), n); - for (i, expected) in strings.iter().enumerate() { - assert_eq!(got[i].as_deref(), Some(expected.as_bytes()), "n={n}, i={i}"); - } - Ok::<_, vortex_error::VortexError>(()) - }) - .unwrap(); + let arr = onpair_compress(&input.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + let canonical = arr.into_array().execute::(&mut ctx)?; + let mask = canonical + .validity()? + .execute_mask(canonical.len(), &mut ctx)?; + let got: Vec>> = (0..canonical.len()) + .map(|i| mask.value(i).then(|| canonical.bytes_at(i).to_vec())) + .collect(); + assert_eq!(got.len(), n); + for (i, expected) in strings.iter().enumerate() { + assert_eq!(got[i].as_deref(), Some(expected.as_bytes()), "n={n}, i={i}"); + } + Ok(()) } /// Empty array — the unroll path must short-circuit cleanly. #[cfg_attr(miri, ignore)] #[test] -fn test_onpair_empty() { +fn test_onpair_empty() -> vortex_error::VortexResult<()> { let input = VarBinArray::from_iter( std::iter::empty::>(), DType::Utf8(Nullability::NonNullable), ); - let len = input.len(); - let dtype = input.dtype().clone(); - let arr = onpair_compress(&input, len, &dtype, DEFAULT_DICT12_CONFIG).unwrap(); - assert_eq!(arr.len(), 0); let mut ctx = SESSION.create_execution_ctx(); - let canonical = arr - .into_array() - .execute::(&mut ctx) - .unwrap(); + let arr = onpair_compress(&input.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; + assert_eq!(arr.len(), 0); + let canonical = arr.into_array().execute::(&mut ctx)?; assert_eq!(canonical.len(), 0); + Ok(()) } /// Filter must share the dictionary — never recompress (this is the @@ -246,7 +225,7 @@ fn test_onpair_empty() { /// and check that the result is bit-exact and still an OnPairArray. #[cfg_attr(miri, ignore)] #[test] -fn test_onpair_filter_shares_dict() { +fn test_onpair_filter_shares_dict() -> vortex_error::VortexResult<()> { let n = 5_000usize; let strings: Vec = (0..n) .map(|i| format!("https://www.example.com/items/{i:08}")) @@ -255,8 +234,8 @@ fn test_onpair_filter_shares_dict() { strings.iter().map(|s| Some(s.as_bytes())), DType::Utf8(Nullability::NonNullable), ); - let arr = - onpair_compress(&varbin, varbin.len(), varbin.dtype(), DEFAULT_DICT12_CONFIG).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; let dict_bytes_before = arr.dict_bytes().clone(); let dict_offsets_len_before = arr.dict_offsets().len(); @@ -270,35 +249,33 @@ fn test_onpair_filter_shares_dict() { .collect(); let mut filter_ctx = SESSION.create_execution_ctx(); - let filtered = ::filter(arr.as_view(), &mask, &mut filter_ctx) - .unwrap() + let filtered = ::filter(arr.as_view(), &mask, &mut filter_ctx)? .expect("OnPair filter must return Some"); assert!( filtered.is::(), "filter dropped OnPair encoding: got {}", filtered.encoding_id() ); - let typed = filtered.try_downcast::().expect("OnPair"); + let typed = filtered + .try_downcast::() + .map_err(|_| vortex_error::vortex_err!("filter result was not OnPair"))?; // Dict must be byte-identical with the input — no retrain, no copy. assert_eq!(typed.dict_bytes().as_slice(), dict_bytes_before.as_slice()); assert_eq!(typed.dict_offsets().len(), dict_offsets_len_before); assert_eq!(typed.len(), expected.len()); - let mut ctx = SESSION.create_execution_ctx(); - let canonical = typed - .into_array() - .execute::(&mut ctx) - .unwrap(); - canonical - .with_iterator(|iter| { - let got: Vec>> = iter.map(|b| b.map(|s| s.to_vec())).collect(); - assert_eq!(got.len(), expected.len()); - for (i, want) in expected.iter().enumerate() { - assert_eq!(got[i].as_deref(), Some(want.as_bytes()), "row {i}"); - } - Ok::<_, vortex_error::VortexError>(()) - }) - .unwrap(); + let canonical = typed.into_array().execute::(&mut ctx)?; + let mask = canonical + .validity()? + .execute_mask(canonical.len(), &mut ctx)?; + let got: Vec>> = (0..canonical.len()) + .map(|i| mask.value(i).then(|| canonical.bytes_at(i).to_vec())) + .collect(); + assert_eq!(got.len(), expected.len()); + for (i, want) in expected.iter().enumerate() { + assert_eq!(got[i].as_deref(), Some(want.as_bytes()), "row {i}"); + } + Ok(()) } /// Rebuild an OnPair array, swapping `codes_offsets` for a narrowed @@ -350,7 +327,7 @@ fn narrow_codes_offsets(arr: &crate::OnPairArray, target: PType) -> crate::OnPai /// of type u16`. The fix dispatches via `match_each_integer_ptype!`. #[cfg_attr(miri, ignore)] #[test] -fn test_onpair_filter_with_narrowed_codes_offsets_u16() { +fn test_onpair_filter_with_narrowed_codes_offsets_u16() -> vortex_error::VortexResult<()> { let n = 200usize; // Short rows so per-row token counts stay small and codes_offsets // values fit in u16. (We narrow manually below regardless — this @@ -361,8 +338,8 @@ fn test_onpair_filter_with_narrowed_codes_offsets_u16() { strings.iter().map(|s| Some(s.as_bytes())), DType::Utf8(Nullability::NonNullable), ); - let arr = - onpair_compress(&varbin, varbin.len(), varbin.dtype(), DEFAULT_DICT12_CONFIG).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; // Force `codes_offsets` to u16 so the panicking pre-fix // `as_slice::()` would fire. @@ -384,27 +361,25 @@ fn test_onpair_filter_with_narrowed_codes_offsets_u16() { let mut filter_ctx = SESSION.create_execution_ctx(); // Pre-fix: this call panics with "Attempted to get slice of type // u32 from array of type u16". Post-fix: succeeds. - let filtered = ::filter(arr.as_view(), &mask, &mut filter_ctx) - .unwrap() + let filtered = ::filter(arr.as_view(), &mask, &mut filter_ctx)? .expect("OnPair filter must return Some"); - let typed = filtered.try_downcast::().expect("OnPair"); + let typed = filtered + .try_downcast::() + .map_err(|_| vortex_error::vortex_err!("filter result was not OnPair"))?; assert_eq!(typed.len(), expected.len()); - let mut ctx = SESSION.create_execution_ctx(); - let canonical = typed - .into_array() - .execute::(&mut ctx) - .unwrap(); - canonical - .with_iterator(|iter| { - let got: Vec>> = iter.map(|b| b.map(|s| s.to_vec())).collect(); - assert_eq!(got.len(), expected.len()); - for (i, want) in expected.iter().enumerate() { - assert_eq!(got[i].as_deref(), Some(want.as_bytes()), "row {i}"); - } - Ok::<_, vortex_error::VortexError>(()) - }) - .unwrap(); + let canonical = typed.into_array().execute::(&mut ctx)?; + let mask = canonical + .validity()? + .execute_mask(canonical.len(), &mut ctx)?; + let got: Vec>> = (0..canonical.len()) + .map(|i| mask.value(i).then(|| canonical.bytes_at(i).to_vec())) + .collect(); + assert_eq!(got.len(), expected.len()); + for (i, want) in expected.iter().enumerate() { + assert_eq!(got[i].as_deref(), Some(want.as_bytes()), "row {i}"); + } + Ok(()) } /// Same regression, narrowed to u8 (smallest possible ptype) — extra @@ -412,25 +387,25 @@ fn test_onpair_filter_with_narrowed_codes_offsets_u16() { /// cascading compressor might pick. #[cfg_attr(miri, ignore)] #[test] -fn test_onpair_filter_with_narrowed_codes_offsets_u8() { +fn test_onpair_filter_with_narrowed_codes_offsets_u8() -> vortex_error::VortexResult<()> { let n = 100usize; let strings: Vec = (0..n).map(|i| format!("{i}")).collect(); let varbin = VarBinArray::from_iter( strings.iter().map(|s| Some(s.as_bytes())), DType::Utf8(Nullability::NonNullable), ); - let arr = - onpair_compress(&varbin, varbin.len(), varbin.dtype(), DEFAULT_DICT12_CONFIG).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; let arr = narrow_codes_offsets(&arr, PType::U8); assert_eq!(arr.as_view().codes_offsets().dtype().as_ptype(), PType::U8); let mask = vortex_mask::Mask::from_iter((0..n).map(|i| i % 2 == 0)); let mut filter_ctx = SESSION.create_execution_ctx(); - let filtered = ::filter(arr.as_view(), &mask, &mut filter_ctx) - .unwrap() + let filtered = ::filter(arr.as_view(), &mask, &mut filter_ctx)? .expect("OnPair filter must return Some"); assert_eq!(filtered.len(), n / 2); + Ok(()) } /// Regression: canonicalising a *sliced* OnPair array. `slice` keeps the full @@ -452,8 +427,8 @@ fn test_onpair_slice_canonicalize() -> vortex_error::VortexResult<()> { strings.iter().map(|s| Some(s.as_bytes())), DType::Utf8(Nullability::NonNullable), ); - let arr = - onpair_compress(&varbin, varbin.len(), varbin.dtype(), DEFAULT_DICT12_CONFIG)?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?.into_array(); // interior (start>0, end0, // end=n), and a near-full window. @@ -466,20 +441,21 @@ fn test_onpair_slice_canonicalize() -> vortex_error::VortexResult<()> { sliced.encoding_id() ); - let mut ctx = SESSION.create_execution_ctx(); let canonical = sliced.execute::(&mut ctx)?; - canonical.with_iterator(|iter| { - let got: Vec>> = iter.map(|b| b.map(|s| s.to_vec())).collect(); - assert_eq!(got.len(), end - start, "window {start}..{end} length"); - for (i, want) in strings[start..end].iter().enumerate() { - assert_eq!( - got[i].as_deref(), - Some(want.as_bytes()), - "window {start}..{end} row {i}" - ); - } - Ok::<_, vortex_error::VortexError>(()) - })?; + let mask = canonical + .validity()? + .execute_mask(canonical.len(), &mut ctx)?; + let got: Vec>> = (0..canonical.len()) + .map(|i| mask.value(i).then(|| canonical.bytes_at(i).to_vec())) + .collect(); + assert_eq!(got.len(), end - start, "window {start}..{end} length"); + for (i, want) in strings[start..end].iter().enumerate() { + assert_eq!( + got[i].as_deref(), + Some(want.as_bytes()), + "window {start}..{end} row {i}" + ); + } } Ok(()) } diff --git a/encodings/experimental/onpair/tests/big_data.rs b/encodings/experimental/onpair/tests/big_data.rs index c6bc163ea9f..e7f74fdd782 100644 --- a/encodings/experimental/onpair/tests/big_data.rs +++ b/encodings/experimental/onpair/tests/big_data.rs @@ -17,7 +17,6 @@ use std::time::Instant; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::accessor::ArrayAccessor; use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinArray; @@ -26,13 +25,11 @@ 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_array::session::ArraySession; use vortex_onpair::DEFAULT_DICT12_CONFIG; use vortex_onpair::onpair_compress; use vortex_session::VortexSession; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn corpus(n: usize) -> Vec { let templates: &[&str] = &[ @@ -65,7 +62,7 @@ fn corpus(n: usize) -> Vec { #[test] #[cfg_attr(miri, ignore)] -fn smoke_100k_rows() { +fn smoke_100k_rows() -> vortex_error::VortexResult<()> { let n = 100_000; let strings = corpus(n); let raw_bytes: usize = strings.iter().map(|s| s.len()).sum(); @@ -75,9 +72,9 @@ fn smoke_100k_rows() { DType::Utf8(Nullability::NonNullable), ); + let mut ctx = SESSION.create_execution_ctx(); let t0 = Instant::now(); - let arr = onpair_compress(&varbin, varbin.len(), varbin.dtype(), DEFAULT_DICT12_CONFIG) - .expect("compress"); + let arr = onpair_compress(&varbin.into_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?; let compress_elapsed = t0.elapsed(); let bits = arr.bits(); eprintln!( @@ -86,26 +83,19 @@ fn smoke_100k_rows() { ); let arr_ref = arr.into_array(); - let mut ctx = SESSION.create_execution_ctx(); // Full canonical round-trip via the pure-Rust decoder. let t0 = Instant::now(); - let decoded = arr_ref - .clone() - .execute::(&mut ctx) - .expect("canonicalize"); + let decoded = arr_ref.clone().execute::(&mut ctx)?; eprintln!("canonicalized in {:?}", t0.elapsed()); assert_eq!(decoded.len(), n); - decoded - .with_iterator(|iter| { - for (i, got) in iter.enumerate() { - let want = strings[i].as_bytes(); - assert_eq!(got, Some(want), "row {} mismatch", i); - } - Ok::<_, vortex_error::VortexError>(()) - }) - .unwrap(); + let mask = decoded.validity()?.execute_mask(decoded.len(), &mut ctx)?; + for i in 0..decoded.len() { + let got = mask.value(i).then(|| decoded.bytes_at(i)); + let want = strings[i].as_bytes(); + assert_eq!(got.as_deref(), Some(want), "row {} mismatch", i); + } eprintln!("roundtrip OK on all {} rows", n); // Equality pushdown: pick a specific row's value and ensure the kernel @@ -117,16 +107,11 @@ fn smoke_100k_rows() { .binary( ConstantArray::new(needle.as_str(), n).into_array(), Operator::Eq, - ) - .unwrap() - .execute::(&mut ctx) - .unwrap() + )? + .execute::(&mut ctx)? .into_array(); - let eq_count = sum(&eq, &mut ctx) - .unwrap() - .as_primitive() - .as_::() - .unwrap(); + let eq_count = sum(&eq, &mut ctx)?.as_primitive().as_::().unwrap(); assert_eq!(eq_count, want_eq); eprintln!("eq pushdown matches reference count ({})", want_eq); + Ok(()) } diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index e0aa5cbb724..4a187f8faae 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -13,6 +13,9 @@ repository = { workspace = true } rust-version = { workspace = true } version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true diff --git a/encodings/fastlanes/benches/bitpack_compare.rs b/encodings/fastlanes/benches/bitpack_compare.rs index c704f2b35ed..739fcd72dc6 100644 --- a/encodings/fastlanes/benches/bitpack_compare.rs +++ b/encodings/fastlanes/benches/bitpack_compare.rs @@ -23,9 +23,11 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar_fn::fns::operators::Operator; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; +use vortex_buffer::Alignment; use vortex_buffer::BufferMut; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArray; use vortex_fastlanes::BitPackedData; use vortex_session::VortexSession; @@ -33,23 +35,50 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); const LENS: &[usize] = &[1024, 64 * 1024]; const BIT_WIDTHS: &[u8] = &[4, 16]; +/// Rebuild the array with its packed buffer copied to a page-aligned allocation. +/// +/// `bitpack_encode` aligns the packed buffer only to the element type, so its cache-line +/// placement depends on allocator state, which shifts with any change to the bench binary. +/// CodSpeed's simulated cache misses are deterministic in those addresses, which made these +/// benches flip ~30% between two layout modes across unrelated commits. Pinning the buffer to a +/// page boundary makes the layout, and therefore the measurement, reproducible. +fn page_aligned(array: BitPackedArray) -> BitPackedArray { + let ptype = array.dtype().as_ptype(); + let parts = BitPacked::into_parts(array); + BitPacked::try_new( + parts.packed.ensure_aligned(Alignment::new(4096)).unwrap(), + ptype, + parts.validity, + parts.patches, + parts.bit_width, + parts.len, + parts.offset, + ) + .unwrap() +} + /// Build a packed array of varied in-range values, plus an out-of-range constant RHS for /// the fast-path benches. fn build_inputs(len: usize) -> (ArrayRef, ArrayRef, ExecutionCtx) { let mut ctx = SESSION.create_execution_ctx(); let buf: BufferMut = (0..len).map(|i| (i as u32) % (1 << BW)).collect(); - let array = BitPackedData::encode( - &PrimitiveArray::new(buf.freeze(), Validity::NonNullable).into_array(), - BW, - &mut ctx, + let array = page_aligned( + BitPackedData::encode( + &PrimitiveArray::new(buf.freeze(), Validity::NonNullable).into_array(), + BW, + &mut ctx, + ) + .unwrap(), ) - .unwrap() .into_array(); // 1 << BW is just past the packable range, so the out-of-range fast path fires. let constant = 1u32 << BW; diff --git a/encodings/fastlanes/benches/bitpack_compare_sweep.rs b/encodings/fastlanes/benches/bitpack_compare_sweep.rs index 0be525566e7..ec7cf9b6892 100644 --- a/encodings/fastlanes/benches/bitpack_compare_sweep.rs +++ b/encodings/fastlanes/benches/bitpack_compare_sweep.rs @@ -27,9 +27,11 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; +use vortex_buffer::Alignment; use vortex_buffer::BufferMut; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArray; use vortex_fastlanes::BitPackedData; use vortex_session::VortexSession; @@ -37,8 +39,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); /// Number of elements per benchmarked array (64 full FastLanes blocks). const LEN: usize = 64 * 1024; @@ -65,6 +70,28 @@ macro_rules! impl_bench_int { impl_bench_int!(u8, u16, u32, u64, i8, i16, i32, i64); +/// Rebuild the array with its packed buffer copied to a page-aligned allocation. +/// +/// `bitpack_encode` aligns the packed buffer only to the element type, so its cache-line +/// placement depends on allocator state, which shifts with any change to the bench binary. +/// CodSpeed's simulated cache misses are deterministic in those addresses, which made the whole +/// sweep flip ~40% between two layout modes across unrelated commits. Pinning the buffer to a +/// page boundary makes the layout, and therefore the measurement, reproducible. +fn page_aligned(array: BitPackedArray) -> BitPackedArray { + let ptype = array.dtype().as_ptype(); + let parts = BitPacked::into_parts(array); + BitPacked::try_new( + parts.packed.ensure_aligned(Alignment::new(4096)).unwrap(), + ptype, + parts.validity, + parts.patches, + parts.bit_width, + parts.len, + parts.offset, + ) + .unwrap() +} + /// Encode `LEN` in-range values of type `T` at the given bit width, returning the packed array, a /// mid-range constant to compare against, and an execution context. fn setup(width: usize) -> (ArrayRef, ArrayRef, ExecutionCtx) { @@ -73,12 +100,14 @@ fn setup(width: usize) -> (ArrayRef, ArrayRef, ExecutionCtx) { let buf: BufferMut = (0..LEN) .map(|i| T::from_counter((i as u64) % cap)) .collect(); - let array = BitPackedData::encode( - &PrimitiveArray::new(buf.freeze(), Validity::NonNullable).into_array(), - width as u8, - &mut ctx, + let array = page_aligned( + BitPackedData::encode( + &PrimitiveArray::new(buf.freeze(), Validity::NonNullable).into_array(), + width as u8, + &mut ctx, + ) + .unwrap(), ) - .unwrap() .into_array(); let rhs = ConstantArray::new(T::from_counter(cap / 2), LEN).into_array(); (array, rhs, ctx) diff --git a/encodings/fastlanes/benches/bitpacking_take.rs b/encodings/fastlanes/benches/bitpacking_take.rs index 2e01f1253e6..eb072017ae3 100644 --- a/encodings/fastlanes/benches/bitpacking_take.rs +++ b/encodings/fastlanes/benches/bitpacking_take.rs @@ -15,7 +15,6 @@ use vortex_array::IntoArray as _; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; @@ -27,8 +26,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); #[divan::bench] fn take_10_stratified(bencher: Bencher) { diff --git a/encodings/fastlanes/benches/canonicalize_bench.rs b/encodings/fastlanes/benches/canonicalize_bench.rs index 7b0ee2453f5..30ee0756960 100644 --- a/encodings/fastlanes/benches/canonicalize_bench.rs +++ b/encodings/fastlanes/benches/canonicalize_bench.rs @@ -12,7 +12,6 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; use vortex_array::builders::ArrayBuilder; use vortex_array::builders::PrimitiveBuilder; -use vortex_array::session::ArraySession; use vortex_error::VortexExpect; use vortex_fastlanes::bitpack_compress::test_harness::make_array; use vortex_session::VortexSession; @@ -35,8 +34,11 @@ const BENCH_ARGS: &[(usize, usize, f64)] = &[ (10000, 1000, 0.00), ]; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); #[cfg(not(codspeed))] #[divan::bench(args = BENCH_ARGS)] diff --git a/encodings/fastlanes/benches/cast_bitpacked.rs b/encodings/fastlanes/benches/cast_bitpacked.rs index 097230930a3..baeff20f5e1 100644 --- a/encodings/fastlanes/benches/cast_bitpacked.rs +++ b/encodings/fastlanes/benches/cast_bitpacked.rs @@ -27,7 +27,6 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; @@ -39,8 +38,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); const U32: DType = DType::Primitive(PType::U32, Nullability::NonNullable); diff --git a/encodings/fastlanes/benches/compute_between.rs b/encodings/fastlanes/benches/compute_between.rs index 5a1bdcee655..2b295157e68 100644 --- a/encodings/fastlanes/benches/compute_between.rs +++ b/encodings/fastlanes/benches/compute_between.rs @@ -16,7 +16,6 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::NativePType; -use vortex_array::session::ArraySession; use vortex_error::VortexExpect; use vortex_fastlanes::bitpack_compress::bitpack_to_best_bit_width; use vortex_session::VortexSession; @@ -25,8 +24,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); fn generate_primitive_array( rng: &mut StdRng, diff --git a/encodings/fastlanes/src/bit_transpose/validity.rs b/encodings/fastlanes/src/bit_transpose/validity.rs index 826e1a2ad35..a6b2ec43cbc 100644 --- a/encodings/fastlanes/src/bit_transpose/validity.rs +++ b/encodings/fastlanes/src/bit_transpose/validity.rs @@ -107,8 +107,10 @@ fn bits_op_with_copy( 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>(); - // We can ignore the spare trailer capacity that can be an artifact of allocator as we requested 128 multiple chunks - let (output_chunks, _) = output.spare_capacity_mut().as_chunks_mut::<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 { @@ -189,4 +191,25 @@ mod tests { 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/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index eba5be1b749..f7a3485c113 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -437,7 +437,6 @@ mod test { use vortex_array::assert_arrays_eq; use vortex_array::builders::ArrayBuilder; use vortex_array::builders::PrimitiveBuilder; - use vortex_array::session::ArraySession; use vortex_buffer::Buffer; use vortex_error::VortexError; use vortex_error::vortex_err; @@ -448,8 +447,11 @@ mod test { use crate::bitpack_compress::test_harness::make_array; use crate::bitpacking::array::BitPackedArrayExt; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_best_bit_width() { @@ -515,14 +517,14 @@ mod test { chunked.append_to_builder(&mut primitive_builder, &mut ctx)?; let ca_into = primitive_builder.finish(); - assert_arrays_eq!(into_ca, ca_into); + assert_arrays_eq!(into_ca, ca_into, &mut ctx); let mut primitive_builder = PrimitiveBuilder::::with_capacity(chunked.dtype().nullability(), 10 * 100); primitive_builder.extend_from_array(&chunked); let ca_into = primitive_builder.finish(); - assert_arrays_eq!(into_ca, ca_into); + assert_arrays_eq!(into_ca, ca_into, &mut ctx); Ok(()) } @@ -555,7 +557,11 @@ mod test { // chunk 1 (1024-2047): no patches -> points to patch index 2 // chunk 2 (2048-3071): patch at 3000 -> starts at patch index 2 // chunk 3 (3072-4095): patch at 3100 -> starts at patch index 3 - assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64, 2, 2, 3])); + assert_arrays_eq!( + chunk_offsets, + PrimitiveArray::from_iter([0u64, 2, 2, 3]), + &mut ctx + ); Ok(()) } @@ -583,7 +589,11 @@ mod test { .clone() .execute::(&mut ctx)?; - assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64, 2, 2])); + assert_arrays_eq!( + chunk_offsets, + PrimitiveArray::from_iter([0u64, 2, 2]), + &mut ctx + ); Ok(()) } @@ -616,7 +626,11 @@ mod test { // chunk 2 (2048-3071): no patches -> points to patch index 3 // chunk 3 (3072-4095): no patches -> points to patch index 3 (remaining chunks filled) // chunk 4 (4096-5119): no patches -> points to patch index 3 (remaining chunks filled) - assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64, 2, 3, 3, 3])); + assert_arrays_eq!( + chunk_offsets, + PrimitiveArray::from_iter([0u64, 2, 3, 3, 3]), + &mut ctx + ); Ok(()) } @@ -645,7 +659,7 @@ mod test { .execute::(&mut ctx)?; // Single chunk starting at patch index 0. - assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64])); + assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64]), &mut ctx); Ok(()) } } diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index a4f8224966a..78aac1aa86e 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -210,7 +210,6 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::dtype::Nullability; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -226,8 +225,11 @@ mod tests { bitpack_encode(array, bit_width, None, &mut SESSION.create_execution_ctx()).unwrap() } - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); fn unpack(bitpacked: &BitPackedArray) -> VortexResult { unpack_array(bitpacked.as_view(), &mut SESSION.create_execution_ctx()) @@ -237,7 +239,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let values = PrimitiveArray::from_iter((0..n).map(|i| (i % 2047) as u16)); let compressed = BitPackedData::encode(&values.clone().into_array(), 11, &mut ctx).unwrap(); - assert_arrays_eq!(compressed, values); + assert_arrays_eq!(compressed, values, &mut ctx); values .as_slice::() @@ -272,7 +274,7 @@ mod tests { .execute::(&mut ctx)?; let bitpacked = encode(&zeros, 0); let actual = unpack(&bitpacked)?; - assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 0, 0, 0])); + assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 0, 0, 0]), &mut ctx); Ok(()) } @@ -284,7 +286,7 @@ mod tests { .execute::(&mut ctx)?; let bitpacked = encode(&zeros, 0); let actual = unpack(&bitpacked)?; - assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 1, 0, 1])); + assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 1, 0, 1]), &mut ctx); Ok(()) } @@ -296,7 +298,7 @@ mod tests { .execute::(&mut ctx)?; let bitpacked = encode(&zeros, 10); let actual = unpack(&bitpacked)?; - assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1024)); + assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1024), &mut ctx); Ok(()) } @@ -311,7 +313,8 @@ mod tests { let actual = unpack(&bitpacked)?; assert_arrays_eq!( actual, - PrimitiveArray::from_iter((5u16..1029).chain(5u16..1029).chain(5u16..1029)) + PrimitiveArray::from_iter((5u16..1029).chain(5u16..1029).chain(5u16..1029)), + &mut ctx ); Ok(()) } @@ -325,7 +328,7 @@ mod tests { let bitpacked = encode(&zeros, 11); assert!(bitpacked.patches().is_none()); let actual = unpack(&bitpacked)?; - assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1025)); + assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1025), &mut ctx); Ok(()) } @@ -339,7 +342,7 @@ mod tests { assert_eq!(bitpacked.len(), 1025); assert!(bitpacked.patches().is_some()); let actual = unpack(&bitpacked)?; - assert_arrays_eq!(actual, PrimitiveArray::from_iter(512u16..1537)); + assert_arrays_eq!(actual, PrimitiveArray::from_iter(512u16..1537), &mut ctx); Ok(()) } @@ -354,7 +357,7 @@ mod tests { assert!(bitpacked.patches().is_some()); let slice_ref = bitpacked.into_array().slice(1023..1025)?; let actual = slice_ref.execute::(&mut ctx)?.into_primitive(); - assert_arrays_eq!(actual, PrimitiveArray::from_iter([1535u16, 1536])); + assert_arrays_eq!(actual, PrimitiveArray::from_iter([1535u16, 1536]), &mut ctx); Ok(()) } @@ -371,7 +374,8 @@ mod tests { let actual = slice_ref.execute::(&mut ctx)?.into_primitive(); assert_arrays_eq!( actual, - PrimitiveArray::from_iter((1023u16..2049).map(|x| x + 512)) + PrimitiveArray::from_iter((1023u16..2049).map(|x| x + 512)), + &mut ctx ); Ok(()) } @@ -459,7 +463,11 @@ mod tests { let result = builder.finish_into_primitive(); // Verify all values were correctly unpacked including patches. - assert_arrays_eq!(result, PrimitiveArray::from_iter(values)); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter(values), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -654,7 +662,11 @@ mod tests { // Verify consistency with unpack_array. let zeros_array = unpack(&zeros_bp)?; assert_eq!(zeros_result.len(), zeros_array.len()); - assert_arrays_eq!(zeros_result, zeros_array); + assert_arrays_eq!( + zeros_result, + zeros_array, + &mut SESSION.create_execution_ctx() + ); // Maximum bit width for u16 (15 bits, since bitpacking requires bit_width < type bit width). let max_values = PrimitiveArray::from_iter([32767u16; 50]); // 2^15 - 1 @@ -681,7 +693,11 @@ mod tests { // Verify consistency. let boundary_unpacked = unpack(&boundary_bp)?; assert_eq!(boundary_result.len(), boundary_unpacked.len()); - assert_arrays_eq!(boundary_result, boundary_unpacked); + assert_arrays_eq!( + boundary_result, + boundary_unpacked, + &mut SESSION.create_execution_ctx() + ); // Single element. let single = PrimitiveArray::from_iter([42u8]); diff --git a/encodings/fastlanes/src/bitpacking/array/mod.rs b/encodings/fastlanes/src/bitpacking/array/mod.rs index e5c64252fbc..772de2ac641 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -327,15 +327,17 @@ mod test { use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; - use vortex_array::session::ArraySession; use vortex_buffer::Buffer; use vortex_session::VortexSession; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_encode() { @@ -357,7 +359,7 @@ mod test { .clone() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(packed_primitive, expected); + assert_arrays_eq!(packed_primitive, expected, &mut ctx); } #[test] @@ -386,7 +388,8 @@ mod test { .unwrap(); assert_arrays_eq!( packed_primitive, - PrimitiveArray::new(values, vortex_array::validity::Validity::NonNullable) + PrimitiveArray::new(values, vortex_array::validity::Validity::NonNullable), + &mut ctx ); } } diff --git a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs index c9158d8bdac..3c77e146ad0 100644 --- a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs +++ b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs @@ -62,7 +62,7 @@ impl> UnpackStrategy for BitPackingStr /// use vortex_fastlanes::BitPackedArrayExt; /// use vortex_fastlanes::unpack_iter::BitUnpackedChunks; /// -/// let mut ctx = vortex_array::LEGACY_SESSION.create_execution_ctx(); +/// 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(); /// @@ -79,7 +79,6 @@ impl> UnpackStrategy for BitPackingStr /// // handle partial trailing chunk /// } /// ``` -/// pub struct UnpackedChunks> { strategy: S, bit_width: usize, diff --git a/encodings/fastlanes/src/bitpacking/compute/between.rs b/encodings/fastlanes/src/bitpacking/compute/between.rs index dcc372ad091..1dd3b61dbf4 100644 --- a/encodings/fastlanes/src/bitpacking/compute/between.rs +++ b/encodings/fastlanes/src/bitpacking/compute/between.rs @@ -142,9 +142,10 @@ where #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; @@ -156,10 +157,17 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::BitPackedArrayExt; use crate::BitPackedData; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn opts(lower: StrictComparison, upper: StrictComparison) -> BetweenOptions { BetweenOptions { lower_strict: lower, @@ -176,7 +184,7 @@ mod tests { #[case] lower_strict: StrictComparison, #[case] upper_strict: StrictComparison, ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: BufferMut = (0..3000u32).map(|i| i % 257).collect(); let prim = PrimitiveArray::new(values.freeze(), Validity::NonNullable); let packed = BitPackedData::encode(&prim.clone().into_array(), 9, &mut ctx)?; @@ -194,13 +202,13 @@ mod tests { .between(lower, upper, options)? .execute::(&mut ctx)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } #[test] fn signed_with_patches_against_primitive_baseline() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Vec = (0..1500) .map(|i| if i % 73 == 0 { 100_000 + i } else { i % 100 }) .collect(); @@ -221,13 +229,13 @@ mod tests { .between(lower, upper, options)? .execute::(&mut ctx)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } #[test] fn nullable_propagates_validity() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let prim = PrimitiveArray::from_option_iter([Some(1u32), None, Some(3), Some(4), None, Some(6)]); let packed = BitPackedData::encode(&prim.clone().into_array(), 3, &mut ctx)?; @@ -244,7 +252,7 @@ mod tests { .into_array() .between(lower, upper, options)? .execute::(&mut ctx)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } } diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index 10060eb57e2..3b917aa36af 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -112,10 +112,11 @@ impl CastKernel for BitPacked { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -128,12 +129,19 @@ mod tests { use vortex_array::match_each_integer_ptype; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::BitPackedArray; use crate::BitPackedData; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn bp(array: &ArrayRef, bit_width: u8) -> BitPackedArray { - BitPackedData::encode(array, bit_width, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() + BitPackedData::encode(array, bit_width, &mut SESSION.create_execution_ctx()).unwrap() } #[test] @@ -151,7 +159,8 @@ mod tests { assert_arrays_eq!( casted, - PrimitiveArray::from_iter([10u32, 20, 30, 40, 50, 60]) + PrimitiveArray::from_iter([10u32, 20, 30, 40, 50, 60]), + &mut SESSION.create_execution_ctx() ); } @@ -213,7 +222,7 @@ mod tests { let source = match_each_integer_ptype!(src, |S| { values::(len) }); let source_ref = source.into_array(); let target = DType::Primitive(tgt, Nullability::NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Reference: plain primitive cast of the same values. let reference = source_ref @@ -227,7 +236,7 @@ mod tests { let casted = packed .cast(target.clone())? .execute::(&mut ctx)?; - assert_arrays_eq!(casted, reference); + assert_arrays_eq!(casted, reference, &mut ctx); // Also exercise the sliced/offset path (offset > 0, trailer present). if len >= 4 { @@ -242,7 +251,7 @@ mod tests { .slice(lo..hi)? .cast(target.clone())? .execute::(&mut ctx)?; - assert_arrays_eq!(casted, reference); + assert_arrays_eq!(casted, reference, &mut ctx); } } } diff --git a/encodings/fastlanes/src/bitpacking/compute/compare.rs b/encodings/fastlanes/src/bitpacking/compute/compare.rs index e622e64a665..51aca62a9aa 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare.rs @@ -118,7 +118,6 @@ mod tests { 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::session::ArraySession; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -126,8 +125,11 @@ mod tests { use crate::BitPackedArrayExt; use crate::BitPackedData; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); /// All six operators on a small in-range input. #[rstest] @@ -148,7 +150,7 @@ mod tests { .unwrap() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter(expected)); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); } /// Sweep every native int type across several bit-widths. 2048 elements spans two @@ -179,7 +181,7 @@ mod tests { .into_array() .binary(rhs.clone(), Operator::from(op))? .execute::(&mut ctx)?; - assert_arrays_eq!(got, want); + assert_arrays_eq!(got, want, &mut SESSION.create_execution_ctx()); } } Ok(()) @@ -217,7 +219,7 @@ mod tests { .into_array() .binary(rhs, Operator::Eq)? .execute::(&mut ctx)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } @@ -261,7 +263,7 @@ mod tests { .slice(start..start + slice_len)? .binary(rhs.clone(), Operator::from(op))? .execute::(&mut ctx)?; - assert_arrays_eq!(got, want); + assert_arrays_eq!(got, want, &mut ctx); } Ok(()) } @@ -298,7 +300,7 @@ mod tests { .slice(start..end)? .binary(rhs, Operator::Eq)? .execute::(&mut ctx)?; - assert_arrays_eq!(got, want); + assert_arrays_eq!(got, want, &mut ctx); Ok(()) } @@ -317,7 +319,7 @@ mod tests { .into_array() .binary(rhs, Operator::Eq)? .execute::(&mut ctx)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } } diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 284ae5bce4c..21184d785a5 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -177,8 +177,9 @@ fn filter_with_indices( #[cfg(test)] mod test { + use std::sync::LazyLock; + use vortex_array::IntoArray as _; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -187,13 +188,20 @@ mod test { use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn take_indices() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a u8 array modulo 63. let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); let bitpacked = BitPackedData::encode(&unpacked.into_array(), 6, &mut ctx).unwrap(); @@ -203,13 +211,14 @@ mod test { let primitive_result = bitpacked.filter(mask).unwrap(); assert_arrays_eq!( primitive_result, - PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]) + PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]), + &mut ctx ); } #[test] fn take_sliced_indices() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a u8 array modulo 63. let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); let bitpacked = BitPackedData::encode(&unpacked.into_array(), 6, &mut ctx).unwrap(); @@ -218,25 +227,30 @@ mod test { let mask = Mask::from_indices(sliced.len(), vec![1919, 1921]); let primitive_result = sliced.filter(mask).unwrap(); - assert_arrays_eq!(primitive_result, PrimitiveArray::from_iter([31u8, 33])); + assert_arrays_eq!( + primitive_result, + PrimitiveArray::from_iter([31u8, 33]), + &mut ctx + ); } #[test] fn filter_bitpacked() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); let bitpacked = BitPackedData::encode(&unpacked.into_array(), 6, &mut ctx).unwrap(); let filtered = bitpacked.filter(Mask::from_indices(4096, 0..1024)).unwrap(); let filtered_prim = filtered.execute::(&mut ctx).unwrap(); assert_arrays_eq!( filtered_prim, - PrimitiveArray::from_iter((0..1024).map(|i| (i % 63) as u8)) + PrimitiveArray::from_iter((0..1024).map(|i| (i % 63) as u8)), + &mut ctx ); } #[test] fn filter_bitpacked_signed() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Buffer = (0..500).collect(); let unpacked = PrimitiveArray::new(values.clone(), Validity::NonNullable); let bitpacked = BitPackedData::encode(&unpacked.into_array(), 9, &mut ctx).unwrap(); @@ -248,13 +262,14 @@ mod test { assert_arrays_eq!( filtered, - PrimitiveArray::from_iter(values[0..250].iter().copied()) + PrimitiveArray::from_iter(values[0..250].iter().copied()), + &mut ctx ); } #[test] fn test_filter_bitpacked_conformance() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test with u8 values let unpacked = buffer![1u8, 2, 3, 4, 5].into_array(); let bitpacked = BitPackedData::encode(&unpacked, 3, &mut ctx).unwrap(); @@ -278,7 +293,7 @@ mod test { /// This test ensures that the type handling is correct. #[test] fn filter_bitpacked_signed_with_patches() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create signed integer values where some exceed the bit width (causing patches). // Values 0-127 fit in 7 bits, but 1000 and 2000 do not. let values: Vec = vec![0, 10, 1000, 20, 30, 2000, 40, 50, 60, 70]; @@ -296,7 +311,11 @@ mod test { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(filtered, PrimitiveArray::from_iter([0i32, 1000, 2000, 70])); + assert_arrays_eq!( + filtered, + PrimitiveArray::from_iter([0i32, 1000, 2000, 70]), + &mut ctx + ); } /// Regression test for signed integers with patches using low selectivity. @@ -305,7 +324,7 @@ mod test { /// that doesn't fully decompress the array first. #[test] fn filter_bitpacked_signed_with_patches_low_selectivity() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a larger array with signed integers and some patches. let values: Vec = (0..1000) .map(|i| { @@ -332,6 +351,6 @@ mod test { .unwrap(); let expected: Vec = values[0..20].to_vec(); - assert_arrays_eq!(filtered, PrimitiveArray::from_iter(expected)); + assert_arrays_eq!(filtered, PrimitiveArray::from_iter(expected), &mut ctx); } } diff --git a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs index 0a743f9269b..3b4e2b9770a 100644 --- a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs +++ b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs @@ -185,9 +185,9 @@ fn apply_patches_idx_typed( #[cfg(test)] mod tests { use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::is_constant::is_constant; + use vortex_array::array_session; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -195,7 +195,7 @@ mod tests { #[test] fn is_constant_with_patches() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = BitPackedData::encode(&buffer![4; 1025].into_array(), 2, &mut ctx)?; assert!(is_constant(&array.into_array(), &mut ctx)?); Ok(()) diff --git a/encodings/fastlanes/src/bitpacking/compute/mod.rs b/encodings/fastlanes/src/bitpacking/compute/mod.rs index 06a4b4597b0..38f86f781bb 100644 --- a/encodings/fastlanes/src/bitpacking/compute/mod.rs +++ b/encodings/fastlanes/src/bitpacking/compute/mod.rs @@ -46,8 +46,8 @@ fn chunked_indices( mod tests { use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; 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; @@ -61,7 +61,7 @@ mod tests { array, bit_width, None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() } @@ -98,7 +98,8 @@ mod tests { #[case::alternating_bits(bp(&PrimitiveArray::from_iter([0u16, 255, 0, 255, 0, 255]), 8))] fn test_bitpacked_consistency(#[case] array: BitPackedArray) { - test_array_consistency(&array.into_array()); + let ctx = &mut array_session().create_execution_ctx(); + test_array_consistency(&array.into_array(), ctx); } #[rstest] @@ -109,6 +110,9 @@ mod tests { #[case::i32_basic(bp(&PrimitiveArray::from_iter([10i32, 20, 30, 40, 50]), 7))] #[case::large_u32(bp(&PrimitiveArray::from_iter((0..100).map(|i| i as u32)), 7))] fn test_bitpacked_binary_numeric(#[case] array: BitPackedArray) { - test_binary_numeric_array(array.into_array()); + test_binary_numeric_array( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/encodings/fastlanes/src/bitpacking/compute/slice.rs b/encodings/fastlanes/src/bitpacking/compute/slice.rs index c019cccd003..996565a2672 100644 --- a/encodings/fastlanes/src/bitpacking/compute/slice.rs +++ b/encodings/fastlanes/src/bitpacking/compute/slice.rs @@ -72,8 +72,8 @@ fn slice_bitpacked( #[cfg(test)] mod tests { use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::SliceArray; use vortex_error::VortexResult; @@ -83,7 +83,7 @@ mod tests { #[test] fn test_reduce_parent_returns_bitpacked_slice() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::from_iter(0u32..2048); let bitpacked = bitpack_encode(&values, 11, None, &mut ctx)?; diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index bbf5bfaf2c8..34aa5a0c3f5 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -150,8 +150,7 @@ fn take_primitive( if let Some(patches) = array.patches() && let Some(patches) = patches.take(&indices.clone().into_array(), ctx)? { - let cast_patches = patches.cast_values(unpatched_taken.dtype())?; - return unpatched_taken.patch(&cast_patches, ctx); + return unpatched_taken.patch(&patches, ctx); } Ok(unpatched_taken) @@ -160,27 +159,35 @@ fn take_primitive( #[cfg(test)] #[expect(clippy::cast_possible_truncation)] mod test { + use std::sync::LazyLock; + use rand::RngExt; use rand::distr::Uniform; use rand::rng; use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; 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_buffer::buffer; + use vortex_session::VortexSession; use crate::BitPackedArray; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; use crate::bitpacking::compute::take::take_primitive; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn take_indices() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![0, 125, 2047, 2049, 2151, 2790].into_array(); // Create a u8 array modulo 63. @@ -190,25 +197,30 @@ mod test { let primitive_result = bitpacked.take(indices).unwrap(); assert_arrays_eq!( primitive_result, - PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]) + PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]), + &mut ctx ); } #[test] fn take_with_patches() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let unpacked = Buffer::from_iter(0u32..1024).into_array(); let bitpacked = BitPackedData::encode(&unpacked, 2, &mut ctx).unwrap(); let indices = buffer![0, 2, 4, 6].into_array(); let primitive_result = bitpacked.take(indices).unwrap(); - assert_arrays_eq!(primitive_result, PrimitiveArray::from_iter([0u32, 2, 4, 6])); + assert_arrays_eq!( + primitive_result, + PrimitiveArray::from_iter([0u32, 2, 4, 6]), + &mut ctx + ); } #[test] fn take_sliced_indices() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![1919, 1921].into_array(); // Create a u8 array modulo 63. @@ -217,13 +229,17 @@ mod test { let sliced = bitpacked.slice(128..2050).unwrap(); let primitive_result = sliced.take(indices).unwrap(); - assert_arrays_eq!(primitive_result, PrimitiveArray::from_iter([31u8, 33])); + assert_arrays_eq!( + primitive_result, + PrimitiveArray::from_iter([31u8, 33]), + &mut ctx + ); } #[test] #[cfg_attr(miri, ignore)] // This test is too slow on miri fn take_random_indices() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let num_patches: usize = 128; let values = (0..u16::MAX as u32 + num_patches as u32).collect::>(); let uncompressed = PrimitiveArray::new(values.clone(), Validity::NonNullable); @@ -256,7 +272,7 @@ mod test { #[test] #[cfg_attr(miri, ignore)] fn take_signed_with_patches() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let start = BitPackedData::encode(&buffer![1i32, 2i32, 3i32, 4i32].into_array(), 1, &mut ctx) .unwrap(); @@ -268,12 +284,16 @@ mod test { &mut ctx, ) .unwrap(); - assert_arrays_eq!(taken_primitive, PrimitiveArray::from_iter([1i32, 2, 3, 4])); + assert_arrays_eq!( + taken_primitive, + PrimitiveArray::from_iter([1i32, 2, 3, 4]), + &mut ctx + ); } #[test] fn take_nullable_with_nullables() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let start = BitPackedData::encode(&buffer![1i32, 2i32, 3i32, 4i32].into_array(), 1, &mut ctx) .unwrap(); @@ -285,19 +305,15 @@ mod test { .unwrap(); assert_arrays_eq!( taken_primitive, - PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4)]) + PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4)]), + &mut ctx ); let taken_primitive_prim = taken_primitive.execute::(&mut ctx).unwrap(); assert_eq!(taken_primitive_prim.invalid_count(&mut ctx).unwrap(), 1); } fn bp(array: vortex_array::ArrayRef, bit_width: u8) -> BitPackedArray { - BitPackedData::encode( - &array, - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap() + BitPackedData::encode(&array, bit_width, &mut SESSION.create_execution_ctx()).unwrap() } #[rstest] diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index bc2304af657..efa0677a91e 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -19,3 +19,7 @@ mod vtable; pub(crate) use plugin::BitPackedPatchedPlugin; pub use vtable::BitPacked; pub use vtable::BitPackedArray; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} diff --git a/encodings/fastlanes/src/bitpacking/plugin.rs b/encodings/fastlanes/src/bitpacking/plugin.rs index d5ecdefd0c7..a621d085514 100644 --- a/encodings/fastlanes/src/bitpacking/plugin.rs +++ b/encodings/fastlanes/src/bitpacking/plugin.rs @@ -100,7 +100,6 @@ mod tests { use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::patched::PatchedArraySlotsExt; use vortex_array::buffer::BufferHandle; - use vortex_array::session::ArraySession; use vortex_array::session::ArraySessionExt; use vortex_buffer::Buffer; use vortex_error::VortexResult; @@ -114,7 +113,7 @@ mod tests { use crate::BitPackedData; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); session.arrays().register(BitPackedPatchedPlugin); session }); diff --git a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs index 87332f736a7..eb0dd9b7a23 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs @@ -1,21 +1,39 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; +use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::between::Between; use vortex_array::scalar_fn::fns::between::BetweenExecuteAdaptor; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; +use vortex_session::VortexSession; use crate::BitPacked; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&BetweenExecuteAdaptor(BitPacked)), - ParentKernelSet::lift(&CastExecuteAdaptor(BitPacked)), - ParentKernelSet::lift(&CompareExecuteAdaptor(BitPacked)), - ParentKernelSet::lift(&FilterExecuteAdaptor(BitPacked)), - ParentKernelSet::lift(&SliceExecuteAdaptor(BitPacked)), - ParentKernelSet::lift(&TakeExecuteAdaptor(BitPacked)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Between.id(), + BitPacked, + BetweenExecuteAdaptor(BitPacked), + ); + kernels.register_execute_parent_kernel(Cast.id(), BitPacked, CastExecuteAdaptor(BitPacked)); + kernels.register_execute_parent_kernel( + Binary.id(), + BitPacked, + CompareExecuteAdaptor(BitPacked), + ); + kernels.register_execute_parent_kernel(Filter.id(), BitPacked, FilterExecuteAdaptor(BitPacked)); + kernels.register_execute_parent_kernel(Slice.id(), BitPacked, SliceExecuteAdaptor(BitPacked)); + kernels.register_execute_parent_kernel(Dict.id(), BitPacked, TakeExecuteAdaptor(BitPacked)); +} diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 52b536e02c5..68fbf1b41d3 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -35,6 +35,7 @@ use vortex_array::vtable::validity_to_child; 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_error::vortex_panic; use vortex_session::VortexSession; @@ -48,7 +49,6 @@ use crate::bitpack_decompress::unpack_into_primitive_builder; use crate::bitpacking::array::BitPackedSlots; use crate::bitpacking::array::BitPackedSlotsView; use crate::bitpacking::array::PATCH_SLOTS; -use crate::bitpacking::vtable::kernels::PARENT_KERNELS; use crate::bitpacking::vtable::rules::RULES; mod kernels; mod operations; @@ -58,6 +58,10 @@ mod validity; /// A [`BitPacked`]-encoded Vortex array. pub type BitPackedArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernels::initialize(session); +} + #[derive(Clone, prost::Message)] pub struct BitPackedMetadata { #[prost(uint32, tag = "1")] @@ -138,6 +142,24 @@ impl VTable for BitPacked { } } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + let mut data = array.data().clone(); + data.packed = buffers[0].clone(); + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, @@ -268,15 +290,6 @@ impl VTable for BitPacked { )) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/encodings/fastlanes/src/bitpacking/vtable/operations.rs b/encodings/fastlanes/src/bitpacking/vtable/operations.rs index 4c277163719..e14b27323c1 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/operations.rs @@ -34,7 +34,6 @@ mod test { use vortex_array::ArrayRef; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::SliceArray; @@ -56,9 +55,10 @@ mod test { use crate::BitPackedArray; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; + use crate::test::SESSION; fn bp(array: &ArrayRef, bit_width: u8) -> BitPackedArray { - BitPackedData::encode(array, bit_width, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() + BitPackedData::encode(array, bit_width, &mut SESSION.create_execution_ctx()).unwrap() } fn slice_via_reduce(array: &BitPackedArray, range: Range) -> BitPackedArray { @@ -78,8 +78,13 @@ mod test { 6, ); let sliced = slice_via_reduce(&arr, 1024..2048); - assert_nth_scalar!(sliced, 0, 1024u32 % 64); - assert_nth_scalar!(sliced, 1023, 2047u32 % 64); + assert_nth_scalar!(sliced, 0, 1024u32 % 64, &mut SESSION.create_execution_ctx()); + assert_nth_scalar!( + sliced, + 1023, + 2047u32 % 64, + &mut SESSION.create_execution_ctx() + ); assert_eq!(sliced.offset(), 0); assert_eq!(sliced.len(), 1024); } @@ -91,8 +96,13 @@ mod test { 6, ); let sliced = slice_via_reduce(&arr, 512..1434); - assert_nth_scalar!(sliced, 0, 512u32 % 64); - assert_nth_scalar!(sliced, 921, 1433u32 % 64); + assert_nth_scalar!(sliced, 0, 512u32 % 64, &mut SESSION.create_execution_ctx()); + assert_nth_scalar!( + sliced, + 921, + 1433u32 % 64, + &mut SESSION.create_execution_ctx() + ); assert_eq!(sliced.offset(), 512); assert_eq!(sliced.len(), 922); } @@ -105,8 +115,18 @@ mod test { ); let compressed = packed.slice(768..9999).unwrap(); - assert_nth_scalar!(compressed, 0, (768 % 63) as u8); - assert_nth_scalar!(compressed, compressed.len() - 1, (9998 % 63) as u8); + assert_nth_scalar!( + compressed, + 0, + (768 % 63) as u8, + &mut SESSION.create_execution_ctx() + ); + assert_nth_scalar!( + compressed, + compressed.len() - 1, + (9998 % 63) as u8, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -117,8 +137,18 @@ mod test { ); let compressed = packed.slice(7168..9216).unwrap(); - assert_nth_scalar!(compressed, 0, (7168 % 63) as u8); - assert_nth_scalar!(compressed, compressed.len() - 1, (9215 % 63) as u8); + assert_nth_scalar!( + compressed, + 0, + (7168 % 63) as u8, + &mut SESSION.create_execution_ctx() + ); + assert_nth_scalar!( + compressed, + compressed.len() - 1, + (9215 % 63) as u8, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -128,20 +158,35 @@ mod test { 6, ); let sliced = slice_via_reduce(&arr, 512..1434); - assert_nth_scalar!(sliced, 0, 512u32 % 64); - assert_nth_scalar!(sliced, 921, 1433u32 % 64); + assert_nth_scalar!(sliced, 0, 512u32 % 64, &mut SESSION.create_execution_ctx()); + assert_nth_scalar!( + sliced, + 921, + 1433u32 % 64, + &mut SESSION.create_execution_ctx() + ); assert_eq!(sliced.offset(), 512); assert_eq!(sliced.len(), 922); let doubly_sliced = slice_via_reduce(&sliced, 127..911); - assert_nth_scalar!(doubly_sliced, 0, (512u32 + 127) % 64); - assert_nth_scalar!(doubly_sliced, 783, (512u32 + 910) % 64); + assert_nth_scalar!( + doubly_sliced, + 0, + (512u32 + 127) % 64, + &mut SESSION.create_execution_ctx() + ); + assert_nth_scalar!( + doubly_sliced, + 783, + (512u32 + 910) % 64, + &mut SESSION.create_execution_ctx() + ); assert_eq!(doubly_sliced.offset(), 639); assert_eq!(doubly_sliced.len(), 784); } #[test] fn slice_empty_patches() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // We create an array that has 1 element that does not fit in the 6-bit range. let array = BitPackedData::encode(&buffer![0u32..=64].into_array(), 6, &mut ctx).unwrap(); @@ -154,12 +199,14 @@ mod test { // reads buffers. The slice range 0..64 excludes the patch at index 64, so the // resulting array should have no patches. let array_ref = array.into_array(); - let slice_array = SliceArray::new(array_ref.clone(), 0..64); - let sliced = array_ref - .execute_parent(&slice_array.into_array(), 0, &mut ctx) - .expect("execute_parent failed") - .expect("expected slice kernel to execute"); - let sliced_bp = sliced.as_::().into_owned(); + let slice_array = SliceArray::new(array_ref, 0..64); + let mut ctx = SESSION.create_execution_ctx(); + let sliced_bp = slice_array + .into_array() + .execute::(&mut ctx) + .expect("slice execution failed") + .as_::() + .into_owned(); assert!(sliced_bp.patches().is_none()); } @@ -214,7 +261,7 @@ mod test { .into_array(); assert_eq!( packed_array - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut SESSION.create_execution_ctx()) .unwrap(), Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable)) ); @@ -222,7 +269,7 @@ mod test { #[test] fn scalar_at() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = (0u32..257).collect::>(); let uncompressed = values.clone().into_array(); let packed = BitPackedData::encode(&uncompressed, 8, &mut ctx).unwrap(); @@ -232,7 +279,7 @@ mod test { assert_eq!( usize::try_from( &patches - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut SESSION.create_execution_ctx()) .unwrap() ) .unwrap(), @@ -240,6 +287,6 @@ mod test { ); let expected = PrimitiveArray::from_iter(values.iter().copied()); - assert_arrays_eq!(packed, expected); + assert_arrays_eq!(packed, expected, &mut ctx); } } diff --git a/encodings/fastlanes/src/delta/array/delta_compress.rs b/encodings/fastlanes/src/delta/array/delta_compress.rs index e35778ad29e..7f00eb172dc 100644 --- a/encodings/fastlanes/src/delta/array/delta_compress.rs +++ b/encodings/fastlanes/src/delta/array/delta_compress.rs @@ -106,7 +106,6 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; - use vortex_array::session::ArraySession; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -116,8 +115,11 @@ mod tests { use crate::delta::array::delta_decompress::delta_decompress; use crate::delta_compress; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[rstest] #[case::u32((0u32..10_000).collect())] @@ -149,7 +151,7 @@ mod tests { let delta = Delta::try_from_primitive_array(&array, &mut SESSION.create_execution_ctx())?; assert_eq!(delta.len(), array.len()); let decompressed = delta_decompress(&delta, &mut SESSION.create_execution_ctx())?; - assert_arrays_eq!(decompressed, array); + assert_arrays_eq!(decompressed, array, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -175,7 +177,7 @@ mod tests { .as_array() .clone() .execute::(&mut ctx)?; - assert_arrays_eq!(packed_delta_prim, array); + assert_arrays_eq!(packed_delta_prim, array, &mut ctx); Ok(()) } diff --git a/encodings/fastlanes/src/delta/array/delta_decompress.rs b/encodings/fastlanes/src/delta/array/delta_decompress.rs index fe2567e63c7..7dcbeb44950 100644 --- a/encodings/fastlanes/src/delta/array/delta_decompress.rs +++ b/encodings/fastlanes/src/delta/array/delta_decompress.rs @@ -70,7 +70,9 @@ where // Allocate a result array. let mut output = BufferMut::with_capacity(deltas.len()); - let (output_chunks, _) = output.spare_capacity_mut().as_chunks_mut::<1024>(); + // Bound to the requested length: `spare_capacity_mut` may expose extra over-aligned capacity + // beyond `deltas.len()`, which would desync the `zip_eq` with `chunks` below and panic. + let (output_chunks, _) = output.spare_capacity_mut()[..deltas.len()].as_chunks_mut::<1024>(); // Loop over all the chunks let mut transposed: [T; 1024] = [T::default(); 1024]; diff --git a/encodings/fastlanes/src/delta/array/mod.rs b/encodings/fastlanes/src/delta/array/mod.rs index 33ece0deddd..7754bb37a59 100644 --- a/encodings/fastlanes/src/delta/array/mod.rs +++ b/encodings/fastlanes/src/delta/array/mod.rs @@ -38,7 +38,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["bases", "deltas"]; /// use vortex_session::VortexSession; /// use vortex_fastlanes::Delta; /// -/// let session = VortexSession::empty().with::(); +/// let session = vortex_array::array_session(); /// let primitive = PrimitiveArray::from_iter([1_u32, 2, 3, 5, 10, 11]); /// let array = Delta::try_from_primitive_array(&primitive, &mut session.create_execution_ctx()).unwrap(); /// ``` @@ -53,7 +53,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["bases", "deltas"]; /// use vortex_session::VortexSession; /// use vortex_fastlanes::Delta; /// -/// let session = VortexSession::empty().with::(); +/// let session = vortex_array::array_session(); /// let primitive = PrimitiveArray::from_iter([-3_i32, -2, -1, 0, 1, 2]); /// let array = Delta::try_from_primitive_array(&primitive, &mut session.create_execution_ctx()).unwrap(); /// ``` diff --git a/encodings/fastlanes/src/delta/compute/cast.rs b/encodings/fastlanes/src/delta/compute/cast.rs index f782eedd524..2ee10e7619a 100644 --- a/encodings/fastlanes/src/delta/compute/cast.rs +++ b/encodings/fastlanes/src/delta/compute/cast.rs @@ -51,14 +51,16 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::session::ArraySession; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::Delta; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_cast_delta_unsigned_widening_wraps() { @@ -77,7 +79,11 @@ mod tests { .cast(DType::Primitive(PType::U32, Nullability::NonNullable)) .unwrap(); - assert_arrays_eq!(casted, PrimitiveArray::from_iter([200u32, 50, 75, 10, 255])); + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([200u32, 50, 75, 10, 255]), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -97,7 +103,8 @@ mod tests { assert_arrays_eq!( casted, - PrimitiveArray::from_iter([10f32, 20.0, 30.0, 40.0, 50.0]) + PrimitiveArray::from_iter([10f32, 20.0, 30.0, 40.0, 50.0]), + &mut SESSION.create_execution_ctx() ); } @@ -117,7 +124,8 @@ mod tests { ); assert_arrays_eq!( casted, - PrimitiveArray::from_option_iter([Some(10u32), Some(20), Some(5), Some(30), Some(15)]) + PrimitiveArray::from_option_iter([Some(10u32), Some(20), Some(5), Some(30), Some(15)]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -136,7 +144,8 @@ mod tests { assert_arrays_eq!( casted, - PrimitiveArray::from_option_iter([Some(10u32), None, Some(30), Some(15), None]) + PrimitiveArray::from_option_iter([Some(10u32), None, Some(30), Some(15), None]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -179,7 +188,11 @@ mod tests { ); // Verify by decoding - assert_arrays_eq!(casted, PrimitiveArray::from_iter([10u32, 20, 30, 40, 50])); + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([10u32, 20, 30, 40, 50]), + &mut SESSION.create_execution_ctx() + ); } #[test] diff --git a/encodings/fastlanes/src/delta/vtable/mod.rs b/encodings/fastlanes/src/delta/vtable/mod.rs index dfb440620ab..6dcfedc0bc9 100644 --- a/encodings/fastlanes/src/delta/vtable/mod.rs +++ b/encodings/fastlanes/src/delta/vtable/mod.rs @@ -108,6 +108,14 @@ impl VTable for Delta { None } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/encodings/fastlanes/src/delta/vtable/operations.rs b/encodings/fastlanes/src/delta/vtable/operations.rs index 785cdf6aca8..7ed57a0886d 100644 --- a/encodings/fastlanes/src/delta/vtable/operations.rs +++ b/encodings/fastlanes/src/delta/vtable/operations.rs @@ -31,11 +31,11 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -44,8 +44,11 @@ mod tests { use crate::Delta; use crate::DeltaArray; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + crate::initialize(&session); + session + }); fn da(array: &PrimitiveArray) -> DeltaArray { Delta::try_from_primitive_array(array, &mut SESSION.create_execution_ctx()) @@ -58,7 +61,7 @@ mod tests { let actual = delta.slice(10..250).unwrap(); let expected = PrimitiveArray::from_iter(10u32..250).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -67,7 +70,7 @@ mod tests { let actual = delta.slice(1024 + 10..1024 + 250).unwrap(); let expected = PrimitiveArray::from_iter((1024 + 10u32)..(1024 + 250)).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -76,7 +79,7 @@ mod tests { let actual = delta.slice(1000..1048).unwrap(); let expected = PrimitiveArray::from_iter(1000u32..1048).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -85,7 +88,7 @@ mod tests { let actual = delta.slice(2040..2050).unwrap(); let expected = PrimitiveArray::from_iter(2040u32..2050).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -94,7 +97,7 @@ mod tests { let actual = delta.slice(0..4096).unwrap(); let expected = PrimitiveArray::from_iter(0u32..4096).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -103,15 +106,15 @@ mod tests { let actual = delta.slice(0..0).unwrap(); let expected = PrimitiveArray::from_iter(Vec::::new()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); let actual = delta.slice(4096..4096).unwrap(); let expected = PrimitiveArray::from_iter(Vec::::new()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); let actual = delta.slice(1024..1024).unwrap(); let expected = PrimitiveArray::from_iter(Vec::::new()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -120,7 +123,7 @@ mod tests { let actual = delta.slice(1024 + 10..1024 + 250).unwrap(); let expected = PrimitiveArray::from_iter((1024 + 10u32)..(1024 + 250)).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -129,15 +132,15 @@ mod tests { let actual = delta.slice(0..0).unwrap(); let expected = PrimitiveArray::from_iter(Vec::::new()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); let actual = delta.slice(4000..4000).unwrap(); let expected = PrimitiveArray::from_iter(Vec::::new()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); let actual = delta.slice(1024..1024).unwrap(); let expected = PrimitiveArray::from_iter(Vec::::new()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -148,7 +151,7 @@ mod tests { let sliced_again = sliced.slice(0..2).unwrap(); let expected = PrimitiveArray::from_iter(vec![10u32, 11]).into_array(); - assert_arrays_eq!(sliced_again, expected); + assert_arrays_eq!(sliced_again, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -159,7 +162,7 @@ mod tests { let sliced_again = sliced.slice(0..2).unwrap(); let expected = PrimitiveArray::from_iter(vec![10u32, 11]).into_array(); - assert_arrays_eq!(sliced_again, expected); + assert_arrays_eq!(sliced_again, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -170,7 +173,7 @@ mod tests { let sliced_again = sliced.slice(0..2).unwrap(); let expected = PrimitiveArray::from_iter(vec![1034u32, 1035]).into_array(); - assert_arrays_eq!(sliced_again, expected); + assert_arrays_eq!(sliced_again, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -181,7 +184,7 @@ mod tests { let sliced_again = sliced.slice(0..2).unwrap(); let expected = PrimitiveArray::from_iter(vec![1034u32, 1035]).into_array(); - assert_arrays_eq!(sliced_again, expected); + assert_arrays_eq!(sliced_again, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -192,7 +195,7 @@ mod tests { let sliced_again = sliced.slice(5..20).unwrap(); let expected = PrimitiveArray::from_iter(1015u32..1030).into_array(); - assert_arrays_eq!(sliced_again, expected); + assert_arrays_eq!(sliced_again, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -203,7 +206,7 @@ mod tests { let sliced_again = sliced.slice(5..20).unwrap(); let expected = PrimitiveArray::from_iter(1015u32..1030).into_array(); - assert_arrays_eq!(sliced_again, expected); + assert_arrays_eq!(sliced_again, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -211,7 +214,7 @@ mod tests { let delta = da(&(0u32..2048).collect()).into_array(); let expected = PrimitiveArray::from_iter(0u32..2048).into_array(); - assert_arrays_eq!(delta, expected); + assert_arrays_eq!(delta, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -227,7 +230,7 @@ mod tests { let delta = da(&(0u32..2000).collect()).into_array(); let expected = PrimitiveArray::from_iter(0u32..2000).into_array(); - assert_arrays_eq!(delta, expected); + assert_arrays_eq!(delta, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -254,7 +257,10 @@ mod tests { #[case::delta_large_i32((-1024i32..1024).collect())] #[case::delta_single_negative(PrimitiveArray::new(buffer![-42i32], Validity::NonNullable))] fn test_delta_consistency(#[case] array: PrimitiveArray) { - test_array_consistency(&da(&array).into_array()); + test_array_consistency( + &da(&array).into_array(), + &mut SESSION.create_execution_ctx(), + ); } #[rstest] @@ -266,6 +272,9 @@ mod tests { #[case::delta_i8_basic(PrimitiveArray::new(buffer![-1i8, -1, -1, -1, -1], Validity::NonNullable))] #[case::delta_i32_basic(PrimitiveArray::new(buffer![-1i32, -1, -1, -1, -1], Validity::NonNullable))] fn test_delta_binary_numeric(#[case] array: PrimitiveArray) { - test_binary_numeric_array(da(&array).into_array()); + test_binary_numeric_array( + &da(&array).into_array(), + &mut SESSION.create_execution_ctx(), + ); } } diff --git a/encodings/fastlanes/src/for/array/for_compress.rs b/encodings/fastlanes/src/for/array/for_compress.rs index ad9566027d2..9c2d5bbe423 100644 --- a/encodings/fastlanes/src/for/array/for_compress.rs +++ b/encodings/fastlanes/src/for/array/for_compress.rs @@ -3,9 +3,8 @@ use num_traits::PrimInt; use num_traits::WrappingSub; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; -use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::NativePType; use vortex_array::expr::stats::Stat; @@ -16,16 +15,17 @@ use vortex_error::vortex_err; use crate::FoR; use crate::FoRArray; use crate::FoRData; + impl FoRData { - pub fn encode(array: PrimitiveArray) -> VortexResult { + pub fn encode(array: PrimitiveArray, ctx: &mut ExecutionCtx) -> VortexResult { let array_ref = array.clone().into_array(); let min = array_ref .statistics() - .compute_stat(Stat::Min, &mut LEGACY_SESSION.create_execution_ctx())? + .compute_stat(Stat::Min, ctx)? .ok_or_else(|| vortex_err!("Min stat not found"))?; let encoded = match_each_integer_ptype!(array.ptype(), |T| { - compress_primitive::(array, T::try_from(&min)?)?.into_array() + compress_primitive::(array, T::try_from(&min)?, ctx)?.into_array() }); FoR::try_new(encoded, min) } @@ -34,10 +34,11 @@ impl FoRData { fn compress_primitive( parray: PrimitiveArray, min: T, + ctx: &mut ExecutionCtx, ) -> VortexResult { // Set null values to the min value, ensuring that decompress into a value in the primitive // range (and stop them wrapping around). - let encoded = parray.map_each_with_validity::(|(v, bool)| { + let encoded = parray.map_each_with_validity::(ctx, |(v, bool)| { if bool { v.wrapping_sub(&min) } else { @@ -53,13 +54,14 @@ mod test { use itertools::Itertools; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::PType; use vortex_array::expr::stats::StatsProvider; use vortex_array::scalar::Scalar; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; + use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -69,31 +71,31 @@ mod test { use crate::r#for::array::for_decompress::decompress; use crate::r#for::array::for_decompress::fused_decompress; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + crate::initialize(&session); + session + }); #[test] fn test_compress_round_trip_small() { - let array = PrimitiveArray::new( - (1i32..10).collect::>(), - Validity::NonNullable, - ); - let compressed = FoRData::encode(array.clone()).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let array = PrimitiveArray::new((1i32..10).collect::>(), Validity::NonNullable); + let compressed = FoRData::encode(array.clone(), &mut ctx).unwrap(); assert_eq!(i32::try_from(compressed.reference_scalar()).unwrap(), 1); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut ctx); } #[test] fn test_compress() { + let mut ctx = SESSION.create_execution_ctx(); // Create a range offset by a million. let array = PrimitiveArray::new( - (0u32..10_000) - .map(|v| v + 1_000_000) - .collect::>(), + (0u32..10_000).map(|v| v + 1_000_000).collect::>(), Validity::NonNullable, ); - let compressed = FoRData::encode(array).unwrap(); + let compressed = FoRData::encode(array, &mut ctx).unwrap(); assert_eq!( u32::try_from(compressed.reference_scalar()).unwrap(), 1_000_000u32 @@ -102,28 +104,27 @@ mod test { #[test] fn test_zeros() { + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::new(buffer![0i32; 100], Validity::NonNullable); assert_eq!(array.statistics().len(), 0); let dtype = array.dtype().clone(); - let compressed = FoRData::encode(array).unwrap(); + let compressed = FoRData::encode(array, &mut ctx).unwrap(); assert_eq!(compressed.reference_scalar().dtype(), &dtype); assert!(compressed.reference_scalar().dtype().is_signed_int()); assert!(compressed.encoded().dtype().is_signed_int()); - let encoded = compressed - .encoded() - .execute_scalar(0, &mut SESSION.create_execution_ctx()) - .unwrap(); + let encoded = compressed.encoded().execute_scalar(0, &mut ctx).unwrap(); assert_eq!(encoded, Scalar::from(0i32)); } #[test] fn test_decompress() { + let mut ctx = SESSION.create_execution_ctx(); // Create a range offset by a million. let array = PrimitiveArray::from_iter((0u32..100_000).step_by(1024).map(|v| v + 1_000_000)); - let compressed = FoRData::encode(array.clone()).unwrap(); - assert_arrays_eq!(compressed, array); + let compressed = FoRData::encode(array.clone(), &mut ctx).unwrap(); + assert_arrays_eq!(compressed, array, &mut ctx); } #[test] @@ -134,7 +135,7 @@ mod test { let array = PrimitiveArray::from_iter((0u32..1024).map(|x| x % 7)); let bp = BitPackedData::encode(&array.into_array(), 3, &mut ctx).unwrap(); let compressed = FoR::try_new(bp.into_array(), 10u32.into()).unwrap(); - assert_arrays_eq!(compressed, expect); + assert_arrays_eq!(compressed, expect, &mut ctx); } #[test] @@ -146,7 +147,7 @@ mod test { let bp = BitPackedData::encode(&array.into_array(), 2, &mut ctx)?; let compressed = FoR::try_new(bp.clone().into_array(), 10u32.into())?; let decompressed = fused_decompress::(&compressed, bp.as_view(), &mut ctx)?; - assert_arrays_eq!(decompressed, expect); + assert_arrays_eq!(decompressed, expect, &mut ctx); Ok(()) } @@ -154,7 +155,7 @@ mod test { fn test_overflow() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter(i8::MIN..=i8::MAX); - let compressed = FoRData::encode(array.clone())?; + let compressed = FoRData::encode(array.clone(), &mut ctx)?; assert_eq!( i8::MIN, compressed @@ -184,7 +185,7 @@ mod test { i8::try_from(&compressed.execute_scalar(i, &mut ctx).unwrap()).unwrap() ); }); - assert_arrays_eq!(decompressed, array); + assert_arrays_eq!(decompressed, array, &mut ctx); Ok(()) } } diff --git a/encodings/fastlanes/src/for/compute/cast.rs b/encodings/fastlanes/src/for/compute/cast.rs index e534f3c76fa..b865af43d12 100644 --- a/encodings/fastlanes/src/for/compute/cast.rs +++ b/encodings/fastlanes/src/for/compute/cast.rs @@ -30,9 +30,12 @@ impl CastReduce for FoR { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; @@ -43,10 +46,17 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexExpect; + use vortex_session::VortexSession; use crate::FoR; use crate::FoRArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn for_arr(encoded: ArrayRef, reference: Scalar) -> FoRArray { FoR::try_new(encoded, reference).vortex_expect("FoR array construction should succeed") } @@ -70,7 +80,8 @@ mod tests { // Verify the values after decoding assert_arrays_eq!( casted, - PrimitiveArray::from_iter([100i64, 110, 120, 130, 140]) + PrimitiveArray::from_iter([100i64, 110, 120, 130, 140]), + &mut SESSION.create_execution_ctx() ); } diff --git a/encodings/fastlanes/src/for/compute/compare.rs b/encodings/fastlanes/src/for/compute/compare.rs index 8cd2fb9c011..384bdee8037 100644 --- a/encodings/fastlanes/src/for/compute/compare.rs +++ b/encodings/fastlanes/src/for/compute/compare.rs @@ -90,18 +90,28 @@ where #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_session::VortexSession; use super::*; use crate::FoR; use crate::FoRArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn for_arr(encoded: ArrayRef, reference: Scalar) -> FoRArray { FoR::try_new(encoded, reference).vortex_expect("FoR array construction should succeed") } @@ -123,7 +133,11 @@ mod tests { ) .unwrap() .unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter([false, true, false].map(Some))); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, true, false].map(Some)), + &mut SESSION.create_execution_ctx() + ); let result = compare_constant( lhs.as_view(), @@ -133,7 +147,11 @@ mod tests { ) .unwrap() .unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter([true, true, false].map(Some))); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, true, false].map(Some)), + &mut SESSION.create_execution_ctx() + ); for op in [ CompareOperator::Lt, @@ -203,7 +221,8 @@ mod tests { .unwrap(); assert_arrays_eq!( result, - BoolArray::from_iter([false, false, false].map(Some)) + BoolArray::from_iter([false, false, false].map(Some)), + &mut SESSION.create_execution_ctx() ); let result = compare_constant( @@ -214,7 +233,11 @@ mod tests { ) .unwrap() .unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter([true, true, true].map(Some))); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, true, true].map(Some)), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -237,7 +260,11 @@ mod tests { ) .unwrap() .unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter([Some(false), Some(true)])); + assert_arrays_eq!( + result, + BoolArray::from_iter([Some(false), Some(true)]), + &mut SESSION.create_execution_ctx() + ); let result = compare_constant( lhs.as_view(), @@ -247,6 +274,10 @@ mod tests { ) .unwrap() .unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter([Some(true), Some(false)])); + assert_arrays_eq!( + result, + BoolArray::from_iter([Some(true), Some(false)]), + &mut SESSION.create_execution_ctx() + ); } } diff --git a/encodings/fastlanes/src/for/compute/is_sorted.rs b/encodings/fastlanes/src/for/compute/is_sorted.rs index 6b72af5902b..7cf43a93da4 100644 --- a/encodings/fastlanes/src/for/compute/is_sorted.rs +++ b/encodings/fastlanes/src/for/compute/is_sorted.rs @@ -60,9 +60,9 @@ impl DynAggregateKernel for FoRIsSortedKernel { #[cfg(test)] mod test { use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::is_sorted::is_sorted; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::validity::Validity; use vortex_buffer::buffer; @@ -72,10 +72,10 @@ mod test { #[test] fn test_sorted() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let a = PrimitiveArray::new(buffer![-1, 0, i8::MAX], Validity::NonNullable); - let b = FoRData::encode(a).unwrap(); + let b = FoRData::encode(a, &mut ctx).unwrap(); assert!( is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), "{}", @@ -83,7 +83,7 @@ mod test { ); let a = PrimitiveArray::new(buffer![i8::MIN, 0, i8::MAX], Validity::NonNullable); - let b = FoRData::encode(a).unwrap(); + let b = FoRData::encode(a, &mut ctx).unwrap(); assert!( is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), "{}", @@ -91,7 +91,7 @@ mod test { ); let a = PrimitiveArray::new(buffer![i8::MIN, 0, 30, 127], Validity::NonNullable); - let b = FoRData::encode(a).unwrap(); + let b = FoRData::encode(a, &mut ctx).unwrap(); assert!( is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), "{}", @@ -99,7 +99,7 @@ mod test { ); let a = PrimitiveArray::new(buffer![i8::MIN, -3, -1], Validity::NonNullable); - let b = FoRData::encode(a).unwrap(); + let b = FoRData::encode(a, &mut ctx).unwrap(); assert!( is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), "{}", @@ -107,7 +107,7 @@ mod test { ); let a = PrimitiveArray::new(buffer![-10, -3, -1], Validity::NonNullable); - let b = FoRData::encode(a).unwrap(); + let b = FoRData::encode(a, &mut ctx).unwrap(); assert!( is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), "{}", @@ -115,7 +115,7 @@ mod test { ); let a = PrimitiveArray::new(buffer![-10, -11, -1], Validity::NonNullable); - let b = FoRData::encode(a).unwrap(); + let b = FoRData::encode(a, &mut ctx).unwrap(); assert!( !is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), "{}", @@ -123,7 +123,7 @@ mod test { ); let a = PrimitiveArray::new(buffer![-10, i8::MIN, -1], Validity::NonNullable); - let b = FoRData::encode(a).unwrap(); + let b = FoRData::encode(a, &mut ctx).unwrap(); assert!( !is_sorted(&b.clone().into_array(), &mut ctx).unwrap(), "{}", diff --git a/encodings/fastlanes/src/for/compute/mod.rs b/encodings/fastlanes/src/for/compute/mod.rs index 7307b50d776..84e8812b598 100644 --- a/encodings/fastlanes/src/for/compute/mod.rs +++ b/encodings/fastlanes/src/for/compute/mod.rs @@ -102,6 +102,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::PrimitiveArray; use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; @@ -145,7 +147,10 @@ mod tests { #[case::for_large_deltas(fa(buffer![100i64, 200, 300, 400, 500].into_array(), Scalar::from(100i64)))] fn test_for_consistency(#[case] array: FoRArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[rstest] @@ -158,6 +163,9 @@ mod tests { Scalar::from(2000i32) ))] fn test_for_binary_numeric(#[case] array: FoRArray) { - test_binary_numeric_array(array.into_array()); + test_binary_numeric_array( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/encodings/fastlanes/src/for/mod.rs b/encodings/fastlanes/src/for/mod.rs index 8a70fe30499..a5d495e3578 100644 --- a/encodings/fastlanes/src/for/mod.rs +++ b/encodings/fastlanes/src/for/mod.rs @@ -10,3 +10,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::FoR; pub use vtable::FoRArray; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} diff --git a/encodings/fastlanes/src/for/vtable/kernels.rs b/encodings/fastlanes/src/for/vtable/kernels.rs index 331b13eceef..6c63292bb9f 100644 --- a/encodings/fastlanes/src/for/vtable/kernels.rs +++ b/encodings/fastlanes/src/for/vtable/kernels.rs @@ -1,13 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::VortexSession; use crate::FoR; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(FoR)), - ParentKernelSet::lift(&TakeExecuteAdaptor(FoR)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), FoR, CompareExecuteAdaptor(FoR)); + kernels.register_execute_parent_kernel(Dict.id(), FoR, TakeExecuteAdaptor(FoR)); +} diff --git a/encodings/fastlanes/src/for/vtable/mod.rs b/encodings/fastlanes/src/for/vtable/mod.rs index 037a858892e..28811bd9582 100644 --- a/encodings/fastlanes/src/for/vtable/mod.rs +++ b/encodings/fastlanes/src/for/vtable/mod.rs @@ -37,7 +37,6 @@ use crate::FoRData; use crate::r#for::array::FoRArrayExt; use crate::r#for::array::SLOT_NAMES; use crate::r#for::array::for_decompress::decompress; -use crate::r#for::vtable::kernels::PARENT_KERNELS; use crate::r#for::vtable::rules::PARENT_RULES; mod kernels; @@ -49,6 +48,10 @@ mod validity; /// A [`FoR`]-encoded Vortex array. pub type FoRArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernels::initialize(session); +} + impl ArrayHash for FoRData { fn array_hash(&self, state: &mut H, _accuracy: EqMode) { self.reference.hash(state); @@ -95,6 +98,14 @@ impl VTable for FoR { None } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { SLOT_NAMES[idx].to_string() } @@ -150,15 +161,6 @@ impl VTable for FoR { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done(decompress(&array, ctx)?.into_array())) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[derive(Clone, Debug)] @@ -179,8 +181,8 @@ impl FoR { } /// Encode a primitive array using Frame of Reference encoding. - pub fn encode(array: PrimitiveArray) -> VortexResult { - FoRData::encode(array) + pub fn encode(array: PrimitiveArray, ctx: &mut ExecutionCtx) -> VortexResult { + FoRData::encode(array, ctx) } } diff --git a/encodings/fastlanes/src/for/vtable/operations.rs b/encodings/fastlanes/src/for/vtable/operations.rs index 5fea6f4bdaa..e549431f973 100644 --- a/encodings/fastlanes/src/for/vtable/operations.rs +++ b/encodings/fastlanes/src/for/vtable/operations.rs @@ -40,15 +40,30 @@ impl OperationsVTable for FoR { #[cfg(test)] mod test { + use std::sync::LazyLock; + + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_session::VortexSession; use crate::FoRData; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn for_scalar_at() { - let for_arr = FoRData::encode(PrimitiveArray::from_iter([-100, 1100, 1500, 1900])).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let for_arr = FoRData::encode( + PrimitiveArray::from_iter([-100, 1100, 1500, 1900]), + &mut ctx, + ) + .unwrap(); let expected = PrimitiveArray::from_iter([-100, 1100, 1500, 1900]); - assert_arrays_eq!(for_arr, expected); + assert_arrays_eq!(for_arr, expected, &mut ctx); } } diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index 9022b7c4e2b..e6192a9cdfb 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -3,6 +3,29 @@ #![expect(clippy::cast_possible_truncation)] +//! FastLanes integer encodings for Vortex arrays. +//! +//! This crate provides SIMD-friendly integer encodings: +//! +//! - [`BitPacked`] stores fixed-width integer values using the minimum bit width plus optional +//! patches. +//! - [`FoR`] stores frame-of-reference deltas from a base value. +//! - [`Delta`] stores adjacent deltas in chunked form. +//! - [`RLE`] stores repeated runs. +//! +//! Call [`initialize`] to register the encodings and encoding-specific aggregate kernels in a +//! session before deserializing or executing arrays that may contain these encodings. +//! +//! ```rust +//! let session = vortex_array::array_session(); +//! vortex_fastlanes::initialize(&session); +//! ``` +//! +//! ## Paper +//! +//! The original encodings are described in the paper [The FastLanes Compression Layout](https://15721.courses.cs.cmu.edu/spring2024/papers/03-data2/p2132-afroozeh.pdf), +//! but are not fully binary compatible. See the underlying [fastlanes](https://github.com/spiraldb/fastlanes) crate for more details. + pub use bitpacking::*; pub use delta::*; pub use r#for::*; @@ -21,7 +44,7 @@ mod delta; mod r#for; mod rle; -pub(crate) const FL_CHUNK_SIZE: usize = 1024; +pub const FL_CHUNK_SIZE: usize = 1024; use bitpacking::compute::is_constant::BitPackedIsConstantKernel; use r#for::compute::is_constant::FoRIsConstantKernel; @@ -47,6 +70,9 @@ pub fn initialize(session: &VortexSession) { session.arrays().register(Delta); session.arrays().register(FoR); session.arrays().register(RLE); + bitpacking::initialize(session); + r#for::initialize(session); + rle::initialize(session); // Register the encoding-specific aggregate kernels. session.aggregate_fns().register_aggregate_kernel( @@ -137,18 +163,14 @@ mod test { use std::sync::LazyLock; use vortex_array::VortexSessionExecute; - use vortex_array::session::ArraySessionExt; use vortex_buffer::BitBufferMut; use vortex_session::VortexSession; use super::*; pub static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty(); - session.arrays().register(BitPacked); - session.arrays().register(Delta); - session.arrays().register(FoR); - session.arrays().register(RLE); + let session = vortex_array::array_session(); + initialize(&session); session }); diff --git a/encodings/fastlanes/src/rle/array/mod.rs b/encodings/fastlanes/src/rle/array/mod.rs index 8793179c7d3..7ddff474836 100644 --- a/encodings/fastlanes/src/rle/array/mod.rs +++ b/encodings/fastlanes/src/rle/array/mod.rs @@ -140,7 +140,6 @@ mod tests { use vortex_array::ArrayContext; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -274,7 +273,7 @@ mod tests { let invalid_slice = rle_array .slice(2..5) .unwrap() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut SESSION.create_execution_ctx()) .unwrap() .into_primitive(); let mut ctx = SESSION.create_execution_ctx(); @@ -311,13 +310,10 @@ mod tests { let validity_mask = sliced_array .validity() .unwrap() - .execute_mask( - sliced_array.len(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) + .execute_mask(sliced_array.len(), &mut SESSION.create_execution_ctx()) .unwrap(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let expected_mask = Validity::from_iter([false, true, false]) .execute_mask(3, &mut ctx) .unwrap(); @@ -347,7 +343,7 @@ mod tests { #[test] fn test_multi_chunk_two_chunks() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]).into_array(); let indices = PrimitiveArray::from_iter([0u16, 1].repeat(1024)).into_array(); let values_idx_offsets = PrimitiveArray::from_iter([0u64, 2]).into_array(); @@ -395,7 +391,11 @@ mod tests { let decoded_data = decoded.execute::(&mut exec_ctx)?; - assert_arrays_eq!(original_data, decoded_data); + assert_arrays_eq!( + original_data, + decoded_data, + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -442,7 +442,11 @@ mod tests { .execute::(&mut exec_ctx)?; let decoded_data = decoded.execute::(&mut exec_ctx)?; - assert_arrays_eq!(original_data, decoded_data); + assert_arrays_eq!( + original_data, + decoded_data, + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -493,7 +497,7 @@ mod tests { .as_array() .clone() .execute::(&mut ctx)?; - assert_arrays_eq!(decoded, original); + assert_arrays_eq!(decoded, original, &mut ctx); Ok(()) } } diff --git a/encodings/fastlanes/src/rle/array/rle_compress.rs b/encodings/fastlanes/src/rle/array/rle_compress.rs index 107c912c23f..b7e04a13699 100644 --- a/encodings/fastlanes/src/rle/array/rle_compress.rs +++ b/encodings/fastlanes/src/rle/array/rle_compress.rs @@ -153,10 +153,11 @@ fn padded_validity(array: &PrimitiveArray, ctx: &mut ExecutionCtx) -> VortexResu #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::MaskedArray; @@ -166,13 +167,20 @@ mod tests { use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use super::*; use crate::rle::array::RLEArrayExt; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_encode_decode() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // u8 let array_u8: Buffer = buffer![1, 1, 2, 2, 3, 3]; let encoded_u8 = RLEData::encode( @@ -184,7 +192,7 @@ mod tests { .clone() .execute::(&mut ctx)?; let expected_u8 = PrimitiveArray::from_iter(vec![1u8, 1, 2, 2, 3, 3]); - assert_arrays_eq!(decoded_u8, expected_u8); + assert_arrays_eq!(decoded_u8, expected_u8, &mut ctx); // u16 let array_u16: Buffer = buffer![100, 100, 200, 200]; @@ -197,7 +205,7 @@ mod tests { .clone() .execute::(&mut ctx)?; let expected_u16 = PrimitiveArray::from_iter(vec![100u16, 100, 200, 200]); - assert_arrays_eq!(decoded_u16, expected_u16); + assert_arrays_eq!(decoded_u16, expected_u16, &mut ctx); // u64 let array_u64: Buffer = buffer![1000, 1000, 2000]; @@ -210,13 +218,13 @@ mod tests { .clone() .execute::(&mut ctx)?; let expected_u64 = PrimitiveArray::from_iter(vec![1000u64, 1000, 2000]); - assert_arrays_eq!(decoded_u64, expected_u64); + assert_arrays_eq!(decoded_u64, expected_u64, &mut ctx); Ok(()) } #[test] fn test_length() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Buffer = buffer![1, 1, 2, 2, 2, 3]; let encoded = RLEData::encode( PrimitiveArray::new(values, Validity::NonNullable).as_view(), @@ -228,7 +236,7 @@ mod tests { #[test] fn test_empty_length() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Buffer = Buffer::empty(); let encoded = RLEData::encode( PrimitiveArray::new(values, Validity::NonNullable).as_view(), @@ -242,7 +250,7 @@ mod tests { #[test] fn test_single_value() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Buffer = vec![42; 2000].into_iter().collect(); let encoded = RLEData::encode( @@ -256,13 +264,13 @@ mod tests { .clone() .execute::(&mut ctx)?; // Verify round-trip let expected = PrimitiveArray::from_iter(vec![42u16; 2000]); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn test_all_different() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Buffer = (0u8..=255).collect(); let encoded = RLEData::encode( @@ -276,13 +284,13 @@ mod tests { .clone() .execute::(&mut ctx)?; // Verify round-trip let expected = PrimitiveArray::from_iter((0u8..=255).collect::>()); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn test_partial_last_chunk() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test array with partial last chunk (not divisible by 1024) let values: Buffer = (0..1500).map(|i| (i / 100) as u32).collect(); let array = PrimitiveArray::new(values, Validity::NonNullable); @@ -290,14 +298,14 @@ mod tests { let encoded = RLEData::encode(array.as_view(), &mut ctx).unwrap(); assert_eq!(encoded.len(), 1500); - assert_arrays_eq!(encoded, array); + assert_arrays_eq!(encoded, array, &mut ctx); // 2 chunks: 1024 + 476 elements assert_eq!(encoded.values_idx_offsets().len(), 2); } #[test] fn test_two_full_chunks() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Array that spans exactly 2 chunks (2048 elements) let values: Buffer = (0..2048).map(|i| (i / 100) as u32).collect(); let array = PrimitiveArray::new(values, Validity::NonNullable); @@ -305,7 +313,7 @@ mod tests { let encoded = RLEData::encode(array.as_view(), &mut ctx).unwrap(); assert_eq!(encoded.len(), 2048); - assert_arrays_eq!(encoded, array); + assert_arrays_eq!(encoded, array, &mut ctx); assert_eq!(encoded.values_idx_offsets().len(), 2); } @@ -324,7 +332,7 @@ mod tests { fn test_roundtrip_primitive_types( #[case] values: Buffer, ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive = values .clone() .into_array() @@ -335,7 +343,7 @@ mod tests { .clone() .execute::(&mut ctx)?; let expected = PrimitiveArray::new(values, primitive.validity()?); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } @@ -366,44 +374,44 @@ mod tests { #[test] fn test_encode_all_null_chunk() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Vec> = vec![None; FL_CHUNK_SIZE]; let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let decoded = with_masked_constant_indices(&rle, &mut ctx)?; - assert_arrays_eq!(decoded, original); + assert_arrays_eq!(decoded, original, &mut ctx); Ok(()) } #[test] fn test_encode_all_null_chunk_then_value_chunk() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // First chunk is entirely null, second chunk has a value preceded by nulls. let mut values: Vec> = vec![None; 2 * FL_CHUNK_SIZE]; values[FL_CHUNK_SIZE + 100] = Some(42); let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let decoded = with_masked_constant_indices(&rle, &mut ctx)?; - assert_arrays_eq!(decoded, original); + assert_arrays_eq!(decoded, original, &mut ctx); Ok(()) } #[test] fn test_encode_one_value_near_end() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Single distinct value near the end of the chunk. let mut values: Vec> = vec![None; FL_CHUNK_SIZE]; values[1000] = Some(42); let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let decoded = with_masked_constant_indices(&rle, &mut ctx)?; - assert_arrays_eq!(decoded, original); + assert_arrays_eq!(decoded, original, &mut ctx); Ok(()) } #[test] fn test_encode_value_chunk_then_all_null_remainder() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // 1085 elements (2 chunks: 1024 + 61 padded to 1024). // Chunk 0 has -1i16 at scattered positions (273..=366), rest null. // Chunk 1 (the remainder) is entirely null. @@ -420,7 +428,7 @@ mod tests { let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let decoded = with_masked_constant_indices(&rle, &mut ctx)?; - assert_arrays_eq!(decoded, original); + assert_arrays_eq!(decoded, original, &mut ctx); Ok(()) } @@ -463,18 +471,18 @@ mod tests { #[test] fn test_random_invalid_indices_all_null_chunk() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Vec> = vec![None; FL_CHUNK_SIZE]; let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let clobbered = with_random_invalid_indices(&rle, &mut ctx)?; - assert_arrays_eq!(clobbered, original); + assert_arrays_eq!(clobbered, original, &mut ctx); Ok(()) } #[test] fn test_random_invalid_indices_sparse_values() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut values: Vec> = vec![None; FL_CHUNK_SIZE]; values[0] = Some(10); values[500] = Some(20); @@ -482,13 +490,13 @@ mod tests { let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let clobbered = with_random_invalid_indices(&rle, &mut ctx)?; - assert_arrays_eq!(clobbered, original); + assert_arrays_eq!(clobbered, original, &mut ctx); Ok(()) } #[test] fn test_random_invalid_indices_multi_chunk() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Two chunks: first has scattered values, second is all null. let mut values: Vec> = vec![None; 2 * FL_CHUNK_SIZE]; values[0] = Some(10); @@ -497,13 +505,13 @@ mod tests { let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let clobbered = with_random_invalid_indices(&rle, &mut ctx)?; - assert_arrays_eq!(clobbered, original); + assert_arrays_eq!(clobbered, original, &mut ctx); Ok(()) } #[test] fn test_random_invalid_indices_partial_last_chunk() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // 1085 elements: chunk 0 has values at scattered positions, chunk 1 is // a partial (61 elements padded to 1024) that is entirely null. let mut values: Vec> = vec![None; 1085]; @@ -513,13 +521,13 @@ mod tests { let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let clobbered = with_random_invalid_indices(&rle, &mut ctx)?; - assert_arrays_eq!(clobbered, original); + assert_arrays_eq!(clobbered, original, &mut ctx); Ok(()) } #[test] fn test_random_invalid_indices_mostly_valid() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Most positions are valid, only a few are null with garbage indices. let mut values: Vec> = (0..FL_CHUNK_SIZE).map(|i| Some((i / 100) as u64)).collect(); @@ -530,7 +538,7 @@ mod tests { let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let clobbered = with_random_invalid_indices(&rle, &mut ctx)?; - assert_arrays_eq!(clobbered, original); + assert_arrays_eq!(clobbered, original, &mut ctx); Ok(()) } @@ -543,11 +551,11 @@ mod tests { fn test_float_zeros( #[case] values: Vec, ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive = PrimitiveArray::from_iter(values); let rle = RLEData::encode(primitive.as_view(), &mut ctx)?; let decoded = rle.as_array().clone().execute::(&mut ctx)?; - assert_arrays_eq!(primitive, decoded); + assert_arrays_eq!(primitive, decoded, &mut ctx); Ok(()) } } diff --git a/encodings/fastlanes/src/rle/compute/cast.rs b/encodings/fastlanes/src/rle/compute/cast.rs index e01d359e59b..92947d927bc 100644 --- a/encodings/fastlanes/src/rle/compute/cast.rs +++ b/encodings/fastlanes/src/rle/compute/cast.rs @@ -49,7 +49,6 @@ mod tests { use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -58,7 +57,6 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -66,8 +64,11 @@ mod tests { use crate::RLEData; use crate::rle::RLEArray; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); fn rle(primitive: &PrimitiveArray, ctx: &mut ExecutionCtx) -> RLEArray { RLEData::encode(primitive.as_view(), ctx).unwrap() @@ -86,7 +87,11 @@ mod tests { .into_array() .cast(DType::Primitive(PType::U16, Nullability::NonNullable)) .unwrap(); - assert_arrays_eq!(casted, PrimitiveArray::from_iter([10u16, 20, 30, 40, 50])); + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([10u16, 20, 30, 40, 50]), + &mut ctx + ); } #[test] @@ -155,7 +160,7 @@ mod tests { ) )] fn test_cast_rle_conformance(#[case] primitive: PrimitiveArray) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let rle_array = rle(&primitive, &mut ctx); test_cast_conformance(&rle_array.into_array()); } diff --git a/encodings/fastlanes/src/rle/kernel.rs b/encodings/fastlanes/src/rle/kernel.rs index f5b654706cd..9dd05ffbae5 100644 --- a/encodings/fastlanes/src/rle/kernel.rs +++ b/encodings/fastlanes/src/rle/kernel.rs @@ -4,20 +4,25 @@ use std::ops::Range; use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::Slice; use vortex_array::arrays::slice::SliceExecuteAdaptor; use vortex_array::arrays::slice::SliceKernel; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_error::VortexResult; +use vortex_session::VortexSession; use crate::FL_CHUNK_SIZE; use crate::RLE; use crate::rle::RLEArrayExt; -pub(crate) static PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&SliceExecuteAdaptor(RLE))]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Slice.id(), RLE, SliceExecuteAdaptor(RLE)); +} impl SliceKernel for RLE { fn slice( diff --git a/encodings/fastlanes/src/rle/mod.rs b/encodings/fastlanes/src/rle/mod.rs index 1c580683908..649d11f307c 100644 --- a/encodings/fastlanes/src/rle/mod.rs +++ b/encodings/fastlanes/src/rle/mod.rs @@ -11,3 +11,7 @@ mod kernel; mod vtable; pub use vtable::RLE; pub use vtable::RLEArray; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + kernel::initialize(session); +} diff --git a/encodings/fastlanes/src/rle/vtable/mod.rs b/encodings/fastlanes/src/rle/vtable/mod.rs index 82c236ae810..cea655e53b4 100644 --- a/encodings/fastlanes/src/rle/vtable/mod.rs +++ b/encodings/fastlanes/src/rle/vtable/mod.rs @@ -38,7 +38,6 @@ use crate::rle::array::SLOT_NAMES; use crate::rle::array::VALUES_IDX_OFFSETS_SLOT; use crate::rle::array::VALUES_SLOT; use crate::rle::array::rle_decompress::rle_decompress; -use crate::rle::kernel::PARENT_KERNELS; use crate::rle::vtable::rules::RULES; mod operations; @@ -122,6 +121,14 @@ impl VTable for RLE { None } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, @@ -193,15 +200,6 @@ impl VTable for RLE { Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done( rle_decompress(&array, ctx)?.into_array(), diff --git a/encodings/fastlanes/src/rle/vtable/operations.rs b/encodings/fastlanes/src/rle/vtable/operations.rs index 0441a1e61b9..b8f79db95a6 100644 --- a/encodings/fastlanes/src/rle/vtable/operations.rs +++ b/encodings/fastlanes/src/rle/vtable/operations.rs @@ -41,20 +41,29 @@ impl OperationsVTable for RLE { #[cfg(test)] mod tests { + + use std::sync::LazyLock; + use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; 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_buffer::buffer; + use vortex_session::VortexSession; use super::*; use crate::RLE; use crate::RLEArray; use crate::RLEData; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + mod fixture { use super::*; @@ -122,7 +131,11 @@ mod tests { let array = fixture::rle_array(); let expected = PrimitiveArray::from_iter([10u32, 10, 20, 20, 20, 30, 10]); - assert_arrays_eq!(array.slice(0..7).unwrap(), expected); + assert_arrays_eq!( + array.slice(0..7).unwrap(), + expected, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -139,7 +152,11 @@ mod tests { Some(30), Some(10), ]); - assert_arrays_eq!(array.slice(0..7).unwrap(), expected); + assert_arrays_eq!( + array.slice(0..7).unwrap(), + expected, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -151,7 +168,7 @@ mod tests { assert_eq!(sliced.len(), 4); let expected = PrimitiveArray::from_iter([20u32, 20, 20, 30]); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -163,12 +180,12 @@ mod tests { assert_eq!(sliced.len(), 4); let expected = PrimitiveArray::from_option_iter([Some(20u32), Some(20), None, Some(30)]); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut SESSION.create_execution_ctx()); } #[test] fn test_scalar_at_multiple_chunks() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test accessing elements around chunk boundaries let values: Buffer = (0..3000).map(|i| (i / 50) as u16).collect(); let expected: Vec = (0..3000).map(|i| (i / 50) as u16).collect(); @@ -197,7 +214,7 @@ mod tests { fn test_scalar_at_out_of_bounds() { let array = fixture::rle_array(); array - .execute_scalar(1025, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1025, &mut SESSION.create_execution_ctx()) .unwrap(); } @@ -206,7 +223,7 @@ mod tests { fn test_scalar_at_slice_out_of_bounds() { let array = fixture::rle_array().slice(0..1).unwrap(); array - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut SESSION.create_execution_ctx()) .unwrap(); } @@ -224,7 +241,11 @@ mod tests { Some(30), Some(10), ]); - assert_arrays_eq!(sliced.into_array(), expected.into_array()); + assert_arrays_eq!( + sliced.into_array(), + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -233,7 +254,7 @@ mod tests { let sliced = array.slice(4..6).unwrap(); // [20, 30] let expected = buffer![20u32, 30].into_array(); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -242,7 +263,7 @@ mod tests { let sliced = array.slice(5..6).unwrap(); // [30] let expected = buffer![30u32].into_array(); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -259,12 +280,16 @@ mod tests { let sliced = array.slice(1..4).unwrap(); // [null, 20, 20] let expected = PrimitiveArray::from_option_iter([Option::::None, Some(20), Some(20)]); - assert_arrays_eq!(sliced.into_array(), expected.into_array()); + assert_arrays_eq!( + sliced.into_array(), + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] fn test_slice_decode_with_nulls() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = fixture::rle_array_with_nulls(); let sliced = array .slice(1..4) @@ -273,7 +298,7 @@ mod tests { .unwrap(); // [null, 20, 20] let expected = PrimitiveArray::from_option_iter([Option::::None, Some(20), Some(20)]); - assert_arrays_eq!(sliced.into_array(), expected.into_array()); + assert_arrays_eq!(sliced.into_array(), expected.into_array(), &mut ctx); } #[test] @@ -286,7 +311,7 @@ mod tests { #[test] fn test_slice_across_chunk_boundaries() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Buffer = (0..2100).map(|i| (i / 100) as u32).collect(); let expected: Vec = (0..2100).map(|i| (i / 100) as u32).collect(); let array = values.into_array(); @@ -298,14 +323,16 @@ mod tests { let slice = encoded.slice(500..1500).unwrap(); assert_arrays_eq!( slice, - PrimitiveArray::from_iter(expected[500..1500].iter().copied()) + PrimitiveArray::from_iter(expected[500..1500].iter().copied()), + &mut ctx ); // Slice across second and third chunk. let slice = encoded.slice(1000..2000).unwrap(); assert_arrays_eq!( slice, - PrimitiveArray::from_iter(expected[1000..2000].iter().copied()) + PrimitiveArray::from_iter(expected[1000..2000].iter().copied()), + &mut ctx ); } } diff --git a/encodings/fsst/Cargo.toml b/encodings/fsst/Cargo.toml index 0d722a131d3..f749f121efa 100644 --- a/encodings/fsst/Cargo.toml +++ b/encodings/fsst/Cargo.toml @@ -18,6 +18,7 @@ workspace = true [dependencies] fsst-rs = { workspace = true } +num-traits = { workspace = true } prost = { workspace = true } rand = { workspace = true, optional = true } vortex-array = { workspace = true } diff --git a/encodings/fsst/benches/chunked_dict_fsst_builder.rs b/encodings/fsst/benches/chunked_dict_fsst_builder.rs index 22131631a6e..8c2c6c2e388 100644 --- a/encodings/fsst/benches/chunked_dict_fsst_builder.rs +++ b/encodings/fsst/benches/chunked_dict_fsst_builder.rs @@ -11,7 +11,6 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; use vortex_array::builders::builder_with_capacity; use vortex_array::dtype::NativePType; -use vortex_array::session::ArraySession; use vortex_error::VortexExpect; use vortex_fsst::test_utils::gen_dict_fsst_test_data; use vortex_session::VortexSession; @@ -29,8 +28,11 @@ const BENCH_ARGS: &[(usize, usize, usize)] = &[ (1000, 1000, 100), ]; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fsst::initialize(&session); + session +}); fn make_dict_fsst_chunks( len: usize, diff --git a/encodings/fsst/benches/fsst_compress.rs b/encodings/fsst/benches/fsst_compress.rs index 88e4e26a903..5e03d1d003d 100644 --- a/encodings/fsst/benches/fsst_compress.rs +++ b/encodings/fsst/benches/fsst_compress.rs @@ -9,6 +9,7 @@ use divan::Bencher; 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::RecursiveCanonical; @@ -23,7 +24,6 @@ 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_array::session::ArraySession; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; use vortex_session::VortexSession; @@ -32,8 +32,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fsst::initialize(&session); + session +}); // [(string_count, avg_len, unique_chars)] const BENCH_ARGS: &[(usize, usize, u8)] = &[ @@ -54,60 +57,50 @@ const BENCH_ARGS: &[(usize, usize, u8)] = &[ #[divan::bench(args = BENCH_ARGS)] fn compress_fsst(bencher: Bencher, (string_count, avg_len, unique_chars): (usize, usize, u8)) { let array = generate_test_data(string_count, avg_len, unique_chars); - let compressor = fsst_train_compressor(&array); + let compressor = fsst_train_compressor(&array, &mut SESSION.create_execution_ctx()).unwrap(); bencher .with_inputs(|| (&array, &compressor, SESSION.create_execution_ctx())) - .bench_refs(|(array, compressor, ctx)| { - fsst_compress(*array, array.len(), array.dtype(), compressor, ctx) - }) + .bench_refs(|(array, compressor, ctx)| fsst_compress(array, compressor, ctx).unwrap()) } #[divan::bench(args = BENCH_ARGS)] fn decompress_fsst(bencher: Bencher, (string_count, avg_len, unique_chars): (usize, usize, u8)) { let array = generate_test_data(string_count, avg_len, unique_chars); - let compressor = fsst_train_compressor(&array); - let len = array.len(); - let dtype = array.dtype().clone(); - let encoded = fsst_compress( - array, - len, - &dtype, - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let encoded = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); bencher .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) - .bench_refs(|(encoded, ctx)| (**encoded).clone().into_array().execute::(ctx)) + .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) - .bench_refs(|array| fsst_train_compressor(array)) + .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); - let compressor = fsst_train_compressor(&array); - let fsst_array = fsst_compress( - &array, - array.len(), - array.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); - let constant = ConstantArray::new(Scalar::from(&b"const"[..]), array.len()); + let len = array.len(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); + let constant = ConstantArray::new(Scalar::from(&b"const"[..]), len); bencher .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, constant, ctx)| { fsst_array .clone() - .into_array() .binary(constant.clone().into_array(), Operator::Eq) .unwrap() .execute::(ctx) @@ -121,22 +114,19 @@ fn canonicalize_compare( (string_count, avg_len, unique_chars): (usize, usize, u8), ) { let array = generate_test_data(string_count, avg_len, unique_chars); - let compressor = fsst_train_compressor(&array); - let fsst_array = fsst_compress( - &array, - array.len(), - array.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); - let constant = ConstantArray::new(Scalar::from(&b"const"[..]), array.len()); + let len = array.len(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); + let constant = ConstantArray::new(Scalar::from(&b"const"[..]), len); bencher .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, constant, ctx)| { - (*fsst_array) + fsst_array .clone() - .into_array() .execute::(ctx) .unwrap() .into_array() @@ -188,11 +178,11 @@ fn chunked_into_canonical( bencher .with_inputs(|| (&array, SESSION.create_execution_ctx())) - .bench_refs(|(array, ctx)| (**array).clone().into_array().execute::(ctx)); + .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); } /// Helper function to generate random string data. -fn generate_test_data(string_count: usize, avg_len: usize, unique_chars: u8) -> VarBinArray { +fn generate_test_data(string_count: usize, avg_len: usize, unique_chars: u8) -> ArrayRef { let mut rng = StdRng::seed_from_u64(0); let mut strings = Vec::with_capacity(string_count); @@ -214,6 +204,7 @@ fn generate_test_data(string_count: usize, avg_len: usize, unique_chars: u8) -> .map(|opt_s| opt_s.map(Vec::into_boxed_slice)), DType::Binary(Nullability::NonNullable), ) + .into_array() } fn generate_chunked_test_data( @@ -221,15 +212,16 @@ fn generate_chunked_test_data( string_count: usize, avg_len: usize, unique_chars: u8, -) -> ChunkedArray { +) -> ArrayRef { let mut ctx = SESSION.create_execution_ctx(); (0..chunk_size) .map(|_| { let array = generate_test_data(string_count, avg_len, unique_chars); - let compressor = fsst_train_compressor(&array); - let len = array.len(); - let dtype = array.dtype().clone(); - fsst_compress(array, len, &dtype, &compressor, &mut ctx).into_array() + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array() }) .collect::() + .into_array() } diff --git a/encodings/fsst/benches/fsst_like.rs b/encodings/fsst/benches/fsst_like.rs index 28aa3b109ae..414c1b7afaf 100644 --- a/encodings/fsst/benches/fsst_like.rs +++ b/encodings/fsst/benches/fsst_like.rs @@ -14,7 +14,6 @@ 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_array::session::ArraySession; use vortex_fsst::FSSTArray; use vortex_fsst::test_utils::NUM_STRINGS; use vortex_fsst::test_utils::make_fsst_clickbench_urls; @@ -30,8 +29,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fsst::initialize(&session); + session +}); const N: usize = NUM_STRINGS; diff --git a/encodings/fsst/benches/fsst_url_compare.rs b/encodings/fsst/benches/fsst_url_compare.rs index 656cd9f1866..9ff4e3debd5 100644 --- a/encodings/fsst/benches/fsst_url_compare.rs +++ b/encodings/fsst/benches/fsst_url_compare.rs @@ -12,13 +12,14 @@ use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::builtins::ArrayBuiltins; use vortex_array::expr::like; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; -use vortex_array::session::ArraySession; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; use vortex_fsst::test_utils::HIGH_MATCH_DOMAIN; @@ -31,39 +32,54 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fsst::initialize(&session); + session +}); const NUM_URLS: usize = NUM_STRINGS; static URL_DATA: LazyLock = LazyLock::new(generate_url_data); +/// `URL_DATA` as a `VarBinViewArray`, used by the view-input bench variants. +static URL_VIEW_DATA: LazyLock = LazyLock::new(|| { + let mut ctx = SESSION.create_execution_ctx(); + URL_DATA + .clone() + .into_array() + .execute::(&mut ctx) + .unwrap() +}); + // --------------------------------------------------------------------------- // Eq compare benchmarks (FSST pushdown vs canonicalize) // --------------------------------------------------------------------------- /// Pick a concrete URL from the dataset that uses the given domain. fn pick_url_with_domain(data: &VarBinArray, domain: &str) -> String { - use vortex_array::accessor::ArrayAccessor; - data.with_iterator(|iter| { - iter.flatten() - .map(|b| std::str::from_utf8(b).unwrap().to_string()) - .find(|u| u.contains(domain)) - .unwrap_or_else(|| format!("http://{domain}/missing")) - }) + let mut ctx = SESSION.create_execution_ctx(); + let mask = data + .validity() + .unwrap() + .execute_mask(data.len(), &mut ctx) + .unwrap(); + (0..data.len()) + .filter(|&i| mask.value(i)) + .map(|i| unsafe { String::from_utf8_unchecked(data.bytes_at(i).to_vec()) }) + .find(|u| u.contains(domain)) + .unwrap_or_else(|| format!("http://{domain}/missing")) } #[divan::bench] fn eq_pushdown_high_match(bencher: Bencher) { let data = &*URL_DATA; - let compressor = fsst_train_compressor(data); - let fsst_array = fsst_compress( - data, - data.len(), - data.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); let match_url = pick_url_with_domain(data, HIGH_MATCH_DOMAIN); let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS); @@ -72,7 +88,6 @@ fn eq_pushdown_high_match(bencher: Bencher) { .bench_refs(|(fsst_array, constant, ctx)| { fsst_array .clone() - .into_array() .binary(constant.clone().into_array(), Operator::Eq) .unwrap() .execute::(ctx) @@ -83,17 +98,61 @@ fn eq_pushdown_high_match(bencher: Bencher) { #[divan::bench] fn eq_pushdown_low_match(bencher: Bencher) { let data = &*URL_DATA; - let compressor = fsst_train_compressor(data); - let fsst_array = fsst_compress( - data, - data.len(), - data.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); let match_url = pick_url_with_domain(data, LOW_MATCH_DOMAIN); let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS); + bencher + .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) + .bench_refs(|(fsst_array, constant, ctx)| { + fsst_array + .clone() + .binary(constant.clone().into_array(), Operator::Eq) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} + +#[divan::bench] +fn eq_pushdown_high_match_view(bencher: Bencher) { + let data = &*URL_VIEW_DATA; + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); + let match_url = pick_url_with_domain(&URL_DATA, HIGH_MATCH_DOMAIN); + let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS); + + bencher + .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) + .bench_refs(|(fsst_array, constant, ctx)| { + fsst_array + .clone() + .binary(constant.clone().into_array(), Operator::Eq) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} + +#[divan::bench] +fn eq_pushdown_low_match_view(bencher: Bencher) { + let data = &*URL_VIEW_DATA; + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx).unwrap(); + let match_url = pick_url_with_domain(&URL_DATA, LOW_MATCH_DOMAIN); + let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS); + bencher .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, constant, ctx)| { @@ -110,23 +169,20 @@ fn eq_pushdown_low_match(bencher: Bencher) { #[divan::bench] fn eq_canonicalize_high_match(bencher: Bencher) { let data = &*URL_DATA; - let compressor = fsst_train_compressor(data); - let fsst_array = fsst_compress( - data, - data.len(), - data.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); let match_url = pick_url_with_domain(data, HIGH_MATCH_DOMAIN); let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS); bencher .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, constant, ctx)| { - (*fsst_array) + fsst_array .clone() - .into_array() .execute::(ctx) .unwrap() .into_array() @@ -140,23 +196,20 @@ fn eq_canonicalize_high_match(bencher: Bencher) { #[divan::bench] fn eq_canonicalize_low_match(bencher: Bencher) { let data = &*URL_DATA; - let compressor = fsst_train_compressor(data); - let fsst_array = fsst_compress( - data, - data.len(), - data.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); let match_url = pick_url_with_domain(data, LOW_MATCH_DOMAIN); let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS); bencher .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, constant, ctx)| { - (*fsst_array) + fsst_array .clone() - .into_array() .execute::(ctx) .unwrap() .into_array() @@ -174,14 +227,12 @@ fn eq_canonicalize_low_match(bencher: Bencher) { #[divan::bench] fn like_substr_high_match(bencher: Bencher) { let data = &*URL_DATA; - let compressor = fsst_train_compressor(data); - let fsst_array = fsst_compress( - data, - data.len(), - data.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); let pattern = format!("%{HIGH_MATCH_DOMAIN}%"); let expr = like(root(), lit(pattern.as_str())); @@ -190,7 +241,6 @@ fn like_substr_high_match(bencher: Bencher) { .bench_refs(|(fsst_array, ctx)| { fsst_array .clone() - .into_array() .apply(&expr) .unwrap() .execute::(ctx) @@ -201,14 +251,12 @@ fn like_substr_high_match(bencher: Bencher) { #[divan::bench] fn like_substr_low_match(bencher: Bencher) { let data = &*URL_DATA; - let compressor = fsst_train_compressor(data); - let fsst_array = fsst_compress( - data, - data.len(), - data.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); let pattern = format!("%{LOW_MATCH_DOMAIN}%"); let expr = like(root(), lit(pattern.as_str())); @@ -217,7 +265,6 @@ fn like_substr_low_match(bencher: Bencher) { .bench_refs(|(fsst_array, ctx)| { fsst_array .clone() - .into_array() .apply(&expr) .unwrap() .execute::(ctx) diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 83360a69e64..79e6dcb2f0c 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -57,7 +57,6 @@ use vortex_session::registry::CachedId; use crate::canonical::canonicalize_fsst; use crate::canonical::fsst_decode_views; -use crate::kernel::PARENT_KERNELS; use crate::rules::RULES; /// A [`FSST`]-encoded Vortex array. @@ -149,6 +148,25 @@ impl VTable for FSST { } } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_ensure!( + buffers.len() == 3, + "Expected 3 buffers, got {}", + buffers.len() + ); + let symbols = Buffer::::from_byte_buffer(buffers[0].clone().try_to_host_sync()?); + let symbol_lengths = Buffer::::from_byte_buffer(buffers[1].clone().try_to_host_sync()?); + let data = FSSTData::try_new(symbols, symbol_lengths, buffers[2].clone(), array.len())?; + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, @@ -313,15 +331,6 @@ impl VTable for FSST { Ok(()) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, @@ -849,9 +858,8 @@ mod test { use prost::Message; use vortex_array::ArrayPlugin; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; - use vortex_array::accessor::ArrayAccessor; + use vortex_array::array_session; use vortex_array::arrays::VarBinViewArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; @@ -859,14 +867,13 @@ mod test { use vortex_array::dtype::PType; use vortex_array::test_harness::check_metadata; use vortex_buffer::Buffer; - use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_err; use crate::FSST; use crate::array::FSSTArrayExt; use crate::array::FSSTMetadata; - use crate::fsst_compress_iter; + use crate::fsst_compress; #[test] fn slice_reuses_initialized_compressor() -> VortexResult<()> { @@ -877,19 +884,9 @@ mod test { let symbol_lengths = Buffer::::copy_from([3, 8]); let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let fsst_array = fsst_compress_iter( - [ - Some(b"abcabcab".as_ref()), - Some(b"defghijk".as_ref()), - Some(b"abcxyz".as_ref()), - ] - .into_iter(), - 3, - DType::Utf8(Nullability::NonNullable), - &compressor, - &mut ctx, - ); + let mut ctx = array_session().create_execution_ctx(); + let strings = VarBinViewArray::from_iter_str(["abcabcab", "defghijk", "abcxyz"]); + let fsst_array = fsst_compress(&strings.into_array(), &compressor, &mut ctx)?; let compressor_ptr = fsst_array.compressor() as *const Compressor; let sliced = fsst_array @@ -923,7 +920,7 @@ mod test { /// This test manually constructs an old-style FSST array and ensures that it can still be /// deserialized. #[test] - fn test_back_compat() { + fn test_back_compat() -> VortexResult<()> { let symbols = Buffer::::copy_from([ Symbol::from_slice(b"abc00000"), Symbol::from_slice(b"defghijk"), @@ -931,14 +928,9 @@ mod test { let symbol_lengths = Buffer::::copy_from([3, 8]); let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let fsst_array = fsst_compress_iter( - [Some(b"abcabcab".as_ref()), Some(b"defghijk".as_ref())].into_iter(), - 2, - DType::Utf8(Nullability::NonNullable), - &compressor, - &mut ctx, - ); + let mut ctx = array_session().create_execution_ctx(); + let input = VarBinViewArray::from_iter_str(["abcabcab", "defghijk"]); + let fsst_array = fsst_compress(&input.into_array(), &compressor, &mut ctx)?; let compressed_codes = fsst_array.codes(); @@ -974,19 +966,18 @@ mod test { .encode_to_vec(), &buffers, &children.as_slice(), - &LEGACY_SESSION, - ) - .unwrap(); - - let decompressed = fsst - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); - decompressed - .with_iterator(|it| { - assert_eq!(it.next().unwrap(), Some(b"abcabcab".as_ref())); - assert_eq!(it.next().unwrap(), Some(b"defghijk".as_ref())); - Ok::<_, VortexError>(()) - }) - .unwrap() + &array_session(), + )?; + + let decompressed = + fsst.execute::(&mut array_session().create_execution_ctx())?; + let mask = decompressed + .validity()? + .execute_mask(decompressed.len(), &mut ctx)?; + assert!(mask.value(0)); + assert_eq!(decompressed.bytes_at(0).as_slice(), b"abcabcab".as_ref()); + assert!(mask.value(1)); + assert_eq!(decompressed.bytes_at(1).as_slice(), b"defghijk".as_ref()); + Ok(()) } } diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index a8b8171b043..43712aad7dd 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -96,7 +96,6 @@ mod tests { use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; @@ -104,15 +103,13 @@ mod tests { use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; - use vortex_array::session::ArraySession; use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::fsst_compress; use crate::fsst_train_compressor; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn make_data() -> (VarBinArray, Vec>>) { const STRING_COUNT: usize = 1000; @@ -153,9 +150,11 @@ mod tests { let (arr_vec, data_vec): (Vec, Vec>>>) = (0..10) .map(|_| { let (array, data) = make_data(); - let compressor = fsst_train_compressor(&array); + let array = array.into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); ( - fsst_compress(&array, array.len(), array.dtype(), &compressor, &mut ctx) + fsst_compress(&array, &compressor, &mut ctx) + .unwrap() .into_array(), data, ) @@ -182,8 +181,10 @@ mod tests { { let arr = builder.finish_into_canonical().into_varbinview(); - let res1 = - arr.with_iterator(|iter| iter.map(|b| b.map(|v| v.to_vec())).collect::>()); + let mask = arr.validity()?.execute_mask(arr.len(), &mut ctx)?; + let res1 = (0..arr.len()) + .map(|i| mask.value(i).then(|| arr.bytes_at(i).to_vec())) + .collect::>(); assert_eq!(data, res1); }; @@ -192,8 +193,10 @@ mod tests { .as_array() .clone() .execute::(&mut ctx)?; - let res2 = - arr2.with_iterator(|iter| iter.map(|b| b.map(|v| v.to_vec())).collect::>()); + let mask = arr2.validity()?.execute_mask(arr2.len(), &mut ctx)?; + let res2 = (0..arr2.len()) + .map(|i| mask.value(i).then(|| arr2.bytes_at(i).to_vec())) + .collect::>(); assert_eq!(data, res2) }; Ok(()) @@ -208,15 +211,14 @@ mod tests { let varbin = VarBinArray::from_iter( [Some(b"long enough too".to_vec().into_boxed_slice())], dtype, - ); + ) + .into_array(); let mut ctx = SESSION.create_execution_ctx(); let fsst_array = fsst_compress( &varbin, - varbin.len(), - varbin.dtype(), - &fsst_train_compressor(&varbin), + &fsst_train_compressor(&varbin, &mut ctx)?, &mut ctx, - ) + )? .into_array(); fsst_array.append_to_builder(&mut builder, &mut ctx)?; diff --git a/encodings/fsst/src/compress.rs b/encodings/fsst/src/compress.rs index 058e0b7e0e1..b9f815b474f 100644 --- a/encodings/fsst/src/compress.rs +++ b/encodings/fsst/src/compress.rs @@ -1,160 +1,362 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -// Compress a set of values into an Array. +//! FSST compression entry points. +//! +//! [`fsst_compress`] and [`fsst_train_compressor`] take an [`ArrayRef`] and dispatch +//! on the input encoding ([`VarBinView`] or [`VarBin`]). Callers don't need to know +//! which string encoding they hold. + +use std::sync::Arc; use fsst::Compressor; -use fsst::Symbol; +use num_traits::AsPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::accessor::ArrayAccessor; +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::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::match_each_integer_ptype; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::AllOr; +use vortex_mask::Mask; -/// Compress a string array using FSST. use crate::FSST; use crate::FSSTArray; -pub fn fsst_compress>( - strings: A, - len: usize, - dtype: &DType, + +/// FSST worst case: every input byte expands to an escape + literal (2x). +const FSST_PER_BYTE_OVERHEAD: usize = 2; + +/// Starting capacity for the per-row `compress_into` scratch buffer; grown monotonically. +const DEFAULT_BUFFER_LEN: usize = 1024 * 1024; + +/// Compress a string array using FSST. +/// +/// Accepts any [`VarBinView`] or [`VarBin`]-encoded array; other encodings error. +pub fn fsst_compress( + array: &ArrayRef, compressor: &Compressor, ctx: &mut ExecutionCtx, -) -> FSSTArray { - strings.with_iterator(|iter| fsst_compress_iter(iter, len, dtype.clone(), compressor, ctx)) +) -> VortexResult { + if let Some(view) = array.as_opt::() { + compress_varbinview(view, compressor, ctx) + } else if let Some(varbin) = array.as_opt::() { + compress_varbin_array(varbin, compressor, ctx) + } else { + vortex_bail!( + "fsst_compress requires VarBinView or VarBin encoding, got {}", + array.encoding_id() + ) + } } -/// Train a compressor from an array. +/// Train an FSST [`Compressor`] from a string array's non-null rows. /// -/// # Panics -/// -/// If the provided array is not FSST compressible. -pub fn fsst_train_compressor>(array: &A) -> Compressor { - array.with_iterator(|iter| fsst_train_compressor_iter(iter)) +/// Accepts any [`VarBinView`] or [`VarBin`]-encoded array; other encodings error. +pub fn fsst_train_compressor(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(view) = array.as_opt::() { + train_varbinview(view, ctx) + } else if let Some(varbin) = array.as_opt::() { + train_varbin_array(varbin, ctx) + } else { + vortex_bail!( + "fsst_train_compressor requires VarBinView or VarBin encoding, got {}", + array.encoding_id() + ) + } } -/// Train a [compressor][Compressor] from an iterator of bytestrings. -fn fsst_train_compressor_iter<'a, I>(iter: I) -> Compressor -where - I: Iterator>, -{ - let mut lines = Vec::with_capacity(8_192); +fn compress_varbinview( + strings: ArrayView, + compressor: &Compressor, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = strings.validity()?.execute_mask(strings.len(), ctx)?; + let views = strings.views(); + + let total_input_bytes = match mask.bit_buffer() { + AllOr::All => views.iter().map(|v| v.len() as usize).sum(), + AllOr::None => 0, + AllOr::Some(bits) => views + .iter() + .zip(bits.iter()) + .filter(|&(_, b)| b) + .map(|(v, _)| v.len() as usize) + .sum(), + }; - for string in iter { - match string { - None => {} - Some(b) => lines.push(b), + if fsst_output_fits_in_i32_offsets(total_input_bytes) { + compress_views::(strings, &mask, compressor, ctx) + } else { + compress_views::(strings, &mask, compressor, ctx) + } +} + +fn compress_varbin_array( + strings: ArrayView, + compressor: &Compressor, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = strings.validity()?.execute_mask(strings.len(), ctx)?; + let offsets = strings.offsets().clone().execute::(ctx)?; + let total_input_bytes = match_each_integer_ptype!(offsets.ptype(), |O| { + let off = offsets.as_slice::(); + let first: usize = off[0].as_(); + let last: usize = off[off.len() - 1].as_(); + last - first + }); + + if fsst_output_fits_in_i32_offsets(total_input_bytes) { + compress_varbin::(strings, &offsets, &mask, compressor, ctx) + } else { + compress_varbin::(strings, &offsets, &mask, compressor, ctx) + } +} + +fn train_varbinview( + strings: ArrayView, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = strings.validity()?.execute_mask(strings.len(), ctx)?; + let views = strings.views(); + let buffers = strings.data_buffers(); + let mut lines: Vec<&[u8]> = Vec::with_capacity(mask.true_count()); + + match mask.bit_buffer() { + AllOr::All => { + for view in views { + lines.push(view_bytes(view, buffers)); + } + } + AllOr::None => {} + AllOr::Some(bits) => { + for (view, valid) in views.iter().zip(bits.iter()) { + if valid { + lines.push(view_bytes(view, buffers)); + } + } } } - Compressor::train(&lines) + Ok(Compressor::train(&lines)) } -/// Most strings are small in practice. If we encounter a larger string, we reallocate -/// the buffer to hold enough capacity for the worst-case compressed value. -const DEFAULT_BUFFER_LEN: usize = 1024 * 1024; +fn train_varbin_array( + strings: ArrayView, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = strings.validity()?.execute_mask(strings.len(), ctx)?; + let offsets = strings.offsets().clone().execute::(ctx)?; + let bytes = strings.bytes().as_slice(); + let mut lines: Vec<&[u8]> = Vec::with_capacity(mask.true_count()); + + match_each_integer_ptype!(offsets.ptype(), |I| { + let off = offsets.as_slice::(); + for_each_varbin_row(off, bytes, &mask, |row| { + if let Some(s) = row { + lines.push(s); + } + }); + }); + + Ok(Compressor::train(&lines)) +} + +#[inline] +fn fsst_output_fits_in_i32_offsets(total_input_bytes: usize) -> bool { + let worst = total_input_bytes.saturating_mul(FSST_PER_BYTE_OVERHEAD); + worst <= i32::MAX as usize +} -/// Compress from an iterator of bytestrings using FSST. -pub fn fsst_compress_iter<'a, I>( - iter: I, - len: usize, - dtype: DType, +#[inline] +fn view_bytes<'a>(view: &'a BinaryView, buffers: &'a Arc<[BufferHandle]>) -> &'a [u8] { + if view.is_inlined() { + view.as_inlined().value() + } else { + let r = view.as_view(); + &buffers[r.buffer_index as usize].as_host()[r.as_range()] + } +} + +fn compress_views( + strings: ArrayView, + mask: &Mask, compressor: &Compressor, ctx: &mut ExecutionCtx, -) -> FSSTArray +) -> VortexResult where - I: Iterator>, + O: IntegerPType + 'static, { - let mut buffer = Vec::with_capacity(DEFAULT_BUFFER_LEN); - - // Offsets are widened to i64 because the cumulative compressed bytes can exceed i32::MAX for - // large inputs (see issue #7833). Per-string sizes still fit in i32. - let mut builder = VarBinBuilder::::with_capacity(len); - let mut uncompressed_lengths: BufferMut = BufferMut::with_capacity(len); - - for string in iter { - match string { - None => { - builder.append_null(); - uncompressed_lengths.push(0); + let mut sink = FsstSink::::with_capacity(strings.len(), compressor); + let views = strings.views(); + let buffers = strings.data_buffers(); + match mask.bit_buffer() { + AllOr::All => { + for view in views { + sink.emit(Some(view_bytes(view, buffers))); } - Some(s) => { - uncompressed_lengths.push( - s.len() - .try_into() - .vortex_expect("string length must fit in i32"), - ); - - // make sure the buffer is 2x+7 larger than the input - let target_size = 2 * s.len() + 7; - if target_size > buffer.len() { - let additional_capacity = target_size - buffer.len(); - buffer.reserve(additional_capacity); - } + } + AllOr::None => { + for _ in 0..mask.len() { + sink.emit(None); + } + } + AllOr::Some(bits) => { + for (view, valid) in views.iter().zip(bits.iter()) { + sink.emit(valid.then(|| view_bytes(view, buffers))); + } + } + } + sink.finish(strings.dtype().clone(), ctx) +} - // SAFETY: buffer is always sized to be large enough - unsafe { compressor.compress_into(s, &mut buffer) }; +fn compress_varbin( + strings: ArrayView, + offsets: &PrimitiveArray, + mask: &Mask, + compressor: &Compressor, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + O: IntegerPType + 'static, +{ + let mut sink = FsstSink::::with_capacity(strings.len(), compressor); + let bytes = strings.bytes().as_slice(); + match_each_integer_ptype!(offsets.ptype(), |I| { + let off = offsets.as_slice::(); + for_each_varbin_row(off, bytes, mask, |row| sink.emit(row)); + }); + sink.finish(strings.dtype().clone(), ctx) +} - builder.append_value(&buffer); +/// Call `f` once per row of a `VarBinArray` with the row bytes or `None`. +/// Validity dispatch is hoisted out of the per-row loop. +#[inline] +fn for_each_varbin_row<'a, I, F>(off: &[I], bytes: &'a [u8], mask: &Mask, mut f: F) +where + I: IntegerPType + 'static, + F: FnMut(Option<&'a [u8]>), +{ + match mask.bit_buffer() { + AllOr::All => { + for w in off.windows(2) { + f(Some(&bytes[w[0].as_()..w[1].as_()])); + } + } + AllOr::None => { + for _ in 0..mask.len() { + f(None); + } + } + AllOr::Some(bits) => { + for (w, valid) in off.windows(2).zip(bits.iter()) { + f(valid.then(|| &bytes[w[0].as_()..w[1].as_()])); } } } +} + +/// Per-row output state for an FSST compression pass. +struct FsstSink<'c, O: IntegerPType + '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 { + Self { + buffer: Vec::with_capacity(DEFAULT_BUFFER_LEN), + builder: VarBinBuilder::::with_capacity(len), + uncompressed_lengths: BufferMut::with_capacity(len), + compressor, + } + } + + #[inline] + fn emit(&mut self, row: Option<&[u8]>) { + let Some(s) = row else { + self.builder.append_null(); + self.uncompressed_lengths.push(0); + return; + }; + + // A single row > i32::MAX (2 GiB) is not supported. + self.uncompressed_lengths.push( + 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()); + } + + // SAFETY: `self.buffer` has capacity for the FSST worst-case output of `s`. + unsafe { self.compressor.compress_into(s, &mut self.buffer) }; - let codes = builder.finish(DType::Binary(dtype.nullability())); - let symbols: Buffer = Buffer::copy_from(compressor.symbol_table()); - let symbol_lengths: Buffer = Buffer::::copy_from(compressor.symbol_lengths()); - - let uncompressed_lengths = uncompressed_lengths.into_array(); - - FSST::try_new( - dtype, - symbols, - symbol_lengths, - codes, - uncompressed_lengths, - ctx, - ) - .vortex_expect("FSST parts must be valid") + 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( + dtype, + Buffer::copy_from(self.compressor.symbol_table()), + Buffer::::copy_from(self.compressor.symbol_lengths()), + codes, + self.uncompressed_lengths.into_array(), + ctx, + ) + } } #[cfg(test)] mod tests { - use fsst::CompressorBuilder; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::scalar::Scalar; + use vortex_array::array_session; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::varbin::VarBinArrayExt; + use vortex_array::dtype::PType; + use vortex_error::VortexResult; - use crate::compress::DEFAULT_BUFFER_LEN; - use crate::fsst_compress_iter; + use super::fsst_compress; + use super::fsst_output_fits_in_i32_offsets; + use super::fsst_train_compressor; + use crate::array::FSSTArrayExt; + /// Regression for #7833: the i32-vs-i64 codes-offsets decision must cross at + /// `i32::MAX` against the worst-case bound `2 * total + 7 * non_null`. #[test] - fn test_large_string() { - let big_string: String = "abc" - .chars() - .cycle() - .take(10 * DEFAULT_BUFFER_LEN) - .collect(); - - let compressor = CompressorBuilder::default().build(); - - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let compressed = fsst_compress_iter( - [Some(big_string.as_bytes())].into_iter(), - 1, - DType::Utf8(Nullability::NonNullable), - &compressor, - &mut ctx, - ); - - let decoded = compressed.execute_scalar(0, &mut ctx).unwrap(); - - let expected = Scalar::utf8(big_string, Nullability::NonNullable); + fn offset_width_boundary() { + let m = i32::MAX as usize; + assert!(fsst_output_fits_in_i32_offsets(m / 2 - 7)); + assert!(fsst_output_fits_in_i32_offsets(m / 2)); + assert!(fsst_output_fits_in_i32_offsets(0)); + assert!(!fsst_output_fits_in_i32_offsets(usize::MAX)); + } - assert_eq!(decoded, expected); + /// Small inputs fit the i32 bound, so `fsst_compress` must pick i32 offsets. + /// The i64 branch is covered by `tests::fsst_compress_offsets_overflow_i32`. + #[test] + fn codes_offsets_dtype_small_input_is_i32() -> VortexResult<()> { + let array = VarBinViewArray::from_iter_str(["hello", "world", "fsst encoded"]); + let mut ctx = array_session().create_execution_ctx(); + let compressor = fsst_train_compressor(array.as_array(), &mut ctx)?; + let fsst = fsst_compress(array.as_array(), &compressor, &mut ctx)?; + assert_eq!(fsst.codes().offsets().dtype().as_ptype(), PType::I32); + Ok(()) } } diff --git a/encodings/fsst/src/compute/cast.rs b/encodings/fsst/src/compute/cast.rs index 47c324fc2a4..bdca68e841b 100644 --- a/encodings/fsst/src/compute/cast.rs +++ b/encodings/fsst/src/compute/cast.rs @@ -98,34 +98,35 @@ mod tests { use vortex_array::compute::conformance::cast::test_cast_conformance; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; - use vortex_array::session::ArraySession; + use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::fsst_compress; use crate::fsst_train_compressor; + use crate::initialize; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + initialize(&session); + session + }); #[test] - fn test_cast_fsst_nullability() { + fn test_cast_fsst_nullability() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let strings = VarBinArray::from_iter( vec![Some("hello"), Some("world"), Some("hello world")], DType::Utf8(Nullability::NonNullable), - ); + ) + .into_array(); - let compressor = fsst_train_compressor(&strings); - let len = strings.len(); - let dtype = strings.dtype().clone(); - let fsst = fsst_compress(strings, len, &dtype, &compressor, &mut ctx); + let compressor = fsst_train_compressor(&strings, &mut ctx)?; + let fsst = fsst_compress(&strings, &compressor, &mut ctx)?; // Cast to nullable - let casted = fsst - .into_array() - .cast(DType::Utf8(Nullability::Nullable)) - .unwrap(); + let casted = fsst.into_array().cast(DType::Utf8(Nullability::Nullable))?; assert_eq!(casted.dtype(), &DType::Utf8(Nullability::Nullable)); + Ok(()) } #[rstest] @@ -141,10 +142,12 @@ mod tests { vec![Some("test")], DType::Utf8(Nullability::NonNullable) ))] - fn test_cast_fsst_conformance(#[case] array: VarBinArray) { + fn test_cast_fsst_conformance(#[case] array: VarBinArray) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let compressor = fsst_train_compressor(&array); - let fsst = fsst_compress(&array, array.len(), array.dtype(), &compressor, &mut ctx); + let array = array.into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx)?; + let fsst = fsst_compress(&array, &compressor, &mut ctx)?; test_cast_conformance(&fsst.into_array()); + Ok(()) } } diff --git a/encodings/fsst/src/compute/compare.rs b/encodings/fsst/src/compute/compare.rs index 553a3608e75..ec63af3735f 100644 --- a/encodings/fsst/src/compute/compare.rs +++ b/encodings/fsst/src/compute/compare.rs @@ -122,8 +122,9 @@ fn compare_fsst_constant( #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; @@ -134,14 +135,22 @@ mod tests { 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::fsst_compress; use crate::fsst_train_compressor; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] #[cfg_attr(miri, ignore)] - fn test_compare_fsst() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + fn test_compare_fsst() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let lhs = VarBinArray::from_iter( [ Some("hello"), @@ -151,11 +160,10 @@ mod tests { Some("this is a very long string"), ], DType::Utf8(Nullability::Nullable), - ); - let compressor = fsst_train_compressor(&lhs); - let len = lhs.len(); - let dtype = lhs.dtype().clone(); - let lhs = fsst_compress(lhs, len, &dtype, &compressor, &mut ctx); + ) + .into_array(); + let compressor = fsst_train_compressor(&lhs, &mut ctx)?; + let lhs = fsst_compress(&lhs, &compressor, &mut ctx)?; let rhs = ConstantArray::new("world", lhs.len()); @@ -163,31 +171,29 @@ mod tests { let equals = lhs .clone() .into_array() - .binary(rhs.clone().into_array(), Operator::Eq) - .unwrap() - .execute::(&mut ctx) - .unwrap(); + .binary(rhs.clone().into_array(), Operator::Eq)? + .execute::(&mut ctx)?; assert_eq!(equals.dtype(), &DType::Bool(Nullability::Nullable)); assert_arrays_eq!( &equals, - &BoolArray::from_iter([Some(false), None, Some(true), None, Some(false)]) + &BoolArray::from_iter([Some(false), None, Some(true), None, Some(false)]), + &mut ctx ); // Ensure fastpath for Eq exists, and returns correct answer let not_equals = lhs .clone() .into_array() - .binary(rhs.into_array(), Operator::NotEq) - .unwrap() - .execute::(&mut ctx) - .unwrap(); + .binary(rhs.into_array(), Operator::NotEq)? + .execute::(&mut ctx)?; assert_eq!(not_equals.dtype(), &DType::Bool(Nullability::Nullable)); assert_arrays_eq!( ¬_equals, - &BoolArray::from_iter([Some(true), None, Some(false), None, Some(true)]) + &BoolArray::from_iter([Some(true), None, Some(false), None, Some(true)]), + &mut ctx ); // Ensure null constants are handled correctly. @@ -196,20 +202,21 @@ mod tests { let equals_null = lhs .clone() .into_array() - .binary(null_rhs.clone().into_array(), Operator::Eq) - .unwrap(); + .binary(null_rhs.clone().into_array(), Operator::Eq)?; assert_arrays_eq!( &equals_null, - &BoolArray::from_iter([None::, None, None, None, None]) + &BoolArray::from_iter([None::, None, None, None, None]), + &mut ctx ); let noteq_null = lhs .into_array() - .binary(null_rhs.into_array(), Operator::NotEq) - .unwrap(); + .binary(null_rhs.into_array(), Operator::NotEq)?; assert_arrays_eq!( ¬eq_null, - &BoolArray::from_iter([None::, None, None, None, None]) + &BoolArray::from_iter([None::, None, None, None, None]), + &mut ctx ); + Ok(()) } } diff --git a/encodings/fsst/src/compute/like.rs b/encodings/fsst/src/compute/like.rs index 3f704f1e034..c922fba088e 100644 --- a/encodings/fsst/src/compute/like.rs +++ b/encodings/fsst/src/compute/like.rs @@ -96,7 +96,6 @@ mod tests { use vortex_array::scalar_fn::fns::like::Like; use vortex_array::scalar_fn::fns::like::LikeKernel; use vortex_array::scalar_fn::fns::like::LikeOptions; - use vortex_array::session::ArraySession; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -105,21 +104,18 @@ mod tests { use crate::fsst_compress; use crate::fsst_train_compressor; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); fn make_fsst(strings: &[Option<&str>], nullability: Nullability) -> FSSTArray { - let varbin = VarBinArray::from_iter(strings.iter().copied(), DType::Utf8(nullability)); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress( - varbin, - len, - &dtype, - &compressor, - &mut SESSION.create_execution_ctx(), - ) + let array = + VarBinArray::from_iter(strings.iter().copied(), DType::Utf8(nullability)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + fsst_compress(&array, &compressor, &mut ctx).unwrap() } fn run_like(array: FSSTArray, pattern: &str, opts: LikeOptions) -> VortexResult { @@ -152,7 +148,8 @@ mod tests { let result = like(fsst, "http%")?; assert_arrays_eq!( &result, - &BoolArray::from_iter([true, true, false, true, false]) + &BoolArray::from_iter([true, true, false, true, false]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -166,7 +163,8 @@ mod tests { let result = like(fsst, "hel%")?; // spellchecker:disable-line assert_arrays_eq!( &result, - &BoolArray::from_iter([Some(true), None, Some(true), None, Some(false)]) + &BoolArray::from_iter([Some(true), None, Some(true), None, Some(false)]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -183,7 +181,11 @@ mod tests { Nullability::NonNullable, ); let result = like(fsst, "%hello%")?; - assert_arrays_eq!(&result, &BoolArray::from_iter([true, true, false, true])); + assert_arrays_eq!( + &result, + &BoolArray::from_iter([true, true, false, true]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -199,7 +201,11 @@ mod tests { Nullability::NonNullable, ); let result = like(fsst, "%lazy dog%")?; - assert_arrays_eq!(&result, &BoolArray::from_iter([true, false, true, false])); + assert_arrays_eq!( + &result, + &BoolArray::from_iter([true, false, true, false]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -214,7 +220,11 @@ mod tests { case_insensitive: false, }; let result = run_like(fsst, "%sdf%", opts)?; - assert_arrays_eq!(&result, &BoolArray::from_iter([false, false, true])); + assert_arrays_eq!( + &result, + &BoolArray::from_iter([false, false, true]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -225,7 +235,11 @@ mod tests { Nullability::NonNullable, ); let result = like(fsst, "%")?; - assert_arrays_eq!(&result, &BoolArray::from_iter([true, true, true])); + assert_arrays_eq!( + &result, + &BoolArray::from_iter([true, true, true]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -244,7 +258,11 @@ mod tests { let fsst = fsst.as_view(); let result = ::like(fsst, &pattern, LikeOptions::default(), &mut ctx)?; assert!(result.is_some(), "FSST LikeKernel should handle prefix%"); - assert_arrays_eq!(result.unwrap(), BoolArray::from_iter([true, false])); + assert_arrays_eq!( + result.unwrap(), + BoolArray::from_iter([true, false]), + &mut ctx + ); Ok(()) } @@ -261,7 +279,11 @@ mod tests { let fsst = fsst.as_view(); let result = ::like(fsst, &pattern, LikeOptions::default(), &mut ctx)?; assert!(result.is_some(), "FSST LikeKernel should handle %needle%"); - assert_arrays_eq!(result.unwrap(), BoolArray::from_iter([true, false])); + assert_arrays_eq!( + result.unwrap(), + BoolArray::from_iter([true, false]), + &mut ctx + ); Ok(()) } @@ -319,7 +341,8 @@ mod tests { assert!(result.is_some(), "escaped percent prefix should use FSST"); assert_arrays_eq!( result.unwrap(), - BoolArray::from_iter([true, false, false, false, false, false, false]) + BoolArray::from_iter([true, false, false, false, false, false, false]), + &mut ctx ); let pattern = ConstantArray::new(r"\_%", fsst.len()).into_array(); @@ -331,7 +354,8 @@ mod tests { ); assert_arrays_eq!( result.unwrap(), - BoolArray::from_iter([false, true, false, false, false, false, false]) + BoolArray::from_iter([false, true, false, false, false, false, false]), + &mut ctx ); let pattern = ConstantArray::new(r"\\%", fsst.len()).into_array(); @@ -340,7 +364,8 @@ mod tests { assert!(result.is_some(), "escaped backslash prefix should use FSST"); assert_arrays_eq!( result.unwrap(), - BoolArray::from_iter([false, false, true, false, false, false, false]) + BoolArray::from_iter([false, false, true, false, false, false, false]), + &mut ctx ); let pattern = ConstantArray::new(r"%\%%", fsst.len()).into_array(); @@ -349,7 +374,8 @@ mod tests { assert!(result.is_some(), "escaped percent contains should use FSST"); assert_arrays_eq!( result.unwrap(), - BoolArray::from_iter([true, false, false, true, false, false, false]) + BoolArray::from_iter([true, false, false, true, false, false, false]), + &mut ctx ); let pattern = ConstantArray::new(r"%\_%", fsst.len()).into_array(); @@ -361,7 +387,8 @@ mod tests { ); assert_arrays_eq!( result.unwrap(), - BoolArray::from_iter([false, true, false, false, true, false, false]) + BoolArray::from_iter([false, true, false, false, true, false, false]), + &mut ctx ); let pattern = ConstantArray::new(r"%\\%", fsst.len()).into_array(); @@ -373,7 +400,8 @@ mod tests { ); assert_arrays_eq!( result.unwrap(), - BoolArray::from_iter([false, false, true, false, false, true, false]) + BoolArray::from_iter([false, false, true, false, false, true, false]), + &mut ctx ); Ok(()) @@ -402,7 +430,11 @@ mod tests { direct.is_some(), "14-byte prefixes are now handled by the flat prefix DFA" ); - assert_arrays_eq!(direct.unwrap(), BoolArray::from_iter([true, false, true])); + assert_arrays_eq!( + direct.unwrap(), + BoolArray::from_iter([true, false, true]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -432,7 +464,11 @@ mod tests { ); let result = like(fsst, &pattern)?; - assert_arrays_eq!(&result, &BoolArray::from_iter([true, false, true])); + assert_arrays_eq!( + &result, + &BoolArray::from_iter([true, false, true]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -459,7 +495,11 @@ mod tests { direct.is_some(), "254-byte contains needle should stay on the DFA path" ); - assert_arrays_eq!(direct.unwrap(), BoolArray::from_iter([true, false, true])); + assert_arrays_eq!( + direct.unwrap(), + BoolArray::from_iter([true, false, true]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } } diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index d95725b9fbc..0065e0f499a 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -61,39 +61,42 @@ mod tests { use rstest::rstest; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::compute::conformance::take::test_take_conformance; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; + use vortex_error::VortexResult; use crate::FSSTArray; use crate::fsst_compress; use crate::fsst_train_compressor; #[test] - fn test_take_null() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let arr = VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::NonNullable)); - let compr = fsst_train_compressor(&arr); - let fsst = fsst_compress(&arr, arr.len(), arr.dtype(), &compr, &mut ctx); + fn test_take_null() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = + VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::NonNullable)).into_array(); + let compr = fsst_train_compressor(&arr, &mut ctx)?; + let fsst = fsst_compress(&arr, &compr, &mut ctx)?; let idx1: PrimitiveArray = (0..1).collect(); assert_eq!( - fsst.take(idx1.into_array()).unwrap().dtype(), + fsst.take(idx1.into_array())?.dtype(), &DType::Utf8(Nullability::NonNullable) ); let idx2: PrimitiveArray = PrimitiveArray::from_option_iter(vec![Some(0)]); assert_eq!( - fsst.take(idx2.into_array()).unwrap().dtype(), + fsst.take(idx2.into_array())?.dtype(), &DType::Utf8(Nullability::Nullable) ); + Ok(()) } #[rstest] @@ -109,11 +112,13 @@ mod tests { ["single element"].map(Some), DType::Utf8(Nullability::NonNullable), ))] - fn test_take_fsst_conformance(#[case] varbin: VarBinArray) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let compressor = fsst_train_compressor(&varbin); - let array = fsst_compress(&varbin, varbin.len(), varbin.dtype(), &compressor, &mut ctx); + fn test_take_fsst_conformance(#[case] varbin: VarBinArray) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let varbin = varbin.into_array(); + let compressor = fsst_train_compressor(&varbin, &mut ctx)?; + let array = fsst_compress(&varbin, &compressor, &mut ctx)?; test_take_conformance(&array.into_array()); + Ok(()) } type FsstBuilder = fn(&mut ExecutionCtx) -> FSSTArray; @@ -121,51 +126,47 @@ mod tests { #[rstest] // Basic string arrays #[case::fsst_simple(|ctx: &mut ExecutionCtx| { - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( ["hello world", "testing fsst", "compression test", "data array", "vortex encoding"].map(Some), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - fsst_compress(&varbin, varbin.len(), varbin.dtype(), &compressor, ctx) + ).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() })] // Nullable strings #[case::fsst_nullable(|ctx: &mut ExecutionCtx| { - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( [Some("hello"), None, Some("world"), Some("test"), None], DType::Utf8(Nullability::Nullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() })] // Repetitive patterns (good for FSST compression) #[case::fsst_repetitive(|ctx: &mut ExecutionCtx| { - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( ["http://example.com", "http://test.com", "http://vortex.dev", "http://data.org"].map(Some), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - fsst_compress(&varbin, varbin.len(), varbin.dtype(), &compressor, ctx) + ).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() })] // Edge cases #[case::fsst_single(|ctx: &mut ExecutionCtx| { - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( ["single element"].map(Some), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - fsst_compress(&varbin, varbin.len(), varbin.dtype(), &compressor, ctx) + ).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() })] #[case::fsst_empty_strings(|ctx: &mut ExecutionCtx| { - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( ["", "test", "", "hello", ""].map(Some), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() })] // Large arrays #[case::fsst_large(|ctx: &mut ExecutionCtx| { @@ -183,16 +184,17 @@ mod tests { _ => "CREATE TABLE data (id INT, value TEXT)", })) .collect(); - let varbin = VarBinArray::from_iter(data, DType::Utf8(Nullability::NonNullable)); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + let array = VarBinArray::from_iter(data, DType::Utf8(Nullability::NonNullable)).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() })] fn test_fsst_consistency(#[case] build: FsstBuilder) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = build(&mut ctx); - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/encodings/fsst/src/dfa/tests.rs b/encodings/fsst/src/dfa/tests.rs index d41fedde330..3ea9a507c20 100644 --- a/encodings/fsst/src/dfa/tests.rs +++ b/encodings/fsst/src/dfa/tests.rs @@ -20,7 +20,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar_fn::fns::like::Like; use vortex_array::scalar_fn::fns::like::LikeOptions; -use vortex_array::session::ArraySession; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -32,8 +31,11 @@ use crate::FSSTArray; use crate::fsst_compress; use crate::fsst_train_compressor; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); /// Helper: make a Symbol from a byte string (up to 8 bytes, zero-padded). fn sym(bytes: &[u8]) -> Symbol { @@ -266,20 +268,14 @@ fn test_contains_pushdown_rejects_len_255() { // --------------------------------------------------------------------------- fn make_fsst_str(strings: &[Option<&str>]) -> FSSTArray { - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( strings.iter().copied(), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress( - varbin, - len, - &dtype, - &compressor, - &mut SESSION.create_execution_ctx(), ) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + fsst_compress(&array, &compressor, &mut ctx).unwrap() } fn run_like(array: FSSTArray, pattern_arr: ArrayRef) -> VortexResult { @@ -364,6 +360,7 @@ fn test_like_edge_cases( ConstantArray::new(pattern, opts.len()).into_array(), )?; let expected_arr = BoolArray::from_iter(expected.iter().copied()); - assert_arrays_eq!(&result, &expected_arr); + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(&result, &expected_arr, &mut ctx); Ok(()) } diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index 386da3b0f60..30f25006195 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -1,24 +1,34 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +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_array::scalar_fn::fns::like::Like; use vortex_array::scalar_fn::fns::like::LikeExecuteAdaptor; +use vortex_session::VortexSession; use crate::FSST; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(FSST)), - ParentKernelSet::lift(&CompareExecuteAdaptor(FSST)), - ParentKernelSet::lift(&FilterExecuteAdaptor(FSST)), - ParentKernelSet::lift(&TakeExecuteAdaptor(FSST)), - ParentKernelSet::lift(&LikeExecuteAdaptor(FSST)), - ParentKernelSet::lift(&ByteLengthExecuteAdaptor(FSST)), -]); +pub(super) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), FSST, CastExecuteAdaptor(FSST)); + kernels.register_execute_parent_kernel(Binary.id(), FSST, CompareExecuteAdaptor(FSST)); + kernels.register_execute_parent_kernel(Filter.id(), FSST, FilterExecuteAdaptor(FSST)); + kernels.register_execute_parent_kernel(Dict.id(), FSST, TakeExecuteAdaptor(FSST)); + kernels.register_execute_parent_kernel(Like.id(), FSST, LikeExecuteAdaptor(FSST)); + kernels.register_execute_parent_kernel(ByteLength.id(), FSST, ByteLengthExecuteAdaptor(FSST)); +} #[cfg(test)] mod tests { @@ -36,7 +46,6 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::expr::byte_length; use vortex_array::expr::root; - use vortex_array::session::ArraySession; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -45,8 +54,11 @@ mod tests { use crate::fsst_compress; use crate::fsst_train_compressor; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); fn build_test_fsst_array() -> ArrayRef { let mut builder = VarBinBuilder::::with_capacity(10); @@ -62,11 +74,12 @@ mod tests { builder.append_value(b"final string"); let input = builder.finish(DType::Utf8(Nullability::NonNullable)); - let compressor = fsst_train_compressor(&input); - let len = input.len(); - let dtype = input.dtype().clone(); let mut ctx = SESSION.create_execution_ctx(); - fsst_compress(input, len, &dtype, &compressor, &mut ctx).into_array() + let arr = input.into_array(); + let compressor = fsst_train_compressor(&arr, &mut ctx).unwrap(); + fsst_compress(&arr, &compressor, &mut ctx) + .unwrap() + .into_array() } #[test] @@ -89,7 +102,7 @@ mod tests { let expected = fsst_array.filter(mask)?; assert_eq!(result.len(), 2); - assert_arrays_eq!(result.into_array(), expected); + assert_arrays_eq!(result.into_array(), expected, &mut ctx); Ok(()) } @@ -109,13 +122,12 @@ mod tests { let expected = fsst_array.filter(mask)?; assert_eq!(result.len(), 5); - assert_arrays_eq!(result.into_array(), expected); + assert_arrays_eq!(result.into_array(), expected, &mut ctx); Ok(()) } #[test] fn issues_6034_test_fsst_filter_with_nulls_and_special_chars() -> VortexResult<()> { - // // 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) @@ -137,17 +149,11 @@ mod tests { builder.append_null(); } let input = builder.finish(DType::Utf8(Nullability::Nullable)); + let array = input.clone().into_array(); - let compressor = fsst_train_compressor(&input); let mut ctx = SESSION.create_execution_ctx(); - let fsst_array: ArrayRef = fsst_compress( - input.clone(), - input.len(), - input.dtype(), - &compressor, - &mut ctx, - ) - .into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx)?; + let fsst_array: ArrayRef = fsst_compress(&array, &compressor, &mut ctx)?.into_array(); // Filter: only select the last element (index 22) let mut mask = vec![false; 22]; @@ -160,7 +166,7 @@ mod tests { let expected = input.filter(mask)?; assert_eq!(result.len(), 1); - assert_arrays_eq!(result.into_array(), expected); + assert_arrays_eq!(result.into_array(), expected, &mut ctx); Ok(()) } @@ -172,17 +178,11 @@ mod tests { builder.append_null(); let input = builder.finish(DType::Utf8(Nullability::Nullable)); + let array = input.clone().into_array(); - let compressor = fsst_train_compressor(&input); let mut ctx = SESSION.create_execution_ctx(); - let fsst_array: ArrayRef = fsst_compress( - input.clone(), - input.len(), - input.dtype(), - &compressor, - &mut ctx, - ) - .into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx)?; + let fsst_array: ArrayRef = fsst_compress(&array, &compressor, &mut ctx)?.into_array(); let mask = Mask::from_iter([true, false, true]); @@ -192,7 +192,7 @@ mod tests { let expected = input.filter(mask)?; assert_eq!(result.len(), 2); - assert_arrays_eq!(result.into_array(), expected); + assert_arrays_eq!(result.into_array(), expected, &mut ctx); Ok(()) } @@ -207,7 +207,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let result = filter_array.execute::(&mut ctx)?.into_array(); - assert_arrays_eq!(result, fsst_array); + assert_arrays_eq!(result, fsst_array, &mut ctx); Ok(()) } @@ -219,15 +219,15 @@ mod tests { builder.append_value("Пуховички"); // 9 characters, 18 bytes builder.append_value(b""); - let varbin = builder.finish(DType::Utf8(Nullability::NonNullable)); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); + let varbin = builder + .finish(DType::Utf8(Nullability::NonNullable)) + .into_array(); let mut ctx = SESSION.create_execution_ctx(); - let fsst = fsst_compress(varbin, len, &dtype, &compressor, &mut ctx).into_array(); + let compressor = fsst_train_compressor(&varbin, &mut ctx)?; + let fsst = fsst_compress(&varbin, &compressor, &mut ctx)?.into_array(); let result = fsst.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(vec![5u64, 7, 18, 0]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } } diff --git a/encodings/fsst/src/lib.rs b/encodings/fsst/src/lib.rs index 3305c0e66fc..70dcc705249 100644 --- a/encodings/fsst/src/lib.rs +++ b/encodings/fsst/src/lib.rs @@ -27,3 +27,11 @@ mod tests; pub use array::*; pub use compress::*; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; + +/// Initialize FSST encoding in the given session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(FSST); + kernel::initialize(session); +} diff --git a/encodings/fsst/src/test_utils.rs b/encodings/fsst/src/test_utils.rs index eeca412e377..7becb34ffbc 100644 --- a/encodings/fsst/src/test_utils.rs +++ b/encodings/fsst/src/test_utils.rs @@ -41,17 +41,18 @@ pub fn gen_fsst_test_data( )); } - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( strings .into_iter() .map(|opt_s| opt_s.map(Vec::into_boxed_slice)), DType::Binary(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx).into_array() + fsst_compress(&array, &compressor, ctx) + .unwrap() + .into_array() } pub fn gen_dict_fsst_test_data( @@ -144,11 +145,9 @@ pub fn generate_url_data_n(n: usize) -> VarBinArray { } pub fn make_fsst_urls(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { - let varbin = generate_url_data_n(n); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + let array = generate_url_data_n(n).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -237,14 +236,13 @@ pub fn generate_clickbench_urls(n: usize) -> Vec { pub fn make_fsst_clickbench_urls(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let urls = generate_clickbench_urls(n); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( urls.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -305,14 +303,13 @@ pub fn generate_short_urls(n: usize) -> Vec { pub fn make_fsst_short_urls(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let urls = generate_short_urls(n); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( urls.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -377,14 +374,13 @@ pub fn generate_log_lines(n: usize) -> Vec { pub fn make_fsst_log_lines(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let lines = generate_log_lines(n); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( lines.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -436,14 +432,13 @@ pub fn generate_json_strings(n: usize) -> Vec { pub fn make_fsst_json_strings(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let jsons = generate_json_strings(n); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( jsons.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -508,14 +503,13 @@ pub fn generate_file_paths(n: usize) -> Vec { pub fn make_fsst_file_paths(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let paths = generate_file_paths(n); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( paths.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -561,14 +555,13 @@ pub fn generate_emails(n: usize) -> Vec { pub fn make_fsst_emails(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let emails = generate_emails(n); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( emails.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -600,12 +593,11 @@ pub fn generate_rare_match_strings(n: usize, match_rate: f64) -> Vec { pub fn make_fsst_rare_match(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let strings = generate_rare_match_strings(n, 0.00001); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( strings.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } diff --git a/encodings/fsst/src/tests.rs b/encodings/fsst/src/tests.rs index 481b84865df..43cdeb1a02c 100644 --- a/encodings/fsst/src/tests.rs +++ b/encodings/fsst/src/tests.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::LazyLock; + use fsst::CompressorBuilder; use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbin::builder::VarBinBuilder; @@ -14,47 +16,58 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_buffer::buffer; use vortex_mask::Mask; +use vortex_session::VortexSession; use crate::FSST; use crate::fsst_compress; use crate::fsst_train_compressor; +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); + /// this function is VERY slow on miri, so we only want to run it once -pub(crate) fn build_fsst_array() -> ArrayRef { +pub(crate) fn build_fsst_array(ctx: &mut ExecutionCtx) -> ArrayRef { let mut input_array = VarBinBuilder::::with_capacity(3); 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)); + let input_array = input_array + .finish(DType::Utf8(Nullability::NonNullable)) + .into_array(); - let compressor = fsst_train_compressor(&input_array); - let len = input_array.len(); - let dtype = input_array.dtype().clone(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - fsst_compress(input_array, len, &dtype, &compressor, &mut ctx).into_array() + let compressor = fsst_train_compressor(&input_array, ctx).unwrap(); + fsst_compress(&input_array, &compressor, ctx) + .unwrap() + .into_array() } #[test] fn test_fsst_array_ops() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // first test the scalar_at values - let fsst_array = build_fsst_array(); + let fsst_array = build_fsst_array(&mut ctx); assert_nth_scalar!( fsst_array, 0, - "The Greeks never said that the limit could not be overstepped" + "The Greeks never said that the limit could not be overstepped", + &mut ctx ); assert_nth_scalar!( fsst_array, 1, - "They said it existed and that whoever dared to exceed it was mercilessly struck down" + "They said it existed and that whoever dared to exceed it was mercilessly struck down", + &mut ctx ); assert_nth_scalar!( fsst_array, 2, - "Nothing in present history can contradict them" + "Nothing in present history can contradict them", + &mut ctx ); // test slice @@ -64,12 +77,14 @@ fn test_fsst_array_ops() { assert_nth_scalar!( fsst_sliced, 0, - "They said it existed and that whoever dared to exceed it was mercilessly struck down" + "They said it existed and that whoever dared to exceed it was mercilessly struck down", + &mut ctx ); assert_nth_scalar!( fsst_sliced, 1, - "Nothing in present history can contradict them" + "Nothing in present history can contradict them", + &mut ctx ); // test take @@ -79,12 +94,14 @@ fn test_fsst_array_ops() { assert_nth_scalar!( fsst_taken, 0, - "The Greeks never said that the limit could not be overstepped" + "The Greeks never said that the limit could not be overstepped", + &mut ctx ); assert_nth_scalar!( fsst_taken, 1, - "Nothing in present history can contradict them" + "Nothing in present history can contradict them", + &mut ctx ); // test filter @@ -96,7 +113,8 @@ fn test_fsst_array_ops() { assert_nth_scalar!( fsst_filtered, 0, - "They said it existed and that whoever dared to exceed it was mercilessly struck down" + "They said it existed and that whoever dared to exceed it was mercilessly struck down", + &mut ctx ); // test to_canonical @@ -106,13 +124,13 @@ fn test_fsst_array_ops() { .unwrap() .into_array(); - assert_arrays_eq!(fsst_array, canonical_array); + assert_arrays_eq!(fsst_array, canonical_array, &mut ctx); } // TODO(someone): ideally CI would run this in release mode as well since debug builds make the // allocation and compression loop substantially slower. /// Regression for #7833: [`fsst_compress`] must accept inputs whose cumulative compressed -/// bytes exceed [`i32::MAX`]. Before the fix, [`fsst_compress_iter`] hardcoded +/// bytes exceed [`i32::MAX`]. Before the fix, the compress path hardcoded /// [`VarBinBuilder`] for the FSST output and panicked in /// [`VarBinBuilder::append_value`] once cumulative compressed bytes crossed the boundary. /// @@ -132,7 +150,7 @@ fn test_fsst_array_ops() { /// CI=1 cargo test --release -p vortex-fsst fsst_compress_offsets /// ``` /// -/// [`fsst_compress_iter`]: crate::compress::fsst_compress_iter +/// [`fsst_compress`]: crate::compress::fsst_compress #[test_with::env(CI)] #[test_with::no_env(VORTEX_SKIP_SLOW_TESTS)] fn fsst_compress_offsets_overflow_i32() { @@ -148,15 +166,16 @@ fn fsst_compress_offsets_overflow_i32() { for _ in 0..N { builder.append_value(&string); } - let array = builder.finish(DType::Utf8(Nullability::NonNullable)); + let array = builder + .finish(DType::Utf8(Nullability::NonNullable)) + .into_array(); let compressor = CompressorBuilder::default().build(); let len = array.len(); - let dtype = array.dtype().clone(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); println!("compressing to FSST"); - let compressed = fsst_compress(array, len, &dtype, &compressor, &mut ctx); + let compressed = fsst_compress(&array, &compressor, &mut ctx).unwrap(); assert_eq!(compressed.len(), len); // Prove the regression condition was exercised: compressed bytes crossed i32::MAX. assert!(compressed.codes_bytes().len() > i32::MAX as usize); diff --git a/encodings/parquet-variant/Cargo.toml b/encodings/parquet-variant/Cargo.toml index 6b3900abb27..d0c0a0d418e 100644 --- a/encodings/parquet-variant/Cargo.toml +++ b/encodings/parquet-variant/Cargo.toml @@ -27,8 +27,9 @@ prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } +vortex-json = { workspace = true } vortex-mask = { workspace = true } -vortex-proto = { workspace = true } +vortex-proto = { workspace = true, features = ["expr"] } vortex-session = { workspace = true } [dev-dependencies] diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index 37247430c37..7e55357c78f 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -19,9 +19,11 @@ use vortex_array::arrays::List; use vortex_array::arrays::ListArray; use vortex_array::arrays::Struct; 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::struct_::StructArrayExt; +use vortex_array::arrays::variant::VariantArrayExt; #[expect( deprecated, reason = "TODO(aduffy): figure out what to do with Parquet Variant" @@ -42,6 +44,7 @@ use vortex_array::vtable::validity_to_child; use vortex_buffer::BitBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::ParquetVariant; use crate::ParquetVariantArray; @@ -307,6 +310,111 @@ fn inferred_shredded_field_validity( Ok(Validity::from_mask(validity, Nullability::Nullable)) } +/// Reconstructs a Parquet `typed_value` tree from the storage-agnostic canonical shredded tree. +/// +/// This is the inverse of [`logical_shredded_from_parquet_typed_value`]: canonicalization strips +/// the Parquet `value`/`typed_value` wrapper shells out of the shredded tree (representing +/// partially shredded fields as nested [`VariantArray`]s), and this re-adds them. The result is a +/// valid Parquet `typed_value` that can be reattached to a [`ParquetVariant`] and serialized to +/// Arrow, where [`parquet_variant_compute::unshred_variant`] merges it back with the residual +/// `value`. +/// +/// `forward` then `inverse` is structure-preserving for the shapes canonicalization produces; +/// fields the forward transform omitted (Parquet wrapper shells with no `typed_value`) are served +/// from the residual `value` and are intentionally not reconstructed here. +pub(crate) fn parquet_typed_value_from_logical_shredded( + shredded: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if let Some(list_array) = shredded.as_opt::() { + let elements = parquet_shredded_field_from_logical(list_array.elements().clone(), ctx)?; + return Ok(ListArray::try_new( + elements, + list_array.offsets().clone(), + list_array.list_validity(), + )? + .into_array()); + } + + let Some(struct_array) = shredded.as_opt::() else { + // A bare typed leaf (a fully shredded scalar) is already a valid Parquet `typed_value`. + return Ok(shredded); + }; + + let mut names = Vec::with_capacity(struct_array.names().len()); + let mut fields = Vec::with_capacity(struct_array.names().len()); + for (name, field) in struct_array + .names() + .iter() + .zip(struct_array.iter_unmasked_fields()) + { + names.push(FieldName::from(name.as_ref())); + fields.push(parquet_shredded_field_from_logical(field.clone(), ctx)?); + } + + Ok(StructArray::try_new( + FieldNames::from_iter(names), + fields, + struct_array.len(), + struct_array.validity()?, + )? + .into_array()) +} + +/// Reconstructs one Parquet shredded field shell (`{value?, typed_value}`) from its canonical +/// representation, the inverse of [`logical_shredded_from_parquet_field`]. +fn parquet_shredded_field_from_logical( + logical_field: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = logical_field.len(); + + // Partially shredded fields canonicalize to a nested Variant whose core storage holds the + // residual `value` and whose own shredded tree holds the typed children. + if let Some(variant) = logical_field.as_opt::() { + let core = variant + .core_storage() + .as_opt::() + .ok_or_else(|| { + vortex_err!( + "cannot rebuild Parquet shredded field: nested Variant lacks Parquet Variant core storage" + ) + })?; + let value = core.value_array().cloned().ok_or_else(|| { + vortex_err!("cannot rebuild Parquet shredded field: partially shredded Variant has no residual value") + })?; + let typed_value = variant + .shredded() + .cloned() + .map(|shredded| parquet_typed_value_from_logical_shredded(shredded, ctx)) + .transpose()?; + + let mut names = vec![FieldName::from("value")]; + let mut fields = vec![value]; + if let Some(typed_value) = typed_value { + names.push(FieldName::from("typed_value")); + fields.push(typed_value); + } + return Ok(StructArray::try_new( + FieldNames::from_iter(names), + fields, + len, + Validity::NonNullable, + )? + .into_array()); + } + + // Fully shredded field: rebuild its typed subtree and wrap it in a `typed_value`-only shell. + let typed_value = parquet_typed_value_from_logical_shredded(logical_field, ctx)?; + Ok(StructArray::try_new( + FieldNames::from_iter([FieldName::from("typed_value")]), + vec![typed_value], + len, + Validity::NonNullable, + )? + .into_array()) +} + /// Accessors and Arrow conversion for Parquet Variant storage arrays. pub trait ParquetVariantArrayExt: TypedArrayRef { /// Returns the non-nullable Parquet Variant metadata child. @@ -385,10 +493,12 @@ impl> ParquetVariantArrayExt for T {} #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use arrow_array::Array as _; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::Int32Array; + use arrow_array::StringArray; use arrow_array::StructArray; use arrow_array::builder::BinaryViewBuilder; use arrow_buffer::NullBuffer; @@ -396,13 +506,17 @@ mod tests { use arrow_schema::Field; use arrow_schema::Fields; use parquet_variant::Variant as PqVariant; + use parquet_variant_compute::ShreddedSchemaBuilder; use parquet_variant_compute::VariantArray as ArrowVariantArray; use parquet_variant_compute::VariantArrayBuilder; + use parquet_variant_compute::json_to_variant; + use parquet_variant_compute::shred_variant; + use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::variant::VariantArrayExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -410,9 +524,17 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; + use vortex_session::VortexSession; use crate::ParquetVariant; use crate::array::ParquetVariantArrayExt; + use crate::array::parquet_typed_value_from_logical_shredded; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); fn assert_arrow_variant_storage_roundtrip(struct_array: StructArray) -> VortexResult<()> { let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; @@ -421,7 +543,7 @@ mod tests { .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant child"))?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let roundtripped = inner.to_arrow(&mut ctx)?; let roundtripped = roundtripped.inner(); @@ -513,10 +635,11 @@ mod tests { .typed_value_array() .ok_or_else(|| vortex_err!("expected typed_value child"))? .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut SESSION.create_execution_ctx())?; assert_arrays_eq!( typed_value, - PrimitiveArray::from_option_iter([Some(10), None, Some(30)]) + PrimitiveArray::from_option_iter([Some(10), None, Some(30)]), + &mut SESSION.create_execution_ctx() ); Ok(()) @@ -528,7 +651,7 @@ mod tests { let value = VarBinViewArray::from_iter_bin([b"\x10", b"\x11"]).into_array(); let pv_array = ParquetVariant::try_new(Validity::NonNullable, metadata, Some(value), None)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let variant_arr = pv_array.to_arrow(&mut ctx)?; let struct_arr = variant_arr.inner(); @@ -550,7 +673,7 @@ mod tests { Some(typed_value), )?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let variant_arr = pv_array.to_arrow(&mut ctx)?; let struct_arr = variant_arr.inner(); @@ -599,8 +722,12 @@ mod tests { .typed_value_array() .ok_or_else(|| vortex_err!("expected typed_value child"))? .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; - assert_arrays_eq!(typed_value, PrimitiveArray::from_iter([10i32, 20, 30])); + .execute::(&mut SESSION.create_execution_ctx())?; + assert_arrays_eq!( + typed_value, + PrimitiveArray::from_iter([10i32, 20, 30]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -629,7 +756,7 @@ mod tests { assert!(parquet_array.value_array().is_some()); assert!(parquet_array.typed_value_array().is_some()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let roundtripped = parquet_array.to_arrow(&mut ctx)?; let roundtripped = roundtripped.inner(); assert_eq!( @@ -678,4 +805,64 @@ mod tests { assert_arrow_variant_storage_roundtrip(struct_array) } + + /// `parquet_typed_value_from_logical_shredded` must invert the wrapper-stripping that + /// canonicalization performs: an object-shredded Parquet variant, once canonicalized and then + /// rebuilt, must produce the same per-row values as the original. + #[test] + fn parquet_typed_value_inverse_roundtrips_object_shredding() -> VortexResult<()> { + // Shred `$.a` as Int32 over conforming, non-conforming, and missing-field rows. + let json: ArrowArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"{"a":1,"b":"x"}"#), + Some(r#"{"a":"not-a-number","b":"y"}"#), + Some(r#"{"b":"z"}"#), + ])); + let shredding = ShreddedSchemaBuilder::new() + .with_path("a", &DataType::Int32)? + .build(); + let shredded = shred_variant(&json_to_variant(&json)?, &shredding)?; + let original = ParquetVariant::from_arrow_variant(&shredded)?; + assert!( + original + .as_opt::() + .ok_or_else(|| vortex_err!("expected parquet variant"))? + .typed_value_array() + .is_some(), + "fixture must be shredded" + ); + + // Canonicalize: the forward transform lifts `typed_value` into a logical shredded child. + let mut ctx = SESSION.create_execution_ctx(); + let Canonical::Variant(canonical) = original.clone().execute::(&mut ctx)? else { + return Err(vortex_err!("expected canonical variant")); + }; + let core = canonical + .core_storage() + .as_opt::() + .ok_or_else(|| vortex_err!("expected parquet variant core storage"))?; + let logical = canonical + .shredded() + .ok_or_else(|| vortex_err!("expected canonical shredded child"))? + .clone(); + + // 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(), + Some(rebuilt_typed_value), + )? + .into_array(); + + assert_eq!(rebuilt.len(), original.len()); + for idx in 0..original.len() { + assert_eq!( + rebuilt.execute_scalar(idx, &mut ctx)?, + original.execute_scalar(idx, &mut ctx)?, + "row {idx}" + ); + } + Ok(()) + } } diff --git a/encodings/parquet-variant/src/arrow.rs b/encodings/parquet-variant/src/arrow.rs index 9bee5c738d3..da4e648a138 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -36,6 +36,7 @@ use vortex_session::registry::Id; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; +use crate::array::parquet_typed_value_from_logical_shredded; /// Arrow canonical extension name for Parquet Variant storage. const PARQUET_VARIANT_ARROW_EXTENSION_NAME: &str = "arrow.parquet.variant"; @@ -58,7 +59,7 @@ fn parquet_variant_storage_request(fields: &Fields) -> Option<(bool, bool)> { (has_metadata && (has_value || has_typed_value)).then_some((has_value, has_typed_value)) } -fn export_storage_to_target( +pub(crate) fn export_storage_to_target( parquet_array: &T, target_fields: &Fields, ctx: &mut ExecutionCtx, @@ -99,7 +100,7 @@ fn export_storage_to_target( )?)) } -fn export_unshredded_storage_to_target( +pub(crate) fn export_unshredded_storage_to_target( parquet_array: &T, target_fields: &Fields, ctx: &mut ExecutionCtx, @@ -115,7 +116,10 @@ fn export_unshredded_storage_to_target( export_storage_to_target(&unshredded_parquet, target_fields, ctx) } -fn parquet_variant_for_export(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { +pub(crate) fn parquet_variant_for_export( + array: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { let executed = array.execute_until::(ctx)?; if executed.is::() { return Ok(executed); @@ -135,11 +139,16 @@ fn parquet_variant_for_export(array: ArrayRef, ctx: &mut ExecutionCtx) -> Vortex return Ok(core_storage); }; + // The canonical shredded child has had its Parquet `value`/`typed_value` wrapper shells + // stripped; rebuild them so the reattached `typed_value` is valid Parquet storage that + // `to_arrow` and `unshred_variant` can consume. + 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(), - Some(shredded.clone()), + Some(typed_value), ) .map(IntoArray::into_array) } @@ -250,7 +259,7 @@ impl ArrowImportVTable for ParquetVariant { field: &Field, dtype: &DType, ) -> VortexResult { - if !matches!(dtype, DType::Variant(_)) + if !dtype.is_variant() || field .metadata() .get(EXTENSION_TYPE_NAME_KEY) @@ -295,7 +304,6 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -306,7 +314,7 @@ mod tests { #[fixture] fn session() -> VortexSession { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); crate::initialize(&session); session } @@ -500,7 +508,7 @@ mod tests { .arrow() .from_arrow_array(Arc::clone(&exported), &field)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } @@ -701,7 +709,7 @@ mod tests { let actual = session.arrow().from_arrow_array(exported, &field)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } } diff --git a/encodings/parquet-variant/src/compute/allnondistinct.rs b/encodings/parquet-variant/src/compute/allnondistinct.rs new file mode 100644 index 00000000000..4f5568fbcf6 --- /dev/null +++ b/encodings/parquet-variant/src/compute/allnondistinct.rs @@ -0,0 +1,161 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Debug; + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct; +use vortex_array::aggregate_fn::fns::all_non_distinct::all_non_distinct; +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; +use vortex_array::arrays::Struct; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; + +use crate::ParquetVariant; +use crate::ParquetVariantArrayExt; + +/// Lets `AllNonDistinct` compare two `ParquetVariant` arrays without canonicalizing them. +/// +/// `AllNonDistinct` accumulates over a `Struct{lhs, rhs}` batch, so this kernel is registered for +/// the struct encoding and inspects the two children. When both are `ParquetVariant`, we compare +/// the typed (`typed_value`) arrays if both sides are shredded, and fall back to the raw `value` +/// arrays otherwise. Comparing these child arrays directly avoids re-canonicalizing the variant +/// (which would recurse through the `Variant` canonical form). +#[derive(Debug)] +pub struct AllNonDistinctParquetVariant; + +impl DynAggregateKernel for AllNonDistinctParquetVariant { + fn aggregate( + &self, + aggregate_fn: &AggregateFnRef, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if !aggregate_fn.is::() { + return Ok(None); + } + + let Some(batch) = batch.as_opt::() else { + return Ok(None); + }; + let lhs = batch.unmasked_field(0); + let rhs = batch.unmasked_field(1); + let (Some(lhs), Some(rhs)) = ( + lhs.as_opt::(), + rhs.as_opt::(), + ) else { + return Ok(None); + }; + + let typed_identical = match (lhs.typed_value_array(), rhs.typed_value_array()) { + (Some(lhs_typed), Some(rhs_typed)) => { + if lhs_typed.dtype().eq_ignore_nullability(rhs_typed.dtype()) { + all_non_distinct(lhs_typed, rhs_typed, ctx)? + } else { + return Ok(None); + } + } + _ => true, + }; + + if typed_identical { + let values_identical = match (lhs.value_array(), rhs.value_array()) { + (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. + _ => return Ok(None), + }; + Ok(Some(Scalar::bool( + values_identical, + Nullability::NonNullable, + ))) + } else { + Ok(Some(Scalar::bool(false, Nullability::NonNullable))) + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::fns::all_non_distinct::all_non_distinct; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::ParquetVariant; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + /// Non-nullable, minimally-valid metadata column of `len` rows. + fn metadata(len: usize) -> ArrayRef { + VarBinViewArray::from_iter_bin(vec![b"\x01\x00"; len]).into_array() + } + + /// Non-nullable binary `value` column. + fn binary>(values: impl IntoIterator) -> ArrayRef { + VarBinViewArray::from_iter_bin(values).into_array() + } + + fn parquet_variant( + len: usize, + value: Option, + typed_value: Option, + ) -> VortexResult { + Ok( + ParquetVariant::try_new(Validity::NonNullable, metadata(len), value, typed_value)? + .into_array(), + ) + } + + #[test] + fn all_non_distinct_matches_equal_unshredded() -> VortexResult<()> { + let lhs = parquet_variant(2, Some(binary([b"\x10", b"\x11"])), None)?; + let rhs = parquet_variant(2, Some(binary([b"\x10", b"\x11"])), None)?; + let mut ctx = SESSION.create_execution_ctx(); + assert!(all_non_distinct(&lhs, &rhs, &mut ctx)?); + Ok(()) + } + + #[test] + fn all_non_distinct_detects_distinct_unshredded() -> VortexResult<()> { + let lhs = parquet_variant(2, Some(binary([b"\x10", b"\x11"])), None)?; + let rhs = parquet_variant(2, Some(binary([b"\x10", b"\x12"])), None)?; + let mut ctx = SESSION.create_execution_ctx(); + assert!(!all_non_distinct(&lhs, &rhs, &mut ctx)?); + Ok(()) + } + + #[test] + fn all_non_distinct_matches_equal_value_and_typed() -> VortexResult<()> { + let typed = || buffer![1i32, 2].into_array(); + let lhs = parquet_variant(2, Some(binary([b"\x10", b"\x11"])), Some(typed()))?; + let rhs = parquet_variant(2, Some(binary([b"\x10", b"\x11"])), Some(typed()))?; + let mut ctx = SESSION.create_execution_ctx(); + assert!(all_non_distinct(&lhs, &rhs, &mut ctx)?); + Ok(()) + } + + #[test] + fn all_non_distinct_empty_is_true() -> VortexResult<()> { + let lhs = parquet_variant(0, Some(binary(Vec::<&[u8]>::new())), None)?; + let rhs = parquet_variant(0, Some(binary(Vec::<&[u8]>::new())), None)?; + let mut ctx = SESSION.create_execution_ctx(); + assert!(all_non_distinct(&lhs, &rhs, &mut ctx)?); + Ok(()) + } +} diff --git a/encodings/parquet-variant/src/compute/mod.rs b/encodings/parquet-variant/src/compute/mod.rs new file mode 100644 index 00000000000..f69f592b89d --- /dev/null +++ b/encodings/parquet-variant/src/compute/mod.rs @@ -0,0 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod allnondistinct; + +pub use allnondistinct::*; diff --git a/encodings/parquet-variant/src/json_to_variant_tests.rs b/encodings/parquet-variant/src/json_to_variant_tests.rs new file mode 100644 index 00000000000..cd3b2f035ee --- /dev/null +++ b/encodings/parquet-variant/src/json_to_variant_tests.rs @@ -0,0 +1,301 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution tests for the `vortex.json_to_variant` scalar function. +//! +//! The function definition lives in `vortex-json`; the JSON->Variant construction is performed +//! by the execute-parent kernel registered here, so these end-to-end tests live in +//! `vortex-parquet-variant` where that kernel is registered. + +use std::sync::LazyLock; + +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::ExtensionArray; +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::assert_arrays_eq; +use vortex_array::assert_nth_scalar_is_null; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::root; +use vortex_array::expr::variant_get; +use vortex_array::scalar_fn::fns::variant_get::VariantPath; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_json::Json; +use vortex_session::VortexSession; + +use crate::ParquetVariant; +use crate::ParquetVariantArrayExt; +use crate::ShreddingSpec; +use crate::json_to_variant; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); + +fn i64_dtype() -> DType { + DType::Primitive(PType::I64, Nullability::Nullable) +} + +fn shred_field_as_i64(field: &str) -> VortexResult { + ShreddingSpec::try_new([(VariantPath::field(field), i64_dtype())]) +} + +fn json_input(storage: ArrayRef) -> VortexResult { + Ok(ExtensionArray::try_new_from_vtable(Json, EmptyMetadata, storage)?.into_array()) +} + +fn execute_json_to_variant(input: ArrayRef, shredding: ShreddingSpec) -> VortexResult { + let expr = json_to_variant(root(), shredding); + input + .apply(&expr)? + .execute::(&mut SESSION.create_execution_ctx()) +} + +fn assert_variant_i64_rows(array: &ArrayRef, expected: &[Option]) -> VortexResult<()> { + assert_eq!(array.len(), expected.len()); + let mut ctx = SESSION.create_execution_ctx(); + for (idx, expected) in expected.iter().enumerate() { + let scalar = array.execute_scalar(idx, &mut ctx)?; + let variant = scalar.as_variant(); + match expected { + Some(expected) => { + let value = variant + .value() + .ok_or_else(|| vortex_err!("expected non-null variant at row {idx}"))? + .cast(&i64_dtype())?; + assert_eq!(value.as_primitive().typed_value::(), Some(*expected)); + } + None => assert!(scalar.is_null(), "expected null row {idx}"), + } + } + Ok(()) +} + +#[test] +fn rejects_bare_utf8_input() { + let input = VarBinViewArray::from_iter_str(["1", "2"]).into_array(); + + let err = execute_json_to_variant(input, ShreddingSpec::empty()).unwrap_err(); + assert!( + err.to_string().contains("Json extension"), + "unexpected error: {err}" + ); +} + +#[test] +fn converts_json_extension_rows() -> VortexResult<()> { + let input = json_input( + VarBinViewArray::from_iter_str([r#"{"a": 1}"#, "2", r#"{"a": 3}"#]).into_array(), + )?; + + let result = execute_json_to_variant(input, ShreddingSpec::empty())?; + + assert_eq!(result.dtype(), &DType::Variant(Nullability::NonNullable)); + let mut ctx = SESSION.create_execution_ctx(); + let row0 = result.execute_scalar(0, &mut ctx)?; + let object = row0 + .as_variant() + .value() + .ok_or_else(|| vortex_err!("expected non-null variant"))?; + let field = object + .as_struct() + .field("a") + .ok_or_else(|| vortex_err!("expected field a"))? + .as_variant() + .value() + .ok_or_else(|| vortex_err!("expected non-null field a"))? + .cast(&i64_dtype())?; + assert_eq!(field.as_primitive().typed_value::(), Some(1)); + + let row1 = result.execute_scalar(1, &mut ctx)?; + let value = row1 + .as_variant() + .value() + .ok_or_else(|| vortex_err!("expected non-null variant"))? + .cast(&i64_dtype())?; + assert_eq!(value.as_primitive().typed_value::(), Some(2)); + Ok(()) +} + +#[test] +fn converts_json_extension_input() -> VortexResult<()> { + let storage = VarBinViewArray::from_iter_str(["1", "2"]).into_array(); + let input = json_input(storage)?; + + let result = execute_json_to_variant(input, ShreddingSpec::empty())?; + + assert_eq!(result.dtype(), &DType::Variant(Nullability::NonNullable)); + assert_variant_i64_rows(&result, &[Some(1), Some(2)]) +} + +#[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 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(); + let input = DictArray::try_new(codes, values)?.into_array(); + + let result = execute_json_to_variant(input, ShreddingSpec::empty())?; + + assert_eq!(result.dtype(), &DType::Variant(Nullability::NonNullable)); + assert_variant_i64_rows(&result, &[Some(1), Some(2), Some(1), Some(2), Some(1)]) +} + +#[test] +fn null_rows_stay_null_and_json_null_becomes_variant_null() -> VortexResult<()> { + let input = json_input( + VarBinViewArray::from_iter_nullable_str([Some("1"), None, Some("null")]).into_array(), + )?; + + let result = execute_json_to_variant(input, ShreddingSpec::empty())?; + + assert_eq!(result.dtype(), &DType::Variant(Nullability::Nullable)); + let mut ctx = SESSION.create_execution_ctx(); + assert!(!result.execute_scalar(0, &mut ctx)?.is_null()); + assert_nth_scalar_is_null!(result, 1, &mut ctx); + let row2 = result.execute_scalar(2, &mut ctx)?; + assert!(!row2.is_null(), "JSON null must not be a row null"); + assert_eq!(row2.as_variant().is_variant_null(), Some(true)); + Ok(()) +} + +#[test] +fn invalid_json_errors() -> VortexResult<()> { + let input = + json_input(VarBinViewArray::from_iter_str([r#"{"a": 1}"#, r#"{"a":"#]).into_array())?; + + let err = execute_json_to_variant(input, ShreddingSpec::empty()).unwrap_err(); + assert!(!err.to_string().is_empty()); + Ok(()) +} + +#[test] +fn shredding_produces_typed_value_child() -> VortexResult<()> { + let input = json_input( + VarBinViewArray::from_iter_str([ + r#"{"a": 1, "b": "x"}"#, + r#"{"a": 2, "b": "y"}"#, + r#"{"a": "not-a-number", "b": "z"}"#, + r#"{"b": "missing-a"}"#, + ]) + .into_array(), + )?; + + let result = execute_json_to_variant(input, shred_field_as_i64("a")?)?; + + assert!( + result.as_::().typed_value_array().is_some(), + "expected shredded typed_value child" + ); + + // The canonical form must expose field `a` through the shredded tree. + let mut ctx = SESSION.create_execution_ctx(); + let Canonical::Variant(canonical) = result.clone().execute::(&mut ctx)? else { + vortex_bail!("expected canonical variant array"); + }; + let shredded = canonical + .shredded() + .ok_or_else(|| vortex_err!("expected canonical shredded child"))? + .clone() + .execute::(&mut ctx)?; + assert!(shredded.unmasked_field_by_name_opt("a").is_some()); + + // Typed extraction must serve shredded rows and fall back for mismatched rows. + let typed = result + .clone() + .apply(&variant_get( + root(), + VariantPath::field("a"), + Some(i64_dtype()), + ))? + .execute::(&mut ctx)?; + assert_arrays_eq!( + typed, + PrimitiveArray::from_option_iter([Some(1i64), Some(2), None, None]), + &mut ctx + ); + + // Mismatched rows keep their original value through the variant fallback. + let untyped = result + .apply(&variant_get( + root(), + VariantPath::field("a"), + Some(DType::Utf8(Nullability::Nullable)), + ))? + .execute::(&mut ctx)?; + let row2 = untyped.execute_scalar(2, &mut ctx)?; + assert_eq!( + row2.as_utf8().value().map(|value| value.to_string()), + Some("not-a-number".to_string()) + ); + Ok(()) +} + +#[test] +fn shredding_preserves_null_rows() -> VortexResult<()> { + let input = json_input( + VarBinViewArray::from_iter_nullable_str([Some(r#"{"a": 1}"#), None, Some(r#"{"a": 3}"#)]) + .into_array(), + )?; + + let result = execute_json_to_variant(input, shred_field_as_i64("a")?)?; + + assert_eq!(result.dtype(), &DType::Variant(Nullability::Nullable)); + let mut ctx = SESSION.create_execution_ctx(); + assert_nth_scalar_is_null!(result, 1, &mut ctx); + let typed = result + .apply(&variant_get( + root(), + VariantPath::field("a"), + Some(i64_dtype()), + ))? + .execute::(&mut ctx)?; + assert_arrays_eq!( + typed, + PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn shredding_root_path_shreds_top_level_values() -> VortexResult<()> { + let input = + json_input(VarBinViewArray::from_iter_str(["1", "2", r#""not-a-number""#]).into_array())?; + let spec = ShreddingSpec::try_new([(VariantPath::root(), i64_dtype())])?; + + let result = execute_json_to_variant(input, spec)?; + + assert!( + result.as_::().typed_value_array().is_some(), + "expected shredded typed_value child" + ); + assert_variant_i64_rows(&result.slice(0..2)?, &[Some(1), Some(2)])?; + let mut ctx = SESSION.create_execution_ctx(); + let row2 = result.execute_scalar(2, &mut ctx)?; + let value = row2 + .as_variant() + .value() + .ok_or_else(|| vortex_err!("expected non-null variant"))?; + assert_eq!( + value.as_utf8().value().map(|value| value.to_string()), + Some("not-a-number".to_string()) + ); + Ok(()) +} diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index a136b483678..bd7d0bfaa7c 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -11,41 +11,83 @@ use arrow_schema::FieldRef; use parquet_variant::VariantPath as PqVariantPath; use parquet_variant::VariantPathElement as PqVariantPathElement; use parquet_variant_compute::GetOptions; +use parquet_variant_compute::ShreddedSchemaBuilder; use parquet_variant_compute::VariantArray as ArrowVariantArray; +use parquet_variant_compute::shred_variant; use parquet_variant_compute::variant_get as arrow_variant_get; use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Extension; +use vortex_array::arrays::Filter; +use vortex_array::arrays::Slice; +use vortex_array::arrays::Struct; use vortex_array::arrays::dict::TakeExecute; use vortex_array::arrays::dict::TakeExecuteAdaptor; +use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::filter::FilterKernel; use vortex_array::arrays::scalar_fn::ExactScalarFn; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; use vortex_array::arrays::slice::SliceExecuteAdaptor; use vortex_array::arrays::slice::SliceKernel; +use vortex_array::arrow::ArrowSessionExt; use vortex_array::arrow::FromArrowArray; use vortex_array::dtype::DType; use vortex_array::kernel::ExecuteParentKernel; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +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_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; +use vortex_json::Json; +use vortex_json::JsonToVariant; use vortex_mask::Mask; +use vortex_session::VortexSession; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; - -pub(crate) static PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&FilterExecuteAdaptor(ParquetVariant)), - ParentKernelSet::lift(&SliceExecuteAdaptor(ParquetVariant)), - ParentKernelSet::lift(&TakeExecuteAdaptor(ParquetVariant)), - ParentKernelSet::lift(&VariantGetKernel), -]); +use crate::compute::AllNonDistinctParquetVariant; + +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Filter.id(), + ParquetVariant, + FilterExecuteAdaptor(ParquetVariant), + ); + kernels.register_execute_parent_kernel( + Slice.id(), + ParquetVariant, + SliceExecuteAdaptor(ParquetVariant), + ); + kernels.register_execute_parent_kernel( + Dict.id(), + ParquetVariant, + TakeExecuteAdaptor(ParquetVariant), + ); + kernels.register_execute_parent_kernel(VariantGet.id(), ParquetVariant, VariantGetKernel); + kernels.register_execute_parent_kernel( + JsonToVariant.id(), + Extension, + JsonExtensionToVariantKernel, + ); + let aggregates = session.aggregate_fns(); + aggregates.register_aggregate_kernel( + Struct.id(), + Some(AllNonDistinct.id()), + &AllNonDistinctParquetVariant, + ); +} #[derive(Default, Debug)] struct VariantGetKernel; @@ -87,7 +129,67 @@ impl ExecuteParentKernel for VariantGetKernel { } } -fn to_parquet_variant_path(path: &VariantPath) -> VortexResult> { +/// Performs the [`JsonToVariant`] conversion (and optional shredding) over JSON string storage. +/// +/// `JsonToVariant`'s definition lives in `vortex-json`; the registered `execute_parent` kernels +/// delegate here to do the actual JSON parsing and optional shredding using +/// `parquet_variant_compute`, producing a [`ParquetVariant`] array. `strings` is the JSON +/// extension's storage array; it is executed to Arrow and parsed. Nullability of the result follows +/// `parent.dtype()`, which equals the input's nullability. +fn json_strings_to_variant( + strings: ArrayRef, + parent: ScalarFnArrayView<'_, JsonToVariant>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let nullable = parent.dtype().is_nullable(); + let session = ctx.session().clone(); + let arrow_strings = session.arrow().execute_arrow(strings, None, ctx)?; + // Any row that fails to parse as JSON fails the whole conversion. + let arrow_variant = parquet_variant_compute::json_to_variant(&arrow_strings)?; + + let arrow_variant = if parent.options.shredding().is_empty() { + arrow_variant + } else { + let mut builder = ShreddedSchemaBuilder::new(); + for (path, dtype) in parent.options.shredding().fields() { + let field: FieldRef = Arc::new(session.arrow().to_arrow_field("shredded", dtype)?); + builder = builder.with_path(to_parquet_variant_path(path)?, field)?; + } + shred_variant(&arrow_variant, &builder.build())? + }; + + if nullable { + ParquetVariant::from_arrow_variant_nullable(&arrow_variant) + } else { + ParquetVariant::from_arrow_variant(&arrow_variant) + } +} + +/// Builds Parquet Variant arrays for [`JsonToVariant`] over a [`Json`] extension input. +/// +/// This kernel unwraps the extension's string storage and runs the JSON conversion. It is keyed on +/// the shared extension encoding, so it declines any non-`Json` extension. +#[derive(Default, Debug)] +struct JsonExtensionToVariantKernel; + +impl ExecuteParentKernel for JsonExtensionToVariantKernel { + type Parent = ExactScalarFn; + + fn execute_parent( + &self, + array: ArrayView<'_, Extension>, + parent: ScalarFnArrayView<'_, JsonToVariant>, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if child_idx != 0 || !array.ext_dtype().is::() { + return Ok(None); + } + json_strings_to_variant(array.storage_array().clone(), parent, ctx).map(Some) + } +} + +pub(crate) fn to_parquet_variant_path(path: &VariantPath) -> VortexResult> { path.elements() .iter() .map(|element| match element { @@ -188,6 +290,7 @@ impl TakeExecute for ParquetVariant { #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use arrow_array::Array as ArrowArray; use arrow_array::ArrayRef as ArrowArrayRef; @@ -207,7 +310,6 @@ mod tests { use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray as VortexStructArray; @@ -232,10 +334,17 @@ mod tests { use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn make_unshredded_array() -> VortexResult { let mut builder = VariantArrayBuilder::new(4); builder.append_variant(PqVariant::from(42i32)); @@ -358,7 +467,7 @@ mod tests { let expr = variant_get(root(), parse_path(path)?, dtype); array .apply(&expr)? - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut SESSION.create_execution_ctx()) } macro_rules! assert_rows_eq { @@ -368,7 +477,7 @@ mod tests { let expected_rows = [$($expected_idx),*]; assert_eq!(actual.len(), expected_rows.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); for (actual_idx, expected_idx) in expected_rows.into_iter().enumerate() { assert_eq!( actual.execute_scalar(actual_idx, &mut ctx)?, @@ -385,7 +494,7 @@ mod tests { let expected = [$($is_null),*]; assert_eq!(array.len(), expected.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); for (idx, is_null) in expected.into_iter().enumerate() { assert_eq!( array.execute_scalar(idx, &mut ctx)?.is_null(), @@ -493,7 +602,8 @@ mod tests { result.dtype(), &VortexDType::Primitive(PType::I32, Nullability::Nullable) ); - assert_arrays_eq!(result, PrimitiveArray::from_option_iter(expected)); + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(result, PrimitiveArray::from_option_iter(expected), &mut ctx); Ok(()) } @@ -509,7 +619,7 @@ mod tests { let result = execute_variant_get(arr, "$.a", None)?; assert_eq!(result.dtype(), &VortexDType::Variant(Nullability::Nullable)); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let row0 = result.execute_scalar(0, &mut ctx)?; assert_eq!( row0.as_variant() @@ -518,7 +628,7 @@ mod tests { .map(|value| value.as_str()), Some("ok") ); - assert_nth_scalar_is_null!(result, 1); + assert_nth_scalar_is_null!(result, 1, &mut ctx); assert_eq!( result .execute_scalar(2, &mut ctx)? @@ -526,7 +636,7 @@ mod tests { .is_variant_null(), Some(true) ); - assert_nth_scalar_is_null!(result, 3); + assert_nth_scalar_is_null!(result, 3, &mut ctx); Ok(()) } @@ -684,7 +794,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 mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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")); }; @@ -722,7 +832,7 @@ mod tests { fn assert_variant_i32_scalars(array: &ArrayRef, expected: &[Option]) -> VortexResult<()> { assert_eq!(array.len(), expected.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); for (idx, expected) in expected.iter().enumerate() { let scalar = array.execute_scalar(idx, &mut ctx)?; let variant = scalar.as_variant(); @@ -748,7 +858,7 @@ mod tests { ) -> VortexResult<()> { assert_eq!(array.len(), expected_a.len()); assert_eq!(array.len(), expected_b.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); for idx in 0..array.len() { let scalar = array.execute_scalar(idx, &mut ctx)?; let object = scalar @@ -807,7 +917,7 @@ mod tests { array: &ArrayRef, expected: impl IntoIterator>, ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let executed = array.clone().execute::(&mut ctx)?; let typed_value = executed .as_::() @@ -816,7 +926,11 @@ mod tests { .clone() .execute::(&mut ctx)?; - assert_arrays_eq!(typed_value, PrimitiveArray::from_option_iter(expected)); + assert_arrays_eq!( + typed_value, + PrimitiveArray::from_option_iter(expected), + &mut ctx + ); Ok(()) } @@ -856,7 +970,8 @@ mod tests { assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10), Some(30), None]) + PrimitiveArray::from_option_iter([Some(10), Some(30), None]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -873,7 +988,8 @@ mod tests { assert_arrays_eq!( result, - VarBinArray::from(vec![Some("left"), Some("right"), Some("missing_a")]) + VarBinArray::from(vec![Some("left"), Some("right"), Some("missing_a")]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -894,7 +1010,7 @@ mod tests { .shredded() .ok_or_else(|| vortex_err!("expected canonical shredded child"))? .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut SESSION.create_execution_ctx())?; assert_eq!( shredded.unmasked_field_by_name("a")?.dtype(), &VortexDType::Variant(Nullability::Nullable) @@ -909,7 +1025,8 @@ mod tests { assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10), Some(30), None]) + PrimitiveArray::from_option_iter([Some(10), Some(30), None]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -926,7 +1043,8 @@ mod tests { assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10), Some(30), None]) + PrimitiveArray::from_option_iter([Some(10), Some(30), None]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -943,7 +1061,8 @@ mod tests { assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(100), Some(20), Some(30)]) + PrimitiveArray::from_option_iter([Some(100), Some(20), Some(30)]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -987,7 +1106,7 @@ mod tests { .is_some() ); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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")); }; diff --git a/encodings/parquet-variant/src/lib.rs b/encodings/parquet-variant/src/lib.rs index 03d2a046442..c71ea81e734 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/encodings/parquet-variant/src/lib.rs @@ -26,6 +26,9 @@ mod array; mod arrow; +mod compute; +#[cfg(test)] +mod json_to_variant_tests; mod kernel; mod operations; mod validity; @@ -36,13 +39,23 @@ use std::sync::Arc; pub use array::ParquetVariantArrayExt; use vortex_array::arrow::ArrowSessionExt; use vortex_array::session::ArraySessionExt; +pub use vortex_json::JsonToVariant; +pub use vortex_json::JsonToVariantOptions; +pub use vortex_json::ShreddingSpec; +pub use vortex_json::json_to_variant; use vortex_session::VortexSession; pub use vtable::ParquetVariant; pub use vtable::ParquetVariantArray; -/// Register Parquet Variant array and Arrow extension support with a session. +/// Register Parquet Variant array, Arrow extension, and scalar function support with a +/// session. +/// +/// This also initializes [`vortex_json`], registering the `Json` extension dtype and the +/// `json_to_variant` scalar function whose execution this crate provides. pub fn initialize(session: &VortexSession) { + vortex_json::initialize(session); session.arrays().register(ParquetVariant); + kernel::initialize(session); session.arrow().register_exporter(Arc::new(ParquetVariant)); session.arrow().register_importer(Arc::new(ParquetVariant)); } diff --git a/encodings/parquet-variant/src/operations.rs b/encodings/parquet-variant/src/operations.rs index 10f11e2bf2c..3cc3d2e5b4a 100644 --- a/encodings/parquet-variant/src/operations.rs +++ b/encodings/parquet-variant/src/operations.rs @@ -387,8 +387,8 @@ mod tests { use parquet_variant::VariantBuilder; use parquet_variant_compute::VariantArray as ArrowVariantArray; use parquet_variant_compute::VariantArrayBuilder; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; @@ -420,7 +420,7 @@ mod tests { Some(ScalarValue::Variant(Box::new(expected_inner))), )?; assert_eq!( - vortex_arr.execute_scalar(index, &mut LEGACY_SESSION.create_execution_ctx())?, + vortex_arr.execute_scalar(index, &mut array_session().create_execution_ctx())?, expected ); } @@ -452,17 +452,17 @@ mod tests { assert!( vortex_arr - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())? + .execute_scalar(1, &mut array_session().create_execution_ctx())? .is_null() ); assert!( !vortex_arr - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())? + .execute_scalar(0, &mut array_session().create_execution_ctx())? .is_null() ); assert!( !vortex_arr - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())? + .execute_scalar(2, &mut array_session().create_execution_ctx())? .is_null() ); @@ -487,7 +487,7 @@ mod tests { let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; let present_variant_null = - vortex_arr.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())?; + vortex_arr.execute_scalar(0, &mut array_session().create_execution_ctx())?; assert!(!present_variant_null.is_null()); assert_eq!( present_variant_null.as_variant().is_variant_null(), @@ -495,12 +495,12 @@ mod tests { ); let outer_null = - vortex_arr.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())?; + vortex_arr.execute_scalar(1, &mut array_session().create_execution_ctx())?; assert!(outer_null.is_null()); assert_eq!(outer_null.as_variant().is_variant_null(), None); let present_value = - vortex_arr.execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())?; + vortex_arr.execute_scalar(2, &mut array_session().create_execution_ctx())?; assert!(!present_value.is_null()); assert_eq!(present_value.as_variant().is_variant_null(), Some(false)); @@ -526,17 +526,17 @@ mod tests { assert_eq!(vortex_arr.dtype(), &DType::Variant(Nullability::Nullable)); assert!( vortex_arr - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())? + .execute_scalar(0, &mut array_session().create_execution_ctx())? .is_null() ); assert!( vortex_arr - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())? + .execute_scalar(1, &mut array_session().create_execution_ctx())? .is_null() ); let inner_pv = vortex_arr.as_opt::().unwrap(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let roundtripped = inner_pv.to_arrow(&mut ctx)?; assert_eq!(roundtripped.inner().null_count(), 2); @@ -558,12 +558,12 @@ mod tests { ); assert!( !vortex_arr - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())? + .execute_scalar(0, &mut array_session().create_execution_ctx())? .is_null() ); assert!( !vortex_arr - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())? + .execute_scalar(1, &mut array_session().create_execution_ctx())? .is_null() ); assert_scalar_at_matches_arrow_try_value(&arrow_variant, [0, 1])?; @@ -673,7 +673,7 @@ mod tests { let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; - let row0 = vortex_arr.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())?; + let row0 = vortex_arr.execute_scalar(0, &mut array_session().create_execution_ctx())?; let row0 = row0.as_variant().value().unwrap().as_list(); assert_eq!(row0.len(), 2); assert_eq!( @@ -697,7 +697,7 @@ mod tests { Some(20) ); - let row1 = vortex_arr.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())?; + let row1 = vortex_arr.execute_scalar(1, &mut array_session().create_execution_ctx())?; let row1 = row1.as_variant().value().unwrap().as_list(); assert_eq!(row1.len(), 1); assert_eq!( @@ -764,7 +764,7 @@ mod tests { let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; - let object = vortex_arr.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())?; + let object = vortex_arr.execute_scalar(0, &mut array_session().create_execution_ctx())?; let object = object.as_variant().value().unwrap().as_struct(); assert_eq!( diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index f2ad4f64743..4fa3ea10690 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -39,7 +39,6 @@ use crate::array::VALIDITY_SLOT; use crate::array::VALUE_SLOT; use crate::array::core_storage_without_typed_value; use crate::array::logical_shredded_from_parquet_typed_value; -use crate::kernel::PARENT_KERNELS; /// VTable for Arrow's canonical `arrow.parquet.variant` extension storage. /// @@ -172,6 +171,14 @@ impl VTable for ParquetVariant { None } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { SLOT_NAMES[idx].to_string() } @@ -280,20 +287,12 @@ impl VTable for ParquetVariant { VariantArray::try_new(core_storage, shredded)?.into_array(), )) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::Int32Array; @@ -310,7 +309,6 @@ mod tests { use vortex_array::Canonical; use vortex_array::EqMode; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; @@ -323,7 +321,6 @@ mod tests { use vortex_array::dtype::PType; use vortex_array::serde::SerializeOptions; use vortex_array::serde::SerializedArray; - use vortex_array::session::ArraySession; use vortex_array::session::ArraySessionExt; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; @@ -344,11 +341,17 @@ mod tests { use crate::ParquetVariant; use crate::array::ParquetVariantArrayExt; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn roundtrip(array: ArrayRef) -> VortexResult { let dtype = array.dtype().clone(); let len = array.len(); - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); session.arrays().register(ParquetVariant); let ctx = ArrayContext::empty(); @@ -387,8 +390,7 @@ mod tests { #[fixture] fn parquet_variant_file_session() -> VortexSession { - let session = VortexSession::empty() - .with::() + let session = vortex_array::array_session() .with::() .with::(); vortex_file::register_default_encodings(&session); @@ -415,7 +417,7 @@ mod tests { let parquet_variant = ParquetVariant::try_new(Validity::NonNullable, metadata, None, Some(typed_value))?; assert!(parquet_variant.typed_value_array().is_some()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let Canonical::Variant(variant) = parquet_variant .into_array() @@ -439,7 +441,8 @@ mod tests { let shredded = shredded.clone().execute::(&mut ctx)?; assert_arrays_eq!( shredded, - PrimitiveArray::from_option_iter([Some(10), None, Some(30)]) + PrimitiveArray::from_option_iter([Some(10), None, Some(30)]), + &mut ctx ); Ok(()) @@ -468,7 +471,7 @@ mod tests { .read_all() .await?; - assert_arrays_eq!(expected, actual); + assert_arrays_eq!(expected, actual, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -496,7 +499,7 @@ mod tests { .read_all() .await?; - assert_arrays_eq!(expected, actual); + assert_arrays_eq!(expected, actual, &mut SESSION.create_execution_ctx()); Ok(()) } diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index df2d79f63f1..d47a935d845 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -164,6 +164,28 @@ impl VTable for Pco { } } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + let mut data = array.data().clone(); + let chunk_metas_len = data.metadata.chunks.len(); + vortex_ensure!(buffers.len() >= chunk_metas_len); + data.chunk_metas = buffers[..chunk_metas_len] + .iter() + .map(|buffer| buffer.clone().try_to_host_sync()) + .collect::>>()?; + data.pages = buffers[chunk_metas_len..] + .iter() + .map(|buffer| buffer.clone().try_to_host_sync()) + .collect::>>()?; + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, @@ -282,6 +304,7 @@ pub(crate) fn vortex_err_from_pco(err: PcoError) -> VortexError { } #[derive(Clone, Debug)] +/// Pco array encoding marker. pub struct Pco; impl Pco { @@ -317,6 +340,7 @@ pub(super) const NUM_SLOTS: usize = 1; pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; #[derive(Clone, Debug)] +/// Encoding-specific data for a [`PcoArray`]. pub struct PcoData { pub(crate) chunk_metas: Vec, pub(crate) pages: Vec, @@ -338,6 +362,7 @@ impl Display for PcoData { } impl PcoData { + /// Validate dtype, validity, slice, and Pco component invariants. pub fn validate(&self, dtype: &DType, len: usize, validity: &Validity) -> VortexResult<()> { let _ = number_type_from_ptype(self.ptype); vortex_ensure!( @@ -391,6 +416,7 @@ impl PcoData { Ok(()) } + /// Construct unsliced Pco data from chunk metadata, pages, and serialized metadata. pub fn new( chunk_metas: Vec, pages: Vec, @@ -409,6 +435,7 @@ impl PcoData { } } + /// Compress a primitive array into Pco data. pub fn from_primitive( parray: ArrayView<'_, Primitive>, level: usize, @@ -497,6 +524,11 @@ impl PcoData { )) } + /// Downcast and compress an array into Pco data. + /// + /// # Errors + /// + /// Returns an error if the input is not a primitive array or compression fails. pub fn from_array( array: ArrayRef, level: usize, @@ -512,6 +544,7 @@ impl PcoData { Self::from_primitive(parray.as_view(), level, nums_per_page, ctx) } + /// Decompress this Pco data into a primitive array. pub fn decompress( &self, unsliced_validity: &Validity, @@ -532,6 +565,7 @@ impl PcoData { self.ptype, unsliced_validity.slice(self.slice_start..self.slice_stop)?, self.slice_stop - self.slice_start, + ctx, )) } @@ -669,8 +703,8 @@ impl OperationsVTable for Pco { #[cfg(test)] mod tests { use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::validity::Validity; @@ -680,7 +714,7 @@ mod tests { #[test] fn test_slice_nullable() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Create a nullable array with some nulls let values = PrimitiveArray::new( buffer![10u32, 20, 30, 40, 50, 60], @@ -696,7 +730,8 @@ mod tests { Some(40), Some(50), None - ]) + ]), + &mut ctx ); // Slice to get only the non-null values in the middle @@ -704,6 +739,6 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(20u32), Some(30), Some(40), Some(50)]) .into_array(); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut ctx); } } diff --git a/encodings/pco/src/compute/cast.rs b/encodings/pco/src/compute/cast.rs index b35526c1016..5f562fafca6 100644 --- a/encodings/pco/src/compute/cast.rs +++ b/encodings/pco/src/compute/cast.rs @@ -65,15 +65,13 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_session::VortexSession; use crate::Pco; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_cast_pco_f32_to_f64() { @@ -92,7 +90,8 @@ mod tests { assert_arrays_eq!( casted, - PrimitiveArray::from_iter([1.0f64, 2.0, 3.0, 4.0, 5.0]) + PrimitiveArray::from_iter([1.0f64, 2.0, 3.0, 4.0, 5.0]), + &mut ctx ); } @@ -109,7 +108,8 @@ mod tests { .unwrap(); assert_arrays_eq!( casted, - PrimitiveArray::new(buffer![10u32, 20, 30, 40], Validity::AllValid,) + PrimitiveArray::new(buffer![10u32, 20, 30, 40], Validity::AllValid,), + &mut ctx ); } @@ -130,7 +130,11 @@ mod tests { &DType::Primitive(PType::U32, Nullability::NonNullable) ); // Verify the values are correct - assert_arrays_eq!(casted, PrimitiveArray::from_iter([20u32, 30, 40, 50])); + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([20u32, 30, 40, 50]), + &mut ctx + ); } #[test] @@ -153,7 +157,11 @@ mod tests { casted.dtype(), &DType::Primitive(PType::U32, Nullability::NonNullable) ); - assert_arrays_eq!(casted, PrimitiveArray::from_iter([20u32, 30, 40, 50])); + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([20u32, 30, 40, 50]), + &mut ctx + ); } #[rstest] diff --git a/encodings/pco/src/compute/mod.rs b/encodings/pco/src/compute/mod.rs index 036a639ae31..d86c3d4b0a3 100644 --- a/encodings/pco/src/compute/mod.rs +++ b/encodings/pco/src/compute/mod.rs @@ -7,8 +7,8 @@ mod cast; mod tests { use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::consistency::test_array_consistency; @@ -21,7 +21,7 @@ mod tests { values.as_view(), 0, 128, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() } @@ -32,7 +32,7 @@ mod tests { values.as_view(), 0, 128, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() } @@ -43,7 +43,7 @@ mod tests { values.as_view(), 0, 128, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() } @@ -54,7 +54,7 @@ mod tests { values.as_view(), 0, 128, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() } @@ -65,7 +65,7 @@ mod tests { values.as_view(), 0, 128, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() } @@ -76,7 +76,7 @@ mod tests { values.as_view(), 0, 128, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() } @@ -87,7 +87,7 @@ mod tests { values.as_view(), 0, 128, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() } @@ -98,7 +98,7 @@ mod tests { values.as_view(), 3, 128, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() } @@ -113,6 +113,9 @@ mod tests { #[case::single(pco_single())] #[case::large(pco_large())] fn test_pco_consistency(#[case] array: PcoArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index fcf9a9397fb..601bd0826bd 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -1,6 +1,23 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Pco-backed numeric compression encoding for Vortex arrays. +//! +//! [`PcoArray`] stores valid primitive numeric values in Pco chunks and pages, while Vortex +//! validity tracks null rows separately. Page metadata lets slices decompress only the components +//! required for the requested row range. +//! +//! Pco supports integer and floating-point primitive dtypes handled by the upstream `pco` crate. +//! It is normally selected through the BtrBlocks compressor when the `pco` feature is enabled. +//! To deserialize arrays manually, register the encoding in the array session: +//! +//! ```rust +//! use vortex_array::session::ArraySessionExt; +//! +//! let session = vortex_array::array_session(); +//! session.arrays().register(vortex_pco::Pco); +//! ``` + mod array; mod compute; mod rules; @@ -9,9 +26,11 @@ mod slice; pub use array::*; #[derive(Clone, prost::Message)] +/// Metadata for one Pco page. pub struct PcoPageInfo { // Since pco limits to 2^24 values per chunk, u32 is sufficient for the // count of values. + /// Number of valid primitive values stored in this page. #[prost(uint32, tag = "1")] pub n_values: u32, } @@ -19,17 +38,22 @@ pub struct PcoPageInfo { // We're calling this Info instead of Metadata because ChunkMeta refers to a specific // component of a Pco file. #[derive(Clone, prost::Message)] +/// Metadata for one Pco chunk. pub struct PcoChunkInfo { + /// Pages contained in this chunk. #[prost(message, repeated, tag = "1")] pub pages: Vec, } #[derive(Clone, prost::Message)] +/// Serialized metadata for a [`PcoArray`]. pub struct PcoMetadata { // would be nice to reuse one header per vortex file, but it's really only 1 byte, so // no issue duplicating it here per PcoArray + /// Pco file header bytes. #[prost(bytes, tag = "1")] pub header: Vec, + /// Metadata for each compressed chunk. #[prost(message, repeated, tag = "2")] pub chunks: Vec, } diff --git a/encodings/pco/src/tests.rs b/encodings/pco/src/tests.rs index 20fec7dd9e0..573bd6a0fb4 100644 --- a/encodings/pco/src/tests.rs +++ b/encodings/pco/src/tests.rs @@ -6,7 +6,6 @@ use std::sync::LazyLock; use vortex_array::ArrayContext; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; @@ -18,7 +17,6 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::serde::SerializeOptions; use vortex_array::serde::SerializedArray; -use vortex_array::session::ArraySession; use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; use vortex_array::vtable::child_to_validity; @@ -33,7 +31,7 @@ use vortex_session::registry::ReadContext; use crate::PcoData; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); session.arrays().register(Pco); session }); @@ -41,7 +39,7 @@ static SESSION: LazyLock = LazyLock::new(|| { use crate::Pco; #[test] fn test_compress_decompress() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let data: Vec = (0..200).collect(); let array = PrimitiveArray::from_iter(data.clone()); let compressed = Pco::from_primitive(array.as_view(), 3, 0, &mut ctx).unwrap(); @@ -54,39 +52,47 @@ fn test_compress_decompress() { compressed.dtype().nullability(), ); let decompressed = compressed.decompress(&unsliced_validity, &mut ctx).unwrap(); - assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data)); + assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data), &mut ctx); // check slicing works let slice = compressed.slice(100..105).unwrap(); for i in 0_i32..5 { - assert_nth_scalar!(slice, i as usize, 100 + i); + assert_nth_scalar!(slice, i as usize, 100 + i, &mut ctx); } - assert_arrays_eq!(slice, PrimitiveArray::from_iter([100, 101, 102, 103, 104])); + assert_arrays_eq!( + slice, + PrimitiveArray::from_iter([100, 101, 102, 103, 104]), + &mut ctx + ); let slice = compressed.slice(200..200).unwrap(); - assert_arrays_eq!(slice, PrimitiveArray::from_iter(Vec::::new())); + assert_arrays_eq!( + slice, + PrimitiveArray::from_iter(Vec::::new()), + &mut ctx + ); } #[test] fn test_compress_decompress_small() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_option_iter([None, Some(1)]); let compressed = Pco::from_primitive(array.as_view(), 3, 0, &mut ctx).unwrap(); let expected = array.into_array(); - assert_arrays_eq!(compressed, expected); + assert_arrays_eq!(compressed, expected, &mut ctx); let unsliced_validity = child_to_validity( compressed.as_ref().slots()[0].as_ref(), compressed.dtype().nullability(), ); let decompressed = compressed.decompress(&unsliced_validity, &mut ctx).unwrap(); - assert_arrays_eq!(decompressed, expected); + assert_arrays_eq!(decompressed, expected, &mut ctx); } #[test] fn test_empty() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let data: Vec = vec![]; let array = PrimitiveArray::from_iter(data.clone()); let compressed = Pco::from_primitive(array.as_view(), 3, 100, &mut ctx).unwrap(); @@ -95,12 +101,12 @@ fn test_empty() { compressed.dtype().nullability(), ); let primitive = compressed.decompress(&unsliced_validity, &mut ctx).unwrap(); - assert_arrays_eq!(primitive, PrimitiveArray::from_iter(data)); + assert_arrays_eq!(primitive, PrimitiveArray::from_iter(data), &mut ctx); } #[test] fn test_validity_and_multiple_chunks_and_pages() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let data: Vec = (0..200).collect(); let mut validity: Vec = vec![true; 200]; validity[7..15].fill(false); @@ -129,18 +135,18 @@ fn test_validity_and_multiple_chunks_and_pages() { assert_eq!(compressed.metadata.chunks.len(), 6); // 191 values / 33 rounds up to 6 assert_eq!(compressed.metadata.chunks[0].pages.len(), 4); // 33 / 10 rounds up to 4 - assert_nth_scalar!(compressed, 0, 0); - assert_nth_scalar!(compressed, 3, 3); - assert_nth_scalar!(compressed, 7, None::); - assert_nth_scalar!(compressed, 14, None::); - assert_nth_scalar!(compressed, 15, 15); - assert_nth_scalar!(compressed, 101, None::); - assert_nth_scalar!(compressed, 199, 199); + assert_nth_scalar!(compressed, 0, 0, &mut ctx); + assert_nth_scalar!(compressed, 3, 3, &mut ctx); + assert_nth_scalar!(compressed, 7, None::, &mut ctx); + assert_nth_scalar!(compressed, 14, None::, &mut ctx); + assert_nth_scalar!(compressed, 15, 15, &mut ctx); + assert_nth_scalar!(compressed, 101, None::, &mut ctx); + assert_nth_scalar!(compressed, 199, 199, &mut ctx); // check slicing works let slice = compressed.slice(100..103).unwrap(); - assert_nth_scalar!(slice, 0, 100); - assert_nth_scalar!(slice, 2, 102); + assert_nth_scalar!(slice, 0, 100, &mut ctx); + assert_nth_scalar!(slice, 2, 102, &mut ctx); let primitive = slice.execute::(&mut ctx).unwrap(); assert!( @@ -158,7 +164,7 @@ fn test_validity_and_multiple_chunks_and_pages() { #[test] fn test_validity_vtable() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let data: Vec = (0..5).collect(); let mask_bools = vec![false, true, true, false, true]; let array = PrimitiveArray::new( @@ -214,14 +220,13 @@ fn test_serde() -> VortexResult<()> { &ReadContext::new(context.to_ids()), &SESSION, )?; - let data_type = LEGACY_SESSION.arrow().to_arrow_field("", data.dtype())?; - let pco_arrow = LEGACY_SESSION + let data_type = SESSION.arrow().to_arrow_field("", data.dtype())?; + let pco_arrow = SESSION .arrow() .execute_arrow(pco, Some(&data_type), &mut ctx)?; - let decoded_arrow = - LEGACY_SESSION - .arrow() - .execute_arrow(decoded, Some(&data_type), &mut ctx)?; + let decoded_arrow = SESSION + .arrow() + .execute_arrow(decoded, Some(&data_type), &mut ctx)?; assert!(pco_arrow == decoded_arrow); Ok(()) } diff --git a/encodings/runend/Cargo.toml b/encodings/runend/Cargo.toml index 01a5b8d7a3e..268aa6ac3ca 100644 --- a/encodings/runend/Cargo.toml +++ b/encodings/runend/Cargo.toml @@ -52,3 +52,7 @@ harness = false [[bench]] name = "run_end_decode" harness = false + +[[bench]] +name = "run_end_take" +harness = false diff --git a/encodings/runend/benches/run_end_compress.rs b/encodings/runend/benches/run_end_compress.rs index a49a5cfcb50..cd241ad42c6 100644 --- a/encodings/runend/benches/run_end_compress.rs +++ b/encodings/runend/benches/run_end_compress.rs @@ -13,7 +13,6 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::IntegerPType; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_runend::RunEnd; @@ -24,8 +23,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + session +}); const BENCH_ARGS: &[(usize, usize)] = &[ (1000, 4), diff --git a/encodings/runend/benches/run_end_decode.rs b/encodings/runend/benches/run_end_decode.rs index f509e55aed3..9256a4b3a53 100644 --- a/encodings/runend/benches/run_end_decode.rs +++ b/encodings/runend/benches/run_end_decode.rs @@ -10,7 +10,6 @@ use divan::Bencher; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::BufferMut; @@ -21,8 +20,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + session +}); /// Distribution types for bool benchmarks #[derive(Clone, Copy)] diff --git a/encodings/runend/benches/run_end_null_count.rs b/encodings/runend/benches/run_end_null_count.rs index 5392e04189b..f214889daae 100644 --- a/encodings/runend/benches/run_end_null_count.rs +++ b/encodings/runend/benches/run_end_null_count.rs @@ -12,7 +12,6 @@ use rand::rngs::StdRng; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; -use vortex_array::session::ArraySession; use vortex_buffer::Buffer; use vortex_runend::RunEnd; use vortex_runend::RunEndArray; @@ -50,8 +49,11 @@ const BENCH_ARGS: &[(usize, usize, f64)] = &[ (100_000, 1024, 0.5), ]; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + session +}); #[divan::bench(args = BENCH_ARGS)] fn null_count_run_end(bencher: Bencher, (n, run_step, valid_density): (usize, usize, f64)) { diff --git a/encodings/runend/benches/run_end_take.rs b/encodings/runend/benches/run_end_take.rs new file mode 100644 index 00000000000..f0dccccfc48 --- /dev/null +++ b/encodings/runend/benches/run_end_take.rs @@ -0,0 +1,170 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use std::fmt; +use std::sync::LazyLock; + +use divan::Bencher; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_runend::RunEnd; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + session +}); + +#[derive(Clone, Copy)] +enum IndexPattern { + SortedEven, + ReverseDense, + Random, +} + +#[derive(Clone, Copy)] +enum IndexValidity { + NonNullable, + EveryFourthNull, + AllNull, +} + +#[derive(Clone, Copy)] +struct TakeBenchArgs { + name: &'static str, + array_len: usize, + run_step: usize, + take_len: usize, + pattern: IndexPattern, + validity: IndexValidity, +} + +impl fmt::Display for TakeBenchArgs { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "{}_len{}_run{}_take{}", + self.name, self.array_len, self.run_step, self.take_len + ) + } +} + +const BENCH_ARGS: &[TakeBenchArgs] = &[ + // Sorted sparse takes should use the single-pass linear scan over run ends. + TakeBenchArgs { + name: "sorted_linear", + array_len: 16_384, + run_step: 8, + take_len: 512, + pattern: IndexPattern::SortedEven, + validity: IndexValidity::NonNullable, + }, + // Dense unsorted takes should build the logical-position-to-run table. + TakeBenchArgs { + name: "dense_table", + array_len: 32_768, + run_step: 4, + take_len: 2_048, + pattern: IndexPattern::ReverseDense, + validity: IndexValidity::NonNullable, + }, + // Sparse unsorted takes below the large-run threshold should stay on binary search. + TakeBenchArgs { + name: "binary_sparse", + array_len: 65_536, + run_step: 4, + take_len: 512, + pattern: IndexPattern::Random, + validity: IndexValidity::NonNullable, + }, + // Nullable indices exercise masked stats and table lookup. + TakeBenchArgs { + name: "nullable_dense_table", + array_len: 32_768, + run_step: 4, + take_len: 2_048, + pattern: IndexPattern::ReverseDense, + validity: IndexValidity::EveryFourthNull, + }, + // All-null indices should return before touching run ends or values. + TakeBenchArgs { + name: "all_null", + array_len: 65_536, + run_step: 4, + take_len: 2_048, + pattern: IndexPattern::Random, + validity: IndexValidity::AllNull, + }, +]; + +#[divan::bench(args = BENCH_ARGS)] +fn take(bencher: Bencher, args: TakeBenchArgs) { + let array = run_end_array(args.array_len, args.run_step); + let indices = take_indices(args); + + bencher + .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, execution_ctx)| { + array + .take(indices.clone()) + .unwrap() + .execute::(execution_ctx) + .unwrap() + }); +} + +fn run_end_array(len: usize, run_step: usize) -> ArrayRef { + let num_runs = len.div_ceil(run_step); + let ends = (0..num_runs) + .map(|run_idx| ((run_idx + 1) * run_step).min(len) as u64) + .collect::>() + .into_array(); + let values = PrimitiveArray::from_iter((0..num_runs).map(|idx| idx as u64)).into_array(); + + RunEnd::new(ends, values, &mut SESSION.create_execution_ctx()).into_array() +} + +fn take_indices(args: TakeBenchArgs) -> ArrayRef { + let values = index_values(args); + let validity = match args.validity { + IndexValidity::NonNullable => Validity::NonNullable, + IndexValidity::EveryFourthNull => { + Validity::from_iter((0..args.take_len).map(|i| !i.is_multiple_of(4))) + } + IndexValidity::AllNull => Validity::AllInvalid, + }; + + PrimitiveArray::new(values, validity).into_array() +} + +fn index_values(args: TakeBenchArgs) -> Buffer { + let values = match args.pattern { + IndexPattern::SortedEven => (0..args.take_len) + .map(|idx| ((idx * args.array_len) / args.take_len) as u64) + .collect::>(), + IndexPattern::ReverseDense => (0..args.take_len).rev().map(|idx| idx as u64).collect(), + IndexPattern::Random => { + let mut rng = StdRng::seed_from_u64(0); + (0..args.take_len) + .map(|_| rng.random_range(0..args.array_len) as u64) + .collect() + } + }; + + Buffer::from(values) +} diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 7bb069bb458..64d0798ebee 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -48,7 +48,6 @@ use crate::compress::runend_decode_primitive; use crate::compress::runend_decode_varbinview; use crate::compress::runend_encode; use crate::decompress_bool::runend_decode_bools; -use crate::kernel::PARENT_KERNELS; use crate::rules::RULES; /// A [`RunEnd`]-encoded Vortex array. @@ -124,6 +123,14 @@ impl VTable for RunEnd { vortex_panic!("RunEndArray buffer_name index {idx} out of bounds") } + 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, @@ -172,15 +179,6 @@ impl VTable for RunEnd { RULES.evaluate(array, parent, child_idx) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { run_end_canonicalize(&array, ctx).map(ExecutionResult::done) } @@ -411,12 +409,14 @@ impl RunEndData { /// ``` /// # use vortex_array::arrays::BoolArray; /// # use vortex_array::IntoArray; - /// # use vortex_array::{LEGACY_SESSION, VortexSessionExecute}; + /// # use vortex_array::VortexSessionExecute; /// # use vortex_buffer::buffer; /// # use vortex_error::VortexResult; /// # use vortex_runend::RunEnd; /// # fn main() -> VortexResult<()> { - /// let mut ctx = LEGACY_SESSION.create_execution_ctx(); + /// let session = vortex_array::array_session(); + /// vortex_runend::initialize(&session); + /// let mut ctx = session.create_execution_ctx(); /// let ends = buffer![2u8, 3u8].into_array(); /// let values = BoolArray::from_iter([false, true]).into_array(); /// let run_end = RunEnd::new(ends, values, &mut ctx); @@ -521,14 +521,16 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::session::ArraySession; use vortex_buffer::buffer; use vortex_session::VortexSession; use crate::RunEnd; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_runend_constructor() { @@ -548,7 +550,7 @@ mod tests { // 2, 3, 4 => 2 // 5, 6, 7, 8, 9 => 3 let expected = buffer![1, 1, 2, 2, 2, 3, 3, 3, 3, 3].into_array(); - assert_arrays_eq!(arr.into_array(), expected); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); } #[test] @@ -562,7 +564,7 @@ mod tests { let expected = VarBinViewArray::from_iter_str(["a", "a", "b", "b", "b", "c", "c", "c", "c", "c"]) .into_array(); - assert_arrays_eq!(arr.into_array(), expected); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); } #[test] @@ -583,6 +585,6 @@ mod tests { let expected = VarBinViewArray::from_iter_str(["x", "x", "y", "y", "y", "z", "z", "z", "z", "z"]) .into_array(); - assert_arrays_eq!(arr.into_array(), expected); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); } } diff --git a/encodings/runend/src/arrow.rs b/encodings/runend/src/arrow.rs index c26674770c0..193937122ef 100644 --- a/encodings/runend/src/arrow.rs +++ b/encodings/runend/src/arrow.rs @@ -81,7 +81,6 @@ mod tests { use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray as _; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute as _; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -95,8 +94,6 @@ mod tests { use vortex_array::scalar::PValue; use vortex_array::search_sorted::SearchSorted; use vortex_array::search_sorted::SearchSortedSide; - use vortex_array::session::ArraySession; - use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; @@ -107,8 +104,8 @@ mod tests { use crate::ops::find_slice_end_index; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); - session.arrays().register(RunEnd); + let session = vortex_array::array_session(); + crate::initialize(&session); session }); @@ -171,7 +168,8 @@ mod tests { assert_arrays_eq!( vortex_array.into_array(), - buffer![10i32, 10, 10, 20, 20, 30, 30, 30].into_array() + buffer![10i32, 10, 10, 20, 20, 30, 30, 30].into_array(), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -195,7 +193,8 @@ mod tests { None, Some(300i32), Some(300i32) - ]) + ]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -210,7 +209,11 @@ mod tests { // Convert to Vortex let vortex_array = decode_run_array(&arrow_run_array, false)?; - assert_arrays_eq!(vortex_array, buffer![1.5f64, 2.5, 2.5, 3.5].into_array()); + assert_arrays_eq!( + vortex_array, + buffer![1.5f64, 2.5, 2.5, 3.5].into_array(), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -231,7 +234,8 @@ mod tests { let vortex_array = decode_run_array(&sliced_array, false)?; assert_arrays_eq!( vortex_array, - buffer![100, 200, 200, 200, 300, 300].into_array() + buffer![100, 200, 200, 200, 300, 300].into_array(), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -262,7 +266,8 @@ mod tests { Some(30), Some(30), Some(40), - ]) + ]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -346,7 +351,7 @@ mod tests { #[case] expected_ends: &[i32], #[case] expected_values: &[i32], ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = RunEnd::encode( PrimitiveArray::from_iter(input.iter().copied()).into_array(), &mut ctx, diff --git a/encodings/runend/src/compress.rs b/encodings/runend/src/compress.rs index ba46588d88b..544b677dfa3 100644 --- a/encodings/runend/src/compress.rs +++ b/encodings/runend/src/compress.rs @@ -320,7 +320,8 @@ pub fn runend_decode_varbinview( #[cfg(test)] mod tests { - use vortex_array::LEGACY_SESSION; + use std::sync::LazyLock; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -328,27 +329,34 @@ mod tests { use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::compress::runend_decode_primitive; use crate::compress::runend_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn encode() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = PrimitiveArray::from_iter([1i32, 1, 2, 2, 2, 3, 3, 3, 3, 3]); let (ends, values) = runend_encode(arr.as_view(), &mut ctx); let values = values.execute::(&mut ctx)?; let expected_ends = PrimitiveArray::from_iter(vec![2u8, 5, 10]); - assert_arrays_eq!(ends, expected_ends); + assert_arrays_eq!(ends, expected_ends, &mut ctx); let expected_values = PrimitiveArray::from_iter(vec![1i32, 2, 3]); - assert_arrays_eq!(values, expected_values); + assert_arrays_eq!(values, expected_values, &mut ctx); Ok(()) } #[test] fn encode_nullable() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = PrimitiveArray::new( buffer![1i32, 1, 2, 2, 2, 3, 3, 3, 3, 3], Validity::from(BitBuffer::from(vec![ @@ -359,16 +367,16 @@ mod tests { let values = values.execute::(&mut ctx)?; let expected_ends = PrimitiveArray::from_iter(vec![2u8, 4, 5, 8, 10]); - assert_arrays_eq!(ends, expected_ends); + assert_arrays_eq!(ends, expected_ends, &mut ctx); let expected_values = PrimitiveArray::from_option_iter(vec![Some(1i32), None, Some(2), Some(3), None]); - assert_arrays_eq!(values, expected_values); + assert_arrays_eq!(values, expected_values, &mut ctx); Ok(()) } #[test] fn encode_all_null() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = PrimitiveArray::new( buffer![0, 0, 0, 0, 0], Validity::from(BitBuffer::new_unset(5)), @@ -377,21 +385,21 @@ mod tests { let values = values.execute::(&mut ctx)?; let expected_ends = PrimitiveArray::from_iter(vec![5u64]); - assert_arrays_eq!(ends, expected_ends); + assert_arrays_eq!(ends, expected_ends, &mut ctx); let expected_values = PrimitiveArray::from_option_iter(vec![Option::::None]); - assert_arrays_eq!(values, expected_values); + assert_arrays_eq!(values, expected_values, &mut ctx); Ok(()) } #[test] fn decode() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let ends = PrimitiveArray::from_iter([2u32, 5, 10]); let values = PrimitiveArray::from_iter([1i32, 2, 3]); let decoded = runend_decode_primitive(ends, values, 0, 10, &mut ctx)?; let expected = PrimitiveArray::from_iter(vec![1i32, 1, 2, 2, 2, 3, 3, 3, 3, 3]); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } } diff --git a/encodings/runend/src/compute/cast.rs b/encodings/runend/src/compute/cast.rs index f70bb0becb5..fe740739ca9 100644 --- a/encodings/runend/src/compute/cast.rs +++ b/encodings/runend/src/compute/cast.rs @@ -46,15 +46,17 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::session::ArraySession; use vortex_buffer::buffer; use vortex_session::VortexSession; use crate::RunEnd; use crate::RunEndArray; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_cast_runend_i32_to_i64() { @@ -132,7 +134,11 @@ mod tests { let sliced = runend.slice(3..8).unwrap(); // Verify the slice is correct before casting - assert_arrays_eq!(sliced, PrimitiveArray::from_iter([200, 200, 300, 300, 300])); + assert_arrays_eq!( + sliced, + PrimitiveArray::from_iter([200, 200, 300, 300, 300]), + &mut ctx + ); // Cast the sliced array let casted = sliced @@ -142,7 +148,8 @@ mod tests { // Verify the cast preserved the offset assert_arrays_eq!( casted, - PrimitiveArray::from_iter([200i64, 200, 300, 300, 300]) + PrimitiveArray::from_iter([200i64, 200, 300, 300, 300]), + &mut ctx ); } diff --git a/encodings/runend/src/compute/compare.rs b/encodings/runend/src/compute/compare.rs index 4da7a453000..ecdd03e3686 100644 --- a/encodings/runend/src/compute/compare.rs +++ b/encodings/runend/src/compute/compare.rs @@ -48,9 +48,10 @@ impl CompareKernel for RunEnd { #[cfg(test)] mod test { + use std::sync::LazyLock; + use vortex_array::ExecutionCtx; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; @@ -58,10 +59,17 @@ mod test { use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_session::VortexSession; use crate::RunEnd; use crate::RunEndArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn ree_array(ctx: &mut ExecutionCtx) -> RunEndArray { RunEnd::encode( PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).into_array(), @@ -72,7 +80,7 @@ mod test { #[test] fn compare_run_end() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = ree_array(&mut ctx); let res = arr .into_array() @@ -81,6 +89,6 @@ mod test { let expected = BoolArray::from_iter([ false, false, false, false, false, false, false, false, true, true, true, true, ]); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut ctx); } } diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index 60644e2bced..ca99fdf1912 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -116,7 +116,6 @@ fn filter_run_end_primitive + AsPrimitiv #[cfg(test)] mod tests { use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -125,11 +124,12 @@ mod tests { use crate::RunEnd; use crate::RunEndArray; + use crate::tests::SESSION; fn ree_array() -> RunEndArray { RunEnd::encode( PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap() } @@ -139,14 +139,15 @@ mod tests { let arr = ree_array().slice(2..7)?; let filtered = arr.filter(Mask::from_iter([true, false, false, true, true]))?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( filtered, RunEnd::new( PrimitiveArray::from_iter([1u8, 2, 3]).into_array(), PrimitiveArray::from_iter([1i32, 4, 2]).into_array(), &mut ctx, - ) + ), + &mut ctx ); Ok(()) } @@ -157,7 +158,7 @@ mod tests { /// Filter unwrap one layer at a time so RunEnd's FilterKernel can fire. #[test] fn filter_sliced_run_end_preserves_encoding() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // 4 runs of 32 each = 128 rows. Large enough that FilterKernel takes // the run-preserving path (true_count >= 25). @@ -186,7 +187,7 @@ mod tests { .chain(std::iter::repeat_n(30, 16)) .chain(std::iter::repeat_n(40, 16)) .collect(); - assert_arrays_eq!(executed, PrimitiveArray::from_iter(expected)); + assert_arrays_eq!(executed, PrimitiveArray::from_iter(expected), &mut ctx); Ok(()) } diff --git a/encodings/runend/src/compute/min_max.rs b/encodings/runend/src/compute/min_max.rs index 2a65d9de2bf..4dd39edc5f3 100644 --- a/encodings/runend/src/compute/min_max.rs +++ b/encodings/runend/src/compute/min_max.rs @@ -28,16 +28,16 @@ impl DynAggregateKernel for RunEndMinMaxKernel { batch: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if !aggregate_fn.is::() { + let Some(options) = aggregate_fn.as_opt::() else { return Ok(None); - } + }; let Some(run_end) = batch.as_opt::() else { return Ok(None); }; let struct_dtype = make_minmax_dtype(batch.dtype()); - match min_max(run_end.values(), ctx)? { + match min_max(run_end.values(), ctx, *options)? { Some(result) => Ok(Some(Scalar::struct_( struct_dtype, vec![result.min, result.max], diff --git a/encodings/runend/src/compute/mod.rs b/encodings/runend/src/compute/mod.rs index 9bdb6d67c00..fc7fc8804ec 100644 --- a/encodings/runend/src/compute/mod.rs +++ b/encodings/runend/src/compute/mod.rs @@ -13,47 +13,55 @@ pub(crate) mod take_from; #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_buffer::buffer; + use vortex_session::VortexSession; use crate::RunEnd; use crate::RunEndArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[rstest] // Simple run-end arrays #[case::runend_i32(RunEnd::encode( buffer![1i32, 1, 1, 2, 2, 3, 3, 3, 3].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::runend_single_run(RunEnd::encode( buffer![5i32, 5, 5, 5, 5].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::runend_alternating(RunEnd::encode( buffer![1i32, 2, 1, 2, 1, 2].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] // Different types #[case::runend_u64(RunEnd::encode( buffer![100u64, 100, 200, 200, 200].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] // Edge cases #[case::runend_single(RunEnd::encode( buffer![42i32].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::runend_large(RunEnd::encode( PrimitiveArray::from_iter((0..1000).map(|i| i / 10)).into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] fn test_runend_consistency(#[case] array: RunEndArray) { - test_array_consistency(&array.into_array()); + test_array_consistency(&array.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/runend/src/compute/take.rs b/encodings/runend/src/compute/take.rs index 7100faf9eac..c61854635e3 100644 --- a/encodings/runend/src/compute/take.rs +++ b/encodings/runend/src/compute/take.rs @@ -7,107 +7,516 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::TakeExecute; +use vortex_array::dtype::UnsignedPType; use vortex_array::match_each_integer_ptype; -use vortex_array::search_sorted::SearchResult; -use vortex_array::search_sorted::SearchSorted; -use vortex_array::search_sorted::SearchSortedSide; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_mask::AllOr; +use vortex_mask::Mask; use crate::RunEnd; use crate::array::RunEndArrayExt; +use crate::iter::trimmed_ends_iter; + +const SORTED_LINEAR_RUNS_PER_INDEX_THRESHOLD: usize = 16; +const UNSORTED_LINEAR_RUNS_PER_INDEX_THRESHOLD: usize = 4; +/// Sorting the indices and merging only beats per-index binary search once the run ends are too +/// large to stay cache-resident; below this run count binary search wins. +const UNSORTED_LINEAR_MIN_RUNS: usize = 1 << 19; +/// Use a dense logical-position-to-run-index table when the array length is at most this many +/// times the number of valid indices: building the table is O(array_len) and each index then +/// resolves with a single unconditional gather. +const TABLE_LEN_PER_INDEX_THRESHOLD: usize = 8; impl TakeExecute for RunEnd { - #[expect( - clippy::cast_possible_truncation, - reason = "index cast to usize inside macro" - )] fn take( array: ArrayView<'_, Self>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { let primitive_indices = indices.clone().execute::(ctx)?; + let indices_validity = primitive_indices.validity()?; + let indices_mask = indices_validity.execute_mask(primitive_indices.len(), ctx)?; - let checked_indices = match_each_integer_ptype!(primitive_indices.ptype(), |P| { - primitive_indices - .as_slice::

() - .iter() - .copied() - .map(|idx| { - let usize_idx = idx as usize; - if usize_idx >= array.len() { - vortex_bail!(OutOfBounds: usize_idx, 0, array.len()); - } - Ok(usize_idx) - }) - .collect::>>()? + let taken = match_each_integer_ptype!(primitive_indices.ptype(), |P| { + take_indices( + array, + primitive_indices.as_slice::

(), + &indices_validity, + &indices_mask, + true, + ctx, + )? }); - let indices_validity = primitive_indices.validity()?; - take_indices_unchecked(array, &checked_indices, &indices_validity, ctx).map(Some) + Ok(Some(taken)) } } -/// Perform a take operation on a RunEndArray by binary searching for each of the indices. +/// Perform a take operation on a RunEndArray without bounds-checking the indices. +/// +/// The caller must guarantee that all valid indices are in bounds for the array. pub fn take_indices_unchecked>( array: ArrayView<'_, RunEnd>, indices: &[T], validity: &Validity, ctx: &mut ExecutionCtx, ) -> VortexResult { + let validity_mask = validity.execute_mask(indices.len(), ctx)?; + take_indices(array, indices, validity, &validity_mask, false, ctx) +} + +fn take_indices>( + array: ArrayView<'_, RunEnd>, + indices: &[T], + validity: &Validity, + validity_mask: &Mask, + check_bounds: bool, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if validity_mask.all_false() { + return Ok( + ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len()) + .into_array(), + ); + } + + let stats = valid_indices_stats(indices, validity_mask, array.len(), check_bounds)?; let ends = array.ends().clone().execute::(ctx)?; - let ends_len = ends.len(); - - // TODO(joe): use the validity mask to skip search sorted. - let physical_indices = match_each_integer_ptype!(ends.ptype(), |I| { - let end_slices = ends.as_slice::(); - let physical_indices_vec: Vec = indices - .iter() - .map(|idx| idx.as_() + array.offset()) - .map(|idx| { - match ::from(idx) { - Some(idx) => end_slices.search_sorted(&idx, SearchSortedSide::Right), - None => { - // The idx is too large for I, therefore it's out of bounds. - Ok(SearchResult::NotFound(ends_len)) - } - } - }) - .map(|result| result.map(|r| r.to_ends_index(ends_len) as u64)) - .collect::>>()?; - let buffer = Buffer::from(physical_indices_vec); - PrimitiveArray::new(buffer, validity.clone()) + let physical_indices = match_each_unsigned_integer_ptype!(ends.ptype(), |I| { + let ends = ends.as_slice::(); + // Run indices fit in u32 for any realistic array; the narrower physical indices halve + // the memory traffic of the downstream take on the values. + if ends.len() <= u32::MAX as usize { + PrimitiveArray::new( + physical_indices_with_stats::<_, _, u32>( + ends, + array.offset(), + array.len(), + indices, + validity_mask, + stats, + ), + validity.clone(), + ) + } else { + PrimitiveArray::new( + physical_indices_with_stats::<_, _, u64>( + ends, + array.offset(), + array.len(), + indices, + validity_mask, + stats, + ), + validity.clone(), + ) + } }); array.values().take(physical_indices.into_array()) } +#[derive(Clone, Copy)] +struct ValidIndicesStats { + count: usize, + sorted: bool, +} + +fn physical_indices_with_stats( + ends: &[I], + offset: usize, + array_len: usize, + indices: &[T], + validity_mask: &Mask, + stats: ValidIndicesStats, +) -> Buffer +where + I: UnsignedPType, + T: AsPrimitive, + O: UnsignedPType, + usize: AsPrimitive, +{ + if stats.count == 0 { + return Buffer::zeroed(indices.len()); + } + + if stats.sorted + && prefer_linear_scan( + ends.len(), + stats.count, + SORTED_LINEAR_RUNS_PER_INDEX_THRESHOLD, + ) + { + return physical_indices_linear_sorted(ends, offset, indices, validity_mask); + } + + // A dense take resolves fastest through the position table regardless of index ordering. + // Sorted indices reach here only when there are too many runs for the sorted linear scan, + // for example a narrow slice of a heavily run-encoded array where runs far exceed array_len. + if array_len <= stats.count.saturating_mul(TABLE_LEN_PER_INDEX_THRESHOLD) { + return physical_indices_table(ends, offset, array_len, indices, validity_mask); + } + + if ends.len() >= UNSORTED_LINEAR_MIN_RUNS + && prefer_linear_scan( + ends.len(), + stats.count, + UNSORTED_LINEAR_RUNS_PER_INDEX_THRESHOLD, + ) + { + return physical_indices_linear_unsorted(ends, offset, indices, validity_mask, stats.count); + } + + physical_indices_binary(ends, offset, indices, validity_mask) +} + +/// Count the valid indices and determine whether they are sorted, bounds-checking each valid +/// index against `array_len` when `check_bounds` is set. +fn valid_indices_stats>( + indices: &[T], + validity_mask: &Mask, + array_len: usize, + check_bounds: bool, +) -> VortexResult { + debug_assert_eq!(indices.len(), validity_mask.len()); + + let count = validity_mask.true_count(); + if count == 0 { + return Ok(ValidIndicesStats { + count, + sorted: true, + }); + } + + let sorted = match validity_mask.bit_buffer() { + AllOr::All => valid_indices_sorted_all(indices, array_len, check_bounds)?, + AllOr::None => true, + AllOr::Some(validity) => { + valid_indices_sorted_masked(indices, validity.iter(), array_len, check_bounds)? + } + }; + + Ok(ValidIndicesStats { count, sorted }) +} + +fn valid_indices_sorted_all>( + indices: &[T], + array_len: usize, + check_bounds: bool, +) -> VortexResult { + // Seed the comparison with the first index; an empty or single-element slice is trivially + // sorted, so the loop below starts from the second element. + let Some((first, rest)) = indices.split_first() else { + return Ok(true); + }; + + let mut previous_idx = first.as_(); + if check_bounds { + check_index(previous_idx, array_len)?; + } + + let mut sorted = true; + for idx in rest { + let idx = idx.as_(); + if check_bounds { + check_index(idx, array_len)?; + } + if previous_idx > idx { + sorted = false; + if !check_bounds { + break; + } + } + previous_idx = idx; + } + + Ok(sorted) +} + +fn valid_indices_sorted_masked>( + indices: &[T], + is_valid: impl Iterator, + array_len: usize, + check_bounds: bool, +) -> VortexResult { + // Invalid positions are skipped without a bounds check, matching the take path that never + // dereferences them. + let mut valid = is_valid + .zip(indices.iter()) + .filter(|(is_valid, _)| *is_valid) + .map(|(_, idx)| idx.as_()); + + // Seed the comparison with the first valid index; zero or one valid index is trivially + // sorted, so the loop below starts from the second valid index. + let Some(mut previous_idx) = valid.next() else { + return Ok(true); + }; + if check_bounds { + check_index(previous_idx, array_len)?; + } + + let mut sorted = true; + for idx in valid { + if check_bounds { + check_index(idx, array_len)?; + } + if previous_idx > idx { + sorted = false; + if !check_bounds { + break; + } + } + previous_idx = idx; + } + + Ok(sorted) +} + +fn prefer_linear_scan( + ends_len: usize, + valid_count: usize, + runs_per_index_threshold: usize, +) -> bool { + ends_len <= valid_count.saturating_mul(runs_per_index_threshold) +} + +fn check_index(index: usize, array_len: usize) -> VortexResult<()> { + if index >= array_len { + vortex_bail!(OutOfBounds: index, 0, array_len); + } + Ok(()) +} + +fn physical_indices_linear_sorted( + ends: &[I], + offset: usize, + indices: &[T], + validity_mask: &Mask, +) -> Buffer +where + I: UnsignedPType, + T: AsPrimitive, + O: UnsignedPType, + usize: AsPrimitive, +{ + let mut run_idx = 0; + + match validity_mask.bit_buffer() { + AllOr::All => Buffer::from_trusted_len_iter(indices.iter().map(|idx| { + advance_run(ends, &mut run_idx, idx.as_() + offset); + run_idx.as_() + })), + AllOr::None => unreachable!("AllInvalid indices have been handled earlier"), + AllOr::Some(validity) => { + // Invalid positions keep physical index zero, which is always in-bounds for the + // values and masked out by the result validity. + let mut physical_indices = BufferMut::zeroed(indices.len()); + for (idx_pos, (is_valid, idx)) in validity.iter().zip(indices.iter()).enumerate() { + if !is_valid { + continue; + } + + advance_run(ends, &mut run_idx, idx.as_() + offset); + physical_indices[idx_pos] = run_idx.as_(); + } + physical_indices.freeze() + } + } +} + +/// Resolve indices through a dense logical-position-to-run-index table. +/// +/// Building the table costs O(array_len), but every index then resolves with an unconditional +/// gather, which beats per-index binary search and sort-then-merge for dense takes. Invalid +/// indices may hold arbitrary values (even out of bounds), so they are redirected to position +/// zero instead of branching; the result validity masks whatever they resolve to. +fn physical_indices_table( + ends: &[I], + offset: usize, + array_len: usize, + indices: &[T], + validity_mask: &Mask, +) -> Buffer +where + I: UnsignedPType, + T: AsPrimitive, + O: UnsignedPType, + usize: AsPrimitive, +{ + let table = run_index_table::(ends, offset, array_len); + let table = table.as_slice(); + + match validity_mask.bit_buffer() { + AllOr::All => Buffer::from_trusted_len_iter(indices.iter().map(|idx| table[idx.as_()])), + AllOr::None => unreachable!("AllInvalid indices have been handled earlier"), + AllOr::Some(validity) => Buffer::from_trusted_len_iter( + validity + .iter() + .zip(indices.iter()) + .map(|(is_valid, idx)| table[if is_valid { idx.as_() } else { 0 }]), + ), + } +} + +/// Materialize the run index of every logical position in `[0, len)`. +fn run_index_table(ends: &[I], offset: usize, len: usize) -> Buffer +where + I: UnsignedPType, + O: UnsignedPType, + usize: AsPrimitive, +{ + let mut table = BufferMut::with_capacity(len); + let mut run_start = 0; + for (run_idx, run_end) in trimmed_ends_iter(ends, offset, len).enumerate() { + table.push_n(run_idx.as_(), run_end - run_start); + run_start = run_end; + } + table.freeze() +} + +fn physical_indices_linear_unsorted( + ends: &[I], + offset: usize, + indices: &[T], + validity_mask: &Mask, + valid_count: usize, +) -> Buffer +where + I: UnsignedPType, + T: AsPrimitive, + O: UnsignedPType, + usize: AsPrimitive, +{ + let mut pairs = Vec::with_capacity(valid_count); + match validity_mask.bit_buffer() { + AllOr::All => { + pairs.extend( + indices + .iter() + .enumerate() + .map(|(idx_pos, idx)| (idx.as_(), idx_pos)), + ); + } + AllOr::None => unreachable!("AllInvalid indices have been handled earlier"), + AllOr::Some(validity) => { + for (idx_pos, (is_valid, idx)) in validity.iter().zip(indices.iter()).enumerate() { + if is_valid { + pairs.push((idx.as_(), idx_pos)); + } + } + } + } + pairs.sort_unstable(); + + let mut physical_indices = BufferMut::zeroed(indices.len()); + let mut run_idx = 0; + + for (idx, idx_pos) in pairs { + advance_run(ends, &mut run_idx, idx + offset); + physical_indices[idx_pos] = run_idx.as_(); + } + + physical_indices.freeze() +} + +fn physical_indices_binary( + ends: &[I], + offset: usize, + indices: &[T], + validity_mask: &Mask, +) -> Buffer +where + I: UnsignedPType, + T: AsPrimitive, + O: UnsignedPType, + usize: AsPrimitive, +{ + match validity_mask.bit_buffer() { + AllOr::All => Buffer::from_trusted_len_iter( + indices + .iter() + .map(|idx| physical_index_binary(ends, idx.as_() + offset).as_()), + ), + AllOr::None => Buffer::zeroed(indices.len()), + AllOr::Some(validity) => { + let mut physical_indices = BufferMut::zeroed(indices.len()); + for (idx_pos, (is_valid, idx)) in validity.iter().zip(indices.iter()).enumerate() { + if !is_valid { + continue; + } + + physical_indices[idx_pos] = physical_index_binary(ends, idx.as_() + offset).as_(); + } + physical_indices.freeze() + } + } +} + +fn physical_index_binary(ends: &[I], logical_idx: usize) -> usize { + let index = match ::from(logical_idx) { + Some(logical_idx) => ends.partition_point(|end| *end <= logical_idx), + None => ends.len(), + }; + index.min(ends.len() - 1) +} + +fn advance_run(ends: &[I], run_idx: &mut usize, logical_idx: usize) { + // A logical index that overflows the run-end type sits past every run, so it lands in the + // final run; otherwise advance while the current run ends at or before it. + let Some(logical_idx) = I::from(logical_idx) else { + *run_idx = ends.len().saturating_sub(1); + return; + }; + while *run_idx + 1 < ends.len() && ends[*run_idx] <= logical_idx { + *run_idx += 1; + } +} + #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::take::test_take_conformance; + use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_mask::Mask; + use vortex_session::VortexSession; + use super::physical_indices_binary; + use super::physical_indices_linear_sorted; + use super::physical_indices_linear_unsorted; + use super::physical_indices_table; use crate::RunEnd; use crate::RunEndArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn ree_array() -> RunEndArray { RunEnd::encode( buffer![1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap() } @@ -116,14 +525,23 @@ mod tests { fn ree_take() { let taken = ree_array().take(buffer![9, 8, 1, 3].into_array()).unwrap(); let expected = PrimitiveArray::from_iter(vec![5i32, 5, 1, 4]).into_array(); - assert_arrays_eq!(taken, expected); + assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx()); } #[test] fn ree_take_end() { let taken = ree_array().take(buffer![11].into_array()).unwrap(); let expected = PrimitiveArray::from_iter(vec![5i32]).into_array(); - assert_arrays_eq!(taken, expected); + assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx()); + } + + #[test] + fn ree_take_sorted_boundaries() { + let taken = ree_array() + .take(buffer![0, 2, 3, 6, 8, 11].into_array()) + .unwrap(); + let expected = PrimitiveArray::from_iter(vec![1i32, 1, 4, 2, 5, 5]).into_array(); + assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -132,7 +550,7 @@ mod tests { let _array = ree_array() .take(buffer![12].into_array()) .unwrap() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut SESSION.create_execution_ctx()) .unwrap(); } @@ -142,7 +560,16 @@ mod tests { let taken = sliced.take(buffer![1, 3, 4].into_array()).unwrap(); let expected = PrimitiveArray::from_iter(vec![4i32, 2, 5]).into_array(); - assert_arrays_eq!(taken, expected); + assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx()); + } + + #[test] + fn sliced_take_unsorted_dense() { + let sliced = ree_array().slice(4..9).unwrap(); + let taken = sliced.take(buffer![4, 0, 2, 1].into_array()).unwrap(); + + let expected = PrimitiveArray::from_iter(vec![5i32, 4, 2, 4]).into_array(); + assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -152,14 +579,143 @@ mod tests { .unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1i32), None]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); + } + + #[test] + fn ree_take_all_null_indices() { + let taken = ree_array() + .take(PrimitiveArray::from_option_iter([None::, None]).into_array()) + .unwrap(); + + let expected = PrimitiveArray::from_option_iter([None::, None]); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); + } + + #[test] + fn ree_take_null_index_skips_out_of_bounds_value() { + let indices = PrimitiveArray::new( + buffer![1u64, 12], + Validity::Array(BoolArray::from_iter([true, false]).into_array()), + ); + let taken = ree_array().take(indices.into_array()).unwrap(); + + let expected = PrimitiveArray::from_option_iter([Some(1i32), None]); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); + } + + #[test] + fn ree_take_unsorted_null_index_skips_out_of_bounds_value() { + let indices = PrimitiveArray::new( + buffer![3u64, 12, 1], + Validity::Array(BoolArray::from_iter([true, false, true]).into_array()), + ); + let taken = ree_array().take(indices.into_array()).unwrap(); + + let expected = PrimitiveArray::from_option_iter([Some(4i32), None, Some(1)]); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); + } + + #[test] + fn ree_take_dense_null_index_skips_out_of_bounds_value() { + let indices = PrimitiveArray::new( + buffer![0u64, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12], + Validity::Array( + BoolArray::from_iter([ + true, true, true, true, true, true, true, true, true, true, true, false, + ]) + .into_array(), + ), + ); + let taken = ree_array().take(indices.into_array()).unwrap(); + + let expected = PrimitiveArray::from_option_iter([ + Some(1i32), + Some(1), + Some(1), + Some(4), + Some(4), + Some(4), + Some(2), + Some(2), + Some(5), + Some(5), + Some(5), + None, + ]); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); + } + + #[rstest] + #[case(vec![3u32, 6, 8, 12], 0, 12, vec![0u64, 11, 3, 3, 7, 2, 9], Mask::new_true(7))] + #[case(vec![3u32, 6, 8, 12], 0, 12, vec![5u64, 100, 2, 11, 0], Mask::from_indices(5, [0, 2, 3, 4]))] + #[case(vec![6u32, 8, 12], 4, 5, vec![4u64, 0, 2, 1, 3], Mask::new_true(5))] + fn unsorted_strategies_agree( + #[case] ends: Vec, + #[case] offset: usize, + #[case] len: usize, + #[case] indices: Vec, + #[case] mask: Mask, + ) { + let binary = physical_indices_binary::(&ends, offset, &indices, &mask); + let table = physical_indices_table::(&ends, offset, len, &indices, &mask); + let sort_merge = physical_indices_linear_unsorted::( + &ends, + offset, + &indices, + &mask, + mask.true_count(), + ); + + assert_eq!(binary.as_slice(), table.as_slice()); + assert_eq!(binary.as_slice(), sort_merge.as_slice()); + } + + #[rstest] + #[case(vec![3u32, 6, 8, 12], 0, 12, vec![0u64, 2, 3, 6, 8, 11], Mask::new_true(6))] + #[case(vec![3u32, 6, 8, 12], 0, 12, vec![1u64, 100, 5, 9], Mask::from_indices(4, [0, 2, 3]))] + #[case(vec![6u32, 8, 12], 4, 5, vec![0u64, 1, 3, 4], Mask::new_true(4))] + fn sorted_strategies_agree( + #[case] ends: Vec, + #[case] offset: usize, + #[case] len: usize, + #[case] indices: Vec, + #[case] mask: Mask, + ) { + let binary = physical_indices_binary::(&ends, offset, &indices, &mask); + let table = physical_indices_table::(&ends, offset, len, &indices, &mask); + let sorted = + physical_indices_linear_sorted::(&ends, offset, &indices, &mask); + + assert_eq!(binary.as_slice(), table.as_slice()); + assert_eq!(binary.as_slice(), sorted.as_slice()); } #[rstest] #[case(ree_array())] #[case(RunEnd::encode( buffer![1u8, 1, 2, 2, 2, 3, 3, 3, 3, 4].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case(RunEnd::encode( PrimitiveArray::from_option_iter([ @@ -172,14 +728,14 @@ mod tests { Some(20), ]) .into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case(RunEnd::encode(buffer![42i32, 42, 42, 42, 42].into_array(), - &mut LEGACY_SESSION.create_execution_ctx()) + &mut SESSION.create_execution_ctx()) .unwrap())] #[case(RunEnd::encode( buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case({ let mut values = Vec::new(); @@ -190,7 +746,7 @@ mod tests { } RunEnd::encode( PrimitiveArray::from_iter(values).into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap() })] @@ -203,7 +759,7 @@ mod tests { #[case({ let array = RunEnd::encode( buffer![1i32, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); array.slice(2..8).unwrap() diff --git a/encodings/runend/src/compute/take_from.rs b/encodings/runend/src/compute/take_from.rs index 32b22eb96c5..0987cd3fc55 100644 --- a/encodings/runend/src/compute/take_from.rs +++ b/encodings/runend/src/compute/take_from.rs @@ -52,9 +52,12 @@ impl ExecuteParentKernel for RunEndTakeFrom { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -68,6 +71,12 @@ mod tests { use crate::array::RunEndArrayExt; use crate::compute::take_from::RunEndTakeFrom; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + /// Build a DictArray whose codes are run-end encoded. /// /// Input: `[2, 2, 2, 3, 3, 2, 2]` @@ -83,7 +92,7 @@ mod tests { #[test] fn test_execute_parent_no_offset() -> VortexResult<()> { - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = SESSION.create_execution_ctx(); let (codes, dict) = make_dict_with_runend_codes(&mut ctx); let result = RunEndTakeFrom @@ -92,13 +101,13 @@ mod tests { let expected = PrimitiveArray::from_iter([2i32, 2, 2, 3, 3, 2, 2]); let canonical = result.execute::(&mut ctx)?.into_array(); - assert_arrays_eq!(canonical, expected); + assert_arrays_eq!(canonical, expected, &mut ctx); Ok(()) } #[test] fn test_execute_parent_with_offset() -> VortexResult<()> { - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = SESSION.create_execution_ctx(); let (codes, dict) = make_dict_with_runend_codes(&mut ctx); // Slice codes to positions 2..5 → logical codes [0, 1, 1] → values [2, 3, 3] let sliced_codes = unsafe { @@ -116,13 +125,13 @@ mod tests { let expected = PrimitiveArray::from_iter([2i32, 3, 3]); let canonical = result.execute::(&mut ctx)?.into_array(); - assert_arrays_eq!(canonical, expected); + assert_arrays_eq!(canonical, expected, &mut ctx); Ok(()) } #[test] fn test_execute_parent_offset_at_run_boundary() -> VortexResult<()> { - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = SESSION.create_execution_ctx(); let (codes, dict) = make_dict_with_runend_codes(&mut ctx); // Slice codes to positions 3..7 → logical codes [1, 1, 0, 0] → values [3, 3, 2, 2] let sliced_codes = unsafe { @@ -140,13 +149,13 @@ mod tests { let expected = PrimitiveArray::from_iter([3i32, 3, 2, 2]); let canonical = result.execute::(&mut ctx)?.into_array(); - assert_arrays_eq!(canonical, expected); + assert_arrays_eq!(canonical, expected, &mut ctx); Ok(()) } #[test] fn test_execute_parent_single_element_offset() -> VortexResult<()> { - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = SESSION.create_execution_ctx(); let (codes, dict) = make_dict_with_runend_codes(&mut ctx); // Slice to single element at position 4 → code=1 → value=3 let sliced_codes = unsafe { @@ -164,13 +173,13 @@ mod tests { let expected = PrimitiveArray::from_iter([3i32]); let canonical = result.execute::(&mut ctx)?.into_array(); - assert_arrays_eq!(canonical, expected); + assert_arrays_eq!(canonical, expected, &mut ctx); Ok(()) } #[test] fn test_execute_parent_returns_none_for_non_codes_child() -> VortexResult<()> { - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = SESSION.create_execution_ctx(); let (codes, dict) = make_dict_with_runend_codes(&mut ctx); let result = RunEndTakeFrom.execute_parent(codes.as_view(), dict.as_view(), 1, &mut ctx)?; diff --git a/encodings/runend/src/decompress_bool.rs b/encodings/runend/src/decompress_bool.rs index 68b6d3d5981..d2b778fdbbf 100644 --- a/encodings/runend/src/decompress_bool.rs +++ b/encodings/runend/src/decompress_bool.rs @@ -246,7 +246,8 @@ fn decode_nullable_sequential( #[cfg(test)] mod tests { - use vortex_array::LEGACY_SESSION; + use std::sync::LazyLock; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; @@ -255,12 +256,19 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use super::runend_decode_bools; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn decode_bools_alternating() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Alternating true/false: [T, T, F, F, F, T, T, T, T, T] let ends = PrimitiveArray::from_iter([2u32, 5, 10]); let values = BoolArray::from(BitBuffer::from(vec![true, false, true])); @@ -269,13 +277,13 @@ mod tests { let expected = BoolArray::from(BitBuffer::from(vec![ true, true, false, false, false, true, true, true, true, true, ])); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn decode_bools_mostly_true() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Mostly true: [T, T, T, T, T, F, T, T, T, T] let ends = PrimitiveArray::from_iter([5u32, 6, 10]); let values = BoolArray::from(BitBuffer::from(vec![true, false, true])); @@ -284,13 +292,13 @@ mod tests { let expected = BoolArray::from(BitBuffer::from(vec![ true, true, true, true, true, false, true, true, true, true, ])); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn decode_bools_mostly_false() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Mostly false: [F, F, F, F, F, T, F, F, F, F] let ends = PrimitiveArray::from_iter([5u32, 6, 10]); let values = BoolArray::from(BitBuffer::from(vec![false, true, false])); @@ -299,13 +307,13 @@ mod tests { let expected = BoolArray::from(BitBuffer::from(vec![ false, false, false, false, false, true, false, false, false, false, ])); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn decode_bools_all_true_single_run() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let ends = PrimitiveArray::from_iter([10u32]); let values = BoolArray::from(BitBuffer::from(vec![true])); let decoded = runend_decode_bools(ends, values, 0, 10, &mut ctx)?; @@ -313,13 +321,13 @@ mod tests { let expected = BoolArray::from(BitBuffer::from(vec![ true, true, true, true, true, true, true, true, true, true, ])); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn decode_bools_all_false_single_run() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let ends = PrimitiveArray::from_iter([10u32]); let values = BoolArray::from(BitBuffer::from(vec![false])); let decoded = runend_decode_bools(ends, values, 0, 10, &mut ctx)?; @@ -327,13 +335,13 @@ mod tests { let expected = BoolArray::from(BitBuffer::from(vec![ false, false, false, false, false, false, false, false, false, false, ])); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn decode_bools_with_offset() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test with offset: [T, T, F, F, F, T, T, T, T, T] -> slice [2..8] = [F, F, F, T, T, T] let ends = PrimitiveArray::from_iter([2u32, 5, 10]); let values = BoolArray::from(BitBuffer::from(vec![true, false, true])); @@ -341,7 +349,7 @@ mod tests { let expected = BoolArray::from(BitBuffer::from(vec![false, false, false, true, true, true])); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } @@ -349,7 +357,7 @@ mod tests { fn decode_bools_nullable() -> VortexResult<()> { use vortex_array::validity::Validity; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // 3 runs: T (valid), F (null), T (valid) -> [T, T, null, null, null, T, T, T, T, T] let ends = PrimitiveArray::from_iter([2u32, 5, 10]); let values = BoolArray::new( @@ -367,13 +375,13 @@ mod tests { true, true, false, false, false, true, true, true, true, true, ])), ); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn decode_bools_nullable_few_runs() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test few runs (uses fast path): 5 runs of length 2000 each let ends = PrimitiveArray::from_iter([2000u32, 4000, 6000, 8000, 10000]); let values = BoolArray::new( diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index 73f6ebc524e..484a4233789 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -4,29 +4,36 @@ use std::ops::Range; use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::kernel::ExecuteParentKernel; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_error::VortexResult; +use vortex_session::VortexSession; use crate::RunEnd; use crate::array::RunEndArrayExt; use crate::compute::take_from::RunEndTakeFrom; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(RunEnd)), - ParentKernelSet::lift(&RunEndSliceKernel), - ParentKernelSet::lift(&FilterExecuteAdaptor(RunEnd)), - ParentKernelSet::lift(&TakeExecuteAdaptor(RunEnd)), - ParentKernelSet::lift(&RunEndTakeFrom), -]); +pub(super) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), RunEnd, CompareExecuteAdaptor(RunEnd)); + kernels.register_execute_parent_kernel(Slice.id(), RunEnd, RunEndSliceKernel); + kernels.register_execute_parent_kernel(Filter.id(), RunEnd, FilterExecuteAdaptor(RunEnd)); + kernels.register_execute_parent_kernel(Dict.id(), RunEnd, TakeExecuteAdaptor(RunEnd)); + kernels.register_execute_parent_kernel(Dict.id(), RunEnd, RunEndTakeFrom); +} /// Kernel to execute slicing on a RunEnd array. /// diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index 8770dbbf58e..8169f5d8dfc 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -38,6 +38,7 @@ use vortex_session::VortexSession; /// Initialize run-end encoding in the given session. pub fn initialize(session: &VortexSession) { session.arrays().register(RunEnd); + kernel::initialize(session); // Register the RunEnd-specific aggregate kernels. session.aggregate_fns().register_aggregate_kernel( @@ -59,12 +60,21 @@ pub fn initialize(session: &VortexSession) { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use prost::Message; use vortex_array::dtype::PType; use vortex_array::test_harness::check_metadata; + use vortex_session::VortexSession; use crate::RunEndMetadata; + pub static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[cfg_attr(miri, ignore)] #[test] fn test_runend_metadata() { diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index 41a9e69419c..56de024f988 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -49,9 +49,9 @@ pub(crate) fn find_slice_end_index(array: &ArrayRef, index: usize) -> VortexResu #[cfg(test)] mod tests { + use std::sync::LazyLock; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::arrays::PrimitiveArray; @@ -60,12 +60,19 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_buffer::buffer; + use vortex_session::VortexSession; use crate::RunEnd; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn slice_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = RunEnd::try_new( buffer![2u32, 5, 10].into_array(), buffer![1i32, 2, 3].into_array(), @@ -81,12 +88,12 @@ mod tests { assert_eq!(arr.len(), 5); let expected = PrimitiveArray::from_iter(vec![2i32, 2, 3, 3, 3]).into_array(); - assert_arrays_eq!(arr, expected); + assert_arrays_eq!(arr, expected, &mut ctx); } #[test] fn double_slice() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = RunEnd::try_new( buffer![2u32, 5, 10].into_array(), buffer![1i32, 2, 3].into_array(), @@ -100,12 +107,12 @@ mod tests { let doubly_sliced = arr.slice(0..3).unwrap(); let expected = PrimitiveArray::from_iter(vec![2i32, 2, 3]).into_array(); - assert_arrays_eq!(doubly_sliced, expected); + assert_arrays_eq!(doubly_sliced, expected, &mut ctx); } #[test] fn slice_end_inclusive() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = RunEnd::try_new( buffer![2u32, 5, 10].into_array(), buffer![1i32, 2, 3].into_array(), @@ -121,12 +128,12 @@ mod tests { assert_eq!(arr.len(), 6); let expected = PrimitiveArray::from_iter(vec![2i32, 3, 3, 3, 3, 3]).into_array(); - assert_arrays_eq!(arr, expected); + assert_arrays_eq!(arr, expected, &mut ctx); } #[test] fn slice_at_end() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let re_array = RunEnd::try_new( buffer![7_u64, 10].into_array(), buffer![2_u64, 3].into_array(), @@ -142,7 +149,7 @@ mod tests { #[test] fn slice_single_end() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let re_array = RunEnd::try_new( buffer![7_u64, 10].into_array(), buffer![2_u64, 3].into_array(), @@ -159,7 +166,7 @@ mod tests { #[test] fn ree_scalar_at_end() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let scalar = RunEnd::encode( buffer![1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5].into_array(), &mut ctx, @@ -172,7 +179,7 @@ mod tests { #[test] fn slice_along_run_boundaries() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a runend array with runs: [1, 1, 1] [4, 4, 4] [2, 2] [5, 5, 5, 5] // Run ends at indices: 3, 6, 8, 12 let arr = RunEnd::try_new( @@ -186,42 +193,42 @@ mod tests { let slice1 = arr.slice(0..3).unwrap(); assert_eq!(slice1.len(), 3); let expected = PrimitiveArray::from_iter(vec![1i32, 1, 1]).into_array(); - assert_arrays_eq!(slice1, expected); + assert_arrays_eq!(slice1, expected, &mut ctx); // Slice from start of second run to end of second run (indices 3..6) let slice2 = arr.slice(3..6).unwrap(); assert_eq!(slice2.len(), 3); let expected = PrimitiveArray::from_iter(vec![4i32, 4, 4]).into_array(); - assert_arrays_eq!(slice2, expected); + assert_arrays_eq!(slice2, expected, &mut ctx); // Slice from start of third run to end of third run (indices 6..8) let slice3 = arr.slice(6..8).unwrap(); assert_eq!(slice3.len(), 2); let expected = PrimitiveArray::from_iter(vec![2i32, 2]).into_array(); - assert_arrays_eq!(slice3, expected); + assert_arrays_eq!(slice3, expected, &mut ctx); // Slice from start of last run to end of last run (indices 8..12) let slice4 = arr.slice(8..12).unwrap(); assert_eq!(slice4.len(), 4); let expected = PrimitiveArray::from_iter(vec![5i32, 5, 5, 5]).into_array(); - assert_arrays_eq!(slice4, expected); + assert_arrays_eq!(slice4, expected, &mut ctx); // Slice spanning exactly two runs (indices 3..8) let slice5 = arr.slice(3..8).unwrap(); assert_eq!(slice5.len(), 5); let expected = PrimitiveArray::from_iter(vec![4i32, 4, 4, 2, 2]).into_array(); - assert_arrays_eq!(slice5, expected); + assert_arrays_eq!(slice5, expected, &mut ctx); // Slice from middle of first run to end of second run (indices 1..6) let slice6 = arr.slice(1..6).unwrap(); assert_eq!(slice6.len(), 5); let expected = PrimitiveArray::from_iter(vec![1i32, 1, 4, 4, 4]).into_array(); - assert_arrays_eq!(slice6, expected); + assert_arrays_eq!(slice6, expected, &mut ctx); // Slice from start of second run to middle of third run (indices 3..7) let slice7 = arr.slice(3..7).unwrap(); assert_eq!(slice7.len(), 4); let expected = PrimitiveArray::from_iter(vec![4i32, 4, 4, 2]).into_array(); - assert_arrays_eq!(slice7, expected); + assert_arrays_eq!(slice7, expected, &mut ctx); } } diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index af8b146cfba..3b8e0065f06 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -49,7 +49,6 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::compress::sequence_decompress; -use crate::kernel::PARENT_KERNELS; use crate::rules::RULES; /// A [`Sequence`]-encoded Vortex array. @@ -259,6 +258,14 @@ impl VTable for Sequence { vortex_panic!("SequenceArray buffer_name index {idx} out of bounds") } + 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, @@ -331,15 +338,6 @@ impl VTable for Sequence { sequence_decompress(&array).map(ExecutionResult::done) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, @@ -452,7 +450,8 @@ impl Sequence { #[cfg(test)] mod tests { - use vortex_array::LEGACY_SESSION; + use std::sync::LazyLock; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -463,16 +462,23 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::Sequence; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_sequence_canonical() { let arr = Sequence::try_new_typed(2i64, 3, Nullability::NonNullable, 4).unwrap(); let canon = PrimitiveArray::from_iter((0..4).map(|i| 2i64 + i * 3)); - assert_arrays_eq!(arr, canon); + assert_arrays_eq!(arr, canon, &mut SESSION.create_execution_ctx()); } #[test] @@ -484,14 +490,14 @@ mod tests { let canon = PrimitiveArray::from_iter((2..3).map(|i| 2i64 + i * 3)); - assert_arrays_eq!(arr, canon); + assert_arrays_eq!(arr, canon, &mut SESSION.create_execution_ctx()); } #[test] fn test_sequence_scalar_at() { let scalar = Sequence::try_new_typed(2i64, 3, Nullability::NonNullable, 4) .unwrap() - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut SESSION.create_execution_ctx()) .unwrap(); assert_eq!( diff --git a/encodings/sequence/src/compress.rs b/encodings/sequence/src/compress.rs index 9f5e62a5217..b20980cf132 100644 --- a/encodings/sequence/src/compress.rs +++ b/encodings/sequence/src/compress.rs @@ -148,18 +148,24 @@ fn encode_primitive_array + CheckedAdd + CheckedSu #[cfg(test)] mod tests { - #[expect(unused_imports)] - use itertools::Itertools; - use vortex_array::LEGACY_SESSION; + use std::sync::LazyLock; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_session::VortexSession; use crate::sequence_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_encode_array_success() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive_array = PrimitiveArray::from_iter([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); let encoded = sequence_encode(primitive_array.as_view(), &mut ctx).unwrap(); assert!(encoded.is_some()); @@ -167,12 +173,12 @@ mod tests { .unwrap() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(decoded, primitive_array); + assert_arrays_eq!(decoded, primitive_array, &mut ctx); } #[test] fn test_encode_array_1_success() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive_array = PrimitiveArray::from_iter([0]); let encoded = sequence_encode(primitive_array.as_view(), &mut ctx).unwrap(); assert!(encoded.is_some()); @@ -180,12 +186,12 @@ mod tests { .unwrap() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(decoded, primitive_array); + assert_arrays_eq!(decoded, primitive_array, &mut ctx); } #[test] fn test_encode_array_fail() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive_array = PrimitiveArray::from_iter([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0]); let encoded = sequence_encode(primitive_array.as_view(), &mut ctx).unwrap(); @@ -194,7 +200,7 @@ mod tests { #[test] fn test_encode_array_fail_oob() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive_array = PrimitiveArray::from_iter(vec![100i8; 1000]); let encoded = sequence_encode(primitive_array.as_view(), &mut ctx).unwrap(); @@ -203,7 +209,7 @@ mod tests { #[test] fn test_encode_all_u8_values() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive_array = PrimitiveArray::from_iter(0u8..=255); let encoded = sequence_encode(primitive_array.as_view(), &mut ctx).unwrap(); assert!(encoded.is_some()); @@ -211,6 +217,6 @@ mod tests { .unwrap() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(decoded, primitive_array); + assert_arrays_eq!(decoded, primitive_array, &mut ctx); } } diff --git a/encodings/sequence/src/compute/cast.rs b/encodings/sequence/src/compute/cast.rs index e6d64fdf7c5..e0a63d0c922 100644 --- a/encodings/sequence/src/compute/cast.rs +++ b/encodings/sequence/src/compute/cast.rs @@ -99,14 +99,16 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::session::ArraySession; use vortex_session::VortexSession; use crate::Sequence; use crate::SequenceArray; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_cast_sequence_nullability() { @@ -139,7 +141,11 @@ mod tests { // Verify the values let decoded = casted.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(decoded, PrimitiveArray::from_iter([100i64, 110, 120, 130])); + assert_arrays_eq!( + decoded, + PrimitiveArray::from_iter([100i64, 110, 120, 130]), + &mut ctx + ); } #[test] @@ -161,7 +167,8 @@ mod tests { let decoded = casted.execute::(&mut ctx).unwrap(); assert_arrays_eq!( decoded, - PrimitiveArray::from_option_iter([Some(5i32), Some(8), Some(11)]) + PrimitiveArray::from_option_iter([Some(5i32), Some(8), Some(11)]), + &mut ctx ); } @@ -185,7 +192,8 @@ mod tests { let decoded = casted.execute::(&mut ctx).unwrap(); assert_arrays_eq!( decoded, - PrimitiveArray::from_iter([0.0f32, 1.0, 2.0, 3.0, 4.0]) + PrimitiveArray::from_iter([0.0f32, 1.0, 2.0, 3.0, 4.0]), + &mut ctx ); } diff --git a/encodings/sequence/src/compute/compare.rs b/encodings/sequence/src/compute/compare.rs index c0c9d1367d6..5895330ea81 100644 --- a/encodings/sequence/src/compute/compare.rs +++ b/encodings/sequence/src/compute/compare.rs @@ -135,7 +135,10 @@ fn find_intersection( #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::assert_arrays_eq; @@ -143,9 +146,16 @@ mod tests { use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::Nullability::Nullable; use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_session::VortexSession; use crate::Sequence; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_compare_match() { let lhs = Sequence::try_new_typed(2i64, 1, NonNullable, 4).unwrap(); @@ -155,7 +165,7 @@ mod tests { .binary(rhs.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([false, false, true, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -167,7 +177,7 @@ mod tests { .binary(rhs.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([Some(false), Some(false), Some(true), Some(false)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -179,6 +189,6 @@ mod tests { .binary(rhs.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([false, false, false, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/sequence/src/compute/list_contains.rs b/encodings/sequence/src/compute/list_contains.rs index fa2cacb5e19..93bd823a5f2 100644 --- a/encodings/sequence/src/compute/list_contains.rs +++ b/encodings/sequence/src/compute/list_contains.rs @@ -52,8 +52,10 @@ impl ListContainsElementReduce for Sequence { #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::Nullability; @@ -62,9 +64,16 @@ mod tests { use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::scalar::Scalar; + use vortex_session::VortexSession; use crate::Sequence; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_list_contains_seq() { let list_scalar = Scalar::list( @@ -84,7 +93,7 @@ mod tests { let expr = list_contains(lit(list_scalar.clone()), root()); let result = array.into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([Some(true), Some(false), Some(true)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } { @@ -98,7 +107,7 @@ mod tests { let expr = list_contains(lit(list_scalar), root()); let result = array.into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([Some(true), Some(true), Some(false)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } } } diff --git a/encodings/sequence/src/compute/mod.rs b/encodings/sequence/src/compute/mod.rs index 9a99ff235c4..22b4e3ac05b 100644 --- a/encodings/sequence/src/compute/mod.rs +++ b/encodings/sequence/src/compute/mod.rs @@ -14,6 +14,8 @@ mod take; mod tests { use rstest::rstest; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::dtype::Nullability; @@ -84,6 +86,9 @@ mod tests { ).unwrap())] // Results in [1000, 1100, 1200, ..., 150900] fn test_sequence_consistency(#[case] array: SequenceArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/encodings/sequence/src/compute/take.rs b/encodings/sequence/src/compute/take.rs index 6e1ebe263df..4b056d0ae7f 100644 --- a/encodings/sequence/src/compute/take.rs +++ b/encodings/sequence/src/compute/take.rs @@ -104,8 +104,8 @@ mod test { use rstest::rstest; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::Nullability; @@ -174,7 +174,7 @@ mod test { let _array = array .take(indices.into_array()) .unwrap() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut array_session().create_execution_ctx()) .unwrap(); } } diff --git a/encodings/sequence/src/kernel.rs b/encodings/sequence/src/kernel.rs index 394a52ec69e..37521a261e7 100644 --- a/encodings/sequence/src/kernel.rs +++ b/encodings/sequence/src/kernel.rs @@ -1,15 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_session::VortexSession; use crate::Sequence; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(Sequence)), - ParentKernelSet::lift(&FilterExecuteAdaptor(Sequence)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Sequence)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), Sequence, CompareExecuteAdaptor(Sequence)); + kernels.register_execute_parent_kernel(Filter.id(), Sequence, FilterExecuteAdaptor(Sequence)); + kernels.register_execute_parent_kernel(Dict.id(), Sequence, TakeExecuteAdaptor(Sequence)); +} diff --git a/encodings/sequence/src/lib.rs b/encodings/sequence/src/lib.rs index bd6ab2f508c..b898963c346 100644 --- a/encodings/sequence/src/lib.rs +++ b/encodings/sequence/src/lib.rs @@ -27,6 +27,7 @@ use vortex_session::VortexSession; /// Initialize sequence encoding in the given session. pub fn initialize(session: &VortexSession) { session.arrays().register(Sequence); + kernel::initialize(session); // Register the Sequence-specific aggregate kernels. session.aggregate_fns().register_aggregate_kernel( diff --git a/encodings/sparse/benches/sparse_canonical.rs b/encodings/sparse/benches/sparse_canonical.rs index 3861e26da06..27f7340fb91 100644 --- a/encodings/sparse/benches/sparse_canonical.rs +++ b/encodings/sparse/benches/sparse_canonical.rs @@ -17,7 +17,6 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType::I32; use vortex_array::scalar::Scalar; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexExpect; @@ -28,8 +27,11 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_sparse::initialize(&session); + session +}); const LIST_ARGS: &[(usize, usize, usize)] = &[ // len, patch_stride, list_size diff --git a/encodings/sparse/benches/sparse_pushdown.rs b/encodings/sparse/benches/sparse_pushdown.rs index 5abf0905f6e..a066e196d76 100644 --- a/encodings/sparse/benches/sparse_pushdown.rs +++ b/encodings/sparse/benches/sparse_pushdown.rs @@ -19,6 +19,7 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::aggregate_fn::fns::min_max::min_max; use vortex_array::aggregate_fn::fns::null_count::null_count; @@ -31,7 +32,6 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; -use vortex_array::session::ArraySession; use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_session::VortexSession; @@ -45,7 +45,7 @@ const LEN: usize = 1_000_000; /// Session with Sparse and its pushdown kernels registered. static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); vortex_sparse::initialize(&session); session }); @@ -107,7 +107,10 @@ fn sparse_min_max(bencher: Bencher) { bencher .with_inputs(|| (make_sparse(40_000, false), SESSION.create_execution_ctx())) .bench_values(|(array, mut ctx)| { - divan::black_box(min_max(&array, &mut ctx).vortex_expect("min_max")) + divan::black_box( + min_max(&array, &mut ctx, NumericalAggregateOpts::default()) + .vortex_expect("min_max"), + ) }); } diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index dd84123851a..afca2b29a18 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -239,6 +239,7 @@ fn execute_sparse_lists_inner( &mut builder, fill_elements.as_ref(), sparse_idx - next_index, + ctx, ); if patch_valid { @@ -246,7 +247,7 @@ fn execute_sparse_lists_inner( .list_elements_at(patch_idx) .vortex_expect("list_elements_at"); builder - .append_array_as_list(&patch_list) + .append_array_as_list(&patch_list, ctx) .vortex_expect("Failed to append sparse value"); } else { builder.append_null(); @@ -255,7 +256,7 @@ fn execute_sparse_lists_inner( next_index = sparse_idx + 1; } - append_list_fill(&mut builder, fill_elements.as_ref(), len - next_index); + append_list_fill(&mut builder, fill_elements.as_ref(), len - next_index, ctx); builder.finish() } @@ -328,6 +329,7 @@ fn execute_sparse_fixed_size_list_inner( &mut builder, fill_elements.as_ref(), sparse_idx - next_index, + ctx, ); // Append the patch value, handling null patches by appending defaults. @@ -336,7 +338,7 @@ fn execute_sparse_fixed_size_list_inner( .fixed_size_list_elements_at(patch_idx) .vortex_expect("fixed_size_list_elements_at"); builder - .append_array_as_list(&patch_list) + .append_array_as_list(&patch_list, ctx) .vortex_expect("Failed to append sparse fixed-size-list value"); } else { builder.append_null(); @@ -346,7 +348,12 @@ fn execute_sparse_fixed_size_list_inner( } // Fill remaining positions after last patch. - append_fixed_size_list_fill(&mut builder, fill_elements.as_ref(), array_len - next_index); + append_fixed_size_list_fill( + &mut builder, + fill_elements.as_ref(), + array_len - next_index, + ctx, + ); builder.finish_into_fixed_size_list() } @@ -367,11 +374,12 @@ 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) + .append_array_as_list(fill_elements, ctx) .vortex_expect("Failed to append sparse fill value"); } } else { @@ -383,11 +391,12 @@ 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) + .append_array_as_list(fill_elements, ctx) .vortex_expect("Failed to append sparse fixed-size-list fill value"); } } else { @@ -588,11 +597,11 @@ fn execute_varbin_inner( #[cfg(test)] mod test { use std::sync::Arc; + use std::sync::LazyLock; use rstest::rstest; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; @@ -622,15 +631,22 @@ mod test { use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::Sparse; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[rstest] #[case(Some(true))] #[case(Some(false))] #[case(None)] fn test_sparse_bool(#[case] fill_value: Option) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![0u64, 1, 7].into_array(); let values = BoolArray::from_iter([Some(true), None, Some(false)]).into_array(); let sparse_bools = Sparse::try_new(indices, values, 10, Scalar::from(fill_value)).unwrap(); @@ -653,7 +669,7 @@ mod test { fill_value, ]); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[rstest] @@ -661,7 +677,7 @@ mod test { #[case(Some(-1i32))] #[case(None)] fn test_sparse_primitive(#[case] fill_value: Option) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![0u64, 1, 7].into_array(); let values = PrimitiveArray::from_option_iter([Some(0i32), None, Some(1)]).into_array(); let sparse_ints = Sparse::try_new(indices, values, 10, Scalar::from(fill_value)).unwrap(); @@ -685,12 +701,12 @@ mod test { fill_value, ]); - assert_arrays_eq!(&flat_ints, &expected); + assert_arrays_eq!(&flat_ints, &expected, &mut ctx); } #[test] fn test_sparse_struct_valid_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let field_names = FieldNames::from_iter(["a", "b"]); let field_types = vec![ DType::Primitive(PType::I32, Nullable), @@ -760,12 +776,12 @@ mod test { .clone() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_struct_invalid_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let field_names = FieldNames::from_iter(["a", "b"]); let field_types = vec![ DType::Primitive(PType::I32, Nullable), @@ -832,12 +848,12 @@ mod test { .clone() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_decimal() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![0u32, 1u32, 7u32, 8u32].into_array(); let decimal_dtype = DecimalDType::new(3, 2); let patch_values = DecimalArray::new( @@ -850,7 +866,7 @@ mod test { let fill_scalar = Scalar::decimal(DecimalValue::I32(123), decimal_dtype, Nullable); let sparse_struct = Sparse::try_new(indices, patch_values, len, fill_scalar).unwrap(); - let expected = LEGACY_SESSION + let expected = SESSION .arrow() .execute_arrow( DecimalArray::new( @@ -865,7 +881,7 @@ mod test { ) .unwrap(); - let actual = LEGACY_SESSION + let actual = SESSION .arrow() .execute_arrow( sparse_struct @@ -885,7 +901,7 @@ mod test { #[test] fn test_sparse_utf8_varbinview_non_null_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let strings = >::from_iter([ Some("hello"), Some("goodbye"), @@ -927,12 +943,12 @@ mod test { ]) .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_utf8_varbinview_null_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let strings = >::from_iter([ Some("hello"), Some("goodbye"), @@ -974,12 +990,12 @@ mod test { ]) .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_utf8_varbinview_non_nullable() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let strings = VarBinViewArray::from_iter_str(["hello", "goodbye", "hello", "bonjour", "你好"]) .into_array(); @@ -1003,12 +1019,12 @@ mod test { ]) .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_utf8_varbin_null_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let strings = >::from_iter([ Some("hello"), Some("goodbye"), @@ -1050,12 +1066,12 @@ mod test { ]) .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_binary_varbinview_non_null_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let binaries = VarBinViewArray::from_iter_nullable_bin([ Some(b"hello" as &[u8]), Some(b"goodbye"), @@ -1097,12 +1113,12 @@ mod test { ]) .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_list_null_fill() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Use ListViewArray consistently let elements = buffer![1i32, 2, 1, 2].into_array(); // Create ListView with offsets and sizes @@ -1160,7 +1176,7 @@ mod test { #[test] fn test_sparse_list_null_fill_sliced_sparse_values() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create ListViewArray with 8 elements forming 8 single-element lists let elements = buffer![1i32, 2, 1, 2, 1, 2, 1, 2].into_array(); let offsets = buffer![0u32, 1, 2, 3, 4, 5, 6, 7].into_array(); @@ -1216,7 +1232,7 @@ mod test { #[test] fn test_sparse_list_non_null_fill() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create ListViewArray with 4 single-element lists let elements = buffer![1i32, 2, 1, 2].into_array(); let offsets = buffer![0u32, 1, 2, 3].into_array(); @@ -1288,7 +1304,7 @@ mod test { #[test] fn test_sparse_binary_varbin_null_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let strings = >::from_iter([ Some(b"hello" as &[u8]), Some(b"goodbye"), @@ -1330,7 +1346,7 @@ mod test { ]) .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] @@ -1348,7 +1364,7 @@ mod test { let sparse = Sparse::try_new(indices, fsl, 5, fill_value)?.into_array(); let actual = sparse - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut SESSION.create_execution_ctx())? .into_array(); // Expected: [1,2,3], null, [4,5,6], [7,8,9], null. @@ -1362,14 +1378,15 @@ mod test { )? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn test_sparse_fixed_size_list_non_null_fill() -> VortexResult<()> { let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); - let fsl = FixedSizeListArray::try_new(elements, 2, Validity::AllValid, 3)?.into_array(); + // Non-nullable values to match the non-nullable fill below. + let fsl = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 3)?.into_array(); let indices = buffer![0u8, 2u8, 4u8].into_array(); let fill_value = Scalar::fixed_size_list( @@ -1383,7 +1400,7 @@ mod test { let sparse = Sparse::try_new(indices, fsl, 6, fill_value)?.into_array(); let actual = sparse - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut SESSION.create_execution_ctx())? .into_array(); // Expected: [1,2], [99,88], [3,4], [99,88], [5,6], [99,88]. @@ -1391,7 +1408,7 @@ mod test { let expected = FixedSizeListArray::try_new(expected_elements, 2, Validity::NonNullable, 6)? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -1419,7 +1436,7 @@ mod test { let sparse = Sparse::try_new(indices, fsl, 6, fill_value)?.into_array(); let actual = sparse - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut SESSION.create_execution_ctx())? .into_array(); // Expected validity: [true, true, true, false, true, true]. @@ -1435,7 +1452,7 @@ mod test { )? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -1446,7 +1463,8 @@ mod test { // Create patch values: only 3 distinct lists out of 100 total positions. let elements = buffer![10i32, 11, 20, 21, 30, 31].into_array(); - let fsl = FixedSizeListArray::try_new(elements, 2, Validity::AllValid, 3)?.into_array(); + // Non-nullable values to match the non-nullable fill below. + let fsl = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 3)?.into_array(); // Patches at positions 5, 50, and 95 out of 100. let indices = buffer![5u32, 50, 95].into_array(); @@ -1464,7 +1482,7 @@ mod test { let sparse = Sparse::try_new(indices, fsl, 100, fill_value)?.into_array(); let actual = sparse - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut SESSION.create_execution_ctx())? .into_array(); // Build expected: 97 copies of [99,99] with patches at positions 5, 50, 95. @@ -1496,7 +1514,7 @@ mod test { FixedSizeListArray::try_new(expected_elements, 2, Validity::NonNullable, 100)? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -1504,7 +1522,8 @@ mod test { fn test_sparse_fixed_size_list_single_element() -> VortexResult<()> { // Test with a single element FSL array. let elements = buffer![42i32, 43].into_array(); - let fsl = FixedSizeListArray::try_new(elements, 2, Validity::AllValid, 1)?.into_array(); + // Non-nullable values to match the non-nullable fill below. + let fsl = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 1)?.into_array(); let indices = buffer![0u32].into_array(); let fill_value = Scalar::fixed_size_list( @@ -1518,7 +1537,7 @@ mod test { let sparse = Sparse::try_new(indices, fsl, 1, fill_value)?.into_array(); let actual = sparse - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut SESSION.create_execution_ctx())? .into_array(); // Expected: just [42, 43]. @@ -1526,13 +1545,13 @@ mod test { let expected = FixedSizeListArray::try_new(expected_elements, 2, Validity::NonNullable, 1)? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn test_sparse_list_grows_offset_type() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let elements = buffer![1i32, 2, 1, 2].into_array(); let offsets = buffer![0u8, 1, 2, 3, 4].into_array(); let lists = ListArray::try_new(elements, offsets, Validity::AllValid)?.into_array(); @@ -1556,19 +1575,16 @@ mod test { actual_listview.offsets().dtype(), &DType::Primitive(PType::U16, NonNullable) ); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &expected, &mut ctx); // Note that the preferred arrow list representation is `List` (not `ListView`). - let arrow_dtype = LEGACY_SESSION - .arrow() - .to_arrow_field("", expected.dtype())?; - let actual = LEGACY_SESSION + let arrow_dtype = SESSION.arrow().to_arrow_field("", expected.dtype())?; + let actual = SESSION .arrow() .execute_arrow(actual, Some(&arrow_dtype), &mut ctx)?; - let expected = - LEGACY_SESSION - .arrow() - .execute_arrow(expected, Some(&arrow_dtype), &mut ctx)?; + let expected = SESSION + .arrow() + .execute_arrow(expected, Some(&arrow_dtype), &mut ctx)?; assert_eq!(actual.data_type(), expected.data_type()); Ok(()) @@ -1576,7 +1592,7 @@ mod test { #[test] fn test_sparse_listview_null_fill_with_gaps() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // This test specifically catches the bug where the old implementation // incorrectly tracked `last_valid_offset` as the START of the last list // instead of properly handling ListView's offset/size pairs. @@ -1658,7 +1674,7 @@ mod test { #[test] fn test_sparse_listview_sliced_values_null_fill() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // This test uses sliced ListView values to ensure proper handling // of non-zero starting offsets in the source data. diff --git a/encodings/sparse/src/compute/between.rs b/encodings/sparse/src/compute/between.rs index 053a6694131..46fcd7f1e59 100644 --- a/encodings/sparse/src/compute/between.rs +++ b/encodings/sparse/src/compute/between.rs @@ -74,7 +74,6 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::between::BetweenOptions; use vortex_array::scalar_fn::fns::between::StrictComparison; - use vortex_array::session::ArraySession; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -82,7 +81,7 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); initialize(&session); session }); @@ -132,6 +131,6 @@ mod tests { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(kernel, baseline); + assert_arrays_eq!(kernel, baseline, &mut ctx); } } diff --git a/encodings/sparse/src/compute/cast.rs b/encodings/sparse/src/compute/cast.rs index 35ce9474589..21af0bc212a 100644 --- a/encodings/sparse/src/compute/cast.rs +++ b/encodings/sparse/src/compute/cast.rs @@ -48,15 +48,17 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; - use vortex_array::session::ArraySession; use vortex_buffer::buffer; use vortex_session::VortexSession; use crate::Sparse; use crate::SparseArray; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_cast_sparse_i32_to_i64() { @@ -80,7 +82,7 @@ mod tests { let expected = PrimitiveArray::from_iter([0i64, 0, 100, 0, 0, 200, 0, 0, 300, 0]); let casted_primitive = casted.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(casted_primitive, expected); + assert_arrays_eq!(casted_primitive, expected, &mut ctx); } #[test] @@ -142,7 +144,9 @@ mod tests { // zero fill, and keeps the result in the Sparse encoding. let sparse = Sparse::try_new( buffer![0u64, 1, 2, 3, 4].into_array(), - buffer![10u64, 20, 30, 40, 50].into_array(), + // Nullable values to match the null (nullable) fill. + PrimitiveArray::from_option_iter([Some(10u64), Some(20), Some(30), Some(40), Some(50)]) + .into_array(), 5, Scalar::null_native::(), )?; @@ -158,7 +162,7 @@ mod tests { let expected = PrimitiveArray::from_iter([10u64, 20, 30, 40, 50]); let casted_primitive = casted.execute::(&mut ctx)?; - assert_arrays_eq!(casted_primitive, expected); + assert_arrays_eq!(casted_primitive, expected, &mut ctx); Ok(()) } @@ -169,7 +173,8 @@ mod tests { // non-nullable, which must not panic. let sparse = Sparse::try_new( buffer![1u64, 3].into_array(), - buffer![10u64, 20].into_array(), + // Nullable values to match the null (nullable) fill. + PrimitiveArray::from_option_iter([Some(10u64), Some(20)]).into_array(), 5, Scalar::null_native::(), )?; diff --git a/encodings/sparse/src/compute/compare.rs b/encodings/sparse/src/compute/compare.rs index c64e8142298..75bb78622df 100644 --- a/encodings/sparse/src/compute/compare.rs +++ b/encodings/sparse/src/compute/compare.rs @@ -65,7 +65,6 @@ mod tests { use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; - use vortex_array::session::ArraySession; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -74,7 +73,7 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); initialize(&session); session }); @@ -112,6 +111,6 @@ mod tests { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(kernel_bool, canonical_bool); + assert_arrays_eq!(kernel_bool, canonical_bool, &mut ctx); } } diff --git a/encodings/sparse/src/compute/fill_null.rs b/encodings/sparse/src/compute/fill_null.rs index fd62926b2f3..f36032265bc 100644 --- a/encodings/sparse/src/compute/fill_null.rs +++ b/encodings/sparse/src/compute/fill_null.rs @@ -55,7 +55,6 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; - use vortex_array::session::ArraySession; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -63,7 +62,7 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); initialize(&session); session }); @@ -106,6 +105,6 @@ mod tests { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(kernel, baseline); + assert_arrays_eq!(kernel, baseline, &mut ctx); } } diff --git a/encodings/sparse/src/compute/filter.rs b/encodings/sparse/src/compute/filter.rs index f7871f124a6..91e67297fd2 100644 --- a/encodings/sparse/src/compute/filter.rs +++ b/encodings/sparse/src/compute/filter.rs @@ -34,10 +34,13 @@ impl FilterKernel for Sparse { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::fixture; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; @@ -49,9 +52,16 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::Sparse; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[fixture] fn array() -> ArrayRef { Sparse::try_new( @@ -82,7 +92,11 @@ mod tests { ) .unwrap(); - assert_arrays_eq!(filtered_array, expected); + assert_arrays_eq!( + filtered_array, + expected, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -111,7 +125,11 @@ mod tests { ) .unwrap(); - assert_arrays_eq!(filtered_array, expected); + assert_arrays_eq!( + filtered_array, + expected, + &mut SESSION.create_execution_ctx() + ); } #[test] diff --git a/encodings/sparse/src/compute/is_constant.rs b/encodings/sparse/src/compute/is_constant.rs index cfd31c6698d..eee8683df0a 100644 --- a/encodings/sparse/src/compute/is_constant.rs +++ b/encodings/sparse/src/compute/is_constant.rs @@ -73,7 +73,6 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::scalar::Scalar; - use vortex_array::session::ArraySession; use vortex_array::session::ArraySessionExt; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -85,14 +84,14 @@ mod tests { /// Session with Sparse + its pushdown kernels. static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); initialize(&session); session }); /// Baseline session: Sparse registered but no pushdown kernels. static CANONICAL_SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); session.arrays().register(Sparse); session }); diff --git a/encodings/sparse/src/compute/min_max.rs b/encodings/sparse/src/compute/min_max.rs index 22ff3e86c53..553b8eaaee0 100644 --- a/encodings/sparse/src/compute/min_max.rs +++ b/encodings/sparse/src/compute/min_max.rs @@ -7,7 +7,6 @@ use vortex_array::IntoArray; use vortex_array::aggregate_fn::Accumulator; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::DynAccumulator; -use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::fns::min_max::MinMax; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::arrays::ConstantArray; @@ -32,9 +31,9 @@ impl DynAggregateKernel for SparseMinMaxKernel { batch: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if !aggregate_fn.is::() { + let Some(options) = aggregate_fn.as_opt::() else { return Ok(None); - } + }; let Some(sparse) = batch.as_opt::() else { return Ok(None); @@ -42,7 +41,7 @@ impl DynAggregateKernel for SparseMinMaxKernel { let patches = sparse.patches(); - let mut acc = Accumulator::try_new(MinMax, EmptyOptions, batch.dtype().clone())?; + let mut acc = Accumulator::try_new(MinMax, *options, batch.dtype().clone())?; if !patches.values().is_empty() { acc.accumulate(patches.values(), ctx)?; @@ -66,10 +65,10 @@ mod tests { use rstest::rstest; 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::scalar::Scalar; - use vortex_array::session::ArraySession; use vortex_array::session::ArraySessionExt; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -79,13 +78,13 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); initialize(&session); session }); static CANONICAL_SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); session.arrays().register(Sparse); session }); @@ -101,10 +100,18 @@ mod tests { #[case(Sparse::try_new(buffer![0u64, 1, 2].into_array(), buffer![7i32, 3, 9].into_array(), 3, Scalar::from(99i32)).unwrap())] fn min_max_matches_canonical(#[case] array: SparseArray) { let arr = array.into_array(); - let kernel: Option = - min_max(&arr, &mut SESSION.create_execution_ctx()).unwrap(); - let canonical: Option = - min_max(&arr, &mut CANONICAL_SESSION.create_execution_ctx()).unwrap(); + let kernel: Option = min_max( + &arr, + &mut SESSION.create_execution_ctx(), + NumericalAggregateOpts::default(), + ) + .unwrap(); + let canonical: Option = min_max( + &arr, + &mut CANONICAL_SESSION.create_execution_ctx(), + NumericalAggregateOpts::default(), + ) + .unwrap(); assert_eq!(kernel, canonical); } } diff --git a/encodings/sparse/src/compute/mod.rs b/encodings/sparse/src/compute/mod.rs index fac33edf47e..bdae87e560d 100644 --- a/encodings/sparse/src/compute/mod.rs +++ b/encodings/sparse/src/compute/mod.rs @@ -14,15 +14,21 @@ pub(crate) mod sum; mod take; #[cfg(test)] -mod test { +mod tests { + use std::f32; + use std::sync::LazyLock; + use rstest::fixture; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; + use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::compute::conformance::mask::test_mask_conformance; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -31,8 +37,16 @@ mod test { use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::Sparse; + use crate::SparseArray; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + crate::initialize(&session); + session + }); #[fixture] fn array() -> ArrayRef { @@ -64,7 +78,11 @@ mod test { ) .unwrap(); - assert_arrays_eq!(filtered_array, expected); + assert_arrays_eq!( + filtered_array, + expected, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -93,12 +111,16 @@ mod test { ) .unwrap(); - assert_arrays_eq!(filtered_array, expected); + assert_arrays_eq!( + filtered_array, + expected, + &mut SESSION.create_execution_ctx() + ); } #[rstest] - fn test_sparse_binary_numeric(array: ArrayRef) { - test_binary_numeric_array(array) + fn test_sparse_binary_numeric_default(array: ArrayRef) { + test_binary_numeric_array(&array, &mut SESSION.create_execution_ctx()); } #[test] @@ -130,26 +152,6 @@ mod test { .into_array(), ) } -} - -#[cfg(test)] -mod tests { - use std::f32; - - use rstest::rstest; - use vortex_array::IntoArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::builtins::ArrayBuiltins; - use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; - use vortex_array::compute::conformance::consistency::test_array_consistency; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_array::scalar::Scalar; - use vortex_buffer::buffer; - - use crate::Sparse; - use crate::SparseArray; #[rstest] // Basic sparse arrays @@ -211,9 +213,8 @@ mod tests { null_fill_value ).unwrap() })] - fn test_sparse_consistency(#[case] array: SparseArray) { - test_array_consistency(&array.into_array()); + test_array_consistency(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[rstest] @@ -254,6 +255,6 @@ mod tests { Scalar::from(0i32) ).unwrap())] fn test_sparse_binary_numeric(#[case] array: SparseArray) { - test_binary_numeric_array(array.into_array()); + test_binary_numeric_array(&array.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/sparse/src/compute/nan_count.rs b/encodings/sparse/src/compute/nan_count.rs index 7fe03e21103..9a8a3a871c1 100644 --- a/encodings/sparse/src/compute/nan_count.rs +++ b/encodings/sparse/src/compute/nan_count.rs @@ -77,7 +77,6 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::nan_count::nan_count; use vortex_array::scalar::Scalar; - use vortex_array::session::ArraySession; use vortex_array::session::ArraySessionExt; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -87,13 +86,13 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); initialize(&session); session }); static CANONICAL_SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); session.arrays().register(Sparse); session }); diff --git a/encodings/sparse/src/compute/null_count.rs b/encodings/sparse/src/compute/null_count.rs index ad5e875ea6c..f63dad80205 100644 --- a/encodings/sparse/src/compute/null_count.rs +++ b/encodings/sparse/src/compute/null_count.rs @@ -67,7 +67,6 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; - use vortex_array::session::ArraySession; use vortex_array::session::ArraySessionExt; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -77,13 +76,13 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); initialize(&session); session }); static CANONICAL_SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); session.arrays().register(Sparse); session }); diff --git a/encodings/sparse/src/compute/sum.rs b/encodings/sparse/src/compute/sum.rs index 4b068aa57b4..8f85124b3b5 100644 --- a/encodings/sparse/src/compute/sum.rs +++ b/encodings/sparse/src/compute/sum.rs @@ -7,7 +7,6 @@ use vortex_array::IntoArray; use vortex_array::aggregate_fn::Accumulator; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::DynAccumulator; -use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::aggregate_fn::kernels::DynAggregateKernel; use vortex_array::arrays::ConstantArray; @@ -34,9 +33,9 @@ impl DynAggregateKernel for SparseSumKernel { batch: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if !aggregate_fn.is::() { + let Some(options) = aggregate_fn.as_opt::() else { return Ok(None); - } + }; let Some(sparse) = batch.as_opt::() else { return Ok(None); @@ -47,8 +46,8 @@ impl DynAggregateKernel for SparseSumKernel { // Build a fresh Sum accumulator over the array dtype and fold in the fill and patch // contributions. The accumulator's existing semantics (checked overflow → null - // partial) are preserved. - let mut acc = Accumulator::try_new(Sum, EmptyOptions, batch.dtype().clone())?; + // partial, NaN handling per the options) are preserved. + let mut acc = Accumulator::try_new(Sum, *options, batch.dtype().clone())?; if n_fill > 0 { let fill_array = ConstantArray::new(sparse.fill_scalar().clone(), n_fill).into_array(); @@ -72,7 +71,6 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::scalar::Scalar; - use vortex_array::session::ArraySession; use vortex_array::session::ArraySessionExt; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -83,13 +81,13 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); initialize(&session); session }); static CANONICAL_SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::empty().with::(); + let session = vortex_array::array_session(); session.arrays().register(Sparse); session }); diff --git a/encodings/sparse/src/compute/take.rs b/encodings/sparse/src/compute/take.rs index ddf9000fac1..e260eb26e9e 100644 --- a/encodings/sparse/src/compute/take.rs +++ b/encodings/sparse/src/compute/take.rs @@ -55,9 +55,12 @@ impl TakeExecute for Sparse { #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -65,10 +68,17 @@ mod test { use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_session::VortexSession; use crate::Sparse; use crate::SparseArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn test_array_fill_value() -> Scalar { // making this const is annoying Scalar::null_native::() @@ -91,7 +101,11 @@ mod test { let sparse = sparse.slice(30..40).unwrap(); let taken = sparse.take(buffer![6, 7, 8].into_array()).unwrap(); let expected = PrimitiveArray::from_option_iter([Option::::None, Some(0.47), None]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -105,7 +119,11 @@ mod test { Some(1.23), Some(3.5), ]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -113,7 +131,7 @@ mod test { let sparse = sparse_array(); let taken = sparse.take(buffer![69].into_array()).unwrap(); let expected = ConstantArray::new(test_array_fill_value(), 1).into_array(); - assert_arrays_eq!(taken, expected); + assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -123,7 +141,11 @@ mod test { let taken = sparse.take(buffer![69, 37].into_array()).unwrap(); // Index 69 is not in sparse array (fill value is null), index 37 has value 0.47 let expected = PrimitiveArray::from_option_iter([Option::::None, Some(0.47f64)]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -144,7 +166,11 @@ mod test { .unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1), Some(10), Option::::None]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -165,7 +191,11 @@ mod test { .unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1), Some(10), Option::::None]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[rstest] diff --git a/encodings/sparse/src/kernel.rs b/encodings/sparse/src/kernel.rs index 0f5d9fd51c0..30e1f9a6d6a 100644 --- a/encodings/sparse/src/kernel.rs +++ b/encodings/sparse/src/kernel.rs @@ -1,21 +1,31 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; +use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::between::Between; use vortex_array::scalar_fn::fns::between::BetweenExecuteAdaptor; +use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::fns::fill_null::FillNull; use vortex_array::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; +use vortex_session::VortexSession; use crate::Sparse; -pub(crate) static PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&BetweenExecuteAdaptor(Sparse)), - ParentKernelSet::lift(&CompareExecuteAdaptor(Sparse)), - ParentKernelSet::lift(&FillNullExecuteAdaptor(Sparse)), - ParentKernelSet::lift(&FilterExecuteAdaptor(Sparse)), - ParentKernelSet::lift(&SliceExecuteAdaptor(Sparse)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Sparse)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Between.id(), Sparse, BetweenExecuteAdaptor(Sparse)); + kernels.register_execute_parent_kernel(Binary.id(), Sparse, CompareExecuteAdaptor(Sparse)); + kernels.register_execute_parent_kernel(FillNull.id(), Sparse, FillNullExecuteAdaptor(Sparse)); + kernels.register_execute_parent_kernel(Filter.id(), Sparse, FilterExecuteAdaptor(Sparse)); + kernels.register_execute_parent_kernel(Slice.id(), Sparse, SliceExecuteAdaptor(Sparse)); + kernels.register_execute_parent_kernel(Dict.id(), Sparse, TakeExecuteAdaptor(Sparse)); +} diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index e74cedb942e..93c0bbe0363 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -7,7 +7,6 @@ use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; -use kernel::PARENT_KERNELS; use prost::Message as _; use vortex_array::AnyCanonical; use vortex_array::Array; @@ -79,11 +78,11 @@ use vortex_array::session::ArraySessionExt; /// Initialize Sparse encoding in the given session. /// -/// Registers the Sparse array vtable and its aggregate kernels (`IsConstant`, `Sum`, -/// `MinMax`, `NullCount`, `NanCount`). Compare/between/fill_null pushdown is wired -/// through `PARENT_KERNELS` (see `kernel.rs`) and does not require registration here. +/// Registers the Sparse array vtable, parent execution kernels, and aggregate kernels +/// (`IsConstant`, `Sum`, `MinMax`, `NullCount`, `NanCount`). pub fn initialize(session: &VortexSession) { session.arrays().register(Sparse); + kernel::initialize(session); let aggregate_fns = session.aggregate_fns(); aggregate_fns.register_aggregate_kernel( @@ -224,6 +223,14 @@ impl VTable for Sparse { } } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::unsupported_buffer_replacement(array, buffers) + } + fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, @@ -294,15 +301,6 @@ impl VTable for Sparse { RULES.evaluate(array, parent, child_idx) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { // Resolve offset first: wrap indices in Binary(indices, offset, Sub) and // reassemble with offset=0. Uses slot children (not data) since the executor @@ -387,6 +385,12 @@ impl Sparse { fill_value: Scalar, ) -> VortexResult { let dtype = fill_value.dtype().clone(); + vortex_ensure!( + values.dtype() == &dtype, + "sparse values dtype {} must match fill value dtype {}", + values.dtype(), + dtype, + ); let patches = Patches::new(len, 0, indices, values, None)?; let slots = SparseData::make_slots(&patches); let data = SparseData::from_patches(&patches, fill_value)?; @@ -426,25 +430,6 @@ impl Sparse { } impl SparseData { - fn normalize_patches_dtype(patches: Patches, fill_value: &Scalar) -> VortexResult { - let fill_dtype = fill_value.dtype(); - let values_dtype = patches.values().dtype(); - - vortex_ensure!( - values_dtype.eq_ignore_nullability(fill_dtype), - "fill value, {:?}, should be instance of values dtype, {} but was {}.", - fill_value, - values_dtype, - fill_dtype, - ); - - if values_dtype == fill_dtype { - Ok(patches) - } else { - patches.cast_values(fill_dtype) - } - } - pub fn validate( patches: &Patches, fill_value: &Scalar, @@ -484,16 +469,23 @@ impl SparseData { .vortex_expect("SparseArray patch slots must be present") } - /// Build a new SparseData from an existing set of patches, normalizing dtypes. + /// Build a new SparseData from an existing set of patches. pub fn try_new_from_patches(patches: Patches, fill_value: Scalar) -> VortexResult { - let patches = Self::normalize_patches_dtype(patches, &fill_value)?; - Ok(Self::from_patches_unchecked(&patches, fill_value)) + Self::from_patches(&patches, fill_value) } - /// Extract metadata from patches to create SparseData, with dtype normalization. + /// Extract metadata from patches to create SparseData. + /// + /// Patch values must already match the fill dtype; callers are expected to construct patches + /// with the correct dtype rather than relying on this to normalize them. fn from_patches(patches: &Patches, fill_value: Scalar) -> VortexResult { - let patches = Self::normalize_patches_dtype(patches.clone(), &fill_value)?; - Ok(Self::from_patches_unchecked(&patches, fill_value)) + vortex_ensure!( + patches.values().dtype() == fill_value.dtype(), + "patch values dtype {} must match fill dtype {}", + patches.values().dtype(), + fill_value.dtype(), + ); + Ok(Self::from_patches_unchecked(patches, fill_value)) } /// Extract metadata from patches to create SparseData, without validation. @@ -698,9 +690,10 @@ impl ValidityVTable for Sparse { #[cfg(test)] mod test { + use std::sync::LazyLock; + use itertools::Itertools; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; @@ -717,6 +710,12 @@ mod test { use super::*; use crate::Sparse; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + initialize(&session); + session + }); + fn nullable_fill() -> Scalar { Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)) } @@ -741,19 +740,19 @@ mod test { assert_eq!( array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut SESSION.create_execution_ctx()) .unwrap(), nullable_fill() ); assert_eq!( array - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut SESSION.create_execution_ctx()) .unwrap(), Scalar::from(Some(100_i32)) ); assert_eq!( array - .execute_scalar(5, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(5, &mut SESSION.create_execution_ctx()) .unwrap(), Scalar::from(Some(200_i32)) ); @@ -764,7 +763,7 @@ mod test { fn test_scalar_at_oob() { let array = sparse_array(nullable_fill()); array - .execute_scalar(10, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(10, &mut SESSION.create_execution_ctx()) .unwrap(); } @@ -779,19 +778,19 @@ mod test { .unwrap(); assert_eq!( - arr.execute_scalar(10, &mut LEGACY_SESSION.create_execution_ctx()) + arr.execute_scalar(10, &mut SESSION.create_execution_ctx()) .unwrap() .as_primitive() .typed_value::(), Some(1234) ); assert!( - arr.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + arr.execute_scalar(0, &mut SESSION.create_execution_ctx()) .unwrap() .is_null() ); assert!( - arr.execute_scalar(99, &mut LEGACY_SESSION.create_execution_ctx()) + arr.execute_scalar(99, &mut SESSION.create_execution_ctx()) .unwrap() .is_null() ); @@ -803,7 +802,7 @@ mod test { assert_eq!( usize::try_from( &sliced - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut SESSION.create_execution_ctx()) .unwrap() ) .unwrap(), @@ -818,7 +817,7 @@ mod test { sliced .validity() .unwrap() - .execute_mask(sliced.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(sliced.len(), &mut SESSION.create_execution_ctx()) .unwrap(), Mask::from_iter(vec![true, false, false, true, false]) ); @@ -844,7 +843,7 @@ mod test { sliced .validity() .unwrap() - .execute_mask(sliced.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(sliced.len(), &mut SESSION.create_execution_ctx()) .unwrap(), Mask::from_iter(vec![false, true, true, false, true]) ); @@ -856,7 +855,7 @@ mod test { assert_eq!( usize::try_from( &sliced_once - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut SESSION.create_execution_ctx()) .unwrap() ) .unwrap(), @@ -867,7 +866,7 @@ mod test { assert_eq!( usize::try_from( &sliced_twice - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(3, &mut SESSION.create_execution_ctx()) .unwrap() ) .unwrap(), @@ -882,7 +881,7 @@ mod test { array .validity() .unwrap() - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(array.len(), &mut SESSION.create_execution_ctx()) .unwrap() .to_bit_buffer() .iter() @@ -900,7 +899,7 @@ mod test { array .validity() .unwrap() - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(array.len(), &mut SESSION.create_execution_ctx()) .unwrap() .all_true() ); @@ -925,7 +924,7 @@ mod test { #[test] fn encode_with_nulls() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let original = PrimitiveArray::new( buffer![0i32, 1, 2, 3, 3, 3, 3, 3, 3, 3, 4, 4], Validity::from_iter(vec![ @@ -945,7 +944,7 @@ mod test { ]) ); let sparse_primitive = sparse.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(sparse_primitive, original); + assert_arrays_eq!(sparse_primitive, original, &mut ctx); } #[test] @@ -957,7 +956,7 @@ mod test { let actual = array .validity() .unwrap() - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(array.len(), &mut SESSION.create_execution_ctx()) .unwrap(); let expected = Mask::from_iter([ true, false, true, false, false, false, false, false, true, false, diff --git a/encodings/sparse/src/ops.rs b/encodings/sparse/src/ops.rs index 172114b7dfb..568d8d377d1 100644 --- a/encodings/sparse/src/ops.rs +++ b/encodings/sparse/src/ops.rs @@ -25,18 +25,26 @@ impl OperationsVTable for Sparse { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_buffer::buffer; + use vortex_session::VortexSession; use crate::Sparse; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn slice_partially_invalid() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = buffer![0u64].into_array(); let indices = buffer![0u8].into_array(); @@ -46,6 +54,6 @@ mod tests { expected[0] = 0; let values = sliced.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(values, PrimitiveArray::from_iter(expected)); + assert_arrays_eq!(values, PrimitiveArray::from_iter(expected), &mut ctx); } } diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index f03315ae575..05fb8a261e0 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -38,7 +38,6 @@ use vortex_session::registry::CachedId; use zigzag::ZigZag as ExternalZigZag; use crate::compute::ZigZagEncoded; -use crate::kernel::PARENT_KERNELS; use crate::rules::RULES; use crate::zigzag_decode; @@ -91,6 +90,14 @@ impl VTable for ZigZag { vortex_panic!("ZigZagArray buffer_name index {idx} out of bounds") } + 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, @@ -143,15 +150,6 @@ impl VTable for ZigZag { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } impl ArrayHash for ZigZagData { @@ -271,8 +269,8 @@ impl ValidityChild for ZigZag { #[cfg(test)] mod test { use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::scalar::Scalar; use vortex_buffer::buffer; @@ -282,7 +280,7 @@ mod test { #[test] fn test_compute_statistics() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = buffer![1i32, -5i32, 2, 3, 4, 5, 6, 7, 8, 9, 10] .into_array() .execute::(&mut ctx)?; diff --git a/encodings/zigzag/src/compress.rs b/encodings/zigzag/src/compress.rs index 260e2e375b5..e70650a56e1 100644 --- a/encodings/zigzag/src/compress.rs +++ b/encodings/zigzag/src/compress.rs @@ -77,25 +77,33 @@ where #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; + use vortex_session::VortexSession; use super::*; use crate::ZigZag; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[rstest] #[case::i8(PrimitiveArray::from_iter(-100_i8..100))] #[case::i16(PrimitiveArray::from_iter(-100_i16..100))] #[case::i32(PrimitiveArray::from_iter(-100_i32..100))] #[case::i64(PrimitiveArray::from_iter(-100_i64..100))] fn test_compress(#[case] input: PrimitiveArray) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let compressed = zigzag_encode(input.as_view()).unwrap().into_array(); assert!(compressed.is::()); let decompressed = compressed.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(decompressed, input); + assert_arrays_eq!(decompressed, input, &mut ctx); } } diff --git a/encodings/zigzag/src/compute/cast.rs b/encodings/zigzag/src/compute/cast.rs index a71eee7a838..c3fae84e346 100644 --- a/encodings/zigzag/src/compute/cast.rs +++ b/encodings/zigzag/src/compute/cast.rs @@ -26,8 +26,11 @@ impl CastReduce for ZigZag { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; @@ -35,10 +38,17 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_session::VortexSession; use crate::ZigZagArray; use crate::zigzag_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_cast_zigzag_i32_to_i64() { let values = PrimitiveArray::from_iter([-100i32, -1, 0, 1, 100]); @@ -61,7 +71,11 @@ mod tests { "Cast should preserve ZigZag encoding" ); - assert_arrays_eq!(casted, PrimitiveArray::from_iter([-100i64, -1, 0, 1, 100])); + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([-100i64, -1, 0, 1, 100]), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -82,7 +96,8 @@ mod tests { assert_arrays_eq!( casted, - PrimitiveArray::from_iter([100i16, -50, 0, 25, -100]) + PrimitiveArray::from_iter([100i16, -50, 0, 25, -100]), + &mut SESSION.create_execution_ctx() ); // Test i16 to i64 (widening) @@ -101,7 +116,8 @@ mod tests { assert_arrays_eq!( casted64, - PrimitiveArray::from_iter([1000i64, -500, 0, 250, -1000]) + PrimitiveArray::from_iter([1000i64, -500, 0, 250, -1000]), + &mut SESSION.create_execution_ctx() ); } diff --git a/encodings/zigzag/src/compute/mod.rs b/encodings/zigzag/src/compute/mod.rs index a90db7eacc7..1da341989c5 100644 --- a/encodings/zigzag/src/compute/mod.rs +++ b/encodings/zigzag/src/compute/mod.rs @@ -70,10 +70,11 @@ impl ZigZagEncoded for u64 { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -85,17 +86,24 @@ mod tests { use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::ZigZagArray; use crate::zigzag_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] pub fn nullable_scalar_at() -> VortexResult<()> { let zigzag = zigzag_encode( PrimitiveArray::new(buffer![-189, -160, 1], Validity::AllValid).as_view(), )?; assert_eq!( - zigzag.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())?, + zigzag.execute_scalar(1, &mut SESSION.create_execution_ctx())?, Scalar::primitive(-160, Nullability::Nullable) ); Ok(()) @@ -112,7 +120,7 @@ mod tests { let expected = zigzag_encode(PrimitiveArray::new(buffer![-189, 1], Validity::AllValid).as_view())? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -127,7 +135,7 @@ mod tests { let expected = zigzag_encode(PrimitiveArray::new(buffer![-189, 1], Validity::AllValid).as_view())? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -187,7 +195,7 @@ mod tests { fn test_take_zigzag_conformance(#[case] array: ArrayRef) -> VortexResult<()> { use vortex_array::compute::conformance::take::test_take_conformance; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array_primitive = array.execute::(&mut ctx)?; let zigzag = zigzag_encode(array_primitive.as_view())?; test_take_conformance(&zigzag.into_array()); @@ -210,7 +218,7 @@ mod tests { #[case::zigzag_large_i32(zigzag_encode(PrimitiveArray::from_iter(-500..500).as_view()).unwrap())] #[case::zigzag_large_i64(zigzag_encode(PrimitiveArray::from_iter((-1000..1000).map(|i| i as i64 * 100)).as_view()).unwrap())] fn test_zigzag_consistency(#[case] array: ZigZagArray) { - test_array_consistency(&array.into_array()); + test_array_consistency(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[rstest] @@ -220,6 +228,6 @@ mod tests { #[case::zigzag_i64_basic(zigzag_encode(PrimitiveArray::from_iter([-10000i64, -5000, 0, 5000, 10000]).as_view()).unwrap())] #[case::zigzag_i32_large(zigzag_encode(PrimitiveArray::from_iter((-50..50).map(|i| i * 10)).as_view()).unwrap())] fn test_zigzag_binary_numeric(#[case] array: ZigZagArray) { - test_binary_numeric_array(array.into_array()); + test_binary_numeric_array(&array.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/zigzag/src/kernel.rs b/encodings/zigzag/src/kernel.rs index d0096abaae1..e98e00e3982 100644 --- a/encodings/zigzag/src/kernel.rs +++ b/encodings/zigzag/src/kernel.rs @@ -1,10 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::kernel::ParentKernelSet; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_session::VortexSession; use crate::ZigZag; -pub(crate) const PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&TakeExecuteAdaptor(ZigZag))]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Dict.id(), ZigZag, TakeExecuteAdaptor(ZigZag)); +} diff --git a/encodings/zigzag/src/lib.rs b/encodings/zigzag/src/lib.rs index 89da8bd6069..01c3f51541b 100644 --- a/encodings/zigzag/src/lib.rs +++ b/encodings/zigzag/src/lib.rs @@ -3,6 +3,8 @@ pub use array::*; pub use compress::*; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; mod array; mod compress; @@ -10,3 +12,9 @@ mod compute; mod kernel; mod rules; mod slice; + +/// Initialize zigzag encoding in the given session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(ZigZag); + kernel::initialize(session); +} diff --git a/encodings/zstd/Cargo.toml b/encodings/zstd/Cargo.toml index 4383f765a2f..e5e415544a2 100644 --- a/encodings/zstd/Cargo.toml +++ b/encodings/zstd/Cargo.toml @@ -13,6 +13,9 @@ repository = { workspace = true } rust-version = { workspace = true } version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true diff --git a/encodings/zstd/benches/listview_rebuild.rs b/encodings/zstd/benches/listview_rebuild.rs index fe4e9f063d1..e3c0e6d5dd1 100644 --- a/encodings/zstd/benches/listview_rebuild.rs +++ b/encodings/zstd/benches/listview_rebuild.rs @@ -11,7 +11,6 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::listview::ListViewRebuildMode; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -19,8 +18,7 @@ use vortex_zstd::Zstd; use vortex_zstd::ZstdData; /// A shared session for the `ListView` rebuild benchmark, used to create execution contexts. -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[divan::bench(sample_size = 1000)] fn rebuild_naive(bencher: Bencher) { diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 0e3bccc4667..dc9f075e44e 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -22,7 +22,6 @@ use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; -use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; @@ -175,6 +174,33 @@ impl VTable for Zstd { } } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + let mut data = array.data().clone(); + if data.dictionary.is_some() { + let Some((dictionary, frames)) = buffers.split_first() else { + vortex_bail!("Expected dictionary buffer"); + }; + data.dictionary = Some(dictionary.clone().try_to_host_sync()?); + data.frames = frames + .iter() + .map(|buffer| buffer.clone().try_to_host_sync()) + .collect::>>()?; + } else { + data.frames = buffers + .iter() + .map(|buffer| buffer.clone().try_to_host_sync()) + .collect::>>()?; + } + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, @@ -252,9 +278,11 @@ impl VTable for Zstd { } #[derive(Clone, Debug)] +/// Zstd array encoding marker. pub struct Zstd; impl Zstd { + /// Construct a [`ZstdArray`] from validated compressed data and validity. pub fn try_new(dtype: DType, data: ZstdData, validity: Validity) -> VortexResult { let len = data.len(); data.validate(&dtype, len, &validity)?; @@ -309,6 +337,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(), @@ -325,6 +354,7 @@ pub(super) const NUM_SLOTS: usize = 1; pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; #[derive(Clone, Debug)] +/// Encoding-specific data for a [`ZstdArray`]. pub struct ZstdData { pub(crate) dictionary: Option, pub(crate) frames: Vec, @@ -344,13 +374,21 @@ impl Display for ZstdData { } } +/// Movable parts of a [`ZstdData`] value plus its validity. pub struct ZstdDataParts { + /// Optional zstd dictionary shared by all frames. pub dictionary: Option, + /// Compressed zstd frames. pub frames: Vec, + /// Serialized frame and dictionary metadata. pub metadata: ZstdMetadata, + /// Unsliced validity for the array. pub validity: Validity, + /// Unsliced row count. pub n_rows: usize, + /// Start of this logical slice in unsliced row coordinates. pub slice_start: usize, + /// End of this logical slice in unsliced row coordinates. pub slice_stop: usize, } @@ -398,17 +436,28 @@ fn collect_valid_vbv( + mask.true_count() * size_of::(), ); let mut value_byte_indices = Vec::new(); - vbv.with_iterator(|iterator| { - // by flattening, we should omit nulls - for value in iterator.flatten() { - value_byte_indices.push(buffer.len()); - // here's where we write the string lengths - buffer - .extend_trusted(ViewLen::try_from(value.len())?.to_le_bytes().into_iter()); - buffer.extend_from_slice(value); + let views = vbv.views(); + let buffers = vbv + .data_buffers() + .iter() + .map(|b| b.as_host()) + .collect::>(); + // skip nulls, writing only valid values + for (i, view) in views.iter().enumerate() { + if !mask.value(i) { + continue; } - Ok::<_, VortexError>(()) - })?; + let value = 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()] + }; + value_byte_indices.push(buffer.len()); + // here's where we write the string lengths + buffer.extend_trusted(ViewLen::try_from(value.len())?.to_le_bytes().into_iter()); + buffer.extend_from_slice(value); + } (buffer.freeze(), value_byte_indices) } }; @@ -466,6 +515,7 @@ pub fn reconstruct_views( } impl ZstdData { + /// Construct unsliced zstd data from raw frames and metadata. pub fn new( dictionary: Option, frames: Vec, @@ -482,6 +532,7 @@ impl ZstdData { } } + /// Validate dtype, slice, validity, frame, and dictionary invariants. pub fn validate(&self, dtype: &DType, len: usize, validity: &Validity) -> VortexResult<()> { vortex_ensure!( matches!( @@ -796,6 +847,9 @@ impl ZstdData { Ok(ZstdData::new(dictionary, frames, metadata, vbv.len())) } + /// Compress a supported canonical array into zstd data. + /// + /// Returns `Ok(None)` for canonical variants that this encoding does not support. pub fn from_canonical( canonical: &Canonical, level: i32, @@ -819,6 +873,11 @@ impl ZstdData { } } + /// Canonicalize and compress an array into zstd data. + /// + /// # Errors + /// + /// Returns an error if the array's canonical form is unsupported or compression fails. pub fn from_array( array: ArrayRef, level: i32, @@ -935,7 +994,6 @@ impl ZstdData { } else if dtype.is_nullable() && matches!(slice_validity, Validity::NonNullable) { slice_validity = Validity::AllValid; } - // // END OF IMPORTANT BLOCK // @@ -950,6 +1008,7 @@ impl ZstdData { dtype.as_ptype(), slice_validity, slice_n_rows, + ctx, ); Ok(primitive.into_array()) @@ -1020,6 +1079,7 @@ impl ZstdData { self.slice_stop == self.slice_start } + /// Split this data into movable parts, attaching the supplied validity. pub fn into_parts(self, validity: Validity) -> ZstdDataParts { ZstdDataParts { dictionary: self.dictionary, diff --git a/encodings/zstd/src/compute/cast.rs b/encodings/zstd/src/compute/cast.rs index 6cbe4f37795..c5f15eba284 100644 --- a/encodings/zstd/src/compute/cast.rs +++ b/encodings/zstd/src/compute/cast.rs @@ -83,15 +83,13 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_session::VortexSession; use crate::Zstd; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_cast_zstd_i32_to_i64() { @@ -109,7 +107,11 @@ mod tests { ); let decoded = casted.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(decoded, PrimitiveArray::from_iter([1i64, 2, 3, 4, 5])); + assert_arrays_eq!( + decoded, + PrimitiveArray::from_iter([1i64, 2, 3, 4, 5]), + &mut ctx + ); } #[test] @@ -146,7 +148,11 @@ mod tests { ); // Verify the values are correct let decoded = casted.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(decoded, PrimitiveArray::from_iter([20u32, 30, 40, 50])); + assert_arrays_eq!( + decoded, + PrimitiveArray::from_iter([20u32, 30, 40, 50]), + &mut ctx + ); } #[test] @@ -171,7 +177,7 @@ mod tests { ); let decoded = casted.execute::(&mut ctx).unwrap(); let expected = PrimitiveArray::from_iter([20u32, 30, 40, 50]); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); } #[rstest] diff --git a/encodings/zstd/src/compute/mod.rs b/encodings/zstd/src/compute/mod.rs index de175eb3154..75ef5273ca5 100644 --- a/encodings/zstd/src/compute/mod.rs +++ b/encodings/zstd/src/compute/mod.rs @@ -7,8 +7,8 @@ mod cast; mod tests { use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_buffer::buffer; @@ -18,23 +18,23 @@ mod tests { fn zstd_i32() -> ZstdArray { let values = PrimitiveArray::from_iter([100i32, 200, 300, 400, 500]); - Zstd::from_primitive(&values, 0, 0, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() + Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() } fn zstd_f64() -> ZstdArray { let values = PrimitiveArray::from_iter([1.1f64, 2.2, 3.3, 4.4, 5.5]); - Zstd::from_primitive(&values, 0, 0, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() + Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() } fn zstd_u32() -> ZstdArray { let values = PrimitiveArray::from_iter([10u32, 20, 30, 40, 50]); - Zstd::from_primitive(&values, 0, 0, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() + Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() } fn zstd_nullable_i64() -> ZstdArray { let values = PrimitiveArray::from_option_iter([Some(1000i64), None, Some(3000), Some(4000), None]); - Zstd::from_primitive(&values, 0, 0, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() + Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() } fn zstd_single() -> ZstdArray { @@ -42,7 +42,7 @@ mod tests { buffer![42i64], vortex_array::validity::Validity::NonNullable, ); - Zstd::from_primitive(&values, 0, 0, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() + Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() } fn zstd_large() -> ZstdArray { @@ -50,7 +50,7 @@ mod tests { buffer![0u32..1000], vortex_array::validity::Validity::NonNullable, ); - Zstd::from_primitive(&values, 3, 0, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() + Zstd::from_primitive(&values, 3, 0, &mut array_session().create_execution_ctx()).unwrap() } fn zstd_all_same() -> ZstdArray { @@ -58,12 +58,12 @@ mod tests { buffer![42i32; 100], vortex_array::validity::Validity::NonNullable, ); - Zstd::from_primitive(&values, 0, 0, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() + Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() } fn zstd_negative() -> ZstdArray { let values = PrimitiveArray::from_iter([-100i32, -50, 0, 50, 100]); - Zstd::from_primitive(&values, 0, 0, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() + Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() } #[rstest] @@ -76,6 +76,9 @@ mod tests { #[case::all_same(zstd_all_same())] #[case::negative(zstd_negative())] fn test_zstd_consistency(#[case] array: ZstdArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/encodings/zstd/src/lib.rs b/encodings/zstd/src/lib.rs index a927bbfdf26..a178fdec142 100644 --- a/encodings/zstd/src/lib.rs +++ b/encodings/zstd/src/lib.rs @@ -1,7 +1,28 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Zstd-backed compression encodings for variable-width Vortex arrays. +//! +//! [`ZstdArray`] stores UTF-8 or binary values as one or more zstd frames, optionally sharing a +//! 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. +//! +//! 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 +//! array session: +//! +//! ```rust +//! use vortex_array::session::ArraySessionExt; +//! +//! let session = vortex_array::array_session(); +//! session.arrays().register(vortex_zstd::Zstd); +//! ``` + pub use array::*; +use vortex_array::dtype::proto::dtype as pb; #[cfg(feature = "unstable_encodings")] pub use zstd_buffers::*; @@ -16,31 +37,49 @@ mod zstd_buffers; mod test; #[derive(Clone, prost::Message)] +/// Metadata for one zstd frame. pub struct ZstdFrameMetadata { + /// Uncompressed byte size of this frame. #[prost(uint64, tag = "1")] pub uncompressed_size: u64, + /// Number of valid values stored in this frame. #[prost(uint64, tag = "2")] pub n_values: u64, } #[derive(Clone, prost::Message)] +/// Serialized metadata for a [`ZstdArray`]. pub struct ZstdMetadata { // optional, will be 0 if there's no dictionary + /// Dictionary size in bytes, or `0` when no dictionary is present. #[prost(uint32, tag = "1")] pub dictionary_size: u32, + /// Metadata for each compressed frame. #[prost(message, repeated, tag = "2")] pub frames: Vec, } #[derive(Clone, prost::Message)] +/// Serialized metadata for the unstable `ZstdBuffers` encoding. pub struct ZstdBuffersMetadata { + /// Encoding id of the inner array whose buffers were compressed. #[prost(string, tag = "1")] pub inner_encoding_id: String, + /// Serialized metadata of the inner array. #[prost(bytes = "vec", tag = "2")] pub inner_metadata: Vec, + /// Uncompressed byte size of each compressed buffer. #[prost(uint64, repeated, tag = "3")] pub uncompressed_sizes: Vec, /// Alignment of each buffer in bytes (must be a power of two). #[prost(uint32, repeated, tag = "4")] pub buffer_alignments: Vec, + /// DType of child arrays. Children belong to inner encodings, and their + /// dtypes don't persist after serialization, so we need to retrieve them + /// from metadata. + #[prost(message, repeated, tag = "5")] + pub child_dtypes: Vec, + /// Length of each child array, ordered as "child_dtypes" + #[prost(uint64, repeated, tag = "6")] + pub child_lens: Vec, } diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index 094c5ef60cc..cf45eb97c04 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -3,8 +3,8 @@ #![expect(clippy::cast_possible_truncation)] use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; @@ -21,7 +21,7 @@ use crate::Zstd; #[test] fn test_zstd_compress_decompress() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let data: Vec = (0..200).collect(); let array = PrimitiveArray::from_iter(data.clone()); @@ -32,22 +32,30 @@ fn test_zstd_compress_decompress() { // check full decompression works let decompressed = Zstd::decompress(&compressed, &mut ctx).unwrap(); - assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data)); + assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data), &mut ctx); // check slicing works let slice = compressed.slice(100..105).unwrap(); for i in 0_i32..5 { - assert_nth_scalar!(slice, i as usize, 100 + i); + assert_nth_scalar!(slice, i as usize, 100 + i, &mut ctx); } - assert_arrays_eq!(slice, PrimitiveArray::from_iter([100, 101, 102, 103, 104])); + assert_arrays_eq!( + slice, + PrimitiveArray::from_iter([100, 101, 102, 103, 104]), + &mut ctx + ); let slice = compressed.slice(200..200).unwrap(); - assert_arrays_eq!(slice, PrimitiveArray::from_iter(Vec::::new())); + assert_arrays_eq!( + slice, + PrimitiveArray::from_iter(Vec::::new()), + &mut ctx + ); } #[test] fn test_zstd_empty() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let data: Vec = vec![]; let array = PrimitiveArray::new( data.iter().cloned().collect::>(), @@ -56,12 +64,12 @@ fn test_zstd_empty() { let compressed = Zstd::from_primitive(&array, 3, 100, &mut ctx).unwrap(); - assert_arrays_eq!(compressed, PrimitiveArray::from_iter(data)); + assert_arrays_eq!(compressed, PrimitiveArray::from_iter(data), &mut ctx); } #[test] fn test_zstd_with_validity_and_multi_frame() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let data: Vec = (0..200).collect(); let mut validity: Vec = vec![false; 200]; validity[3] = true; @@ -73,10 +81,10 @@ fn test_zstd_with_validity_and_multi_frame() { let compressed = Zstd::from_primitive(&array, 0, 30, &mut ctx).unwrap(); assert!(compressed.dictionary.is_none()); - assert_nth_scalar!(compressed, 0, None::); - assert_nth_scalar!(compressed, 3, 3); - assert_nth_scalar!(compressed, 10, None::); - assert_nth_scalar!(compressed, 177, 177); + assert_nth_scalar!(compressed, 0, None::, &mut ctx); + assert_nth_scalar!(compressed, 3, 3, &mut ctx); + assert_nth_scalar!(compressed, 10, None::, &mut ctx); + assert_nth_scalar!(compressed, 177, 177, &mut ctx); let decompressed = Zstd::decompress(&compressed, &mut ctx) .unwrap() @@ -115,7 +123,7 @@ fn test_zstd_with_validity_and_multi_frame() { #[test] fn test_zstd_with_dict() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let data: Vec = (0..200).collect(); let array = PrimitiveArray::new( data.iter().cloned().collect::>(), @@ -124,24 +132,28 @@ fn test_zstd_with_dict() { let compressed = Zstd::from_primitive(&array, 0, 16, &mut ctx).unwrap(); assert!(compressed.dictionary.is_some()); - assert_nth_scalar!(compressed, 0, 0); - assert_nth_scalar!(compressed, 199, 199); + assert_nth_scalar!(compressed, 0, 0, &mut ctx); + assert_nth_scalar!(compressed, 199, 199, &mut ctx); let decompressed = Zstd::decompress(&compressed, &mut ctx) .unwrap() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data)); + assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data), &mut ctx); // check slicing works let slice = compressed.slice(176..179).unwrap(); let primitive = slice.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(primitive, PrimitiveArray::from_iter([176, 177, 178])); + assert_arrays_eq!( + primitive, + PrimitiveArray::from_iter([176, 177, 178]), + &mut ctx + ); } #[test] fn test_validity_vtable() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mask_bools = vec![false, true, true, false, true]; let array = PrimitiveArray::new( (0..5).collect::>(), @@ -152,7 +164,7 @@ fn test_validity_vtable() { assert_eq!( arr.validity() .unwrap() - .execute_mask(arr.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(arr.len(), &mut array_session().create_execution_ctx()) .unwrap(), Mask::from_iter(mask_bools) ); @@ -161,7 +173,7 @@ fn test_validity_vtable() { sliced .validity() .unwrap() - .execute_mask(sliced.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(sliced.len(), &mut array_session().create_execution_ctx()) .unwrap(), Mask::from_iter(vec![true, true, false]) ); @@ -169,7 +181,7 @@ fn test_validity_vtable() { #[test] fn test_zstd_var_bin_view() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let data: [Option<&'static [u8]>; 5] = [ Some(b"foo"), Some(b"bar"), @@ -181,21 +193,21 @@ fn test_zstd_var_bin_view() { let compressed = Zstd::from_var_bin_view(&array, 0, 3, &mut ctx).unwrap(); assert!(compressed.dictionary.is_none()); - assert_nth_scalar!(compressed, 0, "foo"); - assert_nth_scalar!(compressed, 1, "bar"); - assert_nth_scalar!(compressed, 2, None::); - assert_nth_scalar!(compressed, 3, "Lorem ipsum dolor sit amet"); - assert_nth_scalar!(compressed, 4, "baz"); + assert_nth_scalar!(compressed, 0, "foo", &mut ctx); + assert_nth_scalar!(compressed, 1, "bar", &mut ctx); + assert_nth_scalar!(compressed, 2, None::, &mut ctx); + assert_nth_scalar!(compressed, 3, "Lorem ipsum dolor sit amet", &mut ctx); + assert_nth_scalar!(compressed, 4, "baz", &mut ctx); let sliced = compressed.slice(1..4).unwrap(); - assert_nth_scalar!(sliced, 0, "bar"); - assert_nth_scalar!(sliced, 1, None::); - assert_nth_scalar!(sliced, 2, "Lorem ipsum dolor sit amet"); + assert_nth_scalar!(sliced, 0, "bar", &mut ctx); + assert_nth_scalar!(sliced, 1, None::, &mut ctx); + assert_nth_scalar!(sliced, 2, "Lorem ipsum dolor sit amet", &mut ctx); } #[test] fn test_zstd_decompress_var_bin_view() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let data: [Option<&'static [u8]>; 5] = [ Some(b"foo"), Some(b"bar"), @@ -207,26 +219,26 @@ fn test_zstd_decompress_var_bin_view() { let compressed = Zstd::from_var_bin_view(&array, 0, 3, &mut ctx).unwrap(); assert!(compressed.dictionary.is_none()); - assert_nth_scalar!(compressed, 0, "foo"); - assert_nth_scalar!(compressed, 1, "bar"); - assert_nth_scalar!(compressed, 2, None::); - assert_nth_scalar!(compressed, 3, "Lorem ipsum dolor sit amet"); - assert_nth_scalar!(compressed, 4, "baz"); + assert_nth_scalar!(compressed, 0, "foo", &mut ctx); + assert_nth_scalar!(compressed, 1, "bar", &mut ctx); + assert_nth_scalar!(compressed, 2, None::, &mut ctx); + assert_nth_scalar!(compressed, 3, "Lorem ipsum dolor sit amet", &mut ctx); + assert_nth_scalar!(compressed, 4, "baz", &mut ctx); let decompressed = Zstd::decompress(&compressed, &mut ctx) .unwrap() .execute::(&mut ctx) .unwrap(); - assert_nth_scalar!(decompressed, 0, "foo"); - assert_nth_scalar!(decompressed, 1, "bar"); - assert_nth_scalar!(decompressed, 2, None::); - assert_nth_scalar!(decompressed, 3, "Lorem ipsum dolor sit amet"); - assert_nth_scalar!(decompressed, 4, "baz"); + assert_nth_scalar!(decompressed, 0, "foo", &mut ctx); + assert_nth_scalar!(decompressed, 1, "bar", &mut ctx); + assert_nth_scalar!(decompressed, 2, None::, &mut ctx); + assert_nth_scalar!(decompressed, 3, "Lorem ipsum dolor sit amet", &mut ctx); + assert_nth_scalar!(decompressed, 4, "baz", &mut ctx); } #[test] fn test_sliced_array_children() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let data: Vec> = (0..10).map(|v| (v != 5).then_some(v)).collect(); let compressed = Zstd::from_primitive(&PrimitiveArray::from_option_iter(data), 0, 100, &mut ctx).unwrap(); @@ -238,7 +250,7 @@ fn test_sliced_array_children() { /// the buffer alignment when compressing primitive arrays. #[test] fn test_zstd_frame_start_buffer_alignment() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let data = vec![0u8; 2]; let aligned_buffer = Buffer::copy_from_aligned(&data, Alignment::new(8)); // u8 array now has a 8-byte alignment. diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index dd689fd76bd..f40a451caaf 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -43,9 +43,11 @@ use crate::ZstdBuffersMetadata; pub type ZstdBuffersArray = Array; #[derive(Clone, Debug)] +/// Encoding marker for buffer-level zstd compression. pub struct ZstdBuffers; impl ZstdBuffers { + /// Construct a [`ZstdBuffersArray`] from compressed buffer data. pub fn try_new( dtype: DType, len: usize, @@ -54,6 +56,10 @@ impl ZstdBuffers { Array::try_from_parts(ArrayParts::new(ZstdBuffers, dtype, len, data)) } + /// 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. pub fn compress( array: &ArrayRef, level: i32, @@ -96,6 +102,7 @@ impl ZstdBuffers { Ok(compressed) } + /// Rebuild the wrapped array from decompressed buffer handles. pub fn build_inner( array: &ZstdBuffersArray, buffer_handles: &[BufferHandle], @@ -149,6 +156,7 @@ impl Display for ZstdBuffersData { } #[derive(Clone, Debug)] +/// Decode plan for buffer-level zstd decompression. pub struct ZstdBuffersDecodePlan { compressed_buffers: Vec, frame_sizes: Arc<[usize]>, @@ -160,30 +168,37 @@ pub struct ZstdBuffersDecodePlan { } impl ZstdBuffersDecodePlan { + /// Compressed buffers to decode. pub fn compressed_buffers(&self) -> &[BufferHandle] { &self.compressed_buffers } + /// Compressed frame sizes in bytes. pub fn frame_sizes(&self) -> Arc<[usize]> { Arc::clone(&self.frame_sizes) } + /// Decompressed output size for each buffer. pub fn output_sizes(&self) -> Arc<[usize]> { Arc::clone(&self.output_sizes) } + /// Byte offsets of each decompressed buffer in one contiguous output allocation. pub fn output_offsets(&self) -> &[usize] { &self.output_offsets } + /// Total byte size of the planned contiguous output allocation. pub fn output_size_total(&self) -> usize { self.output_size_total } + /// Largest single decompressed buffer size. pub fn output_size_max(&self) -> usize { self.output_size_max } + /// Number of compressed frames/buffers in the plan. pub fn num_frames(&self) -> usize { self.compressed_buffers.len() } @@ -273,6 +288,7 @@ impl ZstdBuffersData { Ok(result) } + /// Build a decode plan for external or device decompression. pub fn decode_plan(&self) -> VortexResult { // If invariants are somehow broken, device decompression could have UB, so ensure // they still hold. @@ -332,6 +348,7 @@ fn compute_output_layout( (offsets, total_size) } +#[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] fn array_id_from_string(s: &str) -> ArrayId { ArrayId::new(s) } @@ -395,6 +412,19 @@ impl VTable for ZstdBuffers { Some(format!("compressed_{idx}")) } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + let mut data = array.data().clone(); + data.compressed_buffers = buffers.to_vec(); + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { format!("child_{idx}") } @@ -403,12 +433,21 @@ impl VTable for ZstdBuffers { array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { + let children: Vec<&ArrayRef> = array.slots().iter().flatten().collect(); + let child_dtypes = children + .iter() + .map(|child| child.dtype().try_into()) + .collect::>>()?; + let child_lens = children.iter().map(|child| child.len() as u64).collect(); + Ok(Some( ZstdBuffersMetadata { inner_encoding_id: array.inner_encoding_id.to_string(), inner_metadata: array.inner_metadata.clone(), uncompressed_sizes: array.uncompressed_sizes.clone(), buffer_alignments: array.buffer_alignments.clone(), + child_dtypes, + child_lens, } .encode_to_vec(), )) @@ -421,13 +460,23 @@ impl VTable for ZstdBuffers { metadata: &[u8], buffers: &[BufferHandle], children: &dyn ArrayChildren, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult> { let metadata = ZstdBuffersMetadata::decode(metadata)?; let compressed_buffers: Vec = buffers.to_vec(); + // Children belong to inner encodings, and serialization doesn't + // preserve their dtypes and values. Check dtypes are recovered from + // metadata. + vortex_ensure_eq!(metadata.child_dtypes.len(), children.len()); + vortex_ensure_eq!(metadata.child_lens.len(), children.len()); + let slots: ArraySlots = (0..children.len()) - .map(|i| children.get(i, dtype, len).map(Some)) + .map(|i| { + let child_dtype = DType::from_proto(&metadata.child_dtypes[i], session)?; + let child_len = usize::try_from(metadata.child_lens[i])?; + children.get(i, &child_dtype, child_len).map(Some) + }) .collect::>>()? .into(); @@ -490,17 +539,23 @@ impl ValidityVTable for ZstdBuffers { #[cfg(test)] mod tests { use rstest::rstest; + use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; use vortex_array::expr::stats::StatsProvider; + use vortex_array::serde::SerializeOptions; + use vortex_array::serde::SerializedArray; + use vortex_array::session::ArraySessionExt; + use vortex_buffer::ByteBufferMut; use vortex_error::VortexResult; + use vortex_session::registry::ReadContext; use super::*; @@ -544,15 +599,47 @@ mod tests { #[case::empty_primitive(make_empty_primitive_array())] #[case::inlined_varbinview(make_inlined_varbinview_array())] fn test_roundtrip(#[case] input: ArrayRef) -> VortexResult<()> { - let compressed = ZstdBuffers::compress(&input, 3, &LEGACY_SESSION)?; + let compressed = ZstdBuffers::compress(&input, 3, &array_session())?; assert_eq!(compressed.len(), input.len()); assert_eq!(compressed.dtype(), input.dtype()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let decompressed = compressed.into_array().execute::(&mut ctx)?; - assert_arrays_eq!(input, decompressed); + assert_arrays_eq!(input, decompressed, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::primitive(make_primitive_array())] + #[case::varbinview(make_varbinview_array())] + #[case::nullable_primitive(make_nullable_primitive_array())] + #[case::nullable_varbinview(make_nullable_varbinview_array())] + #[case::empty_primitive(make_empty_primitive_array())] + #[case::inlined_varbinview(make_inlined_varbinview_array())] + fn test_serde_roundtrip(#[case] input: ArrayRef) -> VortexResult<()> { + let session = array_session(); + session.arrays().register(ZstdBuffers); + + let compressed = ZstdBuffers::compress(&input, 3, &session)?.into_array(); + let dtype = compressed.dtype().clone(); + let len = compressed.len(); + + let array_ctx = ArrayContext::empty(); + let serialized = + compressed.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, len, &ReadContext::new(array_ctx.to_ids()), &session)?; + + let mut ctx = session.create_execution_ctx(); + let decoded = decoded.execute::(&mut ctx)?; + assert_arrays_eq!(input, decoded, &mut ctx); Ok(()) } @@ -561,7 +648,7 @@ mod tests { let input = make_primitive_array(); input.statistics().set(Stat::Min, Precision::exact(0i32)); - let compressed = ZstdBuffers::compress(&input, 3, &LEGACY_SESSION)?; + let compressed = ZstdBuffers::compress(&input, 3, &array_session())?; assert!(!compressed.statistics().get(Stat::Min).is_absent()); Ok(()) @@ -570,9 +657,9 @@ mod tests { #[test] fn test_validity_delegates_for_nullable_input() -> VortexResult<()> { let input = make_nullable_primitive_array(); - let compressed = ZstdBuffers::compress(&input, 3, &LEGACY_SESSION)?.into_array(); + let compressed = ZstdBuffers::compress(&input, 3, &array_session())?.into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(compressed.all_valid(&mut ctx)?, input.all_valid(&mut ctx)?); assert_eq!( compressed.all_invalid(&mut ctx)?, diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index e2d05b706f9..58c728b23a5 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -36,7 +36,7 @@ vortex-array = { workspace = true, features = ["arbitrary", "_test-harness"] } vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } -vortex-fsst = { workspace = true } +vortex-fsst = { workspace = true, features = ["_test-harness"] } vortex-io = { workspace = true } vortex-mask = { workspace = true } vortex-runend = { workspace = true, features = ["arbitrary"] } diff --git a/fuzz/fuzz_targets/compress_gpu.rs b/fuzz/fuzz_targets/compress_gpu.rs index f88df52d2d5..91d7e9b9f9d 100644 --- a/fuzz/fuzz_targets/compress_gpu.rs +++ b/fuzz/fuzz_targets/compress_gpu.rs @@ -4,20 +4,21 @@ #![no_main] #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use libfuzzer_sys::Corpus; use libfuzzer_sys::fuzz_target; +use tokio::runtime::Builder; +use tokio::runtime::Runtime; use vortex_error::vortex_panic; use vortex_fuzz::FuzzCompressGpu; use vortex_fuzz::run_compress_gpu; -fuzz_target!(|fuzz: FuzzCompressGpu| -> Corpus { - // Use tokio runtime to run async GPU fuzzer - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); +static RUNTIME: LazyLock = + LazyLock::new(|| Builder::new_current_thread().enable_all().build().unwrap()); - match rt.block_on(run_compress_gpu(fuzz)) { +fuzz_target!(|fuzz: FuzzCompressGpu| -> Corpus { + match RUNTIME.block_on(run_compress_gpu(fuzz)) { Ok(true) => Corpus::Keep, Ok(false) => Corpus::Reject, Err(e) => { diff --git a/fuzz/src/array/compare.rs b/fuzz/src/array/compare.rs index 53988426822..f594a9b4102 100644 --- a/fuzz/src/array/compare.rs +++ b/fuzz/src/array/compare.rs @@ -4,7 +4,6 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; @@ -20,6 +19,7 @@ use vortex_array::scalar_fn::fns::binary::scalar_cmp; use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; +use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::vortex_panic; @@ -134,32 +134,45 @@ pub fn compare_canonical_array( .clone() .execute::(ctx) .vortex_expect("to varbinview"); - varbinview.with_iterator(|iter| { - let utf8_value = value.as_utf8(); - compare_to( - iter.map(|v| v.map(|b| unsafe { str::from_utf8_unchecked(b) })), - utf8_value.value().vortex_expect("nulls handled before"), - operator, - result_nullability, - ) - }) + let mask = varbinview + .validity() + .vortex_expect("validity") + .execute_mask(varbinview.len(), ctx) + .vortex_expect("varbinview mask"); + let values: Vec> = (0..varbinview.len()) + .map(|i| mask.value(i).then(|| varbinview.bytes_at(i))) + .collect(); + let utf8_value = value.as_utf8(); + compare_to( + values.iter().map(|v| { + v.as_ref() + .map(|b| unsafe { str::from_utf8_unchecked(b.as_slice()) }) + }), + utf8_value.value().vortex_expect("nulls handled before"), + operator, + result_nullability, + ) } DType::Binary(_) => { let varbinview = array .clone() .execute::(ctx) .vortex_expect("to varbinview"); - varbinview.with_iterator(|iter| { - let binary_value = value.as_binary(); - compare_to( - // Don't understand the lifetime problem here but identity map makes it go away - #[expect(clippy::map_identity)] - iter.map(|v| v), - binary_value.value().vortex_expect("nulls handled before"), - operator, - result_nullability, - ) - }) + let mask = varbinview + .validity() + .vortex_expect("validity") + .execute_mask(varbinview.len(), ctx) + .vortex_expect("varbinview mask"); + let values: Vec> = (0..varbinview.len()) + .map(|i| mask.value(i).then(|| varbinview.bytes_at(i))) + .collect(); + let binary_value = value.as_binary(); + compare_to( + values.iter().map(|v| v.as_deref()), + binary_value.value().vortex_expect("nulls handled before"), + operator, + result_nullability, + ) } DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) => { let scalar_vals: Vec = (0..array.len()) diff --git a/fuzz/src/array/fill_null.rs b/fuzz/src/array/fill_null.rs index d19a9f2b360..f3a49b25637 100644 --- a/fuzz/src/array/fill_null.rs +++ b/fuzz/src/array/fill_null.rs @@ -300,8 +300,8 @@ fn fill_varbinview_array( mod tests { use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; @@ -322,25 +322,25 @@ mod tests { fn canonical(array: impl IntoArray) -> Canonical { array .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut array_session().create_execution_ctx()) .unwrap() } #[test] fn test_fill_null_primitive() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]); let fill_value = Scalar::from(42i32); let result = fill_null_canonical_array(canonical(array), &fill_value, &mut ctx).unwrap(); let expected = PrimitiveArray::from_iter([1i32, 42, 3, 42, 5]); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] fn test_fill_null_bool() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let data_buffer = BitBuffer::from(vec![true, false, false, false]); let validity_buffer = BitBuffer::from(vec![true, false, true, false]); let array = BoolArray::new(data_buffer, Validity::from(validity_buffer)); @@ -349,12 +349,12 @@ mod tests { let result = fill_null_canonical_array(canonical(array), &fill_value, &mut ctx).unwrap(); let expected = BoolArray::from(BitBuffer::from(vec![true, true, false, true])); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] fn test_fill_null_string() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = VarBinViewArray::from_iter( [Some("hello"), None, Some("world")].iter().copied(), DType::Utf8(Nullability::Nullable), @@ -364,37 +364,37 @@ mod tests { let result = fill_null_canonical_array(canonical(array), &fill_value, &mut ctx).unwrap(); let expected = VarBinViewArray::from_iter_str(["hello", "default", "world"]); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] fn test_fill_null_all_invalid() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([None::, None, None]); let fill_value = Scalar::from(100i32); let result = fill_null_canonical_array(canonical(array), &fill_value, &mut ctx).unwrap(); let expected = PrimitiveArray::from_iter([100i32, 100, 100]); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] fn test_fill_null_no_nulls() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter([1i32, 2, 3]); let fill_value = Scalar::from(42i32); let result = fill_null_canonical_array(canonical(array), &fill_value, &mut ctx).unwrap(); let expected = PrimitiveArray::from_iter([1i32, 2, 3]); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] #[should_panic] fn test_fill_null_with_null_value_errors() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]); let fill_value = Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)); @@ -411,7 +411,7 @@ mod tests { #[test] fn test_fill_null_decimal_i32() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::from_option_iter( [Some(100i32), None, Some(300i32), None, Some(500i32)], DecimalDType::new(10, 2), @@ -428,12 +428,12 @@ mod tests { [100i32, 999i32, 300i32, 999i32, 500i32], DecimalDType::new(10, 2), ); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] fn test_fill_null_decimal_i64() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::from_option_iter( [Some(1000i64), None, Some(3000i64)], DecimalDType::new(15, 3), @@ -448,12 +448,12 @@ mod tests { let expected = DecimalArray::from_iter([1000i64, 9999i64, 3000i64], DecimalDType::new(15, 3)); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] fn test_fill_null_decimal_i128() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::from_option_iter( [Some(10000i128), None, Some(30000i128), None], DecimalDType::new(20, 4), @@ -470,12 +470,12 @@ mod tests { [10000i128, 99999i128, 30000i128, 99999i128], DecimalDType::new(20, 4), ); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] fn test_fill_null_decimal_all_invalid() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::from_option_iter([None::, None, None], DecimalDType::new(10, 2)); let fill_value = Scalar::decimal( @@ -493,12 +493,12 @@ mod tests { .into_array() .cast(result.dtype().clone()) .unwrap(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] fn test_fill_null_decimal_no_nulls() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::from_option_iter( [Some(100i32), Some(200i32), Some(300i32)], DecimalDType::new(10, 2), @@ -518,6 +518,6 @@ mod tests { .into_array() .cast(result.dtype().clone()) .unwrap(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } } diff --git a/fuzz/src/array/filter.rs b/fuzz/src/array/filter.rs index 8de1786d25f..aa54b0d9488 100644 --- a/fuzz/src/array/filter.rs +++ b/fuzz/src/array/filter.rs @@ -4,7 +4,6 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; @@ -90,12 +89,12 @@ pub fn filter_canonical_array( } DType::Utf8(_) | DType::Binary(_) => { let utf8 = array.clone().execute::(ctx)?; - let values = utf8.with_iterator(|iter| { - iter.zip(filter.iter()) - .filter(|(_, f)| **f) - .map(|(v, _)| v.map(|u| u.to_vec())) - .collect::>() - }); + let mask = utf8.validity()?.execute_mask(utf8.len(), ctx)?; + let values = (0..utf8.len()) + .zip(filter.iter()) + .filter(|(_, f)| **f) + .map(|(i, _)| mask.value(i).then(|| utf8.bytes_at(i).to_vec())) + .collect::>(); Ok(VarBinViewArray::from_iter(values, array.dtype().clone()).into_array()) } DType::List(..) | DType::FixedSizeList(..) => { diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index 49fc2ad40f9..f3cecec2b10 100644 --- a/fuzz/src/array/mask.rs +++ b/fuzz/src/array/mask.rs @@ -156,9 +156,10 @@ pub fn mask_canonical_array( #[cfg(test)] mod tests { use vortex_array::Canonical; + use vortex_array::ExecutionCtx; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::FixedSizeListArray; @@ -175,20 +176,17 @@ mod tests { use super::mask_canonical_array; - fn canonical(array: impl IntoArray) -> Canonical { - array - .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + fn canonical(array: impl IntoArray, ctx: &mut ExecutionCtx) -> Canonical { + array.into_array().execute::(ctx).unwrap() } #[test] fn test_mask_null_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = NullArray::new(5); let mask = Mask::from_iter([true, false, true, false, true]); - let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); + let result = mask_canonical_array(canonical(array, &mut ctx), &mask, &mut ctx).unwrap(); assert_eq!(result.len(), 5); // All values should still be null @@ -199,43 +197,43 @@ mod tests { #[test] fn test_mask_bool_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = BoolArray::from_iter([true, false, true, false, true]); let mask = Mask::from_iter([false, true, true, false, true]); - let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); + let result = mask_canonical_array(canonical(array, &mut ctx), &mask, &mut ctx).unwrap(); let expected = BoolArray::from_iter([None, Some(false), Some(true), None, Some(true)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_mask_primitive_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]); let mask = Mask::from_iter([true, false, true, false, true]); - let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); + let result = mask_canonical_array(canonical(array, &mut ctx), &mask, &mut ctx).unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_mask_primitive_array_with_nulls() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None]); let mask = Mask::from_iter([false, true, true, false, true]); - let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); + let result = mask_canonical_array(canonical(array, &mut ctx), &mask, &mut ctx).unwrap(); let expected = PrimitiveArray::from_option_iter([None, None, Some(3i32), None, None]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_mask_decimal_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DecimalDType::new(10, 2); let array = DecimalArray::from_option_iter( [Some(1i128), Some(2), Some(3), Some(4), Some(5)], @@ -243,29 +241,29 @@ mod tests { ); let mask = Mask::from_iter([true, true, false, true, true]); - let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); + let result = mask_canonical_array(canonical(array, &mut ctx), &mask, &mut ctx).unwrap(); let expected = DecimalArray::from_option_iter([Some(1i128), Some(2), None, Some(4), Some(5)], dtype); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_mask_varbinview_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = VarBinViewArray::from_iter_str(["one", "two", "three", "four", "five"]); let mask = Mask::from_iter([false, true, false, true, false]); - let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); + let result = mask_canonical_array(canonical(array, &mut ctx), &mask, &mut ctx).unwrap(); let expected = VarBinViewArray::from_iter_nullable_str([None, Some("two"), None, Some("four"), None]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_mask_list_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]).into_array(); let offsets = PrimitiveArray::from_iter([0i32, 2, 4]).into_array(); let sizes = PrimitiveArray::from_iter([2i32, 2, 2]).into_array(); @@ -276,7 +274,7 @@ mod tests { let mask = Mask::from_iter([true, false, true]); - let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); + let result = mask_canonical_array(canonical(array, &mut ctx), &mask, &mut ctx).unwrap(); assert_eq!(result.len(), 3); assert!(result.is_valid(0, &mut ctx).unwrap()); @@ -286,14 +284,14 @@ mod tests { #[test] fn test_mask_fixed_size_list_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]).into_array(); let array = FixedSizeListArray::try_new(elements, 2, Nullability::NonNullable.into(), 3).unwrap(); let mask = Mask::from_iter([false, true, false]); - let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); + let result = mask_canonical_array(canonical(array, &mut ctx), &mask, &mut ctx).unwrap(); assert_eq!(result.len(), 3); assert!(!result.is_valid(0, &mut ctx).unwrap()); @@ -303,7 +301,7 @@ mod tests { #[test] fn test_mask_struct_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let field1 = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); let field2 = PrimitiveArray::from_iter([4i32, 5, 6]).into_array(); let fields = vec![field1, field2]; @@ -318,7 +316,7 @@ mod tests { let mask = Mask::from_iter([true, false, true]); - let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); + let result = mask_canonical_array(canonical(array, &mut ctx), &mask, &mut ctx).unwrap(); assert_eq!(result.len(), 3); assert!(result.is_valid(0, &mut ctx).unwrap()); @@ -328,35 +326,36 @@ mod tests { #[test] fn test_mask_all_false() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]); let mask = Mask::AllFalse(5); - let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); + let result = mask_canonical_array(canonical(array, &mut ctx), &mask, &mut ctx).unwrap(); let expected = PrimitiveArray::from_option_iter([None, None, None, None, None::]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_mask_all_true() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]); let mask = Mask::AllTrue(5); - let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); + let result = mask_canonical_array(canonical(array, &mut ctx), &mask, &mut ctx).unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4), Some(5)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_mask_empty_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::from_iter(Vec::::new()); for mask in [Mask::AllFalse(0), Mask::AllTrue(0)] { - let result = mask_canonical_array(canonical(array.clone()), &mask, &mut ctx).unwrap(); + let result = + mask_canonical_array(canonical(array.clone(), &mut ctx), &mask, &mut ctx).unwrap(); assert_eq!(result.len(), 0); } } diff --git a/fuzz/src/array/min_max.rs b/fuzz/src/array/min_max.rs index 5600fe27568..3ac18954c2c 100644 --- a/fuzz/src/array/min_max.rs +++ b/fuzz/src/array/min_max.rs @@ -4,6 +4,7 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray as _; +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_error::VortexResult; @@ -13,5 +14,9 @@ pub fn min_max_canonical_array( canonical: Canonical, ctx: &mut ExecutionCtx, ) -> VortexResult> { - min_max(&canonical.into_array(), ctx) + min_max( + &canonical.into_array(), + ctx, + NumericalAggregateOpts::default(), + ) } diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index e19a38bae86..7c1bbaad091 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -43,6 +43,7 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::all_non_distinct::all_non_distinct; use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; use vortex_array::aggregate_fn::fns::min_max::min_max; @@ -667,8 +668,9 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> VortexFuzzResult { assert_scalar_eq(&expected.scalar(), &sum_result, i)?; } Action::MinMax => { - let min_max_result = min_max(¤t_array, &mut ctx) - .vortex_expect("min_max operation should succeed in fuzz test"); + let min_max_result = + min_max(¤t_array, &mut ctx, NumericalAggregateOpts::default()) + .vortex_expect("min_max operation should succeed in fuzz test"); assert_min_max_eq(expected.min_max().as_ref(), min_max_result.as_ref(), i)?; } Action::FillNull(fill_value) => { diff --git a/fuzz/src/array/search_sorted.rs b/fuzz/src/array/search_sorted.rs index 788e51bac46..762182b35fb 100644 --- a/fuzz/src/array/search_sorted.rs +++ b/fuzz/src/array/search_sorted.rs @@ -6,7 +6,6 @@ use std::fmt::Debug; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; @@ -133,8 +132,10 @@ pub fn search_sorted_canonical_array( } DType::Utf8(_) | DType::Binary(_) => { let utf8 = array.clone().execute::(ctx)?; - let opt_values = - utf8.with_iterator(|iter| iter.map(|v| v.map(|u| u.to_vec())).collect::>()); + let mask = utf8.validity()?.execute_mask(utf8.len(), ctx)?; + let opt_values = (0..utf8.len()) + .map(|i| mask.value(i).then(|| utf8.bytes_at(i).to_vec())) + .collect::>(); let to_find = if matches!(array.dtype(), DType::Utf8(_)) { BufferString::try_from(scalar)?.as_str().as_bytes().to_vec() } else { diff --git a/fuzz/src/array/slice.rs b/fuzz/src/array/slice.rs index 7758d41219c..2a151002805 100644 --- a/fuzz/src/array/slice.rs +++ b/fuzz/src/array/slice.rs @@ -4,7 +4,6 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::FixedSizeListArray; @@ -69,8 +68,10 @@ pub fn slice_canonical_array( } DType::Utf8(_) | DType::Binary(_) => { let utf8 = array.clone().execute::(ctx)?; - let values = - utf8.with_iterator(|iter| iter.map(|v| v.map(|u| u.to_vec())).collect::>()); + let mask = utf8.validity()?.execute_mask(utf8.len(), ctx)?; + let values = (0..utf8.len()) + .map(|i| mask.value(i).then(|| utf8.bytes_at(i).to_vec())) + .collect::>(); Ok(VarBinViewArray::from_iter( values[start..stop].iter().cloned(), array.dtype().clone(), diff --git a/fuzz/src/array/sort.rs b/fuzz/src/array/sort.rs index c77e06a196b..3f00fccd06e 100644 --- a/fuzz/src/array/sort.rs +++ b/fuzz/src/array/sort.rs @@ -6,7 +6,6 @@ use std::cmp::Ordering; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; @@ -86,8 +85,10 @@ pub fn sort_canonical_array(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexR } DType::Utf8(_) | DType::Binary(_) => { let utf8 = array.clone().execute::(ctx)?; - let mut opt_values = - utf8.with_iterator(|iter| iter.map(|v| v.map(|u| u.to_vec())).collect::>()); + let mask = utf8.validity()?.execute_mask(utf8.len(), ctx)?; + let mut opt_values = (0..utf8.len()) + .map(|i| mask.value(i).then(|| utf8.bytes_at(i).to_vec())) + .collect::>(); opt_values.sort(); Ok(VarBinViewArray::from_iter(opt_values, array.dtype().clone()).into_array()) } diff --git a/fuzz/src/array/take.rs b/fuzz/src/array/take.rs index 7bfc2226ba0..c67ae6184d8 100644 --- a/fuzz/src/array/take.rs +++ b/fuzz/src/array/take.rs @@ -4,7 +4,6 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; @@ -103,8 +102,10 @@ pub fn take_canonical_array( } DType::Utf8(_) | DType::Binary(_) => { let utf8 = array.clone().execute::(ctx)?; - let values = - utf8.with_iterator(|iter| iter.map(|v| v.map(|u| u.to_vec())).collect::>()); + let mask = utf8.validity()?.execute_mask(utf8.len(), ctx)?; + let values = (0..utf8.len()) + .map(|i| mask.value(i).then(|| utf8.bytes_at(i).to_vec())) + .collect::>(); Ok(VarBinViewArray::from_iter( indices .iter() diff --git a/fuzz/src/fsst_like.rs b/fuzz/src/fsst_like.rs index 5ca10af310a..9a8c9ae21ac 100644 --- a/fuzz/src/fsst_like.rs +++ b/fuzz/src/fsst_like.rs @@ -21,7 +21,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar_fn::fns::like::Like; use vortex_array::scalar_fn::fns::like::LikeOptions; -use vortex_array::session::ArraySession; use vortex_error::VortexResult; use vortex_fsst::FSSTArray; use vortex_fsst::fsst_compress; @@ -32,8 +31,7 @@ use crate::error::Backtrace; use crate::error::VortexFuzzError; use crate::error::VortexFuzzResult; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); /// A random string from a small alphabet (`a..=h`) with bounded length. #[derive(Debug)] @@ -109,18 +107,15 @@ pub fn run_fsst_like_fuzz(fuzz: FuzzFsstLike) -> VortexFuzzResult { let varbin = VarBinArray::from_iter( strings.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); + ) + .into_array(); // Train FSST compressor and compress. - let compressor = fsst_train_compressor(&varbin); let mut ctx = SESSION.create_execution_ctx(); - let fsst_array: FSSTArray = fsst_compress( - varbin.clone(), - varbin.len(), - varbin.dtype(), - &compressor, - &mut ctx, - ); + let compressor = fsst_train_compressor(&varbin, &mut ctx) + .map_err(|err| VortexFuzzError::VortexError(err, Backtrace::capture()))?; + let fsst_array: FSSTArray = fsst_compress(&varbin, &compressor, &mut ctx) + .map_err(|err| VortexFuzzError::VortexError(err, Backtrace::capture()))?; let opts = LikeOptions { negated, @@ -128,7 +123,7 @@ pub fn run_fsst_like_fuzz(fuzz: FuzzFsstLike) -> VortexFuzzResult { }; // Run LIKE on the uncompressed array. - let expected = run_like_on_array(&varbin.into_array(), &pattern, len, opts) + let expected = run_like_on_array(&varbin, &pattern, len, opts) .map_err(|err| VortexFuzzError::VortexError(err, Backtrace::capture()))?; // Run LIKE on the FSST-compressed array. diff --git a/fuzz/src/gpu/mod.rs b/fuzz/src/gpu/mod.rs index fa6ef110759..d9ac5c4fa0f 100644 --- a/fuzz/src/gpu/mod.rs +++ b/fuzz/src/gpu/mod.rs @@ -12,7 +12,6 @@ use arbitrary::Unstructured; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::dict::ArbitraryDictArray; use vortex_array::dtype::Nullability; @@ -110,13 +109,11 @@ pub async fn run_compress_gpu(fuzz: FuzzCompressGpu) -> VortexFuzzResult { } let FuzzCompressGpu { array } = fuzz; + let mut ctx = SESSION.create_execution_ctx(); let original_len = array.len(); - let cpu_canonical = match array - .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) - { + let cpu_canonical = match array.clone().execute::(&mut ctx) { Ok(c) => c, Err(e) => { return Err(VortexFuzzError::VortexError(e, Backtrace::capture())); @@ -165,10 +162,10 @@ pub async fn run_compress_gpu(fuzz: FuzzCompressGpu) -> VortexFuzzResult { for i in 0..original_len { let cpu_scalar = cpu_array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut ctx) .map_err(|e| VortexFuzzError::VortexError(e, Backtrace::capture()))?; let gpu_scalar = gpu_array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut ctx) .map_err(|e| VortexFuzzError::VortexError(e, Backtrace::capture()))?; if cpu_scalar != gpu_scalar { diff --git a/java/build.gradle.kts b/java/build.gradle.kts index 43cc24b48ed..8afda976316 100644 --- a/java/build.gradle.kts +++ b/java/build.gradle.kts @@ -6,7 +6,7 @@ import net.ltgt.gradle.errorprone.errorprone plugins { id("com.diffplug.spotless") version "8.6.0" id("com.palantir.git-version") version "5.0.0" - id("com.palantir.java-format") version "2.91.0" + id("com.palantir.java-format") version "2.93.0" id("net.ltgt.errorprone") version "5.1.0" apply false id("com.vanniktech.maven.publish") version "0.36.0" apply false } 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 24780785274..93c5df2c61a 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 @@ -128,6 +128,57 @@ public OptionalLong asOptional() { } } + /** + * Sum of the on-storage byte sizes of all files included in this data source along with the precision of that + * estimate. Mirrors the Rust {@code Option>} returned by {@code DataSource::byte_size}: + * {@link ByteSize.Unknown} when no estimate is available (for example when the filesystem listing did not return + * sizes), {@link ByteSize.Estimate} for an inexact hint (some files contribute extrapolated sizes), and + * {@link ByteSize.Exact} when every file has a known size. + */ + public ByteSize byteSize() { + long[] out = new long[2]; + NativeDataSource.byteSize(pointer, out); + return switch ((int) out[1]) { + case 1 -> new ByteSize.Estimate(out[0]); + case 2 -> new ByteSize.Exact(out[0]); + default -> ByteSize.Unknown.INSTANCE; + }; + } + + /** Precision-aware byte size. See {@link #byteSize()}. */ + public sealed interface ByteSize { + /** Returns the byte size as a long, or {@code OptionalLong.empty()} when unknown. */ + OptionalLong asOptional(); + + /** Byte size is not known. */ + final class Unknown implements ByteSize { + public static final Unknown INSTANCE = new Unknown(); + + private Unknown() {} + + @Override + public OptionalLong asOptional() { + return OptionalLong.empty(); + } + } + + /** Estimated byte size; the actual value may differ. */ + record Estimate(long value) implements ByteSize { + @Override + public OptionalLong asOptional() { + return OptionalLong.of(value); + } + } + + /** Exact byte size. */ + record Exact(long value) implements ByteSize { + @Override + public OptionalLong asOptional() { + return OptionalLong.of(value); + } + } + } + /** Submit a scan. */ public Scan scan(ScanOptions options) { Objects.requireNonNull(options, "options"); 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 b7e58d2dc21..cc2aa163cee 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 @@ -33,4 +33,10 @@ private NativeDataSource() {} * {@code 1=estimate}, {@code 2=exact}. */ public static native void rowCount(long pointer, long[] out); + + /** + * Populate {@code out} with {@code [bytes, precision]}, the sum of on-storage file sizes for the data source. + * Precision is one of {@code 0=unknown}, {@code 1=estimate}, {@code 2=exact}. + */ + public static native void byteSize(long pointer, long[] out); } 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 322dc5522c3..aeee8febcf4 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 @@ -11,6 +11,7 @@ import dev.vortex.jni.NativeLoader; import java.io.IOException; import java.math.BigDecimal; +import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; import java.util.HashMap; @@ -132,6 +133,7 @@ public void testFullScan() throws Exception { DataSource ds = DataSource.open(session, writePath); assertEquals(new DataSource.RowCount.Exact(10L), ds.rowCount()); + assertEquals(new DataSource.ByteSize.Exact(Files.size(tempDir.resolve("minimal.vortex"))), ds.byteSize()); var schema = ds.arrowSchema(allocator); assertEquals( diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexBatchExec.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexBatchExec.java index 8df7ce8e1db..198dfd6c77c 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexBatchExec.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexBatchExec.java @@ -4,6 +4,7 @@ package dev.vortex.spark.read; import com.google.common.collect.ImmutableMap; +import dev.vortex.api.Session; import dev.vortex.jni.NativeFiles; import dev.vortex.spark.VortexFilePartition; import dev.vortex.spark.VortexSparkSession; @@ -76,14 +77,19 @@ public PartitionReaderFactory createReaderFactory() { } private List resolvePaths() { - var session = VortexSparkSession.get(formatOptions); + return resolveVortexPaths(VortexSparkSession.get(formatOptions), paths, formatOptions); + } + + /** + * Expands directory-like entries to concrete {@code .vortex} files; entries that already name a {@code .vortex} + * file are kept as-is. Shared with {@link VortexScan#estimateStatistics()} so planning and execution resolve paths + * identically. + */ + static List resolveVortexPaths(Session session, List paths, Map formatOptions) { return paths.stream() - .flatMap(path -> { - if (path.endsWith(".vortex")) { - return Stream.of(path); - } - return NativeFiles.listFiles(session, path, formatOptions).stream(); - }) + .flatMap(path -> path.endsWith(".vortex") + ? Stream.of(path) + : NativeFiles.listFiles(session, path, formatOptions).stream()) .collect(Collectors.toList()); } diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScan.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScan.java index d5949b57a4d..02a7563f925 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScan.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScan.java @@ -3,38 +3,62 @@ package dev.vortex.spark.read; +import dev.vortex.api.DataSource; +import dev.vortex.api.Session; +import dev.vortex.spark.VortexSparkSession; import java.util.Arrays; import java.util.List; import java.util.Map; +import java.util.OptionalLong; import org.apache.spark.sql.connector.catalog.CatalogV2Util; import org.apache.spark.sql.connector.catalog.Column; +import org.apache.spark.sql.connector.expressions.NamedReference; import org.apache.spark.sql.connector.expressions.filter.Predicate; import org.apache.spark.sql.connector.read.Batch; import org.apache.spark.sql.connector.read.Scan; +import org.apache.spark.sql.connector.read.Statistics; +import org.apache.spark.sql.connector.read.SupportsReportStatistics; +import org.apache.spark.sql.connector.read.colstats.ColumnStatistics; +import org.apache.spark.sql.internal.SQLConf; import org.apache.spark.sql.types.StructType; -/** Spark V2 {@link Scan} over a table of Vortex files. */ -public final class VortexScan implements Scan { +/** + * Spark V2 {@link Scan} over a table of Vortex files. + * + *

Implements {@link SupportsReportStatistics} to surface both the row count Vortex records in each file footer and a + * Spark scan-size estimate. The byte estimate starts from the on-storage file sizes collected by + * {@code MultiFileDataSource}, then follows Spark's file scan convention by applying the SQL file-compression factor + * and scaling by the pushed read schema's default size relative to the full table schema's default size. When the + * listing did not return a size for one or more files the file-byte total is extrapolated before Spark scaling is + * applied. + */ +public final class VortexScan implements Scan, SupportsReportStatistics { private final List paths; + private final List tableColumns; private final List readColumns; private final Map formatOptions; private final Predicate[] pushedPredicates; + private volatile Statistics cachedStatistics; + /** * Creates a new VortexScan for the specified file paths and columns. The caller is responsible for passing * immutable collections; the constructor does not copy. * * @param paths the list of Vortex file paths to scan + * @param tableColumns the full table columns before projection pushdown * @param readColumns the list of columns to read from the files * @param pushedPredicates predicates pushed down by Spark; {@code null} or empty means no pushdown */ public VortexScan( List paths, + List tableColumns, List readColumns, - Map formatOptions, - Predicate[] pushedPredicates) { + Predicate[] pushedPredicates, + Map formatOptions) { this.paths = paths; + this.tableColumns = tableColumns; this.readColumns = readColumns; this.formatOptions = formatOptions; this.pushedPredicates = pushedPredicates == null ? new Predicate[0] : pushedPredicates.clone(); @@ -83,4 +107,70 @@ public Batch toBatch() { public ColumnarSupportMode columnarSupportMode() { return ColumnarSupportMode.SUPPORTED; } + + /** + * Returns statistics for this scan. + * + *

Opens the Vortex {@link DataSource} on first invocation and caches the result. The row count is taken from the + * data source (sum of file-footer row counts; extrapolated from the first opened file when other files are + * deferred). {@link Statistics#sizeInBytes()} is derived from the per-file sizes reported by the filesystem + * listing, then adjusted by Spark's compression factor and the ratio between the pushed read schema and the full + * table schema. When a listing did not return a size for some file the file-byte total is extrapolated. When no + * file size is known at all the value is left empty so Spark falls back to its default heuristic. + * + * @return statistics with row-count and Spark scan-size estimates + */ + @Override + public Statistics estimateStatistics() { + Statistics local = cachedStatistics; + if (local != null) { + return local; + } + synchronized (this) { + if (cachedStatistics == null) { + cachedStatistics = computeStatistics(); + } + return cachedStatistics; + } + } + + private Statistics computeStatistics() { + Session session = VortexSparkSession.get(formatOptions); + List resolvedPaths = VortexBatchExec.resolveVortexPaths(session, paths, formatOptions); + if (resolvedPaths.isEmpty()) { + return new VortexStatistics(OptionalLong.empty(), OptionalLong.empty()); + } + + DataSource source = DataSource.open(session, resolvedPaths, formatOptions); + return new VortexStatistics( + source.rowCount().asOptional(), + scaleSizeInBytes(source.byteSize().asOptional())); + } + + private OptionalLong scaleSizeInBytes(OptionalLong fileBytes) { + if (fileBytes.isEmpty()) { + return OptionalLong.empty(); + } + + StructType tableSchema = CatalogV2Util.v2ColumnsToStructType(tableColumns.toArray(new Column[0])); + StructType readSchema = readSchema(); + int tableDefaultSize = tableSchema.defaultSize(); + if (tableDefaultSize <= 0) { + return fileBytes; + } + + double scaled = SQLConf.get().fileCompressionFactor() + * fileBytes.getAsLong() + / tableDefaultSize + * readSchema.defaultSize(); + return OptionalLong.of((long) scaled); + } + + private record VortexStatistics(OptionalLong numRows, OptionalLong sizeInBytes) implements Statistics { + + @Override + public Map columnStats() { + return Map.of(); + } + } } 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 94990432b45..107d22b30f6 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 @@ -31,7 +31,8 @@ public final class VortexScanBuilder implements ScanBuilder, SupportsPushDownRequiredColumns, SupportsPushDownV2Filters { private final ImmutableList.Builder paths; - private final List columns; + private final List tableColumns; + private final List readColumns; private final Map formatOptions; private final Set partitionColumnNames; private Predicate[] pushedPredicates = new Predicate[0]; @@ -48,10 +49,11 @@ public VortexScanBuilder(Map formatOptions) { */ public VortexScanBuilder(Map formatOptions, Transform[] partitionTransforms) { this.paths = ImmutableList.builder(); - this.columns = new ArrayList<>(); Map options = Maps.newHashMap(); options.put("vortex.workerThreads", "4"); options.putAll(formatOptions); + this.tableColumns = new ArrayList<>(); + this.readColumns = new ArrayList<>(); this.formatOptions = options; this.partitionColumnNames = collectPartitionColumnNames(partitionTransforms); } @@ -74,7 +76,8 @@ public VortexScanBuilder addPath(String path) { * @return this builder for method chaining */ public VortexScanBuilder addColumn(Column column) { - this.columns.add(column); + this.tableColumns.add(column); + this.readColumns.add(column); return this; } @@ -97,7 +100,7 @@ public VortexScanBuilder addAllPaths(Iterable paths) { */ public VortexScanBuilder addAllColumns(Iterable columns) { for (Column column : columns) { - this.columns.add(column); + addColumn(column); } return this; } @@ -116,7 +119,12 @@ public Scan build() { // 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, List.copyOf(this.columns), this.formatOptions, pushedPredicates); + return new VortexScan( + paths, + List.copyOf(this.tableColumns), + List.copyOf(this.readColumns), + pushedPredicates, + this.formatOptions); } /** @@ -129,8 +137,8 @@ public Scan build() { */ @Override public void pruneColumns(StructType requiredSchema) { - columns.clear(); - columns.addAll(Arrays.asList(CatalogV2Util.structTypeToV2Columns(requiredSchema))); + readColumns.clear(); + readColumns.addAll(Arrays.asList(CatalogV2Util.structTypeToV2Columns(requiredSchema))); } /** @@ -145,7 +153,7 @@ public void pruneColumns(StructType requiredSchema) { @Override public Predicate[] pushPredicates(Predicate[] predicates) { Map dataColumnTypes = new HashMap<>(); - for (Column column : columns) { + for (Column column : readColumns) { if (!partitionColumnNames.contains(column.name())) { dataColumnTypes.put(column.name(), column.dataType()); } diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceStatsTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceStatsTest.java new file mode 100644 index 00000000000..0595349a49a --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceStatsTest.java @@ -0,0 +1,241 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.vortex.spark.read.VortexScan; +import dev.vortex.spark.read.VortexScanBuilder; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Comparator; +import java.util.Map; +import java.util.stream.Stream; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SaveMode; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.connector.catalog.Column; +import org.apache.spark.sql.connector.read.Statistics; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +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; + +/** + * Integration tests for {@link VortexScan#estimateStatistics()}. + * + *

Verifies that the Spark V2 scan surfaces both the row count Vortex stores in each file footer and the sum of the + * on-storage file sizes reported by the filesystem listing. + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +public final class VortexDataSourceStatsTest { + private static final String FILE_COMPRESSION_FACTOR_KEY = "spark.sql.sources.fileCompressionFactor"; + + private SparkSession spark; + + @TempDir + Path tempDir; + + @BeforeAll + public void setUp() { + spark = SparkSession.builder() + .appName("VortexStatsTest") + .master("local[2]") + .config("spark.driver.host", "127.0.0.1") + .config("spark.sql.shuffle.partitions", "2") + .config("spark.sql.adaptive.enabled", "false") + .config("spark.ui.enabled", "false") + .getOrCreate(); + } + + @AfterAll + public void tearDown() { + if (spark != null) { + spark.stop(); + } + } + + @Test + @DisplayName("VortexScan reports exact row count for single-file scans") + public void testEstimateStatisticsReportsRowCount() throws IOException { + int numRows = 250; + Path outputPath = writeRows(numRows, "single_file"); + + VortexScan scan = buildScan(outputPath); + Statistics stats = scan.estimateStatistics(); + + assertTrue( + stats.numRows().isPresent(), + "VortexScan should report a row count for a Vortex dataset with a populated footer"); + assertEquals(numRows, stats.numRows().getAsLong(), "Row count should match the rows we wrote"); + } + + @Test + @DisplayName("VortexScan reports aggregate row count across multi-file scans") + public void testEstimateStatisticsAcrossMultipleFiles() throws IOException { + int numRows = 400; + Path outputPath = writeRows(numRows, "multi_file", 4); + + VortexScan scan = buildScan(outputPath); + Statistics stats = scan.estimateStatistics(); + + assertTrue(stats.numRows().isPresent(), "Row count should be reported for multi-file Vortex datasets"); + assertEquals(numRows, stats.numRows().getAsLong(), "Row count should sum across all files"); + } + + @Test + @DisplayName("VortexScan reports sizeInBytes equal to the sum of on-storage file sizes") + public void testEstimateStatisticsReportsSizeInBytes() throws IOException { + Path outputPath = writeRows(120, "with_size", 3); + + long fileBytes = totalVortexFileBytes(outputPath); + assertTrue(fileBytes > 0, "Test setup should produce at least one non-empty .vortex file"); + + VortexScan scan = buildScan(outputPath); + Statistics stats = scan.estimateStatistics(); + + assertTrue( + stats.sizeInBytes().isPresent(), + "VortexScan should surface a sizeInBytes when the filesystem listing reports file sizes"); + // Mirror the scan's Spark-convention scaling (factor 1.0, unpruned schema), which divides and + // re-multiplies by the schema default size in double arithmetic before truncating; asserting + // against the raw byte sum would be sensitive to the floating-point round trip. + StructType schema = spark.read() + .format("vortex") + .option("path", outputPath.toUri().toString()) + .load() + .schema(); + long expectedSize = (long) (1.0 * fileBytes / schema.defaultSize() * schema.defaultSize()); + assertEquals( + expectedSize, + stats.sizeInBytes().getAsLong(), + "sizeInBytes should equal the sum of on-storage .vortex file sizes"); + } + + @Test + @DisplayName("VortexScan scales sizeInBytes by the pushed read schema") + public void testEstimateStatisticsScalesSizeInBytesForProjection() throws IOException { + Path outputPath = writeRows(120, "projected_size", 3); + long fileBytes = totalVortexFileBytes(outputPath); + + StructType fullSchema = spark.read() + .format("vortex") + .option("path", outputPath.toUri().toString()) + .load() + .schema(); + StructType idOnlySchema = new StructType(new StructField[] {fullSchema.fields()[0]}); + + String previousCompressionFactor = spark.conf().get(FILE_COMPRESSION_FACTOR_KEY); + spark.conf().set(FILE_COMPRESSION_FACTOR_KEY, "0.5"); + try { + VortexScan scan = buildScan(outputPath, idOnlySchema); + Statistics stats = scan.estimateStatistics(); + + long expectedSize = (long) (0.5 * fileBytes / fullSchema.defaultSize() * idOnlySchema.defaultSize()); + assertTrue(stats.sizeInBytes().isPresent(), "Projected scans should still surface sizeInBytes"); + assertEquals( + expectedSize, + stats.sizeInBytes().getAsLong(), + "sizeInBytes should follow Spark FileScan's compression and schema-width scaling"); + assertTrue( + stats.sizeInBytes().getAsLong() < fileBytes, + "Projected scan stats should be smaller than full file bytes"); + } finally { + spark.conf().set(FILE_COMPRESSION_FACTOR_KEY, previousCompressionFactor); + } + } + + @Test + @DisplayName("VortexScan caches statistics across repeated calls") + public void testEstimateStatisticsIsCached() throws IOException { + Path outputPath = writeRows(50, "cached", 1); + + VortexScan scan = buildScan(outputPath); + Statistics first = scan.estimateStatistics(); + Statistics second = scan.estimateStatistics(); + + // Same instance returned -- the second call hits the cached value. + assertEquals(first, second, "estimateStatistics() should return the same Statistics object on repeat calls"); + assertInstanceOf(Statistics.class, first); + } + + private VortexScan buildScan(Path outputPath) { + return buildScan(outputPath, null); + } + + private VortexScan buildScan(Path outputPath, StructType requiredSchema) { + Dataset readDf = spark.read() + .format("vortex") + .option("path", outputPath.toUri().toString()) + .load(); + StructType readSchema = readDf.schema(); + + VortexScanBuilder builder = new VortexScanBuilder(Map.of()); + builder.addPath(outputPath.toUri().toString()); + for (StructField field : readSchema.fields()) { + builder.addColumn(Column.create(field.name(), field.dataType())); + } + if (requiredSchema != null) { + builder.pruneColumns(requiredSchema); + } + return (VortexScan) builder.build(); + } + + private Path writeRows(int numRows, String name) throws IOException { + return writeRows(numRows, name, 1); + } + + private long totalVortexFileBytes(Path outputPath) throws IOException { + try (Stream paths = Files.walk(outputPath)) { + return paths.filter(Files::isRegularFile) + .filter(path -> path.getFileName().toString().endsWith(".vortex")) + .mapToLong(path -> { + try { + return Files.size(path); + } catch (IOException e) { + throw new RuntimeException(e); + } + }) + .sum(); + } + } + + private Path writeRows(int numRows, String name, int partitions) throws IOException { + Path outputPath = tempDir.resolve(name); + Dataset df = spark.range(0, numRows) + .selectExpr("cast(id as int) as id", "concat('value_', cast(id as string)) as value"); + + df.repartition(partitions) + .write() + .format("vortex") + .option("path", outputPath.toUri().toString()) + .mode(SaveMode.Overwrite) + .save(); + return outputPath; + } + + @AfterEach + public void cleanupTempFiles() throws IOException { + if (tempDir != null && Files.exists(tempDir)) { + try (Stream paths = Files.walk(tempDir)) { + paths.sorted(Comparator.reverseOrder()).forEach(path -> { + try { + Files.deleteIfExists(path); + } catch (IOException e) { + System.err.println("Failed to delete: " + path); + } + }); + } + } + } +} diff --git a/package-lock.json b/package-lock.json deleted file mode 100644 index 0cd0ca8e96d..00000000000 --- a/package-lock.json +++ /dev/null @@ -1,6 +0,0 @@ -{ - "name": "vortex", - "lockfileVersion": 3, - "requires": true, - "packages": {} -} diff --git a/package.json b/package.json deleted file mode 100644 index 0967ef424bc..00000000000 --- a/package.json +++ /dev/null @@ -1 +0,0 @@ -{} diff --git a/pyproject.toml b/pyproject.toml index f0e82ec52c9..53b769f350c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,11 @@ members = ["bench-orchestrator", "vortex-python", "docs"] [tool.uv.sources] bench-orchestrator = { workspace = true } vortex-data = { workspace = true } +# Resolve vortex-data-cuda from its local path so the optional `cuda` extra is +# lockable before the package is published to PyPI. A path source (not +# `workspace = true`) keeps the CUDA extension out of the default `--all-packages` +# set, so it is only built when the `cuda` extra is requested. +vortex-data-cuda = { path = "vortex-python-cuda" } docs = { workspace = true } [tool.ruff] diff --git a/renovate.json b/renovate.json index 80fb4c20238..68a14baa7ca 100644 --- a/renovate.json +++ b/renovate.json @@ -2,34 +2,30 @@ "$schema": "https://docs.renovatebot.com/renovate-schema.json", "extends": [ "config:recommended", - ":automergeStableNonMajor", - ":automergePr", ":automergeRequireAllStatusChecks", - ":combinePatchMinorReleases", - ":dependencyDashboard", ":separateMultipleMajorReleases", ":configMigration", "group:rust-futuresMonorepo", "helpers:pinGitHubActionDigests", - "schedule:earlyMondays", - "docker:disable" + "docker:disable", + ":maintainLockFilesWeekly" ], - "lockFileMaintenance": { - "enabled": true - }, + "schedule": "* 0-3 * * 1#2,1#4", "autoApprove": true, "automergeStrategy": "squash", "prHourlyLimit": 0, "prConcurrentLimit": 20, "rebaseWhen": "conflicted", - "cloneSubmodules": true, - "platformAutomerge": true, - "patch": { - "groupName": "all patch updates" - }, "labels": [ "changelog/chore" ], + "minor": { "enabled": false }, + "patch": { "enabled": false }, + "osvVulnerabilityAlerts": true, + "vulnerabilityAlerts": { + "enabled": true, + "vulnerabilityFixStrategy": "lowest" + }, "packageRules": [ { "matchPackageNames": [ @@ -40,15 +36,38 @@ }, { "groupName": "arrow-rs", + "matchManagers": ["cargo"], "matchUpdateTypes": [ - "major", - "minor", - "patch" + "major" ], "matchSourceUrls": [ "https://github.com/apache/arrow-rs" ] }, + { + "description": "Do not split lance-bench grouped majors by target major", + "matchManagers": [ + "cargo" + ], + "matchFileNames": [ + "benchmarks/lance-bench/Cargo.toml" + ], + "separateMultipleMajor": false + }, + { + "description": "Keep lance-bench dependency majors together because they are tested as one benchmark stack", + "groupName": "lance benchmark dependencies", + "groupSlug": "lance-bench", + "matchManagers": [ + "cargo" + ], + "matchFileNames": [ + "benchmarks/lance-bench/Cargo.toml" + ], + "matchUpdateTypes": [ + "major" + ] + }, { "groupName": "flatbuffers", "matchSourceUrls": [ @@ -61,6 +80,21 @@ "https://github.com/hyperium/tonic" ] }, + { + "description": "Disable routine npm dependency updates; keep security alerts and lock file maintenance", + "matchManagers": [ + "npm" + ], + "matchUpdateTypes": [ + "digest", + "major", + "minor", + "patch", + "pin", + "replacement" + ], + "enabled": false + }, { "groupName": "Rust lock file maintenance", "groupSlug": "rust-lock-file-maintenance", @@ -100,7 +134,7 @@ "/vortex-duckdb/build\\.rs/" ], "matchStrings": [ - "const DUCKDB_VERSION: &str = \"(?[0-9]+\\.[0-9]+\\.[0-9]+)\";" + "const DEFAULT_DUCKDB_VERSION: &str = \"(?[0-9]+\\.[0-9]+\\.[0-9]+)\";" ], "depNameTemplate": "duckdb/duckdb", "datasourceTemplate": "github-releases", diff --git a/rustfmt.toml b/rustfmt.toml index 83fe627dbf6..e6ffec93d96 100644 --- a/rustfmt.toml +++ b/rustfmt.toml @@ -1,8 +1,9 @@ condense_wildcard_suffixes = true -format_macro_matchers = true format_macro_bodies = true +format_macro_matchers = true group_imports = "StdExternalCrate" -unstable_features = true -use_field_init_shorthand = true imports_granularity = "Item" +normalize_comments = true style_edition = "2024" +unstable_features = true +use_field_init_shorthand = true diff --git a/scripts/_measurement_id.py b/scripts/_measurement_id.py new file mode 100644 index 00000000000..e8834f4c097 --- /dev/null +++ b/scripts/_measurement_id.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Python port of the server-internal `measurement_id` xxhash64 functions. + +This is a byte-for-byte port of the `vortex-data/benchmarks-website` repo's +`server/src/db.rs` (`measurement_id_*`, `hasher_for`, +`write_str`, `write_opt_str`, `write_i32`, `write_f64`, `finish`). The v4 ingest +writer computes the same `measurement_id` the v3 Rust server computes, so +re-ingesting an existing `(commit, dim-tuple)` upserts the existing row via +`ON CONFLICT (measurement_id) DO UPDATE` instead of inserting a duplicate. + +The equivalence is NOT assumed -- it is pinned by golden vectors generated from +the Rust source of truth (the `measurement_id_golden_vectors` test in the +`vortex-data/benchmarks-website` repo's `server/src/db.rs`). The Python port is +checked against those vectors; any drift in either implementation fails that +check. + +## The hash, precisely + +`measurement_id` is a canonical xxhash64 (seed 0) over a byte buffer built as: + +1. The per-table tag bytes, then a single `0x00` separator (`hasher_for`). +2. For each dimensional field, in the exact order the Rust function writes it: + - `write_str(s)` -> `len(utf8(s))` as a little-endian u64, then the utf8 + bytes. The length is the BYTE length, not the character count. + - `write_opt_str(o)` -> `0x00` for `None`; `0x01` then `write_str(s)` for + `Some(s)`. + - `write_i32(v)` -> 4 little-endian two's-complement bytes. + - `write_f64(v)` -> the IEEE-754 bit pattern as 8 little-endian bytes + (`v.to_bits()` written little-endian == `struct.pack(" i64` (Postgres/DuckDB `BIGINT` is + signed), matching Rust's `hasher.finish() as i64`. + +Canonical xxhash64 is streaming-equivalent: feeding the bytes incrementally (as +the Rust `Hasher` does via `write_*`) yields the same digest as hashing the +fully concatenated buffer once (as this module does). + +The little-endian integer encodings are load-bearing and assume a little-endian +host. The Rust side serializes integers with native-endian byte order +(twox-hash 2.x `write_u64` / `write_i32` use `to_ne_bytes`, and `write_f64` is +`write_u64(v.to_bits())`), so byte-for-byte compatibility holds only where +`to_ne_bytes == to_le_bytes`. Every target in play (x86_64 / aarch64 CI runners, +dev machines, the RDS Postgres host, the Vercel reader) is little-endian, and the +golden vectors are generated on a little-endian host and pin it there. On a +big-endian host both this module and the Rust `write_*` would have to switch to a +shared explicit endianness. +""" + +import struct + +import xxhash + +# Must match `XxHash64::with_seed(0)` in `db.rs::hasher_for`. +_SEED = 0 + +# Per-table tag literals. These MUST match the `hasher_for("")` argument in +# each `db.rs::measurement_id_*` function verbatim; the tag is the table name. +_TAG_QUERY_MEASUREMENTS = "query_measurements" +_TAG_COMPRESSION_TIMES = "compression_times" +_TAG_COMPRESSION_SIZES = "compression_sizes" +_TAG_RANDOM_ACCESS_TIMES = "random_access_times" +_TAG_VECTOR_SEARCH_RUNS = "vector_search_runs" + + +def _hasher_buf(tag: str) -> bytearray: + """Start a hash buffer seeded with a per-table tag plus a `0x00` separator. + + Mirrors `db.rs::hasher_for`: two fact tables that share the same dim values + still produce distinct `measurement_id`s because the tag differs. + """ + buf = bytearray() + buf += tag.encode("utf-8") + buf.append(0) + return buf + + +def _write_str(buf: bytearray, s: str) -> None: + """Append a length-prefixed string: utf8 BYTE length as LE u64, then bytes.""" + encoded = s.encode("utf-8") + buf += struct.pack(" None: + """Append an optional string: `0x00` for None, `0x01` + the string for Some.""" + if s is None: + buf.append(0) + else: + buf.append(1) + _write_str(buf, s) + + +def _write_i32(buf: bytearray, v: int) -> None: + """Append a 32-bit signed integer as 4 little-endian two's-complement bytes.""" + buf += struct.pack(" None: + """Append a 64-bit float as its 8 little-endian IEEE-754 bytes. + + `struct.pack(" int: + """Hash the buffer (xxhash64, seed 0) and bit-cast the u64 digest to i64.""" + digest = xxhash.xxh64(bytes(buf), seed=_SEED).intdigest() + # Bit-cast u64 -> i64 to match Rust's `hasher.finish() as i64`. + return digest - (1 << 64) if digest >= (1 << 63) else digest + + +def measurement_id_query( + *, + commit_sha: str, + dataset: str, + dataset_variant: str | None, + scale_factor: str | None, + query_idx: int, + storage: str, + engine: str, + format: str, +) -> int: + """`measurement_id` for a `query_measurements` row. Mirrors + `db.rs::measurement_id_query`.""" + buf = _hasher_buf(_TAG_QUERY_MEASUREMENTS) + _write_str(buf, commit_sha) + _write_str(buf, dataset) + _write_opt_str(buf, dataset_variant) + _write_opt_str(buf, scale_factor) + _write_i32(buf, query_idx) + _write_str(buf, storage) + _write_str(buf, engine) + _write_str(buf, format) + return _finish(buf) + + +def measurement_id_compression_time( + *, + commit_sha: str, + dataset: str, + dataset_variant: str | None, + format: str, + op: str, +) -> int: + """`measurement_id` for a `compression_times` row. Mirrors + `db.rs::measurement_id_compression_time`.""" + buf = _hasher_buf(_TAG_COMPRESSION_TIMES) + _write_str(buf, commit_sha) + _write_str(buf, dataset) + _write_opt_str(buf, dataset_variant) + _write_str(buf, format) + _write_str(buf, op) + return _finish(buf) + + +def measurement_id_compression_size( + *, + commit_sha: str, + dataset: str, + dataset_variant: str | None, + format: str, +) -> int: + """`measurement_id` for a `compression_sizes` row. Mirrors + `db.rs::measurement_id_compression_size`.""" + buf = _hasher_buf(_TAG_COMPRESSION_SIZES) + _write_str(buf, commit_sha) + _write_str(buf, dataset) + _write_opt_str(buf, dataset_variant) + _write_str(buf, format) + return _finish(buf) + + +def measurement_id_random_access( + *, + commit_sha: str, + dataset: str, + format: str, +) -> int: + """`measurement_id` for a `random_access_times` row. Mirrors + `db.rs::measurement_id_random_access`. Note: no `dataset_variant`.""" + buf = _hasher_buf(_TAG_RANDOM_ACCESS_TIMES) + _write_str(buf, commit_sha) + _write_str(buf, dataset) + _write_str(buf, format) + return _finish(buf) + + +def measurement_id_vector_search( + *, + commit_sha: str, + dataset: str, + layout: str, + flavor: str, + threshold: float, +) -> int: + """`measurement_id` for a `vector_search_runs` row. Mirrors + `db.rs::measurement_id_vector_search`. `iterations` is intentionally NOT part + of the dim tuple -- it is a side count.""" + buf = _hasher_buf(_TAG_VECTOR_SEARCH_RUNS) + _write_str(buf, commit_sha) + _write_str(buf, dataset) + _write_str(buf, layout) + _write_str(buf, flavor) + _write_f64(buf, threshold) + return _finish(buf) + + +# Dispatch table keyed by the fact-table name, used by the golden-vector test to +# map a vector's `table` field to the matching port function. Keeping it here (vs. +# in the test) means a new fact table is wired in one place alongside the port. +MEASUREMENT_ID_BY_TABLE = { + _TAG_QUERY_MEASUREMENTS: measurement_id_query, + _TAG_COMPRESSION_TIMES: measurement_id_compression_time, + _TAG_COMPRESSION_SIZES: measurement_id_compression_size, + _TAG_RANDOM_ACCESS_TIMES: measurement_id_random_access, + _TAG_VECTOR_SEARCH_RUNS: measurement_id_vector_search, +} diff --git a/scripts/duckdb-r2-prepare.sh b/scripts/duckdb-r2-prepare.sh new file mode 100755 index 00000000000..35f4655c481 --- /dev/null +++ b/scripts/duckdb-r2-prepare.sh @@ -0,0 +1,61 @@ +#!/bin/sh + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors +# +# Produce $ARCHIVE for upload to R2: either download DuckDB release for +# $REF_DIR or build DuckDB from source at commit $REF_DIR and pack +# libraries and headers. +# +# Required env vars: ARCHIVE, REF_DIR, RELEASE, PLATFORM_OS + +set -eu + +if [ "$RELEASE" = "true" ]; then + echo "Mirroring DuckDB release ${REF_DIR}/${ARCHIVE}" + curl -fSL --retry 3 -o "$ARCHIVE" \ + "https://github.com/duckdb/duckdb/releases/download/${REF_DIR}/${ARCHIVE}" +else + echo "Building DuckDB commit ${REF_DIR} from source" + + curl -fSL --retry 3 -o duckdb-src.zip \ + "https://github.com/duckdb/duckdb/archive/${REF_DIR}.zip" + + # macos zip extract error: cannot create + # <...>/issue2628_������.csv Illegal byte sequence + if [ "$PLATFORM_OS" = "osx" ]; then + 7z x duckdb-src.zip + else + unzip -q duckdb-src.zip + fi + + src_dir="duckdb-${REF_DIR}" + extra="" + if [ "$PLATFORM_OS" = "osx" ]; then + extra="OSX_BUILD_UNIVERSAL=1" + fi + + make -C "$src_dir" \ + GEN=ninja \ + DISABLE_SANITIZER=1 \ + THREADSAN=0 \ + BUILD_SHELL=false \ + BUILD_UNITTESTS=false \ + ENABLE_UNITTEST_CPP_TESTS=false \ + BUILD_EXTENSIONS="parquet;tpch;tpcds" \ + $extra + + lib_dir="${src_dir}/build/release/src" + stage="stage" + mkdir -p "$stage" + + cp -a "${lib_dir}/libduckdb.so" "$stage/" 2>/dev/null || true + cp -a "${lib_dir}/libduckdb.dylib" "$stage/" 2>/dev/null || true + cp -a "${lib_dir}/libduckdb_static.a" "$stage/" + cp -a "${src_dir}/src/include/duckdb.h" "$stage/" 2>/dev/null || true + cp -a "${src_dir}/src/include/duckdb.hpp" "$stage/" 2>/dev/null || true + + ( cd "$stage" && zip -r "../${ARCHIVE}" . ) +fi + +ls -la "$ARCHIVE" diff --git a/scripts/duckdb-r2-resolve.sh b/scripts/duckdb-r2-resolve.sh new file mode 100755 index 00000000000..221c1f2bc8d --- /dev/null +++ b/scripts/duckdb-r2-resolve.sh @@ -0,0 +1,76 @@ +#!/bin/sh + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors +# +# Resolve DuckDB version in vortex-duckdb/build.rs and check which archives +# are present in R2. Output "version", "ref_dir", "release", "matrix", and +# "any_missing". +# +# Required env vars: PUBLIC_BASE_URL, GITHUB_OUTPUT + +set -eu + +version=$(grep -oP 'DEFAULT_DUCKDB_VERSION:\s*&str\s*=\s*"\K[^"]+' \ + vortex-duckdb/build.rs) + +# vortex-duckdb/build.rs: >=2 dot-separated numbers are a +# tagged release (ref dir "vX.Y.Z"), anything else is a commit. +ref=${version#v} +if printf '%s' "$ref" | grep -Eq '^[0-9]+(\.[0-9]+)+$'; then + release=true + ref_dir="v$ref" +else + release=false + ref_dir="$ref" +fi + +echo "DuckDB $version release=$release" + +entries=$(mktemp) +trap 'rm -f "$entries"' EXIT + +for archive in \ + libduckdb-linux-amd64.zip \ + libduckdb-linux-arm64.zip \ + libduckdb-osx-universal.zip; do + + url="${PUBLIC_BASE_URL}/${ref_dir}/${archive}" + code=$(curl -o /dev/null -s -w '%{http_code}' --head "$url" || echo 000) + if [ "$code" = "200" ]; then + echo "present in R2: $archive" + continue + fi + + echo "missing in R2 (HTTP $code): $archive" + case "$archive" in + *linux-amd64*) runner=ubuntu-latest; os=linux; arch=amd64 ;; + *linux-arm64*) runner=ubuntu-24.04-arm; os=linux; arch=arm64 ;; + *osx-universal*) runner=macos-14; os=osx; arch=universal ;; + esac + jq -nc \ + --arg archive "$archive" \ + --arg runner "$runner" \ + --arg os "$os" \ + --arg arch "$arch" \ + '{archive: $archive, runner: $runner, os: $os, arch: $arch}' >> "$entries" +done + +if [ -s "$entries" ]; then + include=$(jq -sc '.' < "$entries") + matrix=$(jq -nc --argjson include "$include" '{include: $include}') + any_missing=true +else + matrix='{"include":[]}' + any_missing=false +fi + +echo "any_missing=$any_missing" + +{ + echo "version=$version" + echo "ref_dir=$ref_dir" + echo "release=$release" + echo "matrix=$matrix" + echo "any_missing=$any_missing" +} >> "$GITHUB_OUTPUT" diff --git a/scripts/post-ingest.py b/scripts/post-ingest.py index 2a72f975d95..e349048f04b 100755 --- a/scripts/post-ingest.py +++ b/scripts/post-ingest.py @@ -7,29 +7,44 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright the Vortex contributors -"""Wrap a `--gh-json-v3` JSONL file in an envelope and POST to /api/ingest. +"""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`, -fills the `commit` envelope by shelling out to `git show`, and POSTs the -envelope to `/api/ingest` with a bearer token. +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: -Standard library only -- urllib, json, subprocess. 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 +- `--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`). +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 -`benchmarks-website/AGENTS.md`): +the `vortex-data/benchmarks-website` repo's `AGENTS.md`): -- `benchmarks-website/server/src/records.rs` - envelope + per-record wire +- 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. -- `benchmarks-website/server/src/schema.rs` - `SCHEMA_VERSION` source of +- 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. """ @@ -37,17 +52,19 @@ import argparse import json +import math import os import subprocess import sys import urllib.error import urllib.request +from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime from pathlib import Path -# MUST equal `benchmarks-website/server/src/schema.rs::SCHEMA_VERSION`. +# 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 `benchmarks-website/AGENTS.md` ("Wire shapes are a +# 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 @@ -59,17 +76,45 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( - description="POST a v3 JSONL records file to /api/ingest.", + description=( + "Ingest a v3 JSONL records file: --server POSTs a v3 envelope to " + "/api/ingest; --postgres upserts directly into the RDS Postgres tables." + ), ) parser.add_argument( "jsonl_path", type=Path, help="Path to the JSONL file written by vortex-bench --gh-json-v3.", ) - parser.add_argument( + mode = parser.add_mutually_exclusive_group(required=True) + mode.add_argument( "--server", - required=True, - help="Server base URL, e.g. http://localhost:8080.", + 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( + "--postgres", + metavar="DSN", + help=( + "v4 mode: 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." + ), + ) + parser.add_argument( + "--region", + default=None, + help=( + "AWS region for RDS IAM token minting (--postgres mode). Precedence: " + "this explicit --region, then the boto3 session region, then the region " + "parsed from the RDS hostname." + ), ) parser.add_argument( "--commit-sha", @@ -78,8 +123,10 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--benchmark-id", - required=True, - help="Run identifier echoed back in run_meta.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", @@ -113,9 +160,11 @@ def parse_args() -> argparse.Namespace: def read_records(path: Path) -> list[dict]: + # The JSONL comes from the project's own `vortex-bench` CI on the same run (trusted UTF-8 + # input); a malformed line still fails loud with `path:line` context for debuggability. records: list[dict] = [] - with path.open("r", encoding="utf-8") as fp: - for line_no, line in enumerate(fp, start=1): + with path.open(encoding="utf-8") as f: + for line_no, line in enumerate(f, start=1): line = line.strip() if not line: continue @@ -127,7 +176,10 @@ def read_records(path: Path) -> list[dict]: def git_show_field(sha: str, fmt: str, cwd: Path | None) -> str: - """Run `git show -s --format= ` and return its stdout (stripped).""" + """Run `git show -s --format= ` and return its stdout (stripped). + + The metadata comes from the project's own (UTF-8) git history. + """ result = subprocess.run( ["git", "show", "-s", f"--format={fmt}", sha], cwd=cwd, @@ -265,9 +317,865 @@ def post(server: str, body: bytes, token: str, timeout: float) -> tuple[int, byt return exc.code, exc.read() -def main() -> int: - args = parse_args() +# -------------------------------------------------------------------------- +# 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. +_RECORD_FIELDS: dict[str, tuple[frozenset[str], frozenset[str]]] = { + "query_measurement": ( + frozenset( + { + "commit_sha", + "dataset", + "query_idx", + "storage", + "engine", + "format", + "value_ns", + "all_runtimes_ns", + } + ), + frozenset( + { + "dataset_variant", + "scale_factor", + "peak_physical", + "peak_virtual", + "physical_delta", + "virtual_delta", + "env_triple", + } + ), + ), + "compression_time": ( + frozenset({"commit_sha", "dataset", "format", "op", "value_ns", "all_runtimes_ns"}), + frozenset({"dataset_variant", "env_triple"}), + ), + "compression_size": ( + frozenset({"commit_sha", "dataset", "format", "value_bytes"}), + frozenset({"dataset_variant"}), + ), + "random_access_time": ( + frozenset({"commit_sha", "dataset", "format", "value_ns", "all_runtimes_ns"}), + frozenset({"env_triple"}), + ), + "vector_search_run": ( + frozenset( + { + "commit_sha", + "dataset", + "layout", + "flavor", + "threshold", + "value_ns", + "all_runtimes_ns", + "matches", + "rows_scanned", + "bytes_scanned", + "iterations", + } + ), + frozenset({"env_triple"}), + ), +} + +_MEASUREMENT_ID_MODULE = None + + +def _measurement_id_module(): + """Lazily load `scripts/_measurement_id.py` by path (cached). + + Loaded by file path rather than `import _measurement_id` so it resolves + regardless of cwd / `sys.path` (the test harness loads sibling scripts the + same way). Importing it pulls in `xxhash`, so it only happens here on the + `--postgres` path. + """ + global _MEASUREMENT_ID_MODULE + if _MEASUREMENT_ID_MODULE is None: + import importlib.util + + path = Path(__file__).resolve().parent / "_measurement_id.py" + spec = importlib.util.spec_from_file_location("_measurement_id", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + _MEASUREMENT_ID_MODULE = module + return _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 + required field is a loud error keyed by the record's index. + """ + if not isinstance(record, dict): + raise SystemExit(f"record {index}: expected a JSON object, got {type(record).__name__}") + kind = record.get("kind") + # `isinstance(kind, str)` first: a non-scalar kind (e.g. a list) is unhashable + # and would raise TypeError at the `in` membership check rather than this + # controlled record-indexed error. + if not isinstance(kind, str) or kind not in _RECORD_FIELDS: + raise SystemExit(f"record {index}: unknown kind {kind!r}; expected one of {sorted(_RECORD_FIELDS)}") + required, optional = _RECORD_FIELDS[kind] + present = set(record) - {"kind"} + unknown = present - required - optional + if unknown: + raise SystemExit(f"record {index} ({kind}): unknown field(s) {sorted(unknown)} (deny_unknown_fields)") + missing = required - present + if missing: + raise SystemExit(f"record {index} ({kind}): missing required field(s) {sorted(missing)}") + return kind + + +def _require_finite(value: object, field: str, kind: str, index: int) -> None: + """Raise loudly unless `value` is a finite number. + + Guards every f64 dim that feeds a `measurement_id` hash (currently only + `vector_search_run.threshold`). A NaN/Inf would hash differently across the + Rust `to_bits()` and Python `struct.pack(' None: + """Raise loudly unless `value` is a plain (non-bool) integer within i32/i64 range. + + 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 + `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. + """ + lo, hi = (_INT32_MIN, _INT32_MAX) if bits == 32 else (_INT64_MIN, _INT64_MAX) + if isinstance(value, bool) or not isinstance(value, int): + raise SystemExit(f"record {index} ({kind}): {field} must be an integer, got {value!r}") + if not (lo <= value <= hi): + raise SystemExit(f"record {index} ({kind}): {field}={value!r} is out of int{bits} range") + + +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 + 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. + """ + 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}") + if any(not (_INT64_MIN <= x <= _INT64_MAX) for x in value): + raise SystemExit(f"record {index} ({kind}): {field} has an element out of int64 range") + + +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.""" + if not isinstance(value, str): + raise SystemExit(f"record {index} ({kind}): {field} must be a string, got {value!r}") + + +def _require_opt_str(value: object, field: str, kind: str, index: int) -> None: + """Raise loudly unless `value` is a string or null. Mirrors `Option`.""" + if value is not None and not isinstance(value, str): + raise SystemExit(f"record {index} ({kind}): {field} must be a string or null, got {value!r}") + + +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. + """ + present = [ + r.get("peak_physical") is not None, + r.get("peak_virtual") is not None, + r.get("physical_delta") is not None, + r.get("virtual_delta") is not None, + ] + 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: +# "str" required String "opt_str" Option +# "i32" required i32 "i64" required i64 +# "opt_i64" Option "i64_list" Vec (all_runtimes_ns) +# "f64" finite f64 (threshold). +# Required tokens (every non-opt_*) line up with `_RECORD_FIELDS` required sets, so +# `record[field]` is always present by the time `_validate_record_values` runs. +_FIELD_TYPES: dict[str, tuple[tuple[str, str], ...]] = { + "query_measurement": ( + ("commit_sha", "str"), + ("dataset", "str"), + ("dataset_variant", "opt_str"), + ("scale_factor", "opt_str"), + ("query_idx", "i32"), + ("storage", "str"), + ("engine", "str"), + ("format", "str"), + ("value_ns", "i64"), + ("all_runtimes_ns", "i64_list"), + ("peak_physical", "opt_i64"), + ("peak_virtual", "opt_i64"), + ("physical_delta", "opt_i64"), + ("virtual_delta", "opt_i64"), + ("env_triple", "opt_str"), + ), + "compression_time": ( + ("commit_sha", "str"), + ("dataset", "str"), + ("dataset_variant", "opt_str"), + ("format", "str"), + ("op", "str"), + ("value_ns", "i64"), + ("all_runtimes_ns", "i64_list"), + ("env_triple", "opt_str"), + ), + "compression_size": ( + ("commit_sha", "str"), + ("dataset", "str"), + ("dataset_variant", "opt_str"), + ("format", "str"), + ("value_bytes", "i64"), + ), + "random_access_time": ( + ("commit_sha", "str"), + ("dataset", "str"), + ("format", "str"), + ("value_ns", "i64"), + ("all_runtimes_ns", "i64_list"), + ("env_triple", "opt_str"), + ), + "vector_search_run": ( + ("commit_sha", "str"), + ("dataset", "str"), + ("layout", "str"), + ("flavor", "str"), + ("threshold", "f64"), + ("value_ns", "i64"), + ("all_runtimes_ns", "i64_list"), + ("matches", "i64"), + ("rows_scanned", "i64"), + ("bytes_scanned", "i64"), + ("iterations", "i32"), + ("env_triple", "opt_str"), + ), +} + + +def _validate_record_values(record: dict, kind: str, index: int) -> None: + """Validate every field's type/range against the v3 server's serde boundary. + + 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). + """ + for field, typ in _FIELD_TYPES[kind]: + if typ == "str": + _require_str(record[field], field, kind, index) + elif typ == "opt_str": + _require_opt_str(record.get(field), field, kind, index) + elif typ == "i32": + _require_int(record[field], field, kind, index, bits=32) + elif typ == "i64": + _require_int(record[field], field, kind, index, bits=64) + elif typ == "opt_i64": + if record.get(field) is not None: + _require_int(record[field], field, kind, index, bits=64) + elif typ == "i64_list": + _require_int_list(record[field], field, kind, index) + elif typ == "f64": + _require_finite(record[field], field, kind, index) + + if kind == "query_measurement": + if record["storage"] not in ("nvme", "s3"): + raise SystemExit( + f"record {index} (query_measurement): storage must be 'nvme' or 's3', got {record['storage']!r}" + ) + if not _memory_quartet_consistent(record): + raise SystemExit( + f"record {index} (query_measurement): memory fields must be populated together (all four or none)" + ) + + +def _upsert_returning_was_update(conn, sql: str, params: tuple) -> bool: + """Run an `INSERT ... ON CONFLICT DO UPDATE ... RETURNING (xmax = 0)` upsert + and return whether the row was an update (vs a fresh insert). + + Classifies inserted-vs-updated atomically from the upsert itself: a fresh + INSERT leaves the new tuple's system column `xmax = 0`, while `ON CONFLICT + DO UPDATE` stamps `xmax` with the current transaction id. A preflight SELECT + (the prior approach) would race the upsert under concurrent re-ingest of the + same `measurement_id`, miscounting a concurrent loser as inserted; deriving + 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 + `RETURNING (xmax = 0) AS inserted`. + """ + row = conn.execute(sql, params).fetchone() + return not row[0] + + +def _insert_query_measurement(conn, mid_mod, r: dict) -> bool: + """Upsert a `query_measurements` row. Mirrors `ingest.rs::insert_query_measurement`. + + Record values are validated by `_validate_record_values` before dispatch. + """ + mid = mid_mod.measurement_id_query( + commit_sha=r["commit_sha"], + dataset=r["dataset"], + dataset_variant=r.get("dataset_variant"), + scale_factor=r.get("scale_factor"), + query_idx=r["query_idx"], + storage=r["storage"], + engine=r["engine"], + format=r["format"], + ) + return _upsert_returning_was_update( + conn, + """ + INSERT INTO query_measurements ( + measurement_id, commit_sha, dataset, dataset_variant, scale_factor, + query_idx, storage, engine, format, + value_ns, all_runtimes_ns, + peak_physical, peak_virtual, physical_delta, virtual_delta, + env_triple, commit_timestamp + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s::bigint[], %s, %s, %s, %s, %s, + (SELECT timestamp FROM commits WHERE commit_sha = %s)) + ON CONFLICT (measurement_id) DO UPDATE SET + commit_sha = excluded.commit_sha, + value_ns = excluded.value_ns, + all_runtimes_ns = excluded.all_runtimes_ns, + peak_physical = excluded.peak_physical, + peak_virtual = excluded.peak_virtual, + physical_delta = excluded.physical_delta, + virtual_delta = excluded.virtual_delta, + env_triple = excluded.env_triple, + commit_timestamp = excluded.commit_timestamp + RETURNING (xmax = 0) AS inserted + """, + ( + mid, + r["commit_sha"], + r["dataset"], + r.get("dataset_variant"), + r.get("scale_factor"), + r["query_idx"], + r["storage"], + r["engine"], + r["format"], + r["value_ns"], + r["all_runtimes_ns"], + r.get("peak_physical"), + r.get("peak_virtual"), + r.get("physical_delta"), + r.get("virtual_delta"), + r.get("env_triple"), + # The denormalized `commit_timestamp` (migration 006) is resolved from the + # `commits` row this same transaction upserted first, so the read path's + # latest-per-series summary never sees a NULL from this writer. + r["commit_sha"], + ), + ) + + +def _insert_compression_time(conn, mid_mod, r: dict) -> bool: + """Upsert a `compression_times` row. Mirrors `ingest.rs::insert_compression_time`.""" + mid = mid_mod.measurement_id_compression_time( + commit_sha=r["commit_sha"], + dataset=r["dataset"], + dataset_variant=r.get("dataset_variant"), + format=r["format"], + op=r["op"], + ) + return _upsert_returning_was_update( + conn, + """ + INSERT INTO compression_times ( + measurement_id, commit_sha, dataset, dataset_variant, + format, op, value_ns, all_runtimes_ns, env_triple + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::bigint[], %s) + ON CONFLICT (measurement_id) DO UPDATE SET + commit_sha = excluded.commit_sha, + value_ns = excluded.value_ns, + all_runtimes_ns = excluded.all_runtimes_ns, + env_triple = excluded.env_triple + RETURNING (xmax = 0) AS inserted + """, + ( + mid, + r["commit_sha"], + r["dataset"], + r.get("dataset_variant"), + r["format"], + r["op"], + r["value_ns"], + r["all_runtimes_ns"], + r.get("env_triple"), + ), + ) + + +def _insert_compression_size(conn, mid_mod, r: dict) -> bool: + """Upsert a `compression_sizes` row. Mirrors `ingest.rs::insert_compression_size`.""" + mid = mid_mod.measurement_id_compression_size( + commit_sha=r["commit_sha"], + dataset=r["dataset"], + dataset_variant=r.get("dataset_variant"), + format=r["format"], + ) + return _upsert_returning_was_update( + conn, + """ + INSERT INTO compression_sizes ( + measurement_id, commit_sha, dataset, dataset_variant, + format, value_bytes + ) VALUES (%s, %s, %s, %s, %s, %s) + ON CONFLICT (measurement_id) DO UPDATE SET + commit_sha = excluded.commit_sha, + value_bytes = excluded.value_bytes + RETURNING (xmax = 0) AS inserted + """, + ( + mid, + r["commit_sha"], + r["dataset"], + r.get("dataset_variant"), + r["format"], + r["value_bytes"], + ), + ) + + +def _insert_random_access(conn, mid_mod, r: dict) -> bool: + """Upsert a `random_access_times` row. Mirrors `ingest.rs::insert_random_access`.""" + mid = mid_mod.measurement_id_random_access( + commit_sha=r["commit_sha"], + dataset=r["dataset"], + format=r["format"], + ) + return _upsert_returning_was_update( + conn, + """ + INSERT INTO random_access_times ( + measurement_id, commit_sha, dataset, format, + value_ns, all_runtimes_ns, env_triple + ) VALUES (%s, %s, %s, %s, %s, %s::bigint[], %s) + ON CONFLICT (measurement_id) DO UPDATE SET + commit_sha = excluded.commit_sha, + value_ns = excluded.value_ns, + all_runtimes_ns = excluded.all_runtimes_ns, + env_triple = excluded.env_triple + RETURNING (xmax = 0) AS inserted + """, + ( + mid, + r["commit_sha"], + r["dataset"], + r["format"], + r["value_ns"], + r["all_runtimes_ns"], + r.get("env_triple"), + ), + ) + +def _insert_vector_search(conn, mid_mod, r: dict) -> bool: + """Upsert a `vector_search_runs` row. Mirrors `ingest.rs::insert_vector_search`. + + `threshold` is validated finite by `_validate_record_values` before dispatch. + """ + threshold = float(r["threshold"]) + mid = mid_mod.measurement_id_vector_search( + commit_sha=r["commit_sha"], + dataset=r["dataset"], + layout=r["layout"], + flavor=r["flavor"], + threshold=threshold, + ) + return _upsert_returning_was_update( + conn, + """ + INSERT INTO vector_search_runs ( + measurement_id, commit_sha, dataset, layout, flavor, threshold, + value_ns, all_runtimes_ns, matches, rows_scanned, bytes_scanned, + iterations, env_triple + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::bigint[], %s, %s, %s, %s, %s) + ON CONFLICT (measurement_id) DO UPDATE SET + commit_sha = excluded.commit_sha, + value_ns = excluded.value_ns, + all_runtimes_ns = excluded.all_runtimes_ns, + matches = excluded.matches, + rows_scanned = excluded.rows_scanned, + bytes_scanned = excluded.bytes_scanned, + iterations = excluded.iterations, + env_triple = excluded.env_triple + RETURNING (xmax = 0) AS inserted + """, + ( + mid, + r["commit_sha"], + r["dataset"], + r["layout"], + r["flavor"], + threshold, + r["value_ns"], + r["all_runtimes_ns"], + r["matches"], + r["rows_scanned"], + r["bytes_scanned"], + r["iterations"], + r.get("env_triple"), + ), + ) + + +# Dispatch from a record's `kind` to its per-table upsert. Keyed identically to +# `_RECORD_FIELDS`; the two maps are wired together when adding a fact table. +_APPLY_RECORD = { + "query_measurement": _insert_query_measurement, + "compression_time": _insert_compression_time, + "compression_size": _insert_compression_size, + "random_access_time": _insert_random_access, + "vector_search_run": _insert_vector_search, +} + + +def _upsert_commit(conn, commit: dict) -> None: + """Upsert the `commits` dim row. Mirrors `ingest.rs::upsert_commit`.""" + conn.execute( + """ + INSERT INTO commits ( + commit_sha, timestamp, message, author_name, author_email, + committer_name, committer_email, tree_sha, url + ) VALUES (%s, %s::timestamptz, %s, %s, %s, %s, %s, %s, %s) + ON CONFLICT (commit_sha) DO UPDATE SET + timestamp = excluded.timestamp, + message = excluded.message, + author_name = excluded.author_name, + author_email = excluded.author_email, + committer_name = excluded.committer_name, + committer_email = excluded.committer_email, + tree_sha = excluded.tree_sha, + url = excluded.url + """, + ( + commit["sha"], + commit["timestamp"], + commit["message"], + commit["author_name"], + commit["author_email"], + commit["committer_name"], + commit["committer_email"], + commit["tree_sha"], + commit["url"], + ), + ) + + +# 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. +_WRITE_CONFLICT_ATTEMPTS = 128 + + +def _retry_write_conflicts(op): + """Retry `op` on a Postgres write conflict, mirroring the v3 server's `retry_write_conflicts`. + + 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 + (`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. + """ + from psycopg import errors as pg_errors + + for attempt in range(1, _WRITE_CONFLICT_ATTEMPTS + 1): + try: + return 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). + 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. + """ + 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). + + 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. + """ + inserted = 0 + updated = 0 + with conn.transaction(): + _upsert_commit(conn, commit) + for idx, record in enumerate(records): + kind = _validate_record_fields(record, idx) + 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}" + ) + _validate_record_values(record, kind, idx) + if _APPLY_RECORD[kind](conn, mid_mod, record): + updated += 1 + else: + inserted += 1 + return inserted, updated + + +def _region_from_host(host: str) -> str | None: + """Parse the AWS region out of an RDS endpoint hostname. + + RDS endpoints look like `...rds.amazonaws.com` (instance) + or `.proxy-..rds.amazonaws.com` (proxy); the region is the + label immediately before `rds.amazonaws.com`. Returns None for any other + shape. + """ + parts = host.split(".") + if len(parts) >= 4 and parts[-3:] == ["rds", "amazonaws", "com"]: + return parts[-4] + return None + + +def _rds_iam_token(*, host: str, port: int, user: str, region: str | None) -> str: + """Mint a short-lived RDS IAM auth token to use as the connection password.""" + import boto3 + + session = boto3.session.Session() + resolved = region or session.region_name or _region_from_host(host) + if not resolved: + raise SystemExit( + "could not determine the AWS region for the RDS IAM token; pass --region or set AWS_DEFAULT_REGION." + ) + client = session.client("rds", region_name=resolved) + return client.generate_db_auth_token(DBHostname=host, Port=port, DBUsername=user, Region=resolved) + + +# The least-privilege login role the v4 CI ingest path must authenticate as +# (created by `migrations/004_ingest_role.sql`; SELECT,INSERT,UPDATE only). +# Phase-2 BAN: do not authenticate the ingest write path as `migrator` / +# `GitHubBenchmarkSchemaRole`. +_INGEST_ROLE = "bench_ingest" + + +def connect_postgres(dsn: str, region: str | None): + """Open a psycopg connection to the RDS Postgres ingest target. + + Enforces the ingest contract: verify-full TLS, and authentication only as the + least-privilege `bench_ingest` role -- always, regardless of auth method (IAM + token or password) or host. This is the production RDS ingest connector; the + test suite exercises ingest via `ingest_postgres` directly, so there is no + local-affordance exception to carve out. A host-based "is this local?" + heuristic would be unreliable anyway because libpq can resolve the host from a + DSN `hostaddr=` or the `$PGHOST` environment variable, either of which could + bypass a DSN-host check. Parses `dsn`; if it carries no password, mints an RDS + IAM auth token for the DSN's user and host (the DSN is expected to also supply + `sslrootcert` for verify-full to validate against). + """ + import psycopg + from psycopg import conninfo + + params = conninfo.conninfo_to_dict(dsn) + + # verify-full is the ingest TLS contract. Default an absent value, but + # refuse a DSN that explicitly downgrades it rather than silently weakening + # the internet-reachable ingest connection. + sslmode = params.get("sslmode") + if sslmode is None: + params["sslmode"] = "verify-full" + elif sslmode != "verify-full": + raise SystemExit( + f"--postgres requires sslmode=verify-full for the RDS ingest path; DSN " + f"specified sslmode={sslmode!r}. Omit it (defaults to verify-full) or set " + f"it to verify-full." + ) + + user = params.get("user") + # Least-privilege: always the bench_ingest role. No host heuristic (see the + # docstring): a misconfigured DSN -- by host=, hostaddr=, or $PGHOST -- must + # never ingest as migrator/postgres. + if user != _INGEST_ROLE: + raise SystemExit(f"--postgres must connect as the least-privilege {_INGEST_ROLE!r} role; DSN user is {user!r}.") + + if not params.get("password"): + # IAM-token path (production CI ingest): mint a token for the DSN's user. + host = params.get("host") + if not host: + raise SystemExit("--postgres DSN must specify host for IAM token minting") + try: + port = int(params.get("port", 5432)) + except (TypeError, ValueError) as exc: + raise SystemExit(f"--postgres DSN has a non-numeric port: {params.get('port')!r}") from exc + params["password"] = _rds_iam_token(host=host, port=port, user=user, region=region) + + # Force `search_path=public` so the writer's unqualified table names always resolve to the + # migration-owned `public.*` tables, regardless of any `search_path` baked into the DSN's + # `options=-c search_path=...` or the role's default. libpq applies repeated `-c` settings + # left-to-right (last wins), so appending ours last makes it authoritative even if the DSN + # already set one. + existing_options = params.get("options") or "" + params["options"] = f"{existing_options} -c search_path=public".strip() + + conn = psycopg.connect(**params) + # Verify the RESOLVED transport actually used TLS, not merely that the DSN requested + # verify-full. The `sslmode` check above rejects an explicit downgrade, but it cannot stop a + # hostless / Unix-socket DSN (host omitted, `host=/...`, or libpq resolving via `$PGHOST`) + # from connecting over a local socket, where libpq silently ignores `sslmode`. Checking + # `ssl_in_use` post-connect validates the connection that actually happened rather than + # trusting the DSN string -- the same "verify the resolved state" reasoning the docstring + # uses to reject a DSN-host heuristic -- and so also closes the `$PGHOST` / `hostaddr` bypass. + # NOTE: `ssl_in_use` lives on the low-level libpq wrapper `conn.pgconn` (a `pq.PGconn`); it is + # NOT on the high-level `conn.info` (`ConnectionInfo`), which exposes only host/dbname/user/etc. + if not conn.pgconn.ssl_in_use: + conn.close() + raise SystemExit( + "--postgres requires a verify-full TLS connection, but the established connection is " + "not using TLS (a hostless or Unix-socket DSN bypasses sslmode); connect to the RDS " + "instance over TCP with sslmode=verify-full." + ) + return conn + + +def _http(method: str, url: str, token: str | None, timeout: float) -> bytes: + """Issue one HTTP request and return the body. Raises on any non-2xx or + transport error; callers in `refresh_site_cache` swallow those.""" + headers = {"accept": "application/json"} + if token is not None: + headers["authorization"] = f"Bearer {token}" + req = urllib.request.Request(url, method=method, headers=headers) + with urllib.request.urlopen(req, timeout=timeout) as resp: + return resp.read() + + +def _warm_default_windows(base: str, timeout: float) -> None: + """Best-effort warm pass: prime the freshly invalidated Data Cache for the + landing page and every group's default last-100 bundle, so the first human + request after an ingest is already hot. Each request is independent; one + failure does not abort the others.""" + + def warm(url: str) -> None: + try: + _http("GET", url, None, timeout) + except Exception as exc: # noqa: BLE001 -- warm is best-effort. + print(f"warning: warm {url} failed: {exc}", file=sys.stderr) + + warm(f"{base}/") + try: + groups_body = _http("GET", f"{base}/api/groups", None, timeout) + slugs = [g["slug"] for g in json.loads(groups_body).get("groups", []) if "slug" in g] + except Exception as exc: # noqa: BLE001 -- group discovery is best-effort. + print(f"warning: warm group discovery failed: {exc}", file=sys.stderr) + return + # A whole-bundle recompute is a few seconds cold, so warm with bounded + # concurrency rather than one slow serial pass. + with ThreadPoolExecutor(max_workers=4) as pool: + pool.map(lambda s: warm(f"{base}/api/group/{s}?n=100"), slugs) + + +def refresh_site_cache(base_url: str, token: str, timeout: float) -> None: + """Revalidate the site's Data Cache tag, then warm the default windows. + + BEST-EFFORT: every failure is logged to stderr and swallowed so a cache + refresh can never change the ingest exit code. The warm pass is skipped + when revalidation fails: warming after a failed flush would repopulate the + Data Cache with stale data, which is the opposite of the intent. + """ + base = base_url.rstrip("/") + try: + _http("POST", f"{base}/api/revalidate", token, timeout) + except Exception as exc: # noqa: BLE001 -- refresh must never raise into ingest. + print(f"warning: cache revalidate failed: {exc}", file=sys.stderr) + return # Skip the warm pass: no point warming a cache that was not flushed. + _warm_default_windows(base, timeout) + + +def _main_postgres(args: argparse.Namespace) -> int: + records = read_records(args.jsonl_path) + # `build_commit` runs `git show `, so the SHA must be in the runner's local git + # history. The v4 ingest step inherits the v3 `--server` step's checkout assumption (the default + # checkout provides the head SHA); a shallow checkout missing the SHA fails loud here, and the + # v4 step is best-effort (continue-on-error), so it never fails the job. + commit = build_commit(args.commit_sha, args.repo_url, args.git_dir) + conn = connect_postgres(args.postgres, args.region) + try: + inserted, updated = ingest_postgres(conn, commit, records) + finally: + conn.close() + print( + json.dumps( + {"records": len(records), "inserted": inserted, "updated": updated}, + 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. + base_url = os.environ.get("BENCH_SITE_BASE_URL") + revalidate_token = os.environ.get("BENCH_REVALIDATE_TOKEN") + if base_url and revalidate_token: + refresh_site_cache(base_url, revalidate_token, args.timeout) + 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( @@ -336,5 +1244,12 @@ def main() -> int: return 0 +def main() -> int: + args = parse_args() + if args.postgres is not None: + return _main_postgres(args) + return _main_server(args) + + if __name__ == "__main__": raise SystemExit(main()) diff --git a/scripts/random-access-split.py b/scripts/random-access-split.py new file mode 100755 index 00000000000..92b44aee7e0 --- /dev/null +++ b/scripts/random-access-split.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +""" +Run random-access-bench once per (dataset, format, pattern, open-mode) +then merge the per-combination outputs +""" + +import argparse +import glob +import json +import subprocess +from collections.abc import Callable +from pathlib import Path + +SCRIPT_DIR = Path(__file__).resolve().parent +BINARY = "target/release_debug/random-access-bench" +PARTS_DIR = Path("parts") + +DATASETS = ["taxi", "feature-vectors", "nested-lists", "nested-structs"] +FORMATS = ["parquet", "lance", "vortex"] +PATTERNS = ["correlated", "uniform"] +OPEN_MODES = ["cached", "reopen"] + + +def run_combinations(emit_v3: 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: + args = [ + "bash", + str(SCRIPT_DIR / "bench-taskset.sh"), + BINARY, + "--datasets", + dataset, + "--formats", + fmt, + "--patterns", + pattern, + "--open-mode", + open_mode, + "-d", + "gh-json", + "-o", + str(PARTS_DIR / f"{i}.gh.json"), + ] + if emit_v3: + args += ["--gh-json-v3", str(PARTS_DIR / f"{i}.v3.jsonl")] + print("+", " ".join(args), flush=True) + subprocess.run(args, check=True) + i += 1 + + +""" +This function exists only because of taxi-legacy. + +Every taxi invocation re-emits the pattern-less legacy taxi rows, so we need +the merge to drop the duplicates. Otherwise we could just merge JSONL lines. +""" + + +def merge(pattern: str, key: Callable[[dict], object], out_path: str) -> None: + seen: set[object] = set() + 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 not line: + continue + identity = key(json.loads(line)) + if identity in seen: + continue + seen.add(identity) + 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( + "--v3", + action="store_true", + help="merge --gh-json-v3 records into results.v3.jsonl", + ) + args = parser.parse_args() + + run_combinations(args.v3) + merge(f"{PARTS_DIR}/*.gh.json", lambda record: record["name"], "results.json") + if args.v3: + merge( + f"{PARTS_DIR}/*.v3.jsonl", + lambda record: (record["kind"], record["dataset"], record["format"]), + "results.v3.jsonl", + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/vortex_gdb.py b/scripts/vortex_gdb.py new file mode 100644 index 00000000000..791bbdd868e --- /dev/null +++ b/scripts/vortex_gdb.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +""" +GDB pretty-printers. Usage: + +(gdb) source scripts/vortex_gdb.py +... +(gdb) print values_idx_offsets +$1 = vortex.primitive(u32, len=465) @ 0x20005ab0d20 [strong=2, weak=0] +(gdb) print values +$2 = fastlanes.bitpacked(i16, len=75133) @ 0x20005ab09a0 [strong=2, weak=0] +(gdb) print encoding_id +$3 = Id("vortex.primitive") +(gdb) print plugin +$4 = dyn ArrayPlugin @ 0xaaaab6deb168 +(gdb) print plugin_arc +$5 = Arc @ 0x20000ab0310> [strong=3, weak=0] +""" + +from __future__ import annotations + +import re +from collections.abc import Callable + +import gdb + +_ARRAY_PLUGIN_FQN = "vortex_array::array::plugin::ArrayPlugin" +_DYN_ARRAY_PLUGIN = "dyn " + _ARRAY_PLUGIN_FQN +_AS_STR_NAME = "vortex_session::registry::Id::as_str" +_ARRAY_REF_FQN = "vortex_array::array::erased::ArrayRef" +_ARRAY_TYPED_PREFIX = "vortex_array::array::typed::Array<" +_ID_FQN = "vortex_session::registry::Id" +_ARRAY_INNER_TY_NAME = "vortex_array::array::typed::ArrayInner" +# "data: T" lives after [strong, weak] inside ArcInner. both are AtomicUsize. +_ARC_INNER_DATA_OFFSET = 16 + +# DType variant -> Display template. "{n}" is nullability suffix. +_DTYPE_TEMPLATES = { + "Null": "null", + "Bool": "bool{n}", + "Utf8": "utf8{n}", + "Binary": "binary{n}", + "Union": "union(){n}", + "Variant": "variant{n}", + "Struct": "{{...}}{n}", + "List": "list(...){n}", + "FixedSizeList": "fixed_size_list(...){n}", + "Decimal": "decimal(...){n}", + "Extension": "ext", +} +_DTYPE_VARIANT_RE = re.compile(r"(?:^|::)DType::(\w+)") +_PTYPE_VARIANT_RE = re.compile(r"(?:^|::)PType::(\w+)") +_NULLABILITY_RE = re.compile(r"(?:^|::)Nullability::(\w+)") + + +def _usize_size() -> int: + return gdb.lookup_type("usize").sizeof + + +def _read_usize(addr: int) -> int: + return int(gdb.Value(addr).cast(gdb.lookup_type("usize").pointer()).dereference()) + + +def _concrete_type(val: gdb.Value, deref_path: str) -> str | None: + try: + gdb.set_convenience_variable("vortex_tmp", val) + return gdb.parse_and_eval(f"*$vortex_tmp{deref_path}").type.name + except (gdb.error, AttributeError): + return None + + +def _trait_fat_pointer(val: gdb.Value) -> tuple[int, int] | None: + # field names vary by rustc + for p in ("pointer", "data_ptr", "__0"): + for v in ("vtable", "__1"): + try: + return int(val[p]), int(val[v]) + except (gdb.error, RuntimeError, KeyError): + continue + try: + fields = [f for f in val.type.fields() if f.name and not f.artificial] + except (gdb.error, RuntimeError): + return None + if len(fields) == 2: + try: + return int(val[fields[0].name]), int(val[fields[1].name]) + except (gdb.error, RuntimeError): + return None + return None + + +def _spur_from_id(val: gdb.Value) -> int | None: + # Id(Spur { key: NonZero(T) }) -> T + cur = val + for _ in range(8): + try: + t = cur.type.strip_typedefs() + except gdb.error: + break + if t.code != gdb.TYPE_CODE_STRUCT: + break + try: + fields = [f for f in t.fields() if f.name and not f.artificial and not f.is_base_class] + except (gdb.error, RuntimeError): + break + if len(fields) != 1: + break + cur = cur[fields[0].name] + try: + return int(cur) + except (gdb.error, ValueError, TypeError): + return None + + +def _str_to_python(s: gdb.Value) -> str | None: + for ptr_n, len_n in (("data_ptr", "length"), ("ptr", "len"), ("__0", "__1")): + try: + data_ptr, length = int(s[ptr_n]), int(s[len_n]) + break + except (gdb.error, RuntimeError): + continue + else: + return None + if length == 0: + return "" + if data_ptr == 0 or length > 4096: + return None + try: + return gdb.selected_inferior().read_memory(data_ptr, length).tobytes().decode("utf-8") + except (gdb.error, UnicodeDecodeError): + return None + + +class _PinnedThread: + """Make a pretty-printer's call safe with multiple threads""" + + def __enter__(self): + try: + self._frame = gdb.selected_frame() + except gdb.error: + self._frame = None + try: + self._prev_lock = gdb.parameter("scheduler-locking") + # only this thread will run during infcall + gdb.execute("set scheduler-locking on", to_string=True) + except gdb.error: + self._prev_lock = None + return self + + def __exit__(self, exc_type, exc, tb): + if self._prev_lock is not None: + try: + gdb.execute(f"set scheduler-locking {self._prev_lock}", to_string=True) + except gdb.error: + pass + if self._frame is not None: + try: + self._frame.select() + except gdb.error: + pass + + +def _id_to_string(id_addr: int) -> str | None: + sym = gdb.lookup_global_symbol(_AS_STR_NAME) or gdb.lookup_static_symbol(_AS_STR_NAME) + if sym is None: + return None + try: + id_ptr = gdb.Value(id_addr).cast(gdb.lookup_type(_ID_FQN).pointer()) + with _PinnedThread(): + str_val = sym.value()(id_ptr) + except gdb.error: + return None + return _str_to_python(str_val) + + +def _format_dtype(dtype_val: gdb.Value) -> str: + s = str(dtype_val).strip() + m = _DTYPE_VARIANT_RE.search(s) + if m is None: + return s + variant = m.group(1) + nm = _NULLABILITY_RE.search(s) + n = "?" if nm is not None and nm.group(1) == "Nullable" else "" + if variant == "Primitive": + pm = _PTYPE_VARIANT_RE.search(s) + return f"{pm.group(1).lower() if pm else '?'}{n}" + tmpl = _DTYPE_TEMPLATES.get(variant) + return tmpl.format(n=n) if tmpl is not None else f"{variant}{n}" + + +def _array_inner_at(addr: int) -> gdb.Value | None: + try: + return gdb.Value(addr).cast(gdb.lookup_type(_ARRAY_INNER_TY_NAME).pointer()).dereference() + except gdb.error: + return None + + +def _format_id(val: gdb.Value) -> str: + addr = val.address + if addr is not None: + text = _id_to_string(int(addr)) + if text is not None: + return f'Id("{text}")' + spur = _spur_from_id(val) + return f"Id({spur})" if spur is not None else "Id(?)" + + +def _format_dyn_plugin(val: gdb.Value) -> str: + ty = _concrete_type(val, "") or "?" + fp = _trait_fat_pointer(val) + addr = f"{fp[0]:#x}" if fp else "?" + return f"dyn ArrayPlugin<{ty}> @ {addr}" + + +def _format_arc_dyn_plugin(val: gdb.Value) -> str: + ty = _concrete_type(val, ".ptr.pointer") or "?" + try: + fp = _trait_fat_pointer(val["ptr"]["pointer"]) + except gdb.error: + fp = None + if fp is None: + return f"Arc>" + arc_inner_base, vtable_ptr = fp + usize = _usize_size() + try: + strong = _read_usize(arc_inner_base) + # ArcInner.weak = Arc::weak_count() + 1 + weak = _read_usize(arc_inner_base + usize) - 1 + align = _read_usize(vtable_ptr + 2 * usize) or usize + except gdb.error: + return f"Arc>" + data_offset = (2 * usize + align - 1) & ~(align - 1) + return f"Arc @ {arc_inner_base + data_offset:#x}> [strong={strong}, weak={weak}]" + + +def _format_array_ref(val: gdb.Value) -> str: + try: + arc = val[val.type.fields()[0].name] + fat_ptr = arc["ptr"]["pointer"] + except (gdb.error, RuntimeError, IndexError): + return "ArrayRef(?)" + + fp = _trait_fat_pointer(fat_ptr) or (int(fat_ptr), None) + arc_inner_base = fp[0] + inner = _array_inner_at(arc_inner_base + _ARC_INNER_DATA_OFFSET) + if inner is None: + return f"ArrayRef(?) @ {arc_inner_base:#x}" + + def _try(fn, default): + try: + return fn() + except (gdb.error, RuntimeError): + return default + + len_ = _try(lambda: int(inner["len"]), None) + enc = inner["encoding_id"] + encoding = (_id_to_string(int(enc.address)) if enc.address else None) or "?" + dtype = _try(lambda: _format_dtype(inner["dtype"]), "?") + + head = f"{encoding}({dtype}, len={'?' if len_ is None else len_})" + usize = _usize_size() + try: + strong = _read_usize(arc_inner_base) + weak = _read_usize(arc_inner_base + usize) - 1 + return f"{head} @ {arc_inner_base:#x} [strong={strong}, weak={weak}]" + except gdb.error: + return f"{head} @ {arc_inner_base:#x}" + + +def _format_typed_array(val: gdb.Value) -> str: + try: + return _format_array_ref(val["inner"]) + except (gdb.error, RuntimeError): + return "Array(?)" + + +class _Printer: + def __init__(self, val: gdb.Value, fmt: Callable[[gdb.Value], str]): + self._val, self._fmt = val, fmt + + def to_string(self) -> str: + return self._fmt(self._val) + + +def _matches_dyn_plugin(name: str) -> bool: + return ( + name in (_ARRAY_PLUGIN_FQN, _DYN_ARRAY_PLUGIN) + or name.startswith(_DYN_ARRAY_PLUGIN + " +") + or name.startswith(_DYN_ARRAY_PLUGIN + "+") + ) + + +def _lookup_printer(val: gdb.Value) -> _Printer | None: + try: + t = val.type.strip_typedefs() + except gdb.error: + return None + name = (t.name or "").strip() + + if name == _ARRAY_REF_FQN: + return _Printer(val, _format_array_ref) + if name.startswith(_ARRAY_TYPED_PREFIX): + return _Printer(val, _format_typed_array) + if name == _ID_FQN: + return _Printer(val, _format_id) + if name.startswith("alloc::sync::Arc<") and _DYN_ARRAY_PLUGIN in name: + return _Printer(val, _format_arc_dyn_plugin) + if _matches_dyn_plugin(name): + return _Printer(val, _format_dyn_plugin) + if t.code in (gdb.TYPE_CODE_PTR, gdb.TYPE_CODE_REF): + target = t.target() + if target is not None and _matches_dyn_plugin((target.name or "").strip()): + return _Printer(val, _format_dyn_plugin) + return None + + +def _install_at_front(container) -> None: + if _lookup_printer in container: + container.remove(_lookup_printer) + container.insert(0, _lookup_printer) + + +def register() -> None: + # Install before rust-gdb's StdRcProvider, otherwise Arc from StdRc wins + _install_at_front(gdb.pretty_printers) + try: + _install_at_front(gdb.current_progspace().pretty_printers) + except (AttributeError, gdb.error): + pass + for of in gdb.objfiles(): + _install_at_front(of.pretty_printers) + gdb.events.new_objfile.connect(lambda e: _install_at_front(e.new_objfile.pretty_printers)) + + +register() +print("vortex_gdb: registered pretty-printers") diff --git a/uv.lock b/uv.lock index ef6e5670aaf..8748a91544e 100644 --- a/uv.lock +++ b/uv.lock @@ -286,28 +286,38 @@ wheels = [ [[package]] name = "duckdb" -version = "1.4.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ea/e7/21cf50a3d52ffceee1f0bcc3997fa96a5062e6bab705baee4f6c4e33cce5/duckdb-1.4.1.tar.gz", hash = "sha256:f903882f045d057ebccad12ac69975952832edfe133697694854bb784b8d6c76", size = 18461687, upload-time = "2025-10-07T10:37:28.605Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d9/52/606f13fa9669a24166d2fe523e28982d8ef9039874b4de774255c7806d1f/duckdb-1.4.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:605d563c1d5203ca992497cd33fb386ac3d533deca970f9dcf539f62a34e22a9", size = 29065894, upload-time = "2025-10-07T10:36:29.837Z" }, - { url = "https://files.pythonhosted.org/packages/84/57/138241952ece868b9577e607858466315bed1739e1fbb47205df4dfdfd88/duckdb-1.4.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d3305c7c4b70336171de7adfdb50431f23671c000f11839b580c4201d9ce6ef5", size = 16163720, upload-time = "2025-10-07T10:36:32.241Z" }, - { url = "https://files.pythonhosted.org/packages/a3/81/afa3a0a78498a6f4acfea75c48a70c5082032d9ac87822713d7c2d164af1/duckdb-1.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:a063d6febbe34b32f1ad2e68822db4d0e4b1102036f49aaeeb22b844427a75df", size = 13756223, upload-time = "2025-10-07T10:36:34.673Z" }, - { url = "https://files.pythonhosted.org/packages/47/dd/5f6064fbd9248e37a3e806a244f81e0390ab8f989d231b584fb954f257fc/duckdb-1.4.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d1ffcaaf74f7d1df3684b54685cbf8d3ce732781c541def8e1ced304859733ae", size = 18487022, upload-time = "2025-10-07T10:36:36.759Z" }, - { url = "https://files.pythonhosted.org/packages/a1/10/b54969a1c42fd9344ad39228d671faceb8aa9f144b67cd9531a63551757f/duckdb-1.4.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:685d3d1599dc08160e0fa0cf09e93ac4ff8b8ed399cb69f8b5391cd46b5b207c", size = 20491004, upload-time = "2025-10-07T10:36:39.318Z" }, - { url = "https://files.pythonhosted.org/packages/ed/d5/7332ae8f804869a4e895937821b776199a283f8d9fc775fd3ae5a0558099/duckdb-1.4.1-cp311-cp311-win_amd64.whl", hash = "sha256:78f1d28a15ae73bd449c43f80233732adffa49be1840a32de8f1a6bb5b286764", size = 12327619, upload-time = "2025-10-07T10:36:41.509Z" }, - { url = "https://files.pythonhosted.org/packages/0e/6c/906a3fe41cd247b5638866fc1245226b528de196588802d4df4df1e6e819/duckdb-1.4.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:cd1765a7d180b7482874586859fc23bc9969d7d6c96ced83b245e6c6f49cde7f", size = 29076820, upload-time = "2025-10-07T10:36:43.782Z" }, - { url = "https://files.pythonhosted.org/packages/66/c7/01dd33083f01f618c2a29f6dd068baf16945b8cbdb132929d3766610bbbb/duckdb-1.4.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:8ed7a86725185470953410823762956606693c0813bb64e09c7d44dbd9253a64", size = 16167558, upload-time = "2025-10-07T10:36:46.003Z" }, - { url = "https://files.pythonhosted.org/packages/81/e2/f983b4b7ae1dfbdd2792dd31dee9a0d35f88554452cbfc6c9d65e22fdfa9/duckdb-1.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8a189bdfc64cfb9cc1adfbe4f2dcfde0a4992ec08505ad8ce33c886e4813f0bf", size = 13762226, upload-time = "2025-10-07T10:36:48.55Z" }, - { url = "https://files.pythonhosted.org/packages/ed/34/fb69a7be19b90f573b3cc890961be7b11870b77514769655657514f10a98/duckdb-1.4.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a9090089b6486f7319c92acdeed8acda022d4374032d78a465956f50fc52fabf", size = 18500901, upload-time = "2025-10-07T10:36:52.445Z" }, - { url = "https://files.pythonhosted.org/packages/e4/a5/1395d7b49d5589e85da9a9d7ffd8b50364c9d159c2807bef72d547f0ad1e/duckdb-1.4.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:142552ea3e768048e0e8c832077a545ca07792631c59edaee925e3e67401c2a0", size = 20514177, upload-time = "2025-10-07T10:36:55.358Z" }, - { url = "https://files.pythonhosted.org/packages/c0/21/08f10706d30252753349ec545833fc0cea67c11abd0b5223acf2827f1056/duckdb-1.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:567f3b3a785a9e8650612461893c49ca799661d2345a6024dda48324ece89ded", size = 12336422, upload-time = "2025-10-07T10:36:57.521Z" }, - { url = "https://files.pythonhosted.org/packages/d7/08/705988c33e38665c969f7876b3ca4328be578554aa7e3dc0f34158da3e64/duckdb-1.4.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:46496a2518752ae0c6c5d75d4cdecf56ea23dd098746391176dd8e42cf157791", size = 29077070, upload-time = "2025-10-07T10:36:59.83Z" }, - { url = "https://files.pythonhosted.org/packages/99/c5/7c9165f1e6b9069441bcda4da1e19382d4a2357783d37ff9ae238c5c41ac/duckdb-1.4.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:1c65ae7e9b541cea07d8075343bcfebdecc29a3c0481aa6078ee63d51951cfcd", size = 16167506, upload-time = "2025-10-07T10:37:02.24Z" }, - { url = "https://files.pythonhosted.org/packages/38/46/267f4a570a0ee3ae6871ddc03435f9942884284e22a7ba9b7cb252ee69b6/duckdb-1.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:598d1a314e34b65d9399ddd066ccce1eeab6a60a2ef5885a84ce5ed62dbaf729", size = 13762330, upload-time = "2025-10-07T10:37:04.581Z" }, - { url = "https://files.pythonhosted.org/packages/15/7b/c4f272a40c36d82df20937d93a1780eb39ab0107fe42b62cba889151eab9/duckdb-1.4.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e2f16b8def782d484a9f035fc422bb6f06941ed0054b4511ddcdc514a7fb6a75", size = 18504687, upload-time = "2025-10-07T10:37:06.991Z" }, - { url = "https://files.pythonhosted.org/packages/17/fc/9b958751f0116d7b0406406b07fa6f5a10c22d699be27826d0b896f9bf51/duckdb-1.4.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a5a7d0aed068a5c33622a8848857947cab5cfb3f2a315b1251849bac2c74c492", size = 20513823, upload-time = "2025-10-07T10:37:09.349Z" }, - { url = "https://files.pythonhosted.org/packages/30/79/4f544d73fcc0513b71296cb3ebb28a227d22e80dec27204977039b9fa875/duckdb-1.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:280fd663dacdd12bb3c3bf41f3e5b2e5b95e00b88120afabb8b8befa5f335c6f", size = 12336460, upload-time = "2025-10-07T10:37:12.154Z" }, +version = "1.5.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/69/00/d579dcb2a536b6ea3a2563cdad6844f77d81a9b2d4b22a858097f2468acf/duckdb-1.5.3.tar.gz", hash = "sha256:df39428eb130faa35ae96fd35245bdeae6ecf43936250b116b5fead568eb9f16", size = 18026640, upload-time = "2026-05-20T11:55:31.901Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3c/fc/a8a89c6c73f31c2b58c6abbc2f543e0b736042dd5ef7cc1784c24ec31428/duckdb-1.5.3-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:341a2672e2551ba51c95c1898f0ade983e76675e79038ccb16342c3d6cfb82d7", size = 32583465, upload-time = "2026-05-20T11:54:13.132Z" }, + { url = "https://files.pythonhosted.org/packages/63/f1/3423a2f523dd034e505d4a5dd8e210ae577212e152598dc13b6a5e736e1b/duckdb-1.5.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:c9e8fa408705081160ede7ead238d16e73a36b8561b700f2bf2d650ae48e7b92", size = 17278520, upload-time = "2026-05-20T11:54:16.368Z" }, + { url = "https://files.pythonhosted.org/packages/e1/1a/7bf5ba1b7ea520557e6b2dbee1c85abab016bdac0c1779d9d0ef76c87300/duckdb-1.5.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:70a18f932cf6d87bd0e554613657a515c1443a1724aacfc7ec5137dd28698b03", size = 15424794, upload-time = "2026-05-20T11:54:19.891Z" }, + { url = "https://files.pythonhosted.org/packages/ad/16/ce4b1e386e45fab0268edbf1b85bace20e9437589e9edb2bd5f9a226fa44/duckdb-1.5.3-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e80eb4d0fb59869cb2c7d7ef494c07fb92014fe8e77d96c170cd1ebc1488a708", size = 19306666, upload-time = "2026-05-20T11:54:22.77Z" }, + { url = "https://files.pythonhosted.org/packages/99/1f/651f8453f26931e8061b7e27b3090f868868185814ecb9216d0bd71ec8ef/duckdb-1.5.3-cp311-cp311-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3248b49cd835ea322574bc6aac0ae7a83be85547f49d4f5f5777cb380ee6627f", size = 21418306, upload-time = "2026-05-20T11:54:25.616Z" }, + { url = "https://files.pythonhosted.org/packages/bc/64/e1ffebf010b1631a6fef8d1508f46d4eab3e97c18729af986bb796fa8452/duckdb-1.5.3-cp311-cp311-win_amd64.whl", hash = "sha256:f4eff89c12c3a362efa012262e57b7b4ab904a7f79bad9178fe365510077abe8", size = 13101423, upload-time = "2026-05-20T11:54:28.107Z" }, + { url = "https://files.pythonhosted.org/packages/e7/42/b1d4e34f9658cc0e13d7aae581ab82643f50a548d5aee8767f0c587cc3a4/duckdb-1.5.3-cp311-cp311-win_arm64.whl", hash = "sha256:75d13308c9da3ee431d1e72b8ab720aa74a1b3e9159d4124cb62435924496334", size = 13951740, upload-time = "2026-05-20T11:54:30.886Z" }, + { url = "https://files.pythonhosted.org/packages/e7/c4/2e34929b16c8d544ef664fad8f7f3a2a9db05746aae1e7c8c4ee3a8b23e4/duckdb-1.5.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:ff11a457258148337ef9a392148a8cdbd1069b6c27c21958816c7b67fe6c542d", size = 32626494, upload-time = "2026-05-20T11:54:33.738Z" }, + { url = "https://files.pythonhosted.org/packages/3a/53/3af681793d03771365ae3e2215331151c196a3ac8193f613344840694671/duckdb-1.5.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:5fd25f533cb1b6b2c84cc767a9a9bab7769bb1aa44571a2a0bfc91ac3e4a38ac", size = 17301121, upload-time = "2026-05-20T11:54:36.928Z" }, + { url = "https://files.pythonhosted.org/packages/15/e2/c80af1eac2ab5d35fc2c372ef0a84668842e549fbbf7799277b3fccf3e39/duckdb-1.5.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:10960400ed60cdf0fe05bab2086fa8eb733889cb0ceca18d07ff9a00c0e0be7b", size = 15449283, upload-time = "2026-05-20T11:54:39.777Z" }, + { url = "https://files.pythonhosted.org/packages/2d/9a/c63af233c9f761bf5178a5210437e1bc6bcb30fa8a9073de6398cfb12c03/duckdb-1.5.3-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c5f18e7561403054433706c187589e86629a7af09a7efc23a06a8b308e6acc68", size = 19332762, upload-time = "2026-05-20T11:54:42.51Z" }, + { url = "https://files.pythonhosted.org/packages/21/cc/2d77af4fff86012f334ef82e6d54a995a86c8745e58074f1218ed7d25171/duckdb-1.5.3-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9fb7516255a8764545e30f7efacea408cc847764a3027b3b0b3e7d1a7bebbc5c", size = 21453290, upload-time = "2026-05-20T11:54:45.272Z" }, + { url = "https://files.pythonhosted.org/packages/8d/5e/9bc4817a98feb4dab83e56f2245cd3a30d00ee646d4dec7926464e2b3f28/duckdb-1.5.3-cp312-cp312-win_amd64.whl", hash = "sha256:8001eccbc28be244dfd04d708526f34ddd6460b47a8aeb5d0e39d6f7f9e3fe15", size = 13118308, upload-time = "2026-05-20T11:54:48.058Z" }, + { url = "https://files.pythonhosted.org/packages/81/35/e3f32e4e53e2450ddb1db8312a17d1ce455d60cc4941b6ad2cfc908794b0/duckdb-1.5.3-cp312-cp312-win_arm64.whl", hash = "sha256:6d2835e39bb6af73891f73c0f8d4324f98afe00d0b00c6d34b2a582c2256cbb0", size = 13927187, upload-time = "2026-05-20T11:54:50.584Z" }, + { url = "https://files.pythonhosted.org/packages/cc/9c/a528eb09d8be51954c485864bd06753e616939a080cbc3dd4417e8c94a57/duckdb-1.5.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:e75a6122c12579a99848517f6f00a4e342aebda3590c30fe9b5cc5f39d5e6afc", size = 32626254, upload-time = "2026-05-20T11:54:53.65Z" }, + { url = "https://files.pythonhosted.org/packages/ec/3c/1534c0a6db347c05eb7d0f6ecfb7aefbe74cbff398e4892a8fd1903a20e8/duckdb-1.5.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:fd3963c1cb9d9567777f4a898a9dbe388a2fe9724681801b1e7d6d93eecf1b76", size = 17300917, upload-time = "2026-05-20T11:54:56.628Z" }, + { url = "https://files.pythonhosted.org/packages/23/fa/beafb91e6e152d2161c4a9cbc472334c87607eb61ad7104b5a7fa8d8d7b1/duckdb-1.5.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:3d5db8c0b55e072cf437948ebb5d7e23d7b9d03d905fa5f9145583e65aa447f7", size = 15449411, upload-time = "2026-05-20T11:54:59.089Z" }, + { url = "https://files.pythonhosted.org/packages/50/0a/49b6fe04e2fcd63729eb607dadd44818dde77342a4f5ce086c6c92f1dd4d/duckdb-1.5.3-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ce80aed7a538422129a57eaca9141e3afb51f8bf562b1908b1576c9725b5b22", size = 19333120, upload-time = "2026-05-20T11:55:01.727Z" }, + { url = "https://files.pythonhosted.org/packages/63/4c/0907c3f76adb9dd90e67610b31e0304a35814e65c4c41a354a262c09b885/duckdb-1.5.3-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:787df63824f07bf18022dbc3b8ca4b2bfab0ebe616464f55c6e8cd0f59ea762e", size = 21453266, upload-time = "2026-05-20T11:55:04.5Z" }, + { url = "https://files.pythonhosted.org/packages/6d/9c/d2f23a7803ddbbd9413f7572ecf66a15120ed5ced7ce5c73e698c1406b76/duckdb-1.5.3-cp313-cp313-win_amd64.whl", hash = "sha256:bb5bb5dcdd09d62ee60f0ddbbef918e71cce304ffe28428b1131949d39ffaabf", size = 13118640, upload-time = "2026-05-20T11:55:07.389Z" }, + { url = "https://files.pythonhosted.org/packages/27/d5/7ba2316415bcdab6edd765bbbe35c2ca8a3800f2fe695cd70e3cdb997f09/duckdb-1.5.3-cp313-cp313-win_arm64.whl", hash = "sha256:2fa17ecdd5d3db122836cb71bb93601c2106a3be883c17dffddc02fbf3fa7888", size = 13926409, upload-time = "2026-05-20T11:55:10.166Z" }, + { url = "https://files.pythonhosted.org/packages/a5/c2/d4b6f8a5e4d3bc25773be6da76a99d9661ebbf3552c007c460d2dd59dbf8/duckdb-1.5.3-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:4bfa9a4dadf71e83e2c4eaca2f9421c82a54defecc1b0b4c0be95e2389dec4fe", size = 32636685, upload-time = "2026-05-20T11:55:13.158Z" }, + { url = "https://files.pythonhosted.org/packages/42/58/e835c8298979d29db7a62cb5acc29e9b57aeaca7cdde2fcd3ac980f5cb18/duckdb-1.5.3-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:aea7baf67ad7e1829ac76f67d7dcbd7fb1f57c3eb179d55ac30952df4709ae30", size = 17308134, upload-time = "2026-05-20T11:55:16.194Z" }, + { url = "https://files.pythonhosted.org/packages/c9/46/617b51363f5613418c8b224b3cce16b58e6dde80904566bec232579c1d4e/duckdb-1.5.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:0b0b4f088a65d77e1217ce5d7eff889e63fedc44281200d899ff47c84d8ff836", size = 15449891, upload-time = "2026-05-20T11:55:18.687Z" }, + { url = "https://files.pythonhosted.org/packages/b3/72/354146656e8d9ba3853d3a5ee80a481b8c5f70edfc3d5ae80a8c4479c967/duckdb-1.5.3-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fe8d0c1f6a120aa03fa6e0d03897c71a1842e6cf7afd31d181348391f7108fe1", size = 19338499, upload-time = "2026-05-20T11:55:21.34Z" }, + { url = "https://files.pythonhosted.org/packages/56/8f/65fc623b51448f2bfba1a9ec6ab3debb4664c0876c0113a5e782600b53ac/duckdb-1.5.3-cp314-cp314-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d0405eae18ec6e8210a471c97dbfe87a7e4d605274b7fe572a1f276e92158f13", size = 21455828, upload-time = "2026-05-20T11:55:23.847Z" }, + { url = "https://files.pythonhosted.org/packages/2b/db/d0274cbe9f5fe219f77c0bdf900ac77103569e83c102a4225ce04cbc607d/duckdb-1.5.3-cp314-cp314-win_amd64.whl", hash = "sha256:33ae08b3e818d7613d8936744b67718c2062c2f530376895bfd89efb51b81538", size = 13640011, upload-time = "2026-05-20T11:55:26.276Z" }, + { url = "https://files.pythonhosted.org/packages/07/5d/8f1899b8bef291caf953992fcd6c24df9f29387a35645e58c2504a5ca473/duckdb-1.5.3-cp314-cp314-win_arm64.whl", hash = "sha256:746433e49bbc667b4df283153415fbe37e9083e0eff6c3cd6e54de7536869cd4", size = 14411554, upload-time = "2026-05-20T11:55:29.037Z" }, ] [[package]] @@ -942,11 +952,11 @@ wheels = [ [[package]] name = "pip" -version = "26.1.1" +version = "26.1.2" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b6/48/cb9b7a682f6fe01a4221e1728941dd4ac3cd9090a17db3779d6ff490b602/pip-26.1.1.tar.gz", hash = "sha256:d36762751d156a4ee895de8af39aa0abeeeb577f93a2eca6ab62467bbf0f8a78", size = 1840400, upload-time = "2026-05-04T19:02:21.248Z" } +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" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3a/eb/fea4d1d51c49832120f7f285d07306db3960f423a2612c6057caf3e8196f/pip-26.1.1-py3-none-any.whl", hash = "sha256:99cb1c2899893b075ff56e4ed0af55669a955b49ad7fb8d8603ecdaf4ed653fb", size = 1812777, upload-time = "2026-05-04T19:02:18.9Z" }, + { 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" }, ] [[package]] @@ -1761,15 +1771,15 @@ wheels = [ [[package]] name = "starlette" -version = "1.2.1" +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/25/44/ec35f1b6e83094b997da438a02c8c9b0ade2b1e84cfc48bd4656780760a6/starlette-1.2.1.tar.gz", hash = "sha256:9b9b5ebb992e67d6093741e63c2f59e4f6fff986f81163c087867bd7b924b3f6", size = 2701854, upload-time = "2026-05-31T01:07:51.847Z" } +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/1c/54/196d0c1db10af76baa4f64894448505d60d3cdf70ef92cbb35f46a4e4c71/starlette-1.2.1-py3-none-any.whl", hash = "sha256:4de0082d08c8f6764a85a54cf1120d6939507a19905c7768acad2a9f875d2b89", size = 73350, upload-time = "2026-05-31T01:07:50.09Z" }, + { 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]] @@ -1889,6 +1899,9 @@ dependencies = [ ] [package.optional-dependencies] +cuda = [ + { name = "vortex-data-cuda" }, +] duckdb = [ { name = "duckdb" }, ] @@ -1931,8 +1944,9 @@ requires-dist = [ { name = "ray", marker = "extra == 'ray'", specifier = ">=2.48" }, { name = "substrait", specifier = ">=0.23.0,<0.29.0" }, { name = "typing-extensions", specifier = ">=4.5.0" }, + { name = "vortex-data-cuda", marker = "extra == 'cuda'", directory = "vortex-python-cuda" }, ] -provides-extras = ["duckdb", "numpy", "pandas", "polars", "ray"] +provides-extras = ["cuda", "duckdb", "numpy", "pandas", "polars", "ray"] [package.metadata.requires-dev] dev = [ @@ -1950,6 +1964,16 @@ dev = [ { name = "ray", specifier = ">=2.48" }, ] +[[package]] +name = "vortex-data-cuda" +source = { directory = "vortex-python-cuda" } +dependencies = [ + { name = "vortex-data" }, +] + +[package.metadata] +requires-dist = [{ name = "vortex-data", specifier = "==0.1.0" }] + [[package]] name = "watchfiles" version = "1.1.0" diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 4ec8a83575e..98de7caf95d 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -70,6 +70,7 @@ tracing = { workspace = true } uuid = { workspace = true } vortex-array-macros = { workspace = true } vortex-buffer = { workspace = true, features = ["arrow"] } +vortex-compute = { workspace = true } vortex-error = { workspace = true, features = ["flatbuffers"] } vortex-flatbuffers = { workspace = true, features = ["array", "dtype"] } vortex-mask = { workspace = true } @@ -113,6 +114,10 @@ harness = false name = "cast_primitive" harness = false +[[bench]] +name = "cast_decimal" +harness = false + [[bench]] name = "search_sorted" harness = false @@ -125,6 +130,14 @@ harness = false name = "compare" harness = false +[[bench]] +name = "binary_ops" +harness = false + +[[bench]] +name = "kleene_bool" +harness = false + [[bench]] name = "interleave" harness = false @@ -223,6 +236,10 @@ harness = false name = "filter_bool" harness = false +[[bench]] +name = "list_length" +harness = false + [[bench]] name = "listview_rebuild" harness = false diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index b067314c1d9..fdc43dbcabd 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -4,18 +4,19 @@ #![expect(clippy::cast_possible_truncation)] #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::ArrayRef; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::DynGroupedAccumulator; -use vortex_array::aggregate_fn::EmptyOptions; 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::arrays::ListViewArray; @@ -24,11 +25,14 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + const GROUP_COUNT: usize = 128; const GROUP_SIZE_SEED: u64 = 42; const MIN_VALUES_PER_GROUP: usize = 1; @@ -149,11 +153,15 @@ fn list_element_dtype(list_view: &ArrayRef) -> DType { fn grouped_accumulator(list_view: &ArrayRef, vtable: V) -> ArrayRef where - V: AggregateFnVTable + Clone, + V: AggregateFnVTable + Clone, { - let mut acc = - GroupedAccumulator::try_new(vtable, EmptyOptions, list_element_dtype(list_view)).unwrap(); - acc.accumulate_list(list_view, &mut LEGACY_SESSION.create_execution_ctx()) + let mut acc = GroupedAccumulator::try_new( + vtable, + NumericalAggregateOpts::default(), + list_element_dtype(list_view), + ) + .unwrap(); + acc.accumulate_list(list_view, &mut SESSION.create_execution_ctx()) .unwrap(); divan::black_box(acc.finish().unwrap()) } diff --git a/vortex-array/benches/aggregate_max.rs b/vortex-array/benches/aggregate_max.rs index 14bf13792dd..4da7d1fc653 100644 --- a/vortex-array/benches/aggregate_max.rs +++ b/vortex-array/benches/aggregate_max.rs @@ -8,7 +8,6 @@ use rand::prelude::*; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; -use vortex_array::session::ArraySession; use vortex_session::VortexSession; fn main() { @@ -17,8 +16,7 @@ fn main() { const N: usize = 100_000; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[divan::bench] fn max_i32(bencher: Bencher) { diff --git a/vortex-array/benches/aggregate_sum.rs b/vortex-array/benches/aggregate_sum.rs index b076bb80b5f..81402b24750 100644 --- a/vortex-array/benches/aggregate_sum.rs +++ b/vortex-array/benches/aggregate_sum.rs @@ -9,7 +9,6 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::expr::stats::Stat; -use vortex_array::session::ArraySession; use vortex_session::VortexSession; fn main() { @@ -18,8 +17,7 @@ fn main() { const N: usize = 100_000; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[divan::bench] fn sum_i32(bencher: Bencher) { diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs new file mode 100644 index 00000000000..98c69413193 --- /dev/null +++ b/vortex-array/benches/binary_ops.rs @@ -0,0 +1,247 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] +#![expect( + clippy::cast_possible_truncation, + reason = "benchmark fixtures use indices that fit in the chosen widths" +)] + +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use vortex_array::ArrayRef; +use vortex_array::Executable; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +const LEN: usize = 65_536; + +#[divan::bench] +fn add_i64_nonnull(bencher: Bencher) { + let lhs = primitive_nonnull(0).into_array(); + let rhs = primitive_nonnull(1_000_000).into_array(); + + bench_primitive(bencher, lhs, rhs, Operator::Add); +} + +#[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(); + + bench_primitive(bencher, lhs, rhs, Operator::Add); +} + +#[divan::bench] +fn mul_i64_nonnull(bencher: Bencher) { + let lhs = primitive_small_nonnull(1).into_array(); + let rhs = primitive_small_nonnull(17).into_array(); + + bench_primitive(bencher, lhs, rhs, Operator::Mul); +} + +#[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(); + + bench_primitive(bencher, lhs, rhs, Operator::Mul); +} + +#[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(); + + bench_primitive(bencher, lhs, rhs, Operator::Mul); +} + +#[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(); + + bench_primitive(bencher, lhs, rhs, Operator::Mul); +} + +#[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(); + + bench_primitive(bencher, lhs, rhs, Operator::Mul); +} + +#[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(); + + bench_primitive(bencher, lhs, rhs, Operator::Mul); +} + +#[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(); + + bench_primitive(bencher, lhs, rhs, Operator::Mul); +} + +#[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(); + + bench_primitive(bencher, lhs, rhs, Operator::Mul); +} + +#[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(); + + bench_primitive(bencher, lhs, rhs, Operator::Mul); +} + +#[divan::bench] +fn div_i64_nonnull(bencher: Bencher) { + let lhs = primitive_nonnull(1_000_000).into_array(); + let rhs = primitive_nonzero().into_array(); + + bench_primitive(bencher, lhs, rhs, Operator::Div); +} + +#[divan::bench] +fn sub_i64_constant(bencher: Bencher) { + let lhs = primitive_nonnull(0).into_array(); + let rhs = ConstantArray::new(37i64, LEN).into_array(); + + bench_primitive(bencher, lhs, rhs, Operator::Sub); +} + +#[divan::bench] +fn eq_i64_constant(bencher: Bencher) { + let lhs = primitive_nonnull(0).into_array(); + let rhs = ConstantArray::new(1024i64, LEN).into_array(); + + bench_bool(bencher, lhs, rhs, Operator::Eq); +} + +#[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(); + + bench_bool(bencher, lhs, rhs, Operator::Lt); +} + +#[divan::bench] +fn and_bool_nullable(bencher: Bencher) { + let lhs = bool_nullable(2, 7).into_array(); + let rhs = bool_nullable(3, 5).into_array(); + + 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_bool(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, operator: Operator) { + bench_binary::(bencher, lhs, rhs, operator); +} + +fn bench_binary( + bencher: Bencher, + lhs: ArrayRef, + rhs: ArrayRef, + operator: Operator, +) { + let mut ctx = SESSION.create_execution_ctx(); + + bencher.counter(ItemsCount::new(LEN)).bench_local(|| { + lhs.clone() + .binary(rhs.clone(), operator) + .unwrap() + .execute::(&mut ctx) + .unwrap() + }); +} + +fn primitive_nonnull(base: i64) -> PrimitiveArray { + PrimitiveArray::from_iter((0..LEN as i64).map(|i| base + i)) +} + +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) -> 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) -> PrimitiveArray { + PrimitiveArray::from_iter((0..LEN).map(|i| ((i + offset as usize) % 15 + 1) as u8)) +} + +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) -> PrimitiveArray { + PrimitiveArray::from_iter((0..LEN).map(|i| ((i + offset as usize) % 251 + 1) as u16)) +} + +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) -> PrimitiveArray { + PrimitiveArray::from_iter((0..LEN).map(|i| ((i + offset as usize) % 4096 + 1) as u32)) +} + +fn primitive_nonzero() -> PrimitiveArray { + PrimitiveArray::from_iter((0..LEN as i64).map(|i| (i % 255) + 1)) +} + +fn primitive_nullable(base: i64, null_every: usize) -> PrimitiveArray { + PrimitiveArray::from_option_iter( + (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| { + (!i.is_multiple_of(null_every)) + .then_some((((i as i64 + offset as i64) % 4096) - 2048) as i32) + })) +} + +fn bool_nullable(true_every: usize, null_every: usize) -> BoolArray { + BoolArray::from_iter( + (0..LEN).map(|i| (!i.is_multiple_of(null_every)).then_some(i.is_multiple_of(true_every))), + ) +} diff --git a/vortex-array/benches/bool_zip.rs b/vortex-array/benches/bool_zip.rs index c537d7b074f..40496fd1f0b 100644 --- a/vortex-array/benches/bool_zip.rs +++ b/vortex-array/benches/bool_zip.rs @@ -3,20 +3,24 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use vortex_array::ArrayRef; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::builtins::ArrayBuiltins; use vortex_mask::Mask; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + const LEN: usize = 65_536; /// Fragmented (alternating) mask: the worst case for the generic per-run builder this kernel @@ -53,7 +57,7 @@ fn run(bencher: Bencher, if_true: ArrayRef, if_false: ArrayRef) { if_true.clone(), if_false.clone(), mask.clone().into_array(), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(t, f, m, ctx)| { diff --git a/vortex-array/benches/cast_decimal.rs b/vortex-array/benches/cast_decimal.rs new file mode 100644 index 00000000000..8e0ddb19b6b --- /dev/null +++ b/vortex-array/benches/cast_decimal.rs @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks for decimal-to-decimal casting across the two dimensions that drive its cost: +//! +//! - **validity**: non-nullable vs nullable inputs. Nullable inputs route the rescale through the +//! masked kernel (`try_map_masked_into`) instead of the dense one (`try_map_into`). +//! - **work**: an *in-place* cast that only widens precision at the same scale, so the values +//! buffer is reused untouched (`O(1)`), vs a *copy* cast that must allocate a new buffer and +//! rescan every value (`O(n)`). + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use rand::prelude::*; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::DecimalArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::Nullability; +use vortex_array::validity::Validity; +use vortex_buffer::BufferMut; +use vortex_session::VortexSession; + +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]; + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +/// Builds an `i64`-backed `Decimal(precision, scale)` array of `n` values that all fit in +/// precision 10, optionally with ~50% nulls. +fn decimal_array(n: usize, precision: u8, scale: i8, nullable: bool) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(42); + let values: BufferMut = (0..n) + .map(|_| rng.random_range(0..1_000_000_000i64)) + .collect(); + let validity = if nullable { + Validity::from_iter((0..n).map(|_| rng.random_bool(0.5))) + } else { + Validity::NonNullable + }; + DecimalArray::new( + values.freeze(), + DecimalDType::new(precision, scale), + validity, + ) + .into_array() +} + +/// Casts `array` to `target` and forces execution, the unit measured by every bench below. +fn bench_cast(bencher: Bencher, array: ArrayRef, target: DType) { + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(a, ctx)| a.cast(target.clone()).unwrap().execute::(ctx)); +} + +// In place: widening precision at the same scale keeps the physical type, so the cast reuses the +// values buffer untouched. + +#[divan::bench(args = SIZES)] +fn in_place_non_nullable(bencher: Bencher, n: usize) { + bench_cast( + bencher, + decimal_array(n, 10, 2, false), + DType::Decimal(DecimalDType::new(18, 2), Nullability::NonNullable), + ); +} + +#[divan::bench(args = SIZES)] +fn in_place_nullable(bencher: Bencher, n: usize) { + bench_cast( + bencher, + decimal_array(n, 10, 2, true), + DType::Decimal(DecimalDType::new(18, 2), Nullability::Nullable), + ); +} + +// Copy: narrowing precision at the same scale cannot reuse the buffer; every value is re-scanned +// and range-checked into a freshly allocated buffer. + +#[divan::bench(args = SIZES)] +fn copy_non_nullable(bencher: Bencher, n: usize) { + bench_cast( + bencher, + decimal_array(n, 18, 2, false), + DType::Decimal(DecimalDType::new(10, 2), Nullability::NonNullable), + ); +} + +#[divan::bench(args = SIZES)] +fn copy_nullable(bencher: Bencher, n: usize) { + bench_cast( + bencher, + decimal_array(n, 18, 2, true), + DType::Decimal(DecimalDType::new(10, 2), Nullability::Nullable), + ); +} diff --git a/vortex-array/benches/cast_primitive.rs b/vortex-array/benches/cast_primitive.rs index 060a8290a2a..3fd89270517 100644 --- a/vortex-array/benches/cast_primitive.rs +++ b/vortex-array/benches/cast_primitive.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#![expect(clippy::unwrap_used)] + use std::sync::LazyLock; use divan::Bencher; @@ -14,28 +16,24 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::stats::Stat; -use vortex_array::session::ArraySession; use vortex_session::VortexSession; fn main() { divan::main(); } -const N: usize = 100_000; +// 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]; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); -#[divan::bench] -fn cast_u16_to_u32(bencher: Bencher) { +#[divan::bench(args = SIZES)] +fn cast_u16_to_u32(bencher: Bencher, n: usize) { let mut rng = StdRng::seed_from_u64(42); - #[expect(clippy::cast_possible_truncation)] - let arr = PrimitiveArray::from_option_iter((0..N).map(|i| { - if rng.random_bool(0.5) { - None - } else { - Some(i as u16) - } + let arr = PrimitiveArray::from_option_iter((0..n).map(|i| { + #[expect(clippy::cast_possible_truncation)] + rng.random_bool(0.5).then_some(i as u16) })) .into_array(); // Pre-compute min/max so values_fit_in is a cache hit during the benchmark. @@ -45,7 +43,43 @@ fn cast_u16_to_u32(bencher: Bencher) { bencher .with_inputs(|| (arr.clone(), SESSION.create_execution_ctx())) .bench_refs(|(a, ctx)| { - #[expect(clippy::unwrap_used)] + a.cast(DType::Primitive(PType::U32, Nullability::Nullable)) + .unwrap() + .execute::(ctx) + }); +} + +/// Narrowing fallible cast that goes through `try_map_with_mask`. Inputs are bounded +/// so every value fits, isolating the kernel's per-lane checked-cast overhead. +#[divan::bench(args = SIZES)] +fn cast_u32_to_u8(bencher: Bencher, n: usize) { + let mut rng = StdRng::seed_from_u64(42); + let arr = PrimitiveArray::from_option_iter((0..n).map(|_| { + rng.random_bool(0.7) + .then(|| rng.random_range(0..u8::MAX) as u32) + })) + .into_array(); + bencher + .with_inputs(|| (arr.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(a, ctx)| { + a.cast(DType::Primitive(PType::U8, Nullability::Nullable)) + .unwrap() + .execute::(ctx) + }); +} + +/// Sign-change cast i32 → u32. Values are non-negative so the kernel succeeds +/// but still pays the per-lane `try_from` check. +#[divan::bench(args = SIZES)] +fn cast_i32_to_u32(bencher: Bencher, n: usize) { + let mut rng = StdRng::seed_from_u64(42); + let arr = PrimitiveArray::from_option_iter( + (0..n).map(|_| rng.random_bool(0.7).then(|| rng.random_range(0..i32::MAX))), + ) + .into_array(); + bencher + .with_inputs(|| (arr.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(a, ctx)| { a.cast(DType::Primitive(PType::U32, Nullability::Nullable)) .unwrap() .execute::(ctx) diff --git a/vortex-array/benches/chunk_array_builder.rs b/vortex-array/benches/chunk_array_builder.rs index d74fd309453..6582071a6e1 100644 --- a/vortex-array/benches/chunk_array_builder.rs +++ b/vortex-array/benches/chunk_array_builder.rs @@ -18,7 +18,6 @@ use vortex_array::builders::ArrayBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::builders::builder_with_capacity; use vortex_array::dtype::DType; -use vortex_array::session::ArraySession; use vortex_error::VortexExpect; use vortex_session::VortexSession; @@ -33,8 +32,7 @@ const BENCH_ARGS: &[(usize, usize)] = &[ (1000, 10), ]; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[divan::bench(args = BENCH_ARGS)] fn chunked_bool_canonical_into(bencher: Bencher, (len, chunk_count): (usize, usize)) { diff --git a/vortex-array/benches/chunked_dict_builder.rs b/vortex-array/benches/chunked_dict_builder.rs index 85c8737356f..74722235e54 100644 --- a/vortex-array/benches/chunked_dict_builder.rs +++ b/vortex-array/benches/chunked_dict_builder.rs @@ -11,7 +11,6 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::dict_test::gen_dict_primitive_chunks; use vortex_array::builders::builder_with_capacity; use vortex_array::dtype::NativePType; -use vortex_array::session::ArraySession; use vortex_error::VortexExpect; use vortex_session::VortexSession; @@ -19,8 +18,7 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); const BENCH_ARGS: &[(usize, usize, usize)] = &[ (1000, 10, 10), diff --git a/vortex-array/benches/chunked_fsl_canonicalize.rs b/vortex-array/benches/chunked_fsl_canonicalize.rs index ac9962046de..4277b33382e 100644 --- a/vortex-array/benches/chunked_fsl_canonicalize.rs +++ b/vortex-array/benches/chunked_fsl_canonicalize.rs @@ -18,7 +18,6 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::FixedSizeListArray; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -27,8 +26,7 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); /// Number of lists in each chunk. const LISTS_PER_CHUNK: usize = 1_000; diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index 6fe5be688e9..95702c9f8f5 100644 --- a/vortex-array/benches/compare.rs +++ b/vortex-array/benches/compare.rs @@ -15,7 +15,6 @@ use vortex_array::arrays::BoolArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_buffer::Buffer; -use vortex_session::VortexSession; fn main() { divan::main(); @@ -30,7 +29,7 @@ fn compare_bool(bencher: Bencher) { let arr1 = BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.sample(range) == 0)).into_array(); let arr2 = BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.sample(range) == 0)).into_array(); - let session = VortexSession::empty(); + let session = vortex_array::array_session(); bencher .with_inputs(|| (&arr1, &arr2, session.create_execution_ctx())) @@ -58,7 +57,7 @@ fn compare_int(bencher: Bencher) { .map(|_| rng.sample(range)) .collect::>() .into_array(); - let session = VortexSession::empty(); + let session = vortex_array::array_session(); bencher .with_inputs(|| (&arr1, &arr2, session.create_execution_ctx())) diff --git a/vortex-array/benches/dict_compare.rs b/vortex-array/benches/dict_compare.rs index 401393495dd..446623939f1 100644 --- a/vortex-array/benches/dict_compare.rs +++ b/vortex-array/benches/dict_compare.rs @@ -3,34 +3,31 @@ #![expect(clippy::unwrap_used)] -use std::str::from_utf8; use std::sync::LazyLock; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; -use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::dict_test::gen_primitive_for_dict; use vortex_array::arrays::dict_test::gen_varbin_words; +use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::builders::dict::dict_encode; use vortex_array::builtins::ArrayBuiltins; use vortex_array::expr::eq; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::operators::Operator; -use vortex_array::session::ArraySession; use vortex_session::VortexSession; fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); const LENGTH_AND_UNIQUE_VALUES: &[(usize, usize)] = &[ // length, unique_values @@ -52,17 +49,13 @@ const LENGTH_AND_UNIQUE_VALUES: &[(usize, usize)] = &[ #[divan::bench(args = LENGTH_AND_UNIQUE_VALUES)] fn bench_compare_primitive(bencher: divan::Bencher, (len, uniqueness): (usize, usize)) { + let mut ctx = SESSION.create_execution_ctx(); let primitive_arr = gen_primitive_for_dict::(len, uniqueness); - let dict = dict_encode( - &primitive_arr.clone().into_array(), - &mut SESSION.create_execution_ctx(), - ) - .unwrap(); + let dict = dict_encode(&primitive_arr.clone().into_array(), &mut ctx).unwrap(); let value = primitive_arr.as_slice::()[0]; - let session = VortexSession::empty(); bencher - .with_inputs(|| (&dict, session.create_execution_ctx())) + .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| { dict.clone() .into_array() @@ -75,18 +68,14 @@ fn bench_compare_primitive(bencher: divan::Bencher, (len, uniqueness): (usize, u #[divan::bench(args = LENGTH_AND_UNIQUE_VALUES)] fn bench_compare_varbin(bencher: divan::Bencher, (len, uniqueness): (usize, usize)) { + let mut ctx = SESSION.create_execution_ctx(); let varbin_arr = VarBinArray::from(gen_varbin_words(len, uniqueness)); - let dict = dict_encode( - &varbin_arr.clone().into_array(), - &mut SESSION.create_execution_ctx(), - ) - .unwrap(); - let bytes = varbin_arr.with_iterator(|i| i.next().unwrap().unwrap().to_vec()); - let value = from_utf8(bytes.as_slice()).unwrap(); - let session = VortexSession::empty(); + let dict = dict_encode(&varbin_arr.clone().into_array(), &mut ctx).unwrap(); + let const_bytes = varbin_arr.bytes_at(0); + let value = unsafe { str::from_utf8_unchecked(const_bytes.as_slice()) }; bencher - .with_inputs(|| (&dict, session.create_execution_ctx())) + .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| { dict.clone() .into_array() @@ -99,18 +88,14 @@ fn bench_compare_varbin(bencher: divan::Bencher, (len, uniqueness): (usize, usiz #[divan::bench(args = LENGTH_AND_UNIQUE_VALUES)] fn bench_compare_varbinview(bencher: divan::Bencher, (len, uniqueness): (usize, usize)) { + let mut ctx = SESSION.create_execution_ctx(); let varbinview_arr = VarBinViewArray::from_iter_str(gen_varbin_words(len, uniqueness)); - let dict = dict_encode( - &varbinview_arr.clone().into_array(), - &mut SESSION.create_execution_ctx(), - ) - .unwrap(); - let bytes = varbinview_arr.with_iterator(|i| i.next().unwrap().unwrap().to_vec()); - let value = from_utf8(bytes.as_slice()).unwrap(); - let session = VortexSession::empty(); + let dict = dict_encode(&varbinview_arr.clone().into_array(), &mut ctx).unwrap(); + let const_bytes = varbinview_arr.bytes_at(0); + let value = unsafe { str::from_utf8_unchecked(const_bytes.as_slice()) }; bencher - .with_inputs(|| (&dict, session.create_execution_ctx())) + .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| { dict.clone() .into_array() @@ -138,18 +123,14 @@ fn bench_compare_sliced_dict_primitive( bencher: divan::Bencher, (codes_len, values_len): (usize, usize), ) { + let mut ctx = SESSION.create_execution_ctx(); let primitive_arr = gen_primitive_for_dict::(codes_len.max(values_len), values_len); - let dict = dict_encode( - &primitive_arr.clone().into_array(), - &mut SESSION.create_execution_ctx(), - ) - .unwrap(); + let dict = dict_encode(&primitive_arr.clone().into_array(), &mut ctx).unwrap(); let dict = dict.into_array().slice(0..codes_len).unwrap(); let value = primitive_arr.as_slice::()[0]; - let session = VortexSession::empty(); bencher - .with_inputs(|| (&dict, session.create_execution_ctx())) + .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| { dict.clone() .apply(&eq(root(), lit(value))) @@ -164,19 +145,15 @@ fn bench_compare_sliced_dict_varbinview( bencher: divan::Bencher, (codes_len, values_len): (usize, usize), ) { + let mut ctx = SESSION.create_execution_ctx(); let varbin_arr = VarBinArray::from(gen_varbin_words(codes_len.max(values_len), values_len)); - let dict = dict_encode( - &varbin_arr.clone().into_array(), - &mut SESSION.create_execution_ctx(), - ) - .unwrap(); + let dict = dict_encode(&varbin_arr.clone().into_array(), &mut ctx).unwrap(); let dict = dict.into_array().slice(0..codes_len).unwrap(); - let bytes = varbin_arr.with_iterator(|i| i.next().unwrap().unwrap().to_vec()); - let value = from_utf8(bytes.as_slice()).unwrap(); - let session = VortexSession::empty(); + let const_bytes = varbin_arr.bytes_at(0); + let value = unsafe { str::from_utf8_unchecked(const_bytes.as_slice()) }; bencher - .with_inputs(|| (&dict, session.create_execution_ctx())) + .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| { dict.clone() .apply(&eq(root(), lit(value))) diff --git a/vortex-array/benches/dict_compress.rs b/vortex-array/benches/dict_compress.rs index 21d427bd7f8..2901d82c6b5 100644 --- a/vortex-array/benches/dict_compress.rs +++ b/vortex-array/benches/dict_compress.rs @@ -17,7 +17,6 @@ use vortex_array::arrays::dict_test::gen_primitive_for_dict; use vortex_array::arrays::dict_test::gen_varbin_words; use vortex_array::builders::dict::dict_encode; use vortex_array::dtype::NativePType; -use vortex_array::session::ArraySession; use vortex_session::VortexSession; fn main() { @@ -38,8 +37,7 @@ const BENCH_ARGS: &[(usize, usize)] = &[ (10_000, 512), ]; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[divan::bench(types = [u8, f32, i64], args = BENCH_ARGS)] fn encode_primitives(bencher: Bencher, (len, unique_values): (usize, usize)) @@ -47,29 +45,30 @@ where T: NativePType, StandardUniform: Distribution, { - let primitive_arr = gen_primitive_for_dict::(len, unique_values); + let primitive_arr = gen_primitive_for_dict::(len, unique_values).into_array(); bencher .with_inputs(|| (&primitive_arr, SESSION.create_execution_ctx())) - .bench_refs(|(arr, ctx)| dict_encode(&arr.clone().into_array(), ctx)); + .bench_refs(|(arr, ctx)| dict_encode(arr, ctx)); } #[divan::bench(args = BENCH_ARGS)] fn encode_varbin(bencher: Bencher, (len, unique_values): (usize, usize)) { - let varbin_arr = VarBinArray::from(gen_varbin_words(len, unique_values)); + let varbin_arr = VarBinArray::from(gen_varbin_words(len, unique_values)).into_array(); bencher .with_inputs(|| (&varbin_arr, SESSION.create_execution_ctx())) - .bench_refs(|(arr, ctx)| dict_encode(&arr.clone().into_array(), ctx)); + .bench_refs(|(arr, ctx)| dict_encode(arr, ctx)); } #[divan::bench(args = BENCH_ARGS)] fn encode_varbinview(bencher: Bencher, (len, unique_values): (usize, usize)) { - let varbinview_arr = VarBinViewArray::from_iter_str(gen_varbin_words(len, unique_values)); + let varbinview_arr = + VarBinViewArray::from_iter_str(gen_varbin_words(len, unique_values)).into_array(); bencher .with_inputs(|| (&varbinview_arr, SESSION.create_execution_ctx())) - .bench_refs(|(arr, ctx)| dict_encode(&arr.clone().into_array(), ctx)); + .bench_refs(|(arr, ctx)| dict_encode(arr, ctx)); } #[divan::bench(types = [u8, f32, i64], args = BENCH_ARGS)] @@ -78,13 +77,10 @@ where T: NativePType, StandardUniform: Distribution, { - let primitive_arr = gen_primitive_for_dict::(len, unique_values); - let dict = dict_encode( - &primitive_arr.into_array(), - &mut SESSION.create_execution_ctx(), - ) - .unwrap() - .into_array(); + let primitive_arr = gen_primitive_for_dict::(len, unique_values).into_array(); + let dict = dict_encode(&primitive_arr, &mut SESSION.create_execution_ctx()) + .unwrap() + .into_array(); bencher .with_inputs(|| (&dict, SESSION.create_execution_ctx())) @@ -93,13 +89,10 @@ where #[divan::bench(args = BENCH_ARGS)] fn decode_varbin(bencher: Bencher, (len, unique_values): (usize, usize)) { - let varbin_arr = VarBinArray::from(gen_varbin_words(len, unique_values)); - let dict = dict_encode( - &varbin_arr.into_array(), - &mut SESSION.create_execution_ctx(), - ) - .unwrap() - .into_array(); + let varbin_arr = VarBinArray::from(gen_varbin_words(len, unique_values)).into_array(); + let dict = dict_encode(&varbin_arr, &mut SESSION.create_execution_ctx()) + .unwrap() + .into_array(); bencher .with_inputs(|| (&dict, SESSION.create_execution_ctx())) @@ -108,13 +101,11 @@ fn decode_varbin(bencher: Bencher, (len, unique_values): (usize, usize)) { #[divan::bench(args = BENCH_ARGS)] fn decode_varbinview(bencher: Bencher, (len, unique_values): (usize, usize)) { - let varbinview_arr = VarBinViewArray::from_iter_str(gen_varbin_words(len, unique_values)); - let dict = dict_encode( - &varbinview_arr.into_array(), - &mut SESSION.create_execution_ctx(), - ) - .unwrap() - .into_array(); + let varbinview_arr = + VarBinViewArray::from_iter_str(gen_varbin_words(len, unique_values)).into_array(); + let dict = dict_encode(&varbinview_arr, &mut SESSION.create_execution_ctx()) + .unwrap() + .into_array(); bencher .with_inputs(|| (&dict, SESSION.create_execution_ctx())) diff --git a/vortex-array/benches/dict_mask.rs b/vortex-array/benches/dict_mask.rs index 1951e610a51..04cf18cfd03 100644 --- a/vortex-array/benches/dict_mask.rs +++ b/vortex-array/benches/dict_mask.rs @@ -14,7 +14,6 @@ use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_mask::Mask; -use vortex_session::VortexSession; fn main() { divan::main(); @@ -60,7 +59,7 @@ fn bench_dict_mask(bencher: Bencher, (fraction_valid, fraction_masked): (f64, f6 let values = PrimitiveArray::from_option_iter([None, Some(42i32)]).into_array(); let array = DictArray::try_new(codes, values).unwrap().into_array(); let filter_mask = filter_mask(len, fraction_masked, &mut rng); - let session = VortexSession::empty(); + let session = vortex_array::array_session(); bencher .with_inputs(|| (&array, &filter_mask, session.create_execution_ctx())) .bench_refs(|(array, filter_mask, ctx)| { diff --git a/vortex-array/benches/expr/case_when_bench.rs b/vortex-array/benches/expr/case_when_bench.rs index f37daa2b304..1bada8fa9e1 100644 --- a/vortex-array/benches/expr/case_when_bench.rs +++ b/vortex-array/benches/expr/case_when_bench.rs @@ -22,12 +22,10 @@ use vortex_array::expr::lit; use vortex_array::expr::lt; use vortex_array::expr::nested_case_when; use vortex_array::expr::root; -use vortex_array::session::ArraySession; use vortex_buffer::Buffer; use vortex_session::VortexSession; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn main() { divan::main(); diff --git a/vortex-array/benches/filter_bool.rs b/vortex-array/benches/filter_bool.rs index 6c240db6b60..cd31805bb77 100644 --- a/vortex-array/benches/filter_bool.rs +++ b/vortex-array/benches/filter_bool.rs @@ -19,7 +19,6 @@ use rand_distr::Zipf; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; -use vortex_array::session::ArraySession; use vortex_buffer::BitBuffer; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -28,8 +27,7 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); const SIZES: &[usize] = &[1_000, 10_000, 100_000, 250_000]; const DENSITY_SWEEP_SIZE: usize = 100_000; diff --git a/vortex-array/benches/interleave.rs b/vortex-array/benches/interleave.rs index cba7f6a910e..d92c033c2a4 100644 --- a/vortex-array/benches/interleave.rs +++ b/vortex-array/benches/interleave.rs @@ -1,20 +1,33 @@ // 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: +//! +//! - `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. +//! - `random`, 64 children: the same random gather spread over many value arrays. +//! +//! Each is run nullable and non-nullable. + #![expect(clippy::unwrap_used)] +use std::fmt::Display; +use std::fmt::Formatter; + use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::distr::Uniform; use rand::prelude::StdRng; +use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::InterleaveArray; use vortex_buffer::Buffer; -use vortex_session::VortexSession; fn main() { divan::main(); @@ -22,18 +35,74 @@ fn main() { const ARRAY_SIZE: usize = 8_192; -/// Builds `num_branches` boolean value arrays plus random `(array_indices, row_indices)` selectors -/// describing a full random-access gather of `ARRAY_SIZE` output rows. -fn inputs( - num_branches: usize, +/// The access pattern used to generate the `(array_index, row_index)` selectors. +#[derive(Clone, Copy)] +enum Pattern { + /// A merge: `array_index = i % N`, `row_index = i / N`. + RoundRobin, + /// Fully random `(array_index, row_index)` per output row. + Random, +} + +impl Display for Pattern { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Pattern::RoundRobin => "round_robin", + Pattern::Random => "random", + }) + } +} + +/// A single benchmark configuration: access pattern, branch count, and nullability. +#[derive(Clone, Copy)] +struct Combo { + pattern: Pattern, + branches: usize, nullable: bool, -) -> (Vec, Buffer, Buffer) { +} + +impl Display for Combo { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}/n{}/{}", + self.pattern, + self.branches, + if self.nullable { "null" } else { "nonnull" } + ) + } +} + +/// The configurations the benchmark covers: 2-child round-robin, 2-child random, and 64-child +/// random — each nullable and non-nullable. +fn combos() -> Vec { + let mut out = Vec::new(); + for nullable in [false, true] { + for (pattern, branches) in [ + (Pattern::RoundRobin, 2), + (Pattern::Random, 2), + (Pattern::Random, 64), + ] { + out.push(Combo { + pattern, + branches, + nullable, + }); + } + } + out +} + +/// Builds the Vortex value arrays and the `u32` selector buffers for a [`Combo`]. +/// +/// Seeded only by the combo so a run is deterministic and comparable across revisions. +fn vortex_inputs(combo: Combo) -> (Vec, Buffer, Buffer) { let mut rng = StdRng::seed_from_u64(0); let bit = Uniform::new(0u8, 2).unwrap(); - let values = (0..num_branches) + let values = (0..combo.branches) .map(|_| { - if nullable { + if combo.nullable { BoolArray::from_iter( (0..ARRAY_SIZE).map(|_| (rng.sample(bit) == 0).then_some(rng.sample(bit) == 0)), ) @@ -44,37 +113,27 @@ fn inputs( }) .collect(); - let branch = Uniform::new(0u32, u32::try_from(num_branches).unwrap()).unwrap(); + 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: Buffer = (0..ARRAY_SIZE).map(|_| rng.sample(branch)).collect(); - let row_indices: Buffer = (0..ARRAY_SIZE).map(|_| rng.sample(row)).collect(); - (values, array_indices, row_indices) -} - -#[divan::bench(args = [2, 4])] -fn interleave_bool(bencher: Bencher, num_branches: usize) { - let (values, array_indices, row_indices) = inputs(num_branches, false); - let session = VortexSession::empty(); - 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(), - ) + let array_indices: Buffer = (0..ARRAY_SIZE) + .map(|i| match combo.pattern { + Pattern::Random => rng.sample(branch), + Pattern::RoundRobin => u32::try_from(i % combo.branches).unwrap(), }) - .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); + .collect(); + let row_indices: Buffer = (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 = [2, 4])] -fn interleave_bool_nullable(bencher: Bencher, num_branches: usize) { - let (values, array_indices, row_indices) = inputs(num_branches, true); - let session = VortexSession::empty(); +#[divan::bench(args = combos())] +fn vortex(bencher: Bencher, combo: Combo) { + let (values, array_indices, row_indices) = vortex_inputs(combo); + let session = array_session(); bencher .with_inputs(|| { ( diff --git a/vortex-array/benches/kleene_bool.rs b/vortex-array/benches/kleene_bool.rs new file mode 100644 index 00000000000..b23a0580a5a --- /dev/null +++ b/vortex-array/benches/kleene_bool.rs @@ -0,0 +1,232 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrow::ArrowSession; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::optimizer::kernels::KernelSession; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::session::ArraySession; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(|| { + VortexSession::empty() + .with::() + .with::() + .with::() +}); + +const LEN: usize = 65_536; +const SHIFTED_OFFSET: usize = 1; + +#[divan::bench] +fn and_bool_nonnull_arrays(bencher: Bencher) { + bench_kleene( + bencher, + bool_nonnull(2).into_array(), + bool_nonnull(3).into_array(), + Operator::And, + ); +} + +#[divan::bench] +fn or_bool_nonnull_arrays(bencher: Bencher) { + bench_kleene( + bencher, + bool_nonnull(2).into_array(), + bool_nonnull(3).into_array(), + Operator::Or, + ); +} + +#[divan::bench] +fn and_bool_nullable_arrays(bencher: Bencher) { + and_bool_nullable_arrays_aligned(bencher); +} + +#[divan::bench] +fn and_bool_nullable_arrays_aligned(bencher: Bencher) { + bench_kleene( + bencher, + bool_nullable(2, 7, 0), + bool_nullable(3, 5, 0), + Operator::And, + ); +} + +#[divan::bench] +fn or_bool_nullable_arrays(bencher: Bencher) { + or_bool_nullable_arrays_aligned(bencher); +} + +#[divan::bench] +fn or_bool_nullable_arrays_aligned(bencher: Bencher) { + bench_kleene( + bencher, + bool_nullable(2, 7, 0), + bool_nullable(3, 5, 0), + Operator::Or, + ); +} + +#[divan::bench] +fn and_bool_nullable_arrays_shifted(bencher: Bencher) { + bench_kleene( + bencher, + bool_nullable(2, 7, SHIFTED_OFFSET), + bool_nullable(3, 5, SHIFTED_OFFSET), + Operator::And, + ); +} + +#[divan::bench] +fn or_bool_nullable_arrays_shifted(bencher: Bencher) { + bench_kleene( + bencher, + bool_nullable(2, 7, SHIFTED_OFFSET), + bool_nullable(3, 5, SHIFTED_OFFSET), + Operator::Or, + ); +} + +#[divan::bench] +fn and_true_constant(bencher: Bencher) { + bench_kleene( + bencher, + bool_nullable(2, 7, 0), + ConstantArray::new(true, LEN).into_array(), + Operator::And, + ); +} + +#[divan::bench] +fn or_false_constant(bencher: Bencher) { + bench_kleene( + bencher, + bool_nullable(2, 7, 0), + ConstantArray::new(false, LEN).into_array(), + Operator::Or, + ); +} + +#[divan::bench] +fn and_false_constant(bencher: Bencher) { + bench_kleene( + bencher, + bool_nullable(2, 7, 0), + ConstantArray::new(false, LEN).into_array(), + Operator::And, + ); +} + +#[divan::bench] +fn or_true_constant(bencher: Bencher) { + bench_kleene( + bencher, + bool_nullable(2, 7, 0), + ConstantArray::new(true, LEN).into_array(), + Operator::Or, + ); +} + +#[divan::bench] +fn and_null_constant(bencher: Bencher) { + and_null_constant_aligned(bencher); +} + +#[divan::bench] +fn and_null_constant_aligned(bencher: Bencher) { + bench_kleene( + bencher, + bool_nullable(2, 7, 0), + null_bool_constant(), + Operator::And, + ); +} + +#[divan::bench] +fn or_null_constant(bencher: Bencher) { + or_null_constant_aligned(bencher); +} + +#[divan::bench] +fn or_null_constant_aligned(bencher: Bencher) { + bench_kleene( + bencher, + bool_nullable(2, 7, 0), + null_bool_constant(), + Operator::Or, + ); +} + +#[divan::bench] +fn and_null_constant_shifted(bencher: Bencher) { + bench_kleene( + bencher, + bool_nullable(2, 7, SHIFTED_OFFSET), + null_bool_constant(), + Operator::And, + ); +} + +#[divan::bench] +fn or_null_constant_shifted(bencher: Bencher) { + bench_kleene( + bencher, + bool_nullable(2, 7, SHIFTED_OFFSET), + null_bool_constant(), + Operator::Or, + ); +} + +fn bench_kleene(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, operator: Operator) { + let mut ctx = SESSION.create_execution_ctx(); + + bencher.counter(ItemsCount::new(LEN)).bench_local(|| { + lhs.clone() + .binary(rhs.clone(), operator) + .unwrap() + .execute::(&mut ctx) + .unwrap() + }); +} + +fn bool_nonnull(true_every: usize) -> BoolArray { + BoolArray::from_iter((0..LEN).map(|i| i.is_multiple_of(true_every))) +} + +fn bool_nullable(true_every: usize, null_every: usize, offset: usize) -> ArrayRef { + let len = LEN + offset; + let array = BoolArray::from_iter( + (0..len).map(|i| (!i.is_multiple_of(null_every)).then_some(i.is_multiple_of(true_every))), + ) + .into_array(); + + if offset == 0 { + array + } else { + array.slice(offset..offset + LEN).unwrap() + } +} + +fn null_bool_constant() -> ArrayRef { + ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), LEN).into_array() +} diff --git a/vortex-array/benches/list_length.rs b/vortex-array/benches/list_length.rs new file mode 100644 index 00000000000..6148b54f8f9 --- /dev/null +++ b/vortex-array/benches/list_length.rs @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks for the `list_length` scalar function over `List` and `ListView` inputs. +//! +//! `list_length` reads only the offsets/sizes (never the elements), so its cost scales with the +//! number of lists. + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use std::sync::LazyLock; + +use divan::Bencher; +use rand::RngExt; +use rand::SeedableRng; +use rand::distr::Uniform; +use rand::rngs::StdRng; +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::ListArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::expr::list_length; +use vortex_array::expr::root; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +const BASE_LIST_SIZE: usize = 8; + +const SMALL: usize = 100; +const MEDIUM: usize = 10_000; +const LARGE: usize = 1_000_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. +fn random_lists(num_lists: usize) -> (Vec, Validity) { + let mut rng = StdRng::seed_from_u64(num_lists as u64); + let total = (num_lists * BASE_LIST_SIZE) as i32; + + let cut_dist = Uniform::new_inclusive(0i32, total).unwrap(); + let mut cuts: Vec = (0..num_lists - 1).map(|_| rng.sample(cut_dist)).collect(); + cuts.sort_unstable(); + let mut sizes = Vec::with_capacity(num_lists); + let mut prev = 0i32; + for cut in cuts { + sizes.push(cut - prev); + prev = cut; + } + sizes.push(total - prev); + + let null_dist = Uniform::new(0u32, 8).unwrap(); + let valid = (0..num_lists).map(|_| rng.sample(null_dist) != 0); + ( + sizes, + Validity::Array(BoolArray::from_iter(valid).into_array()), + ) +} + +/// A canonical `List` of `num_lists` variable-length lists, ~1/8 of them null. +fn make_list(num_lists: usize) -> ArrayRef { + let (sizes, validity) = random_lists(num_lists); + let total: i32 = sizes.iter().sum(); + let elements = PrimitiveArray::from_iter(0..total).into_array(); + let offsets: Buffer = std::iter::once(0) + .chain(sizes.iter().scan(0i32, |acc, &s| { + *acc += s; + Some(*acc) + })) + .collect(); + ListArray::try_new(elements, offsets.into_array(), validity) + .unwrap() + .into_array() +} + +/// A gapless `ListView` of `num_lists` variable-length lists, ~1/8 of them null. +fn make_listview(num_lists: usize) -> ArrayRef { + let (sizes, validity) = random_lists(num_lists); + let total: i32 = sizes.iter().sum(); + let elements = PrimitiveArray::from_iter(0..total).into_array(); + let offsets: Buffer = sizes + .iter() + .scan(0i32, |acc, &s| { + let start = *acc; + *acc += s; + Some(start) + }) + .collect(); + let sizes: Buffer = sizes.into_iter().collect(); + ListViewArray::new(elements, offsets.into_array(), sizes.into_array(), validity).into_array() +} + +/// Apply `list_length(root())` and materialize the result. +fn run(bencher: Bencher, array: ArrayRef) { + let expr = list_length(root()); + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| { + array + .clone() + .apply(&expr) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} + +#[divan::bench] +fn list_length_small(bencher: Bencher) { + run(bencher, make_list(SMALL)); +} + +#[divan::bench] +fn list_length_medium(bencher: Bencher) { + run(bencher, make_list(MEDIUM)); +} + +#[divan::bench] +fn list_length_large(bencher: Bencher) { + run(bencher, make_list(LARGE)); +} + +#[divan::bench] +fn listview_length_small(bencher: Bencher) { + run(bencher, make_listview(SMALL)); +} + +#[divan::bench] +fn listview_length_medium(bencher: Bencher) { + run(bencher, make_listview(MEDIUM)); +} + +#[divan::bench] +fn listview_length_large(bencher: Bencher) { + run(bencher, make_listview(LARGE)); +} diff --git a/vortex-array/benches/listview_rebuild.rs b/vortex-array/benches/listview_rebuild.rs index 31fdaf5e890..81157d52b5a 100644 --- a/vortex-array/benches/listview_rebuild.rs +++ b/vortex-array/benches/listview_rebuild.rs @@ -21,7 +21,6 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::listview::ListViewArrayExt; use vortex_array::arrays::listview::ListViewRebuildMode; use vortex_array::dtype::FieldNames; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -31,8 +30,7 @@ fn main() { } /// A shared session for the `ListView` rebuild benchmarks, used to create execution contexts. -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn make_primitive_lv(num_lists: usize, list_size: usize, step: usize) -> ListViewArray { let element_count = step * num_lists + list_size; diff --git a/vortex-array/benches/primitive_zip.rs b/vortex-array/benches/primitive_zip.rs index 0359b464f13..adcabf51d7d 100644 --- a/vortex-array/benches/primitive_zip.rs +++ b/vortex-array/benches/primitive_zip.rs @@ -7,20 +7,24 @@ reason = "benchmark fixtures use indices that fit in the chosen widths" )] +use std::sync::LazyLock; + use divan::Bencher; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_buffer::BufferMut; use vortex_mask::Mask; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + // Sized so the bench stays well under a few hundred microseconds under CodSpeed's instruction-count // simulation, which runs ~10x the local walltime; the branchless value blend is still exercised. const LEN: usize = 16_384; @@ -53,7 +57,7 @@ fn run(bencher: Bencher, if_true: vortex_array::ArrayRef, if_false: vortex_array if_true.clone(), if_false.clone(), mask.clone().into_array(), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(t, f, m, ctx)| { diff --git a/vortex-array/benches/scalar_at_struct.rs b/vortex-array/benches/scalar_at_struct.rs index 5e435b1000a..3921f0a2366 100644 --- a/vortex-array/benches/scalar_at_struct.rs +++ b/vortex-array/benches/scalar_at_struct.rs @@ -15,7 +15,6 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::StructArray; use vortex_array::dtype::FieldNames; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -27,8 +26,7 @@ fn main() { const ARRAY_SIZE: usize = 100_000; const NUM_ACCESSES: usize = 1000; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[divan::bench] fn execute_scalar_struct_simple(bencher: Bencher) { diff --git a/vortex-array/benches/scalar_subtract.rs b/vortex-array/benches/scalar_subtract.rs index 2eee8f3d9ff..ee9aba1055d 100644 --- a/vortex-array/benches/scalar_subtract.rs +++ b/vortex-array/benches/scalar_subtract.rs @@ -17,7 +17,6 @@ use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar_fn::fns::operators::Operator; -use vortex_array::session::ArraySession; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -25,8 +24,7 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[divan::bench] fn scalar_subtract(bencher: Bencher) { diff --git a/vortex-array/benches/take_filter.rs b/vortex-array/benches/take_filter.rs index 0e3ff719a75..b6ec1b4040c 100644 --- a/vortex-array/benches/take_filter.rs +++ b/vortex-array/benches/take_filter.rs @@ -30,7 +30,6 @@ use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; @@ -58,8 +57,7 @@ const INDEX_SEED: u64 = 43; const LIST_SIZE: usize = 4; const NULL_INDEX_INTERVAL: usize = 8; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn primitive_array() -> ArrayRef { PrimitiveArray::from_iter(0..ARRAY_LEN as u32).into_array() diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index 4764f5538be..438faae4983 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -20,7 +20,6 @@ use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::FixedSizeListArray; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -29,8 +28,7 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); /// Number of lists in the source array. const NUM_LISTS: usize = 500; diff --git a/vortex-array/benches/take_patches.rs b/vortex-array/benches/take_patches.rs index 4d87f8b3c63..5dc0c04c325 100644 --- a/vortex-array/benches/take_patches.rs +++ b/vortex-array/benches/take_patches.rs @@ -16,7 +16,6 @@ use vortex_array::IntoArray; use vortex_array::ToCanonical as _; use vortex_array::VortexSessionExecute; use vortex_array::patches::Patches; -use vortex_array::session::ArraySession; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -24,8 +23,7 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); const BENCH_ARGS: &[(f64, f64)] = &[ // patches_sparsity, index_multiple diff --git a/vortex-array/benches/take_primitive.rs b/vortex-array/benches/take_primitive.rs index e211e54d5f7..383b525dcc6 100644 --- a/vortex-array/benches/take_primitive.rs +++ b/vortex-array/benches/take_primitive.rs @@ -19,7 +19,6 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::session::ArraySession; use vortex_session::VortexSession; fn main() { @@ -32,8 +31,7 @@ const NUM_INDICES: &[usize] = &[1_000, 10_000, 100_000]; /// Size of the source vector / dictionary values. const VECTOR_SIZE: &[usize] = &[16, 256, 2048, 8192]; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); // --- DictArray canonicalization benchmarks --- diff --git a/vortex-array/benches/take_struct.rs b/vortex-array/benches/take_struct.rs index 5f15894383a..d49ce9e88a6 100644 --- a/vortex-array/benches/take_struct.rs +++ b/vortex-array/benches/take_struct.rs @@ -15,7 +15,6 @@ use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::StructArray; use vortex_array::dtype::FieldNames; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -24,8 +23,7 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); const ARRAY_SIZE: usize = 100_000; const TAKE_SIZE: usize = 1000; diff --git a/vortex-array/benches/to_arrow.rs b/vortex-array/benches/to_arrow.rs index 0f84feb7baa..48e5e917476 100644 --- a/vortex-array/benches/to_arrow.rs +++ b/vortex-array/benches/to_arrow.rs @@ -25,15 +25,13 @@ use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; -use vortex_array::session::ArraySession; use vortex_session::VortexSession; fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn schema() -> DType { let fields = StructFields::from_iter([ diff --git a/vortex-array/benches/varbinview_zip.rs b/vortex-array/benches/varbinview_zip.rs index d18012d9914..e010291e9dc 100644 --- a/vortex-array/benches/varbinview_zip.rs +++ b/vortex-array/benches/varbinview_zip.rs @@ -13,7 +13,6 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::session::ArraySession; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -21,8 +20,7 @@ fn main() { divan::main(); } -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); /// Benchmarks zip on VarBinView arrays with a highly fragmented mask (worst case for per-slice lookup paths). #[divan::bench] diff --git a/vortex-array/src/accessor.rs b/vortex-array/src/accessor.rs deleted file mode 100644 index 6460f879cb3..00000000000 --- a/vortex-array/src/accessor.rs +++ /dev/null @@ -1,14 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -/// Trait for arrays that support iterative access to their elements. -pub trait ArrayAccessor { - /// Iterate over each element of the array, in-order. - /// - /// The function `f` will be passed an [`Iterator`], it can call [`Iterator::next`] on the - /// iterator `len` times. Iterator elements are `Option` types, - /// regardless of the nullability of the underlying array data. - fn with_iterator(&self, f: F) -> R - where - F: for<'a> FnOnce(&mut dyn Iterator>) -> R; -} diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index c89418e67a6..69bae4e1053 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -275,7 +275,7 @@ mod tests { use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::combined::Combined; use crate::aggregate_fn::combined::PairOptions; use crate::aggregate_fn::fns::mean::Mean; @@ -289,7 +289,6 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::scalar::Scalar; - use crate::session::ArraySession; /// Mean partial sentinel `{sum: 42.0, count: 1}` — distinguishable from the /// natural fan-out result `{sum: 7.0, count: 1}` that `Combined::try_accumulate` @@ -337,7 +336,7 @@ mod tests { } fn fresh_session() -> VortexSession { - VortexSession::empty().with::() + crate::array_session() } fn dict_of_seven() -> ArrayRef { @@ -350,7 +349,10 @@ mod tests { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); Accumulator::try_new( Mean::combined(), - PairOptions(EmptyOptions, EmptyOptions), + PairOptions( + NumericalAggregateOpts::default(), + NumericalAggregateOpts::default(), + ), dtype, ) } diff --git a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs index c33443dd02b..68a58908018 100644 --- a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -33,11 +35,20 @@ impl AggregateFnVTable for AllNan { type Partial = bool; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.all_nan") + static ID: CachedId = CachedId::new("vortex.all_nan"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - Ok(None) + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { @@ -124,12 +135,12 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::all_nan::AllNan; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::Nullability; @@ -137,7 +148,7 @@ mod tests { #[test] fn all_nan_aggregate_fn() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?; @@ -150,7 +161,7 @@ mod tests { #[test] fn all_nan_false_with_non_nan() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?; @@ -170,7 +181,7 @@ mod tests { #[test] fn all_nan_false_with_null() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?; 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 445f4e17b6f..f030cc810e2 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 @@ -12,12 +12,14 @@ mod struct_; #[cfg(test)] mod tests; mod varbin; +mod variant; use std::sync::LazyLock; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_session::registry::CachedId; use self::bool::check_bool_identical; use self::decimal::check_decimal_identical; @@ -38,6 +40,7 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::fns::all_non_distinct::variant::check_variant_identical; use crate::arrays::StructArray; use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; @@ -119,7 +122,8 @@ impl AggregateFnVTable for AllNonDistinct { type Partial = AllNonDistinctPartial; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.all_non_distinct") + static ID: CachedId = CachedId::new("vortex.all_non_distinct"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { @@ -264,8 +268,8 @@ fn check_canonical_identical( (Canonical::Extension(lhs), Canonical::Extension(rhs)) => { check_extension_identical(lhs, rhs, ctx) } - (Canonical::Variant(_), _) | (_, Canonical::Variant(_)) => { - vortex_bail!("Variant arrays don't support AllNonDistinct") + (Canonical::Variant(lhs), Canonical::Variant(rhs)) => { + check_variant_identical(lhs, rhs, ctx) } _ => Err(vortex_err!( "Canonical type mismatch in AllNonDistinct: {:?} vs {:?}", diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs index f11a9a2bc52..1ad0ca6d646 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::dtype::NativePType; use vortex_error::VortexResult; use crate::arrays::primitive::PrimitiveArrayExt; @@ -12,6 +13,10 @@ where R: PrimitiveArrayExt, { match_each_native_ptype!(lhs.ptype(), |P| { - Ok(lhs.as_slice::

() == rhs.as_slice::

()) + Ok(lhs + .as_slice::

() + .iter() + .zip(rhs.as_slice::

()) + .all(|(l, r)| l.is_eq(*r))) }) } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/tests.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/tests.rs index 86fe5b2faf9..ab513670c18 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/tests.rs @@ -9,8 +9,8 @@ use super::all_non_distinct; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::DecimalArray; @@ -43,7 +43,7 @@ fn scalar_baseline(a: &ArrayRef, b: &ArrayRef, ctx: &mut ExecutionCtx) -> Vortex /// Assert that `all_non_distinct` agrees with the scalar baseline. fn assert_matches_baseline(a: &ArrayRef, b: &ArrayRef) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let expected = scalar_baseline(a, b, &mut ctx)?; let actual = all_non_distinct(a, b, &mut ctx)?; assert_eq!( @@ -88,7 +88,7 @@ fn bool_non_nullable( ) -> VortexResult<()> { let a = BoolArray::from_iter(a.iter().copied()).into_array(); let b = BoolArray::from_iter(b.iter().copied()).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(all_non_distinct(&a, &b, &mut ctx)?, expected); assert_matches_baseline(&a, &b) } @@ -121,7 +121,7 @@ fn bool_nullable( ) -> VortexResult<()> { let a = BoolArray::from_iter(a.iter().copied()).into_array(); let b = BoolArray::from_iter(b.iter().copied()).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(all_non_distinct(&a, &b, &mut ctx)?, expected); assert_matches_baseline(&a, &b) } @@ -140,7 +140,7 @@ fn primitive_i32( ) -> VortexResult<()> { let a = PrimitiveArray::from_iter(a).into_array(); let b = PrimitiveArray::from_iter(b).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(all_non_distinct(&a, &b, &mut ctx)?, expected); assert_matches_baseline(&a, &b) } @@ -206,7 +206,7 @@ fn primitive_nullable( ) -> VortexResult<()> { let a = PrimitiveArray::from_option_iter(a.iter().copied()).into_array(); let b = PrimitiveArray::from_option_iter(b.iter().copied()).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(all_non_distinct(&a, &b, &mut ctx)?, expected); assert_matches_baseline(&a, &b) } @@ -235,7 +235,7 @@ fn strings_non_nullable( ) -> VortexResult<()> { let a = VarBinViewArray::from_iter_str(a.iter().copied()).into_array(); let b = VarBinViewArray::from_iter_str(b.iter().copied()).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(all_non_distinct(&a, &b, &mut ctx)?, expected); assert_matches_baseline(&a, &b) } @@ -258,7 +258,7 @@ fn strings_nullable( ) -> VortexResult<()> { let a = VarBinViewArray::from_iter_nullable_str(a.iter().copied()).into_array(); let b = VarBinViewArray::from_iter_nullable_str(b.iter().copied()).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(all_non_distinct(&a, &b, &mut ctx)?, expected); assert_matches_baseline(&a, &b) } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs new file mode 100644 index 00000000000..1576b8c3d6a --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ExecutionCtx; +use crate::arrays::VariantArray; + +/// Checks whether two canonical variant arrays are element-wise non-distinct. +/// +/// Variant values cannot be routed back through [`all_non_distinct`]: canonicalizing a variant +/// value array yields another canonical variant (with no shredded tree), which would recurse +/// forever. The generic fallback therefore compares logical variant scalars row-by-row. Encodings +/// that can compare their typed/value children more cheaply (e.g. `ParquetVariant`) register an +/// aggregate kernel that intercepts the comparison before it reaches this fallback. +/// +/// [`all_non_distinct`]: super::all_non_distinct +pub(super) fn check_variant_identical( + lhs: &VariantArray, + rhs: &VariantArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if lhs.len() != rhs.len() { + return Ok(false); + } + for idx in 0..lhs.len() { + if lhs.execute_scalar(idx, ctx)? != rhs.execute_scalar(idx, ctx)? { + return Ok(false); + } + } + Ok(true) +} diff --git a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs index f8d371cb776..fe8527da966 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -33,11 +35,20 @@ impl AggregateFnVTable for AllNonNan { type Partial = bool; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.all_non_nan") + static ID: CachedId = CachedId::new("vortex.all_non_nan"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - Ok(None) + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { @@ -114,12 +125,12 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::all_non_nan::AllNonNan; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::Nullability; @@ -127,7 +138,7 @@ mod tests { #[test] fn all_non_nan_aggregate_fn() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNonNan, EmptyOptions, dtype)?; @@ -140,7 +151,7 @@ mod tests { #[test] fn all_non_nan_false_with_nan() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNonNan, EmptyOptions, dtype)?; @@ -169,7 +180,7 @@ mod tests { #[test] fn all_non_nan_true_with_nulls() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNonNan, EmptyOptions, dtype)?; diff --git a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs index a297c318fc9..c07dbb907c9 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -25,11 +27,20 @@ impl AggregateFnVTable for AllNonNull { type Partial = bool; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.all_non_null") + static ID: CachedId = CachedId::new("vortex.all_non_null"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - Ok(None) + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option { @@ -104,12 +115,12 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::all_non_null::AllNonNull; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::Nullability; @@ -117,7 +128,7 @@ mod tests { #[test] fn all_non_null_aggregate_fn() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNonNull, EmptyOptions, dtype)?; @@ -130,7 +141,7 @@ mod tests { #[test] fn all_non_null_false_with_nulls() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNonNull, EmptyOptions, dtype)?; @@ -143,7 +154,7 @@ mod tests { #[test] fn all_non_null_true_for_empty_input() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNonNull, EmptyOptions, dtype)?; diff --git a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs index 5476c7d534b..ec64e3d5c43 100644 --- a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -25,11 +27,20 @@ impl AggregateFnVTable for AllNull { type Partial = bool; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.all_null") + static ID: CachedId = CachedId::new("vortex.all_null"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - Ok(None) + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option { @@ -107,12 +118,12 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::all_null::AllNull; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::Nullability; @@ -120,7 +131,7 @@ mod tests { #[test] fn all_null_aggregate_fn() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNull, EmptyOptions, dtype)?; @@ -133,7 +144,7 @@ mod tests { #[test] fn all_null_false_with_non_nulls() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNull, EmptyOptions, dtype)?; @@ -146,7 +157,7 @@ mod tests { #[test] fn all_null_true_for_empty_input() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNull, EmptyOptions, dtype)?; diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs index 998b6032938..bb27d9f9ab1 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs @@ -4,13 +4,16 @@ use std::fmt::Display; use std::fmt::Formatter; use std::num::NonZeroUsize; +use std::sync::LazyLock; use vortex_buffer::BufferString; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -20,16 +23,28 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::max::Max; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::min_max::min_max; +use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::StructFields; use crate::partial_ord::partial_max; use crate::scalar::Scalar; use crate::scalar::ScalarTruncation; use crate::scalar::upper_bound; +/// Field name for the bounded maximum upper-bound value in the partial state. +pub const BOUNDED_MAX_BOUND: &str = "bound"; +/// Field name for whether the partial state represents an unknown upper bound. +pub const BOUNDED_MAX_UNKNOWN: &str = "unknown"; + +static NAMES: LazyLock = + LazyLock::new(|| FieldNames::from([BOUNDED_MAX_BOUND, BOUNDED_MAX_UNKNOWN])); + /// Options for [`BoundedMax`]. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct BoundedMaxOptions { @@ -61,12 +76,8 @@ pub struct BoundedMaxPartial { } impl BoundedMaxPartial { - fn merge(&mut self, max: Scalar) { + fn merge_bound(&mut self, max: Scalar) { if max.is_null() { - // Serialized partials encode both empty input and unknown upper bounds as null. - // Treat null as unknown when merging; this may lose a bound from an empty shard, but - // it preserves pruning soundness. - self.state = BoundedMaxState::Unknown; return; } @@ -82,6 +93,32 @@ impl BoundedMaxPartial { fn unknown(&mut self) { self.state = BoundedMaxState::Unknown; } + + fn final_scalar(&self) -> VortexResult { + let dtype = self.element_dtype.as_nullable(); + match &self.state { + BoundedMaxState::Value(max) => max.cast(&dtype), + BoundedMaxState::Empty | BoundedMaxState::Unknown => Ok(Scalar::null(dtype)), + } + } +} + +/// Return the serialized partial-state dtype for [`BoundedMax`]. +/// +/// A null struct means the partial is empty. A non-null struct with a null `bound` and +/// `unknown = true` means the input has a non-null maximum but no finite upper bound could be +/// represented within the configured byte limit. +pub fn make_bounded_max_partial_dtype(element_dtype: &DType) -> DType { + DType::Struct( + StructFields::new( + NAMES.clone(), + vec![ + element_dtype.as_nullable(), + DType::Bool(Nullability::NonNullable), + ], + ), + Nullability::Nullable, + ) } impl AggregateFnVTable for BoundedMax { @@ -89,7 +126,8 @@ impl AggregateFnVTable for BoundedMax { type Partial = BoundedMaxPartial; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.bounded_max") + static ID: CachedId = CachedId::new("vortex.bounded_max"); + *ID } fn serialize(&self, options: &Self::Options) -> VortexResult>> { @@ -136,7 +174,11 @@ impl AggregateFnVTable for BoundedMax { }; } - if requested.is::() { + // The stored bound skips NaNs, so it cannot stand in for a NaN-including maximum. + if requested + .as_opt::() + .is_some_and(|options| options.skip_nans) + { AggregateFnSatisfaction::Approximate } else { AggregateFnSatisfaction::No @@ -144,7 +186,7 @@ impl AggregateFnVTable for BoundedMax { } fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { - self.return_dtype(options, input_dtype) + supported_dtype(options, input_dtype).map(make_bounded_max_partial_dtype) } fn empty_partial( @@ -160,15 +202,50 @@ impl AggregateFnVTable for BoundedMax { } fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - partial.merge(other); + if other.is_null() { + return Ok(()); + } + + let Some(other) = other.as_struct_opt() else { + vortex_bail!("BoundedMax partial must be a struct, got {}", other.dtype()); + }; + let Some(bound) = other.field_by_idx(0) else { + vortex_bail!("BoundedMax partial is missing its bound field"); + }; + let Some(unknown) = other + .field_by_idx(1) + .and_then(|unknown| unknown.as_bool().value()) + else { + vortex_bail!("BoundedMax partial is missing its non-null unknown field"); + }; + + if unknown { + partial.unknown(); + } else { + partial.merge_bound(bound); + } Ok(()) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - let dtype = partial.element_dtype.as_nullable(); + let dtype = make_bounded_max_partial_dtype(&partial.element_dtype); + let bound_dtype = partial.element_dtype.as_nullable(); match &partial.state { - BoundedMaxState::Value(max) => max.cast(&dtype), - BoundedMaxState::Empty | BoundedMaxState::Unknown => Ok(Scalar::null(dtype)), + BoundedMaxState::Empty => Ok(Scalar::null(dtype)), + BoundedMaxState::Value(max) => Ok(Scalar::struct_( + dtype, + vec![ + max.cast(&bound_dtype)?, + Scalar::bool(false, Nullability::NonNullable), + ], + )), + BoundedMaxState::Unknown => Ok(Scalar::struct_( + dtype, + vec![ + Scalar::null(bound_dtype), + Scalar::bool(true, Nullability::NonNullable), + ], + )), } } @@ -192,28 +269,28 @@ impl AggregateFnVTable for BoundedMax { Columnar::Canonical(canonical) => canonical.clone().into_array(), Columnar::Constant(constant) => constant.clone().into_array(), }; - let Some(result) = min_max(&array, ctx)? else { + let Some(result) = min_max(&array, ctx, NumericalAggregateOpts::default())? else { return Ok(()); }; match truncate_max(result.max, partial.max_bytes.get())? { - Some(bound) => partial.merge(bound), + Some(bound) => partial.merge_bound(bound), None => partial.unknown(), } Ok(()) } fn finalize(&self, partials: ArrayRef) -> VortexResult { - Ok(partials) + partials.get_item(BOUNDED_MAX_BOUND) } fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) + partial.final_scalar() } } fn supported_dtype<'a>(_options: &BoundedMaxOptions, input_dtype: &'a DType) -> Option<&'a DType> { MinMax - .return_dtype(&EmptyOptions, input_dtype) + .return_dtype(&NumericalAggregateOpts::default(), input_dtype) .map(|_| input_dtype) } @@ -246,23 +323,23 @@ mod tests { use vortex_session::VortexSession; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnSatisfaction; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::bounded_max::BoundedMax; use crate::aggregate_fn::fns::bounded_max::BoundedMaxOptions; + use crate::aggregate_fn::fns::bounded_max::make_bounded_max_partial_dtype; use crate::aggregate_fn::fns::max::Max; use crate::aggregate_fn::fns::min::Min; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; use crate::dtype::Nullability; use crate::scalar::Scalar; - use crate::session::ArraySession; use crate::validity::Validity; fn max_bytes(value: usize) -> NonZeroUsize { @@ -270,12 +347,12 @@ mod tests { } fn fresh_session() -> VortexSession { - VortexSession::empty().with::() + array_session() } #[test] fn bounded_max_truncates_utf8_to_upper_bound() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = VarBinViewArray::from_iter_str(["aardvark", "char🪩"]).into_array(); let mut acc = Accumulator::try_new( BoundedMax, @@ -293,7 +370,7 @@ mod tests { #[test] fn bounded_max_unknown_upper_bound_returns_null() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = VarBinViewArray::from_iter_bin([&[255u8, 255, 255][..]]).into_array(); let mut acc = Accumulator::try_new( BoundedMax, @@ -311,7 +388,7 @@ mod tests { #[test] fn bounded_max_empty_does_not_poison_later_values() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let empty = VarBinViewArray::from_iter_bin(Vec::<&[u8]>::new()).into_array(); let values = VarBinViewArray::from_iter_bin([&[1u8][..]]).into_array(); let mut acc = Accumulator::try_new( @@ -334,7 +411,7 @@ mod tests { #[test] fn bounded_max_unknown_poisons_later_values() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let unknown = VarBinViewArray::from_iter_bin([&[255u8, 255, 255][..]]).into_array(); let values = VarBinViewArray::from_iter_bin([&[1u8][..]]).into_array(); let mut acc = Accumulator::try_new( @@ -353,7 +430,29 @@ mod tests { } #[test] - fn bounded_max_null_partial_poisons_existing_bound() -> VortexResult<()> { + fn bounded_max_empty_partial_does_not_poison_existing_bound() -> VortexResult<()> { + let mut ctx = fresh_session().create_execution_ctx(); + let values = VarBinViewArray::from_iter_bin([&[1u8][..]]).into_array(); + let mut acc = Accumulator::try_new( + BoundedMax, + BoundedMaxOptions { + max_bytes: max_bytes(2), + }, + values.dtype().clone(), + )?; + + acc.accumulate(&values, &mut ctx)?; + acc.combine_partials(Scalar::null(make_bounded_max_partial_dtype(values.dtype())))?; + + assert_eq!( + acc.finish()?, + Scalar::binary(buffer![1u8], Nullability::Nullable) + ); + Ok(()) + } + + #[test] + fn bounded_max_unknown_partial_poisons_existing_bound() -> VortexResult<()> { let mut ctx = fresh_session().create_execution_ctx(); let values = VarBinViewArray::from_iter_bin([&[1u8][..]]).into_array(); let mut acc = Accumulator::try_new( @@ -364,8 +463,17 @@ mod tests { values.dtype().clone(), )?; + let partial_dtype = make_bounded_max_partial_dtype(values.dtype()); + let unknown = Scalar::struct_( + partial_dtype, + vec![ + Scalar::null(values.dtype().as_nullable()), + Scalar::bool(true, Nullability::NonNullable), + ], + ); + acc.accumulate(&values, &mut ctx)?; - acc.combine_partials(Scalar::null(values.dtype().as_nullable()))?; + acc.combine_partials(unknown)?; assert_eq!(acc.finish()?, Scalar::null(values.dtype().as_nullable())); Ok(()) @@ -373,7 +481,7 @@ mod tests { #[test] fn bounded_max_keeps_fixed_width_values_exact() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::new(buffer![10i32, 20, 5], Validity::NonNullable).into_array(); let mut acc = Accumulator::try_new( BoundedMax, @@ -417,15 +525,25 @@ mod tests { AggregateFnSatisfaction::No ); assert_eq!( - stored.can_satisfy(&Max.bind(EmptyOptions)), + stored.can_satisfy(&Max.bind(NumericalAggregateOpts::default())), AggregateFnSatisfaction::Approximate ); assert_eq!( - Max.bind(EmptyOptions).can_satisfy(&stored), + stored.can_satisfy(&Max.bind(NumericalAggregateOpts::include_nans())), + AggregateFnSatisfaction::No + ); + assert_eq!( + Max.bind(NumericalAggregateOpts::include_nans()) + .can_satisfy(&stored), + AggregateFnSatisfaction::No + ); + assert_eq!( + Max.bind(NumericalAggregateOpts::default()) + .can_satisfy(&stored), AggregateFnSatisfaction::Approximate ); assert_eq!( - stored.can_satisfy(&Min.bind(EmptyOptions)), + stored.can_satisfy(&Min.bind(NumericalAggregateOpts::default())), AggregateFnSatisfaction::No ); } diff --git a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs index 7d24457e150..7a72442550b 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs @@ -11,6 +11,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -20,7 +21,7 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min::Min; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::min_max::min_max; @@ -79,7 +80,8 @@ impl AggregateFnVTable for BoundedMin { type Partial = BoundedMinPartial; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.bounded_min") + static ID: CachedId = CachedId::new("vortex.bounded_min"); + *ID } fn serialize(&self, options: &Self::Options) -> VortexResult>> { @@ -126,7 +128,11 @@ impl AggregateFnVTable for BoundedMin { }; } - if requested.is::() { + // The stored bound skips NaNs, so it cannot stand in for a NaN-including minimum. + if requested + .as_opt::() + .is_some_and(|options| options.skip_nans) + { AggregateFnSatisfaction::Approximate } else { AggregateFnSatisfaction::No @@ -182,7 +188,7 @@ impl AggregateFnVTable for BoundedMin { Columnar::Canonical(canonical) => canonical.clone().into_array(), Columnar::Constant(constant) => constant.clone().into_array(), }; - let Some(result) = min_max(&array, ctx)? else { + let Some(result) = min_max(&array, ctx, NumericalAggregateOpts::default())? else { return Ok(()); }; if let Some(bound) = truncate_min(result.min, partial.max_bytes.get())? { @@ -202,7 +208,7 @@ impl AggregateFnVTable for BoundedMin { fn supported_dtype<'a>(_options: &BoundedMinOptions, input_dtype: &'a DType) -> Option<&'a DType> { MinMax - .return_dtype(&EmptyOptions, input_dtype) + .return_dtype(&NumericalAggregateOpts::default(), input_dtype) .map(|_| input_dtype) } @@ -234,23 +240,22 @@ mod tests { use vortex_session::VortexSession; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnSatisfaction; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::bounded_min::BoundedMin; use crate::aggregate_fn::fns::bounded_min::BoundedMinOptions; use crate::aggregate_fn::fns::max::Max; use crate::aggregate_fn::fns::min::Min; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; use crate::dtype::Nullability; use crate::scalar::Scalar; - use crate::session::ArraySession; use crate::validity::Validity; fn max_bytes(value: usize) -> NonZeroUsize { @@ -258,12 +263,12 @@ mod tests { } fn fresh_session() -> VortexSession { - VortexSession::empty().with::() + array_session() } #[test] fn bounded_min_truncates_utf8_to_lower_bound() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = VarBinViewArray::from_iter_str(["snowman⛄️snowman", "untruncated"]).into_array(); let mut acc = Accumulator::try_new( @@ -285,7 +290,7 @@ mod tests { #[test] fn bounded_min_keeps_fixed_width_values_exact() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::new(buffer![10i32, 20, 5], Validity::NonNullable).into_array(); let mut acc = Accumulator::try_new( BoundedMin, @@ -351,15 +356,25 @@ mod tests { AggregateFnSatisfaction::No ); assert_eq!( - stored.can_satisfy(&Min.bind(EmptyOptions)), + stored.can_satisfy(&Min.bind(NumericalAggregateOpts::default())), AggregateFnSatisfaction::Approximate ); assert_eq!( - Min.bind(EmptyOptions).can_satisfy(&stored), + stored.can_satisfy(&Min.bind(NumericalAggregateOpts::include_nans())), + AggregateFnSatisfaction::No + ); + assert_eq!( + Min.bind(NumericalAggregateOpts::include_nans()) + .can_satisfy(&stored), + AggregateFnSatisfaction::No + ); + assert_eq!( + Min.bind(NumericalAggregateOpts::default()) + .can_satisfy(&stored), AggregateFnSatisfaction::Approximate ); assert_eq!( - stored.can_satisfy(&Max.bind(EmptyOptions)), + stored.can_satisfy(&Max.bind(NumericalAggregateOpts::default())), AggregateFnSatisfaction::No ); } diff --git a/vortex-array/src/aggregate_fn/fns/count/grouped.rs b/vortex-array/src/aggregate_fn/fns/count/grouped.rs index fb94489dde0..39a957530bf 100644 --- a/vortex-array/src/aggregate_fn/fns/count/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/count/grouped.rs @@ -27,7 +27,12 @@ impl DynGroupedAggregateKernel for CountGroupedKernel { groups: &GroupedArray, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if !aggregate_fn.is::() { + let Some(options) = aggregate_fn.as_opt::() else { + return Ok(None); + }; + // NaN-skipping counts over floats must inspect the element values, which this + // validity-only kernel cannot do; fall back to the per-group accumulator path. + if options.skip_nans && groups.elements().dtype().is_float() { return Ok(None); } try_grouped_count(groups, ctx) @@ -86,13 +91,14 @@ mod tests { use vortex_error::VortexResult; use crate::ArrayRef; + use crate::ExecutionCtx; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::DynGroupedAccumulator; - use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::GroupedAccumulator; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::count::Count; + use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -105,9 +111,17 @@ mod tests { use crate::validity::Validity; /// Run a grouped count through the accumulator. - fn grouped_count_actual(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { - let mut acc = GroupedAccumulator::try_new(Count, EmptyOptions, elem_dtype.clone())?; - acc.accumulate_list(groups, &mut LEGACY_SESSION.create_execution_ctx())?; + fn grouped_count_actual( + groups: &ArrayRef, + elem_dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let mut acc = GroupedAccumulator::try_new( + Count, + NumericalAggregateOpts::default(), + elem_dtype.clone(), + )?; + acc.accumulate_list(groups, ctx)?; acc.finish() } @@ -116,7 +130,7 @@ mod tests { elements: &ArrayRef, ranges: &[(usize, usize)], ) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let counts: Buffer = ranges .iter() .map(|&(offset, size)| { @@ -142,24 +156,26 @@ mod tests { #[test] fn listview_counts_all_valid() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::I32, NonNullable); let ranges = [(0, 2), (2, 1), (3, 3), (6, 0)]; let groups = listview(elements.clone(), &ranges)?; - let actual = grouped_count_actual(&groups, &elem_dtype)?; + let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; let expected = grouped_count_reference(&elements, &ranges)?; let direct = PrimitiveArray::new(buffer![2u64, 1, 3, 0], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct, &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn listview_counts_with_nulls() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) .into_array(); @@ -167,18 +183,19 @@ mod tests { let ranges = [(0, 3), (3, 2), (5, 1)]; let groups = listview(elements.clone(), &ranges)?; - let actual = grouped_count_actual(&groups, &elem_dtype)?; + let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; let expected = grouped_count_reference(&elements, &ranges)?; // Group 0: {1, null, 3} -> 2. Group 1: {null, null} -> 0. Group 2: {9} -> 1. let direct = PrimitiveArray::new(buffer![2u64, 0, 1], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct, &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn listview_counts_varbinview_with_nulls() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements = VarBinViewArray::from_iter_nullable_str([ Some("a"), None, @@ -191,26 +208,51 @@ mod tests { let ranges = [(0, 2), (2, 2), (4, 1)]; let groups = listview(elements.clone(), &ranges)?; - let actual = grouped_count_actual(&groups, &elem_dtype)?; + let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; let expected = grouped_count_reference(&elements, &ranges)?; let direct = PrimitiveArray::new(buffer![1u64, 1, 1], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct, &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) + } + + #[test] + fn fixed_size_counts_float_nans() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = + PrimitiveArray::from_option_iter([Some(1.0f64), Some(f64::NAN), None, Some(2.0)]) + .into_array(); + let elem_dtype = DType::Primitive(PType::F64, Nullable); + let groups = + FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?.into_array(); + + // NaNs are excluded by default and counted otherwise. + let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; + let expected = PrimitiveArray::new(buffer![1u64, 1], Validity::NonNullable).into_array(); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + let mut acc = + GroupedAccumulator::try_new(Count, NumericalAggregateOpts::include_nans(), elem_dtype)?; + acc.accumulate_list(&groups, &mut ctx)?; + let actual = acc.finish()?; + let expected = PrimitiveArray::new(buffer![2u64, 1], Validity::NonNullable).into_array(); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn fixed_size_counts_with_nulls() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(); let elem_dtype = DType::Primitive(PType::I32, Nullable); let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?.into_array(); - let actual = grouped_count_actual(&groups, &elem_dtype)?; + let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; let direct = PrimitiveArray::new(buffer![1u64, 2], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct); + assert_arrays_eq!(&actual, &direct, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index 1fe984fb099..8f7d68027bc 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -5,13 +5,15 @@ mod grouped; pub(crate) use grouped::CountGroupedKernel; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; use crate::ExecutionCtx; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::nan_count::nan_count; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -21,15 +23,27 @@ use crate::scalar::Scalar; /// /// Applies to all types. Returns a `u64` count. /// The identity value is zero. +/// +/// For float inputs, NaN handling is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the +/// default) NaN values are treated as missing and excluded from the count, otherwise they are +/// counted like any other non-null value. #[derive(Clone, Debug)] pub struct Count; +/// Partial accumulator state for the count aggregate. +pub struct CountPartial { + count: u64, + /// Whether NaN values must be excluded from the count (float input with `skip_nans`). + exclude_nans: bool, +} + impl AggregateFnVTable for Count { - type Options = EmptyOptions; - type Partial = u64; + type Options = NumericalAggregateOpts; + type Partial = CountPartial; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.count") + static ID: CachedId = CachedId::new("vortex.count"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { @@ -46,10 +60,13 @@ impl AggregateFnVTable for Count { fn empty_partial( &self, - _options: &Self::Options, - _input_dtype: &DType, + options: &Self::Options, + input_dtype: &DType, ) -> VortexResult { - Ok(0u64) + Ok(CountPartial { + count: 0, + exclude_nans: options.skip_nans && input_dtype.is_float(), + }) } fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { @@ -57,16 +74,16 @@ impl AggregateFnVTable for Count { .as_primitive() .typed_value::() .vortex_expect("count partial should not be null"); - *partial += val; + partial.count += val; Ok(()) } fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(Scalar::primitive(*partial, Nullability::NonNullable)) + Ok(Scalar::primitive(partial.count, Nullability::NonNullable)) } fn reset(&self, partial: &mut Self::Partial) { - *partial = 0; + partial.count = 0; } #[inline] @@ -80,7 +97,12 @@ impl AggregateFnVTable for Count { batch: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { - *state += batch.valid_count(ctx)? as u64; + let mut count = batch.valid_count(ctx)? as u64; + if state.exclude_nans { + // `nan_count` shortcircuits on an exact `Stat::NaNCount` before scanning the batch. + count = count.saturating_sub(nan_count(batch, ctx)? as u64); + } + state.count += count; Ok(true) } @@ -104,19 +126,21 @@ impl AggregateFnVTable for Count { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::count::Count; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -124,11 +148,20 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; 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; + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + pub fn count(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let mut acc = Accumulator::try_new(Count, EmptyOptions, array.dtype().clone())?; + let mut acc = Accumulator::try_new( + Count, + NumericalAggregateOpts::default(), + array.dtype().clone(), + )?; acc.accumulate(array, ctx)?; let result = acc.finish()?; @@ -144,7 +177,7 @@ mod tests { fn count_all_valid() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5], Validity::NonNullable).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_eq!(count(&array, &mut ctx)?, 5); Ok(()) } @@ -153,7 +186,7 @@ mod tests { fn count_with_nulls() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]) .into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_eq!(count(&array, &mut ctx)?, 3); Ok(()) } @@ -161,7 +194,7 @@ mod tests { #[test] fn count_all_null() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_eq!(count(&array, &mut ctx)?, 0); Ok(()) } @@ -169,7 +202,7 @@ mod tests { #[test] fn count_empty() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Count, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Count, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) @@ -177,9 +210,9 @@ mod tests { #[test] fn count_multi_batch() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); - let mut acc = Accumulator::try_new(Count, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Count, NumericalAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -194,9 +227,9 @@ mod tests { #[test] fn count_finish_resets_state() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); - let mut acc = Accumulator::try_new(Count, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Count, NumericalAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::from_option_iter([Some(1i32), None]).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -213,7 +246,7 @@ mod tests { #[test] fn count_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = Count.empty_partial(&EmptyOptions, &dtype)?; + let mut state = Count.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; let scalar1 = Scalar::primitive(5u64, Nullability::NonNullable); Count.combine_partials(&mut state, scalar1)?; @@ -227,10 +260,73 @@ mod tests { Ok(()) } + fn count_with_options( + array: &ArrayRef, + ctx: &mut ExecutionCtx, + options: NumericalAggregateOpts, + ) -> VortexResult { + let mut acc = Accumulator::try_new(Count, options, array.dtype().clone())?; + acc.accumulate(array, ctx)?; + Ok(acc + .finish()? + .as_primitive() + .typed_value::() + .vortex_expect("count result should not be null")) + } + + #[test] + fn count_float_excludes_nans_by_default() -> VortexResult<()> { + let array = + PrimitiveArray::from_option_iter([Some(1.0f64), Some(f64::NAN), None, Some(3.0)]) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!(count(&array, &mut ctx)?, 2); + Ok(()) + } + + #[test] + fn count_float_includes_nans_when_not_skipping() -> VortexResult<()> { + let array = + PrimitiveArray::from_option_iter([Some(1.0f64), Some(f64::NAN), None, Some(3.0)]) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!( + count_with_options(&array, &mut ctx, NumericalAggregateOpts::include_nans())?, + 3 + ); + Ok(()) + } + + #[test] + fn count_float_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> { + // The array has no NaNs; a planted exact NaNCount stat proves the count is derived from + // the stat rather than a scan. + let array = + PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0, 4.0], Validity::NonNullable).into_array(); + array + .statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(3u64))); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!(count(&array, &mut ctx)?, 1); + Ok(()) + } + + #[test] + fn count_constant_nan() -> VortexResult<()> { + let array = ConstantArray::new(f64::NAN, 5).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!(count(&array, &mut ctx)?, 0); + assert_eq!( + count_with_options(&array, &mut ctx, NumericalAggregateOpts::include_nans())?, + 5 + ); + Ok(()) + } + #[test] fn count_constant_non_null() -> VortexResult<()> { let array = ConstantArray::new(42i32, 10); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_eq!(count(&array.into_array(), &mut ctx)?, 10); Ok(()) } @@ -241,7 +337,7 @@ mod tests { Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), 10, ); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_eq!(count(&array.into_array(), &mut ctx)?, 0); Ok(()) } @@ -252,7 +348,7 @@ mod tests { let chunk2 = PrimitiveArray::from_option_iter([None, Some(5i32), None]); let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_eq!(count(&chunked.into_array(), &mut ctx)?, 3); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/first/mod.rs b/vortex-array/src/aggregate_fn/fns/first/mod.rs index 22e9ef7e39d..c41e2575057 100644 --- a/vortex-array/src/aggregate_fn/fns/first/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/first/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -40,7 +41,8 @@ impl AggregateFnVTable for First { type Partial = FirstPartial; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.first") + static ID: CachedId = CachedId::new("vortex.first"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { @@ -130,7 +132,6 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; @@ -138,6 +139,7 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::first::First; use crate::aggregate_fn::fns::first::first; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -152,7 +154,7 @@ mod tests { #[test] fn first_non_null() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![10i32, 20, 30], Validity::NonNullable).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(first(&array, &mut ctx)?, Scalar::primitive(10i32, Nullable)); Ok(()) } @@ -161,7 +163,7 @@ mod tests { fn first_skips_leading_nulls() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter([None, None, Some(7i32), Some(8)]).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(first(&array, &mut ctx)?, Scalar::primitive(7i32, Nullable)); Ok(()) } @@ -169,7 +171,7 @@ mod tests { #[test] fn first_all_null() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullable); assert_eq!(first(&array, &mut ctx)?, Scalar::null(dtype)); Ok(()) @@ -187,7 +189,7 @@ mod tests { #[test] fn first_constant() -> VortexResult<()> { let array = ConstantArray::new(42i32, 10).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(first(&array, &mut ctx)?, Scalar::primitive(42i32, Nullable)); Ok(()) } @@ -196,7 +198,7 @@ mod tests { fn first_constant_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullable); let array = ConstantArray::new(Scalar::null(dtype.clone()), 10).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(first(&array, &mut ctx)?, Scalar::null(dtype)); Ok(()) } @@ -208,14 +210,14 @@ mod tests { DType::Utf8(Nullable), ) .into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(first(&array, &mut ctx)?, Scalar::utf8("hello", Nullable)); Ok(()) } #[test] fn first_multi_batch_picks_earliest_non_null() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullable); let mut acc = Accumulator::try_new(First, EmptyOptions, dtype)?; @@ -240,7 +242,7 @@ mod tests { #[test] fn first_finish_resets_state() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = Accumulator::try_new(First, EmptyOptions, dtype)?; @@ -278,7 +280,7 @@ mod tests { let chunk2 = PrimitiveArray::from_option_iter([None, Some(42i32), Some(100)]); let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!( first(&chunked.into_array(), &mut ctx)?, Scalar::primitive(42i32, Nullable) 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 b14b2a050f8..490ede7f640 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -13,7 +13,7 @@ mod varbin; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_mask::Mask; +use vortex_session::registry::CachedId; use self::bool::check_bool_constant; use self::decimal::check_decimal_constant; @@ -50,7 +50,6 @@ use crate::scalar_fn::fns::operators::Operator; /// Check if two arrays of the same length have equal values at every position (null-safe). /// /// Two positions are considered equal if they are both null, or both non-null with the same value. -/// // TODO(ngates): move this function out when we have any/all aggregate functions. fn arrays_value_equal(a: &ArrayRef, b: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { debug_assert_eq!(a.len(), b.len()); @@ -74,7 +73,7 @@ fn arrays_value_equal(a: &ArrayRef, b: &ArrayRef, ctx: &mut ExecutionCtx) -> Vor // Compare values element-wise. Result is null where both inputs are null, // true/false where both are valid. let eq_result = a.binary(b.clone(), Operator::Eq)?; - let eq_result = eq_result.execute::(ctx)?; + let eq_result = eq_result.null_as_false().execute(ctx)?; Ok(eq_result.true_count() == valid_count) } @@ -260,7 +259,8 @@ impl AggregateFnVTable for IsConstant { type Partial = IsConstantPartial; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.is_constant") + static ID: CachedId = CachedId::new("vortex.is_constant"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { @@ -441,9 +441,9 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray as _; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_constant::is_constant; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::DecimalArray; @@ -461,7 +461,7 @@ mod tests { // Tests migrated from compute/is_constant.rs #[test] fn is_constant_min_max_no_nan() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = buffer![0, 1].into_array(); arr.statistics() @@ -480,7 +480,7 @@ mod tests { #[test] fn is_constant_min_max_with_nan() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::from_iter([0.0, 0.0, f32::NAN]).into_array(); arr.statistics() @@ -507,7 +507,7 @@ mod tests { }, false)] fn test_bool_is_constant(#[case] input: Vec, #[case] expected: bool) -> VortexResult<()> { let array = BoolArray::from_iter(input); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(is_constant(&array.into_array(), &mut ctx)?, expected); Ok(()) } @@ -527,7 +527,7 @@ mod tests { )? .into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert!(is_constant(&chunked, &mut ctx)?); Ok(()) } @@ -535,7 +535,7 @@ mod tests { // Tests migrated from arrays/decimal/compute/is_constant.rs #[test] fn test_decimal_is_constant() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::new( buffer![0i128, 1i128, 2i128], @@ -556,7 +556,7 @@ mod tests { // Tests migrated from arrays/list/compute/is_constant.rs #[test] fn test_is_constant_nested_list() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let xs = ListArray::try_new( buffer![0i32, 1, 0, 1].into_array(), @@ -614,7 +614,7 @@ mod tests { Validity::NonNullable, )?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(is_constant(&list_array.into_array(), &mut ctx)?, expected); Ok(()) } @@ -632,7 +632,7 @@ mod tests { Validity::NonNullable, )?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Both outer lists contain [[1], [2]], so should be constant assert!(is_constant(&outer_list.into_array(), &mut ctx)?); Ok(()) @@ -677,7 +677,7 @@ mod tests { Validity::NonNullable, )?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(is_constant(&list_array.into_array(), &mut ctx)?, expected); Ok(()) } 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 a4a94fc5211..c3ee7f5d4e7 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -13,6 +13,7 @@ use std::fmt::Formatter; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use self::bool::check_bool_sorted; use self::decimal::check_decimal_sorted; @@ -232,7 +233,8 @@ impl AggregateFnVTable for IsSorted { type Partial = IsSortedPartial; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.is_sorted") + static ID: CachedId = CachedId::new("vortex.is_sorted"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { @@ -487,7 +489,7 @@ impl AggregateFnVTable for IsSorted { let batch_is_sorted = match c { Canonical::Primitive(p) => check_primitive_sorted(p, partial.strict, ctx)?, Canonical::Bool(b) => check_bool_sorted(b, partial.strict, ctx)?, - Canonical::VarBinView(v) => check_varbinview_sorted(v, partial.strict)?, + Canonical::VarBinView(v) => check_varbinview_sorted(v, partial.strict, ctx)?, Canonical::Decimal(d) => check_decimal_sorted(d, partial.strict, ctx)?, Canonical::Extension(e) => check_extension_sorted(e, partial.strict, ctx)?, Canonical::Null(_) => !partial.strict, @@ -558,10 +560,10 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_sorted::is_sorted; use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::validity::Validity; @@ -569,7 +571,7 @@ mod tests { // Tests migrated from compute/is_sorted.rs #[test] fn test_is_sorted() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new(buffer!(0, 1, 2, 3), Validity::AllValid).into_array(); assert!(is_sorted(&arr, &mut ctx)?); @@ -603,7 +605,7 @@ mod tests { #[test] fn test_is_strict_sorted() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new(buffer!(0, 1, 2, 3), Validity::AllValid).into_array(); assert!(is_strict_sorted(&arr, &mut ctx)?); @@ -640,7 +642,7 @@ mod tests { #[case(PrimitiveArray::from_option_iter([None, None, Some(1i32), Some(1)]), true)] #[case(PrimitiveArray::from_option_iter([None, Some(5_u8), None]), false)] fn test_primitive_is_sorted(#[case] array: PrimitiveArray, #[case] expected: bool) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!( is_sorted(&array.into_array(), &mut ctx) .vortex_expect("operation should succeed in test"), @@ -655,7 +657,7 @@ mod tests { #[case(PrimitiveArray::from_option_iter([None, None, Some(1i32), Some(1), None]), false)] #[case(PrimitiveArray::from_option_iter([None, Some(5_u8), None]), false)] fn test_primitive_is_strict_sorted(#[case] array: PrimitiveArray, #[case] expected: bool) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!( is_strict_sorted(&array.into_array(), &mut ctx) .vortex_expect("operation should succeed in test"), @@ -672,7 +674,7 @@ mod tests { use crate::arrays::DecimalArray; use crate::dtype::DecimalDType; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DecimalDType::new(19, 2); let i100 = parse_decimal::("100.00", dtype.precision(), dtype.scale())?; let i200 = parse_decimal::("200.00", dtype.precision(), dtype.scale())?; @@ -697,7 +699,7 @@ mod tests { use crate::arrays::DecimalArray; use crate::dtype::DecimalDType; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DecimalDType::new(19, 2); let i100 = parse_decimal::("100.00", dtype.precision(), dtype.scale())?; let i200 = parse_decimal::("200.00", dtype.precision(), dtype.scale())?; diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/varbin.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/varbin.rs index 84529ef16c7..9e17e12e291 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/varbin.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/varbin.rs @@ -4,13 +4,37 @@ use vortex_error::VortexResult; use super::IsSortedIteratorExt; -use crate::accessor::ArrayAccessor; +use crate::ExecutionCtx; use crate::arrays::VarBinViewArray; -pub(super) fn check_varbinview_sorted(array: &VarBinViewArray, strict: bool) -> VortexResult { +pub(super) fn check_varbinview_sorted( + array: &VarBinViewArray, + strict: bool, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = array + .validity()? + .execute_mask(array.len(), ctx)? + .to_bit_buffer(); + let views = array.views(); + let buffers = array + .data_buffers() + .iter() + .map(|b| b.as_host()) + .collect::>(); + let iter = views.iter().zip(mask.iter()).map(|(view, valid)| { + valid.then(|| { + 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()] + } + }) + }); Ok(if strict { - array.with_iterator(|bytes_iter| bytes_iter.is_strict_sorted()) + iter.is_strict_sorted() } else { - array.with_iterator(|bytes_iter| bytes_iter.is_sorted()) + iter.is_sorted() }) } diff --git a/vortex-array/src/aggregate_fn/fns/last/mod.rs b/vortex-array/src/aggregate_fn/fns/last/mod.rs index a0034d0820b..63ee54a6efa 100644 --- a/vortex-array/src/aggregate_fn/fns/last/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/last/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -40,7 +41,8 @@ impl AggregateFnVTable for Last { type Partial = LastPartial; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.last") + static ID: CachedId = CachedId::new("vortex.last"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { @@ -128,7 +130,6 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; @@ -136,6 +137,7 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::last::Last; use crate::aggregate_fn::fns::last::last; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -150,7 +152,7 @@ mod tests { #[test] fn last_non_null() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![10i32, 20, 30], Validity::NonNullable).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(last(&array, &mut ctx)?, Scalar::primitive(30i32, Nullable)); Ok(()) } @@ -159,7 +161,7 @@ mod tests { fn last_skips_trailing_nulls() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter([Some(7i32), Some(8), None, None]).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(last(&array, &mut ctx)?, Scalar::primitive(8i32, Nullable)); Ok(()) } @@ -167,7 +169,7 @@ mod tests { #[test] fn last_all_null() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullable); assert_eq!(last(&array, &mut ctx)?, Scalar::null(dtype)); Ok(()) @@ -185,7 +187,7 @@ mod tests { #[test] fn last_constant() -> VortexResult<()> { let array = ConstantArray::new(42i32, 10).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(last(&array, &mut ctx)?, Scalar::primitive(42i32, Nullable)); Ok(()) } @@ -194,7 +196,7 @@ mod tests { fn last_constant_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullable); let array = ConstantArray::new(Scalar::null(dtype.clone()), 10).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(last(&array, &mut ctx)?, Scalar::null(dtype)); Ok(()) } @@ -206,14 +208,14 @@ mod tests { DType::Utf8(Nullable), ) .into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(last(&array, &mut ctx)?, Scalar::utf8("world", Nullable)); Ok(()) } #[test] fn last_multi_batch_picks_latest_non_null() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullable); let mut acc = Accumulator::try_new(Last, EmptyOptions, dtype)?; @@ -237,7 +239,7 @@ mod tests { #[test] fn last_finish_resets_state() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = Accumulator::try_new(Last, EmptyOptions, dtype)?; @@ -275,7 +277,7 @@ mod tests { let chunk2 = PrimitiveArray::from_option_iter::([None, None]); let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!( last(&chunked.into_array(), &mut ctx)?, Scalar::primitive(100i32, Nullable) diff --git a/vortex-array/src/aggregate_fn/fns/max/mod.rs b/vortex-array/src/aggregate_fn/fns/max/mod.rs index 755e1a4ed35..ccf9f4d899d 100644 --- a/vortex-array/src/aggregate_fn/fns/max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/max/mod.rs @@ -3,6 +3,8 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -12,15 +14,24 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::bounded_max::BoundedMax; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::min_max::min_max; +use crate::aggregate_fn::fns::min_max::nan_scalar; +use crate::aggregate_fn::fns::min_max::scalar_is_nan; use crate::dtype::DType; +use crate::expr::stats::Precision; +use crate::expr::stats::Stat; +use crate::expr::stats::StatsProvider; +use crate::expr::stats::StatsProviderExt; use crate::partial_ord::partial_max; use crate::scalar::Scalar; /// Compute the maximum non-null value of an array. +/// +/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the +/// default) NaN values are ignored, otherwise any NaN value poisons the maximum to NaN. #[derive(Clone, Debug)] pub struct Max; @@ -28,6 +39,7 @@ pub struct Max; pub struct MaxPartial { max: Option, element_dtype: DType, + skip_nans: bool, } impl MaxPartial { @@ -36,39 +48,69 @@ impl MaxPartial { return; } + // NaN scalars are incomparable under `partial_max`; they poison the maximum when NaNs + // participate, and are dropped when they are skipped. + if scalar_is_nan(&max) || self.is_poisoned() { + if !self.skip_nans { + self.poison(); + } + return; + } + self.max = Some(match self.max.take() { Some(current) => partial_max(max, current).vortex_expect("incomparable max scalars"), None => max, }); } + + fn poison(&mut self) { + self.max = Some(nan_scalar(&self.element_dtype)); + } + + fn is_poisoned(&self) -> bool { + self.element_dtype.is_float() && self.max.as_ref().is_some_and(scalar_is_nan) + } } impl AggregateFnVTable for Max { - type Options = EmptyOptions; + type Options = NumericalAggregateOpts; type Partial = MaxPartial; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.max") + static ID: CachedId = CachedId::new("vortex.max"); + *ID } - fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - Ok(None) + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.serialize())) } - fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + NumericalAggregateOpts::deserialize(metadata) + } + + fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { MinMax - .return_dtype(&EmptyOptions, input_dtype) + .return_dtype(options, input_dtype) .map(|_| input_dtype.as_nullable()) } fn can_satisfy( &self, - _options: &Self::Options, + options: &Self::Options, requested: &AggregateFnRef, ) -> AggregateFnSatisfaction { - if requested.is::() { + if requested + .as_opt::() + .is_some_and(|other| other == options) + { AggregateFnSatisfaction::Exact - } else if requested.is::() { + } else if requested.is::() && options.skip_nans { + // A NaN-including maximum may be NaN, which is not a usable upper bound. AggregateFnSatisfaction::Approximate } else { AggregateFnSatisfaction::No @@ -81,12 +123,13 @@ impl AggregateFnVTable for Max { fn empty_partial( &self, - _options: &Self::Options, + options: &Self::Options, input_dtype: &DType, ) -> VortexResult { Ok(MaxPartial { max: None, element_dtype: input_dtype.clone(), + skip_nans: options.skip_nans, }) } @@ -107,8 +150,38 @@ impl AggregateFnVTable for Max { partial.max = None; } - fn is_saturated(&self, _partial: &Self::Partial) -> bool { - false + fn is_saturated(&self, partial: &Self::Partial) -> bool { + // A poisoned NaN-including maximum is fully determined. + partial.is_poisoned() + } + + fn try_accumulate( + &self, + partial: &mut Self::Partial, + batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + // NaN-aware shortcircuits only apply to the NaN-including float maximum; everything else + // takes the default dispatch path. + if partial.skip_nans || !partial.element_dtype.is_float() { + return Ok(false); + } + match batch.statistics().get_as::(Stat::NaNCount) { + Precision::Exact(0) => { + // NaN-free batch: the cached NaN-skipping maximum (if any) is valid. `to_scalar` + // re-casts to the result dtype, so the cached scalar can merge as-is. + if let Some(max) = batch.statistics().get(Stat::Max).as_exact() { + partial.merge(max); + return Ok(true); + } + Ok(false) + } + Precision::Exact(_) => { + partial.poison(); + Ok(true) + } + _ => Ok(false), + } } fn accumulate( @@ -123,7 +196,10 @@ impl AggregateFnVTable for Max { Columnar::Canonical(canonical) => canonical.clone().into_array(), Columnar::Constant(constant) => constant.clone().into_array(), }; - if let Some(result) = min_max(&array, ctx)? { + let options = NumericalAggregateOpts { + skip_nans: partial.skip_nans, + }; + if let Some(result) = min_max(&array, ctx, options)? { partial.merge(result.max); } Ok(()) @@ -144,12 +220,12 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray as _; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::max::Max; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::Nullability; @@ -162,9 +238,9 @@ mod tests { #[test] fn max_aggregate_fn() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Max, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Max, NumericalAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20, 5], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -182,7 +258,7 @@ mod tests { #[test] fn max_empty_group_returns_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Max, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Max, NumericalAggregateOpts::default(), dtype)?; assert_eq!( acc.finish()?, @@ -191,14 +267,90 @@ mod tests { Ok(()) } + #[test] + fn max_with_nan_not_skipping() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(Max, NumericalAggregateOpts::include_nans(), dtype)?; + + let batch = PrimitiveArray::new(buffer![1.0f64, f64::NAN, -5.0], Validity::NonNullable) + .into_array(); + acc.accumulate(&batch, &mut ctx)?; + assert!(acc.is_saturated()); + + let result = acc.finish()?; + assert!( + result + .as_primitive() + .typed_value::() + .is_some_and(f64::is_nan) + ); + Ok(()) + } + + #[test] + fn max_not_skipping_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // The array has no NaNs; a planted exact NaNCount stat proves the poisoning came from + // the stat rather than a scan. + let batch = PrimitiveArray::new(buffer![1.0f64, 2.0], Validity::NonNullable).into_array(); + batch + .statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(1u64))); + let mut acc = Accumulator::try_new( + Max, + NumericalAggregateOpts::include_nans(), + batch.dtype().clone(), + )?; + acc.accumulate(&batch, &mut ctx)?; + let result = acc.finish()?; + assert!( + result + .as_primitive() + .typed_value::() + .is_some_and(f64::is_nan) + ); + Ok(()) + } + + #[test] + fn max_nan_including_nullable_cached_stat() -> VortexResult<()> { + // A nullable float array's cached Max stat is reconstructed as a nullable scalar. The + // NaN-including shortcircuit merges it as-is; `to_scalar` re-casts to the result dtype. + let mut ctx = array_session().create_execution_ctx(); + let array = + PrimitiveArray::from_option_iter([Some(1.0f64), Some(2.0), Some(3.0)]).into_array(); + array + .statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(0u64))); + array + .statistics() + .set(Stat::Max, Precision::Exact(ScalarValue::from(3.0f64))); + let mut acc = Accumulator::try_new( + Max, + NumericalAggregateOpts::include_nans(), + array.dtype().clone(), + )?; + acc.accumulate(&array, &mut ctx)?; + assert_eq!( + acc.finish()?, + Scalar::primitive(3.0f64, Nullability::Nullable) + ); + Ok(()) + } + #[test] fn max_casts_nonnullable_legacy_stat_to_nullable_partial() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let batch = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); batch .statistics() .set(Stat::Max, Precision::Exact(ScalarValue::from(25i32))); - let mut acc = Accumulator::try_new(Max, EmptyOptions, batch.dtype().clone())?; + let mut acc = Accumulator::try_new( + Max, + NumericalAggregateOpts::default(), + batch.dtype().clone(), + )?; acc.accumulate(&batch, &mut ctx)?; diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index 17fb2616d44..b421d69a966 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -3,6 +3,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::ExecutionCtx; @@ -10,7 +11,7 @@ use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::combined::BinaryCombined; use crate::aggregate_fn::combined::Combined; use crate::aggregate_fn::combined::CombinedOptions; @@ -30,7 +31,10 @@ use crate::scalar_fn::fns::operators::Operator; pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let mut acc = Accumulator::try_new( Mean::combined(), - PairOptions(EmptyOptions, EmptyOptions), + PairOptions( + NumericalAggregateOpts::default(), + NumericalAggregateOpts::default(), + ), array.dtype().clone(), )?; acc.accumulate(array, ctx)?; @@ -59,7 +63,8 @@ impl BinaryCombined for Mean { type Right = Count; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.mean") + static ID: CachedId = CachedId::new("vortex.mean"); + *ID } fn left(&self) -> Sum { @@ -165,8 +170,8 @@ mod tests { use super::*; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -177,7 +182,7 @@ mod tests { fn mean_all_valid() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0, 4.0, 5.0], Validity::NonNullable) .into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mean(&array, &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(3.0)); Ok(()) @@ -186,7 +191,7 @@ mod tests { #[test] fn mean_with_nulls() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter([Some(2.0f64), None, Some(4.0)]).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mean(&array, &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(3.0)); Ok(()) @@ -195,7 +200,7 @@ mod tests { #[test] fn mean_integers() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![10i32, 20, 30], Validity::NonNullable).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mean(&array, &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(20.0)); Ok(()) @@ -204,7 +209,7 @@ mod tests { #[test] fn mean_bool() -> VortexResult<()> { let array: BoolArray = [true, false, true, true].into_iter().collect(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mean(&array.into_array(), &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(0.75)); Ok(()) @@ -213,7 +218,7 @@ mod tests { #[test] fn mean_constant_non_null() -> VortexResult<()> { let array = ConstantArray::new(5.0f64, 4); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mean(&array.into_array(), &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(5.0)); Ok(()) @@ -225,16 +230,44 @@ mod tests { let chunk2 = PrimitiveArray::from_option_iter([Some(5.0f64), None]); let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mean(&chunked.into_array(), &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(3.0)); Ok(()) } + #[test] + fn mean_skips_nans_by_default() -> VortexResult<()> { + // NaNs are excluded from both the sum and the count. + let array = + PrimitiveArray::new(buffer![1.0f64, f64::NAN, 3.0], Validity::NonNullable).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let result = mean(&array, &mut ctx)?; + assert_eq!(result.as_primitive().as_::(), Some(2.0)); + Ok(()) + } + + #[test] + fn mean_with_nan_not_skipping() -> VortexResult<()> { + let array = + PrimitiveArray::new(buffer![1.0f64, f64::NAN, 3.0], Validity::NonNullable).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let keep_nans = NumericalAggregateOpts::include_nans(); + let mut acc = Accumulator::try_new( + Mean::combined(), + PairOptions(keep_nans, keep_nans), + array.dtype().clone(), + )?; + acc.accumulate(&array, &mut ctx)?; + let result = acc.finish()?; + assert!(result.as_primitive().as_::().is_some_and(f64::is_nan)); + Ok(()) + } + #[test] fn mean_all_null_returns_nan() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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)); Ok(()) @@ -242,11 +275,14 @@ mod tests { #[test] fn mean_multi_batch() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); let mut acc = Accumulator::try_new( Mean::combined(), - PairOptions(EmptyOptions, EmptyOptions), + PairOptions( + NumericalAggregateOpts::default(), + NumericalAggregateOpts::default(), + ), dtype, )?; diff --git a/vortex-array/src/aggregate_fn/fns/min/mod.rs b/vortex-array/src/aggregate_fn/fns/min/mod.rs index b176bbaf742..488405e8f14 100644 --- a/vortex-array/src/aggregate_fn/fns/min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min/mod.rs @@ -3,6 +3,8 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -12,15 +14,24 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnSatisfaction; use crate::aggregate_fn::AggregateFnVTable; -use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::bounded_min::BoundedMin; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::min_max::min_max; +use crate::aggregate_fn::fns::min_max::nan_scalar; +use crate::aggregate_fn::fns::min_max::scalar_is_nan; use crate::dtype::DType; +use crate::expr::stats::Precision; +use crate::expr::stats::Stat; +use crate::expr::stats::StatsProvider; +use crate::expr::stats::StatsProviderExt; use crate::partial_ord::partial_min; use crate::scalar::Scalar; /// Compute the minimum non-null value of an array. +/// +/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the +/// default) NaN values are ignored, otherwise any NaN value poisons the minimum to NaN. #[derive(Clone, Debug)] pub struct Min; @@ -28,6 +39,7 @@ pub struct Min; pub struct MinPartial { min: Option, element_dtype: DType, + skip_nans: bool, } impl MinPartial { @@ -36,39 +48,69 @@ impl MinPartial { return; } + // NaN scalars are incomparable under `partial_min`; they poison the minimum when NaNs + // participate, and are dropped when they are skipped. + if scalar_is_nan(&min) || self.is_poisoned() { + if !self.skip_nans { + self.poison(); + } + return; + } + self.min = Some(match self.min.take() { Some(current) => partial_min(min, current).vortex_expect("incomparable min scalars"), None => min, }); } + + fn poison(&mut self) { + self.min = Some(nan_scalar(&self.element_dtype)); + } + + fn is_poisoned(&self) -> bool { + self.element_dtype.is_float() && self.min.as_ref().is_some_and(scalar_is_nan) + } } impl AggregateFnVTable for Min { - type Options = EmptyOptions; + type Options = NumericalAggregateOpts; type Partial = MinPartial; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.min") + static ID: CachedId = CachedId::new("vortex.min"); + *ID } - fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - Ok(None) + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.serialize())) } - fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + NumericalAggregateOpts::deserialize(metadata) + } + + fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { MinMax - .return_dtype(&EmptyOptions, input_dtype) + .return_dtype(options, input_dtype) .map(|_| input_dtype.as_nullable()) } fn can_satisfy( &self, - _options: &Self::Options, + options: &Self::Options, requested: &AggregateFnRef, ) -> AggregateFnSatisfaction { - if requested.is::() { + if requested + .as_opt::() + .is_some_and(|other| other == options) + { AggregateFnSatisfaction::Exact - } else if requested.is::() { + } else if requested.is::() && options.skip_nans { + // A NaN-including minimum may be NaN, which is not a usable lower bound. AggregateFnSatisfaction::Approximate } else { AggregateFnSatisfaction::No @@ -81,12 +123,13 @@ impl AggregateFnVTable for Min { fn empty_partial( &self, - _options: &Self::Options, + options: &Self::Options, input_dtype: &DType, ) -> VortexResult { Ok(MinPartial { min: None, element_dtype: input_dtype.clone(), + skip_nans: options.skip_nans, }) } @@ -107,8 +150,38 @@ impl AggregateFnVTable for Min { partial.min = None; } - fn is_saturated(&self, _partial: &Self::Partial) -> bool { - false + fn is_saturated(&self, partial: &Self::Partial) -> bool { + // A poisoned NaN-including minimum is fully determined. + partial.is_poisoned() + } + + fn try_accumulate( + &self, + partial: &mut Self::Partial, + batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + // NaN-aware shortcircuits only apply to the NaN-including float minimum; everything else + // takes the default dispatch path. + if partial.skip_nans || !partial.element_dtype.is_float() { + return Ok(false); + } + match batch.statistics().get_as::(Stat::NaNCount) { + Precision::Exact(0) => { + // NaN-free batch: the cached NaN-skipping minimum (if any) is valid. `to_scalar` + // re-casts to the result dtype, so the cached scalar can merge as-is. + if let Some(min) = batch.statistics().get(Stat::Min).as_exact() { + partial.merge(min); + return Ok(true); + } + Ok(false) + } + Precision::Exact(_) => { + partial.poison(); + Ok(true) + } + _ => Ok(false), + } } fn accumulate( @@ -123,7 +196,10 @@ impl AggregateFnVTable for Min { Columnar::Canonical(canonical) => canonical.clone().into_array(), Columnar::Constant(constant) => constant.clone().into_array(), }; - if let Some(result) = min_max(&array, ctx)? { + let options = NumericalAggregateOpts { + skip_nans: partial.skip_nans, + }; + if let Some(result) = min_max(&array, ctx, options)? { partial.merge(result.min); } Ok(()) @@ -144,12 +220,12 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray as _; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min::Min; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::Nullability; @@ -162,9 +238,9 @@ mod tests { #[test] fn min_aggregate_fn() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Min, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Min, NumericalAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20, 5], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -182,7 +258,7 @@ mod tests { #[test] fn min_empty_group_returns_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Min, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Min, NumericalAggregateOpts::default(), dtype)?; assert_eq!( acc.finish()?, @@ -191,14 +267,90 @@ mod tests { Ok(()) } + #[test] + fn min_with_nan_not_skipping() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(Min, NumericalAggregateOpts::include_nans(), dtype)?; + + let batch = PrimitiveArray::new(buffer![1.0f64, f64::NAN, -5.0], Validity::NonNullable) + .into_array(); + acc.accumulate(&batch, &mut ctx)?; + assert!(acc.is_saturated()); + + let result = acc.finish()?; + assert!( + result + .as_primitive() + .typed_value::() + .is_some_and(f64::is_nan) + ); + Ok(()) + } + + #[test] + fn min_not_skipping_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // The array has no NaNs; a planted exact NaNCount stat proves the poisoning came from + // the stat rather than a scan. + let batch = PrimitiveArray::new(buffer![1.0f64, 2.0], Validity::NonNullable).into_array(); + batch + .statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(1u64))); + let mut acc = Accumulator::try_new( + Min, + NumericalAggregateOpts::include_nans(), + batch.dtype().clone(), + )?; + acc.accumulate(&batch, &mut ctx)?; + let result = acc.finish()?; + assert!( + result + .as_primitive() + .typed_value::() + .is_some_and(f64::is_nan) + ); + Ok(()) + } + + #[test] + fn min_nan_including_nullable_cached_stat() -> VortexResult<()> { + // A nullable float array's cached Min stat is reconstructed as a nullable scalar. The + // NaN-including shortcircuit merges it as-is; `to_scalar` re-casts to the result dtype. + let mut ctx = array_session().create_execution_ctx(); + let array = + PrimitiveArray::from_option_iter([Some(1.0f64), Some(2.0), Some(3.0)]).into_array(); + array + .statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(0u64))); + array + .statistics() + .set(Stat::Min, Precision::Exact(ScalarValue::from(1.0f64))); + let mut acc = Accumulator::try_new( + Min, + NumericalAggregateOpts::include_nans(), + array.dtype().clone(), + )?; + acc.accumulate(&array, &mut ctx)?; + assert_eq!( + acc.finish()?, + Scalar::primitive(1.0f64, Nullability::Nullable) + ); + Ok(()) + } + #[test] fn min_casts_nonnullable_legacy_stat_to_nullable_partial() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let batch = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); batch .statistics() .set(Stat::Min, Precision::Exact(ScalarValue::from(3i32))); - let mut acc = Accumulator::try_new(Min, EmptyOptions, batch.dtype().clone())?; + let mut acc = Accumulator::try_new( + Min, + NumericalAggregateOpts::default(), + batch.dtype().clone(), + )?; acc.accumulate(&batch, &mut ctx)?; diff --git a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs b/vortex-array/src/aggregate_fn/fns/min_max/extension.rs index 60962dd008e..e981c7c6af1 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/extension.rs @@ -7,6 +7,7 @@ use super::MinMaxPartial; use super::MinMaxResult; use super::min_max; use crate::ExecutionCtx; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::arrays::ExtensionArray; use crate::arrays::extension::ExtensionArrayExt; use crate::dtype::Nullability; @@ -18,11 +19,15 @@ pub(super) fn accumulate_extension( ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let non_nullable_ext_dtype = array.ext_dtype().with_nullability(Nullability::NonNullable); - let local = - min_max(array.storage_array(), ctx)?.map(|MinMaxResult { min, max }| MinMaxResult { - min: Scalar::extension_ref(non_nullable_ext_dtype.clone(), min), - max: Scalar::extension_ref(non_nullable_ext_dtype, max), - }); + let local = min_max( + array.storage_array(), + ctx, + NumericalAggregateOpts::default(), + )? + .map(|MinMaxResult { min, max }| MinMaxResult { + min: Scalar::extension_ref(non_nullable_ext_dtype.clone(), min), + max: Scalar::extension_ref(non_nullable_ext_dtype, max), + }); partial.merge(local); 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 540b5608e28..51601ce76fd 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -12,6 +12,8 @@ use std::sync::LazyLock; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_panic; +use vortex_session::registry::CachedId; use self::bool::accumulate_bool; use self::decimal::accumulate_decimal; @@ -26,14 +28,17 @@ use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; +use crate::dtype::PType; use crate::dtype::StructFields; +use crate::dtype::half::f16; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; +use crate::expr::stats::StatsProviderExt; use crate::partial_ord::partial_max; use crate::partial_ord::partial_min; use crate::scalar::Scalar; @@ -42,9 +47,40 @@ static NAMES: LazyLock = LazyLock::new(|| FieldNames::from(["min", " /// The minimum and maximum non-null values of an array, or `None` if there are no non-null values. /// +/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the +/// default) NaN values are ignored and the cached `Stat::Min`/`Stat::Max` statistics are consulted +/// and updated. With `skip_nans=false`, any NaN value in a float array poisons both extrema to +/// NaN; an exact `Stat::NaNCount` statistic shortcircuits the NaN scan in either direction. +/// /// The result scalars have the non-nullable version of the array dtype. /// This will update the stats set of the array as a side effect. -pub fn min_max(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { +pub fn min_max( + array: &ArrayRef, + ctx: &mut ExecutionCtx, + options: NumericalAggregateOpts, +) -> VortexResult> { + if !options.skip_nans && array.dtype().is_float() { + match array.statistics().get_as::(Stat::NaNCount) { + // NaN-free: identical to the NaN-skipping path below, including its stat caching. + Precision::Exact(0) => {} + // At least one NaN value poisons both extrema. + Precision::Exact(_) => return Ok(Some(nan_minmax_result(array.dtype()))), + _ => { + if array.is_empty() || array.valid_count(ctx)? == 0 { + return Ok(None); + } + // Compute with NaN-including options; the NaN-skipping `Stat::Min`/`Stat::Max` + // caches are neither read nor written. + let mut acc = Accumulator::try_new(MinMax, options, array.dtype().clone())?; + acc.accumulate(array, ctx)?; + return MinMaxResult::from_scalar(acc.finish()?); + } + } + } + + // NaN-skipping path. Also reached for NaN-free not-skipping float arrays and all non-float + // arrays, where `skip_nans` has no effect. + // Short-circuit using cached array statistics. let cached_min = array.statistics().get(Stat::Min).as_exact(); let cached_max = array.statistics().get(Stat::Max).as_exact(); @@ -67,7 +103,11 @@ pub fn min_max(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult. - let mut acc = Accumulator::try_new(MinMax, EmptyOptions, array.dtype().clone())?; + let mut acc = Accumulator::try_new( + MinMax, + NumericalAggregateOpts::default(), + array.dtype().clone(), + )?; acc.accumulate(array, ctx)?; let result_scalar = acc.finish()?; let result = MinMaxResult::from_scalar(result_scalar)?; @@ -89,6 +129,34 @@ pub fn min_max(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult MinMaxResult { + let nan = nan_scalar(dtype); + MinMaxResult { + min: nan.clone(), + max: nan, + } +} + +/// A non-nullable NaN scalar of the float `dtype`. +pub(crate) fn nan_scalar(dtype: &DType) -> Scalar { + match dtype.as_ptype() { + PType::F16 => Scalar::primitive(f16::NAN, Nullability::NonNullable), + PType::F32 => Scalar::primitive(f32::NAN, Nullability::NonNullable), + PType::F64 => Scalar::primitive(f64::NAN, Nullability::NonNullable), + _ => vortex_panic!("NaN scalar requested for non-float dtype {dtype}"), + } +} + +/// Whether a scalar holds a primitive float NaN value. +pub(crate) fn scalar_is_nan(scalar: &Scalar) -> bool { + if !scalar.dtype().is_float() { + return false; + } + + scalar.as_primitive_opt().is_some_and(|p| p.is_nan()) +} + /// The minimum and maximum non-null values of an array. #[derive(Debug, Clone, PartialEq, Eq)] pub struct MinMaxResult { @@ -119,6 +187,9 @@ impl MinMaxResult { /// /// Returns a nullable struct scalar `{min: T, max: T}` where `T` is the non-nullable input dtype. /// The struct is null when the array is empty or all-null. +/// +/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the +/// default) NaN values are ignored, otherwise any NaN value poisons both extrema to NaN. #[derive(Clone, Debug)] pub struct MinMax; @@ -127,6 +198,7 @@ pub struct MinMaxPartial { min: Option, max: Option, element_dtype: DType, + skip_nans: bool, } impl MinMaxPartial { @@ -136,6 +208,16 @@ impl MinMaxPartial { return; }; + // NaN scalars are incomparable under `partial_min`/`partial_max`, so they are handled + // explicitly: a NaN extremum poisons the partial state when NaNs participate, and is + // dropped when they are skipped. + if scalar_is_nan(&min) || scalar_is_nan(&max) || self.is_poisoned() { + if !self.skip_nans { + self.poison(); + } + return; + } + self.min = Some(match self.min.take() { Some(current) => partial_min(min, current).vortex_expect("incomparable min scalars"), None => min, @@ -146,6 +228,18 @@ impl MinMaxPartial { None => max, }); } + + /// Poison the partial state to `{min: NaN, max: NaN}`. + fn poison(&mut self) { + let nan = nan_scalar(&self.element_dtype); + self.min = Some(nan.clone()); + self.max = Some(nan); + } + + /// Whether the partial state is poisoned to NaN. + fn is_poisoned(&self) -> bool { + self.element_dtype.is_float() && self.min.as_ref().is_some_and(scalar_is_nan) + } } /// Creates the struct dtype `{min: T, max: T}` (nullable) used for min/max aggregate results. @@ -194,11 +288,12 @@ fn minmax_compute_supported_dtype(input_dtype: &DType) -> bool { } impl AggregateFnVTable for MinMax { - type Options = EmptyOptions; + type Options = NumericalAggregateOpts; type Partial = MinMaxPartial; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.min_max") + static ID: CachedId = CachedId::new("vortex.min_max"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { @@ -215,13 +310,14 @@ impl AggregateFnVTable for MinMax { fn empty_partial( &self, - _options: &Self::Options, + options: &Self::Options, input_dtype: &DType, ) -> VortexResult { Ok(MinMaxPartial { min: None, max: None, element_dtype: input_dtype.clone(), + skip_nans: options.skip_nans, }) } @@ -245,8 +341,46 @@ impl AggregateFnVTable for MinMax { } #[inline] - fn is_saturated(&self, _partial: &Self::Partial) -> bool { - false + fn is_saturated(&self, partial: &Self::Partial) -> bool { + // A poisoned NaN-including min/max is fully determined. + partial.is_poisoned() + } + + fn try_accumulate( + &self, + partial: &mut Self::Partial, + batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + // NaN-aware shortcircuits only apply to NaN-including float min/max; everything else + // takes the default dispatch path. + if partial.skip_nans || !partial.element_dtype.is_float() { + return Ok(false); + } + match batch.statistics().get_as::(Stat::NaNCount) { + Precision::Exact(0) => { + // NaN-free batch: the cached NaN-skipping extrema (if any) are valid. + let cached_min = batch.statistics().get(Stat::Min).as_exact(); + let cached_max = batch.statistics().get(Stat::Max).as_exact(); + if let Some((min, max)) = cached_min.zip(cached_max) { + // Cached float stats carry the (possibly nullable) array dtype; `to_scalar` + // builds a struct with non-nullable fields, so normalise here. + let non_nullable_dtype = partial.element_dtype.as_nonnullable(); + partial.merge(Some(MinMaxResult { + min: min.cast(&non_nullable_dtype)?, + max: max.cast(&non_nullable_dtype)?, + })); + return Ok(true); + } + Ok(false) + } + Precision::Exact(_) => { + // At least one NaN value poisons both extrema without scanning the batch. + partial.poison(); + Ok(true) + } + _ => Ok(false), + } } fn accumulate( @@ -261,8 +395,11 @@ impl AggregateFnVTable for MinMax { if scalar.is_null() { return Ok(()); } - // Skip NaN float constants - if scalar.as_primitive_opt().is_some_and(|p| p.is_nan()) { + // NaN float constants are skipped or poison the extrema, per the options. + if scalar_is_nan(scalar) { + if !partial.skip_nans { + partial.poison(); + } return Ok(()); } let non_nullable_dtype = scalar.dtype().as_nonnullable(); @@ -276,7 +413,7 @@ impl AggregateFnVTable for MinMax { Columnar::Canonical(c) => match c { Canonical::Primitive(p) => accumulate_primitive(partial, p, ctx), Canonical::Bool(b) => accumulate_bool(partial, b, ctx), - Canonical::VarBinView(v) => accumulate_varbinview(partial, v), + Canonical::VarBinView(v) => accumulate_varbinview(partial, v, ctx), Canonical::Decimal(d) => accumulate_decimal(partial, d, ctx), Canonical::Extension(e) => accumulate_extension(partial, e, ctx), Canonical::Null(_) => Ok(()), @@ -302,19 +439,20 @@ impl AggregateFnVTable for MinMax { #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::IntoArray as _; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::min_max::MinMaxResult; use crate::aggregate_fn::fns::min_max::make_minmax_dtype; @@ -332,17 +470,21 @@ mod tests { use crate::dtype::DecimalDType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::expr::stats::Precision; + use crate::expr::stats::Stat; use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::scalar::ScalarValue; use crate::validity::Validity; + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + #[test] fn test_prim_min_max() -> VortexResult<()> { let p = PrimitiveArray::new(buffer![1, 2, 3], Validity::NonNullable).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_eq!( - min_max(&p, &mut ctx)?, + min_max(&p, &mut ctx, NumericalAggregateOpts::default())?, Some(MinMaxResult { min: 1.into(), max: 3.into() @@ -366,9 +508,9 @@ mod tests { Some(7), ]) .into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_eq!( - min_max(&p, &mut ctx)?, + min_max(&p, &mut ctx, NumericalAggregateOpts::default())?, Some(MinMaxResult { min: 1.into(), max: 9.into() @@ -379,7 +521,7 @@ mod tests { #[test] fn test_bool_min_max() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let all_true = BoolArray::new( BitBuffer::from([true, true, true].as_slice()), @@ -387,7 +529,7 @@ mod tests { ) .into_array(); assert_eq!( - min_max(&all_true, &mut ctx)?, + min_max(&all_true, &mut ctx, NumericalAggregateOpts::default())?, Some(MinMaxResult { min: true.into(), max: true.into() @@ -400,7 +542,7 @@ mod tests { ) .into_array(); assert_eq!( - min_max(&all_false, &mut ctx)?, + min_max(&all_false, &mut ctx, NumericalAggregateOpts::default())?, Some(MinMaxResult { min: false.into(), max: false.into() @@ -413,7 +555,7 @@ mod tests { ) .into_array(); assert_eq!( - min_max(&mixed, &mut ctx)?, + min_max(&mixed, &mut ctx, NumericalAggregateOpts::default())?, Some(MinMaxResult { min: false.into(), max: true.into() @@ -425,8 +567,11 @@ mod tests { #[test] fn test_null_array() -> VortexResult<()> { let p = NullArray::new(1).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - assert_eq!(min_max(&p, &mut ctx)?, None); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!( + min_max(&p, &mut ctx, NumericalAggregateOpts::default())?, + None + ); Ok(()) } @@ -436,8 +581,13 @@ mod tests { buffer![f32::NAN, -f32::NAN, -1.0, 1.0], Validity::NonNullable, ); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let result = min_max(&array.into_array(), &mut ctx)?.vortex_expect("should have result"); + let mut ctx = SESSION.create_execution_ctx(); + let result = min_max( + &array.into_array(), + &mut ctx, + NumericalAggregateOpts::default(), + )? + .vortex_expect("should have result"); assert_eq!(f32::try_from(&result.min)?, -1.0); assert_eq!(f32::try_from(&result.max)?, 1.0); Ok(()) @@ -449,8 +599,13 @@ mod tests { buffer![f32::INFINITY, f32::NEG_INFINITY, -1.0, 1.0], Validity::NonNullable, ); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let result = min_max(&array.into_array(), &mut ctx)?.vortex_expect("should have result"); + let mut ctx = SESSION.create_execution_ctx(); + let result = min_max( + &array.into_array(), + &mut ctx, + NumericalAggregateOpts::default(), + )? + .vortex_expect("should have result"); assert_eq!(f32::try_from(&result.min)?, f32::NEG_INFINITY); assert_eq!(f32::try_from(&result.max)?, f32::INFINITY); Ok(()) @@ -458,9 +613,9 @@ mod tests { #[test] fn test_multi_batch() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(MinMax, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(MinMax, NumericalAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20, 5], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -476,9 +631,9 @@ mod tests { #[test] fn test_finish_resets_state() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(MinMax, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(MinMax, NumericalAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -497,7 +652,7 @@ mod tests { #[test] fn test_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = MinMax.empty_partial(&EmptyOptions, &dtype)?; + let mut state = MinMax.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; let struct_dtype = make_minmax_dtype(&dtype); let scalar1 = Scalar::struct_( @@ -520,19 +675,143 @@ mod tests { fn test_constant_nan() -> VortexResult<()> { let scalar = Scalar::primitive(f16::NAN, Nullability::NonNullable); let array = ConstantArray::new(scalar, 2).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - assert_eq!(min_max(&array, &mut ctx)?, None); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!( + min_max(&array, &mut ctx, NumericalAggregateOpts::default())?, + None + ); + Ok(()) + } + + const KEEP_NANS: NumericalAggregateOpts = NumericalAggregateOpts::include_nans(); + + fn assert_poisoned(result: Option) -> VortexResult<()> { + let result = result.vortex_expect("should have result"); + assert!(f64::try_from(&result.min.cast(&result.min.dtype().as_nullable())?)?.is_nan()); + assert!(f64::try_from(&result.max.cast(&result.max.dtype().as_nullable())?)?.is_nan()); Ok(()) } + #[test] + fn test_prim_nan_not_skipping() -> VortexResult<()> { + let array = PrimitiveArray::new( + buffer![f32::NAN, -f32::NAN, -1.0, 1.0], + Validity::NonNullable, + ) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_poisoned(min_max(&array, &mut ctx, KEEP_NANS)?) + } + + #[test] + fn test_prim_no_nan_not_skipping() -> VortexResult<()> { + let array = + PrimitiveArray::new(buffer![3.0f32, -1.0, 1.0], Validity::NonNullable).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let result = min_max(&array, &mut ctx, KEEP_NANS)?.vortex_expect("should have result"); + assert_eq!(f32::try_from(&result.min)?, -1.0); + assert_eq!(f32::try_from(&result.max)?, 3.0); + Ok(()) + } + + #[test] + fn test_constant_nan_not_skipping() -> VortexResult<()> { + let scalar = Scalar::primitive(f64::NAN, Nullability::NonNullable); + let array = ConstantArray::new(scalar, 2).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + assert_poisoned(min_max(&array, &mut ctx, KEEP_NANS)?) + } + + #[test] + fn test_not_skipping_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> { + // The array has no NaNs; a planted exact NaNCount stat proves the poisoning came from + // the stat rather than a scan. + let array = + PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + array + .statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(2u64))); + let mut ctx = SESSION.create_execution_ctx(); + assert_poisoned(min_max(&array, &mut ctx, KEEP_NANS)?) + } + + #[test] + fn test_not_skipping_uses_cached_stats_when_nan_free() -> VortexResult<()> { + // With an exact NaNCount of zero, the planted exact Min/Max stats are usable as-is. + let array = + PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + array + .statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(0u64))); + array + .statistics() + .set(Stat::Min, Precision::Exact(ScalarValue::from(-10.0f64))); + array + .statistics() + .set(Stat::Max, Precision::Exact(ScalarValue::from(10.0f64))); + let mut ctx = SESSION.create_execution_ctx(); + let result = min_max(&array, &mut ctx, KEEP_NANS)?.vortex_expect("should have result"); + assert_eq!(f64::try_from(&result.min)?, -10.0); + assert_eq!(f64::try_from(&result.max)?, 10.0); + Ok(()) + } + + #[test] + fn test_accumulator_nan_including_nullable_cached_stats() -> VortexResult<()> { + // A nullable float array's cached Min/Max stats are reconstructed as nullable scalars. + // The NaN-including accumulator shortcircuit must normalise them to the non-nullable + // struct field dtype before building the result scalar. + let mut ctx = SESSION.create_execution_ctx(); + let array = + PrimitiveArray::from_option_iter([Some(1.0f64), Some(2.0), Some(3.0)]).into_array(); + array + .statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(0u64))); + array + .statistics() + .set(Stat::Min, Precision::Exact(ScalarValue::from(1.0f64))); + array + .statistics() + .set(Stat::Max, Precision::Exact(ScalarValue::from(3.0f64))); + + let mut acc = Accumulator::try_new(MinMax, KEEP_NANS, array.dtype().clone())?; + acc.accumulate(&array, &mut ctx)?; + let result = MinMaxResult::from_scalar(acc.finish()?)?.vortex_expect("should have result"); + assert_eq!(f64::try_from(&result.min)?, 1.0); + assert_eq!(f64::try_from(&result.max)?, 3.0); + Ok(()) + } + + #[test] + fn test_multi_batch_nan_poisoning() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let mut acc = Accumulator::try_new(MinMax, KEEP_NANS, dtype)?; + + let batch1 = PrimitiveArray::new(buffer![1.0f64, 2.0], Validity::NonNullable).into_array(); + acc.accumulate(&batch1, &mut ctx)?; + assert!(!acc.is_saturated()); + + let batch2 = PrimitiveArray::new(buffer![f64::NAN], Validity::NonNullable).into_array(); + acc.accumulate(&batch2, &mut ctx)?; + assert!(acc.is_saturated()); + + assert_poisoned(MinMaxResult::from_scalar(acc.finish()?)?) + } + #[test] fn test_chunked() -> VortexResult<()> { let chunk1 = PrimitiveArray::from_option_iter([Some(5i32), None, Some(1)]); let chunk2 = PrimitiveArray::from_option_iter([Some(10i32), Some(3), None]); let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let result = min_max(&chunked.into_array(), &mut ctx)?.vortex_expect("should have result"); + let mut ctx = SESSION.create_execution_ctx(); + let result = min_max( + &chunked.into_array(), + &mut ctx, + NumericalAggregateOpts::default(), + )? + .vortex_expect("should have result"); assert_eq!(result.min, Scalar::from(1i32)); assert_eq!(result.max, Scalar::from(10i32)); Ok(()) @@ -541,8 +820,11 @@ mod tests { #[test] fn test_all_null() -> VortexResult<()> { let p = PrimitiveArray::from_option_iter::([None, None, None]); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - assert_eq!(min_max(&p.into_array(), &mut ctx)?, None); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!( + min_max(&p.into_array(), &mut ctx, NumericalAggregateOpts::default())?, + None + ); Ok(()) } @@ -557,8 +839,13 @@ mod tests { ], DType::Utf8(Nullability::Nullable), ); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let result = min_max(&array.into_array(), &mut ctx)?.vortex_expect("should have result"); + let mut ctx = SESSION.create_execution_ctx(); + let result = min_max( + &array.into_array(), + &mut ctx, + NumericalAggregateOpts::default(), + )? + .vortex_expect("should have result"); assert_eq!( result.min, Scalar::utf8("hello world", Nullability::NonNullable) @@ -580,8 +867,13 @@ mod tests { DecimalDType::new(4, 2), Validity::from_iter([true, false, true]), ); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let result = min_max(&decimal.into_array(), &mut ctx)?.vortex_expect("should have result"); + let mut ctx = SESSION.create_execution_ctx(); + let result = min_max( + &decimal.into_array(), + &mut ctx, + NumericalAggregateOpts::default(), + )? + .vortex_expect("should have result"); let non_nullable_dtype = DType::Decimal(DecimalDType::new(4, 2), Nullability::NonNullable); let expected_min = Scalar::try_new( @@ -605,18 +897,18 @@ mod tests { DType::FixedSizeList(Arc::new(element_dtype), 1, Nullability::Nullable); assert_eq!( - MinMax.return_dtype(&EmptyOptions, &list_dtype), + MinMax.return_dtype(&NumericalAggregateOpts::default(), &list_dtype), Some(make_minmax_dtype(&list_dtype)) ); assert_eq!( - MinMax.return_dtype(&EmptyOptions, &fixed_size_list_dtype), + MinMax.return_dtype(&NumericalAggregateOpts::default(), &fixed_size_list_dtype), Some(make_minmax_dtype(&fixed_size_list_dtype)) ); } #[test] fn list_and_fixed_size_list_min_max_returns_none() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let list_array = ListArray::try_new( buffer![1i32, 2, 3].into_array(), @@ -624,7 +916,10 @@ mod tests { Validity::NonNullable, )? .into_array(); - assert_eq!(min_max(&list_array, &mut ctx)?, None); + assert_eq!( + min_max(&list_array, &mut ctx, NumericalAggregateOpts::default())?, + None + ); let fixed_size_list_array = FixedSizeListArray::try_new( buffer![1i32, 2, 3, 4].into_array(), @@ -633,7 +928,14 @@ mod tests { 2, )? .into_array(); - assert_eq!(min_max(&fixed_size_list_array, &mut ctx)?, None); + assert_eq!( + min_max( + &fixed_size_list_array, + &mut ctx, + NumericalAggregateOpts::default() + )?, + None + ); Ok(()) } @@ -642,11 +944,12 @@ mod tests { #[test] fn test_bool_with_nulls() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let result = min_max( &BoolArray::from_iter(vec![Some(true), Some(true), None, None]).into_array(), &mut ctx, + NumericalAggregateOpts::default(), )?; assert_eq!( result, @@ -659,6 +962,7 @@ mod tests { let result = min_max( &BoolArray::from_iter(vec![None, Some(true), Some(true)]).into_array(), &mut ctx, + NumericalAggregateOpts::default(), )?; assert_eq!( result, @@ -671,6 +975,7 @@ mod tests { let result = min_max( &BoolArray::from_iter(vec![None, Some(true), Some(true), None]).into_array(), &mut ctx, + NumericalAggregateOpts::default(), )?; assert_eq!( result, @@ -683,6 +988,7 @@ mod tests { let result = min_max( &BoolArray::from_iter(vec![Some(false), Some(false), None, None]).into_array(), &mut ctx, + NumericalAggregateOpts::default(), )?; assert_eq!( result, @@ -701,7 +1007,7 @@ mod tests { /// partial state. #[test] fn test_bool_chunked_with_empty_chunk() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let empty = BoolArray::new(BitBuffer::from([].as_slice()), Validity::NonNullable); let chunk1 = BoolArray::new( @@ -717,7 +1023,11 @@ mod tests { DType::Bool(Nullability::NonNullable), )?; - let result = min_max(&chunked.into_array(), &mut ctx)?; + let result = min_max( + &chunked.into_array(), + &mut ctx, + NumericalAggregateOpts::default(), + )?; assert_eq!( result, Some(MinMaxResult { @@ -736,7 +1046,7 @@ mod tests { /// running min/max. Empty chunks are now skipped during chunked aggregation. #[test] fn test_chunked_with_empty_constant_chunk() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let empty = ConstantArray::new(Scalar::primitive(u32::MAX, Nullability::NonNullable), 0) .into_array(); @@ -748,7 +1058,11 @@ mod tests { )?; assert_eq!( - min_max(&chunked.into_array(), &mut ctx)?, + min_max( + &chunked.into_array(), + &mut ctx, + NumericalAggregateOpts::default() + )?, Some(MinMaxResult { min: Scalar::primitive(0u32, Nullability::NonNullable), max: Scalar::primitive(7631471u32, Nullability::NonNullable), @@ -763,8 +1077,15 @@ mod tests { vec![Option::<&str>::None, None, None], DType::Utf8(Nullability::Nullable), ); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - assert_eq!(min_max(&array.into_array(), &mut ctx)?, None); + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!( + min_max( + &array.into_array(), + &mut ctx, + NumericalAggregateOpts::default() + )?, + None + ); Ok(()) } } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs index 06b9749678d..7ca97712477 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs @@ -20,8 +20,9 @@ pub(super) fn accumulate_primitive( p: &PrimitiveArray, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { + let skip_nans = partial.skip_nans; match_each_native_ptype!(p.ptype(), |T| { - let local = compute_min_max_with_validity::(p, ctx)?; + let local = compute_min_max_with_validity::(p, ctx, skip_nans)?; partial.merge(local); Ok(()) }) @@ -30,6 +31,7 @@ pub(super) fn accumulate_primitive( fn compute_min_max_with_validity( array: &PrimitiveArray, ctx: &mut ExecutionCtx, + skip_nans: bool, ) -> VortexResult> where T: NativePType, @@ -48,7 +50,7 @@ where if T::PTYPE.is_int() { integer_min_max_raw(slice).map(min_max_result) } else { - compute_min_max(slice.iter()) + compute_min_max(slice.iter(), skip_nans) } } Mask::AllFalse(_) => None, @@ -73,6 +75,7 @@ where v.slices() .iter() .flat_map(|&(start, end)| slice[start..end].iter()), + skip_nans, ) } } @@ -110,15 +113,29 @@ where } } -fn compute_min_max<'a, T>(iter: impl Iterator) -> Option +fn compute_min_max<'a, T>( + iter: impl Iterator, + skip_nans: bool, +) -> Option where T: NativePType, PValue: From, { - match iter - .filter(|v| !v.is_nan()) - .minmax_by(|a, b| a.total_compare(**b)) - { + if skip_nans { + minmax_by_total_order(iter.filter(|v| !v.is_nan())) + } else { + // Compute extrema under the total order (where NaNs sort to the ends) and let the + // partial's merge poison the result if either end is NaN. + minmax_by_total_order(iter) + } +} + +fn minmax_by_total_order<'a, T>(iter: impl Iterator) -> Option +where + T: NativePType, + PValue: From, +{ + match iter.minmax_by(|a, b| a.total_compare(**b)) { itertools::MinMaxResult::NoElements => None, itertools::MinMaxResult::OneElement(&x) => { let scalar = Scalar::primitive(x, NonNullable); diff --git a/vortex-array/src/aggregate_fn/fns/min_max/varbin.rs b/vortex-array/src/aggregate_fn/fns/min_max/varbin.rs index a64cf80a240..5d4de12cd11 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/varbin.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/varbin.rs @@ -7,7 +7,7 @@ use vortex_error::vortex_panic; use super::MinMaxPartial; use super::MinMaxResult; -use crate::accessor::ArrayAccessor; +use crate::ExecutionCtx; use crate::arrays::VarBinViewArray; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; @@ -16,16 +16,41 @@ use crate::scalar::Scalar; pub(super) fn accumulate_varbinview( partial: &mut MinMaxPartial, array: &VarBinViewArray, + ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - partial.merge(varbin_compute_min_max(array, array.dtype())); + partial.merge(varbin_compute_min_max(array, array.dtype(), ctx)?); Ok(()) } -fn varbin_compute_min_max>( - array: &T, +fn varbin_compute_min_max( + array: &VarBinViewArray, dtype: &DType, -) -> Option { - array.with_iterator(|iter| match iter.flatten().minmax() { + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let mask = array + .validity()? + .execute_mask(array.len(), ctx)? + .to_bit_buffer(); + let views = array.views(); + let buffers = array + .data_buffers() + .iter() + .map(|b| b.as_host()) + .collect::>(); + let minmax = views + .iter() + .zip(mask.iter()) + .filter(|(_, v)| *v) + .map(|(view, _)| { + 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()] + } + }) + .minmax(); + Ok(match minmax { itertools::MinMaxResult::NoElements => None, itertools::MinMaxResult::OneElement(value) => { let scalar = make_scalar(dtype, value); diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs index 51087f3bb3e..723847160e1 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs @@ -7,6 +7,8 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use self::primitive::accumulate_primitive; use crate::ArrayRef; @@ -83,11 +85,20 @@ impl AggregateFnVTable for NanCount { type Partial = u64; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.nan_count") + static ID: CachedId = CachedId::new("vortex.nan_count"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - unimplemented!("NanCount is not yet serializable"); + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { @@ -176,7 +187,6 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; @@ -184,6 +194,7 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::nan_count::NanCount; use crate::aggregate_fn::fns::nan_count::nan_count; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -195,7 +206,7 @@ mod tests { #[test] fn nan_count_multi_batch() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); let mut acc = Accumulator::try_new(NanCount, EmptyOptions, dtype)?; @@ -215,7 +226,7 @@ mod tests { #[test] fn nan_count_finish_resets_state() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); let mut acc = Accumulator::try_new(NanCount, EmptyOptions, dtype)?; @@ -253,7 +264,7 @@ mod tests { #[test] fn nan_count_constant_nan() -> VortexResult<()> { let array = ConstantArray::new(f64::NAN, 10); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(nan_count(&array.into_array(), &mut ctx)?, 10); Ok(()) } @@ -261,7 +272,7 @@ mod tests { #[test] fn nan_count_constant_non_nan() -> VortexResult<()> { let array = ConstantArray::new(1.0f64, 10); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(nan_count(&array.into_array(), &mut ctx)?, 0); Ok(()) } @@ -281,7 +292,7 @@ mod tests { let chunk2 = PrimitiveArray::from_option_iter([Some(f64::NAN), Some(f64::NAN), None]); let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(nan_count(&chunked.into_array(), &mut ctx)?, 3); Ok(()) } @@ -289,7 +300,7 @@ mod tests { #[test] fn nan_count_all_null() -> VortexResult<()> { let p = PrimitiveArray::from_option_iter::([None, None, None]); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(nan_count(&p.into_array(), &mut ctx)?, 0); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs b/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs index 8d8358235cc..c9da090c7ab 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs @@ -42,9 +42,9 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::nan_count::nan_count; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::validity::Validity; @@ -62,7 +62,7 @@ mod tests { ], Validity::NonNullable, ); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(nan_count(&p.into_array(), &mut ctx)?, 2); Ok(()) } @@ -70,7 +70,7 @@ mod tests { #[test] fn primitive_nan_count_with_nulls() -> VortexResult<()> { let p = PrimitiveArray::from_option_iter([Some(f64::NAN), None, Some(f64::NAN), Some(1.0)]); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(nan_count(&p.into_array(), &mut ctx)?, 2); Ok(()) } @@ -78,7 +78,7 @@ mod tests { #[test] fn primitive_nan_count_all_valid_no_nans() -> VortexResult<()> { let p = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(nan_count(&p.into_array(), &mut ctx)?, 0); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs index 47faabc82ef..031e8f6a09b 100644 --- a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs @@ -4,6 +4,8 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Columnar; @@ -58,11 +60,20 @@ impl AggregateFnVTable for NullCount { type Partial = u64; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.null_count") + static ID: CachedId = CachedId::new("vortex.null_count"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - Ok(None) + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option { @@ -146,13 +157,13 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::null_count::NullCount; use crate::aggregate_fn::fns::null_count::null_count; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::Nullability; @@ -165,7 +176,7 @@ mod tests { fn null_count_with_nulls() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None]).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(null_count(&array, &mut ctx)?, 2); assert_eq!( @@ -177,7 +188,7 @@ mod tests { #[test] fn null_count_multi_batch() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(NullCount, EmptyOptions, dtype)?; diff --git a/vortex-array/src/aggregate_fn/fns/sum/bool.rs b/vortex-array/src/aggregate_fn/fns/sum/bool.rs index 7728d3f64af..b3993840586 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -37,13 +37,13 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; + use crate::array_session; use crate::arrays::BoolArray; use crate::dtype::DType; use crate::dtype::Nullability; @@ -55,7 +55,7 @@ mod tests { let arr: BoolArray = [true, true, true].into_iter().collect(); let result = sum( &arr.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; assert_eq!(result.as_primitive().typed_value::(), Some(3)); Ok(()) @@ -66,7 +66,7 @@ mod tests { let arr: BoolArray = [true, false, true, false, true].into_iter().collect(); let result = sum( &arr.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; assert_eq!(result.as_primitive().typed_value::(), Some(3)); Ok(()) @@ -77,7 +77,7 @@ mod tests { let arr: BoolArray = [false, false, false].into_iter().collect(); let result = sum( &arr.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) @@ -88,7 +88,7 @@ mod tests { let arr = BoolArray::from_iter([Some(true), None, Some(true), Some(false)]); let result = sum( &arr.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; assert_eq!(result.as_primitive().typed_value::(), Some(2)); Ok(()) @@ -99,7 +99,7 @@ mod tests { let arr = BoolArray::from_iter([None::, None, None]); let result = sum( &arr.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) @@ -108,7 +108,7 @@ mod tests { #[test] fn sum_bool_empty_produces_zero() -> VortexResult<()> { let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) @@ -116,9 +116,9 @@ mod tests { #[test] fn sum_bool_finish_resets_state() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Bool(Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let batch1: BoolArray = [true, true, false].into_iter().collect(); acc.accumulate(&batch1.into_array(), &mut ctx)?; @@ -135,7 +135,10 @@ mod tests { #[test] fn sum_bool_return_dtype() -> VortexResult<()> { let dtype = Sum - .return_dtype(&EmptyOptions, &DType::Bool(Nullability::NonNullable)) + .return_dtype( + &NumericalAggregateOpts::default(), + &DType::Bool(Nullability::NonNullable), + ) .unwrap(); assert_eq!(dtype, DType::Primitive(PType::U64, Nullability::Nullable)); Ok(()) @@ -144,7 +147,7 @@ mod tests { #[test] fn sum_boolean_from_iter() -> VortexResult<()> { let arr = BoolArray::from_iter([true, false, false, true]).into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().as_::(), Some(2)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/constant.rs b/vortex-array/src/aggregate_fn/fns/sum/constant.rs index c55b66a71cc..0f366620e5c 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/constant.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/constant.rs @@ -92,9 +92,9 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::sum::sum; + use crate::array_session; use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::dtype::DecimalDType; @@ -109,7 +109,7 @@ mod tests { #[test] fn sum_constant_unsigned() -> VortexResult<()> { let array = ConstantArray::new(5u64, 10).into_array(); - let result = sum(&array, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, 50u64.into()); Ok(()) } @@ -117,7 +117,7 @@ mod tests { #[test] fn sum_constant_signed() -> VortexResult<()> { let array = ConstantArray::new(-5i64, 10).into_array(); - let result = sum(&array, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, (-50i64).into()); Ok(()) } @@ -126,7 +126,7 @@ mod tests { fn sum_constant_nullable_value() -> VortexResult<()> { let array = ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) .into_array(); - let result = sum(&array, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, Scalar::primitive(0u64, Nullable)); Ok(()) } @@ -134,7 +134,7 @@ mod tests { #[test] fn sum_constant_bool_false() -> VortexResult<()> { let array = ConstantArray::new(false, 10).into_array(); - let result = sum(&array, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, 0u64.into()); Ok(()) } @@ -142,7 +142,7 @@ mod tests { #[test] fn sum_constant_bool_true() -> VortexResult<()> { let array = ConstantArray::new(true, 10).into_array(); - let result = sum(&array, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, 10u64.into()); Ok(()) } @@ -150,7 +150,7 @@ mod tests { #[test] fn sum_constant_bool_null() -> VortexResult<()> { let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); - let result = sum(&array, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!(result, Scalar::primitive(0u64, Nullable)); Ok(()) } @@ -168,7 +168,7 @@ mod tests { ) .into_array(); - let result = sum(&array, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result.as_decimal().decimal_value(), @@ -184,7 +184,7 @@ mod tests { let array = ConstantArray::new(Scalar::null(DType::Decimal(decimal_dtype, Nullable)), 10) .into_array(); - let result = sum(&array, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result, Scalar::decimal( @@ -209,7 +209,7 @@ mod tests { ) .into_array(); - let result = sum(&array, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!( result.as_decimal().decimal_value(), Some(DecimalValue::I256(i256::from_i128(99_999_999_900))) diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index e37b671eb46..872e5769a01 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -112,12 +112,12 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::AggregateFnVTable; - use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; + use crate::array_session; use crate::arrays::DecimalArray; use crate::dtype::DType; use crate::dtype::DecimalDType; @@ -139,7 +139,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; let expected = Scalar::try_new( @@ -161,7 +161,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; let expected = Scalar::try_new( @@ -183,7 +183,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; let expected = Scalar::try_new( @@ -206,7 +206,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; let expected_sum = near_max as i64 + 500 + 400; @@ -230,7 +230,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; let expected_sum = (large_val as i128) * 4 + 1; @@ -253,7 +253,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; let expected = Scalar::try_new( @@ -272,7 +272,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; let expected = Scalar::try_new( @@ -294,7 +294,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; let expected = Scalar::try_new( @@ -317,7 +317,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; let expected_sum = @@ -343,7 +343,7 @@ mod tests { assert_eq!( sum( &decimal.into_array(), - &mut LEGACY_SESSION.create_execution_ctx() + &mut array_session().create_execution_ctx() ) .vortex_expect("operation should succeed in test"), Scalar::null(DType::Decimal(decimal_dtype, Nullable)) @@ -357,7 +357,7 @@ mod tests { // Native type for precision 14 is I64 (max precision 18), so 14 < 18. // Use combine_partials to push state near (but under) 10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = Sum.empty_partial(&EmptyOptions, &input_dtype)?; + let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(99_999_999_999_990i64), @@ -387,7 +387,7 @@ mod tests { // i256 arithmetic does not overflow. This tests the precision-based // saturation path in combine_partials. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = Sum.empty_partial(&EmptyOptions, &input_dtype)?; + let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(99_999_999_999_999i64), @@ -414,7 +414,7 @@ mod tests { fn sum_decimal_precision_overflow_negative() -> VortexResult<()> { // Same setup but with negative values: sum reaches -10^14. let input_dtype = DType::Decimal(DecimalDType::new(4, 0), Nullability::NonNullable); - let mut state = Sum.empty_partial(&EmptyOptions, &input_dtype)?; + let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; let near_limit = Scalar::decimal( DecimalValue::from(-99_999_999_999_999i64), @@ -446,7 +446,7 @@ mod tests { // a real array that pushes it over. let input_dtype = DType::Decimal(DecimalDType::new(27, 0), Nullability::NonNullable); let return_dtype = DecimalDType::new(37, 0); - let mut state = Sum.empty_partial(&EmptyOptions, &input_dtype)?; + let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &input_dtype)?; // Set state to 10^37 - 1 via combine_partials. let near_limit_val: i128 = 10i128.pow(37) - 1; @@ -460,7 +460,7 @@ mod tests { // Drive accumulate through the vtable directly. let columnar = crate::Columnar::Canonical(crate::Canonical::Decimal(decimal)); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); Sum.accumulate(&mut state, &columnar, &mut ctx)?; let result = Sum.to_scalar(&state)?; diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 6f00cce7fdb..efe0825d4d6 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -32,10 +32,10 @@ impl DynGroupedAggregateKernel for PrimitiveGroupedSumEncodingKernel { groups: &GroupedArray, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if !aggregate_fn.is::() { + let Some(options) = aggregate_fn.as_opt::() else { return Ok(None); - } - try_grouped_sum(groups, ctx) + }; + try_grouped_sum(groups, ctx, options.skip_nans) } } @@ -48,6 +48,7 @@ impl DynGroupedAggregateKernel for PrimitiveGroupedSumEncodingKernel { pub(super) fn try_grouped_sum( groups: &GroupedArray, ctx: &mut ExecutionCtx, + skip_nans: bool, ) -> VortexResult> { if !groups.elements().is::() { return Ok(None); @@ -61,6 +62,7 @@ pub(super) fn try_grouped_sum( &group_ranges, &group_validity, ctx, + skip_nans, )?)) } @@ -70,6 +72,7 @@ fn grouped_sum( group_ranges: &GroupRanges, group_validity: &Mask, ctx: &mut ExecutionCtx, + skip_nans: bool, ) -> VortexResult { let elem_mask = elements .as_ref() @@ -91,7 +94,7 @@ fn grouped_sum( 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); false }) + |acc, slice| { sum_float_all(acc, slice, skip_nans); false }) } ); @@ -156,13 +159,13 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::DynGroupedAccumulator; - use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::GroupedAccumulator; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; + use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -176,8 +179,12 @@ mod tests { /// Run a grouped sum through the accumulator. fn grouped_sum_actual(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { - let mut acc = GroupedAccumulator::try_new(Sum, EmptyOptions, elem_dtype.clone())?; - acc.accumulate_list(groups, &mut LEGACY_SESSION.create_execution_ctx())?; + let mut acc = GroupedAccumulator::try_new( + Sum, + NumericalAggregateOpts::default(), + elem_dtype.clone(), + )?; + acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; acc.finish() } @@ -191,9 +198,9 @@ mod tests { ) -> VortexResult { use crate::aggregate_fn::AggregateFnVTable; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let sum_dtype = Sum - .partial_dtype(&EmptyOptions, elem_dtype) + .partial_dtype(&NumericalAggregateOpts::default(), elem_dtype) .expect("sum partial dtype"); let mut builder = builder_with_capacity(&sum_dtype, ranges.len()); for (i, &(offset, size)) in ranges.iter().enumerate() { @@ -229,6 +236,7 @@ mod tests { #[test] fn listview_matches_reference_unsigned() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1u32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::U32, NonNullable); @@ -241,13 +249,14 @@ mod tests { // Unsigned input sums to U64. let direct = PrimitiveArray::from_option_iter([Some(3u64), Some(3u64), Some(15u64)]); - assert_arrays_eq!(&actual, &direct.into_array()); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Offsets are not in group order and a group is null: the group validity must be indexed by // group index, not by element offset. let elements = @@ -262,13 +271,14 @@ mod tests { let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; let direct = PrimitiveArray::from_option_iter([Some(110i64), None, Some(70i64)]); - assert_arrays_eq!(&actual, &direct.into_array()); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn listview_interior_and_full_nulls() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Group 1 has an interior null, group 2 is entirely null, group 3 is empty. let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) @@ -283,13 +293,14 @@ mod tests { let direct = PrimitiveArray::from_option_iter([Some(4i64), Some(0i64), Some(0i64), Some(9i64)]); - assert_arrays_eq!(&actual, &direct.into_array()); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn listview_overflow_group_is_null() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::I64, NonNullable); @@ -302,8 +313,8 @@ mod tests { // First group overflows -> null sum; second group sums normally. let direct = PrimitiveArray::from_option_iter([None, Some(5i64)]); - assert_arrays_eq!(&actual, &direct.into_array()); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } @@ -323,7 +334,7 @@ mod tests { // Group 0: NaN skipped -> 3.0. Group 1: INF + -INF = NaN. (Avoid array equality here since // NaN != NaN; compare element scalars against the reference path instead.) - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; let g0 = actual.execute_scalar(0, &mut ctx)?; assert_eq!(g0.as_primitive().typed_value::(), Some(3.0)); @@ -347,8 +358,33 @@ mod tests { Ok(()) } + #[test] + fn listview_float_nan_not_skipping() -> VortexResult<()> { + let elements = PrimitiveArray::new( + buffer![1.0f64, f64::NAN, 2.0, 3.0, 4.0], + Validity::NonNullable, + ) + .into_array(); + let elem_dtype = DType::Primitive(PType::F64, NonNullable); + let groups = listview(elements, &[(0, 3), (3, 2)], &[true, true])?; + + let mut acc = + GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::include_nans(), elem_dtype)?; + acc.accumulate_list(&groups, &mut array_session().create_execution_ctx())?; + let actual = acc.finish()?; + + let mut ctx = array_session().create_execution_ctx(); + // Group 0 contains a NaN -> NaN sum; group 1 sums normally. + let g0 = actual.execute_scalar(0, &mut ctx)?; + assert!(g0.as_primitive().typed_value::().unwrap().is_nan()); + let g1 = actual.execute_scalar(1, &mut ctx)?; + assert_eq!(g1.as_primitive().typed_value::(), Some(7.0)); + Ok(()) + } + #[test] fn fixed_size_overflow_and_nan() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // FixedSize path: first group overflows -> null sum, second sums normally. let elements = PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); @@ -360,8 +396,8 @@ mod tests { let expected = grouped_sum_reference(&elements, &[(0, 2), (2, 2)], &[true, true], &elem_dtype)?; let direct = PrimitiveArray::from_option_iter([None, Some(5i64)]); - assert_arrays_eq!(&actual, &direct.into_array()); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 9d525bec742..001b0f61657 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -12,6 +12,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use self::bool::accumulate_bool; use self::constant::multiply_constant; @@ -25,7 +27,7 @@ use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; -use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::MAX_PRECISION; @@ -34,6 +36,7 @@ use crate::dtype::PType; 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::Scalar; @@ -48,7 +51,11 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { // Compute using Accumulator. // TODO(ngates): we may want to wrap this three-step dance up into an extension crate maybe. - let mut acc = Accumulator::try_new(Sum, EmptyOptions, array.dtype().clone())?; + let mut acc = Accumulator::try_new( + Sum, + NumericalAggregateOpts::default(), + array.dtype().clone(), + )?; acc.accumulate(array, ctx)?; let result = acc.finish()?; @@ -64,19 +71,31 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { /// /// If the sum overflows, a null scalar will be returned. /// If the array is all-invalid, the sum will be zero. +/// +/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the +/// default) NaN values contribute nothing, otherwise any NaN value poisons the sum to NaN. #[derive(Clone, Debug)] pub struct Sum; impl AggregateFnVTable for Sum { - type Options = EmptyOptions; + type Options = NumericalAggregateOpts; type Partial = SumPartial; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.sum") + static ID: CachedId = CachedId::new("vortex.sum"); + *ID } - fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - unimplemented!("Sum is not yet serializable"); + 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 { @@ -130,6 +149,7 @@ impl AggregateFnVTable for Sum { Ok(SumPartial { return_dtype, current: Some(initial), + skip_nans: options.skip_nans, }) } @@ -214,6 +234,43 @@ impl AggregateFnVTable for Sum { } } + fn try_accumulate( + &self, + partial: &mut Self::Partial, + batch: &ArrayRef, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + // NaN-aware shortcircuits only apply to NaN-including float sums; everything else takes + // the default dispatch path. + if partial.skip_nans || !matches!(partial.current, Some(SumState::Float(_))) { + return Ok(false); + } + match batch.statistics().get_as::(Stat::NaNCount) { + Precision::Exact(0) => { + // NaN-free batch: the cached NaN-skipping sum (if any) equals the + // NaN-including sum. + if let Precision::Exact(sum) = batch.statistics().get(Stat::Sum) { + let sum = if sum.dtype() == &partial.return_dtype { + sum + } else { + sum.cast(&partial.return_dtype)? + }; + self.combine_partials(partial, sum)?; + return Ok(true); + } + Ok(false) + } + Precision::Exact(_) => { + // At least one NaN value: the sum is NaN without scanning the batch. + if let Some(SumState::Float(acc)) = partial.current.as_mut() { + *acc = f64::NAN; + } + Ok(true) + } + _ => Ok(false), + } + } + fn accumulate( &self, partial: &mut Self::Partial, @@ -222,12 +279,17 @@ impl AggregateFnVTable for Sum { ) -> VortexResult<()> { // Constants compute scalar * len and combine via combine_partials. if let Columnar::Constant(c) = batch { + // NaN constants are treated as missing when skipping NaNs. + if partial.skip_nans && c.scalar().as_primitive_opt().is_some_and(|p| p.is_nan()) { + return Ok(()); + } if let Some(product) = multiply_constant(c.scalar(), c.len(), &partial.return_dtype)? { self.combine_partials(partial, product)?; } return Ok(()); } + let skip_nans = partial.skip_nans; let mut inner = match partial.current.take() { Some(inner) => inner, None => return Ok(()), @@ -235,7 +297,7 @@ impl AggregateFnVTable for Sum { let result = match batch { Columnar::Canonical(c) => match c { - Canonical::Primitive(p) => accumulate_primitive(&mut inner, p, ctx), + Canonical::Primitive(p) => accumulate_primitive(&mut inner, p, ctx, skip_nans), Canonical::Bool(b) => accumulate_bool(&mut inner, b, ctx), Canonical::Decimal(d) => accumulate_decimal(&mut inner, d, ctx), _ => vortex_bail!("Unsupported canonical type for sum: {}", batch.dtype()), @@ -269,10 +331,11 @@ pub struct SumPartial { return_dtype: DType, /// The current accumulated state, or `None` if saturated (checked overflow). current: Option, + /// Whether NaN values in float inputs are skipped. + skip_nans: bool, } /// The accumulated sum value. -/// // TODO(ngates): instead of an enum, we should use a Box to avoid dispatcher over the // input type every time? Perhaps? pub enum SumState { @@ -333,16 +396,16 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::DynGroupedAccumulator; - use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::GroupedAccumulator; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -367,7 +430,7 @@ mod tests { /// Sum an array with an initial value (test-only helper). fn sum_with_accumulator(array: &ArrayRef, accumulator: &Scalar) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); if accumulator.is_null() { return Ok(accumulator.clone()); } @@ -423,9 +486,9 @@ mod tests { #[test] fn sum_multi_batch() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -440,9 +503,9 @@ mod tests { #[test] fn sum_finish_resets_state() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); acc.accumulate(&batch1, &mut ctx)?; @@ -461,7 +524,7 @@ mod tests { #[test] fn sum_state_merge() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut state = Sum.empty_partial(&EmptyOptions, &dtype)?; + let mut state = Sum.empty_partial(&NumericalAggregateOpts::default(), &dtype)?; let scalar1 = Scalar::primitive(100i64, Nullable); Sum.combine_partials(&mut state, scalar1)?; @@ -491,7 +554,7 @@ mod tests { // compute sum with accumulator to populate stats sum_with_accumulator(&array, &Scalar::primitive(2i64, Nullable))?; - let sum_without_acc = sum(&array, &mut LEGACY_SESSION.create_execution_ctx())?; + let sum_without_acc = sum(&array, &mut array_session().create_execution_ctx())?; assert_eq!(sum_without_acc, Scalar::primitive(9i64, Nullable)); Ok(()) } @@ -514,13 +577,18 @@ mod tests { // Grouped sum tests fn run_grouped_sum(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { - let mut acc = GroupedAccumulator::try_new(Sum, EmptyOptions, elem_dtype.clone())?; - acc.accumulate_list(groups, &mut LEGACY_SESSION.create_execution_ctx())?; + let mut acc = GroupedAccumulator::try_new( + Sum, + NumericalAggregateOpts::default(), + elem_dtype.clone(), + )?; + acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; acc.finish() } #[test] fn grouped_sum_fixed_size_list() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?; @@ -529,12 +597,13 @@ mod tests { let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; let expected = PrimitiveArray::from_option_iter([Some(6i64), Some(15i64)]).into_array(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } #[test] fn grouped_sum_with_null_elements() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5), Some(6)]) .into_array(); @@ -544,12 +613,13 @@ mod tests { let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; let expected = PrimitiveArray::from_option_iter([Some(4i64), Some(11i64)]).into_array(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } #[test] fn grouped_sum_with_null_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9], Validity::NonNullable) .into_array(); @@ -561,12 +631,13 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(6i64), None, Some(24i64)]).into_array(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } #[test] fn grouped_sum_all_null_elements_in_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([None::, None, Some(3), Some(4)]).into_array(); let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?; @@ -575,12 +646,13 @@ mod tests { let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; let expected = PrimitiveArray::from_option_iter([Some(0i64), Some(7i64)]).into_array(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } #[test] fn grouped_sum_bool() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements: BoolArray = [true, false, true, true, true, true].into_iter().collect(); let groups = FixedSizeListArray::try_new(elements.into_array(), 3, Validity::NonNullable, 2)?; @@ -589,15 +661,16 @@ mod tests { let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; let expected = PrimitiveArray::from_option_iter([Some(2u64), Some(3u64)]).into_array(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } #[test] fn grouped_sum_finish_resets() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = GroupedAccumulator::try_new(Sum, EmptyOptions, elem_dtype)?; + let mut acc = + GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::default(), elem_dtype)?; let elements1 = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); @@ -606,7 +679,7 @@ mod tests { let result1 = acc.finish()?; let expected1 = PrimitiveArray::from_option_iter([Some(3i64), Some(7i64)]).into_array(); - assert_arrays_eq!(&result1, &expected1); + assert_arrays_eq!(&result1, &expected1, &mut ctx); let elements2 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); let groups2 = FixedSizeListArray::try_new(elements2, 2, Validity::NonNullable, 1)?; @@ -614,12 +687,13 @@ mod tests { let result2 = acc.finish()?; let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array(); - assert_arrays_eq!(&result2, &expected2); + assert_arrays_eq!(&result2, &expected2, &mut ctx); Ok(()) } #[test] fn grouped_sum_listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![100i32, 200, 300], Validity::NonNullable).into_array(); let offsets = PrimitiveArray::new(buffer![2i32, 0, 1], Validity::NonNullable).into_array(); @@ -633,7 +707,7 @@ mod tests { // group 0 -> elements[2..3] = 300; group 1 -> null; group 2 -> elements[1..2] = 200. let expected = PrimitiveArray::from_option_iter([Some(300i64), None, Some(200i64)]).into_array(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } @@ -657,7 +731,7 @@ mod tests { let result = sum( &chunked.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; assert_eq!(result.as_primitive().as_::(), Some(20.8)); Ok(()) @@ -671,7 +745,7 @@ mod tests { let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; let result = sum( &chunked.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; assert_eq!(result, Scalar::primitive(0f64, Nullable)); Ok(()) @@ -694,7 +768,7 @@ mod tests { let result = sum( &chunked.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; assert_eq!(result.as_primitive().as_::(), Some(36.0)); Ok(()) @@ -709,7 +783,7 @@ mod tests { let result = sum( &chunked.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; assert_eq!(result.as_primitive().as_::(), Some(1)); Ok(()) @@ -741,7 +815,7 @@ mod tests { let result = sum( &chunked.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; let decimal_result = result.as_decimal(); assert_eq!( @@ -777,7 +851,7 @@ mod tests { let result = sum( &chunked.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; let decimal_result = result.as_decimal(); assert_eq!( @@ -811,7 +885,7 @@ mod tests { let result = sum( &chunked.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; let decimal_result = result.as_decimal(); assert_eq!( diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index df7d929d896..87d8da4b143 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -25,16 +25,21 @@ pub(super) fn accumulate_primitive( inner: &mut SumState, p: &PrimitiveArray, ctx: &mut ExecutionCtx, + skip_nans: bool, ) -> VortexResult { let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; match mask.slices() { AllOr::None => Ok(false), - AllOr::All => accumulate_primitive_all(inner, p), - AllOr::Some(slices) => accumulate_primitive_valid(inner, p, slices), + AllOr::All => accumulate_primitive_all(inner, p, skip_nans), + AllOr::Some(slices) => accumulate_primitive_valid(inner, p, slices, skip_nans), } } -fn accumulate_primitive_all(inner: &mut SumState, p: &PrimitiveArray) -> VortexResult { +fn accumulate_primitive_all( + inner: &mut SumState, + p: &PrimitiveArray, + skip_nans: bool, +) -> VortexResult { match inner { SumState::Unsigned(acc) => match_each_native_ptype!(p.ptype(), unsigned: |T| { Ok(sum_unsigned_all(acc, p.as_slice::())) }, @@ -50,7 +55,7 @@ fn accumulate_primitive_all(inner: &mut SumState, p: &PrimitiveArray) -> VortexR unsigned: |_T| { vortex_panic!("float sum state with unsigned input") }, signed: |_T| { vortex_panic!("float sum state with signed input") }, floating: |T| { - sum_float_all(acc, p.as_slice::()); + sum_float_all(acc, p.as_slice::(), skip_nans); Ok(false) } ), @@ -58,11 +63,18 @@ fn accumulate_primitive_all(inner: &mut SumState, p: &PrimitiveArray) -> VortexR } } -/// Sum the non-NaN values of a float slice into an `f64` accumulator. NaNs are skipped to match the -/// scalar `sum` semantics. Floats cannot overflow the accumulator, so this never reports saturation. -pub(super) fn sum_float_all(acc: &mut f64, slice: &[T]) { - for &v in slice { - if !v.is_nan() { +/// 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) { + if skip_nans { + for &v in slice { + if !v.is_nan() { + *acc += ToPrimitive::to_f64(&v).vortex_expect("float to f64"); + } + } + } else { + for &v in slice { *acc += ToPrimitive::to_f64(&v).vortex_expect("float to f64"); } } @@ -122,6 +134,7 @@ fn accumulate_primitive_valid( inner: &mut SumState, p: &PrimitiveArray, slices: &[(usize, usize)], + skip_nans: bool, ) -> VortexResult { match inner { SumState::Unsigned(acc) => match_each_native_ptype!(p.ptype(), @@ -156,7 +169,7 @@ fn accumulate_primitive_valid( floating: |T| { let values = p.as_slice::(); for &(start, end) in slices { - sum_float_all(acc, &values[start..end]); + sum_float_all(acc, &values[start..end], skip_nans); } Ok(false) } @@ -171,25 +184,29 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; + use crate::array_session; + use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::dtype::DType; 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; #[test] fn sum_i32() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(10)); Ok(()) } @@ -197,7 +214,7 @@ mod tests { #[test] fn sum_u8() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![10u8, 20, 30], Validity::NonNullable).into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(60)); Ok(()) } @@ -206,7 +223,7 @@ mod tests { fn sum_f64() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1.5f64, 2.5, 3.0], Validity::NonNullable).into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(7.0)); Ok(()) } @@ -214,7 +231,7 @@ mod tests { #[test] fn sum_with_nulls() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([Some(2i32), None, Some(4)]).into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(6)); Ok(()) } @@ -234,7 +251,7 @@ mod tests { Some(6), ]) .into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(21)); Ok(()) } @@ -242,7 +259,7 @@ mod tests { #[test] fn sum_all_null() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } @@ -250,7 +267,7 @@ mod tests { #[test] fn sum_all_invalid_float() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result, Scalar::primitive(0f64, Nullable)); Ok(()) } @@ -258,7 +275,7 @@ mod tests { #[test] fn sum_buffer_i32() -> VortexResult<()> { let arr = buffer![1, 1, 1, 1].into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().as_::(), Some(4)); Ok(()) } @@ -266,7 +283,7 @@ mod tests { #[test] fn sum_buffer_f64() -> VortexResult<()> { let arr = buffer![1., 1., 1., 1.].into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().as_::(), Some(4.)); Ok(()) } @@ -274,7 +291,7 @@ mod tests { #[test] fn sum_empty_produces_zero() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) @@ -283,7 +300,7 @@ mod tests { #[test] fn sum_empty_f64_produces_zero() -> VortexResult<()> { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; let result = acc.finish()?; assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); Ok(()) @@ -296,7 +313,7 @@ mod tests { Validity::NonNullable, ) .into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); Ok(()) } @@ -305,7 +322,7 @@ mod tests { fn sum_f32_with_nan() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1.0f32, f32::NAN, 4.0], Validity::NonNullable).into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(5.0)); Ok(()) } @@ -314,7 +331,7 @@ mod tests { fn sum_f64_with_nan_and_nulls() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([Some(1.0f64), None, Some(f64::NAN), Some(3.0)]) .into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert_eq!(result.as_primitive().typed_value::(), Some(4.0)); Ok(()) } @@ -323,8 +340,75 @@ mod tests { fn sum_all_nan() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![f64::NAN, f64::NAN], Validity::NonNullable).into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + Ok(()) + } + + /// Sum an array with explicit [`NumericalAggregateOpts`] (test-only helper). + fn sum_with_options( + arr: &crate::ArrayRef, + options: NumericalAggregateOpts, + ) -> VortexResult { + let mut acc = Accumulator::try_new(Sum, options, arr.dtype().clone())?; + acc.accumulate(arr, &mut array_session().create_execution_ctx())?; + acc.finish() + } + + #[test] + fn sum_f64_with_nan_not_skipping() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![1.0f64, f64::NAN, 2.0], Validity::NonNullable).into_array(); + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert!(result.as_primitive().typed_value::().unwrap().is_nan()); + Ok(()) + } + + #[test] + fn sum_f64_without_nan_not_skipping() -> VortexResult<()> { + let arr = + PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); + Ok(()) + } + + #[test] + fn sum_not_skipping_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> { + // The array has no NaNs; a planted exact NaNCount stat proves the NaN poisoning came + // from the stat rather than a scan. + let arr = + PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + arr.statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(1u64))); + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert!(result.as_primitive().typed_value::().unwrap().is_nan()); + Ok(()) + } + + #[test] + fn sum_not_skipping_uses_cached_sum_when_nan_free() -> VortexResult<()> { + // With an exact NaNCount of zero, the planted exact Sum stat is usable as-is. + let arr = + PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array(); + arr.statistics() + .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(0u64))); + arr.statistics() + .set(Stat::Sum, Precision::Exact(ScalarValue::from(42.0f64))); + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert_eq!(result.as_primitive().typed_value::(), Some(42.0)); + Ok(()) + } + + #[test] + fn sum_constant_nan() -> VortexResult<()> { + let arr = ConstantArray::new(f64::NAN, 4).into_array(); + // NaN constants are skipped by default and poison the sum otherwise. + let result = sum_with_options(&arr, NumericalAggregateOpts::default())?; assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); + + let result = sum_with_options(&arr, NumericalAggregateOpts::include_nans())?; + assert!(result.as_primitive().typed_value::().unwrap().is_nan()); Ok(()) } @@ -335,16 +419,16 @@ mod tests { Validity::NonNullable, ) .into_array(); - let acc = sum(&batch, &mut LEGACY_SESSION.create_execution_ctx())?; + let acc = sum(&batch, &mut array_session().create_execution_ctx())?; // INFINITY + NEG_INFINITY = NaN, which is treated as saturated assert!(acc.as_primitive().typed_value::().unwrap().is_nan()); let mut acc = Accumulator::try_new( Sum, - EmptyOptions, + NumericalAggregateOpts::default(), DType::Primitive(PType::F64, Nullability::NonNullable), )?; - acc.accumulate(&batch, &mut LEGACY_SESSION.create_execution_ctx())?; + acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; assert!(acc.is_saturated()); Ok(()) } @@ -352,7 +436,7 @@ mod tests { #[test] fn sum_checked_overflow() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); - let result = sum(&arr, &mut LEGACY_SESSION.create_execution_ctx())?; + let result = sum(&arr, &mut array_session().create_execution_ctx())?; assert!(result.is_null()); Ok(()) } @@ -360,12 +444,12 @@ mod tests { #[test] fn sum_checked_overflow_is_saturated() -> VortexResult<()> { let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); - let mut acc = Accumulator::try_new(Sum, EmptyOptions, dtype)?; + let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; assert!(!acc.is_saturated()); let batch = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); - acc.accumulate(&batch, &mut LEGACY_SESSION.create_execution_ctx())?; + acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; assert!(acc.is_saturated()); // finish resets state, clearing saturation 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 7a77ed9555c..d73e670f50f 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 @@ -27,6 +27,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::Canonical; @@ -100,11 +102,20 @@ impl AggregateFnVTable for UncompressedSizeInBytes { type Partial = u64; fn id(&self) -> AggregateFnId { - AggregateFnId::new("vortex.uncompressed_size_in_bytes") + static ID: CachedId = CachedId::new("vortex.uncompressed_size_in_bytes"); + *ID } fn serialize(&self, _options: &Self::Options) -> VortexResult>> { - unimplemented!("UncompressedSizeInBytes is not yet serializable"); + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { @@ -307,7 +318,6 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; @@ -315,6 +325,7 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -351,7 +362,7 @@ mod tests { } fn aggregate(array: &ArrayRef) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut acc = Accumulator::try_new(UncompressedSizeInBytes, EmptyOptions, array.dtype().clone())?; acc.accumulate(array, &mut ctx)?; @@ -536,7 +547,7 @@ mod tests { fn variant_stat_is_unsupported() -> VortexResult<()> { let child = ConstantArray::new(Scalar::variant(Scalar::from(42i32)), 3).into_array(); let array = VariantArray::try_new(child, None)?.into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!( array @@ -589,7 +600,7 @@ mod tests { #[test] fn helper_caches_result() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let size = uncompressed_size_in_bytes(&array, &mut ctx)?; diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index bff198b3955..50c585233f0 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -35,6 +35,7 @@ impl AggregateFnRef { /// /// Note: the serialization format is not stable and may change between versions. pub fn from_proto(proto: &pb::AggregateFn, session: &VortexSession) -> VortexResult { + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let agg_fn_id: AggregateFnId = AggregateFnId::new(proto.id.as_str()); let agg_fn = if let Some(plugin) = session.aggregate_fns().find_plugin(&agg_fn_id) { plugin.deserialize(proto.metadata(), session)? @@ -59,6 +60,7 @@ impl AggregateFnRef { #[cfg(test)] mod tests { use prost::Message; + use rstest::rstest; use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_proto::expr as pb; @@ -72,6 +74,8 @@ mod tests { use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::EmptyOptions; + use crate::aggregate_fn::NumericalAggregateOpts; + use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::session::AggregateFnSession; use crate::aggregate_fn::session::AggregateFnSessionExt; use crate::dtype::DType; @@ -85,6 +89,7 @@ mod tests { type Options = EmptyOptions; type Partial = (); + #[expect(clippy::disallowed_methods, reason = "test-only id")] fn id(&self) -> AggregateFnId { AggregateFnId::new("vortex.test.proto") } @@ -155,7 +160,7 @@ mod tests { #[test] fn aggregate_fn_serde() { - let session = VortexSession::empty().with::(); + let session = crate::array_session(); session.aggregate_fns().register(TestAgg); let agg_fn = TestAgg.bind(EmptyOptions); @@ -168,6 +173,26 @@ mod tests { assert_eq!(deserialized, agg_fn); } + /// The `skip_nans` option must survive a protobuf serialize/deserialize round-trip for the + /// numeric aggregates, including the non-default NaN-including configuration. + #[rstest] + #[case(NumericalAggregateOpts::skip_nans())] + #[case(NumericalAggregateOpts::include_nans())] + fn numeric_aggregate_options_round_trip( + #[case] options: NumericalAggregateOpts, + ) -> VortexResult<()> { + let session = crate::array_session(); + let agg_fn = Sum.bind(options); + + let proto = agg_fn.serialize_proto()?; + let buf = proto.encode_to_vec(); + let decoded = pb::AggregateFn::decode(buf.as_slice())?; + let round_tripped = AggregateFnRef::from_proto(&decoded, &session)?; + + assert_eq!(round_tripped, agg_fn); + Ok(()) + } + #[test] fn unknown_aggregate_fn_id_allow_unknown() { let session = VortexSession::empty() diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index c6d7542a687..c496f7a9bf8 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -4,8 +4,8 @@ use std::any::Any; use std::sync::Arc; -use vortex_session::Ref; use vortex_session::SessionExt; +use vortex_session::SessionGuard; use vortex_session::SessionVar; use crate::aggregate_fn::AggregateFnId; @@ -50,7 +50,7 @@ use crate::arrays::dict::compute::min_max::DictMinMaxKernel; /// The default session registers the built-in aggregate functions and kernels. Additional /// aggregate functions and kernels may be registered by extensions when they are added to a /// [`VortexSession`](vortex_session::VortexSession). -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct AggregateFnSession { registry: ArcSwapMap, @@ -222,7 +222,7 @@ impl AggregateFnSession { /// Extension trait for accessing aggregate function session data. pub trait AggregateFnSessionExt: SessionExt { /// Returns the aggregate function session data. - fn aggregate_fns(&self) -> Ref<'_, AggregateFnSession> { + fn aggregate_fns(&self) -> SessionGuard<'_, AggregateFnSession> { self.get::() } } diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 28b91d45166..49b28dd26d7 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -7,8 +7,10 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; +use prost::Message; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_proto::expr as pb; use vortex_session::VortexSession; use crate::ArrayRef; @@ -167,6 +169,73 @@ impl Display for EmptyOptions { } } +/// Options for aggregate functions over primitive numeric inputs, controlling how NaN values in +/// floating-point arrays are handled. +/// +/// When `skip_nans` is `true` (the default), NaN values are treated as missing: they contribute +/// nothing to `sum`/`min`/`max`/`mean` and are excluded from `count`. +/// +/// When `skip_nans` is `false`, NaN values participate in the aggregate: `count` includes them, +/// while any NaN value poisons the result of `sum`/`min`/`max`/`mean` to NaN. +/// +/// The option has no effect on non-float inputs. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct NumericalAggregateOpts { + /// Whether NaN values are skipped (treated as missing) during aggregation. + pub skip_nans: bool, +} + +impl NumericalAggregateOpts { + /// Options that skip NaN values, treating them as missing during aggregation. + /// + /// This is the default configuration; see [`NumericalAggregateOpts::include_nans`] for the + /// NaN-including variant. + pub const fn skip_nans() -> Self { + Self { skip_nans: true } + } + + /// Options that include NaN values in the aggregate: `count` counts them, while any NaN + /// poisons the result of `sum`/`min`/`max`/`mean` to NaN. + /// + /// See [`NumericalAggregateOpts::skip_nans`] for the default NaN-skipping variant. + pub const fn include_nans() -> Self { + Self { skip_nans: false } + } + + /// Serialize these options to protobuf-encoded metadata bytes. + pub fn serialize(&self) -> Vec { + pb::NumericalAggregateOpts { + skip_nans: self.skip_nans, + } + .encode_to_vec() + } + + /// Deserialize these options from protobuf-encoded metadata bytes. + pub fn deserialize(metadata: &[u8]) -> VortexResult { + let opts = pb::NumericalAggregateOpts::decode(metadata)?; + Ok(Self { + skip_nans: opts.skip_nans, + }) + } +} + +impl Default for NumericalAggregateOpts { + fn default() -> Self { + Self::skip_nans() + } +} + +impl Display for NumericalAggregateOpts { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + // Only the non-default configuration is displayed, so that aggregates with default + // options render identically to their pre-options form, e.g. `vortex.sum()`. + if !self.skip_nans { + write!(f, "skip_nans=false")?; + } + Ok(()) + } +} + /// Factory functions for aggregate vtables. pub trait AggregateFnVTableExt: AggregateFnVTable { /// Bind this vtable with the given options into an [`AggregateFnRef`]. diff --git a/vortex-array/src/aliases/mod.rs b/vortex-array/src/aliases/mod.rs index d001e0402e2..6d0ab2524c5 100644 --- a/vortex-array/src/aliases/mod.rs +++ b/vortex-array/src/aliases/mod.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Re-exports of third-party crates we use in macros exported from vortex-array. -//! pub mod paste { //! Re-export of [`paste`](https://docs.rs/paste/latest/paste/). diff --git a/vortex-array/src/arc_swap_map.rs b/vortex-array/src/arc_swap_map.rs index 567e025d48c..73108ae5d57 100644 --- a/vortex-array/src/arc_swap_map.rs +++ b/vortex-array/src/arc_swap_map.rs @@ -23,14 +23,27 @@ use vortex_utils::aliases::hash_map::HashMap; /// optimizer-kernel and aggregate-function registries). Because every write /// clones the entire map, it is intended for maps that are written rarely /// (typically only while a session is being configured) and read often. +/// +/// The map is held behind an [`Arc`] so that [`Clone`] shares the same +/// underlying cell: a registry mutated through one clone is observed by all +/// others. Session variables rely on this so that encodings registered after a +/// session is built remain visible to clones of that session. pub(crate) struct ArcSwapMap { - inner: ArcSwap>, + inner: Arc>>, } impl Default for ArcSwapMap { fn default() -> Self { Self { - inner: ArcSwap::from_pointee(HashMap::default()), + inner: Arc::new(ArcSwap::from_pointee(HashMap::default())), + } + } +} + +impl Clone for ArcSwapMap { + fn clone(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), } } } @@ -42,6 +55,11 @@ impl Debug for ArcSwapMap { } impl ArcSwapMap { + /// Return the currently published map snapshot. + pub(crate) fn snapshot(&self) -> Arc> { + self.inner.load_full() + } + /// Read the current snapshot, passing it to `f`. /// /// Every lookup inside `f` observes the same snapshot, which matters when a @@ -50,14 +68,6 @@ impl ArcSwapMap { f(&self.inner.load()) } - /// Load the current snapshot as an owned [`Arc`]. - /// - /// This is zero-copy: a single `Arc` refcount bump returning the same map - /// instance held inside the [`ArcSwap`], with no allocation or map copy. - pub(crate) fn load_full(&self) -> Arc> { - self.inner.load_full() - } - /// Replace the map with the result of applying `f` to a private copy. /// /// Writes are copy-on-write via [`ArcSwap::rcu`], so `f` may run more than @@ -156,4 +166,30 @@ mod tests { map.insert(2, 2); assert_eq!(map.read(|m| m.values().sum::()), 3); } + + #[test] + fn snapshot_keeps_published_view() { + let map = ArcSwapMap::::default(); + map.insert(1, 10); + + let snapshot = map.snapshot(); + map.insert(1, 20); + map.insert(2, 30); + + assert_eq!(snapshot.get(&1), Some(&10)); + assert_eq!(snapshot.get(&2), None); + assert_eq!(map.get(&1), Some(20)); + assert_eq!(map.get(&2), Some(30)); + } + + #[test] + fn clone_shares_the_same_cell() { + let map = ArcSwapMap::::default(); + let clone = map.clone(); + // A write through one handle is observed through the other. + map.insert(1, 10); + assert_eq!(clone.get(&1), Some(10)); + clone.insert(2, 20); + assert_eq!(map.get(&2), Some(20)); + } } diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index bacda957660..7fb3d7baf01 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -460,8 +460,18 @@ impl ArrayRef { /// This is only valid for physical rewrites: the replacement must have the same logical /// `DType` and `len` as the existing slot. /// + /// # Safety + /// + /// If this returns `Ok`, the caller must guarantee that the replacement slot represents the + /// same logical values as the original slot. Only the physical representation may change. + /// Existing parent statistics are preserved and must remain valid. + /// /// Takes ownership to allow in-place mutation when the refcount is 1. - pub fn with_slot(self, slot_idx: usize, replacement: ArrayRef) -> VortexResult { + pub unsafe fn with_slot( + self, + slot_idx: usize, + replacement: ArrayRef, + ) -> VortexResult { let mut slots: ArraySlots = self.slots().iter().cloned().collect(); let nslots = slots.len(); vortex_ensure!( @@ -488,7 +498,8 @@ impl ArrayRef { replacement.len() ); slots[slot_idx] = Some(replacement); - self.with_slots(slots) + // SAFETY: upheld by the caller of this unsafe API. + unsafe { self.with_slots(slots) } } /// Take a slot for executor-owned physical rewrites. @@ -555,7 +566,13 @@ impl ArrayRef { /// /// This is only valid for physical rewrites: slot count, presence, logical `DType`, and /// logical `len` must remain unchanged. - pub fn with_slots(self, slots: ArraySlots) -> VortexResult { + /// + /// # Safety + /// + /// If this returns `Ok`, the caller must guarantee that each replacement slot represents the + /// same logical values as the original slot. Only physical representation may change. Existing + /// parent statistics are preserved and must remain valid. + pub unsafe fn with_slots(self, slots: ArraySlots) -> VortexResult { let old_slots = self.slots(); vortex_ensure!( old_slots.len() == slots.len(), @@ -589,6 +606,47 @@ impl ArrayRef { self.0.data.with_slots(&self, slots) } + /// Returns a new array with the provided top-level buffer handles. + /// + /// This is only valid for physical rewrites: buffer count, logical `DType`, logical `len`, and + /// child slots must remain unchanged. Encoding-specific validation checks buffer shape, + /// alignment, and metadata consistency. + /// + /// # Safety + /// + /// If this returns `Ok`, the caller must guarantee that the replacement buffers represent the + /// same logical values as the original buffers. Only the buffer handle implementation, + /// placement, or backing storage may change. Existing statistics are preserved and must remain + /// valid. + pub unsafe fn with_buffers( + self, + buffers: impl IntoIterator, + ) -> VortexResult { + let buffers = buffers.into_iter().collect::>(); + let nbuffers = self.nbuffers(); + vortex_ensure!( + nbuffers == buffers.len(), + "buffer count changed from {} to {} during physical rewrite", + nbuffers, + buffers.len() + ); + for (idx, (old_buffer, new_buffer)) in self + .buffer_handles() + .into_iter() + .zip(buffers.iter()) + .enumerate() + { + vortex_ensure!( + old_buffer.len() == new_buffer.len(), + "buffer {} length changed from {} to {} during physical rewrite", + idx, + old_buffer.len(), + new_buffer.len() + ); + } + self.0.data.with_buffers(&self, buffers) + } + pub fn reduce(&self) -> VortexResult> { self.0.data.reduce(self) } @@ -623,15 +681,6 @@ impl ArrayRef { unsafe { (&*inner).data.execute_unchecked(self, ctx) } } - pub fn execute_parent( - &self, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - self.0.data.execute_parent(self, parent, child_idx, ctx) - } - // ArrayVisitor delegation methods /// Returns the children of the array. diff --git a/vortex-array/src/array/foreign.rs b/vortex-array/src/array/foreign.rs index e7ea7238346..4a665044f3f 100644 --- a/vortex-array/src/array/foreign.rs +++ b/vortex-array/src/array/foreign.rs @@ -120,6 +120,20 @@ impl VTable for ForeignArray { Some(format!("buffer[{idx}]")) } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + Ok(ArrayParts::new( + self.clone(), + array.dtype().clone(), + array.len(), + ForeignArrayData::new(array.metadata.clone(), buffers.to_vec()), + ) + .with_slots(array.slots().iter().cloned().collect())) + } + fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index b1b22c4c529..0d3bbe9cae3 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -127,6 +127,9 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { /// Returns a new array with the given slots. fn with_slots(&self, this: &ArrayRef, slots: ArraySlots) -> VortexResult; + /// Returns a new array with the given buffers. + fn with_buffers(&self, this: &ArrayRef, buffers: Vec) -> VortexResult; + /// Returns a new array with the given slots, bypassing encoding-level validation. /// /// Used by the executor to temporarily carry an array that has had one of its child slots @@ -177,15 +180,6 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { ctx: &mut ExecutionCtx, ) -> VortexResult; - /// Attempt to execute the parent of this array. - fn execute_parent( - &self, - this: &ArrayRef, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult>; - /// Execute the scalar at the given index. /// /// This method panics if the index is out of bounds for the array. @@ -199,6 +193,7 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { /// Trait for converting a type into a Vortex [`ArrayRef`]. pub trait IntoArray { + /// Convert this value into the erased array handle used by generic APIs. fn into_array(self) -> ArrayRef; } @@ -371,6 +366,16 @@ impl DynArrayData for ArrayData { .into_array()) } + fn with_buffers(&self, this: &ArrayRef, buffers: Vec) -> VortexResult { + let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; + let stats = this.statistics().to_owned(); + Ok( + Array::::try_from_parts(V::with_buffers(&self.vtable, view, &buffers)?)? + .with_stats_set(stats) + .into_array(), + ) + } + unsafe fn with_slots_unchecked(&self, this: &ArrayRef, slots: ArraySlots) -> ArrayRef { // SAFETY: we intentionally skip `V::validate` here. Caller guarantees that the resulting // array is either repaired or not externally observed. @@ -474,32 +479,6 @@ impl DynArrayData for ArrayData { V::execute(typed, ctx) } - fn execute_parent( - &self, - this: &ArrayRef, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; - let Some(result) = V::execute_parent(view, parent, child_idx, ctx)? else { - return Ok(None); - }; - - if cfg!(debug_assertions) { - vortex_ensure!( - result.len() == parent.len(), - "Executed parent canonical length mismatch" - ); - vortex_ensure!( - result.dtype() == parent.dtype(), - "Executed parent canonical dtype mismatch" - ); - } - - Ok(Some(result)) - } - fn execute_scalar( &self, this: &ArrayRef, diff --git a/vortex-array/src/array/typed.rs b/vortex-array/src/array/typed.rs index 180f771bc9c..b88f3792a84 100644 --- a/vortex-array/src/array/typed.rs +++ b/vortex-array/src/array/typed.rs @@ -3,6 +3,11 @@ //! Typed array wrappers: [`ArrayData`] (heap-allocated), [`Array`] (typed handle), //! and [`ArrayView`] (lightweight borrow). +//! +//! Encoding implementors normally construct arrays through [`ArrayParts`] and +//! [`Array::try_from_parts`]. Compute and serialization code should accept [`ArrayRef`] when it can +//! operate over any encoding, and downcast to [`Array`] or [`ArrayView`] only when it needs +//! encoding-specific state. use std::any::Any; use std::fmt::Debug; @@ -48,14 +53,22 @@ pub(crate) struct ArrayInner { /// Construction parameters for typed arrays. pub struct ArrayParts { + /// The vtable value identifying the array encoding. pub vtable: V, + /// Logical dtype of every row in the array. pub dtype: DType, + /// Number of rows in the array. pub len: usize, + /// Encoding-specific, non-child data. pub data: V::TypedArrayData, + /// Optional child arrays owned by this encoding. pub slots: ArraySlots, } impl ArrayParts { + /// Construct array parts with no child slots. + /// + /// The parts are not validated until they are passed to [`Array::try_from_parts`]. pub fn new(vtable: V, dtype: DType, len: usize, data: V::TypedArrayData) -> Self { Self { vtable, @@ -66,6 +79,9 @@ impl ArrayParts { } } + /// Attach child slots to the construction parts. + /// + /// Slot count, names, and meaning are encoding-specific and validated by [`VTable::validate`]. pub fn with_slots(mut self, slots: ArraySlots) -> Self { self.slots = slots; self @@ -180,6 +196,12 @@ impl Debug for ArrayData { /// `Array` holds an [`ArrayRef`] (shared, heap-allocated) and provides typed access /// to the encoding-specific data via [`Deref`] to `V::TypedArrayData`. /// +/// Buffers are intentionally not stored in the common `ArrayInner` or [`ArrayParts`] state. +/// Encodings may expose buffers only when writing or serializing, and those buffers need not be the +/// same representation they keep in memory. For example, an encoding may hold a deserialized +/// in-memory data structure and synthesize serialized buffers at write time; hoisting buffers here +/// would force it to retain both forms or make the serialized layout dictate the runtime layout. +/// /// This is the primary type for working with typed arrays. Convert to [`ArrayRef`] /// via [`into_array()`](IntoArray::into_array) or [`AsRef`]. pub struct Array { @@ -189,6 +211,9 @@ pub struct Array { impl Array { /// Create a typed array from explicit construction parameters. + /// + /// This is the safe construction path for encoding implementors. It calls + /// [`VTable::validate`] before publishing the array as an [`ArrayRef`]. pub fn try_from_parts(new: ArrayParts) -> VortexResult { let store = ArrayInner::>::try_new(new)?; let inner = ArrayRef::from_inner(Arc::new(store)); @@ -232,7 +257,9 @@ impl Array { } } - /// Try to create from an `ArrayRef`, returning `Err` if the type doesn't match. + /// Try to create a typed handle from an [`ArrayRef`]. + /// + /// Returns the original [`ArrayRef`] in `Err` when the encoding id does not match `V`. pub fn try_from_array_ref(array: ArrayRef) -> Result { if array.is::() { Ok(Self { @@ -244,7 +271,7 @@ impl Array { } } - /// Returns the dtype. + /// Returns the logical dtype. pub fn dtype(&self) -> &DType { self.inner.dtype() } @@ -259,12 +286,12 @@ impl Array { self.inner.len() == 0 } - /// Returns the encoding ID. + /// Returns the encoding ID for `V`. pub fn encoding_id(&self) -> ArrayId { self.inner.encoding_id() } - /// Returns the statistics. + /// Returns this array's statistics set. pub fn statistics(&self) -> StatsSetRef<'_> { self.inner.statistics() } @@ -274,7 +301,9 @@ impl Array { &self.downcast_inner().data } - /// Try to fetch a mut ref to the inner ArrayData. + /// Try to fetch mutable access to the encoding-specific data. + /// + /// Returns `None` when this handle is not the unique owner of the backing allocation. pub fn data_mut(&mut self) -> Option<&mut V::TypedArrayData> { let store = self.inner.inner_mut()?; let array_inner = store.data.as_any_mut().downcast_mut::>(); @@ -302,6 +331,7 @@ impl Array { } } + /// Replace the array's statistics set and return the same typed handle. pub fn with_stats_set(self, stats: StatsSet) -> Self { self.statistics().replace(stats); self @@ -343,6 +373,7 @@ impl Array { /// Public API methods that shadow `DynArrayData` / `ArrayRef` methods. impl Array { + /// Lazily or eagerly slice the array to `range`, depending on available kernels. pub fn slice(&self, range: std::ops::Range) -> VortexResult { self.inner.slice(range) } @@ -365,30 +396,37 @@ impl Array { self.inner.execute_scalar(index, ctx) } + /// Filter the array with a selection mask. pub fn filter(&self, mask: vortex_mask::Mask) -> VortexResult { self.inner.filter(mask) } + /// Gather rows from this array by index. pub fn take(&self, indices: ArrayRef) -> VortexResult { self.inner.take(indices) } + /// Returns the array's validity representation. pub fn validity(&self) -> VortexResult { self.inner.validity() } + /// Returns whether `index` is valid using the provided execution context. pub fn is_valid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { self.inner.is_valid(index, ctx) } + /// Returns whether `index` is null using the provided execution context. pub fn is_invalid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { self.inner.is_invalid(index, ctx) } + /// Returns whether every row is valid. pub fn all_valid(&self, ctx: &mut ExecutionCtx) -> VortexResult { self.inner.all_valid(ctx) } + /// Returns whether every row is null. pub fn all_invalid(&self, ctx: &mut ExecutionCtx) -> VortexResult { self.inner.all_invalid(ctx) } @@ -400,26 +438,32 @@ impl Array { result } + /// Returns the estimated physical bytes owned or referenced by this array tree. pub fn nbytes(&self) -> u64 { self.inner.nbytes() } + /// Returns the number of top-level buffers exposed by this encoding. pub fn nbuffers(&self) -> usize { self.inner.nbuffers() } + /// Returns the scalar value when this array is known to be constant. pub fn as_constant(&self) -> Option { self.inner.as_constant() } + /// Counts valid rows, executing validity arrays when necessary. pub fn valid_count(&self, ctx: &mut ExecutionCtx) -> VortexResult { self.inner.valid_count(ctx) } + /// Counts null rows, executing validity arrays when necessary. pub fn invalid_count(&self, ctx: &mut ExecutionCtx) -> VortexResult { self.inner.invalid_count(ctx) } + /// Append this array's logical values to a canonical builder. pub fn append_to_builder( &self, builder: &mut dyn crate::builders::ArrayBuilder, @@ -489,6 +533,8 @@ mod tests { use vortex_buffer::buffer; use super::Array; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -496,17 +542,19 @@ mod tests { #[test] fn typed_array_into_parts_roundtrips() { + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let expected = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let parts = array.try_into_parts().unwrap(); let rebuilt = Array::::try_from_parts(parts).unwrap(); - assert_arrays_eq!(rebuilt, expected); + assert_arrays_eq!(rebuilt, expected, &mut ctx); } #[test] fn typed_array_try_into_parts_requires_unique_owner() { + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let alias = array.clone(); @@ -515,6 +563,6 @@ mod tests { Err(array) => array, }; - assert_arrays_eq!(array, alias); + assert_arrays_eq!(array, alias, &mut ctx); } } diff --git a/vortex-array/src/array/view.rs b/vortex-array/src/array/view.rs index 969fd0484a3..e5e937e0141 100644 --- a/vortex-array/src/array/view.rs +++ b/vortex-array/src/array/view.rs @@ -16,6 +16,10 @@ use crate::stats::StatsSetRef; use crate::validity::Validity; /// A lightweight, `Copy`-able typed view into an [`ArrayRef`]. +/// +/// Vtable methods receive `ArrayView` so they can inspect common array metadata and +/// encoding-specific data without cloning or downcasting an [`ArrayRef`]. The view is only valid for +/// the lifetime of the borrowed array. pub struct ArrayView<'a, V: VTable> { array: &'a ArrayRef, data: &'a V::TypedArrayData, @@ -37,42 +41,52 @@ impl<'a, V: VTable> ArrayView<'a, V> { Self { array, data } } + /// Returns the erased array handle that owns this view. pub fn array(&self) -> &'a ArrayRef { self.array } + /// Returns the encoding-specific data stored by this array. pub fn data(&self) -> &'a V::TypedArrayData { self.data } + /// Returns this array's child slots. pub fn slots(&self) -> &'a [Option] { self.array.slots() } + /// Returns the logical dtype. pub fn dtype(&self) -> &DType { self.array.dtype() } + /// Returns the number of rows. pub fn len(&self) -> usize { self.array.len() } + /// Returns `true` when the array has no rows. pub fn is_empty(&self) -> bool { self.array.len() == 0 } + /// Returns the encoding ID. pub fn encoding_id(&self) -> ArrayId { self.array.encoding_id() } + /// Returns the array's statistics set. pub fn statistics(&self) -> StatsSetRef<'_> { self.array.statistics() } + /// Returns the array's validity representation. pub fn validity(&self) -> VortexResult { self.array.validity() } + /// Clone the underlying [`ArrayRef`] and return it as an owned typed handle. pub fn into_owned(self) -> Array { // SAFETY: we are ourselves type checked as 'V' unsafe { Array::::from_array_ref_unchecked(self.array.clone()) } diff --git a/vortex-array/src/array/vtable/mod.rs b/vortex-array/src/array/vtable/mod.rs index 366eed77ad9..617b211a4c2 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -2,6 +2,14 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! This module contains the VTable definitions for a Vortex encoding. +//! +//! A Vortex array encoding is implemented by a small static vtable type plus an associated +//! `TypedArrayData` value stored in each array instance. The vtable owns behavior such as +//! validation, serialization, execution, child traversal, scalar access, and validity access. +//! +//! The public [`ArrayRef`] API performs common precondition checks before calling +//! into these traits. Implementations should focus on encoding-specific work and uphold the +//! documented postconditions. mod operations; mod validity; @@ -15,6 +23,8 @@ pub use operations::*; pub use validity::*; 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; @@ -54,15 +64,25 @@ use crate::validity::Validity; /// out of bounds). Post-conditions are validated after invocation of the vtable function and will /// panic if violated. pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { + /// Per-array data owned by this encoding, excluding child arrays. + /// + /// Child arrays belong in [`ArrayParts::slots`](crate::ArrayParts::slots) so traversal, + /// serialization, and layout writers can discover them generically. type TypedArrayData: 'static + Send + Sync + Clone + Debug + Display + ArrayHash + ArrayEq; + /// Scalar and element-wise operation hooks for this encoding. type OperationsVTable: OperationsVTable; + /// Validity hook for nullable instances of this encoding. type ValidityVTable: ValidityVTable; /// Returns the ID of the array. fn id(&self) -> ArrayId; /// Validates that externally supplied logical metadata matches the array data. + /// + /// This is called by [`Array::try_from_parts`](crate::Array::try_from_parts) before the array + /// is published. Implementations should check dtype, length, slot count, child dtypes/lengths, + /// metadata bounds, and any buffer shape invariants that unsafe accessors depend on. fn validate( &self, data: &Self::TypedArrayData, @@ -71,7 +91,7 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { slots: &[Option], ) -> VortexResult<()>; - /// Returns the number of buffers in the array. + /// Returns the number of top-level buffers in the array. fn nbuffers(array: ArrayView<'_, Self>) -> usize; /// Returns the buffer at the given index. @@ -83,6 +103,18 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { /// Returns the name of the buffer at the given index, or `None` if unnamed. fn buffer_name(array: ArrayView<'_, Self>, idx: usize) -> Option; + /// Rebuild this array with replacement top-level buffers. + /// + /// This is for physical rewrites that preserve `dtype`, `len`, child slots, buffer count, and + /// buffer lengths. The caller checks the generic invariants before dispatching here; + /// implementations should interpret the replacement buffers for their encoding-specific + /// in-memory representation. + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult>; + /// Returns the number of children in the array. /// /// The default counts non-None slots. @@ -122,14 +154,21 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { .vortex_expect("child_name index out of bounds") } - /// Serialize metadata into a byte buffer for IPC or file storage. - /// Return `None` if the array cannot be serialized. + /// 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 + /// serialized separately through [`buffer`](Self::buffer), [`nbuffers`](Self::nbuffers), and + /// child traversal. fn serialize( array: ArrayView<'_, Self>, session: &VortexSession, ) -> VortexResult>>; - /// Deserialize an array from serialized components. + /// Deserialize an array from serialized metadata, buffers, and children. + /// + /// The returned [`ArrayParts`] are still validated by the generic adapter. + /// Deserializers should use the provided `session` to resolve plugin-owned metadata instead of + /// relying on global state. fn deserialize( &self, dtype: &DType, @@ -138,9 +177,12 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { buffers: &[BufferHandle], children: &dyn ArrayChildren, session: &VortexSession, - ) -> VortexResult>; + ) -> VortexResult>; - /// Writes the array into a canonical builder. + /// Writes the array's logical values into a canonical builder. + /// + /// The default implementation executes the full array to [`Canonical`] and appends that result. + /// Encodings may override this to avoid materializing an intermediate canonical array. fn append_to_builder( array: ArrayView<'_, Self>, builder: &mut dyn ArrayBuilder, @@ -185,24 +227,18 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { /// incorrectly contains null values. fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult; - /// Attempt to execute the parent of this array. - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - _ = (array, parent, child_idx, ctx); - Ok(None) - } - - /// Attempt to reduce the array to a simpler representation. + /// Attempt to reduce the array to a simpler representation without changing logical values. + /// + /// Reductions are opportunistic and may return `Ok(None)` when no cheaper representation is + /// known. fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { _ = array; Ok(None) } - /// Attempt to perform a reduction of the parent of this array. + /// Attempt to reduce `parent` after this array appears as one of its children. + /// + /// This is used by lazy arrays to let child execution unlock parent simplifications. fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, @@ -217,6 +253,7 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { pub use VTable as ArrayVTable; use crate::array::ArrayId; +use crate::array::ArrayParts; /// Empty array metadata struct for encodings with no per-array metadata. #[derive(Clone, Debug, Default)] @@ -237,6 +274,40 @@ impl Display for EmptyArrayData { } } +/// Rebuild an array that has no top-level buffers. +#[inline] +pub fn with_empty_buffers( + vtable: &V, + array: ArrayView<'_, V>, + buffers: &[BufferHandle], +) -> VortexResult> { + vortex_ensure!( + buffers.is_empty(), + "Array {} expects 0 buffers, got {}", + array.encoding_id(), + buffers.len() + ); + Ok(ArrayParts::new( + vtable.clone(), + array.dtype().clone(), + array.len(), + array.data().clone(), + ) + .with_slots(array.slots().iter().cloned().collect())) +} + +/// Reject buffer replacement for encodings whose exposed buffers are not runtime backing buffers. +#[inline] +pub fn unsupported_buffer_replacement( + array: ArrayView<'_, V>, + _buffers: &[BufferHandle], +) -> VortexResult> { + vortex_bail!( + "Array {} does not support in-memory buffer replacement", + array.encoding_id() + ) +} + /// Placeholder type used to indicate when a particular vtable is not supported by the encoding. pub struct NotSupported; diff --git a/vortex-array/src/array/vtable/operations.rs b/vortex-array/src/array/vtable/operations.rs index a002bc488fa..7f49e683640 100644 --- a/vortex-array/src/array/vtable/operations.rs +++ b/vortex-array/src/array/vtable/operations.rs @@ -10,13 +10,23 @@ use crate::array::VTable; use crate::scalar::Scalar; use crate::vtable::NotSupported; +/// Element-level operations for an array encoding. +/// +/// This trait is separated from [`VTable`] so encodings can organize scalar +/// access independently from traversal, serialization, and execution. The erased +/// [`ArrayRef`](crate::ArrayRef) +/// methods perform common checks before dispatching here. pub trait OperationsVTable { /// Fetch the scalar at the given index. /// /// ## Preconditions /// /// Bounds-checking has already been performed by the time this function is called, - /// and the index is guaranteed to be non-null. + /// and the index is guaranteed to be non-null. Implementations may assume `index < len`. + /// + /// ## Postconditions + /// + /// The returned [`Scalar`] must have the same logical dtype as the array's element dtype. fn scalar_at( array: ArrayView<'_, V>, index: usize, diff --git a/vortex-array/src/array/vtable/validity.rs b/vortex-array/src/array/vtable/validity.rs index 2cb67facd20..b111d970381 100644 --- a/vortex-array/src/array/vtable/validity.rs +++ b/vortex-array/src/array/vtable/validity.rs @@ -8,12 +8,22 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::validity::Validity; +/// Validity access for nullable instances of an encoding. +/// +/// Non-nullable arrays bypass this hook and report [`Validity::NonNullable`]. Nullable arrays call +/// into the encoding so it can expose either a constant validity state or a row-aligned boolean +/// child array. pub trait ValidityVTable { /// Returns the [`Validity`] of the array. /// /// ## Pre-conditions /// /// - The array DType is nullable. + /// + /// ## Post-conditions + /// + /// If this returns [`Validity::Array`], the child array must have the same length as `array` + /// and non-nullable boolean dtype. fn validity(array: ArrayView<'_, V>) -> VortexResult; } @@ -21,7 +31,9 @@ pub trait ValidityVTable { /// to a child array. pub struct ValidityVTableFromChild; +/// Helper trait for encodings whose validity is exactly one child slot. pub trait ValidityChild { + /// Returns the child array that carries validity for `array`. fn validity_child(array: ArrayView<'_, V>) -> ArrayRef; } @@ -38,9 +50,12 @@ where /// and a slice into it. pub struct ValidityVTableFromChildSliceHelper; +/// Helper for encodings that keep an unsliced validity child plus a local slice range. pub trait ValidityChildSliceHelper { + /// Returns `(unsliced_validity, start, stop)` for this array's logical slice. fn unsliced_child_and_slice(&self) -> (&ArrayRef, usize, usize); + /// Returns a sliced validity child array for the logical range. fn sliced_child_array(&self) -> VortexResult { let (unsliced_validity, start, stop) = self.unsliced_child_and_slice(); unsliced_validity.slice(start..stop) diff --git a/vortex-array/src/arrays/assertions.rs b/vortex-array/src/arrays/assertions.rs index 389baa23857..aff8d4e32cb 100644 --- a/vortex-array/src/arrays/assertions.rs +++ b/vortex-array/src/arrays/assertions.rs @@ -9,9 +9,8 @@ use vortex_error::VortexExpect; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::RecursiveCanonical; -use crate::VortexSessionExecute; +use crate::aggregate_fn::fns::all_non_distinct::all_non_distinct; fn format_indices>(indices: I) -> impl Display { indices.into_iter().format(",") @@ -28,14 +27,14 @@ fn execute_to_canonical(array: ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { /// Finds indices where two arrays differ based on `scalar_at` comparison. #[expect(clippy::unwrap_used)] -fn find_mismatched_indices(left: &ArrayRef, right: &ArrayRef) -> Vec { +fn find_mismatched_indices( + left: &ArrayRef, + right: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> Vec { assert_eq!(left.len(), right.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); (0..left.len()) - .filter(|i| { - left.execute_scalar(*i, &mut ctx).unwrap() - != right.execute_scalar(*i, &mut ctx).unwrap() - }) + .filter(|i| left.execute_scalar(*i, ctx).unwrap() != right.execute_scalar(*i, ctx).unwrap()) .collect() } @@ -46,22 +45,16 @@ fn find_mismatched_indices(left: &ArrayRef, right: &ArrayRef) -> Vec { /// # Example /// ```ignore /// let arr = PrimitiveArray::from_iter([1, 2, 3]); -/// assert_nth_scalar!(arr, 0, 1); -/// assert_nth_scalar!(arr, 1, 2); +/// assert_nth_scalar!(arr, 0, 1, &mut ctx); +/// assert_nth_scalar!(arr, 1, 2, &mut ctx); /// ``` #[macro_export] macro_rules! assert_nth_scalar { - ($arr:expr, $n:expr, $expected:expr) => {{ + ($arr:expr, $n:expr, $expected:expr, $ctx:expr) => {{ use $crate::IntoArray as _; - use $crate::LEGACY_SESSION; - use $crate::VortexSessionExecute as _; let arr_ref: $crate::ArrayRef = $crate::IntoArray::into_array($arr.clone()); - assert_eq!( - arr_ref - .execute_scalar($n, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - $expected.try_into().unwrap() - ); + let expected = $expected.try_into().unwrap(); + assert_eq!(arr_ref.execute_scalar($n, $ctx).unwrap(), expected); }}; } @@ -71,33 +64,25 @@ macro_rules! assert_nth_scalar { /// /// ```ignore /// let arr = PrimitiveArray::from_option_iter([Some(1), None, Some(3)]); -/// assert_nth_scalar_null!(arr, 1); +/// assert_nth_scalar_null!(arr, 1, &mut ctx); /// ``` #[macro_export] macro_rules! assert_nth_scalar_is_null { - ($arr:expr, $n:expr) => {{ - use $crate::LEGACY_SESSION; - use $crate::VortexSessionExecute as _; + ($arr:expr, $n:expr, $ctx:expr) => {{ let arr_ref: $crate::ArrayRef = $crate::IntoArray::into_array($arr.clone()); + let scalar = arr_ref.execute_scalar($n, $ctx).unwrap(); assert!( - arr_ref - .execute_scalar($n, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .is_null(), + scalar.is_null(), "expected scalar at index {} to be null, but was {:?}", $n, - arr_ref - .execute_scalar($n, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + scalar ); }}; } #[macro_export] macro_rules! assert_arrays_eq { - ($left:expr, $right:expr) => {{ - - + ($left:expr, $right:expr, $ctx:expr) => {{ let left: $crate::ArrayRef = $crate::IntoArray::into_array($left.clone()); let right: $crate::ArrayRef = $crate::IntoArray::into_array($right.clone()); if left.dtype() != right.dtype() { @@ -122,7 +107,7 @@ macro_rules! assert_arrays_eq { let left = left.clone(); let right = right.clone(); - $crate::arrays::assert_arrays_eq_impl(&left, &right); + $crate::arrays::assert_arrays_eq_impl(&left, &right, $ctx); }}; } @@ -130,13 +115,17 @@ macro_rules! assert_arrays_eq { /// `ArrayRef`. #[track_caller] #[expect(clippy::panic)] -pub fn assert_arrays_eq_impl(left: &ArrayRef, right: &ArrayRef) { - let executed = execute_to_canonical(left.clone(), &mut LEGACY_SESSION.create_execution_ctx()); +pub fn assert_arrays_eq_impl(left: &ArrayRef, right: &ArrayRef, ctx: &mut ExecutionCtx) { + let executed = execute_to_canonical(left.clone(), ctx); + + let left_right_the_same = + all_non_distinct(left, right, ctx).vortex_expect("failed to compare left and right"); + let executed_right_the_same = all_non_distinct(&executed, right, ctx) + .vortex_expect("failed to compare executed left and right"); - let left_right = find_mismatched_indices(left, right); - let executed_right = find_mismatched_indices(&executed, right); + if !left_right_the_same || !executed_right_the_same { + let left_right = find_mismatched_indices(left, right, ctx); - if !left_right.is_empty() || !executed_right.is_empty() { let mut msg = String::new(); if !left_right.is_empty() { msg.push_str(&format!( @@ -144,6 +133,8 @@ pub fn assert_arrays_eq_impl(left: &ArrayRef, right: &ArrayRef) { format_indices(left_right) )); } + + let executed_right = find_mismatched_indices(&executed, right, ctx); if !executed_right.is_empty() { msg.push_str(&format!( "\n executed != right at indices: {}", diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index 82df2f9c167..8b83d0f8128 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -52,7 +52,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; /// ``` /// # fn main() -> vortex_error::VortexResult<()> { /// use vortex_array::arrays::BoolArray; -/// use vortex_array::{IntoArray, LEGACY_SESSION, VortexSessionExecute}; +/// use vortex_array::{IntoArray, array_session, VortexSessionExecute}; /// /// // Create from iterator using FromIterator impl /// let array: BoolArray = [true, false, true, false].into_iter().collect(); @@ -62,7 +62,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; /// assert_eq!(sliced.len(), 2); /// /// // Access individual values -/// let mut ctx = LEGACY_SESSION.create_execution_ctx(); +/// let mut ctx = array_session().create_execution_ctx(); /// let value = array.execute_scalar(0, &mut ctx).unwrap(); /// assert_eq!(value, true.into()); /// # Ok(()) @@ -371,8 +371,8 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; @@ -382,76 +382,54 @@ mod tests { #[test] fn bool_array() { + let mut ctx = array_session().create_execution_ctx(); + let arr = BoolArray::from_iter([true, false, true]); - let scalar = bool::try_from( - &arr.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - ) - .unwrap(); + let scalar = bool::try_from(&arr.execute_scalar(0, &mut ctx).unwrap()).unwrap(); assert!(scalar); } #[test] fn test_all_some_iter() { + let mut ctx = array_session().create_execution_ctx(); + let arr = BoolArray::from_iter([Some(true), Some(false)]); assert!(matches!(arr.validity(), Ok(Validity::AllValid))); - let scalar = bool::try_from( - &arr.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - ) - .unwrap(); + let scalar = bool::try_from(&arr.execute_scalar(0, &mut ctx).unwrap()).unwrap(); assert!(scalar); - let scalar = bool::try_from( - &arr.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - ) - .unwrap(); + let scalar = bool::try_from(&arr.execute_scalar(1, &mut ctx).unwrap()).unwrap(); assert!(!scalar); } #[test] fn test_bool_from_iter() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter([Some(true), Some(true), None, Some(false), None]); - let scalar = bool::try_from( - &arr.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - ) - .unwrap(); + let scalar = bool::try_from(&arr.execute_scalar(0, &mut ctx).unwrap()).unwrap(); assert!(scalar); - let scalar = bool::try_from( - &arr.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - ) - .unwrap(); + let scalar = bool::try_from(&arr.execute_scalar(1, &mut ctx).unwrap()).unwrap(); assert!(scalar); - let scalar = arr - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let scalar = arr.execute_scalar(2, &mut ctx).unwrap(); assert!(scalar.is_null()); - let scalar = bool::try_from( - &arr.execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - ) - .unwrap(); + let scalar = bool::try_from(&arr.execute_scalar(3, &mut ctx).unwrap()).unwrap(); assert!(!scalar); - let scalar = arr - .execute_scalar(4, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let scalar = arr.execute_scalar(4, &mut ctx).unwrap(); assert!(scalar.is_null()); } #[test] fn patch_sliced_bools() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.slice(4..12).unwrap(); - assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8])); + assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut ctx); let arr = { let mut builder = BitBufferMut::new_unset(12); @@ -460,7 +438,11 @@ mod tests { }; let sliced = arr.slice(4..12).unwrap(); let expected_slice: Vec = (4..12).map(|i| (1..12).contains(&i)).collect(); - assert_arrays_eq!(sliced, BoolArray::from_iter(expected_slice.clone())); + assert_arrays_eq!( + sliced, + BoolArray::from_iter(expected_slice.clone()), + &mut ctx + ); // patch the underlying array at index 4 to false let patches = Patches::new( @@ -471,26 +453,26 @@ mod tests { None, ) .unwrap(); - let arr = arr - .patch(&patches, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let arr = arr.patch(&patches, &mut ctx).unwrap(); // After patching index 4 to false: indices 1-3 and 5-11 are true, index 0 and 4 are false let expected_patched: Vec = (0..12).map(|i| (1..12).contains(&i) && i != 4).collect(); - assert_arrays_eq!(arr, BoolArray::from_iter(expected_patched)); + assert_arrays_eq!(arr, BoolArray::from_iter(expected_patched), &mut ctx); // the slice should be unchanged (still has original values before patch) - assert_arrays_eq!(sliced, BoolArray::from_iter(expected_slice)); + assert_arrays_eq!(sliced, BoolArray::from_iter(expected_slice), &mut ctx); } #[test] fn slice_array_in_middle() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(16)); let sliced = arr.slice(4..12).unwrap(); - assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8])); + assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut ctx); } #[test] fn patch_bools_owned() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(16)); let buf_ptr = arr.to_bit_buffer().inner().as_ptr(); @@ -502,21 +484,20 @@ mod tests { None, ) .unwrap(); - let arr = arr - .patch(&patches, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let arr = arr.patch(&patches, &mut ctx).unwrap(); // Verify buffer was reused in place assert_eq!(arr.to_bit_buffer().inner().as_ptr(), buf_ptr); // After patching index 0 to false: [false, true, true, ..., true] (16 values) let expected: BoolArray = once(false).chain(repeat_n(true, 15)).collect(); - assert_arrays_eq!(arr, expected); + assert_arrays_eq!(arr, expected, &mut ctx); } #[test] fn patch_sliced_bools_offset() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.slice(4..15).unwrap(); - assert_arrays_eq!(sliced, BoolArray::from_iter([true; 11])); + assert_arrays_eq!(sliced, BoolArray::from_iter([true; 11]), &mut ctx); } } diff --git a/vortex-array/src/arrays/bool/compute/boolean.rs b/vortex-array/src/arrays/bool/compute/boolean.rs new file mode 100644 index 00000000000..28d1bbb7f52 --- /dev/null +++ b/vortex-array/src/arrays/bool/compute/boolean.rs @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::array::ArrayView; +use crate::arrays::Bool; +use crate::arrays::Constant; +use crate::arrays::bool::BoolArrayExt; +use crate::scalar_fn::fns::binary::BooleanKernel; +use crate::scalar_fn::fns::binary::kleene_boolean_buffer_scalar; +use crate::scalar_fn::fns::binary::kleene_boolean_buffers; +use crate::scalar_fn::fns::operators::Operator; + +impl BooleanKernel for Bool { + fn boolean( + lhs: ArrayView<'_, Self>, + rhs: &ArrayRef, + operator: Operator, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); + + if let Some(rhs) = rhs.as_opt::() { + let rhs = rhs + .scalar() + .as_bool_opt() + .ok_or_else(|| vortex_err!("expected boolean scalar"))?; + return kleene_boolean_buffer_scalar( + lhs.to_bit_buffer(), + lhs.validity()?, + &rhs, + operator, + nullability, + ctx, + ) + .map(Some); + } + + let Some(rhs) = rhs.as_opt::() else { + return Ok(None); + }; + + kleene_boolean_buffers( + lhs.to_bit_buffer(), + lhs.validity()?, + rhs.to_bit_buffer(), + rhs.validity()?, + operator, + nullability, + ctx, + ) + .map(Some) + } +} diff --git a/vortex-array/src/arrays/bool/compute/cast.rs b/vortex-array/src/arrays/bool/compute/cast.rs index 0855af21f25..3b47ce62c8f 100644 --- a/vortex-array/src/arrays/bool/compute/cast.rs +++ b/vortex-array/src/arrays/bool/compute/cast.rs @@ -1,6 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use num_traits::One; +use num_traits::Zero; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use crate::ArrayRef; @@ -9,8 +12,10 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Bool; use crate::arrays::BoolArray; +use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::dtype::DType; +use crate::match_each_native_ptype; use crate::scalar_fn::fns::cast::CastKernel; use crate::scalar_fn::fns::cast::CastReduce; @@ -38,17 +43,34 @@ impl CastKernel for Bool { dtype: &DType, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if !dtype.is_boolean() { - return Ok(None); + if dtype.is_boolean() { + let new_validity = + array + .validity()? + .cast_nullability(dtype.nullability(), array.len(), ctx)?; + return Ok(Some( + BoolArray::new(array.to_bit_buffer(), new_validity).into_array(), + )); } + let DType::Primitive(new_ptype, new_nullability) = dtype else { + return Ok(None); + }; + let new_validity = array .validity()? - .cast_nullability(dtype.nullability(), array.len(), ctx)?; - Ok(Some( - BoolArray::new(array.to_bit_buffer(), new_validity).into_array(), - )) + .cast_nullability(*new_nullability, array.len(), ctx)?; + + let bits = array.to_bit_buffer(); + let len = bits.len(); + + Ok(Some(match_each_native_ptype!(*new_ptype, |T| { + let (one, zero) = (::one(), ::zero()); + let mut buffer = BufferMut::::with_capacity(len); + buffer.extend(bits.iter().map(|v| if v { one } else { zero })); + PrimitiveArray::new(buffer.freeze(), new_validity).into_array() + }))) } } @@ -67,10 +89,9 @@ mod tests { use crate::compute::conformance::cast::test_cast_conformance; use crate::dtype::DType; use crate::dtype::Nullability; - use crate::session::ArraySession; + use crate::dtype::PType; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[test] fn try_cast_bool_success() { @@ -104,4 +125,22 @@ mod tests { fn test_cast_bool_conformance(#[case] array: BoolArray) { test_cast_conformance(&array.into_array()); } + + #[rstest] + #[case(PType::I8)] + #[case(PType::I32)] + #[case(PType::I64)] + #[case(PType::U8)] + #[case(PType::U64)] + #[case(PType::F32)] + #[case(PType::F64)] + fn cast_bool_to_primitive(#[case] target: PType) { + let mut ctx = SESSION.create_execution_ctx(); + let arr = BoolArray::from_iter(vec![true, false, true]).into_array(); + let out = arr + .cast(DType::Primitive(target, Nullability::NonNullable)) + .unwrap(); + let out = out.execute::(&mut ctx).unwrap().into_array(); + assert_eq!(out.len(), 3); + } } diff --git a/vortex-array/src/arrays/bool/compute/filter.rs b/vortex-array/src/arrays/bool/compute/filter.rs index 315c1db2474..f463f27f437 100644 --- a/vortex-array/src/arrays/bool/compute/filter.rs +++ b/vortex-array/src/arrays/bool/compute/filter.rs @@ -303,44 +303,50 @@ mod tests { use super::*; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; #[test] fn filter_bool_test() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_iter([true, false, true]); let filtered = arr.filter(mask).unwrap(); - assert_arrays_eq!(filtered, BoolArray::from_iter([true, false])); + assert_arrays_eq!(filtered, BoolArray::from_iter([true, false]), &mut ctx); } #[test] fn filter_bool_sparse_index_mask() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_indices(3, [0, 2]); let filtered = arr.filter(mask).unwrap(); - assert_arrays_eq!(filtered, BoolArray::from_iter([true, false])); + assert_arrays_eq!(filtered, BoolArray::from_iter([true, false]), &mut ctx); } #[test] fn filter_bool_sparse_slice_mask() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_slices(3, vec![(0, 1), (2, 3)]); let filtered = arr.filter(mask).unwrap(); - assert_arrays_eq!(filtered, BoolArray::from_iter([true, false])); + assert_arrays_eq!(filtered, BoolArray::from_iter([true, false]), &mut ctx); } #[test] fn filter_bool_sparse_buffer_mask() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_buffer(BitBuffer::from_iter([true, false, true])); let filtered = arr.filter(mask).unwrap(); - assert_arrays_eq!(filtered, BoolArray::from_iter([true, false])); + assert_arrays_eq!(filtered, BoolArray::from_iter([true, false]), &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/bool/compute/mod.rs b/vortex-array/src/arrays/bool/compute/mod.rs index 9b71578cd84..6037cac8c36 100644 --- a/vortex-array/src/arrays/bool/compute/mod.rs +++ b/vortex-array/src/arrays/bool/compute/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod boolean; mod cast; mod fill_null; pub(crate) mod filter; @@ -15,6 +16,8 @@ mod tests { use rstest::rstest; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::compute::conformance::consistency::test_array_consistency; @@ -40,6 +43,9 @@ mod tests { None, None, Some(true), None, None, None, Some(false), None, None ]))] fn test_bool_consistency(#[case] array: BoolArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index d70579136a4..abff3527bfe 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -90,6 +90,8 @@ mod test { use crate::IntoArray as _; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; @@ -99,6 +101,7 @@ mod test { #[test] fn take_nullable() { + let mut ctx = array_session().create_execution_ctx(); let reference = BoolArray::from_iter(vec![ Some(false), Some(true), @@ -119,11 +122,12 @@ mod test { let all_invalid_indices = PrimitiveArray::from_option_iter([None::, None, None]); let b = reference.take(all_invalid_indices.into_array()).unwrap(); - assert_arrays_eq!(b, BoolArray::from_iter([None, None, None])); + assert_arrays_eq!(b, BoolArray::from_iter([None, None, None]), &mut ctx); } #[test] fn test_bool_array_take_with_null_out_of_bounds_indices() { + let mut ctx = array_session().create_execution_ctx(); let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -132,11 +136,16 @@ mod test { let actual = values.take(indices.into_array()).unwrap(); // position 3 is null, the third index is null - assert_arrays_eq!(actual, BoolArray::from_iter([Some(false), None, None])); + assert_arrays_eq!( + actual, + BoolArray::from_iter([Some(false), None, None]), + &mut ctx + ); } #[test] fn test_non_null_bool_array_take_with_null_out_of_bounds_indices() { + let mut ctx = array_session().create_execution_ctx(); let values = BoolArray::from_iter(vec![false, true, false, true, false]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -146,30 +155,33 @@ mod test { // the third index is null assert_arrays_eq!( actual, - BoolArray::from_iter([Some(false), Some(true), None]) + BoolArray::from_iter([Some(false), Some(true), None]), + &mut ctx ); } #[test] fn test_bool_array_take_all_null_indices() { + let mut ctx = array_session().create_execution_ctx(); let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]); let indices = PrimitiveArray::new( buffer![0, 3, 100], Validity::Array(BoolArray::from_iter([false, false, false]).into_array()), ); let actual = values.take(indices.into_array()).unwrap(); - assert_arrays_eq!(actual, BoolArray::from_iter([None, None, None])); + assert_arrays_eq!(actual, BoolArray::from_iter([None, None, None]), &mut ctx); } #[test] fn test_non_null_bool_array_take_all_null_indices() { + let mut ctx = array_session().create_execution_ctx(); let values = BoolArray::from_iter(vec![false, true, false, true, false]); let indices = PrimitiveArray::new( buffer![0, 3, 100], Validity::Array(BoolArray::from_iter([false, false, false]).into_array()), ); let actual = values.take(indices.into_array()).unwrap(); - assert_arrays_eq!(actual, BoolArray::from_iter([None, None, None])); + assert_arrays_eq!(actual, BoolArray::from_iter([None, None, None]), &mut ctx); } #[rstest] diff --git a/vortex-array/src/arrays/bool/compute/zip.rs b/vortex-array/src/arrays/bool/compute/zip.rs index 2628b457a85..80e7de80c82 100644 --- a/vortex-array/src/arrays/bool/compute/zip.rs +++ b/vortex-array/src/arrays/bool/compute/zip.rs @@ -89,8 +89,8 @@ mod tests { use super::zip_value_bits; use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::assert_arrays_eq; @@ -132,7 +132,7 @@ mod tests { let bits: Vec = (0..len).map(|i| i.is_multiple_of(5) || i == 64).collect(); let mask = Mask::from_iter(bits.iter().copied()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -147,7 +147,7 @@ mod tests { } })) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -167,7 +167,7 @@ mod tests { let bits: Vec = (0..len).map(|i| i.is_multiple_of(2)).collect(); let mask = Mask::from_iter(bits.iter().copied()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -182,7 +182,7 @@ mod tests { } })) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/arrays/bool/mod.rs b/vortex-array/src/arrays/bool/mod.rs index a4aee2c4faa..ea51e6b5f29 100644 --- a/vortex-array/src/arrays/bool/mod.rs +++ b/vortex-array/src/arrays/bool/mod.rs @@ -15,5 +15,9 @@ pub use compute::rules::BoolMaskedValidityRule; pub use vtable::Bool; pub use vtable::BoolArray; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + #[cfg(feature = "_test-harness")] mod test_harness; diff --git a/vortex-array/src/arrays/bool/patch.rs b/vortex-array/src/arrays/bool/patch.rs index 5efdac30231..babf5a0ae1d 100644 --- a/vortex-array/src/arrays/bool/patch.rs +++ b/vortex-array/src/arrays/bool/patch.rs @@ -47,22 +47,26 @@ mod tests { use vortex_buffer::BitBuffer; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; #[test] fn patch_sliced_bools() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.into_array().slice(4..12).unwrap(); let expected = BoolArray::from_iter([true; 8]); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut ctx); } #[test] fn patch_sliced_bools_offset() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.into_array().slice(4..15).unwrap(); let expected = BoolArray::from_iter([true; 11]); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut ctx); } } diff --git a/vortex-array/src/arrays/bool/test_harness.rs b/vortex-array/src/arrays/bool/test_harness.rs index 4e13d62e3ef..1b6aca704a3 100644 --- a/vortex-array/src/arrays/bool/test_harness.rs +++ b/vortex-array/src/arrays/bool/test_harness.rs @@ -4,19 +4,15 @@ use vortex_error::VortexExpect; use vortex_error::vortex_panic; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; +use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; impl BoolArray { - pub fn opt_bool_vec(&self) -> Vec> { + pub fn opt_bool_vec(&self, ctx: &mut ExecutionCtx) -> Vec> { self.validity() .vortex_expect("failed to get validity") - .execute_mask( - self.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) + .execute_mask(self.as_ref().len(), ctx) .vortex_expect("Failed to compute validity mask") .to_bit_buffer() .iter() @@ -25,13 +21,10 @@ impl BoolArray { .collect() } - pub fn bool_vec(&self) -> Vec { + pub fn bool_vec(&self, ctx: &mut ExecutionCtx) -> Vec { self.validity() .vortex_expect("failed to get validity") - .execute_mask( - self.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) + .execute_mask(self.as_ref().len(), ctx) .vortex_expect("Failed to compute validity mask") .to_bit_buffer() .iter() diff --git a/vortex-array/src/arrays/bool/vtable/kernel.rs b/vortex-array/src/arrays/bool/vtable/kernel.rs index 2fe2d974a8f..c19d00d5118 100644 --- a/vortex-array/src/arrays/bool/vtable/kernel.rs +++ b/vortex-array/src/arrays/bool/vtable/kernel.rs @@ -1,16 +1,28 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; use crate::arrays::Bool; +use crate::arrays::Dict; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::fns::binary::BooleanExecuteAdaptor; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; +use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(Bool)), - ParentKernelSet::lift(&FillNullExecuteAdaptor(Bool)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Bool)), - ParentKernelSet::lift(&ZipExecuteAdaptor(Bool)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), Bool, BooleanExecuteAdaptor(Bool)); + kernels.register_execute_parent_kernel(Cast.id(), Bool, CastExecuteAdaptor(Bool)); + kernels.register_execute_parent_kernel(FillNull.id(), Bool, FillNullExecuteAdaptor(Bool)); + kernels.register_execute_parent_kernel(Dict.id(), Bool, TakeExecuteAdaptor(Bool)); + kernels.register_execute_parent_kernel(Zip.id(), Bool, ZipExecuteAdaptor(Bool)); +} diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index 1bbc091b010..2420642d972 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -4,7 +4,6 @@ use std::hash::Hash; use std::hash::Hasher; -use kernel::PARENT_KERNELS; use prost::Message; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -24,6 +23,8 @@ use crate::array::child_to_validity; use crate::arrays::bool::BoolData; use crate::arrays::bool::array::SLOT_NAMES; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::BoolBuilder; use crate::dtype::DType; use crate::serde::ArrayChildren; use crate::validity::Validity; @@ -43,6 +44,10 @@ use crate::hash::ArrayHash; /// A [`Bool`]-encoded Vortex array. pub type BoolArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(prost::Message)] pub struct BoolMetadata { // The offset in bits must be <8 @@ -92,6 +97,24 @@ impl VTable for Bool { } } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + let mut data = array.data().clone(); + data.bits = buffers[0].clone(); + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, @@ -170,17 +193,21 @@ impl VTable for Bool { SLOT_NAMES[idx].to_string() } - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ExecutionResult::done(array)) - } - - fn execute_parent( + fn append_to_builder( array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, + builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) + ) -> VortexResult<()> { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + return builder.append_bool_array(&array.into_owned(), ctx); + } + + builder.extend_from_array(array.as_ref()); + Ok(()) + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) } fn reduce_parent( @@ -202,7 +229,8 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; - use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; @@ -210,15 +238,17 @@ mod tests { #[test] fn test_nullable_bool_serde_roundtrip() { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); let array = BoolArray::from_iter([Some(true), None, Some(false), None]); let dtype = array.dtype().clone(); let len = array.len(); - let ctx = ArrayContext::empty(); + let array_ctx = ArrayContext::empty(); let serialized = array .clone() .into_array() - .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default()) + .serialize(&array_ctx, &session, &SerializeOptions::default()) .unwrap(); let mut concat = ByteBufferMut::empty(); @@ -227,14 +257,9 @@ mod tests { } let parts = SerializedArray::try_from(concat.freeze()).unwrap(); let decoded = parts - .decode( - &dtype, - len, - &ReadContext::new(ctx.to_ids()), - &LEGACY_SESSION, - ) + .decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session) .unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut ctx); } } diff --git a/vortex-array/src/arrays/bool/vtable/operations.rs b/vortex-array/src/arrays/bool/vtable/operations.rs index b4ac4682d23..9bae2b9e19c 100644 --- a/vortex-array/src/arrays/bool/vtable/operations.rs +++ b/vortex-array/src/arrays/bool/vtable/operations.rs @@ -30,6 +30,8 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; @@ -46,13 +48,15 @@ mod tests { #[test] fn test_slice() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter([Some(true), Some(true), None, Some(false), None]); #[expect(deprecated)] let sliced_arr = arr.into_array().slice(1..4).unwrap().to_bool(); assert_arrays_eq!( sliced_arr, - BoolArray::from_iter([Some(true), None, Some(false)]) + BoolArray::from_iter([Some(true), None, Some(false)]), + &mut ctx ); } } diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index a0ed3f98d67..54544ba1b4d 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -278,8 +278,8 @@ mod test { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; @@ -291,6 +291,7 @@ mod test { #[test] fn test_rechunk_one_chunk() { + let mut ctx = array_session().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![buffer![0u64].into_array()], DType::Primitive(PType::U64, Nullability::NonNullable), @@ -299,11 +300,12 @@ mod test { let rechunked = chunked.rechunk(1 << 16, 1 << 16).unwrap(); - assert_arrays_eq!(chunked, rechunked); + assert_arrays_eq!(chunked, rechunked, &mut ctx); } #[test] fn test_rechunk_two_chunks() { + let mut ctx = array_session().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![buffer![0u64].into_array(), buffer![5u64].into_array()], DType::Primitive(PType::U64, Nullability::NonNullable), @@ -313,11 +315,12 @@ mod test { let rechunked = chunked.rechunk(1 << 16, 1 << 16).unwrap(); assert_eq!(rechunked.nchunks(), 1); - assert_arrays_eq!(chunked, rechunked); + assert_arrays_eq!(chunked, rechunked, &mut ctx); } #[test] fn test_rechunk_tiny_target_chunks() { + let mut ctx = array_session().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ buffer![0u64, 1, 2, 3].into_array(), @@ -331,11 +334,12 @@ mod test { assert_eq!(rechunked.nchunks(), 2); assert!(rechunked.iter_chunks().all(|c| c.len() < 5)); - assert_arrays_eq!(chunked, rechunked); + assert_arrays_eq!(chunked, rechunked, &mut ctx); } #[test] fn test_rechunk_with_too_big_chunk() { + let mut ctx = array_session().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ buffer![0u64, 1, 2].into_array(), @@ -352,7 +356,7 @@ mod test { // greedy so should be: [0, 1, 2] [42, 42, 42, 42, 42, 42] [4, 5, 6, 7] [8, 9] assert_eq!(rechunked.nchunks(), 4); - assert_arrays_eq!(chunked, rechunked); + assert_arrays_eq!(chunked, rechunked, &mut ctx); } #[test] @@ -369,11 +373,11 @@ mod test { ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?; // Should be all_valid since all non-empty chunks are all_valid - assert!(chunked.all_valid(&mut LEGACY_SESSION.create_execution_ctx())?); + assert!(chunked.all_valid(&mut array_session().create_execution_ctx())?); assert!( !chunked .into_array() - .all_invalid(&mut LEGACY_SESSION.create_execution_ctx())? + .all_invalid(&mut array_session().create_execution_ctx())? ); Ok(()) @@ -384,20 +388,20 @@ mod test { // Create chunks where some are empty but all non-empty chunks have all invalid values let chunks = vec![ PrimitiveArray::new(buffer![1u64, 2], Validity::AllInvalid).into_array(), - PrimitiveArray::new(buffer![0u64; 0], Validity::AllInvalid).into_array(), // empty chunk + PrimitiveArray::new(buffer![0u64; 0], Validity::AllInvalid).into_array(), /* empty chunk */ PrimitiveArray::new(buffer![3u64, 4, 5], Validity::AllInvalid).into_array(), - PrimitiveArray::new(buffer![0u64; 0], Validity::AllInvalid).into_array(), // empty chunk + PrimitiveArray::new(buffer![0u64; 0], Validity::AllInvalid).into_array(), /* empty chunk */ ]; let chunked = ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?; // Should be all_invalid since all non-empty chunks are all_invalid - assert!(!chunked.all_valid(&mut LEGACY_SESSION.create_execution_ctx())?); + assert!(!chunked.all_valid(&mut array_session().create_execution_ctx())?); assert!( chunked .into_array() - .all_invalid(&mut LEGACY_SESSION.create_execution_ctx())? + .all_invalid(&mut array_session().create_execution_ctx())? ); Ok(()) @@ -410,18 +414,18 @@ mod test { PrimitiveArray::new(buffer![1u64, 2], Validity::AllValid).into_array(), PrimitiveArray::new(buffer![0u64; 0], Validity::AllValid).into_array(), // empty chunk PrimitiveArray::new(buffer![3u64, 4], Validity::AllInvalid).into_array(), - PrimitiveArray::new(buffer![0u64; 0], Validity::AllInvalid).into_array(), // empty chunk + PrimitiveArray::new(buffer![0u64; 0], Validity::AllInvalid).into_array(), /* empty chunk */ ]; let chunked = ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?; // Should be neither all_valid nor all_invalid - assert!(!chunked.all_valid(&mut LEGACY_SESSION.create_execution_ctx())?); + assert!(!chunked.all_valid(&mut array_session().create_execution_ctx())?); assert!( !chunked .into_array() - .all_invalid(&mut LEGACY_SESSION.create_execution_ctx())? + .all_invalid(&mut array_session().create_execution_ctx())? ); Ok(()) diff --git a/vortex-array/src/arrays/chunked/compute/aggregate.rs b/vortex-array/src/arrays/chunked/compute/aggregate.rs index 52c3db5cf2f..149c72f499e 100644 --- a/vortex-array/src/arrays/chunked/compute/aggregate.rs +++ b/vortex-array/src/arrays/chunked/compute/aggregate.rs @@ -43,12 +43,12 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; - use crate::aggregate_fn::EmptyOptions; + 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::PrimitiveArray; @@ -58,8 +58,12 @@ mod tests { use crate::scalar::Scalar; fn run_sum(batch: &crate::ArrayRef) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let mut acc = Accumulator::try_new(Sum, EmptyOptions, batch.dtype().clone())?; + let mut ctx = array_session().create_execution_ctx(); + let mut acc = Accumulator::try_new( + Sum, + NumericalAggregateOpts::default(), + batch.dtype().clone(), + )?; acc.accumulate(batch, &mut ctx)?; acc.finish() } diff --git a/vortex-array/src/arrays/chunked/compute/cast.rs b/vortex-array/src/arrays/chunked/compute/cast.rs index 9eed54d4300..80c6848b973 100644 --- a/vortex-array/src/arrays/chunked/compute/cast.rs +++ b/vortex-array/src/arrays/chunked/compute/cast.rs @@ -35,6 +35,8 @@ mod test { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -46,6 +48,7 @@ mod test { #[test] fn test_cast_chunked() { + let mut ctx = array_session().create_execution_ctx(); let arr0 = buffer![0u32, 1].into_array(); let arr1 = buffer![2u32, 3].into_array(); @@ -67,7 +70,7 @@ mod test { let result = root .cast(DType::Primitive(PType::U64, Nullability::NonNullable)) .unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([0u64, 1, 2, 3])); + assert_arrays_eq!(result, PrimitiveArray::from_iter([0u64, 1, 2, 3]), &mut ctx); } #[rstest] diff --git a/vortex-array/src/arrays/chunked/compute/kernel.rs b/vortex-array/src/arrays/chunked/compute/kernel.rs index 502c1fe8337..db0042105cd 100644 --- a/vortex-array/src/arrays/chunked/compute/kernel.rs +++ b/vortex-array/src/arrays/chunked/compute/kernel.rs @@ -1,18 +1,28 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; use crate::arrays::Chunked; +use crate::arrays::Dict; +use crate::arrays::Filter; +use crate::arrays::Slice; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; use crate::arrays::slice::SliceExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::mask::MaskExecuteAdaptor; +use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(crate) static PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&FilterExecuteAdaptor(Chunked)), - ParentKernelSet::lift(&MaskExecuteAdaptor(Chunked)), - ParentKernelSet::lift(&SliceExecuteAdaptor(Chunked)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Chunked)), - ParentKernelSet::lift(&ZipExecuteAdaptor(Chunked)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Filter.id(), Chunked, FilterExecuteAdaptor(Chunked)); + kernels.register_execute_parent_kernel(Mask.id(), Chunked, MaskExecuteAdaptor(Chunked)); + kernels.register_execute_parent_kernel(Slice.id(), Chunked, SliceExecuteAdaptor(Chunked)); + kernels.register_execute_parent_kernel(Dict.id(), Chunked, TakeExecuteAdaptor(Chunked)); + kernels.register_execute_parent_kernel(Zip.id(), Chunked, ZipExecuteAdaptor(Chunked)); +} diff --git a/vortex-array/src/arrays/chunked/compute/mod.rs b/vortex-array/src/arrays/chunked/compute/mod.rs index bbc9029657a..f7c9a4e9d61 100644 --- a/vortex-array/src/arrays/chunked/compute/mod.rs +++ b/vortex-array/src/arrays/chunked/compute/mod.rs @@ -18,6 +18,8 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::binary_numeric::test_binary_numeric_array; @@ -79,7 +81,10 @@ mod tests { ).unwrap())] fn test_chunked_consistency(#[case] array: ChunkedArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[rstest] @@ -154,6 +159,9 @@ mod tests { DType::Primitive(PType::I32, Nullability::NonNullable), ).unwrap())] fn test_chunked_binary_numeric(#[case] array: ChunkedArray) { - test_binary_numeric_array(array.into_array()) + test_binary_numeric_array( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ) } } diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 7024c2b3875..34f443f1ce1 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -128,6 +128,8 @@ mod test { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; @@ -141,6 +143,7 @@ mod test { #[test] fn test_take() { + let mut ctx = array_session().create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -149,11 +152,12 @@ mod test { let indices = buffer![0u64, 0, 6, 4].into_array(); let result = arr.take(indices).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 1, 1, 2])); + assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 1, 1, 2]), &mut ctx); } #[test] fn test_take_nullable_values() { + let mut ctx = array_session().create_execution_ctx(); let a = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -164,12 +168,14 @@ mod test { let result = arr.take(indices.into_array()).unwrap(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([1i32, 1, 1, 2].map(Some)) + PrimitiveArray::from_option_iter([1i32, 1, 1, 2].map(Some)), + &mut ctx ); } #[test] fn test_take_nullable_indices() { + let mut ctx = array_session().create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -183,12 +189,14 @@ mod test { let result = arr.take(indices.into_array()).unwrap(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(1i32), None, None, Some(2)]) + PrimitiveArray::from_option_iter([Some(1i32), None, None, Some(2)]), + &mut ctx ); } #[test] fn test_take_nullable_struct() { + let mut ctx = array_session().create_execution_ctx(); let struct_array = StructArray::try_new(FieldNames::default(), vec![], 100, Validity::NonNullable) .unwrap(); @@ -209,11 +217,12 @@ mod test { Validity::Array(BoolArray::from_iter(vec![true, false, true]).into_array()), ) .unwrap(); - assert_arrays_eq!(result, expect); + assert_arrays_eq!(result, expect, &mut ctx); } #[test] fn test_empty_take() { + let mut ctx = array_session().create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -227,12 +236,14 @@ mod test { assert_eq!(result.dtype(), arr.dtype()); assert_arrays_eq!( result, - PrimitiveArray::empty::(Nullability::NonNullable) + PrimitiveArray::empty::(Nullability::NonNullable), + &mut ctx ); } #[test] fn test_take_shuffled_indices() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let c0 = buffer![0i32, 1, 2].into_array(); let c1 = buffer![3i32, 4, 5].into_array(); let c2 = buffer![6i32, 7, 8].into_array(); @@ -249,7 +260,8 @@ mod test { assert_arrays_eq!( result, - PrimitiveArray::from_iter([8i32, 0, 5, 3, 2, 7, 1, 6, 4]) + PrimitiveArray::from_iter([8i32, 0, 5, 3, 2, 7, 1, 6, 4]), + &mut ctx ); Ok(()) } @@ -300,6 +312,7 @@ mod test { #[test] fn test_take_null_indices() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let c0 = buffer![10i32, 20, 30].into_array(); let c1 = buffer![40i32, 50, 60].into_array(); let arr = ChunkedArray::try_new( @@ -323,7 +336,8 @@ mod test { Some(40), None, Some(30) - ]) + ]), + &mut ctx ); Ok(()) } diff --git a/vortex-array/src/arrays/chunked/compute/zip.rs b/vortex-array/src/arrays/chunked/compute/zip.rs index 1d7873888a3..4037fd1dc49 100644 --- a/vortex-array/src/arrays/chunked/compute/zip.rs +++ b/vortex-array/src/arrays/chunked/compute/zip.rs @@ -51,10 +51,10 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::chunked::ChunkedArrayExt; @@ -94,7 +94,7 @@ mod tests { // One step of execution will push down the zip. let zipped = zipped .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut array_session().create_execution_ctx()) .unwrap(); let zipped = zipped .as_opt::() diff --git a/vortex-array/src/arrays/chunked/mod.rs b/vortex-array/src/arrays/chunked/mod.rs index 766d9f6575c..4bd6cb43b7e 100644 --- a/vortex-array/src/arrays/chunked/mod.rs +++ b/vortex-array/src/arrays/chunked/mod.rs @@ -12,5 +12,9 @@ pub(crate) mod paired_chunks; mod vtable; pub use vtable::Chunked; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + compute::kernel::initialize(session); +} + #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index 00d113f273a..578b0a7c068 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -6,12 +6,12 @@ use std::sync::LazyLock; use vortex_buffer::Buffer; use vortex_buffer::buffer; +use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::accessor::ArrayAccessor; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::ListArray; @@ -30,11 +30,9 @@ use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::PType::I32; use crate::executor::execute_into_builder; -use crate::session::ArraySession; use crate::validity::Validity; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(crate::array_session); fn chunked_array() -> ChunkedArray { ChunkedArray::try_new( @@ -65,7 +63,8 @@ fn builder_kernel_path_canonicalizes_primitive_chunks() { assert_arrays_eq!( output, - PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6, 7, 8, 9]) + PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6, 7, 8, 9]), + &mut ctx ); } @@ -98,7 +97,11 @@ fn builder_kernel_nested_chunked_of_chunked() { let mut builder = execute_into_builder(outer, builder, &mut ctx).unwrap(); let output = builder.finish(); - assert_arrays_eq!(output, PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6])); + assert_arrays_eq!( + output, + PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6]), + &mut ctx + ); } #[test] @@ -130,8 +133,8 @@ fn builder_kernel_path_repeated_shared_chunked_dict_execution() { drop(keep_alive); - assert_arrays_eq!(first, expected); - assert_arrays_eq!(second, expected); + assert_arrays_eq!(first, expected, &mut ctx); + assert_arrays_eq!(second, expected, &mut ctx); } #[test] @@ -156,8 +159,8 @@ fn execute_path_repeated_shared_chunked_dict_execution() { drop(keep_alive); - assert_arrays_eq!(first, expected); - assert_arrays_eq!(second, expected); + assert_arrays_eq!(first, expected, &mut ctx); + assert_arrays_eq!(second, expected, &mut ctx); } #[test] @@ -194,85 +197,106 @@ fn execute_path_nested_chunked_dict_of_dict_into_canonical() { drop(keep_alive); - assert_arrays_eq!(first, expected); - assert_arrays_eq!(second, expected); + assert_arrays_eq!(first, expected, &mut ctx); + assert_arrays_eq!(second, expected, &mut ctx); } #[test] fn with_slot_rewrites_chunk_and_offsets() { + let mut ctx = SESSION.create_execution_ctx(); let array = chunked_array().into_array(); - let replacement = buffer![10u64, 11, 12].into_array(); - let array = array.with_slot(1, replacement).unwrap(); + let replacement = buffer![1u64, 2, 3].into_array(); + // SAFETY: the replacement chunk has the same logical values as the original chunk; only the + // physical child handle changes. + let array = unsafe { array.with_slot(1, replacement) }.unwrap(); let array = array.as_::(); assert_eq!(array.nchunks(), 3); assert_eq!(array.chunk_offsets(), [0, 3, 6, 9]); assert_arrays_eq!( array.chunk(0).clone(), - PrimitiveArray::from_iter([10u64, 11, 12]) + PrimitiveArray::from_iter([1u64, 2, 3]), + &mut ctx ); assert_arrays_eq!( array.array().clone(), - PrimitiveArray::from_iter([10u64, 11, 12, 4, 5, 6, 7, 8, 9]) + PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6, 7, 8, 9]), + &mut ctx ); } #[test] fn with_slot_rejects_len_mismatch() { - let err = chunked_array() - .into_array() - .with_slot(1, buffer![10u64, 11].into_array()) - .unwrap_err(); + // SAFETY: this call is expected to fail the checked slot length invariant before any rewritten + // array is returned or observed. + let err = unsafe { + chunked_array() + .into_array() + .with_slot(1, buffer![10u64, 11].into_array()) + } + .unwrap_err(); assert!(err.to_string().contains("physical rewrite")); } #[test] fn slice_middle() { + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(2..5).unwrap(), - PrimitiveArray::from_iter([3u64, 4, 5]) + PrimitiveArray::from_iter([3u64, 4, 5]), + &mut ctx ); } #[test] fn slice_begin() { + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(1..3).unwrap(), - PrimitiveArray::from_iter([2u64, 3]) + PrimitiveArray::from_iter([2u64, 3]), + &mut ctx ); } #[test] fn slice_aligned() { + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(3..6).unwrap(), - PrimitiveArray::from_iter([4u64, 5, 6]) + PrimitiveArray::from_iter([4u64, 5, 6]), + &mut ctx ); } #[test] fn slice_many_aligned() { + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(0..6).unwrap(), - PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6]) + PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6]), + &mut ctx ); } #[test] fn slice_end() { + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(7..8).unwrap(), - PrimitiveArray::from_iter([8u64]) + PrimitiveArray::from_iter([8u64]), + &mut ctx ); } #[test] fn slice_exactly_end() { + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(6..9).unwrap(), - PrimitiveArray::from_iter([7u64, 8, 9]) + PrimitiveArray::from_iter([7u64, 8, 9]), + &mut ctx ); } @@ -286,6 +310,7 @@ fn slice_empty() { #[test] fn scalar_at_empty_children_both_sides() { + let mut ctx = SESSION.create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -297,11 +322,12 @@ fn scalar_at_empty_children_both_sides() { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2])); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2]), &mut ctx); } #[test] fn scalar_at_empty_children_trailing() { + let mut ctx = SESSION.create_execution_ctx(); let array = ChunkedArray::try_new( vec![ buffer![1u64, 2].into_array(), @@ -312,11 +338,12 @@ fn scalar_at_empty_children_trailing() { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4])); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]), &mut ctx); } #[test] fn scalar_at_empty_children_leading() { + let mut ctx = SESSION.create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -327,28 +354,26 @@ fn scalar_at_empty_children_leading() { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4])); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]), &mut ctx); } #[test] -pub fn pack_nested_structs() { +pub fn pack_nested_structs() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let struct_array = StructArray::try_new( ["a"].into(), vec![VarBinViewArray::from_iter_str(["foo", "bar", "baz", "quak"]).into_array()], 4, Validity::NonNullable, - ) - .unwrap(); + )?; let dtype = struct_array.dtype().clone(); let chunked = ChunkedArray::try_new( vec![ - ChunkedArray::try_new(vec![struct_array.clone().into_array()], dtype.clone()) - .unwrap() + ChunkedArray::try_new(vec![struct_array.clone().into_array()], dtype.clone())? .into_array(), ], dtype, - ) - .unwrap() + )? .into_array(); #[expect(deprecated)] let canonical_struct = chunked.to_struct(); @@ -356,11 +381,28 @@ pub fn pack_nested_structs() { let canonical_varbin = canonical_struct.unmasked_fields()[0].to_varbinview(); #[expect(deprecated)] let original_varbin = struct_array.unmasked_fields()[0].to_varbinview(); - let orig_values = - original_varbin.with_iterator(|it| it.map(|a| a.map(|v| v.to_vec())).collect::>()); - let canon_values = - canonical_varbin.with_iterator(|it| it.map(|a| a.map(|v| v.to_vec())).collect::>()); + let orig_mask = original_varbin + .validity()? + .execute_mask(original_varbin.len(), &mut ctx)?; + let orig_values = (0..original_varbin.len()) + .map(|i| { + orig_mask + .value(i) + .then(|| original_varbin.bytes_at(i).to_vec()) + }) + .collect::>(); + let canon_mask = canonical_varbin + .validity()? + .execute_mask(canonical_varbin.len(), &mut ctx)?; + let canon_values = (0..canonical_varbin.len()) + .map(|i| { + canon_mask + .value(i) + .then(|| canonical_varbin.bytes_at(i).to_vec()) + }) + .collect::>(); assert_eq!(orig_values, canon_values); + Ok(()) } #[test] diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index adbe447604b..995d90bc0aa 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -297,10 +297,8 @@ mod tests { use crate::ArrayRef; use crate::Canonical; - use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::accessor::ArrayAccessor; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::FixedSizeListArray; @@ -323,12 +321,10 @@ mod tests { use crate::memory::MemorySessionExt; use crate::memory::WritableHostBuffer; use crate::scalar::Scalar; - use crate::session::ArraySession; use crate::validity::Validity; /// A shared session for these chunked-array tests, used to create execution contexts. - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[derive(Debug)] struct CountingAllocator { @@ -439,7 +435,11 @@ mod tests { assert_variant_values(&variant, &[10, 20, 30])?; let shredded = shredded.clone().execute::(&mut ctx)?; - assert_arrays_eq!(shredded, PrimitiveArray::from_iter([10i32, 20, 30])); + assert_arrays_eq!( + shredded, + PrimitiveArray::from_iter([10i32, 20, 30]), + &mut ctx + ); Ok(()) } @@ -549,10 +549,30 @@ mod tests { .clone() .execute::(&mut ctx) .unwrap(); - let orig_values = original_varbin - .with_iterator(|it| it.map(|a| a.map(|v| v.to_vec())).collect::>()); - let canon_values = canonical_varbin - .with_iterator(|it| it.map(|a| a.map(|v| v.to_vec())).collect::>()); + let orig_mask = original_varbin + .validity() + .unwrap() + .execute_mask(original_varbin.len(), &mut ctx) + .unwrap(); + let orig_values = (0..original_varbin.len()) + .map(|i| { + orig_mask + .value(i) + .then(|| original_varbin.bytes_at(i).to_vec()) + }) + .collect::>(); + let canon_mask = canonical_varbin + .validity() + .unwrap() + .execute_mask(canonical_varbin.len(), &mut ctx) + .unwrap(); + let canon_values = (0..canonical_varbin.len()) + .map(|i| { + canon_mask + .value(i) + .then(|| canonical_varbin.bytes_at(i).to_vec()) + }) + .collect::>(); assert_eq!(orig_values, canon_values); } @@ -640,13 +660,10 @@ mod tests { #[test] fn list_canonicalize_uses_memory_session_allocator() { let allocations = Arc::new(AtomicUsize::new(0)); - let session = VortexSession::empty(); - session - .memory_mut() - .set_allocator(Arc::new(CountingAllocator { - allocations: Arc::clone(&allocations), - })); - let mut ctx = ExecutionCtx::new(session); + let session = crate::array_session().with_allocator(Arc::new(CountingAllocator { + allocations: Arc::clone(&allocations), + })); + let mut ctx = session.create_execution_ctx(); let l1 = ListArray::try_new( buffer![1, 2, 3, 4].into_array(), diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index a70ba6869ab..673855cfc0a 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -29,11 +29,11 @@ use crate::array::ArrayId; use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::VTable; +use crate::array::with_empty_buffers; 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::compute::kernel::PARENT_KERNELS; use crate::arrays::chunked::compute::rules::PARENT_RULES; use crate::arrays::chunked::vtable::canonical::_canonicalize; use crate::buffer::BufferHandle; @@ -157,6 +157,14 @@ impl VTable for Chunked { vortex_panic!("ChunkedArray 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, @@ -263,15 +271,6 @@ impl VTable for Chunked { } } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { Ok(match array.nchunks() { 0 => Some(Canonical::empty(array.dtype()).into_array()), diff --git a/vortex-array/src/arrays/chunked/vtable/operations.rs b/vortex-array/src/arrays/chunked/vtable/operations.rs index ce44b8231f8..8f9e0867a88 100644 --- a/vortex-array/src/arrays/chunked/vtable/operations.rs +++ b/vortex-array/src/arrays/chunked/vtable/operations.rs @@ -30,6 +30,8 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -57,9 +59,11 @@ mod tests { #[case::end(7..8, &[8u64])] #[case::exactly_end(6..9, &[7u64, 8, 9])] fn slice(#[case] range: Range, #[case] expected: &[u64]) { + let mut ctx = array_session().create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(range).unwrap(), - PrimitiveArray::from_iter(expected.iter().copied()) + PrimitiveArray::from_iter(expected.iter().copied()), + &mut ctx ); } @@ -73,6 +77,7 @@ mod tests { #[test] fn scalar_at_empty_children_both_sides() { + let mut ctx = array_session().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -84,11 +89,12 @@ mod tests { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2])); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2]), &mut ctx); } #[test] fn scalar_at_empty_children_trailing() { + let mut ctx = array_session().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ buffer![1u64, 2].into_array(), @@ -99,11 +105,12 @@ mod tests { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4])); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]), &mut ctx); } #[test] fn scalar_at_empty_children_leading() { + let mut ctx = array_session().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -114,6 +121,6 @@ mod tests { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4])); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]), &mut ctx); } } diff --git a/vortex-array/src/arrays/constant/compute/cast.rs b/vortex-array/src/arrays/constant/compute/cast.rs index 2fd9b3ded32..db3a41829f0 100644 --- a/vortex-array/src/arrays/constant/compute/cast.rs +++ b/vortex-array/src/arrays/constant/compute/cast.rs @@ -25,8 +25,15 @@ mod tests { use rstest::rstest; use crate::IntoArray; + use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::ConstantArray; + use crate::builtins::ArrayBuiltins; use crate::compute::conformance::cast::test_cast_conformance; + use crate::dtype::DType; + use crate::dtype::DecimalDType; + use crate::dtype::Nullability; + use crate::scalar::DecimalValue; use crate::scalar::Scalar; #[rstest] @@ -39,4 +46,22 @@ mod tests { fn test_cast_constant_conformance(#[case] array: crate::ArrayRef) { test_cast_conformance(&array); } + + #[test] + fn test_cast_constant_i64_to_decimal() { + let target_dtype = DType::Decimal(DecimalDType::new(21, 2), Nullability::NonNullable); + let casted = ConstantArray::new(Scalar::from(42i64), 5) + .into_array() + .cast(target_dtype.clone()) + .unwrap(); + + assert_eq!(casted.dtype(), &target_dtype); + let scalar = casted + .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .unwrap(); + assert_eq!( + scalar.as_decimal().decimal_value(), + Some(DecimalValue::I128(4200)) + ); + } } diff --git a/vortex-array/src/arrays/constant/compute/fill_null.rs b/vortex-array/src/arrays/constant/compute/fill_null.rs index f69a23c7ab0..1c9c2aaf903 100644 --- a/vortex-array/src/arrays/constant/compute/fill_null.rs +++ b/vortex-array/src/arrays/constant/compute/fill_null.rs @@ -22,8 +22,8 @@ impl FillNullReduce for Constant { #[cfg(test)] mod test { use crate::IntoArray as _; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ConstantArray; use crate::arrow::ArrowSessionExt; use crate::builtins::ArrayBuiltins; @@ -34,7 +34,7 @@ mod test { #[test] fn test_null() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let actual = ConstantArray::new(Scalar::null_native::(), 3) .into_array() .fill_null(Scalar::from(1)) @@ -43,11 +43,11 @@ mod test { assert!(!actual.dtype().is_nullable()); - let actual_arrow = LEGACY_SESSION + let actual_arrow = array_session() .arrow() .execute_arrow(actual.clone(), None, &mut ctx) .unwrap(); - let expected_arrow = LEGACY_SESSION + let expected_arrow = array_session() .arrow() .execute_arrow(expected.clone(), None, &mut ctx) .unwrap(); @@ -62,7 +62,7 @@ mod test { #[test] fn test_non_null() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let actual = ConstantArray::new(Scalar::from(Some(1)), 3) .into_array() .fill_null(Scalar::from(1)) @@ -71,11 +71,11 @@ mod test { assert!(!actual.dtype().is_nullable()); - let actual_arrow = LEGACY_SESSION + let actual_arrow = array_session() .arrow() .execute_arrow(actual.clone(), None, &mut ctx) .unwrap(); - let expected_arrow = LEGACY_SESSION + let expected_arrow = array_session() .arrow() .execute_arrow(expected.clone(), None, &mut ctx) .unwrap(); @@ -90,7 +90,7 @@ mod test { #[test] fn test_non_nullable_with_nullable() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let actual = ConstantArray::new(Scalar::from(1), 3) .into_array() .fill_null(Scalar::new( @@ -104,11 +104,11 @@ mod test { assert!(actual.dtype().is_nullable()); - let actual_arrow = LEGACY_SESSION + let actual_arrow = array_session() .arrow() .execute_arrow(actual.clone(), None, &mut ctx) .unwrap(); - let expected_arrow = LEGACY_SESSION + let expected_arrow = array_session() .arrow() .execute_arrow(expected.clone(), None, &mut ctx) .unwrap(); diff --git a/vortex-array/src/arrays/constant/compute/mod.rs b/vortex-array/src/arrays/constant/compute/mod.rs index bc9f25466f6..80fab02dc88 100644 --- a/vortex-array/src/arrays/constant/compute/mod.rs +++ b/vortex-array/src/arrays/constant/compute/mod.rs @@ -18,6 +18,8 @@ mod test { use rstest::rstest; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ConstantArray; use crate::compute::conformance::consistency::test_array_consistency; use crate::compute::conformance::filter::test_filter_conformance; @@ -62,6 +64,9 @@ mod test { #[case::constant_single(ConstantArray::new(Scalar::from(99u64), 1))] #[case::constant_large(ConstantArray::new(Scalar::from("hello"), 1000))] fn test_constant_consistency(#[case] array: ConstantArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/constant/compute/take.rs b/vortex-array/src/arrays/constant/compute/take.rs index eb3d7918207..9fb1373a5cf 100644 --- a/vortex-array/src/arrays/constant/compute/take.rs +++ b/vortex-array/src/arrays/constant/compute/take.rs @@ -73,10 +73,10 @@ mod tests { use vortex_mask::AllOr; use crate::IntoArray; - use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -87,6 +87,7 @@ mod tests { #[test] fn take_nullable_indices() { + let mut ctx = array_session().create_execution_ctx(); let array = ConstantArray::new(42, 10).into_array(); let taken = array .take( @@ -108,13 +109,14 @@ mod tests { PrimitiveArray::new( buffer![42i32, 42, 42], Validity::from_iter([false, true, false]) - ) + ), + &mut ctx ); assert_eq!( taken .validity() .unwrap() - .execute_mask(taken.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(taken.len(), &mut array_session().create_execution_ctx()) .unwrap() .indices(), AllOr::Some(valid_indices) @@ -123,6 +125,7 @@ mod tests { #[test] fn take_all_valid_indices() { + let mut ctx = array_session().create_execution_ctx(); let array = ConstantArray::new(42, 10).into_array(); let taken = array .take(PrimitiveArray::new(buffer![0, 5, 7], Validity::AllValid).into_array()) @@ -134,13 +137,14 @@ mod tests { assert_arrays_eq!( #[expect(deprecated)] taken.to_primitive(), - PrimitiveArray::new(buffer![42i32, 42, 42], Validity::AllValid) + PrimitiveArray::new(buffer![42i32, 42, 42], Validity::AllValid), + &mut ctx ); assert_eq!( taken .validity() .unwrap() - .execute_mask(taken.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(taken.len(), &mut array_session().create_execution_ctx()) .unwrap() .indices(), AllOr::All diff --git a/vortex-array/src/arrays/constant/compute/uncompressed_size.rs b/vortex-array/src/arrays/constant/compute/uncompressed_size.rs index b9654d5840c..f8a4d4f92a4 100644 --- a/vortex-array/src/arrays/constant/compute/uncompressed_size.rs +++ b/vortex-array/src/arrays/constant/compute/uncompressed_size.rs @@ -10,13 +10,13 @@ mod tests { use crate::ArrayRef; use crate::Canonical; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::canonical_uncompressed_size_in_bytes; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -27,7 +27,7 @@ mod tests { use crate::validity::Validity; fn aggregate(array: &ArrayRef) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut acc = Accumulator::try_new(UncompressedSizeInBytes, EmptyOptions, array.dtype().clone())?; acc.accumulate(array, &mut ctx)?; @@ -82,7 +82,7 @@ mod tests { #[test] fn constant_utf8_matches_canonical_size() -> VortexResult<()> { let constant = ConstantArray::new("abcdefghijkl".to_string(), 10).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let canonical = constant.clone().execute::(&mut ctx)?; let expected = canonical_uncompressed_size_in_bytes(&canonical, &mut ctx)?; diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index dd891778949..e160bcc16a9 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -361,12 +361,10 @@ mod tests { use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; use crate::scalar::Scalar; - use crate::session::ArraySession; use crate::validity::Validity; /// A shared session for these constant-array tests, used to create execution contexts. - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[test] fn test_canonicalize_null() { @@ -386,10 +384,11 @@ mod tests { #[test] fn test_canonicalize_const_str() { + let mut ctx = SESSION.create_execution_ctx(); let const_array = ConstantArray::new("four".to_string(), 4); let expected = VarBinArray::from(vec!["four", "four", "four", "four"]); - assert_arrays_eq!(const_array, expected); + assert_arrays_eq!(const_array, expected, &mut ctx); } #[test] @@ -452,21 +451,24 @@ mod tests { .elements() .clone() .execute::(&mut ctx)?, - PrimitiveArray::from_iter([1u64, 2, 1, 2]) + PrimitiveArray::from_iter([1u64, 2, 1, 2]), + &mut ctx ); assert_arrays_eq!( list_array .offsets() .clone() .execute::(&mut ctx)?, - PrimitiveArray::from_iter([0u64, 2]) + PrimitiveArray::from_iter([0u64, 2]), + &mut ctx ); assert_arrays_eq!( list_array .sizes() .clone() .execute::(&mut ctx)?, - PrimitiveArray::from_iter([2u64, 2]) + PrimitiveArray::from_iter([2u64, 2]), + &mut ctx ); Ok(()) } @@ -493,7 +495,8 @@ mod tests { .clone() .execute::(&mut ctx) .unwrap(), - PrimitiveArray::from_iter([0u64, 0]) + PrimitiveArray::from_iter([0u64, 0]), + &mut ctx ); assert_arrays_eq!( canonical_const @@ -501,7 +504,8 @@ mod tests { .clone() .execute::(&mut ctx) .unwrap(), - PrimitiveArray::from_iter([0u64, 0]) + PrimitiveArray::from_iter([0u64, 0]), + &mut ctx ); } @@ -526,7 +530,8 @@ mod tests { .clone() .execute::(&mut ctx) .unwrap(), - PrimitiveArray::from_iter([0u64, 0]) + PrimitiveArray::from_iter([0u64, 0]), + &mut ctx ); assert_arrays_eq!( canonical_const @@ -534,7 +539,8 @@ mod tests { .clone() .execute::(&mut ctx) .unwrap(), - PrimitiveArray::from_iter([0u64, 0]) + PrimitiveArray::from_iter([0u64, 0]), + &mut ctx ); } @@ -595,7 +601,11 @@ mod tests { for i in 0..4 { let list = canonical.fixed_size_list_elements_at(i).unwrap(); let list_primitive = list.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(list_primitive, PrimitiveArray::from_iter([10i32, 20, 30])); + assert_arrays_eq!( + list_primitive, + PrimitiveArray::from_iter([10i32, 20, 30]), + &mut ctx + ); } } @@ -627,7 +637,8 @@ mod tests { .unwrap(); assert_arrays_eq!( elements, - PrimitiveArray::from_iter([1.5f64, 2.5, 1.5, 2.5, 1.5, 2.5]) + PrimitiveArray::from_iter([1.5f64, 2.5, 1.5, 2.5, 1.5, 2.5]), + &mut ctx ); } @@ -740,7 +751,7 @@ mod tests { .clone() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(elements, PrimitiveArray::from_iter([42i16])); + assert_arrays_eq!(elements, PrimitiveArray::from_iter([42i16]), &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 89c571eb827..d3df3a14b36 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -15,6 +15,7 @@ use vortex_session::registry::CachedId; use crate::ArrayEq; use crate::ArrayHash; +use crate::ArrayParts; use crate::ArrayRef; use crate::EqMode; use crate::ExecutionCtx; @@ -24,6 +25,7 @@ use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; +use crate::array::unsupported_buffer_replacement; use crate::arrays::constant::ConstantData; use crate::arrays::constant::compute::rules::PARENT_RULES; use crate::arrays::constant::vtable::canonical::constant_canonicalize; @@ -109,6 +111,14 @@ impl VTable for Constant { } } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + unsupported_buffer_replacement(array, buffers) + } + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { vortex_panic!("ConstantArray slot_name index {idx} out of bounds") } @@ -131,7 +141,7 @@ impl VTable for Constant { buffers: &[BufferHandle], _children: &dyn ArrayChildren, session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult> { vortex_ensure!( buffers.len() == 1, "Expected 1 buffer, got {}", @@ -144,7 +154,7 @@ impl VTable for Constant { let scalar_value = ScalarValue::from_proto_bytes(bytes, dtype, session)?; let scalar = Scalar::try_new(dtype.clone(), scalar_value)?; - Ok(crate::array::ArrayParts::new( + Ok(ArrayParts::new( self.clone(), dtype.clone(), len, @@ -270,7 +280,6 @@ fn append_value_or_nulls( mod tests { use rstest::rstest; use vortex_error::VortexResult; - use vortex_session::VortexSession; use crate::IntoArray; use crate::VortexSessionExecute; @@ -286,7 +295,7 @@ mod tests { /// Appends `array` into a fresh builder and asserts the result matches `constant_canonicalize`. fn assert_append_matches_canonical(array: ConstantArray) -> VortexResult<()> { - let mut ctx = VortexSession::empty().create_execution_ctx(); + 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()); @@ -294,7 +303,7 @@ mod tests { .into_array() .append_to_builder(builder.as_mut(), &mut ctx)?; let result = builder.finish(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } @@ -303,6 +312,23 @@ mod tests { assert_append_matches_canonical(ConstantArray::new(Scalar::null(DType::Null), 5)) } + #[test] + fn test_with_buffers_rejects_serialized_scalar_buffer() { + let array = + ConstantArray::new(Scalar::primitive(42i32, Nullability::NonNullable), 3).into_array(); + let buffers = array.buffer_handles(); + + // SAFETY: the replacement buffers are the array's existing buffers, so the logical values + // would be unchanged if the encoding supported buffer replacement. + let Err(err) = (unsafe { array.with_buffers(buffers) }) else { + panic!("ConstantArray should reject replacing its serialized scalar buffer"); + }; + assert!( + err.to_string() + .contains("does not support in-memory buffer replacement") + ); + } + #[rstest] #[case::bool_true(true, 5)] #[case::bool_false(false, 3)] diff --git a/vortex-array/src/arrays/datetime/test.rs b/vortex-array/src/arrays/datetime/test.rs index 511b8196a79..1ae114874d9 100644 --- a/vortex-array/src/arrays/datetime/test.rs +++ b/vortex-array/src/arrays/datetime/test.rs @@ -1,5 +1,5 @@ // SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributorsuse crate::dtype::Nullability; +// SPDX-FileCopyrightText: Copyright the Vortex contributors use rstest::rstest; use vortex_buffer::buffer; @@ -10,6 +10,8 @@ use crate::EqMode; use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; +use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::datetime::TemporalData; use crate::assert_arrays_eq; @@ -26,12 +28,14 @@ use crate::validity::Validity; macro_rules! test_temporal_roundtrip { ($prim:ty, $constructor:expr, $unit:expr) => {{ + let mut ctx = array_session().create_execution_ctx(); let array = buffer![100 as $prim].into_array(); let temporal: TemporalData = $constructor(array, $unit); assert_arrays_eq!( temporal.temporal_values(), - PrimitiveArray::from_iter([100 as $prim]) + PrimitiveArray::from_iter([100 as $prim]), + &mut ctx ); assert_eq!(temporal.temporal_metadata().time_unit(), $unit); }}; @@ -145,6 +149,7 @@ test_fail_case!( // We test Timestamp explicitly to avoid the macro getting too complex. #[test] fn test_timestamp() { + let mut ctx = array_session().create_execution_ctx(); let ts = buffer![100i64].into_array(); let ts_array = ts.into_array(); @@ -159,7 +164,8 @@ fn test_timestamp() { assert_arrays_eq!( temporal_array.temporal_values(), - PrimitiveArray::from_iter([100i64]) + PrimitiveArray::from_iter([100i64]), + &mut ctx ); assert_eq!( temporal_array.temporal_metadata(), diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index 432313d3cb6..80fafde5059 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -1,11 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use num_traits::CheckedMul; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_compute::lane_kernels::IndexedSourceExt; +use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -13,12 +19,17 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Decimal; use crate::arrays::DecimalArray; +use crate::dtype::BigCast; use crate::dtype::DType; +use crate::dtype::DecimalDType; use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; +use crate::dtype::i256; use crate::match_each_decimal_value_type; +use crate::scalar::DecimalValue; use crate::scalar_fn::fns::cast::CastKernel; use crate::scalar_fn::fns::cast::CastReduce; +use crate::validity::Validity; impl CastReduce for Decimal { fn cast(array: ArrayView<'_, Decimal>, dtype: &DType) -> VortexResult> { @@ -77,57 +88,246 @@ impl CastKernel for Decimal { ); }; - // Scale changes are not yet supported - if from_decimal_dtype.scale() != to_decimal_dtype.scale() { - vortex_bail!( - "Casting decimal with scale {} to scale {} not yet implemented", - from_decimal_dtype.scale(), - to_decimal_dtype.scale() - ); - } - - // Downcasting precision is not yet supported - if to_decimal_dtype.precision() < from_decimal_dtype.precision() { - vortex_bail!( - "Downcasting decimal from precision {} to {} not yet implemented", - from_decimal_dtype.precision(), - to_decimal_dtype.precision() - ); - } - // If the dtype is exactly the same, return self if array.dtype() == dtype { return Ok(Some(array.array().clone())); } + let validity = array.validity()?; + // Cast the validity to the new nullability - let new_validity = array - .validity()? + let new_validity = validity + .clone() .cast_nullability(*to_nullability, array.len(), ctx)?; - // If the target needs a wider physical type, upcast the values + // Reuse the values buffer untouched when no rescale is required, the target precision + // only widens (so every value still fits), and the current physical type is already wide + // enough to hold the target precision. This keeps the common precision-widening cast + // (and pure nullability changes) zero-copy instead of allocating and re-scanning. + if from_decimal_dtype.scale() == to_decimal_dtype.scale() + && to_decimal_dtype.precision() >= from_decimal_dtype.precision() + && array + .values_type() + .is_compatible_decimal_value_type(*to_decimal_dtype) + { + // SAFETY: the source values are bit-identical and remain in range for the wider + // precision, and new_validity has the same length, only its nullability tag changes. + unsafe { + return Ok(Some( + DecimalArray::new_unchecked_handle( + array.buffer_handle().clone(), + array.values_type(), + *to_decimal_dtype, + new_validity, + ) + .into_array(), + )); + } + } + + let valid_values = validity.execute_mask(array.len(), ctx)?; let target_values_type = DecimalType::smallest_decimal_value_type(to_decimal_dtype); - let array = if target_values_type > array.values_type() { - upcast_decimal_values(array, target_values_type)? - } else { - array.array().as_::().into_owned() - }; - // SAFETY: new_validity same length as previous validity, just cast - unsafe { - Ok(Some( - DecimalArray::new_unchecked_handle( - array.buffer_handle().clone(), - array.values_type(), + match_each_decimal_value_type!(array.values_type(), |F| { + match_each_decimal_value_type!(target_values_type, |T| { + cast_decimal_values::( + array, + *from_decimal_dtype, *to_decimal_dtype, new_validity, + &valid_values, ) - .into_array(), - )) + .map(Some) + }) + }) + } +} + +fn cast_decimal_values( + array: ArrayView<'_, Decimal>, + from_decimal_dtype: DecimalDType, + to_decimal_dtype: DecimalDType, + validity: Validity, + valid_values: &Mask, +) -> VortexResult +where + F: NativeDecimalType, + T: NativeDecimalType + CheckedMul, + DecimalValue: From, +{ + let values = array.buffer::(); + let values = values.as_slice(); + let cast_plan = DecimalCastPlan::::new(from_decimal_dtype, to_decimal_dtype); + + let buffer = 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() + } + Mask::AllFalse(_) => 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() + } + }; + + Ok(DecimalArray::new(buffer, to_decimal_dtype, validity).into_array()) +} + +#[cold] +fn decimal_cast_error( + value: F, + from_decimal_dtype: DecimalDType, + to_decimal_dtype: DecimalDType, +) -> VortexError +where + F: NativeDecimalType, + T: NativeDecimalType, + DecimalValue: From, +{ + match DecimalValue::from(value) + .cast_decimal(from_decimal_dtype, to_decimal_dtype) + .and_then(|value| { + value.cast::().ok_or_else(|| { + vortex_err!( + "decimal value cannot be represented as {} after casting to {}", + T::DECIMAL_TYPE, + to_decimal_dtype + ) + }) + }) { + Ok(_) => { + // The fast path only returns `None` for values the slow path also rejects, so this + // arm should be unreachable. If it is hit, the fast and slow paths have drifted and + // we are erroring on a value that is actually representable. + debug_assert!( + false, + "decimal fast-path cast rejected value {value} that the slow path accepts \ + (from {from_decimal_dtype} to {to_decimal_dtype})" + ); + vortex_err!( + "decimal value cannot be represented as {} after casting from {} to {}", + T::DECIMAL_TYPE, + from_decimal_dtype, + to_decimal_dtype + ) + } + Err(error) => error, + } +} + +#[derive(Debug, Clone, Copy)] +enum DecimalCastPlan { + SameScale { min: T, max: T }, + ScaleUp { factor: T, min: T, max: T }, + ScaleUpOverflow, + ScaleDown { factor: i256, min: i256, max: i256 }, + ScaleDownOverflow, +} + +impl DecimalCastPlan +where + T: NativeDecimalType + CheckedMul, +{ + fn new(from_decimal_dtype: DecimalDType, to_decimal_dtype: DecimalDType) -> Self { + let scale_delta = to_decimal_dtype.scale() as i16 - from_decimal_dtype.scale() as i16; + if scale_delta == 0 { + let (min, max) = decimal_precision_range::(to_decimal_dtype); + return Self::SameScale { min, max }; + } + + if scale_delta > 0 { + let Some(factor) = decimal_scale_factor::(scale_delta as u32) else { + return Self::ScaleUpOverflow; + }; + let (min, max) = decimal_precision_range::(to_decimal_dtype); + return Self::ScaleUp { factor, min, max }; + } + + let Some(factor) = decimal_scale_factor::((-scale_delta) as u32) else { + return Self::ScaleDownOverflow; + }; + let (min, max) = decimal_precision_range::(to_decimal_dtype); + Self::ScaleDown { factor, min, max } + } + + #[inline] + fn cast(&self, value: F) -> Option + where + F: NativeDecimalType, + { + match *self { + DecimalCastPlan::SameScale { min, max } => { + let value = ::from(value)?; + (value >= min && value <= max).then_some(value) + } + DecimalCastPlan::ScaleUp { factor, min, max } => { + let value = ::from(value)?; + let value = value.checked_mul(&factor)?; + (value >= min && value <= max).then_some(value) + } + DecimalCastPlan::ScaleUpOverflow | DecimalCastPlan::ScaleDownOverflow => { + (value == F::default()).then_some(T::default()) + } + DecimalCastPlan::ScaleDown { factor, min, max } => { + let value = ::from(value)?; + if value == i256::ZERO { + return Some(T::default()); + } + if value % factor != i256::ZERO { + return None; + } + + let value = value / factor; + if value < min || value > max { + return None; + } + ::from(value) + } } } } +fn decimal_precision_range(decimal_dtype: DecimalDType) -> (T, T) { + let precision = usize::from(decimal_dtype.precision()); + ( + T::MIN_BY_PRECISION[precision], + T::MAX_BY_PRECISION[precision], + ) +} + +fn decimal_scale_factor(exp: u32) -> Option +where + T: NativeDecimalType + CheckedMul, +{ + let ten = ::from(10_i8)?; + let mut factor = ::from(1_i8)?; + for _ in 0..exp { + factor = factor.checked_mul(&ten)?; + } + Some(factor) +} + /// Upcast a DecimalArray to a wider physical representation (e.g., i32 -> i64) while keeping /// the same precision and scale. /// @@ -187,8 +387,8 @@ mod tests { use super::upcast_decimal_values; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::DecimalArray; use crate::builtins::ArrayBuiltins; #[expect(deprecated)] @@ -262,19 +462,59 @@ mod tests { } #[test] - fn cast_different_scale_fails() { + fn cast_different_scale_rescales() { let array = DecimalArray::new( buffer![100i32], DecimalDType::new(10, 2), Validity::NonNullable, ); - // Try to cast to different scale - not supported + // Cast 1.00 to scale 3, where it is stored as 1000. let different_dtype = DType::Decimal(DecimalDType::new(15, 3), Nullability::NonNullable); #[expect(deprecated)] - let result = array + let casted = array .into_array() .cast(different_dtype) + .unwrap() + .to_decimal(); + + assert_eq!(casted.precision(), 15); + assert_eq!(casted.scale(), 3); + assert_eq!(casted.values_type(), DecimalType::I64); + assert_eq!(casted.buffer::().as_ref(), &[1000]); + } + + #[test] + fn cast_downcast_precision_succeeds_when_values_fit() { + let array = DecimalArray::new( + buffer![100i64], + DecimalDType::new(18, 2), + Validity::NonNullable, + ); + + // Downcasting precision is allowed when every value fits. + let smaller_dtype = DType::Decimal(DecimalDType::new(10, 2), Nullability::NonNullable); + #[expect(deprecated)] + let casted = array.into_array().cast(smaller_dtype).unwrap().to_decimal(); + + assert_eq!(casted.precision(), 10); + assert_eq!(casted.scale(), 2); + assert_eq!(casted.buffer::().as_ref(), &[100]); + } + + #[test] + fn cast_downcast_precision_checks_values() { + let array = DecimalArray::new( + buffer![1000i64], + DecimalDType::new(18, 0), + Validity::NonNullable, + ); + + let smaller_dtype = DType::Decimal(DecimalDType::new(3, 0), Nullability::NonNullable); + #[expect(deprecated)] + let result = array + .into_array() + .cast(smaller_dtype) .and_then(|a| a.to_canonical().map(|c| c.into_array())); assert!(result.is_err()); @@ -282,24 +522,23 @@ mod tests { result .unwrap_err() .to_string() - .contains("Casting decimal with scale 2 to scale 3 not yet implemented") + .contains("does not fit in precision") ); } #[test] - fn cast_downcast_precision_fails() { + fn cast_lower_scale_requires_exact_rescale() { let array = DecimalArray::new( - buffer![100i64], - DecimalDType::new(18, 2), + buffer![123456i64], + DecimalDType::new(10, 4), Validity::NonNullable, ); - // Try to downcast precision - not supported - let smaller_dtype = DType::Decimal(DecimalDType::new(10, 2), Nullability::NonNullable); + let lower_scale_dtype = DType::Decimal(DecimalDType::new(10, 2), Nullability::NonNullable); #[expect(deprecated)] let result = array .into_array() - .cast(smaller_dtype) + .cast(lower_scale_dtype) .and_then(|a| a.to_canonical().map(|c| c.into_array())); assert!(result.is_err()); @@ -307,10 +546,40 @@ mod tests { result .unwrap_err() .to_string() - .contains("Downcasting decimal from precision 18 to 10 not yet implemented") + .contains("would lose precision") ); } + #[test] + fn cast_lower_scale_ignores_null_lane_failures() { + let array = DecimalArray::new( + buffer![100i64, 123456], + DecimalDType::new(10, 4), + Validity::from_iter([true, false]), + ); + + let lower_scale_dtype = DType::Decimal(DecimalDType::new(3, 2), Nullability::Nullable); + #[expect(deprecated)] + let casted = array + .into_array() + .cast(lower_scale_dtype) + .unwrap() + .to_decimal(); + + let mask = casted + .as_ref() + .validity() + .unwrap() + .execute_mask( + casted.as_ref().len(), + &mut array_session().create_execution_ctx(), + ) + .unwrap(); + assert!(mask.value(0)); + assert!(!mask.value(1)); + assert_eq!(casted.buffer::().as_ref()[0], 1); + } + #[test] fn cast_upcast_precision_succeeds() { let array = DecimalArray::new( @@ -331,6 +600,33 @@ mod tests { assert_eq!(casted.values_type(), DecimalType::I128); } + #[test] + fn cast_widening_same_physical_type_is_zero_copy() { + // Decimal(10,2) and Decimal(18,2) are both physically i64 with the same scale, so widening + // the precision must reuse the values buffer rather than allocate and re-scan it. + let array = DecimalArray::new( + buffer![100i64, 200, 300], + DecimalDType::new(10, 2), + Validity::NonNullable, + ); + let src_ptr = array.buffer::().as_ptr(); + + let wider_dtype = DType::Decimal(DecimalDType::new(18, 2), Nullability::NonNullable); + #[expect(deprecated)] + let casted = array.into_array().cast(wider_dtype).unwrap().to_decimal(); + + assert_eq!(casted.precision(), 18); + assert_eq!(casted.scale(), 2); + assert_eq!(casted.values_type(), DecimalType::I64); + assert_eq!(casted.buffer::().as_ref(), &[100, 200, 300]); + // The values buffer must be shared with the source (zero-copy), not reallocated. + assert_eq!( + casted.buffer::().as_ptr(), + src_ptr, + "precision-widening cast must reuse the source values buffer" + ); + } + #[test] fn cast_to_non_decimal_returns_err() { let array = DecimalArray::new( @@ -440,7 +736,7 @@ mod tests { .unwrap() .execute_mask( casted.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap(); assert!(mask.value(0)); diff --git a/vortex-array/src/arrays/decimal/compute/fill_null.rs b/vortex-array/src/arrays/decimal/compute/fill_null.rs index 144c90c74f9..1c2904b7b6b 100644 --- a/vortex-array/src/arrays/decimal/compute/fill_null.rs +++ b/vortex-array/src/arrays/decimal/compute/fill_null.rs @@ -89,8 +89,8 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::DecimalArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; @@ -104,6 +104,7 @@ mod tests { #[test] fn fill_null_leading_none() { + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::from_option_iter( [None, Some(800i128), None, Some(1000i128), None], @@ -121,7 +122,8 @@ mod tests { .to_decimal(); assert_arrays_eq!( p, - DecimalArray::from_iter([4200, 800, 4200, 1000, 4200], decimal_dtype) + DecimalArray::from_iter([4200, 800, 4200, 1000, 4200], decimal_dtype), + &mut ctx ); assert_eq!( p.buffer::().as_slice(), @@ -131,7 +133,10 @@ mod tests { p.as_ref() .validity() .unwrap() - .execute_mask(p.as_ref().len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask( + p.as_ref().len(), + &mut array_session().create_execution_ctx() + ) .unwrap() .all_true() ); @@ -139,6 +144,7 @@ mod tests { #[test] fn fill_null_all_none() { + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::from_option_iter( @@ -158,13 +164,15 @@ mod tests { .to_decimal(); assert_arrays_eq!( p, - DecimalArray::from_iter([25500, 25500, 25500, 25500, 25500], decimal_dtype) + DecimalArray::from_iter([25500, 25500, 25500, 25500, 25500], decimal_dtype), + &mut ctx ); } /// fill_null with a value that overflows the array's storage type should upcast the array. #[test] fn fill_null_overflow_upcasts() { + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(3, 0); let arr = DecimalArray::from_option_iter([None, Some(10i8), None], decimal_dtype); // i8 max is 127, so 200 doesn't fit — the array should be widened to i16. @@ -180,12 +188,14 @@ mod tests { .to_decimal(); assert_arrays_eq!( result, - DecimalArray::from_iter([200i16, 10, 200], decimal_dtype) + DecimalArray::from_iter([200i16, 10, 200], decimal_dtype), + &mut ctx ); } #[test] fn fill_null_non_nullable() { + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::new( @@ -205,7 +215,8 @@ mod tests { .to_decimal(); assert_arrays_eq!( p, - DecimalArray::from_iter([800i128, 1000, 1200, 1400, 1600], decimal_dtype) + DecimalArray::from_iter([800i128, 1000, 1200, 1400, 1600], decimal_dtype), + &mut ctx ); } } diff --git a/vortex-array/src/arrays/decimal/compute/mod.rs b/vortex-array/src/arrays/decimal/compute/mod.rs index 09faba913c0..b2f21fa8398 100644 --- a/vortex-array/src/arrays/decimal/compute/mod.rs +++ b/vortex-array/src/arrays/decimal/compute/mod.rs @@ -14,6 +14,8 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::DecimalArray; use crate::compute::conformance::consistency::test_array_consistency; use crate::dtype::DecimalDType; @@ -53,6 +55,9 @@ mod tests { Validity::NonNullable, ))] fn test_decimal_consistency(#[case] array: DecimalArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &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 index 444fe313d15..762ac3d7365 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -53,6 +53,8 @@ mod tests { 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; @@ -62,6 +64,7 @@ mod tests { #[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], @@ -73,11 +76,12 @@ mod tests { let taken = array.take(indices).unwrap(); let expected = DecimalArray::from_iter([10i128, 12, 13], ddtype); - assert_arrays_eq!(expected, taken); + 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], @@ -89,7 +93,7 @@ mod tests { let taken = array.take(indices).unwrap(); let expected = DecimalArray::from_option_iter([None, Some(12i128), Some(13)], ddtype); - assert_arrays_eq!(expected, taken); + assert_arrays_eq!(expected, taken, &mut ctx); } #[rstest] diff --git a/vortex-array/src/arrays/decimal/mod.rs b/vortex-array/src/arrays/decimal/mod.rs index 489a67a0531..e740a75ea91 100644 --- a/vortex-array/src/arrays/decimal/mod.rs +++ b/vortex-array/src/arrays/decimal/mod.rs @@ -13,6 +13,10 @@ mod vtable; pub use compute::rules::DecimalMaskedValidityRule; pub use vtable::Decimal; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + mod utils; pub use utils::*; diff --git a/vortex-array/src/arrays/decimal/vtable/kernel.rs b/vortex-array/src/arrays/decimal/vtable/kernel.rs index 9ad3803be1d..19bb30d8c95 100644 --- a/vortex-array/src/arrays/decimal/vtable/kernel.rs +++ b/vortex-array/src/arrays/decimal/vtable/kernel.rs @@ -1,16 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; use crate::arrays::Decimal; +use crate::arrays::Dict; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::between::BetweenExecuteAdaptor; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&BetweenExecuteAdaptor(Decimal)), - ParentKernelSet::lift(&CastExecuteAdaptor(Decimal)), - ParentKernelSet::lift(&FillNullExecuteAdaptor(Decimal)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Decimal)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Between.id(), Decimal, BetweenExecuteAdaptor(Decimal)); + kernels.register_execute_parent_kernel(Cast.id(), Decimal, CastExecuteAdaptor(Decimal)); + kernels.register_execute_parent_kernel(FillNull.id(), Decimal, FillNullExecuteAdaptor(Decimal)); + kernels.register_execute_parent_kernel(Dict.id(), Decimal, TakeExecuteAdaptor(Decimal)); +} diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index 74e76ab4f51..674f8598256 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -3,7 +3,6 @@ use std::hash::Hasher; -use kernel::PARENT_KERNELS; use prost::Message; use vortex_buffer::Alignment; use vortex_error::VortexResult; @@ -12,6 +11,7 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use crate::ArrayParts; use crate::ArrayRef; use crate::ExecutionCtx; use crate::ExecutionResult; @@ -43,6 +43,10 @@ use crate::hash::ArrayHash; /// A [`Decimal`]-encoded Vortex array. pub type DecimalArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + // The type of the values can be determined by looking at the type info...right? #[derive(prost::Message)] pub struct DecimalMetadata { @@ -92,6 +96,24 @@ impl VTable for Decimal { } } + fn with_buffers( + &self, + 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(); + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, @@ -144,7 +166,7 @@ impl VTable for Decimal { buffers: &[BufferHandle], children: &dyn ArrayChildren, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult> { let metadata = DecimalMetadata::decode(metadata)?; if buffers.len() != 1 { vortex_bail!("Expected 1 buffer, got {}", buffers.len()); @@ -174,7 +196,7 @@ impl VTable for Decimal { ); DecimalData::try_new_handle(values, metadata.values_type(), *decimal_dtype) })?; - Ok(crate::array::ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { @@ -192,15 +214,6 @@ impl VTable for Decimal { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[derive(Clone, Debug)] @@ -214,7 +227,8 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; - use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::assert_arrays_eq; @@ -225,6 +239,7 @@ mod tests { #[test] fn test_array_serde() { + let session = array_session(); let array = DecimalArray::new( buffer![100i128, 200i128, 300i128, 400i128, 500i128], DecimalDType::new(10, 2), @@ -232,10 +247,10 @@ mod tests { ); let dtype = array.dtype().clone(); - let ctx = ArrayContext::empty(); + let array_ctx = ArrayContext::empty(); let out = array .into_array() - .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default()) + .serialize(&array_ctx, &session, &SerializeOptions::default()) .unwrap(); // Concat into a single buffer let mut concat = ByteBufferMut::empty(); @@ -247,13 +262,15 @@ mod tests { let parts = SerializedArray::try_from(concat).unwrap(); let decoded = parts - .decode(&dtype, 5, &ReadContext::new(ctx.to_ids()), &LEGACY_SESSION) + .decode(&dtype, 5, &ReadContext::new(array_ctx.to_ids()), &session) .unwrap(); assert!(decoded.is::()); } #[test] fn test_nullable_decimal_serde_roundtrip() { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); let array = DecimalArray::new( buffer![1234567i32, 0i32, -9999999i32], DecimalDType::new(7, 3), @@ -262,11 +279,11 @@ mod tests { let dtype = array.dtype().clone(); let len = array.len(); - let ctx = ArrayContext::empty(); + let array_ctx = ArrayContext::empty(); let out = array .clone() .into_array() - .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default()) + .serialize(&array_ctx, &session, &SerializeOptions::default()) .unwrap(); let mut concat = ByteBufferMut::empty(); for buf in out { @@ -275,14 +292,9 @@ mod tests { let parts = SerializedArray::try_from(concat.freeze()).unwrap(); let decoded = parts - .decode( - &dtype, - len, - &ReadContext::new(ctx.to_ids()), - &LEGACY_SESSION, - ) + .decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session) .unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut ctx); } } diff --git a/vortex-array/src/arrays/decimal/vtable/operations.rs b/vortex-array/src/arrays/decimal/vtable/operations.rs index aca0f912583..257f26127ae 100644 --- a/vortex-array/src/arrays/decimal/vtable/operations.rs +++ b/vortex-array/src/arrays/decimal/vtable/operations.rs @@ -32,8 +32,8 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::dtype::DecimalDType; @@ -81,7 +81,7 @@ mod tests { assert_eq!( array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::decimal( DecimalValue::I128(100), diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index 59758924a19..f73f4965439 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -4,6 +4,7 @@ use std::fmt::Display; use std::fmt::Formatter; +use num_traits::AsPrimitive; use smallvec::smallvec; use vortex_buffer::BitBuffer; use vortex_error::VortexExpect; @@ -159,13 +160,9 @@ pub trait DictArrayExt: TypedArrayRef + DictArraySlotsExt { match codes_validity.bit_buffer() { AllOr::All => { match_each_integer_ptype!(codes_primitive.ptype(), |P| { - #[allow( - clippy::cast_possible_truncation, - clippy::cast_sign_loss, - reason = "codes are non-negative indices; a negative signed code would wrap to a large usize and panic on the bounds-checked array index" - )] - for &idx in codes_primitive.as_slice::

() { - values_vec[idx as usize] = referenced_value; + for idx in codes_primitive.as_slice::

() { + let idxu: usize = idx.as_(); + values_vec[idxu] = referenced_value; } }); } @@ -173,14 +170,9 @@ pub trait DictArrayExt: TypedArrayRef + DictArraySlotsExt { AllOr::Some(mask) => { match_each_integer_ptype!(codes_primitive.ptype(), |P| { let codes = codes_primitive.as_slice::

(); - - #[allow( - clippy::cast_possible_truncation, - clippy::cast_sign_loss, - reason = "codes are non-negative indices; a negative signed code would wrap to a large usize and panic on the bounds-checked array index" - )] mask.set_indices().for_each(|idx| { - values_vec[codes[idx] as usize] = referenced_value; + let idxu: usize = codes[idx].as_(); + values_vec[idxu] = referenced_value; }); }); } @@ -307,10 +299,10 @@ mod test { use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; @@ -340,7 +332,7 @@ mod test { .unwrap() .execute_mask( dict.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap(); let AllOr::Some(indices) = mask.indices() else { @@ -366,7 +358,7 @@ mod test { .unwrap() .execute_mask( dict.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap(); let AllOr::Some(indices) = mask.indices() else { @@ -396,7 +388,7 @@ mod test { .unwrap() .execute_mask( dict.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap(); let AllOr::Some(indices) = mask.indices() else { @@ -422,7 +414,7 @@ mod test { .unwrap() .execute_mask( dict.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap(); let AllOr::Some(indices) = mask.indices() else { @@ -462,6 +454,7 @@ mod test { #[test] fn test_dict_array_from_primitive_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let len = 2; let chunk_count = 2; let array = make_dict_primitive_chunks::(len, 2, chunk_count); @@ -470,13 +463,16 @@ mod test { &DType::Primitive(PType::U64, NonNullable), len * chunk_count, ); - array.append_to_builder(builder.as_mut(), &mut LEGACY_SESSION.create_execution_ctx())?; + array.append_to_builder( + builder.as_mut(), + &mut array_session().create_execution_ctx(), + )?; #[expect(deprecated)] let into_prim = array.to_primitive(); let prim_into = builder.finish_into_canonical().into_primitive(); - assert_arrays_eq!(into_prim, prim_into); + assert_arrays_eq!(into_prim, prim_into, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/dict/compute/cast.rs b/vortex-array/src/arrays/dict/compute/cast.rs index 490f1f31549..56a433627a2 100644 --- a/vortex-array/src/arrays/dict/compute/cast.rs +++ b/vortex-array/src/arrays/dict/compute/cast.rs @@ -49,14 +49,15 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_session::VortexSession; + use crate::ArrayRef; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::arrays::Dict; + use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::arrays::dict::DictArraySlotsExt; use crate::assert_arrays_eq; @@ -66,15 +67,15 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; - use crate::session::ArraySession; + use crate::validity::Validity; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[test] fn test_cast_dict_to_wider_type() { + let ctx = &mut SESSION.create_execution_ctx(); let values = buffer![1i32, 2, 3, 2, 1].into_array(); - let dict = dict_encode(&values, &mut SESSION.create_execution_ctx()).unwrap(); + let dict = dict_encode(&values, ctx).unwrap(); let casted = dict .into_array() @@ -85,9 +86,8 @@ mod tests { &DType::Primitive(PType::I64, Nullability::NonNullable) ); - #[expect(deprecated)] - let decoded = casted.to_primitive(); - assert_arrays_eq!(decoded, PrimitiveArray::from_iter([1i64, 2, 3, 2, 1])); + let decoded = casted.into_array().execute::(ctx).unwrap(); + assert_arrays_eq!(decoded, PrimitiveArray::from_iter([1i64, 2, 3, 2, 1]), ctx); } #[test] @@ -108,9 +108,10 @@ mod tests { #[test] fn test_cast_dict_allvalid_to_nonnullable_and_back() { + let ctx = &mut SESSION.create_execution_ctx(); // Create an AllValid dict array (no nulls) let values = buffer![10i32, 20, 30, 40].into_array(); - let dict = dict_encode(&values, &mut SESSION.create_execution_ctx()).unwrap(); + let dict = dict_encode(&values, ctx).unwrap(); // Verify initial state - codes should be NonNullable, values should be NonNullable assert_eq!(dict.codes().dtype().nullability(), Nullability::NonNullable); @@ -171,11 +172,10 @@ mod tests { ); // Verify values are unchanged - #[expect(deprecated)] - let original_values = dict.as_array().to_primitive(); - #[expect(deprecated)] - let final_values = back_to_non_nullable.to_primitive(); - assert_arrays_eq!(original_values, final_values); + let original_values = dict.into_array().execute::(ctx).unwrap(); + + let final_values = back_to_non_nullable.execute::(ctx).unwrap(); + assert_arrays_eq!(original_values, final_values, ctx); } #[rstest] @@ -183,21 +183,19 @@ mod tests { #[case(dict_encode(&buffer![100u32, 200, 100, 300, 200].into_array(), &mut SESSION.create_execution_ctx()).unwrap().into_array())] #[case(dict_encode(&PrimitiveArray::from_option_iter([Some(1i32), None, Some(2), Some(1), None]).into_array(), &mut SESSION.create_execution_ctx()).unwrap().into_array())] #[case(dict_encode(&buffer![1.5f32, 2.5, 1.5, 3.5].into_array(), &mut SESSION.create_execution_ctx()).unwrap().into_array())] - fn test_cast_dict_conformance(#[case] array: crate::ArrayRef) { + fn test_cast_dict_conformance(#[case] array: ArrayRef) { test_cast_conformance(&array); } #[test] fn test_cast_dict_with_unreferenced_null_values_to_nonnullable() { - use crate::arrays::DictArray; - use crate::validity::Validity; - + let ctx = &mut SESSION.create_execution_ctx(); // Create a dict with nullable values that have unreferenced null entries. // Values: [1.0, null, 3.0] (index 1 is null but no code points to it) // Codes: [0, 2, 0] (only reference indices 0 and 2, never 1) let values = PrimitiveArray::new( buffer![1.0f64, 0.0f64, 3.0f64], - Validity::from(vortex_buffer::BitBuffer::from(vec![true, false, true])), + Validity::from(BitBuffer::from(vec![true, false, true])), ) .into_array(); let codes = buffer![0u32, 2, 0].into_array(); @@ -222,8 +220,6 @@ mod tests { casted.dtype(), &DType::Primitive(PType::F64, Nullability::NonNullable) ); - #[expect(deprecated)] - let casted_prim = casted.to_primitive(); - assert_arrays_eq!(casted_prim, PrimitiveArray::from_iter([1.0f64, 3.0, 1.0])); + assert_arrays_eq!(casted, PrimitiveArray::from_iter([1.0f64, 3.0, 1.0]), ctx); } } diff --git a/vortex-array/src/arrays/dict/compute/fill_null.rs b/vortex-array/src/arrays/dict/compute/fill_null.rs index 8f146d728a3..290fdd95bda 100644 --- a/vortex-array/src/arrays/dict/compute/fill_null.rs +++ b/vortex-array/src/arrays/dict/compute/fill_null.rs @@ -94,10 +94,8 @@ mod tests { use vortex_error::VortexExpect; use crate::IntoArray; - use crate::LEGACY_SESSION; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -108,6 +106,7 @@ mod tests { #[test] fn nullable_codes_fill_in_values() { + let mut ctx = array_session().create_execution_ctx(); let dict = DictArray::try_new( PrimitiveArray::new( buffer![0u32, 1, 2], @@ -122,13 +121,12 @@ mod tests { .into_array() .fill_null(Scalar::primitive(20, Nullability::NonNullable)) .vortex_expect("operation should succeed in test"); - #[expect(deprecated)] - let filled_primitive = filled.to_primitive(); - assert_arrays_eq!(filled_primitive, PrimitiveArray::from_iter([10, 20, 20])); - assert!( - filled_primitive - .all_valid(&mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + let filled_primitive = filled.execute::(&mut ctx).unwrap(); + assert_arrays_eq!( + filled_primitive, + PrimitiveArray::from_iter([10, 20, 20]), + &mut ctx ); + assert!(filled_primitive.all_valid(&mut ctx).unwrap()); } } diff --git a/vortex-array/src/arrays/dict/compute/like.rs b/vortex-array/src/arrays/dict/compute/like.rs index 40bcea8fcf1..a4ca7c1f76e 100644 --- a/vortex-array/src/arrays/dict/compute/like.rs +++ b/vortex-array/src/arrays/dict/compute/like.rs @@ -56,6 +56,8 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::DictArray; use crate::arrays::VarBinArray; @@ -68,6 +70,7 @@ mod tests { #[test] fn like_reduce_dict() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dict = DictArray::try_new( buffer![0u8, 1, 0, 2].into_array(), VarBinArray::from(vec!["hello", "world", "help"]).into_array(), @@ -79,7 +82,11 @@ mod tests { .try_new_array(4, LikeOptions::default(), [dict, pattern])? .optimize()?; - assert_arrays_eq!(result, BoolArray::from_iter([true, false, true, false])); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, false, true, false]), + &mut ctx + ); Ok(()) } } diff --git a/vortex-array/src/arrays/dict/compute/min_max.rs b/vortex-array/src/arrays/dict/compute/min_max.rs index 47a1f0eb8ee..ee38c0392d4 100644 --- a/vortex-array/src/arrays/dict/compute/min_max.rs +++ b/vortex-array/src/arrays/dict/compute/min_max.rs @@ -30,9 +30,9 @@ impl DynAggregateKernel for DictMinMaxKernel { batch: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if !aggregate_fn.is::() { + let Some(options) = aggregate_fn.as_opt::() else { return Ok(None); - } + }; let Some(dict) = batch.as_opt::() else { return Ok(None); @@ -42,13 +42,13 @@ impl DynAggregateKernel for DictMinMaxKernel { let result = if dict.has_all_values_referenced() { // All values are referenced, compute min/max directly on the values array. - min_max(dict.values(), ctx)? + min_max(dict.values(), ctx, *options)? } else { // Filter to only referenced values, then compute min/max. let referenced_mask = dict.compute_referenced_values_mask(true)?; let mask = Mask::from(referenced_mask); let filtered_values = dict.values().filter(mask)?; - min_max(&filtered_values, ctx)? + min_max(&filtered_values, ctx, *options)? }; match result { @@ -70,18 +70,20 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::min_max; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::builders::dict::dict_encode; - use crate::session::ArraySession; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); fn assert_min_max(array: &ArrayRef, expected: Option<(i32, i32)>) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - match (min_max(array, &mut ctx)?, expected) { + match ( + min_max(array, &mut ctx, NumericalAggregateOpts::default())?, + expected, + ) { (Some(result), Some((expected_min, expected_max))) => { assert_eq!(i32::try_from(&result.min)?, expected_min); assert_eq!(i32::try_from(&result.max)?, expected_max); diff --git a/vortex-array/src/arrays/dict/compute/mod.rs b/vortex-array/src/arrays/dict/compute/mod.rs index 28ba1d84789..e2b0d07b97d 100644 --- a/vortex-array/src/arrays/dict/compute/mod.rs +++ b/vortex-array/src/arrays/dict/compute/mod.rs @@ -56,17 +56,13 @@ impl FilterReduce for Dict { mod test { use std::sync::LazyLock; - #[expect(unused_imports)] - use itertools::Itertools; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::ArrayRef; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::accessor::ArrayAccessor; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; @@ -81,13 +77,12 @@ mod test { use crate::dtype::Nullability; use crate::dtype::PType::I32; use crate::scalar_fn::fns::operators::Operator; - use crate::session::ArraySession; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[test] fn canonicalise_nullable_primitive() { + let mut ctx = SESSION.create_execution_ctx(); let values: Vec> = (0..65) .map(|i| match i % 3 { 0 => Some(42), @@ -102,16 +97,19 @@ mod test { &mut SESSION.create_execution_ctx(), ) .unwrap(); - #[expect(deprecated)] - let actual = dict.as_array().to_primitive(); + let actual = dict + .into_array() + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); let expected = PrimitiveArray::from_option_iter(values); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn canonicalise_non_nullable_primitive_32_unique_values() { + let mut ctx = SESSION.create_execution_ctx(); let unique_values: Vec = (0..32).collect(); let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 32])); @@ -119,15 +117,19 @@ mod test { &expected.clone().into_array(), &mut SESSION.create_execution_ctx(), ) - .unwrap(); - #[expect(deprecated)] - let actual = dict.as_array().to_primitive(); + .unwrap() + .into_array(); + + let actual = dict + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn canonicalise_non_nullable_primitive_100_unique_values() { + let mut ctx = SESSION.create_execution_ctx(); let unique_values: Vec = (0..100).collect(); let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 100])); @@ -135,35 +137,48 @@ mod test { &expected.clone().into_array(), &mut SESSION.create_execution_ctx(), ) - .unwrap(); - #[expect(deprecated)] - let actual = dict.as_array().to_primitive(); + .unwrap() + .into_array(); + + let actual = dict + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] - fn canonicalise_nullable_varbin() { + fn canonicalise_nullable_varbin() -> VortexResult<()> { let reference = VarBinViewArray::from_iter( vec![Some("a"), Some("b"), None, Some("a"), None, Some("b")], DType::Utf8(Nullability::Nullable), ); assert_eq!(reference.len(), 6); - let dict = dict_encode( - &reference.clone().into_array(), - &mut SESSION.create_execution_ctx(), - ) - .unwrap(); - #[expect(deprecated)] - let flattened_dict = dict.as_array().to_varbinview(); - assert_eq!( - flattened_dict.with_iterator(|iter| iter - .map(|slice| slice.map(|s| s.to_vec())) - .collect::>()), - reference.with_iterator(|iter| iter - .map(|slice| slice.map(|s| s.to_vec())) - .collect::>()), - ); + let mut ctx = SESSION.create_execution_ctx(); + let dict = dict_encode(&reference.clone().into_array(), &mut ctx)?; + let flattened_dict = dict.into_array().execute::(&mut ctx)?; + let flattened_mask = flattened_dict + .validity()? + .execute_mask(flattened_dict.len(), &mut ctx)?; + let flattened_values = (0..flattened_dict.len()) + .map(|i| { + flattened_mask + .value(i) + .then(|| flattened_dict.bytes_at(i).to_vec()) + }) + .collect::>(); + let reference_mask = reference + .validity()? + .execute_mask(reference.len(), &mut ctx)?; + let reference_values = (0..reference.len()) + .map(|i| { + reference_mask + .value(i) + .then(|| reference.bytes_at(i).to_vec()) + }) + .collect::>(); + assert_eq!(flattened_values, reference_values); + Ok(()) } fn sliced_dict_array() -> ArrayRef { @@ -182,6 +197,7 @@ mod test { #[test] fn compare_sliced_dict() { + let mut ctx = SESSION.create_execution_ctx(); use crate::arrays::BoolArray; let sliced = sliced_dict_array(); let compared = sliced @@ -189,7 +205,7 @@ mod test { .unwrap(); let expected = BoolArray::from_iter([Some(false), None, Some(true)]); - assert_arrays_eq!(compared, expected.into_array()); + assert_arrays_eq!(compared, expected.into_array(), &mut ctx); } #[test] @@ -332,10 +348,8 @@ mod tests { use crate::compute::conformance::consistency::test_array_consistency; use crate::dtype::DType; use crate::dtype::Nullability; - use crate::session::ArraySession; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[rstest] // Primitive arrays @@ -366,6 +380,6 @@ mod tests { #[case::dict_all_same(dict_encode(&buffer![5i32, 5, 5, 5, 5].into_array(), &mut SESSION.create_execution_ctx()).unwrap())] #[case::dict_large(dict_encode(&PrimitiveArray::from_iter((0..1000).map(|i| i % 10)).into_array(), &mut SESSION.create_execution_ctx()).unwrap())] fn test_dict_consistency(#[case] array: DictArray) { - test_array_consistency(&array.into_array()); + test_array_consistency(&array.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/vortex-array/src/arrays/dict/compute/rules.rs b/vortex-array/src/arrays/dict/compute/rules.rs index 024240e4030..e4d102ee861 100644 --- a/vortex-array/src/arrays/dict/compute/rules.rs +++ b/vortex-array/src/arrays/dict/compute/rules.rs @@ -9,12 +9,16 @@ use crate::EqMode; use crate::IntoArray; use crate::array::ArrayView; use crate::array::VTable; +use crate::arrays::Chunked; +use crate::arrays::ChunkedArray; use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::Dict; use crate::arrays::DictArray; use crate::arrays::ScalarFn; use crate::arrays::ScalarFnArray; +use crate::arrays::chunked::ChunkedArrayExt; +use crate::arrays::dict::DictArrayExt; use crate::arrays::dict::DictArraySlotsExt; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::scalar_fn::AnyScalarFn; @@ -36,11 +40,59 @@ pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&CastReduceAdaptor(Dict)), ParentRuleSet::lift(&MaskReduceAdaptor(Dict)), ParentRuleSet::lift(&LikeReduceAdaptor(Dict)), + ParentRuleSet::lift(&DictionaryChunkedValuesPullUpRule), ParentRuleSet::lift(&DictionaryScalarFnValuesPushDownRule), ParentRuleSet::lift(&DictionaryScalarFnCodesPullUpRule), ParentRuleSet::lift(&SliceReduceAdaptor(Dict)), ]); +/// Pull a common dictionary values array above chunked dictionary codes. +/// +/// Rewrites `Chunked>` into `Dict, values>` only when +/// every child dictionary shares the exact same values array allocation. +#[derive(Debug)] +struct DictionaryChunkedValuesPullUpRule; + +impl ArrayParentReduceRule for DictionaryChunkedValuesPullUpRule { + type Parent = Chunked; + + fn reduce_parent( + &self, + array: ArrayView<'_, Dict>, + parent: ArrayView<'_, Chunked>, + _child_idx: usize, + ) -> VortexResult> { + let values = array.values(); + let codes_dtype = array.codes().dtype().clone(); + let mut code_chunks = Vec::with_capacity(parent.nchunks()); + let mut all_values_referenced = array.has_all_values_referenced(); + + for chunk in parent.iter_chunks() { + let Some(dict) = chunk.as_opt::() else { + return Ok(None); + }; + if dict.codes().dtype() != &codes_dtype { + return Ok(None); + } + if !ArrayRef::ptr_eq(dict.values(), values) { + return Ok(None); + } + all_values_referenced |= dict.has_all_values_referenced(); + code_chunks.push(dict.codes().clone()); + } + + let codes = ChunkedArray::try_new(code_chunks, codes_dtype)?.into_array(); + let dict = DictArray::try_new(codes, values.clone())?; + let dict = if all_values_referenced { + unsafe { dict.set_all_values_referenced(true) } + } else { + dict + }; + + Ok(Some(dict.into_array())) + } +} + /// Push down a scalar function to run only over the values of a dictionary array. #[derive(Debug)] struct DictionaryScalarFnValuesPushDownRule; @@ -136,13 +188,19 @@ impl ArrayParentReduceRule for DictionaryScalarFnValuesPushDownRule { // 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) }.into_array(); + let new_dict = unsafe { + DictArray::new_unchecked(new_codes, new_values) + .set_all_values_referenced(array.has_all_values_referenced()) + } + .into_array(); return Ok(Some(new_dict.cast(parent.dtype().clone())?)); } - Ok(Some( - unsafe { DictArray::new_unchecked(array.codes().clone(), new_values) }.into_array(), - )) + Ok(Some(unsafe { + DictArray::new_unchecked(array.codes().clone(), new_values) + .set_all_values_referenced(array.has_all_values_referenced()) + .into_array() + })) } } @@ -201,3 +259,96 @@ impl ArrayParentReduceRule for DictionaryScalarFnCodesPullUpRule { Ok(Some(new_dict)) } } + +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::ArrayRef; + use crate::IntoArray; + use crate::arrays::BoolArray; + use crate::arrays::Chunked; + use crate::arrays::ChunkedArray; + 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::executor::VortexSessionExecute; + use crate::optimizer::ArrayOptimizer; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::fns::not::Not; + + #[test] + fn chunked_dict_with_shared_values_pulls_values_up() -> VortexResult<()> { + let values = buffer![10u32, 20, 30].into_array(); + let chunk0 = DictArray::try_new(buffer![0u8, 1].into_array(), values.clone())?.into_array(); + let chunk1 = + DictArray::try_new(buffer![2u8, 0, 1].into_array(), values.clone())?.into_array(); + let array = + ChunkedArray::try_new(vec![chunk0, chunk1], values.dtype().clone())?.into_array(); + + let optimized = array.optimize()?; + let dict = optimized.as_::(); + let codes = dict.codes().as_::(); + + assert!(ArrayRef::ptr_eq(dict.values(), &values)); + assert_eq!(codes.nchunks(), 2); + let mut ctx = crate::LEGACY_SESSION.create_execution_ctx(); + assert_arrays_eq!( + optimized, + PrimitiveArray::from_iter([10u32, 20, 30, 10, 20]), + &mut ctx + ); + + Ok(()) + } + + #[test] + fn chunked_dict_with_distinct_values_stays_chunked() -> VortexResult<()> { + let values0 = buffer![10u32, 20, 30].into_array(); + let values1 = buffer![10u32, 20, 30].into_array(); + let chunk0 = + DictArray::try_new(buffer![0u8, 1].into_array(), values0.clone())?.into_array(); + let chunk1 = DictArray::try_new(buffer![2u8, 0, 1].into_array(), values1)?.into_array(); + let array = + ChunkedArray::try_new(vec![chunk0, chunk1], values0.dtype().clone())?.into_array(); + + let optimized = array.optimize()?; + + assert!(optimized.is::()); + let mut ctx = crate::LEGACY_SESSION.create_execution_ctx(); + assert_arrays_eq!( + optimized, + PrimitiveArray::from_iter([10u32, 20, 30, 10, 20]), + &mut ctx + ); + + Ok(()) + } + + #[test] + fn scalar_fn_values_pushdown_preserves_all_values_referenced() -> VortexResult<()> { + let dict = unsafe { + DictArray::try_new( + buffer![0u8, 1, 0, 1].into_array(), + BoolArray::from_iter([true, false]).into_array(), + )? + .set_all_values_referenced(true) + } + .into_array(); + + let result = Not + .try_new_array(dict.len(), EmptyOptions, [dict])? + .optimize()?; + let result = result.as_::(); + + assert!(result.has_all_values_referenced()); + + Ok(()) + } +} diff --git a/vortex-array/src/arrays/dict/compute/slice.rs b/vortex-array/src/arrays/dict/compute/slice.rs index 1552c9c9fea..376621fba2f 100644 --- a/vortex-array/src/arrays/dict/compute/slice.rs +++ b/vortex-array/src/arrays/dict/compute/slice.rs @@ -91,6 +91,8 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::arrays::dict::compute::slice::ConstantArray; @@ -102,18 +104,20 @@ mod tests { #[test] fn slice_constant_valid_code() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dict = DictArray::new( ConstantArray::new(1u8, 5).into_array(), buffer![10i32, 20, 30].into_array(), ); let sliced = dict.slice(1..4)?; let expected = PrimitiveArray::from_iter([20i32, 20, 20]).into_array(); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut ctx); Ok(()) } #[test] fn slice_constant_null_code() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dict = DictArray::new( ConstantArray::new(Scalar::null(DType::Primitive(PType::U8, Nullable)), 5).into_array(), buffer![10i32, 20, 30].into_array(), @@ -121,7 +125,7 @@ mod tests { let sliced = dict.slice(1..4)?; let expected = PrimitiveArray::from_option_iter([Option::::None, None, None]).into_array(); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/arrays/dict/mod.rs b/vortex-array/src/arrays/dict/mod.rs index 0414eea7def..e8a83aa88b3 100644 --- a/vortex-array/src/arrays/dict/mod.rs +++ b/vortex-array/src/arrays/dict/mod.rs @@ -23,5 +23,9 @@ pub use take::*; pub mod vtable; pub use vtable::*; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/dict/tests.rs b/vortex-array/src/arrays/dict/tests.rs index bc69998b0de..7961e2b7037 100644 --- a/vortex-array/src/arrays/dict/tests.rs +++ b/vortex-array/src/arrays/dict/tests.rs @@ -5,6 +5,8 @@ use vortex_buffer::buffer; use super::DictArray; use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; @@ -14,6 +16,7 @@ use crate::validity::Validity; #[test] fn test_scalar_at_null_code() { + let mut ctx = array_session().create_execution_ctx(); let dict = DictArray::try_new( PrimitiveArray::from_option_iter(vec![None, Some(0u32), None]).into_array(), buffer![1i32].into_array(), @@ -21,7 +24,7 @@ fn test_scalar_at_null_code() { .unwrap(); let expected = PrimitiveArray::from_option_iter(vec![None, Some(1i32), None]).into_array(); - assert_arrays_eq!(dict, expected); + assert_arrays_eq!(dict, expected, &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/dict/vtable/kernel.rs b/vortex-array/src/arrays/dict/vtable/kernel.rs index fd1b0ed0a6d..79659af18dd 100644 --- a/vortex-array/src/arrays/dict/vtable/kernel.rs +++ b/vortex-array/src/arrays/dict/vtable/kernel.rs @@ -1,14 +1,23 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + +use crate::ArrayVTable; +use crate::arrays::Chunked; use crate::arrays::Dict; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; +use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(Dict)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Dict)), - ParentKernelSet::lift(&FillNullExecuteAdaptor(Dict)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), Dict, CompareExecuteAdaptor(Dict)); + kernels.register_execute_parent_kernel(Dict.id(), Chunked, TakeExecuteAdaptor(Chunked)); + kernels.register_execute_parent_kernel(Dict.id(), Dict, TakeExecuteAdaptor(Dict)); + kernels.register_execute_parent_kernel(FillNull.id(), Dict, FillNullExecuteAdaptor(Dict)); +} diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index fa8515dd986..d9ced9eeffa 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -3,7 +3,6 @@ use std::hash::Hasher; -use kernel::PARENT_KERNELS; use prost::Message; use smallvec::smallvec; use vortex_error::VortexResult; @@ -26,11 +25,13 @@ use crate::ArrayHash; use crate::ArrayRef; use crate::Canonical; use crate::EqMode; +use crate::IntoArray; use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::VTable; +use crate::array::with_empty_buffers; use crate::arrays::ConstantArray; use crate::arrays::Primitive; use crate::arrays::dict::DictArrayExt; @@ -38,6 +39,7 @@ 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::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -46,7 +48,6 @@ use crate::executor::ExecutionResult; use crate::require_child; use crate::scalar::Scalar; use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; @@ -55,6 +56,10 @@ mod validity; /// A [`Dict`]-encoded Vortex array. pub type DictArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(Clone, Debug)] pub struct Dict; @@ -112,6 +117,14 @@ impl VTable for Dict { None } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, @@ -179,7 +192,7 @@ impl VTable for Dict { let array = require_child!(array, array.codes(), DictSlots::CODES => Primitive); - if matches!(array.codes().validity()?, Validity::AllInvalid) { + if array.codes().validity()?.definitely_all_null() { return Ok(ExecutionResult::done(ConstantArray::new( Scalar::null(array.dtype().as_nullable()), array.codes().len(), @@ -197,20 +210,38 @@ impl VTable for Dict { )?)) } - fn reduce_parent( + fn append_to_builder( array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if !array.is_empty() + && let (Some(codes), Some(values)) = ( + array.codes().as_opt::(), + array.values().as_opt::(), + ) + && !codes.validity()?.definitely_all_null() + { + let codes = codes.into_owned(); + let canonical = take_canonical(values, &codes, ctx)?.into_array(); + canonical.append_to_builder(builder, ctx)?; + return Ok(()); + } + + let canonical = array + .array() + .clone() + .execute::(ctx)? + .into_array(); + canonical.append_to_builder(builder, ctx)?; + Ok(()) } - fn execute_parent( + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, child_idx: usize, - ctx: &mut ExecutionCtx, ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) + PARENT_RULES.evaluate(array, parent, child_idx) } } diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index 0f4746fd4d6..acf12b47070 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -6,8 +6,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; -use crate::ToCanonical; use crate::array::ArrayView; use crate::arrays::Extension; use crate::arrays::ExtensionArray; @@ -19,6 +19,7 @@ use crate::dtype::PType; use crate::extension::datetime::AnyTemporal; use crate::extension::datetime::TemporalMetadata; use crate::extension::datetime::TimeUnit; +use crate::scalar_fn::fns::cast::CastKernel; use crate::scalar_fn::fns::cast::CastReduce; use crate::validity::Validity; @@ -31,10 +32,11 @@ impl CastReduce for Extension { }; if array.ext_dtype().eq_ignore_nullability(ext_dtype) { + // Same extension type: restructure by casting the storage array. This is buffer-free + // (`cast` returns a lazy cast expression), so it stays in the reduce phase. let new_storage = match array .storage_array() .cast(ext_dtype.storage_dtype().clone()) - .and_then(|a| a.to_canonical().map(|c| c.into_array())) { Ok(arr) => arr, Err(e) => { @@ -48,7 +50,25 @@ impl CastReduce for Extension { )); } - if let Some(new_storage) = cast_temporal_date_to_timestamp(&array, dtype)? { + // A different extension type (e.g. Date -> Timestamp) requires per-value conversion that + // reads buffers, so it is deferred to the execute phase in `CastKernel`, where an + // `ExecutionCtx` is available. + Ok(None) + } +} + +impl CastKernel for Extension { + fn cast( + array: ArrayView<'_, Extension>, + dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let DType::Extension(ext_dtype) = dtype else { + // Non-extension targets are restructured buffer-free by `CastReduce`. + return Ok(None); + }; + + if let Some(new_storage) = cast_temporal_date_to_timestamp(&array, dtype, ctx)? { return Ok(Some( ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(), )); @@ -61,6 +81,7 @@ impl CastReduce for Extension { fn cast_temporal_date_to_timestamp( array: &ArrayView<'_, Extension>, target_dtype: &DType, + ctx: &mut ExecutionCtx, ) -> VortexResult> { let DType::Extension(target_ext_dtype) = target_dtype else { return Ok(None); @@ -81,21 +102,25 @@ fn cast_temporal_date_to_timestamp( let source_i64 = array .storage_array() - .cast(DType::Primitive(PType::I64, array.dtype().nullability()))?; - let source_i64 = source_i64.to_primitive(); + .cast(DType::Primitive(PType::I64, array.dtype().nullability()))? + .execute::(ctx)?; - let converted = cast_date_values_to_timestamp(&source_i64, *source_unit, *target_unit)?; + let converted = cast_date_values_to_timestamp(&source_i64, *source_unit, *target_unit, ctx)?; - converted + let new_storage = converted .into_array() - .cast(target_ext_dtype.storage_dtype().clone()) - .map(Some) + .cast(target_ext_dtype.storage_dtype().clone())? + .execute::(ctx)? + .into_array(); + + Ok(Some(new_storage)) } fn cast_date_values_to_timestamp( values: &PrimitiveArray, source_unit: TimeUnit, target_unit: TimeUnit, + ctx: &mut ExecutionCtx, ) -> VortexResult { let (multiply, divide) = date_to_timestamp_scale(source_unit, target_unit)?; @@ -116,22 +141,19 @@ fn cast_date_values_to_timestamp( unsafe { output.push_unchecked(0i64) }; } } - Validity::Array(validity_arr) => { + Validity::Array(_) => { + // Resolve validity to a boolean mask once. Null slots keep a placeholder 0 so a garbage + // source value in a null slot cannot trip the overflow check in + // `convert_temporal_value`; the output re-uses `validity`, so those slots stay null. + let mask = validity.execute_mask(input.len(), ctx)?; for (i, &value) in input.iter().enumerate() { - let is_valid = validity_arr - .scalar_at(i)? - .as_bool() - .value() - .unwrap_or(false); - if is_valid { - // SAFETY: output has sufficient capacity for all pushed values. - unsafe { - output.push_unchecked(convert_temporal_value(value, multiply, divide)?) - }; + let converted = if mask.value(i) { + convert_temporal_value(value, multiply, divide)? } else { - // SAFETY: output has sufficient capacity for all pushed values. - unsafe { output.push_unchecked(0i64) }; - } + 0i64 + }; + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(converted) }; } } } @@ -214,10 +236,8 @@ mod tests { use crate::extension::datetime::Date; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; - use crate::session::ArraySession; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[test] fn cast_same_ext_dtype() { @@ -253,28 +273,28 @@ mod tests { } #[test] - fn cast_date_days_to_timestamp_nanoseconds() { + fn cast_date_days_to_timestamp_nanoseconds() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let source_dtype = Date::new(TimeUnit::Days, Nullability::NonNullable).erased(); let target_dtype = Timestamp::new(TimeUnit::Nanoseconds, Nullability::NonNullable).erased(); let arr = ExtensionArray::new(source_dtype, buffer![0i32, 1, -1].into_array()); let output = arr .into_array() - .cast(DType::Extension(target_dtype.clone())) - .unwrap() - .to_extension(); - - assert_eq!(output.dtype(), &DType::Extension(target_dtype)); + .cast(DType::Extension(target_dtype.clone()))?; - let storage = output.storage_array().to_primitive(); - assert_eq!( - storage.as_slice::(), - &[0, 86_400_000_000_000, -86_400_000_000_000] - ); + let expected = ExtensionArray::new( + target_dtype, + buffer![0i64, 86_400_000_000_000, -86_400_000_000_000].into_array(), + ) + .into_array(); + assert_arrays_eq!(output, expected, &mut ctx); + Ok(()) } #[test] - fn cast_date_days_to_timestamp_seconds_nullable() { + fn cast_date_days_to_timestamp_seconds_nullable() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let source_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); let target_dtype = Timestamp::new(TimeUnit::Seconds, Nullability::Nullable).erased(); @@ -285,22 +305,15 @@ mod tests { let output = arr .into_array() - .cast(DType::Extension(target_dtype.clone())) - .unwrap() - .to_extension(); - - assert_eq!(output.dtype(), &DType::Extension(target_dtype)); + .cast(DType::Extension(target_dtype.clone()))?; - let storage = output.storage_array().to_primitive(); - assert_eq!( - storage.scalar_at(0).unwrap().as_primitive().as_::(), - Some(0) - ); - assert!(storage.scalar_at(1).unwrap().is_null()); - assert_eq!( - storage.scalar_at(2).unwrap().as_primitive().as_::(), - Some(172_800) - ); + let expected = ExtensionArray::new( + target_dtype, + PrimitiveArray::from_option_iter([Some(0i64), None, Some(172_800)]).into_array(), + ) + .into_array(); + assert_arrays_eq!(output, expected, &mut ctx); + Ok(()) } #[test] @@ -325,6 +338,7 @@ mod tests { #[test] fn cast_timestamp_to_i64() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let ext_dtype = Timestamp::new_with_tz( TimeUnit::Nanoseconds, Some("UTC".into()), @@ -339,7 +353,7 @@ mod tests { result.dtype(), &DType::Primitive(PType::I64, Nullability::NonNullable) ); - assert_arrays_eq!(result, buffer![1i64, 2, 3].into_array()); + assert_arrays_eq!(result, buffer![1i64, 2, 3].into_array(), &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/extension/compute/mod.rs b/vortex-array/src/arrays/extension/compute/mod.rs index ce25ab2c0ea..8309ae43be2 100644 --- a/vortex-array/src/arrays/extension/compute/mod.rs +++ b/vortex-array/src/arrays/extension/compute/mod.rs @@ -87,6 +87,8 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ExtensionArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::consistency::test_array_consistency; @@ -114,6 +116,9 @@ mod tests { ExtensionArray::new(ext_dtype, storage) })] fn test_extension_consistency(#[case] array: ExtensionArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/extension/compute/rules.rs b/vortex-array/src/arrays/extension/compute/rules.rs index 59a8c314517..396729b6347 100644 --- a/vortex-array/src/arrays/extension/compute/rules.rs +++ b/vortex-array/src/arrays/extension/compute/rules.rs @@ -111,10 +111,8 @@ mod tests { use crate::scalar::ScalarValue; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::Operator; - use crate::session::ArraySession; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] struct TestExt; @@ -122,6 +120,7 @@ mod tests { type Metadata = EmptyMetadata; type NativeValue<'a> = &'a str; + #[expect(clippy::disallowed_methods, reason = "test-only id")] fn id(&self) -> ExtId { ExtId::new("test_ext") } @@ -246,6 +245,7 @@ mod tests { type Metadata = EmptyMetadata; type NativeValue<'a> = &'a str; + #[expect(clippy::disallowed_methods, reason = "test-only id")] fn id(&self) -> ExtId { ExtId::new("test_ext_2") } diff --git a/vortex-array/src/arrays/extension/mod.rs b/vortex-array/src/arrays/extension/mod.rs index 57bc9df5ead..8dc5972b86f 100644 --- a/vortex-array/src/arrays/extension/mod.rs +++ b/vortex-array/src/arrays/extension/mod.rs @@ -9,3 +9,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::Extension; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} diff --git a/vortex-array/src/arrays/extension/vtable/kernel.rs b/vortex-array/src/arrays/extension/vtable/kernel.rs index 4cc5f9eb4ac..a81863a3ed6 100644 --- a/vortex-array/src/arrays/extension/vtable/kernel.rs +++ b/vortex-array/src/arrays/extension/vtable/kernel.rs @@ -1,12 +1,26 @@ // 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::Extension; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; +use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(Extension)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Extension)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Binary.id(), + Extension, + CompareExecuteAdaptor(Extension), + ); + kernels.register_execute_parent_kernel(Cast.id(), Extension, CastExecuteAdaptor(Extension)); + kernels.register_execute_parent_kernel(Dict.id(), Extension, TakeExecuteAdaptor(Extension)); +} diff --git a/vortex-array/src/arrays/extension/vtable/mod.rs b/vortex-array/src/arrays/extension/vtable/mod.rs index 852593abddb..c27e70476dc 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 kernel::PARENT_KERNELS; use smallvec::smallvec; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -22,6 +21,7 @@ use crate::array::ArrayParts; 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::compute::rules::PARENT_RULES; @@ -78,6 +78,10 @@ pub struct Extension; /// A [`Extension`]-encoded Vortex array. pub type ExtensionArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + impl VTable for Extension { type TypedArrayData = EmptyArrayData; @@ -132,6 +136,14 @@ impl VTable for Extension { None } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + fn serialize( _array: ArrayView<'_, Self>, _session: &VortexSession, @@ -176,15 +188,6 @@ impl VTable for Extension { Ok(ExecutionResult::done(array)) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { RULES.evaluate(array) } 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 3085debf9e1..392a6373f76 100644 --- a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs +++ b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs @@ -124,6 +124,8 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -149,6 +151,7 @@ mod test { #[test] fn filter_fixed_size_list_selects_correct_lists() { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]); let array = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -160,11 +163,12 @@ mod test { let expected = FixedSizeListArray::new(expected_elements.into_array(), 2, Validity::NonNullable, 2); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn filter_degenerate_list_size_zero() { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::empty::(Nullability::NonNullable); let array = FixedSizeListArray::new(elements.into_array(), 0, Validity::NonNullable, 5); @@ -175,11 +179,12 @@ mod test { let expected = FixedSizeListArray::new(expected_elements.into_array(), 0, Validity::NonNullable, 3); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn filter_nested_fixed_size_lists() { + let mut ctx = array_session().create_execution_ctx(); // Inner lists of size 2, outer lists of size 2 (so 2 outer lists, each with 2 inner lists). let inner_elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(); let inner_fsl = FixedSizeListArray::new(inner_elements, 2, Validity::NonNullable, 4); @@ -196,6 +201,6 @@ mod test { let expected_outer = FixedSizeListArray::new(expected_inner.into_array(), 2, Validity::NonNullable, 1); - assert_arrays_eq!(filtered, expected_outer); + assert_arrays_eq!(filtered, expected_outer, &mut ctx); } } diff --git a/vortex-array/src/arrays/filter/execute/listview.rs b/vortex-array/src/arrays/filter/execute/listview.rs index b3f787133ab..0507452c89d 100644 --- a/vortex-array/src/arrays/filter/execute/listview.rs +++ b/vortex-array/src/arrays/filter/execute/listview.rs @@ -65,19 +65,17 @@ mod test { use vortex_session::VortexSession; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::filter::execute::ConstantArray; use crate::arrays::listview::ListViewArrayExt; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; - use crate::session::ArraySession; use crate::validity::Validity; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(array_session); #[test] fn test_filter_listview_conformance() { @@ -137,6 +135,7 @@ mod test { #[test] fn filter_listview_selects_correct_lists() { + let mut ctx = array_session().create_execution_ctx(); // 3 lists: [10,20], [30,40], [50,60] let elements = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]); let offsets = buffer![0u32, 2, 4].into_array(); @@ -159,11 +158,12 @@ mod test { Validity::NonNullable, ); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_filter_preserves_unreferenced_elements() { + let mut ctx = array_session().create_execution_ctx(); // ListView-specific: Test that filter preserves the entire elements array. // // Logical list: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]] @@ -187,7 +187,8 @@ mod test { // Verify the entire elements array is preserved. assert_arrays_eq!( result_list.elements(), - PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]), + &mut ctx ); // Verify offsets are unchanged. @@ -197,6 +198,7 @@ mod test { #[test] fn test_filter_with_gaps() { + let mut ctx = array_session().create_execution_ctx(); // ListView-specific: Test filtering with gaps in elements array. // // Logical list: [[1,2,3], [7,8,9], [11,12], [2,3], [8,9]] @@ -220,7 +222,8 @@ mod test { // Verify the entire elements array is preserved including gaps. assert_arrays_eq!( result_list.elements(), - PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]) + PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]), + &mut ctx ); // Verify offsets are unchanged. @@ -231,7 +234,8 @@ mod test { // Verify the lists still read correctly. assert_arrays_eq!( result_list.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([7i32, 8, 9]) + PrimitiveArray::from_iter([7i32, 8, 9]), + &mut ctx ); } @@ -324,7 +328,7 @@ mod test { let list0 = result_list.list_elements_at(0).unwrap(); assert_eq!( list0 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut SESSION.create_execution_ctx()) .unwrap() .as_primitive() .as_::() @@ -333,7 +337,7 @@ mod test { ); assert_eq!( list0 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut SESSION.create_execution_ctx()) .unwrap() .as_primitive() .as_::() diff --git a/vortex-array/src/arrays/filter/execute/struct_.rs b/vortex-array/src/arrays/filter/execute/struct_.rs index ee209b90fc3..8980c422ef3 100644 --- a/vortex-array/src/arrays/filter/execute/struct_.rs +++ b/vortex-array/src/arrays/filter/execute/struct_.rs @@ -49,6 +49,8 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -86,6 +88,7 @@ mod test { #[test] fn filter_struct_selects_correct_rows() { + let mut ctx = array_session().create_execution_ctx(); let array = StructArray::try_new( ["x", "y"].into(), vec![ @@ -111,11 +114,12 @@ mod test { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn filter_empty_struct() { + let mut ctx = array_session().create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let mask = Mask::from_iter([ @@ -125,18 +129,19 @@ mod test { let expected = StructArray::try_new(FieldNames::empty(), vec![], 5, Validity::NonNullable).unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn filter_empty_struct_with_empty_filter() { + let mut ctx = array_session().create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 0, Validity::NonNullable).unwrap(); let filtered = struct_arr.filter(Mask::from_iter::<[bool; 0]>([])).unwrap(); let expected = StructArray::try_new(FieldNames::empty(), vec![], 0, Validity::NonNullable).unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/filter/execute/take/tests.rs b/vortex-array/src/arrays/filter/execute/take/tests.rs index fd40865bf02..36ab1a63dcb 100644 --- a/vortex-array/src/arrays/filter/execute/take/tests.rs +++ b/vortex-array/src/arrays/filter/execute/take/tests.rs @@ -4,13 +4,15 @@ use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; -use vortex_session::VortexSession; use crate::IntoArray; use crate::RecursiveCanonical; +use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::DecimalArray; +use crate::arrays::Dict; use crate::arrays::DictArray; +use crate::arrays::Filter; use crate::arrays::FilterArray; use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; @@ -18,12 +20,28 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; +use crate::arrays::dict::TakeExecuteAdaptor; use crate::assert_arrays_eq; use crate::dtype::DecimalDType; use crate::dtype::FieldNames; use crate::executor::ExecutionCtx; +use crate::kernel::ExecuteParentKernel; use crate::validity::Validity; +fn execute_parent( + child: &crate::ArrayRef, + parent: &crate::ArrayRef, + child_idx: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + TakeExecuteAdaptor(Filter).execute_parent( + child.as_::(), + parent.as_::(), + child_idx, + ctx, + ) +} + #[test] fn test_take_execute_kernel_maps_indices_through_filter() -> VortexResult<()> { let filter = FilterArray::new( @@ -41,15 +59,15 @@ fn test_take_execute_kernel_maps_indices_through_filter() -> VortexResult<()> { filter.clone(), )? .into_array(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); assert_arrays_eq!( result.execute::(&mut ctx)?.0, - PrimitiveArray::from_option_iter([Some(40i32), None, Some(10)]).into_array() + PrimitiveArray::from_option_iter([Some(40i32), None, Some(10)]).into_array(), + &mut ctx ); Ok(()) } @@ -70,16 +88,16 @@ fn test_take_execute_kernel_nullable_fast_path_maps_indices_through_filter() -> filter.clone(), )? .into_array(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); assert!(result.as_opt::().is_some()); assert_arrays_eq!( result.execute::(&mut ctx)?.0, - PrimitiveArray::from_option_iter([Some(40i32), None, Some(20)]).into_array() + PrimitiveArray::from_option_iter([Some(40i32), None, Some(20)]).into_array(), + &mut ctx ); Ok(()) } @@ -92,16 +110,16 @@ fn test_take_execute_kernel_fast_path_maps_indices_through_filter() -> VortexRes ) .into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 3].into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); assert!(result.as_opt::().is_some()); assert_arrays_eq!( result.execute::(&mut ctx)?.0, - PrimitiveArray::from_iter([50i32, 20, 60]).into_array() + PrimitiveArray::from_iter([50i32, 20, 60]).into_array(), + &mut ctx ); Ok(()) } @@ -113,9 +131,9 @@ fn assert_take_execute_rejects_out_of_bounds_rank( ) -> VortexResult<()> { let filter = FilterArray::new(child, filter_mask).into_array(); let parent = DictArray::try_new(codes, filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); - if let Err(err) = filter.execute_parent(&parent, 1, &mut ctx) { + if let Err(err) = execute_parent(&filter, &parent, 1, &mut ctx) { assert!( err.to_string().contains("out of bounds"), "unexpected error: {err}" @@ -191,10 +209,9 @@ fn test_take_execute_kernel_handles_empty_sequential_take() -> VortexResult<()> filter.clone(), )? .into_array(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); assert_arrays_eq!( @@ -204,7 +221,8 @@ fn test_take_execute_kernel_handles_empty_sequential_take() -> VortexResult<()> buffer![0u32].into_array(), Validity::NonNullable, )? - .into_array() + .into_array(), + &mut ctx ); Ok(()) } @@ -216,13 +234,16 @@ fn assert_take_execute_maps_child_dtype( let filter = FilterArray::new(child, Mask::from_iter([true, false, true, true, false])).into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); - assert_arrays_eq!(result.execute::(&mut ctx)?.0, expected); + assert_arrays_eq!( + result.execute::(&mut ctx)?.0, + expected, + &mut ctx + ); Ok(()) } @@ -234,9 +255,9 @@ fn test_take_execute_kernel_skips_bool_filter_child() -> VortexResult<()> { ) .into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); - let result = filter.execute_parent(&parent, 1, &mut ctx)?; + let result = execute_parent(&filter, &parent, 1, &mut ctx)?; assert!(result.is_none()); Ok(()) @@ -256,9 +277,9 @@ fn execute_primitive_take( .into_array(); let indices = PrimitiveArray::from_iter((0..take_len).map(|idx| (idx % filtered_len) as u64)); let parent = DictArray::try_new(indices.into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); - filter.execute_parent(&parent, 1, &mut ctx) + execute_parent(&filter, &parent, 1, &mut ctx) } #[test] @@ -323,16 +344,17 @@ fn test_take_execute_kernel_handles_nullable_primitive_filter_child() -> VortexR ) .into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); - let result = filter.execute_parent(&parent, 1, &mut ctx)?; + let result = execute_parent(&filter, &parent, 1, &mut ctx)?; assert_arrays_eq!( result .expect("filter child should execute its take parent") .execute::(&mut ctx)? .0, - PrimitiveArray::from_option_iter([Some(40i32), Some(10), None]).into_array() + PrimitiveArray::from_option_iter([Some(40i32), Some(10), None]).into_array(), + &mut ctx ); Ok(()) } @@ -345,16 +367,16 @@ fn test_take_execute_kernel_preserves_nullable_all_valid_fixed_width_child() -> ) .into_array(); let parent = DictArray::try_new(buffer![0u64, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); assert_eq!(result.dtype(), parent.dtype()); assert_arrays_eq!( result.execute::(&mut ctx)?.0, - PrimitiveArray::new(buffer![10i32, 20], Validity::AllValid).into_array() + PrimitiveArray::new(buffer![10i32, 20], Validity::AllValid).into_array(), + &mut ctx ); Ok(()) } @@ -372,9 +394,9 @@ fn test_take_execute_kernel_handles_nullable_decimal_filter_child() -> VortexRes ) .into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); - let result = filter.execute_parent(&parent, 1, &mut ctx)?; + let result = execute_parent(&filter, &parent, 1, &mut ctx)?; assert_arrays_eq!( result @@ -382,7 +404,8 @@ fn test_take_execute_kernel_handles_nullable_decimal_filter_child() -> VortexRes .execute::(&mut ctx)? .0, DecimalArray::from_option_iter([Some(400i128), Some(100), None], decimal_dtype) - .into_array() + .into_array(), + &mut ctx ); Ok(()) } @@ -465,16 +488,16 @@ fn test_take_execute_kernel_preserves_nullable_indices_dtype_fast_path() -> Vort filter.clone(), )? .into_array(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); - let result = filter - .execute_parent(&parent, 1, &mut ctx)? + let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its nullable take parent"); assert_eq!(result.dtype(), parent.dtype()); assert_arrays_eq!( result.execute::(&mut ctx)?.0, - VarBinViewArray::from_iter_nullable_str([Some("a"), Some("b")]).into_array() + VarBinViewArray::from_iter_nullable_str([Some("a"), Some("b")]).into_array(), + &mut ctx ); Ok(()) } diff --git a/vortex-array/src/arrays/filter/kernel.rs b/vortex-array/src/arrays/filter/kernel.rs index 8c65fa99724..dbbe283138c 100644 --- a/vortex-array/src/arrays/filter/kernel.rs +++ b/vortex-array/src/arrays/filter/kernel.rs @@ -10,22 +10,27 @@ use vortex_error::VortexResult; use vortex_mask::Mask; +use vortex_session::VortexSession; use crate::ArrayRef; +use crate::ArrayVTable; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; use crate::array::ArrayView; use crate::array::VTable; +use crate::arrays::Dict; use crate::arrays::Filter; use crate::arrays::dict::TakeExecuteAdaptor; use crate::kernel::ExecuteParentKernel; -use crate::kernel::ParentKernelSet; use crate::matcher::Matcher; +use crate::optimizer::kernels::ArrayKernelsExt; use crate::optimizer::rules::ArrayParentReduceRule; -pub(super) const PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&TakeExecuteAdaptor(Filter))]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Dict.id(), Filter, TakeExecuteAdaptor(Filter)); +} pub trait FilterReduce: VTable { /// Filter an array with the provided mask without reading buffers. diff --git a/vortex-array/src/arrays/filter/mod.rs b/vortex-array/src/arrays/filter/mod.rs index 39859bc5b25..9510ca92f5b 100644 --- a/vortex-array/src/arrays/filter/mod.rs +++ b/vortex-array/src/arrays/filter/mod.rs @@ -19,3 +19,7 @@ mod rules; mod vtable; pub use vtable::Filter; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + kernel::initialize(session); +} diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index 5be6bc6ac35..0f54a97c574 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -16,6 +16,7 @@ use vortex_session::registry::CachedId; use crate::AnyCanonical; use crate::ArrayEq; use crate::ArrayHash; +use crate::ArrayParts; use crate::ArrayRef; use crate::Canonical; use crate::EqMode; @@ -26,13 +27,13 @@ use crate::array::ArrayView; 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::array::FilterData; use crate::arrays::filter::array::SLOT_NAMES; use crate::arrays::filter::execute::execute_filter; use crate::arrays::filter::execute::execute_filter_fast_paths; -use crate::arrays::filter::kernel::PARENT_KERNELS; use crate::arrays::filter::rules::PARENT_RULES; use crate::arrays::filter::rules::RULES; use crate::buffer::BufferHandle; @@ -118,6 +119,14 @@ impl VTable for Filter { 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 { SLOT_NAMES[idx].to_string() } @@ -139,7 +148,7 @@ impl VTable for Filter { _buffers: &[BufferHandle], _children: &dyn ArrayChildren, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult> { vortex_bail!("Filter array is not serializable") } @@ -171,15 +180,6 @@ impl VTable for Filter { PARENT_RULES.evaluate(array, parent, child_idx) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { RULES.evaluate(array) } diff --git a/vortex-array/src/arrays/fixed_size_list/mod.rs b/vortex-array/src/arrays/fixed_size_list/mod.rs index 5d55d22d880..3d90d38e6e9 100644 --- a/vortex-array/src/arrays/fixed_size_list/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/mod.rs @@ -12,5 +12,9 @@ pub(crate) mod compute; mod vtable; pub use vtable::FixedSizeList; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + #[cfg(test)] mod tests; 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 f910bd4a05f..961801094d8 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/basic.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/basic.rs @@ -6,8 +6,8 @@ use std::sync::Arc; use vortex_buffer::buffer; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::dtype::DType; @@ -40,19 +40,19 @@ fn test_basic_fixed_size_list() { let first_list = fsl.fixed_size_list_elements_at(0).unwrap(); assert_eq!( first_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 1i32.into() ); assert_eq!( first_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 2i32.into() ); assert_eq!( first_list - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(), 3i32.into() ); @@ -60,19 +60,19 @@ fn test_basic_fixed_size_list() { let second_list = fsl.fixed_size_list_elements_at(1).unwrap(); assert_eq!( second_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 4i32.into() ); assert_eq!( second_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 5i32.into() ); assert_eq!( second_list - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(), 6i32.into() ); @@ -80,19 +80,19 @@ fn test_basic_fixed_size_list() { let third_list = fsl.fixed_size_list_elements_at(2).unwrap(); assert_eq!( third_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 7i32.into() ); assert_eq!( third_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 8i32.into() ); assert_eq!( third_list - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(), 9i32.into() ); @@ -100,19 +100,19 @@ fn test_basic_fixed_size_list() { let fourth_list = fsl.fixed_size_list_elements_at(3).unwrap(); assert_eq!( fourth_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 10i32.into() ); assert_eq!( fourth_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 11i32.into() ); assert_eq!( fourth_list - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(), 12i32.into() ); @@ -128,7 +128,7 @@ fn test_scalar_at() { // First list: [1, 2, 3]. let first = fsl - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert_eq!( first, @@ -143,26 +143,26 @@ fn test_scalar_at() { let first_list = fsl.fixed_size_list_elements_at(0).unwrap(); assert_eq!( first_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 1i32.into() ); assert_eq!( first_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 2i32.into() ); assert_eq!( first_list - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(), 3i32.into() ); // Second list: [4, 5, 6]. let second = fsl - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(); assert_eq!( second, @@ -177,19 +177,19 @@ fn test_scalar_at() { let second_list = fsl.fixed_size_list_elements_at(1).unwrap(); assert_eq!( second_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 4i32.into() ); assert_eq!( second_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 5i32.into() ); assert_eq!( second_list - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(), 6i32.into() ); @@ -208,13 +208,13 @@ fn test_fixed_size_list_at() { assert_eq!(first_list.len(), list_size as usize); assert_eq!( first_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 1.0f64.into() ); assert_eq!( first_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 2.0f64.into() ); @@ -224,13 +224,13 @@ fn test_fixed_size_list_at() { assert_eq!(third_list.len(), list_size as usize); assert_eq!( third_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 5.0f64.into() ); assert_eq!( third_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 6.0f64.into() ); 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 39a7ada3af3..f609fd72405 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs @@ -6,8 +6,8 @@ use std::sync::Arc; use vortex_buffer::buffer; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; @@ -57,7 +57,7 @@ fn test_fsl_size_0_length_1_non_nullable() { // Get the single empty list. let scalar = fsl - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!scalar.is_null()); assert_eq!( @@ -85,7 +85,7 @@ fn test_fsl_size_0_huge_length_non_nullable() { // Spot check a few lists. let scalar_first = fsl - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!scalar_first.is_null()); assert_eq!( @@ -98,7 +98,7 @@ fn test_fsl_size_0_huge_length_non_nullable() { ); let scalar_middle = fsl - .execute_scalar(500_000_000_000, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(500_000_000_000, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!scalar_middle.is_null()); assert_eq!( @@ -111,7 +111,7 @@ fn test_fsl_size_0_huge_length_non_nullable() { ); let scalar_end = fsl - .execute_scalar(999_999_999_999, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(999_999_999_999, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!scalar_end.is_null()); assert_eq!( @@ -163,7 +163,7 @@ fn test_fsl_size_0_length_1_nullable_valid() { // Get the single empty list (should be valid). let scalar = fsl - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!scalar.is_null()); assert_eq!( @@ -188,7 +188,7 @@ fn test_fsl_size_0_length_1_nullable_null() { // The single list should be null. let scalar = fsl - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert!(scalar.is_null()); } @@ -215,7 +215,7 @@ fn test_fsl_size_0_length_10_nullable_mixed() { ]; for i in 0..len { let scalar = fsl - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(); if expected_valid[i] { assert!(!scalar.is_null()); @@ -256,7 +256,7 @@ fn test_fsl_size_0_nullable_elements() { // All lists should be empty but valid. for i in 0..len { let scalar = fsl - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!scalar.is_null()); } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs index 7efa87eedee..6f609120901 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs @@ -8,6 +8,8 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -51,6 +53,7 @@ fn test_filter_degenerate_list_size_zero( #[case] mask_values: Vec, #[case] expected_len: usize, ) { + let mut ctx = array_session().create_execution_ctx(); let new_validity = if matches!(validity, Validity::NonNullable) { Validity::NonNullable } else { @@ -74,11 +77,12 @@ fn test_filter_degenerate_list_size_zero( new_validity, expected_len, ); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_filter_with_nulls() { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4), Some(5), Some(6)]); let validity = Validity::from_iter([true, false, true]); @@ -97,15 +101,16 @@ fn test_filter_with_nulls() { let expected = FixedSizeListArray::new( expected_elements.into_array(), 2, - Validity::from_iter([true, true]), // Both selected lists are valid, but type is still nullable. + Validity::from_iter([true, true]), /* Both selected lists are valid, but type is still nullable. */ 2, ); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_filter_all_null_array() { + let mut ctx = array_session().create_execution_ctx(); // Create an array where all elements are null. let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let validity = Validity::AllInvalid; @@ -116,12 +121,13 @@ fn test_filter_all_null_array() { // Verify the result is an array of nulls. assert_eq!(filtered.len(), 2, "All-null FSL should produce 2 elements"); - assert_nth_scalar_is_null!(filtered, 0); - assert_nth_scalar_is_null!(filtered, 1); + assert_nth_scalar_is_null!(filtered, 0, &mut ctx); + assert_nth_scalar_is_null!(filtered, 1, &mut ctx); } #[test] fn test_filter_nested_fixed_size_lists() { + let mut ctx = array_session().create_execution_ctx(); // Create nested fixed-size lists: FSL>. // Inner lists are of size 2, outer lists are of size 3. // So we have 2 outer lists, each containing 3 inner lists, each containing 2 i32s. @@ -160,7 +166,7 @@ fn test_filter_nested_fixed_size_lists() { let expected_outer = FixedSizeListArray::new(expected_inner.into_array(), 3, Validity::NonNullable, 1); - assert_arrays_eq!(filtered, expected_outer); + assert_arrays_eq!(filtered, expected_outer, &mut ctx); } // Conformance tests using rstest for various array configurations. @@ -247,6 +253,7 @@ fn create_fsl_empty() -> ArrayRef { #[test] fn test_filter_all_null_various_list_sizes() { + let mut ctx = array_session().create_execution_ctx(); // Test filtering with all-null arrays of different list sizes. // The implementation returns ConstantArray only when validity_mask() is Mask::AllFalse. @@ -257,8 +264,8 @@ fn test_filter_all_null_various_list_sizes() { let filtered0 = fsl0.filter(mask0).unwrap(); assert_eq!(filtered0.len(), 2); // Check that all elements are null (might be ConstantArray or FixedSizeListArray). - assert_nth_scalar_is_null!(filtered0, 0); - assert_nth_scalar_is_null!(filtered0, 1); + assert_nth_scalar_is_null!(filtered0, 0, &mut ctx); + assert_nth_scalar_is_null!(filtered0, 1, &mut ctx); // Case 2: list_size == 1. let elements1 = buffer![1i32, 2, 3].into_array(); @@ -267,8 +274,8 @@ fn test_filter_all_null_various_list_sizes() { let filtered1 = fsl1.filter(mask1).unwrap(); assert_eq!(filtered1.len(), 2); // Check that all elements are null. - assert_nth_scalar_is_null!(filtered1, 0); - assert_nth_scalar_is_null!(filtered1, 1); + assert_nth_scalar_is_null!(filtered1, 0, &mut ctx); + assert_nth_scalar_is_null!(filtered1, 1, &mut ctx); // Case 3: list_size == 10 (large). let elements10 = buffer![0..50i32].into_array(); @@ -277,14 +284,15 @@ fn test_filter_all_null_various_list_sizes() { let filtered10 = fsl10.filter(mask10).unwrap(); assert_eq!(filtered10.len(), 5); // Check that all elements are null. - assert_nth_scalar_is_null!(filtered10, 0); - assert_nth_scalar_is_null!(filtered10, 4); + assert_nth_scalar_is_null!(filtered10, 0, &mut ctx); + assert_nth_scalar_is_null!(filtered10, 4, &mut ctx); } // Note: test_filter_to_empty_degenerate has been consolidated into test_filter_degenerate_list_size_zero above. #[test] fn test_mask_expansion_threshold_boundary() { + let mut ctx = array_session().create_execution_ctx(); // Test with list_size == 8 (the FSL_SPARSE_MASK_LIST_SIZE_THRESHOLD). let list_size = 8u32; let num_lists = 100; @@ -323,7 +331,7 @@ fn test_mask_expansion_threshold_boundary() { 3, ); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); // Test with list_size == 7 (just below threshold). let list_size_7 = 7u32; @@ -355,12 +363,13 @@ fn test_mask_expansion_threshold_boundary() { 3, ); - assert_arrays_eq!(filtered7, expected7); + assert_arrays_eq!(filtered7, expected7, &mut ctx); } // Test FSL-specific behavior with very large list sizes. #[test] fn test_filter_large_list_size() { + let mut ctx = array_session().create_execution_ctx(); // Test with list_size=100, which is significantly larger than typical use cases. let list_size = 100u32; let num_lists = 5; @@ -394,7 +403,7 @@ fn test_filter_large_list_size() { 3, ); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); // Test edge case: filter out all but one large list. let mask_single = Mask::from_iter([false, false, true, false, false]); @@ -414,5 +423,5 @@ fn test_filter_large_list_size() { 1, ); - assert_arrays_eq!(filtered_single, expected_single); + assert_arrays_eq!(filtered_single, expected_single, &mut ctx); } 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 8fca62c0097..44420d22cb3 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs @@ -6,10 +6,10 @@ use std::sync::Arc; use vortex_buffer::buffer; use crate::IntoArray; -use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -76,7 +76,7 @@ fn test_fsl_of_fsl_basic() { // The first outer list should contain 3 inner lists. // We can check by slicing and examining scalars. let first_scalar = outer_fsl - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!first_scalar.is_null()); @@ -92,13 +92,13 @@ fn test_fsl_of_fsl_basic() { .unwrap(); assert_eq!( inner_list_0 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 1i32.into() ); assert_eq!( inner_list_0 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 2i32.into() ); @@ -111,13 +111,13 @@ fn test_fsl_of_fsl_basic() { .unwrap(); assert_eq!( inner_list_1 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 3i32.into() ); assert_eq!( inner_list_1 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 4i32.into() ); @@ -130,13 +130,13 @@ fn test_fsl_of_fsl_basic() { .unwrap(); assert_eq!( inner_list_2 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 5i32.into() ); assert_eq!( inner_list_2 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 6i32.into() ); @@ -152,13 +152,13 @@ fn test_fsl_of_fsl_basic() { .unwrap(); assert_eq!( inner_list_0 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 7i32.into() ); assert_eq!( inner_list_0 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 8i32.into() ); @@ -171,13 +171,13 @@ fn test_fsl_of_fsl_basic() { .unwrap(); assert_eq!( inner_list_1 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 9i32.into() ); assert_eq!( inner_list_1 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 10i32.into() ); @@ -190,13 +190,13 @@ fn test_fsl_of_fsl_basic() { .unwrap(); assert_eq!( inner_list_2 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 11i32.into() ); assert_eq!( inner_list_2 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 12i32.into() ); @@ -246,7 +246,7 @@ fn test_fsl_of_fsl_with_nulls() { // First outer list is valid. assert!( !outer_fsl - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); @@ -254,7 +254,7 @@ fn test_fsl_of_fsl_with_nulls() { // Second outer list is null. assert!( outer_fsl - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); @@ -262,7 +262,7 @@ fn test_fsl_of_fsl_with_nulls() { // Third outer list is valid. assert!( !outer_fsl - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); @@ -326,13 +326,13 @@ fn test_deeply_nested_fsl() { .unwrap(); assert_eq!( level1_0_0 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 1i32.into() ); assert_eq!( level1_0_0 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 2i32.into() ); @@ -344,13 +344,13 @@ fn test_deeply_nested_fsl() { .unwrap(); assert_eq!( level1_0_1 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 3i32.into() ); assert_eq!( level1_0_1 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 4i32.into() ); @@ -363,13 +363,13 @@ fn test_deeply_nested_fsl() { .unwrap(); assert_eq!( level1_1_0 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 5i32.into() ); assert_eq!( level1_1_0 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 6i32.into() ); @@ -381,13 +381,13 @@ fn test_deeply_nested_fsl() { .unwrap(); assert_eq!( level1_1_1 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 7i32.into() ); assert_eq!( level1_1_1 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 8i32.into() ); diff --git a/vortex-array/src/arrays/fixed_size_list/tests/nullability.rs b/vortex-array/src/arrays/fixed_size_list/tests/nullability.rs index 20deef502ed..f5e5339e756 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/nullability.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/nullability.rs @@ -6,8 +6,8 @@ use std::sync::Arc; use vortex_buffer::buffer; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; @@ -33,7 +33,7 @@ fn test_nullable_fsl_with_nulls() { // First list is valid: [1, 2]. let first = fsl - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!first.is_null()); assert_eq!( @@ -49,26 +49,26 @@ fn test_nullable_fsl_with_nulls() { let first_list = fsl.fixed_size_list_elements_at(0).unwrap(); assert_eq!( first_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 1i32.into() ); assert_eq!( first_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 2i32.into() ); // Second list is null. let second = fsl - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(); assert!(second.is_null()); // Third list is valid: [5, 6]. let third = fsl - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!third.is_null()); assert_eq!( @@ -84,20 +84,20 @@ fn test_nullable_fsl_with_nulls() { let third_list = fsl.fixed_size_list_elements_at(2).unwrap(); assert_eq!( third_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 5i32.into() ); assert_eq!( third_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 6i32.into() ); // Fourth list is null. let fourth = fsl - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(3, &mut array_session().create_execution_ctx()) .unwrap(); assert!(fourth.is_null()); } @@ -123,7 +123,7 @@ fn test_nullable_elements_non_nullable_lists() { // First list: [Some(1), None, Some(3)]. let first = fsl - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!first.is_null()); assert_eq!( @@ -137,7 +137,7 @@ fn test_nullable_elements_non_nullable_lists() { // Second list: [Some(4), Some(5), None]. let second = fsl - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!second.is_null()); assert_eq!( @@ -165,7 +165,7 @@ fn test_nullable_elements_and_nullable_lists() { // First list is valid: [Some(10), None]. let first = fsl - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!first.is_null()); assert_eq!( @@ -181,26 +181,26 @@ fn test_nullable_elements_and_nullable_lists() { let first_list = fsl.fixed_size_list_elements_at(0).unwrap(); assert_eq!( first_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Some(10u16).into() ); assert_eq!( first_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), None::.into() ); // Second list is null (but elements would be [Some(20), Some(30)]). let second = fsl - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(); assert!(second.is_null()); // Third list is valid: [None, None]. let third = fsl - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!third.is_null()); assert_eq!( @@ -216,13 +216,13 @@ fn test_nullable_elements_and_nullable_lists() { let third_list = fsl.fixed_size_list_elements_at(2).unwrap(); assert_eq!( third_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), None::.into() ); assert_eq!( third_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), None::.into() ); @@ -243,7 +243,7 @@ fn test_alternating_nulls() { // Check alternating pattern. for i in 0..len { let scalar = fsl - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(); if i % 2 == 0 { assert!(!scalar.is_null()); @@ -275,7 +275,7 @@ fn test_validity_types() { let fsl = FixedSizeListArray::new(elements.clone(), list_size, Validity::AllInvalid, len); for i in 0..len { assert!( - fsl.execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + fsl.execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); @@ -293,22 +293,22 @@ fn test_validity_types() { ); assert!( - !fsl.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + !fsl.execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert!( - !fsl.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + !fsl.execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert!( - fsl.execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + fsl.execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert!( - !fsl.execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) + !fsl.execute_scalar(3, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); @@ -338,31 +338,31 @@ fn test_mixed_nullability_patterns() { // List 0: valid with [Some(1), None]. let list0 = fsl - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!list0.is_null()); // List 1: null. let list1 = fsl - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(); assert!(list1.is_null()); // List 2: valid with [Some(5), Some(6)]. let list2 = fsl - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!list2.is_null()); // List 3: valid with [Some(7), None]. let list3 = fsl - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(3, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!list3.is_null()); // List 4: valid with [None, Some(10)]. let list4 = fsl - .execute_scalar(4, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(4, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!list4.is_null()); } 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 8bccb80c8cb..e75139e5fc8 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -11,8 +11,8 @@ use super::common::create_nullable_fsl; use super::common::create_single_element_fsl; use crate::ArrayRef; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -40,6 +40,7 @@ fn test_take_fsl_conformance(#[case] fsl: FixedSizeListArray) { #[test] fn test_take_basic_smoke_test() { + 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); @@ -53,7 +54,7 @@ fn test_take_basic_smoke_test() { Validity::NonNullable, 3, ); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } // Parameterized test for FSL-specific degenerate (list_size=0) cases. @@ -97,7 +98,7 @@ fn test_take_degenerate_lists( for (i, expected_null) in expected_nulls.iter().enumerate() { assert_eq!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap() .is_null(), *expected_null @@ -107,6 +108,7 @@ fn test_take_degenerate_lists( #[test] fn test_take_large_list_size() { + let mut ctx = array_session().create_execution_ctx(); let elements = buffer![0i32..300].into_array(); let fsl = FixedSizeListArray::new(elements, 100, Validity::NonNullable, 3); @@ -116,11 +118,12 @@ fn test_take_large_list_size() { // Expected: [[200..300], [0..100]] let expected_elems = PrimitiveArray::from_iter((200i32..300).chain(0..100)).into_array(); let expected = FixedSizeListArray::new(expected_elems, 100, Validity::NonNullable, 2); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] fn test_take_fsl_with_null_indices_preserves_elements() { + 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); @@ -135,7 +138,7 @@ fn test_take_fsl_with_null_indices_preserves_elements() { Validity::from_iter([true, false, true]), 3, ); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } // Element index overflow: with u8 indices and list_size=16, data_idx=16 produces element index @@ -167,8 +170,9 @@ fn test_element_index_overflow( #[case] indices: ArrayRef, #[case] expected: FixedSizeListArray, ) { + let mut ctx = array_session().create_execution_ctx(); let result = fsl.take(indices).unwrap(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } // Parameterized test for nullable array scenarios that are specific to FSL's implementation. @@ -231,7 +235,7 @@ fn test_take_nullable_arrays_fsl_specific( for (i, expected_null) in expected_nulls.iter().enumerate() { assert_eq!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap() .is_null(), *expected_null diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs b/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs index 319751070de..47dc3ec0651 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs @@ -1,14 +1,27 @@ // 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::FixedSizeList; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -impl FixedSizeList { - pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(FixedSizeList)), - ParentKernelSet::lift(&TakeExecuteAdaptor(FixedSizeList)), - ]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Cast.id(), + FixedSizeList, + CastExecuteAdaptor(FixedSizeList), + ); + kernels.register_execute_parent_kernel( + Dict.id(), + FixedSizeList, + TakeExecuteAdaptor(FixedSizeList), + ); } 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 8d86e96a174..8f6fef6a909 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs @@ -15,6 +15,7 @@ use vortex_session::registry::CachedId; use crate::ArrayEq; use crate::ArrayHash; +use crate::ArrayParts; use crate::ArrayRef; use crate::EqMode; use crate::ExecutionCtx; @@ -23,6 +24,7 @@ use crate::array::Array; use crate::array::ArrayId; 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; @@ -39,6 +41,10 @@ mod validity; /// A [`FixedSizeList`]-encoded Vortex array. pub type FixedSizeListArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(Clone, Debug)] pub struct FixedSizeList; @@ -76,21 +82,20 @@ impl VTable for FixedSizeList { vortex_panic!("FixedSizeListArray buffer_name index {idx} out of bounds") } - fn reduce_parent( + fn with_buffers( + &self, array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) } - fn execute_parent( + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, child_idx: usize, - ctx: &mut ExecutionCtx, ) -> VortexResult> { - Self::PARENT_KERNELS.execute(array, parent, child_idx, ctx) + PARENT_RULES.evaluate(array, parent, child_idx) } fn serialize( @@ -150,7 +155,7 @@ impl VTable for FixedSizeList { buffers: &[BufferHandle], children: &dyn ArrayChildren, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult> { if !metadata.is_empty() { vortex_bail!( "FixedSizeListArray expects empty metadata, got {} bytes", @@ -186,7 +191,7 @@ impl VTable for FixedSizeList { let data = FixedSizeListData::try_build(elements.clone(), *list_size, validity.clone(), len)?; let slots = FixedSizeListData::make_slots(&elements, &validity, len); - Ok(crate::array::ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { diff --git a/vortex-array/src/arrays/interleave/execute/bool.rs b/vortex-array/src/arrays/interleave/execute/bool.rs index 755921abd59..fde5b161dfd 100644 --- a/vortex-array/src/arrays/interleave/execute/bool.rs +++ b/vortex-array/src/arrays/interleave/execute/bool.rs @@ -8,7 +8,6 @@ use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_mask::Mask; use super::super::Interleave; use super::super::InterleaveArrayExt; @@ -21,62 +20,44 @@ use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::match_each_unsigned_integer_ptype; use crate::require_child; -use crate::validity::Validity; /// Gathers `N` boolean values under unsigned `array_indices` / `row_indices` selectors, scattering -/// each selected bit (and its validity) into the output position it routes to. +/// each selected bit into the output position it routes to. pub(super) fn execute( array: Array, - ctx: &mut ExecutionCtx, + _ctx: &mut ExecutionCtx, ) -> VortexResult { let num_values = array.num_values(); - // Drive every value and both selectors to canonical encodings so we can operate on raw bits. + // Drive both selectors and every value to canonical encodings so we can operate on raw bits. let mut array = array; + 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 => Bool); + array = require_child!(array, array.value(i), i + 2 => Bool); } - array = require_child!(array, array.array_indices(), num_values => Primitive); - array = require_child!(array, array.row_indices(), num_values + 1 => Primitive); - let dtype = array.as_ref().dtype().clone(); - let len = array.as_ref().len(); - let nullable = dtype.is_nullable(); - - // Materialize each value's bits, and its validity mask only when the output can be null. + // Materialize each value's bits; the selectors gather one bit per output below. let mut value_bits = Vec::with_capacity(num_values); - let mut value_validity = Vec::with_capacity(num_values); for i in 0..num_values { - let value = array.value(i).as_::(); - let bits = value.to_bit_buffer(); - let validity = nullable - .then(|| value.validity()?.execute_mask(bits.len(), ctx)) - .transpose()?; - value_bits.push(bits); - value_validity.push(validity); + value_bits.push(array.value(i).as_::().to_bit_buffer()); } + let validity = array.as_ref().validity()?; + // Scatter directly from the typed selector buffers — no intermediate `usize` materialization. let array_indices = array.array_indices().as_::(); let row_indices = array.row_indices().as_::(); - let (values, validity) = match_each_unsigned_integer_ptype!(array_indices.ptype(), |A| { + let values = match_each_unsigned_integer_ptype!(array_indices.ptype(), |A| { match_each_unsigned_integer_ptype!(row_indices.ptype(), |R| { gather( - len, - num_values, &value_bits, - &value_validity, array_indices.as_slice::(), row_indices.as_slice::(), - nullable, )? }) }); - let validity = match validity { - Some(bits) => Validity::from(bits.freeze()), - None => Validity::NonNullable, - }; Ok(ExecutionResult::done(BoolArray::try_new( values.freeze(), validity, @@ -85,43 +66,70 @@ pub(super) fn execute( /// The scatter, monomorphized on the selector integer widths so each `(array_index, row_index)` /// pair is read straight from its packed buffer. -/// -/// Output bits (and validity) are produced with [`BitBufferMut::collect_bool`], which packs 64 -/// results per word: every output bit is written branchlessly, avoiding a per-row `set`/`unset` -/// (each of which would bounds-check and branch on the random bit value). -#[allow(clippy::too_many_arguments)] fn gather, R: AsPrimitive>( - len: usize, - num_values: usize, value_bits: &[BitBuffer], - value_validity: &[Option], branches: &[A], rows: &[R], - nullable: bool, -) -> VortexResult<(BitBufferMut, Option)> { - // Validate the per-row bounds once up front (returning an error rather than panicking), so the - // word-packing passes below are tight branchless loops. +) -> VortexResult { + let len = validate_selectors(value_bits, branches, rows)?; + + // SAFETY: `validate_selectors` proved `branches.len() == rows.len() == len`, and for every + // `i < len` that `branches[i] < value_bits.len()` and `rows[i] < value_bits[branches[i]].len()`. + Ok(unsafe { gather_bits(len, value_bits, branches, rows) }) +} + +/// Validates the per-row selector bounds, returning the output length (`branches.len()`). +/// +/// On success, `rows.len() == branches.len() == len` and, for every `i < len`, +/// `branches[i] < value_bits.len()` and `rows[i] < value_bits[branches[i]].len()` — exactly the +/// preconditions of [`gather_bits`]. Errors (rather than panics) on any out-of-bounds selector. +fn validate_selectors, R: AsPrimitive>( + value_bits: &[BitBuffer], + branches: &[A], + rows: &[R], +) -> VortexResult { + // The two selectors are validated to equal length at construction, which is the output length. + let len = branches.len(); + vortex_ensure!( + rows.len() == len, + "interleave selectors differ in length: array_indices {len}, row_indices {}", + rows.len() + ); + for i in 0..len { let branch = branches[i].as_(); - vortex_ensure!(branch < num_values, "interleave array index out of bounds"); + vortex_ensure!( + branch < value_bits.len(), + "interleave array index out of bounds" + ); vortex_ensure!( rows[i].as_() < value_bits[branch].len(), "interleave row index out of bounds" ); } - let values = - BitBufferMut::collect_bool(len, |i| value_bits[branches[i].as_()].value(rows[i].as_())); - - // A missing per-value mask means every row of that value is valid; only materialized when the - // output can be null. - let validity = nullable.then(|| { - BitBufferMut::collect_bool(len, |i| { - value_validity[branches[i].as_()] - .as_ref() - .is_none_or(|mask| mask.value(rows[i].as_())) - }) - }); + Ok(len) +} - Ok((values, validity)) +/// Gathers one bit per output from `bits[branches[i]]` at position `rows[i]`, packing 64 results per +/// word with [`BitBufferMut::collect_bool`]. +/// +/// The bounds-checked `BitBuffer::value` is slower still. +/// +/// # Safety +/// +/// `branches` and `rows` must both contain at least `len` elements. For every `i < len`, +/// `branches[i] < bits.len()` and `rows[i] < bits[branches[i]].len()`. +unsafe fn gather_bits, R: AsPrimitive>( + len: usize, + bits: &[BitBuffer], + branches: &[A], + rows: &[R], +) -> BitBufferMut { + // SAFETY: `collect_bool` calls this for `i < len`, and the caller guarantees `branches[i]` and + // `rows[i]` are in bounds for `bits` / the selected buffer. + BitBufferMut::collect_bool(len, |i| unsafe { + bits.get_unchecked(branches.get_unchecked(i).as_()) + .value_unchecked(rows.get_unchecked(i).as_()) + }) } diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index dd5619559dc..bff03ab055f 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -6,8 +6,8 @@ //! //! # Specification //! -//! An [`Interleave`] array has `N + 2` children: `N` *values* followed by an `array_indices` -//! selector and a `row_indices` selector. The output has `array_indices.len()` rows, and output +//! An [`Interleave`] array has `N + 2` children: an `array_indices` selector and a `row_indices` +//! selector followed by `N` *values*. The output has `array_indices.len()` rows, and output //! row `i` comes from `values[array_indices[i]][row_indices[i]]`. //! //! Unlike a `Merge`, which consumes each branch in order under a cursor, an [`Interleave`] is @@ -68,6 +68,7 @@ use crate::array::OperationsVTable; use crate::array::TypedArrayRef; use crate::array::VTable; use crate::array::ValidityVTable; +use crate::array::with_empty_buffers; use crate::arrays::ConstantArray; use crate::buffer::BufferHandle; use crate::dtype::DType; @@ -86,8 +87,8 @@ pub struct Interleave; /// Per-array metadata for an [`InterleaveArray`]. /// -/// The values and selectors live in the array's slots; only the value count is stored here so the -/// selector slots can be located (`slots[num_values]` and `slots[num_values + 1]`). +/// The selectors and values live in the array's slots; the selectors occupy `slots[0]` and +/// `slots[1]`, and the `num_values` values follow at `slots[2..2 + num_values]`. #[derive(Clone, Debug)] pub struct InterleaveData { pub(crate) num_values: usize, @@ -120,21 +121,21 @@ pub trait InterleaveArrayExt: TypedArrayRef { /// The `idx`-th value array (holding the rows that `array_indices` routes to it). fn value(&self, idx: usize) -> &ArrayRef { - self.as_ref().slots()[idx] + self.as_ref().slots()[idx + 2] .as_ref() .vortex_expect("validated interleave value slot") } /// The selector routing each output row to a value array. fn array_indices(&self) -> &ArrayRef { - self.as_ref().slots()[self.num_values] + self.as_ref().slots()[0] .as_ref() .vortex_expect("validated interleave array_indices slot") } /// The selector naming each output row's position within its value array. fn row_indices(&self) -> &ArrayRef { - self.as_ref().slots()[self.num_values + 1] + self.as_ref().slots()[1] .as_ref() .vortex_expect("validated interleave row_indices slot") } @@ -215,19 +216,68 @@ impl Array { row_indices: ArrayRef, ) -> VortexResult { let dtype = Interleave::check(&values, &array_indices, &row_indices)?; - let len = array_indices.len(); - let num_values = values.len(); + // SAFETY: `check` just validated every invariant and computed the matching `dtype`. + Ok(unsafe { Self::new_unchecked(values, array_indices, row_indices, dtype) }) + } - let mut slots: ArraySlots = values.into_iter().map(Some).collect(); + /// Constructs an [`InterleaveArray`] without re-validating the spec invariants. + /// + /// This is the assembly half of [`try_new`](Self::try_new): it lays the selectors into + /// `slots[0]`/`slots[1]` and the values into `slots[2..]` and stamps `dtype`, skipping the + /// `Interleave::check` pass. It exists for hot internal paths — notably building the + /// pushed-down validity interleave in [`ValidityVTable::validity`] — where the inputs are known + /// to satisfy the invariants by construction and re-checking them is pure overhead. + /// + /// # Safety + /// + /// The caller must uphold every [module invariant](self): at least two `values` sharing a dtype + /// up to nullability, both selectors non-nullable unsigned integers of equal length, and `dtype` + /// equal to the value returned by `Interleave::check` for these arguments (the shared value + /// type with the union of the values' nullabilities). + pub unsafe fn new_unchecked( + values: Vec, + array_indices: ArrayRef, + row_indices: ArrayRef, + dtype: DType, + ) -> Self { + let mut slots: ArraySlots = ArraySlots::with_capacity(values.len() + 2); slots.push(Some(array_indices)); slots.push(Some(row_indices)); + slots.extend(values.into_iter().map(Some)); + + // SAFETY: the caller of `new_unchecked` upholds every invariant; here we only assemble the + // canonical slot layout (`array_indices`, `row_indices`, then values) that follows. + unsafe { Self::new_unchecked_slots(slots, dtype) } + } + + /// Constructs an [`InterleaveArray`] from a pre-assembled `slots` buffer, skipping both the + /// spec re-check and the slot copy that [`new_unchecked`](Self::new_unchecked) performs. + /// + /// This is the lowest-level assembly path: it stamps `dtype` onto `slots` as-is. Callers that + /// already own a correctly laid-out [`ArraySlots`] — for example reusing a validated parent's + /// selectors while swapping its values — avoid materializing an intermediate `Vec` + /// and re-pushing it into a fresh `ArraySlots`. + /// + /// # Safety + /// + /// The caller must uphold every [module invariant](self) *and* the slot layout: `slots[0]` is + /// the `array_indices` selector, `slots[1]` is the `row_indices` selector, and `slots[2..]` are + /// the value arrays. Every slot must be `Some`, there must be at least two values + /// (`slots.len() >= 4`), and `dtype` must equal the value returned by `Interleave::check` for + /// these arguments. + pub unsafe fn new_unchecked_slots(slots: ArraySlots, dtype: DType) -> Self { + let num_values = slots.len() - 2; + let len = slots[0] + .as_ref() + .vortex_expect("interleave array_indices slot present") + .len(); - Ok(unsafe { + unsafe { Array::from_parts_unchecked( ArrayParts::new(Interleave, dtype, len, InterleaveData { num_values }) .with_slots(slots), ) - }) + } } } @@ -259,16 +309,14 @@ impl VTable for Interleave { "InterleaveArray slots must all be present" ); - let values: Vec = slots[..data.num_values] + let array_indices = slots[0] + .clone() + .vortex_expect("validated array_indices slot"); + let row_indices = slots[1].clone().vortex_expect("validated row_indices slot"); + let values: Vec = slots[2..] .iter() .map(|s| s.clone().vortex_expect("validated value slot")) .collect(); - let array_indices = slots[data.num_values] - .clone() - .vortex_expect("validated array_indices slot"); - let row_indices = slots[data.num_values + 1] - .clone() - .vortex_expect("validated row_indices slot"); // All semantic invariants live in `check`; here we only confirm the array's cached `dtype` // and `len` agree with what the children imply. @@ -300,13 +348,19 @@ impl VTable for Interleave { None } - fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { - if idx == array.num_values() { - "array_indices".to_string() - } else if idx == array.num_values() + 1 { - "row_indices".to_string() - } else { - format!("value_{idx}") + 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 => "array_indices".to_string(), + 1 => "row_indices".to_string(), + _ => format!("value_{}", idx - 2), } } @@ -370,18 +424,18 @@ impl ValidityVTable for Interleave { if !array.as_ref().dtype().is_nullable() { return Ok(Validity::NonNullable); } - // The output validity is itself an interleave — by the same selectors — of the values' - // validities, expressed as non-nullable boolean arrays. This bottoms out immediately - // because the inner interleave is non-nullable. - let mut value_validities: Vec = Vec::with_capacity(array.num_values()); - for i in 0..array.num_values() { - value_validities.push(value_validity_array(array.value(i))?); + let num_values = array.num_values(); + let mut slots: ArraySlots = ArraySlots::with_capacity(num_values + 2); + slots.push(Some(array.array_indices().clone())); + slots.push(Some(array.row_indices().clone())); + for i in 0..num_values { + slots.push(Some(value_validity_array(array.value(i))?)); } - let interleaved = InterleaveArray::try_new( - value_validities, - array.array_indices().clone(), - array.row_indices().clone(), - )?; + // SAFETY: `value_validity_array` yields a non-nullable boolean array per value, + // the selectors already validated. + let interleaved = unsafe { + InterleaveArray::new_unchecked_slots(slots, DType::Bool(Nullability::NonNullable)) + }; Ok(Validity::Array(interleaved.into_array())) } } @@ -404,8 +458,8 @@ mod tests { use super::*; use crate::Canonical; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -496,10 +550,10 @@ mod tests { InterleaveArray::try_new(values.clone(), array_indices.clone(), row_indices.clone())? .into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let reference = interleave_reference(&values, &array_indices, &row_indices, &mut ctx)?; - assert_arrays_eq!(interleaved, reference); + assert_arrays_eq!(interleaved, reference, &mut ctx); Ok(()) } @@ -524,6 +578,29 @@ mod tests { ) } + #[test] + fn interleave_binary_spans_word_boundary() -> VortexResult<()> { + // Exercises a two-value gather across more than one 64-bit packing word, with out-of-order + // routing into both branches so neither buffer is consumed contiguously. + let branch0: Vec> = (0..100).map(|i| Some(i % 3 == 0)).collect(); + let branch1: Vec> = (0..100).map(|i| Some(i % 5 == 0)).collect(); + let indices: Vec<(usize, usize)> = (0..200).map(|i| (i % 2, (i * 7) % 100)).collect(); + check(&[&branch0, &branch1], &indices) + } + + #[test] + fn interleave_binary_nullable_spans_word_boundary() -> VortexResult<()> { + // Same two-value gather, now with nulls so the validity packing is exercised too. + let branch0: Vec> = (0..70) + .map(|i| (i % 4 != 0).then_some(i % 2 == 0)) + .collect(); + let branch1: Vec> = (0..70) + .map(|i| (i % 3 != 0).then_some(i % 2 == 1)) + .collect(); + let indices: Vec<(usize, usize)> = (0..150).map(|i| (i % 2, (i * 11) % 70)).collect(); + check(&[&branch0, &branch1], &indices) + } + #[test] fn interleave_three_values() -> VortexResult<()> { // An unsigned `array_indices` routes among three values with full random access. @@ -602,6 +679,45 @@ mod tests { assert!(err.to_string().contains("equal length"), "{err}"); } + #[test] + fn execute_rejects_out_of_bounds_array_index() -> VortexResult<()> { + let value = BoolArray::from_iter([true]).into_array(); + let array_indices = PrimitiveArray::from_iter([2u32]).into_array(); + let row_indices = PrimitiveArray::from_iter([0u32]).into_array(); + let interleaved = + InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)? + .into_array(); + + let mut ctx = array_session().create_execution_ctx(); + let err = interleaved + .execute::(&mut ctx) + .err() + .vortex_expect("expected execution to reject out-of-bounds array index"); + assert!( + err.to_string().contains("array index out of bounds"), + "{err}" + ); + Ok(()) + } + + #[test] + fn execute_rejects_out_of_bounds_row_index() -> VortexResult<()> { + let value = BoolArray::from_iter([true]).into_array(); + let array_indices = PrimitiveArray::from_iter([0u32]).into_array(); + let row_indices = PrimitiveArray::from_iter([1u32]).into_array(); + let interleaved = + InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)? + .into_array(); + + let mut ctx = array_session().create_execution_ctx(); + let err = interleaved + .execute::(&mut ctx) + .err() + .vortex_expect("expected execution to reject out-of-bounds row index"); + assert!(err.to_string().contains("row index out of bounds"), "{err}"); + Ok(()) + } + #[test] #[should_panic(expected = "only implemented for boolean values")] fn non_boolean_value_execution_panics() { @@ -613,7 +729,7 @@ mod tests { let interleaved = InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices) .vortex_expect("primitive values should construct") .into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); interleaved.execute::(&mut ctx).ok(); } } diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index 7788511aa25..945991b7e03 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -20,6 +20,7 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::min_max; use crate::array::Array; use crate::array::ArrayParts; @@ -212,7 +213,7 @@ impl ListData { // Validate that offsets min is non-negative, and max does not exceed the length of // the elements array. - if let Some(min_max) = min_max(offsets, &mut ctx)? { + if let Some(min_max) = min_max(offsets, &mut ctx, NumericalAggregateOpts::default())? { match_each_integer_ptype!(offsets_ptype, |P| { #[allow(clippy::absurd_extreme_comparisons, unused_comparisons)] { diff --git a/vortex-array/src/arrays/list/compute/cast.rs b/vortex-array/src/arrays/list/compute/cast.rs index e6d71f6e80c..084d9613e75 100644 --- a/vortex-array/src/arrays/list/compute/cast.rs +++ b/vortex-array/src/arrays/list/compute/cast.rs @@ -66,13 +66,11 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; - use vortex_array::session::ArraySession; use vortex_buffer::buffer; use vortex_session::VortexSession; use crate::Canonical; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::RecursiveCanonical; use crate::VortexSessionExecute; use crate::arrays::BoolArray; @@ -86,8 +84,7 @@ mod tests { use crate::dtype::PType; use crate::validity::Validity; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[test] fn test_cast_list_success() { @@ -171,7 +168,7 @@ mod tests { ); let result = list.into_array().cast(target_dtype).and_then(|a| { - a.execute::(&mut LEGACY_SESSION.create_execution_ctx()) + a.execute::(&mut SESSION.create_execution_ctx()) .map(|c| c.0.into_array()) }); assert!(result.is_err()); diff --git a/vortex-array/src/arrays/list/compute/kernels.rs b/vortex-array/src/arrays/list/compute/kernels.rs index 188c83c1bf5..6af858dbf49 100644 --- a/vortex-array/src/arrays/list/compute/kernels.rs +++ b/vortex-array/src/arrays/list/compute/kernels.rs @@ -1,14 +1,22 @@ // 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::List; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(List)), - ParentKernelSet::lift(&FilterExecuteAdaptor(List)), - ParentKernelSet::lift(&TakeExecuteAdaptor(List)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), List, CastExecuteAdaptor(List)); + kernels.register_execute_parent_kernel(Filter.id(), List, FilterExecuteAdaptor(List)); + kernels.register_execute_parent_kernel(Dict.id(), List, TakeExecuteAdaptor(List)); +} diff --git a/vortex-array/src/arrays/list/compute/mod.rs b/vortex-array/src/arrays/list/compute/mod.rs index 65bacbebdf5..d8ddbd37583 100644 --- a/vortex-array/src/arrays/list/compute/mod.rs +++ b/vortex-array/src/arrays/list/compute/mod.rs @@ -9,7 +9,9 @@ pub(crate) mod rules; mod slice; mod take; -pub(crate) use kernels::PARENT_KERNELS; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + kernels::initialize(session); +} #[cfg(test)] mod tests { @@ -17,6 +19,8 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; @@ -76,6 +80,9 @@ mod tests { Validity::NonNullable, ).unwrap())] fn test_list_consistency(#[case] array: ListArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index f96cac5cf32..2c6089e0c58 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -212,10 +212,10 @@ mod test { use vortex_buffer::buffer; use crate::IntoArray as _; - use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; @@ -258,12 +258,12 @@ mod test { assert!( result - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(0, &mut array_session().create_execution_ctx()) .unwrap() ); assert_eq!( result - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::list( Arc::clone(&element_dtype), @@ -274,18 +274,18 @@ mod test { assert!( result - .is_invalid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_invalid(1, &mut array_session().create_execution_ctx()) .unwrap() ); assert!( result - .is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(2, &mut array_session().create_execution_ctx()) .unwrap() ); assert_eq!( result - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::list( Arc::clone(&element_dtype), @@ -296,12 +296,12 @@ mod test { assert!( result - .is_valid(3, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(3, &mut array_session().create_execution_ctx()) .unwrap() ); assert_eq!( result - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(3, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::list(element_dtype, vec![], Nullability::Nullable) ); @@ -361,12 +361,12 @@ mod test { assert!( result - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(0, &mut array_session().create_execution_ctx()) .unwrap() ); assert_eq!( result - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::list( Arc::clone(&element_dtype), @@ -377,12 +377,12 @@ mod test { assert!( result - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(1, &mut array_session().create_execution_ctx()) .unwrap() ); assert_eq!( result - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::list( Arc::clone(&element_dtype), @@ -393,12 +393,12 @@ mod test { assert!( result - .is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(2, &mut array_session().create_execution_ctx()) .unwrap() ); assert_eq!( result - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::list(element_dtype, vec![], Nullability::NonNullable) ); @@ -488,12 +488,12 @@ mod test { assert_eq!(result_view.len(), 2); assert!( result_view - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(0, &mut array_session().create_execution_ctx()) .unwrap() ); assert!( result_view - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(1, &mut array_session().create_execution_ctx()) .unwrap() ); } @@ -518,17 +518,17 @@ mod test { assert_eq!(result_view.len(), 3); assert!( result_view - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(0, &mut array_session().create_execution_ctx()) .unwrap() ); assert!( result_view - .is_invalid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_invalid(1, &mut array_session().create_execution_ctx()) .unwrap() ); assert!( result_view - .is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(2, &mut array_session().create_execution_ctx()) .unwrap() ); } diff --git a/vortex-array/src/arrays/list/mod.rs b/vortex-array/src/arrays/list/mod.rs index c884f343a10..422a9523ff9 100644 --- a/vortex-array/src/arrays/list/mod.rs +++ b/vortex-array/src/arrays/list/mod.rs @@ -12,6 +12,10 @@ pub(crate) mod compute; mod vtable; pub use vtable::List; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + compute::initialize(session); +} + #[cfg(feature = "_test-harness")] mod test_harness; diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 2ed260a5f55..906f4e03588 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -15,6 +15,7 @@ use super::*; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::FilterArray; use crate::arrays::List; use crate::arrays::PrimitiveArray; @@ -25,12 +26,10 @@ use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType::I32; use crate::scalar::Scalar; -use crate::session::ArraySession; use crate::validity::Validity; /// A shared session for `List` tests, used to create execution contexts. -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(array_session); #[test] fn test_empty_list_array() { @@ -115,6 +114,7 @@ fn test_simple_list_filter() { #[test] fn test_list_filter_dense_mask() { + let mut ctx = SESSION.create_execution_ctx(); // Test filtering with a dense mask (high density of true values). let elements = buffer![0..100].into_array(); let offsets = buffer![0, 10, 25, 40, 60, 85, 100].into_array(); @@ -141,11 +141,12 @@ fn test_list_filter_dense_mask() { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_list_filter_sparse_mask() { + let mut ctx = SESSION.create_execution_ctx(); // Test filtering with a sparse mask (low density of true values). let elements = buffer![0..100].into_array(); let offsets = buffer![0, 10, 25, 40, 60, 85, 100].into_array(); @@ -174,11 +175,12 @@ fn test_list_filter_sparse_mask() { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_list_filter_empty_lists() { + let mut ctx = SESSION.create_execution_ctx(); // Test filtering arrays that contain empty lists. let elements = buffer![0..10].into_array(); let offsets = buffer![0, 0, 3, 3, 7, 10, 10].into_array(); // Lists at indices 0, 2, 5 are empty. @@ -202,7 +204,7 @@ fn test_list_filter_empty_lists() { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] @@ -235,6 +237,7 @@ fn test_list_filter_with_nulls() { #[test] fn test_list_filter_all_true() { + let mut ctx = SESSION.create_execution_ctx(); // Test filtering with an all-true mask. let elements = buffer![0..20].into_array(); let offsets = buffer![0, 5, 10, 15, 20].into_array(); @@ -252,7 +255,7 @@ fn test_list_filter_all_true() { assert_eq!(filtered.len(), 4); let expected = ListArray::try_new(elements, offsets, validity).unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] @@ -277,6 +280,7 @@ fn test_list_filter_all_false() { #[test] fn test_list_filter_single_element() { + let mut ctx = SESSION.create_execution_ctx(); // Test filtering to keep only one element. let elements = buffer![0..50].into_array(); let offsets = buffer![0, 10, 20, 30, 40, 50].into_array(); @@ -300,11 +304,12 @@ fn test_list_filter_single_element() { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_list_filter_alternating_pattern() { + let mut ctx = SESSION.create_execution_ctx(); // Test filtering with an alternating pattern. let elements = buffer![0..60].into_array(); let offsets = buffer![0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60].into_array(); @@ -335,11 +340,12 @@ fn test_list_filter_alternating_pattern() { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_list_filter_variable_sizes() { + let mut ctx = SESSION.create_execution_ctx(); // Test filtering lists with highly variable sizes. let elements = buffer![0..100].into_array(); let offsets = buffer![0, 1, 2, 5, 10, 20, 35, 60, 100].into_array(); @@ -375,7 +381,7 @@ fn test_list_filter_variable_sizes() { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] @@ -576,6 +582,7 @@ fn create_list_of_lists_nullable(data: OptVec>>) -> ListArray #[test] #[expect(clippy::cognitive_complexity)] fn test_list_of_lists() { + let mut ctx = SESSION.create_execution_ctx(); let data = vec![ Some(vec![Some(vec![Some(1), Some(2)]), Some(vec![Some(3)])]), Some(vec![Some(vec![Some(4), Some(5), Some(6)])]), @@ -606,11 +613,11 @@ fn test_list_of_lists() { // Check first inner list [1, 2]. let first_inner = first_outer_list.list_elements_at(0).unwrap(); - assert_arrays_eq!(first_inner, PrimitiveArray::from_iter([1, 2])); + assert_arrays_eq!(first_inner, PrimitiveArray::from_iter([1, 2]), &mut ctx); // Check second inner list [3]. let second_inner = first_outer_list.list_elements_at(1).unwrap(); - assert_arrays_eq!(second_inner, PrimitiveArray::from_iter([3])); + assert_arrays_eq!(second_inner, PrimitiveArray::from_iter([3]), &mut ctx); // Check the second list of lists [[4, 5, 6]]. let second_outer = list_of_lists.list_elements_at(1).unwrap(); @@ -618,7 +625,7 @@ fn test_list_of_lists() { assert_eq!(second_outer_list.len(), 1); let inner = second_outer_list.list_elements_at(0).unwrap(); - assert_arrays_eq!(inner, PrimitiveArray::from_iter([4, 5, 6])); + assert_arrays_eq!(inner, PrimitiveArray::from_iter([4, 5, 6]), &mut ctx); // Check the third list of lists (empty). let third_outer = list_of_lists.list_elements_at(2).unwrap(); @@ -631,7 +638,7 @@ fn test_list_of_lists() { assert_eq!(fourth_outer_list.len(), 1); let inner = fourth_outer_list.list_elements_at(0).unwrap(); - assert_arrays_eq!(inner, PrimitiveArray::from_iter([7])); + assert_arrays_eq!(inner, PrimitiveArray::from_iter([7]), &mut ctx); // Test scalar conversion. let scalar = list_of_lists @@ -794,7 +801,7 @@ fn test_list_of_lists_both_nullable() { let third_list = third_outer.as_::(); assert_eq!(third_list.len(), 1); let inner = third_list.list_elements_at(0).unwrap(); - assert_arrays_eq!(inner, PrimitiveArray::from_iter([3])); + assert_arrays_eq!(inner, PrimitiveArray::from_iter([3]), &mut ctx); // Fourth outer list should have a null inner list. let fourth_outer = list_of_lists.list_elements_at(3).unwrap(); diff --git a/vortex-array/src/arrays/list/vtable/mod.rs b/vortex-array/src/arrays/list/vtable/mod.rs index d93c4673302..79316e76be7 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -25,13 +25,13 @@ use crate::array::ArrayId; 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::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::compute::PARENT_KERNELS; use crate::arrays::list::compute::rules::PARENT_RULES; use crate::arrays::listview::list_view_from_list; use crate::buffer::BufferHandle; @@ -85,6 +85,14 @@ impl VTable for List { vortex_panic!("ListArray buffer_name index {idx} out of bounds") } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, @@ -191,15 +199,6 @@ impl VTable for List { list_view_from_list(array, ctx)?.into_array(), )) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[derive(Clone, Debug)] diff --git a/vortex-array/src/arrays/listview/array.rs b/vortex-array/src/arrays/listview/array.rs index e5b6f9aaf99..1d6297ba7c1 100644 --- a/vortex-array/src/arrays/listview/array.rs +++ b/vortex-array/src/arrays/listview/array.rs @@ -22,6 +22,7 @@ use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::min_max; use crate::array::Array; use crate::array::ArrayParts; @@ -567,12 +568,16 @@ pub trait ListViewArrayExt: TypedArrayRef { }); let offsets = self.offsets().cast(wide_dtype.clone())?; let sizes = self.sizes().cast(wide_dtype)?; - let end = min_max(&offsets.binary(sizes, Operator::Add)?, ctx)? - .vortex_expect("non-empty array must report a min/max") - .max - .as_primitive() - .as_::() - .vortex_expect("max `offset + size` must fit in a usize"); + let end = min_max( + &offsets.binary(sizes, Operator::Add)?, + ctx, + NumericalAggregateOpts::default(), + )? + .vortex_expect("non-empty array must report a min/max") + .max + .as_primitive() + .as_::() + .vortex_expect("max `offset + size` must fit in a usize"); Ok((start, end)) } diff --git a/vortex-array/src/arrays/listview/compute/zip.rs b/vortex-array/src/arrays/listview/compute/zip.rs index 1423a14804c..46537f1fafe 100644 --- a/vortex-array/src/arrays/listview/compute/zip.rs +++ b/vortex-array/src/arrays/listview/compute/zip.rs @@ -170,8 +170,8 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; @@ -216,7 +216,7 @@ mod tests { ); let mask = Mask::from_iter([true, false, true]); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -232,7 +232,7 @@ mod tests { buffer![2u32, 2, 3].into_array(), Validity::NonNullable, ); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -256,7 +256,7 @@ mod tests { // true -> if_true, false -> if_false let mask = Mask::from_iter([false, true, true]); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -269,7 +269,7 @@ mod tests { buffer![1u32, 0, 1].into_array(), Validity::Array(BoolArray::from_iter([true, false, true]).into_array()), ); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -293,7 +293,7 @@ mod tests { ); let mask = Mask::from_iter([true, true, false]); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -307,7 +307,7 @@ mod tests { buffer![2u32, 1, 2].into_array(), Validity::AllValid, ); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -336,7 +336,7 @@ mod tests { ); let mask = Mask::from_iter([true, false]); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -362,7 +362,7 @@ mod tests { buffer![2u32, 1].into_array(), Validity::NonNullable, ); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index 7b553041c68..e85567a9ed9 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -331,37 +331,46 @@ mod tests { // Verify structure. assert_eq!(list_view.len(), 4); - assert_arrays_eq!(elements, list_view.elements().clone()); + assert_arrays_eq!(elements, list_view.elements().clone(), &mut ctx); // Verify offsets (should be same but without last element). let expected_offsets = buffer![0u32, 3, 5, 7].into_array(); - assert_arrays_eq!(expected_offsets, list_view.offsets().clone()); + assert_arrays_eq!(expected_offsets, list_view.offsets().clone(), &mut ctx); // Verify sizes. let expected_sizes = buffer![3u32, 2, 2, 3].into_array(); - assert_arrays_eq!(expected_sizes, list_view.sizes().clone()); + assert_arrays_eq!(expected_sizes, list_view.sizes().clone(), &mut ctx); // Verify data integrity. - assert_arrays_eq!(list_array, list_view); + assert_arrays_eq!(list_array, list_view, &mut ctx); Ok(()) } #[test] fn test_listview_to_list_zero_copy() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let list_view = create_basic_listview(); let list_array = list_from_list_view(list_view.clone(), &mut SESSION.create_execution_ctx())?; // Should have same elements. - assert_arrays_eq!(list_view.elements().clone(), list_array.elements().clone()); + assert_arrays_eq!( + list_view.elements().clone(), + list_array.elements().clone(), + &mut ctx + ); // ListArray offsets should have n+1 elements for n lists (add the final offset). // Check that the first n offsets match. let list_array_offsets_without_last = list_array.offsets().slice(0..list_view.len())?; - assert_arrays_eq!(list_view.offsets().clone(), list_array_offsets_without_last); + assert_arrays_eq!( + list_view.offsets().clone(), + list_array_offsets_without_last, + &mut ctx + ); // Verify data integrity. - assert_arrays_eq!(list_view, list_array); + assert_arrays_eq!(list_view, list_array, &mut ctx); Ok(()) } @@ -409,12 +418,13 @@ mod tests { // Round-trip conversion. let converted_back = list_from_list_view(nullable_list_view, &mut ctx)?; - assert_arrays_eq!(nullable_list, converted_back); + assert_arrays_eq!(nullable_list, converted_back, &mut ctx); Ok(()) } #[test] fn test_non_zero_copy_listview_to_list() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // Create ListViewArray with overlapping lists (not zero-copyable). let list_view = create_overlapping_listview(); let list_array = @@ -428,7 +438,7 @@ mod tests { } // The data should still be correct even though it required a rebuild. - assert_arrays_eq!(list_view, list_array); + assert_arrays_eq!(list_view, list_array, &mut ctx); Ok(()) } @@ -448,7 +458,7 @@ mod tests { // Round-trip. let converted_back = list_view_from_list(list_array, &mut ctx)?; - assert_arrays_eq!(empty_lists_view, converted_back); + assert_arrays_eq!(empty_lists_view, converted_back, &mut ctx); Ok(()) } @@ -474,8 +484,8 @@ mod tests { assert_eq!(list_view_i64.sizes().dtype(), i64_offsets.dtype()); // Verify data integrity. - assert_arrays_eq!(list_i32, list_view_i32); - assert_arrays_eq!(list_i64, list_view_i64); + assert_arrays_eq!(list_i32, list_view_i32, &mut ctx); + assert_arrays_eq!(list_i64, list_view_i64, &mut ctx); Ok(()) } @@ -487,20 +497,20 @@ mod tests { let original = create_basic_listview(); let to_list = list_from_list_view(original.clone(), &mut ctx)?; let back_to_view = list_view_from_list(to_list, &mut ctx)?; - assert_arrays_eq!(original, back_to_view); + assert_arrays_eq!(original, back_to_view, &mut ctx); // Test 2: Nullable round-trip. let nullable = create_nullable_listview(); let nullable_to_list = list_from_list_view(nullable.clone(), &mut ctx)?; let nullable_back = list_view_from_list(nullable_to_list, &mut ctx)?; - assert_arrays_eq!(nullable, nullable_back); + assert_arrays_eq!(nullable, nullable_back, &mut ctx); // Test 3: Non-zero-copyable round-trip. let overlapping = create_overlapping_listview(); let overlapping_to_list = list_from_list_view(overlapping.clone(), &mut ctx)?; let overlapping_back = list_view_from_list(overlapping_to_list, &mut ctx)?; - assert_arrays_eq!(overlapping, overlapping_back); + assert_arrays_eq!(overlapping, overlapping_back, &mut ctx); Ok(()) } @@ -517,11 +527,11 @@ mod tests { // Verify sizes are all 1. let expected_sizes = buffer![1u32, 1, 1].into_array(); - assert_arrays_eq!(expected_sizes, list_view.sizes().clone()); + assert_arrays_eq!(expected_sizes, list_view.sizes().clone(), &mut ctx); // Round-trip. let converted_back = list_from_list_view(list_view, &mut ctx)?; - assert_arrays_eq!(single_elem_list, converted_back); + assert_arrays_eq!(single_elem_list, converted_back, &mut ctx); Ok(()) } @@ -538,16 +548,17 @@ mod tests { // Verify sizes. let expected_sizes = buffer![2u32, 0, 1, 0, 3].into_array(); - assert_arrays_eq!(expected_sizes, list_view.sizes().clone()); + assert_arrays_eq!(expected_sizes, list_view.sizes().clone(), &mut ctx); // Round-trip. let converted_back = list_from_list_view(list_view, &mut ctx)?; - assert_arrays_eq!(mixed_list, converted_back); + assert_arrays_eq!(mixed_list, converted_back, &mut ctx); Ok(()) } #[test] fn test_recursive_simple_listview() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let list_view = create_basic_listview(); let result = recursive_list_from_list_view( list_view.clone().into_array(), @@ -555,12 +566,13 @@ mod tests { )?; assert_eq!(result.len(), list_view.len()); - assert_arrays_eq!(list_view.into_array(), result); + assert_arrays_eq!(list_view.into_array(), result, &mut ctx); Ok(()) } #[test] fn test_recursive_nested_listview() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let inner_elements = buffer![1i32, 2, 3].into_array(); let inner_offsets = buffer![0u32, 2].into_array(); let inner_sizes = buffer![2u32, 1].into_array(); @@ -592,12 +604,13 @@ mod tests { )?; assert_eq!(result.len(), 2); - assert_arrays_eq!(outer_listview.into_array(), result); + assert_arrays_eq!(outer_listview.into_array(), result, &mut ctx); Ok(()) } #[test] fn test_recursive_struct_with_listview_fields() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let listview_field = create_basic_listview().into_array(); let primitive_field = buffer![10i32, 20, 30, 40].into_array(); @@ -614,12 +627,13 @@ mod tests { )?; assert_eq!(result.len(), 4); - assert_arrays_eq!(struct_array.into_array(), result); + assert_arrays_eq!(struct_array.into_array(), result, &mut ctx); Ok(()) } #[test] fn test_recursive_fixed_size_list_with_listview_elements() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let lv1_elements = buffer![1i32, 2].into_array(); let lv1_offsets = buffer![0u32].into_array(); let lv1_sizes = buffer![2u32].into_array(); @@ -659,12 +673,13 @@ mod tests { )?; assert_eq!(result.len(), 2); - assert_arrays_eq!(fixed_list.into_array(), result); + assert_arrays_eq!(fixed_list.into_array(), result, &mut ctx); Ok(()) } #[test] fn test_recursive_deep_nesting() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let innermost_elements = buffer![1i32, 2, 3].into_array(); let innermost_offsets = buffer![0u32, 2].into_array(); let innermost_sizes = buffer![2u32, 1].into_array(); @@ -703,7 +718,7 @@ mod tests { )?; assert_eq!(result.len(), 2); - assert_arrays_eq!(outer_listview.into_array(), result); + assert_arrays_eq!(outer_listview.into_array(), result, &mut ctx); Ok(()) } @@ -733,7 +748,7 @@ mod tests { let result = recursive_list_from_list_view(struct_array.clone().into_array(), &mut ctx)?; assert_eq!(result.len(), 4); - assert_arrays_eq!(struct_array.into_array(), result); + assert_arrays_eq!(struct_array.into_array(), result, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/listview/mod.rs b/vortex-array/src/arrays/listview/mod.rs index d5674372b2c..80b2540a897 100644 --- a/vortex-array/src/arrays/listview/mod.rs +++ b/vortex-array/src/arrays/listview/mod.rs @@ -12,6 +12,10 @@ pub(crate) mod compute; mod vtable; pub use vtable::ListView; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + mod conversion; pub use conversion::list_from_list_view; pub use conversion::list_view_from_list; diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 7ef752e4dcb..bae2c4f3fb9 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -455,12 +455,14 @@ mod tests { // Verify the data is correct assert_arrays_eq!( flattened.list_elements_at(0)?, - PrimitiveArray::from_iter([1i32, 2, 3]) + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx ); assert_arrays_eq!( flattened.list_elements_at(1)?, - PrimitiveArray::from_iter([2i32, 3]) + PrimitiveArray::from_iter([2i32, 3]), + &mut ctx ); Ok(()) } @@ -495,12 +497,14 @@ mod tests { // Verify valid lists contain correct data assert_arrays_eq!( flattened.list_elements_at(0)?, - PrimitiveArray::from_iter([1i32, 2]) + PrimitiveArray::from_iter([1i32, 2]), + &mut ctx ); assert_arrays_eq!( flattened.list_elements_at(2)?, - PrimitiveArray::from_iter([3i32]) + PrimitiveArray::from_iter([3i32]), + &mut ctx ); Ok(()) } @@ -536,12 +540,14 @@ mod tests { // Verify the data is correct. assert_arrays_eq!( trimmed.list_elements_at(0)?, - PrimitiveArray::from_iter([1i32, 2]) + PrimitiveArray::from_iter([1i32, 2]), + &mut ctx ); assert_arrays_eq!( trimmed.list_elements_at(1)?, - PrimitiveArray::from_iter([3i32, 4]) + PrimitiveArray::from_iter([3i32, 4]), + &mut ctx ); // Note that element at index 2 (97) is preserved as a gap. @@ -598,12 +604,14 @@ mod tests { // Verify data is preserved assert_arrays_eq!( exact.list_elements_at(0)?, - PrimitiveArray::from_iter([1i32, 2]) + PrimitiveArray::from_iter([1i32, 2]), + &mut ctx ); assert_arrays_eq!( exact.list_elements_at(1)?, - PrimitiveArray::from_iter([3i32, 4]) + PrimitiveArray::from_iter([3i32, 4]), + &mut ctx ); Ok(()) } @@ -626,7 +634,8 @@ mod tests { let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?; assert_arrays_eq!( trimmed.list_elements_at(1)?, - PrimitiveArray::from_iter([30i32, 40]) + PrimitiveArray::from_iter([30i32, 40]), + &mut ctx ); Ok(()) } @@ -647,7 +656,8 @@ mod tests { let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?; assert_arrays_eq!( trimmed.list_elements_at(1)?, - PrimitiveArray::from_iter([30i32, 40]) + PrimitiveArray::from_iter([30i32, 40]), + &mut ctx ); Ok(()) } @@ -670,11 +680,13 @@ mod tests { // Values: [1,2,3] and [2,3]. assert_arrays_eq!( rebuilt.list_elements_at(0)?, - PrimitiveArray::from_iter([1i32, 2, 3]) + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx ); assert_arrays_eq!( rebuilt.list_elements_at(1)?, - PrimitiveArray::from_iter([2i32, 3]) + PrimitiveArray::from_iter([2i32, 3]), + &mut ctx ); // Signed input -> signed result (offsets widened to i32, sizes kept i16). @@ -711,7 +723,7 @@ mod tests { let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?; // min(offsets) = 0, so nothing to trim; output should equal input. - assert_arrays_eq!(trimmed, listview); + assert_arrays_eq!(trimmed, listview, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/arrays/listview/tests/basic.rs b/vortex-array/src/arrays/listview/tests/basic.rs index 61664c2f608..3d4d6db284d 100644 --- a/vortex-array/src/arrays/listview/tests/basic.rs +++ b/vortex-array/src/arrays/listview/tests/basic.rs @@ -8,8 +8,8 @@ use vortex_buffer::buffer; use vortex_error::VortexResult; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::ListArray; @@ -26,6 +26,7 @@ use crate::validity::Validity; #[test] fn test_basic_listview_comprehensive() { + let mut ctx = array_session().create_execution_ctx(); // Comprehensive test for basic ListView functionality including scalar_at. // Logical lists: [[1,2,3], [4,5], [6,7,8,9]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -50,12 +51,13 @@ fn test_basic_listview_comprehensive() { // Check individual list elements. assert_arrays_eq!( listview.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([1i32, 2, 3]) + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx ); // Test scalar_at which returns entire lists as Scalar values. let first_scalar = listview - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert_eq!( first_scalar, @@ -68,17 +70,20 @@ fn test_basic_listview_comprehensive() { assert_arrays_eq!( listview.list_elements_at(1).unwrap(), - PrimitiveArray::from_iter([4i32, 5]) + PrimitiveArray::from_iter([4i32, 5]), + &mut ctx ); assert_arrays_eq!( listview.list_elements_at(2).unwrap(), - PrimitiveArray::from_iter([6i32, 7, 8, 9]) + PrimitiveArray::from_iter([6i32, 7, 8, 9]), + &mut ctx ); } #[test] fn test_out_of_order_offsets() { + let mut ctx = array_session().create_execution_ctx(); // ListView-specific: Tests that offsets can be non-sequential and out-of-order. // Logical lists: [[7,8,9], [1,2,3], [4,5,6]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -92,13 +97,15 @@ fn test_out_of_order_offsets() { // First list starts at offset 6: [7, 8, 9]. assert_arrays_eq!( listview.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([7i32, 8, 9]) + PrimitiveArray::from_iter([7i32, 8, 9]), + &mut ctx ); // Second list starts at offset 0: [1, 2, 3]. assert_arrays_eq!( listview.list_elements_at(1).unwrap(), - PrimitiveArray::from_iter([1i32, 2, 3]) + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx ); } @@ -128,7 +135,7 @@ fn test_from_list_array() -> VortexResult<()> { let validity = Validity::from_iter([true, false, true]); let list_array = ListArray::try_new(elements, offsets, validity)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let list_view = list_view_from_list(list_array, &mut ctx)?; assert_eq!(list_view.len(), 3); @@ -136,18 +143,20 @@ fn test_from_list_array() -> VortexResult<()> { // Check first list. assert_arrays_eq!( list_view.list_elements_at(0)?, - PrimitiveArray::from_iter([1i32, 2]) + PrimitiveArray::from_iter([1i32, 2]), + &mut ctx ); // Check validity is preserved. - assert!(list_view.is_valid(0, &mut LEGACY_SESSION.create_execution_ctx())?); - assert!(list_view.is_invalid(1, &mut LEGACY_SESSION.create_execution_ctx())?); - assert!(list_view.is_valid(2, &mut LEGACY_SESSION.create_execution_ctx())?); + assert!(list_view.is_valid(0, &mut array_session().create_execution_ctx())?); + assert!(list_view.is_invalid(1, &mut array_session().create_execution_ctx())?); + assert!(list_view.is_valid(2, &mut array_session().create_execution_ctx())?); // Check third list. assert_arrays_eq!( list_view.list_elements_at(2)?, - PrimitiveArray::from_iter([5i32, 6, 7]) + PrimitiveArray::from_iter([5i32, 6, 7]), + &mut ctx ); Ok(()) } @@ -158,6 +167,7 @@ fn test_from_list_array() -> VortexResult<()> { #[case::constant_offsets(false, true)] // Varying sizes, constant offsets #[case::both_constant(true, true)] // Both constant fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_offsets: bool) { + let mut ctx = array_session().create_execution_ctx(); // Logical lists vary by case: // - constant_sizes: [[1,2,3], [4,5,6], [7,8,9]] (size 3 each, varying offsets) // - constant_offsets: [[1,2,3], [1,2], [1]] (all start at 0, varying sizes) @@ -190,7 +200,7 @@ fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_o // All lists are identical [1, 2, 3] (overlapping). let expected = PrimitiveArray::from_iter([1i32, 2, 3]); for i in 0..3 { - assert_arrays_eq!(listview.list_elements_at(i).unwrap(), expected); + assert_arrays_eq!(listview.list_elements_at(i).unwrap(), expected, &mut ctx); } } else if const_sizes { // All lists have size 3, different offsets (no overlap). @@ -203,7 +213,7 @@ fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_o listview .list_elements_at(0) .unwrap() - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 1i32.into() ); @@ -211,7 +221,7 @@ fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_o listview .list_elements_at(1) .unwrap() - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 1i32.into() ); @@ -219,7 +229,7 @@ fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_o listview .list_elements_at(2) .unwrap() - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 1i32.into() ); diff --git a/vortex-array/src/arrays/listview/tests/common.rs b/vortex-array/src/arrays/listview/tests/common.rs index be89b586209..1398a8defe6 100644 --- a/vortex-array/src/arrays/listview/tests/common.rs +++ b/vortex-array/src/arrays/listview/tests/common.rs @@ -12,13 +12,11 @@ use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; -use crate::session::ArraySession; use crate::validity::Validity; /// A shared session for `ListView` tests, used to create execution contexts via /// [`create_execution_ctx`](crate::VortexSessionExecute::create_execution_ctx). -pub static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +pub static SESSION: LazyLock = LazyLock::new(crate::array_session); /// Creates a basic ListView for testing: [[0,1,2], [3,4], [5,6], [7,8,9]] pub fn create_basic_listview() -> ListViewArray { diff --git a/vortex-array/src/arrays/listview/tests/density.rs b/vortex-array/src/arrays/listview/tests/density.rs index ad408f650b3..8538a9fe20c 100644 --- a/vortex-array/src/arrays/listview/tests/density.rs +++ b/vortex-array/src/arrays/listview/tests/density.rs @@ -7,7 +7,6 @@ use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; -use vortex_session::VortexSession; use super::common::create_basic_listview; use super::common::create_empty_lists_listview; @@ -23,13 +22,12 @@ use crate::arrays::listview::tests::common::create_empty_elements_listview; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::scalar::ScalarValue; -use crate::session::ArraySession; use crate::validity::Validity; const EPS: f32 = 1e-6; fn test_execution_ctx() -> ExecutionCtx { - let session = VortexSession::empty().with::(); + let session = crate::array_session(); session.create_execution_ctx() } diff --git a/vortex-array/src/arrays/listview/tests/filter.rs b/vortex-array/src/arrays/listview/tests/filter.rs index e218c94926a..25cc0de4aa8 100644 --- a/vortex-array/src/arrays/listview/tests/filter.rs +++ b/vortex-array/src/arrays/listview/tests/filter.rs @@ -14,7 +14,6 @@ use super::common::create_large_listview; use super::common::create_nullable_listview; use super::common::create_overlapping_listview; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; @@ -22,11 +21,9 @@ use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; -use crate::session::ArraySession; use crate::validity::Validity; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(crate::array_session); // Conformance tests for common filter scenarios. #[rstest] @@ -41,6 +38,7 @@ fn test_filter_listview_conformance(#[case] listview: ListViewArray) { #[test] fn test_filter_preserves_unreferenced_elements() { + let mut ctx = SESSION.create_execution_ctx(); // ListView-specific: Test that filter preserves the entire elements array. // // Logical list: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]] @@ -63,7 +61,8 @@ fn test_filter_preserves_unreferenced_elements() { // Verify the entire elements array is preserved. assert_arrays_eq!( result_list.elements(), - PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]), + &mut ctx ); // Verify offsets are unchanged. @@ -73,6 +72,7 @@ fn test_filter_preserves_unreferenced_elements() { #[test] fn test_filter_with_gaps() { + let mut ctx = SESSION.create_execution_ctx(); // ListView-specific: Test filtering with gaps in elements array. // // Logical list: [[1,2,3], [7,8,9], [11,12], [2,3], [8,9]] @@ -95,7 +95,8 @@ fn test_filter_with_gaps() { // Verify the entire elements array is preserved including gaps. assert_arrays_eq!( result_list.elements(), - PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]) + PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]), + &mut ctx ); // Verify offsets are unchanged. @@ -106,7 +107,8 @@ fn test_filter_with_gaps() { // Verify the lists still read correctly. assert_arrays_eq!( result_list.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([7i32, 8, 9]) + PrimitiveArray::from_iter([7i32, 8, 9]), + &mut ctx ); } @@ -198,7 +200,7 @@ fn test_filter_extreme_offsets() { let list0 = result_list.list_elements_at(0).unwrap(); assert_eq!( list0 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut SESSION.create_execution_ctx()) .unwrap() .as_primitive() .as_::() @@ -207,7 +209,7 @@ fn test_filter_extreme_offsets() { ); assert_eq!( list0 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut SESSION.create_execution_ctx()) .unwrap() .as_primitive() .as_::() diff --git a/vortex-array/src/arrays/listview/tests/nested.rs b/vortex-array/src/arrays/listview/tests/nested.rs index 4aa3edd75d5..4b2676f4491 100644 --- a/vortex-array/src/arrays/listview/tests/nested.rs +++ b/vortex-array/src/arrays/listview/tests/nested.rs @@ -4,8 +4,8 @@ use vortex_buffer::buffer; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -74,7 +74,7 @@ fn test_listview_of_listview_with_overlapping() { // inner[0] should be [1, 2, 3]. assert_eq!( inner0 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .as_primitive() .as_::() @@ -83,7 +83,7 @@ fn test_listview_of_listview_with_overlapping() { ); assert_eq!( inner0 - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap() .as_primitive() .as_::() @@ -94,7 +94,7 @@ fn test_listview_of_listview_with_overlapping() { // inner[1] should be [3, 4, 5] - shares element 3 with inner[0]. assert_eq!( inner1 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .as_primitive() .as_::() @@ -103,7 +103,7 @@ fn test_listview_of_listview_with_overlapping() { ); assert_eq!( inner1 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .as_primitive() .as_::() @@ -287,7 +287,7 @@ fn test_listview_zero_and_overlapping() { assert_eq!(inner1.len(), 3); // [1, 2, 3] assert_eq!( inner1 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .as_primitive() .as_::() @@ -306,7 +306,7 @@ fn test_listview_zero_and_overlapping() { assert_eq!(inner3.len(), 3); // [2, 3, 4] assert_eq!( inner3 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .as_primitive() .as_::() @@ -385,7 +385,7 @@ fn test_listview_of_struct_with_nulls() { // The middle element (struct[2]) should be null. assert!( list1 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); diff --git a/vortex-array/src/arrays/listview/tests/nullability.rs b/vortex-array/src/arrays/listview/tests/nullability.rs index ff3f9c44c6c..4889b73de21 100644 --- a/vortex-array/src/arrays/listview/tests/nullability.rs +++ b/vortex-array/src/arrays/listview/tests/nullability.rs @@ -7,8 +7,8 @@ use rstest::rstest; use vortex_buffer::buffer; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -38,17 +38,17 @@ fn test_nullable_listview_comprehensive() { // Check validity. assert!( listview - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(0, &mut array_session().create_execution_ctx()) .unwrap() ); assert!( listview - .is_invalid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_invalid(1, &mut array_session().create_execution_ctx()) .unwrap() ); assert!( listview - .is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(2, &mut array_session().create_execution_ctx()) .unwrap() ); @@ -60,7 +60,7 @@ fn test_nullable_listview_comprehensive() { // Test scalar_at with nulls. let first = listview - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!first.is_null()); assert_eq!( @@ -73,12 +73,12 @@ fn test_nullable_listview_comprehensive() { ); let second = listview - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(); assert!(second.is_null()); let third = listview - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(); assert!(!third.is_null()); assert_eq!( @@ -95,13 +95,13 @@ fn test_nullable_listview_comprehensive() { assert_eq!(null_list_data.len(), 2); assert_eq!( null_list_data - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 3i32.into() ); assert_eq!( null_list_data - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 4i32.into() ); @@ -123,7 +123,7 @@ fn test_nullable_patterns(#[case] validity: Validity, #[case] expected_validity: for (i, &expected) in expected_validity.iter().enumerate() { assert_eq!( listview - .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(i, &mut array_session().create_execution_ctx()) .unwrap(), expected ); @@ -150,19 +150,19 @@ fn test_nullable_elements() { assert_eq!(first_list.len(), 2); assert!( !first_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert_eq!( first_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 1i32.into() ); assert!( first_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); @@ -171,19 +171,19 @@ fn test_nullable_elements() { let second_list = listview.list_elements_at(1).unwrap(); assert!( !second_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert_eq!( second_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 3i32.into() ); assert!( second_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); @@ -192,25 +192,25 @@ fn test_nullable_elements() { let third_list = listview.list_elements_at(2).unwrap(); assert!( !third_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert_eq!( third_list - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 5i32.into() ); assert!( !third_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert_eq!( third_list - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), 6i32.into() ); diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index c911b9ba7a5..f49680a854e 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -13,11 +13,11 @@ use super::common::create_large_listview; use super::common::create_nullable_listview; use crate::ArrayRef; use crate::IntoArray; -use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_constant::is_constant; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::ListView; @@ -64,10 +64,10 @@ fn test_slice_comprehensive() { assert_eq!( full_list .array() - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(), listview - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(), "Mismatch at index {}", i @@ -84,6 +84,7 @@ fn test_slice_comprehensive() { #[test] fn test_slice_out_of_order() { + let mut ctx = array_session().create_execution_ctx(); // ListView-specific: Test slicing with out-of-order offsets. // Logical lists: [[70,80], [10,20,30], [40,50,60], [90], [30]] let elements = buffer![10i32, 20, 30, 40, 50, 60, 70, 80, 90].into_array(); @@ -123,15 +124,18 @@ fn test_slice_out_of_order() { // Verify the actual list contents are correct. assert_arrays_eq!( sliced_list.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([10i32, 20, 30]) + PrimitiveArray::from_iter([10i32, 20, 30]), + &mut ctx ); assert_arrays_eq!( sliced_list.list_elements_at(1).unwrap(), - PrimitiveArray::from_iter([40i32, 50, 60]) + PrimitiveArray::from_iter([40i32, 50, 60]), + &mut ctx ); assert_arrays_eq!( sliced_list.list_elements_at(2).unwrap(), - PrimitiveArray::from_iter([90i32]) + PrimitiveArray::from_iter([90i32]), + &mut ctx ); } @@ -159,13 +163,13 @@ fn test_slice_with_nulls() { assert!( sliced_list .array() - .is_invalid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_invalid(0, &mut array_session().create_execution_ctx()) .unwrap() ); // Original index 1 was null. assert!( sliced_list .array() - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(1, &mut array_session().create_execution_ctx()) .unwrap() ); // Original index 2 was valid. @@ -295,12 +299,12 @@ fn test_cast_with_nulls() { let result_list = result.to_listview(); assert!( result_list - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(0, &mut array_session().create_execution_ctx()) .unwrap() ); assert!( result_list - .is_invalid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_invalid(1, &mut array_session().create_execution_ctx()) .unwrap() ); } @@ -390,6 +394,7 @@ fn test_cast_large_dataset() { #[test] fn test_zip_widens_false_element_nullability() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // [[1, 2], [3], [4]] let if_true = ListViewArray::new( buffer![1i32, 2, 3, 4].into_array(), @@ -411,7 +416,7 @@ fn test_zip_widens_false_element_nullability() -> VortexResult<()> { let result = mask .into_array() .zip(if_true, if_false)? - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut array_session().create_execution_ctx())?; assert!(result.is::()); assert_eq!( result.dtype(), @@ -429,12 +434,13 @@ fn test_zip_widens_false_element_nullability() -> VortexResult<()> { Validity::NonNullable, ) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } #[test] fn test_zip_widens_true_element_nullability() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // [[1, null], [3], [4]] let if_true = ListViewArray::new( PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(), @@ -456,7 +462,7 @@ fn test_zip_widens_true_element_nullability() -> VortexResult<()> { let result = mask .into_array() .zip(if_true, if_false)? - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut array_session().create_execution_ctx())?; assert!(result.is::()); assert_eq!( result.dtype(), @@ -474,7 +480,7 @@ fn test_zip_widens_true_element_nullability() -> VortexResult<()> { Validity::NonNullable, ) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -534,7 +540,7 @@ fn test_is_constant_basic( ) .into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(is_constant(&listview, &mut ctx).unwrap(), expected); } @@ -553,7 +559,7 @@ fn test_constant_with_constant_elements() { .into_array(); // All lists contain [42, 42] so should be constant. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert!(is_constant(&listview, &mut ctx).unwrap()); } @@ -577,7 +583,7 @@ fn test_constant_with_nulls() { .with_zero_copy_to_list(true) } .into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert!(!is_constant(&listview_mixed, &mut ctx).unwrap()); // Case 2: All nulls - should be constant. @@ -587,7 +593,7 @@ fn test_constant_with_nulls() { .with_zero_copy_to_list(true) } .into_array(); - let mut ctx2 = LEGACY_SESSION.create_execution_ctx(); + let mut ctx2 = array_session().create_execution_ctx(); assert!(is_constant(&listview_all_null, &mut ctx2).unwrap()); } @@ -602,7 +608,7 @@ fn test_constant_repeated_same_lists() { let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable).into_array(); // All lists are [10, 20, 30] so should be constant. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert!(is_constant(&listview, &mut ctx).unwrap()); } @@ -644,22 +650,22 @@ fn test_mask_preserves_structure() { // Check validity: true in selection means null. assert!( !result_list - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(0, &mut array_session().create_execution_ctx()) .unwrap() ); // Masked. assert!( result_list - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(1, &mut array_session().create_execution_ctx()) .unwrap() ); // Not masked. assert!( !result_list - .is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(2, &mut array_session().create_execution_ctx()) .unwrap() ); // Masked. assert!( !result_list - .is_valid(3, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(3, &mut array_session().create_execution_ctx()) .unwrap() ); // Masked. @@ -698,17 +704,17 @@ fn test_mask_with_existing_nulls() { // Check combined validity: assert!( result_list - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(0, &mut array_session().create_execution_ctx()) .unwrap() ); // Was valid, mask is false -> valid. assert!( !result_list - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(1, &mut array_session().create_execution_ctx()) .unwrap() ); // Was invalid, mask is true -> invalid. assert!( !result_list - .is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(2, &mut array_session().create_execution_ctx()) .unwrap() ); // Was valid, mask is true -> invalid. } @@ -731,17 +737,17 @@ fn test_mask_with_gaps() { assert_eq!(result_list.len(), 3); assert!( !result_list - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(0, &mut array_session().create_execution_ctx()) .unwrap() ); // Masked assert!( result_list - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(1, &mut array_session().create_execution_ctx()) .unwrap() ); // Not masked assert!( result_list - .is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(2, &mut array_session().create_execution_ctx()) .unwrap() ); // Not masked @@ -776,17 +782,17 @@ fn test_mask_constant_arrays() { assert_eq!(result_list.len(), 3); assert!( result_list - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(0, &mut array_session().create_execution_ctx()) .unwrap() ); assert!( !result_list - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(1, &mut array_session().create_execution_ctx()) .unwrap() ); // Masked assert!( result_list - .is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(2, &mut array_session().create_execution_ctx()) .unwrap() ); diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index 572340259e9..2a68f711d30 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -13,7 +13,6 @@ use super::common::create_large_listview; use super::common::create_nullable_listview; use super::common::create_overlapping_listview; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; @@ -21,11 +20,9 @@ use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; -use crate::session::ArraySession; use crate::validity::Validity; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(crate::array_session); // Conformance tests for common take scenarios. #[rstest] @@ -42,6 +39,7 @@ fn test_take_listview_conformance(#[case] listview: ListViewArray) { #[test] fn test_take_preserves_unreferenced_elements() { + let mut ctx = SESSION.create_execution_ctx(); // ListView-specific: Test that take preserves the entire elements array // even when taking only a subset of lists. let elements = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -62,7 +60,8 @@ fn test_take_preserves_unreferenced_elements() { // Verify the entire elements array is preserved. assert_arrays_eq!( result_list.elements(), - PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]), + &mut ctx ); // Verify offsets are preserved. @@ -72,6 +71,7 @@ fn test_take_preserves_unreferenced_elements() { #[test] fn test_take_with_gaps() { + let mut ctx = SESSION.create_execution_ctx(); // ListView-specific: Test with gaps in elements array. // Elements with gaps (999 values are "gaps" between used ranges). let elements = buffer![1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12].into_array(); @@ -89,13 +89,15 @@ fn test_take_with_gaps() { // Verify the entire elements array is preserved including gaps. assert_arrays_eq!( result_list.elements(), - PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]) + PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]), + &mut ctx ); // Verify the lists still read correctly despite gaps. assert_arrays_eq!( result_list.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([7i32, 8, 9]) + PrimitiveArray::from_iter([7i32, 8, 9]), + &mut ctx ); } @@ -187,7 +189,7 @@ fn test_take_extreme_offsets() { let list0 = result_list.list_elements_at(0).unwrap(); assert_eq!( list0 - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut SESSION.create_execution_ctx()) .unwrap() .as_primitive() .as_::() @@ -196,7 +198,7 @@ fn test_take_extreme_offsets() { ); assert_eq!( list0 - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut SESSION.create_execution_ctx()) .unwrap() .as_primitive() .as_::() diff --git a/vortex-array/src/arrays/listview/vtable/kernel.rs b/vortex-array/src/arrays/listview/vtable/kernel.rs index 1ad98f62a33..5621b23e79e 100644 --- a/vortex-array/src/arrays/listview/vtable/kernel.rs +++ b/vortex-array/src/arrays/listview/vtable/kernel.rs @@ -1,12 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + use crate::arrays::ListView; -use crate::kernel::ParentKernelSet; +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::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(ListView)), - ParentKernelSet::lift(&ZipExecuteAdaptor(ListView)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), ListView, CastExecuteAdaptor(ListView)); + kernels.register_execute_parent_kernel(Zip.id(), ListView, ZipExecuteAdaptor(ListView)); +} diff --git a/vortex-array/src/arrays/listview/vtable/mod.rs b/vortex-array/src/arrays/listview/vtable/mod.rs index 26e54b2d051..f7a7a5486b6 100644 --- a/vortex-array/src/arrays/listview/vtable/mod.rs +++ b/vortex-array/src/arrays/listview/vtable/mod.rs @@ -16,6 +16,7 @@ use vortex_session::registry::CachedId; use crate::ArrayEq; use crate::ArrayHash; +use crate::ArrayParts; use crate::ArrayRef; use crate::EqMode; use crate::ExecutionCtx; @@ -24,6 +25,7 @@ use crate::array::Array; 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::ListViewData; use crate::arrays::listview::array::ELEMENTS_SLOT; @@ -32,7 +34,6 @@ 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::compute::rules::PARENT_RULES; -use crate::arrays::listview::vtable::kernel::PARENT_KERNELS; use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::dtype::Nullability; @@ -45,6 +46,10 @@ mod validity; /// A [`ListView`]-encoded Vortex array. pub type ListViewArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(Clone, Debug)] pub struct ListView; @@ -92,6 +97,14 @@ impl VTable for ListView { vortex_panic!("ListViewArray 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, @@ -154,7 +167,7 @@ impl VTable for ListView { buffers: &[BufferHandle], children: &dyn ArrayChildren, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult> { let metadata = ListViewMetadata::decode(metadata)?; vortex_ensure!( buffers.is_empty(), @@ -201,7 +214,7 @@ impl VTable for ListView { ListViewData::validate(&elements, &offsets, &sizes, &validity)?; let data = ListViewData::try_new()?; let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len); - Ok(crate::array::ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { @@ -219,13 +232,4 @@ impl VTable for ListView { ) -> VortexResult> { PARENT_RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index b26d101eb33..60b045d521d 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -8,10 +8,10 @@ use vortex_error::VortexResult; use super::*; use crate::Canonical; use crate::IntoArray; -use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -51,7 +51,7 @@ fn test_canonical_dtype_matches_array_dtype() -> VortexResult<()> { let canonical = array .clone() .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut array_session().create_execution_ctx())?; assert_eq!(canonical.dtype(), array.dtype()); Ok(()) } @@ -67,34 +67,35 @@ fn test_masked_child_with_validity() { let prim = array.as_array().to_primitive(); // Positions where validity is false should be null in masked_child. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(prim.valid_count(&mut ctx).unwrap(), 3); assert!( - prim.is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + prim.is_valid(0, &mut array_session().create_execution_ctx()) .unwrap() ); assert!( !prim - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(1, &mut array_session().create_execution_ctx()) .unwrap() ); assert!( - prim.is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()) + prim.is_valid(2, &mut array_session().create_execution_ctx()) .unwrap() ); assert!( !prim - .is_valid(3, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(3, &mut array_session().create_execution_ctx()) .unwrap() ); assert!( - prim.is_valid(4, &mut LEGACY_SESSION.create_execution_ctx()) + prim.is_valid(4, &mut array_session().create_execution_ctx()) .unwrap() ); } #[test] fn test_masked_child_all_valid() { + let mut ctx = array_session().create_execution_ctx(); // When validity is AllValid, masked_child should invert to AllInvalid. let child = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); let array = MaskedArray::try_new(child, Validity::AllValid).unwrap(); @@ -102,13 +103,14 @@ fn test_masked_child_all_valid() { assert_eq!(array.len(), 3); assert_eq!( array - .valid_count(&mut LEGACY_SESSION.create_execution_ctx()) + .valid_count(&mut array_session().create_execution_ctx()) .unwrap(), 3 ); assert_arrays_eq!( PrimitiveArray::from_option_iter([10i32, 20, 30].map(Some)), - array + array, + &mut ctx ); } @@ -129,7 +131,7 @@ fn test_masked_child_preserves_length(#[case] validity: Validity) { assert_eq!(array.len(), len); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert!( array .validity() diff --git a/vortex-array/src/arrays/masked/vtable/canonical.rs b/vortex-array/src/arrays/masked/vtable/canonical.rs index 90bd999b8ed..32a5e334afa 100644 --- a/vortex-array/src/arrays/masked/vtable/canonical.rs +++ b/vortex-array/src/arrays/masked/vtable/canonical.rs @@ -8,8 +8,8 @@ mod tests { use crate::Canonical; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; use crate::dtype::Nullability; @@ -41,7 +41,7 @@ mod tests { let canonical = array .clone() .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut array_session().create_execution_ctx())?; assert_eq!(canonical.dtype().nullability(), expected_nullability); assert_eq!(canonical.dtype(), array.dtype()); Ok(()) @@ -56,17 +56,17 @@ mod tests { let canonical = array .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut array_session().create_execution_ctx())?; let prim = canonical.into_primitive(); // Check that null positions match validity. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(prim.valid_count(&mut ctx)?, 3); - assert!(prim.is_valid(0, &mut LEGACY_SESSION.create_execution_ctx())?); - assert!(!prim.is_valid(1, &mut LEGACY_SESSION.create_execution_ctx())?); - assert!(prim.is_valid(2, &mut LEGACY_SESSION.create_execution_ctx())?); - assert!(!prim.is_valid(3, &mut LEGACY_SESSION.create_execution_ctx())?); - assert!(prim.is_valid(4, &mut LEGACY_SESSION.create_execution_ctx())?); + assert!(prim.is_valid(0, &mut array_session().create_execution_ctx())?); + assert!(!prim.is_valid(1, &mut array_session().create_execution_ctx())?); + assert!(prim.is_valid(2, &mut array_session().create_execution_ctx())?); + assert!(!prim.is_valid(3, &mut array_session().create_execution_ctx())?); + assert!(prim.is_valid(4, &mut array_session().create_execution_ctx())?); Ok(()) } @@ -79,12 +79,12 @@ mod tests { let canonical = array .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut array_session().create_execution_ctx())?; assert_eq!(canonical.dtype().nullability(), Nullability::Nullable); assert_eq!( canonical .into_array() - .valid_count(&mut LEGACY_SESSION.create_execution_ctx())?, + .valid_count(&mut array_session().create_execution_ctx())?, 3 ); Ok(()) diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index 4409c8ff3c6..64d39f63531 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -18,6 +18,7 @@ use vortex_session::registry::CachedId; use crate::AnyCanonical; use crate::ArrayEq; use crate::ArrayHash; +use crate::ArrayParts; use crate::ArrayRef; use crate::Canonical; use crate::EqMode; @@ -29,6 +30,7 @@ use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; use crate::array::validity_to_child; +use crate::array::with_empty_buffers; use crate::arrays::ConstantArray; use crate::arrays::masked::MaskedArrayExt; use crate::arrays::masked::MaskedArraySlotsExt; @@ -105,6 +107,14 @@ impl VTable for Masked { None } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + fn serialize( _array: ArrayView<'_, Self>, _session: &VortexSession, @@ -121,7 +131,7 @@ impl VTable for Masked { buffers: &[BufferHandle], children: &dyn ArrayChildren, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult> { if !metadata.is_empty() { vortex_bail!( "MaskedArray expects empty metadata, got {} bytes", @@ -153,10 +163,8 @@ impl VTable for Masked { child.all_valid(&mut LEGACY_SESSION.create_execution_ctx())?, validity, )?; - Ok( - crate::array::ArrayParts::new(self.clone(), dtype.clone(), len, data) - .with_slots(smallvec![Some(child), validity_slot]), - ) + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data) + .with_slots(smallvec![Some(child), validity_slot])) } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { @@ -165,7 +173,7 @@ impl VTable for Masked { let validity = array.masked_validity(); // Fast path: all masked means result is all nulls. - if matches!(validity, Validity::AllInvalid) { + if validity.definitely_all_null() { return Ok(ExecutionResult::done( ConstantArray::new(Scalar::null(array.dtype().as_nullable()), array.len()) .into_array(), @@ -206,8 +214,8 @@ mod tests { use crate::ArrayContext; use crate::Canonical; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Masked; use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; @@ -243,7 +251,7 @@ mod tests { let serialized = array .clone() .into_array() - .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default()) + .serialize(&ctx, &array_session(), &SerializeOptions::default()) .unwrap(); // Concat into a single buffer. @@ -259,7 +267,7 @@ mod tests { &dtype, len, &ReadContext::new(ctx.to_ids()), - &LEGACY_SESSION, + &array_session(), ) .unwrap(); @@ -287,7 +295,7 @@ mod tests { assert_eq!(array.dtype().nullability(), Nullability::Nullable); // Execute the array. This should produce a Canonical with Nullable dtype. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result: Canonical = array.into_array().execute(&mut ctx)?; assert_eq!( diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 148cf209555..15b6b5aa6be 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -1,7 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! All the built-in encoding schemes and arrays. +//! Built-in array encodings. +//! +//! Canonical arrays are the default uncompressed representation for a logical dtype: +//! [`NullArray`], [`BoolArray`], [`PrimitiveArray`], [`DecimalArray`], [`VarBinViewArray`], +//! [`ListViewArray`], [`FixedSizeListArray`], [`StructArray`], [`ExtensionArray`], and +//! [`VariantArray`]. +//! +//! Utility and lazy arrays represent common transformations without immediately materializing +//! their result. Examples include [`ChunkedArray`] for concatenation, [`ConstantArray`] for repeated +//! values, [`DictArray`] for dictionary encoding, [`FilterArray`] for masked rows, [`SliceArray`] +//! for views, and [`ScalarFnArray`] for deferred scalar-function execution. +//! +//! Some public arrays are primarily internal building blocks. Their constructors and extension +//! traits document the stable contract; avoid depending on undocumented slot order or metadata +//! details. #[cfg(any(test, feature = "_test-harness"))] mod assertions; @@ -105,6 +119,25 @@ pub use varbinview::VarBinViewArray; pub mod variant; pub use variant::Variant; pub use variant::VariantArray; +use vortex_session::VortexSession; + +pub(crate) fn initialize(session: &VortexSession) { + bool::initialize(session); + chunked::initialize(session); + decimal::initialize(session); + dict::initialize(session); + extension::initialize(session); + filter::initialize(session); + fixed_size_list::initialize(session); + list::initialize(session); + listview::initialize(session); + patched::initialize(session); + primitive::initialize(session); + struct_::initialize(session); + varbin::initialize(session); + varbinview::initialize(session); + variant::initialize(session); +} #[cfg(feature = "arbitrary")] pub mod arbitrary; diff --git a/vortex-array/src/arrays/null/compute/cast.rs b/vortex-array/src/arrays/null/compute/cast.rs index be6c73ad7d5..a2527e5fb73 100644 --- a/vortex-array/src/arrays/null/compute/cast.rs +++ b/vortex-array/src/arrays/null/compute/cast.rs @@ -32,8 +32,8 @@ mod tests { use rstest::rstest; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::NullArray; use crate::builtins::ArrayBuiltins; use crate::compute::conformance::cast::test_cast_conformance; @@ -68,7 +68,7 @@ mod tests { for i in 0..5 { assert!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); diff --git a/vortex-array/src/arrays/null/compute/mod.rs b/vortex-array/src/arrays/null/compute/mod.rs index df2bc1547ae..5e756da27d0 100644 --- a/vortex-array/src/arrays/null/compute/mod.rs +++ b/vortex-array/src/arrays/null/compute/mod.rs @@ -15,10 +15,10 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; - use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::NullArray; use crate::compute::conformance::consistency::test_array_consistency; use crate::compute::conformance::filter::test_filter_conformance; @@ -38,7 +38,10 @@ mod test { sliced_arr .validity() .unwrap() - .execute_mask(sliced_arr.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask( + sliced_arr.len(), + &mut array_session().create_execution_ctx() + ) .unwrap(), Mask::AllFalse(4) )); @@ -59,7 +62,7 @@ mod test { taken_arr .validity() .unwrap() - .execute_mask(taken_arr.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(taken_arr.len(), &mut array_session().create_execution_ctx()) .unwrap(), Mask::AllFalse(5) )); @@ -70,7 +73,7 @@ mod test { let nulls = NullArray::new(10); let scalar = nulls - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert!(scalar.is_null()); assert_eq!(scalar.dtype().clone(), DType::Null); @@ -104,6 +107,9 @@ mod test { #[case::null_array_large(NullArray::new(1000))] #[case::null_array_empty(NullArray::new(0))] fn test_null_consistency(#[case] array: NullArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs index dd31ef16853..9c8c8e41892 100644 --- a/vortex-array/src/arrays/null/mod.rs +++ b/vortex-array/src/arrays/null/mod.rs @@ -18,6 +18,7 @@ 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::null::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::dtype::DType; @@ -64,6 +65,14 @@ impl VTable for Null { 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 { vortex_panic!("NullArray slot_name index {idx} out of bounds") } @@ -123,7 +132,7 @@ impl VTable for Null { /// ``` /// # fn main() -> vortex_error::VortexResult<()> { /// use vortex_array::arrays::NullArray; -/// use vortex_array::{IntoArray, LEGACY_SESSION, VortexSessionExecute}; +/// use vortex_array::{IntoArray, VortexSessionExecute, array_session}; /// /// // Create a null array with 5 elements /// let array = NullArray::new(5); @@ -133,7 +142,7 @@ impl VTable for Null { /// assert_eq!(sliced.len(), 2); /// /// // All elements are null -/// let mut ctx = LEGACY_SESSION.create_execution_ctx(); +/// let mut ctx = array_session().create_execution_ctx(); /// let scalar = array.execute_scalar(0, &mut ctx).unwrap(); /// assert!(scalar.is_null()); /// # Ok(()) diff --git a/vortex-array/src/arrays/patched/compute/compare.rs b/vortex-array/src/arrays/patched/compute/compare.rs index 0f1e6add00f..a9262b6621e 100644 --- a/vortex-array/src/arrays/patched/compute/compare.rs +++ b/vortex-array/src/arrays/patched/compute/compare.rs @@ -163,9 +163,8 @@ mod tests { use vortex_error::VortexResult; use vortex_error::vortex_err; - use crate::ExecutionCtx; use crate::IntoArray; - use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::Patched; @@ -189,7 +188,7 @@ mod tests { ) .unwrap(); - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx) .unwrap() @@ -207,7 +206,7 @@ mod tests { let expected = BoolArray::from_indices(512, [509, 510, 511], Validity::NonNullable).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -222,7 +221,7 @@ mod tests { ) .unwrap(); - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx).unwrap(); // Slice the array so that the first patch should be skipped. @@ -240,7 +239,7 @@ mod tests { let expected = BoolArray::from_indices(502, [500, 501], Validity::NonNullable).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -259,7 +258,7 @@ mod tests { None, )?; - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx)? .into_array() .try_downcast::() @@ -277,7 +276,7 @@ mod tests { let expected = BoolArray::from_indices(512, [510], Validity::NonNullable).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); Ok(()) } @@ -293,7 +292,7 @@ mod tests { None, )?; - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx)? .into_array() .try_downcast::() @@ -311,7 +310,7 @@ mod tests { let expected = BoolArray::from_indices(10, [7], Validity::NonNullable).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/patched/compute/filter.rs b/vortex-array/src/arrays/patched/compute/filter.rs index d2c60941059..8cf0c186fb6 100644 --- a/vortex-array/src/arrays/patched/compute/filter.rs +++ b/vortex-array/src/arrays/patched/compute/filter.rs @@ -67,9 +67,8 @@ mod tests { use vortex_error::VortexResult; use vortex_mask::Mask; - use crate::ExecutionCtx; use crate::IntoArray; - use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::Patched; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -79,7 +78,7 @@ mod tests { #[test] fn test_filter_noop() -> VortexResult<()> { // Filter that doesn't prune any chunks (all data fits in one chunk). - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let array = buffer![u16::MIN; 5].into_array(); let patched_indices = buffer![3u8, 4].into_array(); @@ -98,7 +97,7 @@ mod tests { // Values at indices 0 and 4: MIN and MAX. let expected = PrimitiveArray::from_iter([u16::MIN, u16::MAX]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } @@ -106,7 +105,7 @@ mod tests { #[test] fn test_filter_with_offset() -> VortexResult<()> { // Test filtering where offset > 0. - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let array = buffer![u16::MIN; 4096].into_array(); let patched_indices = buffer![5u16, 1030].into_array(); @@ -126,7 +125,7 @@ mod tests { .execute::(&mut ctx)?; let expected = PrimitiveArray::from_iter([u16::MAX, u16::MIN, u16::MIN, u16::MAX]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } @@ -134,7 +133,7 @@ mod tests { #[test] fn test_filter_basic() -> VortexResult<()> { // Basic test: filter with mask that crosses boundaries. - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let array = buffer![u16::MIN; 4096].into_array(); let patched_indices = buffer![1024u16, 1025].into_array(); @@ -153,7 +152,7 @@ mod tests { let expected = PrimitiveArray::from_iter([u16::MAX, u16::MAX, u16::MIN]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } @@ -161,7 +160,7 @@ mod tests { #[test] fn test_filter_complex() -> VortexResult<()> { // Filter with mask that crosses boundaries, with patches offset. - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let array = buffer![u16::MIN; 4096].into_array(); let patched_indices = buffer![1024u16, 1025].into_array(); @@ -180,7 +179,7 @@ mod tests { let expected = PrimitiveArray::from_iter([u16::MAX, u16::MIN, u16::MIN]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } @@ -188,7 +187,7 @@ mod tests { #[test] fn test_filter_sliced() -> VortexResult<()> { // Test filter on a sliced PatchedArray to exercise codepath where offset > 0. - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); // Create a larger array (6 chunks) so we can slice and still have room // for the filter to prune chunks. @@ -217,7 +216,7 @@ mod tests { let expected = PrimitiveArray::from_iter([u16::MAX, u16::MAX, u16::MIN]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } @@ -226,7 +225,7 @@ mod tests { fn test_filter_with_offset_nonuniform() -> VortexResult<()> { // Test filtering with offset > 0 using non-uniform base values. // This catches slice_chunks bugs where inner coordinates are miscalculated. - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); // Use non-uniform values so that incorrect slicing is detectable. let base_values: Vec = (0u16..4096).collect(); @@ -254,7 +253,7 @@ mod tests { // Expected: 9999 (patched at logical 0), 6 (original at logical 1), 8888 (patched at logical 1025). let expected = PrimitiveArray::from_iter([9999u16, 6, 8888]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } @@ -263,7 +262,7 @@ mod tests { fn test_filter_with_offset_last_chunk() -> VortexResult<()> { // Test filtering with offset > 0 where the mask touches the last chunk. // This ensures we don't accidentally slice past the end of the array or mask. - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); // Create a 6-chunk array (6144 elements). let array = buffer![u16::MIN; 6144].into_array(); @@ -294,7 +293,7 @@ mod tests { // Expected: patch at 3976 (was 5000), patch at 4976 (was 6000), and MIN at 5119. let expected = PrimitiveArray::from_iter([u16::MAX, u16::MAX, u16::MIN]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/patched/compute/take.rs b/vortex-array/src/arrays/patched/compute/take.rs index 893d18db1ac..993428808d4 100644 --- a/vortex-array/src/arrays/patched/compute/take.rs +++ b/vortex-array/src/arrays/patched/compute/take.rs @@ -136,8 +136,9 @@ mod tests { use vortex_session::VortexSession; use crate::ArrayRef; - use crate::ExecutionCtx; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Patched; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -159,7 +160,7 @@ mod tests { )?; let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); Patched::from_array_and_patches(values, &patches, &mut ctx)? .into_array() @@ -168,6 +169,7 @@ mod tests { #[test] fn test_take_basic() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at indices [1, 3] with values [10, 30] let array = make_patched_array(&[0; 5], &[1, 3], &[10, 30], 0..5)?; @@ -177,13 +179,14 @@ mod tests { let result = array.take(indices)?.to_canonical()?.into_array(); let expected = PrimitiveArray::from_iter([0u16, 10, 0, 30, 0]).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); Ok(()) } #[test] fn test_take_sliced() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let array = make_patched_array(&[0; 10], &[1, 3], &[100, 200], 2..10)?; let indices = buffer![0u32, 1, 2, 3, 7].into_array(); @@ -191,13 +194,14 @@ mod tests { let result = array.take(indices)?.to_canonical()?.into_array(); let expected = PrimitiveArray::from_iter([0u16, 200, 0, 0, 0]).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); Ok(()) } #[test] fn test_take_out_of_order() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at indices [1, 3] with values [10, 30] let array = make_patched_array(&[0; 5], &[1, 3], &[10, 30], 0..5)?; @@ -207,13 +211,14 @@ mod tests { let result = array.take(indices)?.to_canonical()?.into_array(); let expected = PrimitiveArray::from_iter([0u16, 30, 0, 10, 0]).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); Ok(()) } #[test] fn test_take_duplicates() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at index [2] with value [99] let array = make_patched_array(&[0; 5], &[2], &[99], 0..5)?; @@ -227,13 +232,14 @@ mod tests { let _canonical = result.to_canonical()?.into_primitive(); let expected = PrimitiveArray::from_iter([99u16, 99, 0, 99]).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); Ok(()) } #[test] fn test_take_with_null_indices() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); use crate::arrays::BoolArray; use crate::validity::Validity; @@ -271,7 +277,7 @@ mod tests { .into_array(), ), ); - assert_arrays_eq!(expected.into_array(), result); + assert_arrays_eq!(expected.into_array(), result, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/patched/mod.rs b/vortex-array/src/arrays/patched/mod.rs index 0ba25ab7929..a88628e5587 100644 --- a/vortex-array/src/arrays/patched/mod.rs +++ b/vortex-array/src/arrays/patched/mod.rs @@ -43,7 +43,7 @@ //! `indices` and `values` are aligned and accessed together. //! //! ```text -//! +//! //! chunk 0 chunk 0 chunk 0 chunk 0 chunk 0 chunk 0 //! lane 0 lane 1 lane 2 lane 3 lane 4 lane 5 //! ┌────────────┬────────────┬────────────┬────────────┬────────────┬────────────┐ @@ -78,6 +78,10 @@ pub use array::*; use vortex_buffer::ByteBuffer; pub use vtable::*; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + /// Patches that have been transposed into GPU format. struct TransposedPatches { n_lanes: usize, diff --git a/vortex-array/src/arrays/patched/vtable/kernels.rs b/vortex-array/src/arrays/patched/vtable/kernels.rs index 7994b19e02e..51f53d15625 100644 --- a/vortex-array/src/arrays/patched/vtable/kernels.rs +++ b/vortex-array/src/arrays/patched/vtable/kernels.rs @@ -1,12 +1,19 @@ // 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::Patched; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CompareExecuteAdaptor(Patched)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Patched)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Binary.id(), Patched, CompareExecuteAdaptor(Patched)); + kernels.register_execute_parent_kernel(Dict.id(), Patched, TakeExecuteAdaptor(Patched)); +} diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index 4ba5bc542e1..bbe592785ec 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -32,6 +32,7 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::array::ValidityChild; use crate::array::ValidityVTableFromChild; +use crate::array::with_empty_buffers; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::patched::PatchedArrayExt; @@ -40,7 +41,6 @@ use crate::arrays::patched::PatchedData; use crate::arrays::patched::PatchedSlots; use crate::arrays::patched::PatchedSlotsView; use crate::arrays::patched::compute::rules::PARENT_RULES; -use crate::arrays::patched::vtable::kernels::PARENT_KERNELS; use crate::arrays::primitive::PrimitiveDataParts; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; @@ -55,6 +55,10 @@ use crate::serde::ArrayChildren; /// A [`Patched`]-encoded Vortex array. pub type PatchedArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernels::initialize(session); +} + #[derive(Clone, Debug)] pub struct Patched; @@ -126,6 +130,14 @@ impl VTable for Patched { vortex_panic!("invalid buffer index for PatchedArray: {idx}"); } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + fn child(array: ArrayView<'_, Self>, idx: usize) -> ArrayRef { match idx { PatchedSlots::INNER => array.inner().clone(), @@ -301,15 +313,6 @@ impl VTable for Patched { Ok(ExecutionResult::done(patched_values.into_array())) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, @@ -354,19 +357,22 @@ mod tests { use vortex_buffer::buffer; use vortex_buffer::buffer_mut; use vortex_error::VortexResult; - use vortex_session::VortexSession; use vortex_session::registry::ReadContext; + use crate::Array; use crate::ArrayContext; + use crate::ArrayParts; use crate::ArraySlots; use crate::Canonical; - use crate::ExecutionCtx; use crate::IntoArray; - use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Patched; use crate::arrays::PatchedArray; use crate::arrays::PrimitiveArray; + use crate::arrays::patched::PatchedArrayExt; use crate::arrays::patched::PatchedArraySlotsExt; + use crate::arrays::patched::PatchedData; use crate::arrays::patched::PatchedSlots; use crate::arrays::patched::PatchedSlotsView; use crate::assert_arrays_eq; @@ -389,8 +395,8 @@ mod tests { ) .unwrap(); - let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let session = array_session(); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -422,8 +428,8 @@ mod tests { ) .unwrap(); - let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let session = array_session(); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -455,8 +461,8 @@ mod tests { ) .unwrap(); - let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let session = array_session(); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -473,7 +479,7 @@ mod tests { expected[3] = 30; let expected = expected.into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -488,8 +494,8 @@ mod tests { ) .unwrap(); - let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let session = array_session(); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -506,7 +512,7 @@ mod tests { expected[0] = 30; let expected = expected.into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -525,8 +531,8 @@ mod tests { ) .unwrap(); - let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let session = array_session(); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -552,7 +558,7 @@ mod tests { ]) .into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } fn make_patched_array( @@ -569,8 +575,8 @@ mod tests { let patches = Patches::new(len, 0, indices, patch_vals, None)?; - let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let session = array_session(); + let mut ctx = session.create_execution_ctx(); Patched::from_array_and_patches(array, &patches, &mut ctx) } @@ -590,11 +596,12 @@ mod tests { let dtype = array.dtype().clone(); let len = array.len(); - LEGACY_SESSION.arrays().register(Patched); + let session = array_session(); + session.arrays().register(Patched); - let ctx = ArrayContext::empty().with_registry(LEGACY_SESSION.arrays().registry().clone()); + let ctx = ArrayContext::empty().with_registry(session.arrays().registry().clone()); let serialized = array - .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default()) + .serialize(&ctx, &session, &SerializeOptions::default()) .unwrap(); // Concat into a single buffer. @@ -606,12 +613,7 @@ mod tests { let parts = SerializedArray::try_from(concat).unwrap(); let decoded = parts - .decode( - &dtype, - len, - &ReadContext::new(ctx.to_ids()), - &LEGACY_SESSION, - ) + .decode(&dtype, len, &ReadContext::new(ctx.to_ids()), &session) .unwrap(); assert!(decoded.is::()); @@ -639,24 +641,26 @@ mod tests { // Create new PatchedArray with same children using with_slots let array_ref = array.into_array(); - let new_array = array_ref.clone().with_slots(slots.into_slots())?; + // SAFETY: the replacement slots are the original children, preserving logical values and + // parent statistics. + let new_array = unsafe { array_ref.clone().with_slots(slots.into_slots()) }?; assert!(new_array.is::()); assert_eq!(array_ref.len(), new_array.len()); assert_eq!(array_ref.dtype(), new_array.dtype()); // Execute both and compare results - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = array_session().create_execution_ctx(); let original_executed = array_ref.execute::(&mut ctx)?.into_primitive(); let new_executed = new_array.execute::(&mut ctx)?.into_primitive(); - assert_arrays_eq!(original_executed, new_executed); + assert_arrays_eq!(original_executed, new_executed, &mut ctx); Ok(()) } #[test] - fn test_with_slots_modified_inner() -> VortexResult<()> { + fn test_rebuild_modified_inner_from_parts() -> VortexResult<()> { let array = make_patched_array(vec![0u16; 10], &[1, 2, 3], &[10, 20, 30])?; // Create a different inner array (all 5s instead of 0s) @@ -668,16 +672,23 @@ mod tests { patch_values: array.patch_values().clone(), }; - let array_ref = array.into_array(); - let new_array = array_ref.with_slots(slots.into_slots())?; + let data = PatchedData { + n_lanes: array.n_lanes(), + offset: array.offset(), + }; + let new_array = Array::try_from_parts( + ArrayParts::new(Patched, array.dtype().clone(), array.len(), data) + .with_slots(slots.into_slots()), + )? + .into_array(); // Execute and verify the inner values changed (except at patch positions) - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = array_session().create_execution_ctx(); let executed = new_array.execute::(&mut ctx)?.into_primitive(); // Expected: all 5s except indices 1, 2, 3 which are patched to 10, 20, 30 let expected = PrimitiveArray::from_iter([5u16, 10, 20, 30, 5, 5, 5, 5, 5, 5]); - assert_arrays_eq!(expected, executed); + assert_arrays_eq!(expected, executed, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/patched/vtable/operations.rs b/vortex-array/src/arrays/patched/vtable/operations.rs index f0491666e56..51dd1fc9e3c 100644 --- a/vortex-array/src/arrays/patched/vtable/operations.rs +++ b/vortex-array/src/arrays/patched/vtable/operations.rs @@ -60,10 +60,9 @@ mod tests { use vortex_buffer::buffer; use vortex_session::VortexSession; - use crate::ExecutionCtx; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Patched; use crate::dtype::Nullability; use crate::optimizer::ArrayOptimizer; @@ -83,7 +82,7 @@ mod tests { .unwrap(); let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -91,25 +90,25 @@ mod tests { assert_eq!( array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::primitive(0u16, Nullability::NonNullable) ); assert_eq!( array - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::primitive(1u16, Nullability::NonNullable) ); assert_eq!( array - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::primitive(1u16, Nullability::NonNullable) ); assert_eq!( array - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(3, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::primitive(1u16, Nullability::NonNullable) ); @@ -128,7 +127,7 @@ mod tests { .unwrap(); let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -136,7 +135,7 @@ mod tests { for index in 0..array.len() { let value = array - .execute_scalar(index, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(index, &mut array_session().create_execution_ctx()) .unwrap(); if [1, 2, 3].contains(&index) { @@ -160,7 +159,7 @@ mod tests { .unwrap(); let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -174,14 +173,14 @@ mod tests { assert_eq!( array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), 1u16.into() ); for index in 1..array.len() { assert_eq!( array - .execute_scalar(index, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(index, &mut array_session().create_execution_ctx()) .unwrap(), 0u16.into() ); diff --git a/vortex-array/src/arrays/patched/vtable/slice.rs b/vortex-array/src/arrays/patched/vtable/slice.rs index 8655e0c7d84..1f59c2a854d 100644 --- a/vortex-array/src/arrays/patched/vtable/slice.rs +++ b/vortex-array/src/arrays/patched/vtable/slice.rs @@ -65,9 +65,8 @@ mod tests { use vortex_error::VortexResult; use crate::Canonical; - use crate::ExecutionCtx; use crate::IntoArray; - use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::Patched; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -81,7 +80,7 @@ mod tests { let patch_values = buffer![u16::MAX; 3].into_array(); let patches = Patches::new(512, 0, patch_indices, patch_values, None)?; - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx)?; @@ -127,7 +126,7 @@ mod tests { ) .unwrap(); - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let patched_array = Patched::from_array_and_patches(inner.into_array(), &patches, &mut ctx) .unwrap() @@ -148,7 +147,7 @@ mod tests { .slice(range) .unwrap(); - assert_arrays_eq!(slice_first, slice_last); + assert_arrays_eq!(slice_first, slice_last, &mut ctx); } #[test] @@ -159,7 +158,7 @@ mod tests { let patched_values = buffer![0u64, 1, 2, 3, 4, 5].into_array(); let patches = Patches::new(10_000, 0, patched_indices, patched_values, None).unwrap(); - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -181,6 +180,6 @@ mod tests { let expected = expected.into_array(); - assert_arrays_eq!(expected, sliced); + assert_arrays_eq!(expected, sliced, &mut ctx); } } diff --git a/vortex-array/src/arrays/primitive/array/accessor.rs b/vortex-array/src/arrays/primitive/array/accessor.rs deleted file mode 100644 index f26b201ea2b..00000000000 --- a/vortex-array/src/arrays/primitive/array/accessor.rs +++ /dev/null @@ -1,41 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::iter; - -use vortex_error::VortexExpect; - -#[expect(deprecated)] -use crate::ToCanonical as _; -use crate::accessor::ArrayAccessor; -use crate::arrays::PrimitiveArray; -use crate::dtype::NativePType; -use crate::validity::Validity; - -impl ArrayAccessor for PrimitiveArray { - fn with_iterator(&self, f: F) -> R - where - F: for<'a> FnOnce(&mut dyn Iterator>) -> R, - { - match self - .validity() - .vortex_expect("primitive validity should be derivable") - { - Validity::NonNullable | Validity::AllValid => { - let mut iter = self.as_slice::().iter().map(Some); - f(&mut iter) - } - Validity::AllInvalid => f(&mut iter::repeat_n(None, self.len())), - Validity::Array(v) => { - #[expect(deprecated)] - let validity = v.to_bool().into_bit_buffer(); - let mut iter = self - .as_slice::() - .iter() - .zip(validity.iter()) - .map(|(value, valid)| valid.then_some(value)); - f(&mut iter) - } - } - } -} diff --git a/vortex-array/src/arrays/primitive/array/cast.rs b/vortex-array/src/arrays/primitive/array/cast.rs index 0d9db991d26..bdab4f0beb8 100644 --- a/vortex-array/src/arrays/primitive/array/cast.rs +++ b/vortex-array/src/arrays/primitive/array/cast.rs @@ -50,11 +50,9 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; - use crate::session::ArraySession; use crate::validity::Validity; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[test] fn test_downcast_all_invalid() { diff --git a/vortex-array/src/arrays/primitive/array/mod.rs b/vortex-array/src/arrays/primitive/array/mod.rs index e22fc9a4336..867c7d0c783 100644 --- a/vortex-array/src/arrays/primitive/array/mod.rs +++ b/vortex-array/src/arrays/primitive/array/mod.rs @@ -3,7 +3,7 @@ use std::fmt::Display; use std::fmt::Formatter; -use std::iter; +use std::iter::repeat; use smallvec::smallvec; use vortex_buffer::Alignment; @@ -18,11 +18,10 @@ use vortex_error::vortex_panic; use crate::ArraySlots; use crate::ExecutionCtx; -#[expect(deprecated)] -use crate::ToCanonical as _; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; +use crate::arrays::BoolArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::dtype::DType; @@ -32,7 +31,6 @@ use crate::dtype::PType; use crate::match_each_native_ptype; use crate::validity::Validity; -mod accessor; mod cast; mod conversion; mod patch; @@ -42,6 +40,7 @@ pub use patch::chunk_range; pub use patch::patch_chunk; use crate::ArrayRef; +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; @@ -70,7 +69,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; /// ``` /// # fn main() -> vortex_error::VortexResult<()> { /// use vortex_array::arrays::PrimitiveArray; -/// use vortex_array::{LEGACY_SESSION, VortexSessionExecute}; +/// use vortex_array::{VortexSessionExecute, array_session}; /// /// // Create from iterator using FromIterator impl /// let array: PrimitiveArray = [1i32, 2, 3, 4, 5].into_iter().collect(); @@ -79,7 +78,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; /// let sliced = array.slice(1..3)?; /// /// // Access individual values -/// let mut ctx = LEGACY_SESSION.create_execution_ctx(); +/// let mut ctx = array_session().create_execution_ctx(); /// let value = sliced.execute_scalar(0, &mut ctx).unwrap(); /// assert_eq!(value, 2i32.into()); /// @@ -154,7 +153,7 @@ pub trait PrimitiveArrayExt: TypedArrayRef { return Ok(self.to_owned()); } - let Some(min_max) = min_max(self.as_ref(), ctx)? else { + let Some(min_max) = min_max(self.as_ref(), ctx, NumericalAggregateOpts::default())? else { return Ok(PrimitiveArray::new( Buffer::::zeroed(self.len()), self.validity(), @@ -446,12 +445,18 @@ impl Array { ptype: PType, validity: Validity, n_rows: usize, + ctx: &mut ExecutionCtx, ) -> Self { let dtype = DType::Primitive(ptype, validity.nullability()); let len = n_rows; let slots = PrimitiveData::make_slots(&validity, len); - let data = - PrimitiveData::from_values_byte_buffer(valid_elems_buffer, ptype, validity, n_rows); + let data = PrimitiveData::from_values_byte_buffer( + valid_elems_buffer, + ptype, + validity, + n_rows, + ctx, + ); unsafe { Array::from_parts_unchecked( ArrayParts::new(Primitive, dtype, len, data).with_slots(slots), @@ -475,7 +480,7 @@ impl Array { } } - pub fn map_each_with_validity(self, f: F) -> VortexResult + pub fn map_each_with_validity(self, ctx: &mut ExecutionCtx, f: F) -> VortexResult where T: NativePType, R: NativePType, @@ -487,18 +492,17 @@ impl Array { let buffer = match &validity { Validity::NonNullable | Validity::AllValid => { - BufferMut::::from_iter(buf_iter.zip(iter::repeat(true)).map(f)) + Buffer::::from_trusted_len_iter(buf_iter.zip(repeat(true)).map(f)) } Validity::AllInvalid => { - BufferMut::::from_iter(buf_iter.zip(iter::repeat(false)).map(f)) + Buffer::::from_trusted_len_iter(buf_iter.zip(repeat(false)).map(f)) } Validity::Array(val) => { - #[expect(deprecated)] - let val = val.to_bool().into_bit_buffer(); - BufferMut::::from_iter(buf_iter.zip(val.iter()).map(f)) + let val = val.clone().execute::(ctx)?.into_bit_buffer(); + Buffer::::from_trusted_len_iter(buf_iter.zip(val.iter()).map(f)) } }; - Ok(PrimitiveArray::new(buffer.freeze(), validity)) + Ok(PrimitiveArray::new(buffer, validity)) } } @@ -540,6 +544,7 @@ impl PrimitiveData { ptype: PType, validity: Validity, n_rows: usize, + ctx: &mut ExecutionCtx, ) -> Self { let byte_width = ptype.byte_width(); let alignment = Alignment::new(byte_width); @@ -547,8 +552,10 @@ impl PrimitiveData { Validity::AllValid | Validity::NonNullable => valid_elems_buffer.aligned(alignment), Validity::AllInvalid => ByteBuffer::zeroed_aligned(n_rows * byte_width, alignment), Validity::Array(is_valid) => { - #[expect(deprecated)] - let bool_array = is_valid.to_bool(); + let bool_array = is_valid + .clone() + .execute::(ctx) + .vortex_expect("must be a bool array"); let bool_buffer = bool_array.bit_buffer_view(); let mut bytes = ByteBufferMut::zeroed_aligned(n_rows * byte_width, alignment); for (i, valid_i) in bool_buffer.set_indices().enumerate() { @@ -596,6 +603,9 @@ impl PrimitiveData { } /// Try to extract a mutable buffer from the PrimitiveData with zero copy. + /// + /// # Panic + /// If the buffer is not of type T this will panic pub fn try_into_buffer_mut(self) -> Result, Buffer> { if T::PTYPE != self.ptype() { vortex_panic!( diff --git a/vortex-array/src/arrays/primitive/array/patch.rs b/vortex-array/src/arrays/primitive/array/patch.rs index 197db4d3854..d183552b66f 100644 --- a/vortex-array/src/arrays/primitive/array/patch.rs +++ b/vortex-array/src/arrays/primitive/array/patch.rs @@ -137,6 +137,8 @@ mod tests { use super::*; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; use crate::validity::Validity; @@ -173,13 +175,15 @@ mod tests { #[test] fn patch_sliced() { + let mut ctx = array_session().create_execution_ctx(); let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid); let sliced = input.slice(2..8).unwrap(); #[expect(deprecated)] let sliced_primitive = sliced.to_primitive(); assert_arrays_eq!( sliced_primitive, - PrimitiveArray::new(buffer![2u32; 6], Validity::AllValid) + PrimitiveArray::new(buffer![2u32; 6], Validity::AllValid), + &mut ctx ); } } diff --git a/vortex-array/src/arrays/primitive/array/top_value.rs b/vortex-array/src/arrays/primitive/array/top_value.rs index d3ee5eb5a65..7ae45c89405 100644 --- a/vortex-array/src/arrays/primitive/array/top_value.rs +++ b/vortex-array/src/arrays/primitive/array/top_value.rs @@ -17,7 +17,6 @@ use crate::arrays::primitive::NativeValue; use crate::dtype::NativePType; use crate::match_each_native_ptype; use crate::scalar::PValue; -use crate::validity::Validity; impl PrimitiveArray { /// Compute most common present value of this array @@ -26,7 +25,7 @@ impl PrimitiveArray { return Ok(None); } - if matches!(self.validity()?, Validity::AllInvalid) { + if self.validity()?.definitely_all_null() { return Ok(None); } diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index 10c0b8d6eba..36b95b4c554 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -5,6 +5,9 @@ use num_traits::AsPrimitive; 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::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -102,9 +105,7 @@ impl CastKernel for Primitive { } } -/// Cast values from `F` to `T`. For infallible casts this is a pure pass; for fallible casts -/// each valid value goes through a checked `NumCast::from` and the kernel bails if any of them -/// overflow `T`. Invalid positions use the wrapping `as` cast since their values are masked out. +/// Cast Primitive values from `F` to `T`. fn cast_values( array: ArrayView<'_, Primitive>, new_validity: Validity, @@ -114,53 +115,99 @@ where F: NativePType + AsPrimitive, T: NativePType, { - let values = array.as_slice::(); - - // Fast path: statically infallible, or cached min/max prove every valid value fits in `T`. - // The cached check never triggers a stats computation — if the bounds aren't already known - // we fall through to the per-lane loop below. - if values_always_fit(F::PTYPE, T::PTYPE) || values_fit_in(array, T::PTYPE, ctx, false) { - return Ok(PrimitiveArray::new(cast::(values), new_validity).into_array()); - } - - // TODO(joe): if the values source and target have the same bit-width we can - // mutate in place. - - // Fallible: invalid lanes are pre-multiplied to zero so the checked cast always succeeds for - // them; valid lanes go through `NumCast::from` and the whole cast bails on the first overflow. - let mask = array.validity()?.execute_mask(array.len(), ctx)?; let overflow = || { vortex_err!( Compute: "Cannot cast {} to {} — value exceeds target range", F::PTYPE, T::PTYPE, ) }; - let buffer: Buffer = match &mask { - Mask::AllTrue(_) => BufferMut::try_from_trusted_len_iter( + + // 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) + } + + // Skip the fallible kernel when type widening or (cached) min/max prove every value fits. + let target_dtype = DType::Primitive(T::PTYPE, Nullability::NonNullable); + let infallible = casts_losslessly_to(F::PTYPE, T::PTYPE) + || cached_values_fit_in(array, &target_dtype).unwrap_or(false); + + let len = array.len(); + + // If F and T have the same byte width, try to take unique ownership of the buffer. + let same_bit_width = F::PTYPE.byte_width() == T::PTYPE.byte_width(); + let owned: Option> = same_bit_width + .then(|| array.into_owned().try_into_buffer_mut::().ok()) + .flatten(); + let values: &[F] = array.as_slice::(); + + if infallible { + return match owned { + Some(mut buf) => { + ReinterpretSink::::new(buf.as_mut_slice()).map_into_in_place(|v: F| v.as_()); + // SAFETY: same size + alignment for NativePType + let result: BufferMut = unsafe { buf.transmute::() }; + Ok(PrimitiveArray::new(result.freeze(), new_validity).into_array()) + } + None => { + let mut buffer = BufferMut::::with_capacity(len); + values.map_into(&mut buffer.spare_capacity_mut()[..len], |v| v.as_()); + // SAFETY: map_into initializes every lane. + unsafe { buffer.set_len(len) }; + Ok(PrimitiveArray::new(buffer.freeze(), new_validity).into_array()) + } + }; + } + + let mask = array.validity()?.execute_mask(len, ctx)?; + + let buffer: Buffer = match (&mask, owned) { + (Mask::AllTrue(_), Some(mut buf)) => { + ReinterpretSink::::new(buf.as_mut_slice()) + .try_map_in_place(|v: F| ::from(v)) + .map_err(|_| overflow())?; + // SAFETY: same size + alignment for NativePType + let result: BufferMut = unsafe { buf.transmute::() }; + result.freeze() + } + (Mask::AllTrue(_), None) => { + let mut buffer = BufferMut::::with_capacity(len); values - .iter() - .map(|&v| ::from(v).ok_or_else(overflow)), - )? - .freeze(), - Mask::AllFalse(_) => BufferMut::::zeroed(values.len()).freeze(), - Mask::Values(m) => BufferMut::try_from_trusted_len_iter( - values.iter().zip(m.bit_buffer().iter()).map(|(&v, valid)| { - let factor = if valid { F::one() } else { F::zero() }; - ::from(v * factor).ok_or_else(overflow) - }), - )? - .freeze(), + .try_map_into(&mut buffer.spare_capacity_mut()[..len], |v| { + ::from(v) + }) + .map_err(|_| overflow())?; + // SAFETY: initialized every lane. + unsafe { buffer.set_len(len) }; + buffer.freeze() + } + (Mask::AllFalse(_), _) => BufferMut::::zeroed(len).freeze(), + (Mask::Values(m), Some(mut buf)) => { + ReinterpretSink::::new(buf.as_mut_slice()) + .try_map_masked_in_place(m.bit_buffer(), |v: F| ::from(v)) + .map_err(|_| overflow())?; + // SAFETY: same size + alignment for NativePType + let result: BufferMut = unsafe { buf.transmute::() }; + result.freeze() + } + (Mask::Values(m), None) => { + let mut buffer = BufferMut::::with_capacity(len); + values + .try_map_masked_into( + m.bit_buffer(), + &mut buffer.spare_capacity_mut()[..len], + |v| ::from(v), + ) + .map_err(|_| overflow())?; + // SAFETY: initialized every lane. + unsafe { buffer.set_len(len) }; + buffer.freeze() + } }; Ok(PrimitiveArray::new(buffer, new_validity).into_array()) } -/// Out-of-range values at invalid positions are truncated/wrapped by `as`, which is fine because -/// they are masked out by validity. -fn cast, T: NativePType>(array: &[F]) -> Buffer { - BufferMut::from_trusted_len_iter(array.iter().map(|&src| src.as_())).freeze() -} - fn reinterpret( array: ArrayView<'_, Primitive>, new_ptype: PType, @@ -178,23 +225,6 @@ fn reinterpret( .into_array() } -/// Returns `true` if every value of `src` is guaranteed representable in `target` without -/// overflow. Precision may be lost (e.g. large integers cast to `f32`), but the cast can never -/// produce an out-of-range result. -fn values_always_fit(src: PType, target: PType) -> bool { - if src == target { - return true; - } - if src.is_int() && target.is_int() { - return target.byte_width() > src.byte_width() - && (src.is_unsigned_int() || target.is_signed_int()); - } - if src.is_float() && target.is_float() { - return target.byte_width() > src.byte_width(); - } - src.is_int() && matches!(target, PType::F32 | PType::F64) -} - /// Returns `true` if all valid values in `array` are representable as `target_ptype`. /// /// Cached min/max statistics are consulted first. If either bound is missing, the function either @@ -213,10 +243,14 @@ fn values_fit_in( if !compute { return false; } - aggregate_fn::fns::min_max::min_max(array.array(), ctx) - .ok() - .flatten() - .is_none_or(|mm| mm.min.cast(&target_dtype).is_ok() && mm.max.cast(&target_dtype).is_ok()) + aggregate_fn::fns::min_max::min_max( + array.array(), + ctx, + aggregate_fn::NumericalAggregateOpts::default(), + ) + .ok() + .flatten() + .is_none_or(|mm| mm.min.cast(&target_dtype).is_ok() && mm.max.cast(&target_dtype).is_ok()) } /// Cached-only check: returns `Some(fits)` if both `Min` and `Max` are present as `Exact` in the @@ -236,9 +270,10 @@ mod test { use vortex_error::VortexError; use vortex_mask::Mask; + use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; @@ -252,12 +287,13 @@ mod test { #[test] fn cast_u32_u8() { + let mut ctx = array_session().create_execution_ctx(); let arr = buffer![0u32, 10, 200].into_array(); // cast from u32 to u8 #[expect(deprecated)] let p = arr.cast(PType::U8.into()).unwrap().to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200])); + assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200]), &mut ctx); assert!(matches!(p.validity(), Ok(Validity::NonNullable))); // to nullable @@ -269,7 +305,8 @@ mod test { .to_primitive(); assert_arrays_eq!( p, - PrimitiveArray::new(buffer![0u8, 10, 200], Validity::AllValid) + PrimitiveArray::new(buffer![0u8, 10, 200], Validity::AllValid), + &mut ctx ); assert!(matches!(p.validity(), Ok(Validity::AllValid))); @@ -280,7 +317,7 @@ mod test { .cast(DType::Primitive(PType::U8, Nullability::NonNullable)) .unwrap() .to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200])); + assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200]), &mut ctx); assert!(matches!(p.validity(), Ok(Validity::NonNullable))); // to nullable u32 @@ -292,7 +329,8 @@ mod test { .to_primitive(); assert_arrays_eq!( p, - PrimitiveArray::new(buffer![0u32, 10, 200], Validity::AllValid) + PrimitiveArray::new(buffer![0u32, 10, 200], Validity::AllValid), + &mut ctx ); assert!(matches!(p.validity(), Ok(Validity::AllValid))); @@ -303,16 +341,21 @@ mod test { .cast(DType::Primitive(PType::U8, Nullability::NonNullable)) .unwrap() .to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200])); + assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200]), &mut ctx); assert!(matches!(p.validity(), Ok(Validity::NonNullable))); } #[test] fn cast_u32_f32() { + let mut ctx = array_session().create_execution_ctx(); let arr = buffer![0u32, 10, 200].into_array(); #[expect(deprecated)] let u8arr = arr.cast(PType::F32.into()).unwrap().to_primitive(); - assert_arrays_eq!(u8arr, PrimitiveArray::from_iter([0.0f32, 10., 200.])); + assert_arrays_eq!( + u8arr, + PrimitiveArray::from_iter([0.0f32, 10., 200.]), + &mut ctx + ); } #[test] @@ -346,6 +389,7 @@ mod test { #[test] fn cast_with_invalid_nulls() { + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![-1i32, 0, 10], Validity::from_iter([false, true, true]), @@ -358,13 +402,17 @@ mod test { .to_primitive(); assert_arrays_eq!( p, - PrimitiveArray::from_option_iter([None, Some(0u32), Some(10)]) + PrimitiveArray::from_option_iter([None, Some(0u32), Some(10)]), + &mut ctx ); assert_eq!( p.as_ref() .validity() .unwrap() - .execute_mask(p.as_ref().len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask( + p.as_ref().len(), + &mut array_session().create_execution_ctx() + ) .unwrap(), Mask::from(BitBuffer::from(vec![false, true, true])) ); @@ -374,6 +422,7 @@ mod test { /// buffer without allocation (pointer identity). #[test] fn cast_same_width_int_reinterprets_buffer() -> vortex_error::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(); @@ -383,7 +432,7 @@ mod test { // Zero-copy: the data pointer should be identical. assert_eq!(src_ptr as usize, dst_ptr as usize); - assert_arrays_eq!(dst, PrimitiveArray::from_iter([0i32, 10, 100])); + assert_arrays_eq!(dst, PrimitiveArray::from_iter([0i32, 10, 100]), &mut ctx); Ok(()) } @@ -419,6 +468,7 @@ mod test { /// not prevent the cast from succeeding. #[test] fn cast_same_width_int_nullable_with_out_of_range_nulls() -> vortex_error::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. let arr = PrimitiveArray::new( @@ -432,13 +482,15 @@ mod test { .to_primitive(); assert_arrays_eq!( casted, - PrimitiveArray::from_option_iter([None, Some(0i32), Some(42)]) + PrimitiveArray::from_option_iter([None, Some(0i32), Some(42)]), + &mut ctx ); Ok(()) } #[test] fn cast_u32_to_u8_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![1000u32, 10u32, 42u32], Validity::from_iter([false, true, true]), @@ -450,7 +502,8 @@ mod test { .to_primitive(); assert_arrays_eq!( casted, - PrimitiveArray::from_option_iter([None, Some(10u8), Some(42)]) + PrimitiveArray::from_option_iter([None, Some(10u8), Some(42)]), + &mut ctx ); Ok(()) } @@ -471,7 +524,7 @@ mod test { #[case(PrimitiveArray::from_option_iter([Some(1u8), None, Some(255), Some(0), None]).into_array())] #[case(PrimitiveArray::from_option_iter([Some(1i32), None, Some(-100), Some(0), None]).into_array())] #[case(buffer![42u32].into_array())] - fn test_cast_primitive_conformance(#[case] array: crate::ArrayRef) { + fn test_cast_primitive_conformance(#[case] array: ArrayRef) { test_cast_conformance(&array); } } diff --git a/vortex-array/src/arrays/primitive/compute/fill_null.rs b/vortex-array/src/arrays/primitive/compute/fill_null.rs index 67089bd534e..d73442a150e 100644 --- a/vortex-array/src/arrays/primitive/compute/fill_null.rs +++ b/vortex-array/src/arrays/primitive/compute/fill_null.rs @@ -51,8 +51,8 @@ mod test { use vortex_buffer::buffer; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::compute::fill_null::BoolArray; use crate::assert_arrays_eq; @@ -64,6 +64,7 @@ mod test { #[test] fn fill_null_leading_none() { + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([None, Some(8u8), None, Some(10), None]); #[expect(deprecated)] let p = arr @@ -71,12 +72,19 @@ mod test { .fill_null(Scalar::from(42u8)) .unwrap() .to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([42u8, 8, 42, 10, 42])); + assert_arrays_eq!( + p, + PrimitiveArray::from_iter([42u8, 8, 42, 10, 42]), + &mut ctx + ); assert!( p.as_ref() .validity() .unwrap() - .execute_mask(p.as_ref().len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask( + p.as_ref().len(), + &mut array_session().create_execution_ctx() + ) .unwrap() .all_true() ); @@ -84,6 +92,7 @@ mod test { #[test] fn fill_null_all_none() { + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([Option::::None, None, None, None, None]); #[expect(deprecated)] @@ -92,12 +101,19 @@ mod test { .fill_null(Scalar::from(255u8)) .unwrap() .to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([255u8, 255, 255, 255, 255])); + assert_arrays_eq!( + p, + PrimitiveArray::from_iter([255u8, 255, 255, 255, 255]), + &mut ctx + ); assert!( p.as_ref() .validity() .unwrap() - .execute_mask(p.as_ref().len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask( + p.as_ref().len(), + &mut array_session().create_execution_ctx() + ) .unwrap() .all_true() ); @@ -105,6 +121,7 @@ mod test { #[test] fn fill_null_nullable_non_null() { + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![8u8, 10, 12, 14, 16], Validity::Array(BoolArray::from_iter([true, true, true, true, true]).into_array()), @@ -115,12 +132,19 @@ mod test { .fill_null(Scalar::from(255u8)) .unwrap() .to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([8u8, 10, 12, 14, 16])); + assert_arrays_eq!( + p, + PrimitiveArray::from_iter([8u8, 10, 12, 14, 16]), + &mut ctx + ); assert!( p.as_ref() .validity() .unwrap() - .execute_mask(p.as_ref().len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask( + p.as_ref().len(), + &mut array_session().create_execution_ctx() + ) .unwrap() .all_true() ); @@ -128,15 +152,23 @@ mod test { #[test] fn fill_null_non_nullable() { + let mut ctx = array_session().create_execution_ctx(); let arr = buffer![8u8, 10, 12, 14, 16].into_array(); #[expect(deprecated)] let p = arr.fill_null(Scalar::from(255u8)).unwrap().to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([8u8, 10, 12, 14, 16])); + assert_arrays_eq!( + p, + PrimitiveArray::from_iter([8u8, 10, 12, 14, 16]), + &mut ctx + ); assert!( p.as_ref() .validity() .unwrap() - .execute_mask(p.as_ref().len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask( + p.as_ref().len(), + &mut array_session().create_execution_ctx() + ) .unwrap() .all_true() ); diff --git a/vortex-array/src/arrays/primitive/compute/mod.rs b/vortex-array/src/arrays/primitive/compute/mod.rs index bd59efed20b..1ca4b17d3d5 100644 --- a/vortex-array/src/arrays/primitive/compute/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/mod.rs @@ -15,7 +15,10 @@ mod tests { use rstest::rstest; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::compute::conformance::binary_numeric::test_binary_numeric_array; use crate::compute::conformance::consistency::test_array_consistency; #[rstest] @@ -36,7 +39,10 @@ mod tests { #[case::i16_negative(PrimitiveArray::from_iter([-100i16, -50, 0, 50, 100]))] #[case::f64_special(PrimitiveArray::from_iter([0.0f64, 1.0, -1.0, f64::INFINITY, f64::NEG_INFINITY]))] fn test_primitive_consistency(#[case] array: PrimitiveArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[rstest] @@ -46,7 +52,9 @@ mod tests { #[case::f32(PrimitiveArray::from_iter([1.5f32, 2.5, 3.5, 4.5, 5.5]))] #[case::f64(PrimitiveArray::from_iter([10.1f64, 20.2, 30.3, 40.4, 50.5]))] fn test_primitive_binary_numeric(#[case] array: PrimitiveArray) { - use crate::compute::conformance::binary_numeric::test_binary_numeric_array; - test_binary_numeric_array(array.into_array()); + test_binary_numeric_array( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/primitive/compute/slice.rs b/vortex-array/src/arrays/primitive/compute/slice.rs index 4775adba749..9ac45e3baae 100644 --- a/vortex-array/src/arrays/primitive/compute/slice.rs +++ b/vortex-array/src/arrays/primitive/compute/slice.rs @@ -20,7 +20,7 @@ impl SliceReduce for Primitive { let validity = array.validity()?.slice(range)?; // SAFETY: - //slicing an existing PrimitiveArray on element boundaries preserves the buffer + // slicing an existing PrimitiveArray on element boundaries preserves the buffer // alignment, ptype, length, and validity invariants. let array = unsafe { PrimitiveArray::new_unchecked_from_handle(values, array.ptype(), validity).into_array() diff --git a/vortex-array/src/arrays/primitive/compute/take/avx2.rs b/vortex-array/src/arrays/primitive/compute/take/avx2.rs index e92304dc34b..fb64e6748b0 100644 --- a/vortex-array/src/arrays/primitive/compute/take/avx2.rs +++ b/vortex-array/src/arrays/primitive/compute/take/avx2.rs @@ -46,7 +46,6 @@ 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::PType; use crate::dtype::UnsignedPType; use crate::match_each_native_ptype; use crate::match_each_unsigned_integer_ptype; @@ -112,6 +111,16 @@ where /// 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`. @@ -121,60 +130,29 @@ where /// The caller must ensure the `avx2` feature is enabled. #[target_feature(enable = "avx2")] #[doc(hidden)] -unsafe fn take_avx2(buffer: &[V], indices: &[I]) -> Buffer { - macro_rules! dispatch_avx2 { - ($indices:ty, $values:ty) => { - { let result = dispatch_avx2!($indices, $values, cast: $values); result } - }; - ($indices:ty, $values:ty, cast: $cast:ty) => {{ - let indices = unsafe { std::mem::transmute::<&[I], &[$indices]>(indices) }; - let values = unsafe { std::mem::transmute::<&[V], &[$cast]>(buffer) }; - - let result = exec_take::<$cast, $indices, AVX2Gather>(values, indices); - unsafe { result.transmute::() } - }}; - } - +unsafe fn take_avx2(buffer: &[V], indices: &[I]) -> Buffer { if buffer.is_empty() { return Buffer::zeroed(indices.len()); } - match (I::PTYPE, V::PTYPE) { - // Int value types. Only 32 and 64 bit types are supported. - (PType::U8, PType::I32) => dispatch_avx2!(u8, i32), - (PType::U8, PType::U32) => dispatch_avx2!(u8, u32), - (PType::U8, PType::I64) => dispatch_avx2!(u8, i64), - (PType::U8, PType::U64) => dispatch_avx2!(u8, u64), - (PType::U16, PType::I32) => dispatch_avx2!(u16, i32), - (PType::U16, PType::U32) => dispatch_avx2!(u16, u32), - (PType::U16, PType::I64) => dispatch_avx2!(u16, i64), - (PType::U16, PType::U64) => dispatch_avx2!(u16, u64), - (PType::U32, PType::I32) => dispatch_avx2!(u32, i32), - (PType::U32, PType::U32) => dispatch_avx2!(u32, u32), - (PType::U32, PType::I64) => dispatch_avx2!(u32, i64), - (PType::U32, PType::U64) => dispatch_avx2!(u32, u64), - - // Float value types, treat them as if they were corresponding int types. - (PType::U8, PType::F32) => dispatch_avx2!(u8, f32, cast: u32), - (PType::U16, PType::F32) => dispatch_avx2!(u16, f32, cast: u32), - (PType::U32, PType::F32) => dispatch_avx2!(u32, f32, cast: u32), - (PType::U64, PType::F32) => dispatch_avx2!(u64, f32, cast: u32), - - (PType::U8, PType::F64) => dispatch_avx2!(u8, f64, cast: u64), - (PType::U16, PType::F64) => dispatch_avx2!(u16, f64, cast: u64), - (PType::U32, PType::F64) => dispatch_avx2!(u32, f64, cast: u64), - (PType::U64, PType::F64) => dispatch_avx2!(u64, f64, cast: u64), - - // Scalar fallback for unsupported value types. - _ => { - tracing::trace!( - "take AVX2 kernel missing for indices {} values {}, falling back to scalar", - I::PTYPE, - V::PTYPE - ); - - take_primitive_scalar(buffer, indices) - } + // 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), } } @@ -259,17 +237,6 @@ impl_gather!(u8, store: _mm256_storeu_si256, WIDTH = 8, STRIDE = 16 }, - { i32 => - 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 => @@ -282,17 +249,6 @@ impl_gather!(u8, gather: _mm256_mask_i64gather_epi64, store: _mm256_storeu_si256, WIDTH = 4, STRIDE = 16 - }, - { i64 => - 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 } ); @@ -310,17 +266,6 @@ impl_gather!(u16, store: _mm256_storeu_si256, WIDTH = 8, STRIDE = 8 }, - { i32 => - 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 => @@ -333,17 +278,6 @@ impl_gather!(u16, gather: _mm256_mask_i64gather_epi64, store: _mm256_storeu_si256, WIDTH = 4, STRIDE = 8 - }, - { i64 => - 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 } ); @@ -361,17 +295,6 @@ impl_gather!(u32, store: _mm256_storeu_si256, WIDTH = 8, STRIDE = 8 }, - { i32 => - 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 => @@ -384,17 +307,6 @@ impl_gather!(u32, gather: _mm256_mask_i64gather_epi64, store: _mm256_storeu_si256, WIDTH = 4, STRIDE = 4 - }, - { i64 => - 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 } ); @@ -419,25 +331,6 @@ impl_gather!(u64, store: _mm_storeu_si128, WIDTH = 4, STRIDE = 4 }, - { i32 => - 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 => @@ -450,32 +343,36 @@ impl_gather!(u64, gather: _mm256_mask_i64gather_epi64, store: _mm256_storeu_si256, WIDTH = 4, STRIDE = 4 - }, - { i64 => - 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 certain `Idx` and `Value` type. +/// 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: &[Value], indices: &[Idx]) -> Buffer +fn exec_take(values: &[Out], indices: &[Idx]) -> Buffer where - Value: Copy, + Out: Copy, Idx: UnsignedPType, - Gather: GatherFn, + 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>()); + BufferMut::::with_capacity_aligned(indices_len, Alignment::of::<__m256i>()); let buf_uninit = buffer.spare_capacity_mut(); let mut offset = 0; @@ -483,16 +380,18 @@ where // 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_simd` preconditions satisfied: + // 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 + STRIDE` is always valid. + // 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(), - buf_uninit.as_mut_ptr().add(offset).cast(), + values.as_ptr().cast::(), + buf_uninit.as_mut_ptr().add(offset).cast::(), ) }; offset += Gather::WIDTH; @@ -509,11 +408,11 @@ where // SAFETY: All elements have been initialized. unsafe { buffer.set_len(indices_len) }; - // Reset the buffer alignment to the Value type. + // 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 = buffer.aligned(Alignment::of::()); buffer.freeze() } @@ -601,4 +500,37 @@ mod avx2_tests { 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 index 4023991c65d..9294659708e 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -10,20 +10,22 @@ 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::dtype::NativePType; 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 @@ -81,19 +83,36 @@ impl TakeExecute for Primitive { 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.clone().execute::(ctx)? + indices_nulls_zeroed.execute::(ctx)? } else { // This will fail if all values cannot be converted to unsigned - indices - .clone() + indices_nulls_zeroed .cast(DType::Primitive(ptype.to_unsigned(), *null))? .execute::(ctx)? }; let validity = array .validity()? - .take(&unsigned_indices.clone().into_array())?; + .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(); @@ -106,10 +125,7 @@ impl TakeExecute for Primitive { // Compiler may see this as unused based on enabled features #[inline(always)] -fn take_primitive_scalar( - buffer: &[T], - indices: &[I], -) -> Buffer { +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. @@ -136,8 +152,8 @@ mod test { use vortex_error::VortexExpect; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::compute::take::take_primitive_scalar; @@ -165,21 +181,21 @@ mod test { let actual = values.take(indices.into_array()).unwrap(); assert_eq!( actual - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .vortex_expect("no fail"), Scalar::from(Some(1)) ); // position 3 is null assert_eq!( actual - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .vortex_expect("no fail"), Scalar::null_native::() ); // the third index is null assert_eq!( actual - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .vortex_expect("no fail"), Scalar::null_native::() ); @@ -200,3 +216,34 @@ mod test { test_take_conformance(&array.into_array()); } } + +#[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/compute/zip.rs b/vortex-array/src/arrays/primitive/compute/zip.rs index 0dcb7298460..83f03b10614 100644 --- a/vortex-array/src/arrays/primitive/compute/zip.rs +++ b/vortex-array/src/arrays/primitive/compute/zip.rs @@ -152,8 +152,8 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -170,7 +170,7 @@ mod tests { let bits: Vec = (0..len).map(|i| i.is_multiple_of(3) || i == 64).collect(); let mask = Mask::from_iter(bits.iter().copied()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -181,7 +181,7 @@ mod tests { (0..len).map(|i| if bits[i] { i as i64 } else { 1_000 + i as i64 }), ) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -201,7 +201,7 @@ mod tests { let bits: Vec = (0..len).map(|i| i.is_multiple_of(2)).collect(); let mask = Mask::from_iter(bits.iter().copied()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -217,7 +217,7 @@ mod tests { } })) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/arrays/primitive/mod.rs b/vortex-array/src/arrays/primitive/mod.rs index 4d62f1da517..f1566f8174f 100644 --- a/vortex-array/src/arrays/primitive/mod.rs +++ b/vortex-array/src/arrays/primitive/mod.rs @@ -15,8 +15,13 @@ mod vtable; pub use compute::rules::PrimitiveMaskedValidityRule; pub use vtable::Primitive; +pub(crate) fn initialize(session: &VortexSession) { + vtable::initialize(session); +} + mod native_value; pub use native_value::NativeValue; +use vortex_session::VortexSession; #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/primitive/vtable/kernel.rs b/vortex-array/src/arrays/primitive/vtable/kernel.rs index 341d2f21272..6382ea73794 100644 --- a/vortex-array/src/arrays/primitive/vtable/kernel.rs +++ b/vortex-array/src/arrays/primitive/vtable/kernel.rs @@ -1,18 +1,36 @@ // 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::Primitive; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::between::BetweenExecuteAdaptor; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; +use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&BetweenExecuteAdaptor(Primitive)), - ParentKernelSet::lift(&CastExecuteAdaptor(Primitive)), - ParentKernelSet::lift(&FillNullExecuteAdaptor(Primitive)), - ParentKernelSet::lift(&TakeExecuteAdaptor(Primitive)), - ParentKernelSet::lift(&ZipExecuteAdaptor(Primitive)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Between.id(), + Primitive, + BetweenExecuteAdaptor(Primitive), + ); + kernels.register_execute_parent_kernel(Cast.id(), Primitive, CastExecuteAdaptor(Primitive)); + kernels.register_execute_parent_kernel( + FillNull.id(), + Primitive, + FillNullExecuteAdaptor(Primitive), + ); + kernels.register_execute_parent_kernel(Dict.id(), Primitive, TakeExecuteAdaptor(Primitive)); + kernels.register_execute_parent_kernel(Zip.id(), Primitive, ZipExecuteAdaptor(Primitive)); +} diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index 4b472d122f3..a184253527e 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use kernel::PARENT_KERNELS; 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; use crate::ExecutionCtx; use crate::ExecutionResult; @@ -15,8 +15,11 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::arrays::primitive::PrimitiveData; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::PrimitiveBuilder; use crate::dtype::DType; use crate::dtype::PType; +use crate::match_each_native_ptype; use crate::serde::ArrayChildren; use crate::validity::Validity; mod kernel; @@ -39,6 +42,10 @@ use crate::hash::ArrayHash; /// A [`Primitive`]-encoded Vortex array. pub type PrimitiveArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + impl ArrayHash for PrimitiveData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.buffer.array_hash(state, accuracy); @@ -80,6 +87,24 @@ impl VTable for Primitive { } } + fn with_buffers( + &self, + 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(); + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + fn serialize( _array: ArrayView<'_, Self>, _session: &VortexSession, @@ -125,7 +150,7 @@ impl VTable for Primitive { buffers: &[BufferHandle], children: &dyn ArrayChildren, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult> { if !metadata.is_empty() { vortex_bail!( "PrimitiveArray expects empty metadata, got {} bytes", @@ -173,7 +198,7 @@ impl VTable for Primitive { // SAFETY: checked ahead of time let slots = PrimitiveData::make_slots(&validity, len); let data = unsafe { PrimitiveData::new_unchecked_from_handle(buffer, ptype, validity) }; - Ok(crate::array::ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { @@ -184,21 +209,27 @@ impl VTable for Primitive { Ok(ExecutionResult::done(array)) } - fn reduce_parent( + fn append_to_builder( array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - RULES.evaluate(array, parent, child_idx) + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + match_each_native_ptype!(array.ptype(), |P| { + if let Some(builder) = builder.as_any_mut().downcast_mut::>() { + return builder.append_primitive_array(&array.into_owned(), ctx); + } + }); + + builder.extend_from_array(array.as_ref()); + Ok(()) } - fn execute_parent( + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, child_idx: usize, - ctx: &mut ExecutionCtx, ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) + RULES.evaluate(array, parent, child_idx) } } @@ -209,19 +240,24 @@ pub struct Primitive; mod tests { use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_session::registry::ReadContext; use crate::ArrayContext; use crate::IntoArray; - use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; + use crate::buffer::BufferHandle; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; use crate::validity::Validity; #[test] fn test_nullable_primitive_serde_roundtrip() { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); let array = PrimitiveArray::new( buffer![1i32, 2, 3, 4], Validity::from_iter([true, false, true, false]), @@ -229,11 +265,11 @@ mod tests { let dtype = array.dtype().clone(); let len = array.len(); - let ctx = ArrayContext::empty(); + let array_ctx = ArrayContext::empty(); let serialized = array .clone() .into_array() - .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default()) + .serialize(&array_ctx, &session, &SerializeOptions::default()) .unwrap(); let mut concat = ByteBufferMut::empty(); @@ -242,14 +278,35 @@ mod tests { } let parts = SerializedArray::try_from(concat.freeze()).unwrap(); let decoded = parts - .decode( - &dtype, - len, - &ReadContext::new(ctx.to_ids()), - &LEGACY_SESSION, - ) + .decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session) .unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut ctx); + } + + #[test] + fn test_with_buffers_replaces_primitive_buffer_with_equivalent_contents() -> VortexResult<()> { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + + let array = PrimitiveArray::from_iter([1i32, 2, 3, 4]).into_array(); + let replacement = BufferHandle::new_host(buffer![1i32, 2, 3, 4].into_byte_buffer()); + // SAFETY: the replacement buffer contains the same logical values as the original array; + // only the buffer handle changes. + let rewritten = unsafe { array.with_buffers([replacement]) }?; + let expected = PrimitiveArray::from_iter([1i32, 2, 3, 4]); + + assert_arrays_eq!(rewritten, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_with_buffers_rejects_length_change() { + let array = PrimitiveArray::from_iter([1i32, 2, 3, 4]).into_array(); + let replacement = BufferHandle::new_host(buffer![10i32, 20, 30].into_byte_buffer()); + + // SAFETY: this call is expected to fail the checked buffer length invariant before any + // rewritten array is returned or observed. + assert!(unsafe { array.with_buffers([replacement]) }.is_err()); } } diff --git a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs index 6833848deae..d98e00975ea 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs @@ -28,6 +28,7 @@ use crate::array::ArrayId; use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::VTable; +use crate::array::with_empty_buffers; use crate::arrays::scalar_fn::array::ScalarFnArrayExt; use crate::arrays::scalar_fn::array::ScalarFnData; use crate::arrays::scalar_fn::rules::PARENT_RULES; @@ -116,6 +117,14 @@ impl VTable for ScalarFn { None } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + fn serialize( _array: ArrayView<'_, Self>, _session: &VortexSession, diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index 67de19d936f..c19e3e2288c 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -60,8 +60,8 @@ mod tests { use crate::Canonical; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::ScalarFnArray; @@ -76,6 +76,7 @@ mod tests { #[test] fn test_scalar_fn_add() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let lhs = buffer![1i32, 2, 3].into_array(); let rhs = buffer![10i32, 20, 30].into_array(); @@ -86,10 +87,10 @@ mod tests { let result = scalar_fn_array .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut array_session().create_execution_ctx())? .into_array(); let expected = buffer![11i32, 22, 33].into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -130,6 +131,7 @@ mod tests { #[test] fn test_scalar_fn_mul() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let lhs = buffer![2i32, 3, 4].into_array(); let rhs = buffer![5i32, 6, 7].into_array(); @@ -138,16 +140,17 @@ mod tests { let result = scalar_fn_array .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut array_session().create_execution_ctx())? .into_array(); let expected = buffer![10i32, 18, 28].into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } #[test] fn test_scalar_fn_with_nullable() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let lhs = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let rhs = PrimitiveArray::new( buffer![10i32, 20, 30], @@ -160,20 +163,21 @@ mod tests { let result = scalar_fn_array .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut array_session().create_execution_ctx())? .into_array(); let expected = PrimitiveArray::new( buffer![11i32, 0, 33], Validity::from_iter([true, false, true]), ) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } #[test] fn test_scalar_fn_comparison() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let lhs = buffer![1i32, 5, 3].into_array(); let rhs = buffer![2i32, 5, 1].into_array(); @@ -182,10 +186,10 @@ mod tests { let result = scalar_fn_array .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut array_session().create_execution_ctx())? .into_array(); let expected = BoolArray::from_iter([false, true, false]).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/shared/tests.rs b/vortex-array/src/arrays/shared/tests.rs index b8301389074..d1a33ed3e20 100644 --- a/vortex-array/src/arrays/shared/tests.rs +++ b/vortex-array/src/arrays/shared/tests.rs @@ -3,17 +3,15 @@ use vortex_buffer::buffer; use vortex_error::VortexResult; -use vortex_session::VortexSession; use crate::Canonical; -use crate::ExecutionCtx; use crate::IntoArray; +use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::SharedArray; use crate::arrays::shared::SharedArrayExt; use crate::hash::ArrayEq; use crate::hash::EqMode; -use crate::session::ArraySession; use crate::validity::Validity; #[test] @@ -21,8 +19,8 @@ fn shared_array_caches_on_canonicalize() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable).into_array(); let shared = SharedArray::new(array); - let session = VortexSession::empty().with::(); - let mut ctx = ExecutionCtx::new(session); + let session = crate::array_session(); + let mut ctx = session.create_execution_ctx(); let first = shared.get_or_compute(|source| source.clone().execute::(&mut ctx))?; diff --git a/vortex-array/src/arrays/shared/vtable.rs b/vortex-array/src/arrays/shared/vtable.rs index 3c3a09216d2..fb2c67d1d65 100644 --- a/vortex-array/src/arrays/shared/vtable.rs +++ b/vortex-array/src/arrays/shared/vtable.rs @@ -11,6 +11,7 @@ use vortex_session::registry::CachedId; use crate::ArrayEq; use crate::ArrayHash; +use crate::ArrayParts; use crate::ArrayRef; use crate::Canonical; use crate::EqMode; @@ -22,6 +23,7 @@ use crate::array::ArrayView; use crate::array::OperationsVTable; use crate::array::VTable; 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; @@ -84,6 +86,14 @@ impl VTable for Shared { 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 { SLOT_NAMES[idx].to_string() } @@ -104,7 +114,7 @@ impl VTable for Shared { _buffers: &[BufferHandle], _children: &dyn crate::serde::ArrayChildren, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult> { vortex_error::vortex_bail!("Shared array is not serializable") } diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index 1ca91254123..1ff2e5bc84d 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -18,6 +18,7 @@ use vortex_session::registry::CachedId; use crate::AnyCanonical; use crate::ArrayEq; use crate::ArrayHash; +use crate::ArrayParts; use crate::ArrayRef; use crate::EqMode; use crate::array::Array; @@ -26,6 +27,7 @@ use crate::array::ArrayView; 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; @@ -115,6 +117,14 @@ impl VTable for Slice { 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 { SLOT_NAMES[idx].to_string() } @@ -136,7 +146,7 @@ impl VTable for Slice { _buffers: &[BufferHandle], _children: &dyn ArrayChildren, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult> { vortex_bail!("Slice array is not serializable") } @@ -188,18 +198,21 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::SliceArray; use crate::assert_arrays_eq; #[test] fn test_slice_slice() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Slice(1..4, Slice(2..8, base)) combines to Slice(3..6, base) let arr = PrimitiveArray::from_iter(0i32..10).into_array(); let inner_slice = SliceArray::new(arr, 2..8).into_array(); let slice = inner_slice.slice(1..4)?; - assert_arrays_eq!(slice, PrimitiveArray::from_iter([3i32, 4, 5])); + assert_arrays_eq!(slice, PrimitiveArray::from_iter([3i32, 4, 5]), &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/struct_/array.rs b/vortex-array/src/arrays/struct_/array.rs index 7feffb5484c..8b939deb911 100644 --- a/vortex-array/src/arrays/struct_/array.rs +++ b/vortex-array/src/arrays/struct_/array.rs @@ -22,6 +22,7 @@ use crate::array::child_to_validity; use crate::array::validity_to_child; use crate::arrays::ChunkedArray; use crate::arrays::Struct; +use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; @@ -57,7 +58,7 @@ pub(super) const FIELDS_OFFSET: usize = 1; /// use vortex_array::arrays::{StructArray, BoolArray}; /// use vortex_array::validity::Validity; /// use vortex_array::dtype::FieldNames; -/// use vortex_array::{IntoArray, LEGACY_SESSION, VortexSessionExecute}; +/// use vortex_array::{IntoArray, VortexSessionExecute, array_session}; /// use vortex_buffer::buffer; /// /// // Create struct with all non-null fields but struct-level nulls @@ -71,7 +72,7 @@ pub(super) const FIELDS_OFFSET: usize = 1; /// 2, /// Validity::Array(BoolArray::from_iter([true, false]).into_array()), // row 1 is null /// ).unwrap(); -/// let mut ctx = LEGACY_SESSION.create_execution_ctx(); +/// let mut ctx = array_session().create_execution_ctx(); /// /// // Row 0 is valid - returns a struct scalar with field values /// let row0 = struct_array.execute_scalar(0, &mut ctx).unwrap(); @@ -92,7 +93,7 @@ pub(super) const FIELDS_OFFSET: usize = 1; /// use vortex_array::arrays::struct_::StructArrayExt; /// use vortex_array::validity::Validity; /// use vortex_array::dtype::FieldNames; -/// use vortex_array::{IntoArray, LEGACY_SESSION, VortexSessionExecute}; +/// use vortex_array::{IntoArray, VortexSessionExecute, array_session}; /// use vortex_buffer::buffer; /// /// // Create struct with duplicate "data" field names @@ -108,7 +109,7 @@ pub(super) const FIELDS_OFFSET: usize = 1; /// /// // field_by_name returns the FIRST "data" field /// let first_data = struct_array.unmasked_field_by_name("data").unwrap(); -/// let mut ctx = LEGACY_SESSION.create_execution_ctx(); +/// let mut ctx = array_session().create_execution_ctx(); /// assert_eq!(first_data.execute_scalar(0, &mut ctx).unwrap(), 1i32.into()); /// ``` /// @@ -525,4 +526,38 @@ impl Array { // the correct length and dtype harmony. Ok(unsafe { Array::::new_unchecked(field_arrays, struct_fields, len, validity) }) } + + /// Push the struct's top-level validity into each field, so a row null at the struct level + /// becomes null in every field. + /// + /// If `remove_struct_validity` is set the result is non-nullable; otherwise it keeps its + /// top-level validity. + pub fn push_validity_into_children(&self, remove_struct_validity: bool) -> VortexResult { + let struct_validity = self.struct_validity(); + + let new_validity = if remove_struct_validity { + Validity::NonNullable + } else { + struct_validity.clone() + }; + + // Nothing to push down. + if struct_validity.definitely_no_nulls() { + return Self::try_new( + self.names().clone(), + self.unmasked_fields(), + self.len(), + new_validity, + ); + } + + // Null each field where the struct row is null. + let mask = struct_validity.to_array(self.len()); + let fields = self + .iter_unmasked_fields() + .map(|field| field.clone().mask(mask.clone())) + .collect::>>()?; + + Self::try_new(self.names().clone(), fields, 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 f954dbe2fea..a7e1e760fd3 100644 --- a/vortex-array/src/arrays/struct_/compute/cast.rs +++ b/vortex-array/src/arrays/struct_/compute/cast.rs @@ -4,6 +4,7 @@ use itertools::Itertools; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use vortex_session::VortexSession; use crate::ArrayRef; use crate::ArrayView; @@ -13,33 +14,42 @@ use crate::arrays::ConstantArray; use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::scalar_fn::ExactScalarFn; +use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::arrays::struct_::StructArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::StructFields; -use crate::matcher::Matcher; +use crate::kernel::ExecuteParentKernel; +use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar::Scalar; +use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; -pub(crate) fn struct_cast_execute_parent( - child: &ArrayRef, - parent: &ArrayRef, - _child_idx: usize, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let Some(array) = child.as_opt::() else { - return Ok(None); - }; - let Some(parent) = ExactScalarFn::::try_match(parent) else { - return Ok(None); - }; +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), Struct, StructCastKernel); +} - let dtype = parent.options; - if array.dtype() == parent.options { - return Ok(Some(array.array().clone())); - } +#[derive(Debug)] +struct StructCastKernel; + +impl ExecuteParentKernel for StructCastKernel { + type Parent = ExactScalarFn; - struct_cast(array, dtype, ctx) + fn execute_parent( + &self, + array: ArrayView<'_, Struct>, + parent: ScalarFnArrayView<'_, Cast>, + _child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let dtype = parent.options; + if array.dtype() == parent.options { + return Ok(Some(array.array().clone())); + } + + struct_cast(array, dtype, ctx) + } } pub(crate) fn struct_cast( @@ -140,16 +150,14 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; - use crate::optimizer::kernels::ArrayKernels; use crate::optimizer::kernels::ArrayKernelsExt; use crate::optimizer::kernels::ExecuteParentFn; + use crate::optimizer::kernels::KernelSession; use crate::scalar::Scalar; use crate::scalar_fn::fns::cast::Cast; - use crate::session::ArraySession; use crate::validity::Validity; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); fn null_struct_cast_execute_parent( child: &ArrayRef, @@ -183,25 +191,6 @@ mod tests { test_cast_conformance(&array.into_array()); } - #[test] - fn struct_cast_execute_parent_is_not_static_kernel() { - let source = create_simple_struct().into_array(); - let target = DType::struct_( - [( - "value", - DType::Primitive(PType::I64, Nullability::NonNullable), - )], - Nullability::NonNullable, - ); - - let cast = Cast - .try_new_array(source.len(), target, [source.clone()]) - .unwrap(); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); - - assert!(source.execute_parent(&cast, 0, &mut ctx).unwrap().is_none()); - } - #[test] fn struct_cast_execute_parent_uses_session_plugin() { let source = StructArray::try_new( @@ -224,8 +213,9 @@ mod tests { .try_new_array(source.len(), target.clone(), [source]) .unwrap(); let parent_id = cast.encoding_id(); - let session = VortexSession::empty().with::(); - session.kernels().register_execute_parent( + let session = VortexSession::empty().with_some(KernelSession::empty()); + let kernels = session.kernels(); + kernels.register_execute_parent( parent_id, child_id, &[null_struct_cast_execute_parent as ExecuteParentFn], @@ -237,7 +227,8 @@ mod tests { assert_eq!(result.dtype(), &target); assert_arrays_eq!( result.unmasked_field_by_name("b").unwrap(), - ConstantArray::new(Scalar::null(utf8_null), 1) + ConstantArray::new(Scalar::null(utf8_null), 1), + &mut ctx ); } diff --git a/vortex-array/src/arrays/struct_/compute/mod.rs b/vortex-array/src/arrays/struct_/compute/mod.rs index c30ef532008..c1539a638bf 100644 --- a/vortex-array/src/arrays/struct_/compute/mod.rs +++ b/vortex-array/src/arrays/struct_/compute/mod.rs @@ -19,9 +19,9 @@ mod tests { use crate::Canonical; use crate::IntoArray as _; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_constant::is_constant; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -40,6 +40,7 @@ mod tests { #[test] fn take_empty_struct() { + let mut ctx = array_session().create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let indices = PrimitiveArray::from_option_iter([Some(1), None]); @@ -52,7 +53,8 @@ mod tests { vec![], 2, Validity::from_iter([true, false]) - ) + ), + &mut ctx ); } @@ -67,13 +69,13 @@ mod tests { let taken = struct_arr .take(indices.into_array()) .unwrap() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut array_session().create_execution_ctx()) .unwrap(); assert_eq!(taken.len(), 1); assert!( taken .into_array() - .all_invalid(&mut LEGACY_SESSION.create_execution_ctx()) + .all_invalid(&mut array_session().create_execution_ctx()) .unwrap() ); } @@ -85,13 +87,14 @@ mod tests { let taken = arr .take(indices.into_array()) .unwrap() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut array_session().create_execution_ctx()) .unwrap(); assert_eq!(taken.len(), 1); } #[test] fn take_field_struct() { + let mut ctx = array_session().create_execution_ctx(); let struct_arr = StructArray::from_fields(&[("a", PrimitiveArray::from_iter(0..10).into_array())]) .unwrap(); @@ -103,7 +106,8 @@ mod tests { [("a", buffer![1, 0])], Validity::from_iter([true, false]) ) - .unwrap() + .unwrap(), + &mut ctx ); } @@ -254,7 +258,7 @@ mod tests { #[test] fn test_empty_struct_is_constant() { let array = StructArray::new_fieldless_with_len(2); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = is_constant(&array.into_array(), &mut ctx) .vortex_expect("operation should succeed in test"); assert!(result); @@ -415,6 +419,9 @@ mod tests { StructArray::try_new(["xs", "ys"].into(), vec![xs, ys], 100, Validity::NonNullable).unwrap() })] fn test_struct_consistency(#[case] array: StructArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index ae3a67b50e3..01981803239 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -147,14 +147,13 @@ mod tests { use crate::dtype::StructFields; use crate::executor::VortexSessionExecute; use crate::optimizer::ArrayOptimizer; - use crate::optimizer::kernels::ArrayKernels; + use crate::optimizer::kernels::KernelSession; use crate::optimizer::kernels::ReduceParentFn; use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; use crate::validity::Validity; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); fn no_struct_cast_plugin( _child: &ArrayRef, @@ -166,6 +165,7 @@ mod tests { #[test] fn test_struct_cast_field_reorder() { + let mut ctx = SESSION.create_execution_ctx(); // Source: {a, b}, Target: {c, b, a} - reordered + new null field let source = StructArray::try_new( FieldNames::from(["a", "b"]), @@ -197,15 +197,18 @@ mod tests { .unwrap(); assert_arrays_eq!( result.unmasked_field_by_name("a").unwrap(), - VarBinViewArray::from_iter_nullable_str([Some("A")]) + VarBinViewArray::from_iter_nullable_str([Some("A")]), + &mut ctx ); assert_arrays_eq!( result.unmasked_field_by_name("b").unwrap(), - VarBinViewArray::from_iter_nullable_str([Some("B")]) + VarBinViewArray::from_iter_nullable_str([Some("B")]), + &mut ctx ); assert_arrays_eq!( result.unmasked_field_by_name("c").unwrap(), - ConstantArray::new(Scalar::null(utf8_null), 1) + ConstantArray::new(Scalar::null(utf8_null), 1), + &mut ctx ); } @@ -255,8 +258,8 @@ mod tests { ); let cast = source.cast(target).unwrap(); - let kernels = ArrayKernels::empty(); - kernels.register_reduce_parent( + let kernels = KernelSession::empty(); + kernels.kernels().register_reduce_parent( Cast.id(), Struct.id(), &[no_struct_cast_plugin as ReduceParentFn], @@ -299,6 +302,7 @@ mod tests { #[test] fn cast_struct_drop_field() { + let mut ctx = SESSION.create_execution_ctx(); // Casting to a struct with a subset of fields should succeed. let source = StructArray::try_new( FieldNames::from(["a", "b", "c"]), @@ -332,16 +336,19 @@ mod tests { assert_eq!(result.unmasked_fields().len(), 2); assert_arrays_eq!( result.unmasked_field_by_name("a").unwrap(), - buffer![1i32, 2, 3].into_array() + buffer![1i32, 2, 3].into_array(), + &mut ctx ); assert_arrays_eq!( result.unmasked_field_by_name("c").unwrap(), - buffer![100u8, 200, 255].into_array() + buffer![100u8, 200, 255].into_array(), + &mut ctx ); } #[test] fn cast_struct_field_type_widening() { + let mut ctx = SESSION.create_execution_ctx(); // Casting struct fields to wider types (i32 -> i64). let source = StructArray::try_new( FieldNames::from(["val"]), @@ -371,7 +378,8 @@ mod tests { ); assert_arrays_eq!( result.unmasked_field_by_name("val").unwrap(), - buffer![1i64, 2, 3].into_array() + buffer![1i64, 2, 3].into_array(), + &mut ctx ); } diff --git a/vortex-array/src/arrays/struct_/mod.rs b/vortex-array/src/arrays/struct_/mod.rs index 5fccb58f9cd..da49411c3bc 100644 --- a/vortex-array/src/arrays/struct_/mod.rs +++ b/vortex-array/src/arrays/struct_/mod.rs @@ -10,5 +10,10 @@ pub(crate) mod compute; mod vtable; pub use vtable::Struct; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + compute::cast::initialize(session); + vtable::initialize(session); +} + #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/struct_/tests.rs b/vortex-array/src/arrays/struct_/tests.rs index 2db9e62f4f5..efa52a53a9f 100644 --- a/vortex-array/src/arrays/struct_/tests.rs +++ b/vortex-array/src/arrays/struct_/tests.rs @@ -6,8 +6,8 @@ use vortex_error::VortexResult; use crate::Canonical; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -24,6 +24,7 @@ use crate::validity::Validity; #[test] fn test_project() { + let mut ctx = array_session().create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = VarBinArray::from_vec( vec!["a", "b", "c", "d", "e"], @@ -52,15 +53,21 @@ fn test_project() { let bools = struct_b.unmasked_field(0); assert_arrays_eq!( bools, - BoolArray::from_iter([true, true, true, false, false]) + BoolArray::from_iter([true, true, true, false, false]), + &mut ctx ); let prims = struct_b.unmasked_field(1); - assert_arrays_eq!(prims, PrimitiveArray::from_iter([0i64, 1, 2, 3, 4])); + assert_arrays_eq!( + prims, + PrimitiveArray::from_iter([0i64, 1, 2, 3, 4]), + &mut ctx + ); } #[test] fn test_remove_column() { + let mut ctx = array_session().create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = PrimitiveArray::new(buffer![4u64, 5, 6, 7, 8], Validity::NonNullable); @@ -77,7 +84,11 @@ fn test_remove_column() { removed.dtype(), &DType::Primitive(PType::I64, Nullability::NonNullable) ); - assert_arrays_eq!(removed, PrimitiveArray::from_iter([0i64, 1, 2, 3, 4])); + assert_arrays_eq!( + removed, + PrimitiveArray::from_iter([0i64, 1, 2, 3, 4]), + &mut ctx + ); assert_eq!(data.names(), &["ys"]); assert_eq!(data.struct_fields().nfields(), 1); @@ -88,7 +99,8 @@ fn test_remove_column() { ); assert_arrays_eq!( data.unmasked_field(0), - PrimitiveArray::from_iter([4u64, 5, 6, 7, 8]) + PrimitiveArray::from_iter([4u64, 5, 6, 7, 8]), + &mut ctx ); let empty = data.remove_column("non_existent"); @@ -101,6 +113,7 @@ fn test_remove_column() { #[test] fn test_duplicate_field_names() { + let mut ctx = array_session().create_execution_ctx(); // Test that StructArray allows duplicate field names and returns the first match let field1 = buffer![1i32, 2, 3].into_array(); let field2 = buffer![10i32, 20, 30].into_array(); @@ -119,19 +132,21 @@ fn test_duplicate_field_names() { let first_value_field = struct_array.unmasked_field_by_name("value").unwrap(); assert_arrays_eq!( first_value_field, - PrimitiveArray::from_iter([1i32, 2, 3]) // This is field1, not field3 + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx ); // Verify field_by_name_opt also returns the first match let opt_field = struct_array.unmasked_field_by_name_opt("value").unwrap(); - assert_arrays_eq!( - opt_field, - PrimitiveArray::from_iter([1i32, 2, 3]) // First "value" field - ); + assert_arrays_eq!(opt_field, PrimitiveArray::from_iter([1i32, 2, 3]), &mut ctx); // Verify the third field (second "value") can be accessed by index let third_field = struct_array.unmasked_field(2); - assert_arrays_eq!(third_field, PrimitiveArray::from_iter([100i32, 200, 300])); + assert_arrays_eq!( + third_field, + PrimitiveArray::from_iter([100i32, 200, 300]), + &mut ctx + ); } #[test] @@ -146,14 +161,189 @@ fn test_uncompressed_size_in_bytes() -> VortexResult<()> { let canonical_size = struct_array .clone() .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut array_session().create_execution_ctx())? .into_array() .nbytes(); let uncompressed_size = struct_array .statistics() - .compute_uncompressed_size_in_bytes(&mut LEGACY_SESSION.create_execution_ctx()); + .compute_uncompressed_size_in_bytes(&mut array_session().create_execution_ctx()); assert_eq!(canonical_size, 2); assert_eq!(uncompressed_size, Some(4000)); Ok(()) } + +#[test] +fn test_push_validity_into_children_drops_struct_validity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let struct_array = StructArray::try_new( + FieldNames::from(["a", "b"]), + vec![ + buffer![1i32, 2, 3].into_array(), + buffer![10i32, 20, 30].into_array(), + ], + 3, + Validity::from_iter([true, false, true]), + )?; + + let pushed = struct_array.push_validity_into_children(true)?; + + // The struct is now non-nullable; the row-1 null lives in every field instead. + let expected = StructArray::try_new( + FieldNames::from(["a", "b"]), + vec![ + PrimitiveArray::new( + buffer![1i32, 2, 3], + Validity::from_iter([true, false, true]), + ) + .into_array(), + PrimitiveArray::new( + buffer![10i32, 20, 30], + Validity::from_iter([true, false, true]), + ) + .into_array(), + ], + 3, + Validity::NonNullable, + )?; + + assert!(!pushed.dtype().is_nullable()); + assert_arrays_eq!(pushed, expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_push_validity_into_children_preserves_struct_validity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let struct_array = StructArray::try_new( + FieldNames::from(["a", "b"]), + vec![ + buffer![1i32, 2, 3].into_array(), + buffer![10i32, 20, 30].into_array(), + ], + 3, + Validity::from_iter([true, false, true]), + )?; + + let pushed = struct_array.push_validity_into_children(false)?; + + // The null now exists both at the struct level and in every field. + let expected = StructArray::try_new( + FieldNames::from(["a", "b"]), + vec![ + PrimitiveArray::new( + buffer![1i32, 2, 3], + Validity::from_iter([true, false, true]), + ) + .into_array(), + PrimitiveArray::new( + buffer![10i32, 20, 30], + Validity::from_iter([true, false, true]), + ) + .into_array(), + ], + 3, + Validity::from_iter([true, false, true]), + )?; + + assert!(pushed.dtype().is_nullable()); + assert_arrays_eq!(pushed, expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_push_validity_into_children_intersects_field_validity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + // Fields carry their own nulls (a at row 1, b at row 2) and the struct is null at row 1, + // so pushing intersects both levels rather than overwriting the fields. + let struct_array = StructArray::try_new( + FieldNames::from(["a", "b"]), + vec![ + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), + PrimitiveArray::from_option_iter([Some(10i64), Some(20), None]).into_array(), + ], + 3, + Validity::from_iter([true, false, true]), + )?; + + let pushed = struct_array.push_validity_into_children(true)?; + + // a: null at row 1; b: null at rows 1 and 2. + let expected = StructArray::try_new( + FieldNames::from(["a", "b"]), + vec![ + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), + PrimitiveArray::from_option_iter([Some(10i64), None, None]).into_array(), + ], + 3, + Validity::NonNullable, + )?; + + assert_arrays_eq!(pushed, expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_push_validity_into_children_all_invalid() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let struct_array = StructArray::try_new( + FieldNames::from(["a", "b"]), + vec![ + buffer![1i32, 2, 3].into_array(), + buffer![10i32, 20, 30].into_array(), + ], + 3, + Validity::AllInvalid, + )?; + + let pushed = struct_array.push_validity_into_children(true)?; + + // Every row is null at the struct level, so every field becomes all-null. + let expected = StructArray::try_new( + FieldNames::from(["a", "b"]), + vec![ + PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllInvalid).into_array(), + PrimitiveArray::new(buffer![10i32, 20, 30], Validity::AllInvalid).into_array(), + ], + 3, + Validity::NonNullable, + )?; + + assert_arrays_eq!(pushed, expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_push_validity_into_children_no_nulls() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + // No nulls: the fields are untouched, only the top-level nullability changes. + let struct_array = StructArray::try_new( + FieldNames::from(["a", "b"]), + vec![ + buffer![1i32, 2, 3].into_array(), + buffer![10i32, 20, 30].into_array(), + ], + 3, + Validity::AllValid, + )?; + + let dropped = struct_array.push_validity_into_children(true)?; + let expected = StructArray::try_new( + FieldNames::from(["a", "b"]), + vec![ + buffer![1i32, 2, 3].into_array(), + buffer![10i32, 20, 30].into_array(), + ], + 3, + Validity::NonNullable, + )?; + assert!(!dropped.dtype().is_nullable()); + assert_arrays_eq!(dropped, expected, &mut ctx); + + let preserved = struct_array.push_validity_into_children(false)?; + assert!(preserved.dtype().is_nullable()); + assert_arrays_eq!(preserved, struct_array, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/arrays/struct_/vtable/kernel.rs b/vortex-array/src/arrays/struct_/vtable/kernel.rs index eac7158921d..1e58da2d0d0 100644 --- a/vortex-array/src/arrays/struct_/vtable/kernel.rs +++ b/vortex-array/src/arrays/struct_/vtable/kernel.rs @@ -1,9 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_session::VortexSession; + use crate::arrays::Struct; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&ZipExecuteAdaptor(Struct))]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Zip.id(), Struct, ZipExecuteAdaptor(Struct)); +} diff --git a/vortex-array/src/arrays/struct_/vtable/mod.rs b/vortex-array/src/arrays/struct_/vtable/mod.rs index 71ec8a8e2df..1b23ecf804c 100644 --- a/vortex-array/src/arrays/struct_/vtable/mod.rs +++ b/vortex-array/src/arrays/struct_/vtable/mod.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; -use kernel::PARENT_KERNELS; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -18,6 +17,7 @@ use crate::array::ArrayView; 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; @@ -37,6 +37,10 @@ use crate::array::ArrayId; /// A [`Struct`]-encoded Vortex array. pub type StructArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + impl VTable for Struct { type TypedArrayData = EmptyArrayData; @@ -119,6 +123,14 @@ impl VTable for Struct { vortex_panic!("StructArray 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, @@ -192,15 +204,6 @@ impl VTable for Struct { ) -> VortexResult> { PARENT_RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } #[derive(Clone, Debug)] diff --git a/vortex-array/src/arrays/varbin/accessor.rs b/vortex-array/src/arrays/varbin/accessor.rs deleted file mode 100644 index ee68f94bd55..00000000000 --- a/vortex-array/src/arrays/varbin/accessor.rs +++ /dev/null @@ -1,63 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::iter; - -use vortex_error::VortexExpect; - -#[expect(deprecated)] -use crate::ToCanonical as _; -use crate::accessor::ArrayAccessor; -use crate::arrays::VarBinArray; -use crate::arrays::varbin::VarBinArrayExt; -use crate::match_each_integer_ptype; -use crate::validity::Validity; - -impl ArrayAccessor<[u8]> for VarBinArray { - fn with_iterator(&self, f: F) -> R - where - F: for<'a> FnOnce(&mut dyn Iterator>) -> R, - { - #[expect(deprecated)] - let offsets = self.offsets().to_primitive(); - let validity = self - .validity() - .vortex_expect("varbin validity should be derivable"); - - let bytes = self.bytes(); - let bytes = bytes.as_slice(); - - match_each_integer_ptype!(offsets.ptype(), |T| { - let offsets = offsets.as_slice::(); - - #[allow(clippy::cast_possible_truncation)] - match validity { - Validity::NonNullable | Validity::AllValid => { - let mut iter = offsets - .windows(2) - .map(|w| Some(&bytes[w[0] as usize..w[1] as usize])); - f(&mut iter) - } - Validity::AllInvalid => f(&mut iter::repeat_n(None, self.len())), - Validity::Array(v) => { - #[expect(deprecated)] - let validity = v.to_bool().into_bit_buffer(); - let mut iter = offsets - .windows(2) - .zip(validity.iter()) - .map(|(w, valid)| valid.then(|| &bytes[w[0] as usize..w[1] as usize])); - f(&mut iter) - } - } - }) - } -} - -impl ArrayAccessor<[u8]> for &VarBinArray { - fn with_iterator(&self, f: F) -> R - where - F: for<'a> FnOnce(&mut dyn Iterator>) -> R, - { - >::with_iterator(*self, f) - } -} diff --git a/vortex-array/src/arrays/varbin/array.rs b/vortex-array/src/arrays/varbin/array.rs index 4d435471ce0..220fe948a73 100644 --- a/vortex-array/src/arrays/varbin/array.rs +++ b/vortex-array/src/arrays/varbin/array.rs @@ -255,7 +255,7 @@ impl VarBinData { } _ => None, }; - let all_invalid = matches!(validity, Validity::AllInvalid); + let all_invalid = validity.definitely_all_null(); match_each_integer_ptype!(primitive_offsets.dtype().as_ptype(), |O| { let offsets_slice = primitive_offsets.as_slice::(); diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index ad76eaa981b..404f574af52 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -130,8 +130,8 @@ impl VarBinBuilder { mod tests { use vortex_error::VortexResult; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::varbin::VarBinArrayExt; use crate::arrays::varbin::builder::VarBinBuilder; use crate::dtype::DType; @@ -153,13 +153,13 @@ mod tests { assert_eq!(array.dtype().nullability(), Nullable); assert_eq!( array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::utf8("hello".to_string(), Nullable) ); assert!( array - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); diff --git a/vortex-array/src/arrays/varbin/compute/cast.rs b/vortex-array/src/arrays/varbin/compute/cast.rs index ca5dfde56e5..5983886ea30 100644 --- a/vortex-array/src/arrays/varbin/compute/cast.rs +++ b/vortex-array/src/arrays/varbin/compute/cast.rs @@ -81,10 +81,8 @@ mod tests { use crate::compute::conformance::cast::test_cast_conformance; use crate::dtype::DType; use crate::dtype::Nullability; - use crate::session::ArraySession; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[rstest] #[case( diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index 9a5a615c81e..ede3c409d94 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -58,7 +58,7 @@ impl CompareKernel for VarBin { if rhs_is_empty { let buffer = match operator { - CompareOperator::Gte => BitBuffer::new_set(len), // Every possible value is >= "" + CompareOperator::Gte => BitBuffer::new_set(len), /* Every possible value is >= "" */ CompareOperator::Lt => BitBuffer::new_unset(len), // No value is < "" CompareOperator::Eq | CompareOperator::Lte => { let lhs_offsets = lhs.offsets().clone().execute::(ctx)?; @@ -159,10 +159,10 @@ mod test { use vortex_buffer::ByteBuffer; use crate::IntoArray; - use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; @@ -200,7 +200,7 @@ mod test { .unwrap() .execute_mask( result.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx() + &mut array_session().create_execution_ctx() ) .unwrap() .to_bit_buffer(), @@ -236,7 +236,7 @@ mod test { .unwrap() .execute_mask( result.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx() + &mut array_session().create_execution_ctx() ) .unwrap() .to_bit_buffer(), @@ -254,6 +254,8 @@ mod tests { use vortex_buffer::ByteBuffer; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::VarBinArray; @@ -289,6 +291,7 @@ mod tests { /// [`CompareKernel`]: super::CompareKernel #[test] fn varbin_i64_offsets_compare_constant() { + let mut ctx = array_session().create_execution_ctx(); let mut builder = VarBinBuilder::::with_capacity(3); builder.append_value(b"abc"); builder.append_value(b"xyz"); @@ -304,11 +307,12 @@ mod tests { .unwrap(); let expected = BoolArray::from_iter([true, false, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn varbin_i64_offsets_compare_constant_binary() { + let mut ctx = array_session().create_execution_ctx(); let mut builder = VarBinBuilder::::with_capacity(3); builder.append_value(b"abc"); builder.append_value(b"xyz"); @@ -328,6 +332,6 @@ mod tests { .unwrap(); let expected = BoolArray::from_iter([true, false, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } } diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index a0c78a64363..ab969cf3def 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -225,8 +225,8 @@ mod test { use vortex_buffer::buffer; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::VarBinArray; use crate::arrays::varbin::compute::filter::filter_select_var_bin_by_index; @@ -240,6 +240,7 @@ mod test { #[test] fn filter_var_bin_test() { + let mut ctx = array_session().create_execution_ctx(); let arr = VarBinArray::from_vec( vec![ b"hello".as_slice(), @@ -249,19 +250,14 @@ mod test { DType::Utf8(NonNullable), ); let arr = arr.as_view(); - let buf = filter_select_var_bin_by_index( - arr, - &[0, 2], - 2, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let buf = filter_select_var_bin_by_index(arr, &[0, 2], 2, &mut ctx).unwrap(); - assert_arrays_eq!(buf, VarBinArray::from(vec!["hello", "filter"])); + assert_arrays_eq!(buf, VarBinArray::from(vec!["hello", "filter"]), &mut ctx); } #[test] fn filter_var_bin_slice_test() { + let mut ctx = array_session().create_execution_ctx(); let arr = VarBinArray::from_vec( vec![ b"hello".as_slice(), @@ -274,19 +270,19 @@ mod test { ); let arr = arr.as_view(); - let buf = filter_select_var_bin_by_slice( - arr, - &[(0, 1), (2, 3), (4, 5)], - 3, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let buf = + filter_select_var_bin_by_slice(arr, &[(0, 1), (2, 3), (4, 5)], 3, &mut ctx).unwrap(); - assert_arrays_eq!(buf, VarBinArray::from(vec!["hello", "filter", "filter3"])); + assert_arrays_eq!( + buf, + VarBinArray::from(vec!["hello", "filter", "filter3"]), + &mut ctx + ); } #[test] fn filter_var_bin_slice_null() { + let mut ctx = array_session().create_execution_ctx(); let bytes = [ b"one".as_slice(), b"two".as_slice(), @@ -306,13 +302,7 @@ mod test { let arr = VarBinArray::try_new(offsets, bytes, DType::Utf8(Nullable), validity).unwrap(); let arr = arr.as_view(); - let buf = filter_select_var_bin_by_slice( - arr, - &[(0, 3), (4, 6)], - 5, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let buf = filter_select_var_bin_by_slice(arr, &[(0, 3), (4, 6)], 5, &mut ctx).unwrap(); assert_arrays_eq!( buf, @@ -322,12 +312,14 @@ mod test { Some("three"), Some("five"), Some("six") - ]) + ]), + &mut ctx ); } #[test] fn filter_varbin_nulls() { + let mut ctx = array_session().create_execution_ctx(); let bytes = [b"".as_slice(), b"two".as_slice(), b"two".as_slice()] .into_iter() .flat_map(|x| x.iter().cloned()) @@ -338,19 +330,14 @@ mod test { let arr = VarBinArray::try_new(offsets, bytes, DType::Utf8(Nullable), validity).unwrap(); let arr = arr.as_view(); - let buf = filter_select_var_bin_by_slice( - arr, - &[(0, 1), (2, 3)], - 2, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let buf = filter_select_var_bin_by_slice(arr, &[(0, 1), (2, 3)], 2, &mut ctx).unwrap(); - assert_arrays_eq!(buf, VarBinArray::from(vec![None, Some("two")])); + assert_arrays_eq!(buf, VarBinArray::from(vec![None, Some("two")]), &mut ctx); } #[test] fn filter_varbin_all_null() { + let mut ctx = array_session().create_execution_ctx(); let offsets = buffer![0, 0, 0, 0].into_array(); let validity = Validity::Array(BoolArray::from_iter([false, false, false]).into_array()); let arr = VarBinArray::try_new( @@ -362,15 +349,9 @@ mod test { .unwrap(); let arr = arr.as_view(); - let buf = filter_select_var_bin_by_slice( - arr, - &[(0, 1), (2, 3)], - 2, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let buf = filter_select_var_bin_by_slice(arr, &[(0, 1), (2, 3)], 2, &mut ctx).unwrap(); - assert_arrays_eq!(buf, VarBinArray::from(vec![None::<&str>, None])); + assert_arrays_eq!(buf, VarBinArray::from(vec![None::<&str>, None]), &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/varbin/compute/mod.rs b/vortex-array/src/arrays/varbin/compute/mod.rs index ab2de942138..480c07a3031 100644 --- a/vortex-array/src/arrays/varbin/compute/mod.rs +++ b/vortex-array/src/arrays/varbin/compute/mod.rs @@ -15,6 +15,8 @@ mod tests { use rstest::rstest; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::VarBinArray; use crate::compute::conformance::consistency::test_array_consistency; use crate::dtype::DType; @@ -60,6 +62,9 @@ mod tests { DType::Utf8(Nullability::NonNullable), ))] fn test_varbin_consistency(#[case] array: VarBinArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 5826633aea3..74e0bd6af9d 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -251,11 +251,15 @@ fn take_nullable(); + let mut ctx = array_session().create_execution_ctx(); + let scream = iter::once("a").cycle().take(128).collect::(); let bytes = ByteBuffer::copy_from(scream.as_bytes()); let offsets = buffer![0u8, 128u8].into_array(); @@ -322,6 +327,6 @@ mod tests { [Some(scream.clone()), Some(scream.clone()), Some(scream)], DType::Utf8(Nullability::NonNullable), ); - assert_arrays_eq!(expected, taken); + 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 a8f321fec9a..4bca67fad81 100644 --- a/vortex-array/src/arrays/varbin/mod.rs +++ b/vortex-array/src/arrays/varbin/mod.rs @@ -12,9 +12,11 @@ pub(crate) mod compute; mod vtable; pub use vtable::VarBin; -pub mod builder; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} -mod accessor; +pub mod builder; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; diff --git a/vortex-array/src/arrays/varbin/tests.rs b/vortex-array/src/arrays/varbin/tests.rs index 4711c69102e..e2eb79fff54 100644 --- a/vortex-array/src/arrays/varbin/tests.rs +++ b/vortex-array/src/arrays/varbin/tests.rs @@ -8,6 +8,8 @@ use vortex_buffer::buffer; use crate::ArrayRef; use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; @@ -32,17 +34,21 @@ fn binary_array() -> ArrayRef { #[rstest] pub fn test_scalar_at(binary_array: ArrayRef) { + let mut ctx = array_session().create_execution_ctx(); assert_arrays_eq!( binary_array, - VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) + VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]), + &mut ctx ); } #[rstest] pub fn slice_array(binary_array: ArrayRef) { + let mut ctx = array_session().create_execution_ctx(); let binary_arr = binary_array.slice(1..2).unwrap(); assert_arrays_eq!( binary_arr, - VarBinViewArray::from_iter_str(["hello world this is a long string"]) + VarBinViewArray::from_iter_str(["hello world this is a long string"]), + &mut ctx ); } diff --git a/vortex-array/src/arrays/varbin/vtable/canonical.rs b/vortex-array/src/arrays/varbin/vtable/canonical.rs index fc30ee2c94e..4865df5e1a3 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -47,8 +47,8 @@ pub(crate) fn varbin_to_canonical( mod tests { use rstest::rstest; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbin::builder::VarBinBuilder; @@ -79,7 +79,7 @@ mod tests { assert!( !canonical - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(0, &mut array_session().create_execution_ctx()) .unwrap() ); @@ -96,6 +96,7 @@ mod tests { #[case(DType::Utf8(Nullability::NonNullable))] #[case(DType::Binary(Nullability::NonNullable))] fn test_canonical_varbin_unsliced(#[case] dtype: DType) { + let mut ctx = array_session().create_execution_ctx(); let varbin = VarBinArray::from_iter_nonnull(["foo", "bar", "baz"], dtype.clone()); #[expect(deprecated)] let canonical = varbin.as_array().to_varbinview(); @@ -103,7 +104,7 @@ mod tests { DType::Utf8(_) => VarBinViewArray::from_iter_str(["foo", "bar", "baz"]), _ => VarBinViewArray::from_iter_bin(["foo", "bar", "baz"]), }; - assert_arrays_eq!(canonical, expected); + assert_arrays_eq!(canonical, expected, &mut ctx); } // Empty array: offsets has exactly one element; no elements to canonicalize. diff --git a/vortex-array/src/arrays/varbin/vtable/kernel.rs b/vortex-array/src/arrays/varbin/vtable/kernel.rs index 9258e677933..9e80abd1037 100644 --- a/vortex-array/src/arrays/varbin/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbin/vtable/kernel.rs @@ -1,16 +1,25 @@ // 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::VarBin; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(VarBin)), - ParentKernelSet::lift(&CompareExecuteAdaptor(VarBin)), - ParentKernelSet::lift(&FilterExecuteAdaptor(VarBin)), - ParentKernelSet::lift(&TakeExecuteAdaptor(VarBin)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), VarBin, CastExecuteAdaptor(VarBin)); + kernels.register_execute_parent_kernel(Binary.id(), VarBin, CompareExecuteAdaptor(VarBin)); + kernels.register_execute_parent_kernel(Filter.id(), VarBin, FilterExecuteAdaptor(VarBin)); + kernels.register_execute_parent_kernel(Dict.id(), VarBin, TakeExecuteAdaptor(VarBin)); +} diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index dae80b663de..301fdebc084 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_session::registry::CachedId; +use crate::ArrayParts; use crate::ArrayRef; use crate::ExecutionCtx; use crate::ExecutionResult; @@ -36,7 +37,6 @@ mod operations; mod validity; use canonical::varbin_to_canonical; -use kernel::PARENT_KERNELS; use vortex_session::VortexSession; use crate::EqMode; @@ -47,6 +47,10 @@ use crate::hash::ArrayHash; /// A [`VarBin`]-encoded Vortex array. pub type VarBinArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(Clone, prost::Message)] pub struct VarBinMetadata { #[prost(enumeration = "PType", tag = "1")] @@ -121,6 +125,24 @@ impl VTable for VarBin { } } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + let mut data = array.data().clone(); + data.bytes = buffers[0].clone(); + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, @@ -139,11 +161,10 @@ impl VTable for VarBin { dtype: &DType, len: usize, metadata: &[u8], - buffers: &[BufferHandle], children: &dyn ArrayChildren, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult> { let metadata = VarBinMetadata::decode(metadata)?; let validity = if children.len() == 1 { Validity::from(dtype.nullability()) @@ -167,7 +188,7 @@ impl VTable for VarBin { let data = VarBinData::try_build(offsets.clone(), bytes, dtype.clone(), validity.clone())?; let slots = VarBinData::make_slots(offsets, &validity, len); - Ok(crate::array::ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { @@ -182,15 +203,6 @@ impl VTable for VarBin { PARENT_RULES.evaluate(array, parent, child_idx) } - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done( varbin_to_canonical(array.as_view(), ctx)?.into_array(), diff --git a/vortex-array/src/arrays/varbinview/accessor.rs b/vortex-array/src/arrays/varbinview/accessor.rs deleted file mode 100644 index cd42c3f2b58..00000000000 --- a/vortex-array/src/arrays/varbinview/accessor.rs +++ /dev/null @@ -1,72 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::iter; - -use vortex_error::VortexExpect; - -#[expect(deprecated)] -use crate::ToCanonical as _; -use crate::accessor::ArrayAccessor; -use crate::arrays::VarBinViewArray; -use crate::validity::Validity; - -impl ArrayAccessor<[u8]> for VarBinViewArray { - fn with_iterator FnOnce(&mut dyn Iterator>) -> R, R>( - &self, - f: F, - ) -> R { - let bytes = (0..self.data_buffers().len()) - .map(|i| self.buffer(i)) - .collect::>(); - - let views = self.views(); - - match self - .validity() - .vortex_expect("varbinview validity should be derivable") - { - Validity::NonNullable | Validity::AllValid => { - let mut iter = views.iter().map(|view| { - if view.is_inlined() { - Some(view.as_inlined().value()) - } else { - Some( - &bytes[view.as_view().buffer_index as usize][view.as_view().as_range()], - ) - } - }); - f(&mut iter) - } - Validity::AllInvalid => f(&mut iter::repeat_n(None, views.len())), - Validity::Array(v) => { - #[expect(deprecated)] - let validity = v.to_bool().into_bit_buffer(); - let mut iter = views.iter().zip(validity.iter()).map(|(view, valid)| { - if valid { - if view.is_inlined() { - Some(view.as_inlined().value()) - } else { - Some( - &bytes[view.as_view().buffer_index as usize] - [view.as_view().as_range()], - ) - } - } else { - None - } - }); - f(&mut iter) - } - } - } -} - -impl ArrayAccessor<[u8]> for &VarBinViewArray { - fn with_iterator(&self, f: F) -> R - where - F: for<'a> FnOnce(&mut dyn Iterator>) -> R, - { - >::with_iterator(*self, f) - } -} diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index 6effc7c656a..fc3f3477ffc 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -18,6 +18,9 @@ use crate::arrays::varbinview::Ref; use crate::builders::ArrayBuilder; use crate::builders::VarBinViewBuilder; +const DEFAULT_COMPACTION_THRESHOLD: f64 = 0.5; +const MIN_RETAINED_BYTES_PER_ROW_TO_CHECK_COMPACTION: u64 = 128; + impl VarBinViewArray { /// Returns a compacted copy of the input array, where all wasted space has been cleaned up. This /// operation can be very expensive, in the worst case copying all existing string data into @@ -33,8 +36,7 @@ impl VarBinViewArray { return Ok(self.clone()); } - // Use selective compaction with threshold of 1.0 (compact any buffer with any waste) - self.compact_with_threshold(1.0) + self.compact_with_threshold(DEFAULT_COMPACTION_THRESHOLD) } fn should_compact(&self) -> VortexResult { @@ -50,12 +52,18 @@ impl VarBinViewArray { return Ok(true); } - let bytes_referenced: u64 = self.count_referenced_bytes()?; let buffer_total_bytes: u64 = self.buffers.iter().map(|buf| buf.len() as u64).sum(); + if buffer_total_bytes == 0 { + return Ok(true); + } - // If there is any wasted space, we want to repack. - // This is very aggressive. - Ok(bytes_referenced < buffer_total_bytes || buffer_total_bytes == 0) + let len = u64::try_from(self.len()).unwrap_or(u64::MAX); + if len > 0 && buffer_total_bytes / len <= MIN_RETAINED_BYTES_PER_ROW_TO_CHECK_COMPACTION { + return Ok(false); + } + + let bytes_referenced: u64 = self.count_referenced_bytes()?; + Ok((bytes_referenced as f64 / buffer_total_bytes as f64) < DEFAULT_COMPACTION_THRESHOLD) } /// Iterates over all valid, non-inlined views, calling the provided @@ -193,8 +201,8 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; @@ -202,6 +210,7 @@ mod tests { use crate::dtype::Nullability; #[test] fn test_optimize_compacts_buffers() { + let mut ctx = array_session().create_execution_ctx(); // Create a VarBinViewArray with some long strings that will create multiple buffers let original = VarBinViewArray::from_iter_nullable_str([ Some("short"), @@ -219,7 +228,7 @@ mod tests { let indices = buffer![0u32, 4u32].into_array(); let taken = original.take(indices).unwrap(); let taken = taken - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut array_session().create_execution_ctx()) .unwrap(); // The taken array should still have the same number of buffers assert_eq!(taken.data_buffers().len(), original_buffers); @@ -235,12 +244,14 @@ mod tests { // Verify the data is still correct assert_arrays_eq!( optimized_array, - >::from_iter([Some("short"), Some("tiny")]) + >::from_iter([Some("short"), Some("tiny")]), + &mut ctx ); } #[test] fn test_optimize_with_long_strings() { + let mut ctx = array_session().create_execution_ctx(); // Create strings that are definitely longer than 12 bytes let long_string_1 = "this is definitely a very long string that exceeds the inline limit"; let long_string_2 = "another extremely long string that also needs external buffer storage"; @@ -258,11 +269,10 @@ mod tests { let indices = buffer![0u32, 2u32].into_array(); let taken = original.take(indices).unwrap(); let taken_array = taken - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut array_session().create_execution_ctx()) .unwrap(); - // Optimize the taken array - let optimized_array = taken_array.compact_buffers().unwrap(); + let optimized_array = taken_array.compact_with_threshold(1.0).unwrap(); // The optimized array should have exactly 1 buffer (consolidated) assert_eq!(optimized_array.data_buffers().len(), 1); @@ -270,12 +280,14 @@ mod tests { // Verify the data is still correct assert_arrays_eq!( optimized_array, - VarBinArray::from(vec![long_string_1, long_string_3]) + VarBinArray::from(vec![long_string_1, long_string_3]), + &mut ctx ); } #[test] fn test_optimize_no_buffers() { + let mut ctx = array_session().create_execution_ctx(); // Create an array with only short strings (all inlined) let original = VarBinViewArray::from_iter_str(["a", "bb", "ccc", "dddd"]); @@ -287,11 +299,12 @@ mod tests { assert_eq!(optimized_array.data_buffers().len(), 0); - assert_arrays_eq!(optimized_array, original); + assert_arrays_eq!(optimized_array, original, &mut ctx); } #[test] fn test_optimize_single_buffer() { + let mut ctx = array_session().create_execution_ctx(); // Create an array that naturally has only one buffer let str1 = "this is a long string that goes into a buffer"; let str2 = "another long string in the same buffer"; @@ -306,11 +319,12 @@ mod tests { assert_eq!(optimized_array.data_buffers().len(), 1); - assert_arrays_eq!(optimized_array, original); + assert_arrays_eq!(optimized_array, original, &mut ctx); } #[test] fn test_selective_compaction_with_threshold_zero() { + let mut ctx = array_session().create_execution_ctx(); // threshold=0 should keep all buffers (no compaction) let original = VarBinViewArray::from_iter_str([ "this is a longer string that will be stored in a buffer", @@ -324,7 +338,7 @@ mod tests { let indices = buffer![0u32].into_array(); let taken = original.take(indices).unwrap(); let taken = taken - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut array_session().create_execution_ctx()) .unwrap(); // Compact with threshold=0 (should not compact) let compacted = taken.compact_with_threshold(0.0).unwrap(); @@ -333,11 +347,12 @@ mod tests { assert_eq!(compacted.data_buffers().len(), taken.data_buffers().len()); // Verify correctness - assert_arrays_eq!(compacted, taken); + assert_arrays_eq!(compacted, taken, &mut ctx); } #[test] fn test_selective_compaction_with_high_threshold() { + let mut ctx = array_session().create_execution_ctx(); // threshold=1.0 should compact any buffer with waste let original = VarBinViewArray::from_iter_str([ "this is a longer string that will be stored in a buffer", @@ -349,7 +364,7 @@ mod tests { let indices = buffer![0u32, 2u32].into_array(); let taken = original.take(indices).unwrap(); let taken = taken - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut array_session().create_execution_ctx()) .unwrap(); let original_buffers = taken.data_buffers().len(); @@ -361,11 +376,12 @@ mod tests { assert!(compacted.data_buffers().len() <= original_buffers); // Verify correctness - assert_arrays_eq!(compacted, taken); + assert_arrays_eq!(compacted, taken, &mut ctx); } #[test] fn test_selective_compaction_preserves_well_utilized_buffers() { + let mut ctx = array_session().create_execution_ctx(); // Create an array with multiple strings in one buffer (well-utilized) let str1 = "first long string that needs external buffer storage"; let str2 = "second long string also in buffer"; @@ -383,11 +399,12 @@ mod tests { assert_eq!(compacted.data_buffers().len(), 1); // Verify all data is correct - assert_arrays_eq!(compacted, original); + assert_arrays_eq!(compacted, original, &mut ctx); } #[test] fn test_selective_compaction_with_mixed_utilization() { + let mut ctx = array_session().create_execution_ctx(); // Create array with some long strings let strings: Vec = (0..10) .map(|i| { @@ -404,7 +421,7 @@ mod tests { let indices_array = buffer![0u32, 2u32, 4u32, 6u32, 8u32].into_array(); let taken = original.take(indices_array).unwrap(); let taken = taken - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut array_session().create_execution_ctx()) .unwrap(); // Compact with moderate threshold @@ -414,11 +431,12 @@ mod tests { [0, 2, 4, 6, 8].map(|i| Some(strings[i].as_str())), DType::Utf8(Nullability::NonNullable), ); - assert_arrays_eq!(expected, compacted); + assert_arrays_eq!(expected, compacted, &mut ctx); } #[test] fn test_slice_strategy_with_contiguous_range() { + let mut ctx = array_session().create_execution_ctx(); // Create array with strings that will be in one buffer let strings: Vec = (0..20) .map(|i| format!("this is a long string number {} for slice test", i)) @@ -430,7 +448,7 @@ mod tests { let indices_array = buffer![0u32, 1u32, 2u32, 3u32, 4u32].into_array(); let taken = original.take(indices_array).unwrap(); let taken = taken - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut array_session().create_execution_ctx()) .unwrap(); // Get buffer stats before compaction let utils_before = taken.buffer_utilizations().unwrap(); @@ -447,7 +465,7 @@ mod tests { ); // Verify correctness - assert_arrays_eq!(&compacted, taken); + assert_arrays_eq!(&compacted, taken, &mut ctx); // Verify that if there was only one buffer, the compacted version also has one // (it was sliced, not rewritten into multiple buffers) diff --git a/vortex-array/src/arrays/varbinview/compute/cast.rs b/vortex-array/src/arrays/varbinview/compute/cast.rs index 449e735d1fc..476c74e7309 100644 --- a/vortex-array/src/arrays/varbinview/compute/cast.rs +++ b/vortex-array/src/arrays/varbinview/compute/cast.rs @@ -85,10 +85,8 @@ mod tests { use crate::compute::conformance::cast::test_cast_conformance; use crate::dtype::DType; use crate::dtype::Nullability; - use crate::session::ArraySession; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[rstest] #[case( diff --git a/vortex-array/src/arrays/varbinview/compute/mod.rs b/vortex-array/src/arrays/varbinview/compute/mod.rs index b69efad1e88..0358ad01460 100644 --- a/vortex-array/src/arrays/varbinview/compute/mod.rs +++ b/vortex-array/src/arrays/varbinview/compute/mod.rs @@ -11,14 +11,14 @@ mod zip; #[cfg(test)] mod tests { use vortex_buffer::buffer; + use vortex_error::VortexResult; use crate::IntoArray; - use crate::accessor::ArrayAccessor; use crate::arrays::VarBinViewArray; #[expect(deprecated)] use crate::canonical::ToCanonical as _; #[test] - fn take_nullable() { + fn take_nullable() -> VortexResult<()> { let arr = VarBinViewArray::from_iter_nullable_str([ Some("one"), None, @@ -28,19 +28,27 @@ mod tests { Some("six"), ]); - let taken = arr.take(buffer![0, 3].into_array()).unwrap(); + let taken = arr.take(buffer![0, 3].into_array())?; assert!(taken.dtype().is_nullable()); + let mut ctx = array_session().create_execution_ctx(); #[expect(deprecated)] - let result = taken.to_varbinview().with_iterator(|it| { - it.map(|v| v.map(|b| unsafe { String::from_utf8_unchecked(b.to_vec()) })) - .collect::>() - }); + let taken = taken.to_varbinview(); + let mask = taken.validity()?.execute_mask(taken.len(), &mut ctx)?; + let result = (0..taken.len()) + .map(|i| { + mask.value(i) + .then(|| unsafe { String::from_utf8_unchecked(taken.bytes_at(i).to_vec()) }) + }) + .collect::>(); assert_eq!(result, [Some("one".to_string()), Some("four".to_string())]); + Ok(()) } // Consistency tests use rstest::rstest; + use crate::VortexSessionExecute; + use crate::array_session; use crate::compute::conformance::consistency::test_array_consistency; use crate::dtype::DType; use crate::dtype::Nullability; @@ -79,6 +87,9 @@ mod tests { None::<&str>, None, None, None ]))] fn test_varbinview_consistency(#[case] array: VarBinViewArray) { - test_array_consistency(&array.into_array()); + test_array_consistency( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 5c5ff08e2d0..1b7435729f1 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -90,20 +90,20 @@ mod tests { use rstest::rstest; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; + use vortex_error::VortexResult; use crate::IntoArray; - use crate::accessor::ArrayAccessor; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::compute::take::PrimitiveArray; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; use crate::compute::conformance::take::test_take_conformance; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::validity::Validity; #[test] - fn take_nullable() { + fn take_nullable() -> VortexResult<()> { let arr = VarBinViewArray::from_iter_nullable_str([ Some("one"), None, @@ -113,19 +113,24 @@ mod tests { Some("six"), ]); - let taken = arr.take(buffer![0, 3].into_array()).unwrap(); + let taken = arr.take(buffer![0, 3].into_array())?; assert!(taken.dtype().is_nullable()); - #[expect(deprecated)] - let result = taken.to_varbinview().with_iterator(|it| { - it.map(|v| v.map(|b| unsafe { String::from_utf8_unchecked(b.to_vec()) })) - .collect::>() - }); + let mut ctx = array_session().create_execution_ctx(); + let taken = taken.execute::(&mut ctx)?; + let mask = taken.validity()?.execute_mask(taken.len(), &mut ctx)?; + let result = (0..taken.len()) + .map(|i| { + mask.value(i) + .then(|| unsafe { String::from_utf8_unchecked(taken.bytes_at(i).to_vec()) }) + }) + .collect::>(); assert_eq!(result, [Some("one".to_string()), Some("four".to_string())]); + Ok(()) } #[test] - fn take_nullable_indices() { + fn take_nullable_indices() -> VortexResult<()> { let arr = VarBinViewArray::from_iter(["one", "two"].map(Some), DType::Utf8(NonNullable)); let indices = PrimitiveArray::new( @@ -134,15 +139,20 @@ mod tests { Validity::from(BitBuffer::from(vec![true, false])), ); - let taken = arr.take(indices.into_array()).unwrap(); + let taken = arr.take(indices.into_array())?; assert!(taken.dtype().is_nullable()); - #[expect(deprecated)] - let result = taken.to_varbinview().with_iterator(|it| { - it.map(|v| v.map(|b| unsafe { String::from_utf8_unchecked(b.to_vec()) })) - .collect::>() - }); + let mut ctx = array_session().create_execution_ctx(); + let taken = taken.execute::(&mut ctx)?; + let mask = taken.validity()?.execute_mask(taken.len(), &mut ctx)?; + let result = (0..taken.len()) + .map(|i| { + mask.value(i) + .then(|| unsafe { String::from_utf8_unchecked(taken.bytes_at(i).to_vec()) }) + }) + .collect::>(); assert_eq!(result, [Some("two".to_string()), None]); + Ok(()) } #[rstest] diff --git a/vortex-array/src/arrays/varbinview/compute/zip.rs b/vortex-array/src/arrays/varbinview/compute/zip.rs index 9c198fe6f2c..75e30808821 100644 --- a/vortex-array/src/arrays/varbinview/compute/zip.rs +++ b/vortex-array/src/arrays/varbinview/compute/zip.rs @@ -209,10 +209,12 @@ fn push_view( #[cfg(test)] mod tests { + use vortex_error::VortexResult; use vortex_mask::Mask; use crate::IntoArray; - use crate::accessor::ArrayAccessor; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::VarBinViewArray; use crate::builtins::ArrayBuiltins; #[expect(deprecated)] @@ -221,7 +223,7 @@ mod tests { use crate::dtype::Nullability; #[test] - fn zip_varbinview_kernel_zips() { + fn zip_varbinview_kernel_zips() -> VortexResult<()> { let a = VarBinViewArray::from_iter( [ Some("aaaaaaaaaaaaa_long"), // outlined @@ -252,14 +254,18 @@ mod tests { let zipped = mask .clone() .into_array() - .zip(a.into_array(), b.into_array()) - .unwrap() + .zip(a.into_array(), b.into_array())? .to_varbinview(); - let values = zipped.with_iterator(|it| { - it.map(|v| v.map(|bytes| String::from_utf8(bytes.to_vec()).unwrap())) - .collect::>() - }); + let mut ctx = array_session().create_execution_ctx(); + let validity_mask = zipped.validity()?.execute_mask(zipped.len(), &mut ctx)?; + let values = (0..zipped.len()) + .map(|i| { + validity_mask + .value(i) + .then(|| String::from_utf8(zipped.bytes_at(i).to_vec()).unwrap()) + }) + .collect::>(); assert_eq!( values, @@ -274,5 +280,6 @@ mod tests { ); assert_eq!(zipped.len(), mask.len()); assert_eq!(zipped.dtype(), &DType::Utf8(Nullability::Nullable)); + Ok(()) } } diff --git a/vortex-array/src/arrays/varbinview/mod.rs b/vortex-array/src/arrays/varbinview/mod.rs index be82b952f30..a6b8c9db46c 100644 --- a/vortex-array/src/arrays/varbinview/mod.rs +++ b/vortex-array/src/arrays/varbinview/mod.rs @@ -7,7 +7,6 @@ pub use array::VarBinViewData; pub use array::VarBinViewDataParts; pub use vtable::VarBinViewArray; -mod accessor; pub(crate) mod compact; pub(crate) mod compute; @@ -15,6 +14,10 @@ pub(crate) mod compute; mod vtable; pub use vtable::VarBinView; +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + pub mod build_views; mod view; diff --git a/vortex-array/src/arrays/varbinview/tests.rs b/vortex-array/src/arrays/varbinview/tests.rs index 2892ff16283..981d4f602a1 100644 --- a/vortex-array/src/arrays/varbinview/tests.rs +++ b/vortex-array/src/arrays/varbinview/tests.rs @@ -1,42 +1,46 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -#[expect(deprecated)] -use crate::ToCanonical as _; +use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::BinaryView; use crate::assert_arrays_eq; #[test] pub fn varbin_view() { + let mut ctx = array_session().create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]); assert_arrays_eq!( binary_arr, - VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) + VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]), + &mut ctx ); } #[test] pub fn slice_array() { + let mut ctx = array_session().create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) .slice(1..2) .unwrap(); assert_arrays_eq!( binary_arr, - VarBinViewArray::from_iter_str(["hello world this is a long string"]) + VarBinViewArray::from_iter_str(["hello world this is a long string"]), + &mut ctx ); } #[test] pub fn flatten_array() { + let mut ctx = array_session().create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["string1", "string2"]); - #[expect(deprecated)] - let var_bin = binary_arr.as_array().to_varbinview(); assert_arrays_eq!( - var_bin, - VarBinViewArray::from_iter_str(["string1", "string2"]) + binary_arr, + VarBinViewArray::from_iter_str(["string1", "string2"]), + &mut ctx ); } diff --git a/vortex-array/src/arrays/varbinview/vtable/kernel.rs b/vortex-array/src/arrays/varbinview/vtable/kernel.rs index cd9d68010af..e09b381d590 100644 --- a/vortex-array/src/arrays/varbinview/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbinview/vtable/kernel.rs @@ -1,14 +1,22 @@ // 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::VarBinView; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::kernel::ParentKernelSet; +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::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ - ParentKernelSet::lift(&CastExecuteAdaptor(VarBinView)), - ParentKernelSet::lift(&TakeExecuteAdaptor(VarBinView)), - ParentKernelSet::lift(&ZipExecuteAdaptor(VarBinView)), -]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(Cast.id(), VarBinView, CastExecuteAdaptor(VarBinView)); + kernels.register_execute_parent_kernel(Dict.id(), VarBinView, TakeExecuteAdaptor(VarBinView)); + kernels.register_execute_parent_kernel(Zip.id(), VarBinView, ZipExecuteAdaptor(VarBinView)); +} diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index c093b9c0203..9335ed092ca 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -5,7 +5,6 @@ use std::hash::Hasher; use std::mem::size_of; use std::sync::Arc; -use kernel::PARENT_KERNELS; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -15,6 +14,7 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::ArrayParts; use crate::ArrayRef; use crate::EqMode; use crate::ExecutionCtx; @@ -29,6 +29,8 @@ use crate::arrays::varbinview::array::NUM_SLOTS; use crate::arrays::varbinview::array::SLOT_NAMES; use crate::arrays::varbinview::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::hash::ArrayEq; use crate::hash::ArrayHash; @@ -40,6 +42,10 @@ mod validity; /// A [`VarBinView`]-encoded Vortex array. pub type VarBinViewArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(Clone, Debug)] pub struct VarBinView; @@ -126,6 +132,26 @@ impl VTable for VarBinView { } } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + let Some((views, data_buffers)) = buffers.split_last() else { + vortex_bail!("Expected at least 1 buffer, got 0"); + }; + let data = VarBinViewData::try_new_handle( + views.clone(), + Arc::from(data_buffers.to_vec()), + array.dtype().clone(), + array.validity()?, + )?; + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + fn serialize( _array: ArrayView<'_, Self>, _session: &VortexSession, @@ -142,7 +168,7 @@ impl VTable for VarBinView { buffers: &[BufferHandle], children: &dyn ArrayChildren, _session: &VortexSession, - ) -> VortexResult> { + ) -> VortexResult> { if !metadata.is_empty() { vortex_bail!( "VarBinViewArray expects empty metadata, got {} bytes", @@ -183,10 +209,7 @@ impl VTable for VarBinView { validity.clone(), )?; let slots = VarBinViewData::make_slots(&validity, len); - return Ok( - crate::array::ArrayParts::new(self.clone(), dtype.clone(), len, data) - .with_slots(slots), - ); + return Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)); } let data_buffers = data_handles @@ -202,7 +225,7 @@ impl VTable for VarBinView { validity.clone(), )?; let slots = VarBinViewData::make_slots(&validity, len); - Ok(crate::array::ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { @@ -217,13 +240,17 @@ impl VTable for VarBinView { PARENT_RULES.evaluate(array, parent, child_idx) } - fn execute_parent( + fn append_to_builder( array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, + builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) + ) -> VortexResult<()> { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + return builder.append_varbinview_array(&array.into_owned(), ctx); + } + + builder.extend_from_array(array.as_ref()); + Ok(()) } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { @@ -239,7 +266,8 @@ mod tests { use super::*; use crate::ArrayContext; use crate::IntoArray; - use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; @@ -256,11 +284,13 @@ mod tests { let dtype = array.dtype().clone(); let len = array.len(); - let ctx = ArrayContext::empty(); + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); let serialized = array .clone() .into_array() - .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default()) + .serialize(&array_ctx, &session, &SerializeOptions::default()) .unwrap(); let mut concat = ByteBufferMut::empty(); @@ -269,14 +299,9 @@ mod tests { } let parts = SerializedArray::try_from(concat.freeze()).unwrap(); let decoded = parts - .decode( - &dtype, - len, - &ReadContext::new(ctx.to_ids()), - &LEGACY_SESSION, - ) + .decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session) .unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut ctx); } } diff --git a/vortex-array/src/arrays/variant/mod.rs b/vortex-array/src/arrays/variant/mod.rs index 7d7f23e34c5..1a5157e9aa1 100644 --- a/vortex-array/src/arrays/variant/mod.rs +++ b/vortex-array/src/arrays/variant/mod.rs @@ -11,6 +11,11 @@ use vortex_error::vortex_ensure; pub use self::vtable::Variant; pub use self::vtable::VariantArray; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} + use crate::ArrayRef; use crate::array::Array; use crate::array::ArrayParts; @@ -83,8 +88,8 @@ mod tests { use crate::ArrayRef; use crate::Canonical; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -136,7 +141,7 @@ mod tests { } fn execute_variant(array: ArrayRef) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let Canonical::Variant(variant) = array.execute::(&mut ctx)? else { return Err(vortex_err!("expected canonical variant array")); }; @@ -154,7 +159,7 @@ mod tests { let shredded = array .shredded() .ok_or_else(|| vortex_err!("expected shredded child"))?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let shredded = shredded.clone().execute::(&mut ctx)?; let expected_shredded_array = if let Some(values) = expected_shredded .iter() @@ -165,7 +170,7 @@ mod tests { } else { PrimitiveArray::from_option_iter(expected_shredded.iter().copied()) }; - assert_arrays_eq!(shredded, expected_shredded_array); + assert_arrays_eq!(shredded, expected_shredded_array, &mut ctx); Ok(()) } @@ -176,7 +181,7 @@ mod tests { ) -> VortexResult<()> { assert_eq!(array.len(), expected_core.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); for (idx, expected) in expected_core.iter().enumerate() { let scalar = array.core_storage().execute_scalar(idx, &mut ctx)?; let variant = scalar.as_variant(); @@ -268,7 +273,7 @@ mod tests { let shredded = PrimitiveArray::from_option_iter([Some(10i32), Some(20), None]).into_array(); let variant = VariantArray::try_new(core_storage, Some(shredded))?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); for (idx, expected) in [Some(10i32), None, Some(3)].into_iter().enumerate() { let scalar = variant.execute_scalar(idx, &mut ctx)?; let variant = scalar.as_variant(); @@ -374,6 +379,7 @@ mod tests { #[test] fn variant_get_keeps_valid_shredded_rows_for_matching_dtype() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let core_storage = row_storage([1, 2, 3])?; let shredded = StructArray::try_from_iter([( "a", @@ -389,17 +395,19 @@ mod tests { let result = variant .into_array() .apply(&expr)? - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut array_session().create_execution_ctx())?; assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]) + PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), + &mut ctx ); Ok(()) } #[test] fn variant_get_treats_value_and_typed_value_as_logical_field_names() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let core_storage = row_storage([1, 2, 3])?; let shredded = StructArray::try_from_iter([ ( @@ -422,10 +430,11 @@ mod tests { .clone() .into_array() .apply(&value_expr)? - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut array_session().create_execution_ctx())?; assert_arrays_eq!( value_result, - PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]) + PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), + &mut ctx ); let typed_value_expr = variant_get( @@ -436,10 +445,11 @@ mod tests { let typed_value_result = variant .into_array() .apply(&typed_value_expr)? - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut array_session().create_execution_ctx())?; assert_arrays_eq!( typed_value_result, - PrimitiveArray::from_option_iter([Some(40i32), Some(50), Some(60)]) + PrimitiveArray::from_option_iter([Some(40i32), Some(50), Some(60)]), + &mut ctx ); Ok(()) } diff --git a/vortex-array/src/arrays/variant/vtable/kernel.rs b/vortex-array/src/arrays/variant/vtable/kernel.rs index a24eb81fd5d..f289748eded 100644 --- a/vortex-array/src/arrays/variant/vtable/kernel.rs +++ b/vortex-array/src/arrays/variant/vtable/kernel.rs @@ -3,6 +3,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_session::VortexSession; use super::merge_typed_scalar_as_variant; use crate::ArrayRef; @@ -24,14 +25,17 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; use crate::kernel::ExecuteParentKernel; -use crate::kernel::ParentKernelSet; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::variant_get::VariantGet; use crate::scalar_fn::fns::variant_get::VariantGetOptions; use crate::scalar_fn::fns::variant_get::VariantPath; use crate::scalar_fn::fns::variant_get::VariantPathElement; -pub(super) const PARENT_KERNELS: ParentKernelSet = - ParentKernelSet::new(&[ParentKernelSet::lift(&VariantGetKernel)]); +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel(VariantGet.id(), Variant, VariantGetKernel); +} #[derive(Default, Debug)] struct VariantGetKernel; diff --git a/vortex-array/src/arrays/variant/vtable/mod.rs b/vortex-array/src/arrays/variant/vtable/mod.rs index 3c0dd76c575..4d84511ae6a 100644 --- a/vortex-array/src/arrays/variant/vtable/mod.rs +++ b/vortex-array/src/arrays/variant/vtable/mod.rs @@ -5,7 +5,6 @@ mod kernel; mod operations; mod validity; -use kernel::PARENT_KERNELS; use prost::Message; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -25,6 +24,7 @@ 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::variant::CORE_STORAGE_SLOT; use crate::arrays::variant::NUM_SLOTS; use crate::arrays::variant::SHREDDED_SLOT; @@ -43,6 +43,10 @@ use crate::serde::ArrayChildren; /// A [`Variant`]-encoded Vortex array. pub type VariantArray = Array; +pub(crate) fn initialize(session: &VortexSession) { + kernel::initialize(session); +} + #[derive(Clone, Debug)] pub struct Variant; @@ -122,6 +126,14 @@ impl VTable for Variant { None } + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + fn serialize( array: ArrayView<'_, Self>, _session: &VortexSession, @@ -140,7 +152,6 @@ impl VTable for Variant { dtype: &DType, len: usize, metadata: &[u8], - buffers: &[BufferHandle], children: &dyn ArrayChildren, session: &VortexSession, @@ -192,15 +203,6 @@ impl VTable for Variant { ) -> VortexResult> { RULES.evaluate(array, parent, child_idx) } - - fn execute_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - PARENT_KERNELS.execute(array, parent, child_idx, ctx) - } } fn merge_typed_scalar_as_variant( diff --git a/vortex-array/src/arrow/convert.rs b/vortex-array/src/arrow/convert.rs index 3d536402e5c..e62b3a15d22 100644 --- a/vortex-array/src/arrow/convert.rs +++ b/vortex-array/src/arrow/convert.rs @@ -722,8 +722,8 @@ mod tests { use crate::arrays::list::ListArrayExt; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::struct_::StructArrayExt; + use crate::arrow::ArrowSessionExt; use crate::arrow::FromArrowArray as _; - use crate::arrow::executor::ArrowArrayExecutor as _; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -1608,13 +1608,14 @@ mod tests { let struct_elements = list_array.elements().as_::(); assert_eq!(struct_elements.names().len(), 2); // key and value fields - // Convert back to Arrow as a MapArray + // Convert back to Arrow as a MapArray via an ArrowSession. let map_dtype = arrow_map.data_type().clone(); - let arrow_back = vortex_array - .execute_arrow( - Some(&map_dtype), - &mut crate::LEGACY_SESSION.create_execution_ctx(), - ) + let session = crate::array_session(); + let mut ctx = session.create_execution_ctx(); + let target = Field::new("", map_dtype, vortex_array.dtype().is_nullable()); + let arrow_back = session + .arrow() + .execute_arrow(vortex_array, Some(&target), &mut ctx) .unwrap(); let map_back = arrow_back .as_any() diff --git a/vortex-array/src/arrow/executor/byte.rs b/vortex-array/src/arrow/executor/byte.rs index bace6665bca..2db9d3e6494 100644 --- a/vortex-array/src/arrow/executor/byte.rs +++ b/vortex-array/src/arrow/executor/byte.rs @@ -80,10 +80,13 @@ mod tests { use arrow_array::cast::AsArray; use arrow_schema::DataType; use rstest::rstest; + use vortex_error::VortexResult; + use vortex_mask::Mask; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrow::ArrowArrayExecutor; use crate::arrow::executor::byte::VarBinViewArray; use crate::dtype::DType; @@ -119,7 +122,7 @@ mod tests { #[case] vortex_array: VarBinViewArray, #[case] target_dtype: DataType, ) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arrow = vortex_array .into_array() .execute_arrow(Some(&target_dtype), &mut ctx) @@ -166,7 +169,7 @@ mod tests { vortex_dtype, ); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arrow = vortex_array .into_array() .execute_arrow(Some(&target_dtype), &mut ctx) @@ -179,4 +182,27 @@ mod tests { assert!(arrow.is_null(1)); assert!(!arrow.is_null(2)); } + + #[test] + fn filtered_utf8_view_export_does_not_retain_unselected_buffers() -> VortexResult<()> { + let unselected = "x".repeat(1 << 20); + let array = + VarBinViewArray::from_iter_str(["selected", unselected.as_str(), unselected.as_str()]); + let filtered = array + .into_array() + .filter(Mask::from_iter([true, false, false]))?; + + let arrow = filtered.execute_arrow( + Some(&DataType::Utf8View), + &mut LEGACY_SESSION.create_execution_ctx(), + )?; + + assert_eq!(arrow.as_string_view().value(0), "selected"); + assert!( + arrow.get_array_memory_size() < unselected.len(), + "filtered export retained unselected payload: {} bytes", + arrow.get_array_memory_size() + ); + Ok(()) + } } diff --git a/vortex-array/src/arrow/executor/byte_view.rs b/vortex-array/src/arrow/executor/byte_view.rs index b88b1895d53..89f37681220 100644 --- a/vortex-array/src/arrow/executor/byte_view.rs +++ b/vortex-array/src/arrow/executor/byte_view.rs @@ -12,7 +12,6 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::ExecutionCtx; use crate::arrays::VarBinViewArray; -use crate::arrow::executor::validity::to_arrow_null_buffer; use crate::arrow::null_buffer::to_null_buffer; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; @@ -48,19 +47,8 @@ pub fn execute_varbinview_to_arrow( array: &VarBinViewArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - let views = - ScalarBuffer::::from(array.views_handle().as_host().clone().into_arrow_buffer()); - let buffers: Vec<_> = array - .data_buffers() - .iter() - .map(|buffer| buffer.as_host().clone().into_arrow_buffer()) - .collect(); - let nulls = to_arrow_null_buffer(array.validity()?, array.len(), ctx)?; - - // SAFETY: our own VarBinView array is considered safe. - Ok(Arc::new(unsafe { - GenericByteViewArray::::new_unchecked(views, buffers, nulls) - })) + let compacted = array.compact_buffers()?; + canonical_varbinview_to_arrow::(&compacted, ctx) } pub(super) fn to_arrow_byte_view( @@ -73,6 +61,7 @@ pub(super) fn to_arrow_byte_view( // 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.execute::(ctx)?; let varbinview = array.execute::(ctx)?; - canonical_varbinview_to_arrow::(&varbinview, ctx) + execute_varbinview_to_arrow::(&varbinview, ctx) } diff --git a/vortex-array/src/arrow/executor/decimal.rs b/vortex-array/src/arrow/executor/decimal.rs index 077495354cc..a9aa9fdd2f5 100644 --- a/vortex-array/src/arrow/executor/decimal.rs +++ b/vortex-array/src/arrow/executor/decimal.rs @@ -222,9 +222,9 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array::IntoArray; + use crate::array_session; use crate::arrow::ArrowArrayExecutor; use crate::arrow::executor::decimal::DecimalArray; use crate::builders::ArrayBuilder; @@ -235,7 +235,7 @@ mod tests { #[test] fn decimal_to_arrow() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Make a very simple i128 and i256 array. let decimal_vortex = DecimalArray::new( buffer![1i128, 2i128, 3i128, 4i128, 5i128], @@ -264,7 +264,7 @@ mod tests { fn test_to_arrow_decimal128( #[case] _decimal_type: T, ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); decimal.append_value(10); decimal.append_value(11); @@ -292,7 +292,7 @@ mod tests { fn test_to_arrow_decimal32(#[case] _decimal_type: T) -> VortexResult<()> { use arrow_array::Decimal32Array; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); decimal.append_value(10); decimal.append_value(11); @@ -320,7 +320,7 @@ mod tests { fn test_to_arrow_decimal64(#[case] _decimal_type: T) -> VortexResult<()> { use arrow_array::Decimal64Array; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); decimal.append_value(10); decimal.append_value(11); @@ -348,7 +348,7 @@ mod tests { fn test_to_arrow_decimal256( #[case] _decimal_type: T, ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); decimal.append_value(10); decimal.append_value(11); diff --git a/vortex-array/src/arrow/executor/dictionary.rs b/vortex-array/src/arrow/executor/dictionary.rs index 99e315bd8f7..97d25c3f7cc 100644 --- a/vortex-array/src/arrow/executor/dictionary.rs +++ b/vortex-array/src/arrow/executor/dictionary.rs @@ -149,7 +149,7 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; use crate::arrow::ArrowArrayExecutor; @@ -165,7 +165,7 @@ mod tests { } fn execute(array: crate::ArrayRef, dt: &DataType) -> VortexResult { - array.execute_arrow(Some(dt), &mut LEGACY_SESSION.create_execution_ctx()) + array.execute_arrow(Some(dt), &mut array_session().create_execution_ctx()) } fn dict_basic_input() -> crate::ArrayRef { diff --git a/vortex-array/src/arrow/executor/list.rs b/vortex-array/src/arrow/executor/list.rs index e8bf95e14cd..7bf7a16f496 100644 --- a/vortex-array/src/arrow/executor/list.rs +++ b/vortex-array/src/arrow/executor/list.rs @@ -222,12 +222,10 @@ mod tests { use crate::arrow::executor::list::ListViewArray; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; - use crate::session::ArraySession; use crate::validity::Validity; /// A shared session for these list-executor tests, used to create execution contexts. - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[test] fn test_to_arrow_list_i32() -> VortexResult<()> { diff --git a/vortex-array/src/arrow/executor/list_view.rs b/vortex-array/src/arrow/executor/list_view.rs index 5a2cd517b37..7d5666b3762 100644 --- a/vortex-array/src/arrow/executor/list_view.rs +++ b/vortex-array/src/arrow/executor/list_view.rs @@ -124,12 +124,10 @@ mod tests { use crate::arrow::ArrowArrayExecutor; use crate::arrow::executor::list_view::ListViewArray; use crate::arrow::executor::list_view::PrimitiveArray; - use crate::session::ArraySession; use crate::validity::Validity; /// A shared session for these list-view-executor tests, used to create execution contexts. - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[test] fn trims_zero_copy_with_significant_trailing_waste() -> VortexResult<()> { diff --git a/vortex-array/src/arrow/executor/run_end.rs b/vortex-array/src/arrow/executor/run_end.rs index b579cab7e4d..0632f871f29 100644 --- a/vortex-array/src/arrow/executor/run_end.rs +++ b/vortex-array/src/arrow/executor/run_end.rs @@ -217,10 +217,8 @@ mod tests { use crate::dtype::PType; use crate::executor::VortexSessionExecute; use crate::scalar::Scalar; - use crate::session::ArraySession; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); fn ree_type(ends: DataType, values_dtype: DataType) -> DataType { DataType::RunEndEncoded( diff --git a/vortex-array/src/arrow/executor/struct_.rs b/vortex-array/src/arrow/executor/struct_.rs index ebea596e219..d3b7bcda2dc 100644 --- a/vortex-array/src/arrow/executor/struct_.rs +++ b/vortex-array/src/arrow/executor/struct_.rs @@ -208,9 +208,9 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array; + use crate::array_session; use crate::arrays; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -221,7 +221,7 @@ mod tests { #[test] fn struct_nullable_non_null_to_arrow() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::AllValid); let struct_a = StructArray::try_new( @@ -242,7 +242,7 @@ mod tests { #[test] fn struct_nullable_with_nulls_to_arrow() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let xs = PrimitiveArray::from_option_iter(vec![Some(0_i64), Some(1), Some(2), None, Some(3)]); @@ -267,7 +267,7 @@ mod tests { #[test] fn struct_to_arrow_with_schema_mismatch() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::AllValid); let struct_a = StructArray::try_new( @@ -297,7 +297,7 @@ mod tests { #[test] fn test_to_arrow() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = StructArray::from_fields( vec![ ( @@ -340,7 +340,7 @@ mod tests { #[test] fn to_arrow_with_non_nullable_fields() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = StructArray::from_fields( vec![ ( diff --git a/vortex-array/src/arrow/session.rs b/vortex-array/src/arrow/session.rs index 705f7971f90..d3f43443d5d 100644 --- a/vortex-array/src/arrow/session.rs +++ b/vortex-array/src/arrow/session.rs @@ -39,8 +39,8 @@ use tracing::trace; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_session::Ref; use vortex_session::SessionExt; +use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_session::registry::Id; @@ -58,7 +58,7 @@ use crate::dtype::FieldName; use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::StructFields; -use crate::dtype::arrow::FromArrowType; +use crate::dtype::arrow::TryFromArrowType; use crate::dtype::arrow::to_data_type_naive; use crate::dtype::extension::ExtId; use crate::extension::datetime::AnyTemporal; @@ -163,7 +163,7 @@ pub type ArrowImportVTableRef = Arc; /// keyed by Arrow extension name. The default session pre-registers the builtin UUID /// plugin; temporal extensions are handled by the canonical Arrow ↔ Vortex path and do not /// need plugins. -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct ArrowSession { exporters: ArcSwapMap>, exporters_by_vortex: ArcSwapMap>, @@ -314,7 +314,8 @@ impl ArrowSession { /// 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::from_arrow`]. + /// Arrow → Vortex mapping via [`DType::try_from_arrow`]. + #[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() { @@ -346,7 +347,7 @@ impl ArrowSession { .collect::>>()?; DType::Struct(StructFields::from_iter(entries), nullability) } - _ => DType::from_arrow(field), + _ => DType::try_from_arrow(field)?, }) } @@ -407,6 +408,7 @@ impl ArrowSession { /// /// With `target = None` the fallback path picks the array's preferred Arrow physical type /// and executes directly into that, ignoring extension types. + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] pub fn execute_arrow( &self, array: ArrayRef, @@ -475,6 +477,7 @@ impl ArrowSession { /// through to the canonical Arrow → Vortex array conversion. pub fn from_arrow_array(&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)); if !importers.is_empty() { let dtype = self.from_arrow_field(field)?; @@ -610,11 +613,11 @@ impl SessionVar for ArrowSession { /// Extension trait for accessing the [`ArrowSession`] on a Vortex session. pub trait ArrowSessionExt: SessionExt { /// Get the Arrow session. - fn arrow(&self) -> Ref<'_, ArrowSession>; + fn arrow(&self) -> SessionGuard<'_, ArrowSession>; } impl ArrowSessionExt for S { - fn arrow(&self) -> Ref<'_, ArrowSession> { + fn arrow(&self) -> SessionGuard<'_, ArrowSession> { self.get::() } } @@ -631,8 +634,8 @@ mod tests { use vortex_error::VortexResult; use super::*; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::Nullability; @@ -765,8 +768,9 @@ mod tests { #[test] fn execute_arrow_target_none_preserves_top_level_uuid_metadata() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let session = LEGACY_SESSION.arrow(); + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); let mut field = Field::new("id", DataType::FixedSizeBinary(16), false); field.try_with_extension_type(ArrowUuid)?; diff --git a/vortex-array/src/buffer.rs b/vortex-array/src/buffer.rs index 6c59140c9a6..43039eb52ab 100644 --- a/vortex-array/src/buffer.rs +++ b/vortex-array/src/buffer.rs @@ -9,7 +9,6 @@ use std::ops::Range; use std::sync::Arc; use futures::future::BoxFuture; -use vortex_buffer::ALIGNMENT_TO_HOST_COPY; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; @@ -321,7 +320,7 @@ impl BufferHandle { pub fn try_to_host_sync(&self) -> VortexResult { match &self.0 { Inner::Host(b) => Ok(b.clone()), - Inner::Device(device) => device.copy_to_host_sync(ALIGNMENT_TO_HOST_COPY), + Inner::Device(device) => device.copy_to_host_sync(Alignment::HOST_COPY), } } @@ -331,7 +330,7 @@ impl BufferHandle { pub fn try_into_host_sync(self) -> VortexResult { match self.0 { Inner::Host(b) => Ok(b), - Inner::Device(device) => device.copy_to_host_sync(ALIGNMENT_TO_HOST_COPY), + Inner::Device(device) => device.copy_to_host_sync(Alignment::HOST_COPY), } } @@ -352,7 +351,7 @@ impl BufferHandle { let buffer = b.clone(); Ok(Box::pin(async move { Ok(buffer) })) } - Inner::Device(device) => device.copy_to_host(ALIGNMENT_TO_HOST_COPY), + Inner::Device(device) => device.copy_to_host(Alignment::HOST_COPY), } } @@ -370,7 +369,7 @@ impl BufferHandle { pub fn try_into_host(self) -> VortexResult>> { match self.0 { Inner::Host(b) => Ok(Box::pin(async move { Ok(b) })), - Inner::Device(device) => device.copy_to_host(ALIGNMENT_TO_HOST_COPY), + Inner::Device(device) => device.copy_to_host(Alignment::HOST_COPY), } } diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index e829a58f6ec..f26cbf04b57 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_ensure; use vortex_mask::Mask; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; @@ -71,6 +72,17 @@ impl BoolBuilder { self.nulls.finish_with_nullability(self.dtype.nullability()), ) } + + pub(crate) fn append_bool_array( + &mut self, + array: &BoolArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + self.inner.append_buffer(&array.to_bit_buffer()); + self.nulls + .append_validity_mask(&BoolArrayExt::validity(array).execute_mask(array.len(), ctx)?); + Ok(()) + } } impl ArrayBuilder for BoolBuilder { @@ -118,19 +130,8 @@ impl ArrayBuilder for BoolBuilder { unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { #[expect(deprecated)] let bool_array = array.to_bool(); - - self.inner.append_buffer(&bool_array.to_bit_buffer()); - self.nulls.append_validity_mask( - &bool_array - .as_ref() - .validity() - .vortex_expect("validity_mask") - .execute_mask( - bool_array.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("Failed to compute validity mask"), - ); + self.append_bool_array(&bool_array, &mut LEGACY_SESSION.create_execution_ctx()) + .vortex_expect("Failed to append bool array"); } fn reserve_exact(&mut self, additional: usize) { @@ -160,8 +161,8 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; @@ -198,7 +199,7 @@ mod tests { let chunk_count = 10; let chunk = make_opt_bool_chunks(len, chunk_count); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); chunk .clone() @@ -220,6 +221,7 @@ mod tests { #[test] fn test_append_scalar() { + let mut ctx = array_session().create_execution_ctx(); let mut builder = BoolBuilder::with_capacity(Nullability::Nullable, 10); // Test appending true value. @@ -236,7 +238,7 @@ mod tests { let array = builder.finish_into_bool(); let expected = BoolArray::from_iter([Some(true), Some(false), None]); - assert_arrays_eq!(&array, &expected); + assert_arrays_eq!(&array, &expected, &mut ctx); // Test wrong dtype error. let mut builder = BoolBuilder::with_capacity(Nullability::NonNullable, 10); diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index f8d438d477b..aa777051e39 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -309,8 +309,8 @@ impl Default for DecimalBuffer { #[cfg(test)] mod tests { - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::DecimalBuilder; @@ -333,10 +333,10 @@ mod tests { for i in 0..i8s.len() { assert_eq!( - i8s.execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + i8s.execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(), i128s - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap() ); } @@ -344,6 +344,7 @@ mod tests { #[test] fn test_append_scalar() { + let mut ctx = array_session().create_execution_ctx(); use crate::scalar::Scalar; // Simply test that the builder accepts its own finish output via scalar. @@ -357,19 +358,19 @@ mod tests { [Some(1234i64), Some(5678), None], DecimalDType::new(10, 2), ); - assert_arrays_eq!(&array, &expected); + 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()); for i in 0..array.len() { let scalar = array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(); builder2.append_scalar(&scalar).unwrap(); } let array2 = builder2.finish(); - assert_arrays_eq!(&array2, &array); + assert_arrays_eq!(&array2, &array, &mut ctx); // Test wrong dtype error. let mut builder = DecimalBuilder::new::(DecimalDType::new(10, 2), false.into()); diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs index 081cef33f39..33b2e5225f7 100644 --- a/vortex-array/src/builders/dict/bytes.rs +++ b/vortex-array/src/builders/dict/bytes.rs @@ -1,15 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::cell::OnceCell; use std::hash::BuildHasher; use std::mem; use std::sync::Arc; +use num_traits::AsPrimitive; +use vortex_array::ExecutionCtx; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; +use vortex_mask::AllOr; +use vortex_mask::Mask; use vortex_utils::aliases::hash_map::DefaultHashBuilder; use vortex_utils::aliases::hash_map::HashTable; use vortex_utils::aliases::hash_map::HashTableEntry; @@ -18,23 +25,24 @@ use vortex_utils::aliases::hash_map::RandomState; use super::DictConstraints; use super::DictEncoder; use crate::ArrayRef; +use crate::ArrayView; use crate::IntoArray; -use crate::accessor::ArrayAccessor; use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; +use crate::arrays::varbin::VarBinArrayExt; use crate::arrays::varbinview::build_views::BinaryView; -#[expect(deprecated)] -use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::PType; use crate::dtype::UnsignedPType; +use crate::match_each_integer_ptype; use crate::validity::Validity; /// Dictionary encode varbin array. Specializes for primitive byte arrays to avoid double copying -pub struct BytesDictBuilder { - lookup: Option>, +pub struct BytesDictBuilder { + lookup: Option>, + null_code: OnceCell, views: BufferMut, values: ByteBufferMut, values_nulls: BitBufferMut, @@ -58,6 +66,7 @@ impl BytesDictBuilder { Self { lookup: Some(HashTable::new()), views: BufferMut::::empty(), + null_code: OnceCell::new(), values: BufferMut::empty(), values_nulls: BitBufferMut::empty(), hasher: DefaultHashBuilder::default(), @@ -71,18 +80,18 @@ impl BytesDictBuilder { self.views.len() * size_of::() + self.values.len() } - fn lookup_bytes(&self, idx: usize) -> Option<&[u8]> { - self.values_nulls.value(idx).then(|| { - let bin_view = &self.views[idx]; - if bin_view.is_inlined() { - bin_view.as_inlined().value() - } else { - &self.values[bin_view.as_view().as_range()] - } - }) + fn lookup_bytes(&self, idx: usize) -> &[u8] { + let bin_view = &self.views[idx]; + if bin_view.is_inlined() { + bin_view.as_inlined().value() + } else { + &self.values[bin_view.as_view().as_range()] + } } - fn encode_value(&mut self, lookup: &mut HashTable, val: Option<&[u8]>) -> Option { + /// Returns `None` when assigning a code would exceed the dictionary constraints, + /// and callers should stop encoding after the current prefix. + fn encode_value(&mut self, lookup: &mut HashTable, val: &[u8]) -> Option { match lookup.entry( self.hasher.hash_one(val), |idx| val == self.lookup_bytes(idx.as_()), @@ -95,35 +104,25 @@ impl BytesDictBuilder { } let next_code = self.views.len(); - match val { - None => { - // Null value - self.views.push(BinaryView::default()); - self.values_nulls.append_false(); - } - Some(val) => { - let view = BinaryView::make_view( - val, - 0, - u32::try_from(self.values.len()) - .vortex_expect("values length must fit in u32"), - ); - let additional_bytes = if view.is_inlined() { - size_of::() - } else { - size_of::() + val.len() - }; + let view = BinaryView::make_view( + val, + 0, + u32::try_from(self.values.len()).vortex_expect("values length must fit in u32"), + ); + let additional_bytes = if view.is_inlined() { + size_of::() + } else { + size_of::() + val.len() + }; - if self.dict_bytes() + additional_bytes > self.max_dict_bytes { - return None; - } + if self.dict_bytes() + additional_bytes > self.max_dict_bytes { + return None; + } - self.views.push(view); - self.values_nulls.append_true(); - if !view.is_inlined() { - self.values.extend_from_slice(val); - } - } + self.views.push(view); + self.values_nulls.append_true(); + if !view.is_inlined() { + self.values.extend_from_slice(val); } let next_code = Code::from_usize(next_code).unwrap_or_else(|| { @@ -134,29 +133,138 @@ impl BytesDictBuilder { } } - fn encode_bytes>(&mut self, accessor: &A, len: usize) -> ArrayRef { + /// Returns `None` when assigning the null code would exceed the dictionary constraints, + /// and callers should stop encoding after the current prefix. + fn encode_null(&mut self) -> Option { + if let Some(code) = self.null_code.get() { + return Some(*code); + } + + if self.views.len() >= self.max_dict_len + || self.dict_bytes() + size_of::() > self.max_dict_bytes + { + return None; + } + + let code = self.views.len(); + self.views.push(BinaryView::default()); + self.values_nulls.append_false(); + let code = Code::from_usize(code) + .unwrap_or_else(|| vortex_panic!("{} has to fit into {}", code, Code::PTYPE)); + self.null_code + .set(code) + .ok() + .vortex_expect("null code is initialized once"); + Some(code) + } + + /// Encode row values against the dictionary, honoring the supplied validity mask. + /// + /// `value_at` is called only for valid rows. That matters for VarBinView arrays because null + /// rows can hold arbitrary view metadata. + fn encode_validity<'a, F>( + &mut self, + len: usize, + validity_mask: Mask, + mut value_at: F, + ) -> VortexResult + where + 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); - accessor.with_iterator(|it| { - for value in it { - let Some(code) = self.encode_value(&mut local_lookup, value) else { - break; - }; - // SAFETY: we reserved capacity in the buffer for `len` elements - unsafe { codes.push_unchecked(code) } + match validity_mask.bit_buffer() { + AllOr::All => { + for idx in 0..len { + let Some(code) = self.encode_value(&mut local_lookup, value_at(idx)) else { + break; + }; + // SAFETY: we reserved capacity in the buffer for `len` elements + unsafe { codes.push_unchecked(code) } + } + } + AllOr::None => { + if let Some(code) = self.encode_null() { + unsafe { codes.push_n_unchecked(code, len) } + } } - }); + AllOr::Some(b) => { + for (idx, valid) in b.iter().enumerate() { + if !valid { + let Some(code) = self.encode_null() else { + break; + }; + // SAFETY: we reserved capacity in the buffer for `len` elements + unsafe { codes.push_unchecked(code) } + } else { + let Some(code) = self.encode_value(&mut local_lookup, value_at(idx)) else { + break; + }; + // SAFETY: we reserved capacity in the buffer for `len` elements + unsafe { codes.push_unchecked(code) } + } + } + } + } // Restore lookup dictionary back into the struct self.lookup = Some(local_lookup); - PrimitiveArray::new(codes, Validity::NonNullable).into_array() + Ok(PrimitiveArray::new(codes, Validity::NonNullable)) + } + + fn encode_varbin( + &mut self, + var_bin: ArrayView, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let offsets = var_bin.offsets().clone().execute::(ctx)?; + let bytes = var_bin.bytes(); + let validity_mask = var_bin.validity()?.execute_mask(var_bin.len(), ctx)?; + let len = var_bin.len(); + + match_each_integer_ptype!(offsets.ptype(), |P| { + let slice_offsets = offsets.as_slice::

(); + self.encode_validity(len, validity_mask, |idx| { + let start: usize = slice_offsets[idx].as_(); + let end: usize = slice_offsets[idx + 1].as_(); + &bytes[start..end] + }) + }) + } + + fn encode_varbinview( + &mut self, + var_bin_view: ArrayView, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let validity_mask = var_bin_view + .validity()? + .execute_mask(var_bin_view.len(), ctx)?; + let len = var_bin_view.len(); + let views = var_bin_view.views(); + let buffers = var_bin_view + .data_buffers() + .iter() + .map(|b| b.as_host()) + .collect::>(); + + self.encode_validity(len, validity_mask, |idx| view_bytes(&buffers, &views[idx])) + } +} + +fn view_bytes<'a>(buffers: &[&'a ByteBuffer], view: &'a BinaryView) -> &'a [u8] { + if view.is_inlined() { + view.as_inlined().value() + } else { + let view = view.as_view(); + &buffers[view.buffer_index as usize][view.as_range()] } } impl DictEncoder for BytesDictBuilder { - fn encode(&mut self, array: &ArrayRef) -> ArrayRef { + fn encode(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { debug_assert_eq!( &self.dtype, array.dtype(), @@ -165,17 +273,15 @@ impl DictEncoder for BytesDictBuilder { self.dtype ); - let len = array.len(); if let Some(varbinview) = array.as_opt::() { - self.encode_bytes(&varbinview.into_owned(), len) + self.encode_varbinview(varbinview, ctx) } else if let Some(varbin) = array.as_opt::() { - self.encode_bytes(&varbin.into_owned(), len) + self.encode_varbin(varbin, ctx) } else { // NOTE(aduffy): it is very rare that this path would be taken, only e.g. // if we're performing dictionary encoding downstream of some other compression. - #[expect(deprecated)] - let varbinview = array.to_varbinview(); - self.encode_bytes(&varbinview, len) + let vbv_array = array.clone().execute::(ctx)?; + self.encode_varbinview(vbv_array.as_view(), ctx) } } @@ -204,46 +310,49 @@ impl DictEncoder for BytesDictBuilder { #[cfg(test)] mod test { - use std::str; + use std::sync::Arc; use std::sync::LazyLock; + use vortex_buffer::Buffer; + use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::accessor::ArrayAccessor; + use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; + use crate::arrays::VarBinViewArray; use crate::arrays::dict::DictArraySlotsExt; + use crate::arrays::varbinview::BinaryView; + use crate::buffer::BufferHandle; use crate::builders::dict::dict_encode; - use crate::session::ArraySession; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::validity::Validity; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[test] - fn encode_varbin() { - let arr = VarBinArray::from(vec!["hello", "world", "hello", "again", "world"]); - let dict = dict_encode(&arr.into_array(), &mut SESSION.create_execution_ctx()).unwrap(); - #[expect(deprecated)] - let codes = dict.codes().to_primitive(); + fn encode_varbin() -> VortexResult<()> { + let arr = VarBinViewArray::from_iter_str(vec!["hello", "world", "hello", "again", "world"]); + let mut ctx = SESSION.create_execution_ctx(); + let dict = dict_encode(&arr.into_array(), &mut ctx)?; + let codes = dict.codes().clone().execute::(&mut ctx)?; assert_eq!(codes.as_slice::(), &[0, 1, 0, 2, 1]); - #[expect(deprecated)] - let values = dict.values().to_varbinview(); - values.with_iterator(|iter| { - assert_eq!( - iter.flatten() - .map(|b| unsafe { str::from_utf8_unchecked(b) }) - .collect::>(), - vec!["hello", "world", "again"] - ); - }); + let values = dict.values().clone().execute::(&mut ctx)?; + let mask = values.validity()?.execute_mask(values.len(), &mut ctx)?; + let decoded = (0..values.len()) + .filter(|&i| mask.value(i)) + .map(|i| unsafe { String::from_utf8_unchecked(values.bytes_at(i).to_vec()) }) + .collect::>(); + assert_eq!(decoded, vec!["hello", "world", "again"]); + Ok(()) } #[test] - fn encode_varbin_nulls() { - let arr: VarBinArray = vec![ + fn encode_varbin_nulls() -> VortexResult<()> { + let arr: VarBinViewArray = vec![ Some("hello"), None, Some("world"), @@ -255,37 +364,70 @@ mod test { ] .into_iter() .collect(); - let dict = dict_encode(&arr.into_array(), &mut SESSION.create_execution_ctx()).unwrap(); - #[expect(deprecated)] - let codes = dict.codes().to_primitive(); + let mut ctx = SESSION.create_execution_ctx(); + let dict = dict_encode(&arr.into_array(), &mut ctx)?; + let codes = dict.codes().clone().execute::(&mut ctx)?; assert_eq!(codes.as_slice::(), &[0, 1, 2, 0, 1, 3, 2, 1]); - #[expect(deprecated)] - let values = dict.values().to_varbinview(); - values.with_iterator(|iter| { - assert_eq!( - iter.map(|b| b.map(|v| unsafe { str::from_utf8_unchecked(v) })) - .collect::>(), - vec![Some("hello"), None, Some("world"), Some("again")] - ); - }); + let values = dict.values().clone().execute::(&mut ctx)?; + let mask = values.validity()?.execute_mask(values.len(), &mut ctx)?; + let decoded = (0..values.len()) + .map(|i| { + mask.value(i) + .then(|| unsafe { String::from_utf8_unchecked(values.bytes_at(i).to_vec()) }) + }) + .collect::>(); + assert_eq!( + decoded, + vec![ + Some("hello".to_string()), + None, + Some("world".to_string()), + Some("again".to_string()) + ] + ); + Ok(()) + } + + #[test] + fn encode_varbinview_ignores_invalid_null_views() { + let value = b"outlined value"; + let valid_view = BinaryView::make_view(value, 0, 0); + let invalid_null_view = BinaryView::make_view(b"invalid null view", 99, 0); + let views = Buffer::copy_from([valid_view, invalid_null_view, valid_view]); + let buffers = Arc::from([BufferHandle::new_host(ByteBuffer::copy_from(value))]); + let arr = unsafe { + VarBinViewArray::new_handle_unchecked( + BufferHandle::new_host(views.into_byte_buffer()), + buffers, + DType::Utf8(Nullability::Nullable), + Validity::from_iter([true, false, true]), + ) + } + .into_array(); + + let dict = dict_encode(&arr, &mut SESSION.create_execution_ctx()).unwrap(); + let codes = dict + .codes() + .clone() + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); + assert_eq!(codes.as_slice::(), &[0, 1, 0]); } #[test] - fn repeated_values() { + fn repeated_values() -> VortexResult<()> { let arr = VarBinArray::from(vec!["a", "a", "b", "b", "a", "b", "a", "b"]); - let dict = dict_encode(&arr.into_array(), &mut SESSION.create_execution_ctx()).unwrap(); - #[expect(deprecated)] - let values = dict.values().to_varbinview(); - values.with_iterator(|iter| { - assert_eq!( - iter.flatten() - .map(|b| unsafe { str::from_utf8_unchecked(b) }) - .collect::>(), - vec!["a", "b"] - ); - }); - #[expect(deprecated)] - let codes = dict.codes().to_primitive(); + let mut ctx = SESSION.create_execution_ctx(); + let dict = dict_encode(&arr.into_array(), &mut ctx)?; + let values = dict.values().clone().execute::(&mut ctx)?; + let mask = values.validity()?.execute_mask(values.len(), &mut ctx)?; + let decoded = (0..values.len()) + .filter(|&i| mask.value(i)) + .map(|i| unsafe { String::from_utf8_unchecked(values.bytes_at(i).to_vec()) }) + .collect::>(); + assert_eq!(decoded, vec!["a", "b"]); + let codes = dict.codes().clone().execute::(&mut ctx)?; assert_eq!(codes.as_slice::(), &[0, 0, 1, 1, 0, 1, 0, 1]); + Ok(()) } } diff --git a/vortex-array/src/builders/dict/mod.rs b/vortex-array/src/builders/dict/mod.rs index c64e436bb45..94834f83cc4 100644 --- a/vortex-array/src/builders/dict/mod.rs +++ b/vortex-array/src/builders/dict/mod.rs @@ -35,7 +35,7 @@ pub const UNCONSTRAINED: DictConstraints = DictConstraints { pub trait DictEncoder: Send { /// Assign dictionary codes to the given input array. - fn encode(&mut self, array: &ArrayRef) -> ArrayRef; + fn encode(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; /// Clear the encoder state to make it ready for a new round of decoding. fn reset(&mut self) -> ArrayRef; @@ -68,8 +68,7 @@ pub fn dict_encode_with_constraints( ctx: &mut ExecutionCtx, ) -> VortexResult { let mut encoder = dict_encoder(array, constraints); - let encoded = encoder.encode(array); - let codes = encoded.execute::(ctx)?.narrow(ctx)?; + 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 // since we build the dictionary from the codes we observe during encoding diff --git a/vortex-array/src/builders/dict/primitive.rs b/vortex-array/src/builders/dict/primitive.rs index 7c0e68a456a..dea2932367b 100644 --- a/vortex-array/src/builders/dict/primitive.rs +++ b/vortex-array/src/builders/dict/primitive.rs @@ -1,23 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::cell::OnceCell; use std::hash::Hash; use std::mem; use rustc_hash::FxBuildHasher; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; +use vortex_mask::Mask; use vortex_utils::aliases::hash_map::Entry; use vortex_utils::aliases::hash_map::HashMap; use super::DictConstraints; use super::DictEncoder; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; -#[expect(deprecated)] -use crate::ToCanonical as _; -use crate::accessor::ArrayAccessor; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::NativeValue; use crate::dtype::NativePType; @@ -72,6 +74,7 @@ where .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(), nullability, @@ -79,8 +82,10 @@ where } } - fn encode_value(&mut self, v: Option) -> Option { - match self.lookup.entry(v.map(NativeValue)) { + /// Returns `None` when assigning a code would exceed the dictionary constraints, + /// and callers should stop encoding after the current prefix. + fn encode_value(&mut self, v: T) -> Option { + match self.lookup.entry(NativeValue(v)) { Entry::Occupied(o) => Some(*o.get()), Entry::Vacant(vac) => { if self.values.len() >= self.max_dict_len { @@ -89,28 +94,43 @@ where let next_code = Code::from_usize(self.values.len()).unwrap_or_else(|| { vortex_panic!("{} has to fit into {}", self.values.len(), Code::PTYPE) }); - vac.insert(next_code); - match v { - None => { - self.values.push(T::default()); - self.values_nulls.append_false(); - } - Some(v) => { - self.values.push(v); - self.values_nulls.append_true(); - } - } - Some(next_code) + self.values.push(v); + self.values_nulls.append_true(); + Some(*vac.insert(next_code)) } } } + + /// Returns `None` when assigning the null code would exceed the dictionary constraints, + /// and callers should stop encoding after the current prefix. + fn encode_null(&mut self) -> Option { + if let Some(code) = self.null_code.get() { + return Some(*code); + } + + if self.values.len() >= self.max_dict_len { + return None; + } + + let code = Code::from_usize(self.values.len()).unwrap_or_else(|| { + vortex_panic!("{} has to fit into {}", self.values.len(), Code::PTYPE) + }); + self.values.push(T::default()); + self.values_nulls.append_false(); + self.null_code + .set(code) + .ok() + .vortex_expect("null code is initialized once"); + Some(code) + } } /// Dictionary encode primitive array with given PType. /// /// Null values are stored in the values of the dictionary such that codes are always non-null. pub struct PrimitiveDictBuilder { - lookup: HashMap>, Code, FxBuildHasher>, + lookup: HashMap, Code, FxBuildHasher>, + null_code: OnceCell, values: BufferMut, values_nulls: BitBufferMut, nullability: Nullability, @@ -123,21 +143,43 @@ where NativeValue: Hash + Eq, Code: UnsignedPType, { - fn encode(&mut self, array: &ArrayRef) -> ArrayRef { + fn encode(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let mut codes = BufferMut::::with_capacity(array.len()); - #[expect(deprecated)] - let prim = array.to_primitive(); - prim.with_iterator(|it| { - for value in it { - let Some(code) = self.encode_value(value.copied()) else { - break; - }; - unsafe { codes.push_unchecked(code) } + let prim = array.clone().execute::(ctx)?; + match prim.validity()?.execute_mask(array.len(), ctx)? { + Mask::AllTrue(_) => { + for &value in prim.as_slice::() { + let Some(code) = self.encode_value(value) else { + break; + }; + unsafe { codes.push_unchecked(code) } + } } - }); + Mask::AllFalse(_) => { + if let Some(code) = self.encode_null() { + unsafe { codes.push_n_unchecked(code, array.len()) } + } + } + Mask::Values(v) => { + let bit_buff = v.bit_buffer(); + for (&value, valid) in prim.as_slice::().iter().zip(bit_buff) { + if !valid { + let Some(code) = self.encode_null() else { + break; + }; + unsafe { codes.push_unchecked(code) } + } else { + let Some(code) = self.encode_value(value) else { + break; + }; + unsafe { codes.push_unchecked(code) } + } + } + } + } - PrimitiveArray::new(codes, Validity::NonNullable).into_array() + Ok(PrimitiveArray::new(codes, Validity::NonNullable)) } fn reset(&mut self) -> ArrayRef { @@ -157,8 +199,6 @@ where mod test { use std::sync::LazyLock; - #[expect(unused_imports)] - use itertools::Itertools; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -168,25 +208,25 @@ mod test { use crate::assert_arrays_eq; use crate::builders::dict::dict_encode; use crate::builders::dict::primitive::PrimitiveArray; - use crate::session::ArraySession; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[test] fn encode_primitive() { + let mut ctx = SESSION.create_execution_ctx(); let arr = buffer![1, 1, 3, 3, 3].into_array(); let dict = dict_encode(&arr, &mut SESSION.create_execution_ctx()).unwrap(); let expected_codes = buffer![0u8, 0, 1, 1, 1].into_array(); - assert_arrays_eq!(dict.codes(), expected_codes); + assert_arrays_eq!(dict.codes(), expected_codes, &mut ctx); let expected_values = buffer![1i32, 3].into_array(); - assert_arrays_eq!(dict.values(), expected_values); + assert_arrays_eq!(dict.values(), expected_values, &mut ctx); } #[test] fn encode_primitive_nulls() { + let mut ctx = SESSION.create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([ Some(1), Some(1), @@ -196,14 +236,15 @@ mod test { None, Some(3), None, - ]); - let dict = dict_encode(&arr.into_array(), &mut SESSION.create_execution_ctx()).unwrap(); + ]) + .into_array(); + let dict = dict_encode(&arr, &mut SESSION.create_execution_ctx()).unwrap(); let expected_codes = buffer![0u8, 0, 1, 2, 2, 1, 2, 1].into_array(); - assert_arrays_eq!(dict.codes(), expected_codes); + assert_arrays_eq!(dict.codes(), expected_codes, &mut ctx); let expected_values = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); - assert_arrays_eq!(dict.values(), expected_values); + assert_arrays_eq!(dict.values(), expected_values, &mut ctx); } } diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index 7e1da3b8f3e..61d5fe8d7b0 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -127,6 +127,8 @@ impl ArrayBuilder for ExtensionBuilder { #[cfg(test)] mod tests { use super::*; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; @@ -137,6 +139,7 @@ mod tests { #[test] fn test_append_scalar() { + 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()); @@ -165,7 +168,7 @@ mod tests { PrimitiveArray::from_option_iter([Some(42i32), Some(84), None]).into_array(), ); - assert_arrays_eq!(&array, &expected); + assert_arrays_eq!(&array, &expected, &mut ctx); assert_eq!(array.len(), 3); // Test wrong dtype error. diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 7a3c3bb068a..9129d8cb644 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -12,6 +12,7 @@ use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; @@ -83,7 +84,11 @@ impl FixedSizeListBuilder { /// /// Note that the list entry will be non-null but the elements themselves are allowed to be null /// (only if the elements [`DType`] is nullable, of course). - pub fn append_array_as_list(&mut self, array: &ArrayRef) -> VortexResult<()> { + pub fn append_array_as_list( + &mut self, + array: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { vortex_ensure!( array.dtype() == self.element_dtype(), "Array dtype {:?} does not match list element dtype {:?}", @@ -97,7 +102,7 @@ impl FixedSizeListBuilder { self.list_size() ); - self.elements_builder.extend_from_array(array); + array.append_to_builder(self.elements_builder.as_mut(), ctx)?; self.nulls.append_non_null(); Ok(()) @@ -282,10 +287,10 @@ mod tests { use super::FixedSizeListBuilder; use crate::IntoArray as _; - use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::builders::ArrayBuilder; @@ -443,7 +448,7 @@ mod tests { #[test] fn test_nullable_lists_non_nullable_elements() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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); @@ -562,7 +567,7 @@ mod tests { #[test] fn test_append_nulls() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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); @@ -592,7 +597,7 @@ mod tests { #[test] fn test_append_scalar_nulls() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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); @@ -665,7 +670,7 @@ mod tests { #[test] fn test_extend_from_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); // Create a source array. @@ -736,7 +741,7 @@ mod tests { #[test] fn test_extend_degenerate_arrays() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); // Create degenerate source arrays (size = 0). @@ -841,7 +846,7 @@ mod tests { #[test] fn test_mixed_operations() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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); @@ -927,7 +932,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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); @@ -997,12 +1002,13 @@ mod tests { #[test] 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); // Append a primitive array as a single list entry. let arr1 = buffer![1i32, 2, 3].into_array(); - builder.append_array_as_list(&arr1).unwrap(); + builder.append_array_as_list(&arr1, &mut ctx).unwrap(); // Interleave with a list scalar. builder @@ -1018,7 +1024,7 @@ mod tests { // Append another primitive array as a single list entry. let arr2 = buffer![4i32, 5, 6].into_array(); - builder.append_array_as_list(&arr2).unwrap(); + builder.append_array_as_list(&arr2, &mut ctx).unwrap(); // Interleave with another list scalar. builder @@ -1048,11 +1054,19 @@ mod tests { let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 10); let wrong_dtype_arr = buffer![1i64, 2, 3].into_array(); - assert!(builder.append_array_as_list(&wrong_dtype_arr).is_err()); + assert!( + builder + .append_array_as_list(&wrong_dtype_arr, &mut ctx) + .is_err() + ); // Test length mismatch error. let mut builder = FixedSizeListBuilder::with_capacity(dtype, 3, NonNullable, 10); let wrong_len_arr = buffer![1i32, 2].into_array(); - assert!(builder.append_array_as_list(&wrong_len_arr).is_err()); + assert!( + builder + .append_array_as_list(&wrong_len_arr, &mut ctx) + .is_err() + ); } } diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index ac97230daa8..12c74ba53a4 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -13,6 +13,7 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; @@ -95,7 +96,11 @@ impl ListBuilder { /// /// Note that the list entry will be non-null but the elements themselves are allowed to be null /// (only if the elements [`DType`] in nullable, of course). - pub fn append_array_as_list(&mut self, array: &ArrayRef) -> VortexResult<()> { + pub fn append_array_as_list( + &mut self, + array: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { vortex_ensure!( array.dtype() == self.element_dtype(), "Array dtype {:?} does not match list element dtype {:?}", @@ -103,7 +108,7 @@ impl ListBuilder { self.element_dtype() ); - self.elements_builder.extend_from_array(array); + array.append_to_builder(self.elements_builder.as_mut(), ctx)?; self.nulls.append_non_null(); self.offsets_builder.append_value( O::from_usize(self.elements_builder.len()) @@ -324,9 +329,9 @@ mod tests { use vortex_error::VortexExpect; use crate::IntoArray; - use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::arrays::list::ListArrayExt; @@ -451,7 +456,7 @@ mod tests { .unwrap(); assert_eq!(list.len(), 3); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut builder = ListBuilder::::with_capacity(Arc::new(I32.into()), Nullable, 18, 9); builder.extend_from_array(&list); @@ -478,9 +483,9 @@ mod tests { let actual = builder.finish_into_canonical().into_listview(); - assert_arrays_eq!(actual.elements(), expected.elements()); + assert_arrays_eq!(actual.elements(), expected.elements(), &mut ctx); - assert_arrays_eq!(actual.offsets(), expected.offsets()); + assert_arrays_eq!(actual.offsets(), expected.offsets(), &mut ctx); assert!( actual @@ -539,18 +544,18 @@ mod tests { assert_eq!( one_trailing_unused_element - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), canon_values - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() ); assert_eq!( second_array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), canon_values - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() ); } @@ -580,7 +585,7 @@ mod tests { let array = builder.finish_into_list(); assert_eq!(array.len(), 3); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Check actual values using scalar_at. @@ -637,12 +642,13 @@ mod tests { #[test] 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); // Append a primitive array as a single list entry. let arr1 = buffer![1i32, 2, 3].into_array(); - builder.append_array_as_list(&arr1).unwrap(); + builder.append_array_as_list(&arr1, &mut ctx).unwrap(); // Interleave with a list scalar. builder @@ -658,11 +664,11 @@ mod tests { // Append another primitive array as a single list entry. let arr2 = buffer![4i32, 5].into_array(); - builder.append_array_as_list(&arr2).unwrap(); + builder.append_array_as_list(&arr2, &mut ctx).unwrap(); // Append an empty array as a single list entry (empty list). let arr3 = buffer![0i32; 0].into_array(); - builder.append_array_as_list(&arr3).unwrap(); + builder.append_array_as_list(&arr3, &mut ctx).unwrap(); // Interleave with another list scalar (empty list). builder @@ -675,18 +681,24 @@ mod tests { // Verify elements array: [1, 2, 3, 10, 11, 4, 5]. assert_arrays_eq!( list.elements(), - PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]) + PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]), + &mut ctx ); // Verify offsets array. assert_arrays_eq!( list.offsets(), - PrimitiveArray::from_iter([0u32, 3, 5, 7, 7, 7]) + PrimitiveArray::from_iter([0u32, 3, 5, 7, 7, 7]), + &mut ctx ); // Test dtype mismatch error. let mut builder = ListBuilder::::with_capacity(dtype, NonNullable, 20, 10); let wrong_dtype_arr = buffer![1i64, 2, 3].into_array(); - assert!(builder.append_array_as_list(&wrong_dtype_arr).is_err()); + assert!( + builder + .append_array_as_list(&wrong_dtype_arr, &mut ctx) + .is_err() + ); } } diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 4c350f1888d..cb131048339 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -20,6 +20,7 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; +use crate::ExecutionCtx; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array::IntoArray; @@ -117,7 +118,11 @@ impl ListViewBuilder { /// /// Note that the list entry will be non-null but the elements themselves are allowed to be null /// (only if the elements [`DType`] is nullable, of course). - pub fn append_array_as_list(&mut self, array: &ArrayRef) -> VortexResult<()> { + pub fn append_array_as_list( + &mut self, + array: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { vortex_ensure!( array.dtype() == self.element_dtype(), "Array dtype {:?} does not match list element dtype {:?}", @@ -135,7 +140,7 @@ impl ListViewBuilder { "appending this list would cause an offset overflow" ); - self.elements_builder.extend_from_array(array); + array.append_to_builder(self.elements_builder.as_mut(), ctx)?; self.nulls.append_non_null(); self.offsets_builder.append_value( @@ -434,8 +439,8 @@ mod tests { use super::ListViewBuilder; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::listview::ListViewArrayExt; @@ -460,7 +465,7 @@ mod tests { #[test] fn test_basic_append_and_nulls() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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); @@ -498,7 +503,8 @@ mod tests { // Check first list: [1, 2, 3]. assert_arrays_eq!( listview.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([1i32, 2, 3]) + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx ); // Check empty list. @@ -516,12 +522,14 @@ mod tests { // Check last list: [4, 5]. assert_arrays_eq!( listview.list_elements_at(3).unwrap(), - PrimitiveArray::from_iter([4i32, 5]) + PrimitiveArray::from_iter([4i32, 5]), + &mut ctx ); } #[test] fn test_different_offset_size_types() { + let mut ctx = array_session().create_execution_ctx(); // Test u32 offsets with u8 sizes. let dtype: Arc = Arc::new(I32.into()); let mut builder = @@ -555,13 +563,15 @@ mod tests { // Verify first list: [1, 2]. assert_arrays_eq!( listview.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([1i32, 2]) + PrimitiveArray::from_iter([1i32, 2]), + &mut ctx ); // Verify second list: [3, 4, 5]. assert_arrays_eq!( listview.list_elements_at(1).unwrap(), - PrimitiveArray::from_iter([3i32, 4, 5]) + PrimitiveArray::from_iter([3i32, 4, 5]), + &mut ctx ); // Test u64 offsets with u16 sizes. @@ -584,14 +594,15 @@ mod tests { for i in 0..5i32 { assert_arrays_eq!( listview2.list_elements_at(i as usize).unwrap(), - PrimitiveArray::from_iter([i * 10]) + PrimitiveArray::from_iter([i * 10]), + &mut ctx ); } } #[test] fn test_builder_trait_methods() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + 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); @@ -637,13 +648,14 @@ mod tests { // Last is the regular list: [10, 20]. assert_arrays_eq!( listview.list_elements_at(4).unwrap(), - PrimitiveArray::from_iter([10i32, 20]) + PrimitiveArray::from_iter([10i32, 20]), + &mut ctx ); } #[test] fn test_extend_from_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); // Create a source ListArray. @@ -679,13 +691,15 @@ mod tests { // First list: [0] (initial data). assert_arrays_eq!( listview.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([0i32]) + PrimitiveArray::from_iter([0i32]), + &mut ctx ); // Second list: [1, 2, 3] (from source). assert_arrays_eq!( listview.list_elements_at(1).unwrap(), - PrimitiveArray::from_iter([1i32, 2, 3]) + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx ); // Third list: null (from source). @@ -700,13 +714,14 @@ mod tests { // Fourth list: [4, 5] (from source). assert_arrays_eq!( listview.list_elements_at(3).unwrap(), - PrimitiveArray::from_iter([4i32, 5]) + PrimitiveArray::from_iter([4i32, 5]), + &mut ctx ); } #[test] fn test_extend_from_array_overlapping_listview() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); // Non-ZCTL source: @@ -733,7 +748,8 @@ mod tests { assert_arrays_eq!( listview.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([10i32, 20]) + PrimitiveArray::from_iter([10i32, 20]), + &mut ctx ); assert!( !listview @@ -745,7 +761,8 @@ mod tests { assert_eq!(listview.list_elements_at(1).unwrap().len(), 0); assert_arrays_eq!( listview.list_elements_at(2).unwrap(), - PrimitiveArray::from_iter([10i32]) + PrimitiveArray::from_iter([10i32]), + &mut ctx ); } @@ -773,12 +790,13 @@ mod tests { #[test] 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); // Append a primitive array as a single list entry. let arr1 = buffer![1i32, 2, 3].into_array(); - builder.append_array_as_list(&arr1).unwrap(); + builder.append_array_as_list(&arr1, &mut ctx).unwrap(); // Interleave with a list scalar. builder @@ -794,11 +812,11 @@ mod tests { // Append another primitive array as a single list entry. let arr2 = buffer![4i32, 5].into_array(); - builder.append_array_as_list(&arr2).unwrap(); + builder.append_array_as_list(&arr2, &mut ctx).unwrap(); // Append an empty array as a single list entry (empty list). let arr3 = buffer![0i32; 0].into_array(); - builder.append_array_as_list(&arr3).unwrap(); + builder.append_array_as_list(&arr3, &mut ctx).unwrap(); // Interleave with another list scalar. builder @@ -811,24 +829,31 @@ mod tests { // Verify elements array: [1, 2, 3, 10, 11, 4, 5]. assert_arrays_eq!( listview.elements(), - PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]) + PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]), + &mut ctx ); // Verify offsets array. assert_arrays_eq!( listview.offsets(), - PrimitiveArray::from_iter([0u32, 3, 5, 7, 7]) + PrimitiveArray::from_iter([0u32, 3, 5, 7, 7]), + &mut ctx ); // Verify sizes array. assert_arrays_eq!( listview.sizes(), - PrimitiveArray::from_iter([3u32, 2, 2, 0, 0]) + PrimitiveArray::from_iter([3u32, 2, 2, 0, 0]), + &mut ctx ); // Test dtype mismatch error. let mut builder = ListViewBuilder::::with_capacity(dtype, NonNullable, 20, 10); let wrong_dtype_arr = buffer![1i64, 2, 3].into_array(); - assert!(builder.append_array_as_list(&wrong_dtype_arr).is_err()); + assert!( + builder + .append_array_as_list(&wrong_dtype_arr, &mut ctx) + .is_err() + ); } } diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index e9f86fa4332..48a347aa3d5 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -11,7 +11,7 @@ //! ``` //! use vortex_array::builders::{builder_with_capacity, ArrayBuilder}; //! use vortex_array::dtype::{DType, Nullability}; -//! use vortex_array::{LEGACY_SESSION, VortexSessionExecute}; +//! use vortex_array::{VortexSessionExecute, array_session}; //! //! // Create a new builder for string data. //! let mut builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 4); @@ -22,7 +22,7 @@ //! builder.append_scalar(&"d".into()).unwrap(); //! //! let strings = builder.finish(); -//! let mut ctx = LEGACY_SESSION.create_execution_ctx(); +//! let mut ctx = array_session().create_execution_ctx(); //! //! assert_eq!(strings.execute_scalar(0, &mut ctx).unwrap(), "a".into()); //! assert_eq!(strings.execute_scalar(1, &mut ctx).unwrap(), "b".into()); @@ -225,7 +225,7 @@ pub trait ArrayBuilder: Send { /// ``` /// use vortex_array::builders::{builder_with_capacity, ArrayBuilder}; /// use vortex_array::dtype::{DType, Nullability}; -/// use vortex_array::{LEGACY_SESSION, VortexSessionExecute}; +/// use vortex_array::{VortexSessionExecute, array_session}; /// /// // Create a new builder for string data. /// let mut builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 4); @@ -236,7 +236,7 @@ pub trait ArrayBuilder: Send { /// builder.append_scalar(&"d".into()).unwrap(); /// /// let strings = builder.finish(); -/// let mut ctx = LEGACY_SESSION.create_execution_ctx(); +/// let mut ctx = array_session().create_execution_ctx(); /// /// assert_eq!(strings.execute_scalar(0, &mut ctx).unwrap(), "a".into()); /// assert_eq!(strings.execute_scalar(1, &mut ctx).unwrap(), "b".into()); diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index c4b5abfb083..e7e8d3d7512 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_ensure; use vortex_mask::Mask; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; @@ -131,6 +132,28 @@ impl PrimitiveBuilder { self.values.extend(iter); self.nulls.append_validity_mask(mask); } + + pub(crate) fn append_primitive_array( + &mut self, + array: &PrimitiveArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + debug_assert_eq!( + array.ptype(), + T::PTYPE, + "Cannot append primitive array with different ptype" + ); + + self.values.extend_from_slice(array.as_slice::()); + self.nulls.append_validity_mask( + &array + .as_ref() + .validity() + .vortex_expect("validity_mask") + .execute_mask(array.as_ref().len(), ctx)?, + ); + Ok(()) + } } impl ArrayBuilder for PrimitiveBuilder { @@ -181,25 +204,8 @@ impl ArrayBuilder for PrimitiveBuilder { #[expect(deprecated)] let array = array.to_primitive(); - // This should be checked in `extend_from_array` but we can check it again. - debug_assert_eq!( - array.ptype(), - T::PTYPE, - "Cannot extend from array with different ptype" - ); - - self.values.extend_from_slice(array.as_slice::()); - self.nulls.append_validity_mask( - &array - .as_ref() - .validity() - .vortex_expect("validity_mask") - .execute_mask( - array.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("Failed to compute validity mask"), - ); + self.append_primitive_array(&array, &mut LEGACY_SESSION.create_execution_ctx()) + .vortex_expect("Failed to append primitive array"); } fn reserve_exact(&mut self, additional: usize) { @@ -370,6 +376,8 @@ mod tests { use vortex_error::VortexExpect; use super::*; + use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; /// REGRESSION TEST: This test verifies that multiple sequential ranges have correct offsets. @@ -378,6 +386,7 @@ mod tests { /// buffer. #[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); // First range. @@ -407,7 +416,11 @@ mod tests { assert_eq!(builder.values(), &[1, 2, 3, 4, 5]); let array = builder.finish_into_primitive(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1i32, 2, 3, 4, 5])); + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]), + &mut ctx + ); } /// REGRESSION TEST: This test verifies that `append_mask` was correctly moved from @@ -441,19 +454,19 @@ mod tests { // Check validity using scalar_at - nulls will return is_null() = true. assert!( !array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert!( array - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert!( !array - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); @@ -484,6 +497,7 @@ mod tests { /// This verifies the new simplified API without the redundant `len` parameter. #[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 range = builder.uninit_range(6); @@ -498,7 +512,11 @@ mod tests { } let array = builder.finish_into_primitive(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6])); + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]), + &mut ctx + ); } /// Test that `set_bit` uses relative indexing within the range. @@ -546,13 +564,13 @@ mod tests { // Check validity - the first two should be valid (from append_value). assert!( !array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); // initial value 100 assert!( !array - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); // initial value 200 @@ -560,19 +578,19 @@ mod tests { // Check the range items with modified validity. assert!( !array - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); // range index 0 - set to valid assert!( array - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(3, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); // range index 1 - left as null assert!( !array - .execute_scalar(4, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(4, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); // range index 2 - set to valid @@ -588,12 +606,10 @@ mod tests { /// Test that creating an uninit range exceeding capacity panics. #[test] - #[should_panic( - expected = "uninit_range of len 10 exceeds builder with length 0 and capacity 6" - )] + #[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 _range = builder.uninit_range(10); + let _range = builder.uninit_range(261); } /// Test that `copy_from_slice` debug asserts on out-of-bounds access. @@ -670,7 +686,7 @@ mod tests { // values[2] might be any value since it's null. // Check validity - first two should be valid, third should be null. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert!( array .validity() diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index d396cf186a0..016d42508ee 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -211,8 +211,8 @@ impl ArrayBuilder for StructBuilder { #[cfg(test)] mod tests { use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; use crate::assert_arrays_eq; @@ -258,7 +258,7 @@ mod tests { assert_eq!(struct_.dtype(), &dtype); assert_eq!( struct_ - .valid_count(&mut LEGACY_SESSION.create_execution_ctx()) + .valid_count(&mut array_session().create_execution_ctx()) .unwrap(), 1 ); @@ -266,6 +266,7 @@ mod tests { #[test] fn test_append_scalar() { + let mut ctx = array_session().create_execution_ctx(); use crate::scalar::Scalar; let dtype = DType::Struct( @@ -328,7 +329,7 @@ mod tests { Validity::from_iter([true, true, false]), ) .unwrap(); - assert_arrays_eq!(&array, &expected); + assert_arrays_eq!(&array, &expected, &mut ctx); // Test wrong dtype error. let struct_fields = match &dtype { diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index b725a4636b5..02babab4a7d 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -8,8 +8,8 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::builders::ArrayBuilder; use crate::builders::builder_with_capacity; use crate::dtype::DType; @@ -95,10 +95,10 @@ fn test_append_zeros_matches_default_value(#[case] dtype: DType) { // Compare each element. for i in 0..num_elements { let scalar_zeros = array_zeros - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(); let scalar_manual = array_manual - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(); assert_eq!( @@ -197,7 +197,7 @@ fn test_append_defaults_behavior(#[case] dtype: DType, #[case] should_be_null: b for i in 0..3 { let scalar = array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(); if should_be_null { assert!(scalar.is_null(), "Element at index {} should be null", i); @@ -255,10 +255,10 @@ where // Compare each element. for i in 0..array_direct.len() { let scalar_direct = array_direct - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(); let scalar_indirect = array_indirect - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(); assert_eq!( @@ -549,7 +549,7 @@ fn test_append_scalar_comprehensive(#[case] dtype: DType) { // Verify each scalar matches. for (i, expected_scalar) in scalars.iter().enumerate() { let actual_scalar = array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(); assert_scalars_equal(&actual_scalar, expected_scalar, &dtype, i); } @@ -557,7 +557,7 @@ fn test_append_scalar_comprehensive(#[case] dtype: DType) { // If nullable, verify the last element is null. if dtype.is_nullable() { let null_scalar = array - .execute_scalar(num_elements, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(num_elements, &mut array_session().create_execution_ctx()) .unwrap(); assert!( null_scalar.is_null(), @@ -707,31 +707,31 @@ fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { // Check the pattern. assert!( !array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert!( array - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert!( !array - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert!( array - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(3, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); assert!( !array - .execute_scalar(4, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(4, &mut array_session().create_execution_ctx()) .unwrap() .is_null() ); @@ -739,7 +739,7 @@ fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { // Verify non-null values match. assert_scalars_equal( &array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), &test_scalars[0], &dtype, @@ -747,7 +747,7 @@ fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { ); assert_scalars_equal( &array - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap(), &test_scalars[1], &dtype, @@ -755,7 +755,7 @@ fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { ); assert_scalars_equal( &array - .execute_scalar(4, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(4, &mut array_session().create_execution_ctx()) .unwrap(), &test_scalars[2], &dtype, @@ -812,7 +812,7 @@ fn test_append_scalar_repeated_same_instance() { // All values should be 42. for i in 0..5 { let actual = array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(); assert_eq!( actual.as_primitive().typed_value::(), @@ -855,7 +855,7 @@ fn test_set_validity_overrides_validity( builder.set_validity(mask); let validity = builder.finish().validity()?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); for (i, &valid) in expected.iter().enumerate() { assert_eq!( validity.execute_is_valid(i, &mut ctx)?, @@ -877,7 +877,7 @@ fn test_set_validity_noop_when_non_nullable() -> VortexResult<()> { builder.set_validity(Mask::new_false(4)); let validity = builder.finish().validity()?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); for i in 0..4 { assert!( validity.execute_is_valid(i, &mut ctx)?, diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 392bef59997..3aabfd87c31 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -6,6 +6,7 @@ use std::ops::Range; use std::sync::Arc; use itertools::Itertools; +use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; @@ -19,10 +20,12 @@ use vortex_utils::aliases::hash_map::Entry; use vortex_utils::aliases::hash_map::HashMap; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::VarBinViewArray; +use crate::arrays::varbinview::VarBinViewArrayExt; use crate::arrays::varbinview::build_views::BinaryView; use crate::arrays::varbinview::compact::BufferUtilization; use crate::builders::ArrayBuilder; @@ -39,7 +42,7 @@ pub struct VarBinViewBuilder { views_builder: BufferMut, nulls: LazyBitBufferBuilder, completed: CompletedBuffers, - in_progress: ByteBufferMut, + in_progress: Option, growth_strategy: BufferGrowthStrategy, compaction_threshold: f64, } @@ -81,10 +84,14 @@ impl VarBinViewBuilder { "VarBinViewBuilder DType must be Utf8 or Binary." ); Self { - views_builder: BufferMut::::with_capacity(capacity), + views_builder: BufferMut::with_capacity_preferred_aligned( + capacity, + Alignment::of::(), + None, + ), nulls: LazyBitBufferBuilder::new(capacity), completed, - in_progress: ByteBufferMut::empty(), + in_progress: None, dtype, growth_strategy, compaction_threshold, @@ -127,15 +134,14 @@ impl VarBinViewBuilder { } fn flush_in_progress(&mut self) { - if self.in_progress.is_empty() { + let Some(block) = self.in_progress.take() else { return; - } - let block = std::mem::take(&mut self.in_progress).freeze(); + }; assert!(block.len() < u32::MAX as usize, "Block too large"); let initial_len = self.completed.len(); - self.completed.push(block); + self.completed.push(block.freeze()); assert_eq!( self.completed.len(), initial_len + 1, @@ -143,23 +149,41 @@ 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( + to_reserve, + Alignment::of::(), + None, + )); + } + /// append a non inlined value to self.in_progress. fn append_value_to_buffer(&mut self, value: &[u8]) -> (u32, u32) { assert!( value.len() > BinaryView::MAX_INLINED_SIZE, "must inline small strings" ); - let required_cap = self.in_progress.len() + value.len(); - if self.in_progress.capacity() < required_cap { - self.flush_in_progress(); - let next_buffer_size = self.growth_strategy.next_size() as usize; - let to_reserve = next_buffer_size.max(value.len()); - self.in_progress.reserve(to_reserve); - } + + if let Some(in_progress) = &mut self.in_progress { + let required_cap = in_progress.len() + value.len(); + if in_progress.capacity() < required_cap { + self.flush_in_progress(); + self.init_in_progress(value.len()); + } + } else { + self.init_in_progress(value.len()) + }; + + let in_progress = self + .in_progress + .as_mut() + .vortex_expect("in_progress just set"); let buffer_idx = self.completed.len(); - let offset = u32::try_from(self.in_progress.len()).vortex_expect("too many buffers"); - self.in_progress.extend_from_slice(value); + let offset = u32::try_from(in_progress.len()).vortex_expect("too many buffers"); + in_progress.extend_from_slice(value); (buffer_idx, offset) } @@ -171,7 +195,7 @@ impl VarBinViewBuilder { /// 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_empty() + self.in_progress.is_some() } /// Pushes buffers and pre-adjusted views into the builder. @@ -239,6 +263,60 @@ impl VarBinViewBuilder { 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, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + self.flush_in_progress(); + + let mask = array.varbinview_validity().execute_mask(array.len(), ctx)?; + + self.push_only_validity_mask(&mask); + + let view_adjustment = + self.completed + .extend_from_compaction(BuffersWithOffsets::from_array( + array, + self.compaction_threshold, + )); + + match view_adjustment { + ViewAdjustment::Precomputed(adjustment) => self.views_builder.extend_trusted( + array + .views() + .iter() + .map(|view| adjustment.adjust_view(view)), + ), + ViewAdjustment::Rewriting(adjustment) => match mask { + Mask::AllTrue(_) => { + for (idx, &view) in array.views().iter().enumerate() { + let new_view = self.push_view(view, &adjustment, array, idx); + self.views_builder.push(new_view); + } + } + Mask::AllFalse(_) => { + self.views_builder + .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); + } + } + }, + } + + Ok(()) + } } impl ArrayBuilder for VarBinViewBuilder { @@ -297,55 +375,8 @@ impl ArrayBuilder for VarBinViewBuilder { unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { #[expect(deprecated)] let array = array.to_varbinview(); - self.flush_in_progress(); - - let mask = array - .validity() - .vortex_expect("validity_mask") - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) - .vortex_expect("Failed to compute validity mask"); - - self.push_only_validity_mask(&mask); - - let view_adjustment = - self.completed - .extend_from_compaction(BuffersWithOffsets::from_array( - &array, - self.compaction_threshold, - )); - - match view_adjustment { - ViewAdjustment::Precomputed(adjustment) => self.views_builder.extend_trusted( - array - .views() - .iter() - .map(|view| adjustment.adjust_view(view)), - ), - ViewAdjustment::Rewriting(adjustment) => match mask { - Mask::AllTrue(_) => { - for (idx, &view) in array.views().iter().enumerate() { - let new_view = self.push_view(view, &adjustment, &array, idx); - self.views_builder.push(new_view); - } - } - Mask::AllFalse(_) => { - self.views_builder - .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); - } - } - }, - } + self.append_varbinview_array(&array, &mut LEGACY_SESSION.create_execution_ctx()) + .vortex_expect("Failed to append varbinview array"); } fn reserve_exact(&mut self, additional: usize) { @@ -842,8 +873,8 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::VarBinViewBuilder; @@ -853,6 +884,7 @@ mod tests { #[test] fn test_utf8_builder() { + let mut ctx = array_session().create_execution_ctx(); let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); builder.append_value("Hello"); @@ -875,11 +907,12 @@ mod tests { Some(""), Some("test"), ]); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] 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); @@ -903,7 +936,7 @@ mod tests { None, Some("Hello3"), ]); - assert_arrays_eq!(actual.into_array(), expected.into_array()); + assert_arrays_eq!(actual.into_array(), expected.into_array(), &mut ctx); } #[test] @@ -920,7 +953,7 @@ mod tests { let mut builder = VarBinViewBuilder::with_buffer_deduplication(DType::Utf8(Nullability::Nullable), 10); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); array.append_to_builder(&mut builder, &mut ctx)?; assert_eq!(builder.completed_block_count(), 1); @@ -955,6 +988,7 @@ mod tests { #[test] fn test_append_scalar() { + let mut ctx = array_session().create_execution_ctx(); use crate::scalar::Scalar; // Test with Utf8 builder. @@ -976,7 +1010,7 @@ mod tests { let array = utf8_builder.finish(); let expected = >::from_iter([Some("hello"), Some("world"), None]); - assert_arrays_eq!(&array, &expected); + assert_arrays_eq!(&array, &expected, &mut ctx); // Test with Binary builder. let mut binary_builder = @@ -991,7 +1025,7 @@ mod tests { let binary_array = binary_builder.finish(); let expected = >::from_iter([Some(vec![1u8, 2, 3]), None]); - assert_arrays_eq!(&binary_array, &expected); + assert_arrays_eq!(&binary_array, &expected, &mut ctx); // Test wrong dtype error. let mut builder = @@ -1047,7 +1081,7 @@ mod tests { // Verify the value was stored correctly let retrieved = array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .as_binary() .value() diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 3c07650b2ef..02029fc8e3e 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -69,7 +69,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::Binary`]. [`DType::Union`] does not yet have a public canonical array. /// /// # Laziness /// @@ -128,7 +128,9 @@ pub enum Canonical { List(ListViewArray), FixedSizeList(FixedSizeListArray), Struct(StructArray), + /// 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. Variant(VariantArray), } @@ -709,7 +711,9 @@ fn recursively_canonicalize_slots( .transpose() }) .collect::>()?; - array.clone().with_slots(slots) + // SAFETY: recursive canonicalization rewrites child slots to equivalent canonical + // representations, preserving the parent array's logical values and statistics. + unsafe { array.clone().with_slots(slots) } } impl Executable for RecursiveCanonical { fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { @@ -1155,11 +1159,9 @@ mod test { use crate::canonical::StructArray; use crate::dtype::Nullability; use crate::scalar::Scalar; - use crate::session::ArraySession; /// A shared session for these canonical tests, used to create execution contexts. - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); fn variant_core_storage(len: usize) -> ArrayRef { ConstantArray::new( diff --git a/vortex-array/src/compute/conformance/binary_numeric.rs b/vortex-array/src/compute/conformance/binary_numeric.rs index ce3eba940df..1fe409bf648 100644 --- a/vortex-array/src/compute/conformance/binary_numeric.rs +++ b/vortex-array/src/compute/conformance/binary_numeric.rs @@ -21,19 +21,22 @@ //! - Multiplication (`*`) //! - Division (`/`) +use std::fmt::Debug; + use itertools::Itertools; +use num_traits::Bounded; +use num_traits::Float; use num_traits::Num; +use num_traits::Signed; use vortex_error::VortexExpect; use vortex_error::vortex_err; use vortex_error::vortex_panic; use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::RecursiveCanonical; -#[expect(deprecated)] -use crate::ToCanonical as _; -use crate::VortexSessionExecute; use crate::arrays::ConstantArray; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; @@ -43,12 +46,12 @@ use crate::scalar::NumericOperator; use crate::scalar::PrimitiveScalar; use crate::scalar::Scalar; -fn to_vec_of_scalar(array: &ArrayRef) -> Vec { +fn to_vec_of_scalar(array: &ArrayRef, ctx: &mut ExecutionCtx) -> Vec { // Not fast, but obviously correct (0..array.len()) .map(|index| { array - .execute_scalar(index, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(index, ctx) .vortex_expect("scalar_at should succeed in conformance test") }) .collect_vec() @@ -78,24 +81,31 @@ fn to_vec_of_scalar(array: &ArrayRef) -> Vec { /// Panics if: /// - The array cannot be converted to primitive form /// - Results don't match expected values (for operations that don't overflow) -fn test_binary_numeric_conformance(array: ArrayRef) -where +fn test_binary_numeric_conformance( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) where Scalar: From, { // First test with the standard scalar value of 1 - test_standard_binary_numeric::(array.clone()); + test_standard_binary_numeric::(array, ctx); // Then test edge cases - test_binary_numeric_edge_cases(array); + test_binary_numeric_edge_cases(array, ctx); } -fn test_standard_binary_numeric(array: ArrayRef) -where +fn test_standard_binary_numeric( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) where Scalar: From, { - #[expect(deprecated)] - let canonicalized_array = array.to_primitive(); - let original_values = to_vec_of_scalar(&canonicalized_array.into_array()); + 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 one = T::from(1) .ok_or_else(|| vortex_err!("could not convert 1 into array native type")) @@ -119,7 +129,7 @@ where let result = array .binary(rhs_const.clone(), op.into()) .vortex_expect("apply shouldn't fail") - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(ctx) .map(|c| c.0.into_array()); // Skip this operator if the entire operation fails @@ -128,7 +138,7 @@ where continue; }; - let actual_values = to_vec_of_scalar(&result); + let actual_values = to_vec_of_scalar(&result, ctx); // Check each element for overflow/underflow let expected_results: Vec> = original_values @@ -157,7 +167,7 @@ where // Test scalar operator array (e.g., 1 + array) let result = rhs_const.binary(array.clone(), op.into()).and_then(|a| { - a.execute::(&mut LEGACY_SESSION.create_execution_ctx()) + a.execute::(ctx) .map(|c| c.0.into_array()) }); @@ -166,7 +176,7 @@ where continue; }; - let actual_values = to_vec_of_scalar(&result); + let actual_values = to_vec_of_scalar(&result, ctx); // Check each element for overflow/underflow let expected_results: Vec> = original_values @@ -212,23 +222,23 @@ where /// test_binary_numeric_array(array.into_array()); /// } /// ``` -pub fn test_binary_numeric_array(array: ArrayRef) { +pub fn test_binary_numeric_array(array: &ArrayRef, ctx: &mut ExecutionCtx) { match array.dtype() { DType::Primitive(ptype, _) => match ptype { - PType::I8 => test_binary_numeric_conformance::(array), - PType::I16 => test_binary_numeric_conformance::(array), - PType::I32 => test_binary_numeric_conformance::(array), - PType::I64 => test_binary_numeric_conformance::(array), - PType::U8 => test_binary_numeric_conformance::(array), - PType::U16 => test_binary_numeric_conformance::(array), - PType::U32 => test_binary_numeric_conformance::(array), - PType::U64 => test_binary_numeric_conformance::(array), + PType::I8 => test_binary_numeric_conformance::(array, ctx), + PType::I16 => test_binary_numeric_conformance::(array, ctx), + PType::I32 => test_binary_numeric_conformance::(array, ctx), + PType::I64 => test_binary_numeric_conformance::(array, ctx), + PType::U8 => test_binary_numeric_conformance::(array, ctx), + PType::U16 => test_binary_numeric_conformance::(array, ctx), + PType::U32 => test_binary_numeric_conformance::(array, ctx), + PType::U64 => test_binary_numeric_conformance::(array, ctx), PType::F16 => { // F16 not supported in num-traits, skip eprintln!("Skipping f16 binary numeric tests (not supported)"); } - PType::F32 => test_binary_numeric_conformance::(array), - PType::F64 => test_binary_numeric_conformance::(array), + PType::F32 => test_binary_numeric_conformance::(array, ctx), + PType::F64 => test_binary_numeric_conformance::(array, ctx), }, dtype => vortex_panic!( "Binary numeric tests are only supported for primitive numeric types, got {dtype}", @@ -243,22 +253,22 @@ pub fn test_binary_numeric_array(array: ArrayRef) { /// - Negative one (tests signed arithmetic) /// - Maximum value (tests overflow behavior) /// - Minimum value (tests underflow behavior) -fn test_binary_numeric_edge_cases(array: ArrayRef) { +fn test_binary_numeric_edge_cases(array: &ArrayRef, ctx: &mut ExecutionCtx) { match array.dtype() { DType::Primitive(ptype, _) => match ptype { - PType::I8 => test_binary_numeric_edge_cases_signed::(array), - PType::I16 => test_binary_numeric_edge_cases_signed::(array), - PType::I32 => test_binary_numeric_edge_cases_signed::(array), - PType::I64 => test_binary_numeric_edge_cases_signed::(array), - PType::U8 => test_binary_numeric_edge_cases_unsigned::(array), - PType::U16 => test_binary_numeric_edge_cases_unsigned::(array), - PType::U32 => test_binary_numeric_edge_cases_unsigned::(array), - PType::U64 => test_binary_numeric_edge_cases_unsigned::(array), + PType::I8 => test_binary_numeric_edge_cases_signed::(array, ctx), + PType::I16 => test_binary_numeric_edge_cases_signed::(array, ctx), + PType::I32 => test_binary_numeric_edge_cases_signed::(array, ctx), + PType::I64 => test_binary_numeric_edge_cases_signed::(array, ctx), + PType::U8 => test_binary_numeric_edge_cases_unsigned::(array, ctx), + PType::U16 => test_binary_numeric_edge_cases_unsigned::(array, ctx), + PType::U32 => test_binary_numeric_edge_cases_unsigned::(array, ctx), + PType::U64 => test_binary_numeric_edge_cases_unsigned::(array, ctx), PType::F16 => { eprintln!("Skipping f16 edge case tests (not supported)"); } - PType::F32 => test_binary_numeric_edge_cases_float::(array), - PType::F64 => test_binary_numeric_edge_cases_float::(array), + PType::F32 => test_binary_numeric_edge_cases_float::(array, ctx), + PType::F64 => test_binary_numeric_edge_cases_float::(array, ctx), }, dtype => vortex_panic!( "Binary numeric edge case tests are only supported for primitive numeric types, got {dtype}" @@ -266,73 +276,76 @@ fn test_binary_numeric_edge_cases(array: ArrayRef) { } } -fn test_binary_numeric_edge_cases_signed(array: ArrayRef) +fn test_binary_numeric_edge_cases_signed(array: &ArrayRef, ctx: &mut ExecutionCtx) where - T: NativePType + Num + Copy + std::fmt::Debug + num_traits::Bounded + num_traits::Signed, + T: NativePType + Num + Copy + Debug + Bounded + Signed, Scalar: From, { // Test with zero - test_binary_numeric_with_scalar(array.clone(), T::zero()); + test_binary_numeric_with_scalar(array, T::zero(), ctx); // Test with -1 - test_binary_numeric_with_scalar(array.clone(), -T::one()); + test_binary_numeric_with_scalar(array, -T::one(), ctx); // Test with max value - test_binary_numeric_with_scalar(array.clone(), T::max_value()); + test_binary_numeric_with_scalar(array, T::max_value(), ctx); // Test with min value - test_binary_numeric_with_scalar(array, T::min_value()); + test_binary_numeric_with_scalar(array, T::min_value(), ctx); } -fn test_binary_numeric_edge_cases_unsigned(array: ArrayRef) +fn test_binary_numeric_edge_cases_unsigned(array: &ArrayRef, ctx: &mut ExecutionCtx) where - T: NativePType + Num + Copy + std::fmt::Debug + num_traits::Bounded, + T: NativePType + Num + Copy + Debug + Bounded, Scalar: From, { // Test with zero - test_binary_numeric_with_scalar(array.clone(), T::zero()); + test_binary_numeric_with_scalar(array, T::zero(), ctx); // Test with max value - test_binary_numeric_with_scalar(array, T::max_value()); + test_binary_numeric_with_scalar(array, T::max_value(), ctx); } -fn test_binary_numeric_edge_cases_float(array: ArrayRef) +fn test_binary_numeric_edge_cases_float(array: &ArrayRef, ctx: &mut ExecutionCtx) where - T: NativePType + Num + Copy + std::fmt::Debug + num_traits::Float, + T: NativePType + Num + Copy + Debug + Float, Scalar: From, { // Test with zero - test_binary_numeric_with_scalar(array.clone(), T::zero()); + test_binary_numeric_with_scalar(array, T::zero(), ctx); // Test with -1 - test_binary_numeric_with_scalar(array.clone(), -T::one()); + test_binary_numeric_with_scalar(array, -T::one(), ctx); // Test with max value - test_binary_numeric_with_scalar(array.clone(), T::max_value()); + test_binary_numeric_with_scalar(array, T::max_value(), ctx); // Test with min value - test_binary_numeric_with_scalar(array.clone(), T::min_value()); + test_binary_numeric_with_scalar(array, T::min_value(), ctx); // Test with small positive value - test_binary_numeric_with_scalar(array.clone(), T::epsilon()); + test_binary_numeric_with_scalar(array, T::epsilon(), ctx); // Test with min positive value (subnormal) - test_binary_numeric_with_scalar(array.clone(), T::min_positive_value()); + test_binary_numeric_with_scalar(array, T::min_positive_value(), ctx); // Test with special float values (NaN, Infinity) - test_binary_numeric_with_scalar(array.clone(), T::nan()); - test_binary_numeric_with_scalar(array.clone(), T::infinity()); - test_binary_numeric_with_scalar(array, T::neg_infinity()); + test_binary_numeric_with_scalar(array, T::nan(), ctx); + test_binary_numeric_with_scalar(array, T::infinity(), ctx); + test_binary_numeric_with_scalar(array, T::neg_infinity(), ctx); } -fn test_binary_numeric_with_scalar(array: ArrayRef, scalar_value: T) +fn test_binary_numeric_with_scalar(array: &ArrayRef, scalar_value: T, ctx: &mut ExecutionCtx) where - T: NativePType + Num + Copy + std::fmt::Debug, + T: NativePType + Num + Copy + Debug, Scalar: From, { - #[expect(deprecated)] - let canonicalized_array = array.to_primitive(); - let original_values = to_vec_of_scalar(&canonicalized_array.into_array()); + 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::from(scalar_value) .cast(array.dtype()) @@ -363,7 +376,7 @@ where let result = array .binary(rhs_const, op.into()) .vortex_expect("apply failed") - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(ctx) .map(|x| x.0.into_array()); // Skip if the entire operation fails @@ -373,7 +386,7 @@ where } let result = result.vortex_expect("operation should succeed in conformance test"); - let actual_values = to_vec_of_scalar(&result); + let actual_values = to_vec_of_scalar(&result, ctx); // Check each element for overflow/underflow let expected_results: Vec> = original_values diff --git a/vortex-array/src/compute/conformance/cast.rs b/vortex-array/src/compute/conformance/cast.rs index be5142d8f53..c823b2a778a 100644 --- a/vortex-array/src/compute/conformance/cast.rs +++ b/vortex-array/src/compute/conformance/cast.rs @@ -7,11 +7,12 @@ use vortex_error::vortex_panic; use crate::ArrayRef; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::RecursiveCanonical; use crate::VortexSessionExecute; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::MinMaxResult; use crate::aggregate_fn::fns::min_max::min_max; +use crate::array_session; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; @@ -22,7 +23,7 @@ use crate::scalar::Scalar; fn cast_and_execute(array: &ArrayRef, dtype: DType) -> VortexResult { Ok(array .cast(dtype)? - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut array_session().create_execution_ctx())? .0 .into_array()) } @@ -74,10 +75,10 @@ fn test_cast_identity(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -109,7 +110,7 @@ fn test_cast_from_null(array: &ArrayRef) { for i in 0..array.len().min(10) { assert!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .is_null() ); @@ -129,7 +130,7 @@ fn test_cast_from_null(array: &ArrayRef) { fn test_cast_to_non_nullable(array: &ArrayRef) { if array - .invalid_count(&mut LEGACY_SESSION.create_execution_ctx()) + .invalid_count(&mut array_session().create_execution_ctx()) .vortex_expect("invalid_count should succeed in conformance test") == 0 { @@ -141,10 +142,10 @@ fn test_cast_to_non_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), non_nullable - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -157,10 +158,10 @@ fn test_cast_to_non_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), back_to_nullable - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -190,10 +191,10 @@ fn test_cast_to_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), nullable - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -206,9 +207,9 @@ fn test_cast_to_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), - back.execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + back.execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -247,9 +248,9 @@ fn fits(value: &Scalar, ptype: PType) -> bool { } fn test_cast_to_primitive(array: &ArrayRef, target_ptype: PType, test_round_trip: bool) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let maybe_min_max = - min_max(array, &mut ctx).vortex_expect("cast should succeed in conformance test"); + let mut ctx = array_session().create_execution_ctx(); + let maybe_min_max = min_max(array, &mut ctx, NumericalAggregateOpts::default()) + .vortex_expect("cast should succeed in conformance test"); if let Some(MinMaxResult { min, max }) = maybe_min_max && (!fits(&min, target_ptype) || !fits(&max, target_ptype)) @@ -288,20 +289,20 @@ fn test_cast_to_primitive(array: &ArrayRef, target_ptype: PType, test_round_trip array .validity() .vortex_expect("validity_mask should succeed in conformance test") - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(array.len(), &mut array_session().create_execution_ctx()) .vortex_expect("Failed to compute validity mask"), casted .validity() .vortex_expect("validity_mask should succeed in conformance test") - .execute_mask(casted.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(casted.len(), &mut array_session().create_execution_ctx()) .vortex_expect("Failed to compute validity mask") ); for i in 0..array.len().min(10) { let original = array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let casted = casted - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( original diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index 8ae7101c6db..73f98cce4a7 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -27,10 +27,11 @@ use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; @@ -49,7 +50,7 @@ use crate::scalar_fn::fns::operators::Operator; /// - Creates indices array containing positions where mask is true /// - Applies take with these indices /// - Verifies both results are identical -fn test_filter_take_consistency(array: &ArrayRef) { +fn test_filter_take_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len == 0 { return; @@ -89,10 +90,10 @@ fn test_filter_take_consistency(array: &ArrayRef) { for i in 0..filtered.len() { let filtered_val = filtered - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let taken_val = taken - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( filtered_val, taken_val, @@ -117,7 +118,7 @@ fn test_filter_take_consistency(array: &ArrayRef) { /// # Why This Matters /// This test ensures that mask operations compose correctly, which is critical for /// complex query operations that may apply multiple filters. -fn test_double_mask_consistency(array: &ArrayRef) { +fn test_double_mask_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len == 0 { return; @@ -163,10 +164,10 @@ fn test_double_mask_consistency(array: &ArrayRef) { for i in 0..double_masked.len() { let double_val = double_masked - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let direct_val = directly_masked - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( double_val, direct_val, @@ -190,7 +191,7 @@ fn test_double_mask_consistency(array: &ArrayRef) { /// # Why This Matters /// This is an identity operation that should be optimized in implementations /// to avoid unnecessary copying. -fn test_filter_identity(array: &ArrayRef) { +fn test_filter_identity(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len == 0 { return; @@ -213,10 +214,10 @@ fn test_filter_identity(array: &ArrayRef) { for i in 0..len { let original_val = array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let filtered_val = filtered - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( filtered_val, original_val, @@ -239,7 +240,7 @@ fn test_filter_identity(array: &ArrayRef) { /// # Why This Matters /// Masking always produces a nullable array, even when no values are actually masked. /// This test ensures the type system handles this correctly. -fn test_mask_identity(array: &ArrayRef) { +fn test_mask_identity(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len == 0 { return; @@ -269,10 +270,10 @@ fn test_mask_identity(array: &ArrayRef) { for i in 0..len { let original_val = array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let masked_val = masked - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let expected_val = original_val.clone().into_nullable(); assert_eq!( @@ -297,7 +298,7 @@ fn test_mask_identity(array: &ArrayRef) { /// # Why This Matters /// When a filter mask represents a contiguous range, it should be equivalent to /// a slice operation. Some implementations may optimize this case. -fn test_slice_filter_consistency(array: &ArrayRef) { +fn test_slice_filter_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len < 4 { return; // Need at least 4 elements for meaningful test @@ -328,10 +329,10 @@ fn test_slice_filter_consistency(array: &ArrayRef) { for i in 0..filtered.len() { let filtered_val = filtered - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let sliced_val = sliced - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( filtered_val, sliced_val, @@ -354,7 +355,7 @@ fn test_slice_filter_consistency(array: &ArrayRef) { /// /// # Why This Matters /// Sequential takes are a common pattern that can be optimized to slice operations. -fn test_take_slice_consistency(array: &ArrayRef) { +fn test_take_slice_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len < 3 { return; // Need at least 3 elements @@ -383,10 +384,10 @@ fn test_take_slice_consistency(array: &ArrayRef) { for i in 0..taken.len() { let taken_val = taken - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let sliced_val = sliced - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( taken_val, sliced_val, @@ -397,7 +398,7 @@ fn test_take_slice_consistency(array: &ArrayRef) { } /// Tests that filter preserves relative ordering -fn test_filter_preserves_order(array: &ArrayRef) { +fn test_filter_preserves_order(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len < 4 { return; @@ -416,33 +417,33 @@ fn test_filter_preserves_order(array: &ArrayRef) { if len >= 4 { assert_eq!( filtered - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); assert_eq!( filtered - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); assert_eq!( filtered - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(3, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } } /// Tests that take with repeated indices works correctly -fn test_take_repeated_indices(array: &ArrayRef) { +fn test_take_repeated_indices(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len == 0 { return; @@ -458,17 +459,17 @@ fn test_take_repeated_indices(array: &ArrayRef) { for i in 0..3 { assert_eq!( taken - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } } /// Tests mask and filter interaction with nulls -fn test_mask_filter_null_consistency(array: &ArrayRef) { +fn test_mask_filter_null_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len < 3 { return; @@ -498,10 +499,10 @@ fn test_mask_filter_null_consistency(array: &ArrayRef) { for i in 0..filtered.len() { assert_eq!( filtered - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), direct_filtered - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -537,7 +538,7 @@ fn test_empty_operations_consistency(array: &ArrayRef) { } /// Tests that take preserves array properties -fn test_take_preserves_properties(array: &ArrayRef) { +fn test_take_preserves_properties(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len == 0 { return; @@ -555,10 +556,10 @@ fn test_take_preserves_properties(array: &ArrayRef) { for i in 0..len { assert_eq!( taken - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -581,7 +582,7 @@ fn test_take_preserves_properties(array: &ArrayRef) { /// # Why This Matters /// Nullable indices are a powerful feature that allows introducing nulls during /// a take operation, which is useful for outer joins and similar operations. -fn test_nullable_indices_consistency(array: &ArrayRef) { +fn test_nullable_indices_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len < 3 { return; // Need at least 3 elements to test indices 0 and 2 @@ -610,11 +611,11 @@ fn test_nullable_indices_consistency(array: &ArrayRef) { // Check first element (from index 0) let expected_0 = array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable(); let actual_0 = taken - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( actual_0, expected_0, @@ -624,7 +625,7 @@ fn test_nullable_indices_consistency(array: &ArrayRef) { // Check second element (should be null) let actual_1 = taken - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert!( actual_1.is_null(), @@ -633,11 +634,11 @@ fn test_nullable_indices_consistency(array: &ArrayRef) { // Check third element (from index 2) let expected_2 = array - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, ctx) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable(); let actual_2 = taken - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( actual_2, expected_2, @@ -647,7 +648,7 @@ fn test_nullable_indices_consistency(array: &ArrayRef) { } /// Tests large array consistency -fn test_large_array_consistency(array: &ArrayRef) { +fn test_large_array_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len < 1000 { return; @@ -672,10 +673,10 @@ fn test_large_array_consistency(array: &ArrayRef) { for i in 0..taken.len() { assert_eq!( taken - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), filtered - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -697,7 +698,7 @@ fn test_large_array_consistency(array: &ArrayRef) { /// Comparison operations must maintain logical consistency across encodings. /// This test catches bugs where an encoding might implement one comparison /// correctly but fail on its logical inverse. -fn test_comparison_inverse_consistency(array: &ArrayRef) { +fn test_comparison_inverse_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len == 0 { return; @@ -718,19 +719,15 @@ fn test_comparison_inverse_consistency(array: &ArrayRef) { return; } else { array - .execute_scalar(len / 2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(len / 2, ctx) .vortex_expect("scalar_at should succeed in conformance test") }; // Test Eq vs NotEq - let const_array = crate::arrays::ConstantArray::new(test_scalar, len); + let const_array = ConstantArray::new(test_scalar, len).into_array(); if let (Ok(eq_result), Ok(neq_result)) = ( - array - .clone() - .binary(const_array.clone().into_array(), Operator::Eq), - array - .clone() - .binary(const_array.clone().into_array(), Operator::NotEq), + array.binary(const_array.clone(), Operator::Eq), + array.binary(const_array.clone(), Operator::NotEq), ) { let inverted_eq = eq_result .not() @@ -744,10 +741,10 @@ fn test_comparison_inverse_consistency(array: &ArrayRef) { for i in 0..inverted_eq.len() { let inv_val = inverted_eq - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let neq_val = neq_result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( inv_val, neq_val, @@ -759,12 +756,8 @@ fn test_comparison_inverse_consistency(array: &ArrayRef) { // Test Gt vs Lte if let (Ok(gt_result), Ok(lte_result)) = ( - array - .clone() - .binary(const_array.clone().into_array(), Operator::Gt), - array - .clone() - .binary(const_array.clone().into_array(), Operator::Lte), + array.binary(const_array.clone(), Operator::Gt), + array.binary(const_array.clone(), Operator::Lte), ) { let inverted_gt = gt_result .not() @@ -772,10 +765,10 @@ fn test_comparison_inverse_consistency(array: &ArrayRef) { for i in 0..inverted_gt.len() { let inv_val = inverted_gt - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let lte_val = lte_result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( inv_val, lte_val, @@ -787,12 +780,8 @@ fn test_comparison_inverse_consistency(array: &ArrayRef) { // Test Lt vs Gte if let (Ok(lt_result), Ok(gte_result)) = ( - array - .clone() - .binary(const_array.clone().into_array(), Operator::Lt), - array - .clone() - .binary(const_array.into_array(), Operator::Gte), + array.binary(const_array.clone(), Operator::Lt), + array.binary(const_array, Operator::Gte), ) { let inverted_lt = lt_result .not() @@ -800,10 +789,10 @@ fn test_comparison_inverse_consistency(array: &ArrayRef) { for i in 0..inverted_lt.len() { let inv_val = inverted_lt - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let gte_val = gte_result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( inv_val, gte_val, @@ -829,7 +818,7 @@ fn test_comparison_inverse_consistency(array: &ArrayRef) { /// # Why This Matters /// Ensures that comparison operations maintain mathematical ordering properties /// regardless of operand order. -fn test_comparison_symmetry_consistency(array: &ArrayRef) { +fn test_comparison_symmetry_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len == 0 { return; @@ -850,22 +839,17 @@ fn test_comparison_symmetry_consistency(array: &ArrayRef) { return; } else { array - .execute_scalar(len / 2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(len / 2, ctx) .vortex_expect("scalar_at should succeed in conformance test") }; // Create a constant array with the test scalar for reverse comparison - let const_array = crate::arrays::ConstantArray::new(test_scalar, len); + let const_array = ConstantArray::new(test_scalar, len).into_array(); // Test Gt vs Lt symmetry if let (Ok(arr_gt_scalar), Ok(scalar_lt_arr)) = ( - array - .clone() - .binary(const_array.clone().into_array(), Operator::Gt), - const_array - .clone() - .into_array() - .binary(array.clone(), Operator::Lt), + array.binary(const_array.clone(), Operator::Gt), + const_array.binary(array.clone(), Operator::Lt), ) { assert_eq!( arr_gt_scalar.len(), @@ -875,10 +859,10 @@ fn test_comparison_symmetry_consistency(array: &ArrayRef) { for i in 0..arr_gt_scalar.len() { let arr_gt = arr_gt_scalar - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let scalar_lt = scalar_lt_arr - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( arr_gt, scalar_lt, @@ -890,17 +874,15 @@ fn test_comparison_symmetry_consistency(array: &ArrayRef) { // Test Eq symmetry if let (Ok(arr_eq_scalar), Ok(scalar_eq_arr)) = ( - array - .clone() - .binary(const_array.clone().into_array(), Operator::Eq), - const_array.into_array().binary(array.clone(), Operator::Eq), + array.binary(const_array.clone(), Operator::Eq), + const_array.binary(array.clone(), Operator::Eq), ) { for i in 0..arr_eq_scalar.len() { let arr_eq = arr_eq_scalar - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let scalar_eq = scalar_eq_arr - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( arr_eq, scalar_eq, @@ -927,7 +909,7 @@ fn test_comparison_symmetry_consistency(array: &ArrayRef) { /// Boolean operations must maintain logical consistency across encodings. /// This test catches bugs where encodings might optimize boolean operations /// incorrectly, breaking fundamental logical properties. -fn test_boolean_demorgan_consistency(array: &ArrayRef) { +fn test_boolean_demorgan_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { if !matches!(array.dtype(), DType::Bool(_)) { return; } @@ -940,7 +922,7 @@ fn test_boolean_demorgan_consistency(array: &ArrayRef) { // Test first De Morgan's law: NOT(A AND B) = (NOT A) OR (NOT B) if let (Ok(a_and_b), Ok(not_a), Ok(not_b)) = ( - array.clone().binary(bool_mask.clone(), Operator::And), + array.binary(bool_mask.clone(), Operator::And), array.not(), bool_mask.not(), ) { @@ -959,10 +941,10 @@ fn test_boolean_demorgan_consistency(array: &ArrayRef) { for i in 0..not_a_and_b.len() { let left = not_a_and_b - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let right = not_a_or_not_b - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( left, right, @@ -974,7 +956,7 @@ fn test_boolean_demorgan_consistency(array: &ArrayRef) { // Test second De Morgan's law: NOT(A OR B) = (NOT A) AND (NOT B) if let (Ok(a_or_b), Ok(not_a), Ok(not_b)) = ( - array.clone().binary(bool_mask.clone(), Operator::Or), + array.binary(bool_mask.clone(), Operator::Or), array.not(), bool_mask.not(), ) { @@ -987,10 +969,10 @@ fn test_boolean_demorgan_consistency(array: &ArrayRef) { for i in 0..not_a_or_b.len() { let left = not_a_or_b - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let right = not_a_and_not_b - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( left, right, @@ -1016,14 +998,13 @@ fn test_boolean_demorgan_consistency(array: &ArrayRef) { /// # Why This Matters /// Aggregate operations on sliced arrays must produce correct results /// regardless of the underlying encoding's offset handling. -fn test_slice_aggregate_consistency(array: &ArrayRef) { +fn test_slice_aggregate_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { + use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::min_max; use crate::aggregate_fn::fns::nan_count::nan_count; use crate::aggregate_fn::fns::sum::sum; use crate::dtype::DType; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let len = array.len(); if len < 5 { return; // Need enough elements for meaningful slice @@ -1037,8 +1018,10 @@ fn test_slice_aggregate_consistency(array: &ArrayRef) { let sliced = array .slice(start..end) .vortex_expect("slice should succeed in conformance test"); - #[expect(deprecated)] - let canonical = array.to_canonical().vortex_expect("to_canonical failed"); + let canonical = array + .clone() + .execute::(ctx) + .vortex_expect("to_canonical failed"); let canonical_sliced = canonical .into_array() .slice(start..end) @@ -1046,10 +1029,10 @@ fn test_slice_aggregate_consistency(array: &ArrayRef) { // Test null count through invalid_count let sliced_invalid_count = sliced - .invalid_count(&mut ctx) + .invalid_count(ctx) .vortex_expect("invalid_count should succeed in conformance test"); let canonical_invalid_count = canonical_sliced - .invalid_count(&mut ctx) + .invalid_count(ctx) .vortex_expect("invalid_count should succeed in conformance test"); assert_eq!( sliced_invalid_count, canonical_invalid_count, @@ -1062,9 +1045,7 @@ fn test_slice_aggregate_consistency(array: &ArrayRef) { return; } - if let (Ok(slice_sum), Ok(canonical_sum)) = - (sum(&sliced, &mut ctx), sum(&canonical_sliced, &mut ctx)) - { + if let (Ok(slice_sum), Ok(canonical_sum)) = (sum(&sliced, ctx), sum(&canonical_sliced, ctx)) { // Compare sum scalars assert_eq!( slice_sum, canonical_sum, @@ -1075,8 +1056,8 @@ fn test_slice_aggregate_consistency(array: &ArrayRef) { // Test min_max if let (Ok(slice_minmax), Ok(canonical_minmax)) = ( - min_max(&sliced, &mut ctx), - min_max(&canonical_sliced, &mut ctx), + min_max(&sliced, ctx, NumericalAggregateOpts::default()), + min_max(&canonical_sliced, ctx, NumericalAggregateOpts::default()), ) { match (slice_minmax, canonical_minmax) { (Some(s_result), Some(c_result)) => { @@ -1100,10 +1081,8 @@ fn test_slice_aggregate_consistency(array: &ArrayRef) { // Test nan_count for floating point types if array.dtype().is_float() - && let (Ok(slice_nan_count), Ok(canonical_nan_count)) = ( - nan_count(&sliced, &mut ctx), - nan_count(&canonical_sliced, &mut ctx), - ) + && let (Ok(slice_nan_count), Ok(canonical_nan_count)) = + (nan_count(&sliced, ctx), nan_count(&canonical_sliced, ctx)) { assert_eq!( slice_nan_count, canonical_nan_count, @@ -1128,7 +1107,7 @@ fn test_slice_aggregate_consistency(array: &ArrayRef) { /// This test specifically catches bugs where encodings (like RunEndArray) fail to preserve /// offset information during cast operations. Such bugs can lead to incorrect data being /// returned after casting a sliced array. -fn test_cast_slice_consistency(array: &ArrayRef) { +fn test_cast_slice_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len < 5 { return; // Need at least 5 elements for meaningful slice @@ -1139,8 +1118,11 @@ fn test_cast_slice_consistency(array: &ArrayRef) { let end = 7.min(len - 2).max(start + 1); // Ensure we have at least 1 element // Get canonical form of the original array - #[expect(deprecated)] - let canonical = array.to_canonical().vortex_expect("to_canonical failed"); + let canonical = array + .clone() + .execute::(ctx) + .vortex_expect("to_canonical failed") + .into_array(); // Choose appropriate target dtype based on the array's type let target_dtypes = match array.dtype() { @@ -1265,10 +1247,10 @@ fn test_cast_slice_consistency(array: &ArrayRef) { .vortex_expect("slice should succeed in conformance test"); // Try to cast the sliced array (force execution via to_canonical) - let slice_then_cast = match sliced.cast(target_dtype.clone()).and_then(|a| { - #[expect(deprecated)] - a.to_canonical().map(|c| c.into_array()) - }) { + let slice_then_cast = match sliced + .cast(target_dtype.clone()) + .and_then(|a| a.execute::(ctx).map(|c| c.into_array())) + { Ok(result) => result, Err(_) => continue, // Skip if cast fails }; @@ -1285,14 +1267,12 @@ fn test_cast_slice_consistency(array: &ArrayRef) { // Compare each value against the canonical form for i in 0..slice_then_cast.len() { let slice_cast_val = slice_then_cast - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); // Get the corresponding value from the canonical array (adjusted for slice offset) let canonical_val = canonical - .clone() - .into_array() - .execute_scalar(start + i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(start + i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); // Cast the canonical scalar to the target dtype @@ -1317,10 +1297,10 @@ fn test_cast_slice_consistency(array: &ArrayRef) { } // Also test the other way: cast then slice - let casted = match array.clone().cast(target_dtype.clone()).and_then(|a| { - #[expect(deprecated)] - a.to_canonical().map(|c| c.into_array()) - }) { + let casted = match array + .cast(target_dtype.clone()) + .and_then(|a| a.execute::(ctx).map(|c| c.into_array())) + { Ok(result) => result, Err(_) => continue, // Skip if cast fails }; @@ -1337,10 +1317,10 @@ fn test_cast_slice_consistency(array: &ArrayRef) { for i in 0..slice_then_cast.len() { let slice_cast_val = slice_then_cast - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let cast_slice_val = cast_then_slice - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( slice_cast_val, cast_slice_val, @@ -1392,38 +1372,38 @@ fn test_cast_slice_consistency(array: &ArrayRef) { /// ## Large Arrays /// - **Performance**: Operations scale correctly to large arrays (1000+ elements) /// ```text -pub fn test_array_consistency(array: &ArrayRef) { +pub fn test_array_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { // Core operation consistency - test_filter_take_consistency(array); - test_double_mask_consistency(array); - test_slice_filter_consistency(array); - test_take_slice_consistency(array); - test_cast_slice_consistency(array); + test_filter_take_consistency(array, ctx); + test_double_mask_consistency(array, ctx); + test_slice_filter_consistency(array, ctx); + test_take_slice_consistency(array, ctx); + test_cast_slice_consistency(array, ctx); // Boolean operations - test_boolean_demorgan_consistency(array); + test_boolean_demorgan_consistency(array, ctx); // Comparison operations - test_comparison_inverse_consistency(array); - test_comparison_symmetry_consistency(array); + test_comparison_inverse_consistency(array, ctx); + test_comparison_symmetry_consistency(array, ctx); // Aggregate operations - test_slice_aggregate_consistency(array); + test_slice_aggregate_consistency(array, ctx); // Identity operations - test_filter_identity(array); - test_mask_identity(array); - test_take_preserves_properties(array); + test_filter_identity(array, ctx); + test_mask_identity(array, ctx); + test_take_preserves_properties(array, ctx); // Ordering and correctness - test_filter_preserves_order(array); - test_take_repeated_indices(array); + test_filter_preserves_order(array, ctx); + test_take_repeated_indices(array, ctx); // Null handling - test_mask_filter_null_consistency(array); - test_nullable_indices_consistency(array); + test_mask_filter_null_consistency(array, ctx); + test_nullable_indices_consistency(array, ctx); // Edge cases test_empty_operations_consistency(array); - test_large_array_consistency(array); + test_large_array_consistency(array, ctx); } diff --git a/vortex-array/src/compute/conformance/filter.rs b/vortex-array/src/compute/conformance/filter.rs index 7d5f4abcad8..01bfb3bb2fa 100644 --- a/vortex-array/src/compute/conformance/filter.rs +++ b/vortex-array/src/compute/conformance/filter.rs @@ -6,8 +6,8 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -62,12 +62,13 @@ pub fn create_runs_pattern(len: usize, run_length: usize) -> Vec { /// Tests that filtering with an all-true mask returns all elements unchanged fn test_all_filter(array: &ArrayRef) { + let mut ctx = array_session().create_execution_ctx(); let len = array.len(); let mask = Mask::new_true(len); let filtered = array .filter(mask) .vortex_expect("filter should succeed in conformance test"); - assert_arrays_eq!(filtered, array); + assert_arrays_eq!(filtered, array, &mut ctx); } /// Tests that filtering with an all-false mask returns an empty array with the same dtype @@ -100,10 +101,10 @@ fn test_selective_filter(array: &ArrayRef) { for (filtered_idx, i) in (0..len).step_by(2).enumerate() { assert_eq!( filtered - .execute_scalar(filtered_idx, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(filtered_idx, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -120,18 +121,18 @@ fn test_selective_filter(array: &ArrayRef) { assert_eq!(filtered.len(), 2); assert_eq!( filtered - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); assert_eq!( filtered - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(len - 1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(len - 1, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -153,10 +154,10 @@ fn test_single_element_filter(array: &ArrayRef) { assert_eq!(filtered.len(), 1); assert_eq!( filtered - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); @@ -171,10 +172,10 @@ fn test_single_element_filter(array: &ArrayRef) { assert_eq!(filtered.len(), 1); assert_eq!( filtered - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(len - 1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(len - 1, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -237,10 +238,10 @@ fn test_alternating_pattern_filter(array: &ArrayRef) { if keep { assert_eq!( filtered - .execute_scalar(filtered_idx, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(filtered_idx, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); filtered_idx += 1; diff --git a/vortex-array/src/compute/conformance/mask.rs b/vortex-array/src/compute/conformance/mask.rs index 42d7bb9c494..7a002662d38 100644 --- a/vortex-array/src/compute/conformance/mask.rs +++ b/vortex-array/src/compute/conformance/mask.rs @@ -6,8 +6,8 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::builtins::ArrayBuiltins; @@ -54,16 +54,16 @@ fn test_heterogenous_mask(array: &ArrayRef) { if masked_out { assert!( !masked - .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(i, &mut array_session().create_execution_ctx()) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( masked - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -87,10 +87,10 @@ fn test_empty_mask(array: &ArrayRef) { for i in 0..len { assert_eq!( masked - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -113,7 +113,7 @@ fn test_full_mask(array: &ArrayRef) { for i in 0..len { assert!( !masked - .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(i, &mut array_session().create_execution_ctx()) .vortex_expect("is_valid should succeed in conformance test") ); } @@ -135,16 +135,16 @@ fn test_alternating_mask(array: &ArrayRef) { if i % 2 == 0 { assert!( !masked - .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(i, &mut array_session().create_execution_ctx()) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( masked - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -173,7 +173,7 @@ fn test_sparse_mask(array: &ArrayRef) { let valid_count = (0..len) .filter(|&i| { masked - .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(i, &mut array_session().create_execution_ctx()) .vortex_expect("is_valid should succeed in conformance test") }) .count(); @@ -185,7 +185,7 @@ fn test_sparse_mask(array: &ArrayRef) { .filter(|&i| { pattern[i] || !array - .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(i, &mut array_session().create_execution_ctx()) .vortex_expect("is_valid should succeed in conformance test") }) .count(); @@ -208,17 +208,17 @@ fn test_single_element_mask(array: &ArrayRef) { .vortex_expect("mask should succeed in conformance test"); assert!( !masked - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(0, &mut array_session().create_execution_ctx()) .vortex_expect("is_valid should succeed in conformance test") ); for i in 1..len { assert_eq!( masked - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -249,16 +249,16 @@ fn test_double_mask(array: &ArrayRef) { if mask1_pattern[i] || mask2_pattern[i] { assert!( !double_masked - .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(i, &mut array_session().create_execution_ctx()) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( double_masked - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -282,7 +282,7 @@ fn test_nullable_mask_input(array: &ArrayRef) { let nullable_mask = BoolArray::new(bool_array.to_bit_buffer(), validity); let mask_array = - nullable_mask.to_mask_fill_null_false(&mut LEGACY_SESSION.create_execution_ctx()); + nullable_mask.to_mask_fill_null_false(&mut array_session().create_execution_ctx()); let masked = array .clone() .mask((!&mask_array).into_array()) @@ -293,16 +293,16 @@ fn test_nullable_mask_input(array: &ArrayRef) { if bool_values[i] && validity_values[i] { assert!( !masked - .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(i, &mut array_session().create_execution_ctx()) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( masked - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); diff --git a/vortex-array/src/compute/conformance/search_sorted.rs b/vortex-array/src/compute/conformance/search_sorted.rs index 472dfec9dc1..0d0cf9d8cc8 100644 --- a/vortex-array/src/compute/conformance/search_sorted.rs +++ b/vortex-array/src/compute/conformance/search_sorted.rs @@ -10,9 +10,9 @@ use vortex_buffer::Buffer; use vortex_buffer::buffer; use crate::ArrayRef; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array::IntoArray; +use crate::array_session; use crate::arrays::PrimitiveArray; use crate::patches::Patches; use crate::validity::Validity; @@ -32,7 +32,7 @@ pub fn sparse_high_null_fill() -> ArrayRef { None, ) .unwrap(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() .into_array() @@ -49,7 +49,7 @@ pub fn sparse_high_non_null_fill() -> ArrayRef { None, ) .unwrap(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() .into_array() @@ -66,7 +66,7 @@ pub fn sparse_low() -> ArrayRef { None, ) .unwrap(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() .into_array() @@ -83,7 +83,7 @@ pub fn sparse_low_high() -> ArrayRef { None, ) .unwrap(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() .into_array() @@ -100,7 +100,7 @@ pub fn sparse_edge_patch_high() -> ArrayRef { None, ) .unwrap(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() .into_array() @@ -117,7 +117,7 @@ pub fn sparse_edge_patch_low() -> ArrayRef { None, ) .unwrap(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .unwrap() .into_array() diff --git a/vortex-array/src/compute/conformance/take.rs b/vortex-array/src/compute/conformance/take.rs index 911a2ae7295..babaf392eb9 100644 --- a/vortex-array/src/compute/conformance/take.rs +++ b/vortex-array/src/compute/conformance/take.rs @@ -7,8 +7,8 @@ use vortex_error::VortexExpect; use crate::ArrayRef; use crate::Canonical; use crate::IntoArray as _; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::Nullability; @@ -84,10 +84,10 @@ fn test_take_all(array: &ArrayRef) { for i in 0..len { assert_eq!( array - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -125,11 +125,11 @@ fn test_take_selective(array: &ArrayRef) { array .execute_scalar( original_idx as usize, - &mut LEGACY_SESSION.create_execution_ctx() + &mut array_session().create_execution_ctx() ) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(result_idx, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(result_idx, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -145,18 +145,18 @@ fn test_take_first_and_last(array: &ArrayRef) { assert_eq!(result.len(), 2); assert_eq!( array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); assert_eq!( array - .execute_scalar(len - 1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(len - 1, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -190,17 +190,17 @@ fn test_take_with_nullable_indices(array: &ArrayRef) { match idx_opt { Some(idx) => { let expected = array - .execute_scalar(*idx as usize, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(*idx as usize, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); let actual = result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!(expected, actual); } None => { assert!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") .is_null() ); @@ -222,12 +222,12 @@ fn test_take_repeated_indices(array: &ArrayRef) { assert_eq!(result.len(), 3); let first_elem = array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"); for i in 0..3 { assert_eq!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), first_elem ); @@ -258,10 +258,10 @@ fn test_take_reverse(array: &ArrayRef) { for i in 0..len { assert_eq!( array - .execute_scalar(len - 1 - i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(len - 1 - i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -279,10 +279,10 @@ fn test_take_single_middle(array: &ArrayRef) { assert_eq!(result.len(), 1); assert_eq!( array - .execute_scalar(middle_idx, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(middle_idx, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -310,10 +310,10 @@ fn test_take_random_unsorted(array: &ArrayRef) { for (i, &idx) in indices.iter().enumerate() { assert_eq!( array - .execute_scalar(idx as usize, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(idx as usize, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -336,10 +336,10 @@ fn test_take_contiguous_range(array: &ArrayRef) { for i in 0..(end - start) { assert_eq!( array - .execute_scalar(start + i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(start + i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -372,10 +372,10 @@ fn test_take_mixed_repeated(array: &ArrayRef) { for (i, &idx) in indices.iter().enumerate() { assert_eq!( array - .execute_scalar(idx as usize, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(idx as usize, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -404,10 +404,10 @@ fn test_take_large_indices(array: &ArrayRef) { let expected_idx = indices[i] as usize; assert_eq!( array - .execute_scalar(expected_idx, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(expected_idx, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .vortex_expect("scalar_at should succeed in conformance test") ); } diff --git a/vortex-array/src/dtype/arrow.rs b/vortex-array/src/dtype/arrow.rs index e03cf673c21..41d8082a131 100644 --- a/vortex-array/src/dtype/arrow.rs +++ b/vortex-array/src/dtype/arrow.rs @@ -130,36 +130,63 @@ impl FromArrowType for DType { } } +impl TryFromArrowType for DType { + fn try_from_arrow(value: SchemaRef) -> VortexResult { + Self::try_from_arrow(value.as_ref()) + } +} + impl FromArrowType<&Schema> for DType { fn from_arrow(value: &Schema) -> Self { - Self::Struct( - StructFields::from_arrow(value.fields()), + Self::try_from_arrow(value).vortex_expect("arrow schema to dtype") + } +} + +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 { - StructFields::from_iter(value.into_iter().map(|f| { - ( - FieldName::from(f.name().as_str()), - DType::from_arrow(f.as_ref()), - ) - })) + 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((data_type, nullability): (&DataType, Nullability)) -> Self { + 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 DType::Primitive( - PType::try_from_arrow(data_type).vortex_expect("arrow float/integer to ptype"), + return Ok(DType::Primitive( + PType::try_from_arrow(data_type)?, nullability, - ); + )); } - match data_type { + Ok(match data_type { DataType::Null => DType::Null, DataType::Decimal32(precision, scale) | DataType::Decimal64(precision, scale) @@ -189,7 +216,7 @@ impl FromArrowType<(&DataType, Nullability)> for DType { | DataType::LargeList(e) | DataType::ListView(e) | DataType::LargeListView(e) => { - DType::List(Arc::new(Self::from_arrow(e.as_ref())), nullability) + DType::List(Arc::new(Self::try_from_arrow(e.as_ref())?), nullability) } DataType::Map(entries_field, _ordered) => { // Map is logically List>. @@ -200,33 +227,39 @@ impl FromArrowType<(&DataType, Nullability)> for DType { ) } DataType::FixedSizeList(e, size) => DType::FixedSizeList( - Arc::new(Self::from_arrow(e.as_ref())), + Arc::new(Self::try_from_arrow(e.as_ref())?), *size as u32, nullability, ), - DataType::Struct(f) => DType::Struct(StructFields::from_arrow(f), nullability), + DataType::Struct(f) => DType::Struct(StructFields::try_from_arrow(f)?, nullability), DataType::Dictionary(_, value_type) => { - Self::from_arrow((value_type.as_ref(), nullability)) + Self::try_from_arrow((value_type.as_ref(), nullability))? } DataType::RunEndEncoded(_, value_type) => { - Self::from_arrow((value_type.data_type(), nullability)) + Self::try_from_arrow((value_type.data_type(), nullability))? } - _ => unimplemented!("Arrow data type not yet supported: {:?}", data_type), - } + _ => 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") + } +} + +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 DType::Variant(field.is_nullable().into()); + return Ok(DType::Variant(field.is_nullable().into())); } - Self::from_arrow((field.data_type(), field.is_nullable().into())) + Self::try_from_arrow((field.data_type(), field.is_nullable().into())) } } @@ -602,4 +635,30 @@ mod test { assert_eq!(original_dtype, roundtripped_dtype); } + + // Regression test for https://github.com/vortex-data/vortex/issues/8346: unsupported Arrow + // types must return an error instead of panicking with `unimplemented!`. + #[rstest] + #[case::duration(DataType::Duration(ArrowTimeUnit::Microsecond))] + #[case::interval(DataType::Interval(arrow_schema::IntervalUnit::DayTime))] + #[case::fixed_size_binary(DataType::FixedSizeBinary(3))] + fn test_try_from_arrow_unsupported_type_errors(#[case] data_type: DataType) { + let err = DType::try_from_arrow((&data_type, Nullability::NonNullable)) + .expect_err("unsupported Arrow type should not convert") + .to_string(); + assert!(err.contains("not supported"), "unexpected error: {err}"); + + // The same unsupported type nested in a field or schema must also error cleanly. + let field = Field::new("c0", data_type, true); + assert!(DType::try_from_arrow(&field).is_err()); + let schema = Schema::new(vec![field]); + assert!(DType::try_from_arrow(&schema).is_err()); + } + + #[test] + fn test_try_from_arrow_supported_type_succeeds() -> VortexResult<()> { + let dtype = DType::try_from_arrow((&DataType::Int32, Nullability::Nullable))?; + assert_eq!(dtype, DType::Primitive(PType::I32, Nullability::Nullable)); + Ok(()) + } } diff --git a/vortex-array/src/dtype/extension/mod.rs b/vortex-array/src/dtype/extension/mod.rs index 96dfa7e22a3..44a6e7a01a7 100644 --- a/vortex-array/src/dtype/extension/mod.rs +++ b/vortex-array/src/dtype/extension/mod.rs @@ -3,6 +3,15 @@ //! Extension DTypes, and interfaces for working with extension types. //! +//! An extension dtype gives semantic meaning to values stored in an ordinary Vortex storage dtype. +//! For example, a timestamp extension can store an integer while carrying unit/time-zone metadata. +//! Arrays with [`DType::Extension`](crate::dtype::DType::Extension) use +//! [`crate::arrays::ExtensionArray`] to pair the extension dtype with the storage array. +//! +//! Extension dtype implementations provide an [`ExtVTable`] for identity, metadata +//! serialization, storage dtype validation, scalar validation, and typed scalar unpacking. Register +//! extension plugins in a session before deserializing data that may reference them. +//! //! ## File layout convention //! //! Each vtable-backed concept `Foo` follows this module structure: diff --git a/vortex-array/src/dtype/extension/vtable.rs b/vortex-array/src/dtype/extension/vtable.rs index d4a00fbdec4..c89f296a22b 100644 --- a/vortex-array/src/dtype/extension/vtable.rs +++ b/vortex-array/src/dtype/extension/vtable.rs @@ -15,7 +15,13 @@ use crate::scalar::ScalarValue; /// The public API for defining new extension types. /// /// This is the non-object-safe trait that plugin authors implement to define a new extension type. -/// It specifies the type's identity, metadata, serialization, and validation. +/// It specifies the type's identity, metadata, serialization, storage compatibility, and scalar +/// 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) +/// 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 { /// Associated type containing the deserialized metadata for this extension type. type Metadata: 'static + Send + Sync + Clone + Debug + Display + Eq + Hash; @@ -37,6 +43,9 @@ pub trait ExtVTable: 'static + Sized + Send + Sync + Clone + Debug + Eq + Hash { fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult; /// Validate that the given storage type is compatible with this extension type. + /// + /// This is called when constructing an [`ExtDType`] and should check both storage dtype and + /// extension metadata. fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()>; /// Can a value of `other` be implicitly widened into this type? (e.g. GeographyType might diff --git a/vortex-array/src/dtype/field.rs b/vortex-array/src/dtype/field.rs index 84f631b5074..c507b51a093 100644 --- a/vortex-array/src/dtype/field.rs +++ b/vortex-array/src/dtype/field.rs @@ -84,7 +84,6 @@ impl Display for Field { /// let dtype_i32 = DType::Primitive(PType::I32, Nullability::NonNullable); /// assert_eq!(dtype_i32, FieldPath::root().resolve(dtype_i32.clone()).unwrap()); /// ``` -/// // TODO(ngates): we should probably reverse the path. Or better yet, store a Arc<[Field]> along // with a positional index to allow cheap step_into. #[derive(Clone, Default, Debug, PartialEq, Eq, Hash)] diff --git a/vortex-array/src/dtype/mod.rs b/vortex-array/src/dtype/mod.rs index a3a4765c68c..009c2610c73 100644 --- a/vortex-array/src/dtype/mod.rs +++ b/vortex-array/src/dtype/mod.rs @@ -1,10 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! A type system for Vortex +//! A type system for Vortex. //! //! This crate contains the core logical type system for Vortex, including the definition of data types, //! and (optionally) logic for their serialization and deserialization. +//! +//! Vortex dtypes are logical domains, not physical layouts. Encodings are tracked separately on +//! arrays, so the same dtype may be backed by canonical buffers, dictionary codes, compressed +//! children, or a lazy expression array. +//! +//! Every non-null logical dtype carries [`Nullability`] directly. This differs from Apache Arrow, +//! where nullability is usually field metadata. #[cfg(feature = "arbitrary")] mod arbitrary; @@ -101,11 +108,19 @@ pub enum DType { /// `DType`. See [`StructFields`] for more information. Struct(StructFields, Nullability), - // TODO(connor)[Union]: Add more info here! /// A logical union (sum) type. + /// + /// `Union` is reserved for values that are drawn from one of several possible child domains. + /// The exact child-type metadata and canonical storage are not yet part of the stable public + /// Rust API, so most callers should prefer [`DType::Variant`] for dynamically typed values or + /// [`DType::Struct`] for fixed schemas. Union(Nullability), - /// Variant type. + /// Dynamically typed values stored as Vortex scalars. + /// + /// A variant value preserves its full logical value in variant storage and may also carry a + /// typed, row-aligned shredded representation for selected paths when materialized as a + /// [`VariantArray`](crate::arrays::VariantArray). Variant(Nullability), /// A user-defined extension type. @@ -188,8 +203,5 @@ mod test { use vortex_session::VortexSession; - use crate::dtype::session::DTypeSession; - - pub(crate) static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + pub(crate) static SESSION: LazyLock = LazyLock::new(crate::array_session); } diff --git a/vortex-array/src/dtype/ptype.rs b/vortex-array/src/dtype/ptype.rs index ab443618d66..a038eb328c2 100644 --- a/vortex-array/src/dtype/ptype.rs +++ b/vortex-array/src/dtype/ptype.rs @@ -21,6 +21,7 @@ use num_traits::Unsigned; use num_traits::bounds::UpperBounded; use vortex_error::VortexError; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_err; use crate::dtype::DType; @@ -845,7 +846,7 @@ impl TryFrom<&DType> for PType { if let DType::Primitive(p, _) = value { Ok(*p) } else { - Err(vortex_err!("Cannot convert DType {} into PType", value)) + vortex_bail!("Cannot convert DType {value} into PType") } } } diff --git a/vortex-array/src/dtype/serde/flatbuffers.rs b/vortex-array/src/dtype/serde/flatbuffers.rs index f697174e062..a61fc8f5527 100644 --- a/vortex-array/src/dtype/serde/flatbuffers.rs +++ b/vortex-array/src/dtype/serde/flatbuffers.rs @@ -207,6 +207,7 @@ impl TryFrom for DType { let fb_ext = fb .type__as_extension() .ok_or_else(|| vortex_err!("failed to parse extension from flatbuffer"))?; + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let id = ExtId::new(fb_ext.id().ok_or_else(|| { vortex_err!("failed to parse extension id from flatbuffer") diff --git a/vortex-array/src/dtype/serde/proto.rs b/vortex-array/src/dtype/serde/proto.rs index 35e944e0d66..2f992dce591 100644 --- a/vortex-array/src/dtype/serde/proto.rs +++ b/vortex-array/src/dtype/serde/proto.rs @@ -89,6 +89,7 @@ impl DType { DtypeType::Union(u) => Ok(Self::Union(u.nullable.into())), DtypeType::Variant(v) => Ok(Self::Variant(v.nullable.into())), DtypeType::Extension(e) => { + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let id = ExtId::new(e.id.as_str()); let storage_dtype = DType::from_proto( e.storage_dtype @@ -229,8 +230,6 @@ impl TryFrom<&pb::FieldPath> for FieldPath { mod tests { use std::sync::Arc; - use vortex_session::VortexSession; - use super::*; use crate::dtype::DType; use crate::dtype::DecimalDType; @@ -503,7 +502,7 @@ mod tests { #[test] fn test_unknown_extension_allow_unknown() { - let session = VortexSession::empty().allow_unknown(); + let session = crate::array_session().allow_unknown(); let proto = pb::DType { dtype_type: Some(DtypeType::Extension(Box::new(pb::Extension { id: "vortex.test.foreign_ext".to_string(), diff --git a/vortex-array/src/dtype/serde/serde.rs b/vortex-array/src/dtype/serde/serde.rs index 0be3e877f0e..789fd17306b 100644 --- a/vortex-array/src/dtype/serde/serde.rs +++ b/vortex-array/src/dtype/serde/serde.rs @@ -577,6 +577,7 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, ExtDTypeRef> { } let id = id.ok_or_else(|| de::Error::missing_field("id"))?; + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let id = ExtId::new(&id); let storage_dtype = storage_dtype.ok_or_else(|| de::Error::missing_field("storage_dtype"))?; diff --git a/vortex-array/src/dtype/session.rs b/vortex-array/src/dtype/session.rs index 757ffbd0ae5..2e13b923bc3 100644 --- a/vortex-array/src/dtype/session.rs +++ b/vortex-array/src/dtype/session.rs @@ -6,8 +6,8 @@ use std::any::Any; use std::sync::Arc; -use vortex_session::Ref; use vortex_session::SessionExt; +use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_session::registry::Registry; @@ -22,7 +22,7 @@ use crate::extension::uuid::Uuid; pub type ExtDTypeRegistry = Registry; /// Session for managing extension dtypes. -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct DTypeSession { registry: ExtDTypeRegistry, } @@ -69,11 +69,11 @@ impl DTypeSession { /// Extension trait for accessing the DType session. pub trait DTypeSessionExt: SessionExt { /// Get the DType session. - fn dtypes(&self) -> Ref<'_, DTypeSession>; + fn dtypes(&self) -> SessionGuard<'_, DTypeSession>; } impl DTypeSessionExt for S { - fn dtypes(&self) -> Ref<'_, DTypeSession> { + fn dtypes(&self) -> SessionGuard<'_, DTypeSession> { self.get::() } } diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 8fe78f4a0ad..4f41302c91f 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -15,6 +15,7 @@ use std::env::VarError; use std::fmt; use std::fmt::Display; +use std::sync::Arc; use std::sync::LazyLock; #[cfg(debug_assertions)] use std::sync::atomic::AtomicUsize; @@ -26,7 +27,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_session::SessionExt; use vortex_session::VortexSession; use crate::AnyCanonical; @@ -41,8 +41,9 @@ use crate::matcher::Matcher; use crate::memory::HostAllocatorRef; use crate::memory::MemorySessionExt; use crate::optimizer::ArrayOptimizer; -use crate::optimizer::kernels::ArrayKernels; -use crate::optimizer::kernels::KernelSnapshot; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::ParentExecutionKernels; +use crate::optimizer::kernels::execute_parent_key; use crate::stats::ArrayStats; use crate::stats::StatsSet; @@ -140,12 +141,16 @@ impl ArrayRef { /// - yes -> skip Step 2a / 2b /// - no -> try parent kernels /// - /// Step 2a: current_array.execute_parent(stack.top.parent_array) - /// child looks up at the suspended parent from ExecuteSlot + /// Step 2a: if stack.top exists: + /// parent = stack.top.parent_array + /// child = current_array + /// kernels[(parent.encoding_id(), child.encoding_id())] + /// .try_execute_parent(child, parent, stack.top.slot_idx) /// /// Step 2b: for child in current_array.children(): - /// child.execute_parent(current_array) - /// each child looks up at current_array + /// parent = current_array + /// kernels[(parent.encoding_id(), child.encoding_id())] + /// .try_execute_parent(child, parent, child.slot_idx) /// /// Step 3: match current_array.execute() /// ExecuteSlot(i, pred) -> push parent on stack, focus child `i` @@ -162,6 +167,8 @@ impl ArrayRef { let mut current_array = self; let mut current_builder: Option> = None; let mut stack: Vec = Vec::new(); + let execute_parent_kernels = Arc::clone(&ctx.execute_parent_kernels); + let kernels = execute_parent_kernels.as_ref(); let max_iterations = max_iterations(); for _ in 0..max_iterations { @@ -198,8 +205,15 @@ impl ArrayRef { // would be lost when we restore frame.parent_builder. if current_builder.is_none() && let Some(frame) = stack.last() - && let Some(result) = - current_array.execute_parent(&frame.parent_array, frame.slot_idx, ctx)? + && let Some(result) = { + execute_parent_for_child( + &frame.parent_array, + ¤t_array, + frame.slot_idx, + kernels, + ctx, + )? + } { ctx.log(format_args!( "execute_parent (stack) rewrote {} -> {}", @@ -213,7 +227,7 @@ impl ArrayRef { // Step 2b: execute_parent against current_array's own children. if current_builder.is_none() - && let Some(rewritten) = try_execute_parent(¤t_array, ctx)? + && let Some(rewritten) = try_execute_parent(¤t_array, kernels, ctx)? { ctx.log(format_args!( "execute_parent rewrote {} -> {}", @@ -305,16 +319,7 @@ struct StackFrame { #[derive(Debug, Clone)] pub struct ExecutionCtx { session: VortexSession, - /// Snapshot of the session's [`ArrayKernels`] execute-parent kernels, resolved once at - /// construction. - /// - /// The registry is session-scoped and mutable through its public `register_*` methods, so - /// this context sees the kernels as registered when it was created; later registrations are - /// picked up by the next context (contexts are created per evaluation). Caching the snapshot - /// avoids a per-array-node session clone plus a sharded `DashMap` `RwLock` probe in the hot - /// `execute_until` loop, and avoids holding the session-variable read guard across kernel - /// invocation. `None` means the session had no [`ArrayKernels`] when the context was created. - kernels: Option, + execute_parent_kernels: Arc, #[cfg(debug_assertions)] id: usize, #[cfg(debug_assertions)] @@ -323,11 +328,15 @@ pub struct ExecutionCtx { impl ExecutionCtx { /// Create a new execution context with the given session. + /// + /// This captures a snapshot of the session's execute-parent kernel registry. Kernels + /// registered after this context is created are not visible to it; create a new + /// [`ExecutionCtx`] after registration to use newly registered kernels. pub fn new(session: VortexSession) -> Self { - let kernels = session.get_opt::().map(|k| k.snapshot()); + let execute_parent_kernels = session.kernels().execute_parent_snapshot(); Self { session, - kernels, + execute_parent_kernels, #[cfg(debug_assertions)] id: { static EXEC_CTX_ID: AtomicUsize = AtomicUsize::new(0); @@ -343,12 +352,6 @@ impl ExecutionCtx { &self.session } - /// Get the [`KernelSnapshot`] resolved once for this execution context, if the session had an - /// [`ArrayKernels`] registry. Cheap to clone (one `Arc` clone). - pub(crate) fn kernels(&self) -> Option<&KernelSnapshot> { - self.kernels.as_ref() - } - /// Get the session-scoped host allocator for this execution context. pub fn allocator(&self) -> HostAllocatorRef { self.session.allocator() @@ -362,7 +365,7 @@ impl ExecutionCtx { /// Use the [`format_args!`] macro to create the `msg` argument. pub fn log(&mut self, msg: fmt::Arguments<'_>) { #[cfg(debug_assertions)] - if tracing::enabled!(tracing::Level::DEBUG) { + if tracing::enabled!(tracing::Level::TRACE) { let formatted = format!(" - {msg}"); tracing::trace!("exec[{}]: {formatted}", self.id); self.ops.push(formatted); @@ -442,12 +445,13 @@ impl Executable for ArrayRef { } } - let kernels = ctx.kernels().cloned(); + 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.as_ref(), ctx)? + execute_parent_for_child(&array, child, slot_idx, kernels, ctx)? { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", @@ -474,7 +478,9 @@ impl Executable for ArrayRef { ExecutionStep::ExecuteSlot(i, _) => { let child = array.slots()[i].clone().vortex_expect("valid slot index"); let executed_child = child.execute::(ctx)?; - array.with_slot(i, executed_child) + // SAFETY: execution of a child slot produces a logically equivalent array in a + // different physical representation, preserving parent values and statistics. + unsafe { array.with_slot(i, executed_child) } } ExecutionStep::AppendChild(_) => { // Single-step: build the entire parent via the builder path. @@ -559,31 +565,41 @@ fn execute_parent_for_child( parent: &ArrayRef, child: &ArrayRef, slot_idx: usize, - kernels: Option<&KernelSnapshot>, + kernels: &ParentExecutionKernels, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if let Some(kernels) = kernels - && let Some(plugins) = - kernels.find_execute_parent(parent.encoding_id(), child.encoding_id()) - { + let key = execute_parent_key(parent.encoding_id(), child.encoding_id()); + if let Some(plugins) = kernels.get(&key) { for plugin in plugins.as_ref() { - if let Some(result) = plugin(child, parent, slot_idx, ctx)? { + if let Some(result) = plugin.execute_parent(child, parent, slot_idx, ctx)? { + if cfg!(debug_assertions) { + vortex_ensure!( + result.len() == parent.len(), + "Executed parent canonical length mismatch" + ); + vortex_ensure!( + result.dtype() == parent.dtype(), + "Executed parent canonical dtype mismatch" + ); + } return Ok(Some(result)); } } } - child.execute_parent(parent, slot_idx, ctx) + Ok(None) } /// Try execute_parent on each occupied slot of the array. -fn try_execute_parent(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { - let kernels = ctx.kernels().cloned(); - +fn try_execute_parent( + array: &ArrayRef, + kernels: &ParentExecutionKernels, + ctx: &mut ExecutionCtx, +) -> VortexResult> { 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.as_ref(), ctx)? + execute_parent_for_child(array, child, slot_idx, kernels, ctx)? { ctx.log(format_args!( "execute_parent: slot[{}]({}) rewrote {} -> {}", @@ -834,3 +850,55 @@ impl VortexSessionExecute for VortexSession { ExecutionCtx::new(self.clone()) } } + +#[cfg(test)] +mod tests { + use vortex_session::VortexSession; + + use super::*; + use crate::VTable as _; + use crate::VortexSessionExecute; + use crate::arrays::Bool; + use crate::arrays::Primitive; + use crate::optimizer::kernels::ExecuteParentFn; + use crate::optimizer::kernels::KernelSession; + use crate::optimizer::kernels::execute_parent_key; + + fn noop_execute_parent( + _child: &ArrayRef, + _parent: &ArrayRef, + _child_idx: usize, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(None) + } + + #[test] + fn execution_ctx_snapshots_execute_parent_kernels_at_creation() { + let session = VortexSession::empty().with_some(KernelSession::empty()); + let key = execute_parent_key(Bool.id(), Primitive.id()); + + let before_registration = session.create_execution_ctx(); + assert!( + !before_registration + .execute_parent_kernels + .contains_key(&key) + ); + + let kernels = session.kernels(); + kernels.register_execute_parent( + Bool.id(), + Primitive.id(), + &[noop_execute_parent as ExecuteParentFn], + ); + + assert!( + !before_registration + .execute_parent_kernels + .contains_key(&key) + ); + + let after_registration = session.create_execution_ctx(); + assert!(after_registration.execute_parent_kernels.contains_key(&key)); + } +} diff --git a/vortex-array/src/expr/analysis/annotation.rs b/vortex-array/src/expr/analysis/annotation.rs index be842742a5a..32608f38d1e 100644 --- a/vortex-array/src/expr/analysis/annotation.rs +++ b/vortex-array/src/expr/analysis/annotation.rs @@ -42,6 +42,25 @@ pub fn descendent_annotations( let mut visitor = AnnotationVisitor { annotations: Default::default(), annotate, + propagate_up: true, + }; + expr.accept(&mut visitor).vortex_expect("Infallible"); + visitor.annotations +} + +/// Walk the expression tree and annotate each expression with zero or more +/// 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> { + let mut visitor = AnnotationVisitor { + annotations: Default::default(), + annotate, + propagate_up: false, }; expr.accept(&mut visitor).vortex_expect("Infallible"); visitor.annotations @@ -50,6 +69,7 @@ pub fn descendent_annotations( struct AnnotationVisitor<'a, A: AnnotationFn> { annotations: Annotations<'a, A::Annotation>, annotate: A, + propagate_up: bool, } impl<'a, A: AnnotationFn> NodeVisitor<'a> for AnnotationVisitor<'a, A> { @@ -70,6 +90,9 @@ impl<'a, A: AnnotationFn> NodeVisitor<'a> for AnnotationVisitor<'a, A> { } fn visit_up(&mut self, node: &'a Expression) -> VortexResult { + if !self.propagate_up { + return Ok(TraversalOrder::Continue); + } let child_annotations = node .children() .iter() diff --git a/vortex-array/src/expr/analysis/labeling.rs b/vortex-array/src/expr/analysis/labeling.rs index 090e2992fe2..721241a14c9 100644 --- a/vortex-array/src/expr/analysis/labeling.rs +++ b/vortex-array/src/expr/analysis/labeling.rs @@ -29,7 +29,6 @@ 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, diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index cc21fb9a9a6..41cd61a369d 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -15,9 +15,7 @@ use vortex_error::vortex_ensure; use vortex_session::VortexSession; use crate::dtype::DType; -use crate::expr::StatsCatalog; use crate::expr::display::DisplayTreeExpr; -use crate::expr::stats::Stat; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::fns::root::Root; @@ -114,28 +112,6 @@ impl Expression { self.scalar_fn.validity(self) } - /// An expression over zone-statistics which implies all records in the zone evaluate to false. - /// - /// Given an expression, `e`, if `e.stat_falsification(..)` evaluates to true, it is guaranteed - /// that `e` evaluates to false on all records in the zone. However, the inverse is not - /// necessarily true: even if the falsification evaluates to false, `e` need not evaluate to - /// true on all records. - /// - /// The [`StatsCatalog`] can be used to constrain or rename stats used in the final expr. - /// - /// # Examples - /// - /// - An expression over one variable: `x > 0` is false for all records in a zone if the maximum - /// value of the column `x` in that zone is less than or equal to zero: `max(x) <= 0`. - /// - An expression over two variables: `x > y` becomes `max(x) <= min(y)`. - /// - A conjunctive expression: `x > y AND z < x` becomes `max(x) <= min(y) OR min(z) >= max(x). - /// - /// Some expressions, in theory, have falsifications but this function does not support them - /// such as `x < (y < z)` or `x LIKE "needle%"`. - pub fn stat_falsification(&self, catalog: &dyn StatsCatalog) -> Option { - self.scalar_fn().stat_falsification(self, catalog) - } - /// Returns an expression that proves this predicate is definitely false from stats. /// /// `scope` is the dtype of the row this expression evaluates over. @@ -164,28 +140,6 @@ impl Expression { crate::stats::rewrite::StatsRewriteCtx::new(session, scope).satisfy(self) } - /// Returns an expression representing the zoned statistic for the given stat, if available. - /// - /// The [`StatsCatalog`] returns expressions that can be evaluated using the zone map as a - /// scope. Expressions can implement this function to propagate such statistics through the - /// expression tree. For example, the `a + 10` expression could propagate `min: min(a) + 10`. - /// - /// NOTE(gatesn): we currently cannot represent statistics over nested fields. Please file an - /// issue to discuss a solution to this. - pub fn stat_expression(&self, stat: Stat, catalog: &dyn StatsCatalog) -> Option { - self.scalar_fn().stat_expression(self, stat, catalog) - } - - /// Returns an expression representing the zoned maximum statistic, if available. - pub fn stat_min(&self, catalog: &dyn StatsCatalog) -> Option { - self.stat_expression(Stat::Min, catalog) - } - - /// Returns an expression representing the zoned maximum statistic, if available. - pub fn stat_max(&self, catalog: &dyn StatsCatalog) -> Option { - self.stat_expression(Stat::Max, catalog) - } - /// Format the expression as a compact string. /// /// Since this is a recursive formatter, it is exposed on the public Expression type. diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 73f4991d05c..ec12a321552 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -29,6 +29,7 @@ use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::dynamic::DynamicComparison; use crate::scalar_fn::fns::dynamic::DynamicComparisonExpr; use crate::scalar_fn::fns::dynamic::Rhs; +use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::is_not_null::IsNotNull; @@ -36,6 +37,7 @@ use crate::scalar_fn::fns::is_null::IsNull; use crate::scalar_fn::fns::like::Like; use crate::scalar_fn::fns::like::LikeOptions; use crate::scalar_fn::fns::list_contains::ListContains; +use crate::scalar_fn::fns::list_length::ListLength; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::merge::DuplicateHandling; @@ -421,13 +423,13 @@ where /// ``` /// # use vortex_array::IntoArray; /// # use vortex_array::arrow::ArrowArrayExecutor; -/// # use vortex_array::{LEGACY_SESSION, VortexSessionExecute}; +/// # use vortex_array::{VortexSessionExecute, array_session}; /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{checked_add, lit, root}; /// let xs = buffer![1, 2, 3].into_array(); /// let result = xs.apply(&checked_add(root(), lit(5))).unwrap(); /// -/// let mut ctx = LEGACY_SESSION.create_execution_ctx(); +/// let mut ctx = array_session().create_execution_ctx(); /// assert_eq!( /// &result.execute_arrow(None, &mut ctx).unwrap(), /// &buffer![6, 7, 8] @@ -737,3 +739,29 @@ pub fn list_contains(list: Expression, value: Expression) -> Expression { pub fn byte_length(input: Expression) -> Expression { ByteLength.new_expr(EmptyOptions, [input]) } + +// ---- ExtStorage ---- + +/// Creates an expression that extracts the storage values from an extension array. +/// +/// ```rust +/// # use vortex_array::expr::{ext_storage, root}; +/// let expr = ext_storage(root()); +/// ``` +pub fn ext_storage(input: Expression) -> Expression { + ExtStorage.new_expr(EmptyOptions, [input]) +} + +// ---- ListLength ---- + +/// Creates an expression that computes the number of elements in each list +/// for `List` and `FixedSizeList` inputs. This is akin to ANSI SQL `CARDINALITY()`, +/// or DuckDB's `len()`/`array_length()`. +/// +/// ```rust +/// # use vortex_array::expr::{list_length, root}; +/// let expr = list_length(root()); +/// ``` +pub fn list_length(input: Expression) -> Expression { + ListLength.new_expr(EmptyOptions, [input]) +} diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index a5d32510443..4331b2817f0 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -1,9 +1,38 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Vortex's expression language. +//! Vortex's expression language: scalar operations over [arrays](crate::ArrayRef). //! -//! All expressions are serializable, and own their own wire format. +//! An [`Expression`] is a tree of scalar operations rooted at a scope (see [`root`]). Expressions +//! are the common currency of scans: a scan takes a *filter* expression that resolves to a boolean +//! and a *projection* expression that shapes the output. All expressions are serializable and own +//! their own wire format, so they can be pushed down to remote sources and reconstructed on workers. +//! +//! # Scalar functions +//! +//! 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 +//! functions live in [`crate::scalar_fn`]; integration and plugin crates supply additional, +//! use-case-specific functions. +//! +//! # Deferred execution +//! +//! Applying an expression to an array does not compute the result eagerly. Instead it builds a +//! [`ScalarFnArray`](crate::arrays::ScalarFnArray) representing the deferred application, letting +//! downstream encodings push the computation into compressed data, or fuse several expressions +//! together, before any data is materialized. The deferred tree is executed toward canonical form +//! only when a result is actually required. +//! +//! # Typing and coercion +//! +//! 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`. +//! +//! Filter expressions are decomposed into independent conjuncts with [`split_conjunction`] so that +//! scans can evaluate and reorder the most selective predicates first. //! //! The implementation takes inspiration from [Postgres] and [Apache Datafusion]. //! @@ -34,7 +63,6 @@ pub(crate) mod field; pub mod forms; mod optimize; pub mod proto; -pub mod pruning; pub mod stats; pub mod transform; pub mod traversal; @@ -42,7 +70,6 @@ pub mod traversal; pub use analysis::*; pub use expression::*; pub use exprs::*; -pub use pruning::StatsCatalog; pub trait VortexExprExt { /// Accumulate all field references from this expression and its children in a set diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index 27859235e88..9f2b81bde84 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -39,6 +39,7 @@ impl ExprSerializeProtoExt for Expression { impl Expression { pub fn from_proto(expr: &pb::Expr, session: &VortexSession) -> VortexResult { + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] let expr_id = ScalarFnId::new(expr.id.as_str()); let children = expr .children @@ -74,7 +75,7 @@ mod tests { use vortex_session::VortexSession; use super::ExprSerializeProtoExt; - use crate::LEGACY_SESSION; + use crate::array_session; use crate::expr::Expression; use crate::expr::and; use crate::expr::between; @@ -108,7 +109,7 @@ mod tests { let s_expr = expr.serialize_proto().unwrap(); let buf = s_expr.encode_to_vec(); let s_expr = pb::Expr::decode(buf.as_slice()).unwrap(); - let deser_expr = Expression::from_proto(&s_expr, &LEGACY_SESSION).unwrap(); + let deser_expr = Expression::from_proto(&s_expr, &array_session()).unwrap(); assert_eq!(&deser_expr, &expr); } diff --git a/vortex-array/src/expr/pruning/mod.rs b/vortex-array/src/expr/pruning/mod.rs deleted file mode 100644 index 7c20508b7a8..00000000000 --- a/vortex-array/src/expr/pruning/mod.rs +++ /dev/null @@ -1,27 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -pub(crate) mod pruning_expr; -mod relation; - -pub use pruning_expr::RequiredStats; -pub use pruning_expr::checked_pruning_expr; -pub use pruning_expr::field_path_stat_field_name; -pub use relation::Relation; - -use crate::dtype::FieldPath; -use crate::expr::Expression; -use crate::expr::stats::Stat; - -/// A catalog of available stats that are associated with field paths. -pub trait StatsCatalog { - /// Given a field path and statistic, return an expression that when evaluated over the catalog - /// will return that stat for the referenced field. - /// - /// This is likely to be a column expression, or a literal. - /// - /// Returns `None` if the stat is not available for the field path. - fn stats_ref(&self, _field_path: &FieldPath, _stat: Stat) -> Option { - None - } -} diff --git a/vortex-array/src/expr/pruning/pruning_expr.rs b/vortex-array/src/expr/pruning/pruning_expr.rs deleted file mode 100644 index 00d29fbcf99..00000000000 --- a/vortex-array/src/expr/pruning/pruning_expr.rs +++ /dev/null @@ -1,561 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::cell::RefCell; -use std::iter; - -use itertools::Itertools; -use vortex_utils::aliases::hash_map::HashMap; - -use super::relation::Relation; -use crate::dtype::Field; -use crate::dtype::FieldName; -use crate::dtype::FieldPath; -use crate::dtype::FieldPathSet; -use crate::expr::Expression; -use crate::expr::StatsCatalog; -use crate::expr::get_item; -use crate::expr::root; -use crate::expr::stats::Stat; - -pub type RequiredStats = Relation; - -// A catalog that return a stat column whenever it is required, tracking all accessed -// stats and returning them later. -#[derive(Default)] -pub(crate) struct TrackingStatsCatalog { - usage: RefCell>, -} - -impl TrackingStatsCatalog { - /// Consume the catalog, yielding a map of field statistics that were required - /// for each expression. - fn into_usages(self) -> HashMap<(FieldPath, Stat), Expression> { - self.usage.into_inner() - } -} - -// A catalog that return a stat column if it exists in the given scope. -struct ScopeStatsCatalog<'a> { - inner: TrackingStatsCatalog, - available_stats: &'a FieldPathSet, -} - -impl StatsCatalog for ScopeStatsCatalog<'_> { - fn stats_ref(&self, field_path: &FieldPath, stat: Stat) -> Option { - let stat_path = field_path.clone().push(stat.name()); - - if self.available_stats.contains(&stat_path) { - self.inner.stats_ref(field_path, stat) - } else { - None - } - } -} - -impl StatsCatalog for TrackingStatsCatalog { - fn stats_ref(&self, field_path: &FieldPath, stat: Stat) -> Option { - let mut expr = root(); - let name = field_path_stat_field_name(field_path, stat); - expr = get_item(name, expr); - self.usage - .borrow_mut() - .insert((field_path.clone(), stat), expr.clone()); - Some(expr) - } -} - -#[doc(hidden)] -pub fn field_path_stat_field_name(field_path: &FieldPath, stat: Stat) -> FieldName { - field_path - .parts() - .iter() - .map(|f| match f { - Field::Name(n) => n.as_ref(), - Field::ElementType => todo!("element type not currently handled"), - }) - .chain(iter::once(stat.name())) - .join("_") - .into() -} - -/// Build a pruning expr mask, using an existing set of stats. -/// The available stats are provided as a set of [`FieldPath`]. -/// -/// A pruning expression is one that returns `true` for all positions where the original expression -/// cannot hold, and false if it cannot be determined from stats alone whether the positions can -/// be pruned. -/// -/// Some rewrites, such as `is_not_null(...)`, emit -/// [`row_count`][crate::scalar_fn::internal::row_count] placeholders. The evaluation layer must -/// replace those placeholders with the row count for its current scope before -/// executing the returned expression. -/// -/// If the falsification logic attempts to access an unknown stat, -/// this function will return `None`. -pub fn checked_pruning_expr( - expr: &Expression, - available_stats: &FieldPathSet, -) -> Option<(Expression, RequiredStats)> { - let catalog = ScopeStatsCatalog { - inner: Default::default(), - available_stats, - }; - - let expr = expr.stat_falsification(&catalog)?; - - // TODO(joe): filter access by used exprs - let mut relation: Relation = Relation::new(); - for ((field_path, stat), _) in catalog.inner.into_usages() { - relation.insert(field_path, stat) - } - - Some((expr, relation)) -} - -#[cfg(test)] -mod tests { - use rstest::fixture; - use rstest::rstest; - use vortex_utils::aliases::hash_set::HashSet; - - use super::HashMap; - use crate::dtype::DType; - use crate::dtype::FieldName; - use crate::dtype::FieldNames; - use crate::dtype::FieldPath; - use crate::dtype::FieldPathSet; - use crate::dtype::Nullability; - use crate::dtype::StructFields; - use crate::expr::and; - use crate::expr::between; - use crate::expr::cast; - 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_eq; - use crate::expr::or; - use crate::expr::pruning::checked_pruning_expr; - use crate::expr::pruning::field_path_stat_field_name; - use crate::expr::root; - use crate::expr::stats::Stat; - use crate::scalar_fn::fns::between::BetweenOptions; - use crate::scalar_fn::fns::between::StrictComparison; - - // Implement some checked pruning expressions. - #[fixture] - fn available_stats() -> FieldPathSet { - let field_a = FieldPath::from_name("a"); - let field_b = FieldPath::from_name("b"); - - FieldPathSet::from_iter([ - field_a.clone().push(Stat::Min.name()), - field_a.push(Stat::Max.name()), - field_b.clone().push(Stat::Min.name()), - field_b.push(Stat::Max.name()), - ]) - } - - #[rstest] - pub fn pruning_equals(available_stats: FieldPathSet) { - let name = FieldName::from("a"); - let literal_eq = lit(42); - let eq_expr = eq(get_item("a", root()), literal_eq.clone()); - let (converted, _refs) = checked_pruning_expr(&eq_expr, &available_stats).unwrap(); - let expected_expr = or( - gt( - get_item( - field_path_stat_field_name(&FieldPath::from_name(name.clone()), Stat::Min), - root(), - ), - literal_eq.clone(), - ), - gt( - literal_eq, - col(field_path_stat_field_name( - &FieldPath::from_name(name), - Stat::Max, - )), - ), - ); - assert_eq!(&converted, &expected_expr); - } - - #[rstest] - pub fn pruning_equals_column(available_stats: FieldPathSet) { - let column = FieldName::from("a"); - let other_col = FieldName::from("b"); - let eq_expr = eq(col(column.clone()), col(other_col.clone())); - - let (converted, refs) = checked_pruning_expr(&eq_expr, &available_stats).unwrap(); - assert_eq!( - refs.map(), - &HashMap::from_iter([ - ( - FieldPath::from_name(column.clone()), - HashSet::from_iter([Stat::Min, Stat::Max]) - ), - ( - FieldPath::from_name(other_col.clone()), - HashSet::from_iter([Stat::Max, Stat::Min]) - ) - ]) - ); - let expected_expr = or( - gt( - col(field_path_stat_field_name( - &FieldPath::from_name(column.clone()), - Stat::Min, - )), - col(field_path_stat_field_name( - &FieldPath::from_name(other_col.clone()), - Stat::Max, - )), - ), - gt( - col(field_path_stat_field_name( - &FieldPath::from_name(other_col), - Stat::Min, - )), - col(field_path_stat_field_name( - &FieldPath::from_name(column), - Stat::Max, - )), - ), - ); - assert_eq!(&converted, &expected_expr); - } - - #[rstest] - pub fn pruning_not_equals_column(available_stats: FieldPathSet) { - let column = FieldName::from("a"); - let other_col = FieldName::from("b"); - let not_eq_expr = not_eq(col(column.clone()), col(other_col.clone())); - - let (converted, refs) = checked_pruning_expr(¬_eq_expr, &available_stats).unwrap(); - assert_eq!( - refs.map(), - &HashMap::from_iter([ - ( - FieldPath::from_name(column.clone()), - HashSet::from_iter([Stat::Min, Stat::Max]) - ), - ( - FieldPath::from_name(other_col.clone()), - HashSet::from_iter([Stat::Max, Stat::Min]) - ) - ]) - ); - let expected_expr = and( - eq( - col(field_path_stat_field_name( - &FieldPath::from_name(column.clone()), - Stat::Min, - )), - col(field_path_stat_field_name( - &FieldPath::from_name(other_col.clone()), - Stat::Max, - )), - ), - eq( - col(field_path_stat_field_name( - &FieldPath::from_name(column), - Stat::Max, - )), - col(field_path_stat_field_name( - &FieldPath::from_name(other_col), - Stat::Min, - )), - ), - ); - - assert_eq!(&converted, &expected_expr); - } - - #[rstest] - pub fn pruning_gt_column(available_stats: FieldPathSet) { - let column = FieldName::from("a"); - let other_col = FieldName::from("b"); - let other_expr = col(other_col.clone()); - let not_eq_expr = gt(col(column.clone()), other_expr); - - let (converted, refs) = checked_pruning_expr(¬_eq_expr, &available_stats).unwrap(); - assert_eq!( - refs.map(), - &HashMap::from_iter([ - ( - FieldPath::from_name(column.clone()), - HashSet::from_iter([Stat::Max]) - ), - ( - FieldPath::from_name(other_col.clone()), - HashSet::from_iter([Stat::Min]) - ) - ]) - ); - let expected_expr = lt_eq( - col(field_path_stat_field_name( - &FieldPath::from_name(column), - Stat::Max, - )), - col(field_path_stat_field_name( - &FieldPath::from_name(other_col), - Stat::Min, - )), - ); - assert_eq!(&converted, &expected_expr); - } - - #[rstest] - pub fn pruning_gt_value(available_stats: FieldPathSet) { - let column = FieldName::from("a"); - let other_col = lit(42); - let not_eq_expr = gt(col(column.clone()), other_col.clone()); - - let (converted, refs) = checked_pruning_expr(¬_eq_expr, &available_stats).unwrap(); - assert_eq!( - refs.map(), - &HashMap::from_iter([( - FieldPath::from_name(column.clone()), - HashSet::from_iter([Stat::Max]) - ),]) - ); - let expected_expr = lt_eq( - col(field_path_stat_field_name( - &FieldPath::from_name(column), - Stat::Max, - )), - other_col, - ); - assert_eq!(&converted, &(expected_expr)); - } - - #[rstest] - pub fn pruning_lt_column(available_stats: FieldPathSet) { - let column = FieldName::from("a"); - let other_col = FieldName::from("b"); - let other_expr = col(other_col.clone()); - let not_eq_expr = lt(col(column.clone()), other_expr); - - let (converted, refs) = checked_pruning_expr(¬_eq_expr, &available_stats).unwrap(); - assert_eq!( - refs.map(), - &HashMap::from_iter([ - ( - FieldPath::from_name(column.clone()), - HashSet::from_iter([Stat::Min]) - ), - ( - FieldPath::from_name(other_col.clone()), - HashSet::from_iter([Stat::Max]) - ) - ]) - ); - let expected_expr = gt_eq( - col(field_path_stat_field_name( - &FieldPath::from_name(column), - Stat::Min, - )), - col(field_path_stat_field_name( - &FieldPath::from_name(other_col), - Stat::Max, - )), - ); - assert_eq!(&converted, &expected_expr); - } - - #[rstest] - pub fn pruning_lt_value(available_stats: FieldPathSet) { - // expression => a < 42 - // pruning expr => a.min >= 42 - let expr = lt(col("a"), lit(42)); - - let (converted, refs) = checked_pruning_expr(&expr, &available_stats).unwrap(); - assert_eq!( - refs.map(), - &HashMap::from_iter([(FieldPath::from_name("a"), HashSet::from_iter([Stat::Min]))]) - ); - assert_eq!(&converted, >_eq(col("a_min"), lit(42))); - } - - #[rstest] - fn pruning_identity(available_stats: FieldPathSet) { - let expr = or(lt(col("a"), lit(10)), gt(col("a"), lit(50))); - - let (predicate, _) = checked_pruning_expr(&expr, &available_stats).unwrap(); - - let expected_expr = and(gt_eq(col("a_min"), lit(10)), lt_eq(col("a_max"), lit(50))); - assert_eq!(&predicate.to_string(), &expected_expr.to_string()); - } - #[rstest] - pub fn pruning_and_or_operators(available_stats: FieldPathSet) { - // Test case: a > 10 AND a < 50 - let column = FieldName::from("a"); - let and_expr = and(gt(col(column.clone()), lit(10)), lt(col(column), lit(50))); - let (predicate, _) = checked_pruning_expr(&and_expr, &available_stats).unwrap(); - - // Expected: a_max <= 10 OR a_min >= 50 - assert_eq!( - &predicate, - &or( - lt_eq(col(FieldName::from("a_max")), lit(10)), - gt_eq(col(FieldName::from("a_min")), lit(50)), - ), - ); - } - - #[rstest] - fn test_gt_eq_with_booleans(available_stats: FieldPathSet) { - // Consider this unusual, but valid (in Arrow, BooleanArray implements ArrayOrd), filter expression: - // x > (y > z) - // The x column is a Boolean-valued column. The y and z columns are numeric. True > False. - // Suppose we had a Vortex zone whose min/max statistics for each column were: - // x: [True, True] - // y: [1, 2] - // z: [0, 2] - // The pruning predicate will convert the aforementioned expression into: - // x_max <= (y_min > z_min) - // If we evaluate that pruning expression on our zone we get: - // x_max <= (y_min > z_min) - // x_max <= (1 > 0 ) - // x_max <= True - // True <= True - // True - // If a pruning predicate evaluates to true then, as stated in PruningPredicate::evaluate: - // > a true value means the chunk can be pruned. - // But, the following record lies within the above intervals and *passes* the filter expression! We - // cannot prune this zone because we need this record! - // {x: True, y: 1, z: 2} - // x > (y > z) - // True > (1 > 2) - // True > False - // True - let expr = gt_eq(col("x"), gt(col("y"), col("z"))); - assert!(checked_pruning_expr(&expr, &available_stats).is_none()); - // TODO(DK): a sufficiently complex pruner would produce: `x_max <= (y_max > z_min)` - } - - #[fixture] - fn available_stats_with_nans() -> FieldPathSet { - let float_col = FieldPath::from_name("float_col"); - let int_col = FieldPath::from_name("int_col"); - - FieldPathSet::from_iter([ - // Float columns will have a NaNCount. - float_col.clone().push(Stat::Min.name()), - float_col.clone().push(Stat::Max.name()), - float_col.push(Stat::NaNCount.name()), - // int columns will not have a NanCount serialized into the layout - int_col.clone().push(Stat::Min.name()), - int_col.push(Stat::Max.name()), - ]) - } - - #[rstest] - fn pruning_checks_nans(available_stats_with_nans: FieldPathSet) { - let expr = gt_eq(col("float_col"), lit(f32::NAN)); - let (converted, _) = checked_pruning_expr(&expr, &available_stats_with_nans).unwrap(); - assert_eq!( - &converted, - &and( - and( - eq(col("float_col_nan_count"), lit(0u64)), - // NaNCount of NaN is 1 - eq(lit(1u64), lit(0u64)), - ), - // This is the standard conversion of the >= operator. Comparing NAN to a max - // stat is nonsensical, as min/max stats ignore NaNs, but this should be short-circuited - // by the previous check for nan_count anyway. - lt(col("float_col_max"), lit(f32::NAN)), - ) - ); - - // One half of the expression requires NAN count check, the other half does not. - let expr = and( - gt(col("float_col"), lit(10f32)), - lt(col("int_col"), lit(10)), - ); - - let (converted, _) = checked_pruning_expr(&expr, &available_stats_with_nans).unwrap(); - - assert_eq!( - &converted, - &or( - // NaNCount check is enforced for the float column - and( - and( - eq(col("float_col_nan_count"), lit(0u64)), - // NanCount of a non-NaN float literal is 0 - eq(lit(0u64), lit(0u64)), - ), - // We want the opposite: we can prune IF either one is false. - lt_eq(col("float_col_max"), lit(10f32)), - ), - // NanCount check is skipped for the int column - gt_eq(col("int_col_min"), lit(10)), - ) - ) - } - - #[rstest] - fn pruning_between(available_stats: FieldPathSet) { - let expr = between( - col("a"), - lit(10), - lit(50), - BetweenOptions { - lower_strict: StrictComparison::NonStrict, - upper_strict: StrictComparison::NonStrict, - }, - ); - let (converted, refs) = checked_pruning_expr(&expr, &available_stats).unwrap(); - assert_eq!( - refs.map(), - &HashMap::from_iter([( - FieldPath::from_name("a"), - HashSet::from_iter([Stat::Min, Stat::Max]) - )]) - ); - assert_eq!( - &converted, - &or(gt(lit(10), col("a_max")), gt(col("a_min"), lit(50))) - ); - } - - #[rstest] - fn pruning_cast_get_item_eq(available_stats: FieldPathSet) { - // This test verifies that cast properly forwards analysis methods to - // enable pruning. - let struct_dtype = DType::Struct( - StructFields::new( - FieldNames::from([FieldName::from("a"), FieldName::from("b")]), - vec![ - DType::Utf8(Nullability::Nullable), - DType::Utf8(Nullability::Nullable), - ], - ), - Nullability::NonNullable, - ); - let expr = eq(get_item("a", cast(root(), struct_dtype)), lit("value")); - let (converted, refs) = checked_pruning_expr(&expr, &available_stats).unwrap(); - assert_eq!( - refs.map(), - &HashMap::from_iter([( - FieldPath::from_name("a"), - HashSet::from_iter([Stat::Min, Stat::Max]) - )]) - ); - assert_eq!( - &converted, - &or( - gt(col("a_min"), lit("value")), - gt(lit("value"), col("a_max")) - ) - ); - } -} diff --git a/vortex-array/src/expr/pruning/relation.rs b/vortex-array/src/expr/pruning/relation.rs deleted file mode 100644 index 3a33771c46e..00000000000 --- a/vortex-array/src/expr/pruning/relation.rs +++ /dev/null @@ -1,50 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::hash::Hash; - -use vortex_utils::aliases::hash_map::HashMap; -use vortex_utils::aliases::hash_map::IntoIter; -use vortex_utils::aliases::hash_set::HashSet; - -#[derive(Debug, Clone)] -pub struct Relation { - map: HashMap>, -} - -impl Default for Relation { - fn default() -> Self { - Self::new() - } -} - -impl Relation { - pub fn new() -> Self { - Relation { - map: HashMap::new(), - } - } - - pub fn insert(&mut self, k: K, v: V) { - self.map.entry(k).or_default().insert(v); - } - - pub fn map(&self) -> &HashMap> { - &self.map - } -} - -impl From>> for Relation { - fn from(value: HashMap>) -> Self { - Self { map: value } - } -} - -impl IntoIterator for Relation { - type Item = (K, HashSet); - type IntoIter = IntoIter>; - - fn into_iter(self) -> Self::IntoIter { - self.map.into_iter() - } -} diff --git a/vortex-array/src/expr/stats/mod.rs b/vortex-array/src/expr/stats/mod.rs index cd5da7811e9..57e11135ef3 100644 --- a/vortex-array/src/expr/stats/mod.rs +++ b/vortex-array/src/expr/stats/mod.rs @@ -28,6 +28,7 @@ use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; #[derive( Debug, @@ -186,17 +187,20 @@ impl Stat { .return_dtype(&EmptyOptions, data_type); } Self::Sum => { - return aggregate_fn::fns::sum::Sum.return_dtype(&EmptyOptions, data_type); + // Statistics follow NaN-skipping semantics; request it explicitly. + return aggregate_fn::fns::sum::Sum + .return_dtype(&NumericalAggregateOpts::skip_nans(), data_type); } }) } /// Return the built-in aggregate function corresponding to this statistic, if one exists. pub fn aggregate_fn(&self) -> Option { + // Statistics follow NaN-skipping semantics; request it explicitly rather than the default. Some(match self { - Self::Max => aggregate_fn::fns::max::Max.bind(EmptyOptions), - Self::Min => aggregate_fn::fns::min::Min.bind(EmptyOptions), - Self::Sum => aggregate_fn::fns::sum::Sum.bind(EmptyOptions), + Self::Max => aggregate_fn::fns::max::Max.bind(NumericalAggregateOpts::skip_nans()), + Self::Min => aggregate_fn::fns::min::Min.bind(NumericalAggregateOpts::skip_nans()), + Self::Sum => aggregate_fn::fns::sum::Sum.bind(NumericalAggregateOpts::skip_nans()), Self::NullCount => aggregate_fn::fns::null_count::NullCount.bind(EmptyOptions), Self::NaNCount => aggregate_fn::fns::nan_count::NanCount.bind(EmptyOptions), Self::UncompressedSizeInBytes => { @@ -208,9 +212,12 @@ impl Stat { } /// Return the statistic represented by `aggregate_fn`, if it has a legacy stat slot. + /// + /// Min/max/sum statistics skip NaN values, so NaN-including configurations of those + /// aggregates have no stat slot. pub fn from_aggregate_fn(aggregate_fn: &AggregateFnRef) -> Option { - if aggregate_fn.is::() { - return Some(Self::Sum); + if let Some(options) = aggregate_fn.as_opt::() { + return options.skip_nans.then_some(Self::Sum); } if aggregate_fn.is::() { return Some(Self::NaNCount); @@ -218,11 +225,11 @@ impl Stat { if aggregate_fn.is::() { return Some(Self::NullCount); } - if aggregate_fn.is::() { - return Some(Self::Min); + if let Some(options) = aggregate_fn.as_opt::() { + return options.skip_nans.then_some(Self::Min); } - if aggregate_fn.is::() { - return Some(Self::Max); + if let Some(options) = aggregate_fn.as_opt::() { + return options.skip_nans.then_some(Self::Max); } if aggregate_fn .is::() @@ -261,8 +268,8 @@ impl Display for Stat { mod test { use enum_iterator::all; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::expr::stats::Stat; @@ -270,7 +277,7 @@ mod test { fn min_of_nulls_is_not_panic() { let min = PrimitiveArray::from_option_iter::([None, None, None, None]) .statistics() - .compute_as::(Stat::Min, &mut LEGACY_SESSION.create_execution_ctx()); + .compute_as::(Stat::Min, &mut array_session().create_execution_ctx()); assert_eq!(min, None); } diff --git a/vortex-array/src/expr/transform/partition.rs b/vortex-array/src/expr/transform/partition.rs index ef852a6cd7f..92dca145d47 100644 --- a/vortex-array/src/expr/transform/partition.rs +++ b/vortex-array/src/expr/transform/partition.rs @@ -3,6 +3,7 @@ use std::fmt::Display; use std::fmt::Formatter; +use std::hash::Hash; use itertools::Itertools; use vortex_error::VortexExpect; @@ -49,11 +50,22 @@ where { // Annotate each expression with the annotations that any of its descendent expressions have. let annotations = descendent_annotations(&expr, annotate_fn); + partition_annotations(expr.clone(), scope, annotations) +} +pub fn partition_annotations( + expr: Expression, + scope: &DType, + annotations: Annotations, +) -> VortexResult> +where + A: Display + Clone + Eq + Hash, + FieldName: From, +{ // Now we split the original expression into sub-expressions based on the annotations, and // generate a root expression to re-assemble the results. - let mut splitter = StructFieldExpressionSplitter::::new(&annotations); - let root = expr.clone().rewrite(&mut splitter)?.value; + let mut splitter = StructFieldExpressionSplitter::::new(&annotations); + let root = expr.rewrite(&mut splitter)?.value; let mut partitions = Vec::with_capacity(splitter.sub_expressions.len()); let mut partition_annotations = Vec::with_capacity(splitter.sub_expressions.len()); diff --git a/vortex-array/src/extension/datetime/date.rs b/vortex-array/src/extension/datetime/date.rs index 8985c7fe8c5..9eda23dd2d1 100644 --- a/vortex-array/src/extension/datetime/date.rs +++ b/vortex-array/src/extension/datetime/date.rs @@ -9,6 +9,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_session::registry::CachedId; use crate::dtype::DType; use crate::dtype::Nullability; @@ -79,7 +80,8 @@ impl ExtVTable for Date { type NativeValue<'a> = DateValue; fn id(&self) -> ExtId { - ExtId::new("vortex.date") + static ID: CachedId = CachedId::new("vortex.date"); + *ID } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { diff --git a/vortex-array/src/extension/datetime/time.rs b/vortex-array/src/extension/datetime/time.rs index b8fe92b2eeb..88455a121b5 100644 --- a/vortex-array/src/extension/datetime/time.rs +++ b/vortex-array/src/extension/datetime/time.rs @@ -9,6 +9,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_session::registry::CachedId; use crate::dtype::DType; use crate::dtype::Nullability; @@ -80,7 +81,8 @@ impl ExtVTable for Time { type NativeValue<'a> = TimeValue; fn id(&self) -> ExtId { - ExtId::new("vortex.time") + static ID: CachedId = CachedId::new("vortex.time"); + *ID } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { diff --git a/vortex-array/src/extension/datetime/timestamp.rs b/vortex-array/src/extension/datetime/timestamp.rs index 135545468a4..d75820ee181 100644 --- a/vortex-array/src/extension/datetime/timestamp.rs +++ b/vortex-array/src/extension/datetime/timestamp.rs @@ -13,6 +13,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_session::registry::CachedId; use crate::dtype::DType; use crate::dtype::Nullability; @@ -114,7 +115,8 @@ impl ExtVTable for Timestamp { type NativeValue<'a> = TimestampValue<'a>; fn id(&self) -> ExtId { - ExtId::new("vortex.timestamp") + static ID: CachedId = CachedId::new("vortex.timestamp"); + *ID } // NOTE(ngates): unfortunately we're stuck with this hand-rolled serialization format for diff --git a/vortex-array/src/extension/tests/divisible_int.rs b/vortex-array/src/extension/tests/divisible_int.rs index 4ab08f7b9f5..2e42f9fa138 100644 --- a/vortex-array/src/extension/tests/divisible_int.rs +++ b/vortex-array/src/extension/tests/divisible_int.rs @@ -34,6 +34,7 @@ impl ExtVTable for DivisibleInt { type Metadata = Divisor; type NativeValue<'a> = u64; + #[expect(clippy::disallowed_methods, reason = "test-only id")] fn id(&self) -> ExtId { ExtId::new("test.divisible_int") } diff --git a/vortex-array/src/extension/uuid/vtable.rs b/vortex-array/src/extension/uuid/vtable.rs index cd4d4d971ae..c6378c4ead6 100644 --- a/vortex-array/src/extension/uuid/vtable.rs +++ b/vortex-array/src/extension/uuid/vtable.rs @@ -7,6 +7,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; +use vortex_session::registry::CachedId; use crate::dtype::DType; use crate::dtype::PType; @@ -27,7 +28,8 @@ impl ExtVTable for Uuid { type NativeValue<'a> = uuid::Uuid; fn id(&self) -> ExtId { - ExtId::new("vortex.uuid") + static ID: CachedId = CachedId::new("vortex.uuid"); + *ID } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { diff --git a/vortex-array/src/kernel.rs b/vortex-array/src/kernel.rs index f5b75471437..dff9769f139 100644 --- a/vortex-array/src/kernel.rs +++ b/vortex-array/src/kernel.rs @@ -10,12 +10,10 @@ //! run ends rather than decoding the entire array and slicing the result. //! //! Encodings declare their parent kernels by implementing [`ExecuteParentKernel`] and -//! registering them in a [`ParentKernelSet`]. Each kernel specifies which parent types it -//! handles via a [`Matcher`]. +//! registering them with the session's optimizer-kernel registry. Each kernel specifies which +//! parent types it handles via a [`Matcher`]. -use std::any::type_name; use std::fmt::Debug; -use std::marker::PhantomData; use vortex_error::VortexResult; @@ -25,57 +23,6 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::matcher::Matcher; -/// A collection of [`ExecuteParentKernel`]s registered for a specific child encoding. -/// -/// During execution, the scheduler iterates over each child's `ParentKernelSet` looking for -/// a kernel whose [`Matcher`] matches the parent array type. The first matching kernel that -/// returns `Some` wins. -pub struct ParentKernelSet { - kernels: &'static [&'static dyn DynParentKernel], -} - -impl ParentKernelSet { - /// Create a new parent kernel set with the given kernels. - /// - /// Use [`ParentKernelSet::lift`] to lift static rules into dynamic trait objects. - pub const fn new(kernels: &'static [&'static dyn DynParentKernel]) -> Self { - Self { kernels } - } - - /// Lift the given rule into a dynamic trait object. - pub const fn lift>( - kernel: &'static K, - ) -> &'static dyn DynParentKernel { - // Assert that self is zero-sized - const { - assert!( - !(size_of::() != 0), - "Rule must be zero-sized to be lifted" - ); - } - unsafe { &*(kernel as *const K as *const ParentKernelAdapter) } - } - - /// Evaluate the parent kernels on the given child and parent arrays. - pub fn execute( - &self, - child: ArrayView<'_, V>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - for kernel in self.kernels.iter() { - if !kernel.matches(parent) { - continue; - } - if let Some(reduced) = kernel.execute_parent(child, parent, child_idx, ctx)? { - return Ok(Some(reduced)); - } - } - Ok(None) - } -} - /// A kernel that allows a child encoding `V` to execute its parent array in a fused manner. /// /// This is the typed trait that encoding authors implement. The associated `Parent` type @@ -99,55 +46,3 @@ pub trait ExecuteParentKernel: Debug + Send + Sync + 'static { ctx: &mut ExecutionCtx, ) -> VortexResult>; } - -/// Type-erased version of [`ExecuteParentKernel`] used for dynamic dispatch within -/// [`ParentKernelSet`]. -pub trait DynParentKernel: Send + Sync { - /// Returns `true` if this kernel's parent [`Matcher`] matches the given parent array. - fn matches(&self, parent: &ArrayRef) -> bool; - - /// Attempt to execute the parent array fused with the child array. - fn execute_parent( - &self, - child: ArrayView<'_, V>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult>; -} - -/// Bridges a concrete [`ExecuteParentKernel`] to the type-erased [`DynParentKernel`] -/// trait. Created by [`ParentKernelSet::lift`]. -pub struct ParentKernelAdapter { - kernel: K, - _phantom: PhantomData, -} - -impl> Debug for ParentKernelAdapter { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ParentKernelAdapter") - .field("parent", &type_name::()) - .field("kernel", &self.kernel) - .finish() - } -} - -impl> DynParentKernel for ParentKernelAdapter { - fn matches(&self, parent: &ArrayRef) -> bool { - K::Parent::matches(parent) - } - - fn execute_parent( - &self, - child: ArrayView<'_, V>, - parent: &ArrayRef, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let Some(parent_view) = K::Parent::try_match(parent) else { - return Ok(None); - }; - self.kernel - .execute_parent(child, parent_view, child_idx, ctx) - } -} diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 1f3189eecbb..7496c527172 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -10,6 +10,72 @@ //! //! Every data type recognized by Vortex also has a canonical physical encoding format, which //! arrays can be [canonicalized](Canonical) into for ease of access in compute functions. +//! +//! # Core Handles +//! +//! [`ArrayRef`] is the erased, shared handle used by most public APIs. It carries the logical +//! [`DType`], row count, encoding id, children, buffers, and statistics for an +//! array tree. Use it when an API should accept any encoding. +//! +//! [`Array`] is the typed owned handle for a known encoding `V: VTable`. It wraps an +//! [`ArrayRef`] and dereferences to the encoding-specific `V::TypedArrayData`. +//! +//! [`ArrayView`] is the lightweight typed borrow handed to vtable methods. It exposes both the +//! shared [`ArrayRef`] metadata and the encoding-specific data without cloning the handle. +//! +//! [`ArrayParts`] is the construction boundary for typed arrays. It groups externally supplied +//! logical metadata and encoding data, then [`Array::try_from_parts`] validates that they agree. +//! +//! # Logical Types and Physical Encodings +//! +//! A [`DType`] describes the logical values an array may hold. It does not +//! describe the memory layout. For example, a `DType::Primitive(I32, Nullable)` can be stored as a +//! canonical [`PrimitiveArray`], a dictionary, a slice, or a +//! compressed external encoding. +//! +//! The [`Canonical`] enum names the default uncompressed encoding for each logical family. Execution +//! normally moves an array tree toward canonical form, but canonicalization is shallow: children of +//! canonical struct/list arrays may still be encoded. +//! +//! # Built-in, Lazy, and Experimental Arrays +//! +//! Built-in arrays live in [`arrays`]. Some are canonical (`PrimitiveArray`, `StructArray`, +//! `VarBinViewArray`); others are utility or lazy arrays such as [`ChunkedArray`], +//! [`ConstantArray`], [`FilterArray`], [`SliceArray`], and [`ScalarFnArray`]. +//! Lazy arrays defer work so compute kernels can operate on encoded data or prune children +//! before materialization. +//! +//! Experimental arrays are public because they are used inside Vortex, but their storage contracts +//! may still move. Prefer the higher-level constructors and accessors documented on each array +//! module rather than relying on child slot order. +//! +//! # Nulls and Scalars +//! +//! [`Validity`](crate::validity::Validity) separates nullness from values. It can be a cheap +//! constant state (`NonNullable`, `AllValid`, `AllInvalid`) or a boolean array that may itself be +//! encoded. [`Scalar`](crate::scalar::Scalar) is the single-value counterpart: it pairs a +//! [`DType`] with an optional [`ScalarValue`](crate::scalar::ScalarValue). +//! +//! # Extending Vortex +//! +//! New array encodings implement [`VTable`], usually through the local `array_slots!` and +//! `vtable!` patterns used by built-ins. The important extension contracts are: +//! +//! - [`VTable::validate`] checks that externally supplied dtype, length, slots, and data agree. +//! - [`VTable::execute`] returns an [`ExecutionResult`] that makes progress toward canonical form. +//! - [`OperationsVTable`] provides scalar access. +//! - [`ValidityVTable`] exposes validity only for nullable arrays. +//! +//! New logical extension dtypes implement [`ExtVTable`](crate::dtype::extension::ExtVTable) and +//! store values in an ordinary Vortex storage dtype. +//! +//! [`PrimitiveArray`]: crate::arrays::PrimitiveArray +//! [`DType`]: crate::dtype::DType +//! [`ChunkedArray`]: crate::arrays::ChunkedArray +//! [`ConstantArray`]: crate::arrays::ConstantArray +//! [`FilterArray`]: crate::arrays::FilterArray +//! [`SliceArray`]: crate::arrays::SliceArray +//! [`ScalarFnArray`]: crate::arrays::ScalarFnArray extern crate self as vortex_array; @@ -24,12 +90,19 @@ pub use mask_future::*; pub use metadata::*; pub use smallvec; pub use vortex_array_macros::array_slots; +use vortex_session::SessionExt; use vortex_session::VortexSession; use vortex_session::registry::Context; +use crate::aggregate_fn::session::AggregateFnSession; +use crate::arrow::ArrowSession; +use crate::dtype::session::DTypeSession; +use crate::memory::MemorySession; +use crate::optimizer::kernels::KernelSession; +use crate::scalar_fn::session::ScalarFnSession; use crate::session::ArraySession; +use crate::stats::session::StatsSession; -pub mod accessor; pub mod aggregate_fn; #[doc(hidden)] pub mod aliases; @@ -78,10 +151,39 @@ pub mod flatbuffers { pub use vortex_flatbuffers::array::*; } +/// Register vortex-array's built-in session-scoped kernels into the active +/// [`ArrayKernels`](crate::optimizer::kernels::ArrayKernels) registry. +/// +/// If the session contains a [`KernelSession`], this registers into its registry. Sessions that use +/// [`KernelSession::default`] already receive these built-in kernels. +pub fn initialize(session: &VortexSession) { + if session.get_opt::().is_some() { + arrays::initialize(session); + } +} + +/// Builds a fresh [`VortexSession`] registered with all of vortex-array's built-in session +/// variables: arrays, dtypes, scalar functions, stats, optimizer kernels, aggregate functions, +/// Arrow conversion, and memory. +/// +/// Each call returns an independent session (with its own registries), so callers may register +/// additional encodings or kernels into it without affecting any other session. This does not +/// register file, layout, or runtime state — those live in higher-level crates. +pub fn array_session() -> VortexSession { + VortexSession::empty() + .with::() + .with::() + .with::() + .with::() + .with::() + .with::() + .with::() + .with::() +} + // TODO(ngates): canonicalize doesn't currently take a session, therefore we cannot invoke execute // from the new array encodings to support back-compat for legacy encodings. So we hold a session // here... -pub static LEGACY_SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +pub static LEGACY_SESSION: LazyLock = LazyLock::new(array_session); pub type ArrayContext = Context; diff --git a/vortex-array/src/mask.rs b/vortex-array/src/mask.rs index ab035670e24..ee55abc9875 100644 --- a/vortex-array/src/mask.rs +++ b/vortex-array/src/mask.rs @@ -1,8 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::ops::BitAnd; - use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_mask::Mask; @@ -12,19 +10,23 @@ use crate::Executable; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::BoolArray; -use crate::arrays::Constant; use crate::columnar::Columnar; use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::validity::Validity; impl Executable for Mask { + /// Executes a boolean array into a [`Mask`]. + /// + /// The array must have a non-nullable boolean dtype. To execute a nullable boolean array, + /// coercing null elements to `false`, first call + /// [`ArrayRef::fill_null(false)`](crate::builtins::ArrayBuiltins::fill_null). fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - if !matches!(array.dtype(), DType::Bool(_)) { - vortex_bail!("Mask array must have boolean dtype, not {}", array.dtype()); - } - - if let Some(constant) = array.as_opt::() { - let mask_value = constant.scalar().as_bool().value().unwrap_or(false); - return Ok(Mask::new(array.len(), mask_value)); + if !matches!(array.dtype(), DType::Bool(Nullability::NonNullable)) { + vortex_bail!( + "Mask array must have boolean(NonNullable) dtype, not {}", + array.dtype() + ); } let array_len = array.len(); @@ -34,16 +36,149 @@ impl Executable for Mask { } Columnar::Canonical(a) => { let bool = a.into_array().execute::(ctx)?; - let mask = bool - .as_ref() - .validity()? - .execute_mask(bool.as_ref().len(), ctx)?; - let bits = bool.into_bit_buffer(); - // To handle nullable boolean arrays, we treat nulls as false in the mask. - // TODO(ngates): is this correct? Feels like we should just force the caller to - // pass non-nullable boolean arrays. - mask.bitand(&Mask::from(bits)) + Mask::from(bool.into_bit_buffer()) + } + }) + } +} + +/// An adapter that coerces null elements of a boolean array to `false` before executing it into a +/// [`Mask`]. Created by [`ArrayRef::null_as_false`]. +/// +/// Use for filter and pruning predicates over nullable data, where SQL semantics treat `NULL` as +/// not matching. +/// +/// Prefer `array.null_as_false().execute(ctx)` over `array.fill_null(false)?.execute::(ctx)`: +/// `fill_null` on a lazy `ScalarFn` array (e.g. the result of `apply()`) is currently +/// slow because its `validity()` executes the predicate expression. +pub struct NullAsFalse(ArrayRef); + +impl ArrayRef { + /// Returns an adapter that treats null elements of this boolean array as `false` when executed + /// into a [`Mask`]. See [`NullAsFalse`]. + pub fn null_as_false(self) -> NullAsFalse { + NullAsFalse(self) + } +} + +impl NullAsFalse { + /// Executes the boolean array into a [`Mask`], coercing null elements to `false`. + /// + /// Canonicalizes the (possibly lazy) array exactly once and folds validity into the value bits + /// with a single `AND` that reuses the value buffer when it is uniquely owned. + pub fn execute(self, ctx: &mut ExecutionCtx) -> VortexResult { + let array = self.0; + if !matches!(array.dtype(), DType::Bool(_)) { + vortex_bail!("Mask array must have boolean dtype, not {}", array.dtype()); + } + // Non-nullable input needs no coercion; defer to the strict `Mask` execution. + if !array.dtype().is_nullable() { + return array.execute::(ctx); + } + + let len = array.len(); + Ok(match array.execute::(ctx)? { + Columnar::Constant(c) => Mask::new(len, c.scalar().as_bool().value().unwrap_or(false)), + Columnar::Canonical(c) => { + let bool = c.into_array().execute::(ctx)?; + match bool.as_ref().validity()? { + Validity::NonNullable | Validity::AllValid => { + Mask::from_buffer(bool.into_bit_buffer()) + } + Validity::AllInvalid => Mask::new_false(len), + Validity::Array(v) => { + let validity_bits = v.execute::(ctx)?.into_bit_buffer(); + Mask::from_buffer(bool.into_bit_buffer() & &validity_bits) + } + } } }) } } + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + use vortex_mask::Mask; + + use crate::ExecutionCtx; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::BoolArray; + use crate::arrays::ConstantArray; + use crate::builtins::ArrayBuiltins; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::scalar::Scalar; + + fn ctx() -> ExecutionCtx { + array_session().create_execution_ctx() + } + + #[test] + fn mask_non_nullable() -> VortexResult<()> { + let array = BoolArray::from_iter([true, false, true]).into_array(); + let mask = array.execute::(&mut ctx())?; + assert_eq!(mask, Mask::from_iter([true, false, true])); + Ok(()) + } + + #[test] + fn mask_rejects_nullable() { + let array = BoolArray::from_iter([Some(true), None]).into_array(); + assert!(array.execute::(&mut ctx()).is_err()); + } + + #[test] + fn fill_null_then_mask_coerces_nulls() -> VortexResult<()> { + let array = BoolArray::from_iter([Some(true), None, Some(false), None]).into_array(); + let mask = array.fill_null(false)?.execute::(&mut ctx())?; + assert_eq!(mask, Mask::from_iter([true, false, false, false])); + Ok(()) + } + + #[test] + fn fill_null_then_mask_null_constant() -> VortexResult<()> { + let array = + ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 4).into_array(); + let mask = array.fill_null(false)?.execute::(&mut ctx())?; + assert_eq!(mask, Mask::new_false(4)); + Ok(()) + } + + #[test] + fn null_as_false_non_nullable() -> VortexResult<()> { + let array = BoolArray::from_iter([true, false, true]).into_array(); + let mask = array.null_as_false().execute(&mut ctx())?; + assert_eq!(mask, Mask::from_iter([true, false, true])); + Ok(()) + } + + #[test] + fn null_as_false_treats_null_as_false() -> VortexResult<()> { + let array = BoolArray::from_iter([Some(true), None, Some(false), None]).into_array(); + let mask = array.null_as_false().execute(&mut ctx())?; + assert_eq!(mask, Mask::from_iter([true, false, false, false])); + Ok(()) + } + + #[test] + fn null_as_false_null_constant() -> VortexResult<()> { + let array = + ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 4).into_array(); + let mask = array.null_as_false().execute(&mut ctx())?; + assert_eq!(mask, Mask::new_false(4)); + Ok(()) + } + + #[test] + fn null_as_false_matches_fill_null_then_mask() -> VortexResult<()> { + let array = + BoolArray::from_iter([Some(true), None, Some(false), Some(true), None]).into_array(); + let via_fill_null = array.fill_null(false)?.execute::(&mut ctx())?; + let via_coerce = array.null_as_false().execute(&mut ctx())?; + assert_eq!(via_coerce, via_fill_null); + Ok(()) + } +} diff --git a/vortex-array/src/memory.rs b/vortex-array/src/memory.rs index d3669e49f93..b52b3442136 100644 --- a/vortex-array/src/memory.rs +++ b/vortex-array/src/memory.rs @@ -16,10 +16,10 @@ use vortex_buffer::ByteBufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; -use vortex_session::Ref; -use vortex_session::RefMut; use vortex_session::SessionExt; +use vortex_session::SessionGuard; use vortex_session::SessionVar; +use vortex_session::VortexSession; /// Mutable host buffer contract used by [`WritableHostBuffer`]. pub trait HostBufferMut: Send + 'static { @@ -173,7 +173,7 @@ pub trait HostAllocatorExt: HostAllocator { impl HostAllocatorExt for A {} /// Session-scoped memory configuration for Vortex arrays. -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct MemorySession { allocator: HostAllocatorRef, } @@ -214,7 +214,7 @@ impl SessionVar for MemorySession { /// Extension trait for accessing session-scoped memory configuration. pub trait MemorySessionExt: SessionExt { /// Returns the memory session for this execution/session context. - fn memory(&self) -> Ref<'_, MemorySession> { + fn memory(&self) -> SessionGuard<'_, MemorySession> { self.get::() } @@ -223,9 +223,12 @@ pub trait MemorySessionExt: SessionExt { self.memory().allocator() } - /// Returns mutable access to the memory session. - fn memory_mut(&self) -> RefMut<'_, MemorySession> { - self.get_mut::() + /// Configures the session to use `allocator` as its host allocator, mutating it in place and + /// returning it for chaining. + fn with_allocator(self, allocator: HostAllocatorRef) -> VortexSession { + let session = self.session(); + session.get_mut::().set_allocator(allocator); + session } } diff --git a/vortex-array/src/normalize.rs b/vortex-array/src/normalize.rs index 24a8f50d8a1..99a08c1ec1d 100644 --- a/vortex-array/src/normalize.rs +++ b/vortex-array/src/normalize.rs @@ -84,7 +84,9 @@ impl ArrayRef { } if any_slot_changed { - normalized = normalized.with_slots(normalized_slots)?; + // SAFETY: normalization only rewrites child slots to logically equivalent allowed + // encodings, preserving parent logical values and statistics. + normalized = unsafe { normalized.with_slots(normalized_slots) }?; } Ok(normalized) @@ -98,14 +100,13 @@ impl ArrayRef { #[cfg(test)] mod tests { use vortex_error::VortexResult; - use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; use super::NormalizeOptions; use super::Operation; use crate::ArrayRef; - use crate::ExecutionCtx; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::array::VTable; use crate::arrays::Dict; use crate::arrays::DictArray; @@ -128,7 +129,7 @@ mod tests { )? .into_array(); let allowed = HashSet::from_iter([array.encoding_id(), field.encoding_id()]); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); let normalized = array.clone().normalize(&mut NormalizeOptions { allowed: &allowed, @@ -173,7 +174,7 @@ mod tests { )? .into_array(); let allowed = HashSet::from_iter([array.encoding_id(), unchanged.encoding_id()]); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); let normalized = array.clone().normalize(&mut NormalizeOptions { allowed: &allowed, @@ -192,7 +193,11 @@ mod tests { &original_children[1], &normalized_children[1] )); - assert_arrays_eq!(normalized_children[1], PrimitiveArray::from_iter(12i32..16)); + assert_arrays_eq!( + normalized_children[1], + PrimitiveArray::from_iter(12i32..16), + &mut ctx + ); Ok(()) } @@ -208,7 +213,7 @@ mod tests { assert_eq!(sliced.encoding_id(), Slice.id()); let allowed = HashSet::from_iter([Dict.id(), Primitive.id()]); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = crate::array_session().create_execution_ctx(); let normalized = sliced.normalize(&mut NormalizeOptions { allowed: &allowed, @@ -224,7 +229,8 @@ mod tests { let normalized_canonical = normalized.to_canonical()?; assert_arrays_eq!( normalized_canonical, - PrimitiveArray::from_iter(vec![20i32, 10, 20]) + PrimitiveArray::from_iter(vec![20i32, 10, 20]), + &mut ctx ); Ok(()) diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index e16e59abd17..e419e927ca1 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/vortex-array/src/optimizer/kernels.rs @@ -6,10 +6,9 @@ //! [`ArrayKernels`] stores function pointers that participate in array optimization and execution //! without adding rules or kernels to an encoding vtable. The optimizer consults it for //! parent-reduce rewrites before the child encoding's static `PARENT_RULES`, and the executor -//! consults it for parent execution before the child encoding's static parent kernels. A -//! registered function can therefore add support for an extension encoding or take precedence over -//! a built-in rule or kernel. When several functions are registered for the same key and kind, -//! they are tried in registration order until one applies. +//! consults it for parent execution. A registered function can therefore add support for an +//! extension encoding or take precedence over a built-in rule. When several functions are +//! registered for the same key and kind, they are tried in registration order until one applies. //! //! Kernel entries are addressed by `(outer_id, child_id)`. For parent-reduce and execute-parent //! kernels, `outer_id` is the id returned by the parent array's `encoding_id()` and `child_id` is @@ -19,20 +18,24 @@ //! Because registered functions have different signatures for each kernel kind, the registry //! maintains one storage map per function type rather than a single type-erased map. //! -//! Sessions created by the top-level `vortex` crate install the default registry. Other sessions -//! can add it with [`VortexSession::with`](vortex_session::VortexSession::with) or rely on -//! [`ArrayKernelsExt::kernels`] to insert the default value. +//! [`KernelSession`] is the session variable that owns this registry. Its [`Default`] +//! implementation installs vortex-array's built-in parent-reduce and execute-parent kernels, so a +//! session built with [`KernelSession`] participates in the same optimizations and fused execution +//! as the built-in encodings. use std::any::Any; use std::borrow::Borrow; +use std::fmt::Debug; use std::hash::BuildHasher; +use std::ops::Deref; use std::sync::Arc; use std::sync::LazyLock; use vortex_error::VortexResult; -use vortex_session::Ref; use vortex_session::SessionExt; +use vortex_session::SessionGuard; use vortex_session::SessionVar; +use vortex_session::VortexSession; use vortex_session::registry::Id; use vortex_utils::aliases::DefaultHashBuilder; use vortex_utils::aliases::hash_map::HashMap; @@ -42,8 +45,9 @@ use crate::ExecutionCtx; use crate::arc_swap_map::ArcSwapMap; use crate::array::VTable; use crate::arrays::Struct; -use crate::arrays::struct_::compute::cast::struct_cast_execute_parent; use crate::arrays::struct_::compute::rules::struct_cast_reduce_parent; +use crate::kernel::ExecuteParentKernel; +use crate::matcher::Matcher; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; @@ -92,9 +96,69 @@ pub type ExecuteParentFn = fn( ctx: &mut ExecutionCtx, ) -> VortexResult>; +/// Type-erased execute-parent kernel stored in the session registry. +pub trait DynExecuteParentKernel: Debug + Send + Sync + 'static { + /// Attempt to execute the parent array fused with the child array. + fn execute_parent( + &self, + child: &ArrayRef, + parent: &ArrayRef, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; +} + +pub(crate) type ExecuteParentKernelRef = Arc; + +pub(crate) type ParentExecutionKernels = HashMap>; + +#[derive(Debug)] +struct ExecuteParentFnKernel(ExecuteParentFn); + +impl DynExecuteParentKernel for ExecuteParentFnKernel { + fn execute_parent( + &self, + child: &ArrayRef, + parent: &ArrayRef, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + self.0(child, parent, child_idx, ctx) + } +} + +#[derive(Debug)] +struct RegisteredExecuteParentKernel { + _child: V, + kernel: K, +} + +impl DynExecuteParentKernel for RegisteredExecuteParentKernel +where + V: VTable, + K: ExecuteParentKernel, +{ + fn execute_parent( + &self, + child: &ArrayRef, + parent: &ArrayRef, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(child) = child.as_opt::() else { + return Ok(None); + }; + let Some(parent) = K::Parent::try_match(parent) else { + return Ok(None); + }; + + self.kernel.execute_parent(child, parent, child_idx, ctx) + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)] #[repr(transparent)] -struct ExecuteParentFnId(u64); +pub(crate) struct ExecuteParentFnId(u64); impl From for ExecuteParentFnId { fn from(id: u64) -> Self { @@ -112,17 +176,16 @@ impl Borrow for ExecuteParentFnId { /// /// Each kernel kind has its own storage map, keyed by `(outer_id, child_id)`. Registering /// functions for an existing key appends them to that key's ordered list. -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct ArrayKernels { reduce_parent: ArcSwapMap>, - execute_parent: ArcSwapMap>, + execute_parent: ArcSwapMap>, } impl Default for ArrayKernels { fn default() -> ArrayKernels { let this = Self::empty(); this.register_builtin_reduce_parent(); - this.register_builtin_execute_parent(); this } } @@ -144,14 +207,6 @@ impl ArrayKernels { ); } - fn register_builtin_execute_parent(&self) { - self.register_execute_parent( - Cast.id(), - Struct.id(), - &[struct_cast_execute_parent as ExecuteParentFn], - ); - } - /// Register [`ReduceParentFn`]s for `(parent, child)`. /// /// The optimizer invokes these functions in registration order when it sees a parent with @@ -179,63 +234,116 @@ impl ArrayKernels { /// /// The executor invokes these functions in registration order when it sees a parent with /// encoding id `parent` holding a child with encoding id `child` during a parent execution - /// step, before trying the child encoding's static parent kernels. + /// step. /// /// If functions have already been registered for the same pair, these functions are appended /// after them. pub fn register_execute_parent(&self, parent: Id, child: Id, fns: &[ExecuteParentFn]) { + let kernels: Vec = fns + .iter() + .map(|f| Arc::new(ExecuteParentFnKernel(*f)) as ExecuteParentKernelRef) + .collect(); self.execute_parent - .extend(hash_fn_id(parent, child).into(), fns); + .extend(hash_fn_id(parent, child).into(), kernels.as_slice()); } - /// Look up the [`ExecuteParentFn`]s registered for `(parent, child)`. + /// Register a typed [`ExecuteParentKernel`] for `(parent, child.id())`. /// - /// Returns an owned [`Arc`] so the session-variable borrow can be dropped before invoking the - /// functions. - pub fn find_execute_parent(&self, parent: Id, child: Id) -> Option> { - self.execute_parent.get(&hash_fn_id(parent, child)) + /// The executor invokes registered kernels in registration order before falling through to + /// later registered kernels for the same key. `parent` is usually the parent array's encoding + /// id. For `ScalarFnArray`, it is the scalar function id, for example `Cast.id()`. + /// + /// If kernels have already been registered for the same pair, this kernel is appended after + /// them; registering for an existing key cannot override built-in kernels installed earlier. + pub fn register_execute_parent_kernel(&self, parent: Id, child: V, kernel: K) + where + V: VTable, + K: ExecuteParentKernel, + { + let child_id = child.id(); + self.execute_parent.push( + hash_fn_id(parent, child_id).into(), + Arc::new(RegisteredExecuteParentKernel { + _child: child, + kernel, + }) as ExecuteParentKernelRef, + ); } - /// Capture an owned, cheaply-cloneable [`KernelSnapshot`] of the currently-registered - /// execute-parent kernels. - /// - /// The [`ArcSwap`] is loaded once into an [`Arc`], so the snapshot is an `Arc` clone (no map - /// copy) that outlives the session-variable borrow. Registrations made after the snapshot is - /// taken are not visible through it. - pub(crate) fn snapshot(&self) -> KernelSnapshot { - KernelSnapshot { - execute_parent: self.execute_parent.load_full(), - } + /// Returns true when one or more execute-parent kernels are registered for `(parent, child)`. + pub fn has_execute_parent(&self, parent: Id, child: Id) -> bool { + self.execute_parent + .get(&hash_fn_id(parent, child)) + .is_some() + } + + /// Return the currently published execute-parent kernel snapshot. + pub(crate) fn execute_parent_snapshot(&self) -> Arc { + self.execute_parent.snapshot() } } -/// An owned, point-in-time view of the execute-parent kernels registered on an [`ArrayKernels`] -/// registry. +fn hash_fn_id(parent: Id, child: Id) -> u64 { + FN_HASHER.hash_one((parent, child)) +} + +/// Return the registry key for execute-parent kernels registered for `(parent, child)`. +pub(crate) fn execute_parent_key(parent: Id, child: Id) -> u64 { + hash_fn_id(parent, child) +} + +/// Session-scoped holder for the optimizer kernel registry. /// -/// Holding the registry map directly (rather than re-probing the session per array node) lets the -/// executor resolve [`ArrayKernels`] once per execution context. Cloning is one [`Arc`] clone. -#[derive(Debug, Clone)] -pub(crate) struct KernelSnapshot { - execute_parent: Arc>>, +/// `KernelSession` is the session variable that owns an [`ArrayKernels`] registry. Its [`Default`] +/// implementation installs vortex-array's built-in parent-reduce and execute-parent kernels, +/// mirroring how [`ScalarFnSession`](crate::scalar_fn::session::ScalarFnSession) and the other +/// session variables register their built-ins. +#[derive(Clone, Debug)] +pub struct KernelSession { + kernels: ArrayKernels, } -impl KernelSnapshot { - /// Look up the [`ExecuteParentFn`]s registered for `(parent, child)`. - pub(crate) fn find_execute_parent( - &self, - parent: Id, - child: Id, - ) -> Option> { - let id = hash_fn_id(parent, child); - self.execute_parent.get(&id).cloned() +impl KernelSession { + /// Create a [`KernelSession`] with an empty kernel registry. + pub fn empty() -> Self { + Self { + kernels: ArrayKernels::empty(), + } + } + + /// Returns the [`ArrayKernels`] registry held by this session. + pub fn kernels(&self) -> &ArrayKernels { + &self.kernels } } -fn hash_fn_id(parent: Id, child: Id) -> u64 { - FN_HASHER.hash_one((parent, child)) +/// Derefs to the held [`ArrayKernels`] registry, so a [`KernelSession`] (or a +/// [`SessionGuard`](SessionGuard) read from a session) can be used wherever an +/// `&ArrayKernels` is expected. +impl Deref for KernelSession { + type Target = ArrayKernels; + + fn deref(&self) -> &ArrayKernels { + &self.kernels + } } -impl SessionVar for ArrayKernels { +impl Default for KernelSession { + fn default() -> Self { + // `ArrayKernels::default` installs the built-in parent-reduce kernels. The execute-parent + // kernels are registered by the per-encoding `initialize` functions, which operate on a + // session. `KernelSession` clones share their registry storage, so kernels registered into + // the temporary session land in `this.kernels`. + let this = Self { + kernels: ArrayKernels::default(), + }; + let session = VortexSession::empty().with_some(this.clone()); + crate::arrays::initialize(&session); + this + } +} + +impl SessionVar for KernelSession { fn as_any(&self) -> &dyn Any { self } @@ -245,14 +353,58 @@ impl SessionVar for ArrayKernels { } } -/// Extension trait for accessing optimizer kernels from a -/// [`VortexSession`](vortex_session::VortexSession). +/// Extension trait for accessing the optimizer kernel registry from a [`VortexSession`]. pub trait ArrayKernelsExt: SessionExt { - /// Returns the [`ArrayKernels`] session variable, inserting a default-constructed one if - /// none has been registered on the session yet. - fn kernels(&self) -> Ref<'_, ArrayKernels> { - self.get::() + /// Returns the session's [`KernelSession`], inserting a default one (with the built-in + /// kernels) if it does not exist. + /// + /// The returned [`SessionGuard`] borrows the session snapshot it was read from (so the registry + /// stays alive even if the session is concurrently mutated) and derefs through [`KernelSession`] + /// to the [`ArrayKernels`] registry, so it can be used wherever an `&ArrayKernels` is expected. + /// The registry shares its storage with the session, so kernels registered through it remain + /// visible to the session. + fn kernels(&self) -> SessionGuard<'_, KernelSession> { + self.get::() } } impl ArrayKernelsExt for S {} + +#[cfg(test)] +mod tests { + use vortex_session::VortexSession; + + use super::ArrayKernelsExt; + use super::KernelSession; + use crate::ArrayVTable; + use crate::arrays::Bool; + use crate::scalar_fn::ScalarFnVTable; + use crate::scalar_fn::fns::binary::Binary; + + #[test] + fn kernel_session_default_registers_builtin_kernels() { + let session = VortexSession::empty().with::(); + + assert!(session.kernels().has_execute_parent(Binary.id(), Bool.id())); + } + + #[test] + fn initialize_registers_builtin_kernels_into_empty_kernel_session() { + let session = VortexSession::empty().with_some(KernelSession::empty()); + + assert!(!session.kernels().has_execute_parent(Binary.id(), Bool.id())); + + crate::initialize(&session); + + assert!(session.kernels().has_execute_parent(Binary.id(), Bool.id())); + } + + #[test] + fn kernels_inserts_default_kernel_session() { + let session = VortexSession::empty(); + + // `kernels()` uses `get`, so it inserts a default `KernelSession` (with the built-in + // kernels) rather than returning `None`. + assert!(session.kernels().has_execute_parent(Binary.id(), Bool.id())); + } +} diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs index d6e93ca0561..9b51a468ee8 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -10,19 +10,19 @@ //! There are three public entry points on [`ArrayOptimizer`]: //! //! - [`ArrayOptimizer::optimize`] uses only static rules registered on encoding vtables. -//! - [`ArrayOptimizer::optimize_ctx`] also consults session-scoped [`ArrayKernels`] before -//! static parent-reduce rules, so this is the entry point used by execution. +//! - [`ArrayOptimizer::optimize_ctx`] also consults the session's active +//! [`kernels::ArrayKernels`] registry before static parent-reduce rules, so this is the entry +//! point used by execution. //! - [`ArrayOptimizer::optimize_recursive`] applies the session-aware optimizer to the root and //! every descendant. use smallvec::SmallVec; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_session::SessionExt; use vortex_session::VortexSession; use crate::ArrayRef; -use crate::optimizer::kernels::ArrayKernels; +use crate::optimizer::kernels::ArrayKernelsExt; pub mod kernels; pub mod rules; @@ -32,14 +32,16 @@ pub trait ArrayOptimizer { /// Optimize the root array node by running reduce and reduce_parent rules to fixpoint. /// /// This uses only static rules registered on encoding vtables. Use [`Self::optimize_ctx`] - /// when session-registered [`ArrayKernels`] should participate. + /// when a session-scoped [`kernels::ArrayKernels`] registry should participate. fn optimize(&self) -> VortexResult; - /// Optimize the root array node using static rules and any [`ArrayKernels`] on `session`. + /// Optimize the root array node using static rules and the active + /// [`kernels::ArrayKernels`] registry on `session`, if any. /// /// Session kernels are checked for each `(parent_encoding_id, child_encoding_id)` pair before - /// the child's static `PARENT_RULES`. If `session` does not contain [`ArrayKernels`], this - /// behaves like [`Self::optimize`]. + /// the child's static `PARENT_RULES`. The registry comes from the [`kernels::KernelSession`] on + /// `session`, if any. If `session` does not contain a [`kernels::KernelSession`], this behaves + /// like [`Self::optimize`]. fn optimize_ctx(&self, session: &VortexSession) -> VortexResult; /// Optimize the entire array tree recursively (root and all descendants). @@ -69,7 +71,7 @@ fn try_optimize( ) -> VortexResult> { let mut current_array = array.clone(); let mut any_optimizations = false; - let array_ref = session.and_then(|s| s.get_opt::()); + let array_ref = session.map(|s| s.kernels()); // Apply reduction rules to the current array until no more rules apply. let mut loop_counter = 0; @@ -154,7 +156,9 @@ fn try_optimize_recursive( } if any_slot_optimized { - current_array = current_array.with_slots(new_slots)?; + // SAFETY: optimizer rules only replace child slots with logically equivalent arrays, so + // parent logical values and statistics remain valid. + current_array = unsafe { current_array.with_slots(new_slots) }?; any_optimizations = true; } diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index f4ab1cb0f9a..2df53a8e7f9 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -24,9 +24,8 @@ use crate::ArraySlots; use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; -#[expect(deprecated)] -use crate::ToCanonical as _; use crate::VortexSessionExecute; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builtins::ArrayBuiltins; @@ -441,20 +440,6 @@ impl Patches { )) } - pub fn cast_values(self, values_dtype: &DType) -> VortexResult { - // SAFETY: casting does not affect the relationship between the indices and values - unsafe { - Ok(Self::new_unchecked( - self.array_len, - self.offset, - self.indices, - self.values.cast(values_dtype.clone())?, - self.chunk_offsets, - self.offset_within_chunk, - )) - } - } - /// Get the patched value at a given index if it exists. pub fn get_patched(&self, index: usize) -> VortexResult> { self.search_index(index)? @@ -497,9 +482,7 @@ impl Patches { /// [`SearchResult::Found`] with the position if needle exists, or [`SearchResult::NotFound`] /// with the insertion point if not found. fn search_index_binary_search(indices: &ArrayRef, needle: usize) -> VortexResult { - if indices.is_canonical() { - #[expect(deprecated)] - let primitive = indices.to_primitive(); + if let Some(primitive) = indices.as_opt::() { match_each_integer_ptype!(primitive.ptype(), |T| { let Ok(needle) = T::try_from(needle) else { // If the needle is not of type T, then it cannot possibly be in this array. @@ -696,6 +679,10 @@ impl Patches { /// Mask the patches, REMOVING the patches where the mask is true. /// Unlike filter, this preserves the patch indices. /// Unlike mask on a single array, this does not set masked values to null. + /// + /// Masking an array always yields a nullable result (see [`crate::scalar_fn::fns::mask`]), + /// so the surviving patch values are widened to nullable to match the masked parent. Callers + /// therefore receive patches with the correct dtype and do not need to re-cast them. // TODO(joe): make this lazy and remove the ctx. pub fn mask(&self, mask: &Mask, ctx: &mut ExecutionCtx) -> VortexResult> { if mask.len() != self.array_len { @@ -708,7 +695,7 @@ impl Patches { let filter_mask = match mask.bit_buffer() { AllOr::All => return Ok(None), - AllOr::None => return Ok(Some(self.clone())), + AllOr::None => return self.clone().into_nullable_values().map(Some), AllOr::Some(masked) => { let patch_indices = self.indices().clone().execute::(ctx)?; match_each_unsigned_integer_ptype!(patch_indices.ptype(), |P| { @@ -730,7 +717,7 @@ impl Patches { let filtered_indices = self.indices.filter(filter_mask.clone())?; let filtered_values = self.values.filter(filter_mask)?; - Ok(Some(Self { + Self { array_len: self.array_len, offset: self.offset, indices: filtered_indices, @@ -738,7 +725,21 @@ impl Patches { // TODO(0ax1): Chunk offsets are invalid after a filter is applied. chunk_offsets: None, offset_within_chunk: self.offset_within_chunk, - })) + } + .into_nullable_values() + .map(Some) + } + + /// Widen the patch values to their nullable dtype, leaving indices and offsets untouched. + /// + /// The values stay logically unchanged (all currently-valid entries remain valid); only the + /// dtype's nullability flag is set. Used by [`Self::mask`], whose result must be nullable. + fn into_nullable_values(self) -> VortexResult { + if self.values.dtype().is_nullable() { + return Ok(self); + } + let nullable = self.values.dtype().as_nullable(); + self.map_values(|values| values.cast(nullable)) } /// Slice the patches by a range of the patched array. @@ -1213,10 +1214,10 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; - use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; use crate::patches::Patches; use crate::patches::PrimitiveArray; @@ -1225,6 +1226,7 @@ mod test { #[test] fn test_filter() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 100, 0, @@ -1235,19 +1237,25 @@ mod test { .unwrap(); let filtered = patches - .filter( - &Mask::from_indices(100, vec![10, 20, 30]), - &mut LEGACY_SESSION.create_execution_ctx(), - ) + .filter(&Mask::from_indices(100, vec![10, 20, 30]), &mut ctx) .unwrap() .unwrap(); - assert_arrays_eq!(filtered.indices(), PrimitiveArray::from_iter([0u64, 1])); - assert_arrays_eq!(filtered.values(), PrimitiveArray::from_iter([100i32, 200])); + assert_arrays_eq!( + filtered.indices(), + PrimitiveArray::from_iter([0u64, 1]), + &mut ctx + ); + assert_arrays_eq!( + filtered.values(), + PrimitiveArray::from_iter([100i32, 200]), + &mut ctx + ); } #[test] fn take_with_nulls() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1261,7 +1269,7 @@ mod test { .take( &PrimitiveArray::new(buffer![9, 0], Validity::from_iter(vec![true, false])) .into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .unwrap() .unwrap(); @@ -1272,18 +1280,20 @@ mod test { assert_eq!(taken.array_len(), 2); assert_arrays_eq!( primitive_values, - PrimitiveArray::from_option_iter([Some(44i32)]) + PrimitiveArray::from_option_iter([Some(44i32)]), + &mut ctx + ); + assert_arrays_eq!( + primitive_indices, + PrimitiveArray::from_iter([0u64]), + &mut ctx ); - assert_arrays_eq!(primitive_indices, PrimitiveArray::from_iter([0u64])); assert_eq!( primitive_values .as_ref() .validity() .unwrap() - .execute_mask( - primitive_values.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx() - ) + .execute_mask(primitive_values.as_ref().len(), &mut ctx) .unwrap(), Mask::from_iter(vec![true]) ); @@ -1291,6 +1301,7 @@ mod test { #[test] fn take_search_with_nulls_chunked() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1304,7 +1315,7 @@ mod test { .take_search( PrimitiveArray::new(buffer![9, 0], Validity::from_iter([true, false])), true, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .unwrap() .unwrap(); @@ -1314,19 +1325,21 @@ mod test { assert_eq!(taken.array_len(), 2); assert_arrays_eq!( primitive_values, - PrimitiveArray::from_option_iter([Some(44i32), None]) + PrimitiveArray::from_option_iter([Some(44i32), None]), + &mut ctx + ); + assert_arrays_eq!( + taken.indices(), + PrimitiveArray::from_iter([0u64, 1]), + &mut ctx ); - assert_arrays_eq!(taken.indices(), PrimitiveArray::from_iter([0u64, 1])); assert_eq!( primitive_values .as_ref() .validity() .unwrap() - .execute_mask( - primitive_values.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx() - ) + .execute_mask(primitive_values.as_ref().len(), &mut ctx) .unwrap(), Mask::from_iter([true, false]) ); @@ -1334,6 +1347,7 @@ mod test { #[test] fn take_search_chunked_multiple_chunks() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 2048, 0, @@ -1347,7 +1361,7 @@ mod test { .take_search( PrimitiveArray::new(buffer![500, 1200, 999], Validity::AllValid), true, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .unwrap() .unwrap(); @@ -1355,12 +1369,14 @@ mod test { assert_eq!(taken.array_len(), 3); assert_arrays_eq!( taken.values(), - PrimitiveArray::from_option_iter([Some(20i32), Some(30)]) + PrimitiveArray::from_option_iter([Some(20i32), Some(30)]), + &mut ctx ); } #[test] fn take_search_chunked_indices_with_no_patches() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1374,7 +1390,7 @@ mod test { .take_search( PrimitiveArray::new(buffer![3, 4, 5], Validity::AllValid), true, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .unwrap(); @@ -1383,6 +1399,7 @@ mod test { #[test] fn take_search_chunked_interleaved() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 30, 0, @@ -1396,7 +1413,7 @@ mod test { .take_search( PrimitiveArray::new(buffer![10, 15, 20, 99], Validity::AllValid), true, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .unwrap() .unwrap(); @@ -1404,12 +1421,14 @@ mod test { assert_eq!(taken.array_len(), 4); assert_arrays_eq!( taken.values(), - PrimitiveArray::from_option_iter([Some(200i32), Some(300)]) + PrimitiveArray::from_option_iter([Some(200i32), Some(300)]), + &mut ctx ); } #[test] fn test_take_search_multiple_chunk_offsets() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 1500, 0, @@ -1423,7 +1442,7 @@ mod test { .take_search( PrimitiveArray::new(BufferMut::from_iter(0..1500u64), Validity::AllValid), false, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .unwrap() .unwrap(); @@ -1433,6 +1452,7 @@ mod test { #[test] fn test_slice() { + let mut ctx = array_session().create_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); @@ -1440,11 +1460,16 @@ mod test { let sliced = patches.slice(15..100).unwrap().unwrap(); assert_eq!(sliced.array_len(), 100 - 15); - assert_arrays_eq!(sliced.values(), PrimitiveArray::from_iter([13531u32])); + assert_arrays_eq!( + sliced.values(), + PrimitiveArray::from_iter([13531u32]), + &mut ctx + ); } #[test] fn doubly_sliced() { + let mut ctx = array_session().create_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); @@ -1452,17 +1477,23 @@ mod test { let sliced = patches.slice(15..100).unwrap().unwrap(); assert_eq!(sliced.array_len(), 100 - 15); - assert_arrays_eq!(sliced.values(), PrimitiveArray::from_iter([13531u32])); + assert_arrays_eq!( + sliced.values(), + PrimitiveArray::from_iter([13531u32]), + &mut ctx + ); let doubly_sliced = sliced.slice(35..36).unwrap().unwrap(); assert_arrays_eq!( doubly_sliced.values(), - PrimitiveArray::from_iter([13531u32]) + PrimitiveArray::from_iter([13531u32]), + &mut ctx ); } #[test] fn test_mask_all_true() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1473,14 +1504,13 @@ mod test { .unwrap(); let mask = Mask::new_true(10); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap(); assert!(masked.is_none()); } #[test] fn test_mask_all_false() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1491,41 +1521,30 @@ mod test { .unwrap(); let mask = Mask::new_false(10); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); - // No patch values should be masked + // Masking widens the surviving (still valid) values to nullable. + assert!(masked.values().dtype().is_nullable()); assert_arrays_eq!( masked.values(), - PrimitiveArray::from_iter([100i32, 200, 300]) - ); - assert!( - masked - .values() - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ); - assert!( - masked - .values() - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ); - assert!( - masked - .values() - .is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + PrimitiveArray::from_option_iter([Some(100i32), Some(200), Some(300)]), + &mut ctx ); + assert!(masked.values().is_valid(0, &mut ctx).unwrap()); + assert!(masked.values().is_valid(1, &mut ctx).unwrap()); + assert!(masked.values().is_valid(2, &mut ctx).unwrap()); // Indices should remain unchanged - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([2u64, 5, 8])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([2u64, 5, 8]), + &mut ctx + ); } #[test] fn test_mask_partial() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1539,21 +1558,27 @@ mod test { let mask = Mask::from_iter([ false, false, true, false, false, false, false, false, true, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); // Only the patch at index 5 should remain assert_eq!(masked.values().len(), 1); - assert_arrays_eq!(masked.values(), PrimitiveArray::from_iter([200i32])); + assert_arrays_eq!( + masked.values(), + PrimitiveArray::from_option_iter([Some(200i32)]), + &mut ctx + ); // Only index 5 should remain - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([5u64])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([5u64]), + &mut ctx + ); } #[test] fn test_mask_with_offset() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 5, // offset @@ -1568,18 +1593,24 @@ mod test { false, false, true, false, false, false, false, false, false, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); assert_eq!(masked.array_len(), 10); assert_eq!(masked.offset(), 5); - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([10u64, 13])); - assert_arrays_eq!(masked.values(), PrimitiveArray::from_iter([200i32, 300])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([10u64, 13]), + &mut ctx + ); + assert_arrays_eq!( + masked.values(), + PrimitiveArray::from_option_iter([Some(200i32), Some(300)]), + &mut ctx + ); } #[test] fn test_mask_nullable_values() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1593,41 +1624,30 @@ mod test { let mask = Mask::from_iter([ false, false, true, false, false, false, false, false, false, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); // Patches at indices 5 and 8 should remain - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([5u64, 8])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([5u64, 8]), + &mut ctx + ); // Values should be the null and 300 #[expect(deprecated)] let masked_values = masked.values().to_primitive(); assert_eq!(masked_values.len(), 2); - assert!( - !masked_values - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ); // the null value at index 5 - assert!( - masked_values - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ); // the 300 value at index 8 + assert!(!masked_values.is_valid(0, &mut ctx).unwrap()); // the null value at index 5 + assert!(masked_values.is_valid(1, &mut ctx).unwrap()); // the 300 value at index 8 assert_eq!( - i32::try_from( - &masked_values - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ) - .unwrap(), + i32::try_from(&masked_values.execute_scalar(1, &mut ctx).unwrap()).unwrap(), 300i32 ); } #[test] fn test_filter_keep_all() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1639,20 +1659,23 @@ mod test { // Keep all indices (mask with indices 0-9) let mask = Mask::from_indices(10, 0..10); - let filtered = patches - .filter(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let filtered = patches.filter(&mask, &mut ctx).unwrap().unwrap(); - assert_arrays_eq!(filtered.indices(), PrimitiveArray::from_iter([2u64, 5, 8])); + assert_arrays_eq!( + filtered.indices(), + PrimitiveArray::from_iter([2u64, 5, 8]), + &mut ctx + ); assert_arrays_eq!( filtered.values(), - PrimitiveArray::from_iter([100i32, 200, 300]) + PrimitiveArray::from_iter([100i32, 200, 300]), + &mut ctx ); } #[test] fn test_filter_none() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1664,14 +1687,13 @@ mod test { // Filter out all (empty mask means keep nothing) let mask = Mask::from_indices(10, vec![]); - let filtered = patches - .filter(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let filtered = patches.filter(&mask, &mut ctx).unwrap(); assert!(filtered.is_none()); } #[test] fn test_filter_with_indices() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1683,17 +1705,23 @@ mod test { // Keep indices 2, 5, 9 (so patches at 2 and 5 remain) let mask = Mask::from_indices(10, vec![2, 5, 9]); - let filtered = patches - .filter(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let filtered = patches.filter(&mask, &mut ctx).unwrap().unwrap(); - assert_arrays_eq!(filtered.indices(), PrimitiveArray::from_iter([0u64, 1])); // Adjusted indices - assert_arrays_eq!(filtered.values(), PrimitiveArray::from_iter([100i32, 200])); + assert_arrays_eq!( + filtered.indices(), + PrimitiveArray::from_iter([0u64, 1]), + &mut ctx + ); // Adjusted indices + assert_arrays_eq!( + filtered.values(), + PrimitiveArray::from_iter([100i32, 200]), + &mut ctx + ); } #[test] fn test_slice_full_range() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1705,15 +1733,21 @@ mod test { let sliced = patches.slice(0..10).unwrap().unwrap(); - assert_arrays_eq!(sliced.indices(), PrimitiveArray::from_iter([2u64, 5, 8])); + assert_arrays_eq!( + sliced.indices(), + PrimitiveArray::from_iter([2u64, 5, 8]), + &mut ctx + ); assert_arrays_eq!( sliced.values(), - PrimitiveArray::from_iter([100i32, 200, 300]) + PrimitiveArray::from_iter([100i32, 200, 300]), + &mut ctx ); } #[test] fn test_slice_partial() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1726,8 +1760,16 @@ mod test { // Slice from 3 to 8 (includes patch at 5) let sliced = patches.slice(3..8).unwrap().unwrap(); - assert_arrays_eq!(sliced.indices(), PrimitiveArray::from_iter([5u64])); // Index stays the same - assert_arrays_eq!(sliced.values(), PrimitiveArray::from_iter([200i32])); + assert_arrays_eq!( + sliced.indices(), + PrimitiveArray::from_iter([5u64]), + &mut ctx + ); // Index stays the same + assert_arrays_eq!( + sliced.values(), + PrimitiveArray::from_iter([200i32]), + &mut ctx + ); assert_eq!(sliced.array_len(), 5); // 8 - 3 = 5 assert_eq!(sliced.offset(), 3); // New offset } @@ -1750,6 +1792,7 @@ mod test { #[test] fn test_slice_with_offset() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 5, // offset @@ -1762,13 +1805,22 @@ mod test { // Slice from 3 to 8 (includes patch at actual index 5) let sliced = patches.slice(3..8).unwrap().unwrap(); - assert_arrays_eq!(sliced.indices(), PrimitiveArray::from_iter([10u64])); // Index stays the same (offset + 5 = 10) - assert_arrays_eq!(sliced.values(), PrimitiveArray::from_iter([200i32])); + assert_arrays_eq!( + sliced.indices(), + PrimitiveArray::from_iter([10u64]), + &mut ctx + ); // Index stays the same (offset + 5 = 10) + assert_arrays_eq!( + sliced.values(), + PrimitiveArray::from_iter([200i32]), + &mut ctx + ); assert_eq!(sliced.offset(), 8); // New offset = 5 + 3 } #[test] fn test_patch_values() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1781,30 +1833,15 @@ mod test { #[expect(deprecated)] let values = patches.values().to_primitive(); assert_eq!( - i32::try_from( - &values - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ) - .unwrap(), + i32::try_from(&values.execute_scalar(0, &mut ctx).unwrap()).unwrap(), 100i32 ); assert_eq!( - i32::try_from( - &values - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ) - .unwrap(), + i32::try_from(&values.execute_scalar(1, &mut ctx).unwrap()).unwrap(), 200i32 ); assert_eq!( - i32::try_from( - &values - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ) - .unwrap(), + i32::try_from(&values.execute_scalar(2, &mut ctx).unwrap()).unwrap(), 300i32 ); } @@ -1849,6 +1886,7 @@ mod test { #[test] fn test_mask_boundary_patches() { + let mut ctx = array_session().create_execution_ctx(); // Test masking patches at array boundaries let patches = Patches::new( 10, @@ -1862,17 +1900,24 @@ mod test { let mask = Mask::from_iter([ true, false, false, false, false, false, false, false, false, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap(); assert!(masked.is_some()); let masked = masked.unwrap(); - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([9u64])); - assert_arrays_eq!(masked.values(), PrimitiveArray::from_iter([200i32])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([9u64]), + &mut ctx + ); + assert_arrays_eq!( + masked.values(), + PrimitiveArray::from_option_iter([Some(200i32)]), + &mut ctx + ); } #[test] fn test_mask_all_patches_removed() { + let mut ctx = array_session().create_execution_ctx(); // Test when all patches are masked out let patches = Patches::new( 10, @@ -1887,14 +1932,13 @@ mod test { let mask = Mask::from_iter([ false, false, true, false, false, true, false, false, true, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap(); assert!(masked.is_none()); } #[test] fn test_mask_no_patches_removed() { + let mut ctx = array_session().create_execution_ctx(); // Test when no patches are masked let patches = Patches::new( 10, @@ -1909,20 +1953,23 @@ mod test { let mask = Mask::from_iter([ true, false, false, true, false, false, true, false, false, true, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([2u64, 5, 8])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([2u64, 5, 8]), + &mut ctx + ); assert_arrays_eq!( masked.values(), - PrimitiveArray::from_iter([100i32, 200, 300]) + PrimitiveArray::from_option_iter([Some(100i32), Some(200), Some(300)]), + &mut ctx ); } #[test] fn test_mask_single_patch() { + let mut ctx = array_session().create_execution_ctx(); // Test with a single patch let patches = Patches::new( 5, @@ -1935,22 +1982,22 @@ mod test { // Mask that removes the single patch let mask = Mask::from_iter([false, false, true, false, false]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap(); assert!(masked.is_none()); // Mask that keeps the single patch let mask = Mask::from_iter([true, false, false, true, false]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([2u64])); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([2u64]), + &mut ctx + ); } #[test] fn test_mask_contiguous_patches() { + let mut ctx = array_session().create_execution_ctx(); // Test with contiguous patches let patches = Patches::new( 10, @@ -1965,17 +2012,23 @@ mod test { let mask = Mask::from_iter([ false, false, false, false, true, true, false, false, false, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([3u64, 6])); - assert_arrays_eq!(masked.values(), PrimitiveArray::from_iter([100i32, 400])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([3u64, 6]), + &mut ctx + ); + assert_arrays_eq!( + masked.values(), + PrimitiveArray::from_option_iter([Some(100i32), Some(400)]), + &mut ctx + ); } #[test] fn test_mask_with_large_offset() { + let mut ctx = array_session().create_execution_ctx(); // Test with a large offset that shifts all indices let patches = Patches::new( 20, @@ -1991,18 +2044,24 @@ mod test { false, false, true, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([16u64, 19])); - assert_arrays_eq!(masked.values(), PrimitiveArray::from_iter([100i32, 300])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([16u64, 19]), + &mut ctx + ); + assert_arrays_eq!( + masked.values(), + PrimitiveArray::from_option_iter([Some(100i32), Some(300)]), + &mut ctx + ); } #[test] #[should_panic(expected = "Filter mask length 5 does not match array length 10")] fn test_mask_wrong_length() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -2014,9 +2073,7 @@ mod test { // Mask with wrong length let mask = Mask::from_iter([false, false, true, false, false]); - patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + patches.mask(&mask, &mut ctx).unwrap(); } #[test] diff --git a/vortex-array/src/scalar/arrow.rs b/vortex-array/src/scalar/arrow.rs index ea8a9decd34..cca6749cb19 100644 --- a/vortex-array/src/scalar/arrow.rs +++ b/vortex-array/src/scalar/arrow.rs @@ -455,6 +455,7 @@ mod tests { type Metadata = String; type NativeValue<'a> = &'a str; + #[expect(clippy::disallowed_methods, reason = "test-only id")] fn id(&self) -> ExtId { ExtId::new("some_ext") } @@ -557,10 +558,10 @@ mod tests { #[rstest] #[case(TimeUnit::Nanoseconds, "UTC", 1234567890000000000i64)] - #[case(TimeUnit::Microseconds, "EST", 1234567890000000i64)] + #[case(TimeUnit::Microseconds, "America/New_York", 1234567890000000i64)] #[case(TimeUnit::Microseconds, "Asia/Qatar", 1234567890000000i64)] #[case(TimeUnit::Microseconds, "Australia/Sydney", 1234567890000000i64)] - #[case(TimeUnit::Milliseconds, "HST", 1234567890000i64)] + #[case(TimeUnit::Milliseconds, "Pacific/Honolulu", 1234567890000i64)] #[case(TimeUnit::Seconds, "GMT", 1234567890i64)] fn test_temporal_timestamp_tz_to_arrow( #[case] time_unit: TimeUnit, diff --git a/vortex-array/src/scalar/mod.rs b/vortex-array/src/scalar/mod.rs index 8496037ea20..ff8d2c6134b 100644 --- a/vortex-array/src/scalar/mod.rs +++ b/vortex-array/src/scalar/mod.rs @@ -5,9 +5,13 @@ //! //! This crate provides scalar types and values that can be used to represent individual data //! elements in the Vortex array system. [`Scalar`]s are composed of a logical data type ([`DType`]) -//! and an optional (encoding nullablity) value ([`ScalarValue`]). +//! and an optional (encoding nullability) value ([`ScalarValue`]). //! //! Note that the implementations of `Scalar` are split into several different modules. +//! +//! `Scalar` is the single-row counterpart to [`ArrayRef`](crate::ArrayRef): it is logical, not tied +//! to any physical array encoding. A scalar always carries its [`DType`], and null scalars are +//! represented by `value == None`. #[cfg(feature = "arbitrary")] pub mod arbitrary; diff --git a/vortex-array/src/scalar/scalar_impl.rs b/vortex-array/src/scalar/scalar_impl.rs index 6675444be1a..60a675e42b9 100644 --- a/vortex-array/src/scalar/scalar_impl.rs +++ b/vortex-array/src/scalar/scalar_impl.rs @@ -188,10 +188,19 @@ impl Scalar { DType::Utf8(_) => value.as_utf8().is_empty(), DType::Binary(_) => value.as_binary().is_empty(), DType::List(..) => value.as_list().is_empty(), - // TODO(connor): This seems wrong... - DType::FixedSizeList(_, list_size, _) => value.as_list().len() == *list_size as usize, - // TODO(connor): This also seems wrong... - DType::Struct(struct_fields, _) => value.as_list().len() == struct_fields.nfields(), + // A fixed-size list is zero only if it has the expected number of elements and every + // element is itself a non-null zero value. + DType::FixedSizeList(_, list_size, _) => { + let list = self.as_list(); + list.len() == *list_size as usize + && (0..list.len()) + .all(|i| list.element(i).is_some_and(|e| e.is_zero() == Some(true))) + } + // A struct is zero only if every one of its fields is itself a non-null zero value. + DType::Struct(..) => self + .as_struct() + .fields_iter() + .is_some_and(|mut fields| fields.all(|f| f.is_zero() == Some(true))), DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), DType::Variant(_) => self.as_variant().is_zero()?, DType::Extension(_) => self.as_extension().to_storage_scalar().is_zero()?, @@ -428,3 +437,157 @@ fn partial_cmp_struct_values( Some(Ordering::Equal) } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::dtype::StructFields; + use crate::scalar::Scalar; + + fn i32_scalar(value: i32) -> Scalar { + Scalar::primitive::(value, Nullability::NonNullable) + } + + fn nullable_i32(value: Option) -> Scalar { + match value { + Some(value) => Scalar::primitive::(value, Nullability::Nullable), + None => Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), + } + } + + fn ab_struct_dtype(nullability: Nullability) -> DType { + DType::Struct( + StructFields::new( + ["a", "b"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + ], + ), + nullability, + ) + } + + #[rstest] + // A fixed-size list of all-zero elements is itself zero. + #[case(vec![0, 0], Some(true))] + #[case(vec![0], Some(true))] + // A single non-zero element makes the whole list non-zero. On `develop` these incorrectly + // returned `Some(true)` because only the element count was checked. + #[case(vec![0, 5], Some(false))] + #[case(vec![5, 0], Some(false))] + #[case(vec![1, 2], Some(false))] + fn fixed_size_list_is_zero(#[case] values: Vec, #[case] expected: Option) { + let element_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let children: Vec = values.into_iter().map(i32_scalar).collect(); + let scalar = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + assert_eq!(scalar.is_zero(), expected); + } + + #[test] + fn null_fixed_size_list_is_zero_is_none() { + let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)); + let scalar = Scalar::null(DType::FixedSizeList( + element_dtype, + 2, + Nullability::Nullable, + )); + assert_eq!(scalar.is_zero(), None); + } + + #[test] + fn fixed_size_list_with_null_element_is_not_zero() { + // A non-null fixed-size list containing a null element is not a zero value. On `develop` + // this incorrectly returned `Some(true)`. + let element_dtype = DType::Primitive(PType::I32, Nullability::Nullable); + let children = vec![nullable_i32(Some(0)), nullable_i32(None)]; + let scalar = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + assert_eq!(scalar.is_zero(), Some(false)); + } + + #[test] + fn struct_with_all_zero_fields_is_zero() { + let scalar = Scalar::struct_( + ab_struct_dtype(Nullability::NonNullable), + vec![i32_scalar(0), Scalar::utf8("", Nullability::NonNullable)], + ); + assert_eq!(scalar.is_zero(), Some(true)); + } + + #[rstest] + // A non-zero primitive field, a non-empty string field, or both, make the struct non-zero. On + // `develop` all of these incorrectly returned `Some(true)`. + #[case(5, "")] + #[case(0, "x")] + #[case(7, "y")] + fn struct_with_non_zero_field_is_not_zero(#[case] a: i32, #[case] b: &str) { + let scalar = Scalar::struct_( + ab_struct_dtype(Nullability::NonNullable), + vec![i32_scalar(a), Scalar::utf8(b, Nullability::NonNullable)], + ); + assert_eq!(scalar.is_zero(), Some(false)); + } + + #[test] + fn null_struct_is_zero_is_none() { + let scalar = Scalar::null(ab_struct_dtype(Nullability::Nullable)); + assert_eq!(scalar.is_zero(), None); + } + + #[test] + fn struct_with_null_field_is_not_zero() { + // A non-null struct with a null field is not a zero value. On `develop` this incorrectly + // returned `Some(true)`. + let dtype = DType::Struct( + StructFields::new( + ["a", "b"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::Nullable), + DType::Primitive(PType::I32, Nullability::Nullable), + ], + ), + Nullability::NonNullable, + ); + let scalar = Scalar::struct_(dtype, vec![nullable_i32(Some(0)), nullable_i32(None)]); + assert_eq!(scalar.is_zero(), Some(false)); + } + + #[test] + fn nested_struct_of_fixed_size_list_recurses() { + // Zero-checking must recurse through both structs and fixed-size lists. On `develop` the + // non-zero case incorrectly returned `Some(true)`. + let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)); + let fsl_dtype = + DType::FixedSizeList(Arc::clone(&element_dtype), 2, Nullability::NonNullable); + let struct_dtype = DType::Struct( + StructFields::new(["fsl"].into(), vec![fsl_dtype]), + Nullability::NonNullable, + ); + + let all_zero = Scalar::struct_( + struct_dtype.clone(), + vec![Scalar::fixed_size_list( + Arc::clone(&element_dtype), + vec![i32_scalar(0), i32_scalar(0)], + Nullability::NonNullable, + )], + ); + assert_eq!(all_zero.is_zero(), Some(true)); + + let with_non_zero = Scalar::struct_( + struct_dtype, + vec![Scalar::fixed_size_list( + element_dtype, + vec![i32_scalar(0), i32_scalar(9)], + Nullability::NonNullable, + )], + ); + assert_eq!(with_non_zero.is_zero(), Some(false)); + } +} diff --git a/vortex-array/src/scalar/tests/casting.rs b/vortex-array/src/scalar/tests/casting.rs index 8f287d2e87b..20eb442078a 100644 --- a/vortex-array/src/scalar/tests/casting.rs +++ b/vortex-array/src/scalar/tests/casting.rs @@ -32,6 +32,7 @@ mod tests { type Metadata = usize; type NativeValue<'a> = &'a str; + #[expect(clippy::disallowed_methods, reason = "test-only id")] fn id(&self) -> ExtId { ExtId::new("apples") } @@ -247,6 +248,7 @@ mod tests { type Metadata = usize; type NativeValue<'a> = &'a str; + #[expect(clippy::disallowed_methods, reason = "test-only id")] fn id(&self) -> ExtId { ExtId::new("f16_ext") } @@ -304,6 +306,7 @@ mod tests { type Metadata = usize; type NativeValue<'a> = &'a str; + #[expect(clippy::disallowed_methods, reason = "test-only id")] fn id(&self) -> ExtId { ExtId::new("struct_ext") } diff --git a/vortex-array/src/scalar/tests/mod.rs b/vortex-array/src/scalar/tests/mod.rs index a82b1908921..676363c853c 100644 --- a/vortex-array/src/scalar/tests/mod.rs +++ b/vortex-array/src/scalar/tests/mod.rs @@ -14,7 +14,4 @@ use std::sync::LazyLock; use vortex_session::VortexSession; -use crate::dtype::session::DTypeSession; - -pub(crate) static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +pub(crate) static SESSION: LazyLock = LazyLock::new(crate::array_session); diff --git a/vortex-array/src/scalar/typed_view/decimal/dvalue.rs b/vortex-array/src/scalar/typed_view/decimal/dvalue.rs index 4066a77e968..c1a3be7b71c 100644 --- a/vortex-array/src/scalar/typed_view/decimal/dvalue.rs +++ b/vortex-array/src/scalar/typed_view/decimal/dvalue.rs @@ -12,13 +12,34 @@ use num_traits::CheckedDiv; use num_traits::CheckedMul; use num_traits::CheckedSub; use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use crate::dtype::BigCast; use crate::dtype::DecimalDType; use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; use crate::dtype::ToI256; use crate::dtype::i256; use crate::match_each_decimal_value; +use crate::match_each_decimal_value_type; + +/// Widens both operands to the larger of their two decimal types, then applies the checked op. +macro_rules! checked_widening_binary_op { + ($self:expr, $other:expr, $op:path) => {{ + let target = $self.decimal_type().max($other.decimal_type()); + match_each_decimal_value_type!(target, |T| { + let a: T = $self + .cast() + .vortex_expect("widening cast to wider decimal type must always succeed"); + let b: T = $other + .cast() + .vortex_expect("widening cast to wider decimal type must always succeed"); + Some(DecimalValue::from($op(&a, &b)?)) + }) + }}; +} /// A decimal value that can be stored in various integer widths. /// @@ -99,6 +120,80 @@ impl DecimalValue { }) } + /// Rescales a stored decimal value from one scale to another. + /// + /// This preserves the represented numeric value exactly. Reducing scale fails if doing so + /// would discard non-zero fractional digits. + pub(crate) fn rescale_i256(value: i256, from_scale: i8, to_scale: i8) -> VortexResult { + if from_scale == to_scale || value == i256::ZERO { + return Ok(value); + } + + let scale_delta = to_scale as i16 - from_scale as i16; + if scale_delta > 0 { + let factor = decimal_scale_factor(scale_delta as u32)?; + value.checked_mul(&factor).ok_or_else(|| { + vortex_err!( + "Rescaling decimal from scale {} to {} overflows", + from_scale, + to_scale + ) + }) + } else { + let factor = decimal_scale_factor((-scale_delta) as u32)?; + let remainder = value % factor; + if remainder != i256::ZERO { + vortex_bail!( + "Rescaling decimal value {} from scale {} to {} would lose precision", + value, + from_scale, + to_scale + ); + } + Ok(value / factor) + } + } + + /// Rescales this value to `to_decimal_dtype`, checks precision, and stores it in the target + /// decimal value width. + pub(crate) fn cast_decimal( + &self, + from_decimal_dtype: DecimalDType, + to_decimal_dtype: DecimalDType, + ) -> VortexResult { + let rescaled = Self::rescale_i256( + self.as_i256(), + from_decimal_dtype.scale(), + to_decimal_dtype.scale(), + )?; + Self::try_from_i256(rescaled, to_decimal_dtype) + } + + /// Converts an untyped stored decimal integer into the physical value type selected by + /// `decimal_dtype`, after enforcing the dtype precision. + pub(crate) fn try_from_i256(value: i256, decimal_dtype: DecimalDType) -> VortexResult { + let decimal_value = Self::I256(value); + if !decimal_value.fits_in_precision(decimal_dtype) { + vortex_bail!( + "decimal value {} does not fit in precision of {}", + decimal_value, + decimal_dtype + ); + } + + let target_type = DecimalType::smallest_decimal_value_type(&decimal_dtype); + match_each_decimal_value_type!(target_type, |T| { + let value = ::from(value).ok_or_else(|| { + vortex_err!( + "decimal value {} cannot be represented as {}", + decimal_value, + target_type + ) + })?; + Ok(Self::from(value)) + }) + } + /// Returns the 0 value given the [`DecimalType`]. pub fn zero(decimal_type: &DecimalDType) -> Self { let smallest_type = DecimalType::smallest_decimal_value_type(decimal_type); @@ -134,39 +229,36 @@ impl DecimalValue { value_i256 > min_value && value_i256 < max_value } - /// Helper function to perform a checked binary operation on two decimal values. - /// - /// Both values are upcast to i256 before the operation, and the result is returned as I256. - fn checked_binary_op(&self, other: &Self, op: F) -> Option - where - F: FnOnce(i256, i256) -> Option, - { - let self_upcast = self.as_i256(); - let other_upcast = other.as_i256(); - op(self_upcast, other_upcast).map(DecimalValue::I256) - } - /// Checked addition. Returns `None` on overflow. pub fn checked_add(&self, other: &Self) -> Option { - self.checked_binary_op(other, |a, b| a.checked_add(&b)) + checked_widening_binary_op!(self, other, CheckedAdd::checked_add) } /// Checked subtraction. Returns `None` on overflow. pub fn checked_sub(&self, other: &Self) -> Option { - self.checked_binary_op(other, |a, b| a.checked_sub(&b)) + checked_widening_binary_op!(self, other, CheckedSub::checked_sub) } /// Checked multiplication. Returns `None` on overflow. pub fn checked_mul(&self, other: &Self) -> Option { - self.checked_binary_op(other, |a, b| a.checked_mul(&b)) + checked_widening_binary_op!(self, other, CheckedMul::checked_mul) } /// Checked division. Returns `None` on overflow or division by zero. pub fn checked_div(&self, other: &Self) -> Option { - self.checked_binary_op(other, |a, b| a.checked_div(&b)) + checked_widening_binary_op!(self, other, CheckedDiv::checked_div) } } +fn decimal_scale_factor(exp: u32) -> VortexResult { + i256::from_i128(10).checked_pow(exp).ok_or_else(|| { + vortex_err!( + "decimal scale factor 10^{} cannot be represented in i256", + exp + ) + }) +} + // Additional trait implementations for decimal types to ensure consistency. // Comparisons between DecimalValue types should upcast to i256 and operate in the upcast space. diff --git a/vortex-array/src/scalar/typed_view/decimal/scalar.rs b/vortex-array/src/scalar/typed_view/decimal/scalar.rs index b0890177f21..e365a1be02c 100644 --- a/vortex-array/src/scalar/typed_view/decimal/scalar.rs +++ b/vortex-array/src/scalar/typed_view/decimal/scalar.rs @@ -69,20 +69,9 @@ impl<'a> DecimalScalar<'a> { pub(crate) fn cast(&self, dtype: &DType) -> VortexResult { match dtype { DType::Decimal(target_dtype, target_nullability) => { - // Cast between decimal types - if self.decimal_type == *target_dtype { - // Same decimal type, just change nullability if needed - return Scalar::try_new( - dtype.clone(), - self.decimal_value.map(ScalarValue::Decimal), - ); - } - - // TODO(connor): Implement proper decimal scaling logic - whatever that means??? - // Different precision/scale - need to implement scaling logic - // For now, we'll do a simple value preservation without scaling if let Some(value) = &self.decimal_value { - Ok(Scalar::decimal(*value, *target_dtype, *target_nullability)) + let value = value.cast_decimal(self.decimal_type, *target_dtype)?; + Ok(Scalar::decimal(value, *target_dtype, *target_nullability)) } else { Ok(Scalar::null(dtype.clone())) } diff --git a/vortex-array/src/scalar/typed_view/decimal/tests.rs b/vortex-array/src/scalar/typed_view/decimal/tests.rs index ca5b06778d9..63f924b2239 100644 --- a/vortex-array/src/scalar/typed_view/decimal/tests.rs +++ b/vortex-array/src/scalar/typed_view/decimal/tests.rs @@ -115,7 +115,7 @@ fn test_decimal_cast_between_decimal_types() { Nullability::NonNullable, ); - // Cast to different decimal type (currently just preserves value) + // Cast to different decimal type. let result = decimal_scalar .cast(&DType::Decimal( DecimalDType::new(20, 4), @@ -123,9 +123,9 @@ fn test_decimal_cast_between_decimal_types() { )) .unwrap(); - // Value should be preserved (TODO(connor): proper scaling logic - whatever that means???) + // 123.45 with scale 2 is represented as 123.4500 with scale 4. let decimal_value: Option = result.try_into().unwrap(); - assert_eq!(decimal_value, Some(DecimalValue::I32(12345))); + assert_eq!(decimal_value, Some(DecimalValue::I128(1234500))); } #[test] @@ -309,7 +309,6 @@ fn test_decimal_to_decimal_different_scale() { ); // Cast to decimal with scale=4 - // TODO: This should properly rescale, but for now it preserves the raw value let target_dtype = DType::Decimal(DecimalDType::new(10, 4), Nullability::NonNullable); let result = decimal.cast(&target_dtype); assert!(result.is_ok()); @@ -317,7 +316,85 @@ fn test_decimal_to_decimal_different_scale() { let casted = result.unwrap(); assert_eq!( casted.as_decimal().decimal_value(), - Some(DecimalValue::I32(10000)) + Some(DecimalValue::I64(1000000)) + ); +} + +#[test] +fn test_decimal_to_decimal_lower_scale_exact() { + let decimal = Scalar::decimal( + DecimalValue::I64(1234500), // 123.4500 + DecimalDType::new(10, 4), + Nullability::NonNullable, + ); + + let casted = decimal + .cast(&DType::Decimal( + DecimalDType::new(10, 2), + Nullability::NonNullable, + )) + .unwrap(); + + assert_eq!( + casted.as_decimal().decimal_value(), + Some(DecimalValue::I64(12345)) + ); +} + +#[test] +fn test_decimal_to_decimal_lower_scale_lossy_fails() { + let decimal = Scalar::decimal( + DecimalValue::I64(1234567), // 123.4567 + DecimalDType::new(10, 4), + Nullability::NonNullable, + ); + + let result = decimal.cast(&DType::Decimal( + DecimalDType::new(10, 2), + Nullability::NonNullable, + )); + + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("would lose precision") + ); +} + +#[test] +fn test_primitive_i64_to_decimal_rescales() { + let scalar = Scalar::primitive(42i64, Nullability::NonNullable); + + let casted = scalar + .cast(&DType::Decimal( + DecimalDType::new(21, 2), + Nullability::NonNullable, + )) + .unwrap(); + + assert_eq!( + casted.as_decimal().decimal_value(), + Some(DecimalValue::I128(4200)) + ); +} + +#[test] +fn test_primitive_to_decimal_precision_checked() { + let scalar = Scalar::primitive(1000i32, Nullability::NonNullable); + + let result = scalar.cast(&DType::Decimal( + DecimalDType::new(2, 0), + Nullability::NonNullable, + )); + + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("does not fit in precision") ); } @@ -992,7 +1069,7 @@ fn test_decimal_value_checked_add() { let a = DecimalValue::I64(100); let b = DecimalValue::I64(200); let result = a.checked_add(&b).unwrap(); - assert_eq!(result, DecimalValue::I256(i256::from_i128(300))); + assert_eq!(result, DecimalValue::I64(300)); } #[test] @@ -1000,7 +1077,7 @@ fn test_decimal_value_checked_sub() { let a = DecimalValue::I64(500); let b = DecimalValue::I64(200); let result = a.checked_sub(&b).unwrap(); - assert_eq!(result, DecimalValue::I256(i256::from_i128(300))); + assert_eq!(result, DecimalValue::I64(300)); } #[test] @@ -1008,7 +1085,7 @@ fn test_decimal_value_checked_mul() { let a = DecimalValue::I32(50); let b = DecimalValue::I32(10); let result = a.checked_mul(&b).unwrap(); - assert_eq!(result, DecimalValue::I256(i256::from_i128(500))); + assert_eq!(result, DecimalValue::I32(500)); } #[test] @@ -1016,7 +1093,7 @@ fn test_decimal_value_checked_div() { let a = DecimalValue::I64(1000); let b = DecimalValue::I64(10); let result = a.checked_div(&b).unwrap(); - assert_eq!(result, DecimalValue::I256(i256::from_i128(100))); + assert_eq!(result, DecimalValue::I64(100)); } #[test] @@ -1033,7 +1110,69 @@ fn test_decimal_value_mixed_types() { let a = DecimalValue::I8(10); let b = DecimalValue::I128(20); let result = a.checked_add(&b).unwrap(); - assert_eq!(result, DecimalValue::I256(i256::from_i128(30))); + assert_eq!(result, DecimalValue::I128(30)); +} + +#[test] +fn test_checked_ops_preserve_type() { + // Operations should return the wider of the two operand types, not unconditionally upcast to I256 + let add = DecimalValue::I32(5) + .checked_add(&DecimalValue::I32(3)) + .unwrap(); + assert_eq!(add.decimal_type(), DecimalType::I32); + + let sub = DecimalValue::I64(10) + .checked_sub(&DecimalValue::I64(3)) + .unwrap(); + assert_eq!(sub.decimal_type(), DecimalType::I64); + + let mul = DecimalValue::I8(2) + .checked_mul(&DecimalValue::I8(3)) + .unwrap(); + assert_eq!(mul.decimal_type(), DecimalType::I8); + + let div = DecimalValue::I128(10) + .checked_div(&DecimalValue::I128(2)) + .unwrap(); + assert_eq!(div.decimal_type(), DecimalType::I128); + + let add_i256 = DecimalValue::I256(i256::from_i128(1)) + .checked_add(&DecimalValue::I256(i256::from_i128(2))) + .unwrap(); + assert_eq!(add_i256.decimal_type(), DecimalType::I256); +} + +#[test] +fn test_checked_ops_mixed_types_use_wider() { + let add = DecimalValue::I8(1) + .checked_add(&DecimalValue::I64(2)) + .unwrap(); + assert_eq!(add.decimal_type(), DecimalType::I64); + + let sub = DecimalValue::I32(10) + .checked_sub(&DecimalValue::I128(3)) + .unwrap(); + assert_eq!(sub.decimal_type(), DecimalType::I128); +} + +#[test] +fn test_checked_ops_overflow_at_target_width() { + assert_eq!( + DecimalValue::I8(i8::MAX).checked_add(&DecimalValue::I8(1)), + None + ); + assert_eq!( + DecimalValue::I16(i16::MIN).checked_sub(&DecimalValue::I16(1)), + None + ); + assert_eq!( + DecimalValue::I32(i32::MAX).checked_mul(&DecimalValue::I32(2)), + None + ); + assert_eq!( + DecimalValue::I8(i8::MIN).checked_div(&DecimalValue::I8(-1)), + None + ); } #[test] diff --git a/vortex-array/src/scalar/typed_view/extension/tests.rs b/vortex-array/src/scalar/typed_view/extension/tests.rs index 8bf52aa5369..2b08c530af6 100644 --- a/vortex-array/src/scalar/typed_view/extension/tests.rs +++ b/vortex-array/src/scalar/typed_view/extension/tests.rs @@ -20,6 +20,7 @@ impl ExtVTable for TestI32Ext { type Metadata = EmptyMetadata; type NativeValue<'a> = &'a str; + #[expect(clippy::disallowed_methods, reason = "test-only id")] fn id(&self) -> ExtId { ExtId::new("test_ext") } @@ -103,6 +104,7 @@ fn test_ext_scalar_partial_ord_different_types() { type Metadata = EmptyMetadata; type NativeValue<'a> = &'a str; + #[expect(clippy::disallowed_methods, reason = "test-only id")] fn id(&self) -> ExtId { ExtId::new("test_ext_2") } @@ -285,6 +287,7 @@ fn test_ext_scalar_with_metadata() { type Metadata = usize; type NativeValue<'a> = &'a str; + #[expect(clippy::disallowed_methods, reason = "test-only id")] fn id(&self) -> ExtId { ExtId::new("test_ext_metadata") } diff --git a/vortex-array/src/scalar/typed_view/primitive/scalar.rs b/vortex-array/src/scalar/typed_view/primitive/scalar.rs index e3981de8dd3..3ca4d337a45 100644 --- a/vortex-array/src/scalar/typed_view/primitive/scalar.rs +++ b/vortex-array/src/scalar/typed_view/primitive/scalar.rs @@ -18,15 +18,19 @@ use num_traits::CheckedSub; use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use super::pvalue::CoercePValue; use crate::dtype::DType; +use crate::dtype::DecimalDType; use crate::dtype::FromPrimitiveOrF16; use crate::dtype::NativePType; use crate::dtype::PType; +use crate::dtype::i256; use crate::match_each_native_ptype; +use crate::scalar::DecimalValue; use crate::scalar::NumericOperator; use crate::scalar::PValue; use crate::scalar::Scalar; @@ -164,13 +168,21 @@ impl<'a> PrimitiveScalar<'a> { /// Casts this scalar to the given `dtype`. pub(crate) fn cast(&self, dtype: &DType) -> VortexResult { - let ptype = PType::try_from(dtype)?; let pvalue = self .pvalue .vortex_expect("nullness handled in Scalar::cast"); - Ok(match_each_native_ptype!(ptype, |Q| { - Scalar::primitive(pvalue.cast::()?, dtype.nullability()) - })) + + match dtype { + DType::Primitive(ptype, nullability) => Ok(match_each_native_ptype!(*ptype, |Q| { + Scalar::primitive(pvalue.cast::()?, *nullability) + })), + DType::Decimal(decimal_dtype, nullability) => Ok(Scalar::decimal( + pvalue_to_decimal(pvalue, *decimal_dtype)?, + *decimal_dtype, + *nullability, + )), + _ => vortex_bail!("Cannot cast primitive scalar to {dtype}"), + } } /// Returns true if the scalar is nan. @@ -267,6 +279,25 @@ impl<'a> PrimitiveScalar<'a> { } } +fn pvalue_to_decimal(pvalue: PValue, decimal_dtype: DecimalDType) -> VortexResult { + let value = match pvalue { + PValue::U8(v) => i256::from_i128(i128::from(v)), + PValue::U16(v) => i256::from_i128(i128::from(v)), + PValue::U32(v) => i256::from_i128(i128::from(v)), + PValue::U64(v) => i256::from_i128(i128::from(v)), + PValue::I8(v) => i256::from_i128(i128::from(v)), + PValue::I16(v) => i256::from_i128(i128::from(v)), + PValue::I32(v) => i256::from_i128(i128::from(v)), + PValue::I64(v) => i256::from_i128(i128::from(v)), + PValue::F16(_) | PValue::F32(_) | PValue::F64(_) => { + vortex_bail!("Cannot cast floating primitive {pvalue} to decimal {decimal_dtype}") + } + }; + + let scaled = DecimalValue::rescale_i256(value, 0, decimal_dtype.scale())?; + DecimalValue::try_from_i256(scaled, decimal_dtype) +} + impl Sub for PrimitiveScalar<'_> { type Output = Self; diff --git a/vortex-array/src/scalar_fn/erased.rs b/vortex-array/src/scalar_fn/erased.rs index 10e82d25455..6e0011c297a 100644 --- a/vortex-array/src/scalar_fn/erased.rs +++ b/vortex-array/src/scalar_fn/erased.rs @@ -20,8 +20,6 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; use crate::expr::Expression; -use crate::expr::StatsCatalog; -use crate::expr::stats::Stat; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ReduceCtx; @@ -180,25 +178,6 @@ impl ScalarFnRef { pub(crate) fn simplify_untyped(&self, expr: &Expression) -> VortexResult> { self.0.simplify_untyped(expr) } - - /// Compute stat falsification expression. - pub(crate) fn stat_falsification( - &self, - expr: &Expression, - catalog: &dyn StatsCatalog, - ) -> Option { - self.0.stat_falsification(expr, catalog) - } - - /// Compute stat expression. - pub(crate) fn stat_expression( - &self, - expr: &Expression, - stat: Stat, - catalog: &dyn StatsCatalog, - ) -> Option { - self.0.stat_expression(expr, stat, catalog) - } } impl Debug for ScalarFnRef { diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index 0e0d9949195..2151a310dde 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -25,7 +25,6 @@ use crate::arrays::Primitive; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::DType::Bool; -use crate::expr::StatsCatalog; use crate::expr::expression::Expression; use crate::scalar::Scalar; use crate::scalar_fn::Arity; @@ -33,8 +32,6 @@ use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::execute_boolean; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; @@ -298,22 +295,6 @@ impl ScalarFnVTable for Between { between_canonical(&arr, &lower, &upper, options, ctx) } - fn stat_falsification( - &self, - options: &Self::Options, - expr: &Expression, - catalog: &dyn StatsCatalog, - ) -> Option { - let arr = expr.child(0).clone(); - let lower = expr.child(1).clone(); - let upper = expr.child(2).clone(); - - let lhs = Binary.new_expr(options.lower_strict.to_operator(), [lower, arr.clone()]); - let rhs = Binary.new_expr(options.upper_strict.to_operator(), [arr, upper]); - - and(lhs, rhs).stat_falsification(catalog) - } - fn validity( &self, _options: &Self::Options, @@ -357,12 +338,10 @@ mod tests { use crate::expr::root; use crate::scalar::DecimalValue; use crate::scalar::Scalar; - use crate::session::ArraySession; use crate::test_harness::to_int_indices; use crate::validity::Validity; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); #[test] fn test_display() { @@ -402,6 +381,7 @@ mod tests { let lower = buffer![0, 0, 0, 0, 2].into_array(); let array = buffer![1, 0, 1, 0, 1].into_array(); let upper = buffer![2, 1, 1, 0, 0].into_array(); + let ctx = &mut SESSION.create_execution_ctx(); let matches = between_canonical( &array, @@ -411,13 +391,13 @@ mod tests { lower_strict, upper_strict, }, - &mut SESSION.create_execution_ctx(), + ctx, ) .unwrap() - .execute::(&mut SESSION.create_execution_ctx()) + .execute::(ctx) .unwrap(); - let indices = to_int_indices(matches).unwrap(); + let indices = to_int_indices(matches, ctx).unwrap(); assert_eq!(indices, expected); } @@ -425,6 +405,7 @@ mod tests { fn test_constants() { let lower = buffer![0, 0, 2, 0, 2].into_array(); let array = buffer![1, 0, 1, 0, 1].into_array(); + let ctx = &mut SESSION.create_execution_ctx(); // upper is null let upper = ConstantArray::new( @@ -441,13 +422,13 @@ mod tests { lower_strict: StrictComparison::NonStrict, upper_strict: StrictComparison::NonStrict, }, - &mut SESSION.create_execution_ctx(), + ctx, ) .unwrap() - .execute::(&mut SESSION.create_execution_ctx()) + .execute::(ctx) .unwrap(); - let indices = to_int_indices(matches).unwrap(); + let indices = to_int_indices(matches, ctx).unwrap(); assert!(indices.is_empty()); // upper is a fixed constant @@ -460,12 +441,12 @@ mod tests { lower_strict: StrictComparison::NonStrict, upper_strict: StrictComparison::NonStrict, }, - &mut SESSION.create_execution_ctx(), + ctx, ) .unwrap() - .execute::(&mut SESSION.create_execution_ctx()) + .execute::(ctx) .unwrap(); - let indices = to_int_indices(matches).unwrap(); + let indices = to_int_indices(matches, ctx).unwrap(); assert_eq!(indices, vec![0, 1, 3]); // lower is also a constant @@ -479,17 +460,18 @@ mod tests { lower_strict: StrictComparison::NonStrict, upper_strict: StrictComparison::NonStrict, }, - &mut SESSION.create_execution_ctx(), + ctx, ) .unwrap() - .execute::(&mut SESSION.create_execution_ctx()) + .execute::(ctx) .unwrap(); - let indices = to_int_indices(matches).unwrap(); + let indices = to_int_indices(matches, ctx).unwrap(); assert_eq!(indices, vec![0, 1, 2, 3, 4]); } #[test] fn test_between_decimal() { + let ctx = &mut SESSION.create_execution_ctx(); let values = buffer![100i128, 200i128, 300i128, 400i128]; let decimal_type = DecimalDType::new(3, 2); let array = DecimalArray::new(values, decimal_type, Validity::NonNullable).into_array(); @@ -522,12 +504,13 @@ mod tests { lower_strict: StrictComparison::Strict, upper_strict: StrictComparison::NonStrict, }, - &mut SESSION.create_execution_ctx(), + ctx, ) .unwrap(); assert_arrays_eq!( between_strict, - BoolArray::from_iter([false, true, true, true]) + BoolArray::from_iter([false, true, true, true]), + ctx ); // Non-strict lower bound, strict upper bound @@ -539,12 +522,13 @@ mod tests { lower_strict: StrictComparison::NonStrict, upper_strict: StrictComparison::Strict, }, - &mut SESSION.create_execution_ctx(), + ctx, ) .unwrap(); assert_arrays_eq!( between_strict, - BoolArray::from_iter([true, true, true, false]) + BoolArray::from_iter([true, true, true, false]), + ctx ); } @@ -567,6 +551,7 @@ mod tests { #[case] upper_val: DecimalValue, #[case] expected_indices: Vec, ) { + let ctx = &mut SESSION.create_execution_ctx(); // Array uses I16 storage with precision=5 (values fit in i16 even though precision=5 // nominally maps to I32 as the smallest storage type). let decimal_type = DecimalDType::new(5, -67); @@ -596,12 +581,12 @@ mod tests { lower_strict: StrictComparison::NonStrict, upper_strict: StrictComparison::NonStrict, }, - &mut SESSION.create_execution_ctx(), + ctx, ) .unwrap() - .execute::(&mut SESSION.create_execution_ctx()) + .execute::(ctx) .unwrap(); - assert_eq!(to_int_indices(result).unwrap(), expected_indices); + assert_eq!(to_int_indices(result, ctx).unwrap(), expected_indices); } } diff --git a/vortex-array/src/scalar_fn/fns/binary/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/boolean.rs index 218aba806c9..b7f838c6d13 100644 --- a/vortex-array/src/scalar_fn/fns/binary/boolean.rs +++ b/vortex-array/src/scalar_fn/fns/binary/boolean.rs @@ -1,21 +1,103 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::iter::repeat_n; + use arrow_array::cast::AsArray; +use vortex_buffer::BitBuffer; +use vortex_buffer::BufferMut; +use vortex_buffer::read_u64_le; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray; +use crate::array::ArrayView; +use crate::array::VTable; +use crate::arrays::Bool; +use crate::arrays::BoolArray; use crate::arrays::Constant; use crate::arrays::ConstantArray; +use crate::arrays::ScalarFn; +use crate::arrays::scalar_fn::ExactScalarFn; +use crate::arrays::scalar_fn::ScalarFnArrayExt; +use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::arrow::ArrowSessionExt; use crate::arrow::FromArrowArray; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::executor::ExecutionCtx; +use crate::dtype::Nullability; +use crate::kernel::ExecuteParentKernel; +use crate::scalar::BoolScalar; use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::Operator; +use crate::validity::Validity; + +/// Trait for encoding-specific boolean kernels that operate in encoded space. +/// +/// Implementations receive the encoded array as the left operand. `rhs` may be any boolean array +/// encoding or a constant; implementations should return `Ok(None)` when they cannot handle that +/// operand without falling back to ordinary execution. +/// +/// Vortex's boolean [`Operator::And`] and [`Operator::Or`] variants use Kleene semantics; there is +/// no separate two-valued boolean operator path to dispatch here. +pub trait BooleanKernel: VTable { + /// Execute `lhs rhs` using Kleene boolean semantics. + fn boolean( + lhs: ArrayView<'_, Self>, + rhs: &ArrayRef, + operator: Operator, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; +} + +/// Adaptor that bridges [`BooleanKernel`] implementations to [`ExecuteParentKernel`]. +/// +/// When a `ScalarFnArray(Binary, And|Or)` wraps a child implementing [`BooleanKernel`], this +/// adaptor extracts the other operand and delegates to the encoding-specific kernel. +#[derive(Default, Debug)] +pub struct BooleanExecuteAdaptor(pub V); + +impl ExecuteParentKernel for BooleanExecuteAdaptor +where + V: BooleanKernel, +{ + type Parent = ExactScalarFn; + + fn execute_parent( + &self, + array: ArrayView<'_, V>, + parent: ScalarFnArrayView<'_, Binary>, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let op = *parent.options; + if !is_boolean_operator(op) { + return Ok(None); + } + + let Some(scalar_fn_array) = parent.as_opt::() else { + return Ok(None); + }; + let other = match child_idx { + 0 => scalar_fn_array.get_child(1), + 1 => scalar_fn_array.get_child(0), + _ => return Ok(None), + }; + + if let Some(result) = constant_boolean(array.array(), other, op)? { + return Ok(Some(result)); + } + + V::boolean(array, other, op, ctx) + } +} /// Point-wise Kleene logical _and_ between two Boolean arrays. #[deprecated(note = "Use `ArrayBuiltins::binary` instead")] @@ -32,16 +114,35 @@ pub fn or_kleene(lhs: &ArrayRef, rhs: &ArrayRef) -> VortexResult { /// Execute a Kleene boolean operation between two arrays. /// /// This is the entry point for boolean operations from the binary expression. -/// Handles constant-constant directly, otherwise falls back to Arrow. +/// Handles constants and canonical boolean arrays directly, otherwise falls back to Arrow. pub(crate) fn execute_boolean( lhs: &ArrayRef, rhs: &ArrayRef, op: Operator, ctx: &mut ExecutionCtx, ) -> VortexResult { + let nullable = boolean_nullability(lhs, rhs); + + if lhs.is_empty() { + return Ok(Canonical::empty(&DType::Bool(nullable)).into_array()); + } + if let Some(result) = constant_boolean(lhs, rhs, op)? { return Ok(result); } + + if let Some(lhs) = lhs.as_opt::() + && let Some(result) = ::boolean(lhs, rhs, op, ctx)? + { + return Ok(result); + } + + if let Some(rhs) = rhs.as_opt::() + && let Some(result) = ::boolean(rhs, lhs, op, ctx)? + { + return Ok(result); + } + arrow_execute_boolean(lhs.clone(), rhs.clone(), op, ctx) } @@ -52,84 +153,707 @@ fn arrow_execute_boolean( op: Operator, ctx: &mut ExecutionCtx, ) -> VortexResult { - let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable(); + let nullable = boolean_nullability(&lhs, &rhs); let session = ctx.session().clone(); let lhs = session .arrow() .execute_arrow(lhs, None, ctx)? - .as_boolean() + .as_boolean_opt() + .ok_or_else(|| vortex_err!("expected lhs to be boolean"))? .clone(); let rhs = session .arrow() .execute_arrow(rhs, None, ctx)? - .as_boolean() + .as_boolean_opt() + .ok_or_else(|| vortex_err!("expected rhs to be boolean"))? .clone(); let array = match op { Operator::And => arrow_arith::boolean::and_kleene(&lhs, &rhs)?, Operator::Or => arrow_arith::boolean::or_kleene(&lhs, &rhs)?, - other => return Err(vortex_err!("Not a boolean operator: {other}")), + other => vortex_bail!("Not a boolean operator: {other}"), }; - ArrayRef::from_arrow(&array, nullable) + ArrayRef::from_arrow(&array, nullable == Nullability::Nullable) } -/// Constant-folds a boolean operation between two constant arrays. +/// Handles boolean operations where at least one operand is a constant array. fn constant_boolean( lhs: &ArrayRef, rhs: &ArrayRef, op: Operator, ) -> VortexResult> { - let (Some(lhs), Some(rhs)) = (lhs.as_opt::(), rhs.as_opt::()) else { - return Ok(None); - }; + let nullable = boolean_nullability(lhs, rhs); - let length = lhs.len(); - let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable(); - let lhs_val = lhs.scalar().as_bool().value(); - let rhs_val = rhs - .scalar() - .as_bool_opt() - .ok_or_else(|| vortex_err!("expected rhs to be boolean"))? - .value(); + match (lhs.as_opt::(), rhs.as_opt::()) { + (Some(lhs), Some(rhs)) => { + let result = boolean_scalar_scalar( + bool_scalar_value(lhs.scalar())?, + bool_scalar_value(rhs.scalar())?, + op, + )?; + + Ok(Some(constant_bool_result(result, lhs.len(), nullable))) + } + (Some(lhs), None) => constant_array_boolean(lhs.scalar(), rhs, op, nullable), + (None, Some(rhs)) => constant_array_boolean(rhs.scalar(), lhs, op, nullable), + (None, None) => Ok(None), + } +} + +fn constant_array_boolean( + constant: &Scalar, + array: &ArrayRef, + op: Operator, + nullability: Nullability, +) -> VortexResult> { + match (op, bool_scalar_value(constant)?) { + (Operator::And, Some(false)) => Ok(Some(constant_bool_result( + Some(false), + array.len(), + nullability, + ))), + (Operator::And, Some(true)) => Ok(Some(cast_bool_nullability(array, nullability)?)), + (Operator::Or, Some(true)) => Ok(Some(constant_bool_result( + Some(true), + array.len(), + nullability, + ))), + (Operator::Or, Some(false)) => Ok(Some(cast_bool_nullability(array, nullability)?)), + (Operator::And | Operator::Or, None) => Ok(None), + (other, _) => vortex_bail!("Not a boolean operator: {other}"), + } +} - let result = match op { - Operator::And => match (lhs_val, rhs_val) { +fn boolean_scalar_scalar( + lhs: Option, + rhs: Option, + op: Operator, +) -> VortexResult> { + Ok(match op { + Operator::And => match (lhs, rhs) { (Some(false), _) | (_, Some(false)) => Some(false), (None, _) | (_, None) => None, (Some(l), Some(r)) => Some(l & r), }, - Operator::Or => match (lhs_val, rhs_val) { + Operator::Or => match (lhs, rhs) { (Some(true), _) | (_, Some(true)) => Some(true), (None, _) | (_, None) => None, (Some(l), Some(r)) => Some(l | r), }, - other => return Err(vortex_err!("Not a boolean operator: {other}")), + other => vortex_bail!("Not a boolean operator: {other}"), + }) +} + +fn bool_scalar_value(scalar: &Scalar) -> VortexResult> { + Ok(scalar + .as_bool_opt() + .ok_or_else(|| vortex_err!("expected boolean scalar"))? + .value()) +} + +/// Execute a Kleene boolean operation from boolean value bitmaps and validity values. +pub fn kleene_boolean_buffers( + lhs_values: BitBuffer, + lhs_validity: Validity, + rhs_values: BitBuffer, + rhs_validity: Validity, + operator: Operator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs_values.len(); + debug_assert_eq!(rhs_values.len(), len); + + if lhs_validity.definitely_no_nulls() && rhs_validity.definitely_no_nulls() { + let values = match operator { + Operator::And => lhs_values & &rhs_values, + Operator::Or => lhs_values | &rhs_values, + other => vortex_bail!("Not a boolean operator: {other}"), + }; + return Ok(BoolArray::try_new(values, Validity::from(nullability))?.into_array()); + } + + let lhs_valid = lhs_validity.execute_mask(len, ctx)?; + let rhs_valid = rhs_validity.execute_mask(len, ctx)?; + fused_boolean_buffers( + len, + &lhs_values, + &lhs_valid, + &rhs_values, + &rhs_valid, + operator, + nullability, + ) +} + +/// Execute a Kleene boolean operation between boolean value bits and a scalar. +pub fn kleene_boolean_buffer_scalar( + values: BitBuffer, + validity: Validity, + scalar: &BoolScalar<'_>, + operator: Operator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let scalar_value = scalar.value(); + let len = values.len(); + let result = match (operator, scalar_value) { + (Operator::And, Some(false)) => { + return Ok(constant_bool_result(Some(false), len, nullability)); + } + (Operator::And, Some(true)) => { + return Ok( + BoolArray::try_new(values, validity.union_nullability(nullability))?.into_array(), + ); + } + (Operator::Or, Some(true)) => { + return Ok(constant_bool_result(Some(true), len, nullability)); + } + (Operator::Or, Some(false)) => { + return Ok( + BoolArray::try_new(values, validity.union_nullability(nullability))?.into_array(), + ); + } + (Operator::And, None) => { + let valid = validity + .execute_mask(len, ctx)? + .bitand_not(&Mask::from_buffer(values)); + BoolArray::try_new( + BitBuffer::new_unset(len), + Validity::from_mask(valid, nullability), + )? + } + (Operator::Or, None) => { + let valid = validity.execute_mask(len, ctx)? & &Mask::from_buffer(values); + BoolArray::try_new( + BitBuffer::new_set(len), + Validity::from_mask(valid, nullability), + )? + } + (other, _) => vortex_bail!("Not a boolean operator: {other}"), }; - let scalar = result - .map(|b| Scalar::bool(b, nullable.into())) - .unwrap_or_else(|| Scalar::null(DType::Bool(nullable.into()))); + Ok(result.into_array()) +} + +fn fused_boolean_buffers( + len: usize, + lhs_values: &BitBuffer, + lhs_validity: &Mask, + rhs_values: &BitBuffer, + rhs_validity: &Mask, + operator: Operator, + nullability: Nullability, +) -> VortexResult { + if let Some(result) = fused_boolean_buffers_aligned( + len, + lhs_values, + lhs_validity, + rhs_values, + rhs_validity, + operator, + nullability, + )? { + return Ok(result); + } + + let n_words = len.div_ceil(64); + + macro_rules! fuse { + ($lhs_valid_words:expr, $rhs_valid_words:expr) => { + fused_boolean_words( + len, + lhs_values.chunks().iter_padded(), + rhs_values.chunks().iter_padded(), + $lhs_valid_words, + $rhs_valid_words, + operator, + nullability, + ) + }; + } + + match (lhs_validity.bit_buffer(), rhs_validity.bit_buffer()) { + (AllOr::All, AllOr::All) => { + fuse!(repeat_n(u64::MAX, n_words), repeat_n(u64::MAX, n_words)) + } + (AllOr::All, AllOr::None) => { + fuse!(repeat_n(u64::MAX, n_words), repeat_n(0, n_words)) + } + (AllOr::All, AllOr::Some(rhs_validity)) => fuse!( + repeat_n(u64::MAX, n_words), + rhs_validity.chunks().iter_padded() + ), + (AllOr::None, AllOr::All) => { + fuse!(repeat_n(0, n_words), repeat_n(u64::MAX, n_words)) + } + (AllOr::None, AllOr::None) => { + fuse!(repeat_n(0, n_words), repeat_n(0, n_words)) + } + (AllOr::None, AllOr::Some(rhs_validity)) => { + fuse!(repeat_n(0, n_words), rhs_validity.chunks().iter_padded()) + } + (AllOr::Some(lhs_validity), AllOr::All) => fuse!( + lhs_validity.chunks().iter_padded(), + repeat_n(u64::MAX, n_words) + ), + (AllOr::Some(lhs_validity), AllOr::None) => { + fuse!(lhs_validity.chunks().iter_padded(), repeat_n(0, n_words)) + } + (AllOr::Some(lhs_validity), AllOr::Some(rhs_validity)) => fuse!( + lhs_validity.chunks().iter_padded(), + rhs_validity.chunks().iter_padded() + ), + } +} - Ok(Some(ConstantArray::new(scalar, length).into_array())) +#[derive(Clone, Copy)] +enum WordSource<'a> { + Fill(u64), + Bytes(&'a [u8]), +} + +impl WordSource<'_> { + #[inline] + fn word_at(self, byte_offset: usize, len: usize) -> u64 { + match self { + Self::Fill(word) => word, + Self::Bytes(bytes) => read_u64_le(&bytes[byte_offset..byte_offset + len]), + } + } +} + +fn fused_boolean_buffers_aligned( + len: usize, + lhs_values: &BitBuffer, + lhs_validity: &Mask, + rhs_values: &BitBuffer, + rhs_validity: &Mask, + operator: Operator, + nullability: Nullability, +) -> VortexResult> { + let Some(lhs_values) = word_source_from_bit_buffer(lhs_values) else { + return Ok(None); + }; + let Some(rhs_values) = word_source_from_bit_buffer(rhs_values) else { + return Ok(None); + }; + let Some(lhs_validity) = word_source_from_mask(lhs_validity) else { + return Ok(None); + }; + let Some(rhs_validity) = word_source_from_mask(rhs_validity) else { + return Ok(None); + }; + + Ok(Some(fused_boolean_word_sources( + len, + lhs_values, + rhs_values, + lhs_validity, + rhs_validity, + operator, + nullability, + )?)) +} + +fn word_source_from_bit_buffer(buffer: &BitBuffer) -> Option> { + buffer.byte_aligned_bytes().map(WordSource::Bytes) +} + +fn word_source_from_mask(mask: &Mask) -> Option> { + match mask.bit_buffer() { + AllOr::All => Some(WordSource::Fill(u64::MAX)), + AllOr::None => Some(WordSource::Fill(0)), + AllOr::Some(buffer) => word_source_from_bit_buffer(buffer), + } +} + +fn fused_boolean_word_sources( + len: usize, + lhs_words: WordSource<'_>, + rhs_words: WordSource<'_>, + lhs_valid_words: WordSource<'_>, + rhs_valid_words: WordSource<'_>, + operator: Operator, + nullability: Nullability, +) -> VortexResult { + match operator { + Operator::And => fused_boolean_and_word_sources( + len, + lhs_words, + rhs_words, + lhs_valid_words, + rhs_valid_words, + nullability, + ), + Operator::Or => fused_boolean_or_word_sources( + len, + lhs_words, + rhs_words, + lhs_valid_words, + rhs_valid_words, + nullability, + ), + other => vortex_bail!("Not a boolean operator: {other}"), + } +} + +fn fused_boolean_and_word_sources( + len: usize, + lhs_words: WordSource<'_>, + rhs_words: WordSource<'_>, + lhs_valid_words: WordSource<'_>, + rhs_valid_words: WordSource<'_>, + nullability: Nullability, +) -> VortexResult { + let n_bytes = len.div_ceil(8); + let n_words = n_bytes.div_ceil(8); + let full_bytes = n_bytes - n_bytes % 8; + let mut values = BufferMut::::with_capacity(n_words); + let mut validity = BufferMut::::with_capacity(n_words); + + for byte_offset in (0..full_bytes).step_by(8) { + let lhs = lhs_words.word_at(byte_offset, 8); + let rhs = rhs_words.word_at(byte_offset, 8); + let lhs_valid = lhs_valid_words.word_at(byte_offset, 8); + let rhs_valid = rhs_valid_words.word_at(byte_offset, 8); + + // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this + // loop plus the optional tail push emits at most `n_words` words. + unsafe { + values.push_unchecked(lhs & rhs); + validity + .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs)); + } + } + + if full_bytes != n_bytes { + let tail_len = n_bytes - full_bytes; + let lhs = lhs_words.word_at(full_bytes, tail_len); + let rhs = rhs_words.word_at(full_bytes, tail_len); + let lhs_valid = lhs_valid_words.word_at(full_bytes, tail_len); + let rhs_valid = rhs_valid_words.word_at(full_bytes, tail_len); + + // SAFETY: see the loop safety comment above. + unsafe { + values.push_unchecked(lhs & rhs); + validity + .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs)); + } + } + + finish_fused_boolean_words(len, n_bytes, values, validity, nullability) +} + +fn fused_boolean_or_word_sources( + len: usize, + lhs_words: WordSource<'_>, + rhs_words: WordSource<'_>, + lhs_valid_words: WordSource<'_>, + rhs_valid_words: WordSource<'_>, + nullability: Nullability, +) -> VortexResult { + let n_bytes = len.div_ceil(8); + let n_words = n_bytes.div_ceil(8); + let full_bytes = n_bytes - n_bytes % 8; + let mut values = BufferMut::::with_capacity(n_words); + let mut validity = BufferMut::::with_capacity(n_words); + + for byte_offset in (0..full_bytes).step_by(8) { + let lhs = lhs_words.word_at(byte_offset, 8); + let rhs = rhs_words.word_at(byte_offset, 8); + let lhs_valid = lhs_valid_words.word_at(byte_offset, 8); + let rhs_valid = rhs_valid_words.word_at(byte_offset, 8); + + // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this + // loop plus the optional tail push emits at most `n_words` words. + unsafe { + values.push_unchecked(lhs | rhs); + validity + .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs)); + } + } + + if full_bytes != n_bytes { + let tail_len = n_bytes - full_bytes; + let lhs = lhs_words.word_at(full_bytes, tail_len); + let rhs = rhs_words.word_at(full_bytes, tail_len); + let lhs_valid = lhs_valid_words.word_at(full_bytes, tail_len); + let rhs_valid = rhs_valid_words.word_at(full_bytes, tail_len); + + // SAFETY: see the loop safety comment above. + unsafe { + values.push_unchecked(lhs | rhs); + validity + .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs)); + } + } + + finish_fused_boolean_words(len, n_bytes, values, validity, nullability) +} + +fn finish_fused_boolean_words( + len: usize, + n_bytes: usize, + values: BufferMut, + validity: BufferMut, + nullability: Nullability, +) -> VortexResult { + let mut values = values.into_byte_buffer(); + values.truncate(n_bytes); + let mut validity = validity.into_byte_buffer(); + validity.truncate(n_bytes); + Ok(BoolArray::try_new( + BitBuffer::new(values.freeze(), len), + Validity::from_mask( + Mask::from_buffer(BitBuffer::new(validity.freeze(), len)), + nullability, + ), + )? + .into_array()) +} + +fn fused_boolean_words( + len: usize, + lhs_words: L, + rhs_words: R, + lhs_valid_words: LV, + rhs_valid_words: RV, + operator: Operator, + nullability: Nullability, +) -> VortexResult +where + L: Iterator, + R: Iterator, + LV: Iterator, + RV: Iterator, +{ + match operator { + Operator::And => fused_boolean_and_words( + len, + lhs_words, + rhs_words, + lhs_valid_words, + rhs_valid_words, + nullability, + ), + Operator::Or => fused_boolean_or_words( + len, + lhs_words, + rhs_words, + lhs_valid_words, + rhs_valid_words, + nullability, + ), + other => vortex_bail!("Not a boolean operator: {other}"), + } +} + +fn fused_boolean_and_words( + len: usize, + lhs_words: L, + rhs_words: R, + lhs_valid_words: LV, + rhs_valid_words: RV, + nullability: Nullability, +) -> VortexResult +where + L: Iterator, + R: Iterator, + LV: Iterator, + RV: Iterator, +{ + let n_words = len.div_ceil(64); + let mut values = BufferMut::::with_capacity(n_words); + let mut validity = BufferMut::::with_capacity(n_words); + + for (((lhs, rhs), lhs_valid), rhs_valid) in lhs_words + .zip(rhs_words) + .zip(lhs_valid_words) + .zip(rhs_valid_words) + .take(n_words) + { + // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this loop is + // capped at `n_words`. + unsafe { + values.push_unchecked(lhs & rhs); + validity + .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs)); + } + } + + finish_fused_boolean_words(len, len.div_ceil(8), values, validity, nullability) +} + +fn fused_boolean_or_words( + len: usize, + lhs_words: L, + rhs_words: R, + lhs_valid_words: LV, + rhs_valid_words: RV, + nullability: Nullability, +) -> VortexResult +where + L: Iterator, + R: Iterator, + LV: Iterator, + RV: Iterator, +{ + let n_words = len.div_ceil(64); + let mut values = BufferMut::::with_capacity(n_words); + let mut validity = BufferMut::::with_capacity(n_words); + + for (((lhs, rhs), lhs_valid), rhs_valid) in lhs_words + .zip(rhs_words) + .zip(lhs_valid_words) + .zip(rhs_valid_words) + .take(n_words) + { + // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this loop is + // capped at `n_words`. + unsafe { + values.push_unchecked(lhs | rhs); + validity + .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs)); + } + } + + finish_fused_boolean_words(len, len.div_ceil(8), values, validity, nullability) +} + +fn constant_bool_result(value: Option, len: usize, nullability: Nullability) -> ArrayRef { + let scalar = value + .map(|b| Scalar::bool(b, nullability)) + .unwrap_or_else(|| Scalar::null(DType::Bool(nullability))); + + ConstantArray::new(scalar, len).into_array() +} + +fn cast_bool_nullability(array: &ArrayRef, nullability: Nullability) -> VortexResult { + let dtype = DType::Bool(nullability); + if array.dtype() == &dtype { + Ok(array.clone()) + } else { + array.cast(dtype) + } +} + +fn boolean_nullability(lhs: &ArrayRef, rhs: &ArrayRef) -> Nullability { + lhs.dtype().nullability() | rhs.dtype().nullability() +} + +#[inline] +fn is_boolean_operator(operator: Operator) -> bool { + matches!(operator, Operator::And | Operator::Or) } #[cfg(test)] mod tests { use rstest::rstest; + use vortex_error::VortexResult; use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; + use crate::arrays::ConstantArray; + use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; #[expect(deprecated)] use crate::canonical::ToCanonical as _; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::scalar::Scalar; use crate::scalar_fn::fns::operators::Operator; + #[test] + fn test_kleene_truth_table() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let lhs = BoolArray::from_iter([ + Some(true), + Some(true), + Some(true), + Some(false), + Some(false), + Some(false), + None, + None, + None, + ]) + .into_array(); + let rhs = BoolArray::from_iter([ + Some(true), + Some(false), + None, + Some(true), + Some(false), + None, + Some(true), + Some(false), + None, + ]) + .into_array(); + + assert_arrays_eq!( + lhs.binary(rhs.clone(), Operator::And)?, + BoolArray::from_iter([ + Some(true), + Some(false), + None, + Some(false), + Some(false), + Some(false), + None, + Some(false), + None, + ]), + &mut ctx + ); + + assert_arrays_eq!( + lhs.binary(rhs, Operator::Or)?, + BoolArray::from_iter([ + Some(true), + Some(true), + Some(true), + Some(true), + Some(false), + None, + Some(true), + None, + None, + ]), + &mut ctx + ); + + Ok(()) + } + + #[test] + fn test_null_constant_kleene() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let lhs = BoolArray::from_iter([Some(false), Some(true), None]).into_array(); + let null = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), lhs.len()) + .into_array(); + + assert_arrays_eq!( + lhs.binary(null.clone(), Operator::And)?, + BoolArray::from_iter([Some(false), None, None]), + &mut ctx + ); + assert_arrays_eq!( + lhs.binary(null, Operator::Or)?, + BoolArray::from_iter([None, Some(true), None]), + &mut ctx + ); + + Ok(()) + } + #[rstest] #[case( BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(), @@ -145,22 +869,22 @@ mod tests { let r = r.to_bool().into_array(); let v0 = r - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .as_bool() .value(); let v1 = r - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .as_bool() .value(); let v2 = r - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap() .as_bool() .value(); let v3 = r - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(3, &mut array_session().create_execution_ctx()) .unwrap() .as_bool() .value(); @@ -189,22 +913,22 @@ mod tests { .into_array(); let v0 = r - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap() .as_bool() .value(); let v1 = r - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap() .as_bool() .value(); let v2 = r - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut array_session().create_execution_ctx()) .unwrap() .as_bool() .value(); let v3 = r - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(3, &mut array_session().create_execution_ctx()) .unwrap() .as_bool() .value(); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare.rs b/vortex-array/src/scalar_fn/fns/binary/compare.rs index 5d4509fe406..b7a53c21d62 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare.rs @@ -245,14 +245,14 @@ mod tests { use std::sync::Arc; use rstest::rstest; + use vortex_buffer::BitBuffer; use vortex_buffer::buffer; + use vortex_error::VortexExpect; use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; @@ -280,71 +280,70 @@ mod tests { #[test] fn test_bool_basic_comparisons() { - use vortex_buffer::BitBuffer; - + let ctx = &mut array_session().create_execution_ctx(); let arr = BoolArray::new( BitBuffer::from_iter([true, true, false, true, false]), Validity::from_iter([false, true, true, true, true]), ); - #[expect(deprecated)] let matches = arr .clone() .into_array() .binary(arr.clone().into_array(), Operator::Eq) .unwrap() - .to_bool(); - assert_eq!(to_int_indices(matches).unwrap(), [1u64, 2, 3, 4]); + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [1u64, 2, 3, 4]); - #[expect(deprecated)] let matches = arr .clone() .into_array() .binary(arr.clone().into_array(), Operator::NotEq) .unwrap() - .to_bool(); + .execute::(ctx) + .vortex_expect("must be a bool array"); let empty: [u64; 0] = []; - assert_eq!(to_int_indices(matches).unwrap(), empty); + assert_eq!(to_int_indices(matches, ctx).unwrap(), empty); let other = BoolArray::new( BitBuffer::from_iter([false, false, false, true, true]), Validity::from_iter([false, true, true, true, true]), ); - #[expect(deprecated)] let matches = arr .clone() .into_array() .binary(other.clone().into_array(), Operator::Lte) .unwrap() - .to_bool(); - assert_eq!(to_int_indices(matches).unwrap(), [2u64, 3, 4]); + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [2u64, 3, 4]); - #[expect(deprecated)] let matches = arr .clone() .into_array() .binary(other.clone().into_array(), Operator::Lt) .unwrap() - .to_bool(); - assert_eq!(to_int_indices(matches).unwrap(), [4u64]); + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [4u64]); - #[expect(deprecated)] let matches = other .clone() .into_array() .binary(arr.clone().into_array(), Operator::Gte) .unwrap() - .to_bool(); - assert_eq!(to_int_indices(matches).unwrap(), [2u64, 3, 4]); + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [2u64, 3, 4]); - #[expect(deprecated)] let matches = other .into_array() .binary(arr.into_array(), Operator::Gt) .unwrap() - .to_bool(); - assert_eq!(to_int_indices(matches).unwrap(), [4u64]); + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [4u64]); } #[test] @@ -358,7 +357,7 @@ mod tests { .unwrap(); assert_eq!(result.len(), 10); let scalar = result - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert_eq!(scalar.as_bool().value(), Some(false)); } @@ -369,13 +368,15 @@ mod tests { #[case(VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array())] #[case(VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array())] fn arrow_compare_different_encodings(#[case] left: ArrayRef, #[case] right: ArrayRef) { + let mut ctx = array_session().create_execution_ctx(); let res = left.binary(right, Operator::Eq).unwrap(); let expected = BoolArray::from_iter([true, true]); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut ctx); } #[test] fn test_list_array_comparison() { + let mut ctx = array_session().create_execution_ctx(); let values1 = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); let offsets1 = PrimitiveArray::from_iter([0i32, 2, 4, 6]); let list1 = ListArray::try_new( @@ -400,7 +401,7 @@ mod tests { .binary(list2.clone().into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); let result = list1 .clone() @@ -408,18 +409,19 @@ mod tests { .binary(list2.clone().into_array(), Operator::NotEq) .unwrap(); let expected = BoolArray::from_iter([false, false, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); let result = list1 .into_array() .binary(list2.into_array(), Operator::Lt) .unwrap(); let expected = BoolArray::from_iter([false, false, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_list_array_constant_comparison() { + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); let offsets = PrimitiveArray::from_iter([0i32, 2, 4, 6]); let list = ListArray::try_new( @@ -441,11 +443,12 @@ mod tests { .binary(constant.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([false, true, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_struct_array_comparison() { + let mut ctx = array_session().create_execution_ctx(); let bool_field1 = BoolArray::from_iter([Some(true), Some(false), Some(true)]); let int_field1 = PrimitiveArray::from_iter([1i32, 2, 3]); @@ -470,18 +473,19 @@ mod tests { .binary(struct2.clone().into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); let result = struct1 .into_array() .binary(struct2.into_array(), Operator::Gt) .unwrap(); let expected = BoolArray::from_iter([false, false, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_empty_struct_compare() { + let mut ctx = array_session().create_execution_ctx(); let empty1 = StructArray::try_new( FieldNames::from(Vec::::new()), Vec::new(), @@ -503,13 +507,14 @@ mod tests { .binary(empty2.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([true, true, true, true, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } /// Regression test: comparing struct arrays where the same logical field is backed by /// different Vortex encodings (VarBinArray vs VarBinViewArray) must not panic. #[test] fn struct_compare_mixed_binary_encodings() { + let mut ctx = array_session().create_execution_ctx(); // LHS: struct with a VarBinArray (offset-based) binary field let bin_field1 = VarBinArray::from(vec![ "apple".as_bytes(), @@ -531,7 +536,7 @@ mod tests { .binary(struct2.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } /// Regression test: `scalar_cmp` must error when comparing scalars with incompatible @@ -565,6 +570,7 @@ mod tests { #[test] fn test_empty_list() { + let ctx = &mut array_session().create_execution_ctx(); let list = ListViewArray::new( BoolArray::from_iter(Vec::::new()).into_array(), buffer![0i32, 0i32, 0i32].into_array(), @@ -577,23 +583,8 @@ mod tests { .into_array() .binary(list.into_array(), Operator::Eq) .unwrap(); - assert!( - result - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .is_valid() - ); - assert!( - result - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .is_valid() - ); - assert!( - result - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .is_valid() - ); + assert!(result.execute_scalar(0, ctx).unwrap().is_valid()); + assert!(result.execute_scalar(1, ctx).unwrap().is_valid()); + assert!(result.execute_scalar(2, ctx).unwrap().is_valid()); } } diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 1c860cb75b5..156540a2820 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -17,34 +17,33 @@ use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; -use crate::expr::StatsCatalog; +use crate::dtype::Nullability; use crate::expr::and; -use crate::expr::and_collect; -use crate::expr::eq; use crate::expr::expression::Expression; -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::or_collect; -use crate::expr::stats::Stat; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::SimplifyCtx; +use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; -pub(crate) mod boolean; -pub(crate) use boolean::*; +pub mod boolean; +pub use boolean::BooleanExecuteAdaptor; +pub use boolean::BooleanKernel; +pub(crate) use boolean::execute_boolean; +pub use boolean::kleene_boolean_buffer_scalar; +pub use boolean::kleene_boolean_buffers; mod compare; pub use compare::*; mod numeric; pub(crate) use numeric::*; use crate::scalar::NumericOperator; +use crate::scalar::Scalar; #[derive(Clone)] pub struct Binary; @@ -165,108 +164,73 @@ impl ScalarFnVTable for Binary { } } - fn stat_falsification( + fn simplify_untyped( &self, operator: &Operator, expr: &Expression, - catalog: &dyn StatsCatalog, - ) -> Option { - // Wrap another predicate with an optional NaNCount check, if the stat is available. - // - // For example, regular pruning conversion for `A >= B` would be - // - // A.max < B.min - // - // With NaN predicate introduction, we'd conjunct it with a check for NaNCount, resulting - // in: - // - // (A.nan_count = 0) AND (B.nan_count = 0) AND A.max < B.min - // - // Non-floating point column and literal expressions should be unaffected as they do not - // have a nan_count statistic defined. - fn with_nan_predicate( - lhs: &Expression, - rhs: &Expression, - value_predicate: Expression, - catalog: &dyn StatsCatalog, - ) -> Expression { - let nan_predicate = and_collect( - lhs.stat_expression(Stat::NaNCount, catalog) - .into_iter() - .chain(rhs.stat_expression(Stat::NaNCount, catalog)) - .map(|nans| eq(nans, lit(0u64))), - ); - - if let Some(nan_check) = nan_predicate { - and(nan_check, value_predicate) - } else { - value_predicate - } - } - + ) -> VortexResult> { let lhs = expr.child(0); let rhs = expr.child(1); - match operator { - Operator::Eq => { - let min_lhs = lhs.stat_min(catalog); - let max_lhs = lhs.stat_max(catalog); - - let min_rhs = rhs.stat_min(catalog); - let max_rhs = rhs.stat_max(catalog); - let left = min_lhs.zip(max_rhs).map(|(a, b)| gt(a, b)); - let right = min_rhs.zip(max_lhs).map(|(a, b)| gt(a, b)); - - let min_max_check = or_collect(left.into_iter().chain(right))?; - - // NaN is not captured by the min/max stat, so we must check NaNCount before pruning - Some(with_nan_predicate(lhs, rhs, min_max_check, catalog)) - } - Operator::NotEq => { - let min_lhs = lhs.stat_min(catalog)?; - let max_lhs = lhs.stat_max(catalog)?; + let bool_literal = |expr: &Expression| { + expr.as_opt::()? + .as_bool_opt() + .map(|value| value.value()) + }; - let min_rhs = rhs.stat_min(catalog)?; - let max_rhs = rhs.stat_max(catalog)?; - - let min_max_check = and(eq(min_lhs, max_rhs), eq(max_lhs, min_rhs)); - - Some(with_nan_predicate(lhs, rhs, min_max_check, catalog)) - } - Operator::Gt => { - let min_max_check = lt_eq(lhs.stat_max(catalog)?, rhs.stat_min(catalog)?); - - Some(with_nan_predicate(lhs, rhs, min_max_check, catalog)) - } - Operator::Gte => { - // NaN is not captured by the min/max stat, so we must check NaNCount before pruning - let min_max_check = lt(lhs.stat_max(catalog)?, rhs.stat_min(catalog)?); - - Some(with_nan_predicate(lhs, rhs, min_max_check, catalog)) - } - Operator::Lt => { - // NaN is not captured by the min/max stat, so we must check NaNCount before pruning - let min_max_check = gt_eq(lhs.stat_min(catalog)?, rhs.stat_max(catalog)?); + // AND/OR use Kleene three-valued logic. `None` below is a boolean null. + // + // AND: + // - false AND x => false + // - true AND x => x + // - null AND null => null + // + // OR: + // - true OR x => true + // - false OR x => x + // - null OR null => null + // + // Other null cases either fall out of the identity/annihilator rules + // above (`null AND true`, `null OR false`) or cannot be simplified under + // Kleene semantics (`null AND x`, `null OR x` for non-literal `x`). + Ok(match operator { + Operator::And => match (bool_literal(lhs), bool_literal(rhs)) { + (Some(Some(false)), _) | (_, Some(Some(false))) => Some(lit(false)), + (Some(Some(true)), _) => Some(rhs.clone()), + (_, Some(Some(true))) => Some(lhs.clone()), + (Some(None), Some(None)) => Some(lhs.clone()), + _ => None, + }, + Operator::Or => match (bool_literal(lhs), bool_literal(rhs)) { + (Some(Some(true)), _) | (_, Some(Some(true))) => Some(lit(true)), + (Some(Some(false)), _) => Some(rhs.clone()), + (_, Some(Some(false))) => Some(lhs.clone()), + (Some(None), Some(None)) => Some(lhs.clone()), + _ => None, + }, + _ => None, + }) + } - Some(with_nan_predicate(lhs, rhs, min_max_check, catalog)) - } - Operator::Lte => { - // NaN is not captured by the min/max stat, so we must check NaNCount before pruning - let min_max_check = gt(lhs.stat_min(catalog)?, rhs.stat_max(catalog)?); + fn simplify( + &self, + operator: &Operator, + expr: &Expression, + ctx: &dyn SimplifyCtx, + ) -> VortexResult> { + let is_literal_null = + |expr: &Expression| expr.as_opt::().is_some_and(Scalar::is_null); - Some(with_nan_predicate(lhs, rhs, min_max_check, catalog)) - } - Operator::And => or_collect( - lhs.stat_falsification(catalog) - .into_iter() - .chain(rhs.stat_falsification(catalog)), - ), - Operator::Or => Some(and( - lhs.stat_falsification(catalog)?, - rhs.stat_falsification(catalog)?, - )), - Operator::Add | Operator::Sub | Operator::Mul | Operator::Div => None, + if operator.is_comparison() + && (is_literal_null(expr.child(0)) || is_literal_null(expr.child(1))) + { + // Validate the comparison before reducing it. This preserves type + // errors for expressions like `int_col = null_utf8`. + ctx.return_dtype(expr)?; + return Ok(Some(lit(Scalar::null(DType::Bool(Nullability::Nullable))))); } + + Ok(None) } fn validity( @@ -314,19 +278,25 @@ impl ScalarFnVTable for Binary { #[cfg(test)] mod tests { use vortex_error::VortexExpect; + use vortex_error::VortexResult; use super::*; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; + use crate::dtype::PType; use crate::expr::Expression; use crate::expr::and_collect; use crate::expr::col; + use crate::expr::eq; + 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_eq; use crate::expr::or; use crate::expr::or_collect; @@ -442,6 +412,36 @@ mod tests { ); } + #[test] + fn comparison_with_typed_null_simplifies_after_type_check() -> VortexResult<()> { + let dtype = test_harness::struct_dtype(); + + let expr = eq( + col("col1"), + lit(Scalar::null(DType::Primitive( + PType::U16, + Nullability::Nullable, + ))), + ); + + assert_eq!( + expr.optimize_recursive(&dtype)?, + lit(Scalar::null(DType::Bool(Nullability::Nullable))) + ); + Ok(()) + } + + #[test] + fn comparison_with_incompatible_null_still_type_checks() { + let dtype = test_harness::struct_dtype(); + let expr = eq( + col("col1"), + lit(Scalar::null(DType::Utf8(Nullability::Nullable))), + ); + + assert!(expr.optimize_recursive(&dtype).is_err()); + } + #[test] fn test_display_print() { let expr = gt(lit(1), lit(2)); @@ -499,7 +499,7 @@ mod tests { let result_equal = lhs_struct.binary(rhs_struct_equal, Operator::Eq).unwrap(); assert_eq!( result_equal - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .vortex_expect("value"), Scalar::bool(true, Nullability::NonNullable), "Equal structs should be equal" @@ -510,7 +510,7 @@ mod tests { .unwrap(); assert_eq!( result_different - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .vortex_expect("value"), Scalar::bool(false, Nullability::NonNullable), "Different structs should not be equal" @@ -519,6 +519,7 @@ mod tests { #[test] fn test_or_kleene_validity() { + let mut ctx = array_session().create_execution_ctx(); use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::StructArray; @@ -537,11 +538,16 @@ mod tests { let expr = or(col("a"), col("b")); let result = struct_arr.apply(&expr).unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter([Some(true)]).into_array()) + assert_arrays_eq!( + result, + BoolArray::from_iter([Some(true)]).into_array(), + &mut ctx + ) } #[test] fn test_scalar_subtract_unsigned() { + let mut ctx = array_session().create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -551,11 +557,12 @@ mod tests { let values = buffer![1u16, 2, 3].into_array(); let rhs = ConstantArray::new(Scalar::from(1u16), 3).into_array(); let result = values.binary(rhs, Operator::Sub).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2])); + assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2]), &mut ctx); } #[test] fn test_scalar_subtract_signed() { + let mut ctx = array_session().create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -565,11 +572,12 @@ mod tests { let values = buffer![1i64, 2, 3].into_array(); let rhs = ConstantArray::new(Scalar::from(-1i64), 3).into_array(); let result = values.binary(rhs, Operator::Sub).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4])); + assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4]), &mut ctx); } #[test] fn test_scalar_subtract_nullable() { + let mut ctx = array_session().create_execution_ctx(); use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -579,12 +587,14 @@ mod tests { let result = values.into_array().binary(rhs, Operator::Sub).unwrap(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]) + PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]), + &mut ctx ); } #[test] fn test_scalar_subtract_float() { + let mut ctx = array_session().create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -594,7 +604,11 @@ mod tests { let values = buffer![1.0f64, 2.0, 3.0].into_array(); let rhs = ConstantArray::new(Scalar::from(-1f64), 3).into_array(); let result = values.binary(rhs, Operator::Sub).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([2.0f64, 3.0, 4.0])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([2.0f64, 3.0, 4.0]), + &mut ctx + ); } #[test] diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index 2ef58fb5209..e95627155e9 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric.rs @@ -1,77 +1,924 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::Constant; use crate::arrays::ConstantArray; -use crate::arrow::Datum; -use crate::arrow::from_arrow_columnar; -use crate::executor::ExecutionCtx; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::PType; +use crate::dtype::half::f16; +use crate::match_each_native_ptype; use crate::scalar::NumericOperator; +use crate::scalar::Scalar; +use crate::validity::Validity; + +struct CheckedAdd; + +struct CheckedSub; + +struct CheckedMul; + +struct CheckedDiv; + +trait CheckedPrimitiveOp: Sized { + const ERROR: &'static str; + const CHECKED_VALUE_LOOP: bool = false; + + // The vectorizable kernels need an overflowing-style result, not + // `Option`: every lane can write a value unconditionally while reducing + // the error flag separately. `checked` is still available for scalar and + // one-pass paths where early exit is the right shape. + fn apply(lhs: T, rhs: T) -> (T, bool); + + #[inline(always)] + fn checked(lhs: T, rhs: T) -> Option { + let (value, failed) = Self::apply(lhs, rhs); + if failed { None } else { Some(value) } + } +} + +impl CheckedPrimitiveOp for CheckedAdd { + const ERROR: &'static str = "integer overflow in checked add"; + + #[inline(always)] + fn apply(lhs: T, rhs: T) -> (T, bool) { + (lhs.add_value(rhs), lhs.add_error(rhs)) + } +} + +impl CheckedPrimitiveOp for CheckedSub { + const ERROR: &'static str = "integer overflow in checked sub"; + + #[inline(always)] + fn apply(lhs: T, rhs: T) -> (T, bool) { + (lhs.sub_value(rhs), lhs.sub_error(rhs)) + } +} + +impl CheckedPrimitiveOp for CheckedMul { + const ERROR: &'static str = "integer overflow in checked mul"; + + #[inline(always)] + fn apply(lhs: T, rhs: T) -> (T, bool) { + (lhs.mul_value(rhs), lhs.mul_error(rhs)) + } +} + +impl CheckedPrimitiveOp for CheckedDiv { + const ERROR: &'static str = "integer division by zero or overflow in checked div"; + // Integer division still lowers to scalar divides, so the split + // value/error-scan loop used to auto-vectorize add/sub/mul only adds a + // second full scan. Use the generic branchy checked value loop for integer + // division, matching Arrow/Velox. Float division has no checked errors and + // stays on the split/vectorizable default path. + const CHECKED_VALUE_LOOP: bool = T::DIV_CHECKS_IN_VALUE_LOOP; + + #[inline(always)] + fn apply(lhs: T, rhs: T) -> (T, bool) { + let failed = lhs.div_error(rhs); + let value = if failed { + T::default() + } else { + lhs.div_value(rhs) + }; + (value, failed) + } + + #[inline(always)] + fn checked(lhs: T, rhs: T) -> Option { + lhs.div_checked(rhs) + } +} /// Execute a numeric operation between two arrays. -/// -/// This is the entry point for numeric operations from the binary expression. -/// Handles constant-constant directly, otherwise falls back to Arrow. pub(crate) fn execute_numeric( lhs: &ArrayRef, rhs: &ArrayRef, op: NumericOperator, ctx: &mut ExecutionCtx, ) -> VortexResult { - if let Some(result) = constant_numeric(lhs, rhs, op)? { - return Ok(result); + let ptype = PType::try_from(lhs.dtype())?; + if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) { + vortex_bail!( + "numeric operator requires matching primitive types, got {} and {}", + lhs.dtype(), + rhs.dtype() + ); } - arrow_numeric(lhs, rhs, op, ctx) + + match_each_native_ptype!(ptype, |T| { + match op { + NumericOperator::Add => execute_checked_typed::(lhs, rhs, ctx), + NumericOperator::Sub => execute_checked_typed::(lhs, rhs, ctx), + NumericOperator::Mul => execute_checked_typed::(lhs, rhs, ctx), + NumericOperator::Div => execute_checked_typed::(lhs, rhs, ctx), + } + }) } -/// Implementation of numeric operations using the Arrow crate. -pub(crate) fn arrow_numeric( +fn execute_checked_typed( lhs: &ArrayRef, rhs: &ArrayRef, - operator: NumericOperator, ctx: &mut ExecutionCtx, -) -> VortexResult { - let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable(); +) -> VortexResult +where + T: NativePType, + Op: CheckedPrimitiveOp, + Scalar: From, + Scalar: From>, +{ + let result_dtype = lhs + .dtype() + .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); + let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; + let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; let len = lhs.len(); + if len != rhs.len() { + vortex_bail!( + "numeric operator requires equal lengths, got {} and {}", + len, + rhs.len() + ); + } - let left = Datum::try_new(lhs, ctx)?; - let right = Datum::try_new_with_target_datatype(rhs, left.data_type(), ctx)?; + let validity = lhs.validity().and(rhs.validity())?; + let valid_rows = validity.execute_mask(len, ctx)?; - let array = match operator { - NumericOperator::Add => arrow_arith::numeric::add(&left, &right)?, - NumericOperator::Sub => arrow_arith::numeric::sub(&left, &right)?, - NumericOperator::Mul => arrow_arith::numeric::mul(&left, &right)?, - NumericOperator::Div => arrow_arith::numeric::div(&left, &right)?, + let checked = match (&lhs, &rhs) { + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => checked_array_array::(lhs, rhs, &valid_rows), + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => checked_array_constant::(lhs, *rhs, &valid_rows), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => checked_constant_array::(*lhs, rhs, &valid_rows), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => { + let value = Op::checked(*lhs, *rhs) + .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; + return Ok(constant_result_array(value, len, &result_dtype)); + } + (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { + CheckedValues::zeroed(len) + } }; + check_numeric_errors(checked.failed, Op::ERROR)?; - from_arrow_columnar(array.as_ref(), len, nullable, ctx) + primitive_result_array::(checked.values, validity, &result_dtype) } -fn constant_numeric( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: NumericOperator, -) -> VortexResult> { - let (Some(lhs), Some(rhs)) = (lhs.as_opt::(), rhs.as_opt::()) else { - return Ok(None); +fn primitive_result_array( + values: Buffer, + validity: Validity, + dtype: &DType, +) -> VortexResult { + let array = PrimitiveArray::new(values, validity).into_array(); + if array.dtype() == dtype { + return Ok(array); + } + array.cast(dtype.clone()) +} + +fn constant_result_array(value: T, len: usize, dtype: &DType) -> ArrayRef +where + T: NativePType, + Scalar: From + From>, +{ + if dtype.is_nullable() { + ConstantArray::new(Some(value), len).into_array() + } else { + ConstantArray::new(value, len).into_array() + } +} + +enum PrimitiveOperand { + Array { + values: Buffer, + validity: Validity, + }, + Constant { + value: T, + len: usize, + validity: Validity, + }, + Null(usize), +} + +impl PrimitiveOperand { + fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(constant) = array.as_opt::() { + return Ok( + match constant.scalar().as_primitive().try_typed_value::()? { + Some(value) => Self::Constant { + value, + len: array.len(), + validity: if constant.scalar().dtype().is_nullable() { + Validity::AllValid + } else { + Validity::NonNullable + }, + }, + None => Self::Null(array.len()), + }, + ); + } + + let array = array.clone().execute::(ctx)?; + let validity = array.validity()?; + let values = array.into_buffer::(); + Ok(Self::Array { values, validity }) + } + + fn len(&self) -> usize { + match self { + Self::Array { values, .. } => values.len(), + Self::Constant { len, .. } | Self::Null(len) => *len, + } + } + + fn validity(&self) -> Validity { + match self { + Self::Array { validity, .. } => validity.clone(), + Self::Constant { validity, .. } => validity.clone(), + Self::Null(_) => Validity::AllInvalid, + } + } +} + +struct CheckedValues { + values: Buffer, + failed: bool, +} + +impl CheckedValues { + fn zeroed(len: usize) -> Self { + Self { + values: Buffer::::zeroed(len), + failed: false, + } + } + + fn failed(len: usize) -> Self { + Self { + values: Buffer::::zeroed(len), + failed: true, + } + } +} + +fn checked_array_array(lhs: &[T], rhs: &[T], valid_rows: &Mask) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + debug_assert_eq!(lhs.len(), rhs.len()); + + match valid_rows.bit_buffer() { + AllOr::All if Op::CHECKED_VALUE_LOOP => checked_array_array_one_pass::(lhs, rhs), + AllOr::All => checked_array_array_all_lanes::(lhs, rhs), + AllOr::None => CheckedValues::zeroed(lhs.len()), + AllOr::Some(valid_bits) if Op::CHECKED_VALUE_LOOP => { + checked_array_array_valid_lanes_one_pass::(lhs, rhs, valid_bits) + } + AllOr::Some(valid_bits) => checked_array_array_valid_lanes::(lhs, rhs, valid_bits), + } +} + +fn checked_array_constant(lhs: &[T], rhs: T, valid_rows: &Mask) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + match valid_rows.bit_buffer() { + AllOr::All if Op::CHECKED_VALUE_LOOP => checked_array_constant_one_pass::(lhs, rhs), + AllOr::All => checked_array_constant_all_lanes::(lhs, rhs), + AllOr::None => CheckedValues::zeroed(lhs.len()), + AllOr::Some(valid_bits) if Op::CHECKED_VALUE_LOOP => { + checked_array_constant_valid_lanes_one_pass::(lhs, rhs, valid_bits) + } + AllOr::Some(valid_bits) => { + checked_array_constant_valid_lanes::(lhs, rhs, valid_bits) + } + } +} + +fn checked_constant_array(lhs: T, rhs: &[T], valid_rows: &Mask) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + match valid_rows.bit_buffer() { + AllOr::All if Op::CHECKED_VALUE_LOOP => checked_constant_array_one_pass::(lhs, rhs), + AllOr::All => checked_constant_array_all_lanes::(lhs, rhs), + AllOr::None => CheckedValues::zeroed(rhs.len()), + AllOr::Some(valid_bits) if Op::CHECKED_VALUE_LOOP => { + checked_constant_array_valid_lanes_one_pass::(lhs, rhs, valid_bits) + } + AllOr::Some(valid_bits) => { + checked_constant_array_valid_lanes::(lhs, rhs, valid_bits) + } + } +} + +fn checked_array_array_all_lanes(lhs: &[T], rhs: &[T]) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs[idx])) +} + +fn checked_array_array_valid_lanes( + lhs: &[T], + rhs: &[T], + valid_bits: &BitBuffer, +) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + let mut checked = collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs[idx])); + + checked.failed = checked.failed + && any_valid_error(lhs.len(), valid_bits, |idx| Op::apply(lhs[idx], rhs[idx]).1); + checked +} + +fn checked_array_constant_all_lanes(lhs: &[T], rhs: T) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs)) +} + +fn checked_array_constant_valid_lanes( + lhs: &[T], + rhs: T, + valid_bits: &BitBuffer, +) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + let mut checked = collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs)); + + checked.failed = + checked.failed && any_valid_error(lhs.len(), valid_bits, |idx| Op::apply(lhs[idx], rhs).1); + checked +} + +fn checked_constant_array_all_lanes(lhs: T, rhs: &[T]) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + collect_all_lanes(rhs.len(), |idx| Op::apply(lhs, rhs[idx])) +} + +fn checked_constant_array_valid_lanes( + lhs: T, + rhs: &[T], + valid_bits: &BitBuffer, +) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + let mut checked = collect_all_lanes(rhs.len(), |idx| Op::apply(lhs, rhs[idx])); + + checked.failed = + checked.failed && any_valid_error(rhs.len(), valid_bits, |idx| Op::apply(lhs, rhs[idx]).1); + checked +} + +fn collect_all_lanes(len: usize, mut value_and_error_at: F) -> CheckedValues +where + T: NativePType, + F: FnMut(usize) -> (T, bool), +{ + let mut values = BufferMut::::with_capacity(len); + let slots = values.spare_capacity_mut().as_mut_ptr(); + let mut failed = false; + for idx in 0..len { + let (value, is_error) = value_and_error_at(idx); + failed |= is_error; + // SAFETY: `idx < len <= capacity`, and each slot is written once. + unsafe { slots.add(idx).write(std::mem::MaybeUninit::new(value)) }; + } + // SAFETY: every slot in `0..len` was initialized above. + unsafe { values.set_len(len) }; + CheckedValues { + values: values.freeze(), + failed, + } +} + +fn checked_array_array_one_pass(lhs: &[T], rhs: &[T]) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + checked_all_lanes(lhs.len(), |idx| Op::checked(lhs[idx], rhs[idx])) +} + +fn checked_array_array_valid_lanes_one_pass( + lhs: &[T], + rhs: &[T], + valid_bits: &BitBuffer, +) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + checked_valid_lanes(lhs.len(), valid_bits, |idx| Op::checked(lhs[idx], rhs[idx])) +} + +fn checked_array_constant_one_pass(lhs: &[T], rhs: T) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + checked_all_lanes(lhs.len(), |idx| Op::checked(lhs[idx], rhs)) +} + +fn checked_array_constant_valid_lanes_one_pass( + lhs: &[T], + rhs: T, + valid_bits: &BitBuffer, +) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + checked_valid_lanes(lhs.len(), valid_bits, |idx| Op::checked(lhs[idx], rhs)) +} + +fn checked_constant_array_one_pass(lhs: T, rhs: &[T]) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + checked_all_lanes(rhs.len(), |idx| Op::checked(lhs, rhs[idx])) +} + +fn checked_constant_array_valid_lanes_one_pass( + lhs: T, + rhs: &[T], + valid_bits: &BitBuffer, +) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + checked_valid_lanes(rhs.len(), valid_bits, |idx| Op::checked(lhs, rhs[idx])) +} + +// Checked one-pass ops delay early exit until the end of a small block. This +// keeps the loop generic while avoiding a branch-driven exit decision on every +// lane; it is deliberately independent of mask density or input length. +const CHECKED_BLOCK_LANES: usize = 16; + +fn checked_all_lanes(len: usize, mut checked_at: F) -> CheckedValues +where + T: NativePType, + F: FnMut(usize) -> Option, +{ + let mut values = BufferMut::::with_capacity(len); + let mut base = 0; + + while base + CHECKED_BLOCK_LANES <= len { + let mut block_failed = false; + for idx in base..base + CHECKED_BLOCK_LANES { + match checked_at(idx) { + Some(value) => { + // SAFETY: the buffer is allocated with capacity `len`, and + // this loop pushes at most one value for each `idx`. + unsafe { values.push_unchecked(value) }; + } + None => { + block_failed = true; + // SAFETY: the buffer is allocated with capacity `len`, and + // this loop pushes at most one value for each `idx`. + unsafe { values.push_unchecked(T::default()) }; + } + } + } + + if block_failed { + return CheckedValues::failed(len); + } + base += CHECKED_BLOCK_LANES; + } + + for idx in base..len { + let Some(value) = checked_at(idx) else { + return CheckedValues::failed(len); + }; + // SAFETY: the buffer is allocated with capacity `len`, and this loop + // pushes at most one value for each `idx`. + unsafe { values.push_unchecked(value) }; + } + + CheckedValues { + values: values.freeze(), + failed: false, + } +} + +fn checked_valid_lanes( + len: usize, + valid_bits: &BitBuffer, + mut checked_at: F, +) -> CheckedValues +where + T: NativePType, + F: FnMut(usize) -> Option, +{ + let mut values = BufferMut::::zeroed(len); + let mut failed = false; + { + let values = values.as_mut_slice(); + for_each_valid_idx(len, valid_bits, |idx| { + let Some(value) = checked_at(idx) else { + failed = true; + return false; + }; + values[idx] = value; + true + }); + } + + CheckedValues { + values: values.freeze(), + failed, + } +} + +fn any_valid_error(len: usize, valid_bits: &BitBuffer, is_error: F) -> bool +where + F: Fn(usize) -> bool, +{ + !for_each_valid_idx(len, valid_bits, |idx| !is_error(idx)) +} + +fn for_each_valid_idx(len: usize, valid_bits: &BitBuffer, mut f: F) -> bool +where + F: FnMut(usize) -> bool, +{ + debug_assert_eq!(len, valid_bits.len()); + + for (word_idx, valid_word) in valid_bits.chunks().iter_padded().enumerate() { + if valid_word == 0 { + continue; + } + + let offset = word_idx * 64; + let lanes = len.saturating_sub(offset).min(64); + if lanes == 64 && valid_word == u64::MAX { + for bit_idx in 0..64 { + if !f(offset + bit_idx) { + return false; + } + } + continue; + } + + let mut valid_word = if lanes == 64 { + valid_word + } else { + valid_word & ((1u64 << lanes) - 1) + }; + while valid_word != 0 { + let bit_idx = valid_word.trailing_zeros() as usize; + if !f(offset + bit_idx) { + return false; + } + valid_word &= valid_word - 1; + } + } + + true +} + +trait CheckedArithmetic: NativePType { + const DIV_CHECKS_IN_VALUE_LOOP: bool; + + fn add_value(self, rhs: Self) -> Self; + fn add_error(self, rhs: Self) -> bool; + fn sub_value(self, rhs: Self) -> Self; + fn sub_error(self, rhs: Self) -> bool; + fn mul_value(self, rhs: Self) -> Self; + fn mul_error(self, rhs: Self) -> bool; + fn div_value(self, rhs: Self) -> Self; + fn div_error(self, rhs: Self) -> bool; + fn div_checked(self, rhs: Self) -> Option; +} + +macro_rules! impl_checked_unsigned { + ($ty:ty,widening_mul: $wide:ty) => { + impl CheckedArithmetic for $ty { + const DIV_CHECKS_IN_VALUE_LOOP: bool = true; + + #[inline(always)] + fn add_value(self, rhs: Self) -> Self { + self.wrapping_add(rhs) + } + + #[inline(always)] + fn add_error(self, rhs: Self) -> bool { + self > <$ty>::MAX - rhs + } + + #[inline(always)] + fn sub_value(self, rhs: Self) -> Self { + self.wrapping_sub(rhs) + } + + #[inline(always)] + fn sub_error(self, rhs: Self) -> bool { + self < rhs + } + + #[inline(always)] + fn mul_value(self, rhs: Self) -> Self { + self.wrapping_mul(rhs) + } + + #[inline(always)] + fn mul_error(self, rhs: Self) -> bool { + (self as $wide) * (rhs as $wide) > <$ty>::MAX as $wide + } + + #[inline(always)] + fn div_value(self, rhs: Self) -> Self { + self / rhs + } + + #[inline(always)] + fn div_error(self, rhs: Self) -> bool { + rhs == 0 + } + + #[inline(always)] + fn div_checked(self, rhs: Self) -> Option { + self.checked_div(rhs) + } + } + }; + ($ty:ty,overflowing_mul) => { + impl CheckedArithmetic for $ty { + const DIV_CHECKS_IN_VALUE_LOOP: bool = true; + + #[inline(always)] + fn add_value(self, rhs: Self) -> Self { + self.wrapping_add(rhs) + } + + #[inline(always)] + fn add_error(self, rhs: Self) -> bool { + self > <$ty>::MAX - rhs + } + + #[inline(always)] + fn sub_value(self, rhs: Self) -> Self { + self.wrapping_sub(rhs) + } + + #[inline(always)] + fn sub_error(self, rhs: Self) -> bool { + self < rhs + } + + #[inline(always)] + fn mul_value(self, rhs: Self) -> Self { + self.wrapping_mul(rhs) + } + + #[inline(always)] + fn mul_error(self, rhs: Self) -> bool { + self.overflowing_mul(rhs).1 + } + + #[inline(always)] + fn div_value(self, rhs: Self) -> Self { + self / rhs + } + + #[inline(always)] + fn div_error(self, rhs: Self) -> bool { + rhs == 0 + } + + #[inline(always)] + fn div_checked(self, rhs: Self) -> Option { + self.checked_div(rhs) + } + } + }; +} + +macro_rules! impl_checked_signed { + ($ty:ty,widening_mul: $wide:ty) => { + impl CheckedArithmetic for $ty { + const DIV_CHECKS_IN_VALUE_LOOP: bool = true; + + #[inline(always)] + fn add_value(self, rhs: Self) -> Self { + self.wrapping_add(rhs) + } + + #[inline(always)] + fn add_error(self, rhs: Self) -> bool { + let value = self.wrapping_add(rhs); + ((self ^ value) & (rhs ^ value)) < 0 + } + + #[inline(always)] + fn sub_value(self, rhs: Self) -> Self { + self.wrapping_sub(rhs) + } + + #[inline(always)] + fn sub_error(self, rhs: Self) -> bool { + let value = self.wrapping_sub(rhs); + ((self ^ rhs) & (self ^ value)) < 0 + } + + #[inline(always)] + fn mul_value(self, rhs: Self) -> Self { + self.wrapping_mul(rhs) + } + + #[inline(always)] + fn mul_error(self, rhs: Self) -> bool { + let product = (self as $wide) * (rhs as $wide); + product < <$ty>::MIN as $wide || product > <$ty>::MAX as $wide + } + + #[inline(always)] + fn div_value(self, rhs: Self) -> Self { + self / rhs + } + + #[inline(always)] + fn div_error(self, rhs: Self) -> bool { + rhs == 0 || (self == <$ty>::MIN && rhs == -1) + } + + #[inline(always)] + fn div_checked(self, rhs: Self) -> Option { + self.checked_div(rhs) + } + } }; + ($ty:ty,overflowing_mul) => { + impl CheckedArithmetic for $ty { + const DIV_CHECKS_IN_VALUE_LOOP: bool = true; + + #[inline(always)] + fn add_value(self, rhs: Self) -> Self { + self.wrapping_add(rhs) + } + + #[inline(always)] + fn add_error(self, rhs: Self) -> bool { + let value = self.wrapping_add(rhs); + ((self ^ value) & (rhs ^ value)) < 0 + } + + #[inline(always)] + fn sub_value(self, rhs: Self) -> Self { + self.wrapping_sub(rhs) + } + + #[inline(always)] + fn sub_error(self, rhs: Self) -> bool { + let value = self.wrapping_sub(rhs); + ((self ^ rhs) & (self ^ value)) < 0 + } - let Some(result) = lhs - .scalar() - .as_primitive() - .checked_binary_numeric(&rhs.scalar().as_primitive(), op) - else { - // Overflow detected — fall through to arrow_numeric which uses wrapping arithmetic. - return Ok(None); + #[inline(always)] + fn mul_value(self, rhs: Self) -> Self { + self.wrapping_mul(rhs) + } + + #[inline(always)] + fn mul_error(self, rhs: Self) -> bool { + self.overflowing_mul(rhs).1 + } + + #[inline(always)] + fn div_value(self, rhs: Self) -> Self { + self / rhs + } + + #[inline(always)] + fn div_error(self, rhs: Self) -> bool { + rhs == 0 || (self == <$ty>::MIN && rhs == -1) + } + + #[inline(always)] + fn div_checked(self, rhs: Self) -> Option { + self.checked_div(rhs) + } + } }; +} + +macro_rules! impl_checked_float { + ($($ty:ty),+ $(,)?) => { + $( + impl CheckedArithmetic for $ty { + const DIV_CHECKS_IN_VALUE_LOOP: bool = false; + + #[inline(always)] + fn add_value(self, rhs: Self) -> Self { + self + rhs + } + + #[inline(always)] + fn add_error(self, _rhs: Self) -> bool { + false + } - Ok(Some(ConstantArray::new(result, lhs.len()).into_array())) + #[inline(always)] + fn sub_value(self, rhs: Self) -> Self { + self - rhs + } + + #[inline(always)] + fn sub_error(self, _rhs: Self) -> bool { + false + } + + #[inline(always)] + fn mul_value(self, rhs: Self) -> Self { + self * rhs + } + + #[inline(always)] + fn mul_error(self, _rhs: Self) -> bool { + false + } + + #[inline(always)] + fn div_value(self, rhs: Self) -> Self { + self / rhs + } + + #[inline(always)] + fn div_error(self, _rhs: Self) -> bool { + false + } + + #[inline(always)] + fn div_checked(self, rhs: Self) -> Option { + Some(self / rhs) + } + } + )+ + }; } +impl_checked_unsigned!(u8, widening_mul: u16); +impl_checked_unsigned!(u16, widening_mul: u32); +impl_checked_unsigned!(u32, widening_mul: u64); +impl_checked_unsigned!(u64, overflowing_mul); +impl_checked_signed!(i8, widening_mul: i16); +impl_checked_signed!(i16, widening_mul: i32); +impl_checked_signed!(i32, widening_mul: i64); +impl_checked_signed!(i64, overflowing_mul); +impl_checked_float!(f16, f32, f64); + +fn check_numeric_errors(failed: bool, error: &'static str) -> VortexResult<()> { + if failed { + return Err(vortex_err!(InvalidArgument: "{}", error)); + } + + Ok(()) +} #[cfg(test)] mod test { use vortex_buffer::buffer; @@ -79,15 +926,16 @@ mod test { use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::RecursiveCanonical; use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::scalar::Scalar; - use crate::scalar_fn::fns::binary::numeric::ConstantArray; use crate::scalar_fn::fns::operators::Operator; + use crate::validity::Validity; fn sub_scalar(array: &ArrayRef, scalar: impl Into) -> VortexResult { array @@ -96,40 +944,49 @@ mod test { Operator::Sub, ) .and_then(|a| { - a.execute::(&mut LEGACY_SESSION.create_execution_ctx()) + a.execute::(&mut array_session().create_execution_ctx()) }) .map(|a| a.0.into_array()) } #[test] fn test_scalar_subtract_unsigned() { + let mut ctx = array_session().create_execution_ctx(); let values = buffer![1u16, 2, 3].into_array(); let result = sub_scalar(&values, 1u16).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2])); + assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2]), &mut ctx); } #[test] fn test_scalar_subtract_signed() { + let mut ctx = array_session().create_execution_ctx(); let values = buffer![1i64, 2, 3].into_array(); let result = sub_scalar(&values, -1i64).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4])); + assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4]), &mut ctx); } #[test] fn test_scalar_subtract_nullable() { + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]); let result = sub_scalar(&values.into_array(), Some(1u16)).unwrap(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]) + PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]), + &mut ctx ); } #[test] fn test_scalar_subtract_float() { + let mut ctx = array_session().create_execution_ctx(); let values = buffer![1.0f64, 2.0, 3.0].into_array(); let result = sub_scalar(&values, -1f64).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([2.0f64, 3.0, 4.0])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([2.0f64, 3.0, 4.0]), + &mut ctx + ); } #[test] @@ -138,4 +995,140 @@ mod test { let _results = sub_scalar(&values, 1.0f32).unwrap(); let _results = sub_scalar(&values, f32::MAX).unwrap(); } + + #[test] + fn test_float_divide_by_zero_is_ok() { + let mut ctx = array_session().create_execution_ctx(); + let values = buffer![1.0f64, -1.0].into_array(); + let result = values + .binary( + ConstantArray::new(0.0f64, values.len()).into_array(), + Operator::Div, + ) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())) + .unwrap(); + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([f64::INFINITY, f64::NEG_INFINITY]), + &mut ctx + ); + } + + #[test] + fn test_integer_overflow_errors() { + let values = buffer![u8::MAX].into_array(); + let result = values + .binary( + ConstantArray::new(1u8, values.len()).into_array(), + Operator::Add, + ) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); + + assert!(result.is_err()); + } + + #[test] + fn test_integer_divide_by_zero_errors() { + let values = buffer![1i32].into_array(); + let result = values + .binary( + ConstantArray::new(0i32, values.len()).into_array(), + Operator::Div, + ) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); + + assert!(result.is_err()); + } + + #[test] + fn test_integer_divide_overflow_errors() { + let values = buffer![i64::MIN].into_array(); + let result = values + .binary( + ConstantArray::new(-1i64, values.len()).into_array(), + Operator::Div, + ) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); + + assert!(result.is_err()); + } + + #[test] + fn test_integer_divide_errors_ignore_null_lanes() { + let mut ctx = array_session().create_execution_ctx(); + let lhs = PrimitiveArray::new(buffer![10i32, 10], Validity::from_iter([false, true])) + .into_array(); + let rhs = buffer![0i32, 2].into_array(); + let result = lhs + .binary(rhs, Operator::Div) + .and_then(|a| { + a.execute::(&mut array_session().create_execution_ctx()) + }) + .map(|a| a.0.into_array()) + .unwrap(); + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([None, Some(5i32)]), + &mut ctx + ); + } + + #[test] + fn test_integer_errors_ignore_null_lanes() { + let mut ctx = array_session().create_execution_ctx(); + let values = PrimitiveArray::new(buffer![u8::MAX, 1], Validity::from_iter([false, true])) + .into_array(); + let result = values + .binary( + ConstantArray::new(1u8, values.len()).into_array(), + Operator::Add, + ) + .and_then(|a| { + a.execute::(&mut array_session().create_execution_ctx()) + }) + .map(|a| a.0.into_array()) + .unwrap(); + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([None, Some(2u8)]), + &mut ctx + ); + } + + #[test] + fn test_integer_array_array_errors_on_valid_lanes() { + let lhs = PrimitiveArray::new( + buffer![u8::MAX, 1, u8::MAX], + Validity::from_iter([false, true, true]), + ) + .into_array(); + let rhs = buffer![1u8, 1, 1].into_array(); + let result = lhs + .binary(rhs, Operator::Add) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); + + assert!(result.is_err()); + } + + #[test] + fn test_present_nullable_constant_preserves_nullable_output() { + let mut ctx = array_session().create_execution_ctx(); + let values = buffer![1u8, 2].into_array(); + let result = values + .binary( + ConstantArray::new(Some(1u8), values.len()).into_array(), + Operator::Add, + ) + .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())) + .unwrap(); + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(2u8), Some(3)]), + &mut ctx + ); + } } diff --git a/vortex-array/src/scalar_fn/fns/byte_length.rs b/vortex-array/src/scalar_fn/fns/byte_length.rs index 13a4f3158b5..cfab5290a8f 100644 --- a/vortex-array/src/scalar_fn/fns/byte_length.rs +++ b/vortex-array/src/scalar_fn/fns/byte_length.rs @@ -24,6 +24,7 @@ use crate::arrays::varbinview::VarBinViewArrayExt; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::expr::Expression; use crate::kernel::ExecuteParentKernel; use crate::scalar::Scalar; use crate::scalar_fn::Arity; @@ -122,6 +123,14 @@ impl ScalarFnVTable for ByteLength { } } + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + Ok(Some(expression.child(0).validity()?)) + } + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { false } @@ -171,8 +180,8 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; @@ -193,18 +202,20 @@ mod tests { #[case] array: ArrayRef, #[case] expected_lens: Vec, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let result = array.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(expected_lens); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } #[test] fn test_varbinview_byte_length() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let array = VarBinViewArray::from_iter_str(["short", "a longer string here"]).into_array(); let result = array.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(vec![5u64, 20]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -214,16 +225,16 @@ mod tests { .into_array(); let result = array.apply(&byte_length(root()))?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert!(result.is_valid(0, &mut ctx)?); assert!(!result.is_valid(1, &mut ctx)?); assert!(result.is_valid(2, &mut ctx)?); assert_eq!( - result.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())?, + result.execute_scalar(0, &mut array_session().create_execution_ctx())?, Scalar::primitive(5u64, Nullability::Nullable), ); assert_eq!( - result.execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())?, + result.execute_scalar(2, &mut array_session().create_execution_ctx())?, Scalar::primitive(18u64, Nullability::Nullable), ); Ok(()) @@ -234,7 +245,7 @@ mod tests { let null_scalar = Scalar::null(DType::Utf8(Nullability::Nullable)); let array = ConstantArray::new(null_scalar, 2).into_array(); let result = array.apply(&byte_length(root()))?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert!(!result.is_valid(0, &mut ctx)?); assert!(!result.is_valid(1, &mut ctx)?); Ok(()) diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index c178e536b38..d57026e00f5 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -450,7 +450,6 @@ mod tests { use super::*; use crate::Canonical; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; @@ -473,10 +472,8 @@ mod tests { use crate::expr::root; use crate::expr::test_harness; use crate::scalar::Scalar; - use crate::session::ArraySession; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); /// Helper to evaluate an expression using the apply+execute pattern fn evaluate_expr(expr: &Expression, array: &ArrayRef) -> ArrayRef { @@ -785,6 +782,7 @@ mod tests { #[test] fn test_evaluate_simple_condition() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -797,11 +795,16 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![0i32, 0, 100, 100, 100].into_array()); + assert_arrays_eq!( + result, + buffer![0i32, 0, 100, 100, 100].into_array(), + &mut ctx + ); } #[test] fn test_evaluate_nary_multiple_conditions() { + let mut ctx = SESSION.create_execution_ctx(); // Test n-ary via nested_case_when let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) @@ -817,11 +820,12 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![10i32, 0, 30, 0, 0].into_array()); + assert_arrays_eq!(result, buffer![10i32, 0, 30, 0, 0].into_array(), &mut ctx); } #[test] fn test_evaluate_nary_first_match_wins() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -837,11 +841,16 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![0i32, 0, 100, 100, 100].into_array()); + assert_arrays_eq!( + result, + buffer![0i32, 0, 100, 100, 100].into_array(), + &mut ctx + ); } #[test] fn test_evaluate_no_else_returns_null() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -854,12 +863,14 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([None::, None, None, Some(100), Some(100)]) - .into_array() + .into_array(), + &mut ctx ); } #[test] fn test_evaluate_all_conditions_false() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -872,11 +883,12 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![0i32, 0, 0, 0, 0].into_array()); + assert_arrays_eq!(result, buffer![0i32, 0, 0, 0, 0].into_array(), &mut ctx); } #[test] fn test_evaluate_all_conditions_true() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -889,11 +901,16 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![100i32, 100, 100, 100, 100].into_array()); + assert_arrays_eq!( + result, + buffer![100i32, 100, 100, 100, 100].into_array(), + &mut ctx + ); } #[test] fn test_evaluate_all_true_no_else_returns_correct_dtype() { + let mut ctx = SESSION.create_execution_ctx(); // CASE WHEN value > 0 THEN 100 END — condition is always true, no ELSE. // Result must be Nullable because the implicit ELSE is NULL. let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -910,12 +927,14 @@ mod tests { ); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(100i32), Some(100), Some(100)]).into_array() + PrimitiveArray::from_option_iter([Some(100i32), Some(100), Some(100)]).into_array(), + &mut ctx ); } #[test] fn test_merge_case_branches_widens_nullability_of_later_branch() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // When a later THEN branch is Nullable and branches[0] and ELSE are NonNullable, // the result dtype must still be Nullable. // @@ -945,22 +964,25 @@ mod tests { ); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10), Some(20), Some(0)]).into_array() + PrimitiveArray::from_option_iter([Some(10), Some(20), Some(0)]).into_array(), + &mut ctx ); Ok(()) } #[test] fn test_evaluate_with_literal_condition() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = buffer![1i32, 2, 3].into_array(); let expr = case_when(lit(true), lit(100i32), lit(0i32)); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![100i32, 100, 100].into_array()); + assert_arrays_eq!(result, buffer![100i32, 100, 100].into_array(), &mut ctx); } #[test] fn test_evaluate_with_bool_column_result() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -975,12 +997,14 @@ mod tests { let result = evaluate_expr(&expr, &test_array); assert_arrays_eq!( result, - BoolArray::from_iter([false, false, true, true, true]).into_array() + BoolArray::from_iter([false, false, true, true, true]).into_array(), + &mut ctx ); } #[test] fn test_evaluate_with_nullable_condition() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[( "cond", BoolArray::from_iter([Some(true), None, Some(false), None, Some(true)]).into_array(), @@ -991,11 +1015,12 @@ mod tests { let expr = case_when(get_item("cond", root()), lit(100i32), lit(0i32)); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![100i32, 0, 0, 0, 100].into_array()); + assert_arrays_eq!(result, buffer![100i32, 0, 0, 0, 100].into_array(), &mut ctx); } #[test] fn test_evaluate_with_nullable_result_values() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[ ("value", buffer![1i32, 2, 3, 4, 5].into_array()), ( @@ -1017,12 +1042,14 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(0i32), Some(0), Some(30), Some(40), Some(50)]) - .into_array() + .into_array(), + &mut ctx ); } #[test] fn test_evaluate_with_all_null_condition() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[( "cond", BoolArray::from_iter([None, None, None]).into_array(), @@ -1033,13 +1060,14 @@ mod tests { let expr = case_when(get_item("cond", root()), lit(100i32), lit(0i32)); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![0i32, 0, 0].into_array()); + assert_arrays_eq!(result, buffer![0i32, 0, 0].into_array(), &mut ctx); } // ==================== N-ary Evaluate Tests ==================== #[test] fn test_evaluate_nary_no_else_returns_null() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1059,12 +1087,14 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(10i32), None, Some(30), None, None]) - .into_array() + .into_array(), + &mut ctx ); } #[test] fn test_evaluate_nary_many_conditions() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1083,11 +1113,16 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![10i32, 20, 30, 40, 50].into_array()); + assert_arrays_eq!( + result, + buffer![10i32, 20, 30, 40, 50].into_array(), + &mut ctx + ); } #[test] fn test_evaluate_nary_all_false_no_else() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) .unwrap() .into_array(); @@ -1105,12 +1140,14 @@ mod tests { assert!(result.dtype().is_nullable()); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([None::, None, None]).into_array() + PrimitiveArray::from_option_iter([None::, None, None]).into_array(), + &mut ctx ); } #[test] fn test_evaluate_nary_overlapping_conditions_first_wins() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![10i32, 20, 30].into_array())]) .unwrap() @@ -1130,11 +1167,12 @@ mod tests { let result = evaluate_expr(&expr, &test_array); // First matching condition always wins - assert_arrays_eq!(result, buffer![1i32, 1, 1].into_array()); + assert_arrays_eq!(result, buffer![1i32, 1, 1].into_array(), &mut ctx); } #[test] fn test_evaluate_nary_early_exit_when_remaining_empty() { + let mut ctx = SESSION.create_execution_ctx(); // After branch 0 claims all rows, remaining becomes all_false. // The loop breaks before evaluating branch 1's condition. let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -1151,11 +1189,12 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![100i32, 100, 100].into_array()); + assert_arrays_eq!(result, buffer![100i32, 100, 100].into_array(), &mut ctx); } #[test] fn test_evaluate_nary_skips_branch_with_empty_effective_mask() { + let mut ctx = SESSION.create_execution_ctx(); // Branch 0 claims value=1. Branch 1 targets the same rows but they are already // matched → effective_mask is all_false → branch 1 is skipped (THEN not used). let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -1174,7 +1213,7 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![10i32, 20, 0].into_array()); + assert_arrays_eq!(result, buffer![10i32, 20, 0].into_array(), &mut ctx); } #[test] @@ -1197,19 +1236,19 @@ mod tests { let result = evaluate_expr(&expr, &test_array); assert_eq!( - result.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())?, + result.execute_scalar(0, &mut SESSION.create_execution_ctx())?, Scalar::utf8("low", Nullability::NonNullable) ); assert_eq!( - result.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())?, + result.execute_scalar(1, &mut SESSION.create_execution_ctx())?, Scalar::utf8("low", Nullability::NonNullable) ); assert_eq!( - result.execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())?, + result.execute_scalar(2, &mut SESSION.create_execution_ctx())?, Scalar::utf8("high", Nullability::NonNullable) ); assert_eq!( - result.execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx())?, + result.execute_scalar(3, &mut SESSION.create_execution_ctx())?, Scalar::utf8("high", Nullability::NonNullable) ); Ok(()) @@ -1217,6 +1256,7 @@ mod tests { #[test] fn test_evaluate_nary_with_nullable_conditions() { + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[ ( "cond1", @@ -1242,7 +1282,7 @@ mod tests { // row 0: cond1=true → 10 // row 1: cond1=NULL(→false), cond2=true → 20 // row 2: cond1=false, cond2=NULL(→false) → else=0 - assert_arrays_eq!(result, buffer![10i32, 20, 0].into_array()); + assert_arrays_eq!(result, buffer![10i32, 20, 0].into_array(), &mut ctx); } // ==================== Simplify: COALESCE -> fill_null ==================== @@ -1333,6 +1373,7 @@ mod tests { #[test] fn test_simplify_null_fill_semantic_equivalence() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // The collapse-to-input rewrite must preserve values (and `x`'s nullability). let array = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); let scope = DType::Primitive(PType::I64, Nullability::Nullable); @@ -1350,8 +1391,8 @@ mod tests { ); let expected = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); - assert_arrays_eq!(evaluate_expr(&original, &array), expected); - assert_arrays_eq!(evaluate_expr(&optimized, &array), expected); + assert_arrays_eq!(evaluate_expr(&original, &array), expected, &mut ctx); + assert_arrays_eq!(evaluate_expr(&optimized, &array), expected, &mut ctx); Ok(()) } @@ -1385,6 +1426,7 @@ mod tests { #[test] fn test_simplify_semantic_equivalence() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // The optimized expression must produce the same values as the original CASE WHEN. let array = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); let scope = DType::Primitive(PType::I64, Nullability::Nullable); @@ -1400,17 +1442,20 @@ mod tests { // NonNullable because a non-null fill cannot leave any nulls behind. Values match. assert_arrays_eq!( evaluate_expr(&original, &array), - PrimitiveArray::from_option_iter([Some(1i64), Some(0), Some(3)]).into_array() + PrimitiveArray::from_option_iter([Some(1i64), Some(0), Some(3)]).into_array(), + &mut ctx ); assert_arrays_eq!( evaluate_expr(&optimized, &array), - buffer![1i64, 0, 3].into_array() + buffer![1i64, 0, 3].into_array(), + &mut ctx ); Ok(()) } #[test] fn test_merge_case_branches_alternating_mask() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); // Exercises the scalar path: alternating rows produce one slice per row (no runs), // triggering the per-row cursor path in merge_case_branches. let n = 100usize; @@ -1440,7 +1485,8 @@ mod tests { .collect(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter(expected).into_array() + PrimitiveArray::from_option_iter(expected).into_array(), + &mut ctx ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index abc59af2c9a..20852779d42 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -32,11 +32,8 @@ use crate::arrays::VarBinView; use crate::arrays::struct_::compute::cast::struct_cast; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::expr::StatsCatalog; -use crate::expr::cast; use crate::expr::expression::Expression; use crate::expr::lit; -use crate::expr::stats::Stat; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; @@ -152,39 +149,6 @@ impl ScalarFnVTable for Cast { Ok(None) } - fn stat_expression( - &self, - dtype: &DType, - expr: &Expression, - stat: Stat, - catalog: &dyn StatsCatalog, - ) -> Option { - match stat { - Stat::IsConstant - | Stat::IsSorted - | Stat::IsStrictSorted - | Stat::NaNCount - | Stat::Sum - | Stat::UncompressedSizeInBytes => expr.child(0).stat_expression(stat, catalog), - Stat::Max | Stat::Min => { - // We cast min/max to the new type - expr.child(0) - .stat_expression(stat, catalog) - .map(|x| cast(x, dtype.clone())) - } - Stat::NullCount => { - // if !expr.data().is_nullable() { - // NOTE(ngates): we should decide on the semantics here. In theory, the null - // count of something cast to non-nullable will be zero. But if we return - // that we know this to be zero, then a pruning predicate may eliminate data - // that would otherwise have caused the cast to error. - // return Some(lit(0u64)); - // } - None - } - } - } - fn validity(&self, dtype: &DType, expression: &Expression) -> VortexResult> { Ok(Some(if dtype.is_nullable() { expression.child(0).validity()? diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index 7efebf79220..74058aa19f2 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -20,7 +20,6 @@ use crate::IntoArray; use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::expr::Expression; -use crate::expr::StatsCatalog; use crate::expr::traversal::NodeExt; use crate::expr::traversal::NodeVisitor; use crate::expr::traversal::TraversalOrder; @@ -120,50 +119,6 @@ impl ScalarFnVTable for DynamicComparison { .into_array()) } - fn stat_falsification( - &self, - dynamic: &DynamicComparisonExpr, - expr: &Expression, - catalog: &dyn StatsCatalog, - ) -> Option { - let lhs = expr.child(0); - match dynamic.operator { - CompareOperator::Eq | CompareOperator::NotEq => None, - CompareOperator::Gt => Some(DynamicComparison.new_expr( - DynamicComparisonExpr { - operator: CompareOperator::Lte, - rhs: Arc::clone(&dynamic.rhs), - default: !dynamic.default, - }, - vec![lhs.stat_max(catalog)?], - )), - CompareOperator::Gte => Some(DynamicComparison.new_expr( - DynamicComparisonExpr { - operator: CompareOperator::Lt, - rhs: Arc::clone(&dynamic.rhs), - default: !dynamic.default, - }, - vec![lhs.stat_max(catalog)?], - )), - CompareOperator::Lt => Some(DynamicComparison.new_expr( - DynamicComparisonExpr { - operator: CompareOperator::Gte, - rhs: Arc::clone(&dynamic.rhs), - default: !dynamic.default, - }, - vec![lhs.stat_min(catalog)?], - )), - CompareOperator::Lte => Some(DynamicComparison.new_expr( - DynamicComparisonExpr { - operator: CompareOperator::Gt, - rhs: Arc::clone(&dynamic.rhs), - default: !dynamic.default, - }, - vec![lhs.stat_min(catalog)?], - )), - } - } - // Defer to the child fn is_null_sensitive(&self, _instance: &Self::Options) -> bool { false @@ -322,6 +277,8 @@ mod tests { use super::*; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -348,6 +305,7 @@ mod tests { #[test] fn execute_with_value() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -357,12 +315,13 @@ mod tests { root(), ); let result = input.apply(&expr)?; - assert_arrays_eq!(result, BoolArray::from_iter([true, false, false])); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); Ok(()) } #[test] fn execute_without_value_default_true() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -372,12 +331,13 @@ mod tests { root(), ); let result = input.apply(&expr)?; - assert_arrays_eq!(result, BoolArray::from_iter([true, true, true])); + assert_arrays_eq!(result, BoolArray::from_iter([true, true, true]), &mut ctx); Ok(()) } #[test] fn execute_without_value_default_false() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -387,12 +347,17 @@ mod tests { root(), ); let result = input.apply(&expr)?; - assert_arrays_eq!(result, BoolArray::from_iter([false, false, false])); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, false, false]), + &mut ctx + ); Ok(()) } #[test] fn execute_value_flips() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let threshold = Arc::new(AtomicI32::new(5)); let threshold_clone = Arc::clone(&threshold); let expr = dynamic( @@ -405,11 +370,11 @@ mod tests { let input = buffer![1i32, 5, 10].into_array(); let result = input.clone().apply(&expr)?; - assert_arrays_eq!(result, BoolArray::from_iter([true, false, false])); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); threshold.store(10, Ordering::SeqCst); let result = input.apply(&expr)?; - assert_arrays_eq!(result, BoolArray::from_iter([true, true, false])); + assert_arrays_eq!(result, BoolArray::from_iter([true, true, false]), &mut ctx); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/ext_storage.rs b/vortex-array/src/scalar_fn/fns/ext_storage.rs new file mode 100644 index 00000000000..4a014a3325b --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/ext_storage.rs @@ -0,0 +1,202 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::arrays::ExtensionArray; +use crate::arrays::extension::ExtensionArrayExt; +use crate::dtype::DType; +use crate::expr::Expression; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; + +/// Extract the storage values from an extension array. +#[derive(Clone)] +pub struct ExtStorage; + +impl ScalarFnVTable for ExtStorage { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.ext.storage"); + *ID + } + + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("input"), + _ => unreachable!("Invalid child index {child_idx} for ext_storage()"), + } + } + + fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + let DType::Extension(ext_dtype) = &arg_dtypes[0] else { + vortex_bail!("ext_storage() requires Extension, got {}", arg_dtypes[0]); + }; + + Ok(ext_dtype.storage_dtype().clone()) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + + if !matches!(input.dtype(), DType::Extension(_)) { + vortex_bail!("ext_storage() requires Extension, got {}", input.dtype()); + } + + if let Some(scalar) = input.as_constant() { + let storage_scalar = scalar.as_extension().to_storage_scalar(); + return Ok(ConstantArray::new(storage_scalar, args.row_count()).into_array()); + } + + let input = input.execute::(ctx)?; + Ok(input.storage_array().clone()) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + Ok(Some(expression.child(0).validity()?)) + } + + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { + false + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::ConstantArray; + use crate::arrays::ExtensionArray; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::dtype::extension::ExtDTypeRef; + use crate::expr::ext_storage; + use crate::expr::root; + use crate::extension::datetime::TimeUnit; + use crate::extension::datetime::Timestamp; + use crate::scalar::Scalar; + + fn ext_dtype(nullability: Nullability) -> ExtDTypeRef { + Timestamp::new(TimeUnit::Nanoseconds, nullability).erased() + } + + #[test] + fn extracts_extension_storage_array() -> VortexResult<()> { + let storage = buffer![2i64, 4, 6].into_array(); + let array = + ExtensionArray::new(ext_dtype(Nullability::NonNullable), storage.clone()).into_array(); + + let result = array.apply(&ext_storage(root()))?; + + assert_eq!( + result.dtype(), + &DType::Primitive(PType::I64, Nullability::NonNullable) + ); + assert_arrays_eq!( + result, + storage, + &mut crate::array_session().create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn extracts_nullable_extension_storage_array() -> VortexResult<()> { + let storage = PrimitiveArray::from_option_iter([Some(2i64), None, Some(6)]).into_array(); + let array = + ExtensionArray::new(ext_dtype(Nullability::Nullable), storage.clone()).into_array(); + + let result = array.apply(&ext_storage(root()))?; + + assert_eq!( + result.dtype(), + &DType::Primitive(PType::I64, Nullability::Nullable) + ); + assert_arrays_eq!( + result, + storage, + &mut crate::array_session().create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn extracts_constant_extension_storage_scalar() -> VortexResult<()> { + let storage_scalar = Scalar::primitive(4i64, Nullability::NonNullable); + let scalar = + Scalar::extension_ref(ext_dtype(Nullability::NonNullable), storage_scalar.clone()); + let array = ConstantArray::new(scalar, 3).into_array(); + + let result = array.apply(&ext_storage(root()))?; + + assert_eq!( + result.dtype(), + &DType::Primitive(PType::I64, Nullability::NonNullable) + ); + assert_arrays_eq!( + result, + ConstantArray::new(storage_scalar, 3), + &mut crate::array_session().create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn rejects_non_extension_input() { + let dtype = DType::Primitive(PType::U64, Nullability::NonNullable); + let err = ext_storage(root()).return_dtype(&dtype).unwrap_err(); + assert!(err.to_string().contains("requires Extension")); + } + + #[test] + fn test_display() { + assert_eq!(ext_storage(root()).to_string(), "vortex.ext.storage($)"); + } +} diff --git a/vortex-array/src/scalar_fn/fns/fill_null/kernel.rs b/vortex-array/src/scalar_fn/fns/fill_null/kernel.rs index eea3dd6ef7b..147658ddfb2 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/kernel.rs +++ b/vortex-array/src/scalar_fn/fns/fill_null/kernel.rs @@ -79,7 +79,7 @@ pub(super) fn precondition( } // If all values are null, replace the entire array with the fill value. - if matches!(array.validity()?, Validity::AllInvalid) { + if array.validity()?.definitely_all_null() { return Ok(Some( ConstantArray::new(fill_value.clone(), array.len()).into_array(), )); diff --git a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs index b2d07b41255..c9eae72ad08 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs @@ -179,6 +179,8 @@ mod tests { use vortex_error::VortexExpect; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::assert_arrays_eq; @@ -208,6 +210,7 @@ mod tests { #[test] fn evaluate() { + let mut ctx = array_session().create_execution_ctx(); let test_array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]) .into_array(); @@ -219,11 +222,16 @@ mod tests { result.dtype(), &DType::Primitive(PType::I32, Nullability::NonNullable) ); - assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 42, 3, 42, 5])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([1i32, 42, 3, 42, 5]), + &mut ctx + ); } #[test] fn evaluate_struct_field() { + let mut ctx = array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[( "a", PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), @@ -238,15 +246,16 @@ mod tests { result.dtype(), &DType::Primitive(PType::I32, Nullability::NonNullable) ); - assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 0, 3])); + assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 0, 3]), &mut ctx); } #[test] fn evaluate_non_nullable_input() { + let mut ctx = array_session().create_execution_ctx(); let test_array = buffer![1i32, 2, 3].into_array(); let expr = fill_null(root(), lit(0i32)); let result = test_array.apply(&expr).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 2, 3])); + assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 2, 3]), &mut ctx); } #[test] diff --git a/vortex-array/src/scalar_fn/fns/get_item.rs b/vortex-array/src/scalar_fn/fns/get_item.rs index de7e45ca9b0..b9cb9202f38 100644 --- a/vortex-array/src/scalar_fn/fns/get_item.rs +++ b/vortex-array/src/scalar_fn/fns/get_item.rs @@ -18,12 +18,9 @@ use crate::builtins::ArrayBuiltins; use crate::builtins::ExprBuiltins; use crate::dtype::DType; use crate::dtype::FieldName; -use crate::dtype::FieldPath; use crate::dtype::Nullability; use crate::expr::Expression; -use crate::expr::StatsCatalog; use crate::expr::lit; -use crate::expr::stats::Stat; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; @@ -188,24 +185,6 @@ impl ScalarFnVTable for GetItem { Ok(None) } - fn stat_expression( - &self, - field_name: &FieldName, - _expr: &Expression, - stat: Stat, - catalog: &dyn StatsCatalog, - ) -> Option { - // TODO(ngates): I think we can do better here and support stats over nested fields. - // It would be nice if delegating to our child would return a struct of statistics - // matching the nested DType such that we can write: - // `get_item(expr.child(0).stat_expression(...), expr.data().field_name())` - - // TODO(ngates): this is a bug whereby we may return stats for a nested field of the same - // name as a field in the root struct. This should be resolved with upcoming change to - // falsify expressions, but for now I'm preserving the existing buggy behavior. - catalog.stats_ref(&FieldPath::from_name(field_name.clone()), stat) - } - // This will apply struct nullability field. We could add a dtype?? fn is_null_sensitive(&self, _field_name: &FieldName) -> bool { true diff --git a/vortex-array/src/scalar_fn/fns/is_not_null.rs b/vortex-array/src/scalar_fn/fns/is_not_null.rs index 589333304e2..75d86db5ee8 100644 --- a/vortex-array/src/scalar_fn/fns/is_not_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_not_null.rs @@ -3,7 +3,6 @@ use std::fmt::Formatter; -use vortex_array::scalar_fn::internal::row_count::RowCount; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -15,16 +14,12 @@ use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::Expression; -use crate::expr::StatsCatalog; -use crate::expr::eq; -use crate::expr::stats::Stat; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; use crate::validity::Validity; /// Expression that checks for non-null values. @@ -100,49 +95,41 @@ impl ScalarFnVTable for IsNotNull { fn is_fallible(&self, _instance: &Self::Options) -> bool { false } - - fn stat_falsification( - &self, - _options: &Self::Options, - expr: &Expression, - catalog: &dyn StatsCatalog, - ) -> Option { - // is_not_null is falsified when ALL values are null, i.e. null_count == row_count. - let child = expr.child(0); - let null_count_expr = child.stat_expression(Stat::NullCount, catalog)?; - Some(eq(null_count_expr, RowCount.new_expr(EmptyOptions, []))) - } } #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_buffer::buffer; use vortex_error::VortexExpect as _; - use vortex_utils::aliases::hash_map::HashMap; - use vortex_utils::aliases::hash_set::HashSet; + use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::dtype::DType; - use crate::dtype::Field; - use crate::dtype::FieldPath; - use crate::dtype::FieldPathSet; use crate::dtype::Nullability; use crate::expr::col; use crate::expr::eq; use crate::expr::get_item; use crate::expr::is_not_null; - use crate::expr::pruning::checked_pruning_expr; + use crate::expr::or; use crate::expr::root; - use crate::expr::stats::Stat; use crate::expr::test_harness; use crate::scalar::Scalar; use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::internal::row_count::RowCount; - use crate::scalar_fn::vtable::ScalarFnVTableExt; + use crate::stats::StatsSession; + use crate::stats::all_null; + use crate::stats::null_count; + + static STATS_SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); #[test] fn dtype() { @@ -175,7 +162,7 @@ mod tests { for (i, expected_value) in expected.iter().enumerate() { assert_eq!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(*expected_value, Nullability::NonNullable) ); @@ -192,7 +179,7 @@ mod tests { for i in 0..result.len() { assert_eq!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(true, Nullability::NonNullable) ); @@ -211,7 +198,7 @@ mod tests { for i in 0..result.len() { assert_eq!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(false, Nullability::NonNullable) ); @@ -240,7 +227,7 @@ mod tests { for (i, expected_value) in expected.iter().enumerate() { assert_eq!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(*expected_value, Nullability::NonNullable) ); @@ -262,25 +249,16 @@ mod tests { } #[test] - fn test_is_not_null_falsification() { + fn test_is_not_null_falsification() -> VortexResult<()> { let expr = is_not_null(col("a")); - let (pruning_expr, st) = checked_pruning_expr( - &expr, - &FieldPathSet::from_iter([FieldPath::from_iter([ - Field::Name("a".into()), - Field::Name("null_count".into()), - ])]), - ) - .unwrap(); - - assert_eq!( - &pruning_expr, - &eq(col("a_null_count"), RowCount.new_expr(EmptyOptions, [])) - ); assert_eq!( - st.map(), - &HashMap::from_iter([(FieldPath::from_name("a"), HashSet::from([Stat::NullCount]))]) + expr.falsify(&test_harness::struct_dtype(), &STATS_SESSION)?, + Some(or( + eq(null_count(col("a")), RowCount.new_expr(EmptyOptions, []),), + all_null(col("a")), + )) ); + Ok(()) } } diff --git a/vortex-array/src/scalar_fn/fns/is_null.rs b/vortex-array/src/scalar_fn/fns/is_null.rs index 7315fbe8c07..5e605ca6ee0 100644 --- a/vortex-array/src/scalar_fn/fns/is_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_null.rs @@ -12,11 +12,6 @@ use crate::arrays::ConstantArray; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; -use crate::expr::Expression; -use crate::expr::StatsCatalog; -use crate::expr::eq; -use crate::expr::lit; -use crate::expr::stats::Stat; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; @@ -84,16 +79,6 @@ impl ScalarFnVTable for IsNull { } } - fn stat_falsification( - &self, - _options: &Self::Options, - expr: &Expression, - catalog: &dyn StatsCatalog, - ) -> Option { - let null_count_expr = expr.child(0).stat_expression(Stat::NullCount, catalog)?; - Some(eq(null_count_expr, lit(0u64))) - } - fn is_null_sensitive(&self, _instance: &Self::Options) -> bool { true } @@ -105,31 +90,35 @@ impl ScalarFnVTable for IsNull { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_buffer::buffer; use vortex_error::VortexExpect as _; - use vortex_utils::aliases::hash_map::HashMap; - use vortex_utils::aliases::hash_set::HashSet; + use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::dtype::DType; - use crate::dtype::Field; - use crate::dtype::FieldPath; - use crate::dtype::FieldPathSet; use crate::dtype::Nullability; use crate::expr::col; use crate::expr::eq; use crate::expr::get_item; use crate::expr::is_null; use crate::expr::lit; - use crate::expr::pruning::checked_pruning_expr; + use crate::expr::or; use crate::expr::root; - use crate::expr::stats::Stat; use crate::expr::test_harness; use crate::scalar::Scalar; + use crate::stats::StatsSession; + use crate::stats::all_non_null; + use crate::stats::null_count; + + static STATS_SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); #[test] fn dtype() { @@ -162,7 +151,7 @@ mod tests { for (i, expected_value) in expected.iter().enumerate() { assert_eq!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(*expected_value, Nullability::NonNullable) ); @@ -180,7 +169,7 @@ mod tests { for i in 0..result.len() { assert_eq!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(false, Nullability::NonNullable) ); @@ -200,7 +189,7 @@ mod tests { for i in 0..result.len() { assert_eq!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(true, Nullability::NonNullable) ); @@ -229,7 +218,7 @@ mod tests { for (i, expected_value) in expected.iter().enumerate() { assert_eq!( result - .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(i, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(*expected_value, Nullability::NonNullable) ); @@ -246,23 +235,17 @@ mod tests { } #[test] - fn test_is_null_falsification() { + fn test_is_null_falsification() -> VortexResult<()> { let expr = is_null(col("a")); - let (pruning_expr, st) = checked_pruning_expr( - &expr, - &FieldPathSet::from_iter([FieldPath::from_iter([ - Field::Name("a".into()), - Field::Name("null_count".into()), - ])]), - ) - .unwrap(); - - assert_eq!(&pruning_expr, &eq(col("a_null_count"), lit(0u64))); assert_eq!( - st.map(), - &HashMap::from_iter([(FieldPath::from_name("a"), HashSet::from([Stat::NullCount]))]) + expr.falsify(&test_harness::struct_dtype(), &STATS_SESSION)?, + Some(or( + eq(null_count(col("a")), lit(0u64)), + all_non_null(col("a")), + )) ); + Ok(()) } #[test] diff --git a/vortex-array/src/scalar_fn/fns/like/mod.rs b/vortex-array/src/scalar_fn/fns/like/mod.rs index b7f357020f1..63ced07c8e7 100644 --- a/vortex-array/src/scalar_fn/fns/like/mod.rs +++ b/vortex-array/src/scalar_fn/fns/like/mod.rs @@ -21,20 +21,12 @@ use crate::arrow::Datum; use crate::arrow::from_arrow_columnar; use crate::dtype::DType; use crate::expr::Expression; -use crate::expr::StatsCatalog; use crate::expr::and; -use crate::expr::gt; -use crate::expr::gt_eq; -use crate::expr::lit; -use crate::expr::lt; -use crate::expr::or; -use crate::scalar::StringLike; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::literal::Literal; /// Options for SQL LIKE function #[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -166,47 +158,8 @@ impl ScalarFnVTable for Like { false } - fn stat_falsification( - &self, - like_opts: &LikeOptions, - expr: &Expression, - catalog: &dyn StatsCatalog, - ) -> Option { - // Attempt to do min/max pruning for LIKE 'exact' or LIKE 'prefix%' - - // Don't attempt to handle ilike or negated like - if like_opts.negated || like_opts.case_insensitive { - return None; - } - - // Extract the pattern out - let pat = expr.child(1).as_::(); - - // LIKE NULL is nonsensical, don't try to handle it - let pat_str = pat.as_utf8().value()?; - - let src = expr.child(0).clone(); - let src_min = src.stat_min(catalog)?; - let src_max = src.stat_max(catalog)?; - - match LikeVariant::from_str(pat_str)? { - LikeVariant::Exact(text) => { - // col LIKE 'exact' ==> col.min > 'exact' || col.max < 'exact' - Some(or( - gt(src_min, lit(text.as_ref())), - lt(src_max, lit(text.as_ref())), - )) - } - LikeVariant::Prefix(prefix) => { - // col LIKE 'prefix%' ==> col.max < 'prefix' || col.min >= 'prefiy' - let succ = prefix.to_string().increment().ok()?; - - Some(or( - gt_eq(src_min, lit(succ)), - lt(src_max, lit(prefix.as_ref())), - )) - } - } + fn is_fallible(&self, _options: &Self::Options) -> bool { + false } } @@ -291,19 +244,17 @@ mod tests { use std::borrow::Cow; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::Nullability; - use crate::expr::col; use crate::expr::get_item; - use crate::expr::ilike; use crate::expr::like; use crate::expr::lit; use crate::expr::not; use crate::expr::not_ilike; - use crate::expr::not_like; - use crate::expr::pruning::pruning_expr::TrackingStatsCatalog; use crate::expr::root; use crate::scalar_fn::fns::like::LikeVariant; @@ -311,9 +262,11 @@ mod tests { fn invert_booleans() { let not_expr = not(root()); let bools = BoolArray::from_iter([false, true, false, false, true, true]); + let mut ctx = array_session().create_execution_ctx(); assert_arrays_eq!( bools.into_array().apply(¬_expr).unwrap(), - BoolArray::from_iter([true, false, true, true, false, false]) + BoolArray::from_iter([true, false, true, true, false, false]), + &mut ctx ); } @@ -327,6 +280,13 @@ mod tests { ); } + #[test] + fn signature() { + let like_expr = like(root(), lit("%test%")); + assert!(!like_expr.signature().is_null_sensitive()); + assert!(!like_expr.signature().is_fallible()); + } + #[test] fn test_display() { let expr = like(get_item("name", root()), lit("%john%")); @@ -390,50 +350,4 @@ mod tests { assert_eq!(LikeVariant::from_str(r"%\%%"), None); assert_eq!(LikeVariant::from_str("_pattern"), None); } - - #[test] - fn test_like_pushdown() { - // Test that LIKE prefix and exactness filters can be pushed down into stats filtering - // at scan time. - let catalog = TrackingStatsCatalog::default(); - - let pruning_expr = like(col("a"), lit("prefix%")) - .stat_falsification(&catalog) - .expect("LIKE stat falsification"); - - insta::assert_snapshot!(pruning_expr, @r#"(($.a_min >= "prefiy") or ($.a_max < "prefix"))"#); - - let pruning_expr = like(col("a"), lit(r"\%%")) - .stat_falsification(&catalog) - .expect("LIKE stat falsification"); - insta::assert_snapshot!(pruning_expr, @r#"(($.a_min >= "&") or ($.a_max < "%"))"#); - - // Multiple wildcards - let pruning_expr = like(col("a"), lit("pref%ix%")) - .stat_falsification(&catalog) - .expect("LIKE stat falsification"); - insta::assert_snapshot!(pruning_expr, @r#"(($.a_min >= "preg") or ($.a_max < "pref"))"#); - - let pruning_expr = like(col("a"), lit("pref_ix_")) - .stat_falsification(&catalog) - .expect("LIKE stat falsification"); - insta::assert_snapshot!(pruning_expr, @r#"(($.a_min >= "preg") or ($.a_max < "pref"))"#); - - // Exact match - let pruning_expr = like(col("a"), lit("exactly")) - .stat_falsification(&catalog) - .expect("LIKE stat falsification"); - insta::assert_snapshot!(pruning_expr, @r#"(($.a_min > "exactly") or ($.a_max < "exactly"))"#); - - // Suffix search skips pushdown - let pruning_expr = like(col("a"), lit("%suffix")).stat_falsification(&catalog); - assert_eq!(pruning_expr, None); - - // NOT LIKE, ILIKE not supported currently - assert_eq!( - None, - not_like(col("a"), lit("a")).stat_falsification(&catalog) - ); - assert_eq!(None, ilike(col("a"), lit("a")).stat_falsification(&catalog)); - } } diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index 978a1da1caf..5185e513e70 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -33,13 +33,6 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::Nullability; -use crate::expr::Expression; -use crate::expr::StatsCatalog; -use crate::expr::and_collect; -use crate::expr::gt; -use crate::expr::lit; -use crate::expr::lt; -use crate::expr::or; use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; use crate::scalar::ListScalar; @@ -51,7 +44,6 @@ use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; @@ -129,43 +121,6 @@ impl ScalarFnVTable for ListContains { compute_list_contains(&list_array, &value_array, ctx) } - fn stat_falsification( - &self, - _options: &Self::Options, - expr: &Expression, - catalog: &dyn StatsCatalog, - ) -> Option { - let list = expr.child(0); - let needle = expr.child(1); - - // falsification(contains([1,2,5], x)) => - // falsification(x != 1) and falsification(x != 2) and falsification(x != 5) - let min = list.stat_min(catalog)?; - let max = list.stat_max(catalog)?; - // If the list is constant when we can compare each element to the value - if min == max { - let list_ = min - .as_opt::() - .and_then(|l| l.as_list_opt()) - .and_then(|l| l.elements())?; - if list_.is_empty() { - // contains([], x) is always false. - return Some(lit(true)); - } - let value_max = needle.stat_max(catalog)?; - let value_min = needle.stat_min(catalog)?; - - return and_collect(list_.iter().map(move |v| { - or( - lt(value_max.clone(), lit(v.clone())), - gt(value_min.clone(), lit(v.clone())), - ) - })); - } - - None - } - // Nullability matters for contains([], x) where x is false. fn is_null_sensitive(&self, _instance: &Self::Options) -> bool { true @@ -415,7 +370,7 @@ fn list_is_not_empty( ctx: &mut ExecutionCtx, ) -> VortexResult { // Short-circuit for all invalid. - if matches!(list_array.validity()?, Validity::AllInvalid) { + if list_array.validity()?.definitely_all_null() { return Ok(ConstantArray::new( Scalar::null(DType::Bool(Nullability::Nullable)), list_array.len(), @@ -439,30 +394,29 @@ fn list_is_not_empty( #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use itertools::Itertools; use rstest::rstest; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; - use vortex_utils::aliases::hash_map::HashMap; - use vortex_utils::aliases::hash_set::HashSet; + use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ListArray; use crate::arrays::VarBinArray; use crate::assert_arrays_eq; #[expect(deprecated)] use crate::canonical::ToCanonical as _; use crate::dtype::DType; - use crate::dtype::Field; - use crate::dtype::FieldPath; - use crate::dtype::FieldPathSet; use crate::dtype::Nullability; use crate::dtype::PType::I32; use crate::dtype::StructFields; + use crate::expr::Expression; use crate::expr::and; use crate::expr::col; use crate::expr::get_item; @@ -471,7 +425,6 @@ mod tests { use crate::expr::lit; use crate::expr::lt; use crate::expr::or; - use crate::expr::pruning::checked_pruning_expr; use crate::expr::root; use crate::expr::stats::Stat; use crate::scalar::Scalar; @@ -479,8 +432,17 @@ mod tests { use crate::scalar_fn::fns::list_contains::ConstantArray; use crate::scalar_fn::fns::list_contains::ListViewArray; use crate::scalar_fn::fns::list_contains::PrimitiveArray; + use crate::stats::StatsSession; + use crate::stats::stat as stat_expr; use crate::validity::Validity; + static STATS_SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + + fn stat(expr: Expression, stat: Stat) -> Expression { + stat_expr(expr, stat.aggregate_fn().unwrap()) + } + fn test_array() -> ArrayRef { ListArray::try_new( PrimitiveArray::from_iter(vec![1, 1, 2, 2, 2, 2, 2, 3, 3, 3]).into_array(), @@ -499,12 +461,12 @@ mod tests { let item = arr.apply(&expr).unwrap(); assert_eq!( - item.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + item.execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(true, Nullability::Nullable) ); assert_eq!( - item.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + item.execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(false, Nullability::Nullable) ); @@ -518,12 +480,12 @@ mod tests { let item = arr.apply(&expr).unwrap(); assert_eq!( - item.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + item.execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(true, Nullability::Nullable) ); assert_eq!( - item.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + item.execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(true, Nullability::Nullable) ); @@ -537,12 +499,12 @@ mod tests { let item = arr.apply(&expr).unwrap(); assert_eq!( - item.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + item.execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(false, Nullability::Nullable) ); assert_eq!( - item.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + item.execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(false, Nullability::Nullable) ); @@ -562,12 +524,12 @@ mod tests { let item = arr.apply(&expr).unwrap(); assert_eq!( - item.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + item.execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(true, Nullability::Nullable) ); assert_eq!( - item.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + item.execute_scalar(1, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(false, Nullability::Nullable) ); @@ -587,13 +549,13 @@ mod tests { let item = arr.apply(&expr).unwrap(); assert_eq!( - item.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + item.execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(true, Nullability::Nullable) ); assert!( !item - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) + .is_valid(1, &mut array_session().create_execution_ctx()) .unwrap() ); } @@ -621,7 +583,7 @@ mod tests { } #[test] - pub fn list_falsification() { + pub fn list_falsification() -> VortexResult<()> { let expr = list_contains( lit(Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -630,34 +592,34 @@ mod tests { )), col("a"), ); - - let (expr, st) = checked_pruning_expr( - &expr, - &FieldPathSet::from_iter([ - FieldPath::from_iter([Field::Name("a".into()), Field::Name("max".into())]), - FieldPath::from_iter([Field::Name("a".into()), Field::Name("min".into())]), - ]), - ) - .unwrap(); + let scope = DType::Struct( + StructFields::new( + ["a"].into(), + vec![DType::Primitive(I32, Nullability::NonNullable)], + ), + Nullability::NonNullable, + ); assert_eq!( - &expr, - &and( + expr.falsify(&scope, &STATS_SESSION)?, + Some(and( and( - or(lt(col("a_max"), lit(1i32)), gt(col("a_min"), lit(1i32)),), - or(lt(col("a_max"), lit(2i32)), gt(col("a_min"), lit(2i32)),) + or( + lt(stat(col("a"), Stat::Max), lit(1i32)), + gt(stat(col("a"), Stat::Min), lit(1i32)), + ), + or( + lt(stat(col("a"), Stat::Max), lit(2i32)), + gt(stat(col("a"), Stat::Min), lit(2i32)), + ) ), - or(lt(col("a_max"), lit(3i32)), gt(col("a_min"), lit(3i32)),) - ) - ); - - assert_eq!( - st.map(), - &HashMap::from_iter([( - FieldPath::from_name("a"), - HashSet::from([Stat::Min, Stat::Max]) - )]) + or( + lt(stat(col("a"), Stat::Max), lit(3i32)), + gt(stat(col("a"), Stat::Min), lit(3i32)), + ) + )) ); + Ok(()) } #[test] @@ -685,7 +647,7 @@ mod tests { let result = arr.clone().apply(&expr).unwrap(); assert_eq!( result - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(true, Nullability::NonNullable) ); @@ -695,7 +657,7 @@ mod tests { let result = arr.apply(&expr).unwrap(); assert_eq!( result - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(), Scalar::bool(false, Nullability::NonNullable) ); @@ -790,6 +752,7 @@ mod tests { #[case] value: Option<&str>, #[case] expected: BoolArray, ) { + let mut ctx = array_session().create_execution_ctx(); let element_nullability = list_array .dtype() .as_list_element_opt() @@ -802,11 +765,12 @@ mod tests { let elem = ConstantArray::new(scalar, list_array.len()); let expr = list_contains(root(), lit(elem.scalar().clone())); let result = list_array.apply(&expr).unwrap(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_constant_list() { + let mut ctx = array_session().create_execution_ctx(); let list_array = ConstantArray::new( Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -820,11 +784,12 @@ mod tests { let expr = list_contains(root(), lit(2i32)); let contains = list_array.apply(&expr).unwrap(); let expected = BoolArray::from_iter([true, true]); - assert_arrays_eq!(contains, expected); + assert_arrays_eq!(contains, expected, &mut ctx); } #[test] fn test_all_nulls() { + let mut ctx = array_session().create_execution_ctx(); let list_array = ConstantArray::new( Scalar::null(DType::List( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -841,11 +806,12 @@ mod tests { [false, false, false, false, false].into_iter().collect(), Validity::AllInvalid, ); - assert_arrays_eq!(contains, expected); + assert_arrays_eq!(contains, expected, &mut ctx); } #[test] fn test_list_array_element() { + let mut ctx = array_session().create_execution_ctx(); let list_scalar = Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), vec![1.into(), 3.into(), 6.into()], @@ -857,11 +823,12 @@ mod tests { let contains = arr.apply(&expr).unwrap(); let expected = BoolArray::from_iter([false, true, false, true, false, false, true]); - assert_arrays_eq!(contains, expected); + assert_arrays_eq!(contains, expected, &mut ctx); } #[test] fn test_list_contains_empty_listview() { + let mut ctx = array_session().create_execution_ctx(); let empty_elements = PrimitiveArray::empty::(Nullability::NonNullable); let offsets = Buffer::from_iter([0u32, 0, 0, 0]).into_array(); let sizes = Buffer::from_iter([0u32, 0, 0, 0]).into_array(); @@ -880,11 +847,12 @@ mod tests { let result = list_array.into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([false, false, false, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_list_contains_all_null_elements() { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter::([None, None, None, None, None]); let offsets = Buffer::from_iter([0u32, 2, 4]).into_array(); let sizes = Buffer::from_iter([2u32, 2, 1]).into_array(); @@ -908,18 +876,19 @@ mod tests { [false, false, false].into_iter().collect(), Validity::AllInvalid, ); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); // Searching for non-null let expr2 = list_contains(root(), lit(42i32)); let result2 = list_array.into_array().apply(&expr2).unwrap(); let expected2 = BoolArray::from_iter([false, false, false]); - assert_arrays_eq!(result2, expected2); + assert_arrays_eq!(result2, expected2, &mut ctx); } #[test] fn test_list_contains_large_offsets() { + let mut ctx = array_session().create_execution_ctx(); let elements = Buffer::from_iter([1i32, 2, 3, 4, 5]).into_array(); let offsets = Buffer::from_iter([0u32, 1, 4, 0]).into_array(); @@ -932,17 +901,18 @@ mod tests { let result = list_array.clone().into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([false, true, false, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); let expr5 = list_contains(root(), lit(5i32)); let result5 = list_array.into_array().apply(&expr5).unwrap(); let expected5 = BoolArray::from_iter([false, false, true, false]); - assert_arrays_eq!(result5, expected5); + assert_arrays_eq!(result5, expected5, &mut ctx); } #[test] fn test_list_contains_offset_size_boundary() { + let mut ctx = array_session().create_execution_ctx(); let elements = Buffer::from_iter(0..256).into_array(); let offsets = Buffer::from_iter([0u8, 100, 200, 254]).into_array(); let sizes = Buffer::from_iter([50u8, 50, 54, 2]).into_array(); @@ -954,12 +924,12 @@ mod tests { let result = list_array.clone().into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([false, false, false, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); let expr_zero = list_contains(root(), lit(0i32)); let result_zero = list_array.into_array().apply(&expr_zero).unwrap(); let expected_zero = BoolArray::from_iter([true, false, false, false]); - assert_arrays_eq!(result_zero, expected_zero); + assert_arrays_eq!(result_zero, expected_zero, &mut ctx); } } diff --git a/vortex-array/src/scalar_fn/fns/list_length.rs b/vortex-array/src/scalar_fn/fns/list_length.rs new file mode 100644 index 00000000000..ea88a3b27d6 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_length.rs @@ -0,0 +1,396 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::AsPrimitive; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::ConstantArray; +use crate::arrays::FixedSizeList; +use crate::arrays::List; +use crate::arrays::ListView; +use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::list::ListArrayExt; +use crate::arrays::listview::ListViewArrayExt; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::expr::Expression; +use crate::matcher::Matcher; +use crate::scalar::Scalar; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::operators::Operator; + +/// Number of elements in each list of a `List` or `FixedSizeList` typed array. +/// +/// This is computed purely from the list's offsets (`ListArray`), sizes (`ListViewArray`), or +/// dtype (`FixedSizeListArray`) without reading the element *values*. Validity is carried over +/// from the original array. +#[derive(Clone)] +pub struct ListLength; + +impl ScalarFnVTable for ListLength { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.list.length"); + *ID + } + + fn serialize(&self, _instance: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _instance: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("input"), + _ => unreachable!("Invalid child index {child_idx} for list_length()"), + } + } + + fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + match &arg_dtypes[0] { + DType::List(_, nullable) | DType::FixedSizeList(_, _, nullable) => { + Ok(DType::Primitive(PType::U64, *nullable)) + } + other => vortex_bail!("list_length() requires List or FixedSizeList, got {other}"), + } + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let nullability = input.dtype().nullability(); + + if let Some(scalar) = input.as_constant() { + let len_scalar = scalar_list_length(&scalar, nullability)?; + return Ok(ConstantArray::new(len_scalar, args.row_count()).into_array()); + } + + list_length(&input, nullability, ctx) + } + + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + Ok(Some(expression.child(0).validity()?)) + } + + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { + false + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +fn scalar_list_length(scalar: &Scalar, nullability: Nullability) -> VortexResult { + if scalar.is_null() { + let dtype = DType::Primitive(PType::U64, Nullability::Nullable); + return Ok(Scalar::null(dtype)); + } + let len: u64 = scalar.as_list().len().as_(); + Ok(Scalar::primitive(len, nullability)) +} + +pub(crate) fn list_length( + array: &ArrayRef, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let any_list = array.clone().execute_until::(ctx)?; + + let (lengths, validity) = if let Some(fsl) = any_list.as_opt::() { + // The length of fixed-size list is constant, so just need to carry over validity + let size = fsl.list_size() as u64; + let lengths = + ConstantArray::new(Scalar::primitive(size, Nullability::NonNullable), fsl.len()) + .into_array(); + (lengths, fsl.validity()?) + } else if let Some(lv) = any_list.as_opt::() { + // Length array is exactly the sizes child + (lv.sizes().clone(), lv.listview_validity()) + } else if let Some(l) = any_list.as_opt::() { + let lengths = list_length_from_offsets(l)?; + (lengths, l.list_validity()) + } else { + let dtype = any_list.dtype(); + vortex_bail!("list_length() requires List, ListView, or FixedSizeList but got {dtype}") + }; + + // Cast to `U64` + let len = lengths.len(); + let lengths = lengths.cast(DType::Primitive(PType::U64, nullability))?; + + // Carry over validity mask for nullable arrays + if matches!(nullability, Nullability::Nullable) { + lengths.mask(validity.to_array(len)) + } else { + Ok(lengths) + } +} + +/// Calculate the lengths of `ListArray` elements via the `offsets` child: +/// `length[i] = offsets[i + 1] - offsets[i]`. +fn list_length_from_offsets(list: ArrayView<'_, List>) -> VortexResult { + let offsets = list.offsets(); + let n = offsets.len().saturating_sub(1); + + offsets + .slice(1..offsets.len())? + .binary(offsets.slice(0..n)?, Operator::Sub) +} + +/// Matches an `Array`, `Array`, or `Array` +struct AnyList; + +impl Matcher for AnyList { + type Match<'a> = (); + + fn try_match(array: &ArrayRef) -> Option> { + (array.as_opt::().is_some() + || array.as_opt::().is_some() + || array.as_opt::().is_some()) + .then_some(()) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::BoolArray; + use crate::arrays::ConstantArray; + use crate::arrays::FixedSizeListArray; + use crate::arrays::ListArray; + use crate::arrays::ListViewArray; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::expr::cast; + use crate::expr::list_length; + use crate::expr::root; + use crate::scalar::Scalar; + use crate::validity::Validity; + + fn create_list_elements() -> ArrayRef { + PrimitiveArray::from_option_iter::([ + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + None, + ]) + .into_array() + } + + #[rstest] + #[case(buffer![0u32, 2, 5, 5, 7].into_array())] + #[case(buffer![0u64, 2, 5, 5, 7].into_array())] + fn test_list_length(#[case] offsets: ArrayRef) -> VortexResult<()> { + let elements = create_list_elements(); + let list = ListArray::try_new(elements, offsets, Validity::NonNullable)?.into_array(); + let result = list.apply(&list_length(root()))?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 3, 0, 2]), &mut ctx); + Ok(()) + } + + #[rstest] + #[case(buffer![0u32, 2, 5, 5, 7].into_array())] + #[case(buffer![0u64, 2, 5, 5, 7].into_array())] + fn test_nullable_list_length(#[case] offsets: ArrayRef) -> VortexResult<()> { + let elements = create_list_elements(); + let list = ListArray::try_new( + elements, + offsets, + Validity::Array(BoolArray::from_iter([true, false, true, false]).into_array()), + )? + .into_array(); + let result = list.apply(&list_length(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let result = result.execute::(&mut ctx)?; + + let expected = PrimitiveArray::from_option_iter::([Some(2), None, Some(0), None]); + + assert_arrays_eq!(result, expected, &mut ctx); + + Ok(()) + } + + #[test] + fn test_null_scalar_list_length() -> VortexResult<()> { + let null_scalar = Scalar::null(DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::Nullable, + )); + let array = ConstantArray::new(null_scalar, 2).into_array(); + let result = array.apply(&list_length(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + assert!(!result.is_valid(0, &mut ctx)?); + assert!(!result.is_valid(1, &mut ctx)?); + Ok(()) + } + + #[test] + fn test_listview_length() -> VortexResult<()> { + let elements = create_list_elements(); + let lv = ListViewArray::new( + elements, + buffer![5u32, 0, 4, 1].into_array(), + buffer![2u32, 3, 0, 2].into_array(), + Validity::NonNullable, + ) + .into_array(); + let result = lv.apply(&list_length(root()))?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 3, 0, 2]), &mut ctx); + Ok(()) + } + + #[test] + fn test_listview_length_nullable() -> VortexResult<()> { + let elements = create_list_elements(); + let lv = ListViewArray::new( + elements, + buffer![5u32, 0, 4, 1].into_array(), + buffer![2u32, 3, 0, 2].into_array(), + Validity::Array(BoolArray::from_iter([true, false, true, false]).into_array()), + ) + .into_array(); + let result = lv.apply(&list_length(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let result = result.execute::(&mut ctx)?; + + let expected = PrimitiveArray::from_option_iter::([Some(2), None, Some(0), None]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_list_length_take() -> VortexResult<()> { + let elements = create_list_elements(); + let list = ListArray::try_new( + elements, + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::NonNullable, + )? + .into_array(); + let taken = list.take(buffer![3u64, 0, 2].into_array())?; + + let result = taken.apply(&list_length(root()))?; + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 2, 0]), &mut ctx); + Ok(()) + } + + fn create_fixed_size_list(validity: Validity) -> ArrayRef { + // 4 lists of size 2 over 8 primitive elements. + let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6, 7, 8]).into_array(); + FixedSizeListArray::new(elements, 2, validity, 4).into_array() + } + + #[test] + fn test_fixed_size_list_length() -> VortexResult<()> { + let fsl = create_fixed_size_list(Validity::NonNullable); + let result = fsl.apply(&list_length(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 2, 2, 2]), &mut ctx); + Ok(()) + } + + #[test] + fn test_fixed_size_list_length_nullable() -> VortexResult<()> { + let fsl = create_fixed_size_list(Validity::Array( + BoolArray::from_iter([true, false, true, false]).into_array(), + )); + let result = fsl.apply(&list_length(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let result = result.execute::(&mut ctx)?; + + let expected = PrimitiveArray::from_option_iter::([Some(2), None, Some(2), None]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_fallible_child_expression_fails() -> VortexResult<()> { + let fsl = create_fixed_size_list(Validity::Array( + BoolArray::from_iter([true, false, true, false]).into_array(), + )); + let failing_cast_dtype = DType::FixedSizeList( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + 2, + Nullability::NonNullable, + ); + + let lengths = fsl.apply(&list_length(cast(root(), failing_cast_dtype)))?; + + let mut ctx = array_session().create_execution_ctx(); + let result = lengths.execute::(&mut ctx); + + assert!(result.is_err()); + + let err_message = result.unwrap_err().to_string(); + + assert!( + err_message.contains("Cannot cast array with invalid values to non-nullable type.") + ); + + Ok(()) + } + + #[test] + fn test_display() { + let expr = list_length(root()); + assert_eq!(expr.to_string(), "vortex.list.length($)"); + } +} diff --git a/vortex-array/src/scalar_fn/fns/literal.rs b/vortex-array/src/scalar_fn/fns/literal.rs index 16b112e5a78..5181a5250dd 100644 --- a/vortex-array/src/scalar_fn/fns/literal.rs +++ b/vortex-array/src/scalar_fn/fns/literal.rs @@ -16,9 +16,6 @@ use crate::IntoArray; use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::expr::Expression; -use crate::expr::StatsCatalog; -use crate::expr::stats::Stat; -use crate::match_each_float_ptype; use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -96,50 +93,6 @@ impl ScalarFnVTable for Literal { Ok(ConstantArray::new(scalar.clone(), args.row_count()).into_array()) } - fn stat_expression( - &self, - scalar: &Scalar, - _expr: &Expression, - stat: Stat, - _catalog: &dyn StatsCatalog, - ) -> Option { - // NOTE(ngates): we return incorrect `1` values for counts here since we don't have - // row-count information. We could resolve this in the future by introducing a `count()` - // expression that evaluates to the row count of the provided scope. But since this is - // only currently used for pruning, it doesn't change the outcome. - - match stat { - Stat::Min | Stat::Max => Some(lit(scalar.clone())), - Stat::IsConstant => Some(lit(true)), - Stat::NaNCount => { - // The NaNCount for a non-float literal is not defined. - // For floating point types, the NaNCount is 1 for lit(NaN), and 0 otherwise. - let value = scalar.as_primitive_opt()?; - if !value.ptype().is_float() { - return None; - } - - match_each_float_ptype!(value.ptype(), |T| { - if value.typed_value::().is_some_and(|v| v.is_nan()) { - Some(lit(1u64)) - } else { - Some(lit(0u64)) - } - }) - } - Stat::NullCount => { - if scalar.is_null() { - Some(lit(1u64)) - } else { - Some(lit(0u64)) - } - } - Stat::IsSorted | Stat::IsStrictSorted | Stat::Sum | Stat::UncompressedSizeInBytes => { - None - } - } - } - fn validity( &self, scalar: &Scalar, diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index 390d3a2213e..c9f5783de27 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -279,6 +279,8 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; @@ -318,6 +320,7 @@ mod tests { #[test] pub fn test_merge_right_most() { + let mut ctx = array_session().create_execution_ctx(); let expr = merge_opts( vec![ get_item("0", root()), @@ -367,23 +370,28 @@ mod tests { assert_arrays_eq!( primitive_field(&actual_array, &["a"]).unwrap(), - PrimitiveArray::from_iter([0i32, 0, 0]) + PrimitiveArray::from_iter([0i32, 0, 0]), + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array, &["b"]).unwrap(), - PrimitiveArray::from_iter([2i32, 2, 2]) + PrimitiveArray::from_iter([2i32, 2, 2]), + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array, &["c"]).unwrap(), - PrimitiveArray::from_iter([3i32, 3, 3]) + PrimitiveArray::from_iter([3i32, 3, 3]), + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array, &["d"]).unwrap(), - PrimitiveArray::from_iter([4i32, 4, 4]) + PrimitiveArray::from_iter([4i32, 4, 4]), + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array, &["e"]).unwrap(), - PrimitiveArray::from_iter([5i32, 5, 5]) + PrimitiveArray::from_iter([5i32, 5, 5]), + &mut ctx ); } diff --git a/vortex-array/src/scalar_fn/fns/mod.rs b/vortex-array/src/scalar_fn/fns/mod.rs index b7cfda9224c..f7e98e676b0 100644 --- a/vortex-array/src/scalar_fn/fns/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mod.rs @@ -7,12 +7,14 @@ pub mod byte_length; pub mod case_when; pub mod cast; pub mod dynamic; +pub mod ext_storage; pub mod fill_null; pub mod get_item; pub mod is_not_null; pub mod is_null; pub mod like; pub mod list_contains; +pub mod list_length; pub mod literal; pub mod mask; pub mod merge; diff --git a/vortex-array/src/scalar_fn/fns/pack.rs b/vortex-array/src/scalar_fn/fns/pack.rs index 2a87bf069bc..81019e22d6f 100644 --- a/vortex-array/src/scalar_fn/fns/pack.rs +++ b/vortex-array/src/scalar_fn/fns/pack.rs @@ -172,6 +172,8 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; @@ -230,6 +232,7 @@ mod tests { #[test] pub fn test_simple_pack() { + let mut ctx = array_session().create_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), @@ -246,20 +249,24 @@ mod tests { assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["one"]).unwrap(), - PrimitiveArray::from_iter([0i32, 1, 2]) + PrimitiveArray::from_iter([0i32, 1, 2]), + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["two"]).unwrap(), - PrimitiveArray::from_iter([4i32, 5, 6]) + PrimitiveArray::from_iter([4i32, 5, 6]), + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array.into_array(), &["three"]).unwrap(), - PrimitiveArray::from_iter([0i32, 1, 2]) + PrimitiveArray::from_iter([0i32, 1, 2]), + &mut ctx ); } #[test] pub fn test_nested_pack() { + let mut ctx = array_session().create_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), @@ -285,19 +292,23 @@ mod tests { assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["one"]).unwrap(), - PrimitiveArray::from_iter([0i32, 1, 2]) + PrimitiveArray::from_iter([0i32, 1, 2]), + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["two", "two_one"]).unwrap(), - PrimitiveArray::from_iter([4i32, 5, 6]) + PrimitiveArray::from_iter([4i32, 5, 6]), + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["two", "two_two"]).unwrap(), - PrimitiveArray::from_iter([4i32, 5, 6]) + PrimitiveArray::from_iter([4i32, 5, 6]), + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array.into_array(), &["three"]).unwrap(), - PrimitiveArray::from_iter([0i32, 1, 2]) + PrimitiveArray::from_iter([0i32, 1, 2]), + &mut ctx ); } diff --git a/vortex-array/src/scalar_fn/fns/root.rs b/vortex-array/src/scalar_fn/fns/root.rs index 87b8b62ccf4..7bd5b758796 100644 --- a/vortex-array/src/scalar_fn/fns/root.rs +++ b/vortex-array/src/scalar_fn/fns/root.rs @@ -11,10 +11,7 @@ use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; -use crate::dtype::FieldPath; -use crate::expr::StatsCatalog; use crate::expr::expression::Expression; -use crate::expr::stats::Stat; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; @@ -80,16 +77,6 @@ impl ScalarFnVTable for Root { vortex_bail!("Root expression is not executable") } - fn stat_expression( - &self, - _options: &Self::Options, - _expr: &Expression, - stat: Stat, - catalog: &dyn StatsCatalog, - ) -> Option { - catalog.stats_ref(&FieldPath::root(), stat) - } - fn is_null_sensitive(&self, _options: &Self::Options) -> bool { false } diff --git a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs index ac6aa562aa6..1a4854a6a4a 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs @@ -55,16 +55,7 @@ impl ScalarFnVTable for VariantGet { .path() .elements() .iter() - .map(|element| match element { - VariantPathElement::Field(name) => pb::VariantPathElement { - element: Some(variant_path_element::Element::Field( - name.as_ref().to_string(), - )), - }, - VariantPathElement::Index(index) => pb::VariantPathElement { - element: Some(variant_path_element::Element::Index(*index)), - }, - }) + .map(VariantPathElement::to_proto) .collect(); let dtype = options.dtype().map(TryInto::try_into).transpose()?; @@ -355,15 +346,30 @@ impl VariantPathElement { Self::Index(index) } - fn from_proto(value: pb::VariantPathElement) -> VortexResult { + /// Decodes a path element from its protobuf representation. + pub fn from_proto(value: pb::VariantPathElement) -> VortexResult { match value .element - .ok_or_else(|| vortex_err!("VariantGet path element missing value"))? + .ok_or_else(|| vortex_err!("Variant path element missing value"))? { variant_path_element::Element::Field(field) => Ok(Self::field(field)), variant_path_element::Element::Index(index) => Ok(Self::index(index)), } } + + /// Encodes this path element into its protobuf representation. + pub fn to_proto(&self) -> pb::VariantPathElement { + match self { + VariantPathElement::Field(name) => pb::VariantPathElement { + element: Some(variant_path_element::Element::Field( + name.as_ref().to_string(), + )), + }, + VariantPathElement::Index(index) => pb::VariantPathElement { + element: Some(variant_path_element::Element::Index(*index)), + }, + } + } } impl From for VariantPathElement { @@ -395,8 +401,8 @@ mod tests { use crate::ArrayRef; use crate::Canonical; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -539,7 +545,7 @@ mod tests { let expr = variant_get(root(), parse_path(path)?, dtype); array .apply(&expr)? - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut array_session().create_execution_ctx()) } #[test] @@ -630,7 +636,7 @@ mod tests { Some(DType::Primitive(PType::I32, Nullability::NonNullable)), ); let proto = expr.serialize_proto().unwrap(); - let actual = Expression::from_proto(&proto, &LEGACY_SESSION).unwrap(); + let actual = Expression::from_proto(&proto, &array_session()).unwrap(); assert_eq!(actual, expr); } @@ -672,10 +678,12 @@ mod tests { "$.items[1]", Some(DType::Primitive(PType::I32, Nullability::NonNullable)), )?; + let mut ctx = array_session().create_execution_ctx(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(20i32), None, None]) + PrimitiveArray::from_option_iter([Some(20i32), None, None]), + &mut ctx ); Ok(()) } @@ -704,10 +712,12 @@ mod tests { "$.a", Some(DType::Primitive(PType::I32, Nullability::NonNullable)), )?; + let mut ctx = array_session().create_execution_ctx(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10i32), None, Some(30), None]) + PrimitiveArray::from_option_iter([Some(10i32), None, Some(30), None]), + &mut ctx ); Ok(()) } @@ -736,9 +746,11 @@ mod tests { )?; assert!(!result.is::()); + let mut ctx = array_session().create_execution_ctx(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10i32), Some(20), None]) + PrimitiveArray::from_option_iter([Some(10i32), Some(20), None]), + &mut ctx ); Ok(()) } @@ -760,7 +772,7 @@ mod tests { let result = execute_variant_get(array, "$.a", None)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let row0 = result.execute_scalar(0, &mut ctx)?; assert_eq!( row0.as_variant() @@ -769,7 +781,7 @@ mod tests { .map(|value| value.as_str()), Some("ok") ); - assert_nth_scalar_is_null!(result, 1); + assert_nth_scalar_is_null!(result, 1, &mut ctx); assert_eq!( result .execute_scalar(2, &mut ctx)? @@ -777,7 +789,7 @@ mod tests { .is_variant_null(), Some(true) ); - assert_nth_scalar_is_null!(result, 3); + assert_nth_scalar_is_null!(result, 3, &mut ctx); Ok(()) } @@ -797,8 +809,8 @@ mod tests { let result = execute_variant_get(array, "$.a", None)?; let variant = result .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; - let canonical = result.execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut array_session().create_execution_ctx())?; + let canonical = result.execute::(&mut array_session().create_execution_ctx())?; let Canonical::Variant(canonical_variant) = canonical else { vortex_bail!("expected Variant canonical array"); }; @@ -808,7 +820,7 @@ mod tests { assert_eq!(variant.core_storage().dtype(), variant.dtype()); assert_eq!(variant.core_storage().len(), variant.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); for (idx, expected) in [10i32, 20].into_iter().enumerate() { let scalar = variant.execute_scalar(idx, &mut ctx)?; let actual = scalar diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 8a3e1b5bbc6..48449bc8428 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -276,8 +276,8 @@ mod tests { use super::zip_impl; use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::Struct; @@ -317,6 +317,7 @@ mod tests { #[test] fn test_zip_basic() { + let mut ctx = array_session().create_execution_ctx(); let mask = Mask::from_iter([true, false, false, true, false]); let if_true = buffer![10, 20, 30, 40, 50].into_array(); let if_false = buffer![1, 2, 3, 4, 5].into_array(); @@ -324,11 +325,12 @@ mod tests { let result = mask.into_array().zip(if_true, if_false).unwrap(); let expected = buffer![10, 2, 3, 40, 5].into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_zip_all_true() { + let mut ctx = array_session().create_execution_ctx(); let mask = Mask::new_true(4); let if_true = buffer![10, 20, 30, 40].into_array(); let if_false = @@ -338,12 +340,13 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), Some(40)]).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); assert_eq!(result.dtype(), if_false.dtype()) } #[test] fn test_zip_all_false_widens_nullability() { + let mut ctx = array_session().create_execution_ctx(); let mask = Mask::new_false(4); let if_true = PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array(); @@ -353,7 +356,7 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), Some(4)]).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); assert_eq!(result.dtype(), if_true.dtype()); } @@ -364,12 +367,13 @@ mod tests { let if_false = PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), None]).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = zip_impl(&if_true, &if_false, &mask, &mut ctx)?; assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30), Some(40)]) - .into_array() + .into_array(), + &mut ctx ); assert_eq!(result.dtype(), if_false.dtype()); Ok(()) @@ -382,11 +386,12 @@ mod tests { PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array(); let if_false = buffer![1i32, 2, 3, 4].into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = zip_impl(&if_true, &if_false, &mask, &mut ctx)?; assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]).into_array() + PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]).into_array(), + &mut ctx ); assert_eq!(result.dtype(), if_true.dtype()); Ok(()) @@ -422,7 +427,7 @@ mod tests { let mask = Mask::from_indices(len, indices); let mask_array = mask.into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = mask_array .zip(const1.clone(), const2.clone())? .execute::(&mut ctx)? @@ -482,7 +487,7 @@ mod tests { let mask = Mask::from_indices(200, (0..100).filter(|i| i % 3 != 0)); let mask_array = mask.clone().into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let zipped = mask_array .zip(if_true.clone(), if_false.clone()) .unwrap() @@ -491,25 +496,25 @@ mod tests { let zipped = zipped.as_opt::().unwrap(); assert_eq!(zipped.data_buffers().len(), 2); - let mut arrow_ctx = LEGACY_SESSION.create_execution_ctx(); + let mut arrow_ctx = array_session().create_execution_ctx(); let expected = arrow_zip( - LEGACY_SESSION + array_session() .arrow() .execute_arrow(mask.into_array(), None, &mut arrow_ctx) .unwrap() .as_boolean(), - &LEGACY_SESSION + &array_session() .arrow() .execute_arrow(if_true, None, &mut arrow_ctx) .unwrap(), - &LEGACY_SESSION + &array_session() .arrow() .execute_arrow(if_false, None, &mut arrow_ctx) .unwrap(), ) .unwrap(); - let actual = LEGACY_SESSION + let actual = array_session() .arrow() .execute_arrow(zipped.array().clone(), None, &mut arrow_ctx) .unwrap(); diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 590ccb44224..003dbc2ca48 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -9,6 +9,11 @@ use vortex_session::registry::Id; +use crate::scalar_fn::fns::byte_length::ByteLength; +use crate::scalar_fn::fns::ext_storage::ExtStorage; +use crate::scalar_fn::fns::get_item::GetItem; +use crate::scalar_fn::fns::literal::Literal; + mod vtable; pub use vtable::*; @@ -48,3 +53,18 @@ mod sealed { /// This can be the **only** implementor for [`super::typed::DynScalarFn`]. impl Sealed for TypedScalarFnInstance {} } + +/// A scalar function has a negative cost if applying it to an array and +/// canonicalizing is cheaper than canonicalizing an array and applying it. +/// +/// Example of negative cost expressions are byte_length(), ext_storage(), and get_item() since +/// they don't depend on input size. +/// +/// Example of non-negative cost expression is like() as it's linear over +/// individual input. +pub fn is_negative_cost(id: ScalarFnId) -> bool { + id == ScalarFnVTable::id(&ByteLength) + || id == ScalarFnVTable::id(&ExtStorage) + || id == ScalarFnVTable::id(&GetItem) + || id == ScalarFnVTable::id(&Literal) +} diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index 08a5951567b..477d5758990 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -4,8 +4,8 @@ use std::any::Any; use std::sync::Arc; -use vortex_session::Ref; use vortex_session::SessionExt; +use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_session::registry::Registry; @@ -14,12 +14,14 @@ use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::is_not_null::IsNotNull; use crate::scalar_fn::fns::is_null::IsNull; use crate::scalar_fn::fns::like::Like; use crate::scalar_fn::fns::list_contains::ListContains; +use crate::scalar_fn::fns::list_length::ListLength; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::merge::Merge; use crate::scalar_fn::fns::not::Not; @@ -33,7 +35,7 @@ use crate::scalar_fn::fns::variant_get::VariantGet; pub type ScalarFnRegistry = Registry; /// Session state for scalar function vtables and rewrite rules. -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct ScalarFnSession { registry: ScalarFnRegistry, } @@ -60,12 +62,14 @@ impl Default for ScalarFnSession { this.register(Between); this.register(Binary); this.register(Cast); + this.register(ExtStorage); this.register(FillNull); this.register(GetItem); this.register(IsNotNull); this.register(IsNull); this.register(Like); this.register(ListContains); + this.register(ListLength); this.register(Literal); this.register(Merge); this.register(Not); @@ -92,7 +96,7 @@ impl SessionVar for ScalarFnSession { /// Extension trait for accessing scalar function session data. pub trait ScalarFnSessionExt: SessionExt { /// Returns the scalar function vtable registry. - fn scalar_fns(&self) -> Ref<'_, ScalarFnSession> { + fn scalar_fns(&self) -> SessionGuard<'_, ScalarFnSession> { self.get::() } } diff --git a/vortex-array/src/scalar_fn/typed.rs b/vortex-array/src/scalar_fn/typed.rs index fbfc8caace4..1b14a9d613d 100644 --- a/vortex-array/src/scalar_fn/typed.rs +++ b/vortex-array/src/scalar_fn/typed.rs @@ -24,8 +24,6 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; use crate::expr::Expression; -use crate::expr::StatsCatalog; -use crate::expr::stats::Stat; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; @@ -103,17 +101,6 @@ pub(super) trait DynScalarFn: 'static + Send + Sync + super::sealed::Sealed { ) -> VortexResult>; fn simplify_untyped(&self, expression: &Expression) -> VortexResult>; fn validity(&self, expression: &Expression) -> VortexResult>; - fn stat_falsification( - &self, - expression: &Expression, - catalog: &dyn StatsCatalog, - ) -> Option; - fn stat_expression( - &self, - expression: &Expression, - stat: Stat, - catalog: &dyn StatsCatalog, - ) -> Option; // Options operations — self-contained fn options_serialize(&self) -> VortexResult>>; @@ -225,23 +212,6 @@ impl DynScalarFn for TypedScalarFnInstance { V::validity(&self.vtable, &self.options, expression) } - fn stat_falsification( - &self, - expression: &Expression, - catalog: &dyn StatsCatalog, - ) -> Option { - V::stat_falsification(&self.vtable, &self.options, expression, catalog) - } - - fn stat_expression( - &self, - expression: &Expression, - stat: Stat, - catalog: &dyn StatsCatalog, - ) -> Option { - V::stat_expression(&self.vtable, &self.options, expression, stat, catalog) - } - fn options_serialize(&self) -> VortexResult>> { V::serialize(&self.vtable, &self.options) } diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index f4862f6876a..c66afc34932 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -20,8 +20,6 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; use crate::expr::Expression; -use crate::expr::StatsCatalog; -use crate::expr::stats::Stat; use crate::expr::traversal::Node; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnRef; @@ -179,34 +177,6 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { Ok(None) } - /// See [`Expression::stat_falsification`]. - fn stat_falsification( - &self, - options: &Self::Options, - expr: &Expression, - catalog: &dyn StatsCatalog, - ) -> Option { - _ = options; - _ = expr; - _ = catalog; - None - } - - /// See [`Expression::stat_expression`]. - fn stat_expression( - &self, - options: &Self::Options, - expr: &Expression, - stat: Stat, - catalog: &dyn StatsCatalog, - ) -> Option { - _ = options; - _ = expr; - _ = stat; - _ = catalog; - None - } - /// Returns an expression that evaluates to the validity of the result of this expression. /// /// If a validity expression cannot be constructed, returns `None` and the expression will diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 98a7417722b..2cb8414cbfd 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -6,8 +6,8 @@ use std::sync::Arc; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_session::Ref; use vortex_session::SessionExt; +use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_session::registry::Registry; @@ -33,7 +33,7 @@ use crate::arrays::Variant; pub type ArrayRegistry = Registry; -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct ArraySession { /// The set of registered array encodings. registry: ArrayRegistry, @@ -100,7 +100,7 @@ impl SessionVar for ArraySession { /// Session data for Vortex arrays. pub trait ArraySessionExt: SessionExt { /// Returns the array encoding registry. - fn arrays(&self) -> Ref<'_, ArraySession> { + fn arrays(&self) -> SessionGuard<'_, ArraySession> { self.get::() } @@ -118,3 +118,27 @@ pub trait ArraySessionExt: SessionExt { } impl ArraySessionExt for S {} + +#[cfg(test)] +mod tests { + use vortex_session::VortexSession; + + use crate::ArrayVTable; + use crate::arrays::Bool; + use crate::session::ArraySession; + use crate::session::ArraySessionExt; + + #[test] + fn array_session_default_registers_encodings() { + let session = VortexSession::empty().with::(); + + assert!(session.arrays().registry().find(&Bool.id()).is_some()); + } + + #[test] + fn empty_array_session_registers_no_encodings() { + let session = VortexSession::empty().with_some(ArraySession::empty()); + + assert!(session.arrays().registry().find(&Bool.id()).is_none()); + } +} diff --git a/vortex-array/src/stats/array.rs b/vortex-array/src/stats/array.rs index 7dc7a42603a..d3fb7fd11e4 100644 --- a/vortex-array/src/stats/array.rs +++ b/vortex-array/src/stats/array.rs @@ -16,6 +16,7 @@ use super::StatsSet; use super::StatsSetIntoIter; use super::TypedStatsSetRef; use crate::ArrayRef; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::is_constant::is_constant; use crate::aggregate_fn::fns::is_sorted::is_sorted; use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; @@ -162,8 +163,10 @@ impl StatsSetRef<'_> { } Ok(match stat { - Stat::Min => min_max(self.dyn_array_ref, ctx)?.map(|MinMaxResult { min, max: _ }| min), - Stat::Max => min_max(self.dyn_array_ref, ctx)?.map(|MinMaxResult { min: _, max }| max), + Stat::Min => min_max(self.dyn_array_ref, ctx, NumericalAggregateOpts::default())? + .map(|MinMaxResult { min, max: _ }| min), + Stat::Max => min_max(self.dyn_array_ref, ctx, NumericalAggregateOpts::default())? + .map(|MinMaxResult { min: _, max }| max), Stat::Sum => { Stat::Sum .dtype(self.dyn_array_ref.dtype()) diff --git a/vortex-array/src/stats/bind.rs b/vortex-array/src/stats/bind.rs new file mode 100644 index 00000000000..6d921fd0ab1 --- /dev/null +++ b/vortex-array/src/stats/bind.rs @@ -0,0 +1,207 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Bind abstract `vortex.stat` expressions to a concrete stats representation. +//! +//! Stats rewrite rules describe pruning in terms of `vortex.stat(input, aggregate_fn)` placeholders +//! so the rewrite is independent of where statistics are stored. These stat placeholders are +//! abstract because they name the statistic needed for a proof, but not how that statistic is +//! represented by a specific layout or reader. +//! +//! Binding is the later pass that replaces each abstract placeholder with the representation used +//! by a caller: zone-map field references, file-level stat literals, or typed nulls for missing +//! stats. This lets all callers share the same falsification rules while keeping layout-specific +//! stat storage behind [`StatBinder`]. + +use vortex_error::VortexResult; + +use crate::aggregate_fn::AggregateFnRef; +use crate::dtype::DType; +use crate::expr::Expression; +use crate::expr::lit; +use crate::expr::traversal::NodeExt; +use crate::expr::traversal::Transformed; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::stat::StatFn; + +/// A target that can bind abstract statistics to concrete expressions. +/// +/// Implementations define how a pruning proof should read stats from a specific backing +/// representation. For example, a zone-map binder can translate a `max(col)` placeholder into a +/// field reference in the per-zone stats table, while a file-stats binder can translate the same +/// placeholder into a literal value from the file footer. +pub trait StatBinder { + /// The dtype scope used to type-check expressions before stats are bound. + fn scope(&self) -> &DType; + + /// Bind `aggregate_fn(input)` to a concrete expression. + /// + /// Implementations should return `Ok(None)` when the requested aggregate + /// statistic is unavailable in their backing representation. + fn bind_aggregate( + &self, + input: &Expression, + aggregate_fn: &AggregateFnRef, + stat_dtype: &DType, + ) -> VortexResult>; + + /// Expression to use when a stat is unavailable. + /// + /// The default is a nullable null literal, which preserves three-valued + /// pruning semantics for stats-table execution. + fn missing_stat(&self, dtype: DType) -> VortexResult { + Ok(null_expr(dtype)) + } +} + +/// Bind all `vortex.stat` expressions in `predicate`. +/// +/// The predicate is usually the output of a stats rewrite rule. Rewrite rules +/// are responsible for expressing stat semantics; binding maps aggregate-backed +/// stat requests to the concrete stats representation supported by the binder. +pub fn bind_stats( + predicate: Expression, + binder: &B, +) -> VortexResult { + let scope = binder.scope().clone(); + Ok(predicate + .transform_down(|expr| { + if !expr.is::() { + return Ok(Transformed::no(expr)); + } + + match bind_stat_fn(&expr, &scope, binder)? { + Some(bound) => Ok(Transformed::yes(bound)), + None => { + let dtype = expr.return_dtype(&scope)?; + Ok(Transformed::yes(binder.missing_stat(dtype)?)) + } + } + })? + .into_inner()) +} + +fn bind_stat_fn( + expr: &Expression, + scope: &DType, + binder: &(impl StatBinder + ?Sized), +) -> VortexResult> { + let options = expr.as_::(); + let aggregate_fn = options.aggregate_fn(); + // `StatFn` has exactly one child: the expression the aggregate statistic is computed over. + let input = expr.child(0); + + let stat_dtype = expr.return_dtype(scope)?; + binder.bind_aggregate(input, aggregate_fn, &stat_dtype) +} + +fn null_expr(dtype: DType) -> Expression { + lit(Scalar::null(dtype.as_nullable())) +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::*; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::dtype::StructFields; + use crate::expr::and; + use crate::expr::col; + use crate::expr::get_item; + use crate::expr::is_null; + use crate::expr::or; + use crate::expr::root; + use crate::expr::stats::Stat; + use crate::stats::all_non_nan; + use crate::stats::nan_count; + + struct TestBinder { + input_scope: DType, + bind_nan_count: bool, + } + + impl TestBinder { + fn new(bind_nan_count: bool) -> Self { + Self { + input_scope: DType::Struct( + StructFields::from_iter([( + "f", + DType::Primitive(PType::F32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ), + bind_nan_count, + } + } + } + + impl StatBinder for TestBinder { + fn scope(&self) -> &DType { + &self.input_scope + } + + fn bind_aggregate( + &self, + _input: &Expression, + aggregate_fn: &AggregateFnRef, + _stat_dtype: &DType, + ) -> VortexResult> { + let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else { + return Ok(None); + }; + + if stat == Stat::NaNCount && self.bind_nan_count { + Ok(Some(get_item("f_nan_count", root()))) + } else { + Ok(None) + } + } + } + + #[test] + fn nan_count_binds_to_direct_stat_slot() -> VortexResult<()> { + let binder = TestBinder::new(true); + + let bound = bind_stats(nan_count(col("f")), &binder)?; + + assert_eq!(bound, col("f_nan_count")); + Ok(()) + } + + #[test] + fn all_non_nan_does_not_derive_from_nan_count() -> VortexResult<()> { + let binder = TestBinder::new(true); + + let bound = bind_stats(all_non_nan(col("f")), &binder)?; + + assert_eq!(bound, lit(Scalar::null(DType::Bool(Nullability::Nullable)))); + Ok(()) + } + + #[test] + fn missing_stats_bind_to_null_without_reducing() -> VortexResult<()> { + let binder = TestBinder::new(false); + let null_bool = lit(Scalar::null(DType::Bool(Nullability::Nullable))); + + let bound = bind_stats(and(lit(false), all_non_nan(col("f"))), &binder)?; + + assert_eq!(bound, and(lit(false), null_bool.clone())); + + let bound = bind_stats(or(lit(true), all_non_nan(col("f"))), &binder)?; + + assert_eq!(bound, or(lit(true), null_bool)); + Ok(()) + } + + #[test] + fn unrelated_expressions_do_not_request_nan_count() -> VortexResult<()> { + let binder = TestBinder::new(false); + + let bound = bind_stats(is_null(col("f")), &binder)?; + + assert_eq!(bound, is_null(col("f"))); + Ok(()) + } +} diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index 32fd48804de..2038df0b613 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -6,6 +6,7 @@ use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::all_nan::AllNan; use crate::aggregate_fn::fns::all_non_nan::AllNonNan; use crate::aggregate_fn::fns::all_non_null::AllNonNull; @@ -29,12 +30,14 @@ pub fn stat(expr: Expression, aggregate_fn: AggregateFnRef) -> Expression { /// Creates `stat(expr, min_max)`, returning a nullable `{ min, max }` struct statistic. pub fn min_max(expr: Expression) -> Expression { - stat(expr, MinMax.bind(EmptyOptions)) + // Statistics follow NaN-skipping semantics; request it explicitly rather than via the default. + stat(expr, MinMax.bind(NumericalAggregateOpts::skip_nans())) } /// Creates `stat(expr, sum)`, returning a nullable sum statistic. pub fn sum(expr: Expression) -> Expression { - stat(expr, Sum.bind(EmptyOptions)) + // Statistics follow NaN-skipping semantics; request it explicitly rather than via the default. + stat(expr, Sum.bind(NumericalAggregateOpts::skip_nans())) } /// Creates `stat(expr, null_count)`, returning a nullable null-count statistic. @@ -86,6 +89,7 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -100,11 +104,9 @@ mod tests { use crate::expr::stats::Stat; use crate::scalar::Scalar; use crate::scalar::ScalarValue; - use crate::session::ArraySession; use crate::validity::Validity; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(array_session); #[test] fn stat_expr_reads_cached_sum() -> VortexResult<()> { @@ -122,7 +124,7 @@ mod tests { let expected = ConstantArray::new(Scalar::primitive(6i64, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -141,7 +143,7 @@ mod tests { 3, ) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -176,7 +178,7 @@ mod tests { Validity::from_iter([true, true, false, false, false]), ) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -202,7 +204,7 @@ mod tests { let expected = ConstantArray::new(Scalar::primitive(2u64, Nullability::Nullable), 4).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -221,7 +223,7 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -240,7 +242,7 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(false, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -259,7 +261,7 @@ mod tests { let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -278,7 +280,7 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -297,7 +299,7 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -317,7 +319,7 @@ mod tests { let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 4).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -351,7 +353,7 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -372,7 +374,7 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(false, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -393,7 +395,7 @@ mod tests { let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -412,7 +414,7 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -432,7 +434,7 @@ mod tests { let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -459,7 +461,11 @@ mod tests { .into_array(); let expected_min = ConstantArray::new(Scalar::primitive(1i32, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(min_result, expected_min); + assert_arrays_eq!( + min_result, + expected_min, + &mut SESSION.create_execution_ctx() + ); let max_result = array .apply(&stat( @@ -472,7 +478,11 @@ mod tests { .into_array(); let expected_max = ConstantArray::new(Scalar::primitive(3i32, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(max_result, expected_max); + assert_arrays_eq!( + max_result, + expected_max, + &mut SESSION.create_execution_ctx() + ); Ok(()) } diff --git a/vortex-array/src/stats/mod.rs b/vortex-array/src/stats/mod.rs index 8c001ddd2cb..322e8631138 100644 --- a/vortex-array/src/stats/mod.rs +++ b/vortex-array/src/stats/mod.rs @@ -19,9 +19,10 @@ pub use expr::sum; pub use stats_set::*; mod array; +pub mod bind; pub mod expr; pub mod flatbuffers; -pub(crate) mod rewrite; +pub mod rewrite; pub mod session; mod stats_set; diff --git a/vortex-array/src/stats/rewrite.rs b/vortex-array/src/stats/rewrite.rs index 52d354df1a0..2bbeeb00022 100644 --- a/vortex-array/src/stats/rewrite.rs +++ b/vortex-array/src/stats/rewrite.rs @@ -21,7 +21,7 @@ mod builtins; pub(crate) use builtins::register_builtins; /// Shared reference to a stats rewrite rule. -pub(crate) type StatsRewriteRuleRef = Arc; +pub type StatsRewriteRuleRef = Arc; /// A plugin-provided rule for predicates whose root scalar function matches this rule. /// @@ -40,7 +40,7 @@ pub(crate) type StatsRewriteRuleRef = Arc; /// `expr` is the full predicate expression whose root scalar function id is /// [`Self::scalar_fn_id`]. Use [`StatsRewriteCtx`] to resolve dtypes and recursively rewrite child /// predicates. -pub(crate) trait StatsRewriteRule: Debug + Send + Sync + 'static { +pub trait StatsRewriteRule: Debug + Send + Sync + 'static { /// Returns the scalar function id handled by this rule. fn scalar_fn_id(&self) -> ScalarFnId; @@ -83,35 +83,35 @@ pub(crate) trait StatsRewriteRule: Debug + Send + Sync + 'static { } /// Context passed to stats rewrite rules. -pub(crate) struct StatsRewriteCtx<'a> { +pub struct StatsRewriteCtx<'a> { session: &'a VortexSession, scope: &'a DType, } impl<'a> StatsRewriteCtx<'a> { /// Create a rewrite context for `session`. - pub(crate) fn new(session: &'a VortexSession, scope: &'a DType) -> Self { + pub fn new(session: &'a VortexSession, scope: &'a DType) -> Self { Self { session, scope } } /// Returns the session that owns the rewrite registry. - pub(crate) fn session(&self) -> &'a VortexSession { + pub fn session(&self) -> &'a VortexSession { self.session } /// Return the dtype of `expr` within this rewrite scope. - pub(crate) fn return_dtype(&self, expr: &Expression) -> VortexResult { + pub fn return_dtype(&self, expr: &Expression) -> VortexResult { expr.return_dtype(self.scope) } /// Rewrite `expr` into a stats-backed falsifier. - pub(crate) fn falsify(&self, expr: &Expression) -> VortexResult> { + pub fn falsify(&self, expr: &Expression) -> VortexResult> { self.ensure_predicate(expr)?; rewrite(expr, self, StatsRewriteRule::falsify) } /// Rewrite `expr` into a stats-backed satisfier. - pub(crate) fn satisfy(&self, expr: &Expression) -> VortexResult> { + pub fn satisfy(&self, expr: &Expression) -> VortexResult> { self.ensure_predicate(expr)?; rewrite(expr, self, StatsRewriteRule::satisfy) } @@ -156,7 +156,6 @@ fn rewrite( #[cfg(test)] mod tests { use vortex_error::VortexResult; - use vortex_session::VortexSession; use super::StatsRewriteCtx; use super::StatsRewriteRule; @@ -169,7 +168,6 @@ mod tests { use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::literal::Literal; - use crate::stats::session::StatsSession; use crate::stats::session::StatsSessionExt; #[derive(Debug)] @@ -202,7 +200,7 @@ mod tests { #[test] fn combines_multiple_falsifiers_with_or() -> VortexResult<()> { - let session = VortexSession::empty().with::(); + let session = crate::array_session(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); session.stats().register_rewrite(StaticLiteralRule { falsifier: Some(lit(false)), @@ -222,7 +220,7 @@ mod tests { #[test] fn combines_multiple_satisfiers_with_or() -> VortexResult<()> { - let session = VortexSession::empty().with::(); + let session = crate::array_session(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); session.stats().register_rewrite(StaticLiteralRule { falsifier: None, @@ -242,7 +240,7 @@ mod tests { #[test] fn unregistered_expression_has_no_rewrite() -> VortexResult<()> { - let session = VortexSession::empty().with::(); + let session = crate::array_session(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); assert_eq!(lit(true).falsify(&dtype, &session)?, None); @@ -252,7 +250,7 @@ mod tests { #[test] fn non_predicate_expression_errors() { - let session = VortexSession::empty().with::(); + let session = crate::array_session(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); assert!(lit(7).falsify(&dtype, &session).is_err()); diff --git a/vortex-array/src/stats/rewrite/builtins.rs b/vortex-array/src/stats/rewrite/builtins.rs index 2b7316c7d98..ea656f24546 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/vortex-array/src/stats/rewrite/builtins.rs @@ -52,23 +52,26 @@ use crate::stats::session::StatsSession; /// Register built-in stats rewrite rules. pub(crate) fn register_builtins(session: &StatsSession) { - session.register_rewrite(BinaryStatsRewrite); + session.register_rewrite(BinaryNanCountStatsRewrite); + session.register_rewrite(BinaryAllNonNanStatsRewrite); session.register_rewrite(BetweenStatsRewrite); - session.register_rewrite(IsNullLegacyStatsRewrite); + session.register_rewrite(IsNullNullCountStatsRewrite); session.register_rewrite(IsNullAllNonNullStatsRewrite); session.register_rewrite(IsNullAllNullStatsRewrite); - session.register_rewrite(IsNotNullLegacyStatsRewrite); + session.register_rewrite(IsNotNullNullCountStatsRewrite); session.register_rewrite(IsNotNullAllNullStatsRewrite); session.register_rewrite(IsNotNullAllNonNullStatsRewrite); session.register_rewrite(LikeStatsRewrite); - session.register_rewrite(ListContainsStatsRewrite); - session.register_rewrite(DynamicComparisonStatsRewrite); + session.register_rewrite(ListContainsNanCountStatsRewrite); + session.register_rewrite(ListContainsAllNonNanStatsRewrite); + session.register_rewrite(DynamicComparisonNanCountStatsRewrite); + session.register_rewrite(DynamicComparisonAllNonNanStatsRewrite); } #[derive(Debug)] -struct BinaryStatsRewrite; +struct BinaryNanCountStatsRewrite; -impl StatsRewriteRule for BinaryStatsRewrite { +impl StatsRewriteRule for BinaryNanCountStatsRewrite { fn scalar_fn_id(&self) -> ScalarFnId { Binary.id() } @@ -78,60 +81,93 @@ impl StatsRewriteRule for BinaryStatsRewrite { expr: &Expression, ctx: &StatsRewriteCtx<'_>, ) -> VortexResult> { - let operator = expr.as_::(); - let lhs = expr.child(0); - let rhs = expr.child(1); - - Ok(match operator { - Operator::Eq => { - let left = min(lhs, ctx).zip(max(rhs, ctx)).map(|(a, b)| gt(a, b)); - let right = min(rhs, ctx).zip(max(lhs, ctx)).map(|(a, b)| gt(a, b)); - or_collect(left.into_iter().chain(right)) - .map(|value_predicate| with_nan_predicate(ctx, lhs, rhs, value_predicate)) - .transpose()? - } - Operator::NotEq => min(lhs, ctx) - .zip(max(rhs, ctx)) - .zip(max(lhs, ctx).zip(min(rhs, ctx))) - .map(|((min_lhs, max_rhs), (max_lhs, min_rhs))| { - with_nan_predicate( - ctx, - lhs, - rhs, - and(eq(min_lhs, max_rhs), eq(max_lhs, min_rhs)), - ) - }) - .transpose()?, - Operator::Gt => max(lhs, ctx) - .zip(min(rhs, ctx)) - .map(|(a, b)| with_nan_predicate(ctx, lhs, rhs, lt_eq(a, b))) - .transpose()?, - Operator::Gte => max(lhs, ctx) - .zip(min(rhs, ctx)) - .map(|(a, b)| with_nan_predicate(ctx, lhs, rhs, lt(a, b))) - .transpose()?, - Operator::Lt => min(lhs, ctx) - .zip(max(rhs, ctx)) - .map(|(a, b)| with_nan_predicate(ctx, lhs, rhs, gt_eq(a, b))) - .transpose()?, - Operator::Lte => min(lhs, ctx) - .zip(max(rhs, ctx)) - .map(|(a, b)| with_nan_predicate(ctx, lhs, rhs, gt(a, b))) - .transpose()?, - Operator::And => { - let lhs_falsifier = ctx.falsify(lhs)?; - let rhs_falsifier = ctx.falsify(rhs)?; - or_collect(lhs_falsifier.into_iter().chain(rhs_falsifier)) - } - Operator::Or => match (ctx.falsify(lhs)?, ctx.falsify(rhs)?) { - (Some(lhs), Some(rhs)) => Some(and(lhs, rhs)), - _ => None, - }, - Operator::Add | Operator::Sub | Operator::Mul | Operator::Div => None, - }) + binary_falsify::(expr, ctx) + } +} + +#[derive(Debug)] +struct BinaryAllNonNanStatsRewrite; + +impl StatsRewriteRule for BinaryAllNonNanStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + Binary.id() + } + + fn falsify( + &self, + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + binary_falsify::(expr, ctx) } } +fn binary_falsify( + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, +) -> VortexResult> { + let operator = expr.as_::(); + let lhs = expr.child(0); + let rhs = expr.child(1); + + Ok(match operator { + Operator::Eq => { + let left = min(lhs, ctx).zip(max(rhs, ctx)).map(|(a, b)| gt(a, b)); + let right = min(rhs, ctx).zip(max(lhs, ctx)).map(|(a, b)| gt(a, b)); + or_collect(left.into_iter().chain(right)) + .map(|value_predicate| with_non_nan_guards::

(ctx, [lhs, rhs], value_predicate)) + .transpose()? + .flatten() + } + Operator::NotEq => min(lhs, ctx) + .zip(max(rhs, ctx)) + .zip(max(lhs, ctx).zip(min(rhs, ctx))) + .map(|((min_lhs, max_rhs), (max_lhs, min_rhs))| { + with_non_nan_guards::

( + ctx, + [lhs, rhs], + and(eq(min_lhs, max_rhs), eq(max_lhs, min_rhs)), + ) + }) + .transpose()? + .flatten(), + Operator::Gt => max(lhs, ctx) + .zip(min(rhs, ctx)) + .map(|(a, b)| with_non_nan_guards::

(ctx, [lhs, rhs], lt_eq(a, b))) + .transpose()? + .flatten(), + Operator::Gte => max(lhs, ctx) + .zip(min(rhs, ctx)) + .map(|(a, b)| with_non_nan_guards::

(ctx, [lhs, rhs], lt(a, b))) + .transpose()? + .flatten(), + Operator::Lt => min(lhs, ctx) + .zip(max(rhs, ctx)) + .map(|(a, b)| with_non_nan_guards::

(ctx, [lhs, rhs], gt_eq(a, b))) + .transpose()? + .flatten(), + Operator::Lte => min(lhs, ctx) + .zip(max(rhs, ctx)) + .map(|(a, b)| with_non_nan_guards::

(ctx, [lhs, rhs], gt(a, b))) + .transpose()? + .flatten(), + Operator::And => { + if !P::EMIT_UNGUARDED_REWRITES { + return Ok(None); + } + + let lhs_falsifier = ctx.falsify(lhs)?; + let rhs_falsifier = ctx.falsify(rhs)?; + or_collect(lhs_falsifier.into_iter().chain(rhs_falsifier)) + } + Operator::Or => match (ctx.falsify(lhs)?, ctx.falsify(rhs)?) { + (Some(lhs), Some(rhs)) if P::EMIT_UNGUARDED_REWRITES => Some(and(lhs, rhs)), + _ => None, + }, + Operator::Add | Operator::Sub | Operator::Mul | Operator::Div => None, + }) +} + #[derive(Debug)] struct BetweenStatsRewrite; @@ -157,9 +193,9 @@ impl StatsRewriteRule for BetweenStatsRewrite { } #[derive(Debug)] -struct IsNullLegacyStatsRewrite; +struct IsNullNullCountStatsRewrite; -impl StatsRewriteRule for IsNullLegacyStatsRewrite { +impl StatsRewriteRule for IsNullNullCountStatsRewrite { fn scalar_fn_id(&self) -> ScalarFnId { IsNull.id() } @@ -217,9 +253,9 @@ impl StatsRewriteRule for IsNullAllNullStatsRewrite { } #[derive(Debug)] -struct IsNotNullLegacyStatsRewrite; +struct IsNotNullNullCountStatsRewrite; -impl StatsRewriteRule for IsNotNullLegacyStatsRewrite { +impl StatsRewriteRule for IsNotNullNullCountStatsRewrite { fn scalar_fn_id(&self) -> ScalarFnId { IsNotNull.id() } @@ -332,9 +368,9 @@ impl StatsRewriteRule for LikeStatsRewrite { } #[derive(Debug)] -struct ListContainsStatsRewrite; +struct ListContainsNanCountStatsRewrite; -impl StatsRewriteRule for ListContainsStatsRewrite { +impl StatsRewriteRule for ListContainsNanCountStatsRewrite { fn scalar_fn_id(&self) -> ScalarFnId { ListContains.id() } @@ -344,46 +380,71 @@ impl StatsRewriteRule for ListContainsStatsRewrite { expr: &Expression, ctx: &StatsRewriteCtx<'_>, ) -> VortexResult> { - let list = expr.child(0); - let needle = expr.child(1); + list_contains_falsify::(expr, ctx) + } +} - let Some(list_scalar) = literal_stat(list, Stat::Min) else { - return Ok(None); - }; - let elements = list_scalar - .as_opt::() - .and_then(|literal| literal.as_list_opt()) - .and_then(|list| list.elements()); - let Some(elements) = elements else { - return Ok(None); - }; - if elements.is_empty() { - return Ok(Some(lit(true))); - } +#[derive(Debug)] +struct ListContainsAllNonNanStatsRewrite; - let Some(value_max) = max(needle, ctx) else { - return Ok(None); - }; - let Some(value_min) = min(needle, ctx) else { - return Ok(None); - }; +impl StatsRewriteRule for ListContainsAllNonNanStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + ListContains.id() + } - let value_predicate = and_collect(elements.iter().map(|value| { - or( - lt(value_max.clone(), lit(value.clone())), - gt(value_min.clone(), lit(value.clone())), - ) - })); - value_predicate - .map(|value_predicate| with_all_non_nan_predicate(ctx, [needle], value_predicate)) - .transpose() + fn falsify( + &self, + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + list_contains_falsify::(expr, ctx) } } +fn list_contains_falsify( + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, +) -> VortexResult> { + let list = expr.child(0); + let needle = expr.child(1); + + let Some(list_scalar) = literal_stat(list, Stat::Min) else { + return Ok(None); + }; + let elements = list_scalar + .as_opt::() + .and_then(|literal| literal.as_list_opt()) + .and_then(|list| list.elements()); + let Some(elements) = elements else { + return Ok(None); + }; + if elements.is_empty() { + return Ok(P::EMIT_UNGUARDED_REWRITES.then(|| lit(true))); + } + + let Some(value_max) = max(needle, ctx) else { + return Ok(None); + }; + let Some(value_min) = min(needle, ctx) else { + return Ok(None); + }; + + let value_predicate = and_collect(elements.iter().map(|value| { + or( + lt(value_max.clone(), lit(value.clone())), + gt(value_min.clone(), lit(value.clone())), + ) + })); + value_predicate + .map(|value_predicate| with_non_nan_guards::

(ctx, [needle], value_predicate)) + .transpose() + .map(Option::flatten) +} + #[derive(Debug)] -struct DynamicComparisonStatsRewrite; +struct DynamicComparisonNanCountStatsRewrite; -impl StatsRewriteRule for DynamicComparisonStatsRewrite { +impl StatsRewriteRule for DynamicComparisonNanCountStatsRewrite { fn scalar_fn_id(&self) -> ScalarFnId { DynamicComparison.id() } @@ -393,31 +454,55 @@ impl StatsRewriteRule for DynamicComparisonStatsRewrite { expr: &Expression, ctx: &StatsRewriteCtx<'_>, ) -> VortexResult> { - let dynamic = expr.as_::(); - let lhs = expr.child(0); - - let Some((operator, lhs_stat)) = (match dynamic.operator { - CompareOperator::Eq | CompareOperator::NotEq => None, - CompareOperator::Gt => max(lhs, ctx).map(|lhs_stat| (CompareOperator::Lte, lhs_stat)), - CompareOperator::Gte => max(lhs, ctx).map(|lhs_stat| (CompareOperator::Lt, lhs_stat)), - CompareOperator::Lt => min(lhs, ctx).map(|lhs_stat| (CompareOperator::Gte, lhs_stat)), - CompareOperator::Lte => min(lhs, ctx).map(|lhs_stat| (CompareOperator::Gt, lhs_stat)), - }) else { - return Ok(None); - }; + dynamic_comparison_falsify::(expr, ctx) + } +} - let value_predicate = DynamicComparison.new_expr( - DynamicComparisonExpr { - operator, - rhs: Arc::clone(&dynamic.rhs), - default: !dynamic.default, - }, - [lhs_stat], - ); - with_all_non_nan_predicate(ctx, [lhs], value_predicate).map(Some) +#[derive(Debug)] +struct DynamicComparisonAllNonNanStatsRewrite; + +impl StatsRewriteRule for DynamicComparisonAllNonNanStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + DynamicComparison.id() + } + + fn falsify( + &self, + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + dynamic_comparison_falsify::(expr, ctx) } } +fn dynamic_comparison_falsify( + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, +) -> VortexResult> { + let dynamic = expr.as_::(); + let lhs = expr.child(0); + + let Some((operator, lhs_stat)) = (match dynamic.operator { + CompareOperator::Eq | CompareOperator::NotEq => None, + CompareOperator::Gt => max(lhs, ctx).map(|lhs_stat| (CompareOperator::Lte, lhs_stat)), + CompareOperator::Gte => max(lhs, ctx).map(|lhs_stat| (CompareOperator::Lt, lhs_stat)), + CompareOperator::Lt => min(lhs, ctx).map(|lhs_stat| (CompareOperator::Gte, lhs_stat)), + CompareOperator::Lte => min(lhs, ctx).map(|lhs_stat| (CompareOperator::Gt, lhs_stat)), + }) else { + return Ok(None); + }; + + let value_predicate = DynamicComparison.new_expr( + DynamicComparisonExpr { + operator, + rhs: Arc::clone(&dynamic.rhs), + default: !dynamic.default, + }, + [lhs_stat], + ); + with_non_nan_guards::

(ctx, [lhs], value_predicate) +} + fn min(expr: &Expression, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::Min, ctx) } @@ -438,36 +523,77 @@ fn all_non_null(expr: &Expression) -> Expression { stat_fn(expr.clone(), AllNonNull.bind(AggregateEmptyOptions)) } +enum NanCheck { + NotNeeded, + Check(Expression), + Unavailable, +} + +trait NonNanProof { + const EMIT_UNGUARDED_REWRITES: bool; + + fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult; +} + +struct NanCountProof; + +impl NonNanProof for NanCountProof { + const EMIT_UNGUARDED_REWRITES: bool = true; + + fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult { + non_nan_check(ctx, expr, |expr| { + match stat_expr(expr, Stat::NaNCount, ctx) { + Some(nan_count) => NanCheck::Check(eq(nan_count, lit(0u64))), + None => NanCheck::Unavailable, + } + }) + } +} + +struct AllNonNanProof; + +impl NonNanProof for AllNonNanProof { + const EMIT_UNGUARDED_REWRITES: bool = false; + + fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult { + non_nan_check(ctx, expr, |expr| { + NanCheck::Check(stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions))) + }) + } +} + // Min/max do not order NaN values, so comparison rewrites are only sound when every // candidate value is known to be non-NaN. Cast result dtypes are not enough: a cast // from float to non-float still needs a proof about the float source values. -fn all_non_nan_stat( +fn non_nan_check( ctx: &StatsRewriteCtx<'_>, expr: &Expression, -) -> VortexResult> { + proof: impl FnOnce(&Expression) -> NanCheck, +) -> VortexResult { if let Some(scalar) = expr.as_opt::() { let Some(value) = scalar.as_primitive_opt() else { - return Ok(None); + return Ok(NanCheck::NotNeeded); }; - return Ok(value.is_nan().then(|| lit(false))); + return Ok(if value.is_nan() { + NanCheck::Check(lit(false)) + } else { + NanCheck::NotNeeded + }); } if expr.is::() { if !has_nans(&ctx.return_dtype(expr.child(0))?) { - return Ok(None); + return Ok(NanCheck::NotNeeded); } - return all_non_nan_stat(ctx, expr.child(0)); + return non_nan_check(ctx, expr.child(0), proof); } if !has_nans(&ctx.return_dtype(expr)?) { - return Ok(None); + return Ok(NanCheck::NotNeeded); } - Ok(Some(stat_fn( - expr.clone(), - AllNonNan.bind(AggregateEmptyOptions), - ))) + Ok(proof(expr)) } fn has_nans(dtype: &DType) -> bool { @@ -501,33 +627,28 @@ fn stat_expr(expr: &Expression, stat: Stat, ctx: &StatsRewriteCtx<'_>) -> Option .then(|| stat_fn(expr.clone(), aggregate_fn)) } -fn with_nan_predicate( - ctx: &StatsRewriteCtx<'_>, - lhs: &Expression, - rhs: &Expression, - value_predicate: Expression, -) -> VortexResult { - with_all_non_nan_predicate(ctx, [lhs, rhs], value_predicate) -} - -fn with_all_non_nan_predicate<'a>( +fn with_non_nan_guards<'a, P: NonNanProof>( ctx: &StatsRewriteCtx<'_>, exprs: impl IntoIterator, value_predicate: Expression, -) -> VortexResult { +) -> VortexResult> { let mut nan_checks = Vec::new(); for expr in exprs { - if let Some(check) = all_non_nan_stat(ctx, expr)? { - nan_checks.push(check); + match P::check(ctx, expr)? { + NanCheck::NotNeeded => {} + NanCheck::Check(check) => nan_checks.push(check), + NanCheck::Unavailable => return Ok(None), } } let nan_predicate = and_collect(nan_checks); Ok(match nan_predicate { - Some(nan_check) => and(nan_check, value_predicate), + Some(nan_check) => Some(and(nan_check, value_predicate)), // No possible NaN-bearing expression remains, so the value predicate is - // already guarded. - None => value_predicate, + // 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, }) } @@ -616,10 +737,8 @@ mod tests { 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::session::StatsSession; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(crate::array_session); fn stat(expr: Expression, stat: Stat) -> Expression { let aggregate_fn = stat.aggregate_fn().expect("stat should have aggregate fn"); @@ -659,8 +778,17 @@ mod tests { expr.satisfy(&test_scope(), &SESSION) } - fn nan_free(expr: Expression) -> Expression { - stat_fn(expr, AllNonNan.bind(AggregateEmptyOptions)) + fn nan_guarded(expr: Expression, value_predicate: Expression) -> Expression { + or( + and( + eq(stat(expr.clone(), Stat::NaNCount), lit(0u64)), + value_predicate.clone(), + ), + and( + stat_fn(expr, AllNonNan.bind(AggregateEmptyOptions)), + value_predicate, + ), + ) } #[test] @@ -814,6 +942,33 @@ mod tests { )) ); + let expr = like(col("s"), lit(r"\%%")); + assert_eq!( + falsify(&expr)?, + Some(or( + gt_eq(stat(col("s"), Stat::Min), lit("&")), + lt(stat(col("s"), Stat::Max), lit("%")), + )) + ); + + let expr = like(col("s"), lit("pref%ix%")); + assert_eq!( + falsify(&expr)?, + Some(or( + gt_eq(stat(col("s"), Stat::Min), lit("preg")), + lt(stat(col("s"), Stat::Max), lit("pref")), + )) + ); + + let expr = like(col("s"), lit("pref_ix_")); + assert_eq!( + falsify(&expr)?, + Some(or( + gt_eq(stat(col("s"), Stat::Min), lit("preg")), + lt(stat(col("s"), Stat::Max), lit("pref")), + )) + ); + let expr = like(col("s"), lit("exact")); assert_eq!( falsify(&expr)?, @@ -860,8 +1015,8 @@ mod tests { assert_eq!( falsify(&expr)?, - Some(and( - nan_free(col("f")), + Some(nan_guarded( + col("f"), lt_eq(cast(stat(col("f"), Stat::Max), dtype), lit(5i32)), )) ); diff --git a/vortex-array/src/stats/session.rs b/vortex-array/src/stats/session.rs index 2d4325b2cd7..588a99ec372 100644 --- a/vortex-array/src/stats/session.rs +++ b/vortex-array/src/stats/session.rs @@ -7,8 +7,8 @@ use std::any::Any; use std::sync::Arc; use parking_lot::RwLock; -use vortex_session::Ref; use vortex_session::SessionExt; +use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_utils::aliases::hash_map::HashMap; @@ -20,15 +20,15 @@ use crate::stats::rewrite::register_builtins; type StatsRewriteRuleSet = Arc<[StatsRewriteRuleRef]>; /// Session state for stats APIs. -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct StatsSession { - rewrite_rules: RwLock>, + rewrite_rules: Arc>>, } impl Default for StatsSession { fn default() -> Self { let this = Self { - rewrite_rules: RwLock::new(HashMap::default()), + rewrite_rules: Arc::new(RwLock::new(HashMap::default())), }; register_builtins(&this); this @@ -37,14 +37,12 @@ impl Default for StatsSession { impl StatsSession { /// Register a stats rewrite rule. - #[allow(dead_code)] - pub(crate) fn register_rewrite(&self, rule: R) { + pub fn register_rewrite(&self, rule: R) { self.register_rewrite_ref(Arc::new(rule)); } /// Register a shared stats rewrite rule. - #[allow(dead_code)] - pub(crate) fn register_rewrite_ref(&self, rule: StatsRewriteRuleRef) { + pub fn register_rewrite_ref(&self, rule: StatsRewriteRuleRef) { let mut rules = self.rewrite_rules.write(); let rule_id = rule.scalar_fn_id(); let mut updated_rules = rules @@ -75,9 +73,9 @@ impl SessionVar for StatsSession { } /// Extension trait for accessing stats session data. -pub(crate) trait StatsSessionExt: SessionExt { +pub trait StatsSessionExt: SessionExt { /// Returns the stats session state. - fn stats(&self) -> Ref<'_, StatsSession> { + fn stats(&self) -> SessionGuard<'_, StatsSession> { self.get::() } } diff --git a/vortex-array/src/stats/stats_set.rs b/vortex-array/src/stats/stats_set.rs index 7613b746a33..6025afed2c6 100644 --- a/vortex-array/src/stats/stats_set.rs +++ b/vortex-array/src/stats/stats_set.rs @@ -565,8 +565,8 @@ mod test { use itertools::Itertools; use smallvec::smallvec; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::Nullability; @@ -877,7 +877,7 @@ mod test { .collect_vec(); array .statistics() - .compute_all(&all_stats, &mut LEGACY_SESSION.create_execution_ctx()) + .compute_all(&all_stats, &mut array_session().create_execution_ctx()) .unwrap(); let stats = array.statistics().to_owned(); diff --git a/vortex-array/src/test_harness.rs b/vortex-array/src/test_harness.rs index dabeb8bdde4..98be3a85ce4 100644 --- a/vortex-array/src/test_harness.rs +++ b/vortex-array/src/test_harness.rs @@ -8,8 +8,7 @@ use goldenfile::differs::binary_diff; use itertools::Itertools; use vortex_error::VortexResult; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; +use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; @@ -26,12 +25,12 @@ pub fn check_metadata(name: &str, metadata: &[u8]) { } /// Outputs the indices of the true values in a BoolArray -pub fn to_int_indices(indices_bits: BoolArray) -> VortexResult> { +pub fn to_int_indices(indices_bits: BoolArray, ctx: &mut ExecutionCtx) -> VortexResult> { let buffer = indices_bits.to_bit_buffer(); - let mask = indices_bits.as_ref().validity()?.execute_mask( - indices_bits.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let mask = indices_bits + .as_ref() + .validity()? + .execute_mask(indices_bits.as_ref().len(), ctx)?; Ok(buffer .iter() .enumerate() diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 8d36c07405e..9c80e34267e 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -2,6 +2,11 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Array validity and nullability behavior, used by arrays and compute functions. +//! +//! [`Validity`] describes which rows are logically present without forcing every array to carry a +//! materialized boolean bitmap. Constant states (`NonNullable`, `AllValid`, `AllInvalid`) are cheap +//! to clone and inspect. [`Validity::Array`] may itself be encoded or lazy, so APIs that need exact +//! per-row answers take an [`ExecutionCtx`] and execute the validity array as needed. use std::fmt::Debug; use std::ops::Range; @@ -34,14 +39,14 @@ use crate::scalar::Scalar; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::Operator; -/// Validity information for an array +/// Validity information for an array. #[derive(Clone)] pub enum Validity { - /// Items *can't* be null + /// Items cannot be null because the dtype is non-nullable. NonNullable, - /// All items are valid + /// The dtype is nullable, but every item is valid. AllValid, - /// All items are null + /// The dtype is nullable, and every item is null. AllInvalid, /// The validity of each position in the array is determined by a boolean array. /// @@ -123,6 +128,19 @@ impl Validity { matches!(self, Self::NonNullable | Self::AllValid) } + /// Returns `true` if this validity is *definitely* all-null (every value is null), i.e. it + /// is [`Validity::AllInvalid`]. + /// + /// Returning `false` does not prove that any value is valid: a [`Validity::Array`] may still + /// resolve to all-null once executed. Callers must treat `false` as "unknown without + /// compute". For a definitive answer, execute the validity with [`Self::execute_mask`] and + /// check whether the resulting [`Mask`] is all-false (`Mask::all_false`). This is the + /// all-null counterpart to [`Self::definitely_no_nulls`]. + #[inline] + pub fn definitely_all_null(&self) -> bool { + matches!(self, Self::AllInvalid) + } + /// Returns whether this validity contains no null values, executing the validity array if /// necessary. /// @@ -632,8 +650,8 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::Nullability; use crate::validity::BoolArray; @@ -702,7 +720,7 @@ mod tests { let indices = PrimitiveArray::new(Buffer::copy_from(positions), Validity::NonNullable).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert!( validity @@ -716,7 +734,7 @@ mod tests { #[test] #[should_panic] fn out_of_bounds_patch() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); Validity::NonNullable .patch( 2, @@ -768,7 +786,7 @@ mod tests { #[case] indices: ArrayRef, #[case] expected: Validity, ) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert!( validity .take(&indices) @@ -815,7 +833,7 @@ mod tests { #[case] rhs: Validity, #[case] expected: bool, ) -> vortex_error::VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_eq!(lhs.mask_eq(&rhs, 3, &mut ctx)?, expected); Ok(()) } diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index 3b793c6124a..9cd7fc92d30 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -23,6 +23,7 @@ vortex = { workspace = true, features = [ "tokio", "zstd", ] } +vortex-geo = { workspace = true } vortex-tensor = { workspace = true } # TODO(connor): In the future, this might be inside vortex. anyhow = { workspace = true } @@ -48,6 +49,9 @@ regex = { workspace = true } reqwest = { workspace = true, features = ["stream"] } serde = { workspace = true, features = ["derive"] } serde_json = { 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 } diff --git a/vortex-bench/spatialbench.sql b/vortex-bench/spatialbench.sql new file mode 100644 index 00000000000..5dcd7b72235 --- /dev/null +++ b/vortex-bench/spatialbench.sql @@ -0,0 +1,172 @@ +-- SpatialBench queries (DuckDB dialect), from sedona-spatialbench DuckDBSpatialBenchBenchmark +-- (spatialbench-queries/print_queries.py). Query logic is unchanged, only reformatted for readability +-- and numbered Q1..Q12 (canonical order). The harness splits the file on semicolons, so a comment +-- must never contain one. + +-- Q1: trips starting within 50km of Sedona city center, ordered by distance. +SELECT + t.t_tripkey, + ST_X(t.t_pickuploc) AS pickup_lon, + ST_Y(t.t_pickuploc) AS pickup_lat, + t.t_pickuptime, + ST_Distance(t.t_pickuploc, ST_GeomFromText('POINT (-111.7610 34.8697)')) AS distance_to_center +FROM trip t +WHERE ST_DWithin(t.t_pickuploc, ST_GeomFromText('POINT (-111.7610 34.8697)'), 0.45) +ORDER BY distance_to_center ASC, t.t_tripkey ASC; + +-- Q2: count trips starting within the Coconino County (Arizona) zone. +SELECT COUNT(*) AS trip_count_in_coconino_county +FROM trip t +WHERE ST_Intersects( + t.t_pickuploc, + (SELECT z.z_boundary FROM zone z WHERE z.z_name = 'Coconino County' LIMIT 1) +); + +-- Q3: monthly trip statistics within 15km of Sedona city center (10km bbox + 5km buffer). +SELECT + DATE_TRUNC('month', t.t_pickuptime) AS pickup_month, + COUNT(t.t_tripkey) AS total_trips, + AVG(t.t_distance) AS avg_distance, + AVG(t.t_dropofftime - t.t_pickuptime) AS avg_duration, + AVG(t.t_fare) AS avg_fare +FROM trip t +WHERE ST_DWithin( + t.t_pickuploc, + ST_GeomFromText('POLYGON((-111.9060 34.7347, -111.6160 34.7347, -111.6160 35.0047, -111.9060 35.0047, -111.9060 34.7347))'), + 0.045 +) +GROUP BY pickup_month +ORDER BY pickup_month; + +-- Q4: zone distribution of the top 1000 trips by tip amount. +SELECT z.z_zonekey, z.z_name, COUNT(*) AS trip_count +FROM zone z +JOIN ( + SELECT t.t_pickuploc + FROM trip t + ORDER BY t.t_tip DESC, t.t_tripkey ASC + LIMIT 1000 +) top_trips ON ST_Within(top_trips.t_pickuploc, z.z_boundary) +GROUP BY z.z_zonekey, z.z_name +ORDER BY trip_count DESC, z.z_zonekey ASC; + +-- Q5: monthly travel patterns for repeat customers (convex hull of dropoff locations). +SELECT + c.c_custkey, + c.c_name AS customer_name, + DATE_TRUNC('month', t.t_pickuptime) AS pickup_month, + ST_Area(ST_ConvexHull(ST_Collect(ARRAY_AGG(t.t_dropoffloc)))) AS monthly_travel_hull_area, + COUNT(*) AS dropoff_count +FROM trip t +JOIN customer c ON t.t_custkey = c.c_custkey +GROUP BY c.c_custkey, c.c_name, pickup_month +HAVING dropoff_count > 5 +ORDER BY dropoff_count DESC, c.c_custkey ASC; + +-- Q6: zone statistics for trips intersecting a bounding box. +SELECT + z.z_zonekey, + z.z_name, + COUNT(t.t_tripkey) AS total_pickups, + AVG(t.t_totalamount) AS avg_distance, + AVG(t.t_dropofftime - t.t_pickuptime) AS avg_duration +FROM trip t, zone z +WHERE ST_Intersects( + ST_GeomFromText('POLYGON((-112.2110 34.4197, -111.3110 34.4197, -111.3110 35.3197, -112.2110 35.3197, -112.2110 34.4197))'), + z.z_boundary +) +AND ST_Within(t.t_pickuploc, z.z_boundary) +GROUP BY z.z_zonekey, z.z_name +ORDER BY total_pickups DESC, z.z_zonekey ASC; + +-- Q7: detect potential route detours by comparing reported vs. geometric distances. +WITH trip_lengths AS ( + SELECT + t.t_tripkey, + t.t_distance AS reported_distance_m, + ST_Length(ST_MakeLine(t.t_pickuploc, t.t_dropoffloc)) / 0.000009 AS line_distance_m + FROM trip t +) +SELECT + t.t_tripkey, + t.reported_distance_m, + t.line_distance_m, + t.reported_distance_m / NULLIF(t.line_distance_m, 0) AS detour_ratio +FROM trip_lengths t +ORDER BY detour_ratio DESC NULLS LAST, reported_distance_m DESC, t_tripkey ASC; + +-- Q8: count nearby pickups for each building within ~500m. +SELECT b.b_buildingkey, b.b_name, COUNT(*) AS nearby_pickup_count +FROM trip t +JOIN building b ON ST_DWithin(t.t_pickuploc, b.b_boundary, 0.0045) +GROUP BY b.b_buildingkey, b.b_name +ORDER BY nearby_pickup_count DESC, b.b_buildingkey ASC; + +-- Q9: building conflation (duplicate/overlap detection via IoU). +WITH b1 AS ( + SELECT b_buildingkey AS id, b_boundary AS geom FROM building +), +b2 AS ( + SELECT b_buildingkey AS id, b_boundary AS geom FROM building +), +pairs AS ( + SELECT + b1.id AS building_1, + b2.id AS building_2, + ST_Area(b1.geom) AS area1, + ST_Area(b2.geom) AS area2, + ST_Area(ST_Intersection(b1.geom, b2.geom)) AS overlap_area + FROM b1 + JOIN b2 ON b1.id < b2.id AND ST_Intersects(b1.geom, b2.geom) +) +SELECT + building_1, + building_2, + area1, + area2, + overlap_area, + CASE + WHEN overlap_area = 0 THEN 0.0 + WHEN (area1 + area2 - overlap_area) = 0 THEN 1.0 + ELSE overlap_area / (area1 + area2 - overlap_area) + END AS iou +FROM pairs +ORDER BY iou DESC, building_1 ASC, building_2 ASC; + +-- Q10: zone statistics for trips starting within each zone. +SELECT + z.z_zonekey, + z.z_name AS pickup_zone, + AVG(t.t_dropofftime - t.t_pickuptime) AS avg_duration, + AVG(t.t_distance) AS avg_distance, + COUNT(t.t_tripkey) AS num_trips +FROM zone z +LEFT JOIN trip t ON ST_Within(t.t_pickuploc, z.z_boundary) +GROUP BY z.z_zonekey, z.z_name +ORDER BY avg_duration DESC NULLS LAST, z.z_zonekey ASC; + +-- Q11: count trips that cross between different zones. +SELECT COUNT(*) AS cross_zone_trip_count +FROM trip t +JOIN zone pickup_zone ON ST_Within(t.t_pickuploc, pickup_zone.z_boundary) +JOIN zone dropoff_zone ON ST_Within(t.t_dropoffloc, dropoff_zone.z_boundary) +WHERE pickup_zone.z_zonekey != dropoff_zone.z_zonekey; + +-- Q12: five nearest buildings per trip pickup (CROSS JOIN LATERAL, since DuckDB spatial has no ST_KNN). +SELECT + t.t_tripkey, + t.t_pickuploc, + nb.b_buildingkey, + nb.building_name, + nb.distance_to_building +FROM trip t +CROSS JOIN LATERAL ( + SELECT + b.b_buildingkey, + b.b_name AS building_name, + ST_Distance(t.t_pickuploc, b.b_boundary) AS distance_to_building + FROM building b + ORDER BY distance_to_building + LIMIT 5 +) AS nb +ORDER BY nb.distance_to_building, nb.b_buildingkey; diff --git a/vortex-bench/src/benchmark.rs b/vortex-bench/src/benchmark.rs index 2872a02aa64..fe16df5cd3d 100644 --- a/vortex-bench/src/benchmark.rs +++ b/vortex-bench/src/benchmark.rs @@ -8,6 +8,7 @@ use glob::Pattern; use url::Url; use crate::BenchmarkDataset; +use crate::Engine; use crate::Format; /// Specification for a table in a benchmark dataset. @@ -32,6 +33,12 @@ pub trait Benchmark: Send + Sync { /// Get all available queries for this benchmark fn queries(&self) -> anyhow::Result>; + /// SQL an `engine` must run before this benchmark's queries (e.g. loading engine + /// extensions). Runners replay these after every (re)open. Default: none. + fn engine_init_sql(&self, _engine: Engine) -> Vec { + Vec::new() + } + /// Generate or prepare base data for the benchmark (typically Parquet format). /// This is the canonical source data that can be converted to other formats. /// This should be idempotent - safe to call multiple times. diff --git a/vortex-bench/src/clickbench/benchmark.rs b/vortex-bench/src/clickbench/benchmark.rs index a68ce0b54e8..c036d1c084b 100644 --- a/vortex-bench/src/clickbench/benchmark.rs +++ b/vortex-bench/src/clickbench/benchmark.rs @@ -41,21 +41,41 @@ impl ClickBenchBenchmark { } } +/// ClickBench sorted by event time, with shard filenames shuffled to exercise sort pushdown. +pub struct ClickBenchSortedBenchmark { + pub queries_file: Option, + pub data_url: Url, +} + +impl ClickBenchSortedBenchmark { + /// Create the sorted ClickBench benchmark, optionally using a remote data directory. + pub fn new(use_remote_data_dir: Option) -> Result { + Ok(Self { + queries_file: None, + data_url: resolve_data_url(use_remote_data_dir.as_deref(), CLICKBENCH_SORTED_NAME)?, + }) + } +} + +fn read_clickbench_queries(queries_file: Option<&str>) -> Result> { + let queries_filepath = match queries_file { + Some(file) => file.into(), + None => Path::new(env!("CARGO_MANIFEST_DIR")).join("clickbench_queries.sql"), + }; + + Ok(fs::read_to_string(queries_filepath)? + .split(';') + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .enumerate() + .collect()) +} + #[async_trait::async_trait] impl Benchmark for ClickBenchBenchmark { fn queries(&self) -> Result> { - let queries_filepath = match &self.queries_file { - Some(file) => file.into(), - None => Path::new(env!("CARGO_MANIFEST_DIR")).join("clickbench_queries.sql"), - }; - - Ok(fs::read_to_string(queries_filepath)? - .split(';') - .map(|s| s.trim()) - .filter(|s| !s.is_empty()) - .map(|s| s.to_string()) - .enumerate() - .collect()) + read_clickbench_queries(self.queries_file.as_deref()) } async fn generate_base_data(&self) -> Result<()> { @@ -70,10 +90,7 @@ impl Benchmark for ClickBenchBenchmark { } fn expected_row_counts(&self) -> Option> { - Some(vec![ - 1, 1, 1, 1, 1, 1, 1, 18, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 4, 1, 10, 10, 10, - 10, 10, 10, 25, 25, 1, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, - ]) + Some(clickbench_expected_row_counts()) } fn dataset(&self) -> BenchmarkDataset { @@ -99,6 +116,48 @@ impl Benchmark for ClickBenchBenchmark { } } +#[async_trait::async_trait] +impl Benchmark for ClickBenchSortedBenchmark { + fn queries(&self) -> Result> { + Ok(read_clickbench_queries(self.queries_file.as_deref())? + .into_iter() + .filter(|(idx, _)| CLICKBENCH_SORTED_QUERY_IDS.contains(idx)) + .collect()) + } + + async fn generate_base_data(&self) -> Result<()> { + if self.data_url.scheme() != "file" { + return Ok(()); + } + + generate_sorted_clickbench(CLICKBENCH_SORTED_NAME.to_data_path()).await + } + + fn expected_row_counts(&self) -> Option> { + Some(clickbench_expected_row_counts()) + } + + fn dataset(&self) -> BenchmarkDataset { + BenchmarkDataset::ClickBenchSorted + } + + fn dataset_name(&self) -> &str { + CLICKBENCH_SORTED_NAME + } + + fn dataset_display(&self) -> String { + CLICKBENCH_SORTED_NAME.to_string() + } + + fn data_url(&self) -> &Url { + &self.data_url + } + + fn table_specs(&self) -> Vec { + vec![TableSpec::new("hits", Some(HITS_SCHEMA.clone()))] + } +} + fn clickbench_flavor(flavor: Flavor) -> String { format!("clickbench_{flavor}") } diff --git a/vortex-bench/src/clickbench/data.rs b/vortex-bench/src/clickbench/data.rs index ba2a2104192..bd570d6536c 100644 --- a/vortex-bench/src/clickbench/data.rs +++ b/vortex-bench/src/clickbench/data.rs @@ -3,25 +3,48 @@ use std::fmt; use std::fmt::Display; +use std::fs; +use std::fs::File; +use std::io::Write; use std::path::Path; +use std::path::PathBuf; +use std::process::Command; +use std::process::Stdio; use std::str::FromStr; use std::sync::LazyLock; +use anyhow::Context; use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::Schema; use arrow_schema::TimeUnit; use clap::ValueEnum; +use parquet::file::reader::FileReader; +use parquet::file::reader::SerializedFileReader; use serde::Deserialize; use serde::Serialize; use tracing::info; use vortex::error::VortexExpect; use crate::Format; +use crate::IdempotentPath; // Re-export for use by clickbench_benchmark pub use crate::conversions::convert_parquet_directory_to_vortex; use crate::datasets::data_downloads::download_data; use crate::datasets::data_downloads::download_many; +use crate::utils::file::temp_download_filepath; + +/// Benchmark and local data directory name for ClickBench sorted by event time. +pub const CLICKBENCH_SORTED_NAME: &str = "clickbench-sorted"; +const CLICKBENCH_PARTITIONED_NAME: &str = "clickbench_partitioned"; +const SORTED_SHARD_COUNT: usize = 100; +const SORTED_SHARD_COUNT_U64: u64 = 100; +const SORTED_SHARD_FILENAME_WIDTH: usize = 3; +const SORTED_SHARD_PERMUTATION_MULTIPLIER: u64 = 37; +const SORTED_SHARD_PERMUTATION_OFFSET: u64 = 17; + +/// Zero-based ClickBench query IDs that filter by or order/group on `EventDate`/`EventTime`. +pub const CLICKBENCH_SORTED_QUERY_IDS: &[usize] = &[23, 24, 26, 36, 37, 38, 39, 40, 41, 42]; pub static HITS_SCHEMA: LazyLock = LazyLock::new(|| { use DataType::*; @@ -142,6 +165,14 @@ pub static HITS_SCHEMA: LazyLock = LazyLock::new(|| { ]) }); +/// Expected result row counts for the 43 upstream ClickBench queries. +pub fn clickbench_expected_row_counts() -> Vec { + vec![ + 1, 1, 1, 1, 1, 1, 1, 18, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 4, 1, 10, 10, 10, 10, + 10, 10, 25, 25, 1, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, + ] +} + /// Clickbench has two different flavors: /// - Singe - 1 file containing the whole dataset, just under 100 million rows. /// - Partitioned (which we run by default) - 100 files, each containing ~1 million rows, all sharing the same schema. @@ -206,3 +237,294 @@ impl Flavor { Ok(()) } } + +/// Generate sorted ClickBench Parquet shards under `basepath`. +/// +/// Each shard contains a contiguous `EventTime` range, but the shard filenames are deterministically +/// shuffled. That makes ordinary file listing order bad for ascending TopK queries, so file sort +/// pushdown has a visible job to do. +pub async fn generate_sorted_clickbench(basepath: impl AsRef) -> anyhow::Result<()> { + let basepath = basepath.as_ref(); + let source_base = CLICKBENCH_PARTITIONED_NAME.to_data_path(); + Flavor::Partitioned.download(&source_base).await?; + + let source_parquet_dir = source_base.join(Format::Parquet.name()); + let output_parquet_dir = basepath.join(Format::Parquet.name()); + + if output_parquet_dir.exists() { + info!( + "Sorted ClickBench parquet already exists at {}", + output_parquet_dir.display() + ); + return Ok(()); + } + + let temp_root = temp_download_filepath(); + let result = + generate_sorted_clickbench_inner(&source_parquet_dir, &output_parquet_dir, &temp_root); + if result.is_err() { + drop(fs::remove_dir_all(&temp_root)); + } + result +} + +fn generate_sorted_clickbench_inner( + source_parquet_dir: &Path, + output_parquet_dir: &Path, + temp_root: &Path, +) -> anyhow::Result<()> { + let source_rows = parquet_dir_row_count(source_parquet_dir)?; + anyhow::ensure!( + source_rows > 0, + "ClickBench source parquet directory has no rows: {}", + source_parquet_dir.display() + ); + + fs::create_dir_all(temp_root) + .with_context(|| format!("Failed to create temp dir {}", temp_root.display()))?; + + let temp_output_dir = temp_root.join(Format::Parquet.name()); + let duckdb_temp_dir = temp_root.join("duckdb-tmp"); + fs::create_dir_all(&temp_output_dir) + .with_context(|| format!("Failed to create temp dir {}", temp_output_dir.display()))?; + fs::create_dir_all(&duckdb_temp_dir) + .with_context(|| format!("Failed to create temp dir {}", duckdb_temp_dir.display()))?; + + let script = sorted_clickbench_duckdb_script( + source_parquet_dir, + &temp_output_dir, + &duckdb_temp_dir, + source_rows, + ); + let db_path = temp_root.join("sort.duckdb"); + + info!( + "Generating globally sorted ClickBench parquet in {}", + temp_output_dir.display() + ); + + let mut child = Command::new("duckdb") + .arg(&db_path) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .context("Failed to run DuckDB CLI while generating sorted ClickBench data")?; + + let mut stdin = child + .stdin + .take() + .context("Failed to open DuckDB stdin while generating sorted ClickBench data")?; + stdin + .write_all(script.as_bytes()) + .context("Failed to write sorted ClickBench SQL to DuckDB stdin")?; + drop(stdin); + + let output = child + .wait_with_output() + .context("Failed to wait for DuckDB while generating sorted ClickBench data")?; + + if !output.status.success() { + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + anyhow::bail!( + "DuckDB failed generating sorted ClickBench data: stdout=\"{stdout}\", stderr=\"{stderr}\"" + ); + } + + let output_files = parquet_files(&temp_output_dir)?; + anyhow::ensure!( + output_files.len() == SORTED_SHARD_COUNT, + "Expected {SORTED_SHARD_COUNT} sorted ClickBench shards, got {}", + output_files.len() + ); + + let output_rows = parquet_files_row_count(output_files)?; + anyhow::ensure!( + output_rows == source_rows, + "Sorted ClickBench row-count mismatch: source={source_rows}, output={output_rows}" + ); + + if let Some(parent) = output_parquet_dir.parent() { + fs::create_dir_all(parent) + .with_context(|| format!("Failed to create output dir {}", parent.display()))?; + } + fs::rename(&temp_output_dir, output_parquet_dir).with_context(|| { + format!( + "Failed to move sorted ClickBench parquet from {} to {}", + temp_output_dir.display(), + output_parquet_dir.display() + ) + })?; + + drop(fs::remove_dir_all(temp_root)); + Ok(()) +} + +fn sorted_clickbench_duckdb_script( + source_parquet_dir: &Path, + output_parquet_dir: &Path, + duckdb_temp_dir: &Path, + source_rows: u64, +) -> String { + let source_glob = source_parquet_dir.join("hits_*.parquet"); + let rows_per_shard = source_rows.div_ceil(SORTED_SHARD_COUNT_U64); + let columns = HITS_SCHEMA + .fields() + .iter() + .map(|field| quote_identifier(field.name())) + .collect::>() + .join(", "); + + let mut script = format!( + "\ +PRAGMA temp_directory={temp_dir}; +CREATE TABLE hits_sorted AS + SELECT * + FROM read_parquet({source_glob}) + ORDER BY \"EventTime\"; +", + temp_dir = sql_string_literal(&duckdb_temp_dir.display().to_string()), + source_glob = sql_string_literal(&source_glob.display().to_string()), + ); + + for shard_idx in 0..SORTED_SHARD_COUNT_U64 { + let start = shard_idx * rows_per_shard; + let end = (start + rows_per_shard).min(source_rows); + let output_path = output_parquet_dir.join(sorted_shard_file_name(shard_idx)); + script.push_str(&format!( + "\ +COPY ( + SELECT {columns} + FROM hits_sorted + WHERE rowid >= {start} AND rowid < {end} + ORDER BY rowid +) TO {output_path} (FORMAT parquet, COMPRESSION zstd); +", + output_path = sql_string_literal(&output_path.display().to_string()), + )); + } + + script +} + +fn sorted_shard_file_name(sorted_shard_idx: u64) -> String { + debug_assert!(sorted_shard_idx < SORTED_SHARD_COUNT_U64); + let listing_shard_idx = sorted_shard_listing_idx(sorted_shard_idx); + format!( + "hits_{listing_shard_idx:0width$}.parquet", + width = SORTED_SHARD_FILENAME_WIDTH + ) +} + +fn sorted_shard_listing_idx(sorted_shard_idx: u64) -> u64 { + debug_assert!(sorted_shard_idx < SORTED_SHARD_COUNT_U64); + (sorted_shard_idx * SORTED_SHARD_PERMUTATION_MULTIPLIER + SORTED_SHARD_PERMUTATION_OFFSET) + % SORTED_SHARD_COUNT_U64 +} + +fn parquet_dir_row_count(parquet_dir: &Path) -> anyhow::Result { + let files = parquet_files(parquet_dir)?; + anyhow::ensure!( + !files.is_empty(), + "No Parquet files found in {}", + parquet_dir.display() + ); + parquet_files_row_count(files) +} + +fn parquet_files(parquet_dir: &Path) -> anyhow::Result> { + let mut files = fs::read_dir(parquet_dir) + .with_context(|| format!("Failed to read parquet dir {}", parquet_dir.display()))? + .map(|entry| entry.map(|entry| entry.path())) + .collect::, _>>() + .with_context(|| format!("Failed to list parquet dir {}", parquet_dir.display()))?; + + files.retain(|path| path.extension().is_some_and(|ext| ext == "parquet")); + files.sort(); + Ok(files) +} + +fn parquet_files_row_count(files: Vec) -> anyhow::Result { + let mut total = 0_u64; + for file_path in files { + let file = File::open(&file_path) + .with_context(|| format!("Failed to open parquet file {}", file_path.display()))?; + let reader = SerializedFileReader::new(file) + .with_context(|| format!("Failed to read parquet metadata {}", file_path.display()))?; + let rows = reader.metadata().file_metadata().num_rows(); + let rows = u64::try_from(rows).with_context(|| { + format!("Parquet row count was negative in {}", file_path.display()) + })?; + total = total.checked_add(rows).with_context(|| { + format!( + "Parquet row count overflow while reading {}", + file_path.display() + ) + })?; + } + Ok(total) +} + +fn sql_string_literal(value: &str) -> String { + format!("'{}'", value.replace('\'', "''")) +} + +fn quote_identifier(value: &str) -> String { + format!("\"{}\"", value.replace('"', "\"\"")) +} + +#[cfg(test)] +mod tests { + use vortex::utils::aliases::hash_set::HashSet; + + use super::*; + + #[test] + fn sorted_clickbench_shard_names_use_deterministic_shuffle() { + let names = (0..SORTED_SHARD_COUNT_U64) + .map(sorted_shard_file_name) + .collect::>(); + + assert_eq!(names[0], "hits_017.parquet"); + assert_eq!(names[1], "hits_054.parquet"); + assert_eq!(names[2], "hits_091.parquet"); + assert_eq!(names[99], "hits_080.parquet"); + + let unique_names = names.iter().collect::>(); + assert_eq!(unique_names.len(), SORTED_SHARD_COUNT); + + let mut sorted_names = names.clone(); + sorted_names.sort(); + assert_ne!(names, sorted_names); + + let mut reverse_names = sorted_names; + reverse_names.reverse(); + assert_ne!(names, reverse_names); + } + + #[test] + fn sorted_clickbench_script_sorts_for_topk_and_writes_shuffled_shards() { + let script = sorted_clickbench_duckdb_script( + Path::new("source"), + Path::new("out"), + Path::new("tmp"), + 1000, + ); + + assert!(script.contains("ORDER BY \"EventTime\";")); + + let first = script + .find("hits_017.parquet") + .expect("first sorted shard should use its shuffled shard name"); + let second = script + .find("hits_054.parquet") + .expect("second sorted shard should use its shuffled shard name"); + let third = script + .find("hits_091.parquet") + .expect("third sorted shard should use its shuffled shard name"); + + assert!(first < second); + assert!(second < third); + } +} diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 3f21ab30ba0..ad84ef61b1d 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -7,8 +7,11 @@ use std::path::PathBuf; use futures::StreamExt; use futures::TryStreamExt; +use parquet::arrow::AsyncArrowWriter; use parquet::arrow::ParquetRecordBatchStreamBuilder; use parquet::arrow::async_reader::ParquetRecordBatchStream; +use parquet::file::metadata::KeyValue; +use parquet::file::metadata::ParquetMetaData; use sysinfo::System; use tokio::fs::File; use tokio::fs::OpenOptions; @@ -22,16 +25,26 @@ use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ChunkedArray; +use vortex::array::arrays::ExtensionArray; +use vortex::array::arrays::Struct; +use vortex::array::arrays::StructArray; +use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::arrow::FromArrowArray; use vortex::array::builders::builder_with_capacity; use vortex::array::stream::ArrayStreamAdapter; use vortex::array::stream::ArrayStreamExt; use vortex::dtype::DType; +use vortex::dtype::StructFields; use vortex::dtype::arrow::FromArrowType; +use vortex::dtype::extension::ExtDType; +use vortex::dtype::extension::ExtDTypeRef; use vortex::error::VortexResult; use vortex::error::vortex_err; use vortex::file::WriteOptionsSessionExt; use vortex::session::VortexSession; +use vortex::utils::aliases::hash_set::HashSet; +use vortex_geo::extension::GeoMetadata; +use vortex_geo::extension::WellKnownBinary; use crate::CompactionStrategy; use crate::Format; @@ -115,9 +128,12 @@ pub async fn convert_parquet_file_to_vortex( ) -> anyhow::Result<()> { let file = File::open(parquet_path).await?; let builder = ParquetRecordBatchStreamBuilder::new(file).await?; - let dtype = DType::from_arrow(builder.schema().as_ref()); - let stream = parquet_to_vortex_stream(builder.build()?); + // 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 stream = parquet_to_vortex_stream(builder.build()?) + .map(move |chunk| chunk.and_then(|chunk| tag_geo_array(chunk, &geo_columns))); let mut output_file = OpenOptions::new() .write(true) @@ -222,3 +238,104 @@ pub async fn write_parquet_as_vortex( }) .await } + +/// 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<()> { + let builder = ParquetRecordBatchStreamBuilder::new(File::open(parquet_path).await?).await?; + let already_tagged = builder + .metadata() + .file_metadata() + .key_value_metadata() + .is_some_and(|kvs| kvs.iter().any(|kv| kv.key == "geo")); + if already_tagged { + return Ok(()); + } + + let schema = std::sync::Arc::clone(builder.schema()); + let mut reader = builder.build()?; + + let tmp_path = parquet_path.with_extension("parquet.tmp"); + let mut writer = AsyncArrowWriter::try_new(File::create(&tmp_path).await?, schema, None)?; + 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.close().await?; + tokio::fs::rename(&tmp_path, parquet_path).await?; + Ok(()) +} + +/// Column names a parquet file's GeoParquet `geo` metadata marks as geometry. +fn geoparquet_columns(metadata: &ParquetMetaData) -> HashSet { + metadata + .file_metadata() + .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(|value| { + value + .get("columns") + .and_then(serde_json::Value::as_object) + .map(|columns| columns.keys().cloned().collect()) + }) + .unwrap_or_default() +} + +/// The erased `vortex.geo.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()) +} + +/// Re-type the named binary columns of a struct `dtype` as `vortex.geo.wkb`, so the column +/// self-describes as geometry. +fn tag_geo_dtype(dtype: DType, geo_columns: &HashSet) -> VortexResult { + if geo_columns.is_empty() { + return Ok(dtype); + } + let DType::Struct(fields, nullability) = &dtype else { + return Ok(dtype); + }; + let names = fields.names().clone(); + let tagged = names + .iter() + .zip(fields.fields()) + .map(|(name, field)| { + if geo_columns.contains(name.as_ref()) && field.is_binary() { + Ok(DType::Extension(wkb_ext_dtype(&field)?)) + } else { + Ok(field) + } + }) + .collect::>>()?; + Ok(DType::Struct( + StructFields::new(names, tagged), + *nullability, + )) +} + +/// Wrap the named binary columns of a struct `chunk` as `vortex.geo.wkb` extension arrays. +fn tag_geo_array(chunk: ArrayRef, geo_columns: &HashSet) -> VortexResult { + if geo_columns.is_empty() { + return Ok(chunk); + } + let Some(struct_array) = chunk.as_opt::() else { + return Ok(chunk); + }; + let names = struct_array.names().clone(); + let validity = struct_array.struct_validity(); + let len = struct_array.len(); + let tagged = names + .iter() + .zip(struct_array.iter_unmasked_fields()) + .map(|(name, field)| { + if geo_columns.contains(name.as_ref()) && field.dtype().is_binary() { + let ext = wkb_ext_dtype(field.dtype())?; + Ok(ExtensionArray::try_new(ext, field.clone())?.into_array()) + } else { + Ok(field.clone()) + } + }) + .collect::>>()?; + Ok(StructArray::try_new(names, tagged, len, validity)?.into_array()) +} diff --git a/vortex-bench/src/datasets/mod.rs b/vortex-bench/src/datasets/mod.rs index 3e72ba69e7f..a550a712eba 100644 --- a/vortex-bench/src/datasets/mod.rs +++ b/vortex-bench/src/datasets/mod.rs @@ -67,8 +67,12 @@ pub enum BenchmarkDataset { TpcDS { scale_factor: String }, #[serde(rename = "clickbench")] ClickBench { flavor: Flavor }, + #[serde(rename = "clickbench-sorted")] + ClickBenchSorted, #[serde(rename = "public-bi")] PublicBi { name: String }, + #[serde(rename = "spatialbench")] + SpatialBench { scale_factor: String }, #[serde(rename = "statpopgen")] StatPopGen { n_rows: u64 }, #[serde(rename = "polarsignals")] @@ -86,7 +90,9 @@ impl BenchmarkDataset { BenchmarkDataset::TpcH { .. } => "tpch", BenchmarkDataset::TpcDS { .. } => "tpcds", BenchmarkDataset::ClickBench { .. } => "clickbench", + BenchmarkDataset::ClickBenchSorted => "clickbench-sorted", BenchmarkDataset::PublicBi { .. } => "public-bi", + BenchmarkDataset::SpatialBench { .. } => "spatialbench", BenchmarkDataset::StatPopGen { .. } => "statpopgen", BenchmarkDataset::PolarSignals { .. } => "polarsignals", BenchmarkDataset::Fineweb => "fineweb", @@ -105,7 +111,11 @@ impl Display for BenchmarkDataset { Flavor::Partitioned => write!(f, "clickbench-partitioned"), Flavor::Single => write!(f, "clickbench-single"), }, + BenchmarkDataset::ClickBenchSorted => write!(f, "clickbench-sorted"), BenchmarkDataset::PublicBi { name } => write!(f, "public-bi({name})"), + BenchmarkDataset::SpatialBench { scale_factor } => { + write!(f, "spatialbench(sf={scale_factor})") + } BenchmarkDataset::StatPopGen { n_rows } => write!(f, "statpopgen(n_rows={n_rows})"), BenchmarkDataset::PolarSignals { n_rows } => { write!(f, "polarsignals(n_rows={n_rows})") @@ -162,7 +172,9 @@ impl BenchmarkDataset { "customer", "lineitem", "nation", "orders", "part", "partsupp", "region", "supplier", ], - BenchmarkDataset::ClickBench { .. } | BenchmarkDataset::PublicBi { .. } => todo!(), + BenchmarkDataset::ClickBench { .. } | BenchmarkDataset::ClickBenchSorted => &["hits"], + BenchmarkDataset::PublicBi { .. } => todo!(), + BenchmarkDataset::SpatialBench { .. } => &["trip", "building", "customer", "zone"], BenchmarkDataset::StatPopGen { .. } => &["statpopgen"], BenchmarkDataset::PolarSignals { .. } => &["stacktraces"], BenchmarkDataset::Fineweb => &["fineweb"], diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index 30ff45c97a8..daef12c7e78 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -13,6 +13,7 @@ use anyhow::bail; use appian::AppianBenchmark; use clap::ValueEnum; use clickbench::ClickBenchBenchmark; +use clickbench::ClickBenchSortedBenchmark; use clickbench::Flavor; use fineweb::FinewebBenchmark; use itertools::Itertools; @@ -34,6 +35,8 @@ use vortex::file::VortexWriteOptions; use vortex::file::WriteStrategyBuilder; use vortex::utils::aliases::hash_map::HashMap; +use crate::spatialbench::SpatialBenchBenchmark; + pub mod appian; pub mod benchmark; pub mod clickbench; @@ -51,6 +54,7 @@ pub mod public_bi; pub mod random_access; pub mod realnest; pub mod runner; +pub mod spatialbench; pub mod statpopgen; pub mod tpcds; pub mod tpch; @@ -251,6 +255,8 @@ pub enum BenchmarkArg { Appian, #[clap(name = "clickbench")] ClickBench, + #[clap(name = "clickbench-sorted")] + ClickBenchSorted, #[clap(name = "tpch")] TpcH, #[clap(name = "tpcds")] @@ -265,6 +271,8 @@ pub enum BenchmarkArg { PolarSignals, #[clap(name = "public-bi")] PublicBi, + #[clap(name = "spatialbench")] + SpatialBench, } /// Default scale factor for TPC-related benchmarks @@ -287,6 +295,11 @@ pub fn create_benchmark(b: BenchmarkArg, opts: &Opts) -> anyhow::Result { + let remote_data_dir = opts.get_as::(REMOTE_DATA_KEY); + let benchmark = ClickBenchSortedBenchmark::new(remote_data_dir)?; + Ok(Box::new(benchmark) as _) + } BenchmarkArg::TpcH => { let scale_factor = opts.get(SCALE_FACTOR_KEY).unwrap_or(DEFAULT_SCALE_FACTOR); let remote_data_dir = opts.get_as::(REMOTE_DATA_KEY); @@ -326,6 +339,12 @@ pub fn create_benchmark(b: BenchmarkArg, opts: &Opts) -> anyhow::Result { + let scale_factor = opts.get(SCALE_FACTOR_KEY).unwrap_or(DEFAULT_SCALE_FACTOR); + let remote_data_dir = opts.get_as::(REMOTE_DATA_KEY); + let benchmark = SpatialBenchBenchmark::new(scale_factor.to_string(), remote_data_dir)?; + Ok(Box::new(benchmark) as _) + } } } diff --git a/vortex-bench/src/spatialbench/benchmark.rs b/vortex-bench/src/spatialbench/benchmark.rs new file mode 100644 index 00000000000..1f5ae37fc53 --- /dev/null +++ b/vortex-bench/src/spatialbench/benchmark.rs @@ -0,0 +1,153 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! SpatialBench benchmark implementation + +use std::fs; + +use url::Url; + +use crate::Benchmark; +use crate::BenchmarkDataset; +use crate::Engine; +use crate::Format; +use crate::TableSpec; +use crate::spatialbench::datagen; +use crate::spatialbench::datagen::Table; +use crate::utils::file::resolve_data_url; +use crate::workspace_root; + +/// SpatialBench geospatial 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 { + pub scale_factor: String, + pub data_url: Url, +} + +impl SpatialBenchBenchmark { + pub fn new(scale_factor: String, use_remote_data_dir: Option) -> anyhow::Result { + Ok(Self { + data_url: resolve_data_url( + use_remote_data_dir.as_deref(), + &format!("spatialbench/{scale_factor}"), + )?, + scale_factor, + }) + } +} + +#[async_trait::async_trait] +impl Benchmark for SpatialBenchBenchmark { + /// 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. + let queries_file = workspace_root() + .join("vortex-bench") + .join("spatialbench") + .with_extension("sql"); + let contents = fs::read_to_string(queries_file)?; + Ok(contents + .split_terminator(';') + .map(str::trim) + .filter(|stmt| !stmt.is_empty()) + .enumerate() + .map(|(idx, stmt)| (idx + 1, stmt.to_string())) + .collect()) + } + + async fn generate_base_data(&self) -> anyhow::Result<()> { + if self.data_url.scheme() != "file" { + return Ok(()); + } + let base_data_dir = self + .data_url + .to_file_path() + .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) { + 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?; + } + } + Ok(()) + } + + fn data_url(&self) -> &Url { + &self.data_url + } + + fn expected_row_counts(&self) -> Option> { + // Indexed by `query_idx` (1-based), so index 0 is a dummy and Q1's count is at index 1 (TPC-H + // convention). Only SF1.0 and SF10.0 are validated (like TPC-H); other scale factors return + // `None`. Each vec covers Q1..Q9 — the queries that finish — and is identical for Parquet and + // Vortex. Q10..Q12 are heavy spatial joins that time out, so they are left unvalidated (a + // 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]), + _ => None, + } + } + + fn dataset(&self) -> BenchmarkDataset { + BenchmarkDataset::SpatialBench { + scale_factor: self.scale_factor.clone(), + } + } + + fn dataset_name(&self) -> &str { + "spatialbench" + } + + fn dataset_display(&self) -> String { + format!("spatialbench(sf={})", self.scale_factor) + } + + fn table_specs(&self) -> Vec { + // `zone` is externally sourced and optional; register it only when present so queries that + // don't need it don't fail on the missing glob. + let zone_present = match self.data_url.to_file_path() { + Ok(base) => zone_parquet_present(&base.join(Format::Parquet.name())), + Err(()) => true, + }; + Table::ALL + .into_iter() + .filter(|table| !matches!(table, Table::Zone) || zone_present) + .map(|table| TableSpec::new(table.name(), None)) + .collect() + } + + /// Scope each table to its own `{table}_*.{ext}` files; the default globs every file in the + /// format dir, conflating the `trip` and `building` schemas. + fn pattern(&self, table_name: &str, format: Format) -> Option { + Some( + format!("{}_*.{}", table_name, format.ext()) + .parse() + .expect("valid glob pattern"), + ) + } + + /// DuckDB needs the `spatial` extension for `ST_*`; the runner replays it on each (re)open. + /// First INSTALL needs network. + fn engine_init_sql(&self, engine: Engine) -> Vec { + match engine { + Engine::DuckDB => vec!["INSTALL spatial;".to_string(), "LOAD spatial;".to_string()], + _ => Vec::new(), + } + } +} + +/// Whether an externally-sourced `zone_*.parquet` exists under `parquet_dir` (generated by the +/// upstream `spatialbench-cli`; see the module docs). +fn zone_parquet_present(parquet_dir: &std::path::Path) -> bool { + glob::glob(&parquet_dir.join("zone_*.parquet").to_string_lossy()) + .map(|mut paths| paths.next().is_some()) + .unwrap_or(false) +} diff --git a/vortex-bench/src/spatialbench/datagen/mod.rs b/vortex-bench/src/spatialbench/datagen/mod.rs new file mode 100644 index 00000000000..8ebd1b35d86 --- /dev/null +++ b/vortex-bench/src/spatialbench/datagen/mod.rs @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! SpatialBench data preparation. [`wkb`] generates the canonical WKB base tables (Parquet + Vortex); +//! the [`table`] catalog is the single source of truth for the base tables. + +pub mod table; +pub mod wkb; + +pub use table::Table; +pub use wkb::generate_tables; diff --git a/vortex-bench/src/spatialbench/datagen/table.rs b/vortex-bench/src/spatialbench/datagen/table.rs new file mode 100644 index 00000000000..c84de192fdf --- /dev/null +++ b/vortex-bench/src/spatialbench/datagen/table.rs @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The shared SpatialBench table catalog: the single source of truth for the table set, consumed by +//! both [`super::wkb`] data generation and benchmark table registration. + +/// A SpatialBench table. +#[derive(Clone, Copy)] +pub enum Table { + Trip, + Building, + Customer, + Zone, +} + +impl Table { + /// Every SpatialBench table, in registration order. + pub(crate) const ALL: [Table; 4] = [Table::Trip, Table::Building, Table::Customer, Table::Zone]; + + /// File stem under a format directory, e.g. `Trip` → `trip_{part}.parquet`. + pub(crate) fn name(self) -> &'static str { + match self { + Table::Trip => "trip", + Table::Building => "building", + Table::Customer => "customer", + Table::Zone => "zone", + } + } + + /// Whether this table is generated in-process from the scale factor. `Zone` is sourced externally. + pub(crate) fn is_generated(self) -> bool { + !matches!(self, Table::Zone) + } + + /// Geometry columns SpatialBench tables. + pub(crate) fn geometry_columns(self) -> &'static [&'static str] { + match self { + Table::Trip => &["t_pickuploc", "t_dropoffloc"], + Table::Building => &["b_boundary"], + Table::Zone => &["z_boundary"], + Table::Customer => &[], + } + } +} diff --git a/vortex-bench/src/spatialbench/datagen/wkb.rs b/vortex-bench/src/spatialbench/datagen/wkb.rs new file mode 100644 index 00000000000..422505a3623 --- /dev/null +++ b/vortex-bench/src/spatialbench/datagen/wkb.rs @@ -0,0 +1,136 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! SpatialBench WKB base-table generation via the `spatialbench` crates (a tpchgen-rs fork). +//! Geometry is emitted as WKB, which DuckDB reads directly as `GEOMETRY` via `ST_GeomFromWKB`. + +use std::fs; +use std::path::PathBuf; +use std::sync::Arc; + +use anyhow::Result; +// spatialbench emits arrow-56 batches, so they must be written with its matching arrow-56 +// parquet crate, not the workspace's arrow-58 one. The parquet file itself is version-neutral. +use spatialbench::generators::BuildingGenerator; +use spatialbench::generators::CustomerGenerator; +use spatialbench::generators::TripGenerator; +use spatialbench_arrow::BuildingArrow; +use spatialbench_arrow::CustomerArrow; +use spatialbench_arrow::RecordBatchIterator; +use spatialbench_arrow::TripArrow; +use spatialbench_parquet::arrow::AsyncArrowWriter; +use spatialbench_parquet::basic::Compression; +use spatialbench_parquet::file::properties::WriterProperties; +use spatialbench_parquet::format::KeyValue; +use tokio::fs::File as TokioFile; +use tracing::info; + +use super::table::Table; +use crate::Format; +use crate::utils::file::idempotent_async; + +/// Batch size matching the TPC-H generator's streaming batches. +const BATCH_SIZE: usize = 8192 * 64; + +impl Table { + /// Batch iterator for one partition of this table, from the arrow-56 `spatialbench` crates. Only + /// called for generated tables (see [`Table::is_generated`]). + fn batch_iterator( + self, + scale_factor: f64, + part: i32, + part_count: i32, + ) -> Box { + match self { + Table::Trip => Box::new( + TripArrow::new(TripGenerator::new(scale_factor, part, part_count)) + .with_batch_size(BATCH_SIZE), + ), + Table::Building => Box::new( + BuildingArrow::new(BuildingGenerator::new(scale_factor, part, part_count)) + .with_batch_size(BATCH_SIZE), + ), + Table::Customer => Box::new( + CustomerArrow::new(CustomerGenerator::new(scale_factor, part, part_count)) + .with_batch_size(BATCH_SIZE), + ), + Table::Zone => unreachable!("zone is sourced externally, not generated in-process"), + } + } +} + +/// Generate the SpatialBench base tables as parquet under `{output_dir}/parquet/`. +pub async fn generate_tables(scale_factor: &str, output_dir: PathBuf) -> Result<()> { + let scale_factor = scale_factor.parse::()?; + let parquet_dir = output_dir.join(Format::Parquet.name()); + fs::create_dir_all(&parquet_dir)?; + + // One part per unit of scale factor keeps each file near the ~350MB the trip generator + // produces at SF1. + #[expect(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + let num_parts = (scale_factor.ceil() as usize).max(1); + let part_count = i32::try_from(num_parts)?; + + for table in Table::ALL.into_iter().filter(|table| table.is_generated()) { + for part_idx in 0..num_parts { + let output_file = parquet_dir.join(format!("{}_{part_idx}.parquet", table.name())); + let part = i32::try_from(part_idx + 1)?; + + idempotent_async(output_file.to_string_lossy().as_ref(), |path| async move { + info!( + scale_factor, + part, + part_count, + table = table.name(), + "Generating SpatialBench table" + ); + + let iter = table.batch_iterator(scale_factor, part, part_count); + let schema = Arc::clone(iter.schema()); + + let file = TokioFile::create(&path).await?; + let props = WriterProperties::builder() + .set_compression(Compression::SNAPPY) + .build(); + let mut writer = AsyncArrowWriter::try_new(file, schema, Some(props))?; + for batch in iter { + writer.write(&batch).await?; + } + // 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))); + } + writer.close().await?; + + Ok::<(), anyhow::Error>(()) + }) + .await?; + } + } + + Ok(()) +} + +/// GeoParquet metadata for WKB geometry columns, or `None` when it has none. +pub(crate) fn geo_parquet_metadata(table: Table) -> Option { + let geometry_columns = table.geometry_columns(); + let primary = geometry_columns.first()?; + let columns: serde_json::Map = geometry_columns + .iter() + .map(|&column| { + ( + column.to_string(), + serde_json::json!({ "encoding": "WKB", "geometry_types": [] }), + ) + }) + .collect(); + Some( + serde_json::json!({ + "version": "1.0.0", + "primary_column": primary, + "columns": columns, + }) + .to_string(), + ) +} diff --git a/vortex-bench/src/spatialbench/mod.rs b/vortex-bench/src/spatialbench/mod.rs new file mode 100644 index 00000000000..bba06bd7ef9 --- /dev/null +++ b/vortex-bench/src/spatialbench/mod.rs @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! SpatialBench geospatial analytics benchmark. +//! +//! See . + +pub mod benchmark; +pub mod datagen; + +pub use benchmark::SpatialBenchBenchmark; diff --git a/vortex-bench/src/utils/constants.rs b/vortex-bench/src/utils/constants.rs index c9f5ca4c442..959ad63b7f4 100644 --- a/vortex-bench/src/utils/constants.rs +++ b/vortex-bench/src/utils/constants.rs @@ -4,3 +4,4 @@ /// Storage type constants pub const STORAGE_NVME: &str = "nvme"; pub const STORAGE_S3: &str = "s3"; +pub const STORAGE_GCS: &str = "gcs"; diff --git a/vortex-bench/src/utils/file.rs b/vortex-bench/src/utils/file.rs index fb7ecc668fb..f05113963f7 100644 --- a/vortex-bench/src/utils/file.rs +++ b/vortex-bench/src/utils/file.rs @@ -161,11 +161,13 @@ pub fn resolve_data_url(remote_data_dir: Option<&str>, local_subdir: &str) -> Re /// - A storage type string ("s3", "nvme") /// - Or an error if the scheme is unknown pub fn url_scheme_to_storage(url: &Url) -> Result { + use super::constants::STORAGE_GCS; use super::constants::STORAGE_NVME; use super::constants::STORAGE_S3; match url.scheme() { STORAGE_S3 => Ok(STORAGE_S3.to_owned()), + "gs" => Ok(STORAGE_GCS.to_owned()), "file" => Ok(STORAGE_NVME.to_owned()), otherwise => { bail!("unknown URL scheme: {}", otherwise) diff --git a/vortex-bench/src/v3.rs b/vortex-bench/src/v3.rs index 48e8a7f1c94..a7529866f80 100644 --- a/vortex-bench/src/v3.rs +++ b/vortex-bench/src/v3.rs @@ -288,12 +288,14 @@ fn canonical_tpc_scale_factor(scale_factor: &str) -> String { /// | `TpcH { scale_factor }` | `tpch` | `None` | TPC SF as string (`"1"`, `"10"`, `"100"`, `"1000"`) | Run through `canonical_tpc_scale_factor` so `"1.0"` and `"1"` collapse. | /// | `TpcDS { scale_factor }` | `tpcds` | `None` | TPC SF as string | Same canonicalization as TPC-H. | /// | `ClickBench { flavor: _ }` | `clickbench` | `None` | `None` | Migrate path drops flavor; live emitter matches so historical and live merge. | +/// | `ClickBenchSorted` | `clickbench-sorted` | `None` | `None` | New live-only suite; keep separate from unsorted ClickBench history. | /// | `StatPopGen { n_rows: _ }` | `statpopgen` | `None` | `None` | Migrate path carries no SF for this suite; live drops it for the same reason. | /// | `PolarSignals { n_rows: _ }`| `polarsignals` | `None` | `None` | Same as StatPopGen. | /// | `Fineweb` | `fineweb` | `None` | `None` | | /// | `GhArchive` | `gharchive` | `None` | `None` | | /// | `Appian` | `appian` | `None` | `None` | Static dataset; no scale factor. | /// | `PublicBi { name }` | `public-bi` | dataset name (e.g. `cms-provider`) | `None` | Sub-dataset name lives in `dataset_variant`. | +/// | `SpatialBench { scale_factor }` | `spatialbench` | `None` | SF as string | Same canonicalization as TPC-H; no historical v2 records to merge with. | pub fn benchmark_dataset_dims(d: &BenchmarkDataset) -> (String, Option, Option) { match d { BenchmarkDataset::TpcH { scale_factor } => ( @@ -311,6 +313,7 @@ pub fn benchmark_dataset_dims(d: &BenchmarkDataset) -> (String, Option, // same to keep historical and live records in one `clickbench` group. // Flavor is fixed per CI matrix entry and recoverable from there. BenchmarkDataset::ClickBench { .. } => ("clickbench".to_string(), None, None), + BenchmarkDataset::ClickBenchSorted => ("clickbench-sorted".to_string(), None, None), BenchmarkDataset::PublicBi { name } => ("public-bi".to_string(), Some(name.clone()), None), // StatPopGen / PolarSignals: the migrate path (v2 → v3 backfill) does // not carry a per-record scale factor for these suites, so writing one @@ -318,6 +321,11 @@ pub fn benchmark_dataset_dims(d: &BenchmarkDataset) -> (String, Option, // live). Drop it to keep live ingests merging into the migrated // group. The dataset-level `n_rows` is recoverable from the bench // matrix if ever needed. + BenchmarkDataset::SpatialBench { scale_factor } => ( + "spatialbench".to_string(), + None, + Some(canonical_tpc_scale_factor(scale_factor)), + ), BenchmarkDataset::StatPopGen { .. } => ("statpopgen".to_string(), None, None), BenchmarkDataset::PolarSignals { .. } => ("polarsignals".to_string(), None, None), BenchmarkDataset::Fineweb => ("fineweb".to_string(), None, None), @@ -726,6 +734,16 @@ mod tests { } } + #[test] + fn clickbench_sorted_dims_are_distinct_from_clickbench() { + let (dataset, variant, scale_factor) = + benchmark_dataset_dims(&BenchmarkDataset::ClickBenchSorted); + + assert_eq!(dataset, "clickbench-sorted"); + assert_eq!(variant, None); + assert_eq!(scale_factor, None); + } + #[test] fn compression_records_lowercase_dataset_for_v2_history_match() { // The v2 → v3 migrate classifier stores `dataset = series.to_lowercase()` diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 34a98f85b73..5ee2bad5402 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -13,6 +13,9 @@ repository = { workspace = true } rust-version = { workspace = true } version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [dependencies] itertools = { workspace = true } num-traits = { workspace = true } diff --git a/vortex-btrblocks/benches/compress.rs b/vortex-btrblocks/benches/compress.rs index 8bc47d16404..6898b4a9ab4 100644 --- a/vortex-btrblocks/benches/compress.rs +++ b/vortex-btrblocks/benches/compress.rs @@ -17,14 +17,12 @@ mod benchmarks { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; - use vortex_array::session::ArraySession; use vortex_btrblocks::BtrBlocksCompressor; use vortex_buffer::buffer_mut; use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn make_clickbench_window_name() -> ArrayRef { // A test that's meant to mirror the WindowName column from ClickBench. diff --git a/vortex-btrblocks/benches/compress_listview.rs b/vortex-btrblocks/benches/compress_listview.rs index e1b0595ac27..502acf79e0f 100644 --- a/vortex-btrblocks/benches/compress_listview.rs +++ b/vortex-btrblocks/benches/compress_listview.rs @@ -21,7 +21,6 @@ mod benchmarks { use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::FieldNames; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; use vortex_buffer::buffer_mut; @@ -30,8 +29,7 @@ mod benchmarks { const NUM_ROWS: usize = 8192; const SEED: u64 = 42; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); const SHORT_STRINGS: &[&str] = &[ "alpha_one", diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 43df788473a..6489f2d4e52 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -41,6 +41,9 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ &integer::RunEndScheme, &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. //////////////////////////////////////////////////////////////////////////////////////////////// @@ -74,10 +77,10 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ /// Builder for creating configured [`BtrBlocksCompressor`] instances. /// -/// By default, all schemes in [`ALL_SCHEMES`] are enabled. Feature-gated schemes (Pco, Zstd) -/// are not in `ALL_SCHEMES` and must be added explicitly via -/// [`with_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) or -/// [`with_compact`](BtrBlocksCompressorBuilder::with_compact). +/// By default, all schemes in [`ALL_SCHEMES`] are enabled in a deterministic order. Feature-gated +/// schemes (Pco, Zstd) are not in `ALL_SCHEMES` and must be added explicitly via +/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) or `with_compact` when the +/// `zstd` feature is enabled. /// /// # Examples /// @@ -163,6 +166,9 @@ impl BtrBlocksCompressorBuilder { /// 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. + /// + /// 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 @@ -180,6 +186,10 @@ impl BtrBlocksCompressorBuilder { ]; #[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")] + excluded.push(integer::DeltaScheme::default().id()); let builder = self.exclude_schemes(excluded); #[cfg(all(feature = "zstd", feature = "unstable_encodings"))] diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index 752b01d2f6e..9bcee31d990 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -83,7 +83,6 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; @@ -94,8 +93,7 @@ mod tests { #[cfg(feature = "zstd")] use crate::BtrBlocksCompressorBuilder; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[rstest] #[case::zctl( @@ -122,6 +120,7 @@ mod tests { #[case] input: ListViewArray, #[case] expect_list: bool, ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let array_ref = input.clone().into_array(); let result = BtrBlocksCompressor::default() .compress(&array_ref, &mut SESSION.create_execution_ctx())?; @@ -130,12 +129,13 @@ mod tests { } else { assert!(result.as_opt::().is_some()); } - assert_arrays_eq!(result, input); + assert_arrays_eq!(result, input, &mut ctx); Ok(()) } #[test] fn test_constant_all_true() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let array = BoolArray::new(BitBuffer::from(vec![true; 100]), Validity::NonNullable); let btr = BtrBlocksCompressor::default(); let compressed = btr.compress( @@ -143,12 +143,13 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } #[test] fn test_constant_all_false() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let array = BoolArray::new(BitBuffer::from(vec![false; 100]), Validity::NonNullable); let btr = BtrBlocksCompressor::default(); let compressed = btr.compress( @@ -156,12 +157,13 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } #[test] fn test_nullable_all_valid_compressed() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true; 100]), Validity::from(BitBuffer::from(vec![true; 100])), @@ -172,12 +174,13 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } #[test] fn test_nullable_with_nulls_not_compressed() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let validity = Validity::from(BitBuffer::from_iter((0..100).map(|i| i % 3 != 0))); let array = BoolArray::new(BitBuffer::from(vec![true; 100]), validity); let btr = BtrBlocksCompressor::default(); @@ -186,12 +189,13 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(!compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } #[test] fn test_mixed_not_constant() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true, false, true, false, true]), Validity::NonNullable, @@ -202,12 +206,13 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(!compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } #[test] fn test_binary_constant_compressed() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let values = vec![Some(b"constant-bytes".as_slice()); 100]; let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)); let btr = BtrBlocksCompressor::default(); @@ -216,12 +221,13 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } #[test] fn test_binary_dict_compressed() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let distinct_values: [&[u8]; 3] = [b"alpha", b"beta", b"gamma"]; let values = (0..1000) .map(|idx| Some(distinct_values[idx % distinct_values.len()])) @@ -233,7 +239,7 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } @@ -254,17 +260,15 @@ mod tests { ); let compressor = BtrBlocksCompressorBuilder::default().with_compact().build(); - let compressed = compressor.compress( - &array.clone().into_array(), - &mut SESSION.create_execution_ctx(), - )?; + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; assert!( compressed.is::(), "expected Zstd, got {}", compressed.encoding_id() ); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } @@ -287,17 +291,15 @@ mod tests { let compressor = BtrBlocksCompressorBuilder::default() .only_cuda_compatible() .build(); - let compressed = compressor.compress( - &array.clone().into_array(), - &mut SESSION.create_execution_ctx(), - )?; + 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_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } } diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 39db05246a6..69ddb59738c 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -40,16 +40,28 @@ //! # Example //! //! ```rust +//! use vortex_array::{IntoArray, VortexSessionExecute, array_session}; +//! use vortex_array::arrays::PrimitiveArray; +//! use vortex_array::validity::Validity; //! use vortex_btrblocks::{BtrBlocksCompressor, BtrBlocksCompressorBuilder, Scheme, SchemeExt}; //! use vortex_btrblocks::schemes::integer::IntDictScheme; +//! use vortex_buffer::buffer; +//! +//! # fn example() -> vortex_error::VortexResult<()> { +//! let session = array_session(); +//! let array = PrimitiveArray::new(buffer![42u64; 1024], Validity::NonNullable).into_array(); //! -//! // Default compressor with all schemes enabled. //! let compressor = BtrBlocksCompressor::default(); +//! let compressed = compressor.compress(&array, &mut session.create_execution_ctx())?; +//! assert_eq!(compressed.dtype(), array.dtype()); //! //! // Remove specific schemes using the builder. //! let compressor = BtrBlocksCompressorBuilder::default() //! .exclude_schemes([IntDictScheme.id()]) //! .build(); +//! # let _ = compressor; +//! # Ok(()) +//! # } //! ``` //! //! [BtrBlocks]: https://www.cs.cit.tum.de/fileadmin/w00cfj/dis/papers/btrblocks.pdf diff --git a/vortex-btrblocks/src/schemes/float/alprd.rs b/vortex-btrblocks/src/schemes/float/alprd.rs index f97ae077882..969afabad49 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/vortex-btrblocks/src/schemes/float/alprd.rs @@ -66,7 +66,7 @@ impl Scheme for ALPRDScheme { ptype => vortex_panic!("cannot ALPRD compress ptype {ptype}"), }; - let alp_rd = encoder.encode(primitive_array, exec_ctx); + let alp_rd = encoder.encode(primitive_array); let dtype = alp_rd.dtype().clone(); let right_bit_width = alp_rd.right_bit_width(); let mut parts = ALPRDArrayOwnedExt::into_data_parts(alp_rd); @@ -82,7 +82,6 @@ impl Scheme for ALPRDScheme { parts.right_parts, right_bit_width, parts.left_parts_patches, - exec_ctx, )? .into_array()) } diff --git a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs index 9ddddef6df1..4c7a5f85fa6 100644 --- a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs @@ -14,7 +14,6 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::builders::ArrayBuilder; use vortex_array::builders::PrimitiveBuilder; use vortex_array::dtype::Nullability; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; @@ -22,8 +21,7 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_constant_compressed() -> VortexResult<()> { diff --git a/vortex-btrblocks/src/schemes/float/tests.rs b/vortex-btrblocks/src/schemes/float/tests.rs index 31fd88d99d8..2d0a04542cc 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -12,7 +12,6 @@ use vortex_array::builders::ArrayBuilder; use vortex_array::builders::PrimitiveBuilder; use vortex_array::display::DisplayOptions; use vortex_array::dtype::Nullability; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer_mut; @@ -23,9 +22,7 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::float::FloatRLEScheme; - -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_empty() -> VortexResult<()> { @@ -60,6 +57,7 @@ fn test_compress() -> VortexResult<()> { #[test] fn test_rle_compression() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let mut values = Vec::new(); values.extend(iter::repeat_n(1.5f32, 100)); values.extend(iter::repeat_n(2.7f32, 200)); @@ -73,7 +71,7 @@ fn test_rle_compression() -> VortexResult<()> { assert!(compressed.is::()); let expected = Buffer::copy_from(&values).into_array(); - assert_arrays_eq!(compressed, expected); + assert_arrays_eq!(compressed, expected, &mut ctx); Ok(()) } diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs new file mode 100644 index 00000000000..f25bc6ce466 --- /dev/null +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -0,0 +1,199 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! FastLanes Delta integer encoding. + +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_compressor::builtins::BinaryDictScheme; +use vortex_compressor::builtins::FloatDictScheme; +use vortex_compressor::builtins::IntDictScheme; +use vortex_compressor::builtins::StringDictScheme; +use vortex_compressor::estimate::CompressionEstimate; +use vortex_compressor::estimate::DeferredEstimate; +use vortex_compressor::estimate::EstimateScore; +use vortex_compressor::estimate::EstimateVerdict; +use vortex_compressor::scheme::AncestorExclusion; +use vortex_compressor::scheme::ChildSelection; +use vortex_compressor::scheme::DescendantExclusion; +use vortex_error::VortexResult; +use vortex_fastlanes::Delta; + +use crate::ArrayAndStats; +use crate::CascadingCompressor; +use crate::CompressorContext; +use crate::GenerateStatsOptions; +use crate::Scheme; +use crate::SchemeExt; + +/// FastLanes Delta encoding for smooth / near-monotone integers. +/// +/// Delta replaces each value with its difference from an earlier value (at the FastLanes lane +/// stride), so a later cascade layer (FoR / BitPacking) packs the smaller residuals. It only +/// pays off when those residuals span meaningfully fewer bits than the values themselves. +/// +/// The minimum penalized compression ratio required for Delta to be selected is configurable via +/// [`DeltaScheme::new`]; [`DeltaScheme::default`] uses a ratio of `1.25`. +#[derive(Debug, Copy, Clone, PartialEq)] +pub struct DeltaScheme { + min_ratio: f64, +} + +impl DeltaScheme { + /// Creates a Delta scheme requiring `min_ratio` after the delta penalty before it wins. + /// + /// Pass a higher ratio to make Delta more conservative, or a lower one to select it more + /// eagerly. [`DeltaScheme::default`] uses a ratio of `1.25`. + pub const fn new(min_ratio: f64) -> Self { + Self { min_ratio } + } +} + +impl Default for DeltaScheme { + fn default() -> Self { + Self::new(1.25) + } +} + +/// Multiplicative penalty applied to Delta's estimated compression ratio. +/// +/// Unlike FoR/BitPacking, Delta breaks random access and adds a prefix-sum decode pass, and it +/// carries a structural sign bit on its residuals. We therefore require Delta to be meaningfully +/// (~5%) smaller than the best alternative before it wins, rather than picking it for a +/// single-bit gain. This factor encodes that "delta tax". +const DELTA_PENALTY: f64 = 0.95; + +/// Minimum length before Delta is worth considering (one FastLanes chunk). +const MIN_DELTA_LEN: usize = 1024; + +impl Scheme for DeltaScheme { + fn scheme_name(&self) -> &'static str { + "vortex.int.delta" + } + + fn matches(&self, canonical: &Canonical) -> bool { + canonical.dtype().is_int() + } + + fn num_children(&self) -> usize { + 2 + } + + /// Delta-encode the data at most once per path: exclude Delta from the subtrees of both the + /// bases and the deltas children so we never delta-encode data that was already delta-encoded. + fn descendant_exclusions(&self) -> Vec { + vec![DescendantExclusion { + excluded: self.id(), + children: ChildSelection::All, + }] + } + + /// Delta over dictionary codes just adds indirection: codes are compact integers with no + /// monotone structure, so (like FoR/Sequence) skip the codes child. + fn ancestor_exclusions(&self) -> Vec { + vec![ + AncestorExclusion { + ancestor: IntDictScheme.id(), + children: ChildSelection::One(1), + }, + AncestorExclusion { + ancestor: FloatDictScheme.id(), + children: ChildSelection::One(1), + }, + AncestorExclusion { + ancestor: StringDictScheme.id(), + children: ChildSelection::One(1), + }, + AncestorExclusion { + ancestor: BinaryDictScheme.id(), + children: ChildSelection::One(1), + }, + ] + } + + fn expected_compression_ratio( + &self, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + // Delta only pays off if a later cascade layer (FoR/BitPacking) packs the residuals. + if compress_ctx.finished_cascading() { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + } + // Too short to transpose into FastLanes chunks meaningfully. + if data.array_len() < MIN_DELTA_LEN { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + } + + // Estimating Delta needs the real transposed-delta span, so defer to a callback that + // delta-encodes the array and measures the residual range. + let min_ratio = self.min_ratio; + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + 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; + + // 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) { + return Ok(EstimateVerdict::Skip); + } + + // Measure the actual FastLanes transposed-delta span. This is the lane-stride + // difference that gets bit-packed, not the lag-1 difference (which the transpose + // makes optimistic), so it is what truly drives the compressed size. + let (_bases, deltas) = vortex_fastlanes::delta_compress(&primitive, exec_ctx)?; + let delta_stats = + ArrayAndStats::new(deltas.into_array(), GenerateStatsOptions::default()); + let span = delta_stats.integer_stats(exec_ctx).erased().max_minus_min(); + + // Bits needed to FoR-pack the residuals. A zero span means constant deltas, which + // SequenceScheme already captures more cheaply, so defer to it. + let delta_bits = match span.checked_ilog2() { + Some(l) => (l + 1) as f64, + None => return Ok(EstimateVerdict::Skip), + }; + + let ratio = full_width / delta_bits * DELTA_PENALTY; + if ratio <= min_ratio { + return Ok(EstimateVerdict::Skip); + } + Ok(EstimateVerdict::Ratio(ratio)) + }, + ))) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let primitive = data.array().clone().execute::(exec_ctx)?; + let len = primitive.len(); + let (bases, deltas) = vortex_fastlanes::delta_compress(&primitive, exec_ctx)?; + + let compressed_bases = compressor.compress_child( + &bases.into_array(), + &compress_ctx, + self.id(), + 0, + exec_ctx, + )?; + let compressed_deltas = compressor.compress_child( + &deltas.into_array(), + &compress_ctx, + self.id(), + 1, + exec_ctx, + )?; + + Delta::try_new(compressed_bases, compressed_deltas, 0, len).map(IntoArray::into_array) + } +} diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/vortex-btrblocks/src/schemes/integer/for_.rs index 86accdb5e71..ec23f1114db 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -123,7 +123,7 @@ impl Scheme for FoRScheme { exec_ctx: &mut ExecutionCtx, ) -> VortexResult { let primitive = data.array().clone().execute::(exec_ctx)?; - let for_array = FoR::encode(primitive)?; + let for_array = FoR::encode(primitive, exec_ctx)?; let biased = for_array .encoded() .clone() diff --git a/vortex-btrblocks/src/schemes/integer/mod.rs b/vortex-btrblocks/src/schemes/integer/mod.rs index aed29f1ad3d..abe5868f5c8 100644 --- a/vortex-btrblocks/src/schemes/integer/mod.rs +++ b/vortex-btrblocks/src/schemes/integer/mod.rs @@ -4,6 +4,8 @@ //! Integer compression schemes. mod bitpacking; +#[cfg(feature = "unstable_encodings")] +mod delta; mod for_; mod rle; mod runend; @@ -15,6 +17,8 @@ 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; diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index 2e0fb269fda..e4227a472ec 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -17,7 +17,6 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; use vortex_array::expr::stats::StatsProviderExt; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; @@ -29,9 +28,7 @@ use vortex_session::VortexSession; use vortex_sparse::Sparse; use crate::BtrBlocksCompressor; - -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_constant_compressed() -> VortexResult<()> { @@ -143,7 +140,11 @@ fn test_sequence_compressed() -> VortexResult<()> { fn test_rle_compressed() -> VortexResult<()> { let mut values: Vec = Vec::new(); for i in 0..1024 { - values.extend(iter::repeat_n(i, 10)); + // Scramble the per-run value so the data is run-length-dominant but not monotone: this + // keeps RunEnd the winner instead of Delta (whose residuals would be small on a smooth + // ramp). + let v = (i as u32).wrapping_mul(2_654_435_761) as i32; + values.extend(iter::repeat_n(v, 10)); } let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); let btr = BtrBlocksCompressor::default(); @@ -152,3 +153,59 @@ fn test_rle_compressed() -> VortexResult<()> { assert!(compressed.is::()); Ok(()) } + +/// A strictly-increasing column with small, irregular steps: not a perfect arithmetic sequence +/// (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(); + 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..4096) + .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, got tree:\n{}", + compressed.display_tree() + ); + // Delta must appear at most once per tree: no Delta node may be nested under another. + assert!( + !has_nested_delta(&compressed, false), + "Delta was applied more than once in the tree:\n{}", + compressed.display_tree() + ); + assert_arrays_eq!(compressed, array.into_array(), &mut ctx); + 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; + + let is_delta = array.is::(); + if is_delta && under_delta { + return true; + } + array + .children() + .iter() + .any(|child| has_nested_delta(child, under_delta || is_delta)) +} diff --git a/vortex-btrblocks/src/schemes/integer/tests.rs b/vortex-btrblocks/src/schemes/integer/tests.rs index 2f559af8582..a9ef24dc0e0 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -15,7 +15,6 @@ use vortex_array::arrays::Dict; use vortex_array::arrays::Masked; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -28,9 +27,7 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::integer::IntRLEScheme; - -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_empty() -> VortexResult<()> { @@ -74,6 +71,7 @@ fn test_dict_encodable() -> VortexResult<()> { #[test] fn constant_mostly_nulls() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::new( buffer![189u8, 189, 189, 189, 189, 189, 189, 189, 189, 0, 46], Validity::from_iter(vec![ @@ -91,12 +89,13 @@ fn constant_mostly_nulls() -> VortexResult<()> { let decoded = compressed; let expected = PrimitiveArray::new(buffer![0u8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 46], validity).into_array(); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn nullable_sequence() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let values = (0i32..20).step_by(7).collect_vec(); let array = PrimitiveArray::from_option_iter(values.clone().into_iter().map(Some)); @@ -106,12 +105,13 @@ fn nullable_sequence() -> VortexResult<()> { let decoded = compressed; let expected = PrimitiveArray::from_option_iter(values.into_iter().map(Some)).into_array(); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn test_rle_compression() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let mut values = Vec::new(); values.extend(iter::repeat_n(42i32, 100)); values.extend(iter::repeat_n(123i32, 200)); @@ -124,7 +124,7 @@ fn test_rle_compression() -> VortexResult<()> { assert!(compressed.is::()); let expected = Buffer::copy_from(&values).into_array(); - assert_arrays_eq!(compressed, expected); + assert_arrays_eq!(compressed, expected, &mut ctx); Ok(()) } diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index a90fbf4f452..5e6baf73933 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -64,9 +64,9 @@ impl Scheme for FSSTScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let utf8 = data.array_as_varbinview().into_owned(); - let compressor_fsst = fsst_train_compressor(&utf8); - let fsst = fsst_compress(&utf8, utf8.len(), utf8.dtype(), &compressor_fsst, exec_ctx); + let utf8 = data.array_as_varbinview().into_owned().into_array(); + let compressor_fsst = fsst_train_compressor(&utf8, exec_ctx)?; + let fsst = fsst_compress(&utf8, &compressor_fsst, exec_ctx)?; let uncompressed_lengths_primitive = fsst .uncompressed_lengths() diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index 123601f7bb1..7bbe480d059 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -73,7 +73,7 @@ impl Scheme for OnPairScheme { exec_ctx: &mut ExecutionCtx, ) -> VortexResult { let utf8 = data.array_as_varbinview().into_owned(); - let onpair_array = onpair_compress(&utf8, utf8.len(), utf8.dtype(), DEFAULT_DICT12_CONFIG)?; + let onpair_array = onpair_compress(utf8.as_array(), DEFAULT_DICT12_CONFIG, exec_ctx)?; let dict_offsets = compress_offsets_child( compressor, diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index a32ac353a7f..c1473146607 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -12,15 +12,13 @@ use vortex_array::arrays::Dict; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::session::ArraySession; use vortex_error::VortexResult; use vortex_fsst::FSST; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_constant_compressed() -> VortexResult<()> { diff --git a/vortex-btrblocks/src/schemes/string/tests.rs b/vortex-btrblocks/src/schemes/string/tests.rs index 05b256eda6d..ec9c4060f64 100644 --- a/vortex-btrblocks/src/schemes/string/tests.rs +++ b/vortex-btrblocks/src/schemes/string/tests.rs @@ -11,14 +11,12 @@ use vortex_array::builders::VarBinViewBuilder; use vortex_array::display::DisplayOptions; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::session::ArraySession; use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_strings() -> VortexResult<()> { diff --git a/vortex-btrblocks/tests/onpair_roundtrip.rs b/vortex-btrblocks/tests/onpair_roundtrip.rs index 3843e02c319..272fa930bd0 100644 --- a/vortex-btrblocks/tests/onpair_roundtrip.rs +++ b/vortex-btrblocks/tests/onpair_roundtrip.rs @@ -17,16 +17,13 @@ use std::sync::LazyLock; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::accessor::ArrayAccessor; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::session::ArraySession; use vortex_btrblocks::BtrBlocksCompressor; use vortex_session::VortexSession; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); /// Helper: synthetic short-string corpus that the cascading compressor should /// route through OnPair. @@ -74,19 +71,21 @@ fn nonnullable_roundtrip_via_default_compressor() { .execute::(&mut SESSION.create_execution_ctx()) .expect("decompress"); assert_eq!(decoded.len(), n); - decoded - .with_iterator(|iter| { - for (i, got) in iter.enumerate() { - assert_eq!( - got, - Some(strings[i].as_bytes()), - "mismatch at row {i}: got {:?}", - got.map(|b| String::from_utf8_lossy(b).into_owned()), - ); - } - Ok::<_, vortex_error::VortexError>(()) - }) + let mask = decoded + .validity() + .unwrap() + .execute_mask(decoded.len(), &mut SESSION.create_execution_ctx()) .unwrap(); + for i in 0..decoded.len() { + let got = mask.value(i).then(|| decoded.bytes_at(i)); + assert_eq!( + got.as_deref(), + Some(strings[i].as_bytes()), + "mismatch at row {i}: got {:?}", + got.as_deref() + .map(|b| String::from_utf8_lossy(b).into_owned()), + ); + } } #[test] @@ -114,15 +113,16 @@ fn nullable_roundtrip_via_default_compressor() { .execute::(&mut SESSION.create_execution_ctx()) .expect("decompress"); assert_eq!(decoded.len(), n); - decoded - .with_iterator(|iter| { - for (i, got) in iter.enumerate() { - let want = strings[i].as_deref().map(str::as_bytes); - assert_eq!(got, want, "mismatch at row {i}"); - } - Ok::<_, vortex_error::VortexError>(()) - }) + let mask = decoded + .validity() + .unwrap() + .execute_mask(decoded.len(), &mut SESSION.create_execution_ctx()) .unwrap(); + for i in 0..decoded.len() { + let got = mask.value(i).then(|| decoded.bytes_at(i)); + let want = strings[i].as_deref().map(str::as_bytes); + assert_eq!(got.as_deref(), want, "mismatch at row {i}"); + } } /// Larger corpus that exercises the offsets-narrowing / delta-encoding paths @@ -147,14 +147,15 @@ fn large_unique_short_strings_roundtrip() { .execute::(&mut SESSION.create_execution_ctx()) .expect("decompress"); assert_eq!(decoded.len(), n); - decoded - .with_iterator(|iter| { - for (i, got) in iter.enumerate() { - assert_eq!(got, Some(strings[i].as_bytes()), "row {i}"); - } - Ok::<_, vortex_error::VortexError>(()) - }) + let mask = decoded + .validity() + .unwrap() + .execute_mask(decoded.len(), &mut SESSION.create_execution_ctx()) .unwrap(); + for i in 0..decoded.len() { + let got = mask.value(i).then(|| decoded.bytes_at(i)); + assert_eq!(got.as_deref(), Some(strings[i].as_bytes()), "row {i}"); + } } #[test] @@ -173,15 +174,17 @@ fn empty_and_short_string_roundtrip() { let decoded = compressed .execute::(&mut SESSION.create_execution_ctx()) .expect("decompress"); - decoded - .with_iterator(|iter| { - let got: Vec<_> = iter.collect(); - for (i, want) in strings.iter().enumerate() { - assert_eq!(got[i], Some(want.as_bytes()), "row {i}"); - } - Ok::<_, vortex_error::VortexError>(()) - }) + let mask = decoded + .validity() + .unwrap() + .execute_mask(decoded.len(), &mut SESSION.create_execution_ctx()) .unwrap(); + let got: Vec>> = (0..decoded.len()) + .map(|i| mask.value(i).then(|| decoded.bytes_at(i).to_vec())) + .collect(); + for (i, want) in strings.iter().enumerate() { + assert_eq!(got[i].as_deref(), Some(want.as_bytes()), "row {i}"); + } } /// Regression for the Euro2016 compress-bench panic @@ -217,12 +220,13 @@ fn delta_dict_offsets_roundtrip() { .execute::(&mut SESSION.create_execution_ctx()) .expect("decompress"); assert_eq!(decoded.len(), n); - decoded - .with_iterator(|iter| { - for (i, got) in iter.enumerate() { - assert_eq!(got, Some(strings[i].as_bytes()), "row {i}"); - } - Ok::<_, vortex_error::VortexError>(()) - }) + let mask = decoded + .validity() + .unwrap() + .execute_mask(decoded.len(), &mut SESSION.create_execution_ctx()) .unwrap(); + for i in 0..decoded.len() { + let got = mask.value(i).then(|| decoded.bytes_at(i)); + assert_eq!(got.as_deref(), Some(strings[i].as_bytes()), "row {i}"); + } } diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index ae9d7e6cc05..e0397bacc91 100644 --- a/vortex-buffer/Cargo.toml +++ b/vortex-buffer/Cargo.toml @@ -39,6 +39,7 @@ workspace = true [dev-dependencies] divan = { workspace = true } num-traits = { workspace = true } +rand = { workspace = true } rstest = { workspace = true } [[bench]] diff --git a/vortex-buffer/src/alignment.rs b/vortex-buffer/src/alignment.rs index bcd17a64f10..ad5c51761bd 100644 --- a/vortex-buffer/src/alignment.rs +++ b/vortex-buffer/src/alignment.rs @@ -8,26 +8,33 @@ use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::vortex_err; -/// Default alignment for device-to-host buffer copies. -pub const ALIGNMENT_TO_HOST_COPY: Alignment = Alignment::new(256); - /// The alignment of a buffer. /// -/// This type is a wrapper around `usize` that ensures the alignment is a power of 2 and fits into -/// a `u16`. +/// This type is a wrapper around `usize` that ensures the alignment is a non-zero power of 2. #[derive(Clone, Debug, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct Alignment(usize); impl Alignment { + /// Default alignment for device-to-host buffer copies. + pub const HOST_COPY: Self = Alignment::new(256); + + /// Default alignment for all buffers. + /// + /// Chosen to be larger than any SIMD register (e.g. AVX-512's 64-byte + /// registers) so that buffers can be processed with vectorized loads/stores + /// without alignment fixups, and to match the alignment guarantees of the + /// CUDA allocator (256 bytes) so host buffers can be copied to/from device + /// memory without re-alignment. + pub const DEFAULT_ALIGNMENT: Self = Alignment::new(256); + /// Create a new alignment. /// /// ## Panics /// - /// Panics if `align` is not a power of 2, or is greater than `u16::MAX`. + /// Panics if `align` is zero or is not a power of 2. #[inline] pub const fn new(align: usize) -> Self { assert!(align > 0, "Alignment must be greater than 0"); - assert!(align <= u16::MAX as usize, "Alignment must fit into u16"); assert!(align.is_power_of_two(), "Alignment must be a power of 2"); Self(align) } @@ -54,6 +61,9 @@ impl Alignment { Self::new(align_of::()) } + /// The largest valid alignment: the greatest power of 2 representable in a `usize`. + pub const MAX: Alignment = Alignment::new(1 << (usize::BITS - 1)); + /// Check if `self` alignment is a "larger" than `other` alignment. /// /// ## Example @@ -67,23 +77,45 @@ impl Alignment { /// assert!(!b.is_aligned_to(a)); /// ``` #[inline] - pub fn is_aligned_to(&self, other: Alignment) -> bool { - // Since we know alignments are powers of 2, we can compare them by checking if the number - // of trailing zeros in the binary representation of the alignment is greater or equal. - self.0.trailing_zeros() >= other.0.trailing_zeros() + pub const fn is_aligned_to(&self, other: Alignment) -> bool { + // Since both alignments are powers of 2, divisibility is equivalent to ordering. + self.0 >= other.0 + } + + /// Check if the given byte offset (or length) is a multiple of this alignment. + /// + /// ## Example + /// + /// ``` + /// use vortex_buffer::Alignment; + /// + /// let a = Alignment::new(4); + /// assert!(a.is_offset_aligned(8)); + /// assert!(!a.is_offset_aligned(2)); + /// ``` + #[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 + } + + /// Check if the given pointer is aligned to this alignment. + #[inline] + pub fn is_ptr_aligned(&self, ptr: *const T) -> bool { + self.is_offset_aligned(ptr.addr()) } /// Returns the log2 of the alignment. pub fn exponent(&self) -> u8 { u8::try_from(self.0.trailing_zeros()) - .vortex_expect("alignment fits into u16, so exponent fits in u7") + .vortex_expect("alignment is a power of 2 within usize, so its exponent fits in u8") } /// Create from the log2 exponent of the alignment. /// /// ## Panics /// - /// Panics if `alignment` is not a power of 2, or is greater than `u16::MAX`. + /// Panics if `1 << exponent` overflows `usize`. #[inline] pub const fn from_exponent(exponent: u8) -> Self { Self::new(1 << exponent) @@ -126,13 +158,6 @@ impl From for usize { } } -impl From for u16 { - #[inline] - fn from(value: Alignment) -> Self { - u16::try_from(value.0).vortex_expect("Alignment must fit into u16") - } -} - impl From for u32 { #[inline] fn from(value: Alignment) -> Self { @@ -150,9 +175,6 @@ impl TryFrom for Alignment { if value == 0 { return Err(vortex_err!("Alignment must be greater than 0")); } - if value > u16::MAX as usize { - return Err(vortex_err!("Alignment must fit into u16, got {value}")); - } if !value.is_power_of_two() { return Err(vortex_err!("Alignment must be a power of 2, got {value}")); } @@ -172,9 +194,11 @@ mod test { } #[test] - #[should_panic] - fn alignment_overflow() { - Alignment::new(u16::MAX as usize + 1); + 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, Alignment::from_exponent(16)); } #[test] @@ -204,6 +228,10 @@ mod test { Ok(alignment) => assert_eq!(alignment, Alignment::new(8)), Err(err) => panic!("unexpected error for valid alignment: {err}"), } + match Alignment::try_from(1u32 << 16) { + Ok(alignment) => assert_eq!(alignment, Alignment::new(1 << 16)), + Err(err) => panic!("64KiB alignment should be valid: {err}"), + } assert!(Alignment::try_from(0u32).is_err()); assert!(Alignment::try_from(3u32).is_err()); } diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index 457827f4ab9..a8bd9e1929c 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -228,10 +228,11 @@ impl BitBuffer { unsafe { buffer.push_unchecked(packed) } } - buffer.truncate(len.div_ceil(8)); + let mut bytes = buffer.into_byte_buffer(); + bytes.truncate(len.div_ceil(8)); Self { - buffer: buffer.freeze().into_byte_buffer(), + buffer: bytes.freeze(), offset: 0, len, } @@ -271,6 +272,23 @@ impl BitBuffer { &self.buffer } + /// Return the backing bytes for this bit buffer when its logical offset is byte-aligned. + /// + /// The returned slice contains exactly `self.len().div_ceil(8)` bytes. Bits past the logical + /// length in the final byte are outside the buffer's logical range and should be ignored by + /// callers. + #[inline] + pub fn byte_aligned_bytes(&self) -> Option<&[u8]> { + if !self.offset.is_multiple_of(8) { + return None; + } + + let n_bytes = self.len.div_ceil(8); + let start = self.offset / 8; + let end = start + n_bytes; + Some(&self.buffer.as_slice()[start..end]) + } + /// Retrieve the value at the given index. /// /// Panics if the index is out of bounds. @@ -312,7 +330,23 @@ impl BitBuffer { assert!(end <= self.len); let len = end - start; - Self::new_with_offset(self.buffer.clone(), len, self.offset + start) + let offset = self.offset + start; + let byte_offset = offset / 8; + let bit_offset = offset % 8; + + // 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. + let buffer = if byte_offset != 0 { + self.buffer.slice_unaligned(byte_offset..) + } else { + self.buffer.clone().aligned(Alignment::none()) + }; + + Self { + buffer, + offset: bit_offset, + len, + } } /// Slice any full bytes from the buffer, leaving the offset < 8. @@ -675,6 +709,19 @@ mod tests { assert_eq!(sliced.offset(), 2); } + #[test] + fn test_byte_aligned_bytes() { + let bytes: ByteBuffer = buffer![0b1010_0101u8, 0b0000_0011]; + let buf = BitBuffer::new(bytes.clone(), 10); + assert_eq!(buf.byte_aligned_bytes(), Some(bytes.as_slice())); + + let byte_sliced = buf.slice(8..10); + assert_eq!(byte_sliced.byte_aligned_bytes(), Some(&[0b0000_0011][..])); + + let bit_sliced = buf.slice(1..9); + assert!(bit_sliced.byte_aligned_bytes().is_none()); + } + #[test] fn test_from_indices_dense_crosses_words() { let len = 130; diff --git a/vortex-buffer/src/bit/buf_mut.rs b/vortex-buffer/src/bit/buf_mut.rs index 02e4ce66755..df38b0cd434 100644 --- a/vortex-buffer/src/bit/buf_mut.rs +++ b/vortex-buffer/src/bit/buf_mut.rs @@ -267,7 +267,9 @@ impl BitBufferMut { /// Clears the bit buffer (but keeps any allocated memory). pub fn clear(&mut self) { - // Since there are no items we need to drop, we simply set the length to 0. + // 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. + self.buffer.clear(); self.len = 0; self.offset = 0; } @@ -649,6 +651,8 @@ impl FromIterator for BitBufferMut { #[cfg(test)] mod tests { + use rstest::rstest; + use crate::BufferMut; use crate::bit::buf_mut::BitBufferMut; use crate::bitbuffer; @@ -923,6 +927,65 @@ mod tests { assert!(frozen.value(7)); } + #[cfg_attr(miri, ignore)] // bitvec crate uses a ptr cast that Miri doesn't support + #[test] + fn append_after_clear_reads_back_false() { + // `clear` must not leave stale set bits behind: `append_false` and `append_buffer` + // rely on bits beyond `len` being zero. + let mut bools = BitBufferMut::new_set(16); + bools.clear(); + bools.append_false(); + bools.append_buffer(&crate::BitBuffer::new_unset(8)); + + let bools = bools.freeze(); + assert_eq!(bools.len(), 9); + assert_eq!(bools.true_count(), 0); + } + + #[cfg_attr(miri, ignore)] // bitvec crate uses a ptr cast that Miri doesn't support + #[test] + fn test_append_buffer_after_truncate() { + // Truncating leaves stale set bits in the last partial byte; an append after that + // must overwrite them rather than OR into them. + let mut buf = BitBufferMut::new_set(16); + buf.truncate(3); + buf.append_buffer(&crate::BitBuffer::new_unset(8)); + + let frozen = buf.freeze(); + assert_eq!(frozen.len(), 11); + for i in 0..3 { + assert!(frozen.value(i), "bit {i} should be set"); + } + for i in 3..11 { + assert!(!frozen.value(i), "bit {i} should be unset"); + } + } + + #[rstest] + #[case::both_aligned(0, 0)] + #[case::dst_unaligned(3, 0)] + #[case::src_unaligned(0, 5)] + #[case::mismatched(3, 5)] + #[case::equal_nonzero(5, 5)] + #[cfg_attr(miri, ignore)] // bitvec crate uses a ptr cast that Miri doesn't support + fn test_append_buffer_long(#[case] dst_prefix: usize, #[case] src_start: usize) { + // Exercise every alignment combination across many words. + let source = crate::BitBuffer::from_iter((0..301).map(|i| i % 3 == 0)); + let source = source.slice(src_start..301); + + let mut dest = BitBufferMut::with_capacity(512); + dest.append_n(true, dst_prefix); + dest.append_buffer(&source); + + assert_eq!(dest.len(), dst_prefix + source.len()); + for i in 0..dst_prefix { + assert!(dest.value(i), "prefix bit {i}"); + } + for i in 0..source.len() { + assert_eq!(dest.value(dst_prefix + i), source.value(i), "bit {i}"); + } + } + #[cfg_attr(miri, ignore)] // bitvec crate uses a ptr cast that Miri doesn't support #[test] fn test_append_buffer_with_offsets() { diff --git a/vortex-buffer/src/bit/mod.rs b/vortex-buffer/src/bit/mod.rs index 41bb5797266..a39c93e3f25 100644 --- a/vortex-buffer/src/bit/mod.rs +++ b/vortex-buffer/src/bit/mod.rs @@ -74,6 +74,19 @@ where } } +/// Read up to 8 bytes as a little-endian `u64`, zero-padding the high bytes when fewer than 8 +/// bytes are supplied. +/// +/// This preserves Vortex's least-significant-bit-first bitmap numbering on little- and big-endian +/// targets. For a full 8-byte slice it lowers to a single word load. +#[inline] +pub fn read_u64_le(bytes: &[u8]) -> u64 { + debug_assert!(bytes.len() <= 8); + let mut buf = [0u8; 8]; + buf[..bytes.len()].copy_from_slice(bytes); + u64::from_le_bytes(buf) +} + /// Splice a packed word `w` (whose bits above the highest valid bit are zero) into /// `words` at the given bit position. /// @@ -178,6 +191,7 @@ pub unsafe fn unset_bit_unchecked(buf: *mut u8, index: usize) { mod tests { use super::collect_bool_word; use super::pack_bools_into_words; + use super::read_u64_le; #[test] fn collect_bool_word_packs_lsb_first() { @@ -190,6 +204,12 @@ mod tests { assert_eq!(collect_bool_word(0, |_| true), 0); } + #[test] + fn read_u64_le_zero_pads_tail() { + assert_eq!(read_u64_le(&[0x34, 0x12]), 0x1234); + assert_eq!(read_u64_le(&[0xff; 8]), u64::MAX); + } + #[test] #[should_panic(expected = "cannot pack 65 bits into a u64 word")] fn collect_bool_word_rejects_too_many_bits() { diff --git a/vortex-buffer/src/bit/ops.rs b/vortex-buffer/src/bit/ops.rs index be1401d912a..ec922b54bc2 100644 --- a/vortex-buffer/src/bit/ops.rs +++ b/vortex-buffer/src/bit/ops.rs @@ -5,20 +5,9 @@ use std::mem::MaybeUninit; use crate::BitBuffer; use crate::BitBufferMut; -use crate::Buffer; +use crate::BufferMut; use crate::ByteBufferMut; -use crate::trusted_len::TrustedLenExt; - -/// Read up to 8 bytes as a little-endian `u64`, zero-padding the high bytes when fewer than 8 are -/// supplied. Using [`u64::from_le_bytes`] keeps the bit-numbering identical on little- and -/// big-endian targets; for a full 8-byte slice it lowers to a single word load. -#[inline] -fn read_u64_le(bytes: &[u8]) -> u64 { - debug_assert!(bytes.len() <= 8); - let mut buf = [0u8; 8]; - buf[..bytes.len()].copy_from_slice(bytes); - u64::from_le_bytes(buf) -} +use crate::read_u64_le; trait BitWordTarget { fn byte_len(&self) -> usize; @@ -188,36 +177,52 @@ pub(super) fn bitwise_binary_op u64>( mut op: F, ) -> BitBuffer { assert_eq!(left.len(), right.len()); - - // If the buffers are aligned, we can use the fast path. - if left.offset().is_multiple_of(8) && right.offset().is_multiple_of(8) { - let left_chunks = left.unaligned_chunks(); - let right_chunks = right.unaligned_chunks(); - if left_chunks.lead_padding() == 0 - && left_chunks.trailing_padding() == 0 - && right_chunks.lead_padding() == 0 - && right_chunks.trailing_padding() == 0 - { - let iter = left_chunks - .iter() - .zip(right_chunks.iter()) - .map(|(l, r)| op(l, r)); - let iter = unsafe { iter.trusted_len() }; - let result = Buffer::::from_trusted_len_iter(iter).into_byte_buffer(); - return BitBuffer::new(result, left.len()); - } + let len = left.len(); + if len == 0 { + return BitBuffer::empty(); } - let iter = left - .chunks() - .iter_padded() - .zip(right.chunks().iter_padded()) - .map(|(l, r)| op(l, r)); - let iter = unsafe { iter.trusted_len() }; + let n_bytes = len.div_ceil(8); + let out = if left.offset().is_multiple_of(8) && right.offset().is_multiple_of(8) { + // Byte-aligned operands: logical bits map onto physical `u64` words, so read the backing + // bytes straight as words and build the result from a `TrustedLen` iterator. + let l_start = left.offset() / 8; + let r_start = right.offset() / 8; + let lhs = &left.inner().as_slice()[l_start..l_start + n_bytes]; + let rhs = &right.inner().as_slice()[r_start..r_start + n_bytes]; - let result = Buffer::::from_trusted_len_iter(iter).into_byte_buffer(); + let (lhs_words, lhs_tail) = lhs.as_chunks::<8>(); + let (rhs_words, rhs_tail) = rhs.as_chunks::<8>(); + + let mut out = BufferMut::::from_trusted_len_iter( + lhs_words + .iter() + .zip(rhs_words) + .map(|(l, r)| op(u64::from_le_bytes(*l), u64::from_le_bytes(*r))), + ); + if !lhs_tail.is_empty() { + out.push(op(read_u64_le(lhs_tail), read_u64_le(rhs_tail))); + } + out + } else { + // Sub-byte offset: `iter_padded` realigns the bits and appends one pad word, so take + // exactly `ceil(len / 64)` words. + let n_words = len.div_ceil(64); + let mut out = BufferMut::::with_capacity(n_words); + for (l, r) in left + .chunks() + .iter_padded() + .zip(right.chunks().iter_padded()) + .take(n_words) + { + out.push(op(l, r)); + } + out + }; - BitBuffer::new(result, left.len()) + let mut bytes = out.into_byte_buffer(); + bytes.truncate(n_bytes); + BitBuffer::new(bytes.freeze(), len) } #[cfg(test)] @@ -317,6 +322,41 @@ mod tests { assert_eq!(result, bitbuffer![false, true, true, false]); } + /// `bitwise_binary_op` must match a naive per-bit reference for every op, offset and length, + /// independent of the chunked kernels. + #[rstest] + #[case::aligned(0, 0)] + #[case::byte_aligned(8, 16)] + #[case::byte_aligned_mismatch(16, 0)] + #[case::sub_byte(3, 3)] + #[case::sub_byte_mismatch(0, 5)] + fn binary_op_matches_naive(#[case] left_offset: usize, #[case] right_offset: usize) { + #[allow(clippy::cast_possible_truncation)] + let make = |offset: usize, len: usize, salt: u8| -> BitBuffer { + let bytes: ByteBufferMut = (0..(offset + len).div_ceil(8).max(1)) + .map(|i| (i as u8).wrapping_mul(31).wrapping_add(salt)) + .collect(); + BitBufferMut::from_buffer(bytes, offset, len).freeze() + }; + let ops: [fn(u64, u64) -> u64; 4] = + [|a, b| a & b, |a, b| a | b, |a, b| a ^ b, |a, b| a & !b]; + + for len in [1usize, 5, 8, 63, 64, 65, 127, 128, 200, 256] { + let left = make(left_offset, len, 0xC3); + let right = make(right_offset, len, 0x5A); + for op in ops { + let got = bitwise_binary_op(&left, &right, op); + let expected: BitBuffer = (0..len) + .map(|i| op(u64::from(left.value(i)), u64::from(right.value(i))) & 1 == 1) + .collect(); + assert_eq!( + got, expected, + "loff={left_offset} roff={right_offset} len={len}" + ); + } + } + } + /// Regression test for a bug where [`bitwise_unary_op`] produced corrupt results when /// the [`BitBuffer`]'s underlying byte pointer was not u64-aligned. Slicing a buffer by /// a non-multiple-of-8 number of bytes can cause this misalignment. The bug only diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 3b778358577..4b587b94629 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -32,10 +32,20 @@ pub struct Buffer { pub(crate) _marker: PhantomData, } +/// 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) } +}; + impl Default for Buffer { fn default() -> Self { Self { - bytes: Default::default(), + bytes: Bytes::from_static(EMPTY_BACKING), length: 0, alignment: Alignment::of::(), _marker: PhantomData, @@ -85,8 +95,25 @@ impl Buffer { } /// 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 + /// `alignment`. Use [`copy_from_preferred_aligned`] to control the over-alignment. + /// + /// [`copy_from_preferred_aligned`]: Self::copy_from_preferred_aligned pub fn copy_from_aligned(values: impl AsRef<[T]>, alignment: Alignment) -> Self { - BufferMut::copy_from_aligned(values, alignment).freeze() + Self::copy_from_preferred_aligned(values, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + } + + /// Returns a new `Buffer` copied from the provided slice and with the requested alignment. + /// + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. + pub fn copy_from_preferred_aligned( + values: impl AsRef<[T]>, + alignment: Alignment, + preferred_alignment: Option, + ) -> Self { + BufferMut::copy_from_preferred_aligned(values, alignment, preferred_alignment).freeze() } /// Create a new zeroed `Buffer` with the given value. @@ -94,19 +121,51 @@ impl Buffer { Self::zeroed_aligned(len, Alignment::of::()) } - /// Create a new zeroed `Buffer` with the given value. + /// Create a new zeroed `Buffer` with the requested alignment. + /// + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`zeroed_preferred_aligned`] to control the over-alignment. + /// + /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self { - BufferMut::zeroed_aligned(len, alignment).freeze() + Self::zeroed_preferred_aligned(len, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + } + + /// Create a new zeroed `Buffer` with the requested alignment. + /// + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. + pub fn zeroed_preferred_aligned( + len: usize, + alignment: Alignment, + preferred_alignment: Option, + ) -> Self { + BufferMut::zeroed_preferred_aligned(len, alignment, preferred_alignment).freeze() } /// Create a new empty `ByteBuffer` with the provided alignment. pub fn empty() -> Self { - BufferMut::empty().freeze() + Self::empty_aligned(Alignment::of::()) } /// 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`]. pub fn empty_aligned(alignment: Alignment) -> Self { - BufferMut::empty_aligned(alignment).freeze() + if !alignment.is_aligned_to(Alignment::of::()) { + vortex_panic!( + "Alignment {} must align to the scalar type's alignment {}", + alignment, + Alignment::of::(), + ); + } + Self { + bytes: Bytes::from_static(EMPTY_BACKING), + length: 0, + alignment, + _marker: PhantomData, + } } /// Create a new full `ByteBuffer` with the given value. @@ -152,7 +211,7 @@ impl Buffer { Alignment::of::(), ); } - if bytes.as_ptr().align_offset(*alignment) != 0 { + if !alignment.is_ptr_aligned(bytes.as_ptr()) { vortex_panic!( "Bytes alignment must align to the requested alignment {}", alignment, @@ -320,7 +379,7 @@ impl Buffer { let begin_byte = begin * size_of::(); let end_byte = end * size_of::(); - if !begin_byte.is_multiple_of(*alignment) { + if !alignment.is_offset_aligned(begin_byte) { vortex_panic!( "range start must be aligned to {alignment:?}, byte {}", begin_byte @@ -369,7 +428,7 @@ impl Buffer { vortex_panic!("slice_ref subset alignment must at least align to the buffer alignment") } - if subset.as_ptr().align_offset(*alignment) != 0 { + if !alignment.is_ptr_aligned(subset.as_ptr()) { vortex_panic!("slice_ref subset must be aligned to {:?}", alignment); } @@ -435,17 +494,17 @@ impl Buffer { /// 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(&buffer)) + .unwrap_or_else(|buffer| BufferMut::::copy_from_aligned(&buffer, buffer.alignment)) } /// Returns whether a `Buffer` is aligned to the given alignment. pub fn is_aligned(&self, alignment: Alignment) -> bool { - self.bytes.as_ptr().align_offset(*alignment) == 0 + alignment.is_ptr_aligned(self.bytes.as_ptr()) } /// Return a `Buffer` with the given alignment. Where possible, this will be zero-copy. pub fn aligned(mut self, alignment: Alignment) -> Self { - if self.as_ptr().align_offset(*alignment) == 0 { + if alignment.is_ptr_aligned(self.as_ptr()) { self.alignment = alignment; self } else { @@ -462,7 +521,7 @@ impl Buffer { /// Return a `Buffer` with the given alignment. Panics if the buffer is not aligned. pub fn ensure_aligned(mut self, alignment: Alignment) -> Self { - if self.as_ptr().align_offset(*alignment) == 0 { + if alignment.is_ptr_aligned(self.as_ptr()) { self.alignment = alignment; self } else { @@ -634,7 +693,7 @@ impl Buf for ByteBuffer { #[inline] fn advance(&mut self, cnt: usize) { - if !cnt.is_multiple_of(*self.alignment) { + if !self.alignment.is_offset_aligned(cnt) { vortex_panic!( "Cannot advance buffer by {} items, resulting alignment is not {}", cnt, @@ -778,6 +837,29 @@ mod test { assert_eq!(buf.as_slice(), &[0; LEN]); } + #[test] + fn copy_from_over_aligns_to_default() { + let values = [1u32, 2, 3]; + let buf = Buffer::::copy_from(values); + + // The buffer reports the scalar type's alignment, ... + assert_eq!(buf.alignment(), Alignment::of::()); + // ... but the underlying allocation is over-aligned to DEFAULT_ALIGNMENT. + assert!(buf.is_aligned(Alignment::DEFAULT_ALIGNMENT)); + assert_eq!(buf.as_slice(), &values); + } + + #[test] + fn zeroed_over_aligns_to_default() { + const LEN: usize = 17; + + let buf = Buffer::::zeroed(LEN); + + assert_eq!(buf.alignment(), Alignment::of::()); + assert!(buf.is_aligned(Alignment::DEFAULT_ALIGNMENT)); + assert_eq!(buf.as_slice(), &[0; LEN]); + } + #[test] fn from_vec() { let vec = vec![1, 2, 3, 4, 5]; @@ -786,6 +868,31 @@ mod test { assert_eq!(vec, buff.as_ref()); } + #[test] + fn empty_aligned_max_alignment() { + // Empty buffers are backed by a static and must satisfy any valid alignment. + let buf = Buffer::::empty_aligned(Alignment::MAX); + assert!(buf.is_empty()); + assert!(buf.is_aligned(Alignment::MAX)); + } + + #[test] + fn empty_slice_preserves_alignment() { + let buf = Buffer::::zeroed_aligned(8, Alignment::new(64)); + let sliced = buf.slice(0..0); + assert!(sliced.is_empty()); + assert_eq!(sliced.alignment(), Alignment::new(64)); + assert!(sliced.is_aligned(Alignment::new(64))); + } + + #[test] + fn empty_into_mut_preserves_alignment() { + let buf = Buffer::::empty_aligned(Alignment::new(64)); + let buf_mut = buf.into_mut(); + assert_eq!(buf_mut.alignment(), Alignment::new(64)); + assert!(buf_mut.is_empty()); + } + #[test] fn test_slice_unaligned_end_pos() { let data = vec![0u8; 2]; @@ -797,4 +904,12 @@ mod test { // to be aligned. aligned_buffer.slice(0..1); } + + #[test] + fn test_empty_equality() { + let a = Buffer::::empty(); + let b = Buffer::::empty(); + + assert_eq!(a, b); + } } diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index ceb8732d701..a7edfb00293 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -3,6 +3,7 @@ use core::mem::MaybeUninit; use std::any::type_name; +use std::cmp::max; use std::fmt::Debug; use std::fmt::Formatter; use std::io::Write; @@ -38,7 +39,33 @@ impl BufferMut { } /// Create a new `BufferMut` with the requested alignment and capacity. + /// + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`with_capacity_preferred_aligned`] to control the over-alignment. + /// + /// [`with_capacity_preferred_aligned`]: Self::with_capacity_preferred_aligned pub fn with_capacity_aligned(capacity: usize, alignment: Alignment) -> Self { + Self::with_capacity_preferred_aligned( + capacity, + alignment, + Some(Alignment::DEFAULT_ALIGNMENT), + ) + } + + /// Create a new `BufferMut` with the requested alignment and capacity. + /// + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. + pub fn with_capacity_preferred_aligned( + capacity: usize, + alignment: Alignment, + preferred_alignment: Option, + ) -> Self { + let actual = max( + alignment, + preferred_alignment.unwrap_or(Alignment::of::()), + ); + if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!( "Alignment {} must align to the scalar type's alignment {}", @@ -47,8 +74,8 @@ impl BufferMut { ); } - let mut bytes = BytesMut::with_capacity((capacity * size_of::()) + *alignment); - bytes.align_empty(alignment); + let mut bytes = BytesMut::with_capacity((capacity * size_of::()) + *actual); + bytes.align_empty(actual); Self { bytes, @@ -63,10 +90,29 @@ impl BufferMut { Self::zeroed_aligned(len, Alignment::of::()) } - /// Create a new zeroed `BufferMut`. + /// Create a new zeroed `BufferMut` with the requested alignment. + /// + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`zeroed_preferred_aligned`] to control the over-alignment. + /// + /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self { - let mut bytes = BytesMut::zeroed((len * size_of::()) + *alignment); - bytes.advance(bytes.as_ptr().align_offset(*alignment)); + Self::zeroed_preferred_aligned(len, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + } + + /// Create a new zeroed `BufferMut` with the requested alignment. + /// + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. + pub fn zeroed_preferred_aligned( + len: usize, + alignment: Alignment, + preferred_alignment: Option, + ) -> Self { + 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); Self { @@ -83,8 +129,24 @@ impl BufferMut { } /// Create a new empty `BufferMut` with the provided alignment. + /// + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`empty_preferred_aligned`] to control the over-alignment. + /// + /// [`empty_preferred_aligned`]: Self::empty_preferred_aligned pub fn empty_aligned(alignment: Alignment) -> Self { - BufferMut::with_capacity_aligned(0, alignment) + Self::empty_preferred_aligned(alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + } + + /// Create a new empty `BufferMut` with the provided alignment. + /// + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. + pub fn empty_preferred_aligned( + alignment: Alignment, + preferred_alignment: Option, + ) -> Self { + BufferMut::with_capacity_preferred_aligned(0, alignment, preferred_alignment) } /// Create a new full `BufferMut` with the given value. @@ -104,15 +166,37 @@ impl BufferMut { /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. /// + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`copy_from_preferred_aligned`] to control the over-alignment. + /// + /// [`copy_from_preferred_aligned`]: Self::copy_from_preferred_aligned + /// /// ## Panics /// /// 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)) + } + + /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. + /// + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. + /// + /// ## Panics + /// + /// Panics when the requested alignment isn't itself aligned to type T. + pub fn copy_from_preferred_aligned( + other: impl AsRef<[T]>, + alignment: Alignment, + preferred_alignment: Option, + ) -> 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_aligned(other.len(), alignment); + let mut buffer = + Self::with_capacity_preferred_aligned(other.len(), alignment, preferred_alignment); buffer.extend_from_slice(other); debug_assert_eq!(buffer.alignment(), alignment); buffer @@ -214,6 +298,10 @@ impl BufferMut { /// reading from a file) before marking the data as initialized using the /// [`set_len`] method. /// + /// Note that the returned slice may be larger than the capacity requested at + /// construction, since the underlying allocation can be rounded up (e.g. to + /// satisfy alignment requirements). + /// /// [`set_len`]: BufferMut::set_len /// [`Vec::spare_capacity_mut`]: Vec::spare_capacity_mut /// @@ -358,7 +446,7 @@ impl BufferMut { } let bytes_at = at * size_of::(); - if !bytes_at.is_multiple_of(*self.alignment) { + if !self.alignment.is_offset_aligned(bytes_at) { vortex_panic!( "Cannot split buffer at {}, resulting alignment is not {}", at, @@ -742,7 +830,7 @@ impl Buf for ByteBufferMut { } fn advance(&mut self, cnt: usize) { - if !cnt.is_multiple_of(*self.alignment) { + if !self.alignment.is_offset_aligned(cnt) { vortex_panic!( "Cannot advance buffer by {} items, resulting alignment is not {}", cnt, @@ -765,7 +853,7 @@ unsafe impl BufMut for ByteBufferMut { #[inline] unsafe fn advance_mut(&mut self, cnt: usize) { - if !cnt.is_multiple_of(*self.alignment) { + if !self.alignment.is_offset_aligned(cnt) { vortex_panic!( "Cannot advance buffer by {} items, resulting alignment is not {}", cnt, diff --git a/vortex-buffer/src/string.rs b/vortex-buffer/src/string.rs index bcaf8d7c3d2..ca5ae5251d5 100644 --- a/vortex-buffer/src/string.rs +++ b/vortex-buffer/src/string.rs @@ -25,7 +25,7 @@ impl BufferString { /// Creates an empty `BufferString`. pub fn empty() -> Self { - Self(ByteBuffer::from(vec![])) + Self(ByteBuffer::empty()) } /// Return a view of the contents of BufferString as an immutable `&str`. diff --git a/vortex-buffer/src/trusted_len.rs b/vortex-buffer/src/trusted_len.rs index 13cf25d0546..7ef6f682be4 100644 --- a/vortex-buffer/src/trusted_len.rs +++ b/vortex-buffer/src/trusted_len.rs @@ -1,8 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use itertools::ProcessResults; - /// Trait for all types which have a known upper-bound. /// /// Functions that receive a `TrustedLen` iterator can assume that it's `size_hint` is exact, @@ -146,7 +144,7 @@ unsafe impl TrustedLen for crate::Iter<'_, T> {} unsafe impl TrustedLen for crate::BufferIterator {} // ProcessResults -unsafe impl<'a, I, T: 'a, E: 'a> TrustedLen for ProcessResults<'a, I, E> where +unsafe impl<'a, I, T: 'a, E: 'a> TrustedLen for itertools::ProcessResults<'a, I, E> where I: TrustedLen> { } @@ -158,7 +156,7 @@ unsafe impl TrustedLen for std::iter::Enumerate where I: TrustedLen TrustedLen for std::iter::Zip where T: TrustedLen, - U: TrustedLen, + U: Iterator, { } diff --git a/vortex-compressor/benches/dict_encode.rs b/vortex-compressor/benches/dict_encode.rs index 68b7cd98013..2c4e24108a7 100644 --- a/vortex-compressor/benches/dict_encode.rs +++ b/vortex-compressor/benches/dict_encode.rs @@ -11,15 +11,11 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::builders::dict::dict_encode; -use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; -use vortex_compressor::builtins::integer_dictionary_encode; -use vortex_compressor::stats::IntegerStats; use vortex_session::VortexSession; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn make_array() -> PrimitiveArray { let values: BufferMut = (0..50).cycle().take(64_000).collect(); @@ -44,16 +40,6 @@ fn encode_generic(bencher: Bencher) { .bench_refs(|(array, ctx)| dict_encode(array, ctx).unwrap()); } -#[cfg(not(codspeed))] -#[divan::bench] -fn encode_specialized(bencher: Bencher) { - let array = make_array(); - let stats = IntegerStats::generate(&array, &mut SESSION.create_execution_ctx()); - bencher - .with_inputs(|| &stats) - .bench_refs(|stats| integer_dictionary_encode(array.as_view(), stats)); -} - fn main() { divan::main() } diff --git a/vortex-compressor/benches/stats_calc.rs b/vortex-compressor/benches/stats_calc.rs index 69a8f970608..abbc34817f1 100644 --- a/vortex-compressor/benches/stats_calc.rs +++ b/vortex-compressor/benches/stats_calc.rs @@ -9,7 +9,6 @@ mod benchmarks { use divan::Bencher; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -17,8 +16,7 @@ mod benchmarks { use vortex_compressor::stats::IntegerStats; use vortex_session::VortexSession; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn generate_dataset(max_run: u32, distinct: u32) -> Buffer { let mut output = BufferMut::with_capacity(64_000); diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index bf591963db3..5288637a609 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -253,11 +253,9 @@ mod tests { use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::assert_arrays_eq; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; - use vortex_session::VortexSession; use super::dictionary_encode; use crate::stats::FloatStats; @@ -265,9 +263,7 @@ mod tests { #[test] fn test_float_dict_encode() -> VortexResult<()> { - let mut ctx = VortexSession::empty() - .with::() - .create_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let values = buffer![1f32, 2f32, 2f32, 0f32, 1f32]; let validity = Validity::Array(BoolArray::from_iter([true, true, true, false, true]).into_array()); @@ -294,7 +290,7 @@ mod tests { .clone() .execute::(&mut ctx)? .into_array(); - assert_arrays_eq!(undict, expected); + assert_arrays_eq!(undict, expected, &mut ctx); Ok(()) } } diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 82af7f14de7..8ec3bf53345 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -260,20 +260,16 @@ mod tests { use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::assert_arrays_eq; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; - use vortex_session::VortexSession; use super::dictionary_encode; use crate::stats::IntegerStats; #[test] fn test_dict_encode_integer_stats() -> VortexResult<()> { - let mut ctx = VortexSession::empty() - .with::() - .create_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let data = buffer![100i32, 200, 100, 0, 100]; let validity = Validity::Array(BoolArray::from_iter([true, true, true, false, true]).into_array()); @@ -300,7 +296,7 @@ mod tests { .clone() .execute::(&mut ctx)? .into_array(); - assert_arrays_eq!(undict, expected); + assert_arrays_eq!(undict, expected, &mut ctx); Ok(()) } } diff --git a/vortex-compressor/src/compressor.rs b/vortex-compressor/src/compressor.rs index 33502104b13..965c719bf4b 100644 --- a/vortex-compressor/src/compressor.rs +++ b/vortex-compressor/src/compressor.rs @@ -75,6 +75,9 @@ mod root_list_children { /// /// No scheme may appear twice in a cascade chain. The compressor enforces this automatically /// along with push/pull exclusion rules declared by each scheme. +/// +/// Downstream crates usually wrap this type with a preconfigured scheme set. Use it directly when +/// embedding a custom fixed scheme list or testing scheme interactions. #[derive(Debug, Clone)] pub struct CascadingCompressor { /// The enabled compression schemes. @@ -580,7 +583,9 @@ impl CascadingCompressor { }) .collect::>()?; - array.clone().with_slots(slots) + // SAFETY: compression rewrites each child slot to an equivalent physical representation, + // preserving the parent array's logical values and statistics. + unsafe { array.clone().with_slots(slots) } } } @@ -596,7 +601,6 @@ mod tests { use vortex_array::arrays::Constant; use vortex_array::arrays::NullArray; use vortex_array::arrays::PrimitiveArray; - use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_session::VortexSession; @@ -613,8 +617,7 @@ mod tests { use crate::estimate::WinnerEstimate; use crate::scheme::SchemeExt; - static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn compressor() -> CascadingCompressor { CascadingCompressor::new(vec![&IntDictScheme, &FloatDictScheme, &StringDictScheme]) diff --git a/vortex-compressor/src/lib.rs b/vortex-compressor/src/lib.rs index 1ecc9e4d5b4..7e6854eacbe 100644 --- a/vortex-compressor/src/lib.rs +++ b/vortex-compressor/src/lib.rs @@ -16,6 +16,31 @@ //! This crate contains no encoding dependencies. Batteries-included compressors are provided by //! downstream crates like `vortex-btrblocks`, which register different encodings to the compressor. //! +//! # Example +//! +//! A [`CascadingCompressor`] can be created directly with a fixed scheme list. With no schemes it +//! still canonicalizes supported inputs and recursively handles nested structure, but no leaf +//! compression is selected. +//! +//! ```rust +//! use vortex_array::{IntoArray, VortexSessionExecute, array_session}; +//! use vortex_array::arrays::PrimitiveArray; +//! use vortex_array::validity::Validity; +//! use vortex_buffer::buffer; +//! use vortex_compressor::CascadingCompressor; +//! +//! # fn example() -> vortex_error::VortexResult<()> { +//! let session = array_session(); +//! let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable).into_array(); +//! let compressor = CascadingCompressor::new(Vec::new()); +//! +//! let result = compressor.compress(&array, &mut session.create_execution_ctx())?; +//! assert_eq!(result.dtype(), array.dtype()); +//! assert_eq!(result.len(), array.len()); +//! # Ok(()) +//! # } +//! ``` +//! //! # Observability //! //! The compressor emits a small set of `tracing` spans and events on a single target so you can diff --git a/vortex-compressor/src/sample.rs b/vortex-compressor/src/sample.rs index 4821c7c6d95..deb5bbe6f37 100644 --- a/vortex-compressor/src/sample.rs +++ b/vortex-compressor/src/sample.rs @@ -133,6 +133,8 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> #[cfg(test)] mod tests { use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::validity::Validity; @@ -143,6 +145,7 @@ mod tests { #[test] fn sample_is_deterministic() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Create a deterministic array with linear-with-noise pattern let values: Vec = (0i64..100_000).map(|i| i + (i * 7 + 3) % 11).collect(); @@ -153,7 +156,7 @@ mod tests { for _ in 0..10 { let again = sample(&array, SAMPLE_SIZE, SAMPLE_COUNT); assert_eq!(first.nbytes(), again.nbytes()); - assert_arrays_eq!(&first, &again); + assert_arrays_eq!(&first, &again, &mut ctx); } Ok(()) } diff --git a/vortex-compressor/src/scheme.rs b/vortex-compressor/src/scheme.rs index 4c5e7f0e940..6e6899a3ea1 100644 --- a/vortex-compressor/src/scheme.rs +++ b/vortex-compressor/src/scheme.rs @@ -139,6 +139,10 @@ pub struct AncestorExclusion { /// to declare what they require. The compressor merges all eligible schemes' options before /// generating stats, so each stat is always computed at most once for a given array. /// +/// A scheme implementation should be deterministic for a fixed input array and context. The +/// compressor uses scheme order for deterministic tie-breaking, so non-deterministic estimates make +/// compressed output harder to reproduce and compare. +/// /// [`scheme_name`]: Scheme::scheme_name /// [`matches`]: Scheme::matches /// [`compress`]: Scheme::compress diff --git a/vortex-compressor/src/stats/bool.rs b/vortex-compressor/src/stats/bool.rs index 8825ec8a7f6..fbbc4001b90 100644 --- a/vortex-compressor/src/stats/bool.rs +++ b/vortex-compressor/src/stats/bool.rs @@ -92,8 +92,8 @@ impl BoolStats { #[cfg(test)] mod tests { - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; @@ -103,7 +103,7 @@ mod tests { #[test] fn test_all_true() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true, true, true]), Validity::NonNullable, @@ -118,7 +118,7 @@ mod tests { #[test] fn test_all_false() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![false, false, false]), Validity::NonNullable, @@ -133,7 +133,7 @@ mod tests { #[test] fn test_mixed() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true, false, true]), Validity::NonNullable, @@ -148,7 +148,7 @@ mod tests { #[test] fn test_with_nulls() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true, false, true]), Validity::from_iter([true, false, true]), diff --git a/vortex-compressor/src/stats/float.rs b/vortex-compressor/src/stats/float.rs index d968e8d368f..c505993e9e5 100644 --- a/vortex-compressor/src/stats/float.rs +++ b/vortex-compressor/src/stats/float.rs @@ -276,24 +276,25 @@ where #[cfg(test)] mod tests { use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; 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_error::VortexResult; use super::FloatStats; + use crate::stats::GenerateStatsOptions; #[test] fn test_float_stats() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let floats = buffer![0.0f32, 1.0f32, 2.0f32].into_array(); let floats = floats.execute::(&mut ctx)?; let stats = FloatStats::generate_opts( &floats, - crate::stats::GenerateStatsOptions { + GenerateStatsOptions { count_distinct_values: true, }, &mut ctx, @@ -308,7 +309,7 @@ mod tests { #[test] fn test_float_stats_leading_nulls() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let floats = PrimitiveArray::new( buffer![0.0f32, 1.0f32, 2.0f32], Validity::from_iter([false, true, true]), @@ -316,7 +317,7 @@ mod tests { let stats = FloatStats::generate_opts( &floats, - crate::stats::GenerateStatsOptions { + GenerateStatsOptions { count_distinct_values: true, }, &mut ctx, diff --git a/vortex-compressor/src/stats/integer.rs b/vortex-compressor/src/stats/integer.rs index 64345ac5f06..a399c1060f9 100644 --- a/vortex-compressor/src/stats/integer.rs +++ b/vortex-compressor/src/stats/integer.rs @@ -561,8 +561,8 @@ fn inner_loop_naive( mod tests { use std::iter; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; @@ -575,7 +575,7 @@ mod tests { #[test] fn test_naive_count_distinct_values() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::new(buffer![217u8, 0], Validity::NonNullable); let stats = typed_int_stats::(&array, true, &mut ctx)?; assert_eq!(stats.distinct_count().unwrap(), 2); @@ -584,7 +584,7 @@ mod tests { #[test] fn test_naive_count_distinct_values_nullable() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::new( buffer![217u8, 0], Validity::from(BitBuffer::from(vec![true, false])), @@ -596,7 +596,7 @@ mod tests { #[test] fn test_count_distinct_values() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::new((0..128u8).collect::>(), Validity::NonNullable); let stats = typed_int_stats::(&array, true, &mut ctx)?; assert_eq!(stats.distinct_count().unwrap(), 128); @@ -605,7 +605,7 @@ mod tests { #[test] fn test_count_distinct_values_nullable() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::new( (0..128u8).collect::>(), Validity::from(BitBuffer::from_iter( @@ -619,7 +619,7 @@ mod tests { #[test] fn test_integer_stats_leading_nulls() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let ints = PrimitiveArray::new(buffer![0, 1, 2], Validity::from_iter([false, true, true])); let stats = IntegerStats::generate_opts( diff --git a/vortex-compute/Cargo.toml b/vortex-compute/Cargo.toml new file mode 100644 index 00000000000..f223d455395 --- /dev/null +++ b/vortex-compute/Cargo.toml @@ -0,0 +1,37 @@ +[package] +name = "vortex-compute" +authors = { workspace = true } +categories = { workspace = true } +description = "Lane-level compute kernels for Vortex buffers" +edition = { workspace = true } +homepage = { workspace = true } +include = { workspace = true } +keywords = { workspace = true } +license = { workspace = true } +readme = { workspace = true } +repository = { workspace = true } +rust-version = { workspace = true } +version = { workspace = true } + +[package.metadata.docs.rs] +all-features = true + +[dependencies] +vortex-buffer = { workspace = true } + +[dev-dependencies] +arrow-arith = { workspace = true } +arrow-array = { workspace = true } +arrow-buffer = { workspace = true } +arrow-cast = { workspace = true } +arrow-schema = { workspace = true } +divan = { workspace = true } +num-traits = { workspace = true } +rand = { workspace = true } + +[lints] +workspace = true + +[[bench]] +name = "lane_kernels" +harness = false diff --git a/vortex-compute/benches/lane_kernels.rs b/vortex-compute/benches/lane_kernels.rs new file mode 100644 index 00000000000..af8e27d6199 --- /dev/null +++ b/vortex-compute/benches/lane_kernels.rs @@ -0,0 +1,401 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Coverage benchmark for the lane-kernel variants used by primitive casts, +//! bit-packing paths, and `LaneZip` binary kernels. +//! +//! `add_checked` parity assertions (run at startup) verify that the bit-packed +//! fail-tracking scheme: +//! - propagates valid-lane overflow as `Err`, and +//! - suppresses null-lane overflow without the closure ever inspecting `valid`. +//! +//! 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. + +#![expect(clippy::unwrap_used)] +#![expect(clippy::clone_on_ref_ptr)] + +use std::mem::MaybeUninit; +use std::sync::Arc; + +use arrow_arith::numeric::add; +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::Int32Array; +use arrow_array::UInt16Array; +use arrow_array::UInt32Array; +use arrow_array::UInt64Array; +use arrow_buffer::NullBuffer; +use arrow_cast::CastOptions; +use arrow_cast::cast_with_options; +use arrow_schema::DataType; +use divan::Bencher; +use num_traits::AsPrimitive; +use num_traits::NumCast; +use rand::SeedableRng; +use rand::prelude::*; +use rand::rngs::StdRng; +use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_buffer::Buffer; +use vortex_compute::lane_kernels::IndexedSinkExt; +use vortex_compute::lane_kernels::IndexedSourceExt; +use vortex_compute::lane_kernels::LaneZip; +use vortex_compute::lane_kernels::ReinterpretSink; + +fn main() { + assert_overflow_parity(); + assert_null_overflow_suppressed(); + divan::main(); +} + +const SIZES: &[usize] = &[65_536]; + +// ----------------------------------------------------------------------------- +// Cast fixture (u64/u16/i32 lanes + a single validity mask). +// ----------------------------------------------------------------------------- + +struct CastFixture { + values_u64: Buffer, + values_u16: Buffer, + /// Positive `i32` values (always representable as `u32`). Used by the + /// in-place-vs-out-of-place cast bench. + values_i32: Buffer, + mask: BitBuffer, + /// Validity as a plain `Vec` — the source of truth used to build both + /// the Vortex `BitBuffer` mask and the arrow `NullBuffer`. + valid: Vec, +} + +fn cast_fixture(n: usize) -> CastFixture { + let mut rng = StdRng::seed_from_u64(0xC457_1D3E); + + let raw_values: Vec = (0..n) + .map(|_| rng.random_range(0..(u32::MAX as u64))) + .collect(); + let raw_valid: Vec = (0..n).map(|_| rng.random_bool(0.8)).collect(); + + #[expect(clippy::cast_possible_truncation)] + let values_u16 = raw_values + .iter() + .copied() + .map(|v| v as u16) + .collect::>(); + + // Positive i32 values (top bit cleared) — every value fits in u32. + #[expect(clippy::cast_possible_truncation)] + let values_i32 = raw_values + .iter() + .copied() + .map(|v| (v as i32) & i32::MAX) + .collect::>(); + + CastFixture { + values_u64: raw_values.into(), + values_u16, + values_i32, + mask: BitBufferMut::from_iter(raw_valid.iter().copied()).freeze(), + valid: raw_valid, + } +} + +fn uninit_out(n: usize) -> Vec> { + let mut out = Vec::with_capacity(n); + // SAFETY: A `MaybeUninit` does not require initialization. + unsafe { + out.set_len(n); + } + out +} + +// ----------------------------------------------------------------------------- +// Cast benches (single-input, source -> output). +// ----------------------------------------------------------------------------- + +#[divan::bench(args = SIZES)] +fn try_map_into_narrow_u64_u32(bencher: Bencher, n: usize) { + let f = cast_fixture(n); + + bencher + .with_inputs(|| (f.values_u64.clone(), uninit_out::(n))) + .bench_values(|(values, mut out)| { + values + .as_slice() + .try_map_into(out.as_mut_slice(), ::from) + .unwrap(); + out + }); +} + +#[divan::bench(args = SIZES)] +fn arrow_narrow_u64_u32(bencher: Bencher, n: usize) { + let f = cast_fixture(n); + let arr: ArrowArrayRef = Arc::new(UInt64Array::from(f.values_u64.as_slice().to_vec())); + let opts = CastOptions { + safe: false, + ..CastOptions::default() + }; + + bencher + .with_inputs(|| arr.clone()) + .bench_values(|arr| cast_with_options(&arr, &DataType::UInt32, &opts).unwrap()); +} + +#[divan::bench(args = SIZES)] +fn map_with_mask_narrow_u64_u32(bencher: Bencher, n: usize) { + let f = cast_fixture(n); + + bencher + .with_inputs(|| (f.values_u64.clone(), uninit_out::(n))) + .bench_values(|(values, mut out)| { + values.as_slice().map_into(&mut out, |v| v.as_()); + out + }); +} + +/// `try_map_masked_into_widen_u16_u32` and `map_with_mask_widen_u16_u32` have the +/// same runtime — for always-true map operations `try_map_masked_into` is +/// sufficient. +#[divan::bench(args = SIZES)] +fn try_map_masked_into_widen_u16_u32(bencher: Bencher, n: usize) { + let f = cast_fixture(n); + + bencher + .with_inputs(|| (f.values_u16.clone(), f.mask.clone(), uninit_out::(n))) + .bench_values(|(values, mask, mut out)| { + values + .as_slice() + .try_map_masked_into(&mask, out.as_mut_slice(), ::from) + .unwrap(); + out + }); +} + +#[divan::bench(args = SIZES)] +fn map_with_mask_widen_u16_u32(bencher: Bencher, n: usize) { + let f = cast_fixture(n); + + bencher + .with_inputs(|| (f.values_u16.clone(), uninit_out::(n))) + .bench_values(|(values, mut out)| { + values.as_slice().map_into(out.as_mut_slice(), |v| v.as_()); + }); +} + +#[divan::bench(args = SIZES)] +fn arrow_widen_u16_u32(bencher: Bencher, n: usize) { + let f = cast_fixture(n); + let nulls = NullBuffer::from(f.valid.clone()); + let arr: ArrowArrayRef = Arc::new(UInt16Array::new( + f.values_u16.as_slice().to_vec().into(), + Some(nulls), + )); + + bencher.with_inputs(|| arr.clone()).bench_values(|arr| { + cast_with_options(&arr, &DataType::UInt32, &CastOptions::default()).unwrap() + }); +} + +// ----------------------------------------------------------------------------- +// In-place vs out-of-place fallible cast i32 → u32 (same byte width). +// +// `try_map_masked_in_place` mutates the input via `ReinterpretSink` and +// transmutes the wrapper — no output allocation. `try_map_masked_into` allocates +// a fresh `BufferMut` and writes through it. Input values are all positive +// `i32` so every lane succeeds; the two kernels do the same arithmetic, so any +// delta is allocation + memory-traffic overhead. +// ----------------------------------------------------------------------------- + +#[divan::bench(args = SIZES)] +fn try_map_masked_into_narrow_i32_u32(bencher: Bencher, n: usize) { + let f = cast_fixture(n); + + bencher + .with_inputs(|| (f.values_i32.clone(), f.mask.clone(), uninit_out::(n))) + .bench_values(|(values, mask, mut out)| { + values + .as_slice() + .try_map_masked_into(&mask, out.as_mut_slice(), ::from) + .unwrap(); + out + }); +} + +#[divan::bench(args = SIZES)] +fn try_map_masked_in_place_narrow_i32_u32(bencher: Bencher, n: usize) { + let f = cast_fixture(n); + + bencher + .with_inputs(|| (f.values_i32.as_slice().to_vec(), f.mask.clone())) + .bench_values(|(mut values, mask)| { + ReinterpretSink::::new(values.as_mut_slice()) + .try_map_masked_in_place(&mask, ::from) + .unwrap(); + values + }); +} + +#[divan::bench(args = SIZES)] +fn arrow_narrow_i32_u32(bencher: Bencher, n: usize) { + let f = cast_fixture(n); + let nulls = NullBuffer::from(f.valid.clone()); + let arr: ArrowArrayRef = Arc::new(Int32Array::new( + f.values_i32.as_slice().to_vec().into(), + Some(nulls), + )); + let opts = CastOptions { + safe: false, + ..CastOptions::default() + }; + + bencher + .with_inputs(|| arr.clone()) + .bench_values(|arr| cast_with_options(&arr, &DataType::UInt32, &opts).unwrap()); +} + +// ----------------------------------------------------------------------------- +// LaneZip binary kernel: checked `u32 + u32 -> u32` over two nullable columns. +// +// Per-lane `is_none()` flags are bit-packed and AND-ed with the chunk validity +// word, so null-lane overflow is filtered without the closure inspecting `valid`. +// Verified at startup via parity assertions (`assert_overflow_parity` and +// `assert_null_overflow_suppressed`). +// ----------------------------------------------------------------------------- + +const ADD_LHS_VALID_RATE: f64 = 0.7; +const ADD_RHS_VALID_RATE: f64 = 0.8; + +struct AddFixture { + /// Valid lanes carry bounded values; null lanes hold `u32::MAX` so a kernel + /// that ignores validity would `Err` on them. The implementation under test + /// must suppress that. + lhs: Buffer, + rhs: Buffer, + lhs_mask: BitBuffer, + rhs_mask: BitBuffer, + /// Plain `Vec` mirrors of the validity masks — used to build the arrow + /// `NullBuffer`s for the baseline bench. + lhs_valid: Vec, + rhs_valid: Vec, +} + +fn add_fixture(n: usize) -> AddFixture { + let mut lhs_rng = StdRng::seed_from_u64(0); + let mut rhs_rng = StdRng::seed_from_u64(1); + let mut lvr = StdRng::seed_from_u64(2); + let mut rvr = StdRng::seed_from_u64(3); + + let lhs_valid: Vec = (0..n) + .map(|_| lvr.random_bool(ADD_LHS_VALID_RATE)) + .collect(); + let rhs_valid: Vec = (0..n) + .map(|_| rvr.random_bool(ADD_RHS_VALID_RATE)) + .collect(); + + let lhs: Buffer = (0..n) + .map(|i| { + if lhs_valid[i] { + lhs_rng.random_range(0..u16::MAX as u32) + } else { + u32::MAX + } + }) + .collect(); + let rhs: Buffer = (0..n) + .map(|i| { + if rhs_valid[i] { + rhs_rng.random_range(0..u16::MAX as u32) + } else { + u32::MAX + } + }) + .collect(); + + let lhs_mask = BitBufferMut::from_iter(lhs_valid.iter().copied()).freeze(); + let rhs_mask = BitBufferMut::from_iter(rhs_valid.iter().copied()).freeze(); + + AddFixture { + lhs, + rhs, + lhs_mask, + rhs_mask, + lhs_valid, + rhs_valid, + } +} + +#[divan::bench(args = SIZES)] +fn lanezip_checked_add_u32(bencher: Bencher, n: usize) { + let f = add_fixture(n); + bencher + .with_inputs(|| { + ( + f.lhs.clone(), + f.rhs.clone(), + f.lhs_mask.clone(), + f.rhs_mask.clone(), + ) + }) + .bench_refs(|(lhs, rhs, lm, rm)| { + let combined = lm as &BitBuffer & rm as &BitBuffer; + let mut out = uninit_out::(n); + LaneZip::new(lhs.as_slice(), rhs.as_slice()) + .try_map_masked_into(&combined, out.as_mut_slice(), |(a, b)| a.checked_add(b)) + .unwrap(); + (combined, out) + }); +} + +#[divan::bench(args = SIZES)] +fn arrow_checked_add_u32(bencher: Bencher, n: usize) { + let f = add_fixture(n); + let lhs_arr: ArrowArrayRef = Arc::new(UInt32Array::new( + f.lhs.as_slice().to_vec().into(), + Some(NullBuffer::from(f.lhs_valid.clone())), + )); + let rhs_arr: ArrowArrayRef = Arc::new(UInt32Array::new( + f.rhs.as_slice().to_vec().into(), + Some(NullBuffer::from(f.rhs_valid.clone())), + )); + + bencher + .with_inputs(|| (lhs_arr.clone(), rhs_arr.clone())) + .bench_values(|(lhs, rhs)| add(&lhs, &rhs).unwrap()); +} + +// ----------------------------------------------------------------------------- +// Parity assertions — must pass before divan runs benches. +// ----------------------------------------------------------------------------- + +/// Overflow at a valid lane must propagate as `Err`. +fn assert_overflow_parity() { + let lhs: Vec = vec![1, 2, u32::MAX, 4]; + let rhs: Vec = vec![10, 20, 1, 40]; + let valid = vec![true; 4]; + + let mask = BitBufferMut::from_iter(valid).freeze(); + let mut out: Vec> = (0..4).map(|_| MaybeUninit::uninit()).collect(); + let r = LaneZip::new(lhs.as_slice(), rhs.as_slice()).try_map_masked_into( + &mask, + out.as_mut_slice(), + |(a, b)| a.checked_add(b), + ); + assert!(r.is_err(), "bitpack should Err on overflow"); +} + +/// Overflow at a null lane must NOT propagate. +fn assert_null_overflow_suppressed() { + // Lane 2 is null and holds an overflowing value; valid lanes are safe. + let lhs: Vec = vec![1, 2, u32::MAX, 4]; + let rhs: Vec = vec![10, 20, 1, 40]; + let valid = vec![true, true, false, true]; + + let mask = BitBufferMut::from_iter(valid).freeze(); + let mut out = uninit_out::(4); + let r = LaneZip::new(lhs.as_slice(), rhs.as_slice()).try_map_masked_into( + &mask, + out.as_mut_slice(), + |(a, b)| a.checked_add(b), + ); + assert!(r.is_ok(), "bitpack: null-lane overflow leaked"); +} diff --git a/vortex-compute/src/lane_kernels/map_in_place.rs b/vortex-compute/src/lane_kernels/map_in_place.rs new file mode 100644 index 00000000000..335e1c456df --- /dev/null +++ b/vortex-compute/src/lane_kernels/map_in_place.rs @@ -0,0 +1,368 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! In-place lane kernels: read from an [`IndexedSink`] and write back through +//! the same sink (no separate output buffer). +//! +//! [`IndexedSink`]: crate::lane_kernels::sink::IndexedSink + +use vortex_buffer::BitBuffer; + +use crate::lane_kernels::CHUNK_LEN; +use crate::lane_kernels::sink::IndexedSink; + +/// Extension trait providing in-place lane-kernel methods on any [`IndexedSink`]. +/// +/// All methods have default implementations and are inherited via the blanket +/// `impl IndexedSinkExt for S` below. Bring the trait into scope +/// (`use vortex_compute::lane_kernels::IndexedSinkExt;`) to call them with +/// method syntax. +/// +/// [`IndexedSink`]: crate::lane_kernels::sink::IndexedSink +pub trait IndexedSinkExt: IndexedSink + Sized { + /// In-place counterpart of [`IndexedSourceExt::map_into`]. Each lane + /// is replaced with `f(self[i])`. + /// + /// The closure reads `Self::Item` and returns `Self::Write`. For the common + /// case `Self = &mut [T]` both are `T`; for `ReinterpretSink` the read and + /// write types can differ (e.g. read `f32`, write `u32`) over the same + /// backing memory when sizes and alignments match. + /// + /// As with [`IndexedSourceExt::map_into`], use this only when the + /// input is known non-nullable. + /// + /// [`IndexedSourceExt::map_into`]: crate::lane_kernels::map_into::IndexedSourceExt::map_into + #[inline] + fn map_into_in_place(self, mut f: F) + where + F: FnMut(Self::Item) -> Self::Write, + { + #[inline(always)] + fn chunk(values: &mut S, f: &mut F, base: usize, count: usize) + where + S: IndexedSink, + F: FnMut(S::Item) -> S::Write, + { + for bit_idx in 0..count { + let idx = base + bit_idx; + // SAFETY: caller guarantees base + count <= len. + let val = unsafe { values.get_unchecked(idx) }; + let result = f(val); + // SAFETY: caller guarantees base + count <= len. + unsafe { values.set_unchecked(idx, result) }; + } + } + + let mut values = self; + let len = values.len(); + let chunks_count = len / CHUNK_LEN; + let remainder = len % CHUNK_LEN; + + for chunk_idx in 0..chunks_count { + chunk(&mut values, &mut f, chunk_idx * CHUNK_LEN, CHUNK_LEN); + } + if remainder != 0 { + chunk(&mut values, &mut f, chunks_count * CHUNK_LEN, remainder); + } + } + + /// In-place counterpart of [`IndexedSourceExt::try_map_into`]. Each + /// lane is replaced with `f(self[i])`, or `Self::Write::default()` when `f` + /// returns `None`. On failure returns `Err(first_failing_lane)`; the buffer + /// state on `Err` is unspecified. + /// + /// ## Error attribution + /// + /// Per-lane `is_none()` flags are bit-packed into a `u64` at the lane's + /// position — `fail_bits |= (opt.is_none() as u64) << bit_idx`. After the + /// 64-lane loop, `trailing_zeros()` of `fail_bits` recovers the first + /// failing lane index. `OR + shift` per lane is friendlier to the + /// autovectorizer than `min`/`csel` — see [`try_map_masked_in_place`] for + /// the same scheme over a masked variant. + /// + /// [`try_map_masked_in_place`]: IndexedSinkExt::try_map_masked_in_place + /// [`IndexedSourceExt::try_map_into`]: crate::lane_kernels::map_into::IndexedSourceExt::try_map_into + #[inline] + fn try_map_in_place(self, mut f: F) -> Result<(), usize> + where + Self::Write: Default, + F: FnMut(Self::Item) -> Option, + { + #[inline(always)] + fn chunk(values: &mut S, base: usize, count: usize, f: &mut F) -> Option + where + S: IndexedSink, + S::Write: Default, + F: FnMut(S::Item) -> Option, + { + let mut fail_bits: u64 = 0; + for bit_idx in 0..count { + let idx = base + bit_idx; + // SAFETY: caller guarantees base + count <= len. + let val = unsafe { values.get_unchecked(idx) }; + let opt = f(val); + fail_bits |= (opt.is_none() as u64) << bit_idx; + let result = opt.unwrap_or_default(); + // SAFETY: caller guarantees base + count <= len. + unsafe { values.set_unchecked(idx, result) }; + } + (fail_bits != 0).then_some(base + fail_bits.trailing_zeros() as usize) + } + + let mut values = self; + let len = values.len(); + let chunks_count = len / 64; + let remainder = len % 64; + + for chunk_idx in 0..chunks_count { + if let Some(failing) = chunk(&mut values, chunk_idx * 64, 64, &mut f) { + return Err(failing); + } + } + if remainder != 0 + && let Some(failing) = chunk(&mut values, chunks_count * 64, remainder, &mut f) + { + return Err(failing); + } + Ok(()) + } + + /// In-place counterpart of [`IndexedSourceExt::try_map_masked_into`]. Each + /// lane of `self` is replaced with `f(self[i])`, or `Self::Write::default()` + /// if `f` returned `None`. On failure returns `Err(first_failing_lane)`; + /// lanes before that point have been written, and lanes within the failing + /// chunk hold their unwrapped-or-default result. The buffer state on `Err` + /// is intentionally unspecified. + /// + /// **Null-lane failures are filtered automatically** — same semantics as + /// [`IndexedSourceExt::try_map_masked_into`]. The closure has no `valid` + /// parameter; the kernel AND-combines `is_none()` with the chunk's validity + /// bitmap before folding it into the attribution accumulator. + /// + /// ## Why in-place is slower at cache-resident sizes + /// + /// At sizes that fit in L1/L2 the in-place kernel is ~1.5× slower than the + /// out-of-place kernel despite having half the memory traffic, because + /// input and output share memory and the compiler must be conservative + /// reordering loads/stores across iterations. At sizes that exceed L2 the + /// in-place kernel wins back the gap by avoiding the second buffer's DRAM + /// read+write traffic. + /// + /// [`IndexedSourceExt::try_map_masked_into`]: crate::lane_kernels::map_into::IndexedSourceExt::try_map_masked_into + /// + /// # Panics + /// + /// Panics if `self.len() != mask.len()`. + #[inline] + fn try_map_masked_in_place(self, mask: &BitBuffer, mut f: F) -> Result<(), usize> + where + Self::Write: Default, + F: FnMut(Self::Item) -> Option, + { + /// 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. + #[inline(always)] + fn chunk( + values: &mut S, + src_chunk: u64, + base: usize, + count: usize, + f: &mut F, + ) -> Option + where + S: IndexedSink, + S::Write: Default, + F: FnMut(S::Item) -> Option, + { + let mut fail_bits: u64 = 0; + for bit_idx in 0..count { + let idx = base + bit_idx; + // SAFETY: caller guarantees `base + count <= values.len()`. + let val = unsafe { values.get_unchecked(idx) }; + let opt = f(val); + fail_bits |= (opt.is_none() as u64) << bit_idx; + let result = opt.unwrap_or_default(); + unsafe { values.set_unchecked(idx, result) }; + } + let valid_failures = fail_bits & src_chunk; + (valid_failures != 0).then_some(base + valid_failures.trailing_zeros() as usize) + } + + let mut values = self; + let len = values.len(); + assert_eq!(len, mask.len(), "values and mask must have the same length"); + + let chunks = mask.chunks(); + let chunks_count = len / 64; + let remainder = len % 64; + + for (chunk_idx, src_chunk) in chunks.iter().enumerate() { + if let Some(failing) = chunk(&mut values, src_chunk, chunk_idx * 64, 64, &mut f) { + return Err(failing); + } + } + if remainder != 0 + && let Some(failing) = chunk( + &mut values, + chunks.remainder_bits(), + chunks_count * 64, + remainder, + &mut f, + ) + { + return Err(failing); + } + Ok(()) + } +} + +impl IndexedSinkExt for S {} + +#[cfg(test)] +#[allow(clippy::cast_possible_truncation)] +mod tests { + use vortex_buffer::BitBuffer; + use vortex_buffer::BitBufferMut; + + use super::*; + use crate::lane_kernels::sink::ReinterpretSink; + + #[test] + fn try_map_masked_in_place_all_ok() { + let mut values: Vec = (0..200).collect(); + let mask = BitBuffer::new_set(200); + let res = values + .as_mut_slice() + .try_map_masked_in_place(&mask, |v| v.checked_mul(2)); + assert!(res.is_ok()); + let expected: Vec = (0..200u32).map(|v| v * 2).collect(); + assert_eq!(values, expected); + } + + #[test] + fn try_map_masked_in_place_first_failing_chunk_wins() { + let mut values: Vec = (0..200).collect(); + values[83] = u32::MAX; + values[150] = u32::MAX; + let mask = BitBuffer::new_set(200); + let res = values + .as_mut_slice() + .try_map_masked_in_place(&mask, |v| v.checked_mul(2)); + assert_eq!(res, Err(83)); + } + + #[test] + fn try_map_masked_in_place_within_chunk_reports_lowest() { + let mut values: Vec = (0..200).collect(); + values[80] = u32::MAX; + values[100] = u32::MAX; + let mask = BitBuffer::new_set(200); + let res = values + .as_mut_slice() + .try_map_masked_in_place(&mask, |v| v.checked_mul(2)); + assert_eq!(res, Err(80)); + } + + #[test] + fn try_map_masked_in_place_single_failure_lane_exact() { + let mut values: Vec = (0..200).collect(); + values[42] = u32::MAX; + let mask = BitBuffer::new_set(200); + let res = values + .as_mut_slice() + .try_map_masked_in_place(&mask, |v| v.checked_mul(2)); + assert_eq!(res, Err(42)); + } + + #[test] + fn try_map_masked_in_place_null_bypass() { + let mut values: Vec = (0..200).collect(); + values[5] = u32::MAX; + let mask = { + let mut m = BitBufferMut::with_capacity(200); + for i in 0..200 { + m.append(i != 5); + } + m.freeze() + }; + let res = values + .as_mut_slice() + .try_map_masked_in_place(&mask, |v| v.checked_mul(2)); + assert!(res.is_ok(), "null-lane overflow should be filtered"); + assert_eq!(values[5], 0); + assert_eq!(values[6], 12); + } + + #[test] + fn try_map_masked_in_place_remainder_overflow() { + let mut values: Vec = (0..130).collect(); + values[129] = u32::MAX; + let mask = BitBuffer::new_set(130); + let res = values + .as_mut_slice() + .try_map_masked_in_place(&mask, |v| v.checked_mul(2)); + assert_eq!(res, Err(129)); + } + + #[test] + fn try_map_masked_in_place_sliced_mask() { + let big = BitBuffer::new_set(256); + let mask = big.slice(13..143); + assert_eq!(mask.len(), 130); + + let mut values: Vec = (0..130).collect(); + values[77] = u32::MAX; + let res = values + .as_mut_slice() + .try_map_masked_in_place(&mask, |v| v.checked_mul(2)); + assert_eq!(res, Err(77)); + } + + #[test] + fn reinterpret_sink_same_width_f32_u32() { + let mut buf: Vec = (0..130).map(|i| i as f32).collect(); + let mask = BitBuffer::new_set(130); + ReinterpretSink::::new(buf.as_mut_slice()) + .try_map_masked_in_place(&mask, |f| Some(f.to_bits().wrapping_add(1))) + .unwrap(); + // SAFETY: same size + alignment for f32 and u32; every slot now holds a u32 written by + // the closure. + let as_u32: &[u32] = + unsafe { std::slice::from_raw_parts(buf.as_ptr() as *const u32, buf.len()) }; + for (i, &got) in as_u32.iter().enumerate() { + assert_eq!(got, (i as f32).to_bits().wrapping_add(1), "lane {i}"); + } + } + + #[test] + fn reinterpret_sink_failure_reports_lane() { + let mut buf: Vec = (0..200).map(|i| i as f32).collect(); + let mask = BitBuffer::new_set(200); + let res = ReinterpretSink::::new(buf.as_mut_slice()).try_map_masked_in_place( + &mask, + |f| { + if f as u32 == 137 { + None + } else { + Some(f as u32) + } + }, + ); + assert_eq!(res, Err(137)); + } + + #[test] + fn try_map_masked_in_place_partial_chunk_success() { + let mut values: Vec = (0..130).collect(); + let mask = BitBuffer::new_set(130); + let res = values + .as_mut_slice() + .try_map_masked_in_place(&mask, |v| Some(v + 1)); + assert!(res.is_ok()); + assert_eq!(values[0], 1); + assert_eq!(values[63], 64); + assert_eq!(values[64], 65); + assert_eq!(values[129], 130); + } +} diff --git a/vortex-compute/src/lane_kernels/map_into.rs b/vortex-compute/src/lane_kernels/map_into.rs new file mode 100644 index 00000000000..d20b261c274 --- /dev/null +++ b/vortex-compute/src/lane_kernels/map_into.rs @@ -0,0 +1,499 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Out-of-place lane kernels: read from an [`IndexedSource`] and write into a +//! caller-provided `&mut [MaybeUninit]`. + +use std::mem::MaybeUninit; + +use vortex_buffer::BitBuffer; + +use crate::lane_kernels::CHUNK_LEN; +use crate::lane_kernels::source::IndexedSource; + +/// Extension trait providing out-of-place lane-kernel methods on any [`IndexedSource`]. +/// +/// All methods have default implementations and are inherited via the blanket +/// `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)`. +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). + /// + /// **Null-lane failures are filtered automatically.** The closure is called on + /// every lane regardless of validity; if a null lane's stored value causes `f(v)` + /// to return `None`, the kernel does *not* propagate that as `Err`. The per-lane + /// `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 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 + /// the common case, the kernel's automatic filter is sufficient. + /// + /// On failure returns `Err(failing_lane_index)`. Lanes whose `f` returned `None` + /// write `R::default()` into `out`, but the contents of `out` must not be relied + /// upon when this function returns `Err`. + /// + /// [`try_map_into`]: IndexedSourceExt::try_map_into + /// + /// # Panics + /// + /// Panics if `self.len() != mask.len()` or `out.len() != self.len()`. + #[inline] + fn try_map_masked_into( + self, + mask: &BitBuffer, + out: &mut [MaybeUninit], + mut f: F, + ) -> Result<(), usize> + where + R: Copy + Default, + F: FnMut(Self::Item) -> Option, + { + #[inline(always)] + fn chunk( + values: &S, + out: &mut [MaybeUninit], + f: &mut F, + src_chunk: u64, + base: usize, + count: usize, + ) -> Option + where + S: IndexedSource, + R: Copy + Default, + F: FnMut(S::Item) -> Option, + { + let mut fail_bits: u64 = 0; + for bit_idx in 0..count { + let idx = base + bit_idx; + // SAFETY: caller guarantees base + count <= len. + let val = unsafe { values.get_unchecked(idx) }; + let opt = f(val); + fail_bits |= (opt.is_none() as u64) << bit_idx; + let result = opt.unwrap_or_default(); + unsafe { out.get_unchecked_mut(idx).write(result) }; + } + let valid_failures = fail_bits & src_chunk; + (valid_failures != 0).then_some(base + valid_failures.trailing_zeros() as usize) + } + + let values = self; + let len = values.len(); + assert_eq!(len, mask.len(), "values and mask must have the same length"); + assert_eq!(out.len(), len, "out must have the same length as values"); + + let chunks = mask.chunks(); + let chunks_count = len / 64; + 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) { + return Err(idx); + } + } + if remainder != 0 + && let Some(idx) = chunk( + &values, + out, + &mut f, + chunks.remainder_bits(), + chunks_count * 64, + remainder, + ) + { + return Err(idx); + } + Ok(()) + } + + /// Apply `f(value)` lane-by-lane with **no validity awareness at all** — every + /// closure invocation is treated as "happened", regardless of whether the lane + /// is null. Use this only when the input is known non-nullable. + /// + /// # Panics + /// + /// Panics if `out.len() != self.len()`. + #[inline] + fn map_into(self, out: &mut [MaybeUninit], mut f: F) + where + F: FnMut(Self::Item) -> R, + { + #[inline(always)] + fn chunk( + values: &S, + out: &mut [MaybeUninit], + f: &mut F, + base: usize, + count: usize, + ) where + S: IndexedSource, + F: FnMut(S::Item) -> R, + { + for bit_idx in 0..count { + let idx = base + bit_idx; + // SAFETY: caller guarantees base + count <= len. + let val = unsafe { values.get_unchecked(idx) }; + unsafe { out.get_unchecked_mut(idx).write(f(val)) }; + } + } + + let values = self; + let len = values.len(); + assert_eq!(out.len(), len, "out must have the same length as values"); + + let chunks_count = len / CHUNK_LEN; + let remainder = len % CHUNK_LEN; + + for chunk_idx in 0..chunks_count { + chunk(&values, out, &mut f, chunk_idx * CHUNK_LEN, CHUNK_LEN); + } + if remainder != 0 { + chunk(&values, out, &mut f, chunks_count * CHUNK_LEN, remainder); + } + } + + /// Fallible map with **no validity awareness at all** — every `None` returned + /// by the closure is treated as a failure, even at null lanes. + /// + /// # Use this only for non-nullable inputs. + /// + /// For nullable inputs with a fallible closure, use [`try_map_masked_into`] — + /// it has the same value-only closure shape (and the same perf win) but + /// **correctly suppresses null-lane failures** via per-chunk + /// `fail_bits & mask_chunk`. + /// + /// Using this kernel on a nullable input where a null lane's stored value + /// would cause `f` to return `None` will produce a spurious `Err`. This is a + /// correctness footgun on purpose — the name and this doc are how the API + /// signals "you must know your input has no nulls." + /// + /// On failure returns `Err(failing_lane_index)`. + /// + /// [`try_map_masked_into`]: IndexedSourceExt::try_map_masked_into + /// + /// # Panics + /// + /// Panics if `out.len() != self.len()`. + #[inline] + fn try_map_into(self, out: &mut [MaybeUninit], mut f: F) -> Result<(), usize> + where + R: Copy + Default, + F: FnMut(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. + #[inline(always)] + fn chunk( + values: &S, + out: &mut [MaybeUninit], + f: &mut F, + base: usize, + count: usize, + ) -> bool + where + S: IndexedSource, + R: Copy + Default, + F: FnMut(S::Item) -> Option, + { + let mut fail_acc: u64 = 0; + for bit_idx in 0..count { + let idx = base + bit_idx; + // SAFETY: caller guarantees base + count <= len. + let val = unsafe { values.get_unchecked(idx) }; + let opt = f(val); + fail_acc |= opt.is_none() as u64; + let result = opt.unwrap_or_default(); + unsafe { out.get_unchecked_mut(idx).write(result) }; + } + fail_acc != 0 + } + + let values = self; + let len = values.len(); + assert_eq!(out.len(), len, "out must have the same length as values"); + + let chunks_count = len / CHUNK_LEN; + let remainder = len % CHUNK_LEN; + + 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 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)); + } + } + Ok(()) + } +} + +impl IndexedSourceExt for S {} + +/// Shared cold scan: walks a chunk, returns the first lane index where +/// `lane_fails(bit_idx, value)` returns `true`. Used by +/// [`attribute_failure_no_mask`]. +/// +/// Caller guarantees `base + chunk_len <= values.len()`. +#[cold] +#[inline(never)] +fn cold_scan( + values: &S, + base: usize, + chunk_len: usize, + mut lane_fails: impl FnMut(usize /* bit_idx */, S::Item) -> bool, +) -> usize +where + S: IndexedSource, +{ + for bit_idx in 0..chunk_len { + let idx = base + bit_idx; + // SAFETY: caller guarantees idx < values.len(). + let val = unsafe { values.get_unchecked(idx) }; + if lane_fails(bit_idx, val) { + return idx; + } + } + unreachable!("cold_scan called without a failing lane") +} + +/// 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 +where + S: IndexedSource, + F: FnMut(S::Item) -> Option, +{ + cold_scan(values, base, chunk_len, |_bit_idx, val| f(val).is_none()) +} + +#[cfg(test)] +#[allow(clippy::cast_possible_truncation)] +mod tests { + use vortex_buffer::BitBuffer; + use vortex_buffer::BitBufferMut; + + use super::*; + + fn write_t(out: Vec>) -> Vec { + // SAFETY: tests always fully initialize the buffer. + unsafe { std::mem::transmute(out) } + } + + #[test] + fn try_map_masked_into_all_ok() { + let values: Vec = (0..200).collect(); + let mask = BitBuffer::new_set(200); + let mut out = vec![MaybeUninit::::uninit(); 200]; + let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| { + (v <= u32::MAX as u64).then_some(v as u32) + }); + assert!(res.is_ok()); + let got = write_t(out); + assert_eq!(got, (0..200u32).collect::>()); + } + + #[test] + fn try_map_masked_into_overflow_fails() { + let mut values: Vec = (0..200).collect(); + values[137] = (u32::MAX as u64) + 1; + let mask = BitBuffer::new_set(200); + let mut out = vec![MaybeUninit::::uninit(); 200]; + let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| { + (v <= u32::MAX as u64).then_some(v as u32) + }); + assert_eq!(res, Err(137)); + } + + #[test] + fn try_map_masked_into_overflow_reports_first_failing_lane() { + let mut values: Vec = (0..200).collect(); + values[50] = u64::MAX; + values[51] = u64::MAX; + values[137] = u64::MAX; + let mask = BitBuffer::new_set(200); + let mut out = vec![MaybeUninit::::uninit(); 200]; + let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| { + (v <= u32::MAX as u64).then_some(v as u32) + }); + assert_eq!(res, Err(50)); + } + + #[test] + fn try_map_masked_into_value_only_closure_filters_null_overflow() { + let mut values: Vec = (0..200).collect(); + values[5] = u64::MAX; + values[42] = u64::MAX; + let mask = { + let mut m = BitBufferMut::with_capacity(200); + for i in 0..200 { + m.append(i != 5 && i != 42); + } + m.freeze() + }; + let mut out = vec![MaybeUninit::::uninit(); 200]; + let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| { + (v <= u32::MAX as u64).then_some(v as u32) + }); + assert!( + res.is_ok(), + "null-lane overflow should be filtered by the cold path" + ); + } + + #[test] + fn try_map_masked_into_value_only_closure_reports_first_valid_failure() { + let mut values: Vec = (0..200).collect(); + values[5] = u64::MAX; + values[42] = u64::MAX; + values[77] = u64::MAX; + values[100] = u64::MAX; + let mask = { + let mut m = BitBufferMut::with_capacity(200); + for i in 0..200 { + m.append(i != 5 && i != 42); + } + m.freeze() + }; + let mut out = vec![MaybeUninit::::uninit(); 200]; + let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| { + (v <= u32::MAX as u64).then_some(v as u32) + }); + assert_eq!(res, Err(77)); + } + + #[test] + fn try_map_masked_into_null_lane_bypasses_check() { + let mut values: Vec = (0..200).collect(); + values[5] = u64::MAX; + let mask = { + let mut m = BitBufferMut::with_capacity(200); + for i in 0..200 { + m.append(i != 5); + } + m.freeze() + }; + let mut out = vec![MaybeUninit::::uninit(); 200]; + let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| { + (v <= u32::MAX as u64).then_some(v as u32) + }); + assert!(res.is_ok()); + let got = write_t(out); + assert_eq!(got[5], 0); + assert_eq!(got[6], 6); + } + + #[test] + fn try_map_masked_into_branchful_matches_branchless() { + let mut values: Vec = (0..130).map(|i| i as u64 * 7).collect(); + values[2] = u64::MAX; + values[65] = u32::MAX as u64; + let mask = { + let mut m = BitBufferMut::with_capacity(130); + for i in 0..130 { + m.append(!matches!(i, 2 | 17 | 99)); + } + m.freeze() + }; + + let mut branchless = vec![MaybeUninit::::uninit(); 130]; + let mut branchful = vec![MaybeUninit::::uninit(); 130]; + values + .as_slice() + .try_map_masked_into(&mask, &mut branchless, |v| { + (v <= u32::MAX as u64).then_some(v as u32) + }) + .unwrap(); + values + .as_slice() + .try_map_masked_into(&mask, &mut branchful, |v| u32::try_from(v).ok()) + .unwrap(); + + assert_eq!(write_t(branchful), write_t(branchless)); + } + + #[test] + fn try_map_masked_into_partial_chunk() { + let values: Vec = (0..130).collect(); + let mask = BitBuffer::new_set(130); + let mut out = vec![MaybeUninit::::uninit(); 130]; + let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| { + (v <= u32::MAX as u64).then_some(v as u32) + }); + assert!(res.is_ok()); + let got = write_t(out); + assert_eq!(got.len(), 130); + assert_eq!(got[129], 129); + } + + #[test] + fn try_map_masked_into_sliced_mask_unaligned_offset() { + let big = BitBuffer::new_set(256); + let mask = big.slice(13..143); + assert_eq!(mask.len(), 130); + + let values: Vec = (0..130).collect(); + let mut out = vec![MaybeUninit::::uninit(); 130]; + let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| { + (v <= u32::MAX as u64).then_some(v as u32) + }); + assert!(res.is_ok()); + let got = write_t(out); + assert_eq!(got, (0..130u32).collect::>()); + } + + #[test] + fn try_map_masked_into_sliced_mask_with_overflow() { + let big = BitBuffer::new_set(256); + let mask = big.slice(13..143); + assert_eq!(mask.len(), 130); + + let mut values: Vec = (0..130).collect(); + values[77] = u64::MAX; + let mut out = vec![MaybeUninit::::uninit(); 130]; + let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| { + (v <= u32::MAX as u64).then_some(v as u32) + }); + assert_eq!(res, Err(77)); + } + + #[test] + fn try_map_masked_into_sliced_mask_null_lanes() { + let mut m = BitBufferMut::with_capacity(256); + for i in 0..256 { + m.append(i % 3 != 0); + } + let big = m.freeze(); + let mask = big.slice(13..143); + assert_eq!(mask.len(), 130); + + let mut values: Vec = (0..130).collect(); + values[2] = u64::MAX; + let mut out = vec![MaybeUninit::::uninit(); 130]; + let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| { + (v <= u32::MAX as u64).then_some(v as u32) + }); + assert!(res.is_ok(), "null lane should bypass the range check"); + } + + #[test] + fn try_map_masked_into_overflow_in_remainder() { + let mut values: Vec = (0..130).collect(); + values[129] = (u32::MAX as u64) + 1; + let mask = BitBuffer::new_set(130); + let mut out = vec![MaybeUninit::::uninit(); 130]; + let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| { + (v <= u32::MAX as u64).then_some(v as u32) + }); + assert_eq!(res, Err(129)); + } +} diff --git a/vortex-compute/src/lane_kernels/mod.rs b/vortex-compute/src/lane_kernels/mod.rs new file mode 100644 index 00000000000..5fd464f6a85 --- /dev/null +++ b/vortex-compute/src/lane_kernels/mod.rs @@ -0,0 +1,49 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Elementwise lane kernels over indexed sources. +//! +//! Replaces `&[T]` with an [`IndexedSource`] trait: each lane read is +//! `unsafe fn get_unchecked(i) -> Item`, independent across iterations. For `&[T]` +//! this inlines to the same indexed load as the slice kernel; for [`LaneZip`]`(&[A], &[B])` +//! it gives two independent indexed reads per lane — both shapes the auto-vectorizer +//! handles. +//! +//! The module is split into: +//! +//! - [`source`] — the [`IndexedSource`] trait, [`LaneZip`], and read-only adapters. +//! - [`sink`] — the [`IndexedSink`] trait and [`ReinterpretSink`]. +//! - [`map_into`] — out-of-place kernels via [`IndexedSourceExt`] (writes into a +//! caller-provided `&mut [MaybeUninit]`). +//! - [`map_in_place`] — in-place kernels via [`IndexedSinkExt`] (writes back through +//! the sink itself). +//! +//! The kernels never allocate. Both kernel families handle a mask with a non-byte-aligned +//! offset and with a logical `len` shorter than the underlying byte buffer, via +//! `BitBuffer::chunks`. + +pub mod map_in_place; +pub mod map_into; +pub mod sink; +pub mod source; + +pub use map_in_place::IndexedSinkExt; +pub use map_into::IndexedSourceExt; +pub use sink::IndexedSink; +pub use sink::ReinterpretSink; +pub use source::IndexedSource; +pub use source::LaneZip; + +/// Loop-tiling chunk length for the **no-mask** kernels ([`IndexedSourceExt::map_into`], +/// [`IndexedSourceExt::try_map_into`], [`IndexedSinkExt::map_into_in_place`]). +/// +/// This is a pure tuning knob: those kernels split the lane range into +/// `len / CHUNK_LEN` full chunks plus a remainder, so any value yields correct +/// results and only the codegen/tiling changes. Vary it to tune performance. +/// +/// It does **not** apply to the masked or bit-packed kernels: those consume one +/// [`BitBuffer::chunks`] u64 validity word per chunk and pack per-lane fail bits +/// with `<< bit_idx` into a `u64`, so their chunk length is locked to 64. +/// +/// [`BitBuffer::chunks`]: vortex_buffer::BitBuffer::chunks +pub(crate) const CHUNK_LEN: usize = 64; diff --git a/vortex-compute/src/lane_kernels/sink.rs b/vortex-compute/src/lane_kernels/sink.rs new file mode 100644 index 00000000000..86826fd5217 --- /dev/null +++ b/vortex-compute/src/lane_kernels/sink.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Writable lane sink — the [`IndexedSink`] trait and the [`ReinterpretSink`] adapter +//! for in-place type-punning kernels. + +use std::marker::PhantomData; +use std::mem::align_of; +use std::mem::size_of; + +use crate::lane_kernels::source::IndexedSource; + +/// An [`IndexedSource`] that also supports unchecked indexed writes — the binding +/// for in-place kernels. +/// +/// `Write` is the type written by `set_unchecked` and may differ from +/// `IndexedSource::Item` (the read type). For the canonical `&mut [T]` impl +/// both are `T`. The decoupling is what makes [`ReinterpretSink`] possible — +/// a wrapper that reads `F` and writes `T` over the same backing memory when +/// the two have identical size and alignment. +/// +/// Implemented for `&mut [T]`; not implemented for `LaneZip` (you can't write a +/// `(A, B)` pair back to two separate sources via a single index). +pub trait IndexedSink: IndexedSource { + /// The per-lane write type. Equal to `::Item` for + /// `&mut [T]`; different for [`ReinterpretSink`]. + type Write: Copy; + + /// Write `value` into lane `i` without bounds checking. + /// + /// # Safety + /// + /// `i` must be strictly less than `self.len()`. + unsafe fn set_unchecked(&mut self, i: usize, value: Self::Write); +} + +impl IndexedSink for &mut [T] { + type Write = T; + #[inline] + unsafe fn set_unchecked(&mut self, i: usize, value: T) { + // SAFETY: caller guarantees i < self.len(). + unsafe { *<[T]>::get_unchecked_mut(self, i) = value }; + } +} + +/// A sink that reads `F`-values and writes `T`-values over the same backing +/// slice of `F`, reinterpreting each `T` as `F`-bits on write. +/// +/// Requires `size_of::() == size_of::()` and `align_of::() == align_of::()`. +/// Both hold for any pair of `NativePType` primitives with equal byte width +/// (e.g. `u32` ↔ `f32`, `u64` ↔ `i64`, `f64` ↔ `u64`). +/// +/// Use this when an in-place kernel needs to convert lanes between two +/// types of identical width without allocating a second buffer. After the +/// kernel completes every slot holds a valid `T`-bit pattern; the caller +/// can recover a typed view via `BufferMut::transmute::()`. +pub struct ReinterpretSink<'a, F, T> { + slice: &'a mut [F], + _phantom: PhantomData, +} + +impl<'a, F, T> ReinterpretSink<'a, F, T> { + /// Construct a `ReinterpretSink` from `&mut [F]`. + /// + /// # Panics + /// + /// Panics if `size_of::() != size_of::()` or + /// `align_of::() != align_of::()`. + pub fn new(slice: &'a mut [F]) -> Self { + assert_eq!( + size_of::(), + size_of::(), + "ReinterpretSink requires F and T to have the same size", + ); + assert_eq!( + align_of::(), + align_of::(), + "ReinterpretSink requires F and T to have the same alignment", + ); + Self { + slice, + _phantom: PhantomData, + } + } +} + +impl IndexedSource for ReinterpretSink<'_, F, T> { + type Item = F; + #[inline] + fn len(&self) -> usize { + self.slice.len() + } + #[inline] + unsafe fn get_unchecked(&self, i: usize) -> F { + // SAFETY: caller guarantees i < self.slice.len(). Pointer arithmetic + // avoids method-resolution ambiguity between `<[F]>::get_unchecked` and + // `IndexedSource::get_unchecked`. + unsafe { *self.slice.as_ptr().add(i) } + } +} + +impl IndexedSink for ReinterpretSink<'_, F, T> { + type Write = T; + #[inline] + unsafe fn set_unchecked(&mut self, i: usize, value: T) { + // SAFETY: caller guarantees i < self.slice.len(); `new` enforces + // size_of::() == size_of::() and align_of::() == align_of::(), + // so the F-slot can hold a `T` without overflow or misalignment. + unsafe { + let ptr = self.slice.as_mut_ptr().add(i) as *mut T; + ptr.write(value); + } + } +} diff --git a/vortex-compute/src/lane_kernels/source.rs b/vortex-compute/src/lane_kernels/source.rs new file mode 100644 index 00000000000..44c9cb27096 --- /dev/null +++ b/vortex-compute/src/lane_kernels/source.rs @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Read-only lane source — the [`IndexedSource`] trait and the [`LaneZip`] adapter. + +/// A length-known source supporting unchecked indexed reads. +/// +/// Implemented for `&[T]` (with `T: Copy`) and for [`LaneZip`] over two `IndexedSource`s. +/// The kernels in this crate require this trait instead of `Iterator` so that lane +/// 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; + /// Logical lane count. + fn len(&self) -> usize; + /// Returns true when there are no lanes. + fn is_empty(&self) -> bool { + self.len() == 0 + } + /// Read the lane at `i` without bounds checking. + /// + /// # Safety + /// + /// `i` must be strictly less than `self.len()`. + unsafe fn get_unchecked(&self, i: usize) -> Self::Item; +} + +impl IndexedSource for &[T] { + type Item = T; + #[inline] + fn len(&self) -> usize { + <[T]>::len(self) + } + #[inline] + unsafe fn get_unchecked(&self, i: usize) -> T { + // SAFETY: caller guarantees i < self.len(). + unsafe { *<[T]>::get_unchecked(self, i) } + } +} + +impl IndexedSource for &mut [T] { + type Item = T; + #[inline] + fn len(&self) -> usize { + <[T]>::len(self) + } + #[inline] + unsafe fn get_unchecked(&self, i: usize) -> T { + // SAFETY: caller guarantees i < self.len(). + unsafe { *<[T]>::get_unchecked(self, i) } + } +} + +/// 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); + +impl LaneZip { + /// Build a `LaneZip` from two equal-length sources. + /// + /// # Panics + /// + /// Panics if the two operands have different lengths. + pub fn new(a: A, b: B) -> Self { + assert_eq!( + a.len(), + b.len(), + "LaneZip operands must have the same length" + ); + Self(a, b) + } +} + +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] + unsafe fn get_unchecked(&self, i: usize) -> (A::Item, B::Item) { + // SAFETY: caller guarantees i < self.len(); `new` enforces matching lengths. + unsafe { (self.0.get_unchecked(i), self.1.get_unchecked(i)) } + } +} diff --git a/vortex-compute/src/lib.rs b/vortex-compute/src/lib.rs new file mode 100644 index 00000000000..1c813625561 --- /dev/null +++ b/vortex-compute/src/lib.rs @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![deny(missing_docs)] + +//! Lane-level compute kernels over Vortex buffers. +//! +//! Today this crate hosts the [`lane_kernels`] module — an [`IndexedSource`]/[`IndexedSink`] +//! abstraction plus mask-aware map kernels that the autovectorizer can drive through +//! independent lane reads/writes. Additional kernels will land here. +//! +//! [`IndexedSource`]: lane_kernels::IndexedSource +//! [`IndexedSink`]: lane_kernels::IndexedSink + +pub mod lane_kernels; diff --git a/vortex-cuda/Cargo.toml b/vortex-cuda/Cargo.toml index c4c083b97a3..cc34bee74fb 100644 --- a/vortex-cuda/Cargo.toml +++ b/vortex-cuda/Cargo.toml @@ -45,7 +45,7 @@ vortex-error = { workspace = true, features = ["object_store"] } vortex-nvcomp = { path = "nvcomp" } [dev-dependencies] -criterion = { package = "codspeed-criterion-compat-walltime", version = "4.3.0" } +criterion = { workspace = true } futures = { workspace = true, features = ["executor"] } rstest = { workspace = true } tokio = { workspace = true, features = ["rt", "macros"] } diff --git a/vortex-cuda/README.md b/vortex-cuda/README.md index 24b8f420290..4024372adf3 100644 --- a/vortex-cuda/README.md +++ b/vortex-cuda/README.md @@ -27,10 +27,27 @@ cmake -S cpp -B cpp/build \ -DBUILD_TESTS=ON \ -DDISABLE_DEPRECATION_WARNINGS=ON \ -DCMAKE_BUILD_TYPE=Debug \ + -DCUDF_BUILD_STATIC_DEPS=OFF \ -DCUDF_BUILD_STREAMS_TEST_UTIL=OFF \ -DCUDAToolkit_ROOT=/usr/local/cuda \ -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc \ -DCMAKE_C_COMPILER=gcc \ -DCMAKE_CXX_COMPILER=g++ \ + # In large AArch64 Debug links, CALL26/JUMP26 relocations can exceed their branch range. + # Use 1 MiB stub groups so GNU ld emits veneers close enough to each relocation site. + -DCMAKE_SHARED_LINKER_FLAGS="-Wl,--stub-group-size=1048576" \ + -DCMAKE_EXE_LINKER_FLAGS="-Wl,--stub-group-size=1048576" \ -GNinja && cmake --build cpp/build --target INTEROP_TEST --parallel ``` + +## Running the cuDF test harness + +```sh +cargo build -p vortex-test-e2e-cuda +cmake --build /path/to/cudf-test-harness/build --target cudf-test-harness --parallel + +LD_LIBRARY_PATH=/usr/local/cuda-13.1/compat \ + target/debug/cudf_harness_runner \ + /path/to/cudf-test-harness/build/cudf-test-harness \ + target/debug/libvortex_test_e2e_cuda.so +``` diff --git a/vortex-cuda/benches/alp_cuda.rs b/vortex-cuda/benches/alp_cuda.rs index 342f85d4902..5a1fadb7570 100644 --- a/vortex-cuda/benches/alp_cuda.rs +++ b/vortex-cuda/benches/alp_cuda.rs @@ -21,8 +21,8 @@ use criterion::Throughput; use cudarc::driver::DeviceRepr; use futures::executor::block_on; use vortex::array::IntoArray; -use vortex::array::LEGACY_SESSION; use vortex::array::VortexSessionExecute; +use vortex::array::array_session; use vortex::array::arrays::PrimitiveArray; use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; @@ -32,7 +32,6 @@ use vortex::encodings::alp::ALPArrayExt; use vortex::encodings::alp::ALPFloat; use vortex::encodings::alp::alp_encode; use vortex::error::VortexExpect; -use vortex::session::VortexSession; use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; use vortex_cuda::executor::CudaArrayExt; @@ -75,7 +74,7 @@ where let encoded = alp_encode( primitive_array.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), ) .vortex_expect("failed to ALP-encode array"); @@ -110,7 +109,7 @@ where let timer = timed.timer(); let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()) + 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)); diff --git a/vortex-cuda/benches/arrow_binary_cuda.rs b/vortex-cuda/benches/arrow_binary_cuda.rs index bd3154e471a..6668abd3866 100644 --- a/vortex-cuda/benches/arrow_binary_cuda.rs +++ b/vortex-cuda/benches/arrow_binary_cuda.rs @@ -29,7 +29,6 @@ use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::error::VortexExpect; use vortex::error::VortexResult; -use vortex::session::VortexSession; use vortex_cuda::CudaExecutionCtx; use vortex_cuda::CudaSession; use vortex_cuda::arrow::ArrowDeviceArray; @@ -92,9 +91,10 @@ fn benchmark_arrow_binary_export(c: &mut Criterion) { let timed = TimedLaunchStrategy::default(); let timer = timed.timer(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context") - .with_launch_strategy(Arc::new(timed)); + let mut cuda_ctx = + CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) + .vortex_expect("failed to create execution context") + .with_launch_strategy(Arc::new(timed)); let array = block_on(out_of_line_binary(len, &mut cuda_ctx)) .vortex_expect("failed to create binary fixture"); diff --git a/vortex-cuda/benches/arrow_validity_cuda.rs b/vortex-cuda/benches/arrow_validity_cuda.rs index 9221d7c0698..d1142e361f4 100644 --- a/vortex-cuda/benches/arrow_validity_cuda.rs +++ b/vortex-cuda/benches/arrow_validity_cuda.rs @@ -1,5 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#![expect(clippy::cast_possible_truncation)] //! CUDA benchmarks for Arrow validity bitmap repacking. @@ -9,16 +10,27 @@ mod timed_launch_strategy; use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::Duration; +use std::time::Instant; use criterion::BenchmarkId; use criterion::Criterion; use criterion::Throughput; use futures::executor::block_on; +use vortex::array::IntoArray; +use vortex::array::arrays::BoolArray; +use vortex::array::arrays::PrimitiveArray; use vortex::array::buffer::BufferHandle; +use vortex::array::validity::Validity; use vortex::buffer::BitBuffer; +use vortex::buffer::Buffer; +use vortex::dtype::PType; use vortex::error::VortexExpect; +use vortex::error::VortexResult; use vortex::session::VortexSession; +use vortex_cuda::CudaExecutionCtx; use vortex_cuda::CudaSession; +use vortex_cuda::arrow::ArrowDeviceArray; +use vortex_cuda::arrow::DeviceArrayExt; use vortex_cuda::arrow::test_harness; use vortex_cuda_macros::cuda_available; use vortex_cuda_macros::cuda_not_available; @@ -27,12 +39,117 @@ use crate::timed_launch_strategy::TimedLaunchStrategy; const INPUT_OFFSET: usize = 5; const ARROW_OFFSET: usize = 3; +const EXPORT_BENCH_SIZES: &[(usize, &str)] = &[(100_000_000, "100M")]; + +fn validity_bitmap_byte_len(len: usize, bit_offset: usize) -> usize { + (bit_offset + len).div_ceil(8) +} + +unsafe fn release_arrow_device_array(array: &mut ArrowDeviceArray) { + unsafe { + if let Some(release) = array.array.release { + release(&raw mut array.array); + } + } +} + +async fn device_validity_buffer( + len: usize, + validity_offset: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<(usize, BufferHandle)> { + let validity_bits = BitBuffer::collect_bool(len + validity_offset, |idx| idx % 3 != 0) + .slice(validity_offset..validity_offset + len); + let (validity_offset, _, validity_buffer) = validity_bits.into_inner(); + Ok(( + validity_offset, + ctx.ensure_on_device(BufferHandle::new_host(validity_buffer)) + .await?, + )) +} + +async fn primitive_with_device_bool_validity( + len: usize, + validity_offset: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let values = Buffer::::from_iter((0..len).map(|idx| idx as i32)); + let values = ctx + .ensure_on_device(BufferHandle::new_host(values.into_byte_buffer())) + .await?; + + let (validity_offset, validity_buffer) = + device_validity_buffer(len, validity_offset, ctx).await?; + let validity = + BoolArray::new_handle(validity_buffer, validity_offset, len, Validity::NonNullable) + .into_array(); + + Ok( + PrimitiveArray::from_buffer_handle(values, PType::I32, Validity::Array(validity)) + .into_array(), + ) +} + +fn benchmark_arrow_validity_export(c: &mut Criterion) { + let mut group = c.benchmark_group("cuda"); + + for &(len, len_label) in EXPORT_BENCH_SIZES { + for (case, validity_offset) in + [("device_bitmap", 0), ("device_bitmap_repack", INPUT_OFFSET)] + { + group.throughput(Throughput::Bytes( + validity_bitmap_byte_len(len, validity_offset) as u64, + )); + group.bench_with_input( + BenchmarkId::new(format!("cuda/arrow_validity/export/{case}"), len_label), + &len, + |b, &len| { + b.iter_custom(|iters| { + let mut cuda_ctx = + CudaSession::create_execution_ctx(&VortexSession::empty()) + .vortex_expect("failed to create execution context"); + let array = block_on(primitive_with_device_bool_validity( + len, + validity_offset, + &mut cuda_ctx, + )) + .vortex_expect("failed to create primitive fixture"); + + let mut exported_arrays = Vec::with_capacity( + usize::try_from(iters) + .vortex_expect("iteration count does not fit usize"), + ); + + let start = Instant::now(); + for _ in 0..iters { + exported_arrays.push( + block_on(array.clone().export_device_array(&mut cuda_ctx)) + .vortex_expect("failed to export device array"), + ); + } + let elapsed = start.elapsed(); + + for exported in &mut exported_arrays { + unsafe { release_arrow_device_array(exported) }; + } + + elapsed + }); + }, + ); + } + } + + group.finish(); +} fn benchmark_arrow_validity_repack(c: &mut Criterion) { let mut group = c.benchmark_group("cuda"); for &(len, len_label) in bench_config::BENCH_SIZES { - group.throughput(Throughput::Elements(len as u64)); + group.throughput(Throughput::Bytes( + validity_bitmap_byte_len(len, INPUT_OFFSET) as u64, + )); group.bench_with_input( BenchmarkId::new("cuda/arrow_validity/repack", len_label), &len, @@ -40,16 +157,12 @@ fn benchmark_arrow_validity_repack(c: &mut Criterion) { b.iter_custom(|iters| { let timed = TimedLaunchStrategy::default(); let timer = timed.timer(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) .vortex_expect("failed to create execution context") .with_launch_strategy(Arc::new(timed)); - let source = BitBuffer::collect_bool(len + INPUT_OFFSET, |idx| idx % 3 != 0); - let sliced = source.slice(INPUT_OFFSET..INPUT_OFFSET + len); - let (input_offset, _, input_buffer) = sliced.into_inner(); - let input_buffer = - block_on(cuda_ctx.ensure_on_device(BufferHandle::new_host(input_buffer))) - .vortex_expect("failed to copy validity input to device"); + let (input_offset, input_buffer) = + block_on(device_validity_buffer(len, INPUT_OFFSET, &mut cuda_ctx)) + .vortex_expect("failed to create validity fixture"); for _ in 0..iters { let output = test_harness::repack_arrow_validity_buffer( @@ -72,10 +185,51 @@ fn benchmark_arrow_validity_repack(c: &mut Criterion) { group.finish(); } +fn benchmark_arrow_validity_count_nulls(c: &mut Criterion) { + let mut group = c.benchmark_group("cuda"); + + for &(len, len_label) in bench_config::BENCH_SIZES { + group.throughput(Throughput::Bytes( + validity_bitmap_byte_len(len, ARROW_OFFSET) as u64, + )); + group.bench_with_input( + BenchmarkId::new("cuda/arrow_validity/count_nulls", len_label), + &len, + |b, &len| { + b.iter_custom(|iters| { + let timed = TimedLaunchStrategy::default(); + let timer = timed.timer(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + .vortex_expect("failed to create execution context") + .with_launch_strategy(Arc::new(timed)); + let (_, input_buffer) = + block_on(device_validity_buffer(len, ARROW_OFFSET, &mut cuda_ctx)) + .vortex_expect("failed to create validity fixture"); + + for _ in 0..iters { + let null_count = test_harness::count_arrow_validity_nulls( + &input_buffer, + len, + ARROW_OFFSET, + &mut cuda_ctx, + ) + .vortex_expect("failed to count Arrow validity nulls"); + std::hint::black_box(null_count); + } + + Duration::from_nanos(timer.load(Ordering::Relaxed)) + }); + }, + ); + } + + group.finish(); +} + criterion::criterion_group! { name = benches; config = bench_config::cuda_bench_config(); - targets = benchmark_arrow_validity_repack + targets = benchmark_arrow_validity_repack, benchmark_arrow_validity_count_nulls, benchmark_arrow_validity_export } #[cuda_available] diff --git a/vortex-cuda/benches/bitpacked_cuda.rs b/vortex-cuda/benches/bitpacked_cuda.rs index 5f73f02feb5..f216c33e2dd 100644 --- a/vortex-cuda/benches/bitpacked_cuda.rs +++ b/vortex-cuda/benches/bitpacked_cuda.rs @@ -21,8 +21,8 @@ use criterion::Throughput; use cudarc::driver::DeviceRepr; use futures::executor::block_on; use vortex::array::IntoArray; -use vortex::array::LEGACY_SESSION; use vortex::array::VortexSessionExecute; +use vortex::array::array_session; use vortex::array::arrays::PrimitiveArray; use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; @@ -31,7 +31,6 @@ use vortex::encodings::fastlanes::BitPackedArray; use vortex::encodings::fastlanes::BitPackedData; use vortex::encodings::fastlanes::unpack_iter::BitPacked; use vortex::error::VortexExpect; -use vortex::session::VortexSession; use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; use vortex_cuda::executor::CudaArrayExt; @@ -58,7 +57,7 @@ where .collect(); let primitive_array = PrimitiveArray::new(Buffer::from(values), NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); BitPackedData::encode(&primitive_array.into_array(), bit_width, &mut ctx) .vortex_expect("failed to create BitPacked array") } @@ -99,7 +98,7 @@ where .collect(); let primitive_array = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); BitPackedData::encode(&primitive_array, bit_width, &mut ctx) .vortex_expect("failed to create BitPacked array with patches") } @@ -129,10 +128,11 @@ where let timed = TimedLaunchStrategy::default(); let timer = timed.timer(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context") - .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) - .with_launch_strategy(Arc::new(timed)); + 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(array.clone().into_array().execute_cuda(&mut cuda_ctx)).unwrap(); @@ -184,7 +184,7 @@ where let timer = timed.timer(); let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()) + 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)); diff --git a/vortex-cuda/benches/date_time_parts_cuda.rs b/vortex-cuda/benches/date_time_parts_cuda.rs index 06d45dd2354..f5816717a55 100644 --- a/vortex-cuda/benches/date_time_parts_cuda.rs +++ b/vortex-cuda/benches/date_time_parts_cuda.rs @@ -30,7 +30,6 @@ use vortex::encodings::datetime_parts::DateTimePartsArray; use vortex::error::VortexExpect; use vortex::extension::datetime::TimeUnit; use vortex::extension::datetime::Timestamp; -use vortex::session::VortexSession; use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; use vortex_cuda::executor::CudaArrayExt; @@ -68,10 +67,11 @@ fn benchmark_datetimeparts(c: &mut Criterion) { let timed = TimedLaunchStrategy::default(); let timer = timed.timer(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context") - .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) - .with_launch_strategy(Arc::new(timed)); + 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 immediately here diff --git a/vortex-cuda/benches/dict_cuda.rs b/vortex-cuda/benches/dict_cuda.rs index 0cc785c869a..6ed91f4179e 100644 --- a/vortex-cuda/benches/dict_cuda.rs +++ b/vortex-cuda/benches/dict_cuda.rs @@ -27,7 +27,6 @@ use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; use vortex::dtype::NativePType; use vortex::error::VortexExpect; -use vortex::session::VortexSession; use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; use vortex_cuda::executor::CudaArrayExt; @@ -96,10 +95,11 @@ where let timed = TimedLaunchStrategy::default(); let timer = timed.timer(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context") - .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) - .with_launch_strategy(Arc::new(timed)); + 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(dict_array.clone().into_array().execute_cuda(&mut cuda_ctx)) diff --git a/vortex-cuda/benches/dynamic_dispatch_cuda.rs b/vortex-cuda/benches/dynamic_dispatch_cuda.rs index cf89c95efff..81ee3a45167 100644 --- a/vortex-cuda/benches/dynamic_dispatch_cuda.rs +++ b/vortex-cuda/benches/dynamic_dispatch_cuda.rs @@ -27,8 +27,8 @@ use cudarc::driver::sys::CUevent_flags; use futures::executor::block_on; use vortex::array::ArrayRef; use vortex::array::IntoArray; -use vortex::array::LEGACY_SESSION; use vortex::array::VortexSessionExecute; +use vortex::array::array_session; use vortex::array::arrays::DictArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::buffer::BufferHandle; @@ -53,13 +53,13 @@ use vortex::encodings::runend::RunEnd; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_err; -use vortex::session::VortexSession; use vortex_cuda::CudaBufferExt; use vortex_cuda::CudaDeviceBuffer; use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaExecutionCtx; use vortex_cuda::CudaSession; use vortex_cuda::CudaSessionExt; +use vortex_cuda::cuda_session; use vortex_cuda::dynamic_dispatch::CudaDispatchPlan; use vortex_cuda::dynamic_dispatch::DispatchPlan; use vortex_cuda::dynamic_dispatch::MaterializedPlan; @@ -210,7 +210,7 @@ fn bench_for_bitpacked(c: &mut Criterion) { .map(|i| (i as u64 % (max_val + 1)) as u32) .collect(); let prim = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let bp = BitPackedData::encode(&prim.into_array(), bit_width, &mut ctx).vortex_expect("bitpack"); let array = FoR::try_new(bp.into_array(), Scalar::from(reference)) @@ -222,7 +222,7 @@ fn bench_for_bitpacked(c: &mut Criterion) { len, |b, &n| { let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); @@ -255,7 +255,7 @@ fn bench_dict_bp_codes(c: &mut Criterion) { let codes: Vec = (0..*len).map(|i| (i % dict_size) as u32).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let codes_bp = BitPackedData::encode(&codes_prim.into_array(), dict_bit_width, &mut ctx) .vortex_expect("bitpack codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); @@ -267,7 +267,7 @@ fn bench_dict_bp_codes(c: &mut Criterion) { len, |b, &n| { let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); @@ -300,7 +300,7 @@ fn bench_runend(c: &mut Criterion) { let ends: Vec = (1..=num_runs).map(|i| (i * run_len) as u32).collect(); let values: Vec = (0..num_runs).map(|i| (i * 7 + 42) as u32).collect(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let ends_arr = PrimitiveArray::new(Buffer::from(ends), NonNullable).into_array(); let values_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); let re = RunEnd::new(ends_arr, values_arr, &mut ctx); @@ -311,7 +311,7 @@ fn bench_runend(c: &mut Criterion) { len, |b, &n| { let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); @@ -347,7 +347,7 @@ fn bench_dict_bp_codes_alp_for_bp_values_dynanmic_dispatch(c: &mut Criterion) { .map(|i| ::decode_single(10 + i as i32, exponents)) .collect(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // values: ALP → FoR → BitPacked. let float_prim = PrimitiveArray::new(Buffer::from(dict_floats), NonNullable); @@ -359,6 +359,7 @@ fn bench_dict_bp_codes_alp_for_bp_values_dynanmic_dispatch(c: &mut Criterion) { .clone() .execute::(&mut ctx) .vortex_expect("to primitive"), + &mut ctx, ) .vortex_expect("for encode"); let bp = BitPackedData::encode(for_arr.encoded(), values_bit_width, &mut ctx) @@ -388,7 +389,7 @@ fn bench_dict_bp_codes_alp_for_bp_values_dynanmic_dispatch(c: &mut Criterion) { len, |b, &n| { let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); @@ -663,7 +664,7 @@ fn bench_dict_bp_codes_alp_for_bp_values_composed_standalone(c: &mut Criterion) .map(|i| ::decode_single(10 + (i as i32 % 64), exponents)) .collect(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let float_prim = PrimitiveArray::new(Buffer::from(dict_floats), NonNullable); let alp = alp_encode(float_prim.as_view(), Some(exponents), &mut ctx).vortex_expect("alp_encode"); @@ -673,7 +674,7 @@ fn bench_dict_bp_codes_alp_for_bp_values_composed_standalone(c: &mut Criterion) .clone() .execute::(&mut ctx) .vortex_expect("to primitive"); - let for_arr = FoRData::encode(alp_encoded.clone()).vortex_expect("for encode"); + let for_arr = FoRData::encode(alp_encoded.clone(), &mut ctx).vortex_expect("for encode"); let bp = BitPackedData::encode(for_arr.encoded(), values_bit_width, &mut ctx) .vortex_expect("bitpack values"); let values_bp = bp; @@ -702,7 +703,7 @@ fn bench_dict_bp_codes_alp_for_bp_values_composed_standalone(c: &mut Criterion) &(values_bp, values_reference, codes_bp), |b, (values_bp, values_reference, codes_bp)| { b.iter_custom(|iters| { - let session = VortexSession::empty(); + let session = cuda_session(); let cuda_session = session.cuda_session(); let mut cuda_ctx = CudaSession::create_execution_ctx(&session) .vortex_expect("ctx") @@ -735,7 +736,7 @@ fn bench_dict_bp_codes_alp_for_bp_values_composed_standalone(c: &mut Criterion) // Benchmark: ALP(FoR(BitPacked)) — f64 // --------------------------------------------------------------------------- fn bench_alp_for_bitpacked_f64(c: &mut Criterion) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut group = c.benchmark_group("cuda"); let exponents = Exponents { e: 2, f: 0 }; @@ -759,6 +760,7 @@ fn bench_alp_for_bitpacked_f64(c: &mut Criterion) { .clone() .execute::(&mut ctx) .vortex_expect("to primitive"), + &mut ctx, ) .vortex_expect("for encode"); let bp = BitPackedData::encode(for_arr.encoded(), bit_width, &mut ctx) @@ -778,7 +780,7 @@ fn bench_alp_for_bitpacked_f64(c: &mut Criterion) { len, |b, &n| { let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); @@ -810,7 +812,7 @@ fn bench_dict_bp_codes_bp_for_values(c: &mut Criterion) { // Dict values: residuals 0..63 bitpacked, FoR adds 1_000_000 let dict_residuals: Vec = (0..dict_size as u32).collect(); let dict_prim = PrimitiveArray::new(Buffer::from(dict_residuals), NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dict_bp = BitPackedData::encode(&dict_prim.into_array(), dict_bit_width, &mut ctx) .vortex_expect("bitpack dict"); let dict_for = @@ -835,7 +837,7 @@ fn bench_dict_bp_codes_bp_for_values(c: &mut Criterion) { len, |b, &n| { let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); @@ -857,7 +859,7 @@ fn bench_dict_bp_codes_bp_for_values(c: &mut Criterion) { // Benchmark: ALP(FoR(BitPacked)) for f32 // --------------------------------------------------------------------------- fn bench_alp_for_bitpacked(c: &mut Criterion) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut group = c.benchmark_group("cuda"); let exponents = Exponents { e: 2, f: 0 }; @@ -881,6 +883,7 @@ fn bench_alp_for_bitpacked(c: &mut Criterion) { .clone() .execute::(&mut ctx) .vortex_expect("to primitive"), + &mut ctx, ) .vortex_expect("for encode"); let bp = BitPackedData::encode(for_arr.encoded(), bit_width, &mut ctx) @@ -900,7 +903,7 @@ fn bench_alp_for_bitpacked(c: &mut Criterion) { len, |b, &n| { let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); @@ -935,7 +938,7 @@ fn bench_dict_bp_u8_codes_u32_values(c: &mut Criterion) { let codes: Vec = (0..*len).map(|i| (i % dict_size) as u8).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let codes_bp = BitPackedData::encode(&codes_prim.into_array(), bit_width, &mut ctx) .vortex_expect("bitpack u8 codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); @@ -947,7 +950,7 @@ fn bench_dict_bp_u8_codes_u32_values(c: &mut Criterion) { len, |b, &n| { let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); @@ -978,7 +981,7 @@ fn bench_dict_bp_u16_codes_u32_values(c: &mut Criterion) { let codes: Vec = (0..*len).map(|i| (i % dict_size) as u16).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let codes_bp = BitPackedData::encode(&codes_prim.into_array(), bit_width, &mut ctx) .vortex_expect("bitpack u16 codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); @@ -990,7 +993,7 @@ fn bench_dict_bp_u16_codes_u32_values(c: &mut Criterion) { len, |b, &n| { let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); @@ -1021,7 +1024,7 @@ fn bench_dict_bp_u32_codes_u32_values(c: &mut Criterion) { let codes: Vec = (0..*len).map(|i| (i % dict_size) as u32).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let codes_bp = BitPackedData::encode(&codes_prim.into_array(), bit_width, &mut ctx) .vortex_expect("bitpack u32 codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); @@ -1033,7 +1036,7 @@ fn bench_dict_bp_u32_codes_u32_values(c: &mut Criterion) { len, |b, &n| { let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); let bench_runner = BenchRunner::::new(&array, n, &mut cuda_ctx); diff --git a/vortex-cuda/benches/filter_cuda.rs b/vortex-cuda/benches/filter_cuda.rs index f2d9cb74259..a9f55707a7c 100644 --- a/vortex-cuda/benches/filter_cuda.rs +++ b/vortex-cuda/benches/filter_cuda.rs @@ -27,7 +27,6 @@ use futures::executor::block_on; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_err; -use vortex::session::VortexSession; use vortex_cub::filter::CubFilterable; use vortex_cub::filter::cudaStream_t; use vortex_cuda::CudaDeviceBuffer; @@ -165,7 +164,7 @@ where |b, (input_data, bitmask, true_count)| { b.iter_custom(|iters| { let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()) + CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) .vortex_expect("failed to create execution context"); let num_items = input_data.len() as i64; diff --git a/vortex-cuda/benches/for_cuda.rs b/vortex-cuda/benches/for_cuda.rs index 44e60858a48..31acc764141 100644 --- a/vortex-cuda/benches/for_cuda.rs +++ b/vortex-cuda/benches/for_cuda.rs @@ -21,8 +21,8 @@ use criterion::Throughput; use cudarc::driver::DeviceRepr; use futures::executor::block_on; use vortex::array::IntoArray; -use vortex::array::LEGACY_SESSION; use vortex::array::VortexSessionExecute; +use vortex::array::array_session; use vortex::array::arrays::PrimitiveArray; use vortex::array::validity::Validity; use vortex::buffer::Buffer; @@ -33,7 +33,6 @@ use vortex::encodings::fastlanes::FoR; use vortex::encodings::fastlanes::FoRArray; use vortex::error::VortexExpect; use vortex::scalar::Scalar; -use vortex::session::VortexSession; use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; use vortex_cuda::executor::CudaArrayExt; @@ -59,7 +58,7 @@ where PrimitiveArray::new(Buffer::from(data), Validity::NonNullable).into_array(); if bp && T::PTYPE != PType::U8 { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let child = BitPackedData::encode(&primitive_array, 8, &mut ctx).vortex_expect("failed to bitpack"); FoR::try_new(child.into_array(), reference.into()) @@ -90,10 +89,11 @@ where let timed = TimedLaunchStrategy::default(); let timer = timed.timer(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context") - .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) - .with_launch_strategy(Arc::new(timed)); + 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(for_array.clone().into_array().execute_cuda(&mut cuda_ctx)) @@ -129,10 +129,11 @@ where let timed = TimedLaunchStrategy::default(); let timer = timed.timer(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context") - .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) - .with_launch_strategy(Arc::new(timed)); + 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(for_array.clone().into_array().execute_cuda(&mut cuda_ctx)) diff --git a/vortex-cuda/benches/fsst_cuda.rs b/vortex-cuda/benches/fsst_cuda.rs index c7aeed28151..21db8744919 100644 --- a/vortex-cuda/benches/fsst_cuda.rs +++ b/vortex-cuda/benches/fsst_cuda.rs @@ -23,7 +23,6 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::match_each_integer_ptype; use vortex::encodings::fsst::FSSTArrayExt; use vortex::error::VortexExpect; -use vortex::session::VortexSession; use vortex_cuda::CudaSession; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda_macros::cuda_available; @@ -41,7 +40,7 @@ fn benchmark_fsst_cuda_decompress(c: &mut Criterion) { let mut group = c.benchmark_group("cuda"); for &(n, len_str) in BENCH_SIZES { - let mut setup_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut setup_ctx = CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) .vortex_expect("failed to create execution context"); let fsst = make_fsst_clickbench_urls(n, setup_ctx.execution_ctx()); @@ -66,9 +65,10 @@ fn benchmark_fsst_cuda_decompress(c: &mut Criterion) { let timed = TimedLaunchStrategy::default(); let timer = timed.timer(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context") - .with_launch_strategy(Arc::new(timed)); + let mut cuda_ctx = + CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) + .vortex_expect("failed to create execution context") + .with_launch_strategy(Arc::new(timed)); for _ in 0..iters { block_on(fsst_array.clone().execute_cuda(&mut cuda_ctx)).unwrap(); diff --git a/vortex-cuda/benches/list_view_cuda.rs b/vortex-cuda/benches/list_view_cuda.rs index 0904844b900..4ce365210d6 100644 --- a/vortex-cuda/benches/list_view_cuda.rs +++ b/vortex-cuda/benches/list_view_cuda.rs @@ -25,7 +25,6 @@ use vortex::array::validity::Validity; use vortex::dtype::PType; use vortex::error::VortexExpect; use vortex::error::VortexResult; -use vortex::session::VortexSession; use vortex_cuda::CudaExecutionCtx; use vortex_cuda::CudaSession; use vortex_cuda::arrow::ArrowDeviceArray; @@ -90,9 +89,10 @@ fn benchmark_list_view_export(c: &mut Criterion) { let timed = TimedLaunchStrategy::default(); let timer = timed.timer(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context") - .with_launch_strategy(Arc::new(timed)); + let mut cuda_ctx = + CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) + .vortex_expect("failed to create execution context") + .with_launch_strategy(Arc::new(timed)); let array = block_on(contiguous_list_view(len, &mut cuda_ctx)) .vortex_expect("failed to create list-view fixture"); @@ -120,9 +120,10 @@ fn benchmark_list_view_export(c: &mut Criterion) { let timed = TimedLaunchStrategy::default(); let timer = timed.timer(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context") - .with_launch_strategy(Arc::new(timed)); + let mut cuda_ctx = + CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) + .vortex_expect("failed to create execution context") + .with_launch_strategy(Arc::new(timed)); let array = block_on(non_contiguous_primitive_list_view(len, &mut cuda_ctx)) .vortex_expect("failed to create list-view fixture"); diff --git a/vortex-cuda/benches/runend_cuda.rs b/vortex-cuda/benches/runend_cuda.rs index ceb79115604..194923654e0 100644 --- a/vortex-cuda/benches/runend_cuda.rs +++ b/vortex-cuda/benches/runend_cuda.rs @@ -27,7 +27,6 @@ use vortex::buffer::Buffer; use vortex::dtype::NativePType; use vortex::encodings::runend::RunEnd; use vortex::encodings::runend::RunEndArray; -use vortex::session::VortexSession; use vortex_cuda::CudaSession; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda_macros::cuda_available; @@ -75,7 +74,8 @@ where group.throughput(Throughput::Bytes((len * size_of::()) as u64)); for run_len in [10, 1000, 100000] { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()).unwrap(); + let mut cuda_ctx = + CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()).unwrap(); let runend_array = make_runend_array_typed::(len, run_len, cuda_ctx.execution_ctx()); group.bench_with_input( @@ -87,7 +87,7 @@ where let timer = timed.timer(); let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()) + CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) .unwrap() .with_launch_strategy(Arc::new(timed)); diff --git a/vortex-cuda/benches/throughput_cuda.rs b/vortex-cuda/benches/throughput_cuda.rs index bcbaf8058aa..d40b41c5c41 100644 --- a/vortex-cuda/benches/throughput_cuda.rs +++ b/vortex-cuda/benches/throughput_cuda.rs @@ -21,7 +21,6 @@ use cudarc::driver::DevicePtr; use cudarc::driver::DevicePtrMut; use cudarc::driver::sys; use cudarc::driver::sys::CUevent_flags; -use vortex::session::VortexSession; use vortex_cuda::CudaExecutionCtx; use vortex_cuda::CudaSession; use vortex_cuda_macros::cuda_available; @@ -107,7 +106,7 @@ fn benchmark_transfer_throughput(c: &mut Criterion) { &(input_bytes, output_bytes), |b, &(in_bytes, out_bytes)| { b.iter_custom(|iters| { - let session = VortexSession::empty(); + let session = vortex_cuda::cuda_session(); let mut in_ctx = CudaSession::create_execution_ctx(&session).unwrap(); let mut out_ctx = CudaSession::create_execution_ctx(&session).unwrap(); diff --git a/vortex-cuda/benches/zstd_cuda.rs b/vortex-cuda/benches/zstd_cuda.rs index 797e8b76ec8..c0f9afaadd1 100644 --- a/vortex-cuda/benches/zstd_cuda.rs +++ b/vortex-cuda/benches/zstd_cuda.rs @@ -19,7 +19,6 @@ use vortex::encodings::zstd::ZstdDataParts; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_err; -use vortex::session::VortexSession; use vortex_cuda::CudaSession; use vortex_cuda::ZstdKernelPrep; use vortex_cuda::nvcomp::zstd as nvcomp_zstd; @@ -126,7 +125,7 @@ fn benchmark_zstd_cuda_decompress(c: &mut Criterion) { let mut group = c.benchmark_group("cuda"); for (num_strings, label) in BENCH_SIZES { - let mut setup_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut setup_ctx = CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) .vortex_expect("failed to create execution context"); let (zstd_array, uncompressed_size) = make_zstd_array(*num_strings, &mut setup_ctx) .vortex_expect("failed to create ZSTD array"); @@ -137,8 +136,9 @@ fn benchmark_zstd_cuda_decompress(c: &mut Criterion) { &zstd_array, |b, zstd_array| { b.iter_custom(|iters| { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context"); + let mut cuda_ctx = + CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) + .vortex_expect("failed to create execution context"); let mut total_time = Duration::ZERO; diff --git a/vortex-cuda/build.rs b/vortex-cuda/build.rs index e6024d3c275..e201345cd73 100644 --- a/vortex-cuda/build.rs +++ b/vortex-cuda/build.rs @@ -8,6 +8,7 @@ use std::env; use std::fs::File; use std::io; +use std::io::Write; use std::path::Path; use std::path::PathBuf; use std::process::Command; @@ -32,15 +33,6 @@ fn main() { std::fs::create_dir_all(&kernels_gen).expect("Failed to create kernels/gen directory"); - // Always emit the kernels output directory path as a compile-time env var so any binary - // linking against vortex-cuda can find the PTX files. This must be set regardless - // of CUDA availability since the code using env!() is always compiled. - // At runtime, VORTEX_CUDA_KERNELS_DIR can be set to override this path. - println!( - "cargo:rustc-env=VORTEX_CUDA_KERNELS_DIR={}", - kernels_gen.display() - ); - println!("cargo:rerun-if-env-changed=PROFILE"); // Regenerate bit_unpack kernels only when the generator changes @@ -60,6 +52,8 @@ fn main() { 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() { return; } @@ -97,6 +91,57 @@ fn main() { } } } + + // 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"); +} + +/// 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() { + writeln!( + file, + "pub(crate) fn embedded_ptx(_module_name: &str) -> Option<&'static str> {{" + )?; + writeln!(file, " None")?; + writeln!(file, "}}")?; + return Ok(()); + } + + writeln!( + file, + "pub(crate) fn embedded_ptx(module_name: &str) -> Option<&'static str> {{" + )?; + writeln!(file, " match module_name {{")?; + for path in ptx_files { + let Some(module_name) = path.file_stem().and_then(|stem| stem.to_str()) else { + continue; + }; + writeln!( + file, + " {module_name:?} => Some(include_str!({:?})),", + path.to_string_lossy() + )?; + } + writeln!(file, " _ => None,")?; + writeln!(file, " }}")?; + writeln!(file, "}}")?; + Ok(()) } fn generate_unpack(output_dir: &Path, thread_count: usize) -> io::Result { @@ -206,6 +251,7 @@ fn generate_arrow_device_array_bindings(manifest_dir: &Path, out_dir: &Path) { .header(header.to_string_lossy()) .allowlist_type("ArrowArray") .allowlist_type("ArrowDeviceArray") + .allowlist_type("ArrowDeviceArrayStream") .allowlist_type("ArrowDeviceType") .allowlist_var("ARROW_DEVICE_.*") // ArrowArray/ArrowDeviceArray own producer state through release/private_data. diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index ee0f0fb3ca3..0b06a660afa 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -43,6 +43,18 @@ struct ArrowDeviceArray { }; #endif +#if !defined(ARROW_C_DEVICE_STREAM_INTERFACE) && !defined(USE_OWN_ARROW_DEVICE) +#define ARROW_C_DEVICE_STREAM_INTERFACE +struct ArrowDeviceArrayStream { + ArrowDeviceType device_type; + int (*get_schema)(struct ArrowDeviceArrayStream *, struct ArrowSchema *out); + int (*get_next)(struct ArrowDeviceArrayStream *, struct ArrowDeviceArray *out); + const char *(*get_last_error)(struct ArrowDeviceArrayStream *); + void (*release)(struct ArrowDeviceArrayStream *); + void *private_data; +}; +#endif + /** * Create a CUDA Vortex session. * @@ -69,6 +81,25 @@ int vx_cuda_array_export_arrow_device(const vx_session *session, struct ArrowDeviceArray *out_array, vx_error **error_out); +/** + * Consume a Vortex partition and scan it as an Arrow C Device stream. + * + * This function takes ownership of `partition`. Callers must not free or reuse + * it after calling this function, regardless of success or failure. + * + * On success returns 0 and writes an owned `ArrowDeviceArrayStream` to + * `out_stream`. The stream owns the resulting scan iterator. The caller must + * release the stream through its embedded Arrow `release` callback, and must + * release each produced `ArrowDeviceArray` through its embedded + * `ArrowArray.release` callback. + * + * On error returns 1 and writes a `vx_error` to `error_out` when non-NULL. + */ +int vx_cuda_partition_scan_arrow_device_stream(const vx_session *session, + vx_partition *partition, + struct ArrowDeviceArrayStream *out_stream, + vx_error **error_out); + #ifdef __cplusplus } #endif diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 2a74cb2666c..611f644e535 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -18,11 +18,16 @@ use vortex::session::SessionExt; use vortex::session::VortexSession; use vortex_cuda::CudaSession; use vortex_cuda::arrow::ArrowDeviceArray; +use vortex_cuda::arrow::ArrowDeviceArrayStream; use vortex_cuda::arrow::DeviceArrayExt; +use vortex_cuda::arrow::DeviceArrayStreamExt; +use vortex_ffi::ffi_runtime; use vortex_ffi::try_or; use vortex_ffi::vx_array; use vortex_ffi::vx_array_ref; use vortex_ffi::vx_error; +use vortex_ffi::vx_partition; +use vortex_ffi::vx_partition_into_array_stream; use vortex_ffi::vx_session; use vortex_ffi::vx_session_new_with; use vortex_ffi::vx_session_ref; @@ -30,12 +35,13 @@ use vortex_ffi::vx_session_ref; const VX_CUDA_OK: c_int = 0; const VX_CUDA_ERR: c_int = 1; +/// Return a Vortex session with a [`CudaSession`] session variable. +/// +/// If `session` already has CUDA support, this returns a clone of it. Otherwise it +/// returns a new session cloned from `session` with a default [`CudaSession`] attached. fn session_with_cuda(session: &VortexSession) -> VortexResult { - if session.get_opt::().is_some() { - return Ok(session.clone()); - } - - Ok(session.clone().with_some(CudaSession::try_default()?)) + session.get::(); + Ok(session.clone()) } /// Create a CUDA Vortex session. @@ -101,6 +107,46 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_export_arrow_device( }) } +/// Consume a Vortex partition and scan it as an Arrow C Device stream. +/// +/// This function takes ownership of `partition`. Callers must not free or reuse it after calling +/// this function, regardless of success or failure. +/// +/// On success returns `0` and writes an owned `ArrowDeviceArrayStream` to `out_stream`. The stream +/// owns the resulting scan iterator. The caller must release the stream through its embedded Arrow +/// `release` callback, and must release each produced `ArrowDeviceArray` through its embedded +/// `ArrowArray.release` callback. +/// +/// On error returns `1` and, when `error_out` is non-null, writes a `vx_error` (free with +/// `vx_error_free`). +/// +/// # Safety +/// +/// `session` must be a valid borrowed handle created by `vortex-ffi`. `partition` must be an owned +/// partition handle created by `vortex-ffi`. `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_partition_scan_arrow_device_stream( + session: *const vx_session, + partition: *mut vx_partition, + out_stream: *mut ArrowDeviceArrayStream, + error_out: *mut *mut vx_error, +) -> c_int { + try_or(error_out, VX_CUDA_ERR, || { + vortex_ensure!(!partition.is_null(), "null vx_partition"); + + let array_stream = unsafe { vx_partition_into_array_stream(partition) }?; + vortex_ensure!(!out_stream.is_null(), "null ArrowDeviceArrayStream output"); + + let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?; + // Drive the stream on the same runtime the partition's scan spawned its work onto. + let device_stream = array_stream.export_device_array_stream(&session, ffi_runtime())?; + + unsafe { ptr::write(out_stream, device_stream) }; + Ok(VX_CUDA_OK) + }) +} + #[cfg(test)] mod tests { use std::ptr; diff --git a/vortex-cuda/kernels/src/arrow_validity.cu b/vortex-cuda/kernels/src/arrow_validity.cu index 906054c86ee..d139fa66f25 100644 --- a/vortex-cuda/kernels/src/arrow_validity.cu +++ b/vortex-cuda/kernels/src/arrow_validity.cu @@ -7,8 +7,9 @@ namespace { -// Load the `word_idx`-th little-endian u64 of `input`, treating bytes outside -// `[0, input_bytes)` as zero. `input` must be 8-byte aligned. +// Transform up to 8 input bytes into a zero-extended 64-bit word: +// +// [ b0 ][ b1 ][ b2 ] | end -> [ b0 ][ b1 ][ b2 ][ 00 ][ 00 ][ 00 ][ 00 ][ 00 ] __device__ uint64_t load_input_word(const uint8_t *const input, int64_t word_idx, uint64_t input_bytes) { if (word_idx < 0) { return 0; @@ -28,11 +29,17 @@ __device__ uint64_t load_input_word(const uint8_t *const input, int64_t word_idx return word; } -// Build one 64-bit word of the Arrow validity bitmap. +// Build one output word for sliced validity. The row bits are the same, but +// row 0 may live at a different bit position in the source and Arrow bitmaps. +// For example, `input_offset = 5` and `arrow_offset = 0` shifts row0 from bit 5 +// in the input bitmap to bit 0 in the Arrow bitmap. +// +// input bitmap: [ . ][ . ][ . ][ . ][ . ][ row0 ][ row1 ][ row2 ].... +// ^ input_offset +// Arrow bitmap: [ row0 ][ row1 ][ row2 ].... +// ^ arrow_offset // -// Output bit `b` for `b` in `[arrow_offset, validity_bits)` equals input bit `b + shift`; -// all other bits are zero. Two adjacent input words are funnel-shifted to align the input -// bits with the output word, then the leading/trailing edges are masked. +// Padding bits are cleared so word-sized validity readers can safely over-read. __device__ uint64_t repack_word(const uint8_t *const input, uint64_t word_idx, int64_t shift, @@ -56,55 +63,129 @@ __device__ uint64_t repack_word(const uint8_t *const input, return 0; } - // `>> 6` floors also for negative bit positions, unlike `/ 64` which truncates toward zero. - const int64_t input_bit = static_cast(word_start) + shift; - const int64_t input_word = input_bit >> 6; - const uint32_t bit = static_cast(input_bit & 63); + // Each output bit `b` reads source bit `b + shift`. + // `>> 6` floors for negative positions, unlike `/ 64` which truncates toward zero. + const int64_t source_bit_start = static_cast(word_start) + shift; + const int64_t source_word = source_bit_start >> 6; + const uint32_t source_bit = static_cast(source_bit_start & 63); - const uint64_t lo = load_input_word(input, input_word, input_bytes); - if (bit == 0) { + const uint64_t lo = load_input_word(input, source_word, input_bytes); + if (source_bit == 0) { return lo & mask; } - const uint64_t hi = load_input_word(input, input_word + 1, input_bytes); - return ((lo >> bit) | (hi << (64 - bit))) & mask; + const uint64_t hi = load_input_word(input, source_word + 1, input_bytes); + return ((lo >> source_bit) | (hi << (64 - source_bit))) & mask; } -// Rebuild a possibly bit-offset Vortex validity bitmap into an Arrow-compatible bitmap. +constexpr uint32_t WARP_SIZE = 32; +constexpr uint32_t FULL_WARP_MASK = 0xffffffff; + +// First reduction step for the count kernel: sum one value per lane so each +// warp produces a single partial count. // -// `input_offset` is the bit offset into `input`; `arrow_offset` is the logical Arrow array offset -// to preserve in the output. Bits outside `[arrow_offset, arrow_offset + len)` are left unset. -// The output allocation must hold `ceil((len + arrow_offset) / 64)` full 64-bit words; every -// word is written, so no zero-initialization of the output is required. -__device__ void arrow_validity_repack_device(const uint8_t *const input, - uint64_t *const output, - uint64_t len, - uint64_t input_offset, +// lanes: [a][b][c][d]... -> lane 0: a+b+c+d+... +__device__ uint64_t warp_sum(uint64_t value) { + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(FULL_WARP_MASK, value, offset); + } + return value; +} + +// Mask one bitmap byte down to actual rows. This keeps null counting from +// including Arrow offset padding or trailing padding bits. +// +// byte bits: [ pad ][ row ][ row ][ row ][ pad ] +// mask: [ 0 ][ 1 ][ 1 ][ 1 ][ 0 ] +__device__ uint32_t arrow_validity_byte_mask(uint64_t byte_idx, uint64_t arrow_offset, - uint64_t input_bytes) { - // One worker owns a contiguous range of output words. Each word is rebuilt locally so - // there are no cross-thread bit writes or atomics. - const uint64_t worker = blockIdx.x * blockDim.x + threadIdx.x; - const uint64_t validity_bits = len + arrow_offset; - const uint64_t output_words = (validity_bits + 63) / 64; - const uint64_t stride = static_cast(gridDim.x) * blockDim.x; + uint64_t validity_bits) { + const uint64_t byte_start = byte_idx * 8; - // Translate Arrow-visible output bits back to source bitmap bits. The source bitmap may - // start at any bit offset, while Arrow's buffer pointer is byte-addressed. - const int64_t shift = static_cast(input_offset) - static_cast(arrow_offset); + uint32_t mask = 0xff; + if (byte_start < arrow_offset) { + const uint64_t lead = arrow_offset - byte_start; + mask = lead >= 8 ? 0 : mask << lead; + } - for (uint64_t word_idx = worker; word_idx < output_words; word_idx += stride) { - output[word_idx] = repack_word(input, word_idx, shift, arrow_offset, validity_bits, input_bytes); + const uint64_t remaining = validity_bits - byte_start; + if (remaining < 8) { + mask &= (uint32_t {1} << remaining) - 1; + } + return mask; +} + +// Combine warp partial counts into one block total. Only thread 0 returns a +// non-zero value so the count kernel does one global atomic per block. +// +// per-thread counts -> per-warp sums -> block sum -> atomicAdd +__device__ uint64_t block_sum_to_thread_zero(uint64_t value, uint64_t *const warp_counts) { + const uint32_t thread = threadIdx.x; + const uint32_t lane = thread & (WARP_SIZE - 1); + const uint32_t warp = thread / WARP_SIZE; + const uint32_t block_warps = (blockDim.x + WARP_SIZE - 1) / WARP_SIZE; + + value = warp_sum(value); + if (lane == 0) { + warp_counts[warp] = value; } + __syncthreads(); + + value = lane < block_warps ? warp_counts[lane] : 0; + value = warp == 0 ? warp_sum(value) : 0; + return thread == 0 ? value : 0; } } // namespace -// CUDA entry point for validity bitmap repacking used by Arrow Device export. +// Repack sliced validity when the source bitmap offset does not match the +// Arrow array offset. Each thread writes independent output words. +// +// thread 0 -> output word 0, word N, ... +// thread 1 -> output word 1, word N+1, ... extern "C" __global__ void arrow_validity_repack(const uint8_t *const input, uint64_t *const output, uint64_t len, uint64_t input_offset, uint64_t arrow_offset, uint64_t input_bytes) { - arrow_validity_repack_device(input, output, len, input_offset, arrow_offset, input_bytes); + const uint64_t worker = blockIdx.x * blockDim.x + threadIdx.x; + const uint64_t validity_bits = len + arrow_offset; + const uint64_t output_words = (validity_bits + 63) / 64; + const uint64_t stride = static_cast(gridDim.x) * blockDim.x; + const int64_t shift = static_cast(input_offset) - static_cast(arrow_offset); + + for (uint64_t word_idx = worker; word_idx < output_words; word_idx += stride) { + output[word_idx] = repack_word(input, word_idx, shift, arrow_offset, validity_bits, input_bytes); + } +} + +// Count valid rows directly from the device bitmap so Arrow export can provide +// an exact null_count without copying validity to the CPU. +// +// bytes -> mask padding -> popcount -> block sum -> global count +extern "C" __global__ void arrow_validity_count_valid(const uint8_t *const input, + uint64_t *const output, + uint64_t len, + uint64_t arrow_offset) { + __shared__ uint64_t warp_counts[WARP_SIZE]; + + const uint64_t validity_bits = len + arrow_offset; + const uint64_t input_bytes = (validity_bits + 7) / 8; + const uint64_t worker = blockIdx.x * blockDim.x + threadIdx.x; + const uint64_t stride = static_cast(gridDim.x) * blockDim.x; + + // Grid-stride over bitmap bytes. Each byte contributes the popcount of only + // row bits; leading Arrow offset bits and trailing padding bits are masked out. + uint64_t valid_count = 0; + for (uint64_t byte_idx = worker; byte_idx < input_bytes; byte_idx += stride) { + const uint32_t mask = arrow_validity_byte_mask(byte_idx, arrow_offset, validity_bits); + valid_count += __popc(static_cast(input[byte_idx]) & mask); + } + + // Reduce within the block first so global contention is one atomic add per block. + valid_count = block_sum_to_thread_zero(valid_count, warp_counts); + if (threadIdx.x == 0) { + atomicAdd(reinterpret_cast(output), + static_cast(valid_count)); + } } diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index 6c0f53ef249..109e343e424 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -11,13 +11,13 @@ use cudarc::driver::CudaSlice; use cudarc::driver::DeviceRepr; use cudarc::driver::LaunchConfig; use cudarc::driver::PushKernelArg; -use cudarc::driver::result as cuda_driver; use futures::future::BoxFuture; use futures::future::join; use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; +use vortex::array::arrays::Bool; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::Dict; use vortex::array::arrays::DictArray; @@ -59,7 +59,6 @@ use vortex::error::vortex_bail; use vortex::error::vortex_ensure; use vortex::error::vortex_err; use vortex::extension::datetime::AnyTemporal; -use vortex::mask::Mask; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; @@ -77,6 +76,7 @@ use crate::arrow::cuda_decimal_value_type; use crate::arrow::list_view::export_device_list_view; use crate::cub::exclusive_sum_i32; use crate::executor::CudaArrayExt; +use crate::executor::execute_validity_cuda; /// 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 @@ -773,35 +773,52 @@ pub(super) async fn export_arrow_validity_buffer( arrow_offset: usize, ctx: &mut CudaExecutionCtx, ) -> VortexResult<(Option, i64)> { - let mask = validity.execute_mask(len, ctx.execution_ctx())?; - let null_count = i64::try_from(mask.false_count())?; - let validity_bits = len + arrow_offset; - let validity_bytes = validity_bits.div_ceil(8); - - let validity_buffer = match mask { - Mask::AllTrue(_) => return Ok((None, 0)), - Mask::AllFalse(_) => device_zeroed_byte_buffer(validity_bytes, ctx)?, - Mask::Values(values) => { - // Shrinking the offset below 8 bounds the host-to-device copy to the slice's bytes - // instead of the whole backing bitmap. - let bits = values.bit_buffer().clone().shrink_offset(); - if arrow_offset == 0 && bits.offset() == 0 { - // Fast path: the Vortex bitmap already matches Arrow's byte-addressed layout. - let (_, _, buffer) = bits.into_inner(); - ctx.ensure_on_device(BufferHandle::new_host(buffer)).await? + // Validity is exported separately from the array data. Decode it here so Arrow + // gets a device-resident validity buffer alongside the array it belongs to. + let validity = execute_validity_cuda(validity, len, ctx).await?; + match validity { + Validity::NonNullable | Validity::AllValid => Ok((None, 0)), + // For non-Null Arrow layouts, callers still export the normal value buffers. + // This only marks every row null via buffer 0, the validity bitmap. + Validity::AllInvalid => Ok(( + Some(device_zeroed_byte_buffer( + validity_bitmap_byte_len(len, arrow_offset)?, + ctx, + )?), + i64::try_from(len)?, + )), + Validity::Array(array) => { + let array = array.try_downcast::().map_err(|array| { + vortex_err!( + "canonical validity array must be bool, got {}", + array.dtype() + ) + })?; + let BoolDataParts { bits, meta } = array.into_data().into_parts(len); + let bitmap = ctx.ensure_on_device(bits).await?; + // ArrowDeviceArray uses ArrowArray layout with its buffers being device pointers. + // + // Validity is one bit per row, addressed via the Arrow array offset. Reuse the bitmap + // when Vortex's validity offset already matches Arrow's; otherwise repack on the GPU + // so row i is at Arrow bit `arrow_offset + i`. + let bitmap = if meta.offset() == arrow_offset { + bitmap } else { - // Slow path: bit offsets cannot be represented by the Arrow buffer pointer. - // Repack on the GPU so compact/sliced exports keep Arrow offset semantics. - let (input_offset, _, input_buffer) = bits.into_inner(); - let input_buffer = ctx - .ensure_on_device(BufferHandle::new_host(input_buffer)) - .await?; - repack_arrow_validity_buffer(&input_buffer, input_offset, len, arrow_offset, ctx)? - } + repack_arrow_validity_buffer(&bitmap, meta.offset(), len, arrow_offset, ctx)? + }; + // Keep nullable exports self-describing for consumers that require exact null counts. + let null_count = count_arrow_validity_nulls(&bitmap, len, arrow_offset, ctx)?; + Ok((Some(bitmap), null_count)) } - }; + } +} - Ok((Some(validity_buffer), null_count)) +/// Return the byte length needed for `len` validity bits at the given bit offset. +fn validity_bitmap_byte_len(len: usize, arrow_offset: usize) -> VortexResult { + Ok(len + .checked_add(arrow_offset) + .ok_or_else(|| vortex_err!("Arrow validity bit length overflows usize"))? + .div_ceil(8)) } /// Allocate a zeroed device buffer with cuDF-safe padding for Arrow validity masks. @@ -817,24 +834,77 @@ fn device_zeroed_byte_buffer( Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(buffer))).slice(0..byte_len)) } +pub fn count_arrow_validity_nulls( + bitmap: &BufferHandle, + len: usize, + arrow_offset: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + if len == 0 { + return Ok(0); + } + + let expected_bytes = validity_bitmap_byte_len(len, arrow_offset)?; + vortex_ensure!( + bitmap.len() >= expected_bytes, + "Arrow validity bitmap has {} bytes, expected at least {expected_bytes}", + bitmap.len() + ); + + let mut count = ctx.device_alloc::(1)?; + 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)?; + + let kernel = ctx.load_function_with_suffixes("arrow_validity", &["count_valid"])?; + const COUNT_THREADS_PER_BLOCK: u32 = 256; + const MAX_COUNT_BLOCKS: u32 = 4096; + let num_blocks = u32::try_from(expected_bytes.div_ceil(COUNT_THREADS_PER_BLOCK as usize))? + .clamp(1, MAX_COUNT_BLOCKS); + let config = LaunchConfig { + grid_dim: (num_blocks, 1, 1), + block_dim: (COUNT_THREADS_PER_BLOCK, 1, 1), + shared_mem_bytes: 0, + }; + ctx.launch_kernel_config(&kernel, config, expected_bytes, |args| { + args.arg(&input_view) + .arg(&output_view) + .arg(&len) + .arg(&arrow_offset); + })?; + + let valid_count = ctx + .stream() + .clone_dtoh(&output_view) + .map_err(|err| vortex_err!("Failed to copy Arrow validity count to host: {err}"))? + .into_iter() + .next() + .ok_or_else(|| vortex_err!("Arrow validity count kernel returned no output"))?; + + Ok(i64::try_from(len - valid_count)?) +} + /// Repack a validity bitmap into Arrow layout without copying bitmap bits back to the CPU. /// /// Vortex bitmaps may start at any bit offset. Arrow exposes only a byte-addressed validity buffer /// plus an array offset, so sliced compact exports need a GPU rewrite when either side has a /// bit-level offset. The kernel writes the output one 64-bit word at a time, funnel-shifting two /// adjacent input words, so the allocation is padded to whole words (zeroed by the edge masks). -pub(super) fn repack_arrow_validity_buffer( +pub fn repack_arrow_validity_buffer( input_buffer: &BufferHandle, input_offset: usize, len: usize, arrow_offset: usize, ctx: &mut CudaExecutionCtx, ) -> VortexResult { - let validity_bits = len - .checked_add(arrow_offset) - .ok_or_else(|| vortex_err!("Arrow validity bit length overflows usize"))?; - let output_bytes = validity_bits.div_ceil(8); - let output_words = validity_bits.div_ceil(u64::BITS as usize); + let output_bytes = validity_bitmap_byte_len(len, arrow_offset)?; + let output_words = output_bytes.div_ceil(size_of::()); // The kernel loads the input bitmap as 64-bit words. if !input_buffer @@ -1012,7 +1082,7 @@ async fn export_fixed_size_list( } = array.into_data_parts(); let (validity_buffer, null_count) = export_arrow_validity_buffer(validity, len, 0, ctx).await?; - let offsets_buffer = fixed_size_list_offsets(len, list_size, ctx).await?; + let offsets_buffer = fixed_size_list_offsets(len, list_size, ctx)?; export_list_layout( elements, @@ -1026,7 +1096,7 @@ async fn export_fixed_size_list( .await } -async fn fixed_size_list_offsets( +fn fixed_size_list_offsets( len: usize, list_size: u32, ctx: &mut CudaExecutionCtx, @@ -1036,17 +1106,29 @@ async fn fixed_size_list_offsets( "cannot export FixedSizeList with list size {list_size}: Arrow List offsets require i32" ) })?; - let offsets = (0..=i32::try_from(len)?) - .map(|idx| { - idx.checked_mul(list_size) - .ok_or_else(|| vortex_err!("FixedSizeList Arrow List offsets exceed i32 range")) - }) - .collect::>>()?; + let len_i32 = i32::try_from(len)?; + len_i32 + .checked_mul(list_size) + .ok_or_else(|| vortex_err!("FixedSizeList Arrow List offsets exceed i32 range"))?; + + 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 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) + .arg(&base) + .arg(&list_size) + .arg(&output_len_u64); + })?; - ctx.ensure_on_device(BufferHandle::new_host( - Buffer::from(offsets).into_byte_buffer(), - )) - .await + Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( + offsets, + )))) } /// Return Arrow Device `List` offsets as an `i32` device buffer. @@ -1147,9 +1229,9 @@ fn export_fixed_size( } unsafe extern "C" fn release_array(array: *mut ArrowArray) { - // SAFETY: this is only safe if we're dropping an ArrowArray that was created from Rust - // code. This is necessary to ensure that the fields inside the CudaPrivateData - // get dropped to free native/GPU memory. + // SAFETY: this is only safe if we're dropping an ArrowArray that was + // created from Rust code. This is necessary to ensure that the fields + // inside the CudaPrivateData get dropped to free native/GPU memory. unsafe { if array.is_null() || (*array).release.is_none() { return; @@ -1159,13 +1241,11 @@ unsafe extern "C" fn release_array(array: *mut ArrowArray) { if !private_data_ptr.is_null() { let mut private_data = Box::from_raw(private_data_ptr.cast::()); - // Release may run on a foreign thread; bind this array's context before synchronizing - // so async frees cannot race consumer-side reads. - let cuda_context = Arc::clone(private_data.cuda_stream.context()); - match cuda_context.bind_to_thread() { - Ok(()) => cuda_context.record_err(cuda_driver::ctx::synchronize()), - Err(err) => cuda_context.record_err(Err::<(), _>(err)), - } + // Queue a non-blocking wait so CudaSlice::Drop frees after export work. + let wait_result = private_data.cuda_stream.wait(&private_data.export_event); + // Arrow release callbacks cannot return errors, so record wait-enqueue + // failures for later CUDA operations to observe on this context. + private_data.cuda_stream.context().record_err(wait_result); release_children(&mut private_data); release_dictionary(&mut private_data); } @@ -1250,7 +1330,6 @@ mod tests { use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::extension::datetime::TimeUnit; - use vortex::session::VortexSession; use crate::CudaExecutionCtx; use crate::arrow::ARROW_DEVICE_CUDA; @@ -1585,7 +1664,7 @@ mod tests { #[case] expected_len: i64, #[case] expected_data_type: DataType, ) -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let mut exported = array.export_device_array_with_schema(&mut ctx).await?; @@ -1606,7 +1685,7 @@ mod tests { #[crate::test] async fn test_export_null() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = NullArray::new(7).into_array(); @@ -1622,7 +1701,7 @@ mod tests { #[crate::test] async fn test_export_dictionary() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let out_of_line = "a dictionary value stored out-of-line"; @@ -1667,7 +1746,7 @@ mod tests { #[crate::test] async fn test_export_struct_preserves_dictionary_child() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let dictionary = DictArray::try_new( @@ -1709,7 +1788,7 @@ mod tests { #[crate::test] async fn test_export_dictionary_with_nullable_values() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = DictArray::try_new( @@ -1746,7 +1825,7 @@ mod tests { expected_data_type: DataType, expected_values: Vec, ) -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let mut exported = array.export_device_array_with_schema(&mut ctx).await?; @@ -1876,7 +1955,7 @@ mod tests { #[crate::test] async fn test_export_decimal_narrowing_errors() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = DecimalArray::from_iter([i256::from_parts(0, 1)], DecimalDType::new(38, 0)) .into_array(); @@ -1891,7 +1970,7 @@ mod tests { #[crate::test] async fn test_export_decimal_narrowing_from_arrow_import() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = DecimalArray::from_iter([0i128, 1, -2], DecimalDType::new(10, 2)).into_array(); @@ -1934,7 +2013,7 @@ mod tests { #[crate::test] async fn test_export_temporal() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = TemporalArray::new_date( @@ -1957,7 +2036,7 @@ mod tests { #[crate::test] async fn test_export_bool() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = BoolArray::from_iter([true, false, true]).into_array(); @@ -1976,7 +2055,7 @@ mod tests { #[crate::test] async fn test_export_varbinview() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let out_of_line = "this is a longer string for out-of-line storage"; @@ -1992,7 +2071,7 @@ mod tests { #[crate::test] async fn test_export_binary_inline_outline_values() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let out_of_line = b"this binary payload is longer than twelve bytes"; @@ -2023,7 +2102,7 @@ mod tests { #[crate::test] async fn test_export_binary_empty_and_all_null() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let empty = VarBinViewArray::from_iter_nullable_bin(std::iter::empty::>()) @@ -2049,7 +2128,7 @@ mod tests { #[crate::test] async fn test_export_binary_invalid_data_buffer_ref_errors() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let view = BinaryView::make_view(b"this references a missing data buffer", 0, 0); @@ -2075,7 +2154,7 @@ mod tests { #[crate::test] async fn test_export_binary_i32_offset_overflow_errors() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let view = BinaryView::new_ref(i32::MAX as u32 + 1, [0; 4], 0, 0); @@ -2101,7 +2180,7 @@ mod tests { #[crate::test] async fn test_export_list() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = ListArray::try_new( @@ -2140,7 +2219,7 @@ mod tests { #[crate::test] async fn test_export_host_contiguous_list_view() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = ListViewArray::new( @@ -2167,7 +2246,7 @@ mod tests { #[crate::test] async fn test_export_host_non_contiguous_nested_list_view_falls_back_to_cpu() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let elements = StructArray::new( @@ -2209,7 +2288,7 @@ mod tests { #[crate::test] async fn test_export_host_non_contiguous_dictionary_list_view_schema_matches_rebuilt_child() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let elements = DictArray::try_new( @@ -2267,7 +2346,7 @@ mod tests { #[crate::test] async fn test_export_host_large_lists_dictionary_list_view_schema_matches_rebuilt_child() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let elements = DictArray::try_new( @@ -2320,7 +2399,7 @@ mod tests { #[case] offsets_ptype: PType, #[case] sizes_ptype: PType, ) -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let elements = primitive_i32_on_device(0..5, &mut ctx).await?; @@ -2365,7 +2444,7 @@ mod tests { #[case] fixture: (ArrayRef, [usize; 2]), #[case] expected_data_type: DataType, ) -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let (array, expected_data_buffer_lengths) = fixture; @@ -2396,7 +2475,7 @@ mod tests { #[case::u64(PrimitiveArray::from_iter([0u64, 2, 2, 5]).into_array())] #[crate::test] async fn test_export_list_with_non_i32_offsets(#[case] offsets: ArrayRef) -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = ListArray::try_new( @@ -2428,7 +2507,7 @@ mod tests { #[case] offsets_ptype: PType, #[case] sizes_ptype: PType, ) -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let elements = primitive_i32_on_device([10, 11, 12, 13, 14], &mut ctx).await?; @@ -2457,7 +2536,7 @@ mod tests { #[crate::test] async fn test_export_device_non_contiguous_dictionary_list_view() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let codes = primitive_on_device([0u8, 1, 2, 0, 1], &mut ctx).await?; @@ -2505,7 +2584,7 @@ mod tests { #[crate::test] async fn test_export_device_non_contiguous_dictionary_list_view_nullable_values() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let codes = primitive_on_device([0u8, 1, 2, 0, 1], &mut ctx).await?; @@ -2534,7 +2613,7 @@ mod tests { #[crate::test] async fn test_export_device_non_contiguous_dictionary_list_view_nullable_codes_errors() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let codes = PrimitiveArray::from_option_iter([Some(0u8), None, Some(2), Some(0), Some(1)]); @@ -2572,7 +2651,7 @@ mod tests { #[case] sizes_values: &[i64], #[case] expected_error: &str, ) -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let elements = primitive_i32_on_device(0..4, &mut ctx).await?; @@ -2600,7 +2679,7 @@ mod tests { #[crate::test] async fn test_export_device_non_contiguous_nested_list_view_returns_error() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let field = primitive_i32_on_device(0..4, &mut ctx).await?; @@ -2634,7 +2713,7 @@ mod tests { #[crate::test] async fn test_export_device_non_contiguous_nullable_primitive_list_view_returns_error() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let elements = nullable_primitive_i32_on_device( @@ -2664,7 +2743,7 @@ mod tests { #[crate::test] async fn test_export_fixed_size_list_as_list() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = FixedSizeListArray::new( @@ -2709,7 +2788,7 @@ mod tests { // Check device metadata for data-bearing and metadata-only exports. #[crate::test] async fn test_export_device_metadata() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let expected_device_id = ctx.stream().context().ordinal() as i64; @@ -2731,7 +2810,7 @@ mod tests { // Check sliced arrays preserve the expected Arrow length/offset metadata. #[crate::test] async fn test_export_sliced_arrays() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let primitive = PrimitiveArray::from_iter(0u32..10) @@ -2757,7 +2836,7 @@ mod tests { #[crate::test] async fn test_export_sliced_varbinview_arrays() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let utf8 = VarBinViewArray::from_iter_nullable_str([ @@ -2822,7 +2901,7 @@ mod tests { #[case] arrow_offset: usize, #[case] len: usize, ) -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let logical_bits = (0..len).map(|idx| idx % 3 != 0).collect::>(); @@ -2855,7 +2934,7 @@ mod tests { #[crate::test] async fn test_repack_arrow_validity_buffer_zeroes_padding() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let len = 9; @@ -2885,7 +2964,7 @@ mod tests { #[crate::test] async fn test_export_all_false_validity_buffer_is_zeroed_on_device() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let len = 4; @@ -2911,7 +2990,7 @@ mod tests { // Check nullable primitives export Arrow null bitmaps on device. #[crate::test] async fn test_export_nullable_primitive() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let mut primitive = assert_nullable_export( @@ -2938,7 +3017,7 @@ mod tests { // Check nullable bool exports preserve Arrow offset metadata. #[crate::test] async fn test_export_nullable_bool() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let mut bools = assert_nullable_export( @@ -2959,7 +3038,7 @@ mod tests { // Check synthesized all-null bool validity is large enough for Arrow offset-based reads. #[crate::test] async fn test_export_all_null_sliced_bool_validity_covers_arrow_offset() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let mut bools = assert_nullable_export( @@ -2985,7 +3064,7 @@ mod tests { // Check nullable decimal exports include Arrow null bitmaps. #[crate::test] async fn test_export_nullable_decimal() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let mut decimal = assert_nullable_export( @@ -3015,7 +3094,7 @@ mod tests { // Check nullable temporal exports include Arrow null bitmaps. #[crate::test] async fn test_export_nullable_temporal() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let mut temporal = assert_nullable_export( @@ -3037,7 +3116,7 @@ mod tests { // Check nullable string-view exports include Arrow null bitmaps. #[crate::test] async fn test_export_nullable_varbinview() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let mut varbinview = assert_nullable_export( @@ -3060,7 +3139,7 @@ mod tests { // Check nullable struct exports include Arrow null bitmaps. #[crate::test] async fn test_export_nullable_struct() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let mut struct_array = assert_nullable_export( @@ -3081,10 +3160,41 @@ mod tests { Ok(()) } + // Non-canonical row validity should export as a device-resident bitmap. + #[crate::test] + async fn test_export_struct_non_canonical_validity() -> VortexResult<()> { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let validity = DictArray::try_new( + PrimitiveArray::from_iter([0u32, 1, 0]).into_array(), + BoolArray::from_iter([true, false]).into_array(), + )? + .into_array(); + assert!(!validity.is_canonical()); + + let mut struct_array = assert_nullable_export( + StructArray::try_new( + FieldNames::from_iter(["a"]), + vec![PrimitiveArray::from_iter(0u32..3).into_array()], + 3, + Validity::Array(validity), + )? + .into_array(), + 1, + 1, + &mut ctx, + ) + .await?; + unsafe { release_exported_array(&raw mut struct_array.array) }; + + Ok(()) + } + // Check nested struct children expose cuDF-compatible Arrow Device layouts. #[crate::test] async fn test_export_nested_struct_child_layout() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let mut device_array = nested_struct_array().export_device_array(&mut ctx).await?; @@ -3137,7 +3247,7 @@ mod tests { // Check parent release recursively releases children and is safe to repeat. #[crate::test] async fn test_release_is_idempotent_and_releases_children() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let mut device_array = nested_struct_array().export_device_array(&mut ctx).await?; @@ -3185,7 +3295,7 @@ mod tests { #[crate::test] async fn test_export_struct() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = StructArray::new( @@ -3214,7 +3324,7 @@ mod tests { #[crate::test] async fn test_export_struct_with_schema() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = StructArray::new( @@ -3252,7 +3362,7 @@ mod tests { // Check nested struct device exports carry the matching Arrow schema. #[crate::test] async fn test_export_nested_struct_with_schema() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let mut exported = nested_struct_array() @@ -3284,7 +3394,7 @@ mod tests { #[crate::test] async fn test_export_nested_struct_decimal_with_schema() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let nested = StructArray::new( @@ -3345,7 +3455,7 @@ mod tests { #[crate::test] async fn test_export_primitive_with_schema_is_column_shaped() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = PrimitiveArray::from_iter(0u32..5).into_array(); @@ -3364,7 +3474,7 @@ mod tests { #[crate::test] async fn test_export_varbinview_with_schema_uses_utf8_view_layout() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let japanese = "こんにちは"; diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index 7dac0c8631d..9545ba6abee 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -11,8 +11,13 @@ mod canonical; mod list_view; +use std::ffi::CString; +use std::ffi::c_char; +use std::ffi::c_int; use std::ffi::c_void; use std::fmt::Debug; +use std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; use std::ptr; use std::sync::Arc; @@ -38,13 +43,18 @@ use vortex::array::arrays::listview::ListViewArrayExt; use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::arrow::ArrowSessionExt; use vortex::array::buffer::BufferHandle; +use vortex::array::stream::SendableArrayStream; use vortex::dtype::DType; use vortex::dtype::DecimalDType; use vortex::dtype::DecimalType; use vortex::dtype::PType; use vortex::dtype::StructFields; use vortex::error::VortexResult; +use vortex::error::vortex_ensure; use vortex::error::vortex_err; +use vortex::io::runtime::BlockingRuntime; +use vortex::io::runtime::current::CurrentThreadRuntime; +use vortex::session::VortexSession; use crate::CudaBufferExt; use crate::CudaExecutionCtx; @@ -61,27 +71,15 @@ mod arrow_c_abi { pub use arrow_c_abi::ArrowArray; pub use arrow_c_abi::ArrowDeviceArray; +pub use arrow_c_abi::ArrowDeviceArrayStream; pub use arrow_c_abi::ArrowDeviceType; +use arrow_c_abi::ArrowSchema; #[cfg(feature = "_test-harness")] #[doc(hidden)] pub mod test_harness { - use vortex::array::buffer::BufferHandle; - use vortex::error::VortexResult; - - use crate::CudaExecutionCtx; - use crate::arrow::canonical::repack_arrow_validity_buffer as repack_arrow_validity_buffer_impl; - - /// Repack a validity bitmap into Arrow's byte-addressed bitmap layout on the active stream. - pub fn repack_arrow_validity_buffer( - input_buffer: &BufferHandle, - input_offset: usize, - len: usize, - arrow_offset: usize, - ctx: &mut CudaExecutionCtx, - ) -> VortexResult { - repack_arrow_validity_buffer_impl(input_buffer, input_offset, len, arrow_offset, ctx) - } + pub use crate::arrow::canonical::count_arrow_validity_nulls; + pub use crate::arrow::canonical::repack_arrow_validity_buffer; } /// CUDA device memory. @@ -107,24 +105,36 @@ impl ArrowArray { } } +impl ArrowDeviceArray { + /// A zeroed device array: an empty Arrow array with no device. Used as a + /// callback output placeholder and as the basis for the end-of-stream + /// marker. + fn empty() -> Self { + Self { + array: ArrowArray::empty(), + device_id: 0, + device_type: 0, + sync_event: ptr::null_mut(), + reserved: Default::default(), + } + } +} + unsafe impl Send for ArrowArray {} unsafe impl Sync for ArrowArray {} -#[expect( - unused, - reason = "cuda_stream and cuda_buffers need to have deferred drop" -)] pub(crate) struct PrivateData { /// Hold a reference to the CudaStream so that it stays alive even after CudaExecutionCtx /// has been dropped. pub(crate) cuda_stream: Arc, /// The single boxed slice which owns all buffers that the Rust code allocated on the device. + #[allow(dead_code, reason = "buffers are retained for deferred drop")] pub(crate) buffers: Box<[Option]>, /// Boxed slice of buffer pointers. We return a pointer to the start of this allocation over /// the interface, so we hold it here so the Box contents are not freed. pub(crate) buffer_ptrs: Box<[*const c_void]>, - pub(crate) cuda_event: CudaEvent, - pub(crate) cuda_event_ptr: cudaEvent_t, + pub(crate) export_event: CudaEvent, + pub(crate) export_event_ptr: cudaEvent_t, pub(crate) children: Box<[*mut ArrowArray]>, pub(crate) dictionary: *mut ArrowArray, } @@ -168,11 +178,10 @@ impl PrivateData { .map(|array| Box::into_raw(Box::new(array))) .collect::>(); - // generate the synchronization event - let cuda_event = ctx + let export_event = ctx .stream() .record_event(None) - .map_err(|_| vortex_err!("failed to create cudaEvent_t"))?; + .map_err(|_| vortex_err!("failed to create CUDA export event"))?; let dictionary = dictionary .map(|array| Box::into_raw(Box::new(array))) @@ -184,14 +193,14 @@ impl PrivateData { cuda_stream: Arc::clone(ctx.stream()), children, dictionary, - cuda_event_ptr: cuda_event.cu_event().cast(), - cuda_event, + export_event_ptr: export_event.cu_event().cast(), + export_event, })) } - /// Return a stable pointer to the recorded CUDA event handle. + /// Return a stable pointer to the recorded CUDA export event handle. pub(crate) fn sync_event(&mut self) -> SyncEvent { - (&raw mut self.cuda_event_ptr).cast() + (&raw mut self.export_event_ptr).cast() } } @@ -260,6 +269,408 @@ impl DeviceArrayExt for ArrayRef { } } +/// POSIX EIO for Arrow stream producer/export failures. +const LIBC_EIO: c_int = 5; + +/// POSIX EINVAL for invalid Arrow stream callback arguments or released streams. +const LIBC_EINVAL: c_int = 22; + +#[derive(Debug, PartialEq)] +enum ArrowDeviceStreamSchema { + Schema(Schema), + Field(Field), +} + +impl ArrowDeviceStreamSchema { + /// Convert an Arrow C schema into the stream schema shape for `dtype`. + fn from_ffi(schema: &FFI_ArrowSchema, dtype: &DType) -> VortexResult { + if matches!(dtype, DType::Struct(..)) { + Ok(Self::Schema(Schema::try_from(schema)?)) + } else { + Ok(Self::Field(Field::try_from(schema)?)) + } + } + + /// Convert a Vortex dtype into a stream schema when no stream array is available. + /// + /// This uses only the logical dtype, so it can differ from a non-empty stream's first-array + /// schema for encodings the dtype does not capture: a dictionary column reports a plain field + /// here but `DataType::Dictionary` once a concrete array is seen. + fn from_dtype(dtype: &DType, ctx: &mut CudaExecutionCtx) -> VortexResult { + let dtype = arrow_device_export_dtype(dtype); + if let DType::Struct(struct_dtype, _) = &dtype { + Ok(Self::Schema(Schema::new( + arrow_device_export_struct_fields(struct_dtype, ctx)?, + ))) + } else { + Ok(Self::Field(arrow_device_export_field("", &dtype, ctx)?)) + } + } + + /// Export this stream schema as an owned Arrow C schema. + fn to_ffi(&self) -> VortexResult { + match self { + Self::Schema(schema) => Ok(FFI_ArrowSchema::try_from(schema)?), + Self::Field(field) => Ok(FFI_ArrowSchema::try_from(field)?), + } + } +} + +type ArrayStreamIterator = Box>>; + +struct DeviceArrayStreamPrivateData { + array_iter: ArrayStreamIterator, + ctx: CudaExecutionCtx, + /// Runtime used by Arrow stream callbacks to pull from `array_iter` and block on per-array + /// CUDA exports. It must match the runtime that owns the underlying Vortex scan tasks so those + /// tasks are polled while callbacks are producing arrays. + runtime: CurrentThreadRuntime, + dtype: DType, + schema: Option, + pending_array: Option, + device_id: i64, + last_error: Option, +} + +impl DeviceArrayStreamPrivateData { + /// Clear the last stream error before a new callback invocation. + fn clear_error(&mut self) { + self.last_error = None; + } + + /// Store the last stream error and return the requested Arrow callback error code. + /// + /// Interior NUL bytes are replaced so `get_last_error` is never null while a non-zero status + /// is reported. + fn set_error(&mut self, error: impl ToString, code: c_int) -> c_int { + let message = error.to_string().replace('\0', " "); + self.last_error = Some(CString::new(message).unwrap_or_default()); + code + } + + /// Return the stream schema, exporting the first stream array to derive it if needed. + /// + /// A first array is held in `pending_array` so the following `get_next` returns it. + fn get_or_init_schema(&mut self) -> VortexResult<&ArrowDeviceStreamSchema> { + if self.schema.is_none() { + match self.array_iter.next() { + Some(array) => self.pending_array = Some(self.export_stream_array(array?)?), + None => { + self.schema = Some(ArrowDeviceStreamSchema::from_dtype( + &self.dtype, + &mut self.ctx, + )?); + } + } + } + + self.schema + .as_ref() + .ok_or_else(|| vortex_err!("ArrowDeviceArrayStream schema was not initialized")) + } + + /// Export and return the next Arrow device array, or `None` at end of stream. + fn next_array(&mut self) -> VortexResult> { + if let Some(array) = self.pending_array.take() { + return Ok(Some(array)); + } + + match self.array_iter.next() { + Some(array) => self.export_stream_array(array?).map(Some), + None => Ok(None), + } + } + + /// Export one array from the Vortex stream, validating it against the device stream. + fn export_stream_array(&mut self, array: ArrayRef) -> VortexResult { + vortex_ensure!( + array.dtype() == &self.dtype, + "stream array dtype changed from {} to {}", + self.dtype, + array.dtype() + ); + + let ArrowDeviceArrayWithSchema { + schema: mut ffi_schema, + array: mut device_array, + } = self + .runtime + .block_on(array.export_device_array_with_schema(&mut self.ctx))?; + + // Release the schema we no longer need, and on failure release the array we will not + // return. + let checked = self.check_stream_array(&ffi_schema, &device_array); + release_schema(&mut ffi_schema); + let exported_schema = match checked { + Ok(exported_schema) => exported_schema, + Err(error) => { + release_device_array(&mut device_array); + return Err(error); + } + }; + if self.schema.is_none() { + self.schema = Some(exported_schema); + } + Ok(device_array) + } + + /// Check that a freshly exported device array matches the stream schema and CUDA device. + fn check_stream_array( + &self, + ffi_schema: &FFI_ArrowSchema, + device_array: &ArrowDeviceArray, + ) -> VortexResult { + vortex_ensure!( + device_array.device_type == ARROW_DEVICE_CUDA, + "stream array exported on non-CUDA device type {}", + device_array.device_type + ); + vortex_ensure!( + device_array.device_id == self.device_id, + "stream array moved from CUDA device {} to {}", + self.device_id, + device_array.device_id + ); + + let exported_schema = ArrowDeviceStreamSchema::from_ffi(ffi_schema, &self.dtype)?; + if let Some(stream_schema) = &self.schema { + vortex_ensure!( + stream_schema == &exported_schema, + "stream array Arrow schema changed from {:?} to {:?}; an Arrow C device stream \ + requires every array to share one schema, so chunks must not vary their \ + encoding (for example a dictionary-encoded chunk among plain chunks)", + stream_schema, + exported_schema + ); + } + Ok(exported_schema) + } +} + +impl Drop for DeviceArrayStreamPrivateData { + /// Release a first stream array if `get_schema` exported it and `get_next` never returned it. + fn drop(&mut self) { + if let Some(mut array) = self.pending_array.take() { + release_device_array(&mut array); + } + } +} + +/// Extension trait for exporting a Vortex array stream as an Arrow Device stream. +pub trait DeviceArrayStreamExt { + /// Export this stream as an [`ArrowDeviceArrayStream`]. + /// + /// Arrays are exported by reusing one [`CudaExecutionCtx`], and every produced + /// [`ArrowDeviceArray`] must remain on the CUDA device captured at stream construction. The + /// returned [`ArrowDeviceArrayStream`] owns the Vortex stream and must be released through its + /// embedded `release` callback. + /// + /// The Arrow Device stream contract requires all arrays to share the schema reported by + /// `get_schema`. The schema is derived from the first array, or from the logical dtype + /// for an empty stream. Chunks that export to different Arrow types are rejected mid-stream. + /// + /// Drive the returned stream from one thread. `runtime` must be the runtime that owns the + /// underlying scan tasks and per-array exports. + fn export_device_array_stream( + self, + session: &VortexSession, + runtime: &CurrentThreadRuntime, + ) -> VortexResult; +} + +impl DeviceArrayStreamExt for SendableArrayStream { + /// Drive this stream on `runtime` and export it. + fn export_device_array_stream( + self, + session: &VortexSession, + runtime: &CurrentThreadRuntime, + ) -> VortexResult { + let dtype = self.dtype().clone(); + let ctx = crate::CudaSession::create_execution_ctx(session)?; + let array_iter = Box::new(runtime.block_on_stream(self)); + Ok(device_array_stream(array_iter, dtype, ctx, runtime.clone())) + } +} + +/// Build the Arrow Device stream that owns `array_iter` and exports its arrays through `ctx`. +fn device_array_stream( + array_iter: ArrayStreamIterator, + dtype: DType, + ctx: CudaExecutionCtx, + runtime: CurrentThreadRuntime, +) -> ArrowDeviceArrayStream { + let private_data = Box::new(DeviceArrayStreamPrivateData { + device_id: ctx.stream().context().ordinal() as i64, + array_iter, + ctx, + runtime, + dtype, + schema: None, + pending_array: None, + last_error: None, + }); + + ArrowDeviceArrayStream { + device_type: ARROW_DEVICE_CUDA, + get_schema: Some(device_stream_get_schema), + get_next: Some(device_stream_get_next), + get_last_error: Some(device_stream_get_last_error), + release: Some(device_stream_release), + private_data: Box::into_raw(private_data).cast(), + } +} + +/// Returns the stream state stored in `private_data`. +unsafe fn device_stream_private_data<'a>( + stream: *mut ArrowDeviceArrayStream, +) -> Option<&'a mut DeviceArrayStreamPrivateData> { + let stream = unsafe { stream.as_mut()? }; + unsafe { + stream + .private_data + .cast::() + .as_mut() + } +} + +/// Create the Arrow end-of-stream marker for the stream's CUDA device. +fn released_device_array(device_id: i64) -> ArrowDeviceArray { + ArrowDeviceArray { + device_id, + device_type: ARROW_DEVICE_CUDA, + ..ArrowDeviceArray::empty() + } +} + +/// Release an Arrow C schema if it is live. +pub fn release_schema(schema: &mut FFI_ArrowSchema) { + if let Some(release) = schema.release { + unsafe { release(schema) }; + } +} + +/// Release an Arrow device array if it is live. +pub fn release_device_array(array: &mut ArrowDeviceArray) { + if let Some(release) = array.array.release { + unsafe { release(&raw mut array.array) }; + } +} + +/// Runs an Arrow stream callback body. +/// +/// Returns an Arrow callback status code and stores failures in `last_error`. +fn device_stream_callback( + state: &mut DeviceArrayStreamPrivateData, + panic_message: &'static str, + callback: impl FnOnce(&mut DeviceArrayStreamPrivateData) -> VortexResult<()>, +) -> c_int { + let result = catch_unwind(AssertUnwindSafe(|| callback(state))); + match result { + Ok(Ok(())) => 0, + Ok(Err(err)) => state.set_error(err, LIBC_EIO), + Err(_) => state.set_error(panic_message, LIBC_EIO), + } +} + +/// Write the stream's Arrow schema, initializing it from the first stream array if unset. +unsafe extern "C" fn device_stream_get_schema( + stream: *mut ArrowDeviceArrayStream, + out: *mut ArrowSchema, +) -> c_int { + let Some(state) = (unsafe { device_stream_private_data(stream) }) else { + return LIBC_EINVAL; + }; + state.clear_error(); + + if out.is_null() { + return state.set_error("null ArrowSchema output", LIBC_EINVAL); + } + + fn body(state: &mut DeviceArrayStreamPrivateData, out: *mut ArrowSchema) -> VortexResult<()> { + let schema = state.get_or_init_schema()?.to_ffi()?; + unsafe { ptr::write(out.cast::(), schema) }; + Ok(()) + } + + device_stream_callback( + state, + "panic in ArrowDeviceArrayStream::get_schema", + |state| body(state, out), + ) +} + +/// Write the next exported Arrow device array, or a released array at end of stream. +unsafe extern "C" fn device_stream_get_next( + stream: *mut ArrowDeviceArrayStream, + out: *mut ArrowDeviceArray, +) -> c_int { + let Some(state) = (unsafe { device_stream_private_data(stream) }) else { + return LIBC_EINVAL; + }; + state.clear_error(); + + if out.is_null() { + return state.set_error("null ArrowDeviceArray output", LIBC_EINVAL); + } + + // Keep the fallible part in a local function so `device_stream_callback` handles callback + // status and error reporting consistently. + fn body( + state: &mut DeviceArrayStreamPrivateData, + out: *mut ArrowDeviceArray, + ) -> VortexResult<()> { + let array = state + .next_array()? + .unwrap_or_else(|| released_device_array(state.device_id)); + unsafe { ptr::write(out, array) }; + Ok(()) + } + + device_stream_callback( + state, + "panic in ArrowDeviceArrayStream::get_next", + |state| body(state, out), + ) +} + +/// Return the most recent callback error message, or null if no error is stored. +unsafe extern "C" fn device_stream_get_last_error( + stream: *mut ArrowDeviceArrayStream, +) -> *const c_char { + let Some(state) = (unsafe { device_stream_private_data(stream) }) else { + return ptr::null(); + }; + + state + .last_error + .as_ref() + .map_or(ptr::null(), |error| error.as_ptr()) +} + +/// Free the stream state and null its callbacks. The null `release` makes a second call a no-op. +unsafe extern "C" fn device_stream_release(stream: *mut ArrowDeviceArrayStream) { + let Some(stream_ref) = (unsafe { stream.as_mut() }) else { + return; + }; + if stream_ref.release.is_none() { + return; + } + + stream_ref.get_schema = None; + stream_ref.get_next = None; + stream_ref.get_last_error = None; + stream_ref.release = None; + + let private_data = std::mem::replace(&mut stream_ref.private_data, ptr::null_mut()); + if !private_data.is_null() { + drop(catch_unwind(AssertUnwindSafe(|| unsafe { + drop(Box::from_raw( + private_data.cast::(), + )); + }))); + } +} + /// Build the Arrow C schema that describes the exported device array. pub(crate) fn arrow_schema_for_array( array: &ArrayRef, @@ -490,3 +901,178 @@ pub trait ExportDeviceArray: Debug + Send + Sync + 'static { Ok(ArrowDeviceArrayWithSchema { schema, array }) } } + +#[cfg(test)] +mod tests { + use std::ffi::CStr; + use std::ptr; + + use arrow_schema::DataType; + use arrow_schema::ffi::FFI_ArrowSchema; + use futures::stream; + use vortex::VortexSessionDefault; + use vortex::array::ArrayRef; + use vortex::array::IntoArray; + use vortex::array::arrays::PrimitiveArray; + use vortex::array::stream::ArrayStreamAdapter; + use vortex::array::stream::ArrayStreamExt; + use vortex::dtype::DType; + use vortex::dtype::Nullability; + use vortex::dtype::PType; + use vortex::error::VortexResult; + use vortex::error::vortex_err; + use vortex::io::runtime::current::CurrentThreadRuntime; + use vortex::session::VortexSession; + use vortex_cuda_macros::test as cuda_test; + + use crate::CudaSession; + use crate::arrow::ARROW_DEVICE_CUDA; + use crate::arrow::ArrowDeviceArray; + use crate::arrow::ArrowDeviceArrayStream; + use crate::arrow::ArrowSchema; + use crate::arrow::DeviceArrayStreamExt; + use crate::arrow::LIBC_EINVAL; + use crate::arrow::release_device_array; + use crate::arrow::release_schema; + + fn last_error(stream: &mut ArrowDeviceArrayStream) -> VortexResult { + let get_last_error = stream + .get_last_error + .ok_or_else(|| vortex_err!("stream missing get_last_error callback"))?; + let error = unsafe { get_last_error(stream as *mut ArrowDeviceArrayStream) }; + Ok(if error.is_null() { + String::new() + } else { + unsafe { CStr::from_ptr(error) } + .to_string_lossy() + .into_owned() + }) + } + + #[cuda_test] + fn test_export_device_array_stream_schema_next_eos_release() -> VortexResult<()> { + let runtime = CurrentThreadRuntime::new(); + let session = VortexSession::default().with::(); + let array = PrimitiveArray::from_iter(0u32..5).into_array(); + let stream = array.to_array_stream().boxed(); + let mut device_stream = stream.export_device_array_stream(&session, &runtime)?; + assert_eq!(device_stream.device_type, ARROW_DEVICE_CUDA); + + let mut schema = FFI_ArrowSchema::empty(); + let get_schema = device_stream + .get_schema + .ok_or_else(|| vortex_err!("stream missing get_schema callback"))?; + let status = unsafe { + get_schema( + &raw mut device_stream, + (&raw mut schema).cast::(), + ) + }; + assert_eq!(status, 0); + let field = arrow_schema::Field::try_from(&schema)?; + assert_eq!(field.data_type(), &DataType::UInt32); + + let get_next = device_stream + .get_next + .ok_or_else(|| vortex_err!("stream missing get_next callback"))?; + let mut first_array = ArrowDeviceArray::empty(); + let status = unsafe { get_next(&raw mut device_stream, &raw mut first_array) }; + assert_eq!(status, 0); + assert_eq!(first_array.device_type, ARROW_DEVICE_CUDA); + assert_eq!(first_array.array.length, 5); + assert!(first_array.array.release.is_some()); + + let mut eos = ArrowDeviceArray::empty(); + let status = unsafe { get_next(&raw mut device_stream, &raw mut eos) }; + assert_eq!(status, 0); + assert_eq!(eos.device_type, ARROW_DEVICE_CUDA); + assert!(eos.array.release.is_none()); + + unsafe { + release_device_array(&mut first_array); + release_schema(&mut schema); + let release = device_stream + .release + .ok_or_else(|| vortex_err!("stream missing release callback"))?; + release(&raw mut device_stream); + release(&raw mut device_stream); + } + assert!(device_stream.release.is_none()); + Ok(()) + } + + #[cuda_test] + fn test_export_device_array_stream_empty_stream_schema_and_eos() -> VortexResult<()> { + let runtime = CurrentThreadRuntime::new(); + let session = VortexSession::default().with::(); + let dtype = DType::Primitive(PType::U32, Nullability::NonNullable); + let stream = + ArrayStreamAdapter::new(dtype, stream::empty::>()).boxed(); + let mut device_stream = stream.export_device_array_stream(&session, &runtime)?; + + let get_schema = device_stream + .get_schema + .ok_or_else(|| vortex_err!("stream missing get_schema callback"))?; + let mut schema = FFI_ArrowSchema::empty(); + let status = unsafe { + get_schema( + &raw mut device_stream, + (&raw mut schema).cast::(), + ) + }; + assert_eq!(status, 0); + let field = arrow_schema::Field::try_from(&schema)?; + assert_eq!(field.data_type(), &DataType::UInt32); + + let get_next = device_stream + .get_next + .ok_or_else(|| vortex_err!("stream missing get_next callback"))?; + let mut eos = ArrowDeviceArray::empty(); + let status = unsafe { get_next(&raw mut device_stream, &raw mut eos) }; + assert_eq!(status, 0); + assert!(eos.array.release.is_none()); + + unsafe { + release_schema(&mut schema); + let release = device_stream + .release + .ok_or_else(|| vortex_err!("stream missing release callback"))?; + release(&raw mut device_stream); + } + Ok(()) + } + + #[cuda_test] + fn test_export_device_array_stream_null_outputs_report_error() -> VortexResult<()> { + let runtime = CurrentThreadRuntime::new(); + let session = VortexSession::default().with::(); + let array = PrimitiveArray::from_iter(0u32..5).into_array(); + let stream = array.to_array_stream().boxed(); + let mut device_stream = stream.export_device_array_stream(&session, &runtime)?; + + let get_schema = device_stream + .get_schema + .ok_or_else(|| vortex_err!("stream missing get_schema callback"))?; + let status = unsafe { get_schema(&raw mut device_stream, ptr::null_mut()) }; + assert_eq!(status, LIBC_EINVAL); + assert_eq!(last_error(&mut device_stream)?, "null ArrowSchema output"); + + let get_next = device_stream + .get_next + .ok_or_else(|| vortex_err!("stream missing get_next callback"))?; + let status = unsafe { get_next(&raw mut device_stream, ptr::null_mut()) }; + assert_eq!(status, LIBC_EINVAL); + assert_eq!( + last_error(&mut device_stream)?, + "null ArrowDeviceArray output" + ); + + unsafe { + let release = device_stream + .release + .ok_or_else(|| vortex_err!("stream missing release callback"))?; + release(&raw mut device_stream); + } + Ok(()) + } +} diff --git a/vortex-cuda/src/dynamic_dispatch/mod.rs b/vortex-cuda/src/dynamic_dispatch/mod.rs index 186f2297db1..329ea19c061 100644 --- a/vortex-cuda/src/dynamic_dispatch/mod.rs +++ b/vortex-cuda/src/dynamic_dispatch/mod.rs @@ -32,7 +32,6 @@ use vortex::array::buffer::BufferHandle; use vortex::array::buffer::DeviceBufferExt; use vortex::array::match_each_unsigned_integer_ptype; use vortex::array::scalar::Scalar; -use vortex::array::validity::Validity; use vortex::buffer::Alignment; use vortex::buffer::ByteBuffer; use vortex::buffer::ByteBufferMut; @@ -434,7 +433,7 @@ impl MaterializedPlan { let output_ptype = self.dispatch_plan.output_ptype(); // All values are null — no need to touch the GPU. - if matches!(self.validity, Validity::AllInvalid) { + if self.validity.definitely_all_null() { let dtype = DType::Primitive(output_ptype, Nullability::Nullable); return ConstantArray::new(Scalar::null(dtype), len) .into_array() @@ -546,6 +545,7 @@ mod tests { use vortex::array::validity::Validity; use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; + use vortex::buffer::buffer; use vortex::dtype::DType; use vortex::dtype::NativePType; use vortex::dtype::PType; @@ -561,12 +561,16 @@ mod tests { use vortex::encodings::fastlanes::FoR; use vortex::encodings::fastlanes::FoRArrayExt; use vortex::encodings::runend::RunEnd; + use vortex::encodings::sequence::Sequence; use vortex::encodings::zigzag::ZigZag; use vortex::error::VortexExpect; use vortex::error::VortexResult; - use vortex::session::VortexSession; - use vortex_array::LEGACY_SESSION; + use vortex::mask::Mask; + use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::FilterArray; + use vortex_array::assert_arrays_eq; use vortex_array::patches::Patches; use super::*; @@ -574,23 +578,20 @@ mod tests { use crate::CudaBufferExt; use crate::CudaDeviceBuffer; use crate::CudaExecutionCtx; + use crate::cuda_session; use crate::executor::CudaArrayExt; use crate::executor::CudaDispatchMode; use crate::hybrid_dispatch::try_gpu_dispatch; use crate::session::CudaSession; - fn bitpacked_array_u32(bit_width: u8, len: usize) -> BitPackedArray { + fn bitpacked_array_u32(bit_width: u8, len: usize, ctx: &mut ExecutionCtx) -> BitPackedArray { let max_val = (1u64 << bit_width).saturating_sub(1); let values: Vec = (0..len) .map(|i| ((i as u64) % (max_val + 1)) as u32) .collect(); let primitive = PrimitiveArray::new(Buffer::from(values), NonNullable); - BitPacked::encode( - &primitive.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("failed to create BitPacked array") + BitPacked::encode(&primitive.into_array(), bit_width, ctx) + .vortex_expect("failed to create BitPacked array") } async fn dispatch_plan( @@ -626,6 +627,7 @@ mod tests { #[crate::test] fn test_max_scalar_ops() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width: u8 = 6; let len = 2050; let references: [u32; 4] = [1, 2, 4, 8]; @@ -636,8 +638,8 @@ mod tests { .map(|i| ((i as u64) % (max_val + 1)) as u32 + total_reference) .collect(); - let bitpacked = bitpacked_array_u32(bit_width, len); - let cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let bitpacked = bitpacked_array_u32(bit_width, len, &mut ctx); + let cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let packed = bitpacked.packed().clone(); let device_input = futures::executor::block_on(cuda_ctx.ensure_on_device(packed))?; let input_ptr = device_input.cuda_device_ptr()?; @@ -753,7 +755,7 @@ mod tests { }) .collect(); - let cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let (input_ptr, _di) = copy_raw_to_device(&cuda_ctx, &data)?; let plan = CudaDispatchPlan::new( @@ -846,6 +848,7 @@ mod tests { #[crate::test] async fn test_bitpacked() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width: u8 = 10; let len = 3000; let max_val = (1u64 << bit_width).saturating_sub(1); @@ -853,8 +856,8 @@ mod tests { .map(|i| ((i as u64) % (max_val + 1)) as u32) .collect(); - let bp = bitpacked_array_u32(bit_width, len); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let bp = bitpacked_array_u32(bit_width, len, &mut ctx); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&bp.into_array(), &mut cuda_ctx).await?; let actual = @@ -866,6 +869,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width: u8 = 6; let len = 3000; let reference = 42u32; @@ -876,10 +880,10 @@ mod tests { .collect(); let expected: Vec = raw.iter().map(|&v| v + reference).collect(); - let bp = bitpacked_array_u32(bit_width, len); + let bp = bitpacked_array_u32(bit_width, len, &mut ctx); let for_arr = FoR::try_new(bp.into_array(), Scalar::from(reference))?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&for_arr.into_array(), &mut cuda_ctx).await?; let actual = @@ -901,7 +905,7 @@ mod tests { expected.push(values[run]); } - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let ends_arr = PrimitiveArray::new(Buffer::from(ends), NonNullable).into_array(); let values_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); let re = RunEnd::new(ends_arr, values_arr, cuda_ctx.execution_ctx()); @@ -917,6 +921,7 @@ mod tests { #[crate::test] async fn test_dict_for_bp_values_bp_codes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Dict where both codes and values are BitPacked+FoR. let dict_reference = 1_000_000u32; let dict_residuals: Vec = (0..64).collect(); @@ -929,24 +934,16 @@ mod tests { // BitPack+FoR the dict values let dict_prim = PrimitiveArray::new(Buffer::from(dict_residuals), NonNullable); - let dict_bp = BitPacked::encode( - &dict_prim.into_array(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let dict_bp = BitPacked::encode(&dict_prim.into_array(), 6, &mut ctx)?; let dict_for = FoR::try_new(dict_bp.into_array(), Scalar::from(dict_reference))?; // BitPack the codes let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let codes_bp = BitPacked::encode(&codes_prim.into_array(), 6, &mut ctx)?; let dict = DictArray::try_new(codes_bp.into_array(), dict_for.into_array())?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&dict.into_array(), &mut cuda_ctx).await?; let actual = @@ -958,7 +955,7 @@ mod tests { #[crate::test] async fn test_alp_for_bitpacked() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // ALP(FoR(BitPacked)): encode each layer, then reassemble the tree // bottom-up because encode() methods produce flat outputs. let len = 3000; @@ -970,12 +967,11 @@ mod tests { let alp = alp_encode(float_prim.as_view(), Some(exponents), &mut ctx)?; assert!(alp.patches().is_none()); - let for_arr = FoR::encode(alp.encoded().clone().execute::(&mut ctx)?)?; - let bp = BitPacked::encode( - for_arr.encoded(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), + let for_arr = FoR::encode( + alp.encoded().clone().execute::(&mut ctx)?, + &mut ctx, )?; + let bp = BitPacked::encode(for_arr.encoded(), 6, &mut ctx)?; let tree = ALP::new( FoR::try_new(bp.into_array(), for_arr.reference_scalar().clone())?.into_array(), @@ -983,7 +979,7 @@ mod tests { None, ); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&tree.into_array(), &mut cuda_ctx).await?; let actual = @@ -995,6 +991,7 @@ mod tests { #[crate::test] async fn test_zigzag_bitpacked() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // ZigZag(BitPacked): unpack then zigzag-decode. let bit_width: u8 = 4; let len = 3000; @@ -1009,14 +1006,10 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(raw), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; let zz = ZigZag::try_new(bp.into_array())?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&zz.into_array(), &mut cuda_ctx).await?; let actual = @@ -1040,7 +1033,7 @@ mod tests { expected.push(values[run] + reference); } - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let ends_arr = PrimitiveArray::new(Buffer::from(ends), NonNullable).into_array(); let values_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); let re = RunEnd::new(ends_arr, values_arr, cuda_ctx.execution_ctx()); @@ -1074,7 +1067,7 @@ mod tests { let dict = DictArray::try_new(codes_prim.into_array(), values_prim.into_array())?; let for_arr = FoR::try_new(dict.into_array(), Scalar::from(reference))?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&for_arr.into_array(), &mut cuda_ctx).await?; let actual = @@ -1086,6 +1079,7 @@ mod tests { #[crate::test] async fn test_dict_for_bp_codes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Dict(codes=FoR(BitPacked), values=primitive) let dict_values: Vec = (0..8).map(|i| i * 1000 + 7).collect(); let dict_size = dict_values.len(); @@ -1096,17 +1090,13 @@ mod tests { // BitPack codes, then wrap in FoR (reference=0 so values unchanged) let bit_width: u8 = 3; let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let codes_bp = BitPacked::encode(&codes_prim.into_array(), bit_width, &mut ctx)?; let codes_for = FoR::try_new(codes_bp.into_array(), Scalar::from(0u32))?; let values_prim = PrimitiveArray::new(Buffer::from(dict_values), NonNullable); let dict = DictArray::try_new(codes_for.into_array(), values_prim.into_array())?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&dict.into_array(), &mut cuda_ctx).await?; let actual = @@ -1118,6 +1108,7 @@ mod tests { #[crate::test] async fn test_dict_primitive_values_bp_codes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dict_values: Vec = vec![100, 200, 300, 400]; let dict_size = dict_values.len(); let len = 3000; @@ -1126,16 +1117,12 @@ mod tests { let bit_width: u8 = 2; let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let codes_bp = BitPacked::encode(&codes_prim.into_array(), bit_width, &mut ctx)?; let values_prim = PrimitiveArray::new(Buffer::from(dict_values), NonNullable); let dict = DictArray::try_new(codes_bp.into_array(), values_prim.into_array())?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&dict.into_array(), &mut cuda_ctx).await?; let actual = @@ -1147,6 +1134,7 @@ mod tests { #[crate::test] async fn test_dict_mixed_width_u8_codes_u32_values() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dict_values: Vec = vec![100, 200, 300, 400]; let len = 3000; let codes: Vec = (0..len).map(|i| (i % dict_values.len()) as u8).collect(); @@ -1166,19 +1154,20 @@ mod tests { ); // Execute through the hybrid dispatch path (handles widening). - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut ctx); Ok(()) } #[crate::test] async fn test_dict_mixed_width_u16_codes_u32_values() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let dict_values: Vec = vec![1000, 2000, 3000, 4000, 5000]; let len = 2048; let codes: Vec = (0..len).map(|i| (i % dict_values.len()) as u16).collect(); @@ -1198,24 +1187,25 @@ mod tests { ); // Execute through the hybrid dispatch path (handles widening). - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_runend_mixed_width_u64_ends_u32_values() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let ends: Vec = vec![1000, 2000, 3000]; let values: Vec = vec![10, 20, 30]; let len = 3000; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let ends_arr = PrimitiveArray::new(Buffer::from(ends), NonNullable).into_array(); let values_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); let re = RunEnd::new(ends_arr, values_arr, cuda_ctx.execution_ctx()); @@ -1245,7 +1235,7 @@ mod tests { }) .collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -1279,7 +1269,7 @@ mod tests { let expected: Vec = data[slice_start..slice_end].to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1313,6 +1303,7 @@ mod tests { #[case] slice_start: usize, #[case] slice_end: usize, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width = 10u8; let max_val = (1u32 << bit_width) - 1; let len = 5000; @@ -1324,17 +1315,13 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(raw), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; let zz = ZigZag::try_new(bp.into_array())?; let sliced = zz.into_array().slice(slice_start..slice_end)?; let expected: Vec = all_decoded[slice_start..slice_end].to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1384,7 +1371,7 @@ mod tests { .map(|&c| dict_values[c as usize]) .collect(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1418,22 +1405,19 @@ mod tests { #[case] slice_start: usize, #[case] slice_end: usize, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width = 10u8; let max_val = (1u32 << bit_width) - 1; let len = 5000; let data: Vec = (0..len).map(|i| (i as u32) % max_val).collect(); let prim = PrimitiveArray::new(Buffer::from(data.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; let sliced = bp.into_array().slice(slice_start..slice_end)?; let expected: Vec = data[slice_start..slice_end].to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1467,6 +1451,7 @@ mod tests { #[case] slice_start: usize, #[case] slice_end: usize, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let reference = 100u32; let bit_width = 10u8; let max_val = (1u32 << bit_width) - 1; @@ -1474,11 +1459,7 @@ mod tests { let encoded_data: Vec = (0..len).map(|i| (i as u32) % max_val).collect(); let prim = PrimitiveArray::new(Buffer::from(encoded_data.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; let for_arr = FoR::try_new(bp.into_array(), Scalar::from(reference))?; let all_decoded: Vec = encoded_data.iter().map(|&v| v + reference).collect(); @@ -1486,7 +1467,7 @@ mod tests { let sliced = for_arr.into_array().slice(slice_start..slice_end)?; let expected: Vec = all_decoded[slice_start..slice_end].to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1520,6 +1501,7 @@ mod tests { #[case] slice_start: usize, #[case] slice_end: usize, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dict_reference = 1_000_000u32; let dict_residuals: Vec = (0..64).collect(); let dict_expected: Vec = dict_residuals.iter().map(|&r| r + dict_reference).collect(); @@ -1531,27 +1513,19 @@ mod tests { // BitPack+FoR the dict values let dict_prim = PrimitiveArray::new(Buffer::from(dict_residuals), NonNullable); - let dict_bp = BitPacked::encode( - &dict_prim.into_array(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let dict_bp = BitPacked::encode(&dict_prim.into_array(), 6, &mut ctx)?; let dict_for = FoR::try_new(dict_bp.into_array(), Scalar::from(dict_reference))?; // BitPack the codes let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let codes_bp = BitPacked::encode(&codes_prim.into_array(), 6, &mut ctx)?; let dict = DictArray::try_new(codes_bp.into_array(), dict_for.into_array())?; let sliced = dict.into_array().slice(slice_start..slice_end)?; let expected: Vec = all_decoded[slice_start..slice_end].to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1583,7 +1557,7 @@ mod tests { let seq = Sequence::try_new_typed(base, multiplier, Nullability::NonNullable, len)?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&seq.into_array(), &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1616,7 +1590,7 @@ mod tests { let seq = Sequence::try_new_typed(base, multiplier, Nullability::NonNullable, len)?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&seq.into_array(), &mut cuda_ctx).await?; let actual_u32 = run_dynamic_dispatch_plan( @@ -1633,6 +1607,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_u8() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let bit_width: u8 = 4; let len = 3000; let reference = 100u8; @@ -1644,29 +1619,26 @@ mod tests { .collect(); let primitive = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let bp = BitPacked::encode( - &primitive.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack u8"); + let bp = BitPacked::encode(&primitive.into_array(), bit_width, &mut cpu_ctx) + .vortex_expect("bitpack u8"); let for_arr = FoR::try_new( bp.into_array(), Scalar::primitive(reference, Nullability::NonNullable), )?; let array = for_arr.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_for_bitpacked_u16() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let bit_width: u8 = 10; let len = 3000; let reference = 1000u16; @@ -1680,29 +1652,26 @@ mod tests { .collect(); let primitive = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let bp = BitPacked::encode( - &primitive.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack u16"); + let bp = BitPacked::encode(&primitive.into_array(), bit_width, &mut cpu_ctx) + .vortex_expect("bitpack u16"); let for_arr = FoR::try_new( bp.into_array(), Scalar::primitive(reference, Nullability::NonNullable), )?; let array = for_arr.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_for_bitpacked_u64() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let bit_width: u8 = 20; let len = 3000; let reference = 100_000u64; @@ -1714,24 +1683,20 @@ mod tests { .collect(); let primitive = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let bp = BitPacked::encode( - &primitive.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack u64"); + let bp = BitPacked::encode(&primitive.into_array(), bit_width, &mut cpu_ctx) + .vortex_expect("bitpack u64"); let for_arr = FoR::try_new( bp.into_array(), Scalar::primitive(reference, Nullability::NonNullable), )?; let array = for_arr.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -1739,7 +1704,7 @@ mod tests { async fn test_empty_array() -> VortexResult<()> { let values: Vec = vec![]; let primitive = PrimitiveArray::new(Buffer::from(values), NonNullable); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&primitive.into_array(), &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); assert_eq!(result.len(), 0); @@ -1748,31 +1713,29 @@ mod tests { #[crate::test] async fn test_single_element() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let values: Vec = vec![42]; let primitive = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &primitive.into_array(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack"); + let bp = + BitPacked::encode(&primitive.into_array(), 6, &mut cpu_ctx).vortex_expect("bitpack"); let for_arr = FoR::try_new( bp.into_array(), Scalar::primitive(0u32, Nullability::NonNullable), )?; let array = for_arr.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_exactly_elements_per_block() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Exactly 2048 elements — one full block, no remainder let bit_width: u8 = 6; let len = 2048; @@ -1784,24 +1747,20 @@ mod tests { let expected: Vec = residuals.iter().map(|&r| r + reference).collect(); let primitive = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let bp = BitPacked::encode( - &primitive.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack"); + let bp = BitPacked::encode(&primitive.into_array(), bit_width, &mut ctx) + .vortex_expect("bitpack"); let for_arr = FoR::try_new( bp.into_array(), Scalar::primitive(reference, Nullability::NonNullable), )?; let array = for_arr.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut ctx); Ok(()) } @@ -1811,8 +1770,9 @@ mod tests { #[crate::test] fn test_has_standalone_kernel_true_cases() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // BitPacked — leaf encoding, no children. - let bp = bitpacked_array_u32(6, 2048); + let bp = bitpacked_array_u32(6, 2048, &mut ctx); let bp_arr = bp.into_array(); assert!(plan_builder::has_standalone_kernel(&bp_arr)); @@ -1862,14 +1822,12 @@ mod tests { /// a kernel registered in the CUDA session. #[crate::test] fn test_has_standalone_kernel_implies_registered_kernel() -> VortexResult<()> { - use vortex::dtype::Nullability; - use vortex::encodings::sequence::Sequence; - - let session = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut ctx = array_session().create_execution_ctx(); + let session = CudaSession::create_execution_ctx(&cuda_session())?; let cuda_session = session.cuda_session(); // Leaf encodings. - let bp = bitpacked_array_u32(6, 2048); + let bp = bitpacked_array_u32(6, 2048, &mut ctx); let bp_arr = bp.into_array(); let seq = Sequence::try_new_typed(0u32, 1u32, Nullability::NonNullable, 2048)?; let seq_arr = seq.into_array(); @@ -1886,7 +1844,7 @@ mod tests { let patched_bp = BitPacked::encode( &PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(), 4, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, )?; assert!(patched_bp.patches().is_some(), "expected patches"); let patched_bp_arr = patched_bp.into_array(); @@ -1931,7 +1889,7 @@ mod tests { #[crate::test] async fn test_alp_f64_for_bitpacked() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // ALP(FoR(BitPacked)) with f64: same structure as the f32 test. let len = 3000; let exponents = Exponents { e: 2, f: 0 }; @@ -1942,12 +1900,11 @@ mod tests { let alp = alp_encode(float_prim.as_view(), Some(exponents), &mut ctx)?; assert!(alp.patches().is_none()); - let for_arr = FoR::encode(alp.encoded().clone().execute::(&mut ctx)?)?; - let bp = BitPacked::encode( - for_arr.encoded(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), + let for_arr = FoR::encode( + alp.encoded().clone().execute::(&mut ctx)?, + &mut ctx, )?; + let bp = BitPacked::encode(for_arr.encoded(), 6, &mut ctx)?; let tree = ALP::new( FoR::try_new(bp.into_array(), for_arr.reference_scalar().clone())?.into_array(), @@ -1959,15 +1916,15 @@ mod tests { // CPU decode as ground truth. let cpu = array .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut ctx)? .into_array(); // GPU decode. - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let gpu = CanonicalCudaExt::into_host(canonical).await?.into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(cpu, gpu, &mut ctx); Ok(()) } @@ -1980,6 +1937,7 @@ mod tests { #[case] len: usize, #[case] slice_range: Option>, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let mut values: Vec = (0..len).map(|i| (i as f64) * 1.1).collect(); // Insert exception values that ALP can't encode. values[0] = 99.9; @@ -1993,12 +1951,7 @@ mod tests { } let float_prim = PrimitiveArray::new(Buffer::from(values), NonNullable); - let encoded = alp_encode( - float_prim.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - )? - .into_array(); + let encoded = alp_encode(float_prim.as_view(), None, &mut ctx)?.into_array(); let (array, base_offset) = if let Some(range) = &slice_range { (encoded.slice(range.clone())?, range.start) @@ -2007,12 +1960,10 @@ mod tests { }; // Decode on CPU as ground truth (accounts for ALP precision loss + patches). - let cpu_decoded = array - .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + let cpu_decoded = array.clone().execute::(&mut ctx)?; let expected: Vec = cpu_decoded.as_slice::().to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?; let result_prim = result.as_primitive(); @@ -2031,8 +1982,9 @@ mod tests { #[crate::test] async fn alp_slice_device_patches() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); // Regression test for https://github.com/vortex-data/vortex/issues/7838#issuecomment-4452796116. - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let len = 4096; let exponents = Exponents { e: 0, f: 0 }; @@ -2086,20 +2038,21 @@ mod tests { expected[2048 - 100] = std::f64::consts::LN_2; let expected = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected, gpu); + assert_arrays_eq!(expected, gpu, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_runend_u32_ends_u16_values() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); // RunEnd with u32 ends, u16 values. Output type = u16. // Ends (u32) differ from output (u16) → pending subtree. let ends: Vec = vec![500, 1000, 1500, 2000]; let values: Vec = vec![100, 200, 300, 400]; let len = 2000; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let ends_arr = PrimitiveArray::new(Buffer::from(ends), NonNullable).into_array(); let values_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); let re = RunEnd::new(ends_arr, values_arr, cuda_ctx.execution_ctx()); @@ -2130,13 +2083,14 @@ mod tests { }) .collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_dict_bitpacked_u8_codes_u32_values() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); // Dict with BitPacked u8 codes (narrower than u32 output) and u32 values. // The kernel's bitunpack_typed decodes at the source's native width and // widens to T, so this fuses into a single kernel launch. @@ -2146,12 +2100,8 @@ mod tests { let codes_prim = PrimitiveArray::new(Buffer::from(codes.clone()), NonNullable); // BitPack the u8 codes at 2 bits (4 values need 2 bits) - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - 2, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack codes"); + let codes_bp = BitPacked::encode(&codes_prim.into_array(), 2, &mut cpu_ctx) + .vortex_expect("bitpack codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); let dict = DictArray::try_new(codes_bp.into_array(), values_prim.into_array())?; let array = dict.into_array(); @@ -2162,13 +2112,13 @@ mod tests { "expected Fused for mixed-width Dict with BitPacked codes" ); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -2186,16 +2136,13 @@ mod tests { #[case] len: usize, #[case] dict_values: Vec, ) -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let dict_size = dict_values.len(); let codes: Vec = (0..len).map(|i| (i % dict_size) as u8).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes.clone()), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack codes"); + let codes_bp = BitPacked::encode(&codes_prim.into_array(), bit_width, &mut cpu_ctx) + .vortex_expect("bitpack codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); let dict = DictArray::try_new(codes_bp.into_array(), values_prim.into_array())?; let array = dict.into_array(); @@ -2206,13 +2153,13 @@ mod tests { "expected Fused for mixed-width Dict with BitPacked codes" ); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -2229,16 +2176,13 @@ mod tests { #[case] len: usize, #[case] dict_values: Vec, ) -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let dict_size = dict_values.len(); let codes: Vec = (0..len).map(|i| (i % dict_size) as u8).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes.clone()), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack codes"); + let codes_bp = BitPacked::encode(&codes_prim.into_array(), bit_width, &mut cpu_ctx) + .vortex_expect("bitpack codes"); let codes_for = FoR::try_new(codes_bp.into_array(), Scalar::from(0u8))?; let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); let dict = DictArray::try_new(codes_for.into_array(), values_prim.into_array())?; @@ -2250,13 +2194,13 @@ mod tests { "expected Fused for mixed-width Dict with FoR(BitPacked) codes" ); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -2282,19 +2226,15 @@ mod tests { #[case] ends: Vec, #[case] values: Vec, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let ends_u64: Vec = ends.iter().map(|e| (*e).into()).collect(); let len = *ends_u64.last().unwrap() as usize; let bit_width = 64 - ends_u64.iter().max().unwrap().leading_zeros() as u8; let ends_prim = PrimitiveArray::new(Buffer::from(ends.clone()), NonNullable); - let ends_bp = BitPacked::encode( - &ends_prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack ends"); + let ends_bp = BitPacked::encode(&ends_prim.into_array(), bit_width, &mut ctx) + .vortex_expect("bitpack ends"); let values_prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); let re = RunEnd::new(ends_bp.into_array(), values_prim.into_array(), &mut ctx); let array = re.into_array(); @@ -2304,7 +2244,7 @@ mod tests { "expected Fused for mixed-width RunEnd with BitPacked ends" ); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); @@ -2316,7 +2256,7 @@ mod tests { prev = *end; } let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut ctx); Ok(()) } @@ -2326,21 +2266,17 @@ mod tests { /// at native width and widens to T, fusing everything. #[crate::test] async fn test_runend_mixed_width_for_bp_u16_ends_u32_values() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let ends: Vec = vec![500, 1000, 1500, 2000]; let values: Vec = vec![100, 200, 300, 400]; let len = 2000usize; let ends_prim = PrimitiveArray::new(Buffer::from(ends.clone()), NonNullable); - let ends_bp = BitPacked::encode( - &ends_prim.into_array(), - 11, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack ends"); + let ends_bp = + BitPacked::encode(&ends_prim.into_array(), 11, &mut ctx).vortex_expect("bitpack ends"); let ends_for = FoR::try_new(ends_bp.into_array(), Scalar::from(0u16))?; let values_prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); let re = RunEnd::new(ends_for.into_array(), values_prim.into_array(), &mut ctx); let array = re.into_array(); @@ -2350,7 +2286,7 @@ mod tests { "expected Fused for mixed-width RunEnd with FoR(BitPacked) ends" ); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); @@ -2361,13 +2297,14 @@ mod tests { prev = end; } let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut ctx); Ok(()) } #[crate::test] async fn test_sliced_dict_mixed_width() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Sliced Dict with u8 codes and u32 values — combines PartiallyFused + slice handling. let dict_values: Vec = vec![100, 200, 300, 400]; let full_len = 4096; @@ -2382,7 +2319,7 @@ mod tests { // Slice from 1000..3000 let sliced = dict.into_array().slice(1000..3000)?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&sliced, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); @@ -2391,7 +2328,7 @@ mod tests { .map(|&c| dict_values[c as usize]) .collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut ctx); Ok(()) } @@ -2401,7 +2338,7 @@ mod tests { /// (the bit-pattern for i32(-1)), not u32(0x000000FF) = 255. #[crate::test] fn test_load_element_sign_extends_i8_to_u32() -> VortexResult<()> { - let cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let i8_values: Vec = vec![-1, -2, -3, 127, -128, 0, 1, 42]; let len = i8_values.len(); @@ -2439,7 +2376,7 @@ mod tests { /// Same as above but for i16 → u32 widening. #[crate::test] fn test_load_element_sign_extends_i16_to_u32() -> VortexResult<()> { - let cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let i16_values: Vec = vec![-1, -256, -32768, 32767, 0, 1, -100, 12345]; let len = i16_values.len(); @@ -2478,20 +2415,20 @@ mod tests { /// Nullable Primitive array — LOAD source with validity propagated. #[crate::test] async fn test_nullable_primitive() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let array = PrimitiveArray::from_option_iter( (0..2048u32).map(|i| if i % 3 == 0 { None } else { Some(i) }), ); - let cpu = crate::canonicalize_cpu(array.clone())?.into_array(); - let gpu = try_gpu_dispatch(&array.into_array(), &mut cuda_ctx) + let gpu = try_gpu_dispatch(&array.clone().into_array(), &mut cuda_ctx) .await? .into_host() .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(array, gpu, &mut cpu_ctx); Ok(()) } @@ -2501,7 +2438,8 @@ mod tests { /// validity, so this produces a real nullable FoR(BitPacked) tree. #[crate::test] async fn test_nullable_for_bitpacked() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let len = 2048; let reference = 1000u32; @@ -2517,7 +2455,6 @@ mod tests { }) .collect(); let prim = PrimitiveArray::from_option_iter(values.iter().copied()); - let cpu = crate::canonicalize_cpu(prim.clone())?.into_array(); // FoR encoding: subtract reference to get residuals [0..63]. // Null positions get 0 (from from_option_iter), which is fine — @@ -2526,11 +2463,7 @@ mod tests { PrimitiveArray::from_option_iter(values.iter().map(|v| v.map(|x| x - reference))); // BitPacked::encode preserves nullable validity from the input. - let bp = BitPacked::encode( - &residuals.into_array(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&residuals.into_array(), 6, &mut cpu_ctx)?; let for_arr = FoR::try_new(bp.into_array(), reference.into())?; // Verify the plan actually fuses (not just a LOAD). @@ -2548,14 +2481,14 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(prim, gpu, &mut cpu_ctx); Ok(()) } /// AllInvalid array — kernel should be skipped entirely. #[crate::test] async fn test_all_invalid_skips_kernel() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let array = PrimitiveArray::new(Buffer::from(vec![0u32; 2048]), Validity::AllInvalid); @@ -2573,19 +2506,19 @@ mod tests { /// AllValid nullable array — should fuse and produce AllValid output. #[crate::test] async fn test_all_valid_nullable() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let values: Vec = (0..2048).collect(); let array = PrimitiveArray::new(Buffer::from(values.clone()), Validity::AllValid); - let cpu = crate::canonicalize_cpu(array.clone())?.into_array(); - let gpu = try_gpu_dispatch(&array.into_array(), &mut cuda_ctx) + let gpu = try_gpu_dispatch(&array.clone().into_array(), &mut cuda_ctx) .await? .into_host() .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(array, gpu, &mut cpu_ctx); Ok(()) } @@ -2609,16 +2542,15 @@ mod tests { /// Dict with non-nullable codes but nullable values should still fuse. #[crate::test] async fn test_dict_nullable_values_fuses() -> VortexResult<()> { - use vortex::buffer::buffer; - - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let codes = PrimitiveArray::new(buffer![0u32, 1, 2, 2, 1, 0], NonNullable); let values = PrimitiveArray::from_option_iter([Some(10u32), None, Some(30)]); let dict = DictArray::try_new(codes.into_array(), values.into_array())?; - let cpu = crate::canonicalize_cpu(dict.clone())?.into_array(); let gpu = dict + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -2626,7 +2558,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(dict, gpu, &mut cpu_ctx); Ok(()) } @@ -2634,10 +2566,8 @@ mod tests { /// Validity must survive through fused dispatch and into the filter. #[crate::test] async fn test_nullable_fused_then_filter() -> VortexResult<()> { - use vortex::array::arrays::FilterArray; - use vortex::mask::Mask; - - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let len = 2048usize; let values: Vec> = (0..len) @@ -2655,8 +2585,8 @@ mod tests { let mask = Mask::from_iter((0..len).map(|i| i % 2 == 0)); let filter_array = FilterArray::try_new(prim.into_array(), mask)?; - let cpu = crate::canonicalize_cpu(filter_array.clone())?.into_array(); let gpu = filter_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -2664,14 +2594,14 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(filter_array, gpu, &mut cpu_ctx); Ok(()) } /// Empty nullable array should preserve nullability. #[crate::test] async fn test_empty_nullable_array() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let array = PrimitiveArray::new(Buffer::::empty(), Validity::AllValid); let result = try_gpu_dispatch(&array.into_array(), &mut cuda_ctx).await?; @@ -2687,6 +2617,7 @@ mod tests { #[crate::test] async fn test_bitpacked_with_patches() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let len = 3000; let bit_width: u8 = 4; let max_val = (1u32 << bit_width) - 1; @@ -2701,16 +2632,12 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); let array = bp.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&array, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( &cuda_ctx, @@ -2731,6 +2658,7 @@ mod tests { #[case] len: usize, #[case] slice_range: Option>, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width: u8 = 4; let max_val = (1u32 << bit_width) - 1; let values: Vec = (0..len) @@ -2744,11 +2672,7 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); let (array, expected) = if let Some(range) = slice_range { @@ -2758,7 +2682,7 @@ mod tests { (bp.into_array(), values) }; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&array, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( &cuda_ctx, @@ -2772,6 +2696,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_with_patches() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let len = 3000; let bit_width: u8 = 6; let reference = 42u32; @@ -2788,17 +2713,13 @@ mod tests { let all_values: Vec = residuals.iter().map(|&v| v + reference).collect(); let prim = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); let for_arr = FoR::try_new(bp.into_array(), Scalar::from(reference))?; let array = for_arr.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&array, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( &cuda_ctx, @@ -2812,6 +2733,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_with_patches_sliced() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let len = 5000; let bit_width: u8 = 6; let reference = 42u32; @@ -2828,11 +2750,7 @@ mod tests { let all_values: Vec = residuals.iter().map(|&v| v + reference).collect(); let prim = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); let for_arr = FoR::try_new(bp.into_array(), Scalar::from(reference))?; @@ -2840,7 +2758,7 @@ mod tests { let sliced = for_arr.into_array().slice(range.clone())?; let expected = all_values[range].to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( &cuda_ctx, @@ -2861,7 +2779,7 @@ mod tests { #[case] len: usize, #[case] slice_range: Option>, ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut values: Vec = (0..len).map(|i| (i as f32) * 1.1).collect(); // Insert exception values that ALP can't encode. values[0] = 99.9; @@ -2887,7 +2805,7 @@ mod tests { let cpu_decoded = array.clone().execute::(&mut ctx)?; let expected: Vec = cpu_decoded.as_slice::().to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&array, &mut cuda_ctx).await?; let actual = run_dispatch_plan_f32( &cuda_ctx, @@ -2914,6 +2832,7 @@ mod tests { /// u8 BitPacked with patches (bit_width=3, patch values > 7). #[crate::test] async fn test_bitpacked_with_patches_u8() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width: u8 = 3; let len = 3000usize; let max_val = (1u8 << bit_width) - 1; @@ -2928,25 +2847,22 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&bp.into_array(), &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut ctx); Ok(()) } /// u16 BitPacked with patches (bit_width=6, patch values > 63). #[crate::test] async fn test_bitpacked_with_patches_u16() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let bit_width: u8 = 6; let len = 3000usize; let max_val = (1u16 << bit_width) - 1; @@ -2961,25 +2877,22 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut cpu_ctx)?; assert!(bp.patches().is_some(), "expected patches"); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&bp.into_array(), &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } /// u64 BitPacked with patches (bit_width=4, patch values > 15). #[crate::test] async fn test_bitpacked_with_patches_u64() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let bit_width: u8 = 4; let len = 3000usize; let max_val = (1u64 << bit_width) - 1; @@ -2994,25 +2907,22 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut cpu_ctx)?; assert!(bp.patches().is_some(), "expected patches"); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&bp.into_array(), &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } /// Dict where codes are BitPacked u32 with patches exceeding the bit width. #[crate::test] async fn test_dict_bitpacked_codes_with_patches() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dict_values: Vec = (0..256).map(|i| i * 1000 + 42).collect(); let len = 3000; let bit_width: u8 = 4; @@ -3030,17 +2940,13 @@ mod tests { let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let codes_bp = BitPacked::encode(&codes_prim.into_array(), bit_width, &mut ctx)?; assert!(codes_bp.patches().is_some(), "expected patches on codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values), NonNullable); let dict = DictArray::try_new(codes_bp.into_array(), values_prim.into_array())?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&dict.into_array(), &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan(&cuda_ctx, len, &plan.dispatch_plan, plan.shared_mem_bytes)?; @@ -3051,6 +2957,7 @@ mod tests { /// Patches placed exactly at FastLanes chunk boundaries (1024-element chunks). #[crate::test] async fn test_bitpacked_patches_at_chunk_boundaries() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let len = 4096usize; let bit_width: u8 = 4; let max_val = (1u32 << bit_width) - 1; @@ -3062,14 +2969,10 @@ mod tests { values[2048] = 4000; // start of chunk 2 let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&bp.into_array(), &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan(&cuda_ctx, len, &plan.dispatch_plan, plan.shared_mem_bytes)?; @@ -3080,6 +2983,7 @@ mod tests { /// Large array (100k elements) spanning many blocks with sparse patches. #[crate::test] async fn test_bitpacked_large_array_with_patches() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let len = 100_000usize; let bit_width: u8 = 6; let max_val = (1u32 << bit_width) - 1; @@ -3094,14 +2998,10 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&bp.into_array(), &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan(&cuda_ctx, len, &plan.dispatch_plan, plan.shared_mem_bytes)?; @@ -3113,7 +3013,8 @@ mod tests { /// dispatch alongside patch application. #[crate::test] async fn test_nullable_bitpacked_with_patches() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let len = 3000usize; let bit_width: u8 = 4; @@ -3133,13 +3034,8 @@ mod tests { .collect(); let prim = PrimitiveArray::from_option_iter(values.iter().copied()); - let cpu = crate::canonicalize_cpu(prim.clone())?.into_array(); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.clone().into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); let gpu = try_gpu_dispatch(&bp.into_array(), &mut cuda_ctx) @@ -3148,13 +3044,14 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(prim, gpu, &mut ctx); Ok(()) } /// Extreme case: ALL values are patches (bit_width=1, every value > 1). #[crate::test] async fn test_bitpacked_all_patches() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width: u8 = 1; let len = 2000usize; // All values >= 2, so every single element exceeds max storable (1) and @@ -3162,14 +3059,10 @@ mod tests { let values: Vec = (0..len).map(|i| (i as u32) + 2).collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&bp.into_array(), &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan(&cuda_ctx, len, &plan.dispatch_plan, plan.shared_mem_bytes)?; diff --git a/vortex-cuda/src/executor.rs b/vortex-cuda/src/executor.rs index 069941c266c..6a03bfc2dca 100644 --- a/vortex-cuda/src/executor.rs +++ b/vortex-cuda/src/executor.rs @@ -21,16 +21,22 @@ use vortex::array::ArrayVTable; use vortex::array::Canonical; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; +use vortex::array::arrays::BoolArray; use vortex::array::arrays::Extension; use vortex::array::arrays::ExtensionArray; use vortex::array::arrays::Struct; use vortex::array::arrays::StructArray; +use vortex::array::arrays::bool::BoolDataParts; use vortex::array::arrays::extension::ExtensionArrayExt; use vortex::array::arrays::struct_::StructDataParts; use vortex::array::buffer::BufferHandle; +use vortex::array::validity::Validity; +use vortex::dtype::DType; +use vortex::dtype::Nullability; 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::CudaSession; @@ -96,7 +102,7 @@ pub struct CudaExecutionCtx { impl CudaExecutionCtx { /// Creates a new CUDA execution context. pub(crate) fn new(stream: VortexCudaStream, ctx: ExecutionCtx) -> Self { - let cuda_session = ctx.session().cuda_session().clone(); + let cuda_session = (*ctx.session().cuda_session()).clone(); Self { stream, ctx, @@ -375,6 +381,34 @@ pub trait CudaExecute: 'static + Send + Sync + Debug { -> VortexResult; } +pub(crate) async fn execute_validity_cuda( + validity: Validity, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let Validity::Array(array) = validity else { + return Ok(validity); + }; + + vortex_ensure!(array.len() == len, "validity array length mismatch"); + vortex_ensure!( + matches!(array.dtype(), DType::Bool(Nullability::NonNullable)), + "validity array must be non-nullable boolean, got {}", + array.dtype() + ); + + let canonical = array.execute_cuda(ctx).await?; + let Canonical::Bool(bool_array) = canonical else { + vortex_bail!("CUDA validity execution produced {}", canonical.dtype()); + }; + + let BoolDataParts { bits, meta } = bool_array.into_data().into_parts(len); + let bits = ctx.ensure_on_device(bits).await?; + Ok(Validity::Array( + BoolArray::new_handle(bits, meta.offset(), meta.len(), Validity::NonNullable).into_array(), + )) +} + /// Extension trait for executing arrays on CUDA. #[async_trait] pub trait CudaArrayExt { diff --git a/vortex-cuda/src/hybrid_dispatch/mod.rs b/vortex-cuda/src/hybrid_dispatch/mod.rs index 9f5fde8b502..b3a0151a727 100644 --- a/vortex-cuda/src/hybrid_dispatch/mod.rs +++ b/vortex-cuda/src/hybrid_dispatch/mod.rs @@ -171,20 +171,19 @@ mod tests { use vortex::error::VortexResult; use vortex::mask::Mask; use vortex::scalar::Scalar; - use vortex::session::VortexSession; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use crate::CanonicalCudaExt; - use crate::canonicalize_cpu; use crate::executor::CudaArrayExt; use crate::session::CudaSession; fn for_bp>(values: Vec, reference: T) -> ArrayRef { + let mut ctx = array_session().create_execution_ctx(); let bp = BitPacked::encode( &PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(), 7, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .vortex_expect("bp"); FoR::try_new(bp.into_array(), reference.into()) @@ -195,26 +194,27 @@ mod tests { /// FoR(BitPacked) u32 — entire tree compiles into a single fused plan. #[crate::test] async fn test_fused() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let mut ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); let values: Vec = (0..2048).map(|i| (i % 128) as u32).collect(); let bp = BitPacked::encode( &PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(), 7, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut cpu_ctx, ) .vortex_expect("bp"); let arr = FoR::try_new(bp.into_array(), 1000u32.into()).vortex_expect("for"); - let cpu = canonicalize_cpu(arr.clone())?.into_array(); let gpu = arr + .clone() .into_array() .execute_cuda(&mut ctx) .await? .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(arr, gpu, &mut cpu_ctx); Ok(()) } @@ -222,16 +222,17 @@ mod tests { /// Exercises the unsigned type reinterpretation in CudaDispatchPlan::execute. #[crate::test] async fn test_fused_f32() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); use vortex::encodings::alp::ALP; use vortex::encodings::alp::Exponents; let mut ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); let encoded: Vec = (0i32..2048).map(|i| i % 500).collect(); let bp = BitPacked::encode( &PrimitiveArray::new(Buffer::from(encoded), NonNullable).into_array(), 9, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut cpu_ctx, ) .vortex_expect("bp"); let alp = ALP::try_new( @@ -242,21 +243,22 @@ mod tests { None, )?; - let cpu = canonicalize_cpu(alp.clone())?.into_array(); let gpu = alp + .clone() .into_array() .execute_cuda(&mut ctx) .await? .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(alp, gpu, &mut cpu_ctx); Ok(()) } /// ALP with patches — plan builder rejects it, falls back to ALPExecutor. #[crate::test] async fn test_fallback() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); use vortex::array::patches::Patches; use vortex::array::validity::Validity::NonNullable as NN; use vortex::buffer::buffer; @@ -264,7 +266,7 @@ mod tests { use vortex::encodings::alp::Exponents; let mut ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); let encoded = PrimitiveArray::new( Buffer::from((0i32..2048).map(|i| i % 500).collect::>()), NonNullable, @@ -280,15 +282,15 @@ mod tests { .unwrap(); let arr = ALP::try_new(encoded, Exponents { e: 0, f: 2 }, Some(patches))?; - let cpu = canonicalize_cpu(arr.clone())?.into_array(); let gpu = arr + .clone() .into_array() .execute_cuda(&mut ctx) .await? .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(arr, gpu, &mut cpu_ctx); Ok(()) } @@ -298,12 +300,13 @@ mod tests { #[cfg(feature = "unstable_encodings")] #[crate::test] async fn test_partial_fusion() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); use vortex::array::arrays::DictArray; use vortex::array::session::ArraySessionExt; use vortex::encodings::fastlanes; use vortex::encodings::zstd::ZstdBuffers; - let session = VortexSession::empty(); + let session = crate::cuda_session(); fastlanes::initialize(&session); session.arrays().register(ZstdBuffers); let mut ctx = CudaSession::create_execution_ctx(&session).vortex_expect("ctx"); @@ -318,7 +321,7 @@ mod tests { ) .into_array(); let vals = FoR::try_new( - BitPacked::encode(&vals, 6, &mut LEGACY_SESSION.create_execution_ctx()) + BitPacked::encode(&vals, 6, &mut cpu_ctx) .vortex_expect("bp") .into_array(), 0u32.into(), @@ -333,7 +336,7 @@ mod tests { ) .into_array(); let codes = FoR::try_new( - BitPacked::encode(&codes, 6, &mut LEGACY_SESSION.create_execution_ctx()) + BitPacked::encode(&codes, 6, &mut cpu_ctx) .vortex_expect("bp") .into_array(), 0u32.into(), @@ -354,22 +357,23 @@ mod tests { .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(cpu, gpu, &mut cpu_ctx); Ok(()) } /// Filter(FoR(BP), mask) — FoR+BP fuses via dyn dispatch, then CUB filters the result. #[crate::test] async fn test_filter_fused_child() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let mut ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); let len = 2048u32; let data: Vec = (0..len).map(|i| i % 128).collect(); let bp = BitPacked::encode( &PrimitiveArray::new(Buffer::from(data.clone()), NonNullable).into_array(), 7, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut cpu_ctx, ) .vortex_expect("bp"); let for_arr = FoR::try_new(bp.into_array(), 100u32.into()).vortex_expect("for"); @@ -386,7 +390,7 @@ mod tests { .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(cpu, gpu, &mut cpu_ctx); Ok(()) } @@ -404,11 +408,12 @@ mod tests { ))] #[crate::test] async fn test_ext_storage_gpu_decode(#[case] ext: ExtensionArray) -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let mut ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); - let expected_storage = canonicalize_cpu(ext.storage_array().clone())?.into_array(); - let expected = ExtensionArray::new(ext.ext_dtype().clone(), expected_storage).into_array(); + let expected = + ExtensionArray::new(ext.ext_dtype().clone(), ext.storage_array().clone()).into_array(); let actual = ext.into_array().execute_cuda(&mut ctx).await?; let storage = actual.as_extension().storage_array(); @@ -420,15 +425,16 @@ mod tests { assert!(!storage.is_host(), "storage was not decoded on the device"); let actual = actual.into_host().await?.into_array(); - assert_arrays_eq!(expected, actual); + assert_arrays_eq!(expected, actual, &mut cpu_ctx); Ok(()) } /// Extension over already-canonical storage executes unchanged. #[crate::test] async fn test_ext_canonical_storage() -> VortexResult<()> { + let mut cpu_ctx = array_session().create_execution_ctx(); let mut ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx"); + CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); let ext = ExtensionArray::new( Date::new(TimeUnit::Days, Nullability::NonNullable).erased(), @@ -444,7 +450,7 @@ mod tests { .into_host() .await? .into_array(); - assert_arrays_eq!(ext.into_array(), actual); + assert_arrays_eq!(ext.into_array(), actual, &mut cpu_ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/arrays/constant.rs b/vortex-cuda/src/kernel/arrays/constant.rs index 8e5966da243..d385333a024 100644 --- a/vortex-cuda/src/kernel/arrays/constant.rs +++ b/vortex-cuda/src/kernel/arrays/constant.rs @@ -199,7 +199,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::scalar::Scalar; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -224,63 +224,60 @@ mod tests { async fn test_cuda_constant_materialization( #[case] constant_array: ConstantArray, ) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 cpu_result = crate::canonicalize_cpu(constant_array.clone())?; - let gpu_result = ConstantNumericExecutor - .execute(constant_array.into_array(), &mut cuda_ctx) + .execute(constant_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU materialization failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(constant_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_constant_empty_array() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 constant_array = ConstantArray::new(42i32, 0); - let cpu_result = crate::canonicalize_cpu(constant_array.clone())?; - let gpu_result = ConstantNumericExecutor - .execute(constant_array.into_array(), &mut cuda_ctx) + .execute(constant_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU materialization failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(constant_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_constant_small_array() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Test with array smaller than one block (< 2048 elements) let constant_array = ConstantArray::new(99i32, 100); - let cpu_result = crate::canonicalize_cpu(constant_array.clone())?; - let gpu_result = ConstantNumericExecutor - .execute(constant_array.into_array(), &mut cuda_ctx) + .execute(constant_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU materialization failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(constant_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/arrays/dict.rs b/vortex-cuda/src/kernel/arrays/dict.rs index 442b730a037..394de699f49 100644 --- a/vortex-cuda/src/kernel/arrays/dict.rs +++ b/vortex-cuda/src/kernel/arrays/dict.rs @@ -308,7 +308,7 @@ mod tests { use vortex::dtype::DecimalDType; use vortex::dtype::i256; use vortex::error::VortexExpect; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -325,7 +325,8 @@ mod tests { #[crate::test] async fn test_cuda_dict_u32_values_u8_codes() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Dictionary values: [100, 200, 300, 400] @@ -338,26 +339,24 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_u64_values_u16_codes() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Dictionary values: large u64 values @@ -373,26 +372,24 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_i32_values_u32_codes() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Dictionary values: signed integers including negatives @@ -405,25 +402,23 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_large_array() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Dictionary with 256 values @@ -436,26 +431,24 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values_array.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_values_with_validity() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Dictionary values with nulls: [100, null, 300, 400] @@ -469,25 +462,23 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_codes_with_validity() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Dictionary values: [100, 200, 300, 400] @@ -506,25 +497,23 @@ mod tests { let dict_array = DictArray::try_new(codes.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_both_with_validity() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Dictionary values with nulls: [100, null, 300, 400] @@ -550,25 +539,23 @@ mod tests { let dict_array = DictArray::try_new(codes.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_i64_values_with_validity() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Dictionary values with nulls (i64) @@ -595,25 +582,23 @@ mod tests { let dict_array = DictArray::try_new(codes.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_all_valid_matches_baseline() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Non-nullable values @@ -628,19 +613,16 @@ mod tests { let dict_array = DictArray::try_new(codes.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } @@ -656,7 +638,8 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i8_values() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Precision 2 uses i8 backing type @@ -669,22 +652,21 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i16_values() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Precision 4 uses i16 backing type @@ -697,22 +679,21 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i32_values() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Precision 9 uses i32 backing type @@ -725,22 +706,21 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i64_values() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Precision 18 uses i64 backing type @@ -756,22 +736,21 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i128_values() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Precision 38 uses i128 backing type @@ -792,16 +771,14 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } @@ -815,7 +792,8 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_values_u8_codes() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 = VarBinViewArray::from_iter_str(["cat", "dog", "bird", "fish"]); @@ -825,22 +803,21 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_string_values_u16_codes() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 = VarBinViewArray::from_iter_str(["alpha", "beta", "gamma", "delta", "epsilon"]); @@ -850,22 +827,21 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_string_max_inlined_12_bytes() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Exactly 12 bytes — the maximum inlined BinaryView size @@ -877,22 +853,21 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_string_outlined_views() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // 13+ bytes — outlined BinaryViews that reference data buffers @@ -907,22 +882,21 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_string_empty_strings() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 = VarBinViewArray::from_iter_str(["", "a", ""]); @@ -932,22 +906,21 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_string_values_with_validity() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 = VarBinViewArray::from_iter_nullable_str([Some("hello"), None, Some("world")]); @@ -958,22 +931,21 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_string_outlined_with_validity() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Mix of inlined, outlined, and null dictionary values @@ -991,22 +963,21 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i256_values() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Precision 76 uses i256 backing type @@ -1027,16 +998,14 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/alp.rs b/vortex-cuda/src/kernel/encodings/alp.rs index 4a1e7ea931a..d027aa0bbd6 100644 --- a/vortex-cuda/src/kernel/encodings/alp.rs +++ b/vortex-cuda/src/kernel/encodings/alp.rs @@ -141,8 +141,7 @@ mod tests { use std::f64; use vortex::array::IntoArray; - use vortex::array::LEGACY_SESSION; - use vortex::array::VortexSessionExecute; + use vortex::array::array_session; use vortex::array::arrays::PrimitiveArray; use vortex::array::assert_arrays_eq; use vortex::array::patches::Patches; @@ -153,11 +152,10 @@ mod tests { use vortex::encodings::alp::Exponents; use vortex::encodings::alp::alp_encode; use vortex::error::VortexExpect; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; - use crate::canonicalize_cpu; use crate::executor::CudaArrayExt; use crate::session::CudaSession; @@ -171,7 +169,8 @@ mod tests { /// Patches must carry `chunk_offsets` — the fused kernel requires them. #[crate::test] async fn test_cuda_alp_decompression_f32() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // For f32 with exponents (e=0, f=2): decoded = encoded * F10[2] * IF10[0] @@ -197,17 +196,15 @@ mod tests { Some(patches), )?; - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = ALPExecutor - .execute(alp_array.into_array(), &mut cuda_ctx) + .execute(alp_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } @@ -218,7 +215,8 @@ mod tests { /// preserved through the standalone ALP GPU executor. #[crate::test] async fn test_cuda_alp_nullable_with_patches() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 will produce ALP exceptions at non-null positions. @@ -238,12 +236,11 @@ mod tests { let alp_array = alp_encode( prim.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = alp_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -251,7 +248,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } @@ -259,7 +256,8 @@ mod tests { /// elements are actually null. #[crate::test] async fn test_cuda_alp_all_valid_nullable() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); let values = PrimitiveArray::new( @@ -269,12 +267,11 @@ mod tests { let alp_array = alp_encode( values.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = alp_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -282,7 +279,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } @@ -292,7 +289,8 @@ mod tests { /// (zero patches) via the offset math rather than the NULL sentinel. #[crate::test] async fn test_cuda_alp_multi_chunk_sparse_patches() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // 3072 values (3 chunks). Inject exceptions (values ALP can't encode @@ -311,16 +309,15 @@ mod tests { let alp_array = alp_encode( prim.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; assert!( alp_array.patches().is_some(), "expected patches from ALP exceptions" ); - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = alp_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -328,7 +325,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } @@ -337,7 +334,8 @@ mod tests { /// so this guards the fast-path for the (i64, f64) kernel variant. #[crate::test] async fn test_cuda_alp_f64_multi_chunk_with_patches() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // 3072 values (3 chunks). Sprinkle exceptions into each chunk. @@ -354,16 +352,15 @@ mod tests { let alp_array = alp_encode( prim.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; assert!( alp_array.patches().is_some(), "expected patches from ALP exceptions" ); - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = alp_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -371,7 +368,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } @@ -381,7 +378,8 @@ mod tests { /// (existing tests have ≤ 6 patches per chunk). #[crate::test] async fn test_cuda_alp_dense_patches_single_chunk() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Build a 1024-element ALP array manually with exactly 40 patches @@ -410,17 +408,15 @@ mod tests { Some(patches), )?; - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = ALPExecutor - .execute(alp_array.into_array(), &mut cuda_ctx) + .execute(alp_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } @@ -429,7 +425,8 @@ mod tests { /// loop. Includes a patch in the tail. #[crate::test] async fn test_cuda_alp_partial_tail_chunk() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); let values: Buffer = (0u32..1500) @@ -439,16 +436,15 @@ mod tests { let alp_array = alp_encode( prim.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut array_session().create_execution_ctx(), )?; assert!( alp_array.patches().is_some(), "expected patches from ALP exceptions" ); - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = alp_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -456,7 +452,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index 0e528a67c1a..0f00156058b 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -177,14 +177,16 @@ where // Get CUDA view of input let input_view = device_input.cuda_view::()?; + let output_len = offset + len; + // Allocate output buffer - let output_slice = ctx.device_alloc::(len.next_multiple_of(1024))?; + 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 output_width = size_of::() * 8; let cuda_function = bitpacked_cuda_kernel(bit_width, output_width, ctx)?; - let config = bitpacked_cuda_launch_config(output_width, len)?; + let config = bitpacked_cuda_launch_config(output_width, output_len)?; // We hold this here to keep the device buffers alive. let device_patches = if let Some(patches) = patches { @@ -199,7 +201,7 @@ where let patches_arg = build_gpu_patches(device_patches.as_ref())?; - ctx.launch_kernel_config(&cuda_function, config, len, |args| { + ctx.launch_kernel_config(&cuda_function, config, output_len, |args| { args.arg(&input_view) .arg(&output_view) .arg(&reference) @@ -233,8 +235,6 @@ mod tests { use vortex::buffer::buffer; use vortex::encodings::fastlanes::BitPackedArrayExt; use vortex::error::VortexExpect; - use vortex::session::VortexSession; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use super::*; @@ -251,24 +251,19 @@ mod tests { #[case] iter: impl Iterator, #[case] bw: u8, ) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 array = PrimitiveArray::new(iter.collect::>(), NonNullable); // Last two items should be patched - let bp_with_patches = BitPacked::encode( - &array.into_array(), - bw, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp_with_patches = BitPacked::encode(&array.into_array(), bw, &mut ctx)?; assert!(bp_with_patches.patches().is_some()); - let cpu_result = crate::canonicalize_cpu(bp_with_patches.clone())?.into_array(); - let gpu_result = block_on(async { BitPackedExecutor - .execute(bp_with_patches.into_array(), &mut cuda_ctx) + .execute(bp_with_patches.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -276,14 +271,15 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(bp_with_patches, gpu_result, &mut ctx); Ok(()) } #[crate::test] fn test_patches() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 array = PrimitiveArray::new( @@ -292,18 +288,12 @@ mod tests { ); // Last two items should be patched - let bp_with_patches = BitPacked::encode( - &array.into_array(), - 9, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp_with_patches = BitPacked::encode(&array.into_array(), 9, &mut ctx)?; assert!(bp_with_patches.patches().is_some()); - let cpu_result = crate::canonicalize_cpu(bp_with_patches.clone())?.into_array(); - let gpu_result = block_on(async { BitPackedExecutor - .execute(bp_with_patches.into_array(), &mut cuda_ctx) + .execute(bp_with_patches.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -311,7 +301,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(bp_with_patches, gpu_result, &mut ctx); Ok(()) } @@ -326,7 +316,8 @@ mod tests { #[case::bw_7(7)] #[crate::test] fn test_cuda_bitunpack_u8(#[case] bit_width: u8) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 max_val = (1u8 << bit_width).saturating_sub(1); @@ -338,17 +329,11 @@ mod tests { NonNullable, ); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("operation should succeed in test"); - let cpu_result = crate::canonicalize_cpu(bitpacked_array.clone())?; - + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), bit_width, &mut ctx) + .vortex_expect("operation should succeed in test"); let gpu_result = block_on(async { BitPackedExecutor - .execute(bitpacked_array.into_array(), &mut cuda_ctx) + .execute(bitpacked_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -356,7 +341,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(bitpacked_array, gpu_result, &mut ctx); Ok(()) } @@ -379,7 +364,8 @@ mod tests { #[case::bw_15(15)] #[crate::test] fn test_cuda_bitunpack_u16(#[case] bit_width: u8) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 max_val = (1u16 << bit_width).saturating_sub(1); @@ -391,17 +377,11 @@ mod tests { NonNullable, ); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("operation should succeed in test"); - let cpu_result = crate::canonicalize_cpu(bitpacked_array.clone())?; - + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), bit_width, &mut ctx) + .vortex_expect("operation should succeed in test"); let gpu_result = block_on(async { BitPackedExecutor - .execute(bitpacked_array.into_array(), &mut cuda_ctx) + .execute(bitpacked_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -409,7 +389,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(bitpacked_array, gpu_result, &mut ctx); Ok(()) } @@ -448,7 +428,8 @@ mod tests { #[case::bw_31(31)] #[crate::test] fn test_cuda_bitunpack_u32(#[case] bit_width: u8) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 max_val = (1u32 << bit_width).saturating_sub(1); @@ -460,17 +441,11 @@ mod tests { NonNullable, ); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("operation should succeed in test"); - let cpu_result = crate::canonicalize_cpu(bitpacked_array.clone())?; - + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), bit_width, &mut ctx) + .vortex_expect("operation should succeed in test"); let gpu_result = block_on(async { BitPackedExecutor - .execute(bitpacked_array.into_array(), &mut cuda_ctx) + .execute(bitpacked_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -478,7 +453,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(bitpacked_array, gpu_result, &mut ctx); Ok(()) } @@ -549,7 +524,8 @@ mod tests { #[case::bw_63(63)] #[crate::test] fn test_cuda_bitunpack_u64(#[case] bit_width: u8) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 max_val = (1u64 << bit_width).saturating_sub(1); @@ -561,16 +537,11 @@ mod tests { NonNullable, ); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("operation should succeed in test"); - let cpu_result = crate::canonicalize_cpu(bitpacked_array.clone())?; + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), bit_width, &mut ctx) + .vortex_expect("operation should succeed in test"); let gpu_result = block_on(async { BitPackedExecutor - .execute(bitpacked_array.into_array(), &mut cuda_ctx) + .execute(bitpacked_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -578,15 +549,22 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(bitpacked_array, gpu_result, &mut ctx); Ok(()) } + #[rstest] + #[case(67, 3969)] + #[case(1, 1025)] #[crate::test] - fn test_cuda_bitunpack_sliced() -> VortexResult<()> { + fn test_cuda_bitunpack_sliced( + #[case] slice_start: usize, + #[case] slice_end: usize, + ) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); let bit_width = 32; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let max_val = (1u64 << bit_width).saturating_sub(1); @@ -598,17 +576,12 @@ mod tests { NonNullable, ); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("operation should succeed in test"); - let sliced_array = bitpacked_array.into_array().slice(67..3969)?; - let cpu_result = crate::canonicalize_cpu(sliced_array.clone())?; + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), bit_width, &mut ctx) + .vortex_expect("operation should succeed in test"); + let sliced_array = bitpacked_array.into_array().slice(slice_start..slice_end)?; let gpu_result = block_on(async { BitPackedExecutor - .execute(sliced_array, &mut cuda_ctx) + .execute(sliced_array.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -616,7 +589,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(sliced_array, gpu_result, &mut ctx); Ok(()) } @@ -634,17 +607,15 @@ mod tests { #[case] expected_offset: u16, #[case] expected_packed_len: usize, ) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); let values = PrimitiveArray::new( (0u16..4096) .map(|i| if i % 1000 == 0 { 600 } else { i % 512 }) .collect::>(), NonNullable, - ); - let bitpacked = BitPacked::encode( - &values.into_array(), - 9, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + ) + .into_array(); + let bitpacked = BitPacked::encode(&values, 9, &mut ctx)?; assert!(bitpacked.patches().is_some()); let array = if let Some(range) = range { bitpacked.into_array().slice(range)? @@ -668,7 +639,8 @@ mod tests { /// offset_within_chunk. #[crate::test] fn test_cuda_bitunpack_sliced_patches_offset_within_chunk() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Create an array with values that will generate patches. @@ -677,11 +649,7 @@ mod tests { let primitive_array = PrimitiveArray::new(buffer![100u8, 101, 102, 3, 4, 5], NonNullable); // Encode with bit width 4. First 3 elements patched, remainder will pack. - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - 4, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), 4, &mut ctx)?; assert!( bitpacked_array.patches().is_some(), "Expected patches to be present" @@ -689,12 +657,9 @@ mod tests { let sliced_array = bitpacked_array.into_array().slice(2..6)?; - let cpu_result = sliced_array - .clone() - .execute::(cuda_ctx.execution_ctx())?; let gpu_result = block_on(async { BitPackedExecutor - .execute(sliced_array, &mut cuda_ctx) + .execute(sliced_array.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -702,7 +667,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(sliced_array, gpu_result, &mut ctx); Ok(()) } @@ -710,7 +675,8 @@ mod tests { /// Test slicing a bitpacked array multiple times, accumulating offset_within_chunk. #[crate::test] fn test_cuda_bitunpack_double_sliced_patches() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Create an array with values that will generate patches. @@ -726,11 +692,7 @@ mod tests { let primitive_array = PrimitiveArray::new(Buffer::from_iter(values.iter().copied()), NonNullable); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - 9, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), 9, &mut ctx)?; assert!( bitpacked_array.patches().is_some(), "Expected patches to be present" @@ -743,12 +705,9 @@ mod tests { // covers every chunk in the packed buffer. let second_slice = first_slice.slice(50..2900)?; - let cpu_result = second_slice - .clone() - .execute::(cuda_ctx.execution_ctx())?; let gpu_result = block_on(async { BitPackedExecutor - .execute(second_slice, &mut cuda_ctx) + .execute(second_slice.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -756,7 +715,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(second_slice, gpu_result, &mut ctx); Ok(()) } @@ -764,7 +723,8 @@ mod tests { /// Test slicing to skip an entire chunk's worth of patches. #[crate::test] fn test_cuda_bitunpack_sliced_skip_first_chunk_patches() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Create patches in first chunk only, then slice past them all. @@ -782,11 +742,7 @@ mod tests { let primitive_array = PrimitiveArray::new(Buffer::from_iter(values.iter().copied()), NonNullable); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - 9, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), 9, &mut ctx)?; assert!( bitpacked_array.patches().is_some(), "Expected patches to be present" @@ -795,12 +751,9 @@ mod tests { // Slice to skip past all first chunk patches let sliced_array = bitpacked_array.into_array().slice(1024..3072)?; - let cpu_result = sliced_array - .clone() - .execute::(cuda_ctx.execution_ctx())?; let gpu_result = block_on(async { BitPackedExecutor - .execute(sliced_array, &mut cuda_ctx) + .execute(sliced_array.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -808,7 +761,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(sliced_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/date_time_parts.rs b/vortex-cuda/src/kernel/encodings/date_time_parts.rs index d8c655c2558..dcb1fd1f69e 100644 --- a/vortex-cuda/src/kernel/encodings/date_time_parts.rs +++ b/vortex-cuda/src/kernel/encodings/date_time_parts.rs @@ -72,7 +72,7 @@ impl CudaExecute for DateTimePartsExecutor { return Ok(Canonical::empty(array.dtype())); } - if matches!(validity, Validity::AllInvalid) { + if validity.definitely_all_null() { let storage_ptype = ext.storage_dtype().as_ptype(); return Ok(Canonical::Extension( TemporalArray::new_timestamp( @@ -214,7 +214,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::extension::datetime::TimeUnit; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -280,28 +280,28 @@ mod tests { #[case] subseconds: Vec, #[case] time_unit: TimeUnit, ) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 dtp_array = make_datetimeparts_array(days, seconds, subseconds, time_unit); - let cpu_result = crate::canonicalize_cpu(dtp_array.clone())?; - let gpu_result = DateTimePartsExecutor - .execute(dtp_array.into_array(), &mut cuda_ctx) + .execute(dtp_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(dtp_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_datetimeparts_large_array() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 = 2050; @@ -310,24 +310,23 @@ mod tests { let subseconds: Vec = (0..len).map(|i| (i % 1000) as i64).collect(); let dtp_array = make_datetimeparts_array(days, seconds, subseconds, TimeUnit::Milliseconds); - let cpu_result = crate::canonicalize_cpu(dtp_array.clone())?; - let gpu_result = DateTimePartsExecutor - .execute(dtp_array.into_array(), &mut cuda_ctx) + .execute(dtp_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(dtp_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_datetimeparts_with_nulls() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 days_arr = PrimitiveArray::new( @@ -358,17 +357,15 @@ mod tests { ) .vortex_expect("Failed to create DateTimePartsArray"); - let cpu_result = crate::canonicalize_cpu(dtp_array.clone())?; - let gpu_result = DateTimePartsExecutor - .execute(dtp_array.into_array(), &mut cuda_ctx) + .execute(dtp_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(dtp_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs index 70d3c71a485..5f7e639b8a1 100644 --- a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs +++ b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs @@ -70,7 +70,7 @@ mod tests { use vortex::dtype::DecimalDType; use vortex::encodings::decimal_byte_parts::DecimalByteParts; use vortex::error::VortexExpect; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; use super::*; use crate::session::CudaSession; @@ -86,7 +86,8 @@ mod tests { #[case] precision: u8, #[case] scale: i8, ) { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("create execution context"); let decimal_dtype = DecimalDType::new(precision, scale); @@ -96,14 +97,11 @@ mod tests { ) .vortex_expect("create DecimalBytePartsArray"); - let cpu_result = - crate::canonicalize_cpu(dbp_array.clone()).vortex_expect("CPU canonicalize"); - let gpu_result = DecimalBytePartsExecutor - .execute(dbp_array.into_array(), &mut cuda_ctx) + .execute(dbp_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decode"); - assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array()); + assert_arrays_eq!(dbp_array, gpu_result.into_array(), &mut ctx); } } diff --git a/vortex-cuda/src/kernel/encodings/for_.rs b/vortex-cuda/src/kernel/encodings/for_.rs index ce13f4c3b4b..6348f8479b7 100644 --- a/vortex-cuda/src/kernel/encodings/for_.rs +++ b/vortex-cuda/src/kernel/encodings/for_.rs @@ -138,9 +138,8 @@ mod tests { use vortex::encodings::fastlanes::FoRArray; use vortex::error::VortexExpect; use vortex::scalar::Scalar; - use vortex::session::VortexSession; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use super::*; use crate::CanonicalCudaExt; @@ -161,27 +160,27 @@ mod tests { #[case::u64(make_for_array((0..2050).map(|i| (i % 2050) as u64).collect(), 1000000u64))] #[crate::test] async fn test_cuda_for_decompression(#[case] for_array: FoRArray) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); - let cpu_result = crate::canonicalize_cpu(for_array.clone())?; - let gpu_result = FoRExecutor - .execute(for_array.into_array(), &mut cuda_ctx) + .execute(for_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(for_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_signed_ffor() { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); let values = (0i8..8i8) @@ -189,15 +188,13 @@ mod tests { .take(1024) .collect::>() .into_array(); - let packed = BitPacked::encode(&values, 3, &mut LEGACY_SESSION.create_execution_ctx()) + let packed = BitPacked::encode(&values, 3, &mut array_session().create_execution_ctx()) .unwrap() .into_array(); let for_array = FoR::try_new(packed, (-8i8).into()).unwrap(); - let cpu_result = crate::canonicalize_cpu(for_array.clone()).unwrap(); - let gpu_result = FoRExecutor - .execute(for_array.into_array(), &mut cuda_ctx) + .execute(for_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -205,6 +202,6 @@ mod tests { .vortex_expect("copying to host failed") .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(for_array, gpu_result, &mut ctx); } } diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index 5d3d66eaf04..eef0db049b8 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -23,7 +23,6 @@ use vortex::array::arrays::varbinview::build_views::build_views; use vortex::array::buffer::DeviceBuffer; 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::Buffer; use vortex::dtype::NativePType; @@ -62,7 +61,7 @@ impl CudaExecute for FSSTExecutor { let dtype = fsst.dtype().clone(); let validity = fsst.codes().validity()?; - if fsst.is_empty() || matches!(validity, Validity::AllInvalid) { + if fsst.is_empty() || validity.definitely_all_null() { let empty = unsafe { VarBinViewArray::new_unchecked( Buffer::::zeroed(fsst.len()), @@ -210,7 +209,7 @@ mod tests { use vortex::encodings::fsst::fsst_compress; use vortex::encodings::fsst::fsst_train_compressor; use vortex::error::VortexExpect; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -237,49 +236,47 @@ mod tests { #[case] strings: Vec>, #[case] nullability: Nullability, ) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 varbin = VarBinArray::from_iter(strings, DType::Binary(nullability)); - let compressor = fsst_train_compressor(&varbin); - let dtype = varbin.dtype().clone(); - let len = varbin.len(); + let varbin = VarBinArray::from_iter(strings, DType::Binary(nullability)).into_array(); + let compressor = fsst_train_compressor(&varbin, cuda_ctx.execution_ctx())?; let fsst_array = - fsst_compress(&varbin, len, &dtype, &compressor, cuda_ctx.execution_ctx()).into_array(); + fsst_compress(&varbin, &compressor, cuda_ctx.execution_ctx())?.into_array(); - let cpu_result = crate::canonicalize_cpu(fsst_array.clone())?; let gpu_result = FSSTExecutor - .execute(fsst_array, &mut cuda_ctx) + .execute(fsst_array.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(fsst_array, gpu_result, &mut ctx); Ok(()) } /// Exercises the multi-block grid-stride path on a larger dataset. #[crate::test] async fn test_cuda_fsst_decompression_roundtrip_large() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); use vortex_fsst::test_utils::make_fsst_clickbench_urls; - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let fsst_array = make_fsst_clickbench_urls(100_000, cuda_ctx.execution_ctx()).into_array(); - let cpu_result = crate::canonicalize_cpu(fsst_array.clone())?; let gpu_result = FSSTExecutor - .execute(fsst_array, &mut cuda_ctx) + .execute(fsst_array.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(fsst_array, gpu_result, &mut ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/runend.rs b/vortex-cuda/src/kernel/encodings/runend.rs index fca435478d8..4ac042e2972 100644 --- a/vortex-cuda/src/kernel/encodings/runend.rs +++ b/vortex-cuda/src/kernel/encodings/runend.rs @@ -75,7 +75,7 @@ impl CudaExecute for RunEndExecutor { ))); } - if matches!(values.validity()?, Validity::AllInvalid) { + if values.validity()?.definitely_all_null() { return ConstantArray::new(Scalar::null(values.dtype().clone()), output_len) .into_array() .execute::(ctx.execution_ctx()); @@ -174,18 +174,15 @@ mod tests { use vortex::encodings::runend::RunEndArray; use vortex::error::VortexExpect; use vortex::error::VortexResult; - use vortex::session::VortexSession; + use vortex_array::ExecutionCtx; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; use crate::executor::CudaArrayExt; use crate::session::CudaSession; - fn make_runend_array( - ends: Vec, - values: Vec, - ctx: &mut vortex::array::ExecutionCtx, - ) -> RunEndArray + fn make_runend_array(ends: Vec, values: Vec, ctx: &mut ExecutionCtx) -> RunEndArray where V: NativePType, E: NativePType, @@ -197,39 +194,39 @@ mod tests { RunEnd::new(ends_array, values_array, ctx) } - type RunEndBuilder = fn(&mut vortex::array::ExecutionCtx) -> RunEndArray; + type RunEndBuilder = fn(&mut ExecutionCtx) -> RunEndArray; #[rstest] - #[case::u32_ends_u8_values(|ctx: &mut vortex::array::ExecutionCtx| make_runend_array(vec![3u32, 6, 10], vec![10u8, 20, 30], ctx))] - #[case::u32_ends_u32_values(|ctx: &mut vortex::array::ExecutionCtx| make_runend_array(vec![2u32, 5, 10], vec![1u32, 2, 3], ctx))] - #[case::u32_ends_f64_values(|ctx: &mut vortex::array::ExecutionCtx| make_runend_array(vec![2u32, 5, 8], vec![1.5f64, 2.5, 3.5], ctx))] - #[case::u8_ends_i32_values(|ctx: &mut vortex::array::ExecutionCtx| make_runend_array(vec![2u8, 5, 10], vec![1i32, 2, 3], ctx))] - #[case::u32_ends_i32_values(|ctx: &mut vortex::array::ExecutionCtx| make_runend_array(vec![2u32, 5, 10], vec![1i32, 2, 3], ctx))] - #[case::u64_ends_i32_values(|ctx: &mut vortex::array::ExecutionCtx| make_runend_array(vec![2u64, 5, 10], vec![1i32, 2, 3], ctx))] + #[case::u32_ends_u8_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![3u32, 6, 10], vec![10u8, 20, 30], ctx))] + #[case::u32_ends_u32_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u32, 5, 10], vec![1u32, 2, 3], ctx))] + #[case::u32_ends_f64_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u32, 5, 8], vec![1.5f64, 2.5, 3.5], ctx))] + #[case::u8_ends_i32_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u8, 5, 10], vec![1i32, 2, 3], ctx))] + #[case::u32_ends_i32_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u32, 5, 10], vec![1i32, 2, 3], ctx))] + #[case::u64_ends_i32_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u64, 5, 10], vec![1i32, 2, 3], ctx))] #[crate::test] async fn test_cuda_runend_types(#[case] build: RunEndBuilder) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 runend_array = build(cuda_ctx.execution_ctx()); - let cpu_result = crate::canonicalize_cpu(runend_array.clone())?; - let gpu_result = RunEndExecutor - .execute(runend_array.into_array(), &mut cuda_ctx) + .execute(runend_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(runend_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_runend_large_array() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 = 41; @@ -242,46 +239,44 @@ mod tests { let runend_array = make_runend_array(ends, values, cuda_ctx.execution_ctx()); assert_eq!(runend_array.len(), total_len); - let cpu_result = crate::canonicalize_cpu(runend_array.clone())?; - let gpu_result = RunEndExecutor - .execute(runend_array.into_array(), &mut cuda_ctx) + .execute(runend_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(runend_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_runend_single_run() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 runend_array = make_runend_array(vec![100u32], vec![42i32], cuda_ctx.execution_ctx()); - let cpu_result = crate::canonicalize_cpu(runend_array.clone())?; - let gpu_result = RunEndExecutor - .execute(runend_array.into_array(), &mut cuda_ctx) + .execute(runend_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(runend_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_runend_many_small_runs() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Create an array where each run has length 1. @@ -291,24 +286,23 @@ mod tests { let runend_array = make_runend_array(ends, values, cuda_ctx.execution_ctx()); - let cpu_result = crate::canonicalize_cpu(runend_array.clone())?; - let gpu_result = RunEndExecutor - .execute(runend_array.into_array(), &mut cuda_ctx) + .execute(runend_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(runend_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_runend_nullable_values_falls_back_to_cpu() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // Build a RunEnd array whose values have Validity::Array (some nulls). @@ -321,10 +315,9 @@ 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()); - let cpu_result = crate::canonicalize_cpu(runend_array.clone())?.into_array(); - // 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) .await @@ -333,7 +326,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(runend_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/sequence.rs b/vortex-cuda/src/kernel/encodings/sequence.rs index c0721b8d920..e52a706c8b7 100644 --- a/vortex-cuda/src/kernel/encodings/sequence.rs +++ b/vortex-cuda/src/kernel/encodings/sequence.rs @@ -92,7 +92,7 @@ mod tests { use vortex::dtype::Nullability; use vortex::encodings::sequence::Sequence; use vortex::scalar::PValue; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::CudaSession; @@ -125,14 +125,15 @@ mod tests { len: usize, nullability: Nullability, ) { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()).unwrap(); + 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(base, multiplier, nullability, len).unwrap(); - - let cpu_result = crate::canonicalize_cpu(array.clone()).unwrap().into_array(); + let array = Sequence::try_new_typed(base, multiplier, nullability, len) + .unwrap() + .into_array(); let gpu_result = SequenceExecutor - .execute(array.into_array(), &mut cuda_ctx) + .execute(array.clone(), &mut cuda_ctx) .await .unwrap() .into_host() @@ -140,6 +141,6 @@ mod tests { .unwrap() .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(array, gpu_result, &mut ctx); } } diff --git a/vortex-cuda/src/kernel/encodings/zigzag.rs b/vortex-cuda/src/kernel/encodings/zigzag.rs index 98f9534be0a..107ad2b36a1 100644 --- a/vortex-cuda/src/kernel/encodings/zigzag.rs +++ b/vortex-cuda/src/kernel/encodings/zigzag.rs @@ -105,7 +105,7 @@ mod tests { use vortex::buffer::Buffer; use vortex::encodings::zigzag::ZigZag; use vortex::error::VortexExpect; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -113,7 +113,8 @@ mod tests { #[crate::test] async fn test_cuda_zigzag_decompression_u32() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); // ZigZag encoding: 0->0, 1->-1, 2->1, 3->-2, 4->2, ... @@ -122,19 +123,18 @@ mod tests { let zigzag_array = ZigZag::try_new( PrimitiveArray::new(Buffer::from(encoded_data), NonNullable).into_array(), - )?; - - let cpu_result = crate::canonicalize_cpu(zigzag_array.clone())?; + )? + .into_array(); let gpu_result = ZigZagExecutor - .execute(zigzag_array.into_array(), &mut cuda_ctx) + .execute(zigzag_array.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(zigzag_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/zstd.rs b/vortex-cuda/src/kernel/encodings/zstd.rs index abadeb1698f..6d14cc46e86 100644 --- a/vortex-cuda/src/kernel/encodings/zstd.rs +++ b/vortex-cuda/src/kernel/encodings/zstd.rs @@ -356,7 +356,7 @@ mod tests { use vortex::array::assert_arrays_eq; use vortex::encodings::zstd::Zstd; use vortex::error::VortexResult; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -365,7 +365,8 @@ mod tests { #[crate::test] async fn test_cuda_zstd_decompression_utf8() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 = VarBinViewArray::from_iter_str([ @@ -385,13 +386,14 @@ mod tests { .execute(zstd_array.into_array(), &mut cuda_ctx) .await?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array()); + assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array(), &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_zstd_decompression_multiple_frames() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 = VarBinViewArray::from_iter_str([ @@ -421,13 +423,14 @@ mod tests { .execute(zstd_array.into_array(), &mut cuda_ctx) .await?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array()); + assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array(), &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_zstd_decompression_sliced() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 = VarBinViewArray::from_iter_str([ @@ -448,12 +451,11 @@ mod tests { // Slice the array to get a subset (indices 2..7) let sliced_zstd = zstd_array.slice(2..7)?; - let cpu_result = crate::canonicalize_cpu(sliced_zstd.clone())?; let gpu_result = ZstdExecutor .execute(sliced_zstd.clone(), &mut cuda_ctx) .await?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array()); + assert_arrays_eq!(sliced_zstd, gpu_result.into_array(), &mut ctx); Ok(()) } @@ -462,7 +464,8 @@ mod tests { /// correct results instead of panicking. #[crate::test] async fn test_cuda_zstd_nullable_falls_back_to_cpu() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 = VarBinViewArray::from_iter_nullable_str([ @@ -474,20 +477,19 @@ mod tests { Some("another string"), ]); - let zstd_array = Zstd::from_var_bin_view(&strings, 3, 0, cuda_ctx.execution_ctx())?; - - let cpu_result = crate::canonicalize_cpu(zstd_array.clone())?.into_array(); + let zstd_array = + Zstd::from_var_bin_view(&strings, 3, 0, cuda_ctx.execution_ctx())?.into_array(); // execute_cuda should fall back to CPU and still produce the correct result. let gpu_result = zstd_array - .into_array() + .clone() .execute_cuda(&mut cuda_ctx) .await? .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(zstd_array, gpu_result, &mut ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/zstd_buffers.rs b/vortex-cuda/src/kernel/encodings/zstd_buffers.rs index 8a67a338ca5..7de2158ac48 100644 --- a/vortex-cuda/src/kernel/encodings/zstd_buffers.rs +++ b/vortex-cuda/src/kernel/encodings/zstd_buffers.rs @@ -226,7 +226,8 @@ mod tests { use vortex::array::assert_arrays_eq; use vortex::error::VortexExpect; use vortex::error::VortexResult; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use super::*; use crate::CanonicalCudaExt; @@ -234,26 +235,27 @@ mod tests { #[crate::test] async fn test_cuda_zstd_buffers_decompression_primitive() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); let input = PrimitiveArray::from_iter(0i64..1024).into_array(); - let compressed = ZstdBuffers::compress(&input, 3, &VortexSession::empty())?; + let compressed = ZstdBuffers::compress(&input, 3, &crate::cuda_session())?.into_array(); - let cpu_result = crate::canonicalize_cpu(compressed.clone())?; let gpu_result = ZstdBuffersExecutor - .execute(compressed.into_array(), &mut cuda_ctx) + .execute(compressed.clone(), &mut cuda_ctx) .await? .into_host() .await?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array()); + assert_arrays_eq!(compressed, gpu_result.into_array(), &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_zstd_buffers_decompression_varbinview() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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"); let input = VarBinViewArray::from_iter_str([ @@ -265,16 +267,15 @@ mod tests { "baz", ]) .into_array(); - let compressed = ZstdBuffers::compress(&input, 3, &VortexSession::empty())?; + let compressed = ZstdBuffers::compress(&input, 3, &crate::cuda_session())?.into_array(); - let cpu_result = crate::canonicalize_cpu(compressed.clone())?; let gpu_result = ZstdBuffersExecutor - .execute(compressed.into_array(), &mut cuda_ctx) + .execute(compressed.clone(), &mut cuda_ctx) .await? .into_host() .await?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array()); + assert_arrays_eq!(compressed, gpu_result.into_array(), &mut ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/filter/decimal.rs b/vortex-cuda/src/kernel/filter/decimal.rs index 566b1c3313d..dc10b908461 100644 --- a/vortex-cuda/src/kernel/filter/decimal.rs +++ b/vortex-cuda/src/kernel/filter/decimal.rs @@ -48,7 +48,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::mask::Mask; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::FilterExecutor; @@ -89,29 +89,29 @@ mod tests { #[case] input: DecimalArray, #[case] mask: Mask, ) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 CUDA execution context"); let filter_array = FilterArray::try_new(input.clone().into_array(), mask.clone())?; - let cpu_result = crate::canonicalize_cpu(filter_array.clone())?.into_array(); - let gpu_result = FilterExecutor - .execute(filter_array.into_array(), &mut cuda_ctx) + .execute(filter_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU filter failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(filter_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_gpu_filter_decimal_large_array() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 CUDA execution context"); // Create a large array to test multi-block execution @@ -123,18 +123,16 @@ mod tests { let filter_array = FilterArray::try_new(input.into_array(), mask)?; - let cpu_result = crate::canonicalize_cpu(filter_array.clone())?.into_array(); - let gpu_result = FilterExecutor - .execute(filter_array.into_array(), &mut cuda_ctx) + .execute(filter_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU filter failed") .into_host() .await? .into_array(); - assert_eq!(cpu_result.len(), gpu_result.len()); - assert_arrays_eq!(cpu_result, gpu_result); + assert_eq!(filter_array.len(), gpu_result.len()); + assert_arrays_eq!(filter_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/filter/primitive.rs b/vortex-cuda/src/kernel/filter/primitive.rs index 6bff0b20c70..c2a21ccc7d7 100644 --- a/vortex-cuda/src/kernel/filter/primitive.rs +++ b/vortex-cuda/src/kernel/filter/primitive.rs @@ -46,7 +46,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::mask::Mask; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::FilterExecutor; @@ -83,29 +83,29 @@ mod tests { #[case] input: PrimitiveArray, #[case] mask: Mask, ) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 CUDA execution context"); let filter_array = FilterArray::try_new(input.clone().into_array(), mask.clone())?; - let cpu_result = crate::canonicalize_cpu(filter_array.clone())?.into_array(); - let gpu_result = FilterExecutor - .execute(filter_array.into_array(), &mut cuda_ctx) + .execute(filter_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU filter failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(filter_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_gpu_filter_large_array() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 CUDA execution context"); // Create a large array to test multi-block execution @@ -117,18 +117,16 @@ mod tests { let filter_array = FilterArray::try_new(input.into_array(), mask)?; - let cpu_result = crate::canonicalize_cpu(filter_array.clone())?.into_array(); - let gpu_result = FilterExecutor - .execute(filter_array.into_array(), &mut cuda_ctx) + .execute(filter_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU filter failed") .into_host() .await? .into_array(); - assert_eq!(cpu_result.len(), gpu_result.len()); - assert_arrays_eq!(cpu_result, gpu_result); + assert_eq!(filter_array.len(), gpu_result.len()); + assert_arrays_eq!(filter_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/filter/varbinview.rs b/vortex-cuda/src/kernel/filter/varbinview.rs index 23a6c91b048..7bbafff824c 100644 --- a/vortex-cuda/src/kernel/filter/varbinview.rs +++ b/vortex-cuda/src/kernel/filter/varbinview.rs @@ -46,7 +46,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::mask::Mask; - use vortex::session::VortexSession; + use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::FilterExecutor; @@ -69,22 +69,21 @@ mod tests { #[case] input: VarBinViewArray, #[case] mask: Mask, ) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) + 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 CUDA execution context"); let filter_array = FilterArray::try_new(input.into_array(), mask.clone())?; - let cpu_result = crate::canonicalize_cpu(filter_array.clone())?.into_array(); - let gpu_result = FilterExecutor - .execute(filter_array.into_array(), &mut cuda_ctx) + .execute(filter_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU filter failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(filter_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/mod.rs b/vortex-cuda/src/kernel/mod.rs index 5665bc6c3f3..323458db2d8 100644 --- a/vortex-cuda/src/kernel/mod.rs +++ b/vortex-cuda/src/kernel/mod.rs @@ -4,8 +4,6 @@ //! CUDA kernel loading and management. use std::fmt::Debug; -use std::path::Path; -use std::path::PathBuf; use std::sync::Arc; use cudarc::driver::CudaContext; @@ -22,6 +20,10 @@ 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")); +} mod encodings; mod filter; mod patches; @@ -235,24 +237,9 @@ impl KernelLoader { let module = if let Some(entry) = self.modules.get(module_name) { Arc::clone(entry.value()) } else { - let ptx_path = Self::ptx_path_for_module(module_name); - - // Compile and load the CUDA module. - let module = cuda_context - .load_module(Ptx::from_file(&ptx_path)) - .map_err(|e| { - vortex_err!( - "Failed to load CUDA module {}, ptx path {}: {}", - module_name, - ptx_path.display(), - e - ) - })?; - - // Cache the module + let module = Self::load_module(module_name, cuda_context)?; self.modules .insert(module_name.to_string(), Arc::clone(&module)); - module }; @@ -262,22 +249,18 @@ impl KernelLoader { .map_err(|e| vortex_err!("Failed to load kernel function '{}': {}", kernel_name, e)) } - /// Returns the PTX file path for a given module name. - /// - /// Checks for `VORTEX_CUDA_KERNELS_DIR` environment variable at runtime first, - /// falling back to the path baked in at compile time by build.rs. - /// - /// # Arguments - /// - /// * `module_name` - Name of the module - /// - /// # Returns - /// - /// The full path to the PTX file - fn ptx_path_for_module(module_name: &str) -> PathBuf { - let kernels_dir = std::env::var("VORTEX_CUDA_KERNELS_DIR") - .unwrap_or_else(|_| env!("VORTEX_CUDA_KERNELS_DIR").to_string()); - Path::new(&kernels_dir).join(format!("{}.ptx", module_name)) + /// Loads a CUDA module from PTX 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(|| { + vortex_err!("CUDA module {module_name} was not embedded at build time") + })?; + + cuda_context + .load_module(Ptx::from_src(ptx)) + .map_err(|e| vortex_err!("Failed to load embedded CUDA module {module_name}: {e}")) } } diff --git a/vortex-cuda/src/kernel/patches/mod.rs b/vortex-cuda/src/kernel/patches/mod.rs index 3eb3e2ced93..bd6d2b737db 100644 --- a/vortex-cuda/src/kernel/patches/mod.rs +++ b/vortex-cuda/src/kernel/patches/mod.rs @@ -175,8 +175,8 @@ mod tests { use cudarc::driver::DeviceRepr; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; - use vortex::array::LEGACY_SESSION; use vortex::array::VortexSessionExecute; + use vortex::array::array_session; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::assert_arrays_eq; @@ -188,7 +188,6 @@ mod tests { use vortex::dtype::DType; use vortex::dtype::NativePType; use vortex::dtype::Nullability; - use vortex::session::VortexSession; use crate::CanonicalCudaExt; use crate::CudaDeviceBuffer; @@ -217,8 +216,8 @@ mod tests { } async fn full_test_case() { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()).unwrap(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).unwrap(); + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::from_iter(0..128); let values = force_cast::(values, &mut ctx); @@ -265,7 +264,7 @@ mod tests { .unwrap() .into_primitive(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut ctx); } fn force_cast(array: PrimitiveArray, ctx: &mut ExecutionCtx) -> PrimitiveArray { diff --git a/vortex-cuda/src/kernel/patches/types.rs b/vortex-cuda/src/kernel/patches/types.rs index f90d7251c09..3bfe2270b66 100644 --- a/vortex-cuda/src/kernel/patches/types.rs +++ b/vortex-cuda/src/kernel/patches/types.rs @@ -227,8 +227,8 @@ mod tests { use vortex::array::buffer::BufferHandle; use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; - use vortex::session::VortexSession; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::patches::Patches; use vortex_error::VortexResult; @@ -274,7 +274,7 @@ mod tests { #[case] expected_offset: usize, #[case] expected_chunk_offsets: Vec, ) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let indices = PrimitiveArray::from_iter([100u32, 1100, 2100, 3100, 4100]); let values = PrimitiveArray::from_iter([10u32, 11, 12, 13, 14]); let chunk_offsets = PrimitiveArray::from_iter([0u32, 1, 2, 3, 4, 5]); @@ -318,7 +318,8 @@ mod tests { #[case] chunk_offsets: ArrayRef, #[case] expected: ArrayRef, ) -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let indices = PrimitiveArray::from_iter([100u32, 1100, 2100, 3100, 4100]); let values = PrimitiveArray::from_iter([10u32, 11, 12, 13, 14]); let patches = Patches::new( @@ -338,7 +339,7 @@ mod tests { NonNullable, ) .into_array(); - assert_arrays_eq!(expected, actual); + assert_arrays_eq!(expected, actual, &mut ctx); assert_eq!(device_patches.n_chunks, 3); assert_eq!(device_patches.offset, 1024); assert_eq!(device_patches.offset_within_chunk, 0); diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 20aa26e0b3b..1b1d37cd14d 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -40,6 +40,7 @@ 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::LayoutChildType; use vortex::layout::LayoutChildren; use vortex::layout::LayoutEncodingRef; @@ -64,6 +65,7 @@ use vortex::scalar::ScalarTruncation; use vortex::scalar::lower_bound; use vortex::scalar::upper_bound; 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; @@ -130,7 +132,8 @@ impl VTable for CudaFlat { type Metadata = ProstMetadata; fn id(_encoding: &Self::Encoding) -> LayoutId { - LayoutId::new("vortex.cuda_flat") + static ID: CachedId = CachedId::new("vortex.cuda_flat"); + *ID } fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef { @@ -198,7 +201,7 @@ impl VTable for CudaFlat { metadata: &::Output, segment_ids: Vec, _children: &dyn LayoutChildren, - ctx: &ReadContext, + build_ctx: &LayoutBuildContext<'_>, ) -> VortexResult { if segment_ids.len() != 1 { vortex_bail!("CudaFlatLayout must have exactly one segment ID"); @@ -212,7 +215,7 @@ impl VTable for CudaFlat { row_count, dtype: dtype.clone(), segment_id: segment_ids[0], - ctx: ctx.clone(), + ctx: build_ctx.array_read_ctx.clone(), array_tree: ByteBuffer::from(metadata.array_encoding_tree.clone()), host_buffers: Arc::new(host_buffers), }) @@ -331,12 +334,12 @@ impl LayoutReader for CudaFlatReader { let array = array.apply(&expr)?; let array = array.filter(mask.clone())?; let mut ctx = session.create_execution_ctx(); - let array_mask = array.execute::(&mut ctx)?; + let array_mask = array.null_as_false().execute(&mut ctx)?; mask.intersect_by_rank(&array_mask) } else { let array = array.apply(&expr)?; let mut ctx = session.create_execution_ctx(); - let array_mask = array.execute::(&mut ctx)?; + let array_mask = array.null_as_false().execute(&mut ctx)?; mask.bitand(&array_mask) }; diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 9e9de33a288..07cfe40b06f 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -3,8 +3,8 @@ //! CUDA support for Vortex arrays. -use std::process::Command; - +use cudarc::driver::CudaContext; +use cudarc::driver::sys; use tracing::info; pub mod arrow; @@ -25,6 +25,7 @@ mod stream_pool; pub use arrow::ArrowDeviceArrayWithSchema; pub use arrow::DeviceArrayExt; +pub use arrow::DeviceArrayStreamExt; pub use arrow::ExportDeviceArray; pub use canonical::CanonicalCudaExt; pub use device_buffer::CudaBufferExt; @@ -88,24 +89,16 @@ pub use vortex_nvcomp as nvcomp; use crate::kernel::SequenceExecutor; use crate::kernel::SliceExecutor; -#[cfg(test)] -pub(crate) fn canonicalize_cpu( - array: impl vortex::array::IntoArray, -) -> vortex::error::VortexResult { - use vortex::array::LEGACY_SESSION; - use vortex::array::VortexSessionExecute; - - array - .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) -} - -/// Checks if CUDA is available on the system by looking for nvcc. +/// Checks if a CUDA driver and at least one CUDA device are available. +/// +/// cudarc loads `libcuda` lazily and panics if the driver library is absent, so we first probe +/// for it with cudarc's own `is_culib_present`. Creating the context then fails gracefully with +/// `Err`, rather than panicking, when the driver is present but no usable device is. pub fn cuda_available() -> bool { - Command::new("nvcc") - .arg("--version") - .output() - .is_ok_and(|o| o.status.success()) + // SAFETY: `is_culib_present` only tries to dlopen the driver library to test for its + // presence; it upholds no further invariants. + let driver_present = unsafe { sys::is_culib_present() }; + driver_present && CudaContext::new(0).is_ok() } /// Registers CUDA kernels. @@ -131,3 +124,17 @@ pub fn initialize_cuda(session: &CudaSession) { session.register_kernel(Filter.id(), &FilterExecutor); session.register_kernel(Slice.id(), &SliceExecutor); } + +/// Builds a fresh [`VortexSession`](vortex::session::VortexSession) with all array session +/// variables plus a default [`CudaSession`], for use in CUDA tests and benchmarks. +/// +/// Each call returns an independent session with its own CUDA context and stream pool, matching +/// the per-test isolation that lazily-initialized sessions previously provided. +/// +/// # Panics +/// +/// Panics if CUDA device 0 cannot be initialized (the same contract as [`CudaSession::default`]). +#[cfg(any(test, feature = "_test-harness"))] +pub fn cuda_session() -> vortex::session::VortexSession { + vortex::array::array_session().with::() +} diff --git a/vortex-cuda/src/session.rs b/vortex-cuda/src/session.rs index f3f7a2e3a49..5bc4e55de69 100644 --- a/vortex-cuda/src/session.rs +++ b/vortex-cuda/src/session.rs @@ -12,8 +12,8 @@ use vortex::array::ArrayId; use vortex::array::VortexSessionExecute; use vortex::error::VortexResult; use vortex::error::vortex_err; -use vortex::session::Ref; use vortex::session::SessionExt; +use vortex::session::SessionGuard; use vortex::session::SessionVar; use vortex::utils::aliases::dash_map::DashMap; @@ -183,7 +183,7 @@ impl SessionVar for CudaSession { /// Extension trait for accessing the CUDA session from a Vortex session. pub trait CudaSessionExt: SessionExt { /// Returns the CUDA session. - fn cuda_session(&self) -> Ref<'_, CudaSession> { + fn cuda_session(&self) -> SessionGuard<'_, CudaSession> { self.get::() } } diff --git a/vortex-cuda/src/stream.rs b/vortex-cuda/src/stream.rs index 1120ac50e02..6342c39ab28 100644 --- a/vortex-cuda/src/stream.rs +++ b/vortex-cuda/src/stream.rs @@ -251,7 +251,6 @@ fn register_stream_callback(stream: &CudaStream) -> VortexResult VortexResult<()> { - let ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let handle = ctx.stream().copy_to_device(vec![0xab_u8])?.await?; assert_eq!(handle.len(), 1); @@ -281,7 +280,7 @@ mod tests { #[crate::test] async fn test_copy_to_device_sync_preserves_visible_len_with_padding() -> VortexResult<()> { - let ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?; + let ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let handle = ctx.stream().copy_to_device_sync(&[1_u8, 2, 3, 4, 5])?; assert_eq!(handle.len(), 5); diff --git a/vortex-cxx/src/lib.rs b/vortex-cxx/src/lib.rs index 21383fe977f..b6d002513f6 100644 --- a/vortex-cxx/src/lib.rs +++ b/vortex-cxx/src/lib.rs @@ -24,7 +24,6 @@ use write::*; /// By default, the C++ API uses a current-thread runtime, providing control of the threading /// model to the C++ side. -/// // TODO(ngates): in the future, we could expose an API for C++ to spawn threads that can drive // this runtime. pub(crate) static RUNTIME: LazyLock = diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index 21a93b19cf4..86f31d2dfce 100644 --- a/vortex-datafusion/Cargo.toml +++ b/vortex-datafusion/Cargo.toml @@ -14,6 +14,7 @@ rust-version = { workspace = true } version = { workspace = true } [dependencies] +arrow-array = { workspace = true } arrow-schema = { workspace = true } async-trait = { workspace = true } datafusion-catalog = { workspace = true } @@ -23,6 +24,7 @@ datafusion-datasource = { workspace = true, default-features = false } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-functions = { workspace = true } +datafusion-functions-nested = { workspace = true } datafusion-physical-expr = { workspace = true } datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-expr-common = { workspace = true } diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index bf273eabc52..8f3ff81af31 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -4,13 +4,17 @@ use std::sync::Arc; use arrow_schema::DataType; +use arrow_schema::Field; use arrow_schema::Schema; use datafusion_common::Result as DFResult; +use datafusion_common::ScalarValue; use datafusion_common::exec_datafusion_err; use datafusion_common::tree_node::TreeNode; use datafusion_common::tree_node::TreeNodeRecursion; 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::PhysicalExpr; use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::projection::ProjectionExpr; @@ -19,16 +23,18 @@ 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::dtype::DType; +use vortex::VortexSessionDefault; +use vortex::array::arrow::ArrowSessionExt; use vortex::dtype::Nullability; -use vortex::dtype::arrow::FromArrowType; use vortex::expr::Expression; use vortex::expr::and_collect; +use vortex::expr::byte_length; use vortex::expr::cast; use vortex::expr::get_item; use vortex::expr::is_not_null; use vortex::expr::is_null; use vortex::expr::list_contains; +use vortex::expr::list_length; use vortex::expr::lit; use vortex::expr::nested_case_when; use vortex::expr::not; @@ -40,6 +46,7 @@ use vortex::scalar_fn::fns::binary::Binary; use vortex::scalar_fn::fns::like::Like; use vortex::scalar_fn::fns::like::LikeOptions; use vortex::scalar_fn::fns::operators::Operator; +use vortex::session::VortexSession; use crate::convert::FromDataFusion; @@ -109,13 +116,91 @@ pub trait ExpressionConvertor: Send + Sync { } } -/// The default [`ExpressionConvertor`]. -#[derive(Default)] -pub struct DefaultExpressionConvertor {} +/// The default [`ExpressionConvertor`] implementation. +pub struct DefaultExpressionConvertor { + /// Session used to resolve Arrow → Vortex dtypes through the extension + /// plugin registry, so registered extension types (e.g. UUID ⇄ + /// `FixedSizeBinary[16]`) convert correctly instead of hitting the static, + /// non-plugin-aware `DType::from_arrow`. + session: VortexSession, +} + +impl Default for DefaultExpressionConvertor { + fn default() -> Self { + Self { + session: VortexSession::default(), + } + } +} impl DefaultExpressionConvertor { + /// Create a convertor that resolves Arrow extension types using `session`'s + /// dtype registry. + pub fn new(session: VortexSession) -> Self { + Self { session } + } + + /// Attempts to convert DataFusion's `octet_length` function to Vortex `byte_length`. + fn try_convert_octet_length(&self, scalar_fn: &ScalarFunctionExpr) -> DFResult { + let [input] = scalar_fn.args() else { + return Err(exec_datafusion_err!( + "octet_length requires exactly one argument" + )); + }; + + let input = self.convert(input.as_ref())?; + let return_dtype = self + .session + .arrow() + .from_arrow_field(&Field::new( + "", + scalar_fn.return_type().clone(), + scalar_fn.nullable(), + )) + .map_err(|e| exec_datafusion_err!("Failed to convert return type to dtype: {e}"))?; + Ok(cast(byte_length(input), return_dtype)) + } + + /// Attempts to convert DataFusion's `array_length` function (aliased as `list_length`) to + /// Vortex `list_length`. + /// + /// Supports the single-argument form `array_length(arr)` and the equivalent two-argument + /// form with an explicit first dimension `array_length(arr, 1)`. + fn try_convert_array_length(&self, scalar_fn: &ScalarFunctionExpr) -> DFResult { + let Some(input) = array_length_input(scalar_fn) else { + return Err(exec_datafusion_err!( + "array_length pushdown supports only the one-argument form or an explicit first \ + dimension" + )); + }; + + let input = self.convert(input.as_ref())?; + let return_dtype = self + .session + .arrow() + .from_arrow_field(&Field::new( + "", + scalar_fn.return_type().clone(), + scalar_fn.nullable(), + )) + .map_err(|e| exec_datafusion_err!("Failed to convert return type to dtype: {e}"))?; + Ok(cast(list_length(input), return_dtype)) + } + /// Attempts to convert a DataFusion ScalarFunctionExpr to a Vortex expression. fn try_convert_scalar_function(&self, scalar_fn: &ScalarFunctionExpr) -> DFResult { + if let Some(octet_length_fn) = + ScalarFunctionExpr::try_downcast_func::(scalar_fn) + { + return self.try_convert_octet_length(octet_length_fn); + } + + if let Some(array_length_fn) = + ScalarFunctionExpr::try_downcast_func::(scalar_fn) + { + return self.try_convert_array_length(array_length_fn); + } + if let Some(get_field_fn) = ScalarFunctionExpr::try_downcast_func::(scalar_fn) { // DataFusion's GetFieldFunc flattens nested field access into a single call @@ -222,7 +307,11 @@ impl ExpressionConvertor for DefaultExpressionConvertor { } if let Some(cast_expr) = df.downcast_ref::() { - let cast_dtype = DType::from_arrow(cast_expr.target_field().as_ref()); + let cast_dtype = self + .session + .arrow() + .from_arrow_field(cast_expr.target_field().as_ref()) + .map_err(|e| exec_datafusion_err!("Failed to convert cast target to dtype: {e}"))?; let child = self.convert(cast_expr.expr().as_ref())?; return Ok(cast(child, cast_dtype)); } @@ -287,7 +376,7 @@ impl ExpressionConvertor for DefaultExpressionConvertor { let r = projection_expr.expr.apply(|node| { // We only pull column children of scalar functions that we can't push into the scan. if let Some(scalar_fn_expr) = node.downcast_ref::() - && !can_scalar_fn_be_pushed_down(scalar_fn_expr) + && !can_scalar_fn_be_pushed_down(scalar_fn_expr, input_schema) { scan_projection.extend( collect_columns(node) @@ -303,8 +392,8 @@ impl ExpressionConvertor for DefaultExpressionConvertor { // Vortex expects a perfect match so we don't push it down. if let Some(binary_expr) = node.downcast_ref::() && binary_expr.op().is_numerical_operators() - && (is_decimal(&binary_expr.left().data_type(input_schema)?) - && is_decimal(&binary_expr.right().data_type(input_schema)?)) + && binary_expr.left().data_type(input_schema)?.is_decimal() + && binary_expr.right().data_type(input_schema)?.is_decimal() { scan_projection.extend( collect_columns(node) @@ -428,7 +517,7 @@ fn can_be_pushed_down_impl(expr: &Arc, schema: &Schema) -> boo .iter() .all(|e| can_be_pushed_down_impl(e, schema)) } else if let Some(scalar_fn) = expr.downcast_ref::() { - can_scalar_fn_be_pushed_down(scalar_fn) + can_scalar_fn_be_pushed_down(scalar_fn, schema) } else if let Some(case_expr) = expr.downcast_ref::() { can_case_be_pushed_down(case_expr, schema) } else if expr @@ -458,9 +547,11 @@ fn is_convertible_expr(expr: &Arc) -> bool { || expr.downcast_ref::().is_some() || expr.downcast_ref::().is_some() || expr.downcast_ref::().is_some() - || expr - .downcast_ref::() - .is_some_and(|sf| ScalarFunctionExpr::try_downcast_func::(sf).is_some()) + || expr.downcast_ref::().is_some_and(|sf| { + ScalarFunctionExpr::try_downcast_func::(sf).is_some() + || ScalarFunctionExpr::try_downcast_func::(sf).is_some() + || ScalarFunctionExpr::try_downcast_func::(sf).is_some() + }) } fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> bool { @@ -508,20 +599,11 @@ fn supported_data_types(dt: &DataType) -> bool { let is_supported = dt.is_null() || dt.is_numeric() + || dt.is_binary() + || dt.is_string() || matches!( dt, - Boolean - | Utf8 - | LargeUtf8 - | Utf8View - | Binary - | LargeBinary - | BinaryView - | Date32 - | Date64 - | Timestamp(_, _) - | Time32(_) - | Time64(_) + Boolean | Date32 | Date64 | Timestamp(_, _) | Time32(_) | Time64(_) ); if !is_supported { @@ -532,20 +614,72 @@ fn supported_data_types(dt: &DataType) -> bool { } /// Checks if a scalar function can be pushed down. -/// Currently only GetFieldFunc is supported. -fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr) -> bool { - ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() +/// Currently GetFieldFunc, OctetLengthFunc, and ArrayLength are supported. +fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool { + if ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() { + return true; + } + + if ScalarFunctionExpr::try_downcast_func::(scalar_fn) + .is_some_and(|octet_length| can_octet_length_be_pushed_down(octet_length, schema)) + { + return true; + } + + ScalarFunctionExpr::try_downcast_func::(scalar_fn) + .is_some_and(|array_length| can_array_length_be_pushed_down(array_length, schema)) } -// TODO(adam): Replace with `DataType::is_decimal` once its released. -fn is_decimal(dt: &DataType) -> bool { - matches!( - dt, - DataType::Decimal32(_, _) - | DataType::Decimal64(_, _) - | DataType::Decimal128(_, _) - | DataType::Decimal256(_, _) - ) +fn can_octet_length_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool { + let [input] = scalar_fn.args() else { + return false; + }; + + input.data_type(schema).as_ref().is_ok_and(|data_type| { + let dt = if let DataType::Dictionary(_, value_type) = data_type { + value_type.as_ref() + } else { + data_type + }; + + dt.is_binary() || dt.is_string() + }) && can_be_pushed_down_impl(input, schema) +} + +fn can_array_length_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool { + let Some(input) = array_length_input(scalar_fn) else { + return false; + }; + + // The argument must resolve to a list type. We gate on the resolved data type rather than + // `can_be_pushed_down_impl`, since list columns are intentionally rejected there. We still + // require the argument to be a convertible expression (e.g. a column or struct field access). + input.data_type(schema).as_ref().is_ok_and(|data_type| { + matches!( + data_type, + DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) + ) + }) && is_convertible_expr(input) +} + +/// Returns the list argument of an `array_length` call if the call is a form we can rewrite to +/// `list_length`: either the single-argument form `array_length(arr)`, or the two-argument form +/// with an explicit first dimension `array_length(arr, 1)`, which is equivalent. Higher +/// dimensions recurse into nested lists and are not supported. +fn array_length_input(scalar_fn: &ScalarFunctionExpr) -> Option<&Arc> { + match scalar_fn.args() { + [input] => Some(input), + [input, dimension] if is_dimension_one(dimension) => Some(input), + _ => None, + } +} + +/// Returns true if `expr` is an `Int64` literal equal to 1. DataFusion coerces the `array_length` +/// dimension argument to `Int64`, so that is the only form we need to recognize; any other literal +/// simply isn't pushed down. +fn is_dimension_one(expr: &Arc) -> bool { + expr.downcast_ref::() + .is_some_and(|literal| matches!(literal.value(), ScalarValue::Int64(Some(1)))) } #[cfg(test)] @@ -559,7 +693,9 @@ mod tests { use datafusion::arrow::array::AsArray; use datafusion::arrow::datatypes::Int32Type; use datafusion_common::ScalarValue; + use datafusion_common::config::ConfigOptions; use datafusion_expr::Operator as DFOperator; + use datafusion_expr::ScalarUDF; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_plan::expressions as df_expr; use insta::assert_snapshot; @@ -581,13 +717,40 @@ mod tests { true, ), Field::new( - "unsupported_list", + "tags", DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), true, ), ]) } + fn octet_length_expr(input: Arc, schema: &Schema) -> Arc { + Arc::new( + ScalarFunctionExpr::try_new( + Arc::new(ScalarUDF::from(OctetLengthFunc::new())), + vec![input], + schema, + Arc::new(ConfigOptions::new()), + ) + .unwrap(), + ) + } + + fn array_length_expr( + args: Vec>, + schema: &Schema, + ) -> Arc { + Arc::new( + ScalarFunctionExpr::try_new( + Arc::new(ScalarUDF::from(ArrayLength::new())), + args, + schema, + Arc::new(ConfigOptions::new()), + ) + .unwrap(), + ) + } + #[test] fn test_make_vortex_predicate_empty() { let expr_convertor = DefaultExpressionConvertor::default(); @@ -762,6 +925,40 @@ mod tests { assert!(can_be_pushed_down_impl(&case_expr, &schema)); } + #[rstest] + fn test_expr_from_df_octet_length(test_schema: Schema) { + let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; + let octet_length = octet_length_expr(expr, &test_schema); + + let result = DefaultExpressionConvertor::default() + .convert(octet_length.as_ref()) + .unwrap(); + + assert_snapshot!(result.display_tree().to_string(), @r" + vortex.cast(i32?) + └── input: vortex.byte_length() + └── input: vortex.get_item(name) + └── input: vortex.root() + "); + } + + #[rstest] + fn test_expr_from_df_array_length(test_schema: Schema) { + let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + let array_length = array_length_expr(vec![expr], &test_schema); + + let result = DefaultExpressionConvertor::default() + .convert(array_length.as_ref()) + .unwrap(); + + assert_snapshot!(result.display_tree().to_string(), @r" + vortex.cast(u64?) + └── input: vortex.list.length() + └── input: vortex.get_item(tags) + └── input: vortex.root() + "); + } + #[rstest] // Supported types #[case::null(DataType::Null, true)] @@ -825,8 +1022,7 @@ mod tests { #[rstest] fn test_can_be_pushed_down_column_unsupported_type(test_schema: Schema) { - let col_expr = - Arc::new(df_expr::Column::new("unsupported_list", 5)) as Arc; + let col_expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; assert!(!can_be_pushed_down_impl(&col_expr, &test_schema)); } @@ -883,7 +1079,7 @@ mod tests { #[rstest] fn test_can_be_pushed_down_binary_unsupported_operand(test_schema: Schema) { - let left = Arc::new(df_expr::Column::new("unsupported_list", 5)) as Arc; + let left = Arc::new(df_expr::Column::new("tags", 5)) as Arc; let right = Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; let binary_expr = Arc::new(df_expr::BinaryExpr::new(left, DFOperator::Eq, right)) @@ -906,7 +1102,7 @@ mod tests { #[rstest] fn test_can_be_pushed_down_like_unsupported_operand(test_schema: Schema) { - let expr = Arc::new(df_expr::Column::new("unsupported_list", 5)) as Arc; + let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; let pattern = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( "test%".to_string(), )))) as Arc; @@ -916,6 +1112,74 @@ mod tests { assert!(!can_be_pushed_down_impl(&like_expr, &test_schema)); } + #[rstest] + fn test_can_be_pushed_down_octet_length_supported(test_schema: Schema) { + let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; + let octet_length = octet_length_expr(expr, &test_schema); + + assert!(can_be_pushed_down_impl(&octet_length, &test_schema)); + } + + #[rstest] + fn test_can_be_pushed_down_octet_length_unsupported_operand(test_schema: Schema) { + let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + let octet_length = Arc::new(ScalarFunctionExpr::new( + "octet_length", + Arc::new(ScalarUDF::from(OctetLengthFunc::new())), + vec![expr], + Arc::new(Field::new("octet_length", DataType::Int32, true)), + Arc::new(ConfigOptions::new()), + )) as Arc; + + assert!(!can_be_pushed_down_impl(&octet_length, &test_schema)); + } + + #[rstest] + fn test_can_be_pushed_down_array_length_supported(test_schema: Schema) { + let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + let array_length = array_length_expr(vec![expr], &test_schema); + + assert!(can_be_pushed_down_impl(&array_length, &test_schema)); + } + + #[rstest] + fn test_can_be_pushed_down_array_length_unsupported_operand(test_schema: Schema) { + // `array_length` over a non-list column cannot be pushed down. + let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; + let array_length = Arc::new(ScalarFunctionExpr::new( + "array_length", + Arc::new(ScalarUDF::from(ArrayLength::new())), + vec![expr], + Arc::new(Field::new("array_length", DataType::UInt64, true)), + Arc::new(ConfigOptions::new()), + )) as Arc; + + assert!(!can_be_pushed_down_impl(&array_length, &test_schema)); + } + + #[rstest] + fn test_can_be_pushed_down_array_length_dimension_one_supported(test_schema: Schema) { + // `array_length(arr, 1)` is the first-dimension length, equivalent to `list_length`. + let list = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + let dimension = + Arc::new(df_expr::Literal::new(ScalarValue::Int64(Some(1)))) as Arc; + let array_length = array_length_expr(vec![list, dimension], &test_schema); + + assert!(can_be_pushed_down_impl(&array_length, &test_schema)); + } + + #[rstest] + fn test_can_be_pushed_down_array_length_higher_dimension_not_supported(test_schema: Schema) { + // Dimensions other than 1 recurse into nested lists, which `list_length` does not model, + // so they must not be pushed down. + let list = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + let dimension = + Arc::new(df_expr::Literal::new(ScalarValue::Int64(Some(2)))) as Arc; + let array_length = array_length_expr(vec![list, dimension], &test_schema); + + assert!(!can_be_pushed_down_impl(&array_length, &test_schema)); + } + // https://github.com/vortex-data/vortex/issues/6211 #[tokio::test] async fn test_cast_int_to_string() -> anyhow::Result<()> { @@ -949,6 +1213,25 @@ mod tests { Ok(()) } + /// A cast whose target is a UUID-tagged `FixedSizeBinary(16)` must resolve + /// through the dtype extension registry (UUID is registered on the default + /// session) instead of the static, non-plugin-aware `DType::from_arrow`, + /// which does not support `FixedSizeBinary` and previously panicked here. + #[test] + fn test_cast_to_uuid_resolves_via_registry() -> anyhow::Result<()> { + use arrow_schema::extension::Uuid; + + let mut uuid_field = Field::new("id", DataType::FixedSizeBinary(16), true); + uuid_field.try_with_extension_type(Uuid)?; + + let child = Arc::new(df_expr::Column::new("id", 0)) as Arc; + let cast = df_expr::CastExpr::new_with_target_field(child, Arc::new(uuid_field), None); + + // Must convert without panicking — the static path would `unimplemented!()`. + DefaultExpressionConvertor::default().convert(&cast)?; + Ok(()) + } + /// Test that applying a CASE expression to an Arrow RecordBatch using DataFusion /// matches the result of applying the converted Vortex expression. #[test] diff --git a/vortex-datafusion/src/convert/mod.rs b/vortex-datafusion/src/convert/mod.rs index 050987522b7..6a6fb8aef08 100644 --- a/vortex-datafusion/src/convert/mod.rs +++ b/vortex-datafusion/src/convert/mod.rs @@ -1,6 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Utilities and interface to convert DataFusion types to Vortex types. +//! +//! Currently includes: +//! [`ExpressionConvertor`] - Controls the rewrite of DataFusion expressions to Vortex expressions, and whether they can +//! be pushed into the underlying scan. A default implementation is provided. +//! [`FromDataFusion`] - Converts a DataFusion type into a Vortex type infallible. +//! [TryToDataFusion] - Fallibly converts a Vortex type to a DataFusion type. + use vortex::error::VortexResult; pub(crate) mod exprs; @@ -8,12 +16,17 @@ mod scalars; pub(crate) mod schema; pub(crate) mod stats; +pub use exprs::DefaultExpressionConvertor; +pub use exprs::ExpressionConvertor; + /// First-party trait for implementing conversion from DataFusion types to Vortex types. -pub(crate) trait FromDataFusion: Sized { +pub trait FromDataFusion: Sized { + /// Convert to this Vortex type from the input DataFusion type. fn from_df(df: &D) -> Self; } -/// First-party trait for implementing conversion from Vortex to DataFusion types. -pub(crate) trait TryToDataFusion { +/// First-party trait for implementing fallible conversions from Vortex to DataFusion types. +pub trait TryToDataFusion { + /// Try to convert this Vortex type from the input DataFusion type. fn try_to_df(&self) -> VortexResult; } diff --git a/vortex-datafusion/src/convert/scalars.rs b/vortex-datafusion/src/convert/scalars.rs index 95c088742b3..86c1309df57 100644 --- a/vortex-datafusion/src/convert/scalars.rs +++ b/vortex-datafusion/src/convert/scalars.rs @@ -1,6 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + +use arrow_array::Array; +use arrow_array::StructArray; +use arrow_schema::Field; +use arrow_schema::Fields; use datafusion_common::ScalarValue; use vortex::buffer::ByteBuffer; use vortex::dtype::DType; @@ -14,6 +20,8 @@ use vortex::dtype::i256; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::error::vortex_err; +use vortex::error::vortex_panic; use vortex::extension::datetime::AnyTemporal; use vortex::extension::datetime::TemporalMetadata; use vortex::extension::datetime::TimeUnit; @@ -111,10 +119,16 @@ impl TryToDataFusion for Scalar { .cloned() .map(|b| Vec::::from(b.into_inner())), ), - DType::List(..) => todo!("list scalar conversion"), - DType::FixedSizeList(..) => todo!("fixed-size list scalar conversion"), - DType::Struct(..) => todo!("struct scalar conversion"), - DType::Union(..) => todo!("union scalar conversion"), + dtype @ DType::List(..) => vortex_bail!( + "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" + ), + dtype @ DType::FixedSizeList(..) => 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" + ), DType::Variant(_) => vortex_bail!("Variant scalars aren't supported with DF"), DType::Extension(ext) => { let storage_scalar = self.as_extension().to_storage_scalar(); @@ -288,11 +302,74 @@ impl FromDataFusion for Scalar { } } 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"), } } } +/// Converts a Vortex struct scalar to a DataFusion `ScalarValue::Struct`. +fn struct_to_df(scalar: &Scalar) -> VortexResult { + let scalar = scalar.as_struct(); + let struct_fields = scalar.struct_fields(); + let (fields, arrays): (Vec, Vec<_>) = struct_fields + .names() + .iter() + .zip(struct_fields.fields()) + .enumerate() + .map(|(idx, (name, field_dtype))| { + let nullable = field_dtype.is_nullable(); + let child = if scalar.is_null() { + Scalar::null(field_dtype) + } else { + scalar + .field_by_idx(idx) + .ok_or_else(|| vortex_err!("missing struct field {name}"))? + }; + let array = child + .try_to_df()? + .to_array() + .map_err(|e| vortex_err!("failed to build struct field array: {e}"))?; + Ok(( + Field::new(name.as_ref(), array.data_type().clone(), nullable), + array, + )) + }) + .collect::>>()? + .into_iter() + .unzip(); + + let fields = Fields::from(fields); + let struct_array = if scalar.is_null() { + StructArray::new_null(fields, 1) + } else { + StructArray::try_new(fields, arrays, None) + .map_err(|e| vortex_err!("failed to build struct scalar array: {e}"))? + }; + Ok(ScalarValue::Struct(Arc::new(struct_array))) +} + +/// 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)); + if array.is_null(0) { + Scalar::null(dtype) + } else { + let children = array + .columns() + .iter() + .map(|column| { + Scalar::from_df( + &ScalarValue::try_from_array(column.as_ref(), 0).unwrap_or_else(|e| { + vortex_panic!("cannot convert struct field to a Vortex scalar: {e}") + }), + ) + }) + .collect::>(); + Scalar::struct_(dtype, children) + } +} + #[cfg(test)] mod tests { use datafusion_common::ScalarValue; @@ -301,8 +378,10 @@ mod tests { use vortex::buffer::ByteBuffer; use vortex::dtype::DType; use vortex::dtype::DecimalDType; + use vortex::dtype::FieldNames; use vortex::dtype::Nullability; use vortex::dtype::PType; + use vortex::dtype::StructFields; use vortex::dtype::i256; use vortex::scalar::DecimalValue; use vortex::scalar::Scalar; @@ -691,4 +770,66 @@ mod tests { .into(); assert_eq!(result_bytes, vec![1u8, 2, 3, 4, 5]); } + + #[test] + fn struct_scalar_round_trips() -> VortexResult<()> { + let dtype = DType::Struct( + StructFields::new( + FieldNames::from(["x", "y"]), + vec![ + DType::Primitive(PType::F64, Nullability::NonNullable), + DType::Primitive(PType::F64, Nullability::NonNullable), + ], + ), + Nullability::NonNullable, + ); + let original = Scalar::struct_( + dtype, + vec![Scalar::from(-111.7610f64), Scalar::from(34.8697f64)], + ); + + let df = original.try_to_df()?; + assert!(matches!(df, ScalarValue::Struct(_))); + + // Back through `from_df` and out again yields the identical DataFusion struct value. + let back = Scalar::from_df(&df); + assert_eq!(back.try_to_df()?, df); + Ok(()) + } + + #[test] + fn null_struct_scalar_round_trips() -> VortexResult<()> { + let dtype = DType::Struct( + StructFields::new( + FieldNames::from(["x", "y"]), + vec![ + DType::Primitive(PType::F64, Nullability::Nullable), + DType::Primitive(PType::F64, Nullability::Nullable), + ], + ), + Nullability::Nullable, + ); + + let df = Scalar::null(dtype).try_to_df()?; + assert!(matches!(df, ScalarValue::Struct(_))); + assert!(Scalar::from_df(&df).is_null()); + Ok(()) + } + + #[rstest] + #[case::list(Scalar::null(DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), + Nullability::Nullable + )))] + #[case::fixed_size_list(Scalar::null(DType::FixedSizeList( + Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), + 2, + Nullability::Nullable + )))] + #[case::union(Scalar::null(DType::Union(Nullability::Nullable)))] + fn unsupported_vortex_scalars_return_errors(#[case] scalar: Scalar) { + let err = scalar.try_to_df().unwrap_err(); + + assert!(err.to_string().contains("unsupported scalar type"), "{err}"); + } } diff --git a/vortex-datafusion/src/convert/schema.rs b/vortex-datafusion/src/convert/schema.rs index 26d9d86ce18..ecfaa72d01e 100644 --- a/vortex-datafusion/src/convert/schema.rs +++ b/vortex-datafusion/src/convert/schema.rs @@ -85,9 +85,14 @@ fn calculate_physical_field_type( // RunEndEncoded loses its encoding DataType::RunEndEncoded(..) => logical_type.clone(), - // For struct types, recursively check each field + // For struct types, recursively check each field. DataType::Struct(logical_fields) => { - if let DType::Struct(struct_dtype, _) = dtype { + // Walk through any extension layers to reach the underlying struct fields. + let mut inner = dtype; + while let DType::Extension(ext) = inner { + inner = ext.storage_dtype(); + } + if let DType::Struct(struct_dtype, _) = inner { let physical_fields: Vec = struct_dtype .names() .iter() diff --git a/vortex-datafusion/src/lib.rs b/vortex-datafusion/src/lib.rs index 07098b28b49..2858cbe908e 100644 --- a/vortex-datafusion/src/lib.rs +++ b/vortex-datafusion/src/lib.rs @@ -88,7 +88,7 @@ use std::fmt::Debug; use datafusion_common::stats::Precision as DFPrecision; use vortex::expr::stats::Precision; -mod convert; +pub mod convert; mod persistent; pub mod v2; diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 50c7c915eda..3e8f932b272 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -15,8 +15,9 @@ use datafusion_common::GetExt; use datafusion_common::Result as DFResult; use datafusion_common::ScalarValue as DFScalarValue; use datafusion_common::Statistics; +use datafusion_common::config::ConfigExtension; use datafusion_common::config::ConfigField; -use datafusion_common::config_namespace; +use datafusion_common::extensions_options; use datafusion_common::internal_datafusion_err; use datafusion_common::not_impl_err; use datafusion_common::parsers::CompressionTypeVariant; @@ -135,27 +136,62 @@ impl Debug for VortexFormat { } } -config_namespace! { +extensions_options! { /// Options to configure [`VortexFormat`] and [`VortexSource`]. /// - /// These options are usually set on a [`VortexFormatFactory`] and inherited - /// by the `VortexFormat` / `VortexSource` instances created for individual - /// tables. + /// The API follows DataFusion's built-in Parquet and JSON format factories: + /// a format factory may carry customized defaults, the session may carry + /// format defaults, and `CREATE EXTERNAL TABLE ... OPTIONS(...)` can + /// override individual fields for one table. + /// + /// [`FileFormatFactory::create`] builds the `VortexTableOptions` copied into + /// each [`VortexFormat`] as follows: + /// + /// 1. If the factory has explicit options from + /// [`VortexFormatFactory::with_options`] or + /// [`VortexFormatFactory::new_with_options`], start from that complete + /// `VortexTableOptions` value. This matches + /// [`ParquetFormatFactory::new_with_options`] and + /// [`JsonFormatFactory::new_with_options`]: factory options replace + /// session defaults; they are not merged with them field-by-field. + /// 2. If the factory does not have explicit options, read the session's + /// `vortex` extension at the time `create` is called. This is the value + /// changed by `SET vortex.