diff --git a/.clang-format b/.clang-format index 587f746da78..bb77abc1478 100644 --- a/.clang-format +++ b/.clang-format @@ -25,7 +25,7 @@ DerivePointerAlignment: false IncludeBlocks: Regroup InsertBraces: true Language: Cpp -NamespaceIndentation: Inner +NamespaceIndentation: None PointerAlignment: Right SpaceBeforeCpp11BracedList: true SpaceBeforeCtorInitializerColon: true diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index 46e90627eca..2f656d25d63 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -1,6 +1,6 @@ name: Bug Report description: Report something that isn't working correctly -labels: ["needs-triage"] +labels: ["bug"] body: - type: markdown attributes: diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml index 61c4f9a55ed..f6694316c49 100644 --- a/.github/ISSUE_TEMPLATE/config.yml +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -1,5 +1,5 @@ blank_issues_enabled: true contact_links: - - name: Features, Questions - url: https://github.com/vortex-data/vortex/discussions/new/choose - about: Our preferred starting point if you have any questions or suggestions about configuration, features or behavior. + - name: Community Slack + url: https://vortex.dev/slack + about: Chat with the community for quick questions and discussion. diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml new file mode 100644 index 00000000000..e15b7ccca4d --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -0,0 +1,27 @@ +name: Feature Request +description: Suggest a new feature or improvement +labels: ["feature"] +body: + - type: textarea + id: problem + attributes: + label: Problem or motivation + description: What are you trying to do that Vortex doesn't support today? + validations: + required: true + + - type: textarea + id: proposal + attributes: + label: Proposed solution + description: What would you like to see happen? + validations: + required: true + + - type: textarea + id: additional + attributes: + label: Additional context + description: Alternatives considered, links, or anything else helpful + validations: + required: false diff --git a/.github/ISSUE_TEMPLATE/question.yml b/.github/ISSUE_TEMPLATE/question.yml new file mode 100644 index 00000000000..c208dcdb954 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/question.yml @@ -0,0 +1,24 @@ +name: Question +description: Ask a question about using or configuring Vortex +labels: ["question"] +body: + - type: markdown + attributes: + value: | + Please check the [documentation](https://docs.vortex.dev) first. For quick + back-and-forth, try the [Vortex Slack channel](https://vortex.dev/slack). + + - type: textarea + id: question + attributes: + label: What's your question? + validations: + required: true + + - type: textarea + id: context + attributes: + label: Context + description: What you've tried, relevant code, or your environment + validations: + required: false diff --git a/.github/actions/setup-prebuild/action.yml b/.github/actions/setup-prebuild/action.yml index 9cfc837b680..6fdd067fab4 100644 --- a/.github/actions/setup-prebuild/action.yml +++ b/.github/actions/setup-prebuild/action.yml @@ -17,9 +17,9 @@ inputs: description: "optional targets override (e.g. wasm32-unknown-unknown)" required: false enable-sccache: - description: "Should sccache be enabled, true by default." + description: "Should sccache be enabled." required: false - default: "true" + default: "false" runs: using: "composite" @@ -46,4 +46,4 @@ runs: toolchain: ${{ inputs.toolchain }} components: ${{ inputs.components }} targets: ${{ inputs.targets }} - enable-sccache: ${{ inputs.enable-sccache }} + enable-sccache: "false" diff --git a/.github/actions/setup-rust/action.yml b/.github/actions/setup-rust/action.yml index 4f42d38167a..345a2d08f7e 100644 --- a/.github/actions/setup-rust/action.yml +++ b/.github/actions/setup-rust/action.yml @@ -17,9 +17,9 @@ inputs: description: "optional targets override (e.g. wasm32-unknown-unknown)" required: false enable-sccache: - description: "Should sccache be enabled, true by default." + description: "Should sccache be enabled." required: false - default: "true" + default: "false" runs: using: "composite" diff --git a/.github/workflows/approvals.yml b/.github/workflows/approvals.yml deleted file mode 100644 index 5cbfbceeedb..00000000000 --- a/.github/workflows/approvals.yml +++ /dev/null @@ -1,49 +0,0 @@ -name: PR Approval Check - -on: - pull_request: - branches: - - "develop" - -jobs: - check-approvals: - runs-on: ubuntu-latest - timeout-minutes: 10 - steps: - - name: Check required approvals - uses: actions/github-script@450193c5abd4cdb17ba9f3ffcfe8f635c4bb6c2a - with: - script: | - const pr = context.payload.pull_request; - const reviews = await github.rest.pulls.listReviews({ - owner: context.repo.owner, - repo: context.repo.repo, - pull_number: pr.number, - }); - - // Determine if PR author is a bot/GitHub Actions - const authorType = pr.user.type; // 'Bot' vs 'User' - const authorLogin = pr.user.login; // e.g. 'github-actions[bot]' - const isBot = authorType === 'Bot' || authorLogin.endsWith('[bot]'); - const oneApprovalBotAuthors = new Set(['renovate[bot]']); - - // Count unique approvals, including bot approvals. - const latestByUser = {}; - for (const review of reviews.data) { - latestByUser[review.user.login] = review.state; - } - const approvalCount = Object.values(latestByUser) - .filter(state => state === 'APPROVED').length; - - const required = isBot && !oneApprovalBotAuthors.has(authorLogin) ? 2 : 1; - - console.log(`PR author: ${authorLogin} (${authorType}), isBot: ${isBot}`); - console.log(`One-approval bot allowlist: ${oneApprovalBotAuthors.has(authorLogin)}`); - console.log(`Approvals: ${approvalCount} / ${required} required`); - - if (isBot && (approvalCount < required)) { - core.setFailed( - `This PR needs ${required} approval(s) but has ${approvalCount}. ` + - `(Author is ${isBot ? 'a bot' : 'human'})` - ); - } diff --git a/.github/workflows/bench-pr.yml b/.github/workflows/bench-pr.yml index 8223ea80a56..12137726122 100644 --- a/.github/workflows/bench-pr.yml +++ b/.github/workflows/bench-pr.yml @@ -46,10 +46,11 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} + enable-sccache: ${{ github.event.pull_request.head.repo.fork == false && 'true' || 'false' }} - name: Install DuckDB run: | - wget -qO- https://github.com/duckdb/duckdb/releases/download/v1.5.3/duckdb_cli-linux-amd64.zip | funzip > duckdb + wget -qO- https://github.com/duckdb/duckdb/releases/download/v1.5.4/duckdb_cli-linux-amd64.zip | funzip > duckdb chmod +x duckdb echo "$PWD" >> $GITHUB_PATH diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml index fc59a8156ea..b1e29277f00 100644 --- a/.github/workflows/bench.yml +++ b/.github/workflows/bench.yml @@ -60,10 +60,11 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - name: Install DuckDB run: | - wget -qO- https://github.com/duckdb/duckdb/releases/download/v1.5.3/duckdb_cli-linux-amd64.zip | funzip > duckdb + wget -qO- https://github.com/duckdb/duckdb/releases/download/v1.5.4/duckdb_cli-linux-amd64.zip | funzip > duckdb chmod +x duckdb echo "$PWD" >> $GITHUB_PATH @@ -116,47 +117,29 @@ jobs: run: | bash scripts/cat-s3.sh vortex-ci-benchmark-results data.json.gz results.json - - name: Ingest results to v3 server - if: vars.V3_INGEST_URL != '' - shell: bash - env: - INGEST_BEARER_TOKEN: ${{ secrets.INGEST_BEARER_TOKEN }} - run: | - python3 scripts/post-ingest.py results.v3.jsonl \ - --server "${{ vars.V3_INGEST_URL }}" \ - --commit-sha "${{ github.sha }}" \ - --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. + # v4 (Postgres) ingest -- the REQUIRED benchmark-results pipeline feeding the live + # benchmarks website; a failure here fails the job. Gated on the ingest-role ARN var + # (the assume-role input that MUST exist for OIDC to succeed). 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. + # `sync: false` -- the ingest runs `uv run --no-project --with`, which needs only + # the uv binary, never the synced workspace. A full sync would rebuild + # vortex-python via sccache->S3, which fails under the ingest-role creds + # (rds-db:connect only) and is pure waste here. - 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 + with: + sync: false - 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) + - name: Ingest results to v4 Postgres if: vars.GH_BENCH_INGEST_ROLE_ARN != '' - continue-on-error: true shell: bash env: RDS_BENCH_INSTANCE_ENDPOINT: ${{ vars.RDS_BENCH_INSTANCE_ENDPOINT }} @@ -188,3 +171,31 @@ jobs: secrets: inherit with: mode: "develop" + + sql-vortex: + uses: ./.github/workflows/sql-benchmarks.yml + secrets: inherit + with: + mode: "develop" + benchmark_matrix: | + [ + { + "id": "vortex-queries", + "subcommand": "vortex", + "name": "Vortex queries", + "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"} + ], + "iterations": "100" + } + ] diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 44797183cee..6bc9574cb98 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -34,12 +34,12 @@ jobs: duckdb-ready: name: "DuckDB libraries available in R2" needs: duckdb-mirror - if: ${{ !cancelled() }} + if: "!cancelled()" runs-on: ubuntu-latest timeout-minutes: 5 steps: - name: Verify DuckDB mirror - if: ${{ needs.duckdb-mirror.result == 'failure' }} + if: needs.duckdb-mirror.result == 'failure' run: | echo "DuckDB mirror failed; downstream builds would 404" exit 1 @@ -68,7 +68,7 @@ jobs: name: "Python (lint)" runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=python-lint', github.run_id) + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=python-lint', github.run_id) || 'ubuntu-latest' }} timeout-minutes: 10 steps: @@ -78,13 +78,15 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" - name: Install uv if: github.repository != 'vortex-data/vortex' uses: spiraldb/actions/.github/actions/setup-uv@0.18.5 with: sync: false prune-cache: false - # Use uv tool run for ruff to avoid building the Rust extension (saves ~4.5 min) + # Use uvx for ruff to avoid building the Rust extension (saves ~4.5 min) - name: Python Lint - Format run: uv tool run ruff format --check . - name: Python Lint - Ruff @@ -101,7 +103,7 @@ jobs: name: "Python (test)" runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-large/image=ubuntu24-full-x64-pre-v2/tag=python-test', github.run_id) + && format('runs-on={0}/runner=amd64-large/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=python-test', github.run_id) || 'ubuntu-latest' }} timeout-minutes: 30 env: @@ -114,6 +116,8 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" - name: Pytest - Vortex run: | @@ -138,11 +142,30 @@ jobs: uv run --all-packages make html working-directory: docs/ + # Pins the measurement_id hash that scripts/post-ingest.py uses as the primary key of every + # benchmarks-database fact row. The golden vectors are a frozen artifact (see the note inside + # the JSON); this job fails if a change to scripts/_measurement_id.py drifts the hash, which + # would break ON CONFLICT idempotency against the rows already in production. + measurement-id-golden: + name: "measurement_id golden vectors" + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + # sync: false — the test runs via `uv run --no-project` and needs no workspace + # packages; the default `uv sync` builds the vortex-data Rust extension (~6 min). + - uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 + with: + sync: false + - name: Pytest - measurement_id golden vectors + run: | + uv run --no-project --with pytest --with xxhash \ + pytest scripts/tests/test_measurement_id.py + 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) }} + runs-on: runs-on=${{ github.run_id }}/runner=gpu/extras=s3-cache/tag=python-cuda-test timeout-minutes: 30 env: RUST_LOG: "info,maturin=off,uv=debug" @@ -156,6 +179,7 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} components: cargo + enable-sccache: "true" - name: Pin rustup proxy to repository toolchain run: | TOOLCHAIN="$(grep '^channel' rust-toolchain.toml | cut -d '"' -f 2)" @@ -177,7 +201,7 @@ jobs: timeout-minutes: 30 runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-small/image=ubuntu24-full-x64-pre-v2/tag=rust-docs', github.run_id) + && format('runs-on={0}/runner=amd64-small/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=rust-docs', github.run_id) || 'ubuntu-latest' }} steps: - uses: runs-on/action@v2 @@ -186,6 +210,8 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" - name: Docs env: # required to make sure docs build for features @@ -195,14 +221,14 @@ jobs: # including all private docs. 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 --exclude vortex-python-cuda --no-fail-fast + cargo test --profile ci --doc --workspace --all-features --exclude vortex-jni --exclude vortex-ffi --exclude xtask --exclude vortex-python-cuda --no-fail-fast build-rust: name: "Rust build (${{matrix.config.name}})" timeout-minutes: 30 runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner={1}/image=ubuntu24-full-x64-pre-v2/tag={2}', github.run_id, matrix.config.runner, matrix.config.name) + && format('runs-on={0}/runner={1}/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag={2}', github.run_id, matrix.config.runner, matrix.config.name) || 'ubuntu-latest' }} env: # disable lints for build, they will be caught in Rust lint job. @@ -234,8 +260,10 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" - name: Install wasm32 target - if: ${{ matrix.config.target == 'wasm32-unknown-unknown' }} + if: matrix.config.target == 'wasm32-unknown-unknown' run: rustup target add wasm32-unknown-unknown - uses: ./.github/actions/check-rebuild with: @@ -250,7 +278,7 @@ jobs: timeout-minutes: 30 runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=rust-min-deps', github.run_id) + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=rust-min-deps', github.run_id) || 'ubuntu-latest' }} steps: - uses: runs-on/action@v2 @@ -259,6 +287,8 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" - run: cargo minimal-versions check --direct --workspace --ignore-private rust-lint: @@ -267,7 +297,7 @@ jobs: timeout-minutes: 30 runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-large/image=ubuntu24-full-x64-pre-v2/tag=rust-lint', github.run_id) + && format('runs-on={0}/runner=amd64-large/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=rust-lint', github.run_id) || 'ubuntu-latest' }} steps: - uses: runs-on/action@v2 @@ -276,6 +306,8 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" - name: Install nightly for fmt run: rustup toolchain install $NIGHTLY_TOOLCHAIN --component rustfmt - name: Rust Lint - Format @@ -333,7 +365,7 @@ jobs: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: C/C++ Lint - clang-format run: | - git ls-files vortex-cuda vortex-cxx vortex-duckdb vortex-ffi \ + git ls-files lang/cpp vortex-cuda vortex-duckdb vortex-ffi \ | grep -E '\.(cpp|hpp|cu|cuh|h)$' \ | grep -v 'arrow/reference/arrow_c_device\.h$' \ | grep -v 'kernels/src/bit_unpack_.*\.cu$' \ @@ -345,7 +377,7 @@ jobs: timeout-minutes: 30 runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=rust-lint-no-default', github.run_id) + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=rust-lint-no-default', github.run_id) || 'ubuntu-latest' }} steps: - uses: runs-on/action@v2 @@ -354,6 +386,8 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" - name: Install cargo-hack if: github.repository != 'vortex-data/vortex' shell: bash @@ -363,23 +397,14 @@ jobs: run: | cargo hack --ignore-private clippy --profile ci --no-default-features -- -D warnings - rust-test-other: - name: "Rust tests (${{ matrix.os }})" + rust-test-windows: + name: "Rust tests (windows-x64)" needs: duckdb-ready timeout-minutes: 30 - strategy: - fail-fast: false - matrix: - include: - - os: windows-x64 - runner: runs-on=${{ github.run_id }}/pool=windows-x64-pre - fallback_runner: windows-latest - - os: linux-arm64 - runner: runs-on=${{ github.run_id }}/runner=arm64-medium/image=ubuntu24-full-arm64-pre-v2/tag=rust-test-linux-arm64 runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && matrix.runner - || matrix.fallback_runner }} + && format('runs-on={0}/pool=windows-x64-pre/extras=s3-cache', github.run_id) + || 'windows-latest' }} steps: - uses: runs-on/action@v2 if: github.repository == 'vortex-data/vortex' @@ -387,12 +412,12 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - name: Setup (Windows) - if: matrix.os == 'windows-x64' run: | echo "C:\rust\cargo\bin" >> $env:GITHUB_PATH - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" - name: Rust Tests (Windows) - if: matrix.os == 'windows-x64' run: | cargo nextest run --cargo-profile ci --locked --workspace --all-features --no-fail-fast ` --exclude vortex-bench ` @@ -403,12 +428,37 @@ jobs: --exclude lance-bench --exclude datafusion-bench --exclude random-access-bench ` --exclude compress-bench --exclude xtask --exclude vortex-datafusion ` --exclude gpu-scan-cli --exclude vortex-sqllogictest - - name: Rust Tests (Other) - if: matrix.os != 'windows-x64' + + - name: sccache stats + if: always() && github.repository == 'vortex-data/vortex' + run: sccache --show-stats + + - name: Alert incident.io + if: failure() && github.event_name == 'push' && github.ref == 'refs/heads/develop' + uses: ./.github/actions/alert-incident-io + with: + api-key: ${{ secrets.INCIDENT_IO_ALERT_TOKEN }} + alert-title: "Rust tests (windows-x64) failed on develop" + deduplication-key: ci-rust-test-windows-x64-failure + + rust-test-linux-arm64: + name: "Rust tests (linux-arm64)" + needs: duckdb-ready + if: github.repository == 'vortex-data/vortex' + timeout-minutes: 30 + runs-on: runs-on=${{ github.run_id }}/runner=arm64-medium/image=ubuntu24-full-arm64-pre-v2/extras=s3-cache/tag=rust-test-linux-arm64 + steps: + - uses: runs-on/action@v2 + with: + sccache: s3 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" + - name: Rust Tests run: | cargo nextest run --cargo-profile ci --locked --workspace --all-features --no-fail-fast --exclude vortex-bench --exclude xtask --exclude vortex-sqllogictest - uses: ./.github/actions/check-rebuild - if: matrix.os != 'windows-x64' with: command: "cargo test --profile ci --locked --workspace --all-features --no-run --exclude vortex-bench --exclude xtask --exclude vortex-sqllogictest" @@ -417,14 +467,46 @@ jobs: uses: ./.github/actions/alert-incident-io with: api-key: ${{ secrets.INCIDENT_IO_ALERT_TOKEN }} - alert-title: "Rust tests (${{ matrix.os }}) failed on develop" - deduplication-key: ci-rust-test-${{ matrix.os }}-failure + alert-title: "Rust tests (linux-arm64) failed on develop" + deduplication-key: ci-rust-test-linux-arm64-failure + + btrblocks-golden: + name: "Rust (btrblocks golden corpus)" + timeout-minutes: 10 + runs-on: >- + ${{ github.repository == 'vortex-data/vortex' + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=btrblocks-golden', github.run_id) + || 'ubuntu-latest' }} + steps: + - uses: runs-on/action@v2 + if: github.repository == 'vortex-data/vortex' + with: + sccache: s3 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" + # The golden suite pins the compressor's decisions per feature variant, and the + # variants are mutually exclusive at compile time: the `default` variant only exists + # without `unstable_encodings` (which changes ALL_SCHEMES), so the `--all-features` + # workspace test jobs cannot run it. Run each feature combination explicitly. + - name: Golden corpus (default features) + run: | + cargo nextest run --cargo-profile ci --locked --no-fail-fast -p vortex-btrblocks --test golden + - name: Golden corpus (unstable_encodings) + run: | + cargo nextest run --cargo-profile ci --locked --no-fail-fast -p vortex-btrblocks --test golden \ + --features unstable_encodings + - name: Golden corpus (compact, unstable_encodings + zstd + pco) + run: | + cargo nextest run --cargo-profile ci --locked --no-fail-fast -p vortex-btrblocks --test golden \ + --features unstable_encodings,zstd,pco build-java: name: "Java" runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/pool=amd64-medium-pre-v2/tag=java', github.run_id) + && format('runs-on={0}/pool=amd64-medium-pre-v2/extras=s3-cache/tag=java', github.run_id) || 'ubuntu-latest' }} timeout-minutes: 30 steps: @@ -434,6 +516,8 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" - run: ./gradlew javadoc working-directory: ./java - run: ./gradlew check @@ -456,12 +540,12 @@ jobs: with: command: check ${{ matrix.checks }} - cxx-test: - name: "C++ build" - timeout-minutes: 30 + cxx-api: + name: "C++ API" + timeout-minutes: 5 runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=cxx-build', github.run_id) + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=cxx-build', github.run_id) || 'ubuntu-latest' }} steps: - uses: runs-on/action@v2 @@ -470,27 +554,35 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - - name: Build and run C++ unit tests + with: + enable-sccache: "true" + - name: Install Rust nightly toolchain run: | - mkdir -p vortex-cxx/build - cmake -S vortex-cxx -B vortex-cxx/build -DVORTEX_ENABLE_TESTING=ON -DVORTEX_ENABLE_ASAN=ON - cmake --build vortex-cxx/build --parallel $(nproc) - ctest --test-dir vortex-cxx/build -j $(nproc) -V - - name: Build and run the example in release mode + rustup toolchain install $NIGHTLY_TOOLCHAIN + rustup component add --toolchain $NIGHTLY_TOOLCHAIN rust-src rustfmt clippy llvm-tools-preview + - name: Build FFI library (asan) run: | - cmake -S vortex-cxx/examples -B vortex-cxx/examples/build -DCMAKE_BUILD_TYPE=Release - cmake --build vortex-cxx/examples/build --parallel $(nproc) - vortex-cxx/examples/build/hello-vortex vortex-cxx/examples/goldenfiles/example.vortex - - uses: ./.github/actions/check-rebuild - with: - command: "cargo build --profile ci --locked -p vortex-cxx --lib" + RUSTFLAGS="-A warnings -Cunsafe-allow-abi-mismatch=sanitizer \ + -C debuginfo=2 -C opt-level=0 -C strip=none -Zexternal-clangrt \ + -Zsanitizer=address,leak" \ + cargo +$NIGHTLY_TOOLCHAIN build --locked --no-default-features \ + --target x86_64-unknown-linux-gnu -Zbuild-std \ + -p vortex-ffi + - name: Build C++ library (asan) + run: | + mkdir -p build + cmake -S lang/cpp -Bbuild -DSANITIZER=asan -DBUILD_TESTS=1 -DTARGET_TRIPLE="x86_64-unknown-linux-gnu" + cmake --build build --parallel $(nproc) + - name: Run C++ tests + run: | + build/tests/vortex_cxx_test 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) + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=sql-logic-test', github.run_id) || 'ubuntu-latest' }} timeout-minutes: 30 steps: @@ -500,6 +592,8 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" - name: Run sqllogictest tests run: | ./vortex-sqllogictest/slt/tpch/generate_data.sh @@ -511,14 +605,23 @@ jobs: wasm-integration: name: "WASM integration smoke test" - runs-on: ubuntu-latest timeout-minutes: 30 + runs-on: >- + ${{ github.repository == 'vortex-data/vortex' + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=wasm-integration', github.run_id) + || 'ubuntu-latest' }} steps: + - uses: runs-on/action@v2 + if: github.repository == 'vortex-data/vortex' + with: + sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - - uses: ./.github/actions/setup-rust + - uses: ./.github/actions/setup-prebuild with: - repo-token: ${{ secrets.GITHUB_TOKEN }} + enable-sccache: "true" targets: "wasm32-wasip1" + - name: Install wasm32-wasip1 target + run: rustup target add wasm32-wasip1 - name: Setup Wasmer shell: bash run: | @@ -533,7 +636,7 @@ jobs: name: "Check generated source files are up to date" runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=generated-files', github.run_id) + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=generated-files', github.run_id) || 'ubuntu-latest' }} timeout-minutes: 30 steps: @@ -543,6 +646,8 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" - name: Install nightly for cbindgen macro expansion run: rustup toolchain install $NIGHTLY_TOOLCHAIN - name: "regenerate all .fbs/.proto Rust code" @@ -573,7 +678,7 @@ jobs: timeout-minutes: 10 runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=cxx-build', github.run_id) + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=cxx-build', github.run_id) || 'ubuntu-latest' }} steps: - uses: runs-on/action@v2 @@ -582,6 +687,8 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" - name: "regenerate FFI header file" run: | cargo +$NIGHTLY_TOOLCHAIN build --profile ci -p vortex-ffi diff --git a/.github/workflows/claude-write.yml b/.github/workflows/claude-write.yml index acf5add785c..2b1f73d6d51 100644 --- a/.github/workflows/claude-write.yml +++ b/.github/workflows/claude-write.yml @@ -187,8 +187,6 @@ jobs: - name: Setup Rust toolchain uses: ./.github/actions/setup-rust - with: - enable-sccache: "false" - name: Install uv uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 diff --git a/.github/workflows/close-fixed-fuzzer-issues.yml b/.github/workflows/close-fixed-fuzzer-issues.yml index ee203553ac7..9030809f5cd 100644 --- a/.github/workflows/close-fixed-fuzzer-issues.yml +++ b/.github/workflows/close-fixed-fuzzer-issues.yml @@ -27,7 +27,7 @@ jobs: target: [file_io, array_ops, compress_roundtrip] runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=arm64-medium/disk=large/tag=fuzzer-cleanup-{1}', github.run_id, matrix.target) + && format('runs-on={0}/runner=arm64-medium/disk=large/extras=s3-cache/tag=fuzzer-cleanup-{1}', github.run_id, matrix.target) || 'ubuntu-latest' }} timeout-minutes: 30 steps: @@ -42,6 +42,7 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} toolchain: ${{ env.NIGHTLY_TOOLCHAIN }} + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - name: Install llvm uses: aminya/setup-cpp@v1 diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 2cdcbaa9dbf..8c55eb71f7c 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -23,6 +23,25 @@ env: NIGHTLY_TOOLCHAIN: nightly-2026-02-05 jobs: + changes: + name: "Detect CUDA changes" + runs-on: ubuntu-latest + timeout-minutes: 10 + permissions: + pull-requests: read + outputs: + run-cuda-benchmarks: ${{ 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/**" + - ".github/workflows/**" + bench-codspeed: strategy: matrix: @@ -39,7 +58,7 @@ jobs: timeout-minutes: 30 runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=bench-codspeed-{1}', github.run_id, matrix.shard) + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=bench-codspeed-{1}', github.run_id, matrix.shard) || 'ubuntu-latest' }} steps: - uses: runs-on/action@v2 @@ -50,6 +69,8 @@ jobs: - name: Setup benchmark environment run: sudo bash scripts/setup-benchmark.sh - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - uses: ./.github/actions/system-info - run: | sudo apt-get update @@ -72,11 +93,14 @@ jobs: # 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' + needs: [changes] + if: >- + always() && github.repository == 'vortex-data/vortex' && + needs.changes.outputs.run-cuda-benchmarks == 'true' 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 + runs-on=${{ github.run_id }}/family=g5/cpu=8/image=ubuntu24-gpu-x64/extras=s3-cache/tag=bench-codspeed-cuda-build steps: - uses: runs-on/action@v2 with: @@ -85,6 +109,7 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} + enable-sccache: "true" - name: Install Codspeed uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 with: @@ -120,7 +145,7 @@ jobs: name: "Benchmark with Codspeed (CUDA Shard #${{ matrix.shard }} - ${{ matrix.name }})" timeout-minutes: 30 runs-on: >- - runs-on=${{ github.run_id }}/family=g5/cpu=8/image=ubuntu24-gpu-x64/tag=bench-codspeed-cuda-${{ matrix.shard }} + runs-on=${{ github.run_id }}/family=g5/cpu=8/image=ubuntu24-gpu-x64/extras=s3-cache/tag=bench-codspeed-cuda-${{ matrix.shard }} steps: - uses: runs-on/action@v2 with: @@ -129,6 +154,7 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} + enable-sccache: "true" - name: Display NVIDIA SMI details run: | nvidia-smi diff --git a/.github/workflows/commit-metadata.yml b/.github/workflows/commit-metadata.yml index 272bbaad9c9..97f61989d35 100644 --- a/.github/workflows/commit-metadata.yml +++ b/.github/workflows/commit-metadata.yml @@ -1,7 +1,6 @@ # 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. +# every push to develop, to the v4 Postgres, so the `commits` dim stays +# populated even when no benchmark ran. name: Commit metadata @@ -11,7 +10,7 @@ on: workflow_dispatch: { } permissions: - id-token: write # enables AWS-GitHub OIDC for the best-effort v4 ingest step + id-token: write # enables AWS-GitHub OIDC for the v4 ingest step contents: read jobs: @@ -23,45 +22,26 @@ jobs: 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. + # v4 (Postgres) ingest -- REQUIRED (see bench.yml rationale). Empty records: + # post-ingest.py --postgres upserts the commit row only. Gated on the + # ingest-role ARN var (the assume-role input that MUST exist for OIDC to + # succeed). # - # 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. + # `sync: false` -- the ingest runs `uv run --no-project --with`, which needs only + # the uv binary, never the synced workspace (see bench.yml rationale). - 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 + with: + sync: false - 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) + - name: Ingest commit metadata to v4 Postgres if: vars.GH_BENCH_INGEST_ROLE_ARN != '' - continue-on-error: true shell: bash env: RDS_BENCH_INSTANCE_ENDPOINT: ${{ vars.RDS_BENCH_INSTANCE_ENDPOINT }} diff --git a/.github/workflows/compat-gen-upload.yml b/.github/workflows/compat-gen-upload.yml index 8dd4cd1999e..bfc3c580307 100644 --- a/.github/workflows/compat-gen-upload.yml +++ b/.github/workflows/compat-gen-upload.yml @@ -25,7 +25,7 @@ jobs: dry-run: runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=compat-gen-dry-run', github.run_id) + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=compat-gen-dry-run', github.run_id) || 'ubuntu-latest' }} timeout-minutes: 30 permissions: @@ -42,6 +42,8 @@ jobs: with: fetch-depth: 0 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - name: Configure AWS credentials uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6 @@ -76,7 +78,7 @@ jobs: if: inputs.confirm_upload == 'yes' runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=compat-gen-upload', github.run_id) + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=compat-gen-upload', github.run_id) || 'ubuntu-latest' }} timeout-minutes: 30 environment: compat-upload @@ -92,6 +94,8 @@ jobs: with: fetch-depth: 0 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - name: Configure AWS credentials uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6 diff --git a/.github/workflows/compat-validation.yml b/.github/workflows/compat-validation.yml index 8732bdfbf6d..60741d50620 100644 --- a/.github/workflows/compat-validation.yml +++ b/.github/workflows/compat-validation.yml @@ -25,7 +25,7 @@ jobs: compat-test: runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=compat-validation', github.run_id) + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=compat-validation', github.run_id) || 'ubuntu-latest' }} timeout-minutes: 30 steps: @@ -35,6 +35,8 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - name: Run compat tests run: | MODE="${{ inputs.mode || 'last' }}" diff --git a/.github/workflows/cuda.yaml b/.github/workflows/cuda.yaml index 2bd9afbbdf1..004ca7d3a49 100644 --- a/.github/workflows/cuda.yaml +++ b/.github/workflows/cuda.yaml @@ -54,7 +54,7 @@ jobs: 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 + runs-on: runs-on=${{ github.run_id }}/runner=gpu/extras=s3-cache/tag=cuda-build steps: - uses: runs-on/action@v2 with: @@ -63,6 +63,7 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} + enable-sccache: "true" - uses: ./.github/actions/check-rebuild with: command: >- @@ -88,7 +89,7 @@ jobs: 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 + runs-on: runs-on=${{ github.run_id }}/runner=gpu/extras=s3-cache/tag=cuda-tests steps: - uses: runs-on/action@v2 with: @@ -102,6 +103,7 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} + enable-sccache: "true" - name: Install uv uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 with: @@ -144,7 +146,7 @@ jobs: 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 + runs-on: runs-on=${{ github.run_id }}/runner=gpu/extras=s3-cache/tag=cuda-test-sanitizer strategy: fail-fast: false matrix: @@ -171,6 +173,7 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} + enable-sccache: "true" - name: Build tests run: cargo test --profile ci --locked -p vortex-cuda --all-features --target x86_64-unknown-linux-gnu --no-run - name: "CUDA - ${{ matrix.sanitizer }}" @@ -185,7 +188,7 @@ jobs: 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 + runs-on: runs-on=${{ github.run_id }}/runner=gpu/extras=s3-cache/tag=cuda-test-cudf steps: - uses: runs-on/action@v2 with: @@ -199,6 +202,7 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} + enable-sccache: "true" - name: Build cudf test library run: cargo build --profile ci --locked -p vortex-test-e2e-cuda --target x86_64-unknown-linux-gnu - name: Download and run cudf-test-harness diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 1a0d376f319..630383b0311 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -23,9 +23,8 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} - enable-sccache: "false" - name: Set up JDK 17 - uses: actions/setup-java@ad2b38190b15e4d6bdf0c97fb4fca8412226d287 # v5 + uses: actions/setup-java@0f481fcb613427c0f801b606911222b5b6f3083a # v5 with: java-version: "17" distribution: "temurin" diff --git a/.github/workflows/fuzz-coverage.yml b/.github/workflows/fuzz-coverage.yml index fd6ca8aed1b..88bfb7d01b7 100644 --- a/.github/workflows/fuzz-coverage.yml +++ b/.github/workflows/fuzz-coverage.yml @@ -18,7 +18,7 @@ jobs: timeout-minutes: 120 runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=arm64-medium/disk=large', github.run_id) + && format('runs-on={0}/runner=arm64-medium/disk=large/extras=s3-cache', github.run_id) || 'ubuntu-latest' }} steps: - uses: runs-on/action@v2 @@ -33,6 +33,7 @@ jobs: repo-token: ${{ secrets.GITHUB_TOKEN }} toolchain: ${{ env.NIGHTLY_TOOLCHAIN }} components: clippy, rustfmt, llvm-tools + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - name: Ensure llvm-tools are installed run: | diff --git a/.github/workflows/fuzzer-fix-automation.yml b/.github/workflows/fuzzer-fix-automation.yml index 7b8e7012511..0db67516346 100644 --- a/.github/workflows/fuzzer-fix-automation.yml +++ b/.github/workflows/fuzzer-fix-automation.yml @@ -86,7 +86,6 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} toolchain: ${{ env.NIGHTLY_TOOLCHAIN }} - enable-sccache: "false" - name: Install llvm uses: aminya/setup-cpp@v1 diff --git a/.github/workflows/minimize_fuzz_corpus_workflow.yml b/.github/workflows/minimize_fuzz_corpus_workflow.yml index 00f9ba73b84..b0c2fc58d50 100644 --- a/.github/workflows/minimize_fuzz_corpus_workflow.yml +++ b/.github/workflows/minimize_fuzz_corpus_workflow.yml @@ -36,7 +36,7 @@ jobs: name: "Minimize ${{ inputs.fuzz_target }}" runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=arm64-medium/disk=large/tag={1}-minimize', github.run_id, inputs.fuzz_target) + && format('runs-on={0}/runner=arm64-medium/disk=large/extras=s3-cache/tag={1}-minimize', github.run_id, inputs.fuzz_target) || 'ubuntu-latest' }} timeout-minutes: 240 steps: @@ -57,6 +57,7 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} toolchain: ${{ env.NIGHTLY_TOOLCHAIN }} + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - name: Install cargo-fuzz uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 diff --git a/.github/workflows/package.yml b/.github/workflows/package.yml index 3de9ac5c81d..4853e48e551 100644 --- a/.github/workflows/package.yml +++ b/.github/workflows/package.yml @@ -39,7 +39,6 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} targets: ${{ matrix.target.target }} - enable-sccache: "false" - name: Install cargo-edit uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 @@ -79,7 +78,7 @@ jobs: PY - name: Set up Python - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6 + uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6 # Latest macOS doesn't allow maturin to install stuff into the global Python interpreter if: "${{ matrix.target.runs-on == 'macos-latest' }}" with: @@ -117,14 +116,13 @@ jobs: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: fetch-depth: 0 - - uses: actions/setup-java@ad2b38190b15e4d6bdf0c97fb4fca8412226d287 # v5 + - uses: actions/setup-java@0f481fcb613427c0f801b606911222b5b6f3083a # v5 with: distribution: "corretto" java-version: "17" - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} - enable-sccache: "false" - run: cargo build --release --package vortex-jni - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7 with: @@ -140,9 +138,9 @@ jobs: matrix: include: - target: x86_64-unknown-linux-gnu - runner: runs-on=${{ github.run_id }}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=prepare-java-linux-amd64 + runner: runs-on=${{ github.run_id }}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=prepare-java-linux-amd64 - target: aarch64-unknown-linux-gnu - runner: runs-on=${{ github.run_id }}/runner=arm64-medium/image=ubuntu24-full-arm64-pre-v2/tag=prepare-java-linux-arm64 + runner: runs-on=${{ github.run_id }}/runner=arm64-medium/image=ubuntu24-full-arm64-pre-v2/extras=s3-cache/tag=prepare-java-linux-arm64 runs-on: ${{ matrix.runner }} steps: - uses: runs-on/action@v2 @@ -152,11 +150,13 @@ jobs: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: fetch-depth: 0 - - uses: actions/setup-java@ad2b38190b15e4d6bdf0c97fb4fca8412226d287 # v5 + - uses: actions/setup-java@0f481fcb613427c0f801b606911222b5b6f3083a # v5 with: distribution: "corretto" java-version: "17" - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - uses: mlugg/setup-zig@d1434d08867e3ee9daa34448df10607b98908d29 # v2.2.1 - name: Install cargo-zigbuild uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 diff --git a/.github/workflows/publish-benchmarks-website.yml b/.github/workflows/publish-benchmarks-website.yml deleted file mode 100644 index 5555dcf841a..00000000000 --- a/.github/workflows/publish-benchmarks-website.yml +++ /dev/null @@ -1,38 +0,0 @@ -name: Publish Benchmarks Website - -on: - push: - branches: [develop] - paths: - - "benchmarks-website/**" - -jobs: - publish: - runs-on: ubuntu-latest - timeout-minutes: 10 - permissions: - contents: read - packages: write - steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - - - name: Log in to GHCR - uses: docker/login-action@650006c6eb7dba73a995cc03b0b2d7f5ca915bee # v4 - with: - registry: ghcr.io - username: ${{ github.actor }} - password: ${{ secrets.GITHUB_TOKEN }} - - - name: Set up QEMU - uses: docker/setup-qemu-action@06116385d9baf250c9f4dcb4858b16962ea869c3 # v4 - - - name: Set up Docker Buildx - uses: docker/setup-buildx-action@d7f5e7f509e45cec5c76c4d5afdd7de93d0b3df5 # v4 - - - name: Build and push - uses: docker/build-push-action@f9f3042f7e2789586610d6e8b85c8f03e5195baf # v7 - with: - context: ./benchmarks-website - platforms: linux/arm64 - push: true - tags: ghcr.io/${{ github.repository }}/benchmarks-website:latest diff --git a/.github/workflows/publish-dry-runs.yml b/.github/workflows/publish-dry-runs.yml index 49f539a19e4..e381b3d4d6f 100644 --- a/.github/workflows/publish-dry-runs.yml +++ b/.github/workflows/publish-dry-runs.yml @@ -35,7 +35,6 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} - enable-sccache: "false" - uses: mlugg/setup-zig@d1434d08867e3ee9daa34448df10607b98908d29 # v2.2.1 - name: Install uv uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 @@ -70,9 +69,9 @@ jobs: matrix: include: - target: x86_64-unknown-linux-gnu - runner: runs-on=${{ github.run_id }}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=check-java-publish-build-amd64 + runner: runs-on=${{ github.run_id }}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=check-java-publish-build-amd64 - target: aarch64-unknown-linux-gnu - runner: runs-on=${{ github.run_id }}/runner=arm64-medium/image=ubuntu24-full-arm64-pre-v2/tag=check-java-publish-build-arm64 + runner: runs-on=${{ github.run_id }}/runner=arm64-medium/image=ubuntu24-full-arm64-pre-v2/extras=s3-cache/tag=check-java-publish-build-arm64 runs-on: ${{ matrix.runner }} steps: - uses: runs-on/action@v2 @@ -82,11 +81,13 @@ jobs: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 with: fetch-depth: 0 - - uses: actions/setup-java@ad2b38190b15e4d6bdf0c97fb4fca8412226d287 # v5 + - uses: actions/setup-java@0f481fcb613427c0f801b606911222b5b6f3083a # v5 with: distribution: "corretto" java-version: "17" - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - uses: mlugg/setup-zig@d1434d08867e3ee9daa34448df10607b98908d29 # v2.2.1 - name: Install cargo-zigbuild uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 @@ -105,7 +106,7 @@ jobs: timeout-minutes: 30 runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-xsmall/image=ubuntu24-full-x64-pre-v2/tag=rust-publish-dry-run', github.run_id) + && format('runs-on={0}/runner=amd64-xsmall/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=rust-publish-dry-run', github.run_id) || 'ubuntu-latest' }} steps: - uses: runs-on/action@v2 @@ -114,6 +115,8 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - name: Install cargo-edit uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 with: diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index aba77190690..200e6f35a4f 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -29,7 +29,6 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} - enable-sccache: "false" - name: Install cargo-edit uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 with: @@ -89,7 +88,7 @@ jobs: working-directory: ./java steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - - uses: actions/setup-java@ad2b38190b15e4d6bdf0c97fb4fca8412226d287 # v5 + - uses: actions/setup-java@0f481fcb613427c0f801b606911222b5b6f3083a # v5 with: distribution: "corretto" java-version: "17" diff --git a/.github/workflows/release-binaries.yml b/.github/workflows/release-binaries.yml index dcedec293c3..eb933f8a0c8 100644 --- a/.github/workflows/release-binaries.yml +++ b/.github/workflows/release-binaries.yml @@ -31,11 +31,11 @@ jobs: # cargo-zigbuild with glibc 2.31 pinning, matching the Python wheel # and JNI builds. - target: aarch64-unknown-linux-gnu - runner: runs-on=${{ github.run_id }}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=release-vx-aarch64-linux + runner: runs-on=${{ github.run_id }}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=release-vx-aarch64-linux fallback_runner: ubuntu-24.04 archive: tgz - target: x86_64-unknown-linux-gnu - runner: runs-on=${{ github.run_id }}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/tag=release-vx-amd64-linux + runner: runs-on=${{ github.run_id }}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=release-vx-amd64-linux fallback_runner: ubuntu-24.04 archive: tgz steps: @@ -52,7 +52,7 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} targets: ${{ matrix.target }} - enable-sccache: ${{ contains(matrix.target, 'linux') && 'true' || 'false' }} + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && contains(matrix.target, 'linux') && 'true' || 'false' }} - name: Setup Zig if: contains(matrix.target, 'linux') diff --git a/.github/workflows/rerun-approvals.yml b/.github/workflows/rerun-approvals.yml deleted file mode 100644 index 4ec40aea4d4..00000000000 --- a/.github/workflows/rerun-approvals.yml +++ /dev/null @@ -1,40 +0,0 @@ -# Re-runs the PR Approval Check workflow when a review is submitted or dismissed. -# Re-running creates a new attempt on the existing approvals.yml workflow run, which -# updates the `check-approvals` check_run in place. This avoids the gate getting stuck -# behind a stale failed check_run from the original `pull_request` event run. - -name: Re-run PR Approval Check on Review - -on: - pull_request_review: - types: [submitted, dismissed] - -permissions: - actions: write - -concurrency: - group: rerun-approvals-${{ github.event.pull_request.head.sha }} - cancel-in-progress: true - -jobs: - rerun: - if: github.event.pull_request.base.ref == 'develop' - runs-on: ubuntu-latest - timeout-minutes: 5 - steps: - - name: Re-run approvals workflow for this PR's head SHA - env: - GH_TOKEN: ${{ github.token }} - REPO: ${{ github.repository }} - HEAD_SHA: ${{ github.event.pull_request.head.sha }} - run: | - set -euo pipefail - run_id=$(gh api \ - "repos/$REPO/actions/workflows/approvals.yml/runs?head_sha=$HEAD_SHA" \ - --jq '.workflow_runs[0].id // empty') - if [ -z "$run_id" ]; then - echo "No approvals.yml run for $HEAD_SHA - nothing to re-run." - exit 0 - fi - echo "Re-running approvals.yml run $run_id" - gh api -X POST "repos/$REPO/actions/runs/$run_id/rerun" diff --git a/.github/workflows/run-fuzzer.yml b/.github/workflows/run-fuzzer.yml index e563dbefbee..35effdc6e59 100644 --- a/.github/workflows/run-fuzzer.yml +++ b/.github/workflows/run-fuzzer.yml @@ -64,7 +64,7 @@ jobs: timeout-minutes: 370 # 6 hours 10 minutes runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner={1}/disk=large/tag={2}-fuzz', github.run_id, inputs.runner, inputs.fuzz_name || inputs.fuzz_target) + && format('runs-on={0}/runner={1}/disk=large/extras=s3-cache/tag={2}-fuzz', github.run_id, inputs.runner, inputs.fuzz_name || inputs.fuzz_target) || 'ubuntu-latest' }} outputs: crashes_found: ${{ steps.check.outputs.crashes_found }} @@ -82,6 +82,7 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} toolchain: ${{ env.NIGHTLY_TOOLCHAIN }} + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - name: Install llvm uses: aminya/setup-cpp@v1 diff --git a/.github/workflows/rust-instrumented.yml b/.github/workflows/rust-instrumented.yml index bf45fc7be13..f4ff732be74 100644 --- a/.github/workflows/rust-instrumented.yml +++ b/.github/workflows/rust-instrumented.yml @@ -101,7 +101,7 @@ jobs: --ignore 'home/*' --ignore 'xtask/*' --ignore 'target/*' --ignore 'vortex-error/*' \ --ignore 'vortex-python/*' --ignore 'vortex-jni/*' --ignore 'vortex-flatbuffers/*' \ --ignore 'vortex-proto/*' --ignore 'vortex-tui/*' --ignore 'vortex-datafusion/examples/*' \ - --ignore 'vortex-ffi/examples/*' --ignore '*/arbitrary/*' --ignore '*/arbitrary.rs' --ignore 'vortex-cxx/*' \ + --ignore 'vortex-ffi/examples/*' --ignore '*/arbitrary/*' --ignore '*/arbitrary.rs' \ --ignore benchmarks/* --ignore 'vortex-test/*' \ -o ${{ env.GRCOV_OUTPUT_FILE }} test -s ${{ env.GRCOV_OUTPUT_FILE }} @@ -114,6 +114,42 @@ jobs: flags: ${{ matrix.suite }} use_oidc: true + cxx-coverage: + name: "C++ API (coverage)" + timeout-minutes: 15 + permissions: + id-token: write + runs-on: >- + ${{ github.repository == 'vortex-data/vortex' + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=cxx-coverage', github.run_id) + || 'ubuntu-latest' }} + steps: + - uses: runs-on/action@v2 + if: github.repository == 'vortex-data/vortex' + with: + sccache: s3 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: "true" + - name: Install lcov + run: | + sudo apt-get update + sudo apt-get install -y lcov libjson-xs-perl + - name: Build FFI library + run: cargo build -p vortex-ffi + - name: Build and test C++ API with coverage + working-directory: lang/cpp + run: ./gcov-report.sh + - name: Codecov + uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7 + with: + name: cxx-api + files: lang/cpp/coverage.info + disable_search: true + flags: cxx-api + use_oidc: true + rust-test-sanitizer: strategy: fail-fast: false @@ -194,7 +230,7 @@ jobs: - sanitizer: tsan sanitizer_flags: "-Zsanitizer=thread" name: "Rust/C++ FFI tests (${{ matrix.sanitizer }})" - timeout-minutes: 30 + timeout-minutes: 5 env: ASAN_OPTIONS: "symbolize=1:check_initialization_order=1:detect_leaks=1:leak_check_at_exit=1" LSAN_OPTIONS: "report_objects=1" @@ -217,9 +253,6 @@ jobs: sccache: s3 - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 - uses: ./.github/actions/setup-prebuild - - name: Install rustfilt - run: | - cargo install rustfilt - name: Install Rust nightly toolchain run: | rustup toolchain install $NIGHTLY_TOOLCHAIN @@ -241,8 +274,7 @@ jobs: - name: Run tests run: | set -o pipefail -# re-enable once we can fix this hang -# ./vortex-ffi/build/test/vortex_ffi_test 2>&1 | rustfilt + ./vortex-ffi/build/test/vortex_ffi_test - name: Run examples run: | set -o pipefail @@ -251,11 +283,11 @@ jobs: # error: Unable to walk dir: File system loop found rm -fr vortex-ffi/build/_deps/nanoarrow-src/python - ./vortex-ffi/build/examples/write_sample file.vortex 2>&1 | rustfilt - ./vortex-ffi/build/examples/write_sample file2.vortex 2>&1 | rustfilt - ./vortex-ffi/build/examples/dtype '*.vortex' 2>&1 | rustfilt - ./vortex-ffi/build/examples/scan '*.vortex' 2>&1 | rustfilt - ./vortex-ffi/build/examples/scan_to_arrow '*.vortex' 2>&1 | rustfilt + ./vortex-ffi/build/examples/write_sample file.vortex + ./vortex-ffi/build/examples/write_sample file2.vortex + ./vortex-ffi/build/examples/dtype '*.vortex' + ./vortex-ffi/build/examples/scan '*.vortex' + ./vortex-ffi/build/examples/scan_to_arrow '*.vortex' miri: name: "Rust tests (miri)" diff --git a/.github/workflows/sql-benchmarks.yml b/.github/workflows/sql-benchmarks.yml index b8c2d08f64f..243db83aff8 100644 --- a/.github/workflows/sql-benchmarks.yml +++ b/.github/workflows/sql-benchmarks.yml @@ -541,6 +541,7 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} + enable-sccache: ${{ (inputs.mode != 'pr' || github.event.pull_request.head.repo.fork == false) && 'true' || 'false' }} - name: Install uv uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 with: @@ -548,7 +549,7 @@ jobs: - name: Install DuckDB run: | - wget -qO- https://github.com/duckdb/duckdb/releases/download/v1.5.3/duckdb_cli-linux-amd64.zip | funzip > duckdb + wget -qO- https://github.com/duckdb/duckdb/releases/download/v1.5.4/duckdb_cli-linux-amd64.zip | funzip > duckdb chmod +x duckdb echo "$PWD" >> "$GITHUB_PATH" @@ -708,46 +709,22 @@ jobs: run: | bash scripts/cat-s3.sh vortex-ci-benchmark-results data.json.gz results.json - - name: Ingest results to v3 server - if: inputs.mode == 'develop' && vars.V3_INGEST_URL != '' - shell: bash - env: - INGEST_BEARER_TOKEN: ${{ secrets.INGEST_BEARER_TOKEN }} - run: | - python3 scripts/post-ingest.py results.v3.jsonl \ - --server "${{ vars.V3_INGEST_URL }}" \ - --commit-sha "${{ github.sha }}" \ - --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 + # v4 (Postgres) ingest -- the REQUIRED benchmark-results pipeline feeding the live + # benchmarks website (see bench.yml for the full rationale); a failure here fails the + # job. Gated on inputs.mode == 'develop' + the ingest-role ARN var (the assume-role + # input that MUST exist for OIDC to succeed). post-ingest.py mints the RDS IAM token + # (boto3) from the assumed GitHubBenchmarkIngestRole; sslmode=verify-full validates + # the cert. No "Install uv" step here: this job already installed the uv binary + # in its "Install uv" step above, and the ingest runs `uv run --no-project --with`, + # which needs nothing beyond the binary. - 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) + - name: Ingest results to v4 Postgres 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 }} diff --git a/.github/workflows/sql-vortex-pr.yml b/.github/workflows/sql-vortex-pr.yml new file mode 100644 index 00000000000..2b99eb0e4d9 --- /dev/null +++ b/.github/workflows/sql-vortex-pr.yml @@ -0,0 +1,47 @@ +# Run Vortex SQL benchmark on every commit in a PR. +# This is not gated behind action/benchmark-sql + +name: PR Vortex Query Benchmark + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + cancel-in-progress: false + +on: + pull_request: + types: [opened, synchronize, reopened] + branches: ["develop"] + +permissions: + contents: read + pull-requests: write + id-token: write + +jobs: + sql: + uses: ./.github/workflows/sql-benchmarks.yml + secrets: inherit + with: + mode: "pr" + benchmark_matrix: | + [ + { + "id": "vortex-queries", + "subcommand": "vortex", + "name": "Vortex queries", + "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"} + ], + "iterations": "100" + } + ] diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml index 8acf85f33a1..7c26746415f 100644 --- a/.github/workflows/stale.yml +++ b/.github/workflows/stale.yml @@ -12,7 +12,7 @@ jobs: actions: write pull-requests: write steps: - - uses: actions/stale@eb5cf3af3ac0a1aa4c9c45633dd1ae542a27a899 # v10 + - uses: actions/stale@1e223db275d687790206a7acac4d1a11bd6fe629 # v10 with: # After 30 days of no activity, a PR will be marked as stale days-before-pr-stale: 14 diff --git a/.github/workflows/wasm-fuzz.yml b/.github/workflows/wasm-fuzz.yml index d61b3b60598..30a688447f9 100644 --- a/.github/workflows/wasm-fuzz.yml +++ b/.github/workflows/wasm-fuzz.yml @@ -33,7 +33,6 @@ jobs: toolchain: ${{ env.NIGHTLY_TOOLCHAIN }} targets: "wasm32-wasip1" components: "rust-src" - enable-sccache: "false" - name: Build WASM fuzz target run: | diff --git a/.github/workflows/web.yml b/.github/workflows/web.yml index d03a2b45d14..7ba67e2b245 100644 --- a/.github/workflows/web.yml +++ b/.github/workflows/web.yml @@ -48,7 +48,6 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} - enable-sccache: "false" - name: Install wasm-pack run: curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh working-directory: . @@ -75,7 +74,6 @@ jobs: - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} - enable-sccache: "false" - name: Install wasm-pack run: curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh working-directory: . diff --git a/.gitignore b/.gitignore index bcc8ef746ee..f9613807332 100644 --- a/.gitignore +++ b/.gitignore @@ -224,6 +224,7 @@ compile_commands.json # AI Agents .agents/settings.local.json +.agents/scheduled_tasks.lock .claude/settings.local.json .opencode diff --git a/AGENTS.md b/AGENTS.md index e5c3d0cc13b..30ddeeb9515 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -151,6 +151,47 @@ Notes: self-explanatory code. - Keep public APIs small and consistent with neighboring crates. +## Performance: avoid hidden-cost accessors in hot loops + +The most common performance trap in this codebase is calling a *per-element accessor that +hides non-trivial work* inside an `O(n)` loop, instead of doing the work once over the whole +chunk. The call site looks like a cheap getter, but each call re-pays a cost that is constant +(or amortizable) across the loop, making the loop `O(n · k)`. + +Watch for these accessors used inside `for i in 0..n { ... }`: + +| Per-element accessor (in a loop) | Hidden cost | Bulk replacement | +| --- | --- | --- | +| `Validity::is_valid(i)` / `is_null(i)` | for array-backed validity, **allocates an `ExecutionCtx` and runs a scalar lookup per call** | `validity.execute_mask(len, ctx)?` once, then `Mask::value(i)` | +| `array.scalar_at(i)` / `array.execute_scalar(i, ctx)` | per-element execution through the compute stack | canonicalize once (`execute::` / `as_slice`) then index | +| `BitBuffer::value(i)` / `Mask::value(i)` accumulated into a count | recomputes the byte address each call; defeats popcount | `true_count()`, `BitBuffer::count_range(start, end)`, `set_indices()` | +| `BitIterator::next()` to accumulate a running rank/prefix count | bit-at-a-time | `count_range` over each gap (SIMD popcount) | +| re-deriving a value inside the loop (e.g. `self.validity()?` each iteration) | re-runs the derivation `n` times | hoist it above the loop | + +Decide per site — bulk is not always the answer: + +- **Sequential / contiguous access** over an accessor that hides amortizable work → hoist and + go bulk (materialize once, then index or iterate the chunk). +- **Gather over arbitrary indices** → you cannot amortize a per-element *decode*, but you can + still materialize the backing buffer once (e.g. `execute_mask`) and then do cheap `O(1)` + random reads, avoiding per-call context/allocation. +- **The accessor is already genuinely `O(1)`** (e.g. reading an already-materialized `Mask`/ + slice, or a native bitmap) → leave it; bulk would not help. + +Even after materializing into a `Mask`/`BitBuffer`, **do not loop with `value(i)` per +element** to act on each set bit — the per-element branch dominates. Iterate a `u64` word +at a time with all-set / all-unset fast paths: use [`BitBuffer::for_each_set_index`]. It +beats `for i in 0..len { if buf.value(i) {..} }` by 2-45x (more at low density) and beats +collecting `set_indices()` by ~2x at mid/high density, while self-adapting from sparse to +dense — see `vortex-mask/benches/mask_iteration.rs`. Reach for the cached `indices()` / +`slices()` representations when you need them more than once; otherwise `for_each_set_index` +needs no materialization. + +When you touch such a loop, back the change with a benchmark (see +`vortex-array/benches/validity_is_valid.rs` for the `is_valid` case, +`vortex-mask/benches/valid_counts.rs` for the popcount case, and +`vortex-mask/benches/mask_iteration.rs` for the per-element-vs-word-vs-sparse comparison). + ## Tests - Strongly consider `rstest` cases when parameterizing repetitive test logic. @@ -179,6 +220,9 @@ Check new and modified lines against this list before finishing: - Updating expected test output to match buggy behavior without independently verifying the intended semantics. - Silently reducing the scope of an approved plan when implementation is harder than expected. +- Calling a hidden-cost per-element accessor (`Validity::is_valid`, `scalar_at`, `BitBuffer:: + value` accumulation) inside a hot loop instead of materializing once — see + "Performance: avoid hidden-cost accessors in hot loops". ## Summaries diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 6d32374293a..dcdf1cbfe19 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -75,11 +75,13 @@ At the time of writing, the following individuals serve as Committers & Maintain Our CI process enforces an extensive set of linter (e.g., `clippy`) rules, as well as language-specific formatters (e.g., `cargo fmt`). Beyond that, we document additional style guidelines in [STYLE.md](STYLE.md). -## Reporting Issues +## Issues and Questions -Bugs should be filed as [GitHub Issues](https://github.com/vortex-data/vortex/issues). Open-ended -questions and feature requests should be filed as -[GitHub Discussions](https://github.com/vortex-data/vortex/discussions). +Bugs, feature requests, and questions should all be filed as +[GitHub Issues](https://github.com/vortex-data/vortex/issues/new/choose). We strongly prefer that +you use one of the provided issue templates rather than opening a blank issue; templates make sure +we get the information needed to act on your report. For quick questions, the +[Vortex Slack channel](https://vortex.dev/slack) is also a good option. ## Developer Certificate of Origin (DCO) diff --git a/Cargo.lock b/Cargo.lock index cb9e255846b..59ac23a8bc0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -192,9 +192,9 @@ checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" [[package]] name = "arrayvec" -version = "0.7.7" +version = "0.7.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f02882884d3e1bc524fb12c79f107f6ad0e1cfd498c536ffb494301740995dfe" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" [[package]] name = "arrow" @@ -821,9 +821,9 @@ checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "aws-lc-rs" -version = "1.17.0" +version = "1.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ec2f1fc3ec205783a5da9a7e6c1509cc69dedf09a1949e412c1e18469326d00" +checksum = "4342d8937fc7e5dd9b1c60292261c0670c882a2cd1719cfc11b1af41731e32ad" dependencies = [ "aws-lc-sys", "zeroize", @@ -831,14 +831,15 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.41.0" +version = "0.42.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a2f9779ce85b93ab6170dd940ad0169b5766ff848247aff13bb788b832fe3f4" +checksum = "6d9ceb1da931507a12f4fccea479dccd00da1943e1b4ae72d8e502d707361444" dependencies = [ "cc", "cmake", "dunce", "fs_extra", + "pkg-config", ] [[package]] @@ -1196,39 +1197,6 @@ dependencies = [ "libbz2-rs-sys", ] -[[package]] -name = "camino" -version = "1.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f2d30e4173c4026932d51d31d6b0613b1fd3014bf3f9f8943d4ba139c437ba0" -dependencies = [ - "serde_core", -] - -[[package]] -name = "cargo-platform" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd0061da739915fae12ea00e16397555ed4371a6bb285431aab930f61b0aa4ba" -dependencies = [ - "serde", - "serde_core", -] - -[[package]] -name = "cargo_metadata" -version = "0.23.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef987d17b0a113becdd19d3d0022d04d7ef41f9efe4f3fb63ac44ba61df3ade9" -dependencies = [ - "camino", - "cargo-platform", - "semver", - "serde", - "serde_json", - "thiserror 2.0.18", -] - [[package]] name = "cast" version = "0.3.0" @@ -1265,9 +1233,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.65" +version = "1.2.66" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" +checksum = "f5d6cac793997bd970000024b2934968efe83b382de4fdcf4fcb46b6ee4ad996" dependencies = [ "find-msvc-tools", "jobserver", @@ -1435,17 +1403,6 @@ version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" -[[package]] -name = "codespan-reporting" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" -dependencies = [ - "serde", - "termcolor", - "unicode-width 0.2.2", -] - [[package]] name = "codspeed" version = "5.0.1" @@ -1602,6 +1559,7 @@ dependencies = [ "tokio", "tracing", "vortex", + "vortex-arrow", "vortex-bench", ] @@ -1655,9 +1613,9 @@ dependencies = [ [[package]] name = "console" -version = "0.16.3" +version = "0.16.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d64e8af5551369d19cf50138de61f1c42074ab970f74e99be916646777f8fc87" +checksum = "4fe5f465a4f6fee88fad41b85d990f84c835335e85b5d9e6e63e0d06d28cba7c" dependencies = [ "encode_unicode", "libc", @@ -1703,9 +1661,9 @@ dependencies = [ [[package]] name = "const_for" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9c50fcfdf972929aff202c16b80086aa3cfc6a3a820af714096c58c7c1d0582" +checksum = "988d3bd6bf67b6d7ae2b519c55296fa57411c27c312575e47b2975f8d1d8590b" [[package]] name = "const_format" @@ -1860,9 +1818,9 @@ dependencies = [ [[package]] name = "crossbeam-epoch" -version = "0.9.18" +version = "0.9.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" dependencies = [ "crossbeam-utils", ] @@ -1997,68 +1955,6 @@ dependencies = [ "pin-project-lite", ] -[[package]] -name = "cxx" -version = "1.0.194" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "747d8437319e3a2f43d93b341c137927ca70c0f5dabeea7a005a73665e247c7e" -dependencies = [ - "cc", - "cxx-build", - "cxxbridge-cmd", - "cxxbridge-flags", - "cxxbridge-macro", - "foldhash 0.2.0", - "link-cplusplus", -] - -[[package]] -name = "cxx-build" -version = "1.0.194" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b0f4697d190a142477b16aef7da8a99bfdc41e7e8b1687583c0d23a79c7afc1e" -dependencies = [ - "cc", - "codespan-reporting", - "indexmap 2.14.0", - "proc-macro2", - "quote", - "scratch", - "syn 2.0.118", -] - -[[package]] -name = "cxxbridge-cmd" -version = "1.0.194" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0956799fa8678d4c50eed028f2de1c0552ae183c76e976cf7ca8c4e36a7c328" -dependencies = [ - "clap", - "codespan-reporting", - "indexmap 2.14.0", - "proc-macro2", - "quote", - "syn 2.0.118", -] - -[[package]] -name = "cxxbridge-flags" -version = "1.0.194" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23384a836ab4f0ad98ace7e3955ad2de39de42378ab487dc28d3990392cb283a" - -[[package]] -name = "cxxbridge-macro" -version = "1.0.194" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6acc6b5822b9526adfb4fc377b67128fdd60aac757cc4a741a6278603f763cf" -dependencies = [ - "indexmap 2.14.0", - "proc-macro2", - "quote", - "syn 2.0.118", -] - [[package]] name = "darling" version = "0.23.0" @@ -2230,6 +2126,7 @@ dependencies = [ "tokio", "url", "vortex", + "vortex-arrow", "vortex-bench", "vortex-cuda", "vortex-datafusion", @@ -3508,9 +3405,9 @@ checksum = "ac6b926516df9c60bfa16e107b21086399f8285a44ca9711344b9e553c5146e2" [[package]] name = "defmt" -version = "1.1.0" +version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6e524506490a1953d237cb87b1cfc1e46f88c18f10a22dfe0f507dc6bfc7f7f" +checksum = "e2953bfe4f93bbd20cc71198842756f77d161884c99ebbabc41d80231ded88d1" dependencies = [ "bitflags 1.3.2", "defmt-macros", @@ -3518,12 +3415,11 @@ dependencies = [ [[package]] name = "defmt-macros" -version = "1.1.0" +version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0a27770e9c8f719a79d8b638281f4d828f77d8fd61e0bd94451b9b85e576a0b" +checksum = "bad9c72e7ca2137e0dc3813245a0d282fd6daad32fd800af018306a9169b5fe8" dependencies = [ "defmt-parser", - "proc-macro-error2", "proc-macro2", "quote", "syn 2.0.118", @@ -3984,9 +3880,9 @@ dependencies = [ [[package]] name = "fs-err" -version = "3.3.0" +version = "3.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73fde052dbfc920003cfd2c8e2c6e6d4cc7c1091538c3a24226cec0665ab08c0" +checksum = "b91aa448ca50d7e79433bdf3ee8d99215430d2ec02ade5aefab2a073a1822e8a" dependencies = [ "autocfg", ] @@ -4238,6 +4134,19 @@ dependencies = [ "wkt 0.14.0", ] +[[package]] +name = "geoarrow-cast" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41c308d653690a4e8ef3cbba69696056bd819e624766ece66d64cc26a638acc1" +dependencies = [ + "arrow-schema 58.3.0", + "geo-traits", + "geoarrow-array", + "geoarrow-schema", + "wkt 0.14.0", +] + [[package]] name = "geoarrow-schema" version = "0.8.0" @@ -4313,11 +4222,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", - "js-sys", "libc", "r-efi 5.3.0", "wasip2", - "wasm-bindgen", ] [[package]] @@ -4565,9 +4472,9 @@ dependencies = [ [[package]] name = "humantime" -version = "2.3.0" +version = "2.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" +checksum = "15cdd26707701c53297e2fa6afb323d55fbc1d0810c3aec078ae3ef0424c3c15" [[package]] name = "hybrid-array" @@ -4870,11 +4777,11 @@ dependencies = [ [[package]] name = "indicatif" -version = "0.18.4" +version = "0.18.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb" +checksum = "9433806cd6b4ec1aba79c021c7e4c58fb4c3b9977c085062e611ac929998fb0c" dependencies = [ - "console 0.16.3", + "console 0.16.4", "futures-core", "portable-atomic", "unicode-width 0.2.2", @@ -4906,7 +4813,7 @@ version = "1.48.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "86f0f8fee8c926415c58d6ae43a08523a26faccb2323f5e6b644fe7dd4ef6b82" dependencies = [ - "console 0.16.3", + "console 0.16.4", "once_cell", "similar 2.7.0", "tempfile", @@ -5036,9 +4943,9 @@ dependencies = [ [[package]] name = "jiff" -version = "0.2.29" +version = "0.2.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34f877a98676d2fb664698d74cc6a51ce6c484ce8c770f05d0108ec9090aeb46" +checksum = "ccfe6121cbe750cf81efa362d85c0bde7ea298ec43092d3a193baca59cdbd634" dependencies = [ "defmt", "jiff-static", @@ -5052,9 +4959,9 @@ dependencies = [ [[package]] name = "jiff-static" -version = "0.2.29" +version = "0.2.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0666b5ab5ecaca213fc2a85b8c0083d9004e84ee2d5f9a7e0017aaf50986f25f" +checksum = "e165e897f662d428f3cd3828a919dbe067c2d42bb1031eede74ef9d27ecdedd2" dependencies = [ "proc-macro2", "quote", @@ -5063,9 +4970,9 @@ dependencies = [ [[package]] name = "jiff-tzdb" -version = "0.1.6" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c900ef84826f1338a557697dc8fc601df9ca9af4ac137c7fb61d4c6f2dfd3076" +checksum = "6142247df1a93c2b3587402a19710be3e6e942f1581a1702e76408f2c21d6590" [[package]] name = "jiff-tzdb-platform" @@ -5129,11 +5036,11 @@ dependencies = [ [[package]] name = "jobserver" -version = "0.1.34" +version = "0.1.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" dependencies = [ - "getrandom 0.3.4", + "getrandom 0.4.3", "libc", ] @@ -5864,9 +5771,9 @@ dependencies = [ [[package]] name = "libredox" -version = "0.1.17" +version = "0.1.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f02ab6bace2054fb888a3c16f990117b579d14a3088e472d63c6011fa185c9d3" +checksum = "c943259e342f1e06ff2da7a83eabdfe7f92ce10262688dbf1895ff0b3e6e4652" dependencies = [ "libc", ] @@ -5892,15 +5799,6 @@ dependencies = [ "bitflags 2.13.0", ] -[[package]] -name = "link-cplusplus" -version = "1.0.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f78c730aaa7d0b9336a299029ea49f9ee53b0ed06e9202e8cb7db9bae7b8c82" -dependencies = [ - "cc", -] - [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -6010,9 +5908,9 @@ dependencies = [ [[package]] name = "lzma-rust2" -version = "0.16.4" +version = "0.16.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce716bf1a316f47a280fc76295f6495b5bea4752bca01c3b3885e101b1c23c02" +checksum = "ca93e534d1142d1d0dcca6d25fe302508a5dfb40b302802904577725ea0b695b" dependencies = [ "sha2 0.11.0", ] @@ -6357,9 +6255,9 @@ dependencies = [ [[package]] name = "num-bigint" -version = "0.4.6" +version = "0.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" dependencies = [ "num-integer", "num-traits", @@ -6514,7 +6412,7 @@ dependencies = [ "parking_lot", "percent-encoding", "quick-xml", - "rand 0.10.1", + "rand 0.10.2", "reqwest 0.12.28", "ring", "rustls-pki-types", @@ -7142,28 +7040,6 @@ dependencies = [ "toml_edit", ] -[[package]] -name = "proc-macro-error-attr2" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" -dependencies = [ - "proc-macro2", - "quote", -] - -[[package]] -name = "proc-macro-error2" -version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" -dependencies = [ - "proc-macro-error-attr2", - "proc-macro2", - "quote", - "syn 2.0.118", -] - [[package]] name = "proc-macro2" version = "1.0.106" @@ -7426,15 +7302,16 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.15" +version = "0.11.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" +checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" dependencies = [ "aws-lc-rs", "bytes", - "getrandom 0.3.4", + "getrandom 0.4.3", "lru-slab", - "rand 0.9.4", + "rand 0.10.2", + "rand_pcg", "ring", "rustc-hash", "rustls", @@ -7448,16 +7325,16 @@ dependencies = [ [[package]] name = "quinn-udp" -version = "0.5.14" +version = "0.5.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694" dependencies = [ "cfg_aliases", "libc", "once_cell", "socket2", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -7510,9 +7387,9 @@ dependencies = [ [[package]] name = "rand" -version = "0.10.1" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2e8e8bcc7961af1fdac401278c6a831614941f6164ee3bf4ce61b7edb162207" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" dependencies = [ "chacha20", "getrandom 0.4.3", @@ -7580,7 +7457,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4d431c2703ccf129de4d45253c03f49ebb22b97d6ad79ee3ecfc7e3f4862c1d8" dependencies = [ "num-traits", - "rand 0.10.1", + "rand 0.10.2", +] + +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" +dependencies = [ + "rand_core 0.10.1", ] [[package]] @@ -7609,7 +7495,7 @@ dependencies = [ "clap", "indicatif", "lance-bench", - "rand 0.10.1", + "rand 0.10.2", "rand_distr 0.6.0", "tabled", "tokio", @@ -8080,9 +7966,9 @@ dependencies = [ [[package]] name = "rustc-hash" -version = "2.1.2" +version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" [[package]] name = "rustc-hex" @@ -8147,9 +8033,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.14.1" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" dependencies = [ "web-time", "zeroize", @@ -8260,12 +8146,6 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" -[[package]] -name = "scratch" -version = "1.0.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d68f2ec51b097e4c1a75b681a8bec621909b5e91f15bb7b840c4f2f7b01148b2" - [[package]] name = "scroll" version = "0.11.0" @@ -8306,10 +8186,6 @@ name = "semver" version = "1.0.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" -dependencies = [ - "serde", - "serde_core", -] [[package]] name = "seq-macro" @@ -8975,9 +8851,9 @@ dependencies = [ [[package]] name = "taffy" -version = "0.10.1" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aea22054047c16c3f34d3ac473a2170be1424b1115b2a3adcf28cfb067c88859" +checksum = "73afc801dd6bd47529eaa7c7e90557f107527d1b7c9c7ed7d7803c7b8d0c357f" dependencies = [ "arrayvec", "grid", @@ -8991,12 +8867,6 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" -[[package]] -name = "take_mut" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f764005d11ee5f36500a149ace24e00e3da98b0158b3e2d53a7495660d3f4d60" - [[package]] name = "tap" version = "1.0.1" @@ -9042,15 +8912,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "termcolor" -version = "1.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" -dependencies = [ - "winapi-util", -] - [[package]] name = "terminal_size" version = "0.4.4" @@ -9194,9 +9055,9 @@ dependencies = [ [[package]] name = "time" -version = "0.3.51" +version = "0.3.53" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85c17d80feb7334b40c484e45ed1a5273dfd8bfda537c3be2e74a06a6686f327" +checksum = "18dfaaeddcb932337b5e7866ee7d0ce9b76d2fd092997146f187ec09b4558a50" dependencies = [ "deranged", "js-sys", @@ -9217,9 +9078,9 @@ checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" [[package]] name = "time-macros" -version = "0.2.30" +version = "0.2.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dcef1a61bdb119096e153208ec5cbec23944ce8bca13be5c7f60c634f7403935" +checksum = "c431b87111666e491a90baa837f914fb45cd5dc3c268591b0220ff5057f2085f" dependencies = [ "num-conv", "time-core", @@ -9678,9 +9539,9 @@ checksum = "7fcfc827f90e53a02eaef5e535ee14266c1d569214c6aa70133a624d8a3164ba" [[package]] name = "utf8-width" -version = "0.1.8" +version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1292c0d970b54115d14f2492fe0170adf21d68a1de108eebc51c1df4f346a091" +checksum = "159a7cadce548703edd50d24069bc294c5415ecab0a480e0cd1ca06d112dc94a" [[package]] name = "utf8_iter" @@ -9736,12 +9597,9 @@ dependencies = [ "arrow-array 58.3.0", "codspeed-divan-compat", "fastlanes", - "futures", "mimalloc", "parquet 58.3.0", - "paste", - "rand 0.10.1", - "rand_distr 0.6.0", + "rand 0.10.2", "serde_json", "tokio", "tracing", @@ -9749,6 +9607,7 @@ dependencies = [ "vortex", "vortex-alp", "vortex-array", + "vortex-arrow", "vortex-btrblocks", "vortex-buffer", "vortex-bytebool", @@ -9785,7 +9644,7 @@ dependencies = [ "itertools 0.14.0", "num-traits", "prost 0.14.4", - "rand 0.10.1", + "rand 0.10.2", "rstest", "rustc-hash", "vortex-array", @@ -9804,15 +9663,7 @@ dependencies = [ "arbitrary", "arc-swap", "arcref", - "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", @@ -9828,6 +9679,8 @@ dependencies = [ "inventory", "itertools 0.14.0", "jiff", + "memchr", + "mimalloc", "num-traits", "num_enum", "parking_lot", @@ -9835,8 +9688,10 @@ dependencies = [ "pin-project-lite", "primitive-types", "prost 0.14.4", - "rand 0.10.1", + "rand 0.10.2", "rand_distr 0.6.0", + "regex", + "regex-syntax", "rstest", "rstest_reuse", "rustc-hash", @@ -9872,6 +9727,29 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "vortex-arrow" +version = "0.1.0" +dependencies = [ + "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", + "codspeed-divan-compat", + "itertools 0.14.0", + "num-traits", + "rstest", + "tracing", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-runend", + "vortex-session", +] + [[package]] name = "vortex-bench" version = "0.1.0" @@ -9884,6 +9762,10 @@ dependencies = [ "bzip2", "clap", "futures", + "geo", + "geo-traits", + "geoarrow", + "geoarrow-cast", "get_dir", "glob", "humansize", @@ -9896,7 +9778,7 @@ dependencies = [ "parking_lot", "parquet 56.2.1", "parquet 58.3.0", - "rand 0.10.1", + "rand 0.10.2", "regex", "reqwest 0.13.4", "serde", @@ -9917,8 +9799,10 @@ dependencies = [ "url", "uuid", "vortex", + "vortex-arrow", "vortex-geo", "vortex-tensor", + "wkb", ] [[package]] @@ -9926,10 +9810,10 @@ name = "vortex-btrblocks" version = "0.1.0" dependencies = [ "codspeed-divan-compat", + "insta", "itertools 0.14.0", - "num-traits", "pco", - "rand 0.10.1", + "rand 0.10.2", "rstest", "test-with", "vortex-alp", @@ -9963,7 +9847,6 @@ dependencies = [ "itertools 0.14.0", "memmap2", "num-traits", - "rand 0.10.1", "rstest", "serde", "simdutf8", @@ -9980,7 +9863,6 @@ dependencies = [ "vortex-array", "vortex-buffer", "vortex-error", - "vortex-mask", "vortex-session", ] @@ -10005,9 +9887,9 @@ dependencies = [ "tpchgen-arrow", "vortex", "vortex-array", + "vortex-arrow", "vortex-buffer", "vortex-error", - "vortex-fsst", "vortex-session", ] @@ -10019,11 +9901,9 @@ dependencies = [ "itertools 0.14.0", "num-traits", "parking_lot", - "rand 0.10.1", - "rstest", + "rand 0.10.2", "rustc-hash", "tracing", - "tracing-subscriber", "vortex-array", "vortex-buffer", "vortex-error", @@ -10043,7 +9923,7 @@ dependencies = [ "arrow-schema 58.3.0", "codspeed-divan-compat", "num-traits", - "rand 0.10.1", + "rand 0.10.2", "vortex-buffer", ] @@ -10066,7 +9946,6 @@ dependencies = [ "arrow-schema 58.3.0", "async-trait", "bindgen", - "bytes", "codspeed-criterion-compat-walltime", "cudarc", "fastlanes", @@ -10082,6 +9961,7 @@ dependencies = [ "tracing", "vortex", "vortex-array", + "vortex-arrow", "vortex-cub", "vortex-cuda", "vortex-cuda-macros", @@ -10097,7 +9977,6 @@ dependencies = [ "arrow-schema 58.3.0", "futures", "vortex", - "vortex-array", "vortex-cuda", "vortex-cuda-macros", "vortex-ffi", @@ -10111,21 +9990,6 @@ dependencies = [ "syn 2.0.118", ] -[[package]] -name = "vortex-cxx" -version = "0.1.0" -dependencies = [ - "anyhow", - "arrow-array 58.3.0", - "arrow-schema 58.3.0", - "async-fs", - "cxx", - "futures", - "paste", - "take_mut", - "vortex", -] - [[package]] name = "vortex-datafusion" version = "0.1.0" @@ -10160,6 +10024,7 @@ dependencies = [ "url", "uuid", "vortex", + "vortex-arrow", "vortex-utils", ] @@ -10197,11 +10062,11 @@ version = "0.1.0" dependencies = [ "anyhow", "async-fs", - "async-trait", "bindgen", "bitvec", "cbindgen", "cc", + "codspeed-divan-compat", "custom-labels", "futures", "geo-types", @@ -10253,7 +10118,7 @@ dependencies = [ "lending-iterator", "num-traits", "prost 0.14.4", - "rand 0.10.1", + "rand 0.10.2", "rstest", "vortex-alp", "vortex-array", @@ -10274,18 +10139,15 @@ dependencies = [ "bytes", "cbindgen", "futures", - "itertools 0.14.0", "mimalloc", - "object_store", "paste", - "prost 0.14.4", - "rand 0.10.1", + "rand 0.10.2", "tempfile", "tracing", "tracing-subscriber", - "url", "vortex", "vortex-array", + "vortex-arrow", ] [[package]] @@ -10306,6 +10168,7 @@ dependencies = [ "rstest", "tokio", "tracing", + "url", "vortex-alp", "vortex-array", "vortex-btrblocks", @@ -10351,7 +10214,7 @@ dependencies = [ "fsst-rs", "num-traits", "prost 0.14.4", - "rand 0.10.1", + "rand 0.10.2", "rstest", "test-with", "vortex-array", @@ -10397,10 +10260,13 @@ dependencies = [ "geo-traits", "geo-types", "geoarrow", + "geoarrow-cast", "prost 0.14.4", "rstest", "vortex-array", + "vortex-arrow", "vortex-error", + "vortex-layout", "vortex-session", "wkb", ] @@ -10462,6 +10328,8 @@ dependencies = [ "arrow-array 58.3.0", "arrow-schema 58.3.0", "async-fs", + "async-lock", + "async-trait", "futures", "jni", "object_store", @@ -10472,6 +10340,7 @@ dependencies = [ "tracing-subscriber", "url", "vortex", + "vortex-arrow", "vortex-parquet-variant", ] @@ -10483,6 +10352,7 @@ dependencies = [ "arrow-schema 58.3.0", "prost 0.14.4", "vortex-array", + "vortex-arrow", "vortex-error", "vortex-proto", "vortex-session", @@ -10500,6 +10370,7 @@ dependencies = [ "bit-vec", "flatbuffers", "futures", + "insta", "itertools 0.14.0", "kanal", "moka", @@ -10517,6 +10388,7 @@ dependencies = [ "tokio", "tracing", "vortex-array", + "vortex-arrow", "vortex-btrblocks", "vortex-buffer", "vortex-error", @@ -10568,7 +10440,6 @@ name = "vortex-onpair" version = "0.1.0" dependencies = [ "codspeed-divan-compat", - "memchr", "num-traits", "onpair", "prost 0.14.4", @@ -10594,6 +10465,7 @@ dependencies = [ "rstest", "tokio", "vortex-array", + "vortex-arrow", "vortex-buffer", "vortex-error", "vortex-file", @@ -10613,6 +10485,7 @@ dependencies = [ "prost 0.14.4", "rstest", "vortex-array", + "vortex-arrow", "vortex-buffer", "vortex-error", "vortex-mask", @@ -10648,6 +10521,7 @@ dependencies = [ "url", "vortex", "vortex-array", + "vortex-arrow", "vortex-python-abi", "vortex-tui", ] @@ -10675,10 +10549,9 @@ dependencies = [ "arrow-array 58.3.0", "arrow-row 58.3.0", "arrow-schema 58.3.0", - "bytes", "codspeed-divan-compat", "mimalloc", - "rand 0.10.1", + "rand 0.10.2", "rstest", "smallvec", "vortex-array", @@ -10693,13 +10566,11 @@ name = "vortex-runend" version = "0.1.0" dependencies = [ "arbitrary", - "arrow-array 58.3.0", - "arrow-schema 58.3.0", "codspeed-divan-compat", "itertools 0.14.0", "num-traits", "prost 0.14.4", - "rand 0.10.1", + "rand 0.10.2", "rstest", "vortex-array", "vortex-buffer", @@ -10727,7 +10598,6 @@ dependencies = [ name = "vortex-sequence" version = "0.1.0" dependencies = [ - "itertools 0.14.0", "num-traits", "prost 0.14.4", "rstest", @@ -10761,6 +10631,7 @@ dependencies = [ "prost 0.14.4", "rstest", "vortex-array", + "vortex-arrow", "vortex-buffer", "vortex-error", "vortex-mask", @@ -10783,6 +10654,7 @@ dependencies = [ "sqllogictest", "thiserror 2.0.18", "tokio", + "tracing-subscriber", "vortex", "vortex-datafusion", "vortex-duckdb", @@ -10794,22 +10666,17 @@ version = "0.1.0" dependencies = [ "arrow-array 58.3.0", "arrow-schema 58.3.0", - "codspeed-divan-compat", "half", "itertools 0.14.0", - "mimalloc", "num-traits", "prost 0.14.4", - "rand 0.10.1", - "rand_distr 0.6.0", "rstest", "vortex-array", - "vortex-btrblocks", + "vortex-arrow", "vortex-buffer", "vortex-compressor", "vortex-error", "vortex-session", - "vortex-utils", ] [[package]] @@ -10841,7 +10708,6 @@ dependencies = [ "humansize", "indicatif", "itertools 0.14.0", - "js-sys", "parquet 58.3.0", "ratatui", "ratzilla", @@ -10850,6 +10716,7 @@ dependencies = [ "taffy", "tokio", "vortex", + "vortex-arrow", "vortex-datafusion", "wasm-bindgen", "wasm-bindgen-futures", @@ -10862,7 +10729,6 @@ version = "0.1.0" dependencies = [ "dashmap", "hashbrown 0.17.1", - "vortex-error", ] [[package]] @@ -10878,6 +10744,7 @@ dependencies = [ "serde", "serde_json", "vortex", + "vortex-arrow", "wasm-bindgen", "wasm-bindgen-futures", "web-sys", @@ -10900,7 +10767,6 @@ dependencies = [ name = "vortex-zstd" version = "0.1.0" dependencies = [ - "codspeed-divan-compat", "itertools 0.14.0", "prost 0.14.4", "rstest", @@ -11214,7 +11080,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", + "windows-targets", ] [[package]] @@ -11223,16 +11089,7 @@ version = "0.59.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" dependencies = [ - "windows-targets 0.52.6", -] - -[[package]] -name = "windows-sys" -version = "0.60.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" -dependencies = [ - "windows-targets 0.53.5", + "windows-targets", ] [[package]] @@ -11250,31 +11107,14 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", - "windows_i686_gnullvm 0.52.6", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", -] - -[[package]] -name = "windows-targets" -version = "0.53.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" -dependencies = [ - "windows-link", - "windows_aarch64_gnullvm 0.53.1", - "windows_aarch64_msvc 0.53.1", - "windows_i686_gnu 0.53.1", - "windows_i686_gnullvm 0.53.1", - "windows_i686_msvc 0.53.1", - "windows_x86_64_gnu 0.53.1", - "windows_x86_64_gnullvm 0.53.1", - "windows_x86_64_msvc 0.53.1", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] [[package]] @@ -11292,96 +11132,48 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_aarch64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" - [[package]] name = "windows_i686_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" -[[package]] -name = "windows_i686_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" - [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" - [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_i686_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" -[[package]] -name = "windows_x86_64_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" -[[package]] -name = "windows_x86_64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" - [[package]] name = "winnow" version = "0.7.15" @@ -11485,7 +11277,6 @@ name = "xtask" version = "0.1.0" dependencies = [ "anyhow", - "cargo_metadata", "clap", "prost-build", "xshell", @@ -11493,9 +11284,9 @@ dependencies = [ [[package]] name = "xxhash-rust" -version = "0.8.15" +version = "0.8.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" +checksum = "4d93c89cdc2d3a63c3ec48ffe926931bdc069eafa8e4402fe6d8f790c9d1e576" [[package]] name = "yansi" @@ -11644,9 +11435,9 @@ dependencies = [ [[package]] name = "zlib-rs" -version = "0.6.4" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "977347db8caa080403f6b6b7c1cda9479a8e869316f7e13a59b19076a40f94e3" +checksum = "5431d5661c32445236631278f27946e444ddafe4684cac70b185272d4f9c52d5" [[package]] name = "zmij" diff --git a/Cargo.toml b/Cargo.toml index bf5057432a6..12e91c7ef20 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,7 @@ members = [ "vortex-io", "vortex-proto", "vortex-array", + "vortex-arrow", "vortex-row", "vortex-tensor", "vortex-json", @@ -32,7 +33,6 @@ members = [ "vortex-cuda/gpu-scan-cli", "vortex-cuda/macros", "vortex-cuda/nvcomp", - "vortex-cxx", "vortex-ffi", "fuzz", "vortex-jni", @@ -94,7 +94,6 @@ rust-version = "1.91.0" version = "0.1.0" [workspace.dependencies] -aho-corasick = "1.1.3" anyhow = "1.0.97" arbitrary = "1.3.2" arc-swap = "1.9" @@ -105,7 +104,6 @@ arrow-buffer = "58.3" arrow-cast = "58.3" arrow-data = "58.3" arrow-ipc = "58.3" -arrow-ord = "58.3" arrow-row = "58.3" arrow-schema = "58.3" arrow-select = "58.3" @@ -121,7 +119,6 @@ bit-vec = "0.9.0" bitvec = "1.0.1" bytes = "1.11.1" bzip2 = "0.6.0" -cargo_metadata = "0.23.1" cbindgen = "0.29.0" cc = "1.2" cfg-if = "1.0.1" @@ -152,7 +149,6 @@ datafusion-physical-expr-common = { version = "54" } datafusion-physical-plan = { version = "54" } datafusion-pruning = { version = "54" } datafusion-sqllogictest = { version = "54" } -dirs = "6.0.0" divan = { package = "codspeed-divan-compat", version = "5.0.0" } enum-iterator = "2.0.0" env_logger = "0.11" @@ -165,6 +161,7 @@ geo = "0.31.0" geo-traits = "0.3.0" geo-types = "0.7.19" geoarrow = "0.8.0" +geoarrow-cast = "0.8.0" get_dir = "0.5.0" glob = "0.3.2" goldenfile = "1" @@ -175,7 +172,6 @@ indicatif = "0.18.0" insta = "1.43" inventory = "0.3.20" itertools = "0.14.0" -jetscii = "0.5.3" jiff = "0.2.0" jni = { version = "0.22.0" } kanal = "0.1.1" @@ -185,12 +181,10 @@ libfuzzer-sys = "0.4" libloading = "0.8" liblzma = "0.4" log = { version = "0.4.21" } -loom = { version = "0.7", features = ["checkpoint"] } memchr = "2.8.0" memmap2 = "0.9.5" mimalloc = "0.1.42" moka = { version = "0.12.10", default-features = false } -multiversion = "0.8.0" noodles-bgzf = "0.47.0" noodles-vcf = { version = "0.88.0", features = ["async"] } num-traits = "0.2.19" @@ -222,8 +216,9 @@ quote = "1.0.44" rand = "0.10.1" rand_distr = "0.6" ratatui = { version = "0.30", default-features = false } -regex = "1.11.0" +regex = "1.12" regex-automata = "0.4" +regex-syntax = "0.8.9" reqwest = { version = "0.13.0", features = [ "blocking", "charset", @@ -257,14 +252,14 @@ strum = "0.28" syn = { version = "2.0.117", features = ["full"] } sysinfo = "0.38.0" tabled = { version = "0.21.0", default-features = false } -taffy = "0.10.0" +taffy = "0.12" take_mut = "0.2.2" tar = "0.4" target-lexicon = "0.13" temp-env = "0.3" tempfile = "3" termtree = { version = "1.0" } -test-with = "0.16" +test-with = "=0.16.3" thiserror = "2.0.3" tokio = { version = "1.52" } tokio-stream = "0.1.17" @@ -291,6 +286,7 @@ zstd = { version = "0.13.3", default-features = false, features = [ vortex = { version = "0.1.0", path = "./vortex" } vortex-alp = { version = "0.1.0", path = "./encodings/alp", default-features = false } vortex-array = { version = "0.1.0", path = "./vortex-array", default-features = false } +vortex-arrow = { version = "0.1.0", path = "./vortex-arrow", default-features = false } vortex-btrblocks = { version = "0.1.0", path = "./vortex-btrblocks", default-features = false } vortex-buffer = { version = "0.1.0", path = "./vortex-buffer", default-features = false } vortex-bytebool = { version = "0.1.0", path = "./encodings/bytebool", default-features = false } @@ -329,11 +325,9 @@ vortex-zstd = { version = "0.1.0", path = "./encodings/zstd", default-features = # No version constraints for unpublished crates. vortex-bench = { path = "./vortex-bench", default-features = false } -vortex-compat = { path = "./vortex-test/compat-gen" } vortex-cuda = { path = "./vortex-cuda", default-features = false } vortex-cuda-macros = { path = "./vortex-cuda/macros" } vortex-duckdb = { path = "./vortex-duckdb", default-features = false } -vortex-test-e2e-cuda = { path = "./vortex-test/e2e-cuda", default-features = false } vortex-tui = { path = "./vortex-tui" } [workspace.lints.rust] diff --git a/README.md b/README.md index 8279e9651ca..893a6a1cb47 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ [![CodSpeed Badge](https://img.shields.io/endpoint?url=https://codspeed.io/badge.json)](https://codspeed.io/vortex-data/vortex) [![Crates.io](https://img.shields.io/crates/v/vortex.svg)](https://crates.io/crates/vortex) [![PyPI - Version](https://img.shields.io/pypi/v/vortex-data)](https://pypi.org/project/vortex-data/) -[![Maven - Version](https://img.shields.io/maven-central/v/dev.vortex/vortex-spark)](https://central.sonatype.com/artifact/dev.vortex/vortex-spark) +[![Maven - Version](https://img.shields.io/maven-central/v/dev.vortex/vortex-spark_2.13)](https://central.sonatype.com/artifact/dev.vortex/vortex-spark_2.13) [![codecov](https://codecov.io/github/vortex-data/vortex/graph/badge.svg)](https://codecov.io/github/vortex-data/vortex) [![Cite](https://img.shields.io/badge/cite-CITATION.cff-blue)](CITATION.cff) diff --git a/REUSE.toml b/REUSE.toml index fb3079a046d..6ce32411711 100644 --- a/REUSE.toml +++ b/REUSE.toml @@ -12,11 +12,6 @@ path = "vortex-btrblocks/decision-tree.*" SPDX-FileCopyrightText = "Copyright the Vortex contributors" SPDX-License-Identifier = "CC-BY-4.0" -[[annotations]] -path = "benchmarks-website/**" -SPDX-FileCopyrightText = "Copyright the Vortex contributors" -SPDX-License-Identifier = "CC-BY-4.0" - # Golden files are licensed under CC-BY-4.0. [[annotations]] path = "**/goldenfiles/**" diff --git a/_typos.toml b/_typos.toml index e2a5ff330b8..359b89c80fb 100644 --- a/_typos.toml +++ b/_typos.toml @@ -8,7 +8,10 @@ extend-ignore-re = [ ] [files] -extend-exclude = ["/vortex-bench/**", "/docs/references.bib", "benchmarks/**", "vortex-sqllogictest/slt/**", "encodings/fsst/src/dfa/tests.rs", "encodings/fsst/src/dfa/flat_contains.rs", "benchmarks-website/server/static/**", "benchmarks-website/server/tests/snapshots/**"] +# scripts/measurement_id_golden.json is a frozen hash-pin artifact whose vectors include +# deliberate Unicode edge cases (JSON-escaped, so "café" reads as the token "caf"); its bytes +# must never change, so it is excluded rather than "fixed". +extend-exclude = ["/vortex-bench/**", "/docs/references.bib", "benchmarks/**", "vortex-sqllogictest/slt/**", "encodings/fsst/src/dfa/tests.rs", "encodings/fsst/src/dfa/flat_contains.rs", "scripts/measurement_id_golden.json"] [type.py] extend-ignore-identifiers-re = [ diff --git a/bench-orchestrator/bench_orchestrator/cli.py b/bench-orchestrator/bench_orchestrator/cli.py index a9e37e70309..b9d7f9bb2ef 100644 --- a/bench-orchestrator/bench_orchestrator/cli.py +++ b/bench-orchestrator/bench_orchestrator/cli.py @@ -85,9 +85,11 @@ def run_ref_auto_complete() -> list[str]: return list(map(lambda x: x.run_id, ResultStore().list_runs(limit=None))) -def targets_from_axes(engine: str, format: str) -> tuple[list[BenchmarkTarget], list[str]]: +def targets_from_axes( + engine: str, format: str, benchmark: Benchmark | None = None +) -> tuple[list[BenchmarkTarget], list[str]]: """Resolve legacy engine/format axes into explicit benchmark targets.""" - return resolve_axis_targets(parse_engines(engine), parse_formats(format)) + return resolve_axis_targets(parse_engines(engine), parse_formats(format), benchmark) def backends_for_engines(engines: list[Engine]) -> list[Engine]: @@ -260,7 +262,7 @@ def run( targets = parse_targets_json(targets_json) warnings: list[str] = [] else: - targets, warnings = targets_from_axes(engine, format) + targets, warnings = targets_from_axes(engine, format, benchmark) except ValueError as exc: console.print(f"[red]{exc}[/red]") raise typer.Exit(1) from exc diff --git a/bench-orchestrator/bench_orchestrator/config.py b/bench-orchestrator/bench_orchestrator/config.py index c597e84c6be..4c9df79fef1 100644 --- a/bench-orchestrator/bench_orchestrator/config.py +++ b/bench-orchestrator/bench_orchestrator/config.py @@ -35,6 +35,7 @@ class Format(Enum): PARQUET = "parquet" VORTEX = "vortex" VORTEX_COMPACT = "vortex-compact" + VORTEX_NATIVE = "vortex-geo-native" DUCKDB = "duckdb" LANCE = "lance" @@ -52,6 +53,8 @@ class Benchmark(Enum): POLARSIGNALS = "polarsignals" PUBLIC_BI = "public-bi" STATPOPGEN = "statpopgen" + SPATIALBENCH = "spatialbench" + VORTEX_QUERIES = "vortex" # Engine to supported formats mapping. @@ -67,11 +70,25 @@ class Benchmark(Enum): Format.PARQUET, Format.VORTEX, Format.VORTEX_COMPACT, + Format.VORTEX_NATIVE, Format.DUCKDB, ], Engine.LANCE: [Format.LANCE], } +# Engines each benchmark can run on. Benchmarks default to *every* engine; list one here only to +# restrict it. SpatialBench's queries use DuckDB-specific `ST_*` spatial SQL that DataFusion has no +# functions for yet. +BENCHMARK_ENGINES: dict[Benchmark, frozenset[Engine]] = { + Benchmark.SPATIALBENCH: frozenset({Engine.DUCKDB}), +} + + +def engines_for_benchmark(benchmark: Benchmark) -> frozenset[Engine]: + """Return the engines `benchmark` supports, defaulting to every engine when unrestricted.""" + return BENCHMARK_ENGINES.get(benchmark, frozenset(Engine)) + + T = TypeVar("T") @@ -175,13 +192,16 @@ def parse_formats_json(value: str) -> list[Format]: def resolve_axis_targets( - engines: Iterable[Engine], formats: Iterable[Format] + engines: Iterable[Engine], formats: Iterable[Format], benchmark: Benchmark | None = None ) -> tuple[list[BenchmarkTarget], list[str]]: """Expand engine/format axes into supported explicit targets.""" warnings: list[str] = [] targets: list[BenchmarkTarget] = [] for engine in engines: + if benchmark is not None and engine not in engines_for_benchmark(benchmark): + warnings.append(f"Benchmark {benchmark.value} does not support engine {engine.value}") + continue for fmt in formats: target = BenchmarkTarget(engine=engine, format=fmt).normalized() if not target.is_supported(): @@ -200,7 +220,9 @@ def group_targets_by_backend(targets: Iterable[BenchmarkTarget]) -> dict[Engine, return groups -def validate_targets(targets: Iterable[BenchmarkTarget], options: dict[str, str]) -> list[str]: +def validate_targets( + targets: Iterable[BenchmarkTarget], options: dict[str, str], benchmark: Benchmark | None = None +) -> list[str]: """Validate explicit targets against benchmark runner constraints.""" errors: list[str] = [] @@ -208,6 +230,8 @@ def validate_targets(targets: Iterable[BenchmarkTarget], options: dict[str, str] for target in normalized_targets: if not target.is_supported(): errors.append(f"Format {target.format.value} is not supported by engine {target.engine.value}") + if benchmark is not None and target.engine not in engines_for_benchmark(benchmark): + errors.append(f"Benchmark {benchmark.value} does not support engine {target.engine.value}") if options.get("remote-data-dir") and any(target.format == Format.LANCE for target in normalized_targets): errors.append("Lance format is not supported for remote storage benchmarks.") @@ -242,7 +266,7 @@ def backends(self) -> list[Engine]: def validate(self) -> list[str]: """Validate the configuration and return any errors.""" - return validate_targets(self.targets, self.options) + return validate_targets(self.targets, self.options, self.benchmark) @dataclass diff --git a/bench-orchestrator/tests/test_config.py b/bench-orchestrator/tests/test_config.py index c7e2d6bb291..fa3be9a2df6 100644 --- a/bench-orchestrator/tests/test_config.py +++ b/bench-orchestrator/tests/test_config.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright the Vortex contributors from bench_orchestrator.config import ( + Benchmark, BenchmarkTarget, Engine, Format, @@ -25,6 +26,23 @@ def test_parse_formats_json_accepts_ci_format_arrays() -> None: assert formats == [Format.PARQUET, Format.VORTEX, Format.DUCKDB] +def test_parse_formats_json_accepts_vortex_native() -> None: + formats = parse_formats_json('["parquet","vortex","vortex-geo-native"]') + + assert formats == [Format.PARQUET, Format.VORTEX, Format.VORTEX_NATIVE] + + +def test_resolve_axis_targets_offers_vortex_native_on_duckdb_only() -> None: + # vortex-geo-native is a DuckDB-only lane; the DataFusion axis is dropped as unsupported. + targets, warnings = resolve_axis_targets( + [Engine.DATAFUSION, Engine.DUCKDB], + [Format.VORTEX_NATIVE], + ) + + assert targets == [BenchmarkTarget(engine=Engine.DUCKDB, format=Format.VORTEX_NATIVE)] + assert warnings == ["Format vortex-geo-native is not supported by engine datafusion"] + + def test_resolve_axis_targets_filters_unsupported_combinations() -> None: targets, warnings = resolve_axis_targets( [Engine.DATAFUSION, Engine.DUCKDB], @@ -39,6 +57,48 @@ def test_resolve_axis_targets_filters_unsupported_combinations() -> None: assert warnings == ["Format arrow is not supported by engine duckdb"] +def test_resolve_axis_targets_skips_engines_a_benchmark_cannot_run() -> None: + # SpatialBench is DuckDB-only (ST_* spatial SQL), so the DataFusion axis is dropped with a warning. + targets, warnings = resolve_axis_targets( + [Engine.DATAFUSION, Engine.DUCKDB], + [Format.PARQUET, Format.VORTEX], + Benchmark.SPATIALBENCH, + ) + + assert targets == [ + BenchmarkTarget(engine=Engine.DUCKDB, format=Format.PARQUET), + BenchmarkTarget(engine=Engine.DUCKDB, format=Format.VORTEX), + ] + assert warnings == ["Benchmark spatialbench does not support engine datafusion"] + + +def test_resolve_axis_targets_expands_spatialbench_three_lanes() -> None: + # The single-command three-lane comparison: parquet, WKB vortex, and native-geometry vortex, all + # on DuckDB. + targets, warnings = resolve_axis_targets( + [Engine.DUCKDB], + [Format.PARQUET, Format.VORTEX, Format.VORTEX_NATIVE], + Benchmark.SPATIALBENCH, + ) + + assert targets == [ + BenchmarkTarget(engine=Engine.DUCKDB, format=Format.PARQUET), + BenchmarkTarget(engine=Engine.DUCKDB, format=Format.VORTEX), + BenchmarkTarget(engine=Engine.DUCKDB, format=Format.VORTEX_NATIVE), + ] + assert warnings == [] + + +def test_validate_targets_rejects_engine_a_benchmark_cannot_run() -> None: + errors = validate_targets( + [BenchmarkTarget(engine=Engine.DATAFUSION, format=Format.PARQUET)], + {}, + Benchmark.SPATIALBENCH, + ) + + assert errors == ["Benchmark spatialbench does not support engine datafusion"] + + def test_validate_targets_rejects_remote_lance() -> None: errors = validate_targets( [BenchmarkTarget(engine=Engine.DATAFUSION, format=Format.LANCE)], diff --git a/bench-orchestrator/tests/test_executor.py b/bench-orchestrator/tests/test_executor.py index dd3253a22ff..86e92b22a3b 100644 --- a/bench-orchestrator/tests/test_executor.py +++ b/bench-orchestrator/tests/test_executor.py @@ -33,6 +33,19 @@ def test_build_command_adds_duckdb_cleanup_flag() -> None: assert "scale-factor=1.0" in cmd +def test_build_command_serializes_vortex_native_format() -> None: + executor = BenchmarkExecutor(Path("/tmp/duckdb-bench"), Engine.DUCKDB) + + cmd = executor.build_command( + benchmark=Benchmark.SPATIALBENCH, + formats=[Format.PARQUET, Format.VORTEX, Format.VORTEX_NATIVE], + iterations=1, + options={"scale-factor": "1.0"}, + ) + + assert "parquet,vortex,vortex-geo-native" in cmd + + def test_build_command_omits_formats_for_lance_backend() -> None: executor = BenchmarkExecutor(Path("/tmp/lance-bench"), Engine.LANCE) diff --git a/benchmarks-website/.dockerignore b/benchmarks-website/.dockerignore deleted file mode 100644 index a21f178d3ab..00000000000 --- a/benchmarks-website/.dockerignore +++ /dev/null @@ -1,3 +0,0 @@ -node_modules -dist -.git diff --git a/benchmarks-website/.gitignore b/benchmarks-website/.gitignore deleted file mode 100644 index 76add878f8d..00000000000 --- a/benchmarks-website/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -node_modules -dist \ No newline at end of file diff --git a/benchmarks-website/Dockerfile b/benchmarks-website/Dockerfile deleted file mode 100644 index 1f87a7148b5..00000000000 --- a/benchmarks-website/Dockerfile +++ /dev/null @@ -1,16 +0,0 @@ -FROM node:24-alpine AS build -WORKDIR /app -COPY package.json package-lock.json ./ -RUN npm ci -COPY . . -RUN npm run build - -FROM node:24-alpine -WORKDIR /app -COPY package.json package-lock.json ./ -RUN npm ci --omit=dev -COPY --from=build /app/dist ./dist -COPY server.js . -COPY src/config.js ./src/config.js -EXPOSE 3000 -CMD ["node", "server.js"] diff --git a/benchmarks-website/README.md b/benchmarks-website/README.md deleted file mode 100644 index 01bbcc265fd..00000000000 --- a/benchmarks-website/README.md +++ /dev/null @@ -1,119 +0,0 @@ - - -# bench.vortex.dev - -The website behind `bench.vortex.dev`. The directory currently houses **two -implementations side by side**, run together until the v3 cutover lands: - -- **v2** (top-level files: `server.js`, `src/`, `index.html`, `vite.config.js`, - `package.json`, `Dockerfile`, `docker-compose.yml`, `public/`). The Node + - React stack that has shipped to production for the life of the site. Built - and published by - [`.github/workflows/publish-benchmarks-website.yml`](../.github/workflows/publish-benchmarks-website.yml). -- **v3** (`server/` + `migrate/` + `ops/`). A single Rust binary - - [`vortex-bench-server`](server/) - that owns a DuckDB file on local disk, - serves the API, and renders the HTML. Compiles all static assets - (`chart.umd.js`, `chart-init.js`, `style.css`) into the binary so deploys - are one file plus a database. Built directly on the EC2 host by - [`ops/deploy.sh`](ops/deploy.sh) - see [`ops/README.md`](ops/README.md). - [`migrate/`](migrate/) is a one-shot tool that loads v2's S3 dataset into a - v3 DuckDB; it is throwaway and goes away after cutover. - -Live results are produced by -[`.github/workflows/bench.yml`](../.github/workflows/bench.yml) and -[`.github/workflows/sql-benchmarks.yml`](../.github/workflows/sql-benchmarks.yml), -which CI runs after every push to `develop`. Until cutover the same payload is -emitted to both stacks (v2 via the legacy `--gh-json` path appended to a public -S3 bucket; v3 via `--gh-json-v3` POSTed to `/api/ingest`). - -## v3 architecture in one paragraph - -`axum` (HTTP) + `maud` (compile-time HTML) + embedded `duckdb-rs` over a single -local DB file. Five fact tables (`query_measurements`, `compression_times`, -`compression_sizes`, `random_access_times`, `vector_search_runs`) plus a -`commits` dim table - see [`server/src/schema.rs`](server/src/schema.rs) for -the column contracts. Three HTML routes (`/`, `/chart/{slug}`, -`/group/{slug}`) and four stable JSON routes (`GET /api/groups`, -`GET /api/chart/{slug}`, `GET /api/group/{slug}`, `GET /health`), plus -versioned group shard artifacts and bearer-gated `POST /api/ingest`. The hot -website path serves precomputed, precompressed latest-100 artifacts from an -in-memory read model; pages render chart shells and hydrate groups via shard -artifacts, while full history warms in the background. See -[`server/ARCHITECTURE.md`](server/ARCHITECTURE.md). - -For the per-module crate map and the request-flow walkthrough, see the -`//!` doc on [`server/src/lib.rs`](server/src/lib.rs). The producer side of -the ingest contract lives in -[`vortex-bench/src/v3.rs`](../vortex-bench/src/v3.rs); the historical-data -side in [`migrate/src/classifier.rs`](migrate/src/classifier.rs). - -## Local dev - -```bash -# v3 server (DuckDB lives at ./bench.duckdb by default). -INGEST_BEARER_TOKEN=dev cargo run -p vortex-bench-server -# server logs: "bench server listening addr=127.0.0.1:3000 db=bench.duckdb" - -# v3 historical migrator (writes a fully populated DuckDB the server can open). -cargo run -p vortex-bench-migrate -- run --output ./bench.duckdb -``` - -Ingest fixture data via the snapshot tests' envelopes (see -[`server/tests/common/mod.rs`](server/tests/common/mod.rs)) or by hand-rolling -a JSONL file and POSTing through `scripts/post-ingest.py`. - -```bash -cargo nextest run -p vortex-bench-server -p vortex-bench-migrate -INSTA_UPDATE=auto cargo nextest run -p vortex-bench-server # update snapshots -``` - -For the v2 stack: - -```bash -cd benchmarks-website -npm install -npm run dev -``` - -## Deployment - -v3 runs as a systemd service on a single EC2 host. The full operator -runbook (first-time install, day-to-day, failure modes) is in -[`ops/README.md`](ops/README.md). Summary: - -- A `vortex-bench-deploy.timer` polls `origin/develop` every 60s. If commits - in the range touch `benchmarks-website/server/`, `benchmarks-website/migrate/`, - `Cargo.toml`, or `Cargo.lock`, it builds and atomically swaps the binary, - then verifies `/health`. Otherwise it fast-forwards the working tree and - exits silently. -- A `vortex-bench-backup.timer` fires hourly: it asks the server to write a - per-table Vortex snapshot (`schema.sql` plus one `.vortex` file per - table) via the bearer-gated `/api/admin/snapshot` endpoint, `tar czf`s the - snapshot directory into `.tar.gz`, uploads it to - `s3://vortex-benchmark-results-database/v3-backups/`, and deletes the local - copies. -- For ad-hoc reads against the live DB, `ops/inspect.sh` calls a - bearer-gated `/api/admin/sql` endpoint - no server stop required. - -The v3 server is throwaway-friendly: every request runs against the local -DuckDB file, and a fresh boot reapplies the schema DDL idempotently. The -migrator deletes the target file (and its `.wal`) before populating it, so -re-running `vortex-bench-migrate run --output ...` is safe. - -## Cutover plan (in flight) - -The work to flip `bench.vortex.dev` from v2 to v3 is tracked outside this -repo. The relevant code-side bits: - -- v3 runs alongside v2 on the same EC2 host today and is fed by CI's - dual-write `--gh-json-v3` path. -- v2 keeps shipping unchanged until DNS flips. **Do not touch the top-level - v2 files unless you are doing the cleanup PR opened post-flip.** -- The v2 cleanup PR removes everything top-level under `benchmarks-website/` - that belongs to v2 (`server.js`, `src/`, `index.html`, `vite.config.js`, - `package.json`, `package-lock.json`, `public/`, the top-level `Dockerfile`, - `docker-compose.yml`, and the `publish-benchmarks-website.yml` workflow). - The v3 tree under `server/`, `migrate/`, and `ops/` is untouched. diff --git a/benchmarks-website/docker-compose.yml b/benchmarks-website/docker-compose.yml deleted file mode 100644 index b97482a230a..00000000000 --- a/benchmarks-website/docker-compose.yml +++ /dev/null @@ -1,29 +0,0 @@ -services: - benchmarks-website: - image: ghcr.io/vortex-data/vortex/benchmarks-website:latest - ports: - - "80:3000" - restart: unless-stopped - - vortex-bench-server: - image: ghcr.io/vortex-data/vortex/vortex-bench-server:latest - ports: - - "3001:3000" - environment: - VORTEX_BENCH_DB: "/app/data/bench.duckdb" - VORTEX_BENCH_BIND: "0.0.0.0:3000" - VORTEX_BENCH_LOG: "info,vortex_bench_server=debug" - env_file: - - /etc/vortex-bench/secrets.env - volumes: - - /opt/benchmarks-website/data:/app/data - restart: unless-stopped - - watchtower: - image: containrrr/watchtower - volumes: - - /var/run/docker.sock:/var/run/docker.sock - environment: - - WATCHTOWER_POLL_INTERVAL=60 - - WATCHTOWER_CLEANUP=true - restart: unless-stopped diff --git a/benchmarks-website/index.html b/benchmarks-website/index.html deleted file mode 100644 index e475f3ad254..00000000000 --- a/benchmarks-website/index.html +++ /dev/null @@ -1,36 +0,0 @@ - - - - - - Vortex Benchmarks - - - - - - - - - - - - - - - - - -
- - - diff --git a/benchmarks-website/package-lock.json b/benchmarks-website/package-lock.json deleted file mode 100644 index 85c5abe0a1f..00000000000 --- a/benchmarks-website/package-lock.json +++ /dev/null @@ -1,1294 +0,0 @@ -{ - "name": "vortex-benchmarks-website", - "version": "2.0.0", - "lockfileVersion": 3, - "requires": true, - "packages": { - "": { - "name": "vortex-benchmarks-website", - "version": "2.0.0", - "dependencies": { - "chart.js": "^4.5.1", - "chartjs-plugin-zoom": "^2.2.0", - "downsample": "^1.4.0", - "hammerjs": "^2.0.8", - "lucide-react": "^1.11.0", - "react": "^19.2.5", - "react-chartjs-2": "^5.3.1", - "react-dom": "^19.2.5" - }, - "devDependencies": { - "@types/react": "^19.2.14", - "@types/react-dom": "^19.2.3", - "@vitejs/plugin-react": "^6.0.1", - "concurrently": "^10.0.0", - "vite": "^8.0.10" - }, - "engines": { - "node": ">=18.0.0" - } - }, - "node_modules/@emnapi/core": { - "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.2", - "tslib": "^2.4.0" - } - }, - "node_modules/@emnapi/runtime": { - "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, - "dependencies": { - "tslib": "^2.4.0" - } - }, - "node_modules/@emnapi/wasi-threads": { - "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, - "dependencies": { - "tslib": "^2.4.0" - } - }, - "node_modules/@kurkle/color": { - "version": "0.3.4", - "resolved": "https://registry.npmjs.org/@kurkle/color/-/color-0.3.4.tgz", - "integrity": "sha512-M5UknZPHRu3DEDWoipU6sE8PdkZ6Z/S+v4dD+Ke8IaNlpdSQah50lz1KtcFBa2vsdOnwbbnxJwVM4wty6udA5w==", - "license": "MIT" - }, - "node_modules/@napi-rs/wasm-runtime": { - "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.3" - }, - "funding": { - "type": "github", - "url": "https://github.com/sponsors/Brooooooklyn" - }, - "peerDependencies": { - "@emnapi/core": "^1.7.1", - "@emnapi/runtime": "^1.7.1" - } - }, - "node_modules/@oxc-project/types": { - "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": { - "url": "https://github.com/sponsors/Boshen" - } - }, - "node_modules/@rolldown/binding-android-arm64": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "android" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-darwin-arm64": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-darwin-x64": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-freebsd-x64": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "freebsd" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-linux-arm-gnueabihf": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-linux-arm64-gnu": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-linux-arm64-musl": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-linux-ppc64-gnu": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-linux-s390x-gnu": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-linux-x64-gnu": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-linux-x64-musl": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-openharmony-arm64": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "openharmony" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-wasm32-wasi": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "dependencies": { - "@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.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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/binding-win32-x64-msvc": { - "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" - ], - "dev": true, - "license": "MIT", - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": "^20.19.0 || >=22.12.0" - } - }, - "node_modules/@rolldown/pluginutils": { - "version": "1.0.1", - "resolved": "https://registry.npmjs.org/@rolldown/pluginutils/-/pluginutils-1.0.1.tgz", - "integrity": "sha512-2j9bGt5Jh8hj+vPtgzPtl72j0yRxHAyumoo6TNfAjsLB04UtpSvPbPcDcBMxz7n+9CYB0c1GxQFxYRg2jimqGw==", - "dev": true, - "license": "MIT" - }, - "node_modules/@tybys/wasm-util": { - "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, - "dependencies": { - "tslib": "^2.4.0" - } - }, - "node_modules/@types/hammerjs": { - "version": "2.0.46", - "resolved": "https://registry.npmjs.org/@types/hammerjs/-/hammerjs-2.0.46.tgz", - "integrity": "sha512-ynRvcq6wvqexJ9brDMS4BnBLzmr0e14d6ZJTEShTBWKymQiHwlAyGu0ZPEFI2Fh1U53F7tN9ufClWM5KvqkKOw==", - "license": "MIT" - }, - "node_modules/@types/react": { - "version": "19.2.17", - "resolved": "https://registry.npmjs.org/@types/react/-/react-19.2.17.tgz", - "integrity": "sha512-MXfmqaVPEVgkBT/aY0aGCkRWWtByiYQXo3xdQ8r5RzuFrPiRn8Gar2tQdXSUQ2GKV3bkXckek89V8wQBY2Q/Aw==", - "dev": true, - "license": "MIT", - "dependencies": { - "csstype": "^3.2.2" - } - }, - "node_modules/@types/react-dom": { - "version": "19.2.3", - "resolved": "https://registry.npmjs.org/@types/react-dom/-/react-dom-19.2.3.tgz", - "integrity": "sha512-jp2L/eY6fn+KgVVQAOqYItbF0VY/YApe5Mz2F0aykSO8gx31bYCZyvSeYxCHKvzHG5eZjc+zyaS5BrBWya2+kQ==", - "dev": true, - "license": "MIT", - "peerDependencies": { - "@types/react": "^19.2.0" - } - }, - "node_modules/@vitejs/plugin-react": { - "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.1" - }, - "engines": { - "node": "^20.19.0 || >=22.12.0" - }, - "peerDependencies": { - "@rolldown/plugin-babel": "^0.1.7 || ^0.2.0", - "babel-plugin-react-compiler": "^1.0.0", - "vite": "^8.0.0" - }, - "peerDependenciesMeta": { - "@rolldown/plugin-babel": { - "optional": true - }, - "babel-plugin-react-compiler": { - "optional": true - } - } - }, - "node_modules/ansi-regex": { - "version": "6.2.2", - "resolved": "https://registry.npmjs.org/ansi-regex/-/ansi-regex-6.2.2.tgz", - "integrity": "sha512-Bq3SmSpyFHaWjPk8If9yc6svM8c56dB5BAtW4Qbw5jHTwwXXcTLoRMkpDJp6VL0XzlWaCHTXrkFURMYmD0sLqg==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">=12" - }, - "funding": { - "url": "https://github.com/chalk/ansi-regex?sponsor=1" - } - }, - "node_modules/ansi-styles": { - "version": "6.2.3", - "resolved": "https://registry.npmjs.org/ansi-styles/-/ansi-styles-6.2.3.tgz", - "integrity": "sha512-4Dj6M28JB+oAH8kFkTLUo+a2jwOFkuqb3yucU0CANcRRUbxS0cP0nZYCGjcc3BNXwRIsUVmDGgzawme7zvJHvg==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">=12" - }, - "funding": { - "url": "https://github.com/chalk/ansi-styles?sponsor=1" - } - }, - "node_modules/chalk": { - "version": "5.6.2", - "resolved": "https://registry.npmjs.org/chalk/-/chalk-5.6.2.tgz", - "integrity": "sha512-7NzBL0rN6fMUW+f7A6Io4h40qQlG+xGmtMxfbnH/K7TAtt8JQWVQK+6g0UXKMeVJoyV5EkkNsErQ8pVD3bLHbA==", - "dev": true, - "license": "MIT", - "engines": { - "node": "^12.17.0 || ^14.13 || >=16.0.0" - }, - "funding": { - "url": "https://github.com/chalk/chalk?sponsor=1" - } - }, - "node_modules/chart.js": { - "version": "4.5.1", - "resolved": "https://registry.npmjs.org/chart.js/-/chart.js-4.5.1.tgz", - "integrity": "sha512-GIjfiT9dbmHRiYi6Nl2yFCq7kkwdkp1W/lp2J99rX0yo9tgJGn3lKQATztIjb5tVtevcBtIdICNWqlq5+E8/Pw==", - "license": "MIT", - "dependencies": { - "@kurkle/color": "^0.3.0" - }, - "engines": { - "pnpm": ">=8" - } - }, - "node_modules/chartjs-plugin-zoom": { - "version": "2.2.0", - "resolved": "https://registry.npmjs.org/chartjs-plugin-zoom/-/chartjs-plugin-zoom-2.2.0.tgz", - "integrity": "sha512-in6kcdiTlP6npIVLMd4zXZ08PDUXC52gZ4FAy5oyjk1zX3gKarXMAof7B9eFiisf9WOC3bh2saHg+J5WtLXZeA==", - "license": "MIT", - "dependencies": { - "@types/hammerjs": "^2.0.45", - "hammerjs": "^2.0.8" - }, - "peerDependencies": { - "chart.js": ">=3.2.0" - } - }, - "node_modules/cliui": { - "version": "9.0.1", - "resolved": "https://registry.npmjs.org/cliui/-/cliui-9.0.1.tgz", - "integrity": "sha512-k7ndgKhwoQveBL+/1tqGJYNz097I7WOvwbmmU2AR5+magtbjPWQTS1C5vzGkBC8Ym8UWRzfKUzUUqFLypY4Q+w==", - "dev": true, - "license": "ISC", - "dependencies": { - "string-width": "^7.2.0", - "strip-ansi": "^7.1.0", - "wrap-ansi": "^9.0.0" - }, - "engines": { - "node": ">=20" - } - }, - "node_modules/concurrently": { - "version": "10.0.3", - "resolved": "https://registry.npmjs.org/concurrently/-/concurrently-10.0.3.tgz", - "integrity": "sha512-hc3LH4UaKWd/bbyDK/IGVa4RB6PtQ3CUYwtrkzqHn+wIG3Hr5fhpRlk0L/gCa8ZE1L/Ufj50Zho69cI5w8SQBA==", - "dev": true, - "license": "MIT", - "dependencies": { - "chalk": "5.6.2", - "rxjs": "7.8.2", - "shell-quote": "1.8.4", - "supports-color": "10.2.2", - "tree-kill": "1.2.2", - "yargs": "18.0.0" - }, - "bin": { - "conc": "dist/bin/index.js", - "concurrently": "dist/bin/index.js" - }, - "engines": { - "node": ">=22" - }, - "funding": { - "url": "https://github.com/open-cli-tools/concurrently?sponsor=1" - } - }, - "node_modules/csstype": { - "version": "3.2.3", - "resolved": "https://registry.npmjs.org/csstype/-/csstype-3.2.3.tgz", - "integrity": "sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==", - "dev": true, - "license": "MIT" - }, - "node_modules/detect-libc": { - "version": "2.1.2", - "resolved": "https://registry.npmjs.org/detect-libc/-/detect-libc-2.1.2.tgz", - "integrity": "sha512-Btj2BOOO83o3WyH59e8MgXsxEQVcarkUOpEYrubB0urwnN10yQ364rsiByU11nZlqWYZm05i/of7io4mzihBtQ==", - "dev": true, - "license": "Apache-2.0", - "engines": { - "node": ">=8" - } - }, - "node_modules/downsample": { - "version": "1.4.0", - "resolved": "https://registry.npmjs.org/downsample/-/downsample-1.4.0.tgz", - "integrity": "sha512-teYPhUPxqwtyICt47t1mP/LjhbRV/ghuKb/LmFDbcZ0CjqFD31tn6rVLZoeCEa1xr8+f2skW8UjRiLiGIKQE4w==", - "license": "MIT" - }, - "node_modules/emoji-regex": { - "version": "10.6.0", - "resolved": "https://registry.npmjs.org/emoji-regex/-/emoji-regex-10.6.0.tgz", - "integrity": "sha512-toUI84YS5YmxW219erniWD0CIVOo46xGKColeNQRgOzDorgBi1v4D71/OFzgD9GO2UGKIv1C3Sp8DAn0+j5w7A==", - "dev": true, - "license": "MIT" - }, - "node_modules/escalade": { - "version": "3.2.0", - "resolved": "https://registry.npmjs.org/escalade/-/escalade-3.2.0.tgz", - "integrity": "sha512-WUj2qlxaQtO4g6Pq5c29GTcWGDyd8itL8zTlipgECz3JesAiiOKotd8JU6otB3PACgG6xkJUyVhboMS+bje/jA==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">=6" - } - }, - "node_modules/fdir": { - "version": "6.5.0", - "resolved": "https://registry.npmjs.org/fdir/-/fdir-6.5.0.tgz", - "integrity": "sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">=12.0.0" - }, - "peerDependencies": { - "picomatch": "^3 || ^4" - }, - "peerDependenciesMeta": { - "picomatch": { - "optional": true - } - } - }, - "node_modules/fsevents": { - "version": "2.3.3", - "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz", - "integrity": "sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==", - "dev": true, - "hasInstallScript": true, - "license": "MIT", - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": "^8.16.0 || ^10.6.0 || >=11.0.0" - } - }, - "node_modules/get-caller-file": { - "version": "2.0.5", - "resolved": "https://registry.npmjs.org/get-caller-file/-/get-caller-file-2.0.5.tgz", - "integrity": "sha512-DyFP3BM/3YHTQOCUL/w0OZHR0lpKeGrxotcHWcqNEdnltqFwXVfhEBQ94eIo34AfQpo0rGki4cyIiftY06h2Fg==", - "dev": true, - "license": "ISC", - "engines": { - "node": "6.* || 8.* || >= 10.*" - } - }, - "node_modules/get-east-asian-width": { - "version": "1.6.0", - "resolved": "https://registry.npmjs.org/get-east-asian-width/-/get-east-asian-width-1.6.0.tgz", - "integrity": "sha512-QRbvDIbx6YklUe6RxeTeleMR0yv3cYH6PsPZHcnVn7xv7zO1BHN8r0XETu8n6Ye3Q+ahtSarc3WgtNWmehIBfA==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">=18" - }, - "funding": { - "url": "https://github.com/sponsors/sindresorhus" - } - }, - "node_modules/hammerjs": { - "version": "2.0.8", - "resolved": "https://registry.npmjs.org/hammerjs/-/hammerjs-2.0.8.tgz", - "integrity": "sha512-tSQXBXS/MWQOn/RKckawJ61vvsDpCom87JgxiYdGwHdOa0ht0vzUWDlfioofFCRU0L+6NGDt6XzbgoJvZkMeRQ==", - "license": "MIT", - "engines": { - "node": ">=0.8.0" - } - }, - "node_modules/lightningcss": { - "version": "1.32.0", - "resolved": "https://registry.npmjs.org/lightningcss/-/lightningcss-1.32.0.tgz", - "integrity": "sha512-NXYBzinNrblfraPGyrbPoD19C1h9lfI/1mzgWYvXUTe414Gz/X1FD2XBZSZM7rRTrMA8JL3OtAaGifrIKhQ5yQ==", - "dev": true, - "license": "MPL-2.0", - "dependencies": { - "detect-libc": "^2.0.3" - }, - "engines": { - "node": ">= 12.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/parcel" - }, - "optionalDependencies": { - "lightningcss-android-arm64": "1.32.0", - "lightningcss-darwin-arm64": "1.32.0", - "lightningcss-darwin-x64": "1.32.0", - "lightningcss-freebsd-x64": "1.32.0", - "lightningcss-linux-arm-gnueabihf": "1.32.0", - "lightningcss-linux-arm64-gnu": "1.32.0", - "lightningcss-linux-arm64-musl": "1.32.0", - "lightningcss-linux-x64-gnu": "1.32.0", - "lightningcss-linux-x64-musl": "1.32.0", - "lightningcss-win32-arm64-msvc": "1.32.0", - "lightningcss-win32-x64-msvc": "1.32.0" - } - }, - "node_modules/lightningcss-android-arm64": { - "version": "1.32.0", - "resolved": "https://registry.npmjs.org/lightningcss-android-arm64/-/lightningcss-android-arm64-1.32.0.tgz", - "integrity": "sha512-YK7/ClTt4kAK0vo6w3X+Pnm0D2cf2vPHbhOXdoNti1Ga0al1P4TBZhwjATvjNwLEBCnKvjJc2jQgHXH0NEwlAg==", - "cpu": [ - "arm64" - ], - "dev": true, - "license": "MPL-2.0", - "optional": true, - "os": [ - "android" - ], - "engines": { - "node": ">= 12.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/parcel" - } - }, - "node_modules/lightningcss-darwin-arm64": { - "version": "1.32.0", - "resolved": "https://registry.npmjs.org/lightningcss-darwin-arm64/-/lightningcss-darwin-arm64-1.32.0.tgz", - "integrity": "sha512-RzeG9Ju5bag2Bv1/lwlVJvBE3q6TtXskdZLLCyfg5pt+HLz9BqlICO7LZM7VHNTTn/5PRhHFBSjk5lc4cmscPQ==", - "cpu": [ - "arm64" - ], - "dev": true, - "license": "MPL-2.0", - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": ">= 12.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/parcel" - } - }, - "node_modules/lightningcss-darwin-x64": { - "version": "1.32.0", - "resolved": "https://registry.npmjs.org/lightningcss-darwin-x64/-/lightningcss-darwin-x64-1.32.0.tgz", - "integrity": "sha512-U+QsBp2m/s2wqpUYT/6wnlagdZbtZdndSmut/NJqlCcMLTWp5muCrID+K5UJ6jqD2BFshejCYXniPDbNh73V8w==", - "cpu": [ - "x64" - ], - "dev": true, - "license": "MPL-2.0", - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": ">= 12.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/parcel" - } - }, - "node_modules/lightningcss-freebsd-x64": { - "version": "1.32.0", - "resolved": "https://registry.npmjs.org/lightningcss-freebsd-x64/-/lightningcss-freebsd-x64-1.32.0.tgz", - "integrity": "sha512-JCTigedEksZk3tHTTthnMdVfGf61Fky8Ji2E4YjUTEQX14xiy/lTzXnu1vwiZe3bYe0q+SpsSH/CTeDXK6WHig==", - "cpu": [ - "x64" - ], - "dev": true, - "license": "MPL-2.0", - "optional": true, - "os": [ - "freebsd" - ], - "engines": { - "node": ">= 12.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/parcel" - } - }, - "node_modules/lightningcss-linux-arm-gnueabihf": { - "version": "1.32.0", - "resolved": "https://registry.npmjs.org/lightningcss-linux-arm-gnueabihf/-/lightningcss-linux-arm-gnueabihf-1.32.0.tgz", - "integrity": "sha512-x6rnnpRa2GL0zQOkt6rts3YDPzduLpWvwAF6EMhXFVZXD4tPrBkEFqzGowzCsIWsPjqSK+tyNEODUBXeeVHSkw==", - "cpu": [ - "arm" - ], - "dev": true, - "license": "MPL-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 12.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/parcel" - } - }, - "node_modules/lightningcss-linux-arm64-gnu": { - "version": "1.32.0", - "resolved": "https://registry.npmjs.org/lightningcss-linux-arm64-gnu/-/lightningcss-linux-arm64-gnu-1.32.0.tgz", - "integrity": "sha512-0nnMyoyOLRJXfbMOilaSRcLH3Jw5z9HDNGfT/gwCPgaDjnx0i8w7vBzFLFR1f6CMLKF8gVbebmkUN3fa/kQJpQ==", - "cpu": [ - "arm64" - ], - "dev": true, - "license": "MPL-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 12.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/parcel" - } - }, - "node_modules/lightningcss-linux-arm64-musl": { - "version": "1.32.0", - "resolved": "https://registry.npmjs.org/lightningcss-linux-arm64-musl/-/lightningcss-linux-arm64-musl-1.32.0.tgz", - "integrity": "sha512-UpQkoenr4UJEzgVIYpI80lDFvRmPVg6oqboNHfoH4CQIfNA+HOrZ7Mo7KZP02dC6LjghPQJeBsvXhJod/wnIBg==", - "cpu": [ - "arm64" - ], - "dev": true, - "license": "MPL-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 12.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/parcel" - } - }, - "node_modules/lightningcss-linux-x64-gnu": { - "version": "1.32.0", - "resolved": "https://registry.npmjs.org/lightningcss-linux-x64-gnu/-/lightningcss-linux-x64-gnu-1.32.0.tgz", - "integrity": "sha512-V7Qr52IhZmdKPVr+Vtw8o+WLsQJYCTd8loIfpDaMRWGUZfBOYEJeyJIkqGIDMZPwPx24pUMfwSxxI8phr/MbOA==", - "cpu": [ - "x64" - ], - "dev": true, - "license": "MPL-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 12.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/parcel" - } - }, - "node_modules/lightningcss-linux-x64-musl": { - "version": "1.32.0", - "resolved": "https://registry.npmjs.org/lightningcss-linux-x64-musl/-/lightningcss-linux-x64-musl-1.32.0.tgz", - "integrity": "sha512-bYcLp+Vb0awsiXg/80uCRezCYHNg1/l3mt0gzHnWV9XP1W5sKa5/TCdGWaR/zBM2PeF/HbsQv/j2URNOiVuxWg==", - "cpu": [ - "x64" - ], - "dev": true, - "license": "MPL-2.0", - "optional": true, - "os": [ - "linux" - ], - "engines": { - "node": ">= 12.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/parcel" - } - }, - "node_modules/lightningcss-win32-arm64-msvc": { - "version": "1.32.0", - "resolved": "https://registry.npmjs.org/lightningcss-win32-arm64-msvc/-/lightningcss-win32-arm64-msvc-1.32.0.tgz", - "integrity": "sha512-8SbC8BR40pS6baCM8sbtYDSwEVQd4JlFTOlaD3gWGHfThTcABnNDBda6eTZeqbofalIJhFx0qKzgHJmcPTnGdw==", - "cpu": [ - "arm64" - ], - "dev": true, - "license": "MPL-2.0", - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": ">= 12.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/parcel" - } - }, - "node_modules/lightningcss-win32-x64-msvc": { - "version": "1.32.0", - "resolved": "https://registry.npmjs.org/lightningcss-win32-x64-msvc/-/lightningcss-win32-x64-msvc-1.32.0.tgz", - "integrity": "sha512-Amq9B/SoZYdDi1kFrojnoqPLxYhQ4Wo5XiL8EVJrVsB8ARoC1PWW6VGtT0WKCemjy8aC+louJnjS7U18x3b06Q==", - "cpu": [ - "x64" - ], - "dev": true, - "license": "MPL-2.0", - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": ">= 12.0.0" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/parcel" - } - }, - "node_modules/lucide-react": { - "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.15", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.15.tgz", - "integrity": "sha512-y7Wygv/7mEOvxTuEQDB8StXdMRBWf1kR/tlhAzBRUFkB2jfcLOAxO/SHmOO2zgz1pVgK29/kyupn059/bCHdjA==", - "dev": true, - "funding": [ - { - "type": "github", - "url": "https://github.com/sponsors/ai" - } - ], - "license": "MIT", - "bin": { - "nanoid": "bin/nanoid.cjs" - }, - "engines": { - "node": "^10 || ^12 || ^13.7 || ^14 || >=15.0.1" - } - }, - "node_modules/picocolors": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/picocolors/-/picocolors-1.1.1.tgz", - "integrity": "sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA==", - "dev": true, - "license": "ISC" - }, - "node_modules/picomatch": { - "version": "4.0.4", - "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.4.tgz", - "integrity": "sha512-QP88BAKvMam/3NxH6vj2o21R6MjxZUAd6nlwAS/pnGvN9IVLocLHxGYIzFhg6fUQ+5th6P4dv4eW9jX3DSIj7A==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">=12" - }, - "funding": { - "url": "https://github.com/sponsors/jonschlinkert" - } - }, - "node_modules/postcss": { - "version": "8.5.16", - "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.16.tgz", - "integrity": "sha512-vuwillviilfKZsg0VGj5R/YwwcHx4SLsIOI/7K6mQkWx+l5cUHTjj5g0AasTBcyXsbfTgrwsUNmVUb5xVwyPwg==", - "dev": true, - "funding": [ - { - "type": "opencollective", - "url": "https://opencollective.com/postcss/" - }, - { - "type": "tidelift", - "url": "https://tidelift.com/funding/github/npm/postcss" - }, - { - "type": "github", - "url": "https://github.com/sponsors/ai" - } - ], - "license": "MIT", - "dependencies": { - "nanoid": "^3.3.12", - "picocolors": "^1.1.1", - "source-map-js": "^1.2.1" - }, - "engines": { - "node": "^10 || ^12 || >=14" - } - }, - "node_modules/react": { - "version": "19.2.7", - "resolved": "https://registry.npmjs.org/react/-/react-19.2.7.tgz", - "integrity": "sha512-HNe9WslTbXmFK8o8cmwgAeJFSBvt1bPdHCVKtaaV+WlAN36mpT4hcRpwbf3fY56ar2oIXzsBpOAiIRHAdY0OlQ==", - "license": "MIT", - "engines": { - "node": ">=0.10.0" - } - }, - "node_modules/react-chartjs-2": { - "version": "5.3.1", - "resolved": "https://registry.npmjs.org/react-chartjs-2/-/react-chartjs-2-5.3.1.tgz", - "integrity": "sha512-h5IPXKg9EXpjoBzUfyWJvllMjG2mQ4EiuHQFhms/AjUm0XSZHhyRy2xVmLXHKrtcdrPO4mnGqRtYoD0vp95A0A==", - "license": "MIT", - "peerDependencies": { - "chart.js": "^4.1.1", - "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" - } - }, - "node_modules/react-dom": { - "version": "19.2.7", - "resolved": "https://registry.npmjs.org/react-dom/-/react-dom-19.2.7.tgz", - "integrity": "sha512-t0BRVXvbiE/o20Hfw669rLbMCDWtYZLvmJigy2f0MxsXF+71pxhR3xOkspmsO8h3ZlNzyibAmtCa3l4lYKk6gQ==", - "license": "MIT", - "dependencies": { - "scheduler": "^0.27.0" - }, - "peerDependencies": { - "react": "^19.2.7" - } - }, - "node_modules/rolldown": { - "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.137.0", - "@rolldown/pluginutils": "^1.0.0" - }, - "bin": { - "rolldown": "bin/cli.mjs" - }, - "engines": { - "node": "^20.19.0 || >=22.12.0" - }, - "optionalDependencies": { - "@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": { - "version": "7.8.2", - "resolved": "https://registry.npmjs.org/rxjs/-/rxjs-7.8.2.tgz", - "integrity": "sha512-dhKf903U/PQZY6boNNtAGdWbG85WAbjT/1xYoZIC7FAY0yWapOBQVsVrDl58W86//e1VpMNBtRV4MaXfdMySFA==", - "dev": true, - "license": "Apache-2.0", - "dependencies": { - "tslib": "^2.1.0" - } - }, - "node_modules/scheduler": { - "version": "0.27.0", - "resolved": "https://registry.npmjs.org/scheduler/-/scheduler-0.27.0.tgz", - "integrity": "sha512-eNv+WrVbKu1f3vbYJT/xtiF5syA5HPIMtf9IgY/nKg0sWqzAUEvqY/xm7OcZc/qafLx/iO9FgOmeSAp4v5ti/Q==", - "license": "MIT" - }, - "node_modules/shell-quote": { - "version": "1.8.4", - "resolved": "https://registry.npmjs.org/shell-quote/-/shell-quote-1.8.4.tgz", - "integrity": "sha512-VsC6n6vz1ihYYyZZwX7YZSF5l5x36ca17OC+a69h94YqB7X6XLwf+5MOgynYir2SLFUbl8gIYvBo8K8RoNQ6bQ==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">= 0.4" - }, - "funding": { - "url": "https://github.com/sponsors/ljharb" - } - }, - "node_modules/source-map-js": { - "version": "1.2.1", - "resolved": "https://registry.npmjs.org/source-map-js/-/source-map-js-1.2.1.tgz", - "integrity": "sha512-UXWMKhLOwVKb728IUtQPXxfYU+usdybtUrK/8uGE8CQMvrhOpwvzDBwj0QhSL7MQc7vIsISBG8VQ8+IDQxpfQA==", - "dev": true, - "license": "BSD-3-Clause", - "engines": { - "node": ">=0.10.0" - } - }, - "node_modules/string-width": { - "version": "7.2.0", - "resolved": "https://registry.npmjs.org/string-width/-/string-width-7.2.0.tgz", - "integrity": "sha512-tsaTIkKW9b4N+AEj+SVA+WhJzV7/zMhcSu78mLKWSk7cXMOSHsBKFWUs0fWwq8QyK3MgJBQRX6Gbi4kYbdvGkQ==", - "dev": true, - "license": "MIT", - "dependencies": { - "emoji-regex": "^10.3.0", - "get-east-asian-width": "^1.0.0", - "strip-ansi": "^7.1.0" - }, - "engines": { - "node": ">=18" - }, - "funding": { - "url": "https://github.com/sponsors/sindresorhus" - } - }, - "node_modules/strip-ansi": { - "version": "7.2.0", - "resolved": "https://registry.npmjs.org/strip-ansi/-/strip-ansi-7.2.0.tgz", - "integrity": "sha512-yDPMNjp4WyfYBkHnjIRLfca1i6KMyGCtsVgoKe/z1+6vukgaENdgGBZt+ZmKPc4gavvEZ5OgHfHdrazhgNyG7w==", - "dev": true, - "license": "MIT", - "dependencies": { - "ansi-regex": "^6.2.2" - }, - "engines": { - "node": ">=12" - }, - "funding": { - "url": "https://github.com/chalk/strip-ansi?sponsor=1" - } - }, - "node_modules/supports-color": { - "version": "10.2.2", - "resolved": "https://registry.npmjs.org/supports-color/-/supports-color-10.2.2.tgz", - "integrity": "sha512-SS+jx45GF1QjgEXQx4NJZV9ImqmO2NPz5FNsIHrsDjh2YsHnawpan7SNQ1o8NuhrbHZy9AZhIoCUiCeaW/C80g==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">=18" - }, - "funding": { - "url": "https://github.com/chalk/supports-color?sponsor=1" - } - }, - "node_modules/tinyglobby": { - "version": "0.2.17", - "resolved": "https://registry.npmjs.org/tinyglobby/-/tinyglobby-0.2.17.tgz", - "integrity": "sha512-wXR/dYpcqKmfWpEdZjiKJOwCNFndD0DMnrW/cYjVGttEkBfVgcLFHoNrlj47mjOVic9yyNu65alsgF4NQyTa2g==", - "dev": true, - "license": "MIT", - "dependencies": { - "fdir": "^6.5.0", - "picomatch": "^4.0.4" - }, - "engines": { - "node": ">=12.0.0" - }, - "funding": { - "url": "https://github.com/sponsors/SuperchupuDev" - } - }, - "node_modules/tree-kill": { - "version": "1.2.2", - "resolved": "https://registry.npmjs.org/tree-kill/-/tree-kill-1.2.2.tgz", - "integrity": "sha512-L0Orpi8qGpRG//Nd+H90vFB+3iHnue1zSSGmNOOCh1GLJ7rUKVwV2HvijphGQS2UmhUZewS9VgvxYIdgr+fG1A==", - "dev": true, - "license": "MIT", - "bin": { - "tree-kill": "cli.js" - } - }, - "node_modules/tslib": { - "version": "2.8.1", - "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", - "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", - "dev": true, - "license": "0BSD" - }, - "node_modules/vite": { - "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.1.2", - "tinyglobby": "^0.2.17" - }, - "bin": { - "vite": "bin/vite.js" - }, - "engines": { - "node": "^20.19.0 || >=22.12.0" - }, - "funding": { - "url": "https://github.com/vitejs/vite?sponsor=1" - }, - "optionalDependencies": { - "fsevents": "~2.3.3" - }, - "peerDependencies": { - "@types/node": "^20.19.0 || >=22.12.0", - "@vitejs/devtools": "^0.3.0", - "esbuild": "^0.27.0 || ^0.28.0", - "jiti": ">=1.21.0", - "less": "^4.0.0", - "sass": "^1.70.0", - "sass-embedded": "^1.70.0", - "stylus": ">=0.54.8", - "sugarss": "^5.0.0", - "terser": "^5.16.0", - "tsx": "^4.8.1", - "yaml": "^2.4.2" - }, - "peerDependenciesMeta": { - "@types/node": { - "optional": true - }, - "@vitejs/devtools": { - "optional": true - }, - "esbuild": { - "optional": true - }, - "jiti": { - "optional": true - }, - "less": { - "optional": true - }, - "sass": { - "optional": true - }, - "sass-embedded": { - "optional": true - }, - "stylus": { - "optional": true - }, - "sugarss": { - "optional": true - }, - "terser": { - "optional": true - }, - "tsx": { - "optional": true - }, - "yaml": { - "optional": true - } - } - }, - "node_modules/wrap-ansi": { - "version": "9.0.2", - "resolved": "https://registry.npmjs.org/wrap-ansi/-/wrap-ansi-9.0.2.tgz", - "integrity": "sha512-42AtmgqjV+X1VpdOfyTGOYRi0/zsoLqtXQckTmqTeybT+BDIbM/Guxo7x3pE2vtpr1ok6xRqM9OpBe+Jyoqyww==", - "dev": true, - "license": "MIT", - "dependencies": { - "ansi-styles": "^6.2.1", - "string-width": "^7.0.0", - "strip-ansi": "^7.1.0" - }, - "engines": { - "node": ">=18" - }, - "funding": { - "url": "https://github.com/chalk/wrap-ansi?sponsor=1" - } - }, - "node_modules/y18n": { - "version": "5.0.8", - "resolved": "https://registry.npmjs.org/y18n/-/y18n-5.0.8.tgz", - "integrity": "sha512-0pfFzegeDWJHJIAmTLRP2DwHjdF5s7jo9tuztdQxAhINCdvS+3nGINqPd00AphqJR/0LhANUS6/+7SCb98YOfA==", - "dev": true, - "license": "ISC", - "engines": { - "node": ">=10" - } - }, - "node_modules/yargs": { - "version": "18.0.0", - "resolved": "https://registry.npmjs.org/yargs/-/yargs-18.0.0.tgz", - "integrity": "sha512-4UEqdc2RYGHZc7Doyqkrqiln3p9X2DZVxaGbwhn2pi7MrRagKaOcIKe8L3OxYcbhXLgLFUS3zAYuQjKBQgmuNg==", - "dev": true, - "license": "MIT", - "dependencies": { - "cliui": "^9.0.1", - "escalade": "^3.1.1", - "get-caller-file": "^2.0.5", - "string-width": "^7.2.0", - "y18n": "^5.0.5", - "yargs-parser": "^22.0.0" - }, - "engines": { - "node": "^20.19.0 || ^22.12.0 || >=23" - } - }, - "node_modules/yargs-parser": { - "version": "22.0.0", - "resolved": "https://registry.npmjs.org/yargs-parser/-/yargs-parser-22.0.0.tgz", - "integrity": "sha512-rwu/ClNdSMpkSrUb+d6BRsSkLUq1fmfsY6TOpYzTwvwkg1/NRG85KBy3kq++A8LKQwX6lsu+aWad+2khvuXrqw==", - "dev": true, - "license": "ISC", - "engines": { - "node": "^20.19.0 || ^22.12.0 || >=23" - } - } - } -} diff --git a/benchmarks-website/package.json b/benchmarks-website/package.json deleted file mode 100644 index a795e0fb22c..00000000000 --- a/benchmarks-website/package.json +++ /dev/null @@ -1,32 +0,0 @@ -{ - "name": "vortex-benchmarks-website", - "version": "2.0.0", - "type": "module", - "scripts": { - "dev": "concurrently \"npm run server\" \"npm run vite\"", - "vite": "vite", - "server": "node server.js", - "build": "vite build", - "preview": "vite preview" - }, - "engines": { - "node": ">=18.0.0" - }, - "dependencies": { - "chart.js": "^4.5.1", - "chartjs-plugin-zoom": "^2.2.0", - "downsample": "^1.4.0", - "hammerjs": "^2.0.8", - "lucide-react": "^1.11.0", - "react": "^19.2.5", - "react-chartjs-2": "^5.3.1", - "react-dom": "^19.2.5" - }, - "devDependencies": { - "@types/react": "^19.2.14", - "@types/react-dom": "^19.2.3", - "@vitejs/plugin-react": "^6.0.1", - "concurrently": "^10.0.0", - "vite": "^8.0.10" - } -} diff --git a/benchmarks-website/public/Vortex_Black_NoBG.png b/benchmarks-website/public/Vortex_Black_NoBG.png deleted file mode 100644 index b4caa0e0287..00000000000 Binary files a/benchmarks-website/public/Vortex_Black_NoBG.png and /dev/null differ diff --git a/benchmarks-website/public/Vortex_White_NoBG.png b/benchmarks-website/public/Vortex_White_NoBG.png deleted file mode 100644 index 702d378306b..00000000000 Binary files a/benchmarks-website/public/Vortex_White_NoBG.png and /dev/null differ diff --git a/benchmarks-website/public/android-chrome-192x192.png b/benchmarks-website/public/android-chrome-192x192.png deleted file mode 100644 index 1089ab6a060..00000000000 Binary files a/benchmarks-website/public/android-chrome-192x192.png and /dev/null differ diff --git a/benchmarks-website/public/android-chrome-512x512.png b/benchmarks-website/public/android-chrome-512x512.png deleted file mode 100644 index 3cfec717ee4..00000000000 Binary files a/benchmarks-website/public/android-chrome-512x512.png and /dev/null differ diff --git a/benchmarks-website/public/apple-touch-icon.png b/benchmarks-website/public/apple-touch-icon.png deleted file mode 100644 index 2b34da9ee0c..00000000000 Binary files a/benchmarks-website/public/apple-touch-icon.png and /dev/null differ diff --git a/benchmarks-website/public/favicon-16x16.png b/benchmarks-website/public/favicon-16x16.png deleted file mode 100644 index 81dd776858a..00000000000 Binary files a/benchmarks-website/public/favicon-16x16.png and /dev/null differ diff --git a/benchmarks-website/public/favicon-32x32.png b/benchmarks-website/public/favicon-32x32.png deleted file mode 100644 index 8d2cb9edabc..00000000000 Binary files a/benchmarks-website/public/favicon-32x32.png and /dev/null differ diff --git a/benchmarks-website/public/favicon.ico b/benchmarks-website/public/favicon.ico deleted file mode 100644 index 3d4e2ff9fe0..00000000000 Binary files a/benchmarks-website/public/favicon.ico and /dev/null differ diff --git a/benchmarks-website/public/site.webmanifest b/benchmarks-website/public/site.webmanifest deleted file mode 100644 index ba9480d90de..00000000000 --- a/benchmarks-website/public/site.webmanifest +++ /dev/null @@ -1,19 +0,0 @@ -{ - "name": "Vortex Benchmarks", - "short_name": "Benchmarks", - "icons": [ - { - "src": "/android-chrome-192x192.png", - "sizes": "192x192", - "type": "image/png" - }, - { - "src": "/android-chrome-512x512.png", - "sizes": "512x512", - "type": "image/png" - } - ], - "theme_color": "#ffffff", - "background_color": "#ffffff", - "display": "standalone" -} diff --git a/benchmarks-website/server.js b/benchmarks-website/server.js deleted file mode 100644 index d6d3f98eef9..00000000000 --- a/benchmarks-website/server.js +++ /dev/null @@ -1,775 +0,0 @@ -import http from "http"; -import fs from "fs"; -import path from "path"; -import { fileURLToPath } from "url"; -import zlib from "zlib"; -import readline from "readline"; -import { Readable } from "stream"; -import { LTTB } from "downsample"; -import { QUERY_SUITES, FAN_OUT_GROUPS, ENGINE_RENAMES } from "./src/config.js"; - -const __filename = fileURLToPath(import.meta.url); -const __dirname = path.dirname(__filename); - -// Configuration -const PORT = process.env.PORT || 3000; -const DATA_URL = - process.env.DATA_URL || - "https://vortex-ci-benchmark-results.s3.amazonaws.com/data.json.gz"; -const COMMITS_URL = - process.env.COMMITS_URL || - "https://vortex-ci-benchmark-results.s3.amazonaws.com/commits.json"; -const REFRESH_INTERVAL = process.env.REFRESH_INTERVAL || 5 * 60 * 1000; -const MAX_POINTS = 200; -const USE_LOCAL_DATA = process.env.USE_LOCAL_DATA === "true"; - -// Benchmark groups: non-query groups + simple suites + fan-out suites -const GROUPS = [ - "Compression", - "Compression Size", - ...QUERY_SUITES.filter((s) => !s.skip && !s.fanOut).map((s) => s.displayName), - "Random Access", - ...FAN_OUT_GROUPS, -]; - -const MIME = { - ".html": "text/html", - ".js": "application/javascript", - ".css": "text/css", - ".json": "application/json", - ".png": "image/png", - ".jpg": "image/jpeg", - ".svg": "image/svg+xml", - ".ico": "image/x-icon", - ".woff": "font/woff", - ".woff2": "font/woff2", - ".webmanifest": "application/manifest+json", -}; - -let store = { - commits: [], - groups: {}, - metadata: null, - downsampled: {}, - lastUpdated: null, -}; - -// Utilities -const rename = (s) => ENGINE_RENAMES[s.toLowerCase()] || ENGINE_RENAMES[s] || s; -const geoMean = (arr) => - arr.length - ? Math.pow( - arr.reduce((a, v) => a * v, 1), - 1 / arr.length, - ) - : null; - -// Categorize benchmarks based on name patterns and metadata -function getGroup(benchmark) { - const name = benchmark.name; - const lower = name.toLowerCase(); - - // Random Access: "random-access/..." or "random access/..." - if ( - lower.startsWith("random-access/") || - lower.startsWith("random access/") - ) { - return "Random Access"; - } - - // Compression Size: size measurements - if ( - lower.startsWith("vortex size/") || - lower.startsWith("vortex-file-compressed size/") || - lower.startsWith("parquet size/") || - lower.startsWith("lance size/") || - lower.includes(":raw size/") || - lower.includes(":parquet-zstd size/") || - lower.includes(":lance size/") - ) { - return "Compression Size"; - } - - // Compression: compress/decompress time and ratio measurements - if ( - lower.startsWith("compress time/") || - lower.startsWith("decompress time/") || - lower.startsWith("parquet_rs-zstd compress") || - lower.startsWith("parquet_rs-zstd decompress") || - lower.startsWith("lance compress") || - lower.startsWith("lance decompress") || - lower.startsWith("vortex:lance ratio") || - lower.startsWith("vortex:parquet-zstd ratio") || - lower.startsWith("vortex:raw ratio") - ) { - return "Compression"; - } - - // SQL query suites: match "{prefix}_q..." or "{prefix}/..." - for (const suite of QUERY_SUITES) { - if ( - !lower.startsWith(suite.prefix + "_q") && - !lower.startsWith(suite.prefix + "/") - ) - continue; - if (suite.skip) return null; - if (!suite.fanOut) return suite.displayName; - // Fan-out suites: expand by storage and scale factor - const storage = benchmark.storage?.toUpperCase() === "S3" ? "S3" : "NVMe"; - const rawSf = benchmark.dataset?.[suite.datasetKey]?.scale_factor; - const sf = rawSf ? Math.round(parseFloat(rawSf)) : 1; - return `${suite.displayName} (${storage}) (SF=${sf})`; - } - - return null; -} - -// Format query name for display: "{prefix}_q00" -> "{QUERY_PREFIX} Q0" -function formatQuery(q) { - const lower = q.toLowerCase(); - for (const suite of QUERY_SUITES) { - if (suite.skip) continue; - const m = lower.match(new RegExp(`^${suite.prefix}[_ ]?q(\\d+)`, "i")); - if (m) return `${suite.queryPrefix} Q${parseInt(m[1], 10)}`; - } - return q.toUpperCase().replace(/[_-]/g, " "); -} - -function normalizeChartName(group, chartName) { - if (group === "Compression Size" && chartName === "VORTEX FILE COMPRESSED SIZE") { - return "VORTEX SIZE"; - } - return chartName; -} - -// LTTB downsampling -function lttbIndices(seriesMap, target) { - const keys = [...seriesMap.keys()]; - if (!keys.length) return []; - const len = seriesMap.get(keys[0])?.length || 0; - if (len <= target) return [...Array(len).keys()]; - - const avg = Array(len); - for (let i = 0; i < len; i++) { - let sum = 0, - n = 0; - for (const arr of seriesMap.values()) { - const v = arr[i]?.value ?? arr[i]; - if (v != null && !isNaN(v)) { - sum += v; - n++; - } - } - avg[i] = [i, n ? sum / n : 0]; - } - - const idx = LTTB(avg, target).map((p) => Math.round(p[0])); - if (!idx.includes(0)) idx.unshift(0); - if (!idx.includes(len - 1)) idx.push(len - 1); - return idx.sort((a, b) => a - b); -} - -function downsample(data, factor) { - const target = Math.ceil(data.commits.length / factor); - if (target >= data.commits.length) return data; - - const idx = lttbIndices(data.series, target); - const series = new Map(); - for (const [k, v] of data.series) - series.set( - k, - idx.map((i) => v[i]), - ); - - return { - ...data, - commits: idx.map((i) => data.commits[i]), - series, - originalLength: data.commits.length, - }; -} - -// Data fetching — streams response body directly instead of buffering -async function fetchJsonl(url) { - const res = await fetch(url); - if (!res.ok) throw new Error(`Fetch failed: ${url} ${res.status}`); - return new Promise((resolve, reject) => { - const results = []; - const rl = readline.createInterface({ - input: Readable.fromWeb(res.body), - crlfDelay: Infinity, - }); - rl.on("line", (l) => { - if (l.trim()) - try { - results.push(JSON.parse(l)); - } catch {} - }); - rl.on("close", () => resolve(results)); - rl.on("error", reject); - }); -} - -function readLocalJsonl(fp) { - return new Promise((resolve, reject) => { - const results = []; - const rl = readline.createInterface({ - input: fs.createReadStream(fp), - crlfDelay: Infinity, - }); - rl.on("line", (l) => { - if (l.trim()) - try { - results.push(JSON.parse(l)); - } catch {} - }); - rl.on("close", () => resolve(results)); - rl.on("error", reject); - }); -} - -// Stream benchmark data record-by-record without buffering the entire dataset -async function forEachBenchmark(callback) { - let stream; - if (USE_LOCAL_DATA) { - stream = fs.createReadStream(path.join(__dirname, "sample/data.json")); - } else { - const res = await fetch(DATA_URL); - if (!res.ok) throw new Error(`Fetch failed: ${DATA_URL} ${res.status}`); - stream = Readable.fromWeb(res.body).pipe(zlib.createGunzip()); - } - return new Promise((resolve, reject) => { - const rl = readline.createInterface({ input: stream, crlfDelay: Infinity }); - rl.on("line", (l) => { - if (l.trim()) - try { - callback(JSON.parse(l)); - } catch {} - }); - rl.on("close", resolve); - rl.on("error", reject); - }); -} - -// Main data processing -async function refresh() { - console.log("Refreshing data..."); - const t0 = Date.now(); - - try { - // Load commits first (small dataset, must be fully in memory for indexing) - const commitsArr = USE_LOCAL_DATA - ? await readLocalJsonl(path.join(__dirname, "sample/commits.json")) - : await fetchJsonl(COMMITS_URL); - - // Build commit index (O(1) lookup) - const commitMap = new Map(commitsArr.map((c) => [c.id, c])); - const commits = commitsArr.sort( - (a, b) => new Date(a.timestamp) - new Date(b.timestamp), - ); - const commitIdx = new Map(commits.map((c, i) => [c.id, i])); - - const groups = Object.fromEntries(GROUPS.map((g) => [g, new Map()])); - let missing = 0; - let benchmarkCount = 0; - const uncategorized = new Set(); - - // Stream benchmarks one record at a time to avoid loading all into memory - await forEachBenchmark((b) => { - benchmarkCount++; - const commit = b.commit || commitMap.get(b.commit_id); - if (!commit) { - missing++; - return; - } - - const group = getGroup(b); - if (!group) { - uncategorized.add(b.name.split("/")[0]); - return; - } - if (!groups[group]) return; - - // Random access names have the form: random-access/{dataset}/{pattern}/{format} - // Historical random access names: random-access/{format} - // Other benchmarks use: {query}/{series} - let seriesName, chartName; - const parts = b.name.split("/"); - if (group === "Random Access" && parts.length === 4) { - chartName = `${parts[1]}/${parts[2]}`.toUpperCase().replace(/[_-]/g, " "); - seriesName = rename(parts[3] || "default"); - } else if (group === "Random Access" && parts.length === 2) { - chartName = "RANDOM ACCESS"; - seriesName = rename(parts[1] || "default"); - } else { - seriesName = rename(parts[1] || "default"); - chartName = formatQuery(parts[0]); - } - chartName = normalizeChartName(group, chartName); - if (chartName.includes("PARQUET-UNC")) return; - - // Skip throughput metrics (keep only time/size) - if (b.name.includes(" throughput")) return; - - let unit = b.unit; - if (!unit) { - if (b.name.toLowerCase().includes(" size/")) unit = "bytes"; - else if (b.name.toLowerCase().includes(" ratio ")) unit = "ratio"; - else unit = "ns"; - } - - const sortPos = parts[0].match(/q(\d+)$/i)?.[1] - ? parseInt(RegExp.$1, 10) - : 0; - const idx = commitIdx.get(commit.id); - if (idx === undefined) return; - - let chart = groups[group].get(chartName); - if (!chart) { - let displayUnit = unit; - if (unit === "ns") displayUnit = "ms/iter"; - else if (unit === "bytes") displayUnit = "MiB"; - chart = { - sort_position: sortPos, - commits, - unit: displayUnit, - series: new Map(), - }; - groups[group].set(chartName, chart); - } - - if (!chart.series.has(seriesName)) { - chart.series.set(seriesName, Array(commits.length).fill(null)); - } - - // Convert values: ns -> ms, bytes -> MiB - let val = b.value; - if (unit === "ns" && typeof val === "number") { - val = val / 1e6; // ns to ms - } else if (unit === "bytes" && typeof val === "number") { - val = val / (1024 * 1024); // bytes to MiB - } - - chart.series.get(seriesName)[idx] = { value: val }; - }); - - console.log( - `Processed ${benchmarkCount} benchmarks, ${commitsArr.length} commits`, - ); - - // Log uncategorized benchmarks for debugging - if (uncategorized.size > 0) { - console.log( - `Uncategorized benchmark prefixes (${uncategorized.size}):`, - [...uncategorized].slice(0, 20).join(", "), - ); - } - - // Trim leading empty commits - let firstIdx = commits.length; - for (const gc of Object.values(groups)) { - for (const cd of gc.values()) { - for (const sd of cd.series.values()) { - const i = sd.findIndex((d) => d !== null); - if (i !== -1 && i < firstIdx) firstIdx = i; - } - } - } - - if (firstIdx > 0 && firstIdx < commits.length) { - console.log(`Trimming ${firstIdx} empty commits`); - commits.splice(0, firstIdx); - for (const gc of Object.values(groups)) { - for (const cd of gc.values()) { - cd.commits = commits; - for (const [k, v] of cd.series) cd.series.set(k, v.slice(firstIdx)); - } - } - } - - // Sort charts within groups - for (const gc of Object.values(groups)) { - const sorted = [...gc.entries()].sort( - (a, b) => - a[1].sort_position - b[1].sort_position || a[0].localeCompare(b[0]), - ); - gc.clear(); - for (const [k, v] of sorted) gc.set(k, v); - } - - // Precompute downsampled versions - const downsampled = {}; - for (const [gn, gc] of Object.entries(groups)) { - downsampled[gn] = {}; - for (const [cn, cd] of gc) { - downsampled[gn][cn] = { - "1x": cd, - "2x": downsample(cd, 2), - "4x": downsample(cd, 4), - "8x": downsample(cd, 8), - }; - } - } - - // Count charts per group for logging - const groupCounts = Object.entries(groups) - .map(([n, g]) => `${n}: ${g.size}`) - .filter((s) => !s.endsWith(": 0")); - console.log("Charts per group:", groupCounts.join(", ")); - - store = { - commits, - groups, - metadata: buildMeta(groups, commits), - downsampled, - lastUpdated: new Date().toISOString(), - }; - console.log( - `Refresh done in ${Date.now() - t0}ms (${missing} missing commits)`, - ); - } catch (e) { - console.error("Refresh error:", e); - } -} - -// Summary calculations -function latestIdx(chart) { - for (let i = chart.commits.length - 1; i >= 0; i--) { - for (const s of chart.series.values()) if (s[i]?.value != null) return i; - } - return -1; -} - -function calcSummary(name, charts) { - if (name === "Random Access") { - for (const q of charts.values()) { - const i = latestIdx(q); - if (i === -1) continue; - const vals = new Map(); - for (const [n, d] of q.series) - if (d[i]?.value != null) vals.set(n, d[i].value); - if (!vals.size) continue; - const min = Math.min(...vals.values()); - return { - type: "randomAccess", - title: "Random Access Performance", - rankings: [...vals] - .map(([n, t]) => ({ name: n, time: t, ratio: t / min })) - .sort((a, b) => a.time - b.time), - explanation: "Random access time | Ratio to fastest (lower is better)", - }; - } - return null; - } - - if (name === "Compression") { - const cc = charts.get("VORTEX:PARQUET ZSTD RATIO COMPRESS TIME"); - const dc = charts.get("VORTEX:PARQUET ZSTD RATIO DECOMPRESS TIME"); - if (!cc && !dc) return null; - const i = latestIdx(cc || dc); - if (i === -1) return null; - const collect = (c) => - c - ? [...c.series] - .filter(([n]) => !n.toLowerCase().includes("wide table")) - .map(([, d]) => d[i]?.value) - .filter((v) => v > 0) - .map((v) => 1 / v) - : []; - return { - type: "compression", - title: "Compression Throughput vs Parquet", - compressRatio: geoMean(collect(cc)), - decompressRatio: geoMean(collect(dc)), - datasetCount: collect(cc).length, - explanation: - "Inverse geomean of Vortex/Parquet ratios (higher is better)", - }; - } - - if (name === "Compression Size") { - const c = charts.get("VORTEX:PARQUET ZSTD SIZE"); - if (!c) return null; - const i = latestIdx(c); - if (i === -1) return null; - const ratios = [...c.series] - .filter(([n]) => !n.toLowerCase().includes("wide table")) - .map(([, d]) => d[i]?.value) - .filter((v) => v > 0); - return ratios.length - ? { - type: "compressionSize", - title: "Compression Size Summary", - minRatio: Math.min(...ratios), - meanRatio: geoMean(ratios), - maxRatio: Math.max(...ratios), - datasetCount: ratios.length, - explanation: - "Geomean of Vortex/Parquet size ratios (lower is better)", - } - : null; - } - - if ( - QUERY_SUITES.some( - (s) => - !s.skip && - (name === s.displayName || name.startsWith(s.displayName + " ")), - ) - ) { - const all = new Map(); - for (const q of charts.values()) - for (const n of q.series.keys()) if (!all.has(n)) all.set(n, new Map()); - for (const [qn, qd] of charts) { - for (const [sn, sd] of qd.series) { - for (let i = sd.length - 1; i >= 0; i--) { - if (sd[i]?.value != null) { - all.get(sn).set(qn, sd[i].value); - break; - } - } - } - } - if (!all.size) return null; - - const scores = new Map(); - for (const [sn, qr] of all) { - let total = 0, - max = 0; - for (const v of qr.values()) { - total += v; - max = Math.max(max, v); - } - const penalty = Math.max(300000, max) * 2; - const ratios = []; - for (const qn of charts.keys()) { - let base = Infinity; - for (const m of all.values()) - if (m.has(qn)) base = Math.min(base, m.get(qn)); - if (base < Infinity) - ratios.push((10 + (qr.get(qn) ?? penalty)) / (10 + base)); - } - if (ratios.length) - scores.set(sn, { score: geoMean(ratios), totalRuntime: total }); - } - - return scores.size - ? { - type: "queryBenchmark", - title: "Performance Summary", - rankings: [...scores] - .map(([n, d]) => ({ name: n, ...d })) - .sort((a, b) => a.score - b.score), - explanation: - "Geomean of query time ratio to fastest (lower is better)", - } - : null; - } - return null; -} - -function buildMeta(groups, commits) { - const meta = {}; - for (const [gn, gc] of Object.entries(groups)) { - const charts = [...gc].map(([cn, cd]) => { - const latest = {}; - for (const [sn, sd] of cd.series) { - for (let i = sd.length - 1; i >= 0; i--) - if (sd[i]?.value != null) { - latest[sn] = sd[i].value; - break; - } - } - return { - name: cn, - unit: cd.unit, - series: [...cd.series.keys()], - sortPosition: cd.sort_position, - totalPoints: cd.commits.length, - latestValues: latest, - }; - }); - meta[gn] = { - charts, - totalCharts: charts.length, - hasData: charts.length > 0, - summary: calcSummary(gn, gc), - }; - } - return { - groups: meta, - totalCommits: commits.length, - commits: commits.map((c) => ({ - id: c.id, - message: c.message?.split("\n")[0] || "", - timestamp: c.timestamp, - author: c.author?.name || "Unknown", - })), - lastUpdated: new Date().toISOString(), - }; -} - -// HTTP handlers -const json = (res, code, data) => { - res.writeHead(code, { - "Content-Type": "application/json", - "Access-Control-Allow-Origin": "*", - }); - res.end(JSON.stringify(data)); -}; - -function serveFile(res, fp) { - fs.readFile(fp, (err, data) => { - if (err) { - res.writeHead(err.code === "ENOENT" ? 404 : 500); - return res.end(err.code === "ENOENT" ? "Not Found" : "Error"); - } - const ext = path.extname(fp).toLowerCase(); - const hdrs = { "Content-Type": MIME[ext] || "application/octet-stream" }; - if (ext === ".js") { - hdrs["Cache-Control"] = "no-cache"; - hdrs["Pragma"] = "no-cache"; - } - res.writeHead(200, hdrs); - res.end(data); - }); -} - -function handleData(res, group, chart, start, end, last, startIdx, endIdx) { - if (!store.downsampled) return json(res, 503, { error: "Loading" }); - const gd = store.downsampled[group]; - if (!gd) return json(res, 404, { error: "Group not found" }); - const cv = gd[chart]; - if (!cv) return json(res, 404, { error: "Chart not found" }); - - const full = cv["1x"]; - const ts = (c) => - typeof c?.timestamp === "number" - ? c.timestamp - : new Date(c?.timestamp).getTime(); - - let si = 0, - ei = full.commits.length - 1; - - // Support "last=N" parameter to get the last N commits - if (last && !start && !end && startIdx === null && endIdx === null) { - const n = parseInt(last, 10); - if (n > 0 && n < full.commits.length) { - si = full.commits.length - n; - } - } else if (startIdx !== null || endIdx !== null) { - // Support index-based range (startIdx, endIdx) - if (startIdx !== null) si = Math.max(0, parseInt(startIdx, 10)); - if (endIdx !== null) - ei = Math.min(full.commits.length - 1, parseInt(endIdx, 10)); - } else { - // Timestamp-based range - if (start) { - const t = +start, - i = full.commits.findIndex((c) => ts(c) >= t); - if (i !== -1) si = i; - } - if (end) { - const t = +end; - for (let i = ei; i >= 0; i--) - if (ts(full.commits[i]) <= t) { - ei = i; - break; - } - } - } - - const len = ei - si + 1; - const ver = - len <= MAX_POINTS - ? "1x" - : len <= MAX_POINTS * 2 - ? "2x" - : len <= MAX_POINTS * 4 - ? "4x" - : "8x"; - const cd = cv[ver]; - const val = (d) => d?.value ?? (typeof d === "number" ? d : null); - - let commits, series; - if (ver === "1x") { - commits = full.commits.slice(si, ei + 1); - series = Object.fromEntries( - [...full.series].map(([n, d]) => [n, d.slice(si, ei + 1).map(val)]), - ); - } else { - const s = +ver[0], - dsi = Math.floor(si / s), - dei = Math.min(Math.ceil(ei / s), cd.commits.length - 1); - commits = cd.commits.slice(dsi, dei + 1); - series = Object.fromEntries( - [...cd.series].map(([n, d]) => [n, d.slice(dsi, dei + 1).map(val)]), - ); - } - - json(res, 200, { - group, - chart, - unit: cd.unit, - downsampleLevel: ver, - originalLength: full.commits.length, - requestedRange: { startIndex: si, endIndex: ei, length: len }, - commits: commits.map((c) => ({ - id: c.id, - message: c.message?.split("\n")[0] || "", - timestamp: c.timestamp, - author: c.author?.name || "Unknown", - url: c.url, - })), - series, - }); -} - -const server = http.createServer((req, res) => { - const [path_, qs] = req.url.split("?"); - const params = new URLSearchParams(qs || ""); - - if (req.method === "OPTIONS") { - res.writeHead(204, { - "Access-Control-Allow-Origin": "*", - "Access-Control-Allow-Methods": "GET", - "Access-Control-Allow-Headers": "Content-Type", - }); - return res.end(); - } - - if (path_ === "/api/metadata") - return store.metadata - ? json(res, 200, store.metadata) - : json(res, 503, { error: "Loading" }); - - if (path_.startsWith("/api/data/")) { - const p = path_.slice(10).split("/"); - return handleData( - res, - decodeURIComponent(p[0] || ""), - decodeURIComponent(p.slice(1).join("/") || ""), - params.get("start"), - params.get("end"), - params.get("last"), - params.has("startIdx") ? params.get("startIdx") : null, - params.has("endIdx") ? params.get("endIdx") : null, - ); - } - - const fp = path.join(__dirname, "dist", path_ === "/" ? "index.html" : path_); - if (!fp.startsWith(__dirname) || fp.includes("/sample/")) { - res.writeHead(403); - return res.end("Forbidden"); - } - serveFile(res, fp); -}); - -async function start() { - console.log("Starting server..."); - await refresh(); - setInterval(refresh, REFRESH_INTERVAL); - server.listen(PORT, () => console.log(`Server at http://localhost:${PORT}`)); -} - -start().catch(console.error); diff --git a/benchmarks-website/src/App.jsx b/benchmarks-website/src/App.jsx deleted file mode 100644 index 0df05bebf01..00000000000 --- a/benchmarks-website/src/App.jsx +++ /dev/null @@ -1,295 +0,0 @@ -import React, { useState, useEffect, useCallback, useMemo, useRef } from 'react'; -import Header from './components/Header'; -import Sidebar from './components/Sidebar'; -import BenchmarkSection from './components/BenchmarkSection'; -import Modal from './components/Modal'; -import { fetchMetadata } from './api'; -import { BENCHMARK_CONFIGS, CATEGORY_TAGS } from './config'; - -export default function App() { - const [metadata, setMetadata] = useState(null); - const [loading, setLoading] = useState(true); - const [error, setError] = useState(null); - const [expandedGroups, setExpandedGroups] = useState(new Set()); - const [sidebarOpen, setSidebarOpen] = useState(false); - const [categoryFilter, setCategoryFilter] = useState('all'); - const [searchFilter, setSearchFilter] = useState(''); - const [viewMode, setViewMode] = useState('grid'); - const [modalChart, setModalChart] = useState(null); - const [showBackToTop, setShowBackToTop] = useState(false); - const metadataFetched = useRef(false); - - useEffect(() => { - if (metadataFetched.current) return; - metadataFetched.current = true; - - async function loadMetadata() { - try { - const data = await fetchMetadata(); - setMetadata(data); - const params = new URLSearchParams(window.location.search); - if (params.get('expanded') === 'true' && data?.groups) { - setExpandedGroups(new Set(Object.keys(data.groups))); - } - } catch (err) { - setError(err.message); - } finally { - setLoading(false); - } - } - loadMetadata(); - }, []); - - useEffect(() => { - const handleScroll = () => { - setShowBackToTop(window.scrollY > 400); - }; - window.addEventListener('scroll', handleScroll, { passive: true }); - return () => window.removeEventListener('scroll', handleScroll); - }, []); - - // Handle hash-based navigation on page load - useEffect(() => { - if (!metadata || loading) return; - - const hash = window.location.hash; - if (hash && hash.startsWith('#group-')) { - const groupId = hash.slice(1); // Remove the '#' - const groupName = groupId.replace('group-', '').replace(/-/g, ' '); - - // Find the matching group (case-insensitive, handle hyphenated names) - const matchingGroup = Object.keys(metadata.groups).find(name => - name.replace(/\s+/g, '-') === groupId.replace('group-', '') - ); - - if (matchingGroup) { - // Expand the group - setExpandedGroups(prev => new Set([...prev, matchingGroup])); - - // Scroll to the element after a short delay to allow rendering - setTimeout(() => { - const element = document.getElementById(groupId); - if (element) { - const headerHeight = 72; - const y = element.getBoundingClientRect().top + window.scrollY - headerHeight - 16; - window.scrollTo({ top: y, behavior: 'smooth' }); - } - }, 100); - } - } - }, [metadata, loading]); - - // Get benchmark config by group name - const getBenchmarkConfig = useCallback((groupName) => { - return BENCHMARK_CONFIGS.find(c => c.name === groupName) || {}; - }, []); - - // Filter groups based on category and search - const filteredGroups = useMemo(() => { - if (!metadata?.groups) return []; - - return Object.keys(metadata.groups).filter(groupName => { - // Category filter - if (categoryFilter !== 'all') { - const tags = CATEGORY_TAGS[groupName] || []; - if (!tags.includes(categoryFilter)) return false; - } - - // Search filter - if (searchFilter) { - const searchLower = searchFilter.toLowerCase(); - const matchesGroup = groupName.toLowerCase().includes(searchLower); - const groupData = metadata.groups[groupName]; - const charts = groupData?.charts || []; - const matchesChart = charts.some(c => - c.name.toLowerCase().includes(searchLower) - ); - if (!matchesGroup && !matchesChart) return false; - } - - return true; - }); - }, [metadata, categoryFilter, searchFilter]); - - // Toggle group expansion - const toggleGroup = useCallback((groupName) => { - setExpandedGroups(prev => { - const next = new Set(prev); - if (next.has(groupName)) { - next.delete(groupName); - } else { - next.add(groupName); - } - return next; - }); - }, []); - - // Expand all groups - const expandAll = useCallback(() => { - if (metadata?.groups) { - setExpandedGroups(new Set(Object.keys(metadata.groups))); - const url = new URL(window.location); - url.searchParams.set('expanded', 'true'); - window.history.replaceState(null, '', url); - } - }, [metadata]); - - // Collapse all groups - const collapseAll = useCallback(() => { - setExpandedGroups(new Set()); - const url = new URL(window.location); - url.searchParams.delete('expanded'); - window.history.replaceState(null, '', url); - }, []); - - // Scroll to group - const scrollToGroup = useCallback((groupName) => { - const element = document.getElementById(`group-${groupName.replace(/\s+/g, '-')}`); - if (element) { - const headerHeight = 72; - const y = element.getBoundingClientRect().top + window.scrollY - headerHeight - 16; - window.scrollTo({ top: y, behavior: 'smooth' }); - } - setSidebarOpen(false); - }, []); - - // Back to top - const scrollToTop = useCallback(() => { - window.scrollTo({ top: 0, behavior: 'smooth' }); - }, []); - - // Clear search - const clearFilter = useCallback(() => { - setSearchFilter(''); - setCategoryFilter('all'); - }, []); - - if (loading) { - return ( -
-
setSidebarOpen(!sidebarOpen)} - categoryFilter={categoryFilter} - onCategoryChange={setCategoryFilter} - searchFilter={searchFilter} - onSearchChange={setSearchFilter} - viewMode={viewMode} - onViewModeChange={setViewMode} - onExpandAll={expandAll} - onCollapseAll={collapseAll} - /> -
-
-
-
-

Loading benchmarks...

-
-
-
-
- ); - } - - if (error) { - return ( -
-
setSidebarOpen(!sidebarOpen)} - categoryFilter={categoryFilter} - onCategoryChange={setCategoryFilter} - searchFilter={searchFilter} - onSearchChange={setSearchFilter} - viewMode={viewMode} - onViewModeChange={setViewMode} - onExpandAll={expandAll} - onCollapseAll={collapseAll} - /> -
-
-
-

Error loading benchmarks: {error}

-
-
-
-
- ); - } - - return ( -
-
setSidebarOpen(!sidebarOpen)} - categoryFilter={categoryFilter} - onCategoryChange={setCategoryFilter} - searchFilter={searchFilter} - onSearchChange={setSearchFilter} - viewMode={viewMode} - onViewModeChange={setViewMode} - onExpandAll={expandAll} - onCollapseAll={collapseAll} - /> - -
- setSidebarOpen(false)} - onGroupClick={scrollToGroup} - onClearFilter={clearFilter} - showClearFilter={categoryFilter !== 'all' || searchFilter !== ''} - /> - -
setSidebarOpen(false)} - /> - -
- {filteredGroups.map(groupName => { - const groupData = metadata.groups[groupName] || {}; - const charts = groupData.charts || []; - const config = getBenchmarkConfig(groupName); - const isExpanded = expandedGroups.has(groupName); - - if (config.hidden) return null; - - return ( - toggleGroup(groupName)} - viewMode={viewMode} - onFullscreen={(chartData) => setModalChart(chartData)} - commitRange={metadata.totalCommits} - summary={groupData.summary} - /> - ); - })} -
-
- - {showBackToTop && ( - - )} - - {modalChart && ( - setModalChart(null)} - /> - )} -
- ); -} diff --git a/benchmarks-website/src/api.js b/benchmarks-website/src/api.js deleted file mode 100644 index 042ea7a6f0b..00000000000 --- a/benchmarks-website/src/api.js +++ /dev/null @@ -1,41 +0,0 @@ -const API_BASE = ''; - -export async function fetchMetadata() { - const response = await fetch(`${API_BASE}/api/metadata`); - if (!response.ok) throw new Error(`Failed to fetch metadata: ${response.status}`); - return response.json(); -} - -export async function fetchChartData(groupName, chartName, options = {}) { - const { startTimestamp, endTimestamp, last, startIdx, endIdx } = options; - let url = `${API_BASE}/api/data/${encodeURIComponent(groupName)}/${encodeURIComponent(chartName)}`; - const params = new URLSearchParams(); - - if (last) { - params.set('last', last); - } else if (startIdx !== undefined || endIdx !== undefined) { - // Index-based range - if (startIdx !== undefined) params.set('startIdx', startIdx); - if (endIdx !== undefined) params.set('endIdx', endIdx); - } else { - // Timestamp-based range - if (startTimestamp) { - const ts = typeof startTimestamp === 'number' - ? startTimestamp - : new Date(startTimestamp).getTime(); - params.set('start', ts); - } - if (endTimestamp) { - const ts = typeof endTimestamp === 'number' - ? endTimestamp - : new Date(endTimestamp).getTime(); - params.set('end', ts); - } - } - - if (params.toString()) url += '?' + params.toString(); - - const response = await fetch(url); - if (!response.ok) throw new Error(`Failed to fetch chart data: ${response.status}`); - return response.json(); -} diff --git a/benchmarks-website/src/components/BenchmarkSection.jsx b/benchmarks-website/src/components/BenchmarkSection.jsx deleted file mode 100644 index 706a9c27fe6..00000000000 --- a/benchmarks-website/src/components/BenchmarkSection.jsx +++ /dev/null @@ -1,147 +0,0 @@ -import React, { useState, useCallback, useMemo } from 'react'; -import { Info, Link2 } from 'lucide-react'; -import ChartContainer from './ChartContainer'; -import BenchmarkSummary from './BenchmarkSummary'; -import { getBenchmarkDescription, remapChartName } from '../utils'; - -export default function BenchmarkSection({ - groupName, - charts, - config, - isExpanded, - onToggle, - viewMode, - onFullscreen, - commitRange, - summary, -}) { - const [engineFilter, setEngineFilter] = useState('all'); - const [copiedLink, setCopiedLink] = useState(false); - - // Get unique engines from chart series - const engines = useMemo(() => { - const engineSet = new Set(); - charts.forEach(chart => { - chart.series?.forEach(seriesName => { - if (seriesName.includes(':')) { - const engine = seriesName.split(':')[0].toLowerCase(); - engineSet.add(engine); - } - }); - }); - return Array.from(engineSet).sort(); - }, [charts]); - - // Filter and sort charts based on config - const filteredCharts = useMemo(() => { - if (!charts) return []; - - let result = charts.filter(chart => { - // Apply keptCharts filter - if (config.keptCharts) { - const upperName = chart.name.toUpperCase(); - return config.keptCharts.some(kept => upperName === kept.toUpperCase()); - } - return true; - }); - - // Sort by keptCharts order if specified - if (config.keptCharts) { - const orderMap = new Map(config.keptCharts.map((name, idx) => [name.toUpperCase(), idx])); - result.sort((a, b) => { - const aIdx = orderMap.get(a.name.toUpperCase()) ?? 999; - const bIdx = orderMap.get(b.name.toUpperCase()) ?? 999; - return aIdx - bIdx; - }); - } - - return result; - }, [charts, config]); - - // Copy link to clipboard - const handleCopyLink = useCallback((e) => { - e.stopPropagation(); - const url = `${window.location.origin}${window.location.pathname}#group-${groupName.replace(/\s+/g, '-')}`; - navigator.clipboard.writeText(url); - setCopiedLink(true); - setTimeout(() => setCopiedLink(false), 2000); - }, [groupName]); - - const description = getBenchmarkDescription(groupName); - const hasData = filteredCharts.length > 0; - const chartCount = filteredCharts.length; - - return ( -
-
-
- {isExpanded ? '▼' : '▶'} -
-

- {groupName} - -

- {description && ( - - - - )} -
- {chartCount} {chartCount === 1 ? 'CHART' : 'CHARTS'} -
-
-
- - {isExpanded && engines.length > 0 && ( -
- Filter by engine: - - {engines.map(engine => ( - - ))} -
- )} -
- - - - {isExpanded && ( -
- {filteredCharts.map(chart => ( - - ))} -
- )} -
- ); -} diff --git a/benchmarks-website/src/components/BenchmarkSummary.jsx b/benchmarks-website/src/components/BenchmarkSummary.jsx deleted file mode 100644 index 6bc2d4f1790..00000000000 --- a/benchmarks-website/src/components/BenchmarkSummary.jsx +++ /dev/null @@ -1,129 +0,0 @@ -import React from 'react'; -import { formatTime } from '../utils'; - -// BenchmarkSummary now uses pre-computed summary from metadata (passed via props) -// instead of fetching all chart data -export default function BenchmarkSummary({ groupName, charts, summary }) { - // Use pre-computed summary from metadata - const summaryData = summary; - - if (!summaryData) return null; - - // Query benchmarks (Clickbench, TPC-H, TPC-DS, etc.) - if (summaryData.type === 'queryBenchmark' && summaryData.rankings?.length > 0) { - return ( -
-

{summaryData.title || 'Performance Summary'}

-
- {summaryData.rankings.map((item, idx) => ( -
- #{idx + 1} - {item.name} - - {item.score.toFixed(2)}x - {formatTime(item.totalRuntime)} - -
- ))} -
-
- {summaryData.explanation || 'Score: geometric mean of query time ratio to fastest (lower is better)'} -
-
- ); - } - - if (summaryData.type === 'randomAccess' && summaryData.rankings?.length > 0) { - return ( -
-

{summaryData.title || 'Random Access Performance'}

-
- {summaryData.rankings.map((item, idx) => ( -
- #{idx + 1} - {item.name} - - {formatTime(item.time)} - {item.ratio.toFixed(2)}x - -
- ))} -
-
- {summaryData.explanation || 'Random access time | Ratio to fastest (lower is better)'} -
-
- ); - } - - if (summaryData.type === 'compression') { - return ( -
-

{summaryData.title || 'Compression Throughput vs Parquet'}

-
- {summaryData.compressRatio && ( -
- ⚡ - Write Speed (Compression) - - {summaryData.compressRatio.toFixed(2)}x - -
- )} - {summaryData.decompressRatio && ( -
- 📤 - Scan Speed (Decompression) - - {summaryData.decompressRatio.toFixed(2)}x - -
- )} -
-
- {summaryData.explanation || `Inverse geometric mean of Vortex/Parquet ratios across ${summaryData.datasetCount || 'multiple'} datasets (higher is better)`} -
-
- ); - } - - if (summaryData.type === 'compressionSize' && summaryData.meanRatio) { - return ( -
-

{summaryData.title || 'Compression Size Summary'}

-
- {summaryData.minRatio && ( -
- ⬇️ - Min Size Ratio - - {summaryData.minRatio.toFixed(2)}x - -
- )} -
- 📊 - Mean Size Ratio - - {summaryData.meanRatio.toFixed(2)}x - -
- {summaryData.maxRatio && ( -
- ⬆️ - Max Size Ratio - - {summaryData.maxRatio.toFixed(2)}x - -
- )} -
-
- {summaryData.explanation || `Geometric mean of Vortex/Parquet size ratios across ${summaryData.datasetCount || 'multiple'} datasets (lower is better)`} -
-
- ); - } - - return null; -} diff --git a/benchmarks-website/src/components/ChartContainer.jsx b/benchmarks-website/src/components/ChartContainer.jsx deleted file mode 100644 index 1dfb193e836..00000000000 --- a/benchmarks-website/src/components/ChartContainer.jsx +++ /dev/null @@ -1,664 +0,0 @@ -import { - CategoryScale, - Chart as ChartJS, - Legend, - LinearScale, - LineElement, - PointElement, - Title, - Tooltip, -} from 'chart.js'; -import zoomPlugin from 'chartjs-plugin-zoom'; -import { - ChevronLeft, - ChevronRight, - Expand, - MoveHorizontal, - SkipBack, - SkipForward, - ZoomIn, - ZoomOut, -} from 'lucide-react'; -import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; -import { Line } from 'react-chartjs-2'; -import { fetchChartData } from '../api'; -import { formatDate, stringToColor } from '../utils'; - -ChartJS.register( - CategoryScale, - LinearScale, - PointElement, - LineElement, - Title, - Tooltip, - Legend, - zoomPlugin -); - -// Custom tooltip positioner - 50px from cursor, 50px above nearest point -Tooltip.positioners.topCorner = function(elements, eventPosition) { - const chart = this.chart; - const chartCenter = (chart.chartArea.left + chart.chartArea.right) / 2; - const chartVerticalCenter = (chart.chartArea.top + chart.chartArea.bottom) / 2; - const isOnRightSide = eventPosition.x > chartCenter; - const isOnTopSide = eventPosition.y < chartVerticalCenter; - - let x = isOnRightSide ? eventPosition.x - 150 : eventPosition.x + 150; - let y = isOnTopSide ? chart.chartArea.top + 100 : chart.chartArea.bottom - 100; - return { - x, - y, - xAlign: isOnRightSide ? 'right' : 'left', - yAlign: isOnTopSide ? 'top' : 'bottom', - }; -}; - -const DEFAULT_RANGE_SIZE = 100; - -export default function ChartContainer({ - groupName, - chartName, - displayName, - unit, - config, - engineFilter, - onFullscreen, -}) { - const [totalCommits, setTotalCommits] = useState(null); - const [chartData, setChartData] = useState(null); - const [loading, setLoading] = useState(true); - const [error, setError] = useState(null); - // viewRange stores the requested range: either { last: N } or { startIdx, endIdx } - const [viewRange, setViewRange] = useState({ last: DEFAULT_RANGE_SIZE }); - const chartRef = useRef(null); - const isResettingZoom = useRef(false); - - // Fetch data for the current view range - useEffect(() => { - let cancelled = false; - - async function loadData() { - setLoading(true); - setError(null); - - try { - let options = {}; - if (viewRange.last) { - // Initial load: get last N commits - options = { last: viewRange.last }; - } else if (viewRange.startIdx !== undefined && viewRange.endIdx !== undefined) { - // Navigation: use index-based range - options = { startIdx: viewRange.startIdx, endIdx: viewRange.endIdx }; - } - - const data = await fetchChartData(groupName, chartName, options); - if (!cancelled && data) { - setChartData(data); - if (data.originalLength) { - setTotalCommits(data.originalLength); - } else if (data.commits) { - setTotalCommits(data.commits.length); - } - } - } catch (err) { - if (!cancelled) { - setError(err.message); - } - } finally { - if (!cancelled) { - setLoading(false); - } - } - } - - loadData(); - - return () => { - cancelled = true; - }; - }, [groupName, chartName, viewRange]); - - // Compute display range info from chartData (for rendering only) - const displayRangeInfo = useMemo(() => { - if (!chartData) return { startIdx: 0, endIdx: 0, total: 0, rangeSize: 0 }; - const total = chartData.originalLength || chartData.commits?.length || 0; - const rangeSize = chartData.commits?.length || 0; - const req = chartData.requestedRange || {}; - const startIdx = req.startIndex ?? (total - rangeSize); - const endIdx = req.endIndex ?? (total - 1); - return { startIdx, endIdx, total, rangeSize }; - }, [chartData]); - - // Compute current range from viewRange state (for navigation calculations) - const getCurrentRange = useCallback(() => { - const total = totalCommits || 0; - if (viewRange.last) { - const rangeSize = Math.min(viewRange.last, total); - return { - startIdx: Math.max(0, total - rangeSize), - endIdx: total - 1, - total, - rangeSize, - }; - } - const startIdx = viewRange.startIdx ?? 0; - const endIdx = viewRange.endIdx ?? (total - 1); - return { - startIdx, - endIdx, - total, - rangeSize: endIdx - startIdx + 1, - }; - }, [viewRange, totalCommits]); - - const isAtStart = displayRangeInfo.startIdx === 0; - const isAtEnd = displayRangeInfo.endIdx >= displayRangeInfo.total - 1; - const currentRangeSize = displayRangeInfo.rangeSize; - - // Navigation handlers - use functional updates to get latest state - const handleGoToStart = useCallback(() => { - const range = getCurrentRange(); - if (range.startIdx === 0 || !range.total) return; - setViewRange({ - startIdx: 0, - endIdx: Math.min(range.rangeSize - 1, range.total - 1), - }); - }, [getCurrentRange]); - - const handleGoToEnd = useCallback(() => { - const range = getCurrentRange(); - if (range.endIdx >= range.total - 1 || !range.total) return; - setViewRange({ last: range.rangeSize }); - }, [getCurrentRange]); - - const handleMoveBackward = useCallback(() => { - const range = getCurrentRange(); - if (range.startIdx === 0 || !range.total) return; - const moveAmount = Math.max(1, Math.floor(range.rangeSize / 2)); - const newStartIdx = Math.max(0, range.startIdx - moveAmount); - const newEndIdx = newStartIdx + range.rangeSize - 1; - setViewRange({ - startIdx: newStartIdx, - endIdx: Math.min(newEndIdx, range.total - 1), - }); - }, [getCurrentRange]); - - const handleMoveForward = useCallback(() => { - const range = getCurrentRange(); - if (range.endIdx >= range.total - 1 || !range.total) return; - const moveAmount = Math.max(1, Math.floor(range.rangeSize / 2)); - const newEndIdx = Math.min(range.total - 1, range.endIdx + moveAmount); - const newStartIdx = Math.max(0, newEndIdx - range.rangeSize + 1); - setViewRange({ startIdx: newStartIdx, endIdx: newEndIdx }); - }, [getCurrentRange]); - - const handleZoomIn = useCallback(() => { - const range = getCurrentRange(); - if (!range.total || range.rangeSize <= 10) return; - const center = Math.floor((range.startIdx + range.endIdx) / 2); - const newRangeSize = Math.max(10, Math.floor(range.rangeSize / 2)); - const halfRange = Math.floor(newRangeSize / 2); - let newStartIdx = center - halfRange; - let newEndIdx = newStartIdx + newRangeSize - 1; - - // Clamp to bounds - if (newStartIdx < 0) { - newStartIdx = 0; - newEndIdx = newRangeSize - 1; - } - if (newEndIdx >= range.total) { - newEndIdx = range.total - 1; - newStartIdx = Math.max(0, newEndIdx - newRangeSize + 1); - } - - setViewRange({ startIdx: newStartIdx, endIdx: newEndIdx }); - }, [getCurrentRange]); - - const handleZoomOut = useCallback(() => { - const range = getCurrentRange(); - if (!range.total) return; - const center = Math.floor((range.startIdx + range.endIdx) / 2); - const newRangeSize = Math.min(range.total, range.rangeSize * 2); - const halfRange = Math.floor(newRangeSize / 2); - let newStartIdx = center - halfRange; - let newEndIdx = newStartIdx + newRangeSize - 1; - - // Clamp to bounds - if (newStartIdx < 0) { - newStartIdx = 0; - newEndIdx = Math.min(newRangeSize - 1, range.total - 1); - } - if (newEndIdx >= range.total) { - newEndIdx = range.total - 1; - newStartIdx = Math.max(0, newEndIdx - newRangeSize + 1); - } - - setViewRange({ startIdx: newStartIdx, endIdx: newEndIdx }); - }, [getCurrentRange]); - - const handleShowFullRange = useCallback(() => { - const range = getCurrentRange(); - if (!range.total) return; - setViewRange({ startIdx: 0, endIdx: range.total - 1 }); - }, [getCurrentRange]); - - const isFullRange = isAtStart && isAtEnd; - - // Handle drag selection zoom - const handleDragZoom = useCallback((startDataIdx, endDataIdx) => { - if (!chartData?.commits || !chartData.requestedRange) return; - - const numCommits = chartData.commits.length; - if (numCommits < 2) return; - - const rangeStart = chartData.requestedRange.startIndex; - const rangeEnd = chartData.requestedRange.endIndex; - const total = chartData.originalLength || rangeEnd + 1; - - // Map chart indices to original dataset indices using linear interpolation - // This correctly handles downsampled data where numCommits < (rangeEnd - rangeStart + 1) - const minIdx = Math.min(startDataIdx, endDataIdx); - const maxIdx = Math.max(startDataIdx, endDataIdx); - const globalStartIdx = rangeStart + Math.round(minIdx / (numCommits - 1) * (rangeEnd - rangeStart)); - const globalEndIdx = rangeStart + Math.round(maxIdx / (numCommits - 1) * (rangeEnd - rangeStart)); - - // Ensure minimum range - if (globalEndIdx - globalStartIdx < 5) return; - - setViewRange({ - startIdx: Math.max(0, globalStartIdx), - endIdx: Math.min(total - 1, globalEndIdx), - }); - }, [chartData]); - - // Process series data with filters and renaming - const processedData = useMemo(() => { - if (!chartData?.series || !chartData?.commits) return null; - - const { series, commits } = chartData; - const datasets = []; - const labels = commits.map(c => formatDate(c.timestamp)); - - Object.entries(series).forEach(([seriesName, points]) => { - // Apply removed datasets filter - if (config.removedDatasets?.has(seriesName)) return; - - // Apply engine filter - if (engineFilter !== 'all') { - const engine = seriesName.split(':')[0].toLowerCase(); - if (engine !== engineFilter && !seriesName.toLowerCase().includes(engineFilter)) { - return; - } - } - - // Rename series if needed - let displaySeriesName = seriesName; - if (config.renamedDatasets) { - const caseInsensitive = {}; - Object.entries(config.renamedDatasets).forEach(([k, v]) => { - caseInsensitive[k.toLowerCase()] = v; - }); - displaySeriesName = caseInsensitive[seriesName.toLowerCase()] || seriesName; - } - - // Check if hidden by default - const hidden = config.hiddenDatasets?.has(seriesName) || - config.hiddenDatasets?.has(displaySeriesName); - - datasets.push({ - label: displaySeriesName, - data: points, - borderColor: stringToColor(displaySeriesName), - backgroundColor: stringToColor(displaySeriesName) + '20', - pointRadius: 2, - pointHoverRadius: 5, - pointStyle: 'cross', - borderWidth: 1.5, - tension: 0, - spanGaps: true, - hidden, - }); - }); - - return { labels, datasets, commits }; - }, [chartData, config, engineFilter]); - - // Handle click on chart point to open commit on GitHub - const handleChartClick = useCallback((event, elements) => { - if (!elements.length || !processedData?.commits) return; - const dataIndex = elements[0].index; - const commit = processedData.commits[dataIndex]; - if (commit?.id) { - window.open(`https://github.com/vortex-data/vortex/commit/${commit.id}`, '_blank'); - } - }, [processedData]); - - // Chart.js options with drag zoom - const options = useMemo(() => ({ - responsive: true, - maintainAspectRatio: false, - animation: false, - onClick: handleChartClick, - onHover: (event, elements) => { - event.native.target.style.cursor = elements.length ? 'pointer' : 'default'; - }, - interaction: { - mode: 'index', - intersect: true, - }, - plugins: { - legend: { - position: 'top', - align: 'start', - labels: { - boxWidth: 12, - padding: 8, - font: { - size: 12, - family: 'Geist, sans-serif', - }, - usePointStyle: true, - pointStyle: 'rectRounded', - }, - }, - tooltip: { - backgroundColor: 'rgba(16, 16, 16, 0.9)', - titleFont: { family: 'Geist, sans-serif', size: 13 }, - bodyFont: { family: 'Geist Mono, monospace', size: 12 }, - padding: 12, - cornerRadius: 4, - position: 'topCorner', - caretSize: 0, - itemSort: (a, b) => b.parsed.y - a.parsed.y, - // Limit to top 10 items by value to prevent tooltip from getting too large - filter: (item, _index, items) => { - if (items.length <= 10) return item.parsed.y != null; - const validItems = items.filter(i => i.parsed.y != null); - if (validItems.length <= 10) return item.parsed.y != null; - const sorted = [...validItems].sort((a, b) => (b.parsed.y ?? 0) - (a.parsed.y ?? 0)); - const top10 = sorted.slice(0, 10); - return top10.some(i => i.datasetIndex === item.datasetIndex); - }, - callbacks: { - title: (items) => { - if (!items.length || !processedData?.commits) return ''; - const commit = processedData.commits[items[0].dataIndex]; - if (!commit) return items[0].label; - const author = commit.author || 'Unknown'; - return `${formatDate(commit.timestamp)} — ${author}\n(${commit.id?.slice(0, 7) || ''}) ${commit.message || ''}`; - }, - label: (item) => { - const value = item.parsed.y; - if (value == null) return null; - const formattedValue = value < 1 ? value.toFixed(4) : value.toFixed(2); - return `${item.dataset.label}: ${formattedValue} ${unit || ''}`; - }, - labelTextColor: (tooltipItem) => { - const chart = tooltipItem.chart; - const activeElements = chart._active || []; - if (activeElements.length === 0) return '#ffffff'; - const lastEvent = chart._lastEvent; - if (!lastEvent) return '#ffffff'; - // Find which dataset point is closest to the cursor - let hoveredDatasetIndex = activeElements[0].datasetIndex; - let closestDist = Infinity; - for (const el of activeElements) { - const point = chart.getDatasetMeta(el.datasetIndex).data[el.index]; - if (point) { - const dist = Math.hypot(point.x - lastEvent.x, point.y - lastEvent.y); - if (dist < closestDist) { - closestDist = dist; - hoveredDatasetIndex = el.datasetIndex; - } - } - } - if (tooltipItem.datasetIndex === hoveredDatasetIndex) { - return '#ffffff'; - } - return 'rgba(255, 255, 255, 0.45)'; - }, - }, - }, - zoom: { - zoom: { - drag: { - enabled: true, - backgroundColor: 'rgba(99, 102, 241, 0.2)', - borderColor: 'rgba(99, 102, 241, 0.8)', - borderWidth: 1, - }, - mode: 'x', - onZoomComplete: ({ chart }) => { - // Prevent infinite loop from resetZoom triggering onZoomComplete - if (isResettingZoom.current) { - isResettingZoom.current = false; - return; - } - - const { min, max } = chart.scales.x; - const startIdx = Math.floor(min); - const endIdx = Math.ceil(max); - if (startIdx >= 0 && endIdx > startIdx) { - handleDragZoom(startIdx, endIdx); - } - // Reset chart zoom state - isResettingZoom.current = true; - chart.resetZoom(); - }, - }, - }, - }, - scales: { - x: { - display: true, - grid: { - display: true, - color: 'rgba(166, 166, 166, 0.12)', - }, - ticks: { - maxRotation: 45, - minRotation: 45, - font: { - size: 11, - family: 'Geist, sans-serif', - }, - maxTicksLimit: 10, - callback: function(value, index, ticks) { - // Always show first and last tick - if (index === 0 || index === ticks.length - 1) { - return this.getLabelForValue(value); - } - // Show intermediate ticks based on maxTicksLimit - const step = Math.ceil(ticks.length / 10); - if (index % step === 0) { - return this.getLabelForValue(value); - } - return null; - }, - }, - }, - y: { - display: true, - beginAtZero: true, - grid: { - color: 'rgba(166, 166, 166, 0.12)', - }, - ticks: { - font: { - size: 12, - family: 'Geist Mono, monospace', - }, - }, - title: { - display: !!unit, - text: unit || '', - font: { - size: 12, - family: 'Geist, sans-serif', - }, - }, - }, - }, - }), [unit, processedData, handleDragZoom, handleChartClick]); - - // Fullscreen handler - const handleFullscreen = useCallback(() => { - if (processedData) { - onFullscreen({ - title: displayName, - groupName, - chartName, - unit, - config, - initialData: processedData, - totalCommits, - currentRange: getCurrentRange(), - }); - } - }, [processedData, displayName, groupName, chartName, unit, config, totalCommits, getCurrentRange, onFullscreen]); - - // Show placeholder only on initial load (no data yet) - const showPlaceholder = !processedData && (loading || error); - const showOverlay = loading && processedData; - - if (showPlaceholder) { - return ( -
-
- {displayName} -
-
- {error ? ( -

Error loading chart

- ) : ( -
- )} -
-
- ); - } - - if (!loading && error) { - return ( -
-
- {displayName} -
-
-

Error loading chart

-
-
- ); - } - - if (!processedData || processedData.datasets.length === 0) { - return ( -
-
- {displayName} -
-
-

No data available

-
-
- ); - } - - return ( -
-
- - {displayName} - {chartData?.downsampleLevel && chartData.downsampleLevel !== '1x' && ( - - {chartData.downsampleLevel} downsampled - - )} - -
-
- - - - - - - -
- -
-
-
- - {showOverlay && ( -
-
-
- )} -
-
- ); -} diff --git a/benchmarks-website/src/components/Header.jsx b/benchmarks-website/src/components/Header.jsx deleted file mode 100644 index b9624794119..00000000000 --- a/benchmarks-website/src/components/Header.jsx +++ /dev/null @@ -1,110 +0,0 @@ -import React from 'react'; - -export default function Header({ - sidebarOpen, - onMenuToggle, - categoryFilter, - onCategoryChange, - searchFilter, - onSearchChange, - viewMode, - onViewModeChange, - onExpandAll, - onCollapseAll, -}) { - return ( -
-
-
- - - - - - Vortex - - -
- -

Vortex Benchmarks

- -
-
- - - - onSearchChange(e.target.value)} - /> -
-
- -
-
- - -
- - - - - GitHub - -
-
-
- ); -} diff --git a/benchmarks-website/src/components/Modal.jsx b/benchmarks-website/src/components/Modal.jsx deleted file mode 100644 index b9bb7d419b1..00000000000 --- a/benchmarks-website/src/components/Modal.jsx +++ /dev/null @@ -1,476 +0,0 @@ -import React, { useState, useEffect, useCallback, useMemo, useRef } from 'react'; -import { Line } from 'react-chartjs-2'; -import { - SkipBack, - ChevronLeft, - ChevronRight, - SkipForward, - ZoomIn, - ZoomOut, - MoveHorizontal, - X, -} from 'lucide-react'; -import { fetchChartData } from '../api'; -import { stringToColor, formatDate } from '../utils'; - -export default function Modal({ chartData, onClose }) { - const [loading, setLoading] = useState(false); - const [viewRange, setViewRange] = useState(null); - const [currentData, setCurrentData] = useState(null); - const [totalCommits, setTotalCommits] = useState(chartData?.totalCommits || 0); - const chartRef = useRef(null); - const isResettingZoom = useRef(false); - - // Initialize with the data passed from parent - useEffect(() => { - if (chartData?.initialData) { - setCurrentData(chartData.initialData); - setTotalCommits(chartData.totalCommits || chartData.initialData.commits?.length || 0); - if (chartData.currentRange) { - setViewRange({ - startIdx: chartData.currentRange.startIdx, - endIdx: chartData.currentRange.endIdx, - }); - } - } - }, [chartData]); - - // Close on escape key - useEffect(() => { - const handleKeyDown = (e) => { - if (e.key === 'Escape') { - onClose(); - } - }; - document.addEventListener('keydown', handleKeyDown); - return () => document.removeEventListener('keydown', handleKeyDown); - }, [onClose]); - - // Prevent body scroll when modal is open - useEffect(() => { - document.body.style.overflow = 'hidden'; - return () => { - document.body.style.overflow = ''; - }; - }, []); - - const handleBackdropClick = useCallback((e) => { - if (e.target === e.currentTarget) { - onClose(); - } - }, [onClose]); - - // Fetch data for a new range - const fetchRange = useCallback(async (range) => { - if (!chartData?.groupName || !chartData?.chartName) return; - - setLoading(true); - try { - let options = {}; - if (range.last) { - options = { last: range.last }; - } else if (range.startIdx !== undefined && range.endIdx !== undefined) { - options = { startIdx: range.startIdx, endIdx: range.endIdx }; - } - - const data = await fetchChartData(chartData.groupName, chartData.chartName, options); - if (data) { - // Process the data similar to ChartContainer - const processedData = processChartData(data, chartData.config); - setCurrentData(processedData); - if (data.originalLength) { - setTotalCommits(data.originalLength); - } - } - } catch (err) { - console.error('Error fetching range:', err); - } finally { - setLoading(false); - } - }, [chartData]); - - // Process chart data (similar to ChartContainer) - const processChartData = useCallback((data, config) => { - if (!data?.series || !data?.commits) return null; - - const { series, commits } = data; - const datasets = []; - const labels = commits.map(c => formatDate(c.timestamp)); - - Object.entries(series).forEach(([seriesName, points]) => { - if (config?.removedDatasets?.has(seriesName)) return; - - let displaySeriesName = seriesName; - if (config?.renamedDatasets) { - const caseInsensitive = {}; - Object.entries(config.renamedDatasets).forEach(([k, v]) => { - caseInsensitive[k.toLowerCase()] = v; - }); - displaySeriesName = caseInsensitive[seriesName.toLowerCase()] || seriesName; - } - - const dataPoints = points.map(p => p); - const hidden = config?.hiddenDatasets?.has(seriesName) || - config?.hiddenDatasets?.has(displaySeriesName); - - datasets.push({ - label: displaySeriesName, - data: dataPoints, - borderColor: stringToColor(displaySeriesName), - backgroundColor: stringToColor(displaySeriesName) + '20', - pointRadius: 2, - pointHoverRadius: 5, - pointStyle: 'cross', - borderWidth: 1.5, - tension: 0, - spanGaps: true, - hidden, - }); - }); - - return { labels, datasets, commits }; - }, []); - - // Get current range info - const getCurrentRange = useCallback(() => { - if (!viewRange) { - return { startIdx: 0, endIdx: totalCommits - 1, total: totalCommits, rangeSize: totalCommits }; - } - const startIdx = viewRange.startIdx ?? 0; - const endIdx = viewRange.endIdx ?? (totalCommits - 1); - return { - startIdx, - endIdx, - total: totalCommits, - rangeSize: endIdx - startIdx + 1, - }; - }, [viewRange, totalCommits]); - - const range = getCurrentRange(); - const isAtStart = range.startIdx === 0; - const isAtEnd = range.endIdx >= range.total - 1; - const currentRangeSize = range.rangeSize; - const isFullRange = isAtStart && isAtEnd; - - // Navigation handlers - const handleGoToStart = useCallback(() => { - if (isAtStart || !range.total) return; - const newRange = { - startIdx: 0, - endIdx: Math.min(currentRangeSize - 1, range.total - 1), - }; - setViewRange(newRange); - fetchRange(newRange); - }, [isAtStart, currentRangeSize, range.total, fetchRange]); - - const handleGoToEnd = useCallback(() => { - if (isAtEnd || !range.total) return; - const newRange = { last: currentRangeSize }; - setViewRange({ - startIdx: range.total - currentRangeSize, - endIdx: range.total - 1, - }); - fetchRange(newRange); - }, [isAtEnd, currentRangeSize, range.total, fetchRange]); - - const handleMoveBackward = useCallback(() => { - if (isAtStart || !range.total) return; - const moveAmount = Math.max(1, Math.floor(currentRangeSize / 2)); - const newStartIdx = Math.max(0, range.startIdx - moveAmount); - const newEndIdx = newStartIdx + currentRangeSize - 1; - const newRange = { - startIdx: newStartIdx, - endIdx: Math.min(newEndIdx, range.total - 1), - }; - setViewRange(newRange); - fetchRange(newRange); - }, [isAtStart, range, currentRangeSize, fetchRange]); - - const handleMoveForward = useCallback(() => { - if (isAtEnd || !range.total) return; - const moveAmount = Math.max(1, Math.floor(currentRangeSize / 2)); - const newEndIdx = Math.min(range.total - 1, range.endIdx + moveAmount); - const newStartIdx = Math.max(0, newEndIdx - currentRangeSize + 1); - const newRange = { startIdx: newStartIdx, endIdx: newEndIdx }; - setViewRange(newRange); - fetchRange(newRange); - }, [isAtEnd, range, currentRangeSize, fetchRange]); - - const handleZoomIn = useCallback(() => { - if (!range.total || currentRangeSize <= 10) return; - const center = Math.floor((range.startIdx + range.endIdx) / 2); - const newRangeSize = Math.max(10, Math.floor(currentRangeSize / 2)); - const halfRange = Math.floor(newRangeSize / 2); - let newStartIdx = center - halfRange; - let newEndIdx = newStartIdx + newRangeSize - 1; - - if (newStartIdx < 0) { - newStartIdx = 0; - newEndIdx = newRangeSize - 1; - } - if (newEndIdx >= range.total) { - newEndIdx = range.total - 1; - newStartIdx = Math.max(0, newEndIdx - newRangeSize + 1); - } - - const newRange = { startIdx: newStartIdx, endIdx: newEndIdx }; - setViewRange(newRange); - fetchRange(newRange); - }, [range, currentRangeSize, fetchRange]); - - const handleZoomOut = useCallback(() => { - if (!range.total) return; - const center = Math.floor((range.startIdx + range.endIdx) / 2); - const newRangeSize = Math.min(range.total, currentRangeSize * 2); - const halfRange = Math.floor(newRangeSize / 2); - let newStartIdx = center - halfRange; - let newEndIdx = newStartIdx + newRangeSize - 1; - - if (newStartIdx < 0) { - newStartIdx = 0; - newEndIdx = Math.min(newRangeSize - 1, range.total - 1); - } - if (newEndIdx >= range.total) { - newEndIdx = range.total - 1; - newStartIdx = Math.max(0, newEndIdx - newRangeSize + 1); - } - - const newRange = { startIdx: newStartIdx, endIdx: newEndIdx }; - setViewRange(newRange); - fetchRange(newRange); - }, [range, currentRangeSize, fetchRange]); - - const handleShowFullRange = useCallback(() => { - if (!range.total) return; - const newRange = { startIdx: 0, endIdx: range.total - 1 }; - setViewRange(newRange); - fetchRange(newRange); - }, [range.total, fetchRange]); - - // Handle drag selection zoom - const handleDragZoom = useCallback((startDataIdx, endDataIdx) => { - if (!currentData?.commits) return; - - const numCommits = currentData.commits.length; - if (numCommits < 2) return; - - const rangeStart = range.startIdx; - const rangeEnd = range.endIdx; - const total = range.total; - - const minIdx = Math.min(startDataIdx, endDataIdx); - const maxIdx = Math.max(startDataIdx, endDataIdx); - const globalStartIdx = rangeStart + Math.round(minIdx / (numCommits - 1) * (rangeEnd - rangeStart)); - const globalEndIdx = rangeStart + Math.round(maxIdx / (numCommits - 1) * (rangeEnd - rangeStart)); - - if (globalEndIdx - globalStartIdx < 5) return; - - const newRange = { - startIdx: Math.max(0, globalStartIdx), - endIdx: Math.min(total - 1, globalEndIdx), - }; - setViewRange(newRange); - fetchRange(newRange); - }, [currentData, range, fetchRange]); - - // Chart options - const options = useMemo(() => ({ - responsive: true, - maintainAspectRatio: false, - animation: false, - interaction: { - mode: 'index', - intersect: false, - }, - plugins: { - legend: { - position: 'top', - align: 'start', - labels: { - boxWidth: 12, - padding: 8, - font: { size: 11, family: 'Geist, sans-serif' }, - usePointStyle: true, - pointStyle: 'rectRounded', - }, - }, - tooltip: { - backgroundColor: 'rgba(16, 16, 16, 0.9)', - titleFont: { family: 'Geist, sans-serif', size: 12 }, - bodyFont: { family: 'Geist Mono, monospace', size: 11 }, - padding: 12, - cornerRadius: 4, - caretSize: 0, - position: 'topCorner', - itemSort: (a, b) => b.parsed.y - a.parsed.y, - callbacks: { - title: (items) => { - if (!items.length || !currentData?.commits) return ''; - const commit = currentData.commits[items[0].dataIndex]; - if (!commit) return items[0].label; - const author = commit.author || 'Unknown'; - return `${formatDate(commit.timestamp)} — ${author}\n(${commit.id?.slice(0, 7) || ''}) ${commit.message || ''}`; - }, - label: (item) => { - const value = item.parsed.y; - if (value == null) return null; - const formattedValue = value < 1 ? value.toFixed(4) : value.toFixed(2); - return `${item.dataset.label}: ${formattedValue} ${chartData?.unit || ''}`; - }, - }, - }, - zoom: { - zoom: { - drag: { - enabled: true, - backgroundColor: 'rgba(99, 102, 241, 0.2)', - borderColor: 'rgba(99, 102, 241, 0.8)', - borderWidth: 1, - }, - mode: 'x', - onZoomComplete: ({ chart }) => { - if (isResettingZoom.current) { - isResettingZoom.current = false; - return; - } - const { min, max } = chart.scales.x; - const startIdx = Math.floor(min); - const endIdx = Math.ceil(max); - if (startIdx >= 0 && endIdx > startIdx) { - handleDragZoom(startIdx, endIdx); - } - isResettingZoom.current = true; - chart.resetZoom(); - }, - }, - }, - }, - scales: { - x: { - display: true, - grid: { display: true, color: 'rgba(0, 0, 0, 0.12)' }, - ticks: { - maxRotation: 45, - minRotation: 45, - font: { size: 10, family: 'Geist, sans-serif' }, - maxTicksLimit: 15, - callback: function(value, index, ticks) { - // Always show first and last tick - if (index === 0 || index === ticks.length - 1) { - return this.getLabelForValue(value); - } - // Show intermediate ticks based on maxTicksLimit - const step = Math.ceil(ticks.length / 15); - if (index % step === 0) { - return this.getLabelForValue(value); - } - return null; - }, - }, - }, - y: { - display: true, - beginAtZero: true, - grid: { color: 'rgba(0, 0, 0, 0.12)' }, - ticks: { font: { size: 11, family: 'Geist Mono, monospace' } }, - title: { - display: !!chartData?.unit, - text: chartData?.unit || '', - font: { size: 11, family: 'Geist, sans-serif' }, - }, - }, - }, - }), [currentData, chartData, handleDragZoom]); - - if (!chartData) return null; - - return ( -
-
-
-

{chartData.title}

-
-
- - - - - - - -
- -
-
-
- {currentData && ( - - )} - {loading && ( -
-
-
- )} -
-
-
- ); -} diff --git a/benchmarks-website/src/components/Sidebar.jsx b/benchmarks-website/src/components/Sidebar.jsx deleted file mode 100644 index 3a53b801f8e..00000000000 --- a/benchmarks-website/src/components/Sidebar.jsx +++ /dev/null @@ -1,59 +0,0 @@ -import React from 'react'; - -export default function Sidebar({ - isOpen, - groups, - onClose, - onGroupClick, - onClearFilter, - showClearFilter, -}) { - return ( - - ); -} diff --git a/benchmarks-website/src/config.js b/benchmarks-website/src/config.js deleted file mode 100644 index dfe5555019d..00000000000 --- a/benchmarks-website/src/config.js +++ /dev/null @@ -1,285 +0,0 @@ -// ============================================================================= -// SQL query benchmark suites — single source of truth. -// To add a new SQL query benchmark, add one entry to QUERY_SUITES. -// The server (routing/formatting) and frontend (UI config) both derive from this. -// ============================================================================= - -export const QUERY_SUITES = [ - { - prefix: "clickbench", - displayName: "Clickbench", - queryPrefix: "CLICKBENCH", - description: - "ClickHouse's analytical benchmark suite testing real-world query patterns on web analytics data", - 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", - queryPrefix: "STATPOPGEN", - description: - "A suite of Statistical and Population genetics queries using the gnomAD dataset", - tags: ["Queries (NVMe)", "StatPopGen"], - }, - { - prefix: "polarsignals", - displayName: "PolarSignals Profiling", - queryPrefix: "POLARSIGNALS", - description: - "Profiling data benchmark modeled on PolarSignals/Parca, exercising scan-layer performance with projection and filter pushdown on deeply nested schemas", - tags: ["Queries (NVMe)", "PolarSignals"], - }, - { - prefix: "tpch", - displayName: "TPC-H", - queryPrefix: "TPC-H", - datasetKey: "tpch", - fanOut: true, - hiddenDatasets: ["datafusion:lance"], - }, - { - prefix: "tpcds", - displayName: "TPC-DS", - queryPrefix: "TPC-DS", - datasetKey: "tpcds", - fanOut: true, - }, - { prefix: "fineweb", skip: true }, -]; - -// Pre-registered fan-out groups (storage x scale factor). -export const FAN_OUT_GROUPS = [ - "TPC-H (NVMe) (SF=1)", - "TPC-H (S3) (SF=1)", - "TPC-H (NVMe) (SF=10)", - "TPC-H (S3) (SF=10)", - "TPC-H (NVMe) (SF=100)", - "TPC-H (S3) (SF=100)", - "TPC-H (NVMe) (SF=1000)", - "TPC-H (S3) (SF=1000)", - "TPC-DS (NVMe) (SF=1)", - "TPC-DS (NVMe) (SF=10)", -]; - -// Canonical engine:format renaming used by all query suites. -export const ENGINE_RENAMES = { - "datafusion:vortex-file-compressed": "datafusion:vortex", - "datafusion:parquet": "datafusion:parquet", - "datafusion:arrow": "datafusion:in-memory-arrow", - "datafusion:lance": "datafusion:lance", - "datafusion:vortex-compact": "datafusion:vortex-compact", - "duckdb:vortex-file-compressed": "duckdb:vortex", - "duckdb:parquet": "duckdb:parquet", - "duckdb:duckdb": "duckdb:duckdb", - "duckdb:vortex-compact": "duckdb:vortex-compact", - "vortex-tokio-local-disk": "vortex-nvme", - "vortex-compact-tokio-local-disk": "vortex-compact-nvme", - "lance-tokio-local-disk": "lance-nvme", - "parquet-tokio-local-disk": "parquet-nvme", - lance: "lance", -}; - -// ============================================================================= -// Below: frontend UI config, derived from QUERY_SUITES where possible. -// ============================================================================= - -// Build BENCHMARK_CONFIGS: bespoke non-query groups + generated query group entries. -const BESPOKE_CONFIGS = [ - { - name: "Random Access", - renamedDatasets: { - "vortex-tokio-local-disk": "vortex-nvme", - "vortex-compact-tokio-local-disk": "vortex-compact-nvme", - "lance-tokio-local-disk": "lance-nvme", - "parquet-tokio-local-disk": "parquet-nvme", - }, - }, - { - name: "Compression", - keptCharts: [ - "COMPRESS TIME", - "DECOMPRESS TIME", - "PARQUET RS ZSTD COMPRESS TIME", - "PARQUET RS ZSTD DECOMPRESS TIME", - "LANCE COMPRESS TIME", - "LANCE DECOMPRESS TIME", - "VORTEX:PARQUET ZSTD RATIO COMPRESS TIME", - "VORTEX:PARQUET ZSTD RATIO DECOMPRESS TIME", - "VORTEX:LANCE RATIO COMPRESS TIME", - "VORTEX:LANCE RATIO DECOMPRESS TIME", - ], - hiddenDatasets: new Set([ - "wide table cols=1000 chunks=1 rows=1000", - "wide table cols=1000 chunks=50 rows=1000", - ]), - removedDatasets: new Set([ - "TPC-H l_comment canonical", - "TPC-H l_comment chunked without fsst", - "wide table cols=10 chunks=1 rows=1000", - "wide table cols=100 chunks=1 rows=1000", - "wide table cols=10 chunks=50 rows=1000", - "wide table cols=100 chunks=50 rows=1000", - ]), - renamedDatasets: { lance: "lance", Lance: "lance", LANCE: "lance" }, - }, - { - name: "Compression Size", - keptCharts: [ - "VORTEX SIZE", - "PARQUET SIZE", - "LANCE SIZE", - "VORTEX:PARQUET ZSTD SIZE", - "VORTEX:LANCE SIZE", - ], - hiddenDatasets: new Set(["wide table cols=1000"]), - removedDatasets: new Set([ - "wide table cols=10 chunks=1 rows=1000", - "wide table cols=100 chunks=1 rows=1000", - "wide table cols=10 chunks=50 rows=1000", - "wide table cols=100 chunks=50 rows=1000", - ]), - renamedDatasets: { lance: "lance", Lance: "lance", LANCE: "lance" }, - }, -]; - -function querySuiteConfig(name, suite) { - const cfg = { name, renamedDatasets: { ...ENGINE_RENAMES } }; - if (suite?.hiddenDatasets?.length) - cfg.hiddenDatasets = new Set(suite.hiddenDatasets); - return cfg; -} - -function buildQueryConfigs() { - const configs = []; - for (const s of QUERY_SUITES) { - if (s.skip) continue; - if (!s.fanOut) { - configs.push(querySuiteConfig(s.displayName, s)); - } - } - for (const g of FAN_OUT_GROUPS) { - const suite = QUERY_SUITES.find( - (s) => s.fanOut && g.startsWith(s.displayName), - ); - const cfg = querySuiteConfig(g, suite); - if (g.includes("SF=1000") || (g.includes("TPC-DS") && g.includes("SF=10)"))) - cfg.hidden = true; - configs.push(cfg); - } - return configs; -} - -export const BENCHMARK_CONFIGS = [...BESPOKE_CONFIGS, ...buildQueryConfigs()]; - -// Chart name remapping (compression benchmarks only) -export const CHART_NAME_MAP = { - "COMPRESS TIME": "VORTEX WRITE TIME (COMPRESSION)", - "DECOMPRESS TIME": "VORTEX SCAN TIME (DECOMPRESSION)", - "PARQUET RS ZSTD COMPRESS TIME": "PARQUET WRITE TIME (COMPRESSION)", - "PARQUET RS ZSTD DECOMPRESS TIME": "PARQUET SCAN TIME (DECOMPRESSION)", - "LANCE COMPRESS TIME": "LANCE WRITE TIME (COMPRESSION)", - "LANCE DECOMPRESS TIME": "LANCE SCAN TIME (DECOMPRESSION)", - "VORTEX SIZE": "VORTEX SIZE", - "PARQUET ZSTD SIZE": "PARQUET SIZE", - "LANCE SIZE": "LANCE SIZE", - "VORTEX:RAW SIZE": "VORTEX vs RAW SIZE RATIO", - "VORTEX:PARQUET ZSTD SIZE": "VORTEX vs PARQUET SIZE RATIO", - "VORTEX:LANCE SIZE": "VORTEX vs LANCE SIZE RATIO", - "VORTEX:PARQUET ZSTD RATIO COMPRESS TIME": - "VORTEX vs PARQUET WRITE TIME RATIO", - "VORTEX:PARQUET ZSTD RATIO DECOMPRESS TIME": - "VORTEX vs PARQUET SCAN TIME RATIO", - "VORTEX:LANCE RATIO COMPRESS TIME": "VORTEX vs LANCE WRITE TIME RATIO", - "VORTEX:LANCE RATIO DECOMPRESS TIME": "VORTEX vs LANCE SCAN TIME RATIO", -}; - -// Category tags for sidebar filtering -export const CATEGORY_TAGS = { - "Random Access": ["Read/Write"], - Compression: ["Read/Write"], - "Compression Size": ["Read/Write"], -}; -for (const s of QUERY_SUITES) { - if (!s.skip && !s.fanOut && s.tags) CATEGORY_TAGS[s.displayName] = s.tags; -} -for (const g of FAN_OUT_GROUPS) { - const m = g.match(/^(.+?) \((NVMe|S3)\) \((SF=\d+)\)$/); - CATEGORY_TAGS[g] = [ - m[2] === "S3" ? "Queries (S3)" : "Queries (NVMe)", - `${m[1]} (${m[3]})`, - ]; -} - -// Benchmark descriptions -export const BENCHMARK_DESCRIPTIONS = { - "Random Access": - "Tests performance of selecting arbitrary row indices from a file on NVMe storage", - Compression: - "Measures encoding and decoding throughput (MB/s) for Vortex files and Parquet files (with zstd page compression)", - "Compression Size": - "Compares compressed file sizes and compression ratios across different encoding strategies", -}; -for (const s of QUERY_SUITES) { - if (s.description) BENCHMARK_DESCRIPTIONS[s.displayName] = s.description; -} - -// Scale factor descriptions -export const SCALE_FACTOR_DESCRIPTIONS = { - 1: "SF=1 (~1GB of data)", - 10: "SF=10 (~10GB of data)", - 100: "SF=100 (~100GB of data)", - 1000: "SF=1000 (~1TB of data)", -}; - -// Engine filter labels -export const ENGINE_LABELS = { - all: "All", - duckdb: "DuckDB", - datafusion: "DataFusion", - vortex: "Vortex", - parquet: "Parquet", -}; - -// Series color map -export const SERIES_COLOR_MAP = { - "vortex-nvme": "#19a508", - "vortex-compact-nvme": "#15850a", - "parquet-nvme": "#ef7f1d", - "lance-nvme": "#3B82F6", - "datafusion:arrow": "#7a27b1", - "datafusion:in-memory-arrow": "#7a27b1", - "datafusion:parquet": "#ef7f1d", - "datafusion:vortex": "#19a508", - "datafusion:vortex-compact": "#15850a", - "datafusion:lance": "#2D936C", - "duckdb:parquet": "#985113", - "duckdb:vortex": "#0e5e04", - "duckdb:vortex-compact": "#0b4a03", - "duckdb:duckdb": "#87752e", - "vortex:lance": "#FF8787", -}; - -// Fallback color palette -export const FALLBACK_PALETTE = [ - "#5971FD", - "#CEE562", - "#EEB3E1", - "#FF8C42", - "#B8336A", - "#726DA8", - "#2D936C", - "#E9B44C", -]; - -// Default visible commits -export const DEFAULT_COMMIT_RANGE = 100; diff --git a/benchmarks-website/src/main.jsx b/benchmarks-website/src/main.jsx deleted file mode 100644 index 32fcc3f979f..00000000000 --- a/benchmarks-website/src/main.jsx +++ /dev/null @@ -1,10 +0,0 @@ -import React from 'react'; -import ReactDOM from 'react-dom/client'; -import App from './App'; -import './styles/index.css'; - -ReactDOM.createRoot(document.getElementById('root')).render( - - - -); diff --git a/benchmarks-website/src/styles/index.css b/benchmarks-website/src/styles/index.css deleted file mode 100644 index 910670ec28f..00000000000 --- a/benchmarks-website/src/styles/index.css +++ /dev/null @@ -1,1319 +0,0 @@ -/* CSS Variables for consistent theming */ -:root { - /* Vortex Brand Colors */ - --vortex-black: #101010; - --vortex-gray: #ECECEC; - --vortex-green: #CEE562; - --vortex-blue: #5971FD; - --vortex-pink: #EEB3E1; - - /* Theme Colors */ - --primary-color: var(--vortex-blue); - --primary-hover: #4A5FE5; - --accent-color: var(--vortex-green); - --bg-color: #ffffff; - --bg-secondary: #FAFAFA; - --text-color: var(--vortex-black); - --text-secondary: #666666; - --border-color: var(--vortex-gray); - - /* Layout */ - --header-height: 72px; - --sidebar-width: 280px; - --chart-spacing: 24px; - --mobile-breakpoint: 768px; - --tablet-breakpoint: 1024px; - - /* Shadows */ - --shadow-sm: 0 1px 3px rgba(16,16,16,0.08); - --shadow-md: 0 4px 8px rgba(16,16,16,0.08); - --shadow-lg: 0 12px 24px rgba(16,16,16,0.12); - - /* Border Radius */ - --radius-sm: 4px; - --radius-md: 8px; - --radius-lg: 12px; -} - -@media (prefers-color-scheme: dark) { :root { - --primary-color: var(--vortex-blue); - --primary-hover: #8a9cff; - --accent-color: var(--vortex-green); - --bg-color: #050507; - --bg-secondary: #121219; - --text-color: #F5F5F7; - --text-secondary: #A0A0B0; - --border-color: #30303a; - --shadow-sm: 0 1px 3px rgba(0, 0, 0, 0.6); - --shadow-md: 0 4px 8px rgba(0, 0, 0, 0.7); - --shadow-lg: 0 12px 24px rgba(0, 0, 0, 0.8); -}} - -/* Reset and base styles */ -* { - box-sizing: border-box; -} - -html { - font-family: "Geist", -apple-system, BlinkMacSystemFont, "Segoe UI", "SF Pro Display", Roboto, sans-serif; - -webkit-font-smoothing: antialiased; - -moz-osx-font-smoothing: grayscale; - background-color: var(--bg-color); - font-size: 16px; - scroll-behavior: smooth; - overflow-x: hidden; -} - -body { - color: var(--text-color); - margin: 0; - padding: 0; - font-size: 1em; - font-weight: 400; - padding-top: var(--header-height); - line-height: 1.6; - letter-spacing: -0.01em; - overflow-x: hidden; - position: relative; -} - -/* Typography */ -h1, h2, h3, h4, h5, h6 { - font-family: "Funnel Display", sans-serif; - font-weight: 600; - letter-spacing: -0.02em; -} - -code, pre { - font-family: "Geist Mono", monospace; - font-size: 0.9em; -} - -/* Sticky Header */ -.sticky-header { - position: fixed; - top: 0; - left: 0; - right: 0; - height: var(--header-height); - background: rgba(255, 255, 255, 0.95); - -webkit-backdrop-filter: blur(10px); - backdrop-filter: blur(10px); - border-bottom: 1px solid var(--border-color); - box-shadow: var(--shadow-sm); - z-index: 1000; - display: flex; - align-items: center; -} - -.header-content { - display: flex; - align-items: center; - justify-content: space-between; - height: 100%; - padding: 0 32px; - width: 100%; - gap: 24px; - background: var(--bg-secondary); -} - -.header-left { - display: flex; - align-items: center; - gap: 12px; -} - -.menu-toggle { - display: block; - background: none; - border: none; - font-size: 20px; - cursor: pointer; - padding: 8px; - border-radius: var(--radius-sm); - transition: all 0.2s; - color: var(--text-color); -} - -.menu-toggle:hover { - background-color: var(--bg-secondary); -} - -.logo-link { - display: flex; - align-items: center; - text-decoration: none; - transition: opacity 0.2s; -} - -.logo-link:hover { - opacity: 0.8; -} - -.site-logo { - height: 24px; - width: auto; - display: block; -} - -.site-title { - font-family: "Funnel Display", sans-serif; - font-size: 1.5rem; - font-weight: 600; - margin: 0; - margin-left: calc(var(--sidebar-width) - 156px); - color: var(--text-color); - display: none; - white-space: nowrap; -} - -@media (min-width: 1400px) { - .site-title { - display: block; - } -} - -.header-center { - flex: 1; - display: flex; - justify-content: center; - padding: 0 20px; -} - -.filter-controls { - display: flex; - align-items: center; - gap: 16px; - max-width: 600px; -} - -.control-btn, .view-btn { - font-family: "Geist", sans-serif; - padding: 8px 20px; - border: 1px solid var(--border-color); - background: var(--bg-color); - color: var(--text-color); - border-radius: var(--radius-md); - cursor: pointer; - font-size: 14px; - font-weight: 500; - transition: all 0.2s; - white-space: nowrap; -} - -.control-btn:hover, .view-btn:hover { - background-color: var(--bg-secondary); - border-color: var(--primary-color); - transform: translateY(-1px); - box-shadow: var(--shadow-sm); -} - -.view-btn.active { - background-color: var(--primary-color); - color: white; - border-color: var(--primary-color); -} - -.category-filter, .search-filter { - font-family: "Geist", sans-serif; - padding: 8px 12px; - border: 1px solid var(--border-color); - border-radius: var(--radius-md); - font-size: 14px; - font-weight: 500; - background: var(--bg-color); - color: var(--text-color); - transition: border-color 0.2s; -} - -.category-filter:focus, .search-filter:focus { - outline: none; - border-color: var(--primary-color); -} - -.search-filter { - width: 200px; -} - -.header-right { - display: flex; - align-items: center; - gap: 16px; -} - -.view-controls { - display: flex; - gap: 4px; -} - -.repo-link { - display: flex; - align-items: center; - gap: 8px; - color: var(--text-color); - text-decoration: none; - font-weight: 600; - font-size: 14px; - padding: 8px 16px; - border: 1px solid var(--border-color); - border-radius: var(--radius-md); - transition: all 0.2s; -} - -.github-logo { - flex-shrink: 0; -} - -.repo-link:hover { - background-color: var(--bg-secondary); - border-color: var(--primary-color); - color: var(--primary-color); -} - -/* Main Container */ -.main-container { - display: flex; - min-height: calc(100vh - var(--header-height)); - overflow-x: hidden; - width: 100%; -} - -/* Sidebar Navigation */ -.sidebar { - width: var(--sidebar-width); - background: var(--bg-secondary); - -webkit-backdrop-filter: blur(10px); - backdrop-filter: blur(10px); - border-right: 1px solid var(--border-color); - position: fixed; - top: var(--header-height); - left: 0; - height: calc(100vh - var(--header-height)); - overflow-y: auto; - transition: transform 0.3s ease; - z-index: 998; - transform: translateX(-100%); - box-shadow: 2px 0 12px rgba(0, 0, 0, 0.15); -} - -.sidebar.active, -.sidebar.open { - transform: translateX(0); -} - -.sidebar-nav { - display: flex; - flex-direction: column; - height: 100%; -} - -.sidebar-header { - padding: 20px; - border-bottom: 1px solid var(--border-color); - display: flex; - justify-content: space-between; - align-items: center; -} - -.sidebar-header h2 { - font-family: "Funnel Display", sans-serif; - margin: 0; - font-size: 1.2rem; - color: var(--text-color); - font-weight: 600; -} - -.sidebar-close { - background: none; - border: none; - font-size: 24px; - cursor: pointer; - padding: 4px; - line-height: 1; -} - -.clear-filter-btn { - font-family: "Geist", sans-serif; - width: calc(100% - 40px); - margin: 16px 20px; - padding: 10px 16px; - background-color: var(--primary-color); - color: white; - border: 1px solid var(--primary-color); - border-radius: var(--radius-md); - font-size: 14px; - font-weight: 500; - cursor: pointer; - transition: all 0.2s; -} - -.clear-filter-btn:hover { - background-color: var(--primary-hover); - border-color: var(--primary-hover); -} - -.toc-list { - list-style: none; - padding: 0; - margin: 0; - flex: 1; - overflow-y: auto; -} - -.toc-list li { - border-bottom: 1px solid rgba(0,0,0,0.05); -} - -.toc-list a { - display: block; - padding: 12px 20px; - color: var(--text-color); - text-decoration: none; - transition: all 0.2s; - position: relative; -} - -.toc-list a:hover { - background-color: rgba(89, 113, 253, 0.08); - color: var(--primary-color); - padding-left: 24px; -} - -.sidebar-footer { - padding: 20px; - border-top: 1px solid var(--border-color); -} - -.download-btn { - display: block; - width: 100%; - padding: 10px 16px; - background-color: transparent; - color: var(--text-secondary); - text-align: center; - text-decoration: none; - border: 1px solid var(--border-color); - border-radius: var(--radius-md); - font-size: 13px; - font-weight: 500; - transition: all 0.2s; -} - -.download-btn:hover { - background-color: var(--bg-secondary); - color: var(--text-color); - border-color: var(--text-secondary); -} - -/* Sidebar overlay */ -.sidebar-overlay { - display: none; - position: fixed; - top: 0; - left: 0; - right: 0; - bottom: 0; - background: rgba(0, 0, 0, 0.3); - z-index: 997; -} - -.sidebar-overlay.active { - display: block; -} - -/* Main Content */ -.main-content { - flex: 1; - padding: 32px; - width: 100%; - position: relative; -} - -@media (min-width: 1600px) { - .main-content { - padding: 40px 60px; - } - .header-content { - padding: 0 60px; - } -} - -@media (min-width: 1920px) { - .main-content { - padding: 48px 80px; - } - .header-content { - padding: 0 80px; - } -} - -/* Loading Indicator */ -.loading-indicator, -.error-indicator { - display: flex; - flex-direction: column; - align-items: center; - justify-content: center; - min-height: 400px; - gap: 16px; -} - -.loading-indicator p, -.error-indicator p { - font-family: "Geist", sans-serif; - color: var(--text-secondary); - font-size: 14px; - font-weight: 500; -} - -.spinner { - width: 48px; - height: 48px; - border: 3px solid var(--border-color); - border-top: 3px solid var(--primary-color); - border-radius: 50%; - animation: spin 1s linear infinite; -} - -@keyframes spin { - 0% { transform: rotate(0deg); } - 100% { transform: rotate(360deg); } -} - -/* Benchmark Sections */ -.benchmark-set { - margin-bottom: 48px; - background: var(--bg-color); - border-radius: var(--radius-lg); - border: 1px solid var(--border-color); - overflow: visible; - transition: opacity 0.3s ease, border-color 0.3s ease; - position: relative; - z-index: 1; -} - -.benchmark-set:first-child { - margin-top: 0; -} - -.benchmark-set.no-data { - opacity: 0.5; -} - -.benchmark-set.no-data .benchmark-header { - cursor: not-allowed; -} - -.benchmark-set.no-data .collapse-icon { - visibility: hidden; -} - -.sticky-header-container { - position: relative; - z-index: 50; - background: var(--bg-secondary); - border-radius: var(--radius-lg) var(--radius-lg) 0 0; - margin: 0; - overflow: visible; - padding: 0; -} - -.benchmark-header { - display: flex; - align-items: center; - gap: 12px; - padding: 16px 32px; - background: var(--bg-secondary); - border-bottom: 1px solid var(--border-color); - border-radius: var(--radius-lg); - cursor: pointer; - -webkit-user-select: none; - user-select: none; - transition: background-color 0.2s; - overflow: visible; -} - -.benchmark-header:hover { - background: var(--bg-color); -} - -.title-wrapper { - display: flex; - align-items: center; - gap: 12px; - flex: 1; - min-width: 0; - overflow: visible; -} - -.benchmark-title { - font-family: "Funnel Display", sans-serif; - font-size: 1.5rem; - font-weight: 600; - margin: 0; - color: var(--text-color); - display: flex; - align-items: center; - gap: 12px; - letter-spacing: -0.02em; - line-height: 1.2; -} - -.group-link-btn { - font-size: 18px; - background: none; - border: none; - cursor: pointer; - opacity: 0.5; - padding: 4px 8px; - border-radius: var(--radius-sm); - transition: opacity 0.2s, background-color 0.2s, color 0.2s; - color: var(--text-secondary); -} - -.benchmark-header:hover .group-link-btn { - opacity: 1; -} - -.group-link-btn:hover { - background-color: var(--bg-secondary); - color: var(--primary-color); -} - -.group-link-btn.copied { - color: var(--accent-color); - opacity: 1; -} - -.collapse-icon { - font-size: 1rem; - flex-shrink: 0; - width: 1rem; - text-align: center; - display: inline-block; -} - -.benchmark-secondary-info { - display: flex; - align-items: center; - gap: 12px; - flex-shrink: 0; -} - -.benchmark-meta { - font-family: "Geist Mono", monospace; - display: flex; - gap: 16px; - font-size: 11px; - font-weight: 500; - color: var(--text-secondary); - letter-spacing: 0.02em; - text-transform: uppercase; - flex-shrink: 0; - line-height: 1; -} - -.info-icon { - display: inline-flex; - align-items: center; - justify-content: center; - width: 20px; - height: 20px; - border-radius: 50%; - background: var(--bg-secondary); - color: var(--text-secondary); - font-size: 16px; - cursor: help; - position: relative; - transition: background-color 0.2s, color 0.2s; - flex-shrink: 0; -} - -.info-icon:hover { - background: var(--primary-color); - color: white; -} - -.info-icon::after { - content: attr(data-tooltip); - position: absolute; - bottom: calc(100% + 8px); - left: 50%; - transform: translateX(-50%); - padding: 8px 12px; - background: rgba(16, 16, 16, 0.95); - color: white; - font-size: 13px; - line-height: 1.4; - border-radius: var(--radius-sm); - white-space: nowrap; - opacity: 0; - visibility: hidden; - transition: opacity 0.2s, visibility 0.2s; - pointer-events: none; - z-index: 10000; - font-weight: normal; -} - -.info-icon:hover::after { - opacity: 1; - visibility: visible; -} - -/* Generic Tooltip */ -[data-tooltip] { - position: relative; -} - -[data-tooltip]::after { - content: attr(data-tooltip); - position: absolute; - bottom: 100%; - left: 50%; - transform: translateX(-50%); - margin-bottom: 6px; - padding: 6px 10px; - background: rgba(16, 16, 16, 0.95); - color: white; - font-family: "Geist", sans-serif; - font-size: 12px; - font-weight: 500; - line-height: 1.3; - border-radius: var(--radius-sm); - white-space: nowrap; - opacity: 0; - visibility: hidden; - transition: all 0.15s; - pointer-events: none; - z-index: 1000; -} - -[data-tooltip]:hover::after { - opacity: 1; - visibility: visible; -} - -[data-tooltip]:disabled::after, -[data-tooltip][disabled]::after { - display: none; -} - -/* Engine Filter Controls */ -.engine-filter-container { - padding: 12px 32px; - background: var(--bg-secondary); - border-bottom: 1px solid var(--border-color); - display: flex; - align-items: center; - gap: 12px; - flex-wrap: wrap; -} - -.engine-filter-label { - font-family: "Geist", sans-serif; - font-size: 14px; - font-weight: 500; - color: var(--text-secondary); -} - -.engine-filter-btn { - font-family: "Geist", sans-serif; - padding: 6px 14px; - border: 1px solid var(--border-color); - background: var(--bg-color); - border-radius: var(--radius-sm); - font-size: 13px; - font-weight: 500; - color: var(--text-color); - cursor: pointer; - transition: all 0.2s; -} - -.engine-filter-btn:hover { - background-color: var(--bg-secondary); - border-color: var(--primary-color); -} - -.engine-filter-btn.active { - background-color: var(--primary-color); - color: white; - border-color: var(--primary-color); -} - -/* Chart Grid */ -.benchmark-graphs { - display: grid; - grid-template-columns: 1fr; - gap: 20px; - padding: 20px; - border-radius: 0 0 var(--radius-lg) var(--radius-lg); - background: var(--bg-color); -} - -@media (min-width: 1200px) { - .benchmark-graphs { - grid-template-columns: repeat(2, 1fr); - gap: 24px; - padding: 24px; - } - - .benchmark-graphs.single-chart { - grid-template-columns: 1fr; - max-width: 1400px; - margin: 0 auto; - } -} - -@media (min-width: 1600px) { - .benchmark-graphs { - padding: 28px 32px; - } -} - -.benchmark-graphs.list-view { - grid-template-columns: 1fr; - max-width: 1200px; - margin: 0 auto; -} - -.chart-container { - background: var(--bg-color); - border: 1px solid var(--border-color); - border-radius: var(--radius-md); - padding: 20px; - position: relative; - transition: all 0.3s ease; - display: flex; - flex-direction: column; -} - -.chart-container:hover { - box-shadow: var(--shadow-md); -} - -.chart-header { - display: flex; - justify-content: space-between; - align-items: center; - margin-bottom: 8px; -} - -.chart-title { - font-family: "Funnel Display", sans-serif; - font-size: 16px; - font-weight: 500; - color: var(--text-color); - letter-spacing: -0.01em; - display: flex; - align-items: center; - gap: 8px; -} - -.downsample-indicator { - font-family: "Geist Mono", monospace; - font-size: 10px; - font-weight: 600; - min-width: 100px; - padding: 2px 6px; - background-color: var(--vortex-pink); - color: var(--vortex-black); - border-radius: var(--radius-sm); - text-transform: uppercase; - letter-spacing: 0.02em; -} - -.chart-actions { - display: flex; - gap: 8px; -} - -.chart-action-btn { - background: var(--bg-color); - border: 1px solid var(--border-color); - padding: 6px 12px; - border-radius: var(--radius-sm); - cursor: pointer; - font-size: 12px; - font-weight: 500; - transition: all 0.2s; - display: flex; - align-items: center; - gap: 6px; -} - -.chart-action-btn:hover { - background-color: #f5f5f5; - border-color: var(--primary-color); -} - -.chart-zoom-controls { - display: flex; - gap: 2px; - margin-right: 8px; -} - -.chart-zoom-btn { - background: var(--bg-color); - border: 1px solid var(--border-color); - padding: 4px 8px; - border-radius: var(--radius-sm); - cursor: pointer; - font-size: 12px; - font-weight: 600; - min-width: 28px; - transition: all 0.15s; - color: var(--text-secondary); - display: flex; - align-items: center; - justify-content: center; -} - -.chart-zoom-btn:hover:not(:disabled) { - background-color: var(--primary-color); - border-color: var(--primary-color); - color: white; -} - -.chart-zoom-btn:disabled { - opacity: 0.4; - cursor: not-allowed; - background: var(--bg-secondary); -} - -/* Chart Canvas */ -.chart-container canvas { - max-height: 450px; - min-height: 320px; - width: 100% !important; - height: auto !important; - display: block; -} - -.benchmark-graphs.list-view .chart-container canvas { - max-height: 600px; - min-height: 400px; -} - -.chart-canvas-wrapper { - position: relative; - height: 100%; - min-height: 320px; -} - -.chart-canvas-wrapper.loading canvas { - filter: blur(2px); - pointer-events: none; - transition: filter 0.2s; -} - -.chart-canvas-placeholder { - display: flex; - align-items: center; - justify-content: center; - min-height: 320px; - max-height: 450px; - background: var(--bg-secondary); - border-radius: var(--radius-sm); - margin-top: 8px; - position: relative; - overflow: hidden; -} - -.chart-loading-overlay { - position: absolute; - top: 0; - left: 0; - right: 0; - bottom: 0; - display: flex; - align-items: center; - justify-content: center; - z-index: 10; -} - -.chart-loading-spinner { - width: 32px; - height: 32px; - border: 3px solid var(--border-color); - border-top-color: var(--primary-color); - border-radius: 50%; - animation: spin 0.8s linear infinite; -} - -@keyframes spin { - to { - transform: rotate(360deg); - } -} - -/* Back to Top Button */ -.back-to-top { - position: fixed; - bottom: 32px; - right: 32px; - width: 48px; - height: 48px; - background-color: var(--primary-color); - color: white; - border: none; - border-radius: 50%; - font-size: 20px; - cursor: pointer; - box-shadow: 0 4px 12px rgba(89, 113, 253, 0.3); - opacity: 0; - visibility: hidden; - transition: all 0.3s ease; - z-index: 999; -} - -.back-to-top.visible { - opacity: 1; - visibility: visible; -} - -.back-to-top:hover { - background-color: var(--primary-hover); - transform: translateY(-4px); - box-shadow: 0 8px 20px rgba(89, 113, 253, 0.4); -} - -/* Chart Modal */ -.chart-modal { - display: none; - position: fixed; - top: 0; - left: 0; - width: 100%; - height: 100%; - background-color: rgba(0, 0, 0, 0.8); - z-index: 2000; -} - -.chart-modal.active { - display: flex; - align-items: center; - justify-content: center; -} - -.modal-content { - background: var(--bg-color); - padding: 32px; - border-radius: var(--radius-lg); - width: 95%; - max-width: 1800px; - height: 90vh; - position: relative; - box-shadow: 0 20px 40px rgba(0, 0, 0, 0.2); -} - -.modal-header { - display: flex; - justify-content: space-between; - align-items: center; - margin-bottom: 16px; -} - -.modal-header h2 { - font-family: 'Funnel Display', sans-serif; - font-size: 20px; - font-weight: 600; - margin: 0; -} - -.modal-controls { - display: flex; - align-items: center; - gap: 12px; -} - -.modal-close-btn { - background: var(--bg-color); - border: 1px solid var(--border-color); - padding: 6px; - border-radius: var(--radius-sm); - cursor: pointer; - color: var(--text-secondary); - display: flex; - align-items: center; - justify-content: center; - transition: all 0.15s; -} - -.modal-close-btn:hover { - background-color: var(--primary-color); - border-color: var(--primary-color); - color: white; -} - -.modal-chart-container { - width: 100%; - height: calc(100% - 60px); - position: relative; -} - -.modal-chart-container.loading canvas { - filter: blur(2px); - pointer-events: none; - transition: filter 0.2s; -} - -/* Benchmark Scores Summary */ -.benchmark-scores-summary { - background: var(--bg-secondary); - border-bottom: 1px solid var(--border-color); - padding: 12px 32px; - margin: 0; - margin-top: 0 !important; -} - -.scores-title { - font-size: 12px; - font-weight: 600; - margin: 0 0 8px 0; - color: var(--text-secondary); - text-transform: uppercase; - letter-spacing: 0.05em; -} - -.scores-list { - column-count: 2; - column-gap: 20px; -} - -.scores-list:has(.score-item:only-child) { - column-count: 1; -} - -@media (max-width: 780px) { - .scores-list { - column-count: 1; - } -} - -.score-item { - display: flex; - align-items: center; - background: transparent; - padding: 6px 0; - margin-bottom: 4px; - border-radius: 0; - border: none; - transition: none; - font-size: 14px; - break-inside: avoid; -} - -.score-rank { - font-weight: 500; - color: var(--primary-color); - min-width: 24px; - font-size: 14px; -} - -.score-series { - flex: 1; - font-weight: 500; - color: var(--text-color); - margin: 0 8px; - font-size: 14px; - white-space: nowrap; - overflow: hidden; - text-overflow: ellipsis; -} - -.score-metrics { - display: flex; - gap: 6px; - align-items: center; -} - -.score-value { - font-family: 'Geist Mono', monospace; - font-weight: 600; - color: var(--primary-color); - font-size: 13px; -} - -.score-runtime { - font-family: 'Geist Mono', monospace; - font-weight: 600; - color: var(--text-secondary); - background: transparent; - padding: 0 4px; - border-radius: 0; - font-size: 13px; -} - -.scores-explanation { - margin-top: 8px; - font-size: 11px; - color: var(--text-secondary); - font-style: italic; - text-align: center; -} - -/* Utility Classes */ -.hidden { - display: none !important; -} - -/* Mobile Styles */ -@media (max-width: 768px) { - :root { - --header-height: 56px; - } - - .header-content { - padding: 0 8px; - gap: 8px; - } - - .header-left { - gap: 8px; - } - - .site-logo { - height: 16px; - } - - .site-title { - display: none; - } - - .header-center { - display: none; - } - - .view-controls { - display: none; - } - - .repo-link span { - display: none; - } - - .main-content { - padding: 16px; - max-width: 100vw; - } - - .benchmark-set { - margin-bottom: 24px; - border-radius: var(--radius-md); - overflow: hidden; - } - - .benchmark-header { - padding: 12px 16px; - flex-wrap: wrap; - gap: 8px; - overflow: hidden; - } - - .title-wrapper { - flex-wrap: wrap; - gap: 8px; - overflow: hidden; - } - - .benchmark-title { - font-size: 1.1rem; - flex: 1 1 auto; - min-width: 0; - word-break: break-word; - } - - .benchmark-meta { - flex-shrink: 1; - font-size: 10px; - } - - .benchmark-secondary-info { - flex-shrink: 1; - min-width: 0; - } - - .benchmark-graphs { - grid-template-columns: 1fr; - gap: 16px; - padding: 12px; - } - - .chart-container { - padding: 12px; - overflow: hidden; - } - - .chart-header { - flex-wrap: wrap; - gap: 8px; - } - - .chart-title { - flex: 1 1 100%; - min-width: 0; - font-size: 14px; - } - - .chart-actions { - width: 100%; - justify-content: flex-start; - } - - .chart-zoom-controls { - flex-wrap: wrap; - gap: 4px; - } - - .chart-zoom-btn { - padding: 6px 8px; - min-width: 32px; - } - - .chart-container canvas { - max-height: 350px; - min-height: 250px; - } - - .chart-container:hover { - transform: none; - box-shadow: none; - } - - .chart-action-btn { - display: none; - } - - .engine-filter-container { - padding: 12px 16px; - flex-wrap: wrap; - } - - .back-to-top { - bottom: 16px; - right: 16px; - width: 40px; - height: 40px; - font-size: 16px; - } - - .modal-content { - padding: 16px; - height: 90vh; - } - - .modal-header { - flex-wrap: wrap; - gap: 8px; - } - - .modal-header h2 { - font-size: 16px; - min-width: 0; - word-break: break-word; - } - - .benchmark-scores-summary { - padding: 8px 16px; - } -} diff --git a/benchmarks-website/src/utils.js b/benchmarks-website/src/utils.js deleted file mode 100644 index 140ebebeb6c..00000000000 --- a/benchmarks-website/src/utils.js +++ /dev/null @@ -1,104 +0,0 @@ -import { SERIES_COLOR_MAP, FALLBACK_PALETTE, CHART_NAME_MAP } from './config'; - -// Simple hash function for color selection -function simpleHash(str) { - let hash = 0; - for (let i = 0; i < str.length; i++) { - const char = str.charCodeAt(i); - hash = ((hash << 5) - hash) + char; - hash = hash & hash; - } - return Math.abs(hash); -} - -export function stringToColor(str) { - if (SERIES_COLOR_MAP[str]) { - return SERIES_COLOR_MAP[str]; - } - - const lowerStr = str - .replace(/^DataFusion:/i, 'datafusion:') - .replace(/^DuckDB:/i, 'duckdb:') - .replace(/^Vortex:/i, 'vortex:') - .replace(/^Arrow:/i, 'arrow:'); - - if (lowerStr !== str && SERIES_COLOR_MAP[lowerStr]) { - return SERIES_COLOR_MAP[lowerStr]; - } - - const index = simpleHash(str) % FALLBACK_PALETTE.length; - return FALLBACK_PALETTE[index]; -} - -export function remapChartName(name) { - if (CHART_NAME_MAP[name]) { - return CHART_NAME_MAP[name]; - } - // Convert dashes to spaces for readability - return name.replace(/-/g, ' '); -} - -export function formatDate(timestamp) { - if (!timestamp) return ''; - const date = new Date(timestamp); - const months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']; - return `${months[date.getMonth()]} ${date.getDate()}, ${date.getFullYear()}`; -} - -export function formatTime(ms) { - if (ms < 1) return `${(ms * 1000).toFixed(0)}μs`; - if (ms < 1000) return `${ms.toFixed(1)}ms`; - if (ms < 60000) return `${(ms / 1000).toFixed(1)}s`; - return `${(ms / 60000).toFixed(1)}m`; -} - -export function debounce(func, wait) { - let timeout; - return function (...args) { - clearTimeout(timeout); - timeout = setTimeout(() => func.apply(this, args), wait); - }; -} - -export function throttle(func, limit) { - let inThrottle; - return function (...args) { - if (!inThrottle) { - func.apply(this, args); - inThrottle = true; - setTimeout(() => (inThrottle = false), limit); - } - }; -} - -export function isMobile() { - return window.innerWidth <= 768; -} - -export function getBenchmarkDescription(categoryName) { - if (categoryName.startsWith('TPC-H')) { - const match = categoryName.match(/SF=(\d+)/); - const sf = match ? match[1] : null; - const sfDesc = sf ? `at SF=${sf} (~${sf === '1' ? '1GB' : sf === '10' ? '10GB' : sf === '100' ? '100GB' : '1TB'} of data)` : ''; - if (categoryName.includes('NVMe')) { - return `TPC-H benchmark queries on local NVMe storage ${sfDesc}`; - } else if (categoryName.includes('S3')) { - return `TPC-H benchmark queries against S3 storage ${sfDesc}`; - } - } - if (categoryName.startsWith('TPC-DS')) { - const match = categoryName.match(/SF=(\d+)/); - const sf = match ? match[1] : null; - const sfDesc = sf ? `at SF=${sf}` : ''; - return `TPC-DS benchmark queries on local NVMe storage ${sfDesc}`; - } - const descriptions = { - 'Random Access': 'Tests performance of selecting arbitrary row indices from a file on NVMe storage', - '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-website/vite.config.js b/benchmarks-website/vite.config.js deleted file mode 100644 index ad0cc7446f6..00000000000 --- a/benchmarks-website/vite.config.js +++ /dev/null @@ -1,19 +0,0 @@ -import { defineConfig } from 'vite'; -import react from '@vitejs/plugin-react'; - -export default defineConfig({ - plugins: [react()], - publicDir: 'public', - server: { - port: 5173, - proxy: { - '/api': { - target: 'http://localhost:3000', - changeOrigin: true, - }, - }, - }, - build: { - outDir: 'dist', - }, -}); diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index baff1daad74..d214af23511 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -30,6 +30,7 @@ regex = { workspace = true } tokio = { workspace = true, features = ["full"] } tracing = { workspace = true } vortex = { workspace = true } +vortex-arrow = { workspace = true } vortex-bench = { workspace = true } [features] diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index 87416cdf924..fae0cc44189 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -18,6 +18,7 @@ use vortex::expr::root; use vortex::expr::select; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; +use vortex_arrow::ToArrowType; use vortex_bench::Format; use vortex_bench::SESSION; use vortex_bench::compress::Compressor; diff --git a/benchmarks/datafusion-bench/Cargo.toml b/benchmarks/datafusion-bench/Cargo.toml index d0015dd2995..a195bbaa03f 100644 --- a/benchmarks/datafusion-bench/Cargo.toml +++ b/benchmarks/datafusion-bench/Cargo.toml @@ -14,6 +14,9 @@ rust-version = { workspace = true } version = { workspace = true } publish = false +[package.metadata.cargo-shear] +ignored = ["vortex-cuda"] + [dependencies] anyhow = { workspace = true } clap = { workspace = true, features = ["derive"] } @@ -36,6 +39,7 @@ parking_lot = { workspace = true } tokio = { workspace = true, features = ["full"] } url = { workspace = true } vortex = { workspace = true, features = ["object_store", "files", "tokio"] } +vortex-arrow = { workspace = true } vortex-bench = { workspace = true } vortex-cuda = { workspace = true, optional = true } vortex-datafusion = { workspace = true } diff --git a/benchmarks/datafusion-bench/src/lib.rs b/benchmarks/datafusion-bench/src/lib.rs index 1100d38d24e..c45aa99be2c 100644 --- a/benchmarks/datafusion-bench/src/lib.rs +++ b/benchmarks/datafusion-bench/src/lib.rs @@ -111,10 +111,9 @@ 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_with_options( - SESSION.clone(), - vortex_table_options(), - )), + Format::OnDiskVortex | Format::VortexCompact | Format::VortexNative => 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`") } diff --git a/benchmarks/datafusion-bench/src/main.rs b/benchmarks/datafusion-bench/src/main.rs index b8f9ac42df6..67521a8147e 100644 --- a/benchmarks/datafusion-bench/src/main.rs +++ b/benchmarks/datafusion-bench/src/main.rs @@ -29,6 +29,7 @@ use parking_lot::Mutex; use tokio::fs::File; use vortex::io::filesystem::FileSystemRef; use vortex::scan::DataSourceRef; +use vortex_arrow::ToArrowType; use vortex_bench::Benchmark; use vortex_bench::BenchmarkArg; use vortex_bench::CompactionStrategy; diff --git a/benchmarks/duckdb-bench/Cargo.toml b/benchmarks/duckdb-bench/Cargo.toml index b70375cf963..77c3d5280fc 100644 --- a/benchmarks/duckdb-bench/Cargo.toml +++ b/benchmarks/duckdb-bench/Cargo.toml @@ -14,6 +14,9 @@ rust-version.workspace = true version.workspace = true publish = false +[package.metadata.cargo-shear] +ignored = ["vortex-cuda"] + [dependencies] anyhow = { workspace = true } clap = { workspace = true, features = ["derive"] } diff --git a/benchmarks/duckdb-bench/src/lib.rs b/benchmarks/duckdb-bench/src/lib.rs index 4ec1efd0993..4f3bdd388fc 100644 --- a/benchmarks/duckdb-bench/src/lib.rs +++ b/benchmarks/duckdb-bench/src/lib.rs @@ -78,6 +78,12 @@ impl DuckClient { for stmt in &statements { self.connection().query(stmt)?; } + // After `LOAD spatial`, shadow the overridden spatial functions so that their filters + // push down. No-op without it. + self.db + .as_ref() + .vortex_expect("DuckClient database accessed after close") + .register_spatial_overrides()?; self.init_sql = statements; Ok(()) } @@ -127,6 +133,11 @@ impl DuckClient { .vortex_expect("connection just opened") .query(stmt)?; } + // Re-shadow the overridden spatial functions against the fresh instance. + self.db + .as_ref() + .vortex_expect("database just opened") + .register_spatial_overrides()?; Ok(()) } @@ -169,7 +180,10 @@ impl DuckClient { file_format: Format, ) -> Result<()> { let object_type = match file_format { - Format::Parquet | Format::OnDiskVortex | Format::VortexCompact => "VIEW", + Format::Parquet + | Format::OnDiskVortex + | Format::VortexCompact + | Format::VortexNative => "VIEW", Format::OnDiskDuckDB => "TABLE", Format::Lance => { anyhow::bail!( diff --git a/benchmarks/duckdb-bench/src/main.rs b/benchmarks/duckdb-bench/src/main.rs index cf4fa071067..c496929f6a5 100644 --- a/benchmarks/duckdb-bench/src/main.rs +++ b/benchmarks/duckdb-bench/src/main.rs @@ -142,6 +142,7 @@ fn main() -> anyhow::Result<()> { // OnDiskDuckDB tables are created during register_tables by loading from Parquet _ => {} } + benchmark.prepare_format(format, &base_path).await?; } anyhow::Ok(()) diff --git a/clippy.toml b/clippy.toml index 9f6d5bd21b8..73f74ca35f5 100644 --- a/clippy.toml +++ b/clippy.toml @@ -17,4 +17,14 @@ disallowed-methods = [ { 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 }, + { path = "vortex_array::legacy_session", reason = "Relies on the hidden global session; thread an explicit `VortexSession`/`ExecutionCtx` through instead" }, + { path = "vortex_array::ToCanonical::to_null", reason = "This function doesn't take the context, use `array.execute::(ctx)` instead", allow-invalid = true }, + { path = "vortex_array::ToCanonical::to_bool", reason = "This function doesn't take the context, use `array.execute::(ctx)` instead", allow-invalid = true }, + { path = "vortex_array::ToCanonical::to_primitive", reason = "This function doesn't take the context, use `array.execute::(ctx)` instead", allow-invalid = true }, + { path = "vortex_array::ToCanonical::to_decimal", reason = "This function doesn't take the context, use `array.execute::(ctx)` instead", allow-invalid = true }, + { path = "vortex_array::ToCanonical::to_struct", reason = "This function doesn't take the context, use `array.execute::(ctx)` instead", allow-invalid = true }, + { path = "vortex_array::ToCanonical::to_listview", reason = "This function doesn't take the context, use `array.execute::(ctx)` instead", allow-invalid = true }, + { path = "vortex_array::ToCanonical::to_fixed_size_list", reason = "This function doesn't take the context, use `array.execute::(ctx)` instead", allow-invalid = true }, + { path = "vortex_array::ToCanonical::to_varbinview", reason = "This function doesn't take the context, use `array.execute::(ctx)` instead", allow-invalid = true }, + { path = "vortex_array::ToCanonical::to_extension", reason = "This function doesn't take the context, use `array.execute::(ctx)` instead", allow-invalid = true }, ] diff --git a/deny.toml b/deny.toml index b9f9c79a796..21240f1f7e7 100644 --- a/deny.toml +++ b/deny.toml @@ -15,8 +15,9 @@ feature-depth = 1 ignore = [ # Paste is no longer maintained because its essentially "finished". "RUSTSEC-2024-0436", - # proc-macro-error-2 is unmaintained, only used by the `test_with` test dependency - "RUSTSEC-2026-0173" + # quick-xml is a transitive dependency that we cannot upgrade ourselves + "RUSTSEC-2026-0195", + "RUSTSEC-2026-0194" ] [licenses] diff --git a/docs/Doxyfile.cpp b/docs/Doxyfile.cpp index 553c07e6dae..23072f7479a 100644 --- a/docs/Doxyfile.cpp +++ b/docs/Doxyfile.cpp @@ -5,7 +5,7 @@ PROJECT_NAME = "Vortex C++" OUTPUT_DIRECTORY = _build/doxygen-cpp # Input sources -INPUT = ../vortex-cxx/cpp/include/vortex +INPUT = ../lang/cpp/include/vortex FILE_PATTERNS = *.hpp RECURSIVE = NO @@ -26,6 +26,3 @@ MACRO_EXPANSION = NO # Suppress warnings about undocumented members (WIP API) WARN_IF_UNDOCUMENTED = NO - -# Exclude cxx bridge internals from documentation -EXCLUDE_SYMBOLS = ffi::* diff --git a/docs/api/c/index.rst b/docs/api/c/index.rst index 3d14a1a5d60..81d142d16c0 100644 --- a/docs/api/c/index.rst +++ b/docs/api/c/index.rst @@ -83,31 +83,5 @@ responsible for freeing them. .. c:autofunction:: vx_error_free :file: vortex.h -.. c:autofunction:: vx_error_get_message +.. c:autofunction:: vx_error_message :file: vortex.h - -Strings -------- - -Vortex strings wrap a Rust `Arc`, and therefore are reference-counted, UTF-8 encoded, and not null-terminated. - -.. c:autotype:: vx_string - :file: vortex.h - -.. c:autofunction:: vx_string_clone - :file: vortex.h - -.. c:autofunction:: vx_string_free - :file: vortex.h - -.. c:autofunction:: vx_string_new - :file: vortex.h - -.. c:autofunction:: vx_string_new_from_cstr - :file: vortex.h - -.. c:autofunction:: vx_string_len - :file: vortex.h - -.. c:autofunction:: vx_string_ptr - :file: vortex.h diff --git a/docs/api/cpp/dtypes.rst b/docs/api/cpp/dtypes.rst deleted file mode 100644 index ae026af604d..00000000000 --- a/docs/api/cpp/dtypes.rst +++ /dev/null @@ -1,17 +0,0 @@ -Data Types -========== - -Logical data types for Vortex arrays. ``DType`` represents a logical type, and ``PType`` -enumerates the physical primitive types. - -PType ------ - -.. doxygenenum:: vortex::PType - -DType ------ - -.. doxygennamespace:: vortex::dtype - :members: - :undoc-members: diff --git a/docs/api/cpp/expr.rst b/docs/api/cpp/expr.rst deleted file mode 100644 index 4e904fff558..00000000000 --- a/docs/api/cpp/expr.rst +++ /dev/null @@ -1,9 +0,0 @@ -Expressions -=========== - -Expression trees used for filter and projection pushdowns in the scan API. Build expressions -from columns, literals, and comparison operators. - -.. doxygennamespace:: vortex::expr - :members: - :undoc-members: diff --git a/docs/api/cpp/file.rst b/docs/api/cpp/file.rst deleted file mode 100644 index 26736a88979..00000000000 --- a/docs/api/cpp/file.rst +++ /dev/null @@ -1,17 +0,0 @@ -File I/O -======== - -Read and write Vortex files from C++. ``VortexFile`` opens an existing file for scanning, while -``VortexWriteOptions`` writes an Arrow array stream to a new file. - -VortexFile ----------- - -.. doxygenclass:: vortex::VortexFile - :members: - -VortexWriteOptions ------------------- - -.. doxygenclass:: vortex::VortexWriteOptions - :members: diff --git a/docs/api/cpp/index.rst b/docs/api/cpp/index.rst index befb980f6a9..642b3d18ce9 100644 --- a/docs/api/cpp/index.rst +++ b/docs/api/cpp/index.rst @@ -1,82 +1,255 @@ C++ API ======= -The Vortex C++ API provides an idiomatic C++ wrapper around the Vortex C FFI, built using -`cxx `_. It currently supports reading and writing Vortex files and integrates -with the Arrow C Data Interface via `nanoarrow `_. - -In the future we will expand the C++ API to cover Vortex's plugin and extension points. Please -reach out if you are interested in extending Vortex from C++ so we can prioritize these features. +The Vortex C++ API allows you to read and write ``.vortex`` files directly or +via an Arrow compatibility layer like `nanoarrow +`_. The only dependency apart from Vortex +is ``nanoarrow``. .. note:: - Both the C++ API and this documentation are a work in progress. The API surface may change - significantly. Please reach out if you are interested in using Vortex from C++ so we can - prioritize stabilization. - + C++ API is a work in progress. Please reach out to us if you are interested + in using Vortex from C++ or you want a feature not covered yet e.g. + extension support. Installation ------------ -The C++ bindings are built using CMake. Requirements: +We don't provide prebuilt library files (yet) so you will need to build Vortex +from source, and for that you will need: -* CMake 3.22 or higher -* C++20 compatible compiler -* Rust toolchain (for building the underlying Rust library) +- C++20, +- Rust toolchain, +- and CMake 3.10. .. code-block:: bash + git clone --depth 1 https://github.com/vortex-data/vortex + cd vortex + cargo build --release -p vortex-ffi cd vortex-cxx - mkdir build && cd build - cmake .. - make -j$(nproc) + mkdir build + cmake -Bbuild -DCMAKE_BUILD_TYPE=Release + # To build the examples, pass -DBUILD_EXAMPLES=1 + # cmake -Bbuild -DBUILD_EXAMPLES=1 -Compatibility -------------- + cmake --build build -j -The C++ bindings are supported on the following architectures: +This produces a shared and a static library which you can use directly or via -* x86_64 Linux -* ARM64 Linux -* Apple Silicon macOS +.. code-block:: cmake -They support any Linux distribution with a GLIBC version >= 2.31. This includes + # static library + target_link_libraries(target PRIVATE vortex_cxx) + # shared library + target_link_libraries(target PRIVATE vortex_cxx_shared) -* Amazon Linux 2022 or newer -* Ubuntu 20.04 or newer +Have a look at the `examples +`_ +directory as well. - -Usage Example +Reading files ------------- -Here's a basic example of reading a Vortex file into an Arrow array stream: +Full source code for this example is `reader.cpp +`_. +For brevity we omit ``main()``, system includes and ``using`` directives. + +Assuming you have Vortex files ``people0``, ``people1``, and ``me`` in a local folder, +each containing U8 column "age" and U16 column "height", this is how you +print all ages for specific heights: .. code-block:: cpp - #include "vortex/file.hpp" - #include "vortex/scan.hpp" + Session session; + DataSource ds = DataSource::open(session, {"people*.vortex", "me.vortex"}); + Scan scan = ds.scan({.filter = col("height") >= lit(50)}); + + for (Partition &partition : scan.partitions()) { + for (Array &array : partition.batches()) { + Array age = array.field("age"); + PrimitiveView age_view = age.values(session); + std::span age_values = age_view.values(); + for (uint8_t value : age_values) { + std::cout << int(value) << " "; + } + } + } + std::cout << "\n"; + +DataSource and Scan +^^^^^^^^^^^^^^^^^^^ + +First, you need to create a Vortex session, which does extension bookkeeping and +may also hold metadata like object store credentials (we'll get back to it +later). Further, a DataSource provides a view over multiple files which may also +be remote. You can specify globs for every item as it's shown. + +Once you have a DataSource, you can create Scans out of it. A Scan is a single +traversal of a DataSource which projects columns and filters on them. Our Scan +is consumed by following calls so it needs to be mutable. ScanOptions, which is +passed to Scan, is a simple C++ aggregate so you can initialize any fields you +want or avoid them altogether (``auto scan = ds.scan()``). + +Expressions +^^^^^^^^^^^ + +If you omit ``.projection`` all columns are read as part of a scan. If we wanted +to return only the "age" column we'd write ``.projection = col("age")``. We pass +an Expression to a filter which returns false for some "height" values, and the +scan filters them out. + +Two additional things to look for in ``.filter``: first, we allow overloading +Expression operators which produce Expressions themselves à la Eigen. This is +opt-in via ``using namespace vortex::expr::ops``. If you prefer, you can use +``eq(col("height"), lit(180))`` instead. The second thing is that we +explicitly pass the type to the ``lit`` expression, which creates a literal +constant. We don't do any type coercion in Vortex so if you were to write +``lit(180)``, C++ would likely deduce the type to ``int`` and fail at runtime. + +Once we're done with the scan, we need to consume the data it provides. Vortex +has Arrow interoperability, but now let's focus on Partitions and Arrays. + +Partitions and Arrays +^^^^^^^^^^^^^^^^^^^^^ + +A Partition is an independent unit of work. Assuming your Vortex file +processing will likely be multithreaded, each thread can pull Partitions from +a Scan and handle them in parallel. For simplicity we don't cover it here, but +look for the next sections! Each Partition produces Arrays which are batches of +rows and columns. - // Open a Vortex file and scan with a row limit - auto stream = vortex::VortexFile::Open("data.vortex") - .CreateScanBuilder() - .WithLimit(1000) - .IntoStream(); +.. note:: + Arrays, DataTypes, and Expressions are reference-counted so copying them is + cheap. - // Consume the Arrow C Data stream - ArrowArray array; - while (stream.get_next(&stream, &array) == 0 && array.release != nullptr) { - // Process each batch... - } +Once we have an Array, we can get access to its values. First, we need to +extract the "age" field. + +.. note:: + Default projection returns the root Array which is a Struct so to get + a field from it we need to use ``field()``. If we were to use ``col()`` in + projection, the array we got would be "age" column already, so we wouldn't + need to use ``field()``. -API Reference +Then we want to get access to the raw bytes. However, +Array likely references compressed data, so now we need to learn about +canonicalization. + +Canonicalization +^^^^^^^^^^^^^^^^ + +Vortex files hold trees of compressed data where each node is a specific +encoding (zstd, FSST, delta) over another node. This is good for performance +because we defer decompression and we can also pass compressed data directly to +other systems. Say, if a reader knows how to deal with bitpacked integers but +not RLE and we have ``RLE(Bitpacked(U64))``, we can decompress just RLE and +pass bitpacked array directly to the reader. + +In this example we want to remove all encodings and uncompress all data fully. +This form is called a canonical Array, and the process is canonicalization. + +When we request ``.values(session)``, we canonicalize the Array and get a +PrimitiveView because we already know the type of the column. As PrimitiveView +holds uncompressed data, we can get raw ``uint64_t`` numbers by calling +``.values()`` on the view. + +Writing files ------------- -.. toctree:: - :maxdepth: 2 +Now let's write the first files to be read by our previous example. +Source code for this example is `writer.cpp +`_. + +.. code-block:: cpp + + Session session; + DataType dtype = dtype::struct_({ + {"age", dtype::uint8()}, + {"height", dtype::uint16(Nullable)}, + }); + + constexpr size_t SAMPLE_ROWS = 100; + std::vector age_buffer(SAMPLE_ROWS); + std::vector height_buffer(SAMPLE_ROWS); + for (size_t i = 0; i < SAMPLE_ROWS; ++i) { + age_buffer[i] = static_cast(i); + height_buffer[i] = static_cast((i + 1) % 200); + } - dtypes - scalars - expr - file - scan + Array age = Array::primitive(age_buffer); + Array array = make_struct({ + {"age", age}, + {"height", Array::primitive(height_buffer, AllValid)}, + }); + + Expression age_gt_10 = expr::gt(expr::col("age"), expr::lit(10)); + Array validity_array = array.apply(age_gt_10); + + Validity validity = Validity::from_array(validity_array); + Array array2 = make_struct({ + {"age", age}, + {"height", Array::primitive(height_buffer, validity)}, + }); + + Writer writer = Writer::open(session, argv[1], dtype); + writer.push({array, array2}); + writer.finish(); + +DataType +^^^^^^^^ + +First, you want to create a schema. DataType is a logical representation of +Vortex's type system. "Logical" as opposed to "physical" means DataType doesn't +know what encodings the data use. As another example, Arrow and Parquet's types +are physical. There's a convenience function ``struct_`` which creates a +``DataTypeVariant::Struct`` with fields passed as a ``initializer_list``. + +Each DataType has a notion of nullability: whether some items in the row can be +invalid (Vortex uses the terms "null" and "invalid" interchangeably). DataTypes +are non-nullable by default, so "age" column is not nullable, but "height" is. + +Once we create Arrays for columns we can merge them into a Struct Array with +``make_struct``, and we need to say a word about Validity. + +Nullability and Validity +^^^^^^^^^^^^^^^^^^^^^^^^ + +Nullability (and ``Nullable`` flag) is a type-system note that we `may` have +invalid elements in the column. A non-``Nullable`` column can't have nulls in +it, but a ``Nullable`` column's items may be all valid as well. Validity, on the +other hand, tells us whether we `do` have nulls in a particular Array. + +"height" in the first Array is AllValid which means in this Array we don't have +invalid items. For the second Array we reuse age by copying it (which is a +reference clone so it's cheap). But for "height" we say Validity is determined +by another Array. This array consists of Bools and we obtained it by applying a +comparison Expression to the "age" column of the first array. So, "height" field +in ``array2`` will have null/invalid values if "age"'s item at the same index +was less than 11. + +Applying Expressions +^^^^^^^^^^^^^^^^^^^^ + +One important thing to know is that applying an Expression is a constant-time +operation, and the returned array contains the original values along with some +meta-information about the expression. Real application happens when you execute +the array. One example of executing the array is canonicalization, which in this +example happens later when we write both arrays to the file. + +Writing to a file +^^^^^^^^^^^^^^^^^ + +Once we're done with the arrays, we need to initialize the Writer and push +arrays into it. ``finish()`` call writes the file footer. If you don't call it and +let Writer go out of scope, the file will be visibly corrupted. + +Now you can build the example and read back the generated files: + +.. code-block:: + + ./build/examples/writer people0.vortex + ./build/examples/writer people1.vortex + ./build/examples/writer me.vortex + ./build/examples/reader diff --git a/docs/api/cpp/scalars.rst b/docs/api/cpp/scalars.rst deleted file mode 100644 index ab52cc22af0..00000000000 --- a/docs/api/cpp/scalars.rst +++ /dev/null @@ -1,9 +0,0 @@ -Scalars -======= - -A scalar is a single typed value. Factory functions create scalars of each primitive type, and -``cast`` converts between types. - -.. doxygennamespace:: vortex::scalar - :members: - :undoc-members: diff --git a/docs/api/cpp/scan.rst b/docs/api/cpp/scan.rst deleted file mode 100644 index dbbd9e3bd36..00000000000 --- a/docs/api/cpp/scan.rst +++ /dev/null @@ -1,18 +0,0 @@ -Scanning -======== - -The scan API provides a builder pattern for reading data from a Vortex file with optional -filter, projection, row range, and limit pushdowns. The resulting stream exposes the -Arrow C Data Interface (``ArrowArrayStream``). - -ScanBuilder ------------ - -.. doxygenclass:: vortex::ScanBuilder - :members: - -StreamDriver ------------- - -.. doxygenclass:: vortex::StreamDriver - :members: diff --git a/docs/concepts/index.md b/docs/concepts/index.md index 878296b3d09..5a2741655cf 100644 --- a/docs/concepts/index.md +++ b/docs/concepts/index.md @@ -42,7 +42,7 @@ segment retrieval, FlatBuffer metadata for O(1) schema access, and support for m ## Integrations **Language bindings:** [Rust](https://docs.rs/vortex), [Python](../api/python/index.rst), -[Java](../api/java/index.rst), [C](../api/c/index.rst), [C++](../api/cpp/index.rst) +[Java](../api/java/index.rst), [C](../api/c/index.rst) **Query engines:** [DataFusion](../user-guide/datafusion.md), [DuckDB](../user-guide/duckdb.md), [Spark](../user-guide/spark.md), [Polars](../user-guide/polars.md), [Ray](../user-guide/ray.md) diff --git a/docs/conf.py b/docs/conf.py index adcd94f262f..a2cadbb255f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,8 +1,6 @@ import doctest import os import re -import shutil -import subprocess from pathlib import Path import hawkmoth.docstring @@ -120,40 +118,13 @@ os.makedirs(os.path.dirname(_doxygen_xml_dir), exist_ok=True) -if not shutil.which("doxygen"): - raise RuntimeError("doxygen is required to build the docs but was not found on PATH") -subprocess.run(["doxygen", "Doxyfile.cpp"], cwd=Path(__file__).parent, check=True) +# if not shutil.which("doxygen"): +# raise RuntimeError("doxygen is required to build the docs but was not found on PATH") +# subprocess.run(["doxygen", "Doxyfile.cpp"], cwd=Path(__file__).parent, check=True) breathe_projects = {"vortex-cpp": _doxygen_xml_dir} breathe_default_project = "vortex-cpp" -# C++ types from cxx bridge and standard library that Sphinx cannot resolve. -nitpick_ignore += [ - ("cpp:identifier", t) - for t in [ - "vortex", - "rust", - "ffi", - "uint8_t", - "uint16_t", - "uint32_t", - "uint64_t", - "int8_t", - "int16_t", - "int32_t", - "int64_t", - "size_t", - "std::size_t", - ] -] -nitpick_ignore_regex = [ - # cxx bridge internals that will never be resolvable in Sphinx. - (r"cpp:identifier", r"rust::.*"), - (r"cpp:identifier", r"ffi::.*"), - # Doxygen file-level labels (e.g. "dtype_8hpp") that we don't generate pages for. - (r"ref", r".*_8hpp"), -] - # -- Options for hawkmoth C API gen ---------------------------- hawkmoth_root = str(git_root / "vortex-ffi/cinclude") @@ -171,6 +142,7 @@ ("c:identifier", "int16_t"), ("c:identifier", "uint8_t"), ("c:identifier", "int8_t"), + ("c:identifier", "vx_view"), ] hawkmoth_transform_default = "c_to_rust" diff --git a/docs/developer-guide/embedding/index.md b/docs/developer-guide/embedding/index.md index e204939431b..8c8154718ca 100644 --- a/docs/developer-guide/embedding/index.md +++ b/docs/developer-guide/embedding/index.md @@ -3,7 +3,7 @@ :::{warning} This section is under construction. For guidance on embedding Vortex, please join the [Vortex Slack channel](https://vortex.dev/slack) -or start a [GitHub Discussion](https://github.com/vortex-data/vortex/discussions). +or open a [GitHub Issue](https://github.com/vortex-data/vortex/issues/new/choose). ::: Vortex can be embedded into applications and services via its C FFI, C++ wrapper, or the Scan API. diff --git a/docs/developer-guide/extending/index.md b/docs/developer-guide/extending/index.md index ff004a69bc5..c8e862c27be 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/vortex-data/vortex/discussions). +or open a [GitHub Issue](https://github.com/vortex-data/vortex/issues/new/choose). ::: Vortex is designed to be extended with custom types, encodings, layouts, and compute functions. diff --git a/docs/project/contributing.md b/docs/project/contributing.md index fa6ea57d926..4721f335dfc 100644 --- a/docs/project/contributing.md +++ b/docs/project/contributing.md @@ -7,15 +7,17 @@ The full contributing guide lives in the repository: Below is a brief summary. -## Reporting Issues +## Issues and Questions -Bugs should be filed as [GitHub Issues](https://github.com/vortex-data/vortex/issues). Open-ended -questions and feature requests should be filed as -[GitHub Discussions](https://github.com/vortex-data/vortex/discussions). +Bugs, feature requests, and questions should all be filed as +[GitHub Issues](https://github.com/vortex-data/vortex/issues/new/choose). We strongly prefer that +you use one of the provided issue templates rather than opening a blank issue; templates make sure +we get the information needed to act on your report. For quick questions, the +[Vortex Slack channel](https://vortex.dev/slack) is also a good option. ## Code Contributions -1. Start a discussion on GitHub (unless the change is trivial). +1. Start a discussion by creating or commenting on a GitHub Issue (unless the change is trivial). 2. Implement the change, including tests for new functionality or bug fixes. 3. Open a pull request — ensure CI passes and that you sign off your commits (see below). CI requires approval from a committer for first-time contributors. diff --git a/docs/user-guide/spark.md b/docs/user-guide/spark.md index 676c11f332b..5b1f9b44b8d 100644 --- a/docs/user-guide/spark.md +++ b/docs/user-guide/spark.md @@ -1,15 +1,52 @@ # Spark Vortex provides a Spark DataSource V2 connector for reading and writing Vortex files. The -connector is published to Maven Central as `dev.vortex:vortex-spark`. +connector is published to Maven Central in two flavors: -## Installation +- `dev.vortex:vortex-spark_2.13` for Spark 4.x (Scala 2.13) +- `dev.vortex:vortex-spark_2.12` for Spark 3.5.x (Scala 2.12) -Add the dependency to your build. The connector is built against Spark 4.x with Scala 2.13. +Use the `all` classifier JAR (e.g. `vortex-spark_2.13-0.78.0-all.jar`). It is self-contained: +it bundles the Vortex JNI bindings, native libraries for Linux (x86_64 and aarch64) and macOS +(aarch64), and relocates its Arrow, Guava, and Jackson dependencies to avoid classpath +conflicts with Spark. The thin (unclassified) JAR does not work on its own because it +references relocated classes that only ship in the `all` JAR. + +## Getting Vortex into Spark + +For `spark-shell`, `spark-submit`, or `pyspark`, pass the `all` JAR with `--jars`. Spark +accepts either a local path or a URL, so you can point directly at Maven Central: + +```shell +spark-shell --jars https://repo1.maven.org/maven2/dev/vortex/vortex-spark_2.13/0.78.0/vortex-spark_2.13-0.78.0-all.jar +``` + +Or equivalently when building a session programmatically, e.g. in PySpark: + +```python +spark = ( + SparkSession.builder + .config("spark.jars", "/path/to/vortex-spark_2.13-0.78.0-all.jar") + .getOrCreate() +) +``` + +```{note} +`--packages dev.vortex:vortex-spark_2.13:0.78.0` does not work: `--packages` cannot select +the `all` classifier and resolves the thin JAR, which fails at runtime with +`NoClassDefFoundError: dev/vortex/relocated/...`. +``` + +Once the JAR is on the classpath, the connector registers itself automatically under the +format name `vortex` — no session configuration is required. + +## Installation as a Build Dependency + +To depend on the connector from a JVM project, add the `all` classifier to the dependency: ````{tab} Gradle (Kotlin) ```kotlin -implementation("dev.vortex:vortex-spark:") +implementation("dev.vortex:vortex-spark_2.13:0.78.0:all") ``` ```` @@ -17,18 +54,18 @@ implementation("dev.vortex:vortex-spark:") ```xml dev.vortex - vortex-spark - ${vortex.version} + vortex-spark_2.13 + 0.78.0 + all ``` ```` -The connector ships as a shadow JAR that relocates its Arrow, Guava, and Protobuf dependencies -to avoid classpath conflicts with Spark. - ## Reading Vortex Files -Use the `vortex` format to read a single file or a directory of Vortex files: +Paths may be local filesystem paths (`/path/to/data`) or URLs (`file:///path/to/data`, +`s3://bucket/path/to/data`). Use the `vortex` format to read a single file or a directory of +Vortex files: ```java Dataset df = spark.read() @@ -65,6 +102,72 @@ Each Spark partition produces one output file named `part-{partitionId}-{taskId} The connector supports all standard Spark save modes: `Overwrite`, `Append`, `Ignore`, and `ErrorIfExists`. +## Spark SQL + +The connector can also be used from pure SQL. To query existing Vortex files, register them +as a temporary view: + +```sql +CREATE TEMPORARY VIEW people +USING vortex +OPTIONS (path '/path/to/data'); + +SELECT name, age FROM people WHERE age > 30; +``` + +Tables can be created with `USING vortex`, then written to and read back with plain SQL. +With a `LOCATION` clause the table is external, backed by the files at that path; without +one the table is managed, and Spark stores its data under the warehouse directory (and +deletes it on `DROP TABLE`): + +```sql +CREATE TABLE student (id INT, name STRING, age INT) +USING vortex; + +INSERT INTO student VALUES (1, 'Alice', 20), (2, 'Bob', 21); + +SELECT * FROM student; +``` + +`CREATE TABLE ... AS SELECT` works the same way: + +```sql +CREATE TABLE adults +USING vortex +AS SELECT * FROM people WHERE age >= 18; +``` + +```{note} +On Spark 3.5, `CREATE TABLE ... USING vortex` additionally requires replacing the session +catalog, because Spark 3.5's built-in catalog cannot read tables backed by a DataSource +V2-only connector: + + spark.sql.catalog.spark_catalog=dev.vortex.spark.VortexSessionCatalog + +The extension delegates everything to the built-in session catalog (including the Hive +metastore, if configured) and only changes how `vortex` tables are resolved; tables of other +providers are untouched. It is not needed on Spark 4, though setting it is harmless. +``` + +## Direct File Queries + +Spark's built-in ``SELECT * FROM format.`path` `` syntax only works for built-in file +formats, so the connector ships a path-based catalog that provides the equivalent for +Vortex. Register it in the session configuration under the name `vortex`: + +```shell +spark-sql --conf spark.sql.catalog.vortex=dev.vortex.spark.VortexCatalog +``` + +Then query a Vortex file, or a directory of Vortex files, directly by path — no view or +table required: + +```sql +SELECT * FROM vortex.`/path/to/data`; + +INSERT INTO vortex.`/path/to/data` VALUES (1, 'Alice', 20); +``` + ## Supported Types | Spark Type | Vortex Type | diff --git a/encodings/alp/src/alp/compute/cast.rs b/encodings/alp/src/alp/compute/cast.rs index d77a7257544..324500889ee 100644 --- a/encodings/alp/src/alp/compute/cast.rs +++ b/encodings/alp/src/alp/compute/cast.rs @@ -154,7 +154,7 @@ mod tests { let array_primitive = array.execute::(&mut ctx)?; let alp = alp_encode(array_primitive.as_view(), None, &mut ctx).vortex_expect("cannot fail"); - test_cast_conformance(&alp.into_array()); + test_cast_conformance(&alp.into_array(), &mut ctx); Ok(()) } diff --git a/encodings/alp/src/alp/compute/filter.rs b/encodings/alp/src/alp/compute/filter.rs index f38a87f19b0..8c7212cdeab 100644 --- a/encodings/alp/src/alp/compute/filter.rs +++ b/encodings/alp/src/alp/compute/filter.rs @@ -65,6 +65,6 @@ mod test { 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()); + test_filter_conformance(&alp.into_array(), &mut ctx); } } diff --git a/encodings/alp/src/alp/compute/mask.rs b/encodings/alp/src/alp/compute/mask.rs index 2a528987143..37f912e5eca 100644 --- a/encodings/alp/src/alp/compute/mask.rs +++ b/encodings/alp/src/alp/compute/mask.rs @@ -76,7 +76,7 @@ mod test { 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()); + test_mask_conformance(&alp.into_array(), &mut ctx); } #[test] @@ -90,7 +90,7 @@ mod test { let array = PrimitiveArray::from_iter(values); let alp = alp_encode(array.as_view(), None, &mut ctx).unwrap(); assert!(alp.patches().is_some(), "expected patches"); - test_mask_conformance(&alp.into_array()); + test_mask_conformance(&alp.into_array(), &mut ctx); } #[test] diff --git a/encodings/alp/src/alp/compute/take.rs b/encodings/alp/src/alp/compute/take.rs index 2e071eaa2aa..2c34b73c6e4 100644 --- a/encodings/alp/src/alp/compute/take.rs +++ b/encodings/alp/src/alp/compute/take.rs @@ -51,6 +51,6 @@ mod test { 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()); + test_take_conformance(&alp.into_array(), &mut ctx); } } diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index 28ca30f909c..dcaeb0fde79 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -159,6 +159,7 @@ impl VTable for ALPRD { )) } + #[allow(clippy::disallowed_methods)] fn deserialize( &self, dtype: &DType, @@ -470,6 +471,7 @@ fn patches_from_slots( PatchesData::patches_from_slots(patches_data, len, slots, LP_PATCH_SLOTS) } +#[allow(clippy::disallowed_methods)] fn validate_parts( dtype: &DType, len: usize, diff --git a/encodings/alp/src/alp_rd/compute/cast.rs b/encodings/alp/src/alp_rd/compute/cast.rs index 68ed39dd390..9a08f00e298 100644 --- a/encodings/alp/src/alp_rd/compute/cast.rs +++ b/encodings/alp/src/alp_rd/compute/cast.rs @@ -28,8 +28,6 @@ impl CastReduce for ALPRD { .with_nullability(dtype.nullability()), )?; - // NOTE: `CastReduce::cast` has a fixed trait signature without `ExecutionCtx`, so we - // construct a legacy ctx locally at this trait boundary. Ok(Some( unsafe { ALPRD::new_unchecked( @@ -149,6 +147,9 @@ mod tests { encoder.encode(arr.as_view()) })] fn test_cast_alprd_conformance(#[case] alprd: crate::alp_rd::ALPRDArray) { - test_cast_conformance(&alprd.into_array()); + test_cast_conformance( + &alprd.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/encodings/alp/src/alp_rd/compute/filter.rs b/encodings/alp/src/alp_rd/compute/filter.rs index 7aa8c247f2a..69c26c53e3b 100644 --- a/encodings/alp/src/alp_rd/compute/filter.rs +++ b/encodings/alp/src/alp_rd/compute/filter.rs @@ -86,10 +86,12 @@ mod test { #[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 = SESSION.create_execution_ctx(); test_filter_conformance( &RDEncoder::new(&[a, b]) .encode(PrimitiveArray::from_iter([a, b, outlier, b, outlier]).as_view()) .into_array(), + &mut ctx, ); } @@ -97,6 +99,7 @@ mod test { #[case(0.1f32, 3e25f32)] #[case(0.5f64, 1e100f64)] fn test_filter_with_nulls(#[case] a: T, #[case] outlier: T) { + let mut ctx = SESSION.create_execution_ctx(); test_filter_conformance( &RDEncoder::new(&[a]) .encode( @@ -104,6 +107,7 @@ mod test { .as_view(), ) .into_array(), + &mut ctx, ); } } diff --git a/encodings/alp/src/alp_rd/compute/mask.rs b/encodings/alp/src/alp_rd/compute/mask.rs index 2773387c738..25bdc8f9ab4 100644 --- a/encodings/alp/src/alp_rd/compute/mask.rs +++ b/encodings/alp/src/alp_rd/compute/mask.rs @@ -14,6 +14,7 @@ use crate::ALPRD; use crate::ALPRDArrayExt; impl MaskReduce for ALPRD { + #[allow(clippy::disallowed_methods)] fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { let masked_left_parts = MaskExpr.try_new_array( array.left_parts().len(), @@ -36,22 +37,35 @@ impl MaskReduce for ALPRD { #[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::PrimitiveArray; use vortex_array::compute::conformance::mask::test_mask_conformance; + use vortex_session::VortexSession; 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_mask_simple(#[case] a: T, #[case] b: T, #[case] outlier: T) { + let mut ctx = SESSION.create_execution_ctx(); test_mask_conformance( &RDEncoder::new(&[a, b]) .encode(PrimitiveArray::from_iter([a, b, outlier, b, outlier]).as_view()) .into_array(), + &mut ctx, ); } @@ -59,6 +73,7 @@ mod tests { #[case(0.1f32, 3e25f32)] #[case(0.5f64, 1e100f64)] fn test_mask_with_nulls(#[case] a: T, #[case] outlier: T) { + let mut ctx = SESSION.create_execution_ctx(); test_mask_conformance( &RDEncoder::new(&[a]) .encode( @@ -66,6 +81,7 @@ mod tests { .as_view(), ) .into_array(), + &mut ctx, ); } } diff --git a/encodings/alp/src/alp_rd/compute/take.rs b/encodings/alp/src/alp_rd/compute/take.rs index 8d12bfbb316..18d5e2b5ad4 100644 --- a/encodings/alp/src/alp_rd/compute/take.rs +++ b/encodings/alp/src/alp_rd/compute/take.rs @@ -142,10 +142,12 @@ 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 = SESSION.create_execution_ctx(); test_take_conformance( &RDEncoder::new(&[a, b]) .encode(PrimitiveArray::from_iter([a, b, outlier, b, outlier]).as_view()) .into_array(), + &mut ctx, ); } @@ -153,6 +155,7 @@ 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 = SESSION.create_execution_ctx(); test_take_conformance( &RDEncoder::new(&[a]) .encode( @@ -160,6 +163,7 @@ mod test { .as_view(), ) .into_array(), + &mut ctx, ); } } diff --git a/encodings/bytebool/Cargo.toml b/encodings/bytebool/Cargo.toml index bca2ed71aa9..53197452ff5 100644 --- a/encodings/bytebool/Cargo.toml +++ b/encodings/bytebool/Cargo.toml @@ -21,7 +21,6 @@ num-traits = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } -vortex-mask = { workspace = true } vortex-session = { workspace = true } [dev-dependencies] diff --git a/encodings/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs index 112abc921a4..7f9a4e803de 100644 --- a/encodings/bytebool/src/compute.rs +++ b/encodings/bytebool/src/compute.rs @@ -158,7 +158,12 @@ impl BooleanKernel for ByteBool { fn truthy_bit_buffer(array: ArrayView<'_, ByteBool>) -> BitBuffer { let bytes = array.truthy_bytes(); - BitBuffer::collect_bool(bytes.len(), |idx| bytes[idx] != 0) + // SAFETY: `collect_bool_multiversioned` invokes the predicate with indices `0..bytes.len()` only; + // the trivially cheap unchecked gather vectorizes inside the wide pack loop. + BitBuffer::collect_bool_multiversioned( + bytes.len(), + |idx| unsafe { *bytes.get_unchecked(idx) } != 0, + ) } #[cfg(test)] @@ -297,17 +302,25 @@ mod tests { #[test] fn test_mask_byte_bool() { - test_mask_conformance(&bb(vec![true, false, true, true, false]).into_array()); + test_mask_conformance( + &bb(vec![true, false, true, true, false]).into_array(), + &mut SESSION.create_execution_ctx(), + ); test_mask_conformance( &bb_opt(vec![Some(true), Some(true), None, Some(false), None]).into_array(), + &mut SESSION.create_execution_ctx(), ); } #[test] fn test_filter_byte_bool() { - test_filter_conformance(&bb(vec![true, false, true, true, false]).into_array()); + test_filter_conformance( + &bb(vec![true, false, true, true, false]).into_array(), + &mut SESSION.create_execution_ctx(), + ); test_filter_conformance( &bb_opt(vec![Some(true), Some(true), None, Some(false), None]).into_array(), + &mut SESSION.create_execution_ctx(), ); } @@ -317,7 +330,7 @@ mod tests { #[case(bb(vec![true, false]))] #[case(bb(vec![true]))] fn test_take_byte_bool_conformance(#[case] array: ByteBoolArray) { - test_take_conformance(&array.into_array()); + test_take_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] @@ -338,7 +351,7 @@ mod tests { #[case(bb(vec![true]))] #[case(bb_opt(vec![Some(true), None]))] fn test_cast_bytebool_conformance(#[case] array: ByteBoolArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[rstest] diff --git a/encodings/datetime-parts/src/compute/cast.rs b/encodings/datetime-parts/src/compute/cast.rs index 90aba343ae3..fbfb6f87e39 100644 --- a/encodings/datetime-parts/src/compute/cast.rs +++ b/encodings/datetime-parts/src/compute/cast.rs @@ -140,6 +140,9 @@ mod tests { ), &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()); + test_cast_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/encodings/datetime-parts/src/compute/filter.rs b/encodings/datetime-parts/src/compute/filter.rs index 696660510d5..6578467eddd 100644 --- a/encodings/datetime-parts/src/compute/filter.rs +++ b/encodings/datetime-parts/src/compute/filter.rs @@ -54,7 +54,7 @@ mod test { TemporalArray::new_timestamp(timestamps, TimeUnit::Milliseconds, Some("UTC".into())); let array = DateTimeParts::try_from_temporal(temporal, &mut ctx).unwrap(); - test_filter_conformance(&array.into_array()); + test_filter_conformance(&array.into_array(), &mut ctx); // Test with nullable values let timestamps = PrimitiveArray::from_option_iter([ @@ -70,6 +70,6 @@ mod test { TemporalArray::new_timestamp(timestamps, TimeUnit::Milliseconds, Some("UTC".into())); let array = DateTimeParts::try_from_temporal(temporal, &mut ctx).unwrap(); - test_filter_conformance(&array.into_array()); + test_filter_conformance(&array.into_array(), &mut ctx); } } diff --git a/encodings/datetime-parts/src/compute/take.rs b/encodings/datetime-parts/src/compute/take.rs index 9f5a50d3a9b..2b49cfaa15d 100644 --- a/encodings/datetime-parts/src/compute/take.rs +++ b/encodings/datetime-parts/src/compute/take.rs @@ -136,6 +136,9 @@ mod tests { Some("UTC".into()) ), &mut array_session().create_execution_ctx()).unwrap())] fn test_take_datetime_parts_conformance(#[case] array: DateTimePartsArray) { - test_take_conformance(&array.into_array()); + test_take_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } 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 376ad12f564..303e9d3247b 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 @@ -118,6 +118,9 @@ mod tests { DecimalDType::new(10, 2), ).unwrap())] fn test_cast_decimal_byte_parts_conformance(#[case] array: DecimalBytePartsArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } 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 04f321cf88f..af6e98bddc2 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 @@ -47,10 +47,7 @@ impl CompareKernel for DecimalByteParts { .decimal_value() .vortex_expect("checked for null in entry func"); - match decimal_value_wrapper_to_primitive( - rhs_decimal, - lhs.msp().as_primitive_typed().ptype(), - ) { + match decimal_value_wrapper_to_primitive(rhs_decimal, lhs.msp().dtype().as_ptype()) { Ok(value) => { let encoded_scalar = Scalar::try_new(scalar_type, Some(value))?; let encoded_const = ConstantArray::new(encoded_scalar, rhs.len()); diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs index c82171b0e13..a1b50ad788a 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs @@ -27,6 +27,8 @@ impl FilterReduce for DecimalByteParts { #[cfg(test)] mod test { use vortex_array::IntoArray; + 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_array::dtype::DecimalDType; @@ -41,7 +43,10 @@ mod test { let decimal_dtype = DecimalDType::new(8, 2); let array = DecimalByteParts::try_new(msp, decimal_dtype).unwrap(); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); // Test with nullable values let msp = PrimitiveArray::from_option_iter([Some(10i64), None, Some(30), Some(40), None]) @@ -49,6 +54,9 @@ mod test { let decimal_dtype = DecimalDType::new(18, 4); let array = DecimalByteParts::try_new(msp, decimal_dtype).unwrap(); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/encodings/experimental/onpair/Cargo.toml b/encodings/experimental/onpair/Cargo.toml index 258568cc75a..2f6b19147fb 100644 --- a/encodings/experimental/onpair/Cargo.toml +++ b/encodings/experimental/onpair/Cargo.toml @@ -17,7 +17,6 @@ version = { workspace = true } workspace = true [dependencies] -memchr = { workspace = true } num-traits = { workspace = true } onpair = { workspace = true } prost = { workspace = true } diff --git a/encodings/experimental/onpair/src/array.rs b/encodings/experimental/onpair/src/array.rs index 58fd0332ac3..519243b100e 100644 --- a/encodings/experimental/onpair/src/array.rs +++ b/encodings/experimental/onpair/src/array.rs @@ -493,14 +493,12 @@ impl VTable for OnPair { ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let Some(builder) = builder.as_any_mut().downcast_mut::() else { - builder.extend_from_array( - &array - .array() - .clone() - .execute::(ctx)? - .into_array(), - ); - return Ok(()); + return array + .array() + .clone() + .execute::(ctx)? + .into_array() + .append_to_builder(builder, ctx); }; let next_buffer_index = builder.completed_block_count() + u32::from(builder.in_progress()); diff --git a/encodings/experimental/onpair/src/compress.rs b/encodings/experimental/onpair/src/compress.rs index adf9e73461c..3616ec7e717 100644 --- a/encodings/experimental/onpair/src/compress.rs +++ b/encodings/experimental/onpair/src/compress.rs @@ -27,7 +27,10 @@ use crate::OnPair; use crate::OnPairArray; /// Default OnPair training configuration: 12-bit codes ("dict-12"). -pub const DEFAULT_DICT12_CONFIG: Config = onpair::DEFAULT_CONFIG; +pub const DEFAULT_DICT12_CONFIG: Config = Config { + seed: Some(42), + ..onpair::DEFAULT_CONFIG +}; fn onpair_compress_varbinview( array: VarBinViewArray, diff --git a/encodings/fastlanes/src/bit_transpose/mod.rs b/encodings/fastlanes/src/bit_transpose/mod.rs index 99b01bc242c..26e8b595195 100644 --- a/encodings/fastlanes/src/bit_transpose/mod.rs +++ b/encodings/fastlanes/src/bit_transpose/mod.rs @@ -28,6 +28,10 @@ mod x86; mod validity; pub use validity::*; +use vortex_buffer::CpuKernel; + +/// Signature shared by all 1024-bit transpose kernels. +type TransposeKernel = unsafe fn(&[u8; 128], &mut [u8; 128]); /// Base indices for the first 64 output bytes (lanes 0-7). /// Each entry indicates the starting input byte index for that output byte group. @@ -45,56 +49,64 @@ const TRANSPOSE_8X8: u64 = 0x0000_0000_F0F0_F0F0; /// Transpose 1024-bits into FastLanes layout. /// -/// Dispatch to the best available implementation at runtime. +/// Dispatch to the best available implementation, selected once on first call. #[inline] pub fn transpose_bits(input: &[u8; 128], output: &mut [u8; 128]) { - #[cfg(target_arch = "x86_64")] - { - // VBMI is fastest - if x86::has_vbmi() { - unsafe { x86::transpose_bits_vbmi(input, output) }; - return; + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + // VBMI is fastest + if x86::has_vbmi() { + return x86::transpose_bits_vbmi; + } + if x86::has_bmi2() { + return x86::transpose_bits_bmi2; + } } - if x86::has_bmi2() { - unsafe { x86::transpose_bits_bmi2(input, output) }; - return; + // NEON is architecturally guaranteed on aarch64, so it needs no probe. + #[cfg(target_arch = "aarch64")] + return aarch64::transpose_bits_neon; + // The aarch64 arm above returns unconditionally, making this portable default + // unreachable there. + #[allow(unreachable_code)] + { + scalar::transpose_bits_scalar } - // Fall back to scalar - scalar::transpose_bits_scalar(input, output); - } - #[cfg(target_arch = "aarch64")] - { - unsafe { aarch64::transpose_bits_neon(input, output) }; - } - #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))] - scalar::transpose_bits_scalar(input, output); + }); + // SAFETY: the selector only returns kernels that are safe or whose required CPU + // features were probed before selection. + unsafe { KERNEL.get()(input, output) } } /// Untranspose 1024-bits from FastLanes layout. /// -/// Dispatch untranspose to the best available implementation at runtime. +/// Dispatch untranspose to the best available implementation, selected once on first call. #[inline] pub fn untranspose_bits(input: &[u8; 128], output: &mut [u8; 128]) { - #[cfg(target_arch = "x86_64")] - { - // VBMI is fastest - if x86::has_vbmi() { - unsafe { x86::untranspose_bits_vbmi(input, output) }; - return; + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + // VBMI is fastest + if x86::has_vbmi() { + return x86::untranspose_bits_vbmi; + } + if x86::has_bmi2() { + return x86::untranspose_bits_bmi2; + } } - if x86::has_bmi2() { - unsafe { x86::untranspose_bits_bmi2(input, output) }; - return; + // NEON is architecturally guaranteed on aarch64, so it needs no probe. + #[cfg(target_arch = "aarch64")] + return aarch64::untranspose_bits_neon; + // The aarch64 arm above returns unconditionally, making this portable default + // unreachable there. + #[allow(unreachable_code)] + { + scalar::untranspose_bits_scalar } - // Fall back to scalar - scalar::untranspose_bits_scalar(input, output); - } - #[cfg(target_arch = "aarch64")] - { - unsafe { aarch64::untranspose_bits_neon(input, output) }; - } - #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))] - scalar::untranspose_bits_scalar(input, output); + }); + // SAFETY: the selector only returns kernels that are safe or whose required CPU + // features were probed before selection. + unsafe { KERNEL.get()(input, output) } } #[cfg(test)] diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index f7a3485c113..a393db6ecc8 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -521,7 +521,7 @@ mod test { let mut primitive_builder = PrimitiveBuilder::::with_capacity(chunked.dtype().nullability(), 10 * 100); - primitive_builder.extend_from_array(&chunked); + chunked.append_to_builder(&mut primitive_builder, &mut ctx)?; let ca_into = primitive_builder.finish(); assert_arrays_eq!(into_ca, ca_into, &mut ctx); diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index 78aac1aa86e..692e7dcdd7f 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -203,13 +203,20 @@ pub fn count_exceptions(bit_width: u8, bit_width_freq: &[usize]) -> usize { #[cfg(test)] mod tests { + use std::sync::Arc; use std::sync::LazyLock; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::ListViewArray; use vortex_array::assert_arrays_eq; + use vortex_array::builders::ListBuilder; + use vortex_array::builders::ListViewBuilder; + use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -266,6 +273,64 @@ mod tests { compression_roundtrip(10_240); } + /// List builders bulk-append the source's elements into their internal elements builder. + /// Fixed-width builder appends assume the caller reserved capacity, so the list builders + /// must grow the elements builder themselves; encoded elements exercise that because + /// BitPacked's `append_to_builder` writes through `uninit_range`. + #[test] + fn test_list_append_grows_elements_builder() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + let elements = encode(&PrimitiveArray::from_iter((0..3072u32).map(|i| i % 256)), 8); + let element_dtype: Arc = Arc::new(PType::U32.into()); + + // 48 lists of 64 elements each. + let offsets = Buffer::from_iter((0..=48u64).map(|i| i * 64)).into_array(); + let list = ListArray::try_new( + elements.clone().into_array(), + offsets, + Validity::NonNullable, + )?; + + let mut listview_builder = ListViewBuilder::::with_capacity( + Arc::clone(&element_dtype), + Nullability::NonNullable, + 0, + 0, + ); + list.clone() + .into_array() + .append_to_builder(&mut listview_builder, &mut ctx)?; + assert_arrays_eq!(listview_builder.finish(), list, &mut ctx); + + let mut list_builder = ListBuilder::::with_capacity( + Arc::clone(&element_dtype), + Nullability::NonNullable, + 0, + 0, + ); + list.clone() + .into_array() + .append_to_builder(&mut list_builder, &mut ctx)?; + assert_arrays_eq!(list_builder.finish(), list, &mut ctx); + + // A `ListViewArray` source appended into a `ListBuilder` appends elements list by list. + let listview = ListViewArray::try_new( + elements.into_array(), + Buffer::from_iter((0..48u64).map(|i| i * 64)).into_array(), + Buffer::from_iter(std::iter::repeat_n(64u32, 48)).into_array(), + Validity::NonNullable, + )?; + let mut list_builder = + ListBuilder::::with_capacity(element_dtype, Nullability::NonNullable, 0, 0); + listview + .into_array() + .append_to_builder(&mut list_builder, &mut ctx)?; + assert_arrays_eq!(list_builder.finish(), list, &mut ctx); + + Ok(()) + } + #[test] fn test_all_zeros() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index 3b917aa36af..f7cfc56acf0 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -266,6 +266,6 @@ mod tests { #[case(bp(&buffer![0u32, 1000, 2000, 3000, 4000].into_array(), 12))] #[case(bp(&PrimitiveArray::from_option_iter([Some(1u32), None, Some(7), Some(15), None]).into_array(), 4))] fn test_cast_bitpacked_conformance(#[case] array: BitPackedArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 21184d785a5..1530c33874d 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -273,17 +273,17 @@ mod test { // Test with u8 values let unpacked = buffer![1u8, 2, 3, 4, 5].into_array(); let bitpacked = BitPackedData::encode(&unpacked, 3, &mut ctx).unwrap(); - test_filter_conformance(&bitpacked.into_array()); + test_filter_conformance(&bitpacked.into_array(), &mut ctx); // Test with u32 values let unpacked = buffer![100u32, 200, 300, 400, 500].into_array(); let bitpacked = BitPackedData::encode(&unpacked, 9, &mut ctx).unwrap(); - test_filter_conformance(&bitpacked.into_array()); + test_filter_conformance(&bitpacked.into_array(), &mut ctx); // Test with nullable values let unpacked = PrimitiveArray::from_option_iter([Some(1u16), None, Some(3), Some(4), None]); let bitpacked = BitPackedData::encode(&unpacked.into_array(), 3, &mut ctx).unwrap(); - test_filter_conformance(&bitpacked.into_array()); + test_filter_conformance(&bitpacked.into_array(), &mut ctx); } /// Regression test for signed integers with patches. diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 34aa5a0c3f5..77f6b9d3127 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -169,6 +169,7 @@ mod test { use vortex_array::VortexSessionExecute; 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_buffer::buffer; @@ -327,7 +328,6 @@ mod test { #[case(bp(buffer![42u32].into_array(), 6))] #[case(bp(PrimitiveArray::from_iter((0..1024).map(|i| i as u32)).into_array(), 8))] fn test_take_bitpacked_conformance(#[case] bitpacked: BitPackedArray) { - use vortex_array::compute::conformance::take::test_take_conformance; - test_take_conformance(&bitpacked.into_array()); + test_take_conformance(&bitpacked.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/fastlanes/src/delta/compute/cast.rs b/encodings/fastlanes/src/delta/compute/cast.rs index 2ee10e7619a..c6de602bb0b 100644 --- a/encodings/fastlanes/src/delta/compute/cast.rs +++ b/encodings/fastlanes/src/delta/compute/cast.rs @@ -244,6 +244,9 @@ mod tests { let delta_array = Delta::try_from_primitive_array(&primitive, &mut SESSION.create_execution_ctx()) .unwrap(); - test_cast_conformance(&delta_array.into_array()); + test_cast_conformance( + &delta_array.into_array(), + &mut SESSION.create_execution_ctx(), + ); } } diff --git a/encodings/fastlanes/src/delta/vtable/validity.rs b/encodings/fastlanes/src/delta/vtable/validity.rs index d6a87ff5956..c8d8b7b4dc0 100644 --- a/encodings/fastlanes/src/delta/vtable/validity.rs +++ b/encodings/fastlanes/src/delta/vtable/validity.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::ArrayView; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; +use vortex_array::legacy_session; use vortex_array::validity::Validity; use vortex_array::vtable::ValidityVTable; use vortex_error::VortexResult; @@ -13,13 +13,14 @@ use crate::bit_transpose::untranspose_validity; use crate::delta::array::DeltaArrayExt; impl ValidityVTable for Delta { + #[allow(clippy::disallowed_methods)] fn validity(array: ArrayView<'_, Delta>) -> VortexResult { let start = array.offset(); let end = start + array.len(); let validity = untranspose_validity( &array.deltas().validity()?, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut legacy_session().create_execution_ctx(), )?; validity.slice(start..end) } diff --git a/encodings/fastlanes/src/for/compute/cast.rs b/encodings/fastlanes/src/for/compute/cast.rs index b865af43d12..e56188d9e00 100644 --- a/encodings/fastlanes/src/for/compute/cast.rs +++ b/encodings/fastlanes/src/for/compute/cast.rs @@ -118,6 +118,6 @@ mod tests { Scalar::from(-100i32) ))] fn test_cast_for_conformance(#[case] array: FoRArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/fastlanes/src/for/compute/mod.rs b/encodings/fastlanes/src/for/compute/mod.rs index 84e8812b598..168fcbc4192 100644 --- a/encodings/fastlanes/src/for/compute/mod.rs +++ b/encodings/fastlanes/src/for/compute/mod.rs @@ -49,8 +49,11 @@ mod test { 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::filter::test_filter_conformance; + use vortex_array::compute::conformance::take::test_take_conformance; use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -68,18 +71,27 @@ mod test { buffer![100i32, 101, 102, 103, 104].into_array(), Scalar::from(100i32), ); - test_filter_conformance(&for_array.into_array()); + test_filter_conformance( + &for_array.into_array(), + &mut array_session().create_execution_ctx(), + ); let for_array = fa( buffer![1000u64, 1001, 1002, 1003, 1004].into_array(), Scalar::from(1000u64), ); - test_filter_conformance(&for_array.into_array()); + test_filter_conformance( + &for_array.into_array(), + &mut array_session().create_execution_ctx(), + ); let values = PrimitiveArray::from_option_iter([Some(50i16), None, Some(52), Some(53), None]); let for_array = fa(values.into_array(), Scalar::from(50i16)); - test_filter_conformance(&for_array.into_array()); + test_filter_conformance( + &for_array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[rstest] @@ -92,8 +104,10 @@ mod test { #[case(fa(buffer![-100i32, -99, -98, -97, -96].into_array(), Scalar::from(-100i32)))] #[case(fa(buffer![42i64].into_array(), Scalar::from(40i64)))] fn test_take_for_conformance(#[case] for_array: FoRArray) { - use vortex_array::compute::conformance::take::test_take_conformance; - test_take_conformance(&for_array.into_array()); + test_take_conformance( + &for_array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/encodings/fastlanes/src/rle/compute/cast.rs b/encodings/fastlanes/src/rle/compute/cast.rs index 92947d927bc..38e7927f8e6 100644 --- a/encodings/fastlanes/src/rle/compute/cast.rs +++ b/encodings/fastlanes/src/rle/compute/cast.rs @@ -162,6 +162,6 @@ mod tests { fn test_cast_rle_conformance(#[case] primitive: PrimitiveArray) { let mut ctx = SESSION.create_execution_ctx(); let rle_array = rle(&primitive, &mut ctx); - test_cast_conformance(&rle_array.into_array()); + test_cast_conformance(&rle_array.into_array(), &mut ctx); } } diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 79e6dcb2f0c..de0cf3157ba 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -25,7 +25,6 @@ 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::VarBin; @@ -37,6 +36,7 @@ use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::legacy_session; use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; @@ -114,6 +114,7 @@ impl VTable for FSST { *ID } + #[allow(clippy::disallowed_methods)] fn validate( &self, data: &Self::TypedArrayData, @@ -122,7 +123,7 @@ impl VTable for FSST { slots: &[Option], ) -> VortexResult<()> { // TODO(ctx): trait fixes - VTable::validate has a fixed signature. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = legacy_session().create_execution_ctx(); data.validate(dtype, len, slots, &mut ctx) } @@ -304,14 +305,12 @@ impl VTable for FSST { ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let Some(builder) = builder.as_any_mut().downcast_mut::() else { - builder.extend_from_array( - &array - .array() - .clone() - .execute::(ctx)? - .into_array(), - ); - return Ok(()); + return array + .array() + .clone() + .execute::(ctx)? + .into_array() + .append_to_builder(builder, ctx); }; // Decompress the whole block of data into a new buffer, and create some views diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 43712aad7dd..5c58dbbd501 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -180,7 +180,7 @@ mod tests { .append_to_builder(&mut builder, &mut ctx)?; { - let arr = builder.finish_into_canonical().into_varbinview(); + let arr = builder.finish_into_canonical(&mut ctx).into_varbinview(); 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())) diff --git a/encodings/fsst/src/compute/cast.rs b/encodings/fsst/src/compute/cast.rs index bdca68e841b..a1469681590 100644 --- a/encodings/fsst/src/compute/cast.rs +++ b/encodings/fsst/src/compute/cast.rs @@ -147,7 +147,7 @@ mod tests { 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()); + test_cast_conformance(&fsst.into_array(), &mut ctx); Ok(()) } } diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index 0065e0f499a..d447b4b4958 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -117,7 +117,7 @@ mod tests { 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()); + test_take_conformance(&array.into_array(), &mut ctx); Ok(()) } diff --git a/encodings/parquet-variant/Cargo.toml b/encodings/parquet-variant/Cargo.toml index d0c0a0d418e..6cdc2b3fa8b 100644 --- a/encodings/parquet-variant/Cargo.toml +++ b/encodings/parquet-variant/Cargo.toml @@ -25,6 +25,7 @@ parquet-variant = { workspace = true } parquet-variant-compute = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } +vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } vortex-json = { workspace = true } diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index 7e55357c78f..bc7fd2c22ec 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -24,13 +24,6 @@ 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" -)] -use vortex_array::arrow::ArrowArrayExecutor; -use vortex_array::arrow::FromArrowArray; -use vortex_array::arrow::to_arrow_null_buffer; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; @@ -41,6 +34,13 @@ use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; +#[expect( + deprecated, + reason = "TODO(aduffy): figure out what to do with Parquet Variant" +)] +use vortex_arrow::ArrowArrayExecutor; +use vortex_arrow::FromArrowArray; +use vortex_arrow::to_arrow_null_buffer; use vortex_buffer::BitBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; diff --git a/encodings/parquet-variant/src/arrow.rs b/encodings/parquet-variant/src/arrow.rs index da4e648a138..65efd6c571c 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -20,14 +20,14 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::Variant; use vortex_array::arrays::variant::VariantArrayExt; -use vortex_array::arrow::ArrowExport; -use vortex_array::arrow::ArrowExportVTable; -use vortex_array::arrow::ArrowImport; -use vortex_array::arrow::ArrowImportVTable; -use vortex_array::arrow::ArrowSession; -use vortex_array::arrow::ArrowSessionExt; -use vortex_array::arrow::to_arrow_null_buffer; use vortex_array::dtype::DType; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::to_arrow_null_buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -300,11 +300,11 @@ mod tests { use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::VariantArray; - use vortex_array::arrow::ArrowSessionExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::validity::Validity; + use vortex_arrow::ArrowSessionExt; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index bd7d0bfaa7c..c783e6600ac 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -37,8 +37,6 @@ 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::optimizer::kernels::ArrayKernelsExt; @@ -46,6 +44,9 @@ use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::variant_get::VariantGet; use vortex_array::scalar_fn::fns::variant_get::VariantPath; use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; +use vortex_arrow::ToArrowType; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; @@ -318,7 +319,6 @@ mod tests { use vortex_array::arrays::VariantArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::variant::VariantArrayExt; - use vortex_array::arrow::FromArrowArray; use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar_is_null; use vortex_array::dtype::DType as VortexDType; @@ -329,6 +329,7 @@ mod tests { use vortex_array::scalar_fn::fns::variant_get::VariantPath; use vortex_array::scalar_fn::fns::variant_get::VariantPathElement; use vortex_array::validity::Validity; + use vortex_arrow::FromArrowArray; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; diff --git a/encodings/parquet-variant/src/lib.rs b/encodings/parquet-variant/src/lib.rs index c71ea81e734..88125af04fd 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/encodings/parquet-variant/src/lib.rs @@ -37,8 +37,8 @@ mod vtable; use std::sync::Arc; pub use array::ParquetVariantArrayExt; -use vortex_array::arrow::ArrowSessionExt; use vortex_array::session::ArraySessionExt; +use vortex_arrow::ArrowSessionExt; pub use vortex_json::JsonToVariant; pub use vortex_json::JsonToVariantOptions; pub use vortex_json::ShreddingSpec; diff --git a/encodings/pco/Cargo.toml b/encodings/pco/Cargo.toml index 8d560d5f068..1683d4875e9 100644 --- a/encodings/pco/Cargo.toml +++ b/encodings/pco/Cargo.toml @@ -28,4 +28,5 @@ vortex-session = { workspace = true } [dev-dependencies] rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-arrow = { workspace = true } diff --git a/encodings/pco/src/compute/cast.rs b/encodings/pco/src/compute/cast.rs index 5f562fafca6..c59e205e2e1 100644 --- a/encodings/pco/src/compute/cast.rs +++ b/encodings/pco/src/compute/cast.rs @@ -188,6 +188,6 @@ mod tests { fn test_cast_pco_conformance(#[case] values: PrimitiveArray) { let mut ctx = SESSION.create_execution_ctx(); let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx).unwrap(); - test_cast_conformance(&pco.into_array()); + test_cast_conformance(&pco.into_array(), &mut ctx); } } diff --git a/encodings/pco/src/tests.rs b/encodings/pco/src/tests.rs index 573bd6a0fb4..5cc771601f0 100644 --- a/encodings/pco/src/tests.rs +++ b/encodings/pco/src/tests.rs @@ -9,7 +9,6 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrow::ArrowSessionExt; use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar; use vortex_array::dtype::DType; @@ -20,6 +19,7 @@ use vortex_array::serde::SerializedArray; use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; use vortex_array::vtable::child_to_validity; +use vortex_arrow::ArrowSessionExt; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; diff --git a/encodings/runend/Cargo.toml b/encodings/runend/Cargo.toml index 268aa6ac3ca..4d36d0c1548 100644 --- a/encodings/runend/Cargo.toml +++ b/encodings/runend/Cargo.toml @@ -15,7 +15,6 @@ version = { workspace = true } [dependencies] arbitrary = { workspace = true, optional = true } -arrow-array = { workspace = true, optional = true } itertools = { workspace = true } num-traits = { workspace = true } prost = { workspace = true } @@ -29,8 +28,6 @@ vortex-session = { workspace = true } workspace = true [dev-dependencies] -arrow-array = { workspace = true } -arrow-schema = { workspace = true } divan = { workspace = true } itertools = { workspace = true } rand = { workspace = true } @@ -39,7 +36,6 @@ vortex-array = { workspace = true, features = ["_test-harness"] } [features] arbitrary = ["dep:arbitrary", "vortex-array/arbitrary"] -arrow = ["dep:arrow-array"] [[bench]] name = "run_end_null_count" @@ -56,3 +52,7 @@ harness = false [[bench]] name = "run_end_take" harness = false + +[[bench]] +name = "run_end_filter" +harness = false diff --git a/encodings/runend/benches/run_end_filter.rs b/encodings/runend/benches/run_end_filter.rs new file mode 100644 index 00000000000..395b41063b1 --- /dev/null +++ b/encodings/runend/benches/run_end_filter.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks for the run-end filter inner loop (`filter_run_end_primitive`). +//! +//! This measures the kernel directly rather than going through the lazy +//! `ArrayRef::filter` (which only builds a `FilterArray` node and does not run +//! the kernel). The hot work is a per-run popcount of the predicate mask, which +//! now uses `BitBuffer::count_range` (SIMD) instead of a bit-by-bit walk. + +#![expect(clippy::cast_possible_truncation)] +#![expect(clippy::cast_precision_loss)] +#![expect(clippy::cast_sign_loss)] +#![expect(clippy::expect_used)] + +use std::fmt; + +use divan::Bencher; +use rand::SeedableRng; +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use vortex_buffer::BitBuffer; +use vortex_runend::_benchmarking::filter_run_end_primitive; + +fn main() { + divan::main(); +} + +#[derive(Clone, Copy)] +struct FilterBenchArgs { + /// Total logical length of the decoded array. + length: usize, + /// Average run length used when building the run-end array. + run_length: usize, + /// Fraction of mask bits that are set to `true`. + density: f64, +} + +impl fmt::Display for FilterBenchArgs { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "len={}_run={}_density={:.1}", + self.length, self.run_length, self.density + ) + } +} + +const FILTER_ARGS: &[FilterBenchArgs] = &[ + FilterBenchArgs { + length: 4_096, + run_length: 16, + density: 0.1, + }, + FilterBenchArgs { + length: 4_096, + run_length: 16, + density: 0.5, + }, + FilterBenchArgs { + length: 4_096, + run_length: 16, + density: 0.9, + }, + FilterBenchArgs { + length: 16_384, + run_length: 16, + density: 0.1, + }, + FilterBenchArgs { + length: 16_384, + run_length: 16, + density: 0.5, + }, + FilterBenchArgs { + length: 16_384, + run_length: 16, + density: 0.9, + }, +]; + +/// Build the run-end boundaries (cumulative run lengths) for `length` rows. +fn build_run_ends(length: usize, run_length: usize) -> Vec { + let n_runs = length.div_ceil(run_length); + (0..n_runs) + .map(|r| (((r + 1) * run_length).min(length)) as u32) + .collect() +} + +/// Build a predicate mask of `length` bits with approximately `density` set bits, +/// shuffled so the set bits are spread across runs. +fn build_mask(length: usize, density: f64) -> BitBuffer { + let n_true = (length as f64 * density).round() as usize; + let mut bits = vec![false; length]; + for b in bits.iter_mut().take(n_true) { + *b = true; + } + let mut rng = StdRng::seed_from_u64(0x5eed); + bits.shuffle(&mut rng); + BitBuffer::from(bits) +} + +#[divan::bench(args = FILTER_ARGS)] +fn filter_run_end(bencher: Bencher, args: FilterBenchArgs) { + let run_ends = build_run_ends(args.length, args.run_length); + let mask = build_mask(args.length, args.density); + let length = args.length as u64; + bencher + .with_inputs(|| (run_ends.clone(), mask.clone())) + .bench_refs(|(run_ends, mask)| { + filter_run_end_primitive::(run_ends, 0, length, mask).expect("filter") + }); +} diff --git a/encodings/runend/src/arbitrary.rs b/encodings/runend/src/arbitrary.rs index baa15e0fd13..d913de378be 100644 --- a/encodings/runend/src/arbitrary.rs +++ b/encodings/runend/src/arbitrary.rs @@ -4,9 +4,9 @@ use arbitrary::Arbitrary; use arbitrary::Result; use arbitrary::Unstructured; +use arbitrary::unstructured::Int; +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::arbitrary::ArbitraryArray; use vortex_array::arrays::arbitrary::ArbitraryArrayConfig; @@ -14,6 +14,7 @@ use vortex_array::arrays::arbitrary::ArbitraryWith; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::dtype::UnsignedPType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexExpect; @@ -31,7 +32,7 @@ impl<'a> Arbitrary<'a> for ArbitraryRunEndArray { let ptype: PType = u.arbitrary()?; let nullability: Nullability = u.arbitrary()?; let dtype = DType::Primitive(ptype, nullability); - Self::with_dtype(u, &dtype, None) + Self::with_dtype(u, &dtype) } } @@ -39,15 +40,16 @@ impl ArbitraryRunEndArray { /// Generate an arbitrary RunEndArray with the given dtype for values. /// /// The dtype must be a primitive or boolean type. - pub fn with_dtype(u: &mut Unstructured, dtype: &DType, len: Option) -> Result { + pub fn with_dtype(u: &mut Unstructured, dtype: &DType) -> Result { // Number of runs (values/ends pairs) let num_runs = u.int_in_range(0..=20)?; - // TODO(ctx): trait fixes - Arbitrary::arbitrary has a fixed signature. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); if num_runs == 0 { // Empty RunEndArray - let ends = PrimitiveArray::from_iter(Vec::::new()).into_array(); + let ends = unsafe { + PrimitiveArray::new_unchecked(Buffer::::empty(), Validity::NonNullable) + .into_array() + }; let values = ArbitraryArray::arbitrary_with_config( u, &ArbitraryArrayConfig { @@ -56,8 +58,7 @@ impl ArbitraryRunEndArray { }, )? .0; - let runend_array = RunEnd::try_new(ends, values, &mut ctx) - .vortex_expect("Empty RunEndArray creation should succeed"); + let runend_array = unsafe { RunEnd::new_unchecked(ends, values, 0, 0) }; return Ok(ArbitraryRunEndArray(runend_array)); } @@ -71,12 +72,12 @@ impl ArbitraryRunEndArray { )? .0; + let len = u.int_in_range(0..=2048)?; // Generate strictly increasing ends // Each end must be > previous end, and first end must be >= 1 let ends = random_strictly_sorted_ends(u, num_runs, len)?; - let runend_array = RunEnd::try_new(ends, values, &mut ctx) - .vortex_expect("RunEndArray creation should succeed in arbitrary impl"); + let runend_array = unsafe { RunEnd::new_unchecked(ends, values, 0, len) }; Ok(ArbitraryRunEndArray(runend_array)) } @@ -89,67 +90,75 @@ impl ArbitraryRunEndArray { fn random_strictly_sorted_ends( u: &mut Unstructured, num_runs: usize, - target_len: Option, -) -> Result { + target_len: usize, +) -> Result { // Choose a random unsigned PType for ends - let ends_ptype = *u.choose(&[PType::U8, PType::U16, PType::U32, PType::U64])?; + let mut ends_ptypes = vec![PType::U8, PType::U16, PType::U32, PType::U64]; + if target_len >= u8::MAX as usize { + ends_ptypes.remove(0); + } + if target_len >= u16::MAX as usize { + ends_ptypes.remove(0); + } + if target_len >= u32::MAX as usize { + ends_ptypes.remove(0); + } + let ends_ptype = *u.choose(&ends_ptypes)?; + match ends_ptype { + PType::U8 => random_strictly_sorted( + u, + num_runs, + u8::try_from(target_len).vortex_expect("must fit in u8"), + ), + PType::U16 => random_strictly_sorted( + u, + num_runs, + u16::try_from(target_len).vortex_expect("must fit in u16"), + ), + PType::U32 => random_strictly_sorted( + u, + num_runs, + u32::try_from(target_len).vortex_expect("must fit in u32"), + ), + PType::U64 => random_strictly_sorted( + u, + num_runs, + u64::try_from(target_len).vortex_expect("must fit in u64"), + ), + _ => unreachable!("Only unsigned integer types are valid for ends"), + } +} + +fn random_strictly_sorted( + u: &mut Unstructured, + num_runs: usize, + target: T, +) -> Result { // Generate strictly increasing values // Start from 0, increment by at least 1 each time - let mut ends: Vec = Vec::with_capacity(num_runs); - let mut current: u64 = 0; + let mut ends: Vec = Vec::with_capacity(num_runs); + let mut current = T::zero(); for i in 0..num_runs { // Each run must have at least length 1, so increment by at least 1 - let increment = match (i == num_runs - 1, target_len) { - (true, Some(target)) => { + let increment = match i == num_runs - 1 { + true => { // Last element should reach target_len - let target = target as u64; if target > current { target - current } else { - 1 + T::one() } } - _ => { + false => { // Random increment between 1 and 10 - u.int_in_range(1..=10)? + u.int_in_range(T::one()..=T::from(10).vortex_expect("10 will fit in all T"))? } }; current += increment; ends.push(current); } - // Convert to the chosen PType - // The values are bounded: max is num_runs (20) * max_increment (10) = 200 - // This fits in all unsigned types - let ends_array = match ends_ptype { - PType::U8 => { - let ends_typed: Vec = ends - .iter() - .map(|&e| u8::try_from(e).vortex_expect("end value fits in u8")) - .collect(); - PrimitiveArray::new(Buffer::copy_from(ends_typed), Validity::NonNullable).into_array() - } - PType::U16 => { - let ends_typed: Vec = ends - .iter() - .map(|&e| u16::try_from(e).vortex_expect("end value fits in u16")) - .collect(); - PrimitiveArray::new(Buffer::copy_from(ends_typed), Validity::NonNullable).into_array() - } - PType::U32 => { - let ends_typed: Vec = ends - .iter() - .map(|&e| u32::try_from(e).vortex_expect("end value fits in u32")) - .collect(); - PrimitiveArray::new(Buffer::copy_from(ends_typed), Validity::NonNullable).into_array() - } - PType::U64 => { - PrimitiveArray::new(Buffer::copy_from(ends), Validity::NonNullable).into_array() - } - _ => unreachable!("Only unsigned integer types are valid for ends"), - }; - - Ok(ends_array) + Ok(PrimitiveArray::new(Buffer::copy_from(ends), Validity::NonNullable).into_array()) } diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 64d0798ebee..2a9659efdd8 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -19,7 +19,6 @@ 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; @@ -28,9 +27,7 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::scalar::PValue; -use vortex_array::search_sorted::SearchSorted; -use vortex_array::search_sorted::SearchSortedSide; +use vortex_array::legacy_session; use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; @@ -48,6 +45,8 @@ 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::ops::find_physical_index; +use crate::ops::find_slice_end_index; use crate::rules::RULES; /// A [`RunEnd`]-encoded Vortex array. @@ -86,6 +85,7 @@ impl VTable for RunEnd { *ID } + #[allow(clippy::disallowed_methods)] fn validate( &self, data: &Self::TypedArrayData, @@ -100,7 +100,7 @@ impl VTable for RunEnd { .as_ref() .vortex_expect("RunEndArray values slot"); // TODO(ctx): trait fixes - VTable::validate has a fixed signature. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = legacy_session().create_execution_ctx(); RunEndData::validate_parts(ends, values, data.offset, len, &mut ctx)?; vortex_ensure!( values.dtype() == dtype, @@ -229,17 +229,15 @@ pub trait RunEndArrayExt: TypedArrayRef { self.values().dtype() } - fn find_physical_index(&self, index: usize) -> VortexResult { - Ok(self - .ends() - .as_primitive_typed() - .search_sorted( - &PValue::from(index + self.offset()), - SearchSortedSide::Right, - )? - .to_ends_index(self.ends().len())) + fn find_physical_index(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + find_physical_index(self.ends(), index + self.offset(), ctx) + } + + fn find_slice_end_index(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + find_slice_end_index(self.ends(), index + self.offset(), ctx) } } + impl> RunEndArrayExt for T {} #[derive(Clone, Debug)] @@ -325,7 +323,9 @@ impl RunEndData { } } - pub(crate) fn validate_parts( + /// Validate that `ends` and `values` form a well-formed run-end array covering + /// `offset..offset + length`. + pub fn validate_parts( ends: &ArrayRef, values: &ArrayRef, offset: usize, diff --git a/encodings/runend/src/compute/cast.rs b/encodings/runend/src/compute/cast.rs index fe740739ca9..4722f9e5549 100644 --- a/encodings/runend/src/compute/cast.rs +++ b/encodings/runend/src/compute/cast.rs @@ -184,6 +184,6 @@ mod tests { fn test_cast_runend_conformance(#[case] build: RunEndBuilder) { let mut ctx = SESSION.create_execution_ctx(); let array = build(&mut ctx); - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut ctx); } } diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index ca99fdf1912..6d8d4aab6ea 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -5,6 +5,7 @@ use std::cmp::min; use std::ops::AddAssign; use num_traits::AsPrimitive; +use num_traits::NumCast; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -74,7 +75,7 @@ impl FilterKernel for RunEnd { } // Code adapted from apache arrow-rs https://github.com/apache/arrow-rs/blob/b1f5c250ebb6c1252b4e7c51d15b8e77f4c361fa/arrow-select/src/filter.rs#L425 -fn filter_run_end_primitive + AsPrimitive>( +pub fn filter_run_end_primitive + AsPrimitive>( run_ends: &[R], offset: u64, length: u64, @@ -87,16 +88,21 @@ fn filter_run_end_primitive + AsPrimitiv let mut count = R::zero(); let new_mask: Mask = BitBuffer::collect_bool(run_ends.len(), |i| { - let mut keep = false; let end = min(run_ends[i].as_() - offset, length); - // Safety: predicate must be the same length as the array the ends have been taken from - for pred in (start..end).map(|i| unsafe { - mask.value_unchecked(i.try_into().vortex_expect("index must fit in usize")) - }) { - count += >::from(pred); - keep |= pred - } + // SIMD popcount of the predicate bits in this run. The range matches the + // bit-by-bit `value_unchecked` read it replaces, so `end <= mask.len()`. + let start_usize = start + .try_into() + .vortex_expect("run start index must fit in usize"); + let end_usize = end + .try_into() + .vortex_expect("run end index must fit in usize"); + let run_trues = mask.count_range(start_usize, end_usize); + count += ::from(run_trues) + .vortex_expect("run popcount must fit in run-end native type"); + let keep = run_trues > 0; + // this is to avoid branching new_run_ends[j] = count; j += keep as usize; diff --git a/encodings/runend/src/compute/take.rs b/encodings/runend/src/compute/take.rs index c61854635e3..4ec601e71a6 100644 --- a/encodings/runend/src/compute/take.rs +++ b/encodings/runend/src/compute/take.rs @@ -751,7 +751,7 @@ mod tests { .unwrap() })] fn test_take_runend_conformance(#[case] array: RunEndArray) { - test_take_conformance(&array.into_array()); + test_take_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[rstest] @@ -765,6 +765,6 @@ mod tests { array.slice(2..8).unwrap() })] fn test_take_sliced_runend_conformance(#[case] sliced: ArrayRef) { - test_take_conformance(&sliced); + test_take_conformance(&sliced, &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index 484a4233789..13f32df07a1 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -63,8 +63,8 @@ fn slice( ) -> VortexResult { let new_length = range.len(); - let slice_begin = array.find_physical_index(range.start)?; - let slice_end = crate::ops::find_slice_end_index(array.ends(), range.end + array.offset())?; + let slice_begin = array.find_physical_index(range.start, ctx)?; + let slice_end = array.find_slice_end_index(range.end, ctx)?; // If the sliced range contains only a single run, opt to return a ConstantArray. if slice_begin + 1 == slice_end { diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index 8169f5d8dfc..ed3ff946d31 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -9,18 +9,17 @@ pub use array::*; pub use iter::trimmed_ends_iter; mod array; -#[cfg(feature = "arrow")] -mod arrow; pub mod compress; mod compute; pub mod decompress_bool; mod iter; mod kernel; -mod ops; +pub mod ops; mod rules; #[doc(hidden)] pub mod _benchmarking { + pub use compute::filter::filter_run_end_primitive; pub use compute::take::take_indices_unchecked; use super::*; diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index 56de024f988..39fb004c3e7 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -4,10 +4,11 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; -use vortex_array::scalar::PValue; +use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::scalar::Scalar; use vortex_array::search_sorted::SearchResult; use vortex_array::search_sorted::SearchSorted; +use vortex_array::search_sorted::SearchSortedPrimitiveArray; use vortex_array::search_sorted::SearchSortedSide; use vortex_array::vtable::OperationsVTable; use vortex_error::VortexResult; @@ -21,9 +22,8 @@ impl OperationsVTable for RunEnd { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - array - .values() - .execute_scalar(array.find_physical_index(index)?, ctx) + let physical_index = array.find_physical_index(index, ctx)?; + array.values().execute_scalar(physical_index, ctx) } } @@ -31,10 +31,15 @@ impl OperationsVTable for RunEnd { /// /// If the index exists in the array we want to take that position (as we are searching from the right) /// otherwise we want to take the next one -pub(crate) fn find_slice_end_index(array: &ArrayRef, index: usize) -> VortexResult { - let result = array - .as_primitive_typed() - .search_sorted(&PValue::from(index), SearchSortedSide::Right)?; +pub fn find_slice_end_index( + array: &ArrayRef, + index: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let result = match_each_unsigned_integer_ptype!(array.dtype().as_ptype(), |T| { + SearchSortedPrimitiveArray::::new(array, ctx) + .search_sorted(&index, SearchSortedSide::Right)? + }); Ok(match result { SearchResult::Found(i) => i, SearchResult::NotFound(i) => { @@ -47,6 +52,19 @@ pub(crate) fn find_slice_end_index(array: &ArrayRef, index: usize) -> VortexResu }) } +/// Find the physical index (the run) that contains the logical `index`, given the run `ends`. +pub fn find_physical_index( + array: &ArrayRef, + index: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match_each_unsigned_integer_ptype!(array.dtype().as_ptype(), |T| { + Ok(SearchSortedPrimitiveArray::::new(array, ctx) + .search_sorted(&index, SearchSortedSide::Right)? + .to_ends_index(array.len())) + }) +} + #[cfg(test)] mod tests { use std::sync::LazyLock; diff --git a/encodings/sequence/Cargo.toml b/encodings/sequence/Cargo.toml index 5dd9c5c187d..e60d6be29b7 100644 --- a/encodings/sequence/Cargo.toml +++ b/encodings/sequence/Cargo.toml @@ -25,7 +25,6 @@ vortex-proto = { workspace = true } vortex-session = { workspace = true } [dev-dependencies] -itertools = { workspace = true } rstest = { workspace = true } vortex-array = { path = "../../vortex-array", features = ["_test-harness"] } diff --git a/encodings/sequence/src/compute/cast.rs b/encodings/sequence/src/compute/cast.rs index e0a63d0c922..fb18cf11181 100644 --- a/encodings/sequence/src/compute/cast.rs +++ b/encodings/sequence/src/compute/cast.rs @@ -213,6 +213,6 @@ mod tests { 5, ).unwrap())] fn test_cast_sequence_conformance(#[case] sequence: SequenceArray) { - test_cast_conformance(&sequence.into_array()); + test_cast_conformance(&sequence.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/sequence/src/compute/filter.rs b/encodings/sequence/src/compute/filter.rs index 3ef24188fbf..aae06610d82 100644 --- a/encodings/sequence/src/compute/filter.rs +++ b/encodings/sequence/src/compute/filter.rs @@ -48,6 +48,8 @@ fn filter_impl(mul: T, base: T, mask: &Mask, validity: Validity) mod tests { use rstest::rstest; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::compute::conformance::filter::LARGE_SIZE; use vortex_array::compute::conformance::filter::MEDIUM_SIZE; use vortex_array::compute::conformance::filter::test_filter_conformance; @@ -70,6 +72,9 @@ mod tests { #[case(Sequence::try_new_typed(0u32, 5, Nullability::NonNullable, MEDIUM_SIZE).unwrap())] #[case(Sequence::try_new_typed(0u64, 1, Nullability::NonNullable, LARGE_SIZE).unwrap())] fn test_filter_sequence_conformance(#[case] array: SequenceArray) { - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &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 4b056d0ae7f..c59842224e2 100644 --- a/encodings/sequence/src/compute/take.rs +++ b/encodings/sequence/src/compute/take.rs @@ -107,6 +107,7 @@ mod test { 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_array::dtype::Nullability; use crate::Sequence; @@ -161,9 +162,11 @@ mod test { Nullability::Nullable, 1000 ).unwrap())] - fn test_take_conformance(#[case] sequence: SequenceArray) { - use vortex_array::compute::conformance::take::test_take_conformance; - test_take_conformance(&sequence.into_array()); + fn sequence_take_conformance(#[case] sequence: SequenceArray) { + test_take_conformance( + &sequence.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] diff --git a/encodings/sparse/Cargo.toml b/encodings/sparse/Cargo.toml index acf99a74ccf..56e87745b77 100644 --- a/encodings/sparse/Cargo.toml +++ b/encodings/sparse/Cargo.toml @@ -31,6 +31,7 @@ divan = { workspace = true } itertools = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-arrow = { workspace = true } [[bench]] name = "sparse_canonical" diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index afca2b29a18..0be6d0a4162 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -613,7 +613,6 @@ mod test { use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::listview::ListViewArrayExt; - use vortex_array::arrow::ArrowSessionExt; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; @@ -625,6 +624,7 @@ mod test { use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; + use vortex_arrow::ArrowSessionExt; use vortex_buffer::ByteBuffer; use vortex_buffer::buffer; use vortex_buffer::buffer_mut; diff --git a/encodings/sparse/src/compute/cast.rs b/encodings/sparse/src/compute/cast.rs index 21af0bc212a..f9596aa9976 100644 --- a/encodings/sparse/src/compute/cast.rs +++ b/encodings/sparse/src/compute/cast.rs @@ -131,7 +131,7 @@ mod tests { Scalar::from(0u8) ).unwrap())] fn test_cast_sparse_conformance(#[case] array: SparseArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] diff --git a/encodings/sparse/src/compute/filter.rs b/encodings/sparse/src/compute/filter.rs index 91e67297fd2..c1818262b9b 100644 --- a/encodings/sparse/src/compute/filter.rs +++ b/encodings/sparse/src/compute/filter.rs @@ -147,6 +147,7 @@ mod tests { ) .unwrap() .into_array(), + &mut SESSION.create_execution_ctx(), ); let ten_fill_value = Scalar::from(10i32); @@ -159,6 +160,7 @@ mod tests { ) .unwrap() .into_array(), + &mut SESSION.create_execution_ctx(), ) } } diff --git a/encodings/sparse/src/compute/mod.rs b/encodings/sparse/src/compute/mod.rs index bdae87e560d..81dfa89080a 100644 --- a/encodings/sparse/src/compute/mod.rs +++ b/encodings/sparse/src/compute/mod.rs @@ -138,6 +138,7 @@ mod tests { ) .unwrap() .into_array(), + &mut SESSION.create_execution_ctx(), ); let ten_fill_value = Scalar::from(10i32); @@ -150,6 +151,7 @@ mod tests { ) .unwrap() .into_array(), + &mut SESSION.create_execution_ctx(), ) } diff --git a/encodings/sparse/src/compute/take.rs b/encodings/sparse/src/compute/take.rs index e260eb26e9e..3babc19e57b 100644 --- a/encodings/sparse/src/compute/take.rs +++ b/encodings/sparse/src/compute/take.rs @@ -64,6 +64,7 @@ mod test { use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::compute::conformance::take::test_take_conformance; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; @@ -227,7 +228,6 @@ mod test { Scalar::from(-1i32), ).unwrap())] fn test_take_sparse_conformance(#[case] sparse: SparseArray) { - use vortex_array::compute::conformance::take::test_take_conformance; - test_take_conformance(&sparse.into_array()); + test_take_conformance(&sparse.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index 93c0bbe0363..4d4ec2623de 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -590,7 +590,7 @@ impl SparseData { // TODO(robert): Support other dtypes, only thing missing is getting most common value out of the array let primitive = array.clone().execute::(ctx)?; let (top_pvalue, _) = primitive - .top_value()? + .top_value(ctx)? .vortex_expect("Non empty or all null array"); Scalar::primitive_value(top_pvalue, top_pvalue.ptype(), array.dtype().nullability()) diff --git a/encodings/zigzag/src/compute/cast.rs b/encodings/zigzag/src/compute/cast.rs index c3fae84e346..94123e4498d 100644 --- a/encodings/zigzag/src/compute/cast.rs +++ b/encodings/zigzag/src/compute/cast.rs @@ -143,6 +143,6 @@ mod tests { #[case(zigzag_encode(PrimitiveArray::from_option_iter([Some(-5i16), None, Some(0), Some(5), None]).as_view()).unwrap())] #[case(zigzag_encode(PrimitiveArray::from_iter([i32::MIN, -1, 0, 1, i32::MAX]).as_view()).unwrap())] fn test_cast_zigzag_conformance(#[case] array: ZigZagArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/zigzag/src/compute/mod.rs b/encodings/zigzag/src/compute/mod.rs index 1da341989c5..e1c2f27f5d7 100644 --- a/encodings/zigzag/src/compute/mod.rs +++ b/encodings/zigzag/src/compute/mod.rs @@ -147,7 +147,7 @@ mod tests { let zigzag = zigzag_encode( PrimitiveArray::new(buffer![-189i32, -160, 1, 42, -73], Validity::AllValid).as_view(), )?; - test_filter_conformance(&zigzag.into_array()); + test_filter_conformance(&zigzag.into_array(), &mut SESSION.create_execution_ctx()); // Test with i64 values let zigzag = zigzag_encode( @@ -157,13 +157,13 @@ mod tests { ) .as_view(), )?; - test_filter_conformance(&zigzag.into_array()); + test_filter_conformance(&zigzag.into_array(), &mut SESSION.create_execution_ctx()); // Test with nullable values let array = PrimitiveArray::from_option_iter([Some(-10i16), None, Some(20), Some(-30), None]); let zigzag = zigzag_encode(array.as_view())?; - test_filter_conformance(&zigzag.into_array()); + test_filter_conformance(&zigzag.into_array(), &mut SESSION.create_execution_ctx()); Ok(()) } @@ -176,13 +176,13 @@ mod tests { PrimitiveArray::new(buffer![-100i32, 200, -300, 400, -500], Validity::AllValid) .as_view(), )?; - test_mask_conformance(&zigzag.into_array()); + test_mask_conformance(&zigzag.into_array(), &mut SESSION.create_execution_ctx()); // Test with i8 values let zigzag = zigzag_encode( PrimitiveArray::new(buffer![-127i8, 0, 127, -1, 1], Validity::AllValid).as_view(), )?; - test_mask_conformance(&zigzag.into_array()); + test_mask_conformance(&zigzag.into_array(), &mut SESSION.create_execution_ctx()); Ok(()) } @@ -198,7 +198,7 @@ mod tests { 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()); + test_take_conformance(&zigzag.into_array(), &mut ctx); Ok(()) } diff --git a/encodings/zstd/Cargo.toml b/encodings/zstd/Cargo.toml index e5e415544a2..1b2cb4f2dd0 100644 --- a/encodings/zstd/Cargo.toml +++ b/encodings/zstd/Cargo.toml @@ -34,10 +34,5 @@ vortex-session = { workspace = true } zstd = { workspace = true } [dev-dependencies] -divan = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } - -[[bench]] -name = "listview_rebuild" -harness = false diff --git a/encodings/zstd/benches/listview_rebuild.rs b/encodings/zstd/benches/listview_rebuild.rs deleted file mode 100644 index e3c0e6d5dd1..00000000000 --- a/encodings/zstd/benches/listview_rebuild.rs +++ /dev/null @@ -1,58 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#![expect(clippy::unwrap_used)] - -use std::sync::LazyLock; - -use divan::Bencher; -use vortex_array::IntoArray; -use vortex_array::VortexSessionExecute; -use vortex_array::arrays::ListViewArray; -use vortex_array::arrays::VarBinViewArray; -use vortex_array::arrays::listview::ListViewRebuildMode; -use vortex_array::validity::Validity; -use vortex_buffer::Buffer; -use vortex_session::VortexSession; -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(vortex_array::array_session); - -#[divan::bench(sample_size = 1000)] -fn rebuild_naive(bencher: Bencher) { - let dudes = VarBinViewArray::from_iter_str(["Washington", "Adams", "Jefferson", "Madison"]) - .into_array(); - let dtype = dudes.dtype().clone(); - let validity = dudes.validity().unwrap(); - let dudes = Zstd::try_new( - dtype, - ZstdData::from_array(dudes, 9, 1024, &mut SESSION.create_execution_ctx()).unwrap(), - validity, - ) - .unwrap() - .into_array(); - - let offsets = std::iter::repeat_n(0u32, 1024) - .collect::>() - .into_array(); - let sizes = [0u64, 1, 2, 3, 4] - .into_iter() - .cycle() - .take(1024) - .collect::>() - .into_array(); - - let list_view = ListViewArray::new(dudes, offsets, sizes, Validity::NonNullable); - - bencher - .with_inputs(|| (&list_view, SESSION.create_execution_ctx())) - .bench_refs(|(list_view, ctx)| { - list_view.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx) - }) -} - -fn main() { - divan::main() -} diff --git a/encodings/zstd/src/compute/cast.rs b/encodings/zstd/src/compute/cast.rs index c5f15eba284..dba59198c88 100644 --- a/encodings/zstd/src/compute/cast.rs +++ b/encodings/zstd/src/compute/cast.rs @@ -200,6 +200,6 @@ mod tests { fn test_cast_zstd_conformance(#[case] values: PrimitiveArray) { let zstd = Zstd::from_primitive(&values, 0, 0, &mut SESSION.create_execution_ctx()).unwrap(); - test_cast_conformance(&zstd.into_array()); + test_cast_conformance(&zstd.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index f40a451caaf..73361a593a1 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -25,6 +25,7 @@ use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; use vortex_array::serde::ArrayChildren; use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; @@ -512,25 +513,22 @@ impl OperationsVTable for ZstdBuffers { // TODO(os): maybe we should not support scalar_at, it is really slow, and adding a cache // layer here is weird. Valid use of zstd buffers array would be by executing it first into // canonical - let inner_array = ZstdBuffers::decompress_and_build_inner( - &array.into_owned(), - &vortex_array::LEGACY_SESSION, - )?; + let inner_array = + ZstdBuffers::decompress_and_build_inner(&array.into_owned(), ctx.session())?; inner_array.execute_scalar(index, ctx) } } impl ValidityVTable for ZstdBuffers { - fn validity( - array: ArrayView<'_, ZstdBuffers>, - ) -> VortexResult { + #[allow(clippy::disallowed_methods)] + fn validity(array: ArrayView<'_, ZstdBuffers>) -> VortexResult { if !array.dtype().is_nullable() { - return Ok(vortex_array::validity::Validity::NonNullable); + return Ok(Validity::NonNullable); } let inner_array = ZstdBuffers::decompress_and_build_inner( &array.into_owned(), - &vortex_array::LEGACY_SESSION, + vortex_array::legacy_session(), )?; inner_array.validity() } diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index 7c1bbaad091..b6402cb1e1d 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -41,6 +41,7 @@ use strum::IntoEnumIterator; use tracing::debug; use vortex_array::ArrayRef; use vortex_array::Canonical; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::NumericalAggregateOpts; @@ -553,25 +554,32 @@ fn random_action_from_list( /// Compress an array using the given strategy. #[cfg(feature = "zstd")] -pub fn compress_array(array: &ArrayRef, strategy: CompressorStrategy) -> ArrayRef { - let mut ctx = SESSION.create_execution_ctx(); +pub fn compress_array( + array: &ArrayRef, + strategy: CompressorStrategy, + ctx: &mut ExecutionCtx, +) -> ArrayRef { match strategy { CompressorStrategy::Default => BtrBlocksCompressor::default() - .compress(array, &mut ctx) + .compress(array, ctx) .vortex_expect("BtrBlocksCompressor compress should succeed in fuzz test"), CompressorStrategy::Compact => BtrBlocksCompressorBuilder::default() .with_compact() .build() - .compress(array, &mut ctx) + .compress(array, ctx) .vortex_expect("Compact compress should succeed in fuzz test"), } } /// Compress an array using the given strategy (only Default). #[cfg(not(feature = "zstd"))] -pub fn compress_array(array: &ArrayRef, _strategy: CompressorStrategy) -> ArrayRef { +pub fn compress_array( + array: &ArrayRef, + _strategy: CompressorStrategy, + ctx: &mut ExecutionCtx, +) -> ArrayRef { BtrBlocksCompressor::default() - .compress(array, &mut SESSION.create_execution_ctx()) + .compress(array, ctx) .vortex_expect("BtrBlocksCompressor compress should succeed in fuzz test") } @@ -602,14 +610,14 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> VortexFuzzResult { .clone() .execute::(&mut ctx) .vortex_expect("execute canonical should succeed in fuzz test"); - current_array = compress_array(&canonical.into_array(), strategy); - assert_array_eq(&expected.array(), ¤t_array, i)?; + current_array = compress_array(&canonical.into_array(), strategy, &mut ctx); + assert_array_eq(&expected.array(), ¤t_array, i, &mut ctx)?; } Action::Slice(range) => { current_array = current_array .slice(range) .vortex_expect("slice operation should succeed in fuzz test"); - assert_array_eq(&expected.array(), ¤t_array, i)?; + assert_array_eq(&expected.array(), ¤t_array, i, &mut ctx)?; } Action::Take(indices) => { if indices.is_empty() { @@ -618,14 +626,14 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> VortexFuzzResult { current_array = current_array .take(indices) .vortex_expect("take operation should succeed in fuzz test"); - assert_array_eq(&expected.array(), ¤t_array, i)?; + assert_array_eq(&expected.array(), ¤t_array, i, &mut ctx)?; } Action::SearchSorted(s, side) => { let mut sorted = sort_canonical_array(¤t_array, &mut ctx) .vortex_expect("sort_canonical_array should succeed in fuzz test"); if !current_array.is_canonical() { - sorted = compress_array(&sorted, CompressorStrategy::Default); + sorted = compress_array(&sorted, CompressorStrategy::Default, &mut ctx); } assert_search_sorted(sorted, s, side, expected.search(), i)?; } @@ -633,7 +641,7 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> VortexFuzzResult { current_array = current_array .filter(mask_val) .vortex_expect("filter operation should succeed in fuzz test"); - assert_array_eq(&expected.array(), ¤t_array, i)?; + assert_array_eq(&expected.array(), ¤t_array, i, &mut ctx)?; } Action::Compare(v, op) => { let compare_result = current_array @@ -642,7 +650,7 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> VortexFuzzResult { Operator::from(op), ) .vortex_expect("compare operation should succeed in fuzz test"); - if let Err(e) = assert_array_eq(&expected.array(), &compare_result, i) { + if let Err(e) = assert_array_eq(&expected.array(), &compare_result, i, &mut ctx) { vortex_panic!( "Failed to compare {}with {op} {v}\nError: {e}", current_array.display_tree() @@ -654,7 +662,7 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> VortexFuzzResult { let cast_result = current_array .cast(to.clone()) .vortex_expect("cast operation should succeed in fuzz test"); - if let Err(e) = assert_array_eq(&expected.array(), &cast_result, i) { + if let Err(e) = assert_array_eq(&expected.array(), &cast_result, i, &mut ctx) { vortex_panic!( "Failed to cast {} to dtype {to}\nError: {e}", current_array.display_tree() @@ -677,13 +685,13 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> VortexFuzzResult { current_array = current_array .fill_null(fill_value.clone()) .vortex_expect("fill_null operation should succeed in fuzz test"); - assert_array_eq(&expected.array(), ¤t_array, i)?; + assert_array_eq(&expected.array(), ¤t_array, i, &mut ctx)?; } Action::Mask(mask_val) => { current_array = current_array .mask(mask_val.into_array()) .vortex_expect("mask operation should succeed in fuzz test"); - assert_array_eq(&expected.array(), ¤t_array, i)?; + assert_array_eq(&expected.array(), ¤t_array, i, &mut ctx)?; } Action::ScalarAt(indices) => { let expected_scalars = expected.scalar_vec(); @@ -730,7 +738,12 @@ fn assert_search_sorted( /// Uses `all_non_distinct` for an efficient buffer-level comparison on the happy path. /// Falls back to element-wise scalar comparison only on mismatch to produce a detailed error. #[expect(clippy::result_large_err)] -pub fn assert_array_eq(lhs: &ArrayRef, rhs: &ArrayRef, step: usize) -> VortexFuzzResult<()> { +pub fn assert_array_eq( + lhs: &ArrayRef, + rhs: &ArrayRef, + step: usize, + ctx: &mut ExecutionCtx, +) -> VortexFuzzResult<()> { if lhs.dtype() != rhs.dtype() { return Err(VortexFuzzError::DTypeMismatch( lhs.clone(), @@ -751,10 +764,8 @@ pub fn assert_array_eq(lhs: &ArrayRef, rhs: &ArrayRef, step: usize) -> VortexFuz )); } - let mut ctx = SESSION.create_execution_ctx(); - // Fast path: buffer-level comparison. - let identical = all_non_distinct(lhs, rhs, &mut ctx) + let identical = all_non_distinct(lhs, rhs, ctx) .map_err(|e| VortexFuzzError::VortexError(e, Backtrace::capture()))?; if identical { return Ok(()); @@ -762,8 +773,8 @@ pub fn assert_array_eq(lhs: &ArrayRef, rhs: &ArrayRef, step: usize) -> VortexFuz // Slow path: find the first differing element for a detailed error message. for idx in 0..lhs.len() { - let l = lhs.execute_scalar(idx, &mut ctx).vortex_expect("scalar_at"); - let r = rhs.execute_scalar(idx, &mut ctx).vortex_expect("scalar_at"); + let l = lhs.execute_scalar(idx, ctx).vortex_expect("scalar_at"); + let r = rhs.execute_scalar(idx, ctx).vortex_expect("scalar_at"); if l != r { return Err(VortexFuzzError::ArrayNotEqual( diff --git a/fuzz/src/compress.rs b/fuzz/src/compress.rs index 1284a7be45e..9953de79b0a 100644 --- a/fuzz/src/compress.rs +++ b/fuzz/src/compress.rs @@ -6,6 +6,8 @@ //! This module generates arbitrary instances of compressed encodings (DictArray, etc.), //! then verifies that `to_canonical()` works and produces correct `len` and `dtype`. +use std::backtrace::Backtrace; + use arbitrary::Arbitrary; use arbitrary::Unstructured; use vortex_array::ArrayRef; @@ -17,6 +19,8 @@ use vortex_array::arrays::dict::ArbitraryDictArray; use vortex_runend::ArbitraryRunEndArray; use crate::SESSION; +use crate::error::VortexFuzzError; +use crate::error::VortexFuzzResult; /// Which compressed encoding to generate. #[derive(Debug, Clone, Copy)] @@ -64,10 +68,7 @@ impl<'a> Arbitrary<'a> for FuzzCompressRoundtrip { /// - `Ok(false)` - reject from corpus /// - `Err(_)` - a bug was found #[expect(clippy::result_large_err)] -pub fn run_compress_roundtrip(fuzz: FuzzCompressRoundtrip) -> crate::error::VortexFuzzResult { - use crate::error::Backtrace; - use crate::error::VortexFuzzError; - +pub fn run_compress_roundtrip(fuzz: FuzzCompressRoundtrip) -> VortexFuzzResult { let FuzzCompressRoundtrip { array } = fuzz; // Store original properties diff --git a/java/README.md b/java/README.md index d9554420998..f45f2474f11 100644 --- a/java/README.md +++ b/java/README.md @@ -3,7 +3,9 @@ We provide two interfaces for working with Vortex from Java: - `vortex-java` - a low-level interface JNI for working with Vortex files and arrays on cloud and local storage -- `vortex-spark` - A Spark connector for working with datasets of Vortex files +- `vortex-spark` - A Spark connector for working with datasets of Vortex files. See + [vortex-spark/README.md](vortex-spark/README.md) for how to load the connector into Spark + and query Vortex files from the DataFrame API or Spark SQL. ## Publishing diff --git a/java/vortex-jni/build.gradle.kts b/java/vortex-jni/build.gradle.kts index bb652d0d86e..8f5106a1262 100644 --- a/java/vortex-jni/build.gradle.kts +++ b/java/vortex-jni/build.gradle.kts @@ -2,7 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar -import org.gradle.kotlin.dsl.support.serviceOf +import org.gradle.api.tasks.Copy +import org.gradle.api.tasks.Exec plugins { `java-library` @@ -111,51 +112,52 @@ tasks.build { dependsOn("shadowJar") } -tasks.register("makeTestFiles") { - description = "Generate files used by unit tests" +val rustWorkspaceDir = rootProject.projectDir.absoluteFile.parentFile +val osName = System.getProperty("os.name").lowercase() +val osArch = System.getProperty("os.arch").lowercase() +val osShortName = + when { + osName.contains("mac") -> "darwin" + osName.contains("nix") || osName.contains("nux") -> "linux" + osName.contains("win") -> "win" + else -> throw GradleException("Unsupported OS for makeTestFiles: $osName") + } +val libExt = + when (osShortName) { + "darwin" -> ".dylib" + "linux" -> ".so" + "win" -> ".dll" + else -> throw GradleException("Unsupported OS short name: $osShortName") + } +val nativeLibrary = rustWorkspaceDir.resolve("target/debug/libvortex_jni$libExt") +val nativeLibraryDir = "src/main/resources/native/$osShortName-$osArch" + +val buildJniLibrary = + tasks.register("buildJniLibrary") { + description = "Build the JNI library for the host architecture" + group = "verification" + + // The publish workflow places release, cross-compiled libs for every supported + // architecture before invoking shadowJar; rebuilding the host-arch debug lib + // here would overwrite them (linux-aarch64 ends up holding a linux-amd64 .so). + onlyIf { System.getenv("VORTEX_SKIP_MAKE_TEST_FILES") != "true" } + + workingDir = rustWorkspaceDir + executable = "cargo" + args("build", "--package", "vortex-jni") + } + +tasks.register("makeTestFiles") { + description = "Stage the JNI library used by unit tests" group = "verification" - // The publish workflow places release, cross-compiled libs for every supported - // architecture before invoking shadowJar; rebuilding the host-arch debug lib - // here would overwrite them (linux-aarch64 ends up holding a linux-amd64 .so). onlyIf { System.getenv("VORTEX_SKIP_MAKE_TEST_FILES") != "true" } + dependsOn(buildJniLibrary) - doLast { - println("makeTestFiles executed") - - val execOps = serviceOf() - - // Build the JNI lib for the host architecture only. - execOps.exec { - workingDir = rootProject.projectDir.absoluteFile.parentFile - executable = "cargo" - args("build", "--package", "vortex-jni") - } - - val osName = System.getProperty("os.name").lowercase() - val osArch = System.getProperty("os.arch").lowercase() - val osShortName = - when { - osName.contains("mac") -> "darwin" - osName.contains("nix") || osName.contains("nux") -> "linux" - osName.contains("win") -> "win" - else -> throw GradleException("Unsupported OS for makeTestFiles: $osName") - } - val libExt = - when (osShortName) { - "darwin" -> ".dylib" - "linux" -> ".so" - "win" -> ".dll" - else -> throw GradleException("Unsupported OS short name: $osShortName") - } - - // Only populate the host-arch directory so cross-compiled libs for other - // architectures (placed by the publish workflow) are preserved. - copy { - from("${rootProject.projectDir.absoluteFile.parentFile}/target/debug/libvortex_jni$libExt") - into("$projectDir/src/main/resources/native/$osShortName-$osArch") - } - } + // Only populate the host-arch directory so cross-compiled libs for other + // architectures (placed by the publish workflow) are preserved. + from(nativeLibrary) + into(nativeLibraryDir) } tasks.named("processResources").configure { 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 93c5df2c61a..9a888f224f9 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 @@ -5,6 +5,7 @@ import com.google.common.base.Preconditions; import dev.vortex.VortexCleaner; +import dev.vortex.io.NativeReadable; import dev.vortex.jni.NativeDataSource; import java.util.Arrays; import java.util.Collections; @@ -71,6 +72,51 @@ public static DataSource open(Session session, List uris, Map readables) { + return open(session, readables, 0); + } + + /** + * Open one or more files through caller-provided byte sources instead of native storage clients. Every read the + * scan performs becomes an upcall into the corresponding {@link NativeReadable}, so this is the integration point + * for external I/O abstractions (for example Iceberg's {@code FileIO}). Each file is identified by + * {@link NativeReadable#name()}, which must be unique within {@code readables}. + * + *

The readables must remain open until this data source and all scans created from it are closed; native code + * never closes them. + * + * @param session open session + * @param readables byte sources + * @param readConcurrency maximum in-flight {@code readFully} calls across all files of this data source; {@code 0} + * selects the default. Each in-flight read occupies one native thread and typically one stream on its readable. + */ + public static DataSource open(Session session, List readables, int readConcurrency) { + Objects.requireNonNull(session, "session"); + Objects.requireNonNull(readables, "readables"); + Preconditions.checkArgument(!readables.isEmpty(), "at least one readable is required"); + Preconditions.checkArgument(readConcurrency >= 0, "readConcurrency must not be negative"); + Object[] readableArray = readables.toArray(); + String[] names = new String[readableArray.length]; + long[] lengths = new long[readableArray.length]; + for (int i = 0; i < readableArray.length; i++) { + Preconditions.checkArgument(readableArray[i] != null, "readables must not contain null values"); + NativeReadable readable = readables.get(i); + names[i] = readable.name(); + Preconditions.checkArgument(names[i] != null, "readable at index %s returned a null name", i); + lengths[i] = readable.length(); + Preconditions.checkArgument(lengths[i] >= 0, "readable for %s reported negative length", names[i]); + } + long pointer = + NativeDataSource.openFiles(session.nativePointer(), readableArray, names, lengths, readConcurrency); + return new DataSource(session, pointer); + } + /** Arrow schema of the data source (and of scans produced from it). */ public Schema arrowSchema(BufferAllocator allocator) { try (ArrowSchema schema = ArrowSchema.allocateNew(allocator)) { diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/VortexColumnStatistics.java b/java/vortex-jni/src/main/java/dev/vortex/api/VortexColumnStatistics.java new file mode 100644 index 00000000000..308ca64a0df --- /dev/null +++ b/java/vortex-jni/src/main/java/dev/vortex/api/VortexColumnStatistics.java @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.api; + +import java.util.Optional; +import java.util.OptionalLong; + +/** Statistics and physical size reported for one top-level column after a Vortex write. */ +public final class VortexColumnStatistics { + private final int columnIndex; + private final long compressedSize; + private final long valueCount; + private final long nullValueCount; + private final long nanValueCount; + private final Object lowerBound; + private final Object upperBound; + + /** + * Construct native write statistics. + * + *

This constructor is public so the JNI implementation can instantiate the value without reflective access. + * Applications should obtain instances from {@link VortexWriteSummary#columnStatistics()}. + */ + public VortexColumnStatistics( + int columnIndex, + long compressedSize, + long valueCount, + long nullValueCount, + long nanValueCount, + Object lowerBound, + Object upperBound) { + this.columnIndex = columnIndex; + this.compressedSize = compressedSize; + this.valueCount = valueCount; + this.nullValueCount = nullValueCount; + this.nanValueCount = nanValueCount; + this.lowerBound = lowerBound; + this.upperBound = upperBound; + } + + /** Zero-based position of this column in the writer's Arrow schema. */ + public int columnIndex() { + return columnIndex; + } + + /** Compressed bytes referenced by this column's physical layout. */ + public long compressedSize() { + return compressedSize; + } + + /** Number of values in this top-level column. */ + public long valueCount() { + return valueCount; + } + + /** Exact null count, or empty when Vortex did not compute one for this data type. */ + public OptionalLong nullValueCount() { + return nullValueCount >= 0 ? OptionalLong.of(nullValueCount) : OptionalLong.empty(); + } + + /** Exact NaN count, or empty for non-floating-point columns. */ + public OptionalLong nanValueCount() { + return nanValueCount >= 0 ? OptionalLong.of(nanValueCount) : OptionalLong.empty(); + } + + /** + * Lower bound represented using the corresponding Arrow scalar's Java type. + * + *

Integers use {@link Integer}, {@link Long}, or {@link java.math.BigInteger}; floating-point values use + * {@link Float} or {@link Double}; decimal values use {@link java.math.BigDecimal}; strings use {@link String}; and + * binary values use {@code byte[]}. + */ + public Optional lowerBound() { + return Optional.ofNullable(lowerBound); + } + + /** Upper bound represented using the corresponding Arrow scalar's Java type. */ + public Optional upperBound() { + return Optional.ofNullable(upperBound); + } +} diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriteSummary.java b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriteSummary.java new file mode 100644 index 00000000000..69f042d2ff4 --- /dev/null +++ b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriteSummary.java @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.api; + +import java.util.List; + +/** Immutable metadata returned after a Vortex writer has finalized its file. */ +public final class VortexWriteSummary { + private final long fileSize; + private final long rowCount; + private final List columnStatistics; + + /** + * Construct a native write summary. + * + *

This constructor is public so the JNI implementation can instantiate the value without reflective access. + * Applications should obtain summaries from {@link VortexWriter#finish()}. + */ + public VortexWriteSummary(long fileSize, long rowCount, VortexColumnStatistics[] columnStatistics) { + this.fileSize = fileSize; + this.rowCount = rowCount; + this.columnStatistics = List.of(columnStatistics.clone()); + } + + /** Exact size of the completed file in bytes. */ + public long fileSize() { + return fileSize; + } + + /** Exact number of rows written to the file. */ + public long rowCount() { + return rowCount; + } + + /** Per-column statistics in Arrow schema order. */ + public List columnStatistics() { + return columnStatistics; + } +} diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java index 7761be59cdc..a62243aab07 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java +++ b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java @@ -5,6 +5,7 @@ import com.google.common.base.Preconditions; import dev.vortex.VortexCleaner; +import dev.vortex.io.NativeWritable; import dev.vortex.jni.NativeWriter; import java.io.IOException; import java.util.Map; @@ -29,6 +30,7 @@ public final class VortexWriter implements AutoCloseable { private final long pointer; private final AtomicBoolean closed = new AtomicBoolean(false); + private volatile VortexWriteSummary summary; private VortexWriter(long pointer) { Preconditions.checkArgument(pointer != 0, "invalid writer pointer"); @@ -65,6 +67,33 @@ public static VortexWriter create( } } + /** + * Create a writer that streams the file into a caller-provided byte sink instead of a native storage client. This + * is the integration point for external I/O abstractions (for example Iceberg's {@code PositionOutputStream}). + * + *

The native side writes and flushes the sink but never closes it: after {@link #close()} returns, all bytes + * have been written and flushed, and the caller must close the sink to finalize the file. + */ + public static VortexWriter create( + Session session, NativeWritable writable, Schema arrowSchema, BufferAllocator allocator) + throws IOException { + Objects.requireNonNull(session, "session"); + Objects.requireNonNull(writable, "writable"); + Objects.requireNonNull(arrowSchema, "arrowSchema"); + Objects.requireNonNull(allocator, "allocator"); + ArrowSchema ffi = ArrowSchema.allocateNew(allocator); + try { + Data.exportSchema(allocator, arrowSchema, null, ffi); + long ptr = NativeWriter.createStream(session.nativePointer(), writable, ffi.memoryAddress()); + if (ptr <= 0) { + throw new IOException("failed to create stream writer (ptr=" + ptr + ")"); + } + return new VortexWriter(ptr); + } finally { + ffi.close(); + } + } + /** Write a batch directly from Arrow C Data Interface addresses. */ public void writeBatch(long arrowArrayAddr, long arrowSchemaAddr) throws IOException { Preconditions.checkState(!closed.get(), "writer already closed"); @@ -79,15 +108,45 @@ public void writeBatch(long arrowArrayAddr, long arrowSchemaAddr) throws IOExcep } } - /** Flush any pending batches and finalize the file. Idempotent. */ - @Override - public void close() throws IOException { + /** + * Return the number of bytes successfully written to the underlying sink so far. + * + *

This count does not include queued batches or data still buffered by layout strategies, so it may lag the + * amount of input accepted by {@link #writeBatch(long, long)}. After {@link #finish()}, it is the exact completed + * file size and is equal to {@link VortexWriteSummary#fileSize()}. + */ + public synchronized long bytesWritten() { + if (summary != null) { + return summary.fileSize(); + } + Preconditions.checkState(!closed.get(), "writer closed without a write summary"); + long bytesWritten = NativeWriter.bytesWritten(pointer); + Preconditions.checkState(bytesWritten >= 0, "native writer returned an invalid byte count"); + return bytesWritten; + } + + /** + * Flush pending batches, finalize the file, and return its statistics and physical sizes. + * + *

This method is idempotent. Later calls return the same immutable summary. + */ + public synchronized VortexWriteSummary finish() throws IOException { if (closed.compareAndSet(false, true)) { try { - NativeWriter.close(pointer); + summary = NativeWriter.finish(pointer); } catch (RuntimeException e) { throw new IOException("failed to close writer", e); } } + if (summary == null) { + throw new IOException("writer was closed without retaining its write summary"); + } + return summary; + } + + /** Flush any pending batches and finalize the file. Idempotent. */ + @Override + public void close() throws IOException { + finish(); } } diff --git a/java/vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java b/java/vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java new file mode 100644 index 00000000000..84f264cc952 --- /dev/null +++ b/java/vortex-jni/src/main/java/dev/vortex/io/NativeReadable.java @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.io; + +import java.io.Closeable; +import java.io.IOException; +import java.nio.ByteBuffer; + +/** + * A random-access byte source supplied by the caller and read from native code. + * + *

Implementations bridge external I/O abstractions — for example Iceberg's {@code FileIO} input streams — into the + * Vortex native reader. The native side invokes {@link #readFully(long, ByteBuffer)} from its own worker threads, + * potentially many calls concurrently, so implementations must be safe for concurrent positional reads (typically by + * opening one underlying stream per concurrent call, or delegating to a positional-read API). + * + *

Lifecycle: the creator owns this object. Native code never calls {@link #close()}; close it only after every data + * source or scan built on top of it has been closed. + */ +public interface NativeReadable extends Closeable { + /** + * Unique name of this source, typically its original location (URI or file path). Used to key per-session caches + * and in native error messages, so it must be stable and unique across sources; it must not contain glob characters + * ({@code *?[}). + */ + String name(); + + /** Total length of the source in bytes. Must be cheap; called once when the source is registered. */ + long length(); + + /** + * Read exactly {@code buffer.remaining()} bytes starting at absolute {@code position} into {@code buffer}. + * + *

Unlike {@link java.io.InputStream#read}, short reads are not permitted: implementations must either fill the + * buffer completely ({@code buffer.remaining() == 0} on return) or throw. + * + *

{@code buffer} is a direct {@link ByteBuffer} wrapping native memory owned by the caller, so bytes + * written to it land in the native reader without an extra copy. It is valid only for the duration of this call: + * implementations must not retain a reference to it after returning. Implementations backed by {@code byte[]} APIs + * can bulk-{@link ByteBuffer#put(byte[], int, int) put} into it. + * + * @throws java.io.EOFException if the range extends past the end of the source + * @throws IOException if the underlying storage fails + */ + void readFully(long position, ByteBuffer buffer) throws IOException; +} diff --git a/java/vortex-jni/src/main/java/dev/vortex/io/NativeWritable.java b/java/vortex-jni/src/main/java/dev/vortex/io/NativeWritable.java new file mode 100644 index 00000000000..a86d127dad6 --- /dev/null +++ b/java/vortex-jni/src/main/java/dev/vortex/io/NativeWritable.java @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.io; + +import java.io.Closeable; +import java.io.IOException; + +/** + * A sequential byte sink supplied by the caller and written to from native code. + * + *

Implementations bridge external output abstractions — for example Iceberg's {@code PositionOutputStream} — into + * the Vortex native writer. Writes are sequential and single-threaded: the native side never issues concurrent + * {@code write} calls against the same instance. + * + *

Lifecycle: the creator owns this object. Native code calls {@link #write} and {@link #flush} but never + * {@link #close()}; after {@link dev.vortex.api.VortexWriter#close()} returns, all bytes have been written and flushed, + * and the creator must close the underlying stream to finalize it. + */ +public interface NativeWritable extends Closeable { + /** Append {@code length} bytes from {@code buffer} starting at {@code offset}. */ + void write(byte[] buffer, int offset, int length) throws IOException; + + /** Flush buffered bytes to the underlying storage. */ + void flush() throws IOException; +} 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 cc2aa163cee..76d09baf42e 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 @@ -22,6 +22,21 @@ private NativeDataSource() {} */ public static native long open(long sessionPointer, String[] uris, Map options); + /** + * Open a data source over caller-provided {@link dev.vortex.io.NativeReadable} objects. All reads become upcalls + * into the supplied readables; no native storage client is created. + * + * @param sessionPointer pointer from {@link NativeSession#newSession()} + * @param readables one {@link dev.vortex.io.NativeReadable} per file + * @param names unique name per file (from {@link dev.vortex.io.NativeReadable#name()}), parallel to + * {@code readables} + * @param lengths file sizes in bytes, parallel to {@code readables} + * @param readConcurrency maximum in-flight {@code readFully} upcalls across all files of this data source; + * {@code <= 0} selects the default + */ + public static native long openFiles( + long sessionPointer, Object[] readables, String[] names, long[] lengths, int readConcurrency); + /** Free a data source pointer. */ public static native void free(long pointer); diff --git a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java index d1f2e725f8c..4f61a381438 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java +++ b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java @@ -3,6 +3,7 @@ package dev.vortex.jni; +import dev.vortex.api.VortexWriteSummary; import java.util.Map; /** JNI boundary for {@link dev.vortex.api.VortexWriter}. */ @@ -17,6 +18,12 @@ private NativeWriter() {} public static native long create( long sessionPointer, String uri, long arrowSchemaAddress, Map options); + /** + * Open a writer that streams the file into a caller-provided {@link dev.vortex.io.NativeWritable}. The native side + * writes and flushes but never closes the writable; the caller must close it after {@link #close}. + */ + public static native long createStream(long sessionPointer, Object writable, long arrowSchemaAddress); + /** * Write a batch directly from Arrow C Data Interface addresses. * @@ -27,6 +34,12 @@ public static native long create( */ public static native boolean writeBatch(long writerPointer, long arrowArrayAddress, long arrowSchemaAddress); + /** Number of bytes successfully written to the underlying sink so far. */ + public static native long bytesWritten(long writerPointer); + /** Flush and close the writer. Must be called exactly once. */ public static native void close(long writerPointer); + + /** Flush and close the writer, returning its completed-file summary. Must be called exactly once. */ + public static native VortexWriteSummary finish(long writerPointer); } 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 aeee8febcf4..3a6f9ca0883 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 @@ -23,7 +23,6 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.FieldVector; -import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.ipc.ArrowReader; @@ -155,8 +154,8 @@ public void testProjectedScan() throws Exception { List people = readAll(ds, options, allocator, batch -> { List results = new ArrayList<>(); - VarCharVector names = (VarCharVector) batch.getVector("Name"); - VarCharVector states = (VarCharVector) batch.getVector("State"); + ViewVarCharVector names = (ViewVarCharVector) batch.getVector("Name"); + ViewVarCharVector states = (ViewVarCharVector) batch.getVector("State"); for (int i = 0; i < batch.getRowCount(); i++) { String name = names.isNull(i) ? null : new String(names.get(i), UTF_8); String state = states.isNull(i) ? null : new String(states.get(i), UTF_8); @@ -272,9 +271,9 @@ private static List readAll( private static List readFullBatch(VectorSchemaRoot root) { List result = new ArrayList<>(); - VarCharVector names = (VarCharVector) root.getVector("Name"); + ViewVarCharVector names = (ViewVarCharVector) root.getVector("Name"); FieldVector salaries = root.getVector("Salary"); - VarCharVector states = (VarCharVector) root.getVector("State"); + ViewVarCharVector states = (ViewVarCharVector) root.getVector("State"); for (int i = 0; i < root.getRowCount(); i++) { String name = names.isNull(i) ? null : new String(names.get(i), UTF_8); diff --git a/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java b/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java new file mode 100644 index 00000000000..5a6d40e05db --- /dev/null +++ b/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java @@ -0,0 +1,332 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.io; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.vortex.api.DataSource; +import dev.vortex.api.Partition; +import dev.vortex.api.Scan; +import dev.vortex.api.ScanOptions; +import dev.vortex.api.Session; +import dev.vortex.api.VortexWriteSummary; +import dev.vortex.api.VortexWriter; +import dev.vortex.arrow.ArrowAllocation; +import dev.vortex.jni.NativeLoader; +import java.io.EOFException; +import java.io.IOException; +import java.io.OutputStream; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.List; +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ViewVarCharVector; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +/** Round-trip tests for the caller-provided I/O bridge ({@link NativeReadable} / {@link NativeWritable}). */ +public final class NativeIOBridgeTest { + @TempDir + Path tempDir; + + @BeforeAll + public static void loadLibrary() { + NativeLoader.loadJni(); + } + + private static Schema personSchema() { + return new Schema(List.of( + Field.notNullable("name", new ArrowType.Utf8()), + Field.notNullable("age", new ArrowType.Int(32, true)))); + } + + /** A {@link NativeReadable} over a local file, safe for concurrent positional reads. */ + private static final class FileChannelReadable implements NativeReadable { + private final String name; + private final FileChannel channel; + private final long length; + + FileChannelReadable(Path path) throws IOException { + this(path.toString(), path); + } + + FileChannelReadable(String name, Path path) throws IOException { + this.name = name; + this.channel = FileChannel.open(path, StandardOpenOption.READ); + this.length = channel.size(); + } + + @Override + public String name() { + return name; + } + + @Override + public long length() { + return length; + } + + @Override + public void readFully(long position, ByteBuffer buffer) throws IOException { + int len = buffer.remaining(); + long pos = position; + while (buffer.hasRemaining()) { + int read = channel.read(buffer, pos); + if (read < 0) { + throw new EOFException("EOF reading " + len + " bytes at position " + position); + } + pos += read; + } + } + + @Override + public void close() throws IOException { + channel.close(); + } + } + + /** A {@link NativeWritable} over a local file. */ + private static final class StreamWritable implements NativeWritable { + private final OutputStream out; + + StreamWritable(Path path) throws IOException { + this.out = Files.newOutputStream( + path, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING, StandardOpenOption.WRITE); + } + + @Override + public void write(byte[] buffer, int offset, int length) throws IOException { + out.write(buffer, offset, length); + } + + @Override + public void flush() throws IOException { + out.flush(); + } + + @Override + public void close() throws IOException { + out.close(); + } + } + + private void writePeopleFile(Session session, BufferAllocator allocator, Path path) throws IOException { + Schema schema = personSchema(); + VortexWriteSummary summary; + try (StreamWritable writable = new StreamWritable(path)) { + try (VortexWriter writer = VortexWriter.create(session, writable, schema, allocator); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + VarCharVector nameVec = (VarCharVector) root.getVector("name"); + IntVector ageVec = (IntVector) root.getVector("age"); + nameVec.allocateNew(3); + ageVec.allocateNew(3); + nameVec.setSafe(0, "Alice".getBytes(UTF_8)); + nameVec.setSafe(1, "Bob".getBytes(UTF_8)); + nameVec.setSafe(2, "Carol".getBytes(UTF_8)); + ageVec.setSafe(0, 30); + ageVec.setSafe(1, 25); + ageVec.setSafe(2, 40); + root.setRowCount(3); + + try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator); + ArrowSchema arrowSchemaFfi = ArrowSchema.allocateNew(allocator)) { + Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchemaFfi); + writer.writeBatch(arrowArray.memoryAddress(), arrowSchemaFfi.memoryAddress()); + } + + // The byte counter must be live for stream writers too: bounded while + // the write task is still flushing, exact once finished. + long bytesWhileOpen = writer.bytesWritten(); + summary = writer.finish(); + assertTrue(bytesWhileOpen >= 0 && bytesWhileOpen <= summary.fileSize()); + assertEquals(summary.fileSize(), writer.bytesWritten()); + } + } + // Everything the writer counted must have reached the caller-provided sink. + assertEquals(Files.size(path), summary.fileSize()); + } + + @Test + public void testWritableThenReadableRoundTrip() throws IOException { + Path path = tempDir.resolve("bridge_roundtrip.vortex"); + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + + writePeopleFile(session, allocator, path); + assertTrue(Files.size(path) > 0, "stream writer should have produced bytes"); + + try (FileChannelReadable readable = new FileChannelReadable(path)) { + DataSource ds = DataSource.open(session, readable); + assertEquals(new DataSource.RowCount.Exact(3L), ds.rowCount()); + + Scan scan = ds.scan(ScanOptions.of()); + int rows = 0; + while (scan.hasNext()) { + Partition p = scan.next(); + try (ArrowReader reader = p.scanArrow(allocator)) { + while (reader.loadNextBatch()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + ViewVarCharVector nameOut = (ViewVarCharVector) root.getVector("name"); + IntVector ageOut = (IntVector) root.getVector("age"); + if (rows == 0) { + assertEquals("Alice", nameOut.getObject(0).toString()); + assertEquals(30, ageOut.get(0)); + } + rows += root.getRowCount(); + } + } + } + assertEquals(3, rows); + } + } + + @Test + public void testMultipleReadables() throws IOException { + Path first = tempDir.resolve("bridge_a.vortex"); + Path second = tempDir.resolve("bridge_b.vortex"); + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + + writePeopleFile(session, allocator, first); + writePeopleFile(session, allocator, second); + + try (FileChannelReadable firstReadable = new FileChannelReadable(first); + FileChannelReadable secondReadable = new FileChannelReadable(second)) { + DataSource ds = DataSource.open(session, List.of(firstReadable, secondReadable)); + // Only the first file is opened eagerly, so the count is an estimate until scanned. + assertEquals(6L, ds.rowCount().asOptional().orElseThrow()); + + Scan scan = ds.scan(ScanOptions.of()); + long rows = 0; + while (scan.hasNext()) { + Partition p = scan.next(); + try (ArrowReader reader = p.scanArrow(allocator)) { + while (reader.loadNextBatch()) { + rows += reader.getVectorSchemaRoot().getRowCount(); + } + } + } + assertEquals(6, rows); + } + } + + @Test + public void testDuplicateReadableNamesAreRejected() throws IOException { + Path path = tempDir.resolve("bridge_duplicate.vortex"); + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + + writePeopleFile(session, allocator, path); + + try (FileChannelReadable first = new FileChannelReadable("/bridge_duplicate.vortex", path); + FileChannelReadable second = new FileChannelReadable("bridge_duplicate.vortex", path)) { + RuntimeException thrown = + assertThrows(RuntimeException.class, () -> DataSource.open(session, List.of(first, second))); + assertTrue( + thrown.getMessage().contains("multiple Java readables normalize to path"), + "error should identify the normalized path collision, got: " + thrown.getMessage()); + } + } + + @Test + public void testLargeRoundTrip() throws IOException { + Path path = tempDir.resolve("bridge_large.vortex"); + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + Schema schema = personSchema(); + + int batches = 20; + int rowsPerBatch = 5_000; + try (StreamWritable writable = new StreamWritable(path)) { + try (VortexWriter writer = VortexWriter.create(session, writable, schema, allocator); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + VarCharVector nameVec = (VarCharVector) root.getVector("name"); + IntVector ageVec = (IntVector) root.getVector("age"); + for (int batch = 0; batch < batches; batch++) { + nameVec.allocateNew(rowsPerBatch); + ageVec.allocateNew(rowsPerBatch); + for (int row = 0; row < rowsPerBatch; row++) { + nameVec.setSafe(row, ("name-" + batch + "-" + row).getBytes(UTF_8)); + ageVec.setSafe(row, batch * rowsPerBatch + row); + } + root.setRowCount(rowsPerBatch); + try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator); + ArrowSchema arrowSchemaFfi = ArrowSchema.allocateNew(allocator)) { + Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchemaFfi); + writer.writeBatch(arrowArray.memoryAddress(), arrowSchemaFfi.memoryAddress()); + } + } + } + } + + try (FileChannelReadable readable = new FileChannelReadable(path)) { + DataSource ds = DataSource.open(session, List.of(readable), 4); + Scan scan = ds.scan(ScanOptions.of()); + long rows = 0; + long ageSum = 0; + while (scan.hasNext()) { + Partition p = scan.next(); + try (ArrowReader reader = p.scanArrow(allocator)) { + while (reader.loadNextBatch()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + IntVector ageOut = (IntVector) root.getVector("age"); + for (int i = 0; i < root.getRowCount(); i++) { + ageSum += ageOut.get(i); + } + rows += root.getRowCount(); + } + } + } + long total = (long) batches * rowsPerBatch; + assertEquals(total, rows); + assertEquals(total * (total - 1) / 2, ageSum, "ages should be 0..N-1 exactly once"); + } + } + + @Test + public void testReadableExceptionPropagates() { + Session session = Session.create(); + NativeReadable failing = new NativeReadable() { + @Override + public String name() { + return "memory/failing.vortex"; + } + + @Override + public long length() { + return 1024; + } + + @Override + public void readFully(long position, ByteBuffer buffer) throws IOException { + throw new IOException("boom: injected read failure"); + } + + @Override + public void close() {} + }; + + RuntimeException thrown = assertThrows(RuntimeException.class, () -> DataSource.open(session, failing)); + assertTrue( + thrown.getMessage().contains("boom: injected read failure"), + "error should carry the Java exception message, got: " + thrown.getMessage()); + } +} diff --git a/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java b/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java index 03869ae8598..1bb51f8799f 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java +++ b/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java @@ -14,6 +14,7 @@ import dev.vortex.api.Scan; import dev.vortex.api.ScanOptions; import dev.vortex.api.Session; +import dev.vortex.api.VortexWriteSummary; import dev.vortex.api.VortexWriter; import dev.vortex.arrow.ArrowAllocation; import java.io.IOException; @@ -30,6 +31,7 @@ import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.types.pojo.ArrowType; @@ -157,6 +159,8 @@ public void testWriteBatch() throws IOException { Schema schema = personSchema(); Session session = Session.create(); + VortexWriteSummary summary; + long bytesWhileOpen; try (VortexWriter writer = VortexWriter.create(session, writePath, schema, new HashMap<>(), allocator); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { VarCharVector nameVec = (VarCharVector) root.getVector("name"); @@ -179,9 +183,25 @@ public void testWriteBatch() throws IOException { Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchemaFfi); writer.writeBatch(arrowArray.memoryAddress(), arrowSchemaFfi.memoryAddress()); } + bytesWhileOpen = writer.bytesWritten(); + summary = writer.finish(); + assertEquals(summary.fileSize(), writer.bytesWritten()); } assertTrue(Files.exists(outputPath), "output file should exist"); + assertTrue(bytesWhileOpen >= 0); + assertTrue(bytesWhileOpen <= summary.fileSize()); + assertEquals(Files.size(outputPath), summary.fileSize()); + assertEquals(3L, summary.rowCount()); + assertEquals(2, summary.columnStatistics().size()); + assertEquals(0, summary.columnStatistics().get(0).columnIndex()); + assertTrue(summary.columnStatistics().get(0).compressedSize() > 0); + assertEquals(3L, summary.columnStatistics().get(0).valueCount()); + assertEquals(0L, summary.columnStatistics().get(0).nullValueCount().orElseThrow()); + assertEquals("Alice", summary.columnStatistics().get(0).lowerBound().orElseThrow()); + assertEquals("Carol", summary.columnStatistics().get(0).upperBound().orElseThrow()); + assertEquals(25, summary.columnStatistics().get(1).lowerBound().orElseThrow()); + assertEquals(40, summary.columnStatistics().get(1).upperBound().orElseThrow()); DataSource ds = DataSource.open(session, writePath); assertEquals(new DataSource.RowCount.Exact(3L), ds.rowCount()); @@ -192,7 +212,7 @@ public void testWriteBatch() throws IOException { try (ArrowReader reader = p.scanArrow(allocator)) { reader.loadNextBatch(); VectorSchemaRoot resultRoot = reader.getVectorSchemaRoot(); - VarCharVector nameOut = (VarCharVector) resultRoot.getVector("name"); + ViewVarCharVector nameOut = (ViewVarCharVector) resultRoot.getVector("name"); IntVector ageOut = (IntVector) resultRoot.getVector("age"); assertEquals("Alice", nameOut.getObject(0).toString()); assertEquals("Bob", nameOut.getObject(1).toString()); diff --git a/java/vortex-spark/README.md b/java/vortex-spark/README.md new file mode 100644 index 00000000000..c6e8c9345bd --- /dev/null +++ b/java/vortex-spark/README.md @@ -0,0 +1,115 @@ +# vortex-spark + +A Spark DataSource V2 connector for reading and writing [Vortex](https://vortex.dev) files. +It registers itself under the format name `vortex` and supports both the DataFrame API and +Spark SQL. + +Two flavors are published to Maven Central: + +| Artifact | Spark | Scala | +|-------------------------------|-----------|-------| +| `dev.vortex:vortex-spark_2.13` | Spark 4.x | 2.13 | +| `dev.vortex:vortex-spark_2.12` | Spark 3.5.x | 2.12 | + +Use the `all` classifier JAR (e.g. `vortex-spark_2.13-0.78.0-all.jar`). It is self-contained: +it bundles the Vortex JNI bindings, native libraries for Linux (x86_64 and aarch64) and macOS +(aarch64), and relocates its Arrow, Guava, and Jackson dependencies to avoid classpath +conflicts with Spark. The thin (unclassified) JAR does not work on its own because it +references relocated classes that only ship in the `all` JAR. + +## Getting Vortex into Spark + +Pass the `all` JAR to `spark-shell`, `spark-submit`, or `pyspark` with `--jars`. Spark accepts +either a local path or a URL, so you can point directly at Maven Central: + +```shell +spark-shell --jars https://repo1.maven.org/maven2/dev/vortex/vortex-spark_2.13/0.78.0/vortex-spark_2.13-0.78.0-all.jar +``` + +Or configure it on the session builder, e.g. in PySpark: + +```python +spark = ( + SparkSession.builder + .config("spark.jars", "/path/to/vortex-spark_2.13-0.78.0-all.jar") + .getOrCreate() +) +``` + +Note that `--packages dev.vortex:vortex-spark_2.13:0.78.0` does not work: `--packages` cannot +select the `all` classifier and resolves the thin JAR, which fails at runtime with +`NoClassDefFoundError: dev/vortex/relocated/...`. + +To depend on the connector from a JVM project instead, add the `all` classifier to the +dependency: + +```kotlin +implementation("dev.vortex:vortex-spark_2.13:0.78.0:all") +``` + +## Usage + +Paths may be local filesystem paths (`/path/to/data`) or URLs (`file:///path/to/data`, +`s3://bucket/path/to/data`). + +### DataFrame API + +```java +// Write +df.write() + .format("vortex") + .option("path", "/path/to/output") + .mode(SaveMode.Overwrite) + .save(); + +// Read a single file or a directory of .vortex files +Dataset df = spark.read() + .format("vortex") + .option("path", "/path/to/output") + .load(); +``` + +### Spark SQL + +```sql +-- Query existing Vortex files through a temporary view +CREATE TEMPORARY VIEW people +USING vortex +OPTIONS (path '/path/to/data'); + +SELECT name, age FROM people WHERE age > 30; + +-- Create a table and write to it. With a LOCATION clause the table is external, +-- backed by the files at that path; without one it is managed by Spark. +CREATE TABLE student (id INT, name STRING, age INT) +USING vortex; + +INSERT INTO student VALUES (1, 'Alice', 20), (2, 'Bob', 21); + +SELECT * FROM student; +``` + +On Spark 3.5, `CREATE TABLE ... USING vortex` additionally requires replacing the session +catalog with `spark.sql.catalog.spark_catalog=dev.vortex.spark.VortexSessionCatalog`, +because Spark 3.5's built-in catalog cannot read tables backed by a DataSource V2-only +connector. The extension delegates everything else to the built-in session catalog and +leaves tables of other providers untouched; it is not needed on Spark 4. + +### Direct file queries + +Spark's built-in ``SELECT * FROM format.`path` `` syntax only works for built-in file +formats, so the connector ships a path-based catalog that provides the equivalent for +Vortex. Register it under the name `vortex` with this session config: + +``` +spark.sql.catalog.vortex=dev.vortex.spark.VortexCatalog +``` + +Then query (or insert into) Vortex files directly by path: + +```sql +SELECT * FROM vortex.`/path/to/data`; +``` + +See the [Spark user guide](https://docs.vortex.dev/user-guide/spark.html) for the full +documentation, including supported types, write options, and S3 configuration. diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/ArrowUtils.java b/java/vortex-spark/src/main/java/dev/vortex/spark/ArrowUtils.java index 6b3c130c467..6e0166a2873 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/ArrowUtils.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/ArrowUtils.java @@ -4,7 +4,6 @@ package dev.vortex.spark; import dev.vortex.relocated.org.apache.arrow.vector.types.DateUnit; -import dev.vortex.relocated.org.apache.arrow.vector.types.FloatingPointPrecision; import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit; import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType; import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.Field; @@ -18,8 +17,18 @@ * Utility class for converting Arrow types to Spark SQL data types. * *

This class provides static methods to convert Arrow field definitions and type definitions into their - * corresponding Spark SQL DataType representations. It handles the mapping between Arrow's type system and Spark's type - * system, including complex types like structs and arrays. + * corresponding Spark SQL DataType representations. The mapping matches Spark 4.1's own {@code ArrowUtils}, extended + * where Vortex produces types Spark's mapping does not cover: + * + *

    + *
  • the Arrow view types map to their logical Spark types: {@code Utf8View} to {@code StringType}, + * {@code BinaryView} to {@code BinaryType}, and {@code ListView} to {@code ArrayType} + *
  • timestamps of every unit map to {@code TimestampType}/{@code TimestampNTZType}; + * {@link dev.vortex.spark.read.VortexArrowColumnVector} normalizes non-microsecond values on read + *
+ * + *

The one intentional divergence from Spark 4.1 is Arrow's Time type: Spark's {@code TimeType} only exists in Spark + * 4.1+, and this module compiles a single source set against both Spark 3.5 and 4.1, so Time is rejected. */ public final class ArrowUtils { private ArrowUtils() {} @@ -27,8 +36,8 @@ private ArrowUtils() {} /** * Converts an Arrow Field to a Spark SQL DataType. * - *

This method handles complex types like structs and arrays by recursively converting their child fields. For - * primitive types, it delegates to {@link #fromArrowType(ArrowType)}. + *

This method handles nested types like structs, lists and maps by recursively converting their child fields. + * For non-nested types, it delegates to {@link #fromArrowType(ArrowType)}. * * @param field the Arrow field to convert * @return the corresponding Spark SQL DataType @@ -43,11 +52,19 @@ public static DataType fromArrowField(Field field) { return new StructField(child.getName(), dt, child.isNullable(), Metadata.empty()); }) .collect(Collectors.toList())); - case List: { + case List: + case ListView: { Field elementField = field.getChildren().get(0); DataType elementType = fromArrowField(elementField); return DataTypes.createArrayType(elementType, elementField.isNullable()); } + case Map: { + Field entries = field.getChildren().get(0); + Field keyField = entries.getChildren().get(0); + Field valueField = entries.getChildren().get(1); + return DataTypes.createMapType( + fromArrowField(keyField), fromArrowField(valueField), valueField.isNullable()); + } default: return fromArrowType(field.getType()); } @@ -56,13 +73,12 @@ public static DataType fromArrowField(Field field) { /** * Converts an Arrow type to a Spark SQL DataType. * - *

This method maps primitive Arrow types to their corresponding Spark SQL types. It supports most common Arrow - * types including integers, floating point numbers, strings, binary data, dates, timestamps, decimals, and nulls. + *

This method maps non-nested Arrow types to their corresponding Spark SQL types, following Spark 4.1's own + * Arrow type mapping plus the view types (see the class documentation). * * @param dt the Arrow type to convert * @return the corresponding Spark SQL DataType - * @throws UnsupportedOperationException if the Arrow type configuration is not supported - * @throws RuntimeException if the Arrow type is not recognized + * @throws UnsupportedOperationException if the Arrow type has no Spark representation */ public static DataType fromArrowType(ArrowType dt) { switch (dt.getTypeID()) { @@ -70,26 +86,31 @@ public static DataType fromArrowType(ArrowType dt) { return DataTypes.BooleanType; case Int: { ArrowType.Int intType = (ArrowType.Int) dt; - if (intType.getIsSigned() && intType.getBitWidth() == 8) { - return DataTypes.ByteType; - } else if (intType.getIsSigned() && intType.getBitWidth() == 16) { - return DataTypes.ShortType; - } else if (intType.getIsSigned() && intType.getBitWidth() == 32) { - return DataTypes.IntegerType; - } else if (intType.getIsSigned() && intType.getBitWidth() == 64) { - return DataTypes.LongType; - } else { - throw new UnsupportedOperationException("Unsupported Arrow type: " + dt); + if (!intType.getIsSigned()) { + throw new UnsupportedOperationException("Unsupported Arrow unsigned integer type: " + dt); + } + switch (intType.getBitWidth()) { + case 8: + return DataTypes.ByteType; + case 16: + return DataTypes.ShortType; + case 32: + return DataTypes.IntegerType; + case 64: + return DataTypes.LongType; + default: + throw new UnsupportedOperationException("Unsupported Arrow integer bit width: " + dt); } } case FloatingPoint: { ArrowType.FloatingPoint floatType = (ArrowType.FloatingPoint) dt; - if (floatType.getPrecision() == FloatingPointPrecision.SINGLE) { - return DataTypes.FloatType; - } else if (floatType.getPrecision() == FloatingPointPrecision.DOUBLE) { - return DataTypes.DoubleType; - } else { - throw new UnsupportedOperationException("Unsupported Arrow type: " + dt); + switch (floatType.getPrecision()) { + case SINGLE: + return DataTypes.FloatType; + case DOUBLE: + return DataTypes.DoubleType; + default: + throw new UnsupportedOperationException("Unsupported Arrow float precision: " + dt); } } case Decimal: { @@ -98,30 +119,52 @@ public static DataType fromArrowType(ArrowType dt) { } case Utf8: case LargeUtf8: + case Utf8View: return DataTypes.StringType; case Binary: case LargeBinary: + case BinaryView: return DataTypes.BinaryType; case Date: { ArrowType.Date dateType = (ArrowType.Date) dt; if (dateType.getUnit() == DateUnit.DAY) { return DataTypes.DateType; - } else { - throw new UnsupportedOperationException("Unsupported Arrow type: " + dt); } + throw new UnsupportedOperationException("Unsupported Arrow date unit: " + dt); } case Timestamp: { + // Spark timestamps are physically microseconds; VortexArrowColumnVector's accessor + // normalizes second/millisecond/nanosecond values on read. ArrowType.Timestamp ts = (ArrowType.Timestamp) dt; - if (ts.getUnit() == TimeUnit.MICROSECOND) { - return ts.getTimezone() != null ? DataTypes.TimestampType : DataTypes.TimestampNTZType; - } else { - throw new UnsupportedOperationException("Unsupported Arrow type: " + dt); - } + return ts.getTimezone() != null ? DataTypes.TimestampType : DataTypes.TimestampNTZType; } case Null: return DataTypes.NullType; + case Interval: { + ArrowType.Interval interval = (ArrowType.Interval) dt; + switch (interval.getUnit()) { + case YEAR_MONTH: + return DataTypes.createYearMonthIntervalType(); + case MONTH_DAY_NANO: + return DataTypes.CalendarIntervalType; + default: + throw new UnsupportedOperationException("Unsupported Arrow interval unit: " + dt); + } + } + case Duration: { + ArrowType.Duration duration = (ArrowType.Duration) dt; + if (duration.getUnit() != TimeUnit.MICROSECOND) { + throw new UnsupportedOperationException("Unsupported Arrow duration unit: " + dt); + } + return DataTypes.createDayTimeIntervalType(); + } + case Time: + // Spark 3.5 has no TIME type (TimeType only exists in Spark 4.1+), and this + // module compiles a single source set against both Spark versions. + throw new UnsupportedOperationException("Arrow Time type has no Spark 3.5 representation: " + dt); default: - throw new IllegalArgumentException("Unsupported Arrow type: " + dt); + throw new UnsupportedOperationException( + "Unsupported Arrow type: " + dt + " (type id: " + dt.getTypeID() + ")"); } } } diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexCatalog.java b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexCatalog.java new file mode 100644 index 00000000000..3d496a4cee0 --- /dev/null +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexCatalog.java @@ -0,0 +1,119 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark; + +import java.util.Map; +import org.apache.spark.sql.catalyst.analysis.NoSuchTableException; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.connector.catalog.Table; +import org.apache.spark.sql.connector.catalog.TableCatalog; +import org.apache.spark.sql.connector.catalog.TableChange; +import org.apache.spark.sql.connector.expressions.Transform; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; + +/** + * A path-based Spark catalog for querying Vortex files directly from SQL. + * + *

Spark only supports {@code SELECT * FROM format.`path`} syntax for built-in file formats, so this catalog provides + * the equivalent for Vortex. Register it under the name {@code vortex}: + * + * 

spark.sql.catalog.vortex=dev.vortex.spark.VortexCatalog
+ * + *

then query a Vortex file, or a directory of Vortex files, directly by path: + * + * 

SELECT * FROM vortex.`/path/to/data`;
+ * + *

The table identifier must look like a path — contain a {@code /} — and resolves to the same table a + * {@code spark.read.format("vortex")} load of that path would produce, so reads, writes ({@code INSERT INTO + * vortex.`/path/to/data`}), and pushdown all behave identically. The catalog holds no state and supports no DDL. + */ +public final class VortexCatalog implements TableCatalog { + private static final String PATH_KEY = "path"; + + private String name = "vortex"; + + /** + * Creates a new catalog instance. + * + *

This no-argument constructor is required for Spark to instantiate the catalog through reflection from the + * {@code spark.sql.catalog.} configuration. + */ + public VortexCatalog() {} + + @Override + public void initialize(String name, CaseInsensitiveStringMap options) { + this.name = name; + } + + @Override + public String name() { + return name; + } + + /** + * Returns no identifiers: this catalog holds no state, tables are addressed by path. + * + * @param namespace the namespace to list, ignored + * @return an empty array + */ + @Override + public Identifier[] listTables(String[] namespace) { + return new Identifier[0]; + } + + /** + * Loads the Vortex file or directory of Vortex files at the path given by the identifier name. + * + * @param ident identifier whose name is a filesystem path or URL, e.g. {@code vortex.`/path/to/data`} + * @return a table backed by the Vortex files at the path + * @throws NoSuchTableException if the identifier does not look like a path, or the path cannot be read + */ + @SuppressWarnings("deprecation") + @Override + public Table loadTable(Identifier ident) throws NoSuchTableException { + String path = ident.name(); + if (ident.namespace().length != 0 || !path.contains("/")) { + throw new NoSuchTableException(ident); + } + var options = new CaseInsensitiveStringMap(Map.of(PATH_KEY, path)); + var provider = new VortexDataSourceV2(); + StructType schema; + Transform[] partitioning; + try { + schema = provider.inferSchema(options); + partitioning = provider.inferPartitioning(options); + } catch (RuntimeException e) { + // Missing or unreadable paths surface as "table not found" to SQL users. + throw new NoSuchTableException(ident); + } + return provider.getTable(schema, partitioning, Map.of(PATH_KEY, path)); + } + + /** Unsupported: tables are addressed by path, create them by writing data with the {@code vortex} format. */ + @Override + public Table createTable( + Identifier ident, StructType schema, Transform[] partitions, Map properties) { + throw new UnsupportedOperationException( + "VortexCatalog does not support CREATE TABLE, write data to the path instead"); + } + + /** Unsupported: this catalog holds no table metadata to alter. */ + @Override + public Table alterTable(Identifier ident, TableChange... changes) { + throw new UnsupportedOperationException("VortexCatalog does not support ALTER TABLE"); + } + + /** Unsupported: this catalog never drops data, returns false. */ + @Override + public boolean dropTable(Identifier ident) { + return false; + } + + /** Unsupported: this catalog holds no table metadata to rename. */ + @Override + public void renameTable(Identifier oldIdent, Identifier newIdent) { + throw new UnsupportedOperationException("VortexCatalog does not support RENAME TABLE"); + } +} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexDataSourceV2.java b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexDataSourceV2.java index b3d7d637504..a5ce4e3b1ce 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexDataSourceV2.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexDataSourceV2.java @@ -126,10 +126,14 @@ public StructType inferSchema(CaseInsensitiveStringMap options) { * explicit partitioning. Returning identity transforms here lets downstream components (notably * {@link dev.vortex.spark.read.VortexScanBuilder}) tell which schema columns are encoded in the directory layout * rather than stored inside the Vortex files, which matters for predicate pushdown. + * + *

The options may contain no path at all: when a managed table is created ({@code CREATE TABLE ... USING vortex} + * without a {@code LOCATION} clause), Spark's session catalog calls this before it has assigned the table's + * warehouse location. No transforms are inferred in that case. */ @Override public Transform[] inferPartitioning(CaseInsensitiveStringMap options) { - var paths = getPaths(options); + var paths = getPathsOrEmpty(options); if (paths.isEmpty()) { return new Transform[0]; } @@ -157,16 +161,20 @@ public Transform[] inferPartitioning(CaseInsensitiveStringMap options) { *

This method creates a VortexWritableTable that can be used to both read from and write to Vortex files. The * partitioning parameter is currently ignored. * + *

The properties may contain no path at all: when a managed table is created ({@code CREATE TABLE ... USING + * vortex} without a {@code LOCATION} clause), Spark's session catalog validates the table before it has assigned + * the table's warehouse location. The returned table then has no paths; once the table is loaded for reads or + * writes, Spark always supplies the resolved table location as the {@code path} property. + * * @param schema the table schema * @param partitioning table partitioning transforms * @param properties the table properties containing file paths and other options * @return a VortexTable instance for reading and writing data - * @throws RuntimeException if required path properties are missing */ @Override public Table getTable(StructType schema, Transform[] partitioning, Map properties) { var uncased = new CaseInsensitiveStringMap(properties); - ImmutableList paths = getPaths(uncased); + ImmutableList paths = getPathsOrEmpty(uncased); return new VortexTable(paths, schema, buildDataSourceOptions(properties), partitioning); } @@ -210,6 +218,13 @@ private Map buildDataSourceOptions(Map propertie return options.build(); } + private static ImmutableList getPathsOrEmpty(CaseInsensitiveStringMap uncased) { + if (!uncased.containsKey(PATH_KEY) && !uncased.containsKey(PATHS_KEY)) { + return ImmutableList.of(); + } + return getPaths(uncased); + } + private static ImmutableList getPaths(CaseInsensitiveStringMap uncased) { if (uncased.containsKey(PATH_KEY)) { return ImmutableList.of(uncased.get(PATH_KEY)); diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionCatalog.java b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionCatalog.java new file mode 100644 index 00000000000..1d546880cc1 --- /dev/null +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/VortexSessionCatalog.java @@ -0,0 +1,79 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark; + +import java.util.Map; +import org.apache.spark.sql.catalyst.analysis.NoSuchNamespaceException; +import org.apache.spark.sql.catalyst.analysis.NoSuchTableException; +import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException; +import org.apache.spark.sql.connector.catalog.DelegatingCatalogExtension; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.connector.catalog.Table; +import org.apache.spark.sql.connector.catalog.TableCatalog; +import org.apache.spark.sql.connector.expressions.Transform; +import org.apache.spark.sql.types.StructType; + +/** + * A session catalog extension that resolves {@code USING vortex} tables through the Vortex DataSource V2 connector. + * + *

Spark 3.5's built-in session catalog resolves the tables it stores through the V1 {@code DataSource} path, which + * rejects DataSource-V2-only connectors like Vortex, so {@code CREATE TABLE ... USING vortex} tables cannot be read + * back. (Spark 4 resolves them through the V2 provider directly and needs none of this.) Registering this extension as + * the session catalog fixes that on Spark 3.5: + * + * 

spark.sql.catalog.spark_catalog=dev.vortex.spark.VortexSessionCatalog
+ * + *

All operations are delegated to the built-in session catalog — table metadata lives wherever it normally would, + * including the Hive metastore — but any table whose provider is {@code vortex} is loaded as a Vortex DataSource V2 + * table, backed by the files at the table's location. Tables of every other provider are untouched. + */ +public final class VortexSessionCatalog extends DelegatingCatalogExtension { + + /** + * Creates a new session catalog extension. + * + *

This no-argument constructor is required for Spark to instantiate the catalog through reflection from the + * {@code spark.sql.catalog.spark_catalog} configuration. + */ + public VortexSessionCatalog() {} + + @Override + public Table loadTable(Identifier ident) throws NoSuchTableException { + return asVortexTableIfVortex(super.loadTable(ident)); + } + + /** + * Creates the table in the delegate session catalog, then returns it resolved through the Vortex connector when its + * provider is {@code vortex}. + * + *

The conversion matters for {@code CREATE TABLE ... AS SELECT}: Spark writes the query result into the table + * returned here, which must therefore support V2 writes. + */ + @SuppressWarnings("deprecation") + @Override + public Table createTable( + Identifier ident, StructType schema, Transform[] partitions, Map properties) + throws TableAlreadyExistsException, NoSuchNamespaceException { + return asVortexTableIfVortex(super.createTable(ident, schema, partitions, properties)); + } + + /** + * Rebuilds a session-catalog table as a Vortex DataSource V2 table when its provider is {@code vortex} and it has a + * location; returns every other table unchanged. The schema and partitioning stored in the catalog are used as-is, + * no file needs to be opened. + */ + @SuppressWarnings("deprecation") + private static Table asVortexTableIfVortex(Table table) { + Map properties = table.properties(); + VortexDataSourceV2 provider = new VortexDataSourceV2(); + if (!provider.shortName().equalsIgnoreCase(properties.get(TableCatalog.PROP_PROVIDER))) { + return table; + } + String location = properties.get(TableCatalog.PROP_LOCATION); + if (location == null) { + return table; + } + return provider.getTable(table.schema(), table.partitioning(), Map.of("path", location)); + } +} diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java index 99e45feab72..be191a0ed43 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java @@ -7,12 +7,16 @@ import com.jakewharton.nopen.annotation.Open; import dev.vortex.relocated.org.apache.arrow.vector.*; import dev.vortex.relocated.org.apache.arrow.vector.complex.ListVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.ListViewVector; import dev.vortex.relocated.org.apache.arrow.vector.complex.MapVector; import dev.vortex.relocated.org.apache.arrow.vector.complex.StructVector; import dev.vortex.relocated.org.apache.arrow.vector.holders.NullableLargeVarCharHolder; import dev.vortex.relocated.org.apache.arrow.vector.holders.NullableVarCharHolder; +import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit; +import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType; import dev.vortex.spark.ArrowUtils; import org.apache.spark.sql.types.DataType; +import org.apache.spark.sql.types.DataTypes; import org.apache.spark.sql.types.Decimal; import org.apache.spark.sql.vectorized.ColumnVector; import org.apache.spark.sql.vectorized.ColumnarArray; @@ -23,12 +27,17 @@ * Spark ColumnVector implementation that wraps Apache Arrow vectors from Vortex data. *

* This class provides a bridge between Vortex's Arrow-based data representation and Spark's - * ColumnVector interface. It supports all major Arrow data types including primitives, strings, - * binary data, decimals, dates, timestamps, arrays, maps, and structs. + * ColumnVector interface. It supports the same Arrow vector types as Spark 4.1's own + * {@code ArrowColumnVector} — booleans, signed integers, single/double floats, 128-bit decimals, + * strings and binary (regular and large), day dates, timestamps of every unit with and without + * timezone (normalized to microseconds), microsecond durations, year-month and month-day-nano + * intervals, nulls, lists, maps, and structs — plus the Arrow view types (Utf8View, BinaryView, + * ListView) that Vortex produces natively. *

- * The implementation uses type-specific accessors to efficiently retrieve values from the - * underlying Arrow vectors while maintaining Spark's expected API contract. - * + * Arrow types outside that set (unsigned integers, half floats, 256-bit decimals, + * non-microsecond durations, Time, fixed-size binary and lists, large lists, unions, run-end and + * dictionary encodings) are rejected with a descriptive {@link UnsupportedOperationException}. + * * @see ColumnVector * @see ValueVector */ @@ -39,7 +48,7 @@ public class VortexArrowColumnVector extends ColumnVector { /** * Returns the underlying Apache Arrow ValueVector wrapped by this column vector. - * + * * @return the Arrow ValueVector containing the actual data */ public ValueVector getValueVector() { @@ -48,7 +57,7 @@ public ValueVector getValueVector() { /** * Returns whether this column contains any null values. - * + * * @return true if the column contains at least one null value, false otherwise */ @Override @@ -58,7 +67,7 @@ public boolean hasNull() { /** * Returns the total number of null values in this column. - * + * * @return the count of null values */ @Override @@ -76,7 +85,7 @@ public void close() {} /** * Returns whether the value at the specified row is null. - * + * * @param rowId the row index to check * @return true if the value at rowId is null, false otherwise */ @@ -87,7 +96,7 @@ public boolean isNullAt(int rowId) { /** * Returns the boolean value at the specified row. - * + * * @param rowId the row index * @return the boolean value at rowId * @throws UnsupportedOperationException if this column is not of boolean type @@ -99,7 +108,7 @@ public boolean getBoolean(int rowId) { /** * Returns the byte value at the specified row. - * + * * @param rowId the row index * @return the byte value at rowId * @throws UnsupportedOperationException if this column is not of byte type @@ -111,7 +120,7 @@ public byte getByte(int rowId) { /** * Returns the short value at the specified row. - * + * * @param rowId the row index * @return the short value at rowId * @throws UnsupportedOperationException if this column is not of short type @@ -123,7 +132,7 @@ public short getShort(int rowId) { /** * Returns the int value at the specified row. - * + * * @param rowId the row index * @return the int value at rowId * @throws UnsupportedOperationException if this column is not of int type @@ -135,7 +144,7 @@ public int getInt(int rowId) { /** * Returns the long value at the specified row. - * + * * @param rowId the row index * @return the long value at rowId * @throws UnsupportedOperationException if this column is not of long type @@ -147,7 +156,7 @@ public long getLong(int rowId) { /** * Returns the float value at the specified row. - * + * * @param rowId the row index * @return the float value at rowId * @throws UnsupportedOperationException if this column is not of float type @@ -159,7 +168,7 @@ public float getFloat(int rowId) { /** * Returns the double value at the specified row. - * + * * @param rowId the row index * @return the double value at rowId * @throws UnsupportedOperationException if this column is not of double type @@ -171,7 +180,7 @@ public double getDouble(int rowId) { /** * Returns the decimal value at the specified row with the given precision and scale. - * + * * @param rowId the row index * @param precision the precision of the decimal * @param scale the scale of the decimal @@ -186,7 +195,7 @@ public Decimal getDecimal(int rowId, int precision, int scale) { /** * Returns the UTF8String value at the specified row. - * + * * @param rowId the row index * @return the UTF8String value at rowId, or null if the value is null * @throws UnsupportedOperationException if this column is not of string type @@ -199,7 +208,7 @@ public UTF8String getUTF8String(int rowId) { /** * Returns the binary data (byte array) at the specified row. - * + * * @param rowId the row index * @return the byte array at rowId, or null if the value is null * @throws UnsupportedOperationException if this column is not of binary type @@ -212,7 +221,7 @@ public byte[] getBinary(int rowId) { /** * Returns the array value at the specified row. - * + * * @param rowId the row index * @return the ColumnarArray at rowId, or null if the value is null * @throws UnsupportedOperationException if this column is not of array type @@ -225,7 +234,7 @@ public ColumnarArray getArray(int rowId) { /** * Returns the map value at the specified row. - * + * * @param rowId the row index * @return the ColumnarMap at rowId, or null if the value is null * @throws UnsupportedOperationException if this column is not of map type @@ -240,8 +249,9 @@ public ColumnarMap getMap(int rowId) { * Returns the child column at the specified ordinal. *

* This is used for complex types like structs where each field is represented - * as a child column. - * + * as a child column. For month-day-nano interval columns, the children follow the + * months/days/microseconds protocol documented on {@link ColumnVector#getInterval(int)}. + * * @param ordinal the index of the child column * @return the child VortexArrowColumnVector at the specified ordinal * @throws ArrayIndexOutOfBoundsException if ordinal is out of bounds @@ -256,7 +266,7 @@ public VortexArrowColumnVector getChild(int ordinal) { *

* This constructor is used internally for creating column vectors before * the underlying Arrow vector is available. - * + * * @param type the Spark DataType for this column */ VortexArrowColumnVector(DataType type) { @@ -268,7 +278,7 @@ public VortexArrowColumnVector getChild(int ordinal) { *

* This constructor automatically determines the appropriate Spark DataType from * the Arrow field and initializes the type-specific accessor. - * + * * @param vector the Arrow ValueVector to wrap * @throws UnsupportedOperationException if the vector type is not supported */ @@ -277,59 +287,78 @@ public VortexArrowColumnVector(ValueVector vector) { initAccessor(vector); } + private static VortexArrowColumnVector withAccessor(DataType type, ArrowVectorAccessor accessor) { + VortexArrowColumnVector column = new VortexArrowColumnVector(type); + column.accessor = accessor; + return column; + } + void initAccessor(ValueVector vector) { - if (vector instanceof BitVector) { - accessor = new VortexArrowColumnVector.BooleanAccessor((BitVector) vector); - } else if (vector instanceof TinyIntVector) { - accessor = new VortexArrowColumnVector.ByteAccessor((TinyIntVector) vector); - } else if (vector instanceof SmallIntVector) { - accessor = new VortexArrowColumnVector.ShortAccessor((SmallIntVector) vector); - } else if (vector instanceof IntVector) { - accessor = new VortexArrowColumnVector.IntAccessor((IntVector) vector); - } else if (vector instanceof BigIntVector) { - accessor = new VortexArrowColumnVector.LongAccessor((BigIntVector) vector); - } else if (vector instanceof Float4Vector) { - accessor = new VortexArrowColumnVector.FloatAccessor((Float4Vector) vector); - } else if (vector instanceof Float8Vector) { - accessor = new VortexArrowColumnVector.DoubleAccessor((Float8Vector) vector); - } else if (vector instanceof DecimalVector) { - accessor = new VortexArrowColumnVector.DecimalAccessor((DecimalVector) vector); - } else if (vector instanceof VarCharVector) { - accessor = new VortexArrowColumnVector.StringAccessor((VarCharVector) vector); - } else if (vector instanceof LargeVarCharVector) { - accessor = new VortexArrowColumnVector.LargeStringAccessor((LargeVarCharVector) vector); - } else if (vector instanceof VarBinaryVector) { - accessor = new VortexArrowColumnVector.BinaryAccessor((VarBinaryVector) vector); - } else if (vector instanceof LargeVarBinaryVector) { - accessor = new VortexArrowColumnVector.LargeBinaryAccessor((LargeVarBinaryVector) vector); - } else if (vector instanceof DateDayVector) { - accessor = new VortexArrowColumnVector.DateAccessor((DateDayVector) vector); - } else if (vector instanceof TimeStampMicroTZVector) { - accessor = new VortexArrowColumnVector.TimestampAccessor((TimeStampMicroTZVector) vector); - } else if (vector instanceof TimeStampMicroVector) { - accessor = new VortexArrowColumnVector.TimestampNTZAccessor((TimeStampMicroVector) vector); - } else if (vector instanceof MapVector) { - MapVector mapVector = (MapVector) vector; + if (vector instanceof BitVector bitVector) { + accessor = new VortexArrowColumnVector.BooleanAccessor(bitVector); + } else if (vector instanceof TinyIntVector tinyIntVector) { + accessor = new VortexArrowColumnVector.ByteAccessor(tinyIntVector); + } else if (vector instanceof SmallIntVector smallIntVector) { + accessor = new VortexArrowColumnVector.ShortAccessor(smallIntVector); + } else if (vector instanceof IntVector intVector) { + accessor = new VortexArrowColumnVector.IntAccessor(intVector); + } else if (vector instanceof BigIntVector bigIntVector) { + accessor = new VortexArrowColumnVector.LongAccessor(bigIntVector); + } else if (vector instanceof Float4Vector float4Vector) { + accessor = new VortexArrowColumnVector.FloatAccessor(float4Vector); + } else if (vector instanceof Float8Vector float8Vector) { + accessor = new VortexArrowColumnVector.DoubleAccessor(float8Vector); + } else if (vector instanceof DecimalVector decimalVector) { + accessor = new VortexArrowColumnVector.DecimalAccessor(decimalVector); + } else if (vector instanceof VarCharVector varCharVector) { + accessor = new VortexArrowColumnVector.StringAccessor(varCharVector); + } else if (vector instanceof LargeVarCharVector largeVarCharVector) { + accessor = new VortexArrowColumnVector.LargeStringAccessor(largeVarCharVector); + } else if (vector instanceof ViewVarCharVector viewVarCharVector) { + accessor = new VortexArrowColumnVector.StringViewAccessor(viewVarCharVector); + } else if (vector instanceof VarBinaryVector varBinaryVector) { + accessor = new VortexArrowColumnVector.BinaryAccessor(varBinaryVector); + } else if (vector instanceof LargeVarBinaryVector largeVarBinaryVector) { + accessor = new VortexArrowColumnVector.LargeBinaryAccessor(largeVarBinaryVector); + } else if (vector instanceof ViewVarBinaryVector viewVarBinaryVector) { + accessor = new VortexArrowColumnVector.BinaryViewAccessor(viewVarBinaryVector); + } else if (vector instanceof DateDayVector dateDayVector) { + accessor = new VortexArrowColumnVector.DateAccessor(dateDayVector); + } else if (vector instanceof TimeStampVector timeStampVector) { + // Covers all eight unit/timezone variants; values are normalized to microseconds. + accessor = new VortexArrowColumnVector.TimestampAccessor(timeStampVector); + } else if (vector instanceof MapVector mapVector) { + // MapVector extends ListVector, so this check must come first. accessor = new VortexArrowColumnVector.MapAccessor(mapVector); - } else if (vector instanceof ListVector) { - ListVector listVector = (ListVector) vector; + } else if (vector instanceof ListVector listVector) { accessor = new VortexArrowColumnVector.ArrayAccessor(listVector); - } else if (vector instanceof StructVector) { - StructVector structVector = (StructVector) vector; + } else if (vector instanceof ListViewVector listViewVector) { + accessor = new VortexArrowColumnVector.ListViewAccessor(listViewVector); + } else if (vector instanceof StructVector structVector) { accessor = new VortexArrowColumnVector.StructAccessor(structVector); childColumns = new VortexArrowColumnVector[structVector.size()]; for (int i = 0; i < childColumns.length; ++i) { childColumns[i] = new VortexArrowColumnVector(structVector.getVectorById(i)); } - } else if (vector instanceof NullVector) { - accessor = new VortexArrowColumnVector.NullAccessor((NullVector) vector); - } else if (vector instanceof IntervalYearVector) { - accessor = new VortexArrowColumnVector.IntervalYearAccessor((IntervalYearVector) vector); - } else if (vector instanceof DurationVector) { - accessor = new VortexArrowColumnVector.DurationAccessor((DurationVector) vector); + } else if (vector instanceof NullVector nullVector) { + accessor = new VortexArrowColumnVector.NullAccessor(nullVector); + } else if (vector instanceof IntervalYearVector intervalYearVector) { + accessor = new VortexArrowColumnVector.IntervalYearAccessor(intervalYearVector); + } else if (vector instanceof IntervalMonthDayNanoVector intervalMonthDayNanoVector) { + accessor = new VortexArrowColumnVector.IntervalMonthDayNanoAccessor(intervalMonthDayNanoVector); + // CalendarInterval values are read through ColumnVector.getInterval, which uses the + // months/days/microseconds child column protocol. + childColumns = new VortexArrowColumnVector[] { + withAccessor(DataTypes.IntegerType, new IntervalMonthsAccessor(intervalMonthDayNanoVector)), + withAccessor(DataTypes.IntegerType, new IntervalDaysAccessor(intervalMonthDayNanoVector)), + withAccessor(DataTypes.LongType, new IntervalMicrosAccessor(intervalMonthDayNanoVector)), + }; + } else if (vector instanceof DurationVector durationVector) { + accessor = new VortexArrowColumnVector.DurationAccessor(durationVector); } else { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException("Unsupported Arrow vector type: " + + vector.getClass().getSimpleName() + " for field " + vector.getField()); } } @@ -578,6 +607,24 @@ final UTF8String getUTF8String(int rowId) { } } + @Open + static class StringViewAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + + private final ViewVarCharVector accessor; + + StringViewAccessor(ViewVarCharVector vector) { + super(vector); + this.accessor = vector; + } + + @Override + final UTF8String getUTF8String(int rowId) { + // View values may be inlined in the view buffer rather than contiguous in a data + // buffer, so copy out rather than aliasing vector memory. + return UTF8String.fromBytes(accessor.get(rowId)); + } + } + @Open static class BinaryAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { @@ -610,6 +657,22 @@ final byte[] getBinary(int rowId) { } } + @Open + static class BinaryViewAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + + private final ViewVarBinaryVector accessor; + + BinaryViewAccessor(ViewVarBinaryVector vector) { + super(vector); + this.accessor = vector; + } + + @Override + final byte[] getBinary(int rowId) { + return accessor.getObject(rowId); + } + } + @Open static class DateAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { @@ -626,45 +689,63 @@ final int getInt(int rowId) { } } + /** + * Reads any of the eight timestamp vector variants (four units, with or without timezone), + * normalizing values to the microseconds Spark expects for TimestampType and TimestampNTZType. + * Seconds and milliseconds fail on overflow rather than silently wrapping; nanoseconds floor + * towards negative infinity, matching Spark's nanosecond-to-microsecond conversions. + */ @Open static class TimestampAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { - private final TimeStampMicroTZVector accessor; + private final TimeStampVector accessor; + private final TimeUnit unit; - TimestampAccessor(TimeStampMicroTZVector vector) { + TimestampAccessor(TimeStampVector vector) { super(vector); this.accessor = vector; + this.unit = ((ArrowType.Timestamp) vector.getField().getType()).getUnit(); } @Override final long getLong(int rowId) { - return accessor.get(rowId); + long value = accessor.get(rowId); + return switch (unit) { + case SECOND -> Math.multiplyExact(value, 1_000_000L); + case MILLISECOND -> Math.multiplyExact(value, 1_000L); + case MICROSECOND -> value; + case NANOSECOND -> Math.floorDiv(value, 1_000L); + }; } } @Open - static class TimestampNTZAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + static class ArrayAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { - private final TimeStampMicroVector accessor; + private final ListVector accessor; + private final VortexArrowColumnVector arrayData; - TimestampNTZAccessor(TimeStampMicroVector vector) { + ArrayAccessor(ListVector vector) { super(vector); this.accessor = vector; + this.arrayData = new VortexArrowColumnVector(vector.getDataVector()); } @Override - final long getLong(int rowId) { - return accessor.get(rowId); + final ColumnarArray getArray(int rowId) { + int start = accessor.getElementStartIndex(rowId); + int end = accessor.getElementEndIndex(rowId); + return new ColumnarArray(arrayData, start, end - start); } } @Open - static class ArrayAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + static class ListViewAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { - private final ListVector accessor; + private final ListViewVector accessor; private final VortexArrowColumnVector arrayData; - ArrayAccessor(ListVector vector) { + ListViewAccessor(ListViewVector vector) { super(vector); this.accessor = vector; this.arrayData = new VortexArrowColumnVector(vector.getDataVector()); @@ -740,6 +821,68 @@ int getInt(int rowId) { } } + /** + * Accessor for the {@code MONTH_DAY_NANO} interval vector itself. Values are read through + * {@link ColumnVector#getInterval(int)}, which uses the months/days/microseconds child + * columns; this accessor only provides null tracking. + */ + @Open + static class IntervalMonthDayNanoAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + + IntervalMonthDayNanoAccessor(IntervalMonthDayNanoVector vector) { + super(vector); + } + } + + @Open + static class IntervalMonthsAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + + private final IntervalMonthDayNanoVector accessor; + + IntervalMonthsAccessor(IntervalMonthDayNanoVector vector) { + super(vector); + this.accessor = vector; + } + + @Override + final int getInt(int rowId) { + return IntervalMonthDayNanoVector.getMonths(accessor.getDataBuffer(), rowId); + } + } + + @Open + static class IntervalDaysAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + + private final IntervalMonthDayNanoVector accessor; + + IntervalDaysAccessor(IntervalMonthDayNanoVector vector) { + super(vector); + this.accessor = vector; + } + + @Override + final int getInt(int rowId) { + return IntervalMonthDayNanoVector.getDays(accessor.getDataBuffer(), rowId); + } + } + + @Open + static class IntervalMicrosAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + + private final IntervalMonthDayNanoVector accessor; + + IntervalMicrosAccessor(IntervalMonthDayNanoVector vector) { + super(vector); + this.accessor = vector; + } + + @Override + final long getLong(int rowId) { + // Truncating division matches Spark's ArrowColumnVector month-day-nano handling. + return IntervalMonthDayNanoVector.getNanoseconds(accessor.getDataBuffer(), rowId) / 1_000L; + } + } + @Open static class DurationAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexBatchWrite.java b/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexBatchWrite.java index 7f96bac4404..03a70c2f676 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexBatchWrite.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexBatchWrite.java @@ -90,9 +90,15 @@ public DataWriterFactory createBatchWriterFactory(PhysicalWriteInfo info) { if (overwrite) { var session = VortexSparkSession.get(options); var uris = NativeFiles.listFiles(session, outputPath, options); - log.info("truncating table with {} files", uris.size()); + // Deleting the existing files is destructive and happens before the new data is written: + // if the subsequent write fails, abort() only removes the newly written files and cannot + // restore what was deleted here. Log loudly so operators can see what was removed. + log.warn( + "Deleting {} existing file(s) under {} because of overwrite, before writing new data; " + + "this cannot be undone if the subsequent write fails", + uris.size(), + outputPath); NativeFiles.delete(session, uris.toArray(new String[0]), options); - log.warn("overwrite currently does not do anything for vortex format"); } return new VortexDataWriterFactory(outputPath, schema, options, resolvedTransforms); diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/ArrowUtilsTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/ArrowUtilsTest.java new file mode 100644 index 00000000000..50c74eb41f8 --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/ArrowUtilsTest.java @@ -0,0 +1,257 @@ +// 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.assertThrows; + +import dev.vortex.relocated.org.apache.arrow.vector.types.DateUnit; +import dev.vortex.relocated.org.apache.arrow.vector.types.FloatingPointPrecision; +import dev.vortex.relocated.org.apache.arrow.vector.types.IntervalUnit; +import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit; +import dev.vortex.relocated.org.apache.arrow.vector.types.UnionMode; +import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType; +import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.Field; +import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.FieldType; +import java.util.List; +import org.apache.spark.sql.types.DataType; +import org.apache.spark.sql.types.DataTypes; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link ArrowUtils}, which maps Arrow types to Spark SQL {@link DataType}s. + * + *

Characterizes both the supported mappings and the type configurations that are explicitly rejected, so that + * regressions in either direction are caught. + */ +final class ArrowUtilsTest { + @Test + @DisplayName("Bool maps to BooleanType") + void boolMapsToBoolean() { + assertEquals(DataTypes.BooleanType, ArrowUtils.fromArrowType(new ArrowType.Bool())); + } + + @Test + @DisplayName("Signed integers map to the matching Spark integral types by bit width") + void signedIntegersMapByWidth() { + assertEquals(DataTypes.ByteType, ArrowUtils.fromArrowType(new ArrowType.Int(8, true))); + assertEquals(DataTypes.ShortType, ArrowUtils.fromArrowType(new ArrowType.Int(16, true))); + assertEquals(DataTypes.IntegerType, ArrowUtils.fromArrowType(new ArrowType.Int(32, true))); + assertEquals(DataTypes.LongType, ArrowUtils.fromArrowType(new ArrowType.Int(64, true))); + } + + @Test + @DisplayName("Single/double floating point map to Float/Double") + void floatingPointMapsByPrecision() { + assertEquals( + DataTypes.FloatType, + ArrowUtils.fromArrowType(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE))); + assertEquals( + DataTypes.DoubleType, + ArrowUtils.fromArrowType(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE))); + } + + @Test + @DisplayName("Decimal preserves precision and scale regardless of bit width") + void decimalPreservesPrecisionAndScale() { + assertEquals(DataTypes.createDecimalType(20, 4), ArrowUtils.fromArrowType(new ArrowType.Decimal(20, 4, 128))); + assertEquals(DataTypes.createDecimalType(20, 4), ArrowUtils.fromArrowType(new ArrowType.Decimal(20, 4, 256))); + } + + @Test + @DisplayName("Utf8, LargeUtf8 and Utf8View map to StringType") + void utf8MapsToString() { + assertEquals(DataTypes.StringType, ArrowUtils.fromArrowType(new ArrowType.Utf8())); + assertEquals(DataTypes.StringType, ArrowUtils.fromArrowType(new ArrowType.LargeUtf8())); + assertEquals(DataTypes.StringType, ArrowUtils.fromArrowType(new ArrowType.Utf8View())); + } + + @Test + @DisplayName("Binary, LargeBinary and BinaryView map to BinaryType") + void binaryMapsToBinary() { + assertEquals(DataTypes.BinaryType, ArrowUtils.fromArrowType(new ArrowType.Binary())); + assertEquals(DataTypes.BinaryType, ArrowUtils.fromArrowType(new ArrowType.LargeBinary())); + assertEquals(DataTypes.BinaryType, ArrowUtils.fromArrowType(new ArrowType.BinaryView())); + } + + @Test + @DisplayName("Date(DAY) maps to DateType") + void dateDayMapsToDate() { + assertEquals(DataTypes.DateType, ArrowUtils.fromArrowType(new ArrowType.Date(DateUnit.DAY))); + } + + @Test + @DisplayName("Timestamps of every unit map to Timestamp with tz, TimestampNTZ without") + void timestampMapsByTimezonePresence() { + for (TimeUnit unit : + new TimeUnit[] {TimeUnit.SECOND, TimeUnit.MILLISECOND, TimeUnit.MICROSECOND, TimeUnit.NANOSECOND}) { + assertEquals(DataTypes.TimestampType, ArrowUtils.fromArrowType(new ArrowType.Timestamp(unit, "UTC"))); + assertEquals(DataTypes.TimestampNTZType, ArrowUtils.fromArrowType(new ArrowType.Timestamp(unit, null))); + } + } + + @Test + @DisplayName("Null maps to NullType") + void nullMapsToNull() { + assertEquals(DataTypes.NullType, ArrowUtils.fromArrowType(new ArrowType.Null())); + } + + @Test + @DisplayName("Interval(YEAR_MONTH) maps to YearMonthIntervalType") + void yearMonthIntervalMapsToYearMonthIntervalType() { + assertEquals( + DataTypes.createYearMonthIntervalType(), + ArrowUtils.fromArrowType(new ArrowType.Interval(IntervalUnit.YEAR_MONTH))); + } + + @Test + @DisplayName("Interval(MONTH_DAY_NANO) maps to CalendarIntervalType") + void monthDayNanoIntervalMapsToCalendarIntervalType() { + assertEquals( + DataTypes.CalendarIntervalType, + ArrowUtils.fromArrowType(new ArrowType.Interval(IntervalUnit.MONTH_DAY_NANO))); + } + + @Test + @DisplayName("Duration(MICROSECOND) maps to DayTimeIntervalType") + void microsecondDurationMapsToDayTimeIntervalType() { + assertEquals( + DataTypes.createDayTimeIntervalType(), + ArrowUtils.fromArrowType(new ArrowType.Duration(TimeUnit.MICROSECOND))); + } + + @Test + @DisplayName("fromArrowField builds a StructType from nested children") + void structFieldBuildsStructType() { + Field struct = new Field( + "s", + FieldType.nullable(new ArrowType.Struct()), + List.of( + new Field("a", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field("b", FieldType.notNullable(new ArrowType.Utf8()), null))); + + DataType expected = DataTypes.createStructType(new org.apache.spark.sql.types.StructField[] { + DataTypes.createStructField("a", DataTypes.IntegerType, true), + DataTypes.createStructField("b", DataTypes.StringType, false) + }); + assertEquals(expected, ArrowUtils.fromArrowField(struct)); + } + + @Test + @DisplayName("fromArrowField builds an ArrayType carrying the element's nullability") + void listFieldBuildsArrayType() { + Field list = new Field( + "l", + FieldType.nullable(new ArrowType.List()), + List.of(new Field("element", FieldType.nullable(new ArrowType.Int(32, true)), null))); + + assertEquals(DataTypes.createArrayType(DataTypes.IntegerType, true), ArrowUtils.fromArrowField(list)); + } + + @Test + @DisplayName("fromArrowField builds an ArrayType from a ListView field") + void listViewFieldBuildsArrayType() { + Field listView = new Field( + "lv", + FieldType.nullable(new ArrowType.ListView()), + List.of(new Field("element", FieldType.notNullable(new ArrowType.Utf8View()), null))); + + assertEquals(DataTypes.createArrayType(DataTypes.StringType, false), ArrowUtils.fromArrowField(listView)); + } + + @Test + @DisplayName("fromArrowField builds a MapType carrying the value's nullability") + void mapFieldBuildsMapType() { + assertEquals( + DataTypes.createMapType(DataTypes.StringType, DataTypes.LongType, true), + ArrowUtils.fromArrowField(mapField(FieldType.nullable(new ArrowType.Int(64, true))))); + assertEquals( + DataTypes.createMapType(DataTypes.StringType, DataTypes.LongType, false), + ArrowUtils.fromArrowField(mapField(FieldType.notNullable(new ArrowType.Int(64, true))))); + } + + private static Field mapField(FieldType valueType) { + Field key = new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null); + Field value = new Field("value", valueType, null); + Field entries = new Field("entries", FieldType.notNullable(new ArrowType.Struct()), List.of(key, value)); + return new Field("m", FieldType.nullable(new ArrowType.Map(false)), List.of(entries)); + } + + @Test + @DisplayName("Unsigned integers are unsupported") + void unsignedIntegerIsUnsupported() { + assertThrows(UnsupportedOperationException.class, () -> ArrowUtils.fromArrowType(new ArrowType.Int(32, false))); + } + + @Test + @DisplayName("Half-precision floating point is unsupported") + void halfPrecisionFloatIsUnsupported() { + assertThrows( + UnsupportedOperationException.class, + () -> ArrowUtils.fromArrowType(new ArrowType.FloatingPoint(FloatingPointPrecision.HALF))); + } + + @Test + @DisplayName("Non-DAY date units are unsupported") + void millisecondDateIsUnsupported() { + assertThrows( + UnsupportedOperationException.class, + () -> ArrowUtils.fromArrowType(new ArrowType.Date(DateUnit.MILLISECOND))); + } + + @Test + @DisplayName("Non-microsecond duration units are unsupported") + void nonMicrosecondDurationsAreUnsupported() { + for (TimeUnit unit : new TimeUnit[] {TimeUnit.SECOND, TimeUnit.MILLISECOND, TimeUnit.NANOSECOND}) { + assertThrows( + UnsupportedOperationException.class, () -> ArrowUtils.fromArrowType(new ArrowType.Duration(unit))); + } + } + + @Test + @DisplayName("Time is unsupported: Spark's TimeType only exists in Spark 4.1+") + void timeIsUnsupported() { + assertThrows( + UnsupportedOperationException.class, + () -> ArrowUtils.fromArrowType(new ArrowType.Time(TimeUnit.NANOSECOND, 64))); + assertThrows( + UnsupportedOperationException.class, + () -> ArrowUtils.fromArrowType(new ArrowType.Time(TimeUnit.MILLISECOND, 32))); + } + + @Test + @DisplayName("Interval(DAY_TIME) is unsupported") + void dayTimeIntervalIsUnsupported() { + assertThrows( + UnsupportedOperationException.class, + () -> ArrowUtils.fromArrowType(new ArrowType.Interval(IntervalUnit.DAY_TIME))); + } + + @Test + @DisplayName("FixedSizeBinary is unsupported") + void fixedSizeBinaryIsUnsupported() { + assertThrows( + UnsupportedOperationException.class, () -> ArrowUtils.fromArrowType(new ArrowType.FixedSizeBinary(16))); + } + + @Test + @DisplayName("Union is unsupported") + void unionIsUnsupported() { + assertThrows( + UnsupportedOperationException.class, + () -> ArrowUtils.fromArrowType(new ArrowType.Union(UnionMode.Sparse, new int[0]))); + } + + @Test + @DisplayName("LargeList and FixedSizeList fields are unsupported") + void largeAndFixedSizeListFieldsAreUnsupported() { + Field element = new Field("element", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field largeList = new Field("ll", FieldType.nullable(new ArrowType.LargeList()), List.of(element)); + Field fixedSizeList = new Field("fsl", FieldType.nullable(new ArrowType.FixedSizeList(2)), List.of(element)); + + assertThrows(UnsupportedOperationException.class, () -> ArrowUtils.fromArrowField(largeList)); + assertThrows(UnsupportedOperationException.class, () -> ArrowUtils.fromArrowField(fixedSizeList)); + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java index 325da0bc98c..90fb1959ade 100644 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceWriteTest.java @@ -129,6 +129,45 @@ public void testWriteAndReadVortexFiles() throws IOException { verifyDataContent(originalDf, readDf); } + @Test + @DisplayName("Write and read Vortex files with a bare path (no URI scheme)") + public void testWriteAndReadWithBarePath() throws IOException { + int numRows = 50; + Dataset originalDf = createTestDataFrame(numRows); + + // A bare filesystem path, without a file:// scheme. + String barePath = tempDir.resolve("bare_path_output").toString(); + originalDf + .write() + .format("vortex") + .option("path", barePath) + .mode(SaveMode.Overwrite) + .save(); + + assertFalse(findVortexFiles(tempDir.resolve("bare_path_output")).isEmpty(), "Write should create files"); + + // Reading a bare directory path exercises schema inference via file listing. + Dataset readDf = + spark.read().format("vortex").option("path", barePath).load(); + + assertSchemaEquals(originalDf.schema(), readDf.schema()); + assertEquals(numRows, readDf.count(), "Read DataFrame should have same number of rows as original"); + verifyDataContent(originalDf, readDf); + + // Overwriting a bare path exercises file listing and deletion of the existing files. + Dataset replacementDf = createTestDataFrame(25); + replacementDf + .write() + .format("vortex") + .option("path", barePath) + .mode(SaveMode.Overwrite) + .save(); + + Dataset reread = + spark.read().format("vortex").option("path", barePath).load(); + assertEquals(25, reread.count(), "Should have data from second write after overwrite"); + } + @Test @DisplayName("Write empty DataFrame as Vortex") public void testWriteEmptyDataFrame() throws IOException { diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexSessionCatalogTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexSessionCatalogTest.java new file mode 100644 index 00000000000..b31d674d884 --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexSessionCatalogTest.java @@ -0,0 +1,113 @@ +// 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.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.io.TempDir; + +/** + * Integration tests for {@link VortexSessionCatalog}, the session catalog extension that makes {@code CREATE TABLE ... + * USING vortex} work on Spark 3.5 as well as Spark 4. + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +public final class VortexSessionCatalogTest { + + private SparkSession spark; + private Path warehouseDir; + + @TempDir + Path tempDir; + + @BeforeAll + public void setUp() throws IOException { + warehouseDir = Files.createTempDirectory("vortex-warehouse"); + spark = SparkSession.builder() + .appName("VortexSessionCatalogTest") + .master("local[2]") + .config("spark.driver.host", "127.0.0.1") + .config("spark.sql.warehouse.dir", warehouseDir.toUri().toString()) + .config("spark.sql.catalog.spark_catalog", VortexSessionCatalog.class.getName()) + .config("spark.ui.enabled", "false") + .getOrCreate(); + } + + @AfterAll + public void tearDown() { + if (spark != null) { + spark.stop(); + } + } + + @Test + @DisplayName("Managed table lifecycle: CREATE, SELECT while empty, INSERT, INSERT OVERWRITE, DROP") + public void testManagedTableLifecycle() { + spark.sql("CREATE TABLE managed_students (id INT, name STRING, age INT) USING vortex"); + + assertEquals(0, spark.sql("SELECT * FROM managed_students").count(), "New managed table should be empty"); + + spark.sql("INSERT INTO managed_students VALUES (1, 'Alice', 20), (2, 'Bob', 21)"); + List rows = spark.sql("SELECT name FROM managed_students WHERE age > 20 ORDER BY name") + .collectAsList(); + assertEquals(1, rows.size()); + assertEquals("Bob", rows.get(0).getString(0)); + + spark.sql("INSERT OVERWRITE managed_students VALUES (3, 'Carol', 22)"); + assertEquals(1, spark.sql("SELECT * FROM managed_students").count(), "Overwrite should replace all rows"); + + Path tableDir = warehouseDir.resolve("managed_students"); + assertTrue(Files.exists(tableDir), "Managed table data should live under the warehouse dir"); + spark.sql("DROP TABLE managed_students"); + assertFalse(Files.exists(tableDir), "Dropping a managed table should remove its data"); + } + + @Test + @DisplayName("CREATE TABLE AS SELECT without a LOCATION clause") + public void testCreateManagedTableAsSelect() { + spark.sql("CREATE TABLE ctas_source (id INT, name STRING) USING vortex"); + spark.sql("INSERT INTO ctas_source VALUES (1, 'Alice'), (2, 'Bob')"); + + spark.sql("CREATE TABLE ctas_target USING vortex AS SELECT * FROM ctas_source WHERE id > 1"); + List rows = spark.sql("SELECT name FROM ctas_target").collectAsList(); + assertEquals(1, rows.size()); + assertEquals("Bob", rows.get(0).getString(0)); + + spark.sql("DROP TABLE ctas_target"); + spark.sql("DROP TABLE ctas_source"); + } + + @Test + @DisplayName("External table with a LOCATION clause") + public void testExternalTable() { + Path location = tempDir.resolve("ext_students"); + spark.sql( + String.format("CREATE TABLE ext_students (id INT, name STRING) USING vortex LOCATION '%s'", location)); + spark.sql("INSERT INTO ext_students VALUES (1, 'Alice')"); + assertEquals(1, spark.sql("SELECT * FROM ext_students").count()); + spark.sql("DROP TABLE ext_students"); + } + + @Test + @DisplayName("Tables of other providers pass through the extension untouched") + public void testOtherProviderPassthrough() { + spark.sql("CREATE TABLE pq_table (id INT) USING parquet"); + spark.sql("INSERT INTO pq_table VALUES (7)"); + assertEquals( + 7, spark.sql("SELECT * FROM pq_table").collectAsList().get(0).getInt(0)); + spark.sql("DROP TABLE pq_table"); + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/VortexSqlTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexSqlTest.java new file mode 100644 index 00000000000..1a9b21ef779 --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/VortexSqlTest.java @@ -0,0 +1,176 @@ +// 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.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assumptions.assumeTrue; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.List; +import org.apache.spark.sql.AnalysisException; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SaveMode; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.io.TempDir; + +/** + * Integration tests for Spark SQL access to Vortex: managed tables created without a {@code LOCATION} clause, and + * direct file queries through {@link VortexCatalog}. + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +public final class VortexSqlTest { + + private SparkSession spark; + private Path warehouseDir; + + @TempDir + Path tempDir; + + @BeforeAll + public void setUp() throws IOException { + warehouseDir = Files.createTempDirectory("vortex-warehouse"); + spark = SparkSession.builder() + .appName("VortexSqlTest") + .master("local[2]") + .config("spark.driver.host", "127.0.0.1") + .config("spark.sql.warehouse.dir", warehouseDir.toUri().toString()) + .config("spark.sql.catalog.vortex", VortexCatalog.class.getName()) + .config("spark.ui.enabled", "false") + .getOrCreate(); + } + + @AfterAll + public void tearDown() { + if (spark != null) { + spark.stop(); + } + } + + /** + * Spark 3.5's built-in session catalog cannot read tables backed by a DataSource-V2-only provider: its + * {@code FindDataSourceTable} rule falls back to the V1 {@code DataSource} path, which rejects the provider with + * "vortex is not a valid Spark SQL Data Source". Spark 4 resolves such tables through the provider directly. On + * Spark 3.5 the {@link VortexSessionCatalog} extension provides the same support — see + * {@link VortexSessionCatalogTest}, which runs these scenarios on both versions. + */ + private void assumeSupportsSqlTables() { + assumeTrue( + spark.version().startsWith("4."), + "CREATE TABLE ... USING vortex requires Spark 4 or the VortexSessionCatalog extension"); + } + + @Test + @DisplayName("Managed table lifecycle: CREATE, SELECT while empty, INSERT, INSERT OVERWRITE, DROP") + public void testManagedTableLifecycle() { + assumeSupportsSqlTables(); + spark.sql("CREATE TABLE managed_students (id INT, name STRING, age INT) USING vortex"); + + assertEquals(0, spark.sql("SELECT * FROM managed_students").count(), "New managed table should be empty"); + + spark.sql("INSERT INTO managed_students VALUES (1, 'Alice', 20), (2, 'Bob', 21)"); + List rows = spark.sql("SELECT name FROM managed_students WHERE age > 20 ORDER BY name") + .collectAsList(); + assertEquals(1, rows.size()); + assertEquals("Bob", rows.get(0).getString(0)); + + spark.sql("INSERT OVERWRITE managed_students VALUES (3, 'Carol', 22)"); + assertEquals(1, spark.sql("SELECT * FROM managed_students").count(), "Overwrite should replace all rows"); + + Path tableDir = warehouseDir.resolve("managed_students"); + assertTrue(Files.exists(tableDir), "Managed table data should live under the warehouse dir"); + spark.sql("DROP TABLE managed_students"); + assertFalse(Files.exists(tableDir), "Dropping a managed table should remove its data"); + } + + @Test + @DisplayName("CREATE TABLE AS SELECT without a LOCATION clause") + public void testCreateManagedTableAsSelect() { + assumeSupportsSqlTables(); + spark.sql("CREATE TABLE ctas_source (id INT, name STRING) USING vortex"); + spark.sql("INSERT INTO ctas_source VALUES (1, 'Alice'), (2, 'Bob')"); + + spark.sql("CREATE TABLE ctas_target USING vortex AS SELECT * FROM ctas_source WHERE id > 1"); + List rows = spark.sql("SELECT name FROM ctas_target").collectAsList(); + assertEquals(1, rows.size()); + assertEquals("Bob", rows.get(0).getString(0)); + + spark.sql("DROP TABLE ctas_target"); + spark.sql("DROP TABLE ctas_source"); + } + + @Test + @DisplayName("Reading the vortex format without a path option still fails") + public void testReadWithoutPathStillThrows() { + assertThrows( + IllegalArgumentException.class, + () -> spark.read().format("vortex").load(), + "A read with no path should not silently return an empty DataFrame"); + } + + @Test + @DisplayName("Direct file query through the vortex catalog") + public void testDirectPathQuery() { + Path dataDir = tempDir.resolve("direct_query"); + writeTestData(dataDir); + + List rows = spark.sql(String.format("SELECT name FROM vortex.`%s` WHERE age > 30 ORDER BY name", dataDir)) + .collectAsList(); + assertEquals(2, rows.size()); + assertEquals("Alice", rows.get(0).getString(0)); + assertEquals("Carol", rows.get(1).getString(0)); + } + + @Test + @DisplayName("INSERT INTO a direct path through the vortex catalog") + public void testDirectPathInsert() { + Path dataDir = tempDir.resolve("direct_insert"); + writeTestData(dataDir); + + spark.sql(String.format("INSERT INTO vortex.`%s` VALUES (4, 'Dave', 50)", dataDir)); + assertEquals( + 4, + spark.sql(String.format("SELECT * FROM vortex.`%s`", dataDir)).count()); + } + + @Test + @DisplayName("Direct queries of missing paths and non-path names fail as table-not-found") + public void testDirectPathNotFound() { + assertThrows(AnalysisException.class, () -> spark.sql( + String.format("SELECT * FROM vortex.`%s`", tempDir.resolve("no_such_dir"))) + .collect()); + assertThrows(AnalysisException.class, () -> spark.sql("SELECT * FROM vortex.not_a_path") + .collect()); + } + + private void writeTestData(Path dataDir) { + StructType schema = DataTypes.createStructType(new StructField[] { + DataTypes.createStructField("id", DataTypes.IntegerType, false), + DataTypes.createStructField("name", DataTypes.StringType, false), + DataTypes.createStructField("age", DataTypes.IntegerType, false) + }); + Dataset df = spark.createDataFrame( + Arrays.asList( + RowFactory.create(1, "Alice", 34), + RowFactory.create(2, "Bob", 27), + RowFactory.create(3, "Carol", 45)), + schema); + df.write().format("vortex").mode(SaveMode.Overwrite).save(dataDir.toString()); + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/read/PartitionPathUtilsTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/read/PartitionPathUtilsTest.java new file mode 100644 index 00000000000..445fad2f56c --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/read/PartitionPathUtilsTest.java @@ -0,0 +1,199 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.read; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; +import java.util.Map; +import org.apache.spark.sql.execution.vectorized.ConstantColumnVector; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.unsafe.types.UTF8String; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link PartitionPathUtils}, which discovers and materializes Hive-style partition columns from file + * paths. + * + *

Characterizes partition-value parsing (including URL decoding and the {@code __HIVE_DEFAULT_PARTITION__} + * sentinel), partition column type inference, and constant-vector materialization for each supported type. + */ +final class PartitionPathUtilsTest { + // --- parsePartitionValues --- + + @Test + @DisplayName("Parses key=value segments from a Hive-style partition path") + void parsesKeyValueSegments() { + Map values = + PartitionPathUtils.parsePartitionValues("/data/warehouse/year=2024/month=12/part-0.vortex"); + assertEquals(Map.of("year", "2024", "month", "12"), values); + } + + @Test + @DisplayName("Preserves the order of partition segments in the path") + void preservesSegmentOrder() { + Map values = PartitionPathUtils.parsePartitionValues("/tbl/b=2/a=1/c=3/file.vortex"); + assertEquals(List.of("b", "a", "c"), List.copyOf(values.keySet())); + } + + @Test + @DisplayName("Ignores segments without key=value shape") + void ignoresNonPartitionSegments() { + Map values = PartitionPathUtils.parsePartitionValues("/data/plain/dir/file.vortex"); + assertTrue(values.isEmpty()); + } + + @Test + @DisplayName("Ignores segments with an empty key or empty value") + void ignoresEmptyKeyOrValue() { + assertTrue(PartitionPathUtils.parsePartitionValues("/tbl/=v/file").isEmpty()); + assertTrue(PartitionPathUtils.parsePartitionValues("/tbl/k=/file").isEmpty()); + } + + @Test + @DisplayName("URL-decodes both keys and values") + void urlDecodesKeysAndValues() { + Map values = PartitionPathUtils.parsePartitionValues("/tbl/city=New%20York/file.vortex"); + assertEquals(Map.of("city", "New York"), values); + } + + @Test + @DisplayName("Splits on the first '=' so values may contain '='") + void splitsOnFirstEquals() { + Map values = PartitionPathUtils.parsePartitionValues("/tbl/k=a=b/file.vortex"); + assertEquals(Map.of("k", "a=b"), values); + } + + // --- inferPartitionColumnType --- + + @Test + @DisplayName("Infers IntegerType for values that fit in an int") + void infersInteger() { + assertEquals(DataTypes.IntegerType, PartitionPathUtils.inferPartitionColumnType("42")); + assertEquals(DataTypes.IntegerType, PartitionPathUtils.inferPartitionColumnType("-7")); + } + + @Test + @DisplayName("Infers LongType for integral values wider than an int") + void infersLong() { + assertEquals(DataTypes.LongType, PartitionPathUtils.inferPartitionColumnType("9999999999")); + } + + @Test + @DisplayName("Infers DoubleType for decimal-point values") + void infersDouble() { + assertEquals(DataTypes.DoubleType, PartitionPathUtils.inferPartitionColumnType("3.14")); + } + + @Test + @DisplayName("Infers BooleanType for true/false in any case") + void infersBoolean() { + assertEquals(DataTypes.BooleanType, PartitionPathUtils.inferPartitionColumnType("true")); + assertEquals(DataTypes.BooleanType, PartitionPathUtils.inferPartitionColumnType("FALSE")); + } + + @Test + @DisplayName("Falls back to StringType for everything else") + void fallsBackToString() { + assertEquals(DataTypes.StringType, PartitionPathUtils.inferPartitionColumnType("hello")); + assertEquals(DataTypes.StringType, PartitionPathUtils.inferPartitionColumnType("2024-12-01")); + } + + @Test + @DisplayName("Null and __HIVE_DEFAULT_PARTITION__ infer StringType") + void nullAndDefaultPartitionInferString() { + assertEquals(DataTypes.StringType, PartitionPathUtils.inferPartitionColumnType(null)); + assertEquals(DataTypes.StringType, PartitionPathUtils.inferPartitionColumnType("__HIVE_DEFAULT_PARTITION__")); + } + + // --- createConstantVector --- + + @Test + @DisplayName("Null and __HIVE_DEFAULT_PARTITION__ materialize as null vectors") + void nullValuesMaterializeAsNull() { + ConstantColumnVector fromNull = PartitionPathUtils.createConstantVector(3, DataTypes.StringType, null); + assertTrue(fromNull.isNullAt(0)); + + ConstantColumnVector fromSentinel = + PartitionPathUtils.createConstantVector(3, DataTypes.StringType, "__HIVE_DEFAULT_PARTITION__"); + assertTrue(fromSentinel.isNullAt(0)); + } + + @Test + @DisplayName("String values materialize as UTF8 strings") + void stringMaterializes() { + ConstantColumnVector vec = PartitionPathUtils.createConstantVector(2, DataTypes.StringType, "us-east"); + assertFalse(vec.isNullAt(0)); + assertEquals(UTF8String.fromString("us-east"), vec.getUTF8String(0)); + } + + @Test + @DisplayName("Integral values materialize with the width of the target type") + void integralTypesMaterialize() { + assertEquals( + 42, + PartitionPathUtils.createConstantVector(1, DataTypes.IntegerType, "42") + .getInt(0)); + assertEquals( + 9999999999L, + PartitionPathUtils.createConstantVector(1, DataTypes.LongType, "9999999999") + .getLong(0)); + assertEquals( + (short) 7, + PartitionPathUtils.createConstantVector(1, DataTypes.ShortType, "7") + .getShort(0)); + assertEquals( + (byte) 3, + PartitionPathUtils.createConstantVector(1, DataTypes.ByteType, "3") + .getByte(0)); + } + + @Test + @DisplayName("DateType parses the value as days since epoch (int)") + void dateMaterializesAsInt() { + assertEquals( + 19000, + PartitionPathUtils.createConstantVector(1, DataTypes.DateType, "19000") + .getInt(0)); + } + + @Test + @DisplayName("Timestamp types parse the value as micros since epoch (long)") + void timestampMaterializesAsLong() { + assertEquals( + 1700000000000000L, + PartitionPathUtils.createConstantVector(1, DataTypes.TimestampType, "1700000000000000") + .getLong(0)); + assertEquals( + 1700000000000000L, + PartitionPathUtils.createConstantVector(1, DataTypes.TimestampNTZType, "1700000000000000") + .getLong(0)); + } + + @Test + @DisplayName("Boolean, float, and double values materialize with their target types") + void booleanAndFloatingPointMaterialize() { + assertTrue(PartitionPathUtils.createConstantVector(1, DataTypes.BooleanType, "true") + .getBoolean(0)); + assertEquals( + 1.5f, + PartitionPathUtils.createConstantVector(1, DataTypes.FloatType, "1.5") + .getFloat(0)); + assertEquals( + 2.25, + PartitionPathUtils.createConstantVector(1, DataTypes.DoubleType, "2.25") + .getDouble(0)); + } + + @Test + @DisplayName("Unrecognized target types fall back to UTF8 strings") + void unrecognizedTypeFallsBackToString() { + ConstantColumnVector vec = + PartitionPathUtils.createConstantVector(1, DataTypes.CalendarIntervalType, "whatever"); + assertEquals(UTF8String.fromString("whatever"), vec.getUTF8String(0)); + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexArrowColumnVectorTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexArrowColumnVectorTest.java new file mode 100644 index 00000000000..9ae0ba86fd8 --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexArrowColumnVectorTest.java @@ -0,0 +1,523 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.read; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.vortex.relocated.org.apache.arrow.memory.BufferAllocator; +import dev.vortex.relocated.org.apache.arrow.memory.RootAllocator; +import dev.vortex.relocated.org.apache.arrow.vector.BigIntVector; +import dev.vortex.relocated.org.apache.arrow.vector.BitVector; +import dev.vortex.relocated.org.apache.arrow.vector.DateDayVector; +import dev.vortex.relocated.org.apache.arrow.vector.DateMilliVector; +import dev.vortex.relocated.org.apache.arrow.vector.Decimal256Vector; +import dev.vortex.relocated.org.apache.arrow.vector.DecimalVector; +import dev.vortex.relocated.org.apache.arrow.vector.DurationVector; +import dev.vortex.relocated.org.apache.arrow.vector.Float2Vector; +import dev.vortex.relocated.org.apache.arrow.vector.Float4Vector; +import dev.vortex.relocated.org.apache.arrow.vector.Float8Vector; +import dev.vortex.relocated.org.apache.arrow.vector.IntVector; +import dev.vortex.relocated.org.apache.arrow.vector.IntervalMonthDayNanoVector; +import dev.vortex.relocated.org.apache.arrow.vector.IntervalYearVector; +import dev.vortex.relocated.org.apache.arrow.vector.LargeVarBinaryVector; +import dev.vortex.relocated.org.apache.arrow.vector.LargeVarCharVector; +import dev.vortex.relocated.org.apache.arrow.vector.NullVector; +import dev.vortex.relocated.org.apache.arrow.vector.SmallIntVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeMicroVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampMicroTZVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampMicroVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampMilliTZVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampMilliVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampNanoTZVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampNanoVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampSecTZVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampSecVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampVector; +import dev.vortex.relocated.org.apache.arrow.vector.TinyIntVector; +import dev.vortex.relocated.org.apache.arrow.vector.UInt4Vector; +import dev.vortex.relocated.org.apache.arrow.vector.VarBinaryVector; +import dev.vortex.relocated.org.apache.arrow.vector.VarCharVector; +import dev.vortex.relocated.org.apache.arrow.vector.ViewVarBinaryVector; +import dev.vortex.relocated.org.apache.arrow.vector.ViewVarCharVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.ListVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.ListViewVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.MapVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.StructVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.impl.NullableStructWriter; +import dev.vortex.relocated.org.apache.arrow.vector.complex.impl.UnionListViewWriter; +import dev.vortex.relocated.org.apache.arrow.vector.complex.impl.UnionListWriter; +import dev.vortex.relocated.org.apache.arrow.vector.complex.impl.UnionMapWriter; +import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit; +import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType; +import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.FieldType; +import dev.vortex.relocated.org.apache.arrow.vector.util.Text; +import java.math.BigDecimal; +import java.nio.charset.StandardCharsets; +import org.apache.spark.sql.types.ArrayType; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.MapType; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.vectorized.ColumnarArray; +import org.apache.spark.sql.vectorized.ColumnarMap; +import org.apache.spark.unsafe.types.CalendarInterval; +import org.apache.spark.unsafe.types.UTF8String; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link VortexArrowColumnVector} covering every Arrow vector type it supports (Spark 4.1's own + * ArrowColumnVector set plus the Arrow view types): the Spark type mapping, the value conversion, and null handling. + */ +final class VortexArrowColumnVectorTest { + + private static final BufferAllocator ALLOCATOR = new RootAllocator(); + + @AfterAll + static void closeAllocator() { + ALLOCATOR.close(); + } + + @Test + @DisplayName("BitVector maps to BooleanType") + void booleanVector() { + try (BitVector vector = new BitVector("bool", ALLOCATOR)) { + vector.allocateNew(3); + vector.set(0, 1); + vector.set(2, 0); + vector.setValueCount(3); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.BooleanType, column.dataType()); + assertTrue(column.getBoolean(0)); + assertTrue(column.isNullAt(1)); + assertFalse(column.getBoolean(2)); + assertTrue(column.hasNull()); + assertEquals(1, column.numNulls()); + } + } + + @Test + @DisplayName("Signed integer vectors map to Byte/Short/Integer/LongType") + void signedIntegers() { + try (TinyIntVector i8 = new TinyIntVector("i8", ALLOCATOR); + SmallIntVector i16 = new SmallIntVector("i16", ALLOCATOR); + IntVector i32 = new IntVector("i32", ALLOCATOR); + BigIntVector i64 = new BigIntVector("i64", ALLOCATOR)) { + i8.allocateNew(2); + i8.set(0, Byte.MIN_VALUE); + i8.setValueCount(2); + i16.allocateNew(2); + i16.set(0, Short.MIN_VALUE); + i16.setValueCount(2); + i32.allocateNew(2); + i32.set(0, Integer.MIN_VALUE); + i32.setValueCount(2); + i64.allocateNew(2); + i64.set(0, Long.MIN_VALUE); + i64.setValueCount(2); + + VortexArrowColumnVector byteColumn = new VortexArrowColumnVector(i8); + assertEquals(DataTypes.ByteType, byteColumn.dataType()); + assertEquals(Byte.MIN_VALUE, byteColumn.getByte(0)); + assertTrue(byteColumn.isNullAt(1)); + + VortexArrowColumnVector shortColumn = new VortexArrowColumnVector(i16); + assertEquals(DataTypes.ShortType, shortColumn.dataType()); + assertEquals(Short.MIN_VALUE, shortColumn.getShort(0)); + + VortexArrowColumnVector intColumn = new VortexArrowColumnVector(i32); + assertEquals(DataTypes.IntegerType, intColumn.dataType()); + assertEquals(Integer.MIN_VALUE, intColumn.getInt(0)); + + VortexArrowColumnVector longColumn = new VortexArrowColumnVector(i64); + assertEquals(DataTypes.LongType, longColumn.dataType()); + assertEquals(Long.MIN_VALUE, longColumn.getLong(0)); + } + } + + @Test + @DisplayName("Float vectors map to Float/DoubleType") + void floats() { + try (Float4Vector f32 = new Float4Vector("f32", ALLOCATOR); + Float8Vector f64 = new Float8Vector("f64", ALLOCATOR)) { + f32.allocateNew(2); + f32.set(0, 2.5f); + f32.setValueCount(2); + f64.allocateNew(2); + f64.set(0, 3.5d); + f64.setValueCount(2); + + VortexArrowColumnVector floatColumn = new VortexArrowColumnVector(f32); + assertEquals(DataTypes.FloatType, floatColumn.dataType()); + assertEquals(2.5f, floatColumn.getFloat(0)); + assertTrue(floatColumn.isNullAt(1)); + + VortexArrowColumnVector doubleColumn = new VortexArrowColumnVector(f64); + assertEquals(DataTypes.DoubleType, doubleColumn.dataType()); + assertEquals(3.5d, doubleColumn.getDouble(0)); + } + } + + @Test + @DisplayName("Decimal128 vectors map to DecimalType") + void decimal() { + try (DecimalVector d128 = new DecimalVector("d128", ALLOCATOR, 10, 2)) { + d128.allocateNew(2); + d128.set(0, new BigDecimal("12345678.90")); + d128.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(d128); + assertEquals(DataTypes.createDecimalType(10, 2), column.dataType()); + assertEquals( + new BigDecimal("12345678.90"), column.getDecimal(0, 10, 2).toJavaBigDecimal()); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("String vectors (regular, large, view) map to StringType") + void strings() { + try (VarCharVector utf8 = new VarCharVector("utf8", ALLOCATOR); + LargeVarCharVector largeUtf8 = new LargeVarCharVector("large_utf8", ALLOCATOR); + ViewVarCharVector utf8View = new ViewVarCharVector("utf8_view", ALLOCATOR)) { + utf8.allocateNew(2); + utf8.setSafe(0, "hello".getBytes(StandardCharsets.UTF_8)); + utf8.setValueCount(2); + largeUtf8.allocateNew(2); + largeUtf8.setSafe(0, "world".getBytes(StandardCharsets.UTF_8)); + largeUtf8.setValueCount(2); + utf8View.allocateNew(2); + utf8View.setSafe(0, new Text("a string long enough to not be inlined")); + utf8View.setValueCount(2); + + VortexArrowColumnVector utf8Column = new VortexArrowColumnVector(utf8); + assertEquals(DataTypes.StringType, utf8Column.dataType()); + assertEquals(UTF8String.fromString("hello"), utf8Column.getUTF8String(0)); + assertTrue(utf8Column.isNullAt(1)); + + VortexArrowColumnVector largeUtf8Column = new VortexArrowColumnVector(largeUtf8); + assertEquals(DataTypes.StringType, largeUtf8Column.dataType()); + assertEquals(UTF8String.fromString("world"), largeUtf8Column.getUTF8String(0)); + + VortexArrowColumnVector utf8ViewColumn = new VortexArrowColumnVector(utf8View); + assertEquals(DataTypes.StringType, utf8ViewColumn.dataType()); + assertEquals( + UTF8String.fromString("a string long enough to not be inlined"), utf8ViewColumn.getUTF8String(0)); + assertTrue(utf8ViewColumn.isNullAt(1)); + } + } + + @Test + @DisplayName("Binary vectors (regular, large, view) map to BinaryType") + void binary() { + byte[] payload = new byte[] {1, 2, 3, 4}; + try (VarBinaryVector bin = new VarBinaryVector("bin", ALLOCATOR); + LargeVarBinaryVector largeBin = new LargeVarBinaryVector("large_bin", ALLOCATOR); + ViewVarBinaryVector binView = new ViewVarBinaryVector("bin_view", ALLOCATOR)) { + bin.allocateNew(2); + bin.setSafe(0, payload); + bin.setValueCount(2); + largeBin.allocateNew(2); + largeBin.setSafe(0, payload); + largeBin.setValueCount(2); + binView.allocateNew(2); + binView.setSafe(0, payload); + binView.setValueCount(2); + + for (VortexArrowColumnVector column : new VortexArrowColumnVector[] { + new VortexArrowColumnVector(bin), + new VortexArrowColumnVector(largeBin), + new VortexArrowColumnVector(binView), + }) { + assertEquals(DataTypes.BinaryType, column.dataType()); + assertArrayEquals(payload, column.getBinary(0)); + assertTrue(column.isNullAt(1)); + } + } + } + + @Test + @DisplayName("Day-unit date vectors map to DateType") + void date() { + try (DateDayVector vector = new DateDayVector("date_day", ALLOCATOR)) { + vector.allocateNew(2); + vector.set(0, 19000); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.DateType, column.dataType()); + assertEquals(19000, column.getInt(0)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Timestamp vectors of every unit normalize to microseconds") + void timestampsWithoutTimezone() { + try (TimeStampSecVector sec = new TimeStampSecVector("ts_s", ALLOCATOR); + TimeStampMilliVector milli = new TimeStampMilliVector("ts_ms", ALLOCATOR); + TimeStampMicroVector micro = new TimeStampMicroVector("ts_us", ALLOCATOR); + TimeStampNanoVector nano = new TimeStampNanoVector("ts_ns", ALLOCATOR)) { + sec.allocateNew(3); + sec.set(0, 1_700_000_000L); + sec.setValueCount(3); + milli.allocateNew(3); + milli.set(0, 1_700_000_000_123L); + milli.setValueCount(3); + micro.allocateNew(3); + micro.set(0, 1_700_000_000_123_456L); + micro.setValueCount(3); + nano.allocateNew(3); + nano.set(0, 1_700_000_000_123_456_789L); + // Negative nanos floor towards negative infinity when reduced to micros. + nano.set(2, -1_500L); + nano.setValueCount(3); + + assertTimestamp(sec, DataTypes.TimestampNTZType, 1_700_000_000_000_000L); + assertTimestamp(milli, DataTypes.TimestampNTZType, 1_700_000_000_123_000L); + assertTimestamp(micro, DataTypes.TimestampNTZType, 1_700_000_000_123_456L); + VortexArrowColumnVector nanoColumn = + assertTimestamp(nano, DataTypes.TimestampNTZType, 1_700_000_000_123_456L); + assertEquals(-2L, nanoColumn.getLong(2)); + } + } + + @Test + @DisplayName("Timezone-aware timestamp vectors of every unit map to TimestampType") + void timestampsWithTimezone() { + try (TimeStampSecTZVector sec = new TimeStampSecTZVector("ts_s", ALLOCATOR, "UTC"); + TimeStampMilliTZVector milli = new TimeStampMilliTZVector("ts_ms", ALLOCATOR, "UTC"); + TimeStampMicroTZVector micro = new TimeStampMicroTZVector("ts_us", ALLOCATOR, "UTC"); + TimeStampNanoTZVector nano = new TimeStampNanoTZVector("ts_ns", ALLOCATOR, "UTC")) { + sec.allocateNew(1); + sec.set(0, 1_700_000_000L); + sec.setValueCount(1); + milli.allocateNew(1); + milli.set(0, 1_700_000_000_123L); + milli.setValueCount(1); + micro.allocateNew(1); + micro.set(0, 1_700_000_000_123_456L); + micro.setValueCount(1); + nano.allocateNew(1); + nano.set(0, 1_700_000_000_123_456_789L); + nano.setValueCount(1); + + assertTimestamp(sec, DataTypes.TimestampType, 1_700_000_000_000_000L); + assertTimestamp(milli, DataTypes.TimestampType, 1_700_000_000_123_000L); + assertTimestamp(micro, DataTypes.TimestampType, 1_700_000_000_123_456L); + assertTimestamp(nano, DataTypes.TimestampType, 1_700_000_000_123_456L); + } + } + + private static VortexArrowColumnVector assertTimestamp( + TimeStampVector vector, org.apache.spark.sql.types.DataType expectedType, long expectedMicros) { + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(expectedType, column.dataType()); + assertEquals(expectedMicros, column.getLong(0)); + return column; + } + + @Test + @DisplayName("Microsecond duration vectors map to DayTimeIntervalType") + void duration() { + try (DurationVector vector = new DurationVector( + "dur_us", FieldType.nullable(new ArrowType.Duration(TimeUnit.MICROSECOND)), ALLOCATOR)) { + vector.allocateNew(2); + vector.set(0, 1_500L); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.createDayTimeIntervalType(), column.dataType()); + assertEquals(1_500L, column.getLong(0)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Year-month interval vectors map to YearMonthIntervalType as months") + void intervalYear() { + try (IntervalYearVector vector = new IntervalYearVector("interval_ym", ALLOCATOR)) { + vector.allocateNew(2); + vector.set(0, 14); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.createYearMonthIntervalType(), column.dataType()); + assertEquals(14, column.getInt(0)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Month-day-nano interval vectors map to CalendarIntervalType") + void intervalMonthDayNano() { + try (IntervalMonthDayNanoVector vector = new IntervalMonthDayNanoVector("interval_mdn", ALLOCATOR)) { + vector.allocateNew(2); + vector.set(0, 1, 2, 3_500L); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.CalendarIntervalType, column.dataType()); + // Nanoseconds truncate to microseconds, matching Spark's ArrowColumnVector. + assertEquals(new CalendarInterval(1, 2, 3L), column.getInterval(0)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Null vectors map to NullType") + void nullVector() { + try (NullVector vector = new NullVector("null_col")) { + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.NullType, column.dataType()); + assertTrue(column.isNullAt(0)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("List vectors map to ArrayType") + void list() { + try (ListVector vector = ListVector.empty("list", ALLOCATOR)) { + UnionListWriter writer = vector.getWriter(); + writer.allocate(); + writer.setPosition(0); + writer.startList(); + writer.writeInt(1); + writer.writeInt(2); + writer.writeInt(3); + writer.endList(); + writer.setPosition(2); + writer.startList(); + writer.endList(); + vector.setValueCount(3); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.IntegerType, ((ArrayType) column.dataType()).elementType()); + ColumnarArray array = column.getArray(0); + assertEquals(3, array.numElements()); + assertEquals(1, array.getInt(0)); + assertEquals(3, array.getInt(2)); + assertTrue(column.isNullAt(1)); + assertEquals(0, column.getArray(2).numElements()); + } + } + + @Test + @DisplayName("List view vectors map to ArrayType") + void listView() { + try (ListViewVector vector = ListViewVector.empty("list_view", ALLOCATOR)) { + UnionListViewWriter writer = vector.getWriter(); + writer.allocate(); + writer.setPosition(0); + writer.startListView(); + writer.writeInt(7); + writer.writeInt(8); + writer.endListView(); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.IntegerType, ((ArrayType) column.dataType()).elementType()); + ColumnarArray array = column.getArray(0); + assertEquals(2, array.numElements()); + assertEquals(7, array.getInt(0)); + assertEquals(8, array.getInt(1)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Map vectors map to MapType") + void map() { + try (MapVector vector = MapVector.empty("map", ALLOCATOR, false)) { + UnionMapWriter writer = vector.getWriter(); + writer.allocate(); + writer.setPosition(0); + writer.startMap(); + writer.startEntry(); + writer.key().integer().writeInt(1); + writer.value().bigInt().writeBigInt(100L); + writer.endEntry(); + writer.startEntry(); + writer.key().integer().writeInt(2); + writer.value().bigInt().writeBigInt(200L); + writer.endEntry(); + writer.endMap(); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + MapType mapType = (MapType) column.dataType(); + assertEquals(DataTypes.IntegerType, mapType.keyType()); + assertEquals(DataTypes.LongType, mapType.valueType()); + + ColumnarMap columnarMap = column.getMap(0); + assertEquals(2, columnarMap.numElements()); + assertEquals(1, columnarMap.keyArray().getInt(0)); + assertEquals(100L, columnarMap.valueArray().getLong(0)); + assertEquals(2, columnarMap.keyArray().getInt(1)); + assertEquals(200L, columnarMap.valueArray().getLong(1)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Struct vectors map to StructType with child columns") + void struct() { + try (StructVector vector = StructVector.empty("struct", ALLOCATOR)) { + NullableStructWriter writer = vector.getWriter(); + writer.allocate(); + writer.setPosition(0); + writer.start(); + writer.integer("a").writeInt(5); + writer.bigInt("b").writeBigInt(7L); + writer.end(); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + StructType structType = (StructType) column.dataType(); + assertEquals(2, structType.fields().length); + assertEquals(DataTypes.IntegerType, structType.fields()[0].dataType()); + assertEquals(DataTypes.LongType, structType.fields()[1].dataType()); + assertEquals(5, column.getChild(0).getInt(0)); + assertEquals(7L, column.getChild(1).getLong(0)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Arrow types outside the supported set are rejected with descriptive errors") + void unsupportedTypes() { + try (UInt4Vector unsignedInt = new UInt4Vector("u32", ALLOCATOR); + Float2Vector halfFloat = new Float2Vector("f16", ALLOCATOR); + Decimal256Vector decimal256 = new Decimal256Vector("d256", ALLOCATOR, 38, 2); + DateMilliVector dateMilli = new DateMilliVector("date_ms", ALLOCATOR); + TimeMicroVector time = new TimeMicroVector("time_us", ALLOCATOR); + DurationVector durationSec = new DurationVector( + "dur_s", FieldType.nullable(new ArrowType.Duration(TimeUnit.SECOND)), ALLOCATOR)) { + assertUnsupported(unsignedInt, "unsigned"); + assertUnsupported(halfFloat, "float precision"); + // Decimal256 maps to DecimalType but has no accessor, matching Spark's ArrowColumnVector. + assertUnsupported(decimal256, "Decimal256Vector"); + assertUnsupported(dateMilli, "date unit"); + assertUnsupported(time, "Time"); + assertUnsupported(durationSec, "duration unit"); + } + } + + private static void assertUnsupported( + dev.vortex.relocated.org.apache.arrow.vector.ValueVector vector, String expectedMessagePart) { + UnsupportedOperationException e = + assertThrows(UnsupportedOperationException.class, () -> new VortexArrowColumnVector(vector)); + assertTrue( + e.getMessage().contains(expectedMessagePart), + "expected error message to contain \"" + expectedMessagePart + "\" but was: " + e.getMessage()); + } +} diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java new file mode 100644 index 00000000000..cd9b3ca184c --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/write/SparkToArrowSchemaTest.java @@ -0,0 +1,176 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.write; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.vortex.relocated.org.apache.arrow.vector.types.DateUnit; +import dev.vortex.relocated.org.apache.arrow.vector.types.FloatingPointPrecision; +import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit; +import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType; +import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.Field; +import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.Schema; +import java.util.List; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link SparkToArrowSchema}, which converts Spark SQL schemas to Arrow schemas on the write path. + * + *

Characterizes the supported conversions (including field-level nullability and nested types) and the Spark types + * that are explicitly rejected, mirroring the read-path coverage in {@code ArrowUtilsTest}. + */ +final class SparkToArrowSchemaTest { + private static Field single(String name, org.apache.spark.sql.types.DataType type, boolean nullable) { + StructType schema = new StructType(new StructField[] {new StructField(name, type, nullable, Metadata.empty())}); + Schema arrowSchema = SparkToArrowSchema.convert(schema); + assertEquals(1, arrowSchema.getFields().size()); + return arrowSchema.getFields().get(0); + } + + @Test + @DisplayName("BooleanType converts to Arrow Bool") + void booleanConvertsToBool() { + assertEquals( + new ArrowType.Bool(), single("b", DataTypes.BooleanType, true).getType()); + } + + @Test + @DisplayName("Integral types convert to signed Arrow Ints of matching width") + void integralTypesConvertBySignedWidth() { + assertEquals( + new ArrowType.Int(8, true), + single("v", DataTypes.ByteType, true).getType()); + assertEquals( + new ArrowType.Int(16, true), + single("v", DataTypes.ShortType, true).getType()); + assertEquals( + new ArrowType.Int(32, true), + single("v", DataTypes.IntegerType, true).getType()); + assertEquals( + new ArrowType.Int(64, true), + single("v", DataTypes.LongType, true).getType()); + } + + @Test + @DisplayName("Float/Double convert to single/double precision Arrow FloatingPoint") + void floatingPointConvertsByPrecision() { + assertEquals( + new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), + single("v", DataTypes.FloatType, true).getType()); + assertEquals( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), + single("v", DataTypes.DoubleType, true).getType()); + } + + @Test + @DisplayName("String and Binary convert to Utf8 and Binary") + void stringAndBinaryConvert() { + assertEquals( + new ArrowType.Utf8(), single("v", DataTypes.StringType, true).getType()); + assertEquals( + new ArrowType.Binary(), single("v", DataTypes.BinaryType, true).getType()); + } + + @Test + @DisplayName("DecimalType preserves precision and scale as 128-bit Arrow Decimal") + void decimalPreservesPrecisionAndScale() { + assertEquals( + new ArrowType.Decimal(20, 4, 128), + single("v", DataTypes.createDecimalType(20, 4), true).getType()); + } + + @Test + @DisplayName("DateType converts to Date(DAY)") + void dateConvertsToDayUnit() { + assertEquals( + new ArrowType.Date(DateUnit.DAY), + single("v", DataTypes.DateType, true).getType()); + } + + @Test + @DisplayName("Timestamp converts to microsecond precision, with UTC tz only for the tz-aware type") + void timestampConvertsByTimezoneAwareness() { + assertEquals( + new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"), + single("v", DataTypes.TimestampType, true).getType()); + assertEquals( + new ArrowType.Timestamp(TimeUnit.MICROSECOND, null), + single("v", DataTypes.TimestampNTZType, true).getType()); + } + + @Test + @DisplayName("Field nullability is carried over to the Arrow field") + void nullabilityIsPreserved() { + assertTrue(single("v", DataTypes.IntegerType, true).isNullable()); + assertFalse(single("v", DataTypes.IntegerType, false).isNullable()); + } + + @Test + @DisplayName("Field names and ordering are preserved") + void fieldNamesAndOrderArePreserved() { + StructType schema = new StructType() + .add("first", DataTypes.IntegerType) + .add("second", DataTypes.StringType) + .add("third", DataTypes.DoubleType); + + Schema arrowSchema = SparkToArrowSchema.convert(schema); + assertEquals( + List.of("first", "second", "third"), + arrowSchema.getFields().stream().map(Field::getName).toList()); + } + + @Test + @DisplayName("Nested StructType converts to a Struct field with converted children") + void structConvertsWithChildren() { + StructType inner = + new StructType().add("a", DataTypes.IntegerType, true).add("b", DataTypes.StringType, false); + Field structField = single("s", inner, true); + + assertEquals(new ArrowType.Struct(), structField.getType()); + assertEquals(2, structField.getChildren().size()); + + Field a = structField.getChildren().get(0); + assertEquals("a", a.getName()); + assertEquals(new ArrowType.Int(32, true), a.getType()); + assertTrue(a.isNullable()); + + Field b = structField.getChildren().get(1); + assertEquals("b", b.getName()); + assertEquals(new ArrowType.Utf8(), b.getType()); + assertFalse(b.isNullable()); + } + + @Test + @DisplayName("ArrayType converts to a List field whose element carries containsNull") + void arrayConvertsWithElementNullability() { + Field withNulls = single("l", DataTypes.createArrayType(DataTypes.IntegerType, true), true); + assertEquals(new ArrowType.List(), withNulls.getType()); + assertEquals(1, withNulls.getChildren().size()); + + Field element = withNulls.getChildren().get(0); + assertEquals("element", element.getName()); + assertEquals(new ArrowType.Int(32, true), element.getType()); + assertTrue(element.isNullable()); + + Field withoutNulls = single("l", DataTypes.createArrayType(DataTypes.StringType, false), true); + assertFalse(withoutNulls.getChildren().get(0).isNullable()); + } + + @Test + @DisplayName("Unsupported Spark types raise UnsupportedOperationException naming the type") + void unsupportedTypeIsRejected() { + StructType schema = new StructType().add("v", DataTypes.CalendarIntervalType); + UnsupportedOperationException e = + assertThrows(UnsupportedOperationException.class, () -> SparkToArrowSchema.convert(schema)); + assertTrue(e.getMessage().contains("CalendarIntervalType")); + } +} diff --git a/lang/cpp/.gitignore b/lang/cpp/.gitignore new file mode 100644 index 00000000000..3ea99d6c913 --- /dev/null +++ b/lang/cpp/.gitignore @@ -0,0 +1 @@ +coverage* diff --git a/lang/cpp/CMakeLists.txt b/lang/cpp/CMakeLists.txt new file mode 100644 index 00000000000..8949c1977cc --- /dev/null +++ b/lang/cpp/CMakeLists.txt @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +cmake_minimum_required(VERSION 3.10) +project(VortexCXX + VERSION 0.0.1 + LANGUAGES CXX) + +include(FetchContent) + +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +find_program(SCCACHE_PROGRAM sccache) +if (SCCACHE_PROGRAM) + set(CMAKE_C_COMPILER_LAUNCHER "${SCCACHE_PROGRAM}") + set(CMAKE_CXX_COMPILER_LAUNCHER "${SCCACHE_PROGRAM}") + message(STATUS "Sccache found: ${SCCACHE_PROGRAM}") +else () + message(STATUS "Sccache not found") +endif () + +option(BUILD_TESTS "Build tests" OFF) + +set(SANITIZER "" CACHE STRING "Build with sanitizers") +set(TARGET_TRIPLE "" CACHE STRING "Rust target triple for FFI library") +set(RUST_BUILD_PROFILE "" CACHE STRING "Cargo profile name for Rust FFI library") + +if (NOT CMAKE_BUILD_TYPE) + set(CMAKE_BUILD_TYPE Debug) +endif() + +if (NOT SANITIZER STREQUAL "") + message(NOTICE "Sanitizer: ${SANITIZER}") + if (NOT CMAKE_BUILD_TYPE STREQUAL "Debug") + message(FATAL_ERROR "Only debug build is supported for sanitizer builds") + endif() + + if (SANITIZER STREQUAL "asan") + set(SANITIZER_FLAGS "-fsanitize=address,undefined,leak") + elseif (SANITIZER STREQUAL "tsan") + set(SANITIZER_FLAGS "-fsanitize=thread") + else() + message(FATAL_ERROR "Unknown sanitizer ${SANITIZER}") + endif() + set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${SANITIZER_FLAGS}") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${SANITIZER_FLAGS}") + set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} ${SANITIZER_FLAGS}") +endif() + +if (RUST_BUILD_PROFILE STREQUAL "") + if (CMAKE_BUILD_TYPE STREQUAL "RelWithDebInfo") + set(RUST_BUILD_PROFILE "release_debug") + else() + string(TOLOWER "${CMAKE_BUILD_TYPE}" RUST_BUILD_PROFILE) + endif() +endif() + +set(VORTEX_FFI_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../vortex-ffi") +set(VORTEX_FFI_LIB_DIR "${VORTEX_FFI_DIR}/../target/${TARGET_TRIPLE}/${RUST_BUILD_PROFILE}") +set(VORTEX_FFI_HEADERS "${VORTEX_FFI_DIR}/cinclude") + +if(APPLE) + set(VORTEX_FFI_STATIC "${VORTEX_FFI_LIB_DIR}/libvortex_ffi.a") + set(VORTEX_FFI_SHARED "${VORTEX_FFI_LIB_DIR}/libvortex_ffi.dylib") +elseif(WIN32) + set(VORTEX_FFI_STATIC "${VORTEX_FFI_LIB_DIR}/libvortex_ffi.lib") + set(VORTEX_FFI_SHARED "${VORTEX_FFI_LIB_DIR}/libvortex_ffi.dll") +else() + set(VORTEX_FFI_STATIC "${VORTEX_FFI_LIB_DIR}/libvortex_ffi.a") + set(VORTEX_FFI_SHARED "${VORTEX_FFI_LIB_DIR}/libvortex_ffi.so") +endif() + +if(NOT EXISTS "${VORTEX_FFI_STATIC}") + message(FATAL_ERROR + "vortex-ffi static library not found at ${VORTEX_FFI_STATIC}. " + "Run: cargo build --profile -p vortex-ffi") +endif() + +add_library(vortex_ffi STATIC IMPORTED) +set_target_properties(vortex_ffi PROPERTIES + IMPORTED_LOCATION "${VORTEX_FFI_STATIC}" + INTERFACE_INCLUDE_DIRECTORIES "${VORTEX_FFI_HEADERS}") + +if(EXISTS "${VORTEX_FFI_SHARED}") + add_library(vortex_ffi_shared SHARED IMPORTED) + set_target_properties(vortex_ffi_shared PROPERTIES + IMPORTED_LOCATION "${VORTEX_FFI_SHARED}" + INTERFACE_INCLUDE_DIRECTORIES "${VORTEX_FFI_HEADERS}" + INTERFACE_LINK_OPTIONS "LINKER:-rpath,${VORTEX_FFI_LIB_DIR}") +endif() + +file(GLOB VORTEX_CXX_SOURCES CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/src/*.cpp") +add_library(vortex_cxx STATIC ${VORTEX_CXX_SOURCES}) +target_include_directories(vortex_cxx PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR} + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/../vortex-ffi/cinclude +) +target_link_libraries(vortex_cxx PUBLIC vortex_ffi) +add_library(vortex::cxx ALIAS vortex_cxx) +target_compile_features(vortex_cxx PUBLIC cxx_std_20) + +set(APPLE_LINK_FLAGS "-framework CoreFoundation -framework Security") + +if(APPLE) + target_link_libraries(vortex_cxx PRIVATE ${APPLE_LINK_FLAGS}) +endif() + +if(TARGET vortex_ffi_shared) + add_library(vortex_cxx_shared SHARED ${VORTEX_CXX_SOURCES}) + set_target_properties(vortex_cxx_shared PROPERTIES POSITION_INDEPENDENT_CODE ON) + target_include_directories(vortex_cxx_shared PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR} + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/../vortex-ffi/cinclude + ) + target_link_libraries(vortex_cxx_shared PUBLIC vortex_ffi_shared) + add_library(vortex::cxx_shared ALIAS vortex_cxx_shared) + target_compile_features(vortex_cxx_shared PUBLIC cxx_std_20) + + if(APPLE) + target_link_libraries(vortex_cxx_shared PRIVATE ${APPLE_LINK_FLAGS}) + endif() +endif() + +if (BUILD_TESTS) + FetchContent_Declare( + Nanoarrow + GIT_REPOSITORY https://github.com/apache/arrow-nanoarrow + GIT_TAG apache-arrow-nanoarrow-0.8.0 + ) + FetchContent_MakeAvailable(Nanoarrow) + + FetchContent_Declare( + Catch + GIT_REPOSITORY https://github.com/catchorg/Catch2.git + GIT_TAG v3.8.1 + ) + FetchContent_MakeAvailable(Catch) + include(Catch) + target_compile_definitions(Catch2 PRIVATE CATCH_CONFIG_NO_POSIX_SIGNALS) + + include(CTest) + add_subdirectory(tests) +endif() diff --git a/lang/cpp/README.md b/lang/cpp/README.md new file mode 100644 index 00000000000..fe9d2527e63 --- /dev/null +++ b/lang/cpp/README.md @@ -0,0 +1,45 @@ +# Vortex C++ bindings + +For a usage guide, see docs/api/cpp/index.rst. + +## Requirements + +- CMake 3.10+ +- C++20 compiler (C++23 compiler for tests). +- Rust toolchain for building `vortex-ffi`. + +## Build + +```sh +cargo build --release -p vortex-ffi +cmake -Bbuild -DCMAKE_BUILD_TYPE=Release +cmake --build build -j +``` + +This will generate `libvortex_cxx` shared and static libraries. +You can use `vortex_cxx` and `vortex_cxx_shared` CMake targets. + +## Test + +```sh +cargo build --release -p vortex-ffi +cmake -Bbuild -DBUILD_TESTS=ON +cmake --build build -j +ctest --test-dir build -j "$(nproc)" +``` + +## Run examples + +```sh +cmake -Bbuild -DBUILD_EXAMPLES=ON +cmake --build build -j +./build/examples/hello-vortex +``` + +## Check coverage + +This will generate an LCOV directory `coverage`: + +```sh +./gcov-report.sh generate +``` diff --git a/lang/cpp/gcov-report.sh b/lang/cpp/gcov-report.sh new file mode 100755 index 00000000000..6c1365df57a --- /dev/null +++ b/lang/cpp/gcov-report.sh @@ -0,0 +1,18 @@ +#!/bin/sh + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +set -eu +cmake -Bbuild -DBUILD_TESTS=1 -DCMAKE_CXX_FLAGS='-fprofile-arcs -ftest-coverage' +cmake --build build -j +ctest --test-dir build --output-on-failure + +geninfo build/CMakeFiles/vortex_cxx_shared.dir/ \ + build/tests/CMakeFiles/vortex_cxx_test.dir/ \ + --rc geninfo_unexecuted_blocks=1 \ + --exclude /usr --exclude build/_deps --exclude tests \ + -j -b src -o coverage.info +if [ $# -gt 0 ]; then + genhtml coverage.info -o coverage +fi diff --git a/lang/cpp/include/vortex/array.hpp b/lang/cpp/include/vortex/array.hpp new file mode 100644 index 00000000000..43b986f4b42 --- /dev/null +++ b/lang/cpp/include/vortex/array.hpp @@ -0,0 +1,401 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "vortex/common.hpp" +#include "vortex/dtype.hpp" +#include "vortex/error.hpp" +#include "vortex/expression.hpp" +#include "vortex/session.hpp" + +#include + +#include +#include +#include +#include +#include +#include + +namespace vortex { + +// Types that a PrimitiveView can hold +template +concept primitive_view = primitive_type || std::is_same_v; + +template +class PrimitiveView; + +class Array; +class StringView; +class BytesView; + +/* + * Validity type tells us whether there are null/invalid values in an Array. + */ +enum class ValidityType { + // Items can't be null + NonNullable = VX_VALIDITY_NON_NULLABLE, + // All items are valid + AllValid = VX_VALIDITY_ALL_VALID, + // All items are invalid + AllInvalid = VX_VALIDITY_ALL_INVALID, + // Item validity is set from a boolean array: true = valid, false = invalid + FromArray = VX_VALIDITY_ARRAY, +}; + +/** + * Array per-element validity of type ValidityType. + * If ValidityType is ValidityType::Array, holds a boolean array + * with validity items. + * + * You can use shortcut constants NonNullable/AllValid/AllInvalid and + * function ValidityArray(validity_bools); + */ +class Validity { +public: + // NonNullable/AllValid/AllInvalid constructor + // NOLINTNEXTLINE(google-explicit-constructor) + Validity(ValidityType type); + // Validity determined by a boolean array, true = valid, false = invalid. + static Validity from_array(const Array &bools); + + Validity(const Validity &other); + Validity(Validity &&other) noexcept; + Validity &operator=(const Validity &other); + Validity &operator=(Validity &&other) noexcept; + ~Validity(); + + ValidityType type() const { + return type_; + } + + // Boolean validity array. Throws if type() != ValidityType::Array. + Array array() const; + +private: + friend struct detail::Access; + friend Validity ValidityArray(const Array &bools); + Validity(ValidityType type, const vx_array *owned) : type_(type), array_(owned) { + } + + ValidityType type_; + const vx_array *array_; +}; + +namespace detail { +// Validity bitmap for typed views. Owns the arrays that back the bits +class ValidityBits { +public: + ValidityBits(ValidityBits &&other) noexcept; + ValidityBits &operator=(ValidityBits &&other) noexcept; + ValidityBits(const ValidityBits &) = delete; + ValidityBits &operator=(const ValidityBits &) = delete; + ~ValidityBits(); + + bool is_null(size_t index) const; + +private: + friend class vortex::Array; + // Materialize validity of "canonical" + ValidityBits(const Session &session, const vx_array *canonical); + + const vx_array *owner_ = nullptr; + const uint8_t *bits_ = nullptr; + size_t bit_offset_ = 0; + bool all_invalid_ = false; +}; +} // namespace detail + +// A reference-counted handle to columnar data in some encoding +class Array { +public: + Array(const Array &other); + Array(Array &&) noexcept = default; + Array &operator=(const Array &other); + Array &operator=(Array &&) noexcept = default; + + // An all-null array with DataType Null. + static Array null(size_t len); + + /** + * A Primitive array copied from a typed buffer. + * + * Example: + * + * std::array buffer = {0, 1, 2}; + * auto array = Array::primitive(buffer); + */ + template + static Array primitive(std::span data, const Validity &validity = ValidityType::NonNullable) { + return primitive_raw(detail::to_ptype(), data.data(), data.size(), validity); + } + + /** + * Import an Arrow array. Consumes both "array" and "schema", do not use + * or release them afterwards. For a record batch pass nullable = false. + */ + static Array from_arrow(ArrowArray *array, ArrowSchema *schema, bool nullable); + + size_t size() const; + bool nullable() const; + bool has_dtype(DataTypeVariant variant) const; + bool is_primitive(PType ptype) const; + DataType dtype() const; + Validity validity() const; + + // Number of null/invalid elements in Array + size_t null_count() const; + + /** + * Get a Struct field by index. Throws if Array is not a Struct or if index + * is out of bounds. + */ + Array field(size_t index) const; + + /** + * Get a Struct field by name. Throws if Array is not a Struct or doesn't + * have this named field. + */ + Array field(std::string_view name) const; + + /* + * Create a new Array slicing [begin; end) rows from original. + * Doesn't copy the original buffer or sliced buffer. + * + * Example: + * + * std::array buffer = {0, 1, 2}; + * Array array = Array::primitive(buffer); + * Array sliced = array.slice(1, 2); + */ + Array slice(size_t begin, size_t end) const; + + /** + * Apply an expression to an array. + * + * This function operates in constant time and doesn't execute the result + * array. To execute the array, canonicalise it. + * + * Example: + * + * using namespace vortex::expr::ops; + * + * std::array buffer = {0, 1, 2}; + * Array array = Array::primitive(buffer); + * Expression expr = expr::root() > expr::lit(0); + * Array result = array.apply(expr); + */ + Array apply(const Expression &expr) const; + + /** + * Bulk view over values. Canonicalizes the array. + * Throws if T does not match Array's ptype. + * + * Example: + * + * Session session; + * std::array buffer = {0, 1, 2}; + * Array array = Array::primitive(buffer); + * auto view = array.values(session); + */ + template + PrimitiveView values(const Session &session) const; + + // Bulk view over Bool values + PrimitiveView bools(const Session &session) const; + + // Bulk view over Utf8 values. + StringView strings(const Session &session) const; + + // Bulk view over Binary values. + BytesView bytes(const Session &session) const; + +private: + friend struct detail::Access; + friend class StringView; + friend class BytesView; + template + friend class PrimitiveView; + + explicit Array(const vx_array *owned) : handle_(owned) { + } + const vx_array *release() && { + return handle_.release(); + } + + static Array primitive_raw(vx_ptype ptype, const void *data, size_t len, const Validity &validity); + Array canonicalize(const Session &session) const; + + struct Deleter { + void operator()(const vx_array *ptr) const noexcept; + }; + std::unique_ptr handle_; +}; + +// Column field of a Struct Array +struct ColumnField { + std::string name; + Array column; +}; + +/** + * Create a Struct array from named columns of equal length. + * + * Example: + * + * using enum ValidityType; + * std::array age_buffer = {0, 1, 2}; + * std::array height_buffer = {0, 1, 2}; + * Array ages = Array::primitive(age_buffer); + * Array heights = Array::primitive(height_buffer); + * Array result = make_struct( + * {{"age", ages}, {"height", heights}}, + * NonNullable); + */ +Array make_struct(std::span fields, const Validity &validity = ValidityType::NonNullable); +Array make_struct(std::initializer_list fields, + const Validity &validity = ValidityType::NonNullable); + +/** + * Typed read-only view over a Primitive array. + * + * Owns a canonicalized copy of Array. values() and anything derived from + * it are valid as long as the view lives. + */ +template +class PrimitiveView { +public: + /* + * Get raw values from this view. Values at null/invalid positions are + * unspecified. + */ + std::span values() const { + return {data_, size_}; + } + bool is_null(size_t index) const { + return validity_.is_null(index); + } + size_t size() const { + return size_; + } + +private: + friend class Array; + PrimitiveView(Array canonical, detail::ValidityBits validity, const T *data, size_t size) + : canonical_(std::move(canonical)), validity_(std::move(validity)), data_(data), size_(size) { + } + + Array canonical_; + detail::ValidityBits validity_; + const T *data_; + size_t size_; +}; + +/** + * Read-only view over a Bool array. As Bool values are bit-packed, there's no + * span. Read individual values with value(i). + */ +template <> +class PrimitiveView { +public: + /* + * Get raw value from this view. Values at null/invalid positions are + * unspecified. + */ + bool value(size_t index) const; + bool is_null(size_t index) const { + return validity_.is_null(index); + } + size_t size() const { + return size_; + } + +private: + friend class Array; + PrimitiveView(Array canonical, detail::ValidityBits validity, size_t size) + : canonical_(std::move(canonical)), validity_(std::move(validity)), size_(size) { + } + + Array canonical_; + detail::ValidityBits validity_; + size_t size_; +}; + +/** + * Read-only view over a Utf8 array. + * + * operator[] is O(1) and borrows from the view's canonical copy. Returned + * string_views are valid as long as the view lives. + */ +class StringView { +public: + /* + * Get raw value from this view. Values at null/invalid positions are + * unspecified. + */ + std::string_view operator[](size_t index) const; + bool is_null(size_t index) const { + return validity_.is_null(index); + } + size_t size() const { + return size_; + } + +private: + friend class Array; + StringView(Array canonical, detail::ValidityBits validity, size_t size) + : canonical_(std::move(canonical)), validity_(std::move(validity)), size_(size) { + } + + Array canonical_; + detail::ValidityBits validity_; + size_t size_; +}; + +/** + * Read-only view over a Bytes array. + * + * Byte spans borrow from the view's canonical copy and are valid as long as + * the view lives. + */ +class BytesView { +public: + /* + * Get raw value from this view. Values at null/invalid positions are + * unspecified. + */ + BinaryView operator[](size_t index) const; + bool is_null(size_t index) const { + return validity_.is_null(index); + } + size_t size() const { + return size_; + } + +private: + friend class Array; + BytesView(Array canonical, detail::ValidityBits validity, size_t size) + : canonical_(std::move(canonical)), validity_(std::move(validity)), size_(size) { + } + + Array canonical_; + detail::ValidityBits validity_; + size_t size_; +}; + +template +PrimitiveView Array::values(const Session &session) const { + Array canonical = canonicalize(session); + const vx_array *raw = detail::Access::c_ptr(canonical); + if (!vx_array_is_primitive(raw, detail::to_ptype())) { + throw VortexException("values: T does not match the array's ptype", ErrorCode::MismatchedTypes); + } + vx_error *error = nullptr; + const void *data = vx_array_data_ptr_primitive(raw, &error); + detail::throw_on_error(error); + detail::ValidityBits validity(session, raw); + const size_t n = vx_array_len(raw); + return PrimitiveView(std::move(canonical), std::move(validity), static_cast(data), n); +} +} // namespace vortex diff --git a/lang/cpp/include/vortex/common.hpp b/lang/cpp/include/vortex/common.hpp new file mode 100644 index 00000000000..cc8559e96c6 --- /dev/null +++ b/lang/cpp/include/vortex/common.hpp @@ -0,0 +1,134 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include +#include +#include +#include +#include +#include + +#if __STDCPP_FLOAT16_T__ != 1 + +namespace vortex { +struct float16_t { + uint16_t bits; + constexpr friend bool operator==(float16_t, float16_t) = default; + // NOLINTNEXTLINE + constexpr operator float() const { + float result; + const uint32_t sign = (bits >> 15) & 1; + const uint32_t exponent = (bits >> 10) & 0x1F; + const uint32_t mantissa = bits & 0x3FF; + + uint32_t out; + if (exponent == 0x1F) { + out = (sign << 31) | 0x7F800000 | (mantissa << 13); + } else if (exponent == 0) { + if (mantissa == 0) { + out = sign << 31; + } else { + uint32_t m = mantissa; + int e = -1; + do { + m <<= 1; + ++e; + } while ((m & 0x400) == 0); + out = (sign << 31) | ((127 - 15 - e) << 23) | ((m & 0x3FF) << 13); + } + } else { + out = (sign << 31) | ((exponent - 15 + 127) << 23) | (mantissa << 13); + } + return std::bit_cast(out); + } +}; +static_assert(sizeof(float16_t) == 2 && std::is_trivially_copyable_v); +} // namespace vortex + +#else + +#include + +namespace vortex { +using float16_t = ::std::float16_t; +} // namespace vortex + +#endif + +namespace vortex::detail { +struct Access { + template + static T adopt(Args &&...args) { + return T(std::forward(args)...); + } + template + static auto release(T &&t) { + return std::forward(t).release(); + } + template + static auto c_ptr(const T &t) { + return t.handle_.get(); + } +}; + +template +constexpr vx_ptype to_ptype() { + if constexpr (std::is_same_v) { + return PTYPE_U8; + } else if constexpr (std::is_same_v) { + return PTYPE_U16; + } else if constexpr (std::is_same_v) { + return PTYPE_U32; + } else if constexpr (std::is_same_v) { + return PTYPE_U64; + } else if constexpr (std::is_same_v) { + return PTYPE_I8; + } else if constexpr (std::is_same_v) { + return PTYPE_I16; + } else if constexpr (std::is_same_v) { + return PTYPE_I32; + } else if constexpr (std::is_same_v) { + return PTYPE_I64; +#if __STDCPP_FLOAT16_T__ == 1 + } else if constexpr (std::is_same_v) { + return PTYPE_F16; +#endif + } else if constexpr (std::is_same_v) { + return PTYPE_F16; + } else if constexpr (std::is_same_v) { + return PTYPE_F32; + } else { + static_assert(std::is_same_v); + return PTYPE_F64; + } +} + +template +inline constexpr bool is_numeric_element = + std::is_same_v || std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || std::is_same_v || + std::is_same_v +#if __STDCPP_FLOAT16_T__ == 1 + || std::is_same_v +#else + || std::is_same_v +#endif + ; +} // namespace vortex::detail + +namespace vortex { + +// View over a single Binary byte range. +using BinaryView = std::span; + +// Types that can be stored in a Primitive array +template +concept primitive_type = detail::is_numeric_element; + +// Types constructible as scalars or literals. +template +concept element_type = primitive_type || std::is_same_v || std::is_same_v || + std::is_same_v; +} // namespace vortex diff --git a/lang/cpp/include/vortex/data_source.hpp b/lang/cpp/include/vortex/data_source.hpp new file mode 100644 index 00000000000..f3221f37068 --- /dev/null +++ b/lang/cpp/include/vortex/data_source.hpp @@ -0,0 +1,60 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "vortex/common.hpp" +#include "vortex/estimate.hpp" +#include "vortex/scan.hpp" +#include "vortex/session.hpp" + +#include + +#include +#include +#include +#include +#include + +namespace vortex { + +/** + * A reference to one or more (possibly remote, possibly glob) paths. + * + * Creating a DataSource opens the first matched path to read the schema. All + * other IO is deferred until a scan is requested. Multiple scans may be + * requested from one data source. + */ +class DataSource { +public: + static DataSource open(const Session &session, std::initializer_list paths); + static DataSource open(const Session &session, std::span paths); + static DataSource open(const Session &session, std::span paths); + + /** + * Create a DataSource from an in-memory Vortex file. Borrows the buffer: + * caller must keep it alive and unmodified while DataSource, Scans, + * Partitions, and Arrays obtained from this buffer are alive. + */ + static DataSource from_buffer(const Session &session, std::span data); + + DataSource(const DataSource &other); + DataSource(DataSource &&) noexcept = default; + DataSource &operator=(const DataSource &other); + DataSource &operator=(DataSource &&) noexcept = default; + + Estimate row_count() const; + DataType dtype() const; + Scan scan(const ScanOptions &options = {}) const; + +private: + friend struct detail::Access; + DataSource(const vx_data_source *owned, Session session) : handle_(owned), session_(std::move(session)) { + } + + struct Deleter { + void operator()(const vx_data_source *ptr) const noexcept; + }; + std::unique_ptr handle_; + Session session_; +}; +} // namespace vortex diff --git a/lang/cpp/include/vortex/dtype.hpp b/lang/cpp/include/vortex/dtype.hpp new file mode 100644 index 00000000000..1be5007c6d6 --- /dev/null +++ b/lang/cpp/include/vortex/dtype.hpp @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "vortex/common.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace vortex { + +struct StructField; + +enum class DataTypeVariant { + Null = DTYPE_NULL, + Bool = DTYPE_BOOL, + // Primitives e.g., u8, i16, f32 + Primitive = DTYPE_PRIMITIVE, + // Variable-length UTF-8 string + Utf8 = DTYPE_UTF8, + // Variable-length binary + Binary = DTYPE_BINARY, + // Nested struct + Struct = DTYPE_STRUCT, + // Nested list + List = DTYPE_LIST, + // User-defined extension + Extension = DTYPE_EXTENSION, + // Decimal with fixed precision and scale + Decimal = DTYPE_DECIMAL, + // Nested fixed-size list + FixedSizeList = DTYPE_FIXED_SIZE_LIST, +}; + +// Primitive type +enum class PType { + U8 = PTYPE_U8, + U16 = PTYPE_U16, + U32 = PTYPE_U32, + U64 = PTYPE_U64, + I8 = PTYPE_I8, + I16 = PTYPE_I16, + I32 = PTYPE_I32, + I64 = PTYPE_I64, + F16 = PTYPE_F16, + F32 = PTYPE_F32, + F64 = PTYPE_F64, +}; + +/** + * A Vortex data type. Data types are logical: they say nothing about physical + * representation. + */ +class DataType { +public: + DataType(const DataType &other); + DataType(DataType &&) noexcept = default; + DataType &operator=(const DataType &other); + DataType &operator=(DataType &&) noexcept = default; + + /** + * Consume an ArrowSchema and convert it into a DataType. The schema must + * not be used after this call. + */ + static DataType from_arrow(ArrowSchema *schema); + + /** + * Convert dtype to an Arrow C schema. Caller is responsible for invoking + * schema's release() callback. + */ + ArrowSchema to_arrow() const; + + DataTypeVariant variant() const; + bool nullable() const; + + PType primitive_type() const; + uint8_t decimal_precision() const; + int8_t decimal_scale() const; + + /** + * For a Struct dtype, return its fields in order. + * Throws if DataType is not Struct. + */ + std::vector fields() const; + + // List accessors. Valid only on List and FixedSizeList dtypes + + DataType list_element() const; + DataType fixed_size_list_element() const; + uint32_t fixed_size_list_size() const; + +private: + friend struct detail::Access; + explicit DataType(const vx_dtype *owned); + const vx_dtype *release() && { + return handle_.release(); + } + + struct Deleter { + void operator()(const vx_dtype *ptr) const noexcept; + }; + std::unique_ptr handle_; +}; + +// Field of a Struct DataType. +struct StructField { + std::string name; + DataType dtype; +}; + +namespace dtype { + +inline constexpr bool Nullable = true; + +DataType null(); +DataType boolean(bool nullable = false); +DataType primitive(PType ptype, bool nullable = false); +DataType int8(bool nullable = false); +DataType int16(bool nullable = false); +DataType int32(bool nullable = false); +DataType int64(bool nullable = false); +DataType uint8(bool nullable = false); +DataType uint16(bool nullable = false); +DataType uint32(bool nullable = false); +DataType uint64(bool nullable = false); +DataType float16(bool nullable = false); +DataType float32(bool nullable = false); +DataType float64(bool nullable = false); +DataType utf8(bool nullable = false); +DataType binary(bool nullable = false); +DataType decimal(uint8_t precision, int8_t scale, bool nullable = false); +DataType list(DataType element, bool nullable = false); +DataType fixed_size_list(DataType element, uint32_t size, bool nullable = false); + +/** + * Create a DataTypeVariant::Struct from a field list. + * + * Example: + * + * using dtype::Nullable; + * DataType dtype = dtype::struct_({ + * {"age", dtype::uint8()}, + * {"height", dtype::uint16(Nullable)}} + * ); + */ +DataType struct_(std::span fields, bool nullable = false); +DataType struct_(std::initializer_list fields, bool nullable = false); +} // namespace dtype +} // namespace vortex diff --git a/lang/cpp/include/vortex/error.hpp b/lang/cpp/include/vortex/error.hpp new file mode 100644 index 00000000000..8793be76dbb --- /dev/null +++ b/lang/cpp/include/vortex/error.hpp @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include + +#include +#include +#include + +namespace vortex { +enum class ErrorCode { + Other = VX_ERROR_CODE_OTHER, + OutOfBounds = VX_ERROR_CODE_OUT_OF_BOUNDS, + Compute = VX_ERROR_CODE_COMPUTE, + InvalidArgument = VX_ERROR_CODE_INVALID_ARGUMENT, + Serialization = VX_ERROR_CODE_SERIALIZATION, + NotImplemented = VX_ERROR_CODE_NOT_IMPLEMENTED, + MismatchedTypes = VX_ERROR_CODE_MISMATCHED_TYPES, + AssertionFailed = VX_ERROR_CODE_ASSERTION_FAILED, + Io = VX_ERROR_CODE_IO, + Panic = VX_ERROR_CODE_PANIC, +}; + +class VortexException : public std::runtime_error { +public: + VortexException(const std::string &message, ErrorCode code) : std::runtime_error(message), code_(code) { + } + + ErrorCode code() const { + return code_; + } + +private: + ErrorCode code_; +}; + +namespace detail { +// Throw VortexException and free "error" if it is non-nullptr. +inline void throw_on_error(vx_error *error) { + if (error == nullptr) { + return; + } + const vx_view str = vx_error_message(error); + const std::string message {str.ptr, str.len}; + const auto code = static_cast(vx_error_get_code(error)); + vx_error_free(error); + throw VortexException(message, code); +} + +inline vx_view to_view(std::string_view view) { + return {view.data(), view.size()}; +} +} // namespace detail +} // namespace vortex diff --git a/lang/cpp/include/vortex/estimate.hpp b/lang/cpp/include/vortex/estimate.hpp new file mode 100644 index 00000000000..1b412555b0f --- /dev/null +++ b/lang/cpp/include/vortex/estimate.hpp @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "vortex/error.hpp" +#include + +#include + +namespace vortex { + +enum class EstimateType { + Unknown = VX_ESTIMATE_UNKNOWN, + Exact = VX_ESTIMATE_EXACT, + Inexact = VX_ESTIMATE_INEXACT, +}; + +// Estimated count (rows in a partition, partitions in a scan) +class Estimate { +public: + explicit Estimate(vx_estimate raw) : raw_(raw) { + } + + inline EstimateType type() const { + return static_cast(raw_.type); + } + + /** + * Estimated count. Throws if type() is Unknown. For inexact estimates this + * is an upper bound. + */ + inline uint64_t value() const { + if (type() == EstimateType::Unknown) { + throw VortexException("estimate is unknown", ErrorCode::InvalidArgument); + } + return raw_.estimate; + } + + inline uint64_t value_or(uint64_t fallback) const { + return type() == EstimateType::Unknown ? fallback : raw_.estimate; + } + +private: + vx_estimate raw_; +}; + +} // namespace vortex diff --git a/lang/cpp/include/vortex/expression.hpp b/lang/cpp/include/vortex/expression.hpp new file mode 100644 index 00000000000..01547450ae5 --- /dev/null +++ b/lang/cpp/include/vortex/expression.hpp @@ -0,0 +1,182 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "vortex/common.hpp" +#include "vortex/scalar.hpp" + +#include + +#include +#include +#include +#include + +namespace vortex { + +// A node in an expression tree used for scan filters and projections. +class Expression { +public: + Expression(const Expression &); + Expression &operator=(const Expression &); + Expression(Expression &&) noexcept = default; + Expression &operator=(Expression &&) noexcept = default; + + /** + * Extract field from a Struct. Output DataType is field's DataType. + * + * Errors at scan/apply time if field does not exist or if root() is + * not a Struct. + * + * Example: + * + * Expression field = root()["age"]["nested"]; + */ + Expression operator[](std::string_view field) const; + + Expression is_null() const; + + /* + * Extract fields from a Struct. Output DataType is a Struct. + * + * Errors at scan/apply time if any of fields does not exist or if root() + * is not a Struct. + */ + Expression select(std::span names) const; + Expression select(std::span names) const; + Expression select(std::initializer_list names) const; + +private: + friend struct detail::Access; + explicit Expression(const vx_expression *owned) : handle_(owned) { + } + const vx_expression *release() && { + return handle_.release(); + } + + struct Deleter { + void operator()(const vx_expression *ptr) const noexcept; + }; + std::unique_ptr handle_; +}; + +// Add, Sub, Mul, and Div error in runtime on overflow and underflow +enum class BinaryOperator { + // x == y + Eq = VX_OPERATOR_EQ, + // x != y + NotEq = VX_OPERATOR_NOT_EQ, + // x > y + Gt = VX_OPERATOR_GT, + // x >= y + Gte = VX_OPERATOR_GTE, + // x < y + Lt = VX_OPERATOR_LT, + // x <= y + Lte = VX_OPERATOR_LTE, + // boolean x AND y, Kleene semantics + KleeneAnd = VX_OPERATOR_KLEENE_AND, + // boolean x OR y, Kleene semantics + KleeneOr = VX_OPERATOR_KLEENE_OR, + // x + y + Add = VX_OPERATOR_ADD, + // x - y + Sub = VX_OPERATOR_SUB, + // x * y + Mul = VX_OPERATOR_MUL, + // x / y + Div = VX_OPERATOR_DIV, +}; + +namespace expr { + +// scanned/applied array. +Expression root(); + +// root()'s named column. +Expression col(std::string_view name); + +// Literal expression +Expression lit(const Scalar &value); + +/* + * Literal expression. + * + * Literal's DataType must match column it's compared against, otherwise scan + * fails at runtime. No type coercion is performed. + */ +template +Expression lit(T value) { + return lit(scalar::of(value)); +} + +Expression eq(const Expression &l, const Expression &r); +Expression neq(const Expression &l, const Expression &r); +Expression lt(const Expression &l, const Expression &r); +Expression lte(const Expression &l, const Expression &r); +Expression gt(const Expression &l, const Expression &r); +Expression gte(const Expression &l, const Expression &r); +Expression add(const Expression &l, const Expression &r); +Expression sub(const Expression &l, const Expression &r); +Expression mul(const Expression &l, const Expression &r); +Expression div(const Expression &l, const Expression &r); +Expression binary_op(BinaryOperator op, const Expression &l, const Expression &r); + +// Kleene AND of children. Errors on an empty list. +Expression and_all(std::span children); +// Kleene OR of children. Errors on an empty list. +Expression or_all(std::span children); + +Expression logical_not(const Expression &child); +Expression is_null(const Expression &child); +Expression list_contains(const Expression &list, const Expression &value); + +/** + * Opt-in operator overloads like in Eigen. + * Note && and || don't short-circuit. + */ +namespace ops { + +inline Expression operator==(const Expression &l, const Expression &r) { + return eq(l, r); +} +inline Expression operator!=(const Expression &l, const Expression &r) { + return neq(l, r); +} +inline Expression operator<(const Expression &l, const Expression &r) { + return lt(l, r); +} +inline Expression operator<=(const Expression &l, const Expression &r) { + return lte(l, r); +} +inline Expression operator>(const Expression &l, const Expression &r) { + return gt(l, r); +} +inline Expression operator>=(const Expression &l, const Expression &r) { + return gte(l, r); +} +inline Expression operator&&(const Expression &l, const Expression &r) { + return binary_op(BinaryOperator::KleeneAnd, l, r); +} +inline Expression operator||(const Expression &l, const Expression &r) { + return binary_op(BinaryOperator::KleeneOr, l, r); +} +inline Expression operator!(const Expression &e) { + return logical_not(e); +} +inline Expression operator+(const Expression &l, const Expression &r) { + return add(l, r); +} +inline Expression operator-(const Expression &l, const Expression &r) { + return sub(l, r); +} +inline Expression operator*(const Expression &l, const Expression &r) { + return mul(l, r); +} +inline Expression operator/(const Expression &l, const Expression &r) { + return div(l, r); +} + +} // namespace ops +} // namespace expr +} // namespace vortex diff --git a/lang/cpp/include/vortex/scalar.hpp b/lang/cpp/include/vortex/scalar.hpp new file mode 100644 index 00000000000..f454d918df3 --- /dev/null +++ b/lang/cpp/include/vortex/scalar.hpp @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "vortex/common.hpp" +#include "vortex/dtype.hpp" +#include "vortex/error.hpp" + +#include + +#include +#include +#include +#include + +namespace vortex { + +// A single value with an associated DataType +class Scalar { +public: + Scalar(const Scalar &other); + Scalar(Scalar &&) noexcept = default; + Scalar &operator=(const Scalar &other); + Scalar &operator=(Scalar &&) noexcept = default; + + bool is_null() const; + DataType dtype() const; + +private: + friend struct detail::Access; + explicit Scalar(vx_scalar *owned) : handle_(owned) { + } + vx_scalar *release() && { + return handle_.release(); + } + + struct Deleter { + void operator()(vx_scalar *ptr) const noexcept; + }; + std::unique_ptr handle_; +}; + +namespace detail { +vx_scalar *make_bool(bool value, bool nullable); +vx_scalar *make_primitive(vx_ptype ptype, const void *value, bool nullable); +vx_scalar *make_utf8(std::string_view value, bool nullable); +vx_scalar *make_binary(BinaryView value, bool nullable); +Scalar adopt(vx_scalar *raw); +} // namespace detail + +namespace scalar { +/** + * A scalar of DataType selected by T: bool, primitive, string_view (utf8), + * or a BinaryView (binary). + * + * Bytes are copied for utf8 and binary scalars. + */ +template +Scalar of(T value, bool nullable = false) { + if constexpr (std::is_same_v) { + return detail::adopt(detail::make_bool(value, nullable)); + } else if constexpr (std::is_same_v) { + return detail::adopt(detail::make_utf8(value, nullable)); + } else if constexpr (std::is_same_v) { + return detail::adopt(detail::make_binary(value, nullable)); + } else { + return detail::adopt(detail::make_primitive(detail::to_ptype(), &value, nullable)); + } +} + +template +Scalar decimal(T value, uint8_t precision, int8_t scale, bool nullable = false) { + vx_error *error = nullptr; + vx_scalar *out = nullptr; + if constexpr (std::is_same_v) { + out = vx_scalar_new_decimal_i8(value, precision, scale, nullable, &error); + } else if constexpr (std::is_same_v) { + out = vx_scalar_new_decimal_i16(value, precision, scale, nullable, &error); + } else if constexpr (std::is_same_v) { + out = vx_scalar_new_decimal_i32(value, precision, scale, nullable, &error); + } else if constexpr (std::is_same_v) { + out = vx_scalar_new_decimal_i64(value, precision, scale, nullable, &error); + } else { + static_assert(false, "can't construct decimal of following scale"); + } + detail::throw_on_error(error); + return detail::Access::adopt(out); +} + +// A typed null of (a nullable copy of) a given DataType. +Scalar null(const DataType &dtype); +} // namespace scalar +} // namespace vortex diff --git a/lang/cpp/include/vortex/scan.hpp b/lang/cpp/include/vortex/scan.hpp new file mode 100644 index 00000000000..1b8cc2ef7a4 --- /dev/null +++ b/lang/cpp/include/vortex/scan.hpp @@ -0,0 +1,239 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "vortex/array.hpp" +#include "vortex/common.hpp" +#include "vortex/estimate.hpp" +#include "vortex/expression.hpp" +#include "vortex/session.hpp" + +#include + +#include +#include +#include +#include +#include +#include + +namespace vortex { + +namespace detail { + +// range-for support for Scan and Partition +template (Source::*Next)()> +class PullRange { +public: + class iterator { + public: + using value_type = Item; + using difference_type = std::ptrdiff_t; + + iterator() = default; + iterator(Source *src, std::optional first) : src_(src), cur_(std::move(first)) { + } + + Item &operator*() const { + return *cur_; + } + iterator &operator++() { + cur_ = (src_->*Next)(); + return *this; + } + void operator++(int) { + ++*this; + } + bool operator==(std::default_sentinel_t) const { + return !cur_.has_value(); + } + + private: + Source *src_ = nullptr; + mutable std::optional cur_; + }; + + explicit PullRange(Source &src) : src_(&src) { + } + iterator begin() { + return iterator(src_, (src_->*Next)()); + } + std::default_sentinel_t end() { + return std::default_sentinel; + } + +private: + Source *src_; +}; +} // namespace detail + +// Wrapper around ArrowArrayStream which releases stream in destructor +class ArrowStream { +public: + ArrowStream(const ArrowStream &) = delete; + ArrowStream &operator=(const ArrowStream &) = delete; + ArrowStream(ArrowStream &&other) noexcept; + ArrowStream &operator=(ArrowStream &&other) noexcept; + ~ArrowStream(); + + ArrowArrayStream *raw() { + return &stream_; + } + +private: + friend struct detail::Access; + ArrowStream(Session session, ArrowArrayStream stream) noexcept + : session_(std::move(session)), stream_(stream) { + } + + Session session_; + ArrowArrayStream stream_ {}; +}; + +/** + * An independent unit of scan work. + * + * Partition's methods are thread-unsafe: drive each partition from + * one worker thread. + * Calling methods of a moved-out Partition is UB. + */ +class Partition { +public: + Partition(const Partition &) = delete; + Partition &operator=(const Partition &) = delete; + Partition(Partition &&) noexcept = default; + Partition &operator=(Partition &&) noexcept = default; + + // Estimated row count. Throws if called after next(). + Estimate row_count() const; + + // Return next Array or nullopt when partition is exhausted. + std::optional next(); + + // range-for over Arrays + auto batches() & { + return detail::PullRange(*this); + } + auto batches() && = delete; + + /* + * Consume the partition into an Arrow stream. + * Blocks until partition is drained. + */ + ArrowStream into_arrow_stream() &&; + +private: + friend struct detail::Access; + Partition(vx_partition *owned, Session session) : handle_(owned), session_(std::move(session)) { + } + + struct Deleter { + void operator()(vx_partition *ptr) const noexcept; + }; + std::unique_ptr handle_; + Session session_; +}; + +/* + * Row range [begin; end) to apply over filtering. + * [0; 0) or convenience constant AllRows means "return all rows". + */ +struct RowRange { + uint64_t begin = 0; + uint64_t end = 0; +}; + +// Return all rows +constexpr RowRange AllRows = RowRange {0, 0}; + +struct Selection { + enum class Kind { + Include = VX_SELECTION_INCLUDE_RANGE, + Exclude = VX_SELECTION_EXCLUDE_RANGE, + }; + Kind kind = Kind::Include; + std::vector indices; +}; + +/** + * Scan configuration. Fields are append-only and must be set via designated + * initializers. + * Default fields have reasonable behaviour: default projection returns all + * fields, default filter, row_range, and selection don't filter etc. + * + * Example: + * + * DataSource ds = DataSource::open(session, {"file.vortex"}); + * Scan scan = ds.scan({.limit = 100}); + */ +struct ScanOptions { + std::optional projection; + std::optional filter; + /* + * Row range [begin; end) to apply over filtering. + * [0; 0) or convenience constant AllRows means "return all rows". + */ + std::optional row_range; + // Row-index filter applied after row_range. + std::optional selection; + /* + * Maximum number of rows to return. 0 means no limit. + * You can either pass a limit or a filter but not both. + */ + uint64_t limit = 0; + // If true, return rows in storage order. + bool ordered = false; +}; + +/** + * A single traversal of a DataSource. A scan can be consumed only once. + * + * next_partition() is internally synchronized, give each partition to its own + * worker thread. + * + * Calling methods of a moved-out Scan is UB. + */ +class Scan { +public: + Scan(const Scan &) = delete; + Scan &operator=(const Scan &) = delete; + Scan(Scan &&) noexcept = default; + Scan &operator=(Scan &&) noexcept = default; + + Estimate partition_count() const { + return estimate_; + } + + /** + * Scan's dtype. + * + * Throws if called after next_partition(). + * UB if called in parallel with next_partition(). + */ + DataType dtype() const; + + // Next partition or nullopt when the scan is exhausted. Thread-safe. + std::optional next_partition(); + + // range-for over partitions + auto partitions() & { + return detail::PullRange(*this); + } + auto partitions() && = delete; + +private: + friend struct detail::Access; + Scan(vx_scan *owned, Estimate estimate, Session session) + : handle_(owned), mutex_(std::make_unique()), estimate_(estimate), + session_(std::move(session)) { + } + + struct Deleter { + void operator()(vx_scan *ptr) const noexcept; + }; + std::unique_ptr handle_; + std::unique_ptr mutex_; + Estimate estimate_; + Session session_; +}; +} // namespace vortex diff --git a/lang/cpp/include/vortex/session.hpp b/lang/cpp/include/vortex/session.hpp new file mode 100644 index 00000000000..29522e7fd74 --- /dev/null +++ b/lang/cpp/include/vortex/session.hpp @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "vortex/common.hpp" + +#include + +#include + +namespace vortex { + +/** + * A handle to a Vortex session, registry of encodings and compute kernels. + * Copying shares the underlying session. + */ +class Session { +public: + Session(); + + Session(const Session &other); + Session(Session &&) noexcept = default; + Session &operator=(const Session &other); + Session &operator=(Session &&) noexcept = default; + +private: + friend struct detail::Access; + + struct Deleter { + void operator()(vx_session *ptr) const noexcept; + }; + std::unique_ptr handle_; +}; + +} // namespace vortex diff --git a/lang/cpp/include/vortex/writer.hpp b/lang/cpp/include/vortex/writer.hpp new file mode 100644 index 00000000000..8826423e995 --- /dev/null +++ b/lang/cpp/include/vortex/writer.hpp @@ -0,0 +1,56 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "vortex/array.hpp" +#include "vortex/dtype.hpp" +#include "vortex/session.hpp" + +#include + +#include +#include + +namespace vortex { + +/** + * Writes arrays into a Vortex file. + * + * finish() writes the footer and finalizes the file. + * Not calling finish() leaves file corrupted. + * + * Writer methods are thread-unsafe. + */ +class Writer { +public: + static Writer open(const Session &session, std::string_view path, const DataType &dtype); + + Writer(const Writer &) = delete; + Writer &operator=(const Writer &) = delete; + Writer(Writer &&) noexcept = default; + Writer &operator=(Writer &&) noexcept = default; + + /* + * Append Array to output file. + * Throws if "array"'s DataType doesn't match writer's DataType. + */ + void push(std::span arrays); + void push(const Array &array); + void push(std::initializer_list arrays); + + /* + * Write footer and finalize the file. + * Throws on failure. Writer is closed afterwards and further uses throws. + */ + void finish(); + +private: + explicit Writer(vx_array_sink *sink) : handle_(sink) { + } + + struct Deleter { + void operator()(vx_array_sink *ptr) const noexcept; + }; + std::unique_ptr handle_; +}; +} // namespace vortex diff --git a/lang/cpp/src/array.cpp b/lang/cpp/src/array.cpp new file mode 100644 index 00000000000..b4adf69fe5c --- /dev/null +++ b/lang/cpp/src/array.cpp @@ -0,0 +1,344 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "vortex/array.hpp" +#include "vortex/common.hpp" +#include "vortex/dtype.hpp" +#include "vortex/error.hpp" + +#include + +#include +#include +#include +#include + +namespace vortex { + +using detail::Access; +using detail::throw_on_error; + +Validity::Validity(const Validity &other) + : type_(other.type_), array_(other.array_ != nullptr ? vx_array_clone(other.array_) : nullptr) { +} + +Validity::Validity(ValidityType type) : type_(type), array_(nullptr) { + if (type == ValidityType::FromArray) { + throw VortexException("Validity(ValidityType) called with ValidityType::FromArray", + ErrorCode::InvalidArgument); + } +} + +Validity Validity::from_array(const Array &bools) { + if (!bools.has_dtype(DataTypeVariant::Bool)) { + throw VortexException("Validity array isn't a Bool array", ErrorCode::InvalidArgument); + } + return {ValidityType::FromArray, vx_array_clone(Access::c_ptr(bools))}; +} + +Validity::Validity(Validity &&other) noexcept : type_(other.type_), array_(other.array_) { + other.array_ = nullptr; + other.type_ = ValidityType::NonNullable; +} + +Validity &Validity::operator=(const Validity &other) { + if (this != &other) { + *this = Validity(other); + } + return *this; +} + +Validity &Validity::operator=(Validity &&other) noexcept { + if (this != &other) { + vx_array_free(array_); + type_ = other.type_; + array_ = other.array_; + other.array_ = nullptr; + other.type_ = ValidityType::NonNullable; + } + return *this; +} + +Validity::~Validity() { + vx_array_free(array_); +} + +Array Validity::array() const { + if (type_ != ValidityType::FromArray || array_ == nullptr) { + throw VortexException("validity has no backing array", ErrorCode::InvalidArgument); + } + return Access::adopt(vx_array_clone(array_)); +} + +namespace detail { + +bool ValidityBits::is_null(size_t index) const { + if (all_invalid_) { + return true; + } + if (bits_ == nullptr) { + return false; + } + const size_t bit = bit_offset_ + index; + return (bits_[bit / 8] >> (bit % 8) & 1) == 0; +} + +ValidityBits::ValidityBits(const Session &session, const vx_array *canonical) { + vx_validity raw {}; + vx_error *error = nullptr; + vx_array_get_validity(canonical, &raw, &error); + throw_on_error(error); + + switch (static_cast(raw.type)) { + case ValidityType::NonNullable: + case ValidityType::AllValid: + return; + case ValidityType::AllInvalid: + all_invalid_ = true; + return; + case ValidityType::FromArray: + break; + } + + owner_ = vx_array_canonicalize(Access::c_ptr(session), raw.array, &error); + vx_array_free(raw.array); + throw_on_error(error); + + bits_ = static_cast(vx_array_data_ptr_bool(owner_, &bit_offset_, &error)); + if (error != nullptr) { + vx_array_free(owner_); + } + throw_on_error(error); +} + +ValidityBits::ValidityBits(ValidityBits &&other) noexcept + : owner_(other.owner_), bits_(other.bits_), bit_offset_(other.bit_offset_), + all_invalid_(other.all_invalid_) { + other.owner_ = nullptr; + other.bits_ = nullptr; +} + +ValidityBits &ValidityBits::operator=(ValidityBits &&other) noexcept { + if (this != &other) { + vx_array_free(owner_); + owner_ = other.owner_; + bits_ = other.bits_; + bit_offset_ = other.bit_offset_; + all_invalid_ = other.all_invalid_; + other.owner_ = nullptr; + other.bits_ = nullptr; + } + return *this; +} + +ValidityBits::~ValidityBits() { + vx_array_free(owner_); +} +} // namespace detail + +static const vx_struct_fields *struct_fields_or_throw(const vx_dtype *dtype) { + const vx_struct_fields *fields = vx_dtype_struct_dtype(dtype); + if (fields == nullptr) { + throw VortexException("dtype is not a struct", ErrorCode::MismatchedTypes); + } + return fields; +} + +void Array::Deleter::operator()(const vx_array *ptr) const noexcept { + vx_array_free(ptr); +} + +Array::Array(const Array &other) : handle_(vx_array_clone(other.handle_.get())) { +} + +Array &Array::operator=(const Array &other) { + if (this != &other) { + handle_.reset(vx_array_clone(other.handle_.get())); + } + return *this; +} + +Array Array::null(size_t len) { + return Access::adopt(vx_array_new_null(len)); +} + +Array Array::primitive_raw(vx_ptype ptype, const void *data, size_t len, const Validity &validity) { + std::optional keep_alive; + vx_validity raw {}; + raw.type = static_cast(validity.type()); + if (validity.type() == ValidityType::FromArray) { + keep_alive = validity.array(); + raw.array = Access::c_ptr(*keep_alive); + } + + vx_error *error = nullptr; + const vx_array *out = vx_array_new_primitive(ptype, data, len, &raw, &error); + throw_on_error(error); + return Access::adopt(out); +} + +Array Array::from_arrow(ArrowArray *array, ArrowSchema *schema, bool nullable) { + vx_error *error = nullptr; + const vx_array *out = vx_array_from_arrow(array, schema, nullable, &error); + throw_on_error(error); + return Access::adopt(out); +} + +size_t Array::size() const { + return vx_array_len(handle_.get()); +} + +bool Array::nullable() const { + return vx_array_is_nullable(handle_.get()); +} + +bool Array::has_dtype(DataTypeVariant v) const { + return vx_array_has_dtype(handle_.get(), static_cast(v)); +} + +bool Array::is_primitive(PType p) const { + return vx_array_is_primitive(handle_.get(), static_cast(p)); +} + +DataType Array::dtype() const { + return Access::adopt(vx_array_dtype(handle_.get())); +} + +Validity Array::validity() const { + vx_validity raw {}; + vx_error *error = nullptr; + vx_array_get_validity(handle_.get(), &raw, &error); + throw_on_error(error); + return Access::adopt(static_cast(raw.type), raw.array); +} + +size_t Array::null_count() const { + vx_error *error = nullptr; + const size_t count = vx_array_invalid_count(handle_.get(), &error); + throw_on_error(error); + return count; +} + +Array Array::field(size_t index) const { + vx_error *error = nullptr; + const vx_array *out = vx_array_get_field(handle_.get(), index, &error); + throw_on_error(error); + return Access::adopt(out); +} + +Array Array::field(std::string_view name) const { + const DataType dt = dtype(); + const std::unique_ptr fields( + struct_fields_or_throw(detail::Access::c_ptr(dt)), + &vx_struct_fields_free); + const uint64_t fields_size = vx_struct_fields_nfields(fields.get()); + for (uint64_t i = 0; i < fields_size; ++i) { + const vx_view field = vx_struct_fields_field_name(fields.get(), i); + if (std::string_view {field.ptr, field.len} == name) { + return this->field(i); + } + } + throw VortexException("no field named \"" + std::string(name) + "\"", ErrorCode::InvalidArgument); +} + +Array Array::slice(size_t begin, size_t end) const { + vx_error *error = nullptr; + const vx_array *out = vx_array_slice(handle_.get(), begin, end, &error); + throw_on_error(error); + return Access::adopt(out); +} + +Array Array::apply(const Expression &expr) const { + vx_error *error = nullptr; + const vx_array *out = vx_array_apply(handle_.get(), Access::c_ptr(expr), &error); + throw_on_error(error); + return Access::adopt(out); +} + +Array Array::canonicalize(const Session &session) const { + vx_error *error = nullptr; + const vx_array *out = vx_array_canonicalize(Access::c_ptr(session), handle_.get(), &error); + throw_on_error(error); + return Access::adopt(out); +} + +Array make_struct(std::span fields, const Validity &validity) { + vx_validity raw {}; + raw.type = static_cast(validity.type()); + + std::optional keep_alive; + if (validity.type() == ValidityType::FromArray) { + keep_alive = validity.array(); + raw.array = Access::c_ptr(*keep_alive); + } + + std::unique_ptr handle( + vx_struct_column_builder_new(&raw, fields.size()), + &vx_struct_column_builder_free); + + vx_error *error = nullptr; + for (const auto &[name, array] : fields) { + vx_struct_column_builder_add_field(handle.get(), detail::to_view(name), Access::c_ptr(array), &error); + throw_on_error(error); + } + + const vx_array *out = vx_struct_column_builder_finalize(handle.release(), &error); + throw_on_error(error); + return Access::adopt(out); +} + +Array make_struct(std::initializer_list fields, const Validity &validity) { + return make_struct({fields.begin(), fields.end()}, validity); +} + +PrimitiveView Array::bools(const Session &session) const { + Array canonical = canonicalize(session); + const vx_array *raw = Access::c_ptr(canonical); + if (!vx_array_has_dtype(raw, DTYPE_BOOL)) { + throw VortexException("bools(): array is not a Bool array", ErrorCode::MismatchedTypes); + } + detail::ValidityBits validity(session, raw); + const size_t len = vx_array_len(raw); + return PrimitiveView(std::move(canonical), std::move(validity), len); +} + +StringView Array::strings(const Session &session) const { + Array canonical = canonicalize(session); + const vx_array *raw = Access::c_ptr(canonical); + if (!vx_array_has_dtype(raw, DTYPE_UTF8)) { + throw VortexException("strings(): array is not a Utf8 array", ErrorCode::MismatchedTypes); + } + detail::ValidityBits validity(session, raw); + const size_t len = vx_array_len(raw); + return StringView(std::move(canonical), std::move(validity), len); +} + +BytesView Array::bytes(const Session &session) const { + Array canonical = canonicalize(session); + const vx_array *raw = Access::c_ptr(canonical); + if (!vx_array_has_dtype(raw, DTYPE_BINARY)) { + throw VortexException("bytes(): array is not a Binary array", ErrorCode::MismatchedTypes); + } + detail::ValidityBits validity(session, raw); + const size_t len = vx_array_len(raw); + return BytesView(std::move(canonical), std::move(validity), len); +} + +bool PrimitiveView::value(size_t i) const { + return vx_array_get_bool(Access::c_ptr(canonical_), i); +} + +std::string_view StringView::operator[](size_t i) const { + vx_error *error = nullptr; + const vx_view out = vx_array_utf8_at(Access::c_ptr(canonical_), i, &error); + throw_on_error(error); + return {out.ptr, out.len}; +} + +BinaryView BytesView::operator[](size_t i) const { + vx_error *error = nullptr; + const vx_view out = vx_array_binary_at(Access::c_ptr(canonical_), i, &error); + throw_on_error(error); + return {reinterpret_cast(out.ptr), out.len}; +} +} // namespace vortex diff --git a/lang/cpp/src/data_source.cpp b/lang/cpp/src/data_source.cpp new file mode 100644 index 00000000000..6aa8d3e3107 --- /dev/null +++ b/lang/cpp/src/data_source.cpp @@ -0,0 +1,112 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "vortex/common.hpp" +#include "vortex/data_source.hpp" +#include "vortex/dtype.hpp" +#include "vortex/error.hpp" + +#include + +#include +#include +#include +#include + +namespace vortex { + +using detail::Access; +using detail::throw_on_error; +using detail::to_view; + +void DataSource::Deleter::operator()(const vx_data_source *ptr) const noexcept { + vx_data_source_free(ptr); +} + +DataSource::DataSource(const DataSource &other) + : handle_(vx_data_source_clone(other.handle_.get())), session_(other.session_) { +} + +DataSource &DataSource::operator=(const DataSource &other) { + if (this != &other) { + handle_.reset(vx_data_source_clone(other.handle_.get())); + session_ = other.session_; + } + return *this; +} + +template +static DataSource open_paths(const Session &session, std::span paths) { + std::vector raw; + raw.reserve(paths.size()); + for (const auto &path : paths) { + raw.push_back(to_view(path)); + } + + vx_data_source_options options {}; + options.paths = raw.data(); + options.paths_len = raw.size(); + + vx_error *error = nullptr; + const vx_data_source *ds = vx_data_source_new(Access::c_ptr(session), &options, &error); + throw_on_error(error); + return Access::adopt(ds, session); +} + +DataSource DataSource::open(const Session &session, std::initializer_list paths) { + std::span span {paths.begin(), paths.end()}; + return open_paths(session, span); +} + +DataSource DataSource::open(const Session &session, std::span paths) { + return open_paths(session, paths); +} + +DataSource DataSource::open(const Session &session, std::span paths) { + return open_paths(session, paths); +} + +DataSource DataSource::from_buffer(const Session &session, std::span data) { + vx_error *error = nullptr; + const vx_data_source *ds = + vx_data_source_new_buffer(Access::c_ptr(session), data.data(), data.size(), &error); + throw_on_error(error); + return Access::adopt(ds, session); +} + +Estimate DataSource::row_count() const { + vx_estimate raw {}; + vx_data_source_get_row_count(handle_.get(), &raw); + return Estimate(raw); +} + +DataType DataSource::dtype() const { + return Access::adopt(vx_data_source_dtype(handle_.get())); +} + +Scan DataSource::scan(const ScanOptions &options) const { + vx_scan_options raw {}; + raw.projection = options.projection.has_value() ? Access::c_ptr(*options.projection) : nullptr; + raw.filter = options.filter.has_value() ? Access::c_ptr(*options.filter) : nullptr; + if (options.row_range.has_value()) { + raw.row_range_begin = options.row_range->begin; + raw.row_range_end = options.row_range->end; + } + if (options.selection.has_value()) { + raw.selection.idx = options.selection->indices.data(); + raw.selection.idx_len = options.selection->indices.size(); + raw.selection.include = static_cast(options.selection->kind); + } else { + raw.selection.include = VX_SELECTION_INCLUDE_ALL; + } + raw.limit = options.limit; + raw.ordered = options.ordered; + + vx_estimate estimate {}; + vx_error *error = nullptr; + vx_scan *scan = vx_data_source_scan(handle_.get(), &raw, &estimate, &error); + throw_on_error(error); + return Access::adopt(scan, Estimate(estimate), session_); +} + +} // namespace vortex diff --git a/lang/cpp/src/dtype.cpp b/lang/cpp/src/dtype.cpp new file mode 100644 index 00000000000..808ead6fd50 --- /dev/null +++ b/lang/cpp/src/dtype.cpp @@ -0,0 +1,210 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "vortex/common.hpp" +#include "vortex/dtype.hpp" +#include "vortex/error.hpp" + +#include + +#include +#include +#include +#include +#include + +namespace vortex { + +using detail::Access; +using detail::throw_on_error; +using detail::to_view; +using namespace std::string_literals; + +void DataType::Deleter::operator()(const vx_dtype *ptr) const noexcept { + vx_dtype_free(ptr); +} + +DataType::DataType(const vx_dtype *owned) : handle_(owned) { +} + +DataType::DataType(const DataType &other) : handle_(vx_dtype_clone(other.handle_.get())) { +} + +DataType &DataType::operator=(const DataType &other) { + if (this != &other) { + handle_.reset(vx_dtype_clone(other.handle_.get())); + } + return *this; +} + +DataType DataType::from_arrow(ArrowSchema *schema) { + vx_error *error = nullptr; + const vx_dtype *dtype = vx_dtype_from_arrow_schema(schema, &error); + throw_on_error(error); + return DataType(dtype); +} + +ArrowSchema DataType::to_arrow() const { + ArrowSchema schema {}; + vx_error *error = nullptr; + vx_dtype_to_arrow_schema(handle_.get(), &schema, &error); + throw_on_error(error); + return schema; +} + +DataTypeVariant DataType::variant() const { + return static_cast(vx_dtype_get_variant(handle_.get())); +} + +bool DataType::nullable() const { + return vx_dtype_is_nullable(handle_.get()); +} + +PType DataType::primitive_type() const { + return static_cast(vx_dtype_primitive_ptype(handle_.get())); +} + +uint8_t DataType::decimal_precision() const { + return vx_dtype_decimal_precision(handle_.get()); +} + +int8_t DataType::decimal_scale() const { + return vx_dtype_decimal_scale(handle_.get()); +} + +namespace { +const vx_struct_fields *struct_fields_or_throw(const vx_dtype *dtype) { + const vx_struct_fields *fields = vx_dtype_struct_dtype(dtype); + if (fields == nullptr) { + throw VortexException("dtype is not a struct", ErrorCode::MismatchedTypes); + } + return fields; +} +} // namespace + +std::vector DataType::fields() const { + const std::unique_ptr fields( + struct_fields_or_throw(handle_.get()), + &vx_struct_fields_free); + const uint64_t fields_size = vx_struct_fields_nfields(fields.get()); + std::vector out; + out.reserve(fields_size); + for (uint64_t idx = 0; idx < fields_size; ++idx) { + const vx_view name = vx_struct_fields_field_name(fields.get(), idx); + if (name.ptr == nullptr) { + throw VortexException("error getting field name at index "s + std::to_string(idx), + ErrorCode::Other); + } + const vx_dtype *dtype = vx_struct_fields_field_dtype(fields.get(), idx); + if (dtype == nullptr) { + throw VortexException("error getting dtype at index "s + std::to_string(idx), ErrorCode::Other); + } + out.push_back(StructField {{name.ptr, name.len}, DataType(dtype)}); + } + return out; +} + +DataType DataType::list_element() const { + const vx_dtype *element = vx_dtype_list_element(handle_.get()); + if (element == nullptr) { + throw VortexException("dtype is not a list", ErrorCode::MismatchedTypes); + } + return DataType(element); +} + +DataType DataType::fixed_size_list_element() const { + const vx_dtype *element = vx_dtype_fixed_size_list_element(handle_.get()); + if (element == nullptr) { + throw VortexException("dtype is not a fixed-size list", ErrorCode::MismatchedTypes); + } + return DataType(element); +} + +uint32_t DataType::fixed_size_list_size() const { + return vx_dtype_fixed_size_list_size(handle_.get()); +} + +namespace dtype { + +DataType null() { + return Access::adopt(vx_dtype_new_null()); +} +DataType boolean(bool nullable) { + return Access::adopt(vx_dtype_new_bool(nullable)); +} +DataType primitive(PType ptype, bool nullable) { + return Access::adopt(vx_dtype_new_primitive(static_cast(ptype), nullable)); +} +DataType int8(bool nullable) { + return primitive(PType::I8, nullable); +} +DataType int16(bool nullable) { + return primitive(PType::I16, nullable); +} +DataType int32(bool nullable) { + return primitive(PType::I32, nullable); +} +DataType int64(bool nullable) { + return primitive(PType::I64, nullable); +} +DataType uint8(bool nullable) { + return primitive(PType::U8, nullable); +} +DataType uint16(bool nullable) { + return primitive(PType::U16, nullable); +} +DataType uint32(bool nullable) { + return primitive(PType::U32, nullable); +} +DataType uint64(bool nullable) { + return primitive(PType::U64, nullable); +} +DataType float16(bool nullable) { + return primitive(PType::F16, nullable); +} +DataType float32(bool nullable) { + return primitive(PType::F32, nullable); +} +DataType float64(bool nullable) { + return primitive(PType::F64, nullable); +} +DataType utf8(bool nullable) { + return Access::adopt(vx_dtype_new_utf8(nullable)); +} +DataType binary(bool nullable) { + return Access::adopt(vx_dtype_new_binary(nullable)); +} +DataType decimal(uint8_t precision, int8_t scale, bool nullable) { + return Access::adopt(vx_dtype_new_decimal(precision, scale, nullable)); +} +DataType list(DataType element, bool nullable) { + return Access::adopt(vx_dtype_new_list(Access::release(std::move(element)), nullable)); +} +DataType fixed_size_list(DataType element, uint32_t size, bool nullable) { + return Access::adopt( + vx_dtype_new_fixed_size_list(Access::release(std::move(element)), size, nullable)); +} + +DataType struct_(std::span fields, bool nullable) { + vx_error *error = nullptr; + std::unique_ptr handle( + vx_struct_fields_builder_new(), + vx_struct_fields_builder_free); + + for (const auto &[name, dtype] : fields) { + vx_struct_fields_builder_add_field(handle.get(), + to_view(name), + vx_dtype_clone(Access::c_ptr(dtype)), + &error); + throw_on_error(error); + } + vx_struct_fields *ffi_fields = vx_struct_fields_builder_finalize(handle.release()); + return Access::adopt(vx_dtype_new_struct(ffi_fields, nullable)); +} + +DataType struct_(std::initializer_list fields, bool nullable) { + return struct_({fields.begin(), fields.end()}, nullable); +} + +} // namespace dtype +} // namespace vortex diff --git a/lang/cpp/src/expression.cpp b/lang/cpp/src/expression.cpp new file mode 100644 index 00000000000..d314ae26293 --- /dev/null +++ b/lang/cpp/src/expression.cpp @@ -0,0 +1,159 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "vortex/common.hpp" +#include "vortex/error.hpp" +#include "vortex/expression.hpp" + +#include + +#include +#include +#include +#include + +namespace vortex { + +using detail::Access; +using detail::throw_on_error; +using detail::to_view; + +void Expression::Deleter::operator()(const vx_expression *ptr) const noexcept { + vx_expression_free(ptr); +} + +Expression::Expression(const Expression &other) : handle_(vx_expression_clone(other.handle_.get())) { +} + +Expression &Expression::operator=(const Expression &other) { + if (this != &other) { + handle_.reset(vx_expression_clone(other.handle_.get())); + } + return *this; +} + +Expression Expression::operator[](std::string_view field) const { + vx_expression *out = vx_expression_get_item(to_view(field), handle_.get()); + if (out == nullptr) { + throw VortexException("get_item: field name is not valid UTF-8", ErrorCode::InvalidArgument); + } + return Access::adopt(out); +} + +Expression Expression::is_null() const { + return expr::is_null(*this); +} + +template +static Expression select_impl(std::span names, const vx_expression *expr) { + std::vector raw; + raw.reserve(names.size()); + for (const auto &name : names) { + raw.push_back(to_view(name)); + } + return Access::adopt(vx_expression_select(raw.data(), raw.size(), expr)); +} + +Expression Expression::select(std::span names) const { + return select_impl(names, handle_.get()); +} + +Expression Expression::select(std::span names) const { + return select_impl(names, handle_.get()); +} + +Expression Expression::select(std::initializer_list names) const { + std::span span {names.begin(), names.end()}; + return select_impl(span, handle_.get()); +} + +namespace expr { + +Expression root() { + return Access::adopt(vx_expression_root()); +} + +Expression col(std::string_view name) { + return root()[name]; +} + +Expression lit(const Scalar &value) { + vx_error *error = nullptr; + vx_expression *out = vx_expression_literal(Access::c_ptr(value), &error); + throw_on_error(error); + return Access::adopt(out); +} + +Expression binary_op(BinaryOperator op, const Expression &l, const Expression &r) { + return Access::adopt( + vx_expression_binary(static_cast(op), Access::c_ptr(l), Access::c_ptr(r))); +} + +Expression eq(const Expression &l, const Expression &r) { + return binary_op(BinaryOperator::Eq, l, r); +} +Expression neq(const Expression &l, const Expression &r) { + return binary_op(BinaryOperator::NotEq, l, r); +} +Expression lt(const Expression &l, const Expression &r) { + return binary_op(BinaryOperator::Lt, l, r); +} +Expression lte(const Expression &l, const Expression &r) { + return binary_op(BinaryOperator::Lte, l, r); +} +Expression gt(const Expression &l, const Expression &r) { + return binary_op(BinaryOperator::Gt, l, r); +} +Expression gte(const Expression &l, const Expression &r) { + return binary_op(BinaryOperator::Gte, l, r); +} +Expression add(const Expression &l, const Expression &r) { + return binary_op(BinaryOperator::Add, l, r); +} +Expression sub(const Expression &l, const Expression &r) { + return binary_op(BinaryOperator::Sub, l, r); +} +Expression mul(const Expression &l, const Expression &r) { + return binary_op(BinaryOperator::Mul, l, r); +} +Expression div(const Expression &l, const Expression &r) { + return binary_op(BinaryOperator::Div, l, r); +} + +static Expression combine(vx_expression *(*combiner)(const vx_expression *const *, size_t), + std::span children) { + std::vector raw; + raw.reserve(children.size()); + for (const auto &child : children) { + raw.push_back(Access::c_ptr(child)); + } + vx_expression *out = combiner(raw.data(), raw.size()); + if (out == nullptr) { + throw VortexException("empty expression list", ErrorCode::InvalidArgument); + } + return Access::adopt(out); +} + +Expression and_all(std::span children) { + return combine(vx_expression_and, children); +} + +Expression or_all(std::span children) { + return combine(vx_expression_or, children); +} + +Expression logical_not(const Expression &child) { + return Access::adopt(vx_expression_not(Access::c_ptr(child))); +} + +Expression is_null(const Expression &child) { + return Access::adopt(vx_expression_is_null(Access::c_ptr(child))); +} + +Expression list_contains(const Expression &list, const Expression &value) { + return Access::adopt(vx_expression_list_contains(Access::c_ptr(list), Access::c_ptr(value))); +} + +} // namespace expr + +} // namespace vortex diff --git a/lang/cpp/src/scalar.cpp b/lang/cpp/src/scalar.cpp new file mode 100644 index 00000000000..2bef2699bba --- /dev/null +++ b/lang/cpp/src/scalar.cpp @@ -0,0 +1,111 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "vortex/common.hpp" +#include "vortex/error.hpp" +#include "vortex/scalar.hpp" + +#include + +#if __STDCPP_FLOAT16_T__ == 1 +#include +#endif +#include + +namespace vortex { + +using detail::Access; +using detail::throw_on_error; + +void Scalar::Deleter::operator()(vx_scalar *ptr) const noexcept { + vx_scalar_free(ptr); +} + +Scalar::Scalar(const Scalar &other) : handle_(vx_scalar_clone(other.handle_.get())) { +} + +Scalar &Scalar::operator=(const Scalar &other) { + if (this != &other) { + handle_.reset(vx_scalar_clone(other.handle_.get())); + } + return *this; +} + +bool Scalar::is_null() const { + return vx_scalar_is_null(handle_.get()); +} + +DataType Scalar::dtype() const { + return Access::adopt(vx_scalar_dtype(handle_.get())); +} + +namespace detail { + +Scalar adopt(vx_scalar *raw) { + return Access::adopt(raw); +} + +vx_scalar *make_bool(bool value, bool nullable) { + return vx_scalar_new_bool(value, nullable); +} + +vx_scalar *make_primitive(vx_ptype ptype, const void *value, bool nullable) { + switch (ptype) { + case PTYPE_U8: + return vx_scalar_new_u8(*static_cast(value), nullable); + case PTYPE_U16: + return vx_scalar_new_u16(*static_cast(value), nullable); + case PTYPE_U32: + return vx_scalar_new_u32(*static_cast(value), nullable); + case PTYPE_U64: + return vx_scalar_new_u64(*static_cast(value), nullable); + case PTYPE_I8: + return vx_scalar_new_i8(*static_cast(value), nullable); + case PTYPE_I16: + return vx_scalar_new_i16(*static_cast(value), nullable); + case PTYPE_I32: + return vx_scalar_new_i32(*static_cast(value), nullable); + case PTYPE_I64: + return vx_scalar_new_i64(*static_cast(value), nullable); +#if __STDCPP_FLOAT16_T__ != 1 + case PTYPE_F16: + return vx_scalar_new_f16_bits(static_cast(value)->bits, nullable); +#else + case PTYPE_F16: + return vx_scalar_new_f16_bits(std::bit_cast(*static_cast(value)), + nullable); +#endif + case PTYPE_F32: + return vx_scalar_new_f32(*static_cast(value), nullable); + case PTYPE_F64: + return vx_scalar_new_f64(*static_cast(value), nullable); + } + throw VortexException("unsupported ptype", ErrorCode::InvalidArgument); +} + +vx_scalar *make_utf8(std::string_view value, bool nullable) { + vx_error *error = nullptr; + vx_scalar *out = vx_scalar_new_utf8(to_view(value), nullable, &error); + throw_on_error(error); + return out; +} + +vx_scalar *make_binary(BinaryView value, bool nullable) { + vx_error *error = nullptr; + vx_scalar *out = + vx_scalar_new_binary(reinterpret_cast(value.data()), value.size(), nullable, &error); + throw_on_error(error); + return out; +} + +} // namespace detail + +namespace scalar { +Scalar null(const DataType &dtype) { + vx_error *error = nullptr; + vx_scalar *out = vx_scalar_new_null(Access::c_ptr(dtype), &error); + throw_on_error(error); + return Access::adopt(out); +} +} // namespace scalar +} // namespace vortex diff --git a/lang/cpp/src/scan.cpp b/lang/cpp/src/scan.cpp new file mode 100644 index 00000000000..8c7967b010d --- /dev/null +++ b/lang/cpp/src/scan.cpp @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "vortex/array.hpp" +#include "vortex/common.hpp" +#include "vortex/dtype.hpp" +#include "vortex/error.hpp" +#include "vortex/scan.hpp" + +#include + +#include +#include + +namespace vortex { + +using detail::Access; +using detail::throw_on_error; + +ArrowStream::ArrowStream(ArrowStream &&other) noexcept + : session_(std::move(other.session_)), stream_(other.stream_) { + other.stream_ = ArrowArrayStream {}; +} + +ArrowStream &ArrowStream::operator=(ArrowStream &&other) noexcept { + if (this != &other) { + if (stream_.release != nullptr) { + stream_.release(&stream_); + } + session_ = std::move(other.session_); + stream_ = other.stream_; + other.stream_ = ArrowArrayStream {}; + } + return *this; +} + +ArrowStream::~ArrowStream() { + if (stream_.release != nullptr) { + stream_.release(&stream_); + } +} + +void Partition::Deleter::operator()(vx_partition *ptr) const noexcept { + vx_partition_free(ptr); +} + +Estimate Partition::row_count() const { + vx_estimate raw {}; + vx_error *error = nullptr; + vx_partition_row_count(handle_.get(), &raw, &error); + throw_on_error(error); + return Estimate(raw); +} + +std::optional Partition::next() { + vx_error *error = nullptr; + const vx_array *array = vx_partition_next(handle_.get(), &error); + throw_on_error(error); + if (array == nullptr) { + return std::nullopt; + } + return Access::adopt(array); +} + +ArrowStream Partition::into_arrow_stream() && { + ArrowArrayStream stream {}; + vx_error *error = nullptr; + // Consumes handle even on error + vx_partition_scan_arrow(Access::c_ptr(session_), handle_.release(), &stream, &error); + throw_on_error(error); + return Access::adopt(std::move(session_), stream); +} + +void Scan::Deleter::operator()(vx_scan *ptr) const noexcept { + vx_scan_free(ptr); +} + +DataType Scan::dtype() const { + vx_error *error = nullptr; + const vx_dtype *out = vx_scan_dtype(handle_.get(), &error); + throw_on_error(error); + return Access::adopt(out); +} + +std::optional Scan::next_partition() { + vx_partition *partition = nullptr; + { + const std::lock_guard guard(*mutex_); + vx_error *error = nullptr; + partition = vx_scan_next_partition(handle_.get(), &error); + throw_on_error(error); + } + if (partition == nullptr) { + return std::nullopt; + } + return Access::adopt(partition, session_); +} +} // namespace vortex diff --git a/lang/cpp/src/session.cpp b/lang/cpp/src/session.cpp new file mode 100644 index 00000000000..fb2753476bd --- /dev/null +++ b/lang/cpp/src/session.cpp @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "vortex/session.hpp" + +#include + +namespace vortex { + +void Session::Deleter::operator()(vx_session *ptr) const noexcept { + vx_session_free(ptr); +} + +Session::Session() : handle_(vx_session_new()) { +} + +Session::Session(const Session &other) : handle_(vx_session_clone(other.handle_.get())) { +} + +Session &Session::operator=(const Session &other) { + if (this != &other) { + handle_.reset(vx_session_clone(other.handle_.get())); + } + return *this; +} + +} // namespace vortex diff --git a/lang/cpp/src/writer.cpp b/lang/cpp/src/writer.cpp new file mode 100644 index 00000000000..57035171a14 --- /dev/null +++ b/lang/cpp/src/writer.cpp @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "vortex/array.hpp" +#include "vortex/common.hpp" +#include "vortex/dtype.hpp" +#include "vortex/error.hpp" +#include "vortex/writer.hpp" +#include "vortex/session.hpp" + +#include +#include + +#include + +namespace vortex { + +using detail::Access; +using detail::throw_on_error; +using detail::to_view; + +void Writer::Deleter::operator()(vx_array_sink *ptr) const noexcept { + vx_array_sink_abort(ptr); +} + +Writer Writer::open(const Session &session, std::string_view path, const DataType &dtype) { + vx_error *error = nullptr; + vx_array_sink *sink = + vx_array_sink_open_file(Access::c_ptr(session), to_view(path), Access::c_ptr(dtype), &error); + throw_on_error(error); + return Writer(sink); +} + +void Writer::push(std::span arrays) { + if (handle_ == nullptr) { + throw VortexException("null handle_", ErrorCode::InvalidArgument); + } + vx_error *error = nullptr; + for (const Array &array : arrays) { + vx_array_sink_push(handle_.get(), Access::c_ptr(array), &error); + throw_on_error(error); + } +} + +void Writer::push(const Array &array) { + push(std::span {&array, 1}); +} + +void Writer::push(std::initializer_list arrays) { + push({arrays.begin(), arrays.end()}); +} + +void Writer::finish() { + if (handle_ == nullptr) { + throw VortexException("finish() called twice", ErrorCode::InvalidArgument); + } + vx_error *error = nullptr; + vx_array_sink_close(handle_.release(), &error); + throw_on_error(error); +} +} // namespace vortex diff --git a/lang/cpp/tests/CMakeLists.txt b/lang/cpp/tests/CMakeLists.txt new file mode 100644 index 00000000000..16a6e0290b9 --- /dev/null +++ b/lang/cpp/tests/CMakeLists.txt @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +# need this for std::float16_t +set(CMAKE_CXX_STANDARD 23) + +FetchContent_Declare( + magic_enum + GIT_REPOSITORY https://github.com/Neargye/magic_enum.git + GIT_TAG v0.9.7 +) +FetchContent_MakeAvailable(magic_enum) + +file(GLOB TEST_FILES CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp") + +add_executable(vortex_cxx_test ${TEST_FILES}) +target_include_directories(vortex_cxx_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) +target_link_libraries(vortex_cxx_test PRIVATE + vortex_cxx_shared + Catch2::Catch2WithMain + magic_enum::magic_enum + nanoarrow_shared) + +catch_discover_tests(vortex_cxx_test) diff --git a/lang/cpp/tests/array.cpp b/lang/cpp/tests/array.cpp new file mode 100644 index 00000000000..998cb945833 --- /dev/null +++ b/lang/cpp/tests/array.cpp @@ -0,0 +1,184 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include +#include +#include +#include +#include + +using namespace vortex; +using namespace vortex::expr::ops; + +namespace { +using enum vortex::PType; +using enum ValidityType; + +TEST_CASE("Null array", "[array]") { + Array a = Array::null(1999); + REQUIRE(a.size() == 1999); + REQUIRE(a.nullable()); + REQUIRE(a.has_dtype(DataTypeVariant::Null)); +} + +TEST_CASE("Empty array", "[array]") { + Session session; + + auto empty = Array::primitive({}); + REQUIRE(empty.size() == 0); + REQUIRE(empty.is_primitive(I32)); + + REQUIRE(empty.null_count() == 0); + + auto view = empty.values(session); + REQUIRE(view.size() == 0); + auto values = view.values(); + REQUIRE(values.empty()); +} + +void test_primitive_array(Array array, const int32_t *begin) { + Session session; + REQUIRE(array.size() == 3); + REQUIRE(array.is_primitive(I32)); + REQUIRE_FALSE(array.nullable()); + REQUIRE(array.null_count() == 0); + + auto view = array.values(session); + REQUIRE(view.size() == 3); + REQUIRE(std::equal(view.values().begin(), view.values().end(), begin)); + REQUIRE_FALSE(view.is_null(1)); +} + +TEST_CASE("Primitive array", "[array]") { + int32_t c_array[3] = {10, 20, 30}; + test_primitive_array(Array::primitive(c_array), c_array); + + const int32_t const_c_array[3] = {10, 20, 30}; + test_primitive_array(Array::primitive(const_c_array), const_c_array); + + const std::array cpp_array = {10, 20, 30}; + test_primitive_array(Array::primitive(cpp_array), cpp_array.begin()); + + std::vector cpp_vector = {10, 20, 30}; + test_primitive_array(Array::primitive(cpp_vector), cpp_vector.data()); +} + +TEST_CASE("values with wrong type", "[array]") { + Session session; + std::vector data = {1}; + Array a = Array::primitive(data); + REQUIRE_THROWS_AS(a.values(session), VortexException); +} + +TEST_CASE("Validity from a boolean mask", "[array]") { + Session session; + std::vector data = {10, 20, 30}; + std::vector mask_bytes = {1, 0, 1}; + + Array mask_u8 = Array::primitive(std::span(mask_bytes)); + Array mask = mask_u8.apply(expr::root() == expr::lit(1)); + + Array a = Array::primitive(std::span(data), Validity::from_array(mask)); + REQUIRE(a.nullable()); + REQUIRE(a.null_count() == 1); + + auto view = a.values(session); + REQUIRE_FALSE(view.is_null(0)); + REQUIRE(view.is_null(1)); + REQUIRE_FALSE(view.is_null(2)); + REQUIRE(view.values()[0] == 10); + REQUIRE(view.values()[2] == 30); + + Validity validity = a.validity(); + REQUIRE(validity.type() == FromArray); + REQUIRE(validity.array().size() == 3); +} + +TEST_CASE("Invalid validity", "[array]") { + std::vector invalid_mask = {1, 2, 3}; + Array mask = Array::primitive(invalid_mask); + REQUIRE_THROWS_AS(Validity::from_array(mask), VortexException); +} + +TEST_CASE("AllInvalid", "[array]") { + Session session; + std::vector data = {1, 2}; + Array a = Array::primitive(std::span(data), AllInvalid); + REQUIRE(a.null_count() == 2); + auto view = a.values(session); + REQUIRE(view.is_null(0)); + REQUIRE(view.is_null(1)); +} + +TEST_CASE("make_struct and fields", "[array]") { + Array empty = make_struct({}); + REQUIRE(empty.size() == 0); + REQUIRE(empty.has_dtype(DataTypeVariant::Struct)); + REQUIRE(empty.dtype().fields().size() == 0); + + std::vector ages = {10, 20, 30}; + std::vector heights = {150, 160, 170}; + + Array s = make_struct({ + {"age", Array::primitive(ages)}, + {"height", Array::primitive(heights, AllValid)}, + }); + + REQUIRE(s.size() == 3); + REQUIRE(s.has_dtype(DataTypeVariant::Struct)); + REQUIRE(s.dtype().fields().size() == 2); + + Array by_index = s.field(0); + REQUIRE(by_index.is_primitive(U8)); + + Session session; + Array by_name = s.field("height"); + REQUIRE(by_name.is_primitive(U16)); + auto view = by_name.values(session); + REQUIRE(view.values()[2] == 170); + + REQUIRE_THROWS_AS(s.field(2), VortexException); + REQUIRE_THROWS_AS(s.field("nope"), VortexException); + + std::vector fields_vec; + fields_vec.emplace_back("age", Array::primitive(ages)); + Array other = make_struct(fields_vec); + REQUIRE(other.size() == 3); + REQUIRE(other.has_dtype(DataTypeVariant::Struct)); +} + +TEST_CASE("Mismatched field length", "[array]") { + std::vector a = {1, 2}; + std::vector b = {1, 2, 3}; + REQUIRE_THROWS_AS(make_struct({ + {"a", Array::primitive(a)}, + {"b", Array::primitive(b)}, + }), + VortexException); +} + +TEST_CASE("Slice", "[array]") { + Session session; + std::vector data = {0, 1, 2, 3, 4, 5}; + Array a = Array::primitive(data); + Array sliced = a.slice(2, 5); + REQUIRE(sliced.size() == 3); + auto view = sliced.values(session); + REQUIRE(view.values()[0] == 2); + REQUIRE(view.values()[2] == 4); + + REQUIRE_THROWS_AS(a.slice(2, 100), VortexException); +} + +TEST_CASE("Error with a code", "[array]") { + std::vector data = {0}; + Array a = Array::primitive(data); + try { + (void)a.slice(2, 100); + FAIL("expected exception"); + } catch (const VortexException &e) { + REQUIRE_FALSE(std::string(e.what()).empty()); + (void)e.code(); + } +} +} // namespace diff --git a/lang/cpp/tests/arrow.cpp b/lang/cpp/tests/arrow.cpp new file mode 100644 index 00000000000..8675d816067 --- /dev/null +++ b/lang/cpp/tests/arrow.cpp @@ -0,0 +1,115 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include + +#include +#include + +typedef struct ArrowSchema FFI_ArrowSchema; +typedef struct ArrowArray FFI_ArrowArray; +typedef struct ArrowArrayStream FFI_ArrowArrayStream; +#define USE_OWN_ARROW 1 + +#include + +#include "common.hpp" + +using namespace vortex; +using vortex_test::sample_dtype; +using vortex_test::TempPath; +using vortex_test::write_sample; + +namespace { +using enum vortex::PType; + +TEST_CASE("dtype to ArrowSchema", "[arrow]") { + DataType d = sample_dtype(); + ArrowSchema schema = d.to_arrow(); + + nanoarrow::UniqueSchema unique_schema; + ArrowSchemaMove(&schema, unique_schema.get()); + REQUIRE(unique_schema->format != nullptr); + REQUIRE(unique_schema->n_children == 2); +} + +TEST_CASE("dtype from ArrowSchema", "[arrow]") { + nanoarrow::UniqueSchema schema; + REQUIRE(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRUCT) == NANOARROW_OK); + REQUIRE(ArrowSchemaAllocateChildren(schema.get(), 1) == NANOARROW_OK); + REQUIRE(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT64) == NANOARROW_OK); + REQUIRE(ArrowSchemaSetName(schema->children[0], "n") == NANOARROW_OK); + + ArrowSchema raw = {}; + ArrowSchemaMove(schema.get(), &raw); + DataType d = DataType::from_arrow(&raw); + REQUIRE(d.variant() == DataTypeVariant::Struct); + const std::vector fields = d.fields(); + REQUIRE(fields.size() == 1); + REQUIRE(fields[0].name == "n"); + REQUIRE(fields[0].dtype.primitive_type() == I64); +} + +TEST_CASE("Import Arrow array as Vortex array", "[arrow]") { + Session session; + nanoarrow::UniqueSchema schema; + REQUIRE(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRUCT) == NANOARROW_OK); + REQUIRE(ArrowSchemaAllocateChildren(schema.get(), 1) == NANOARROW_OK); + REQUIRE(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32) == NANOARROW_OK); + REQUIRE(ArrowSchemaSetName(schema->children[0], "a") == NANOARROW_OK); + + nanoarrow::UniqueArray arr; + REQUIRE(ArrowArrayInitFromSchema(arr.get(), schema.get(), nullptr) == NANOARROW_OK); + REQUIRE(ArrowArrayStartAppending(arr.get()) == NANOARROW_OK); + for (int i : {10, 20, 30}) { + REQUIRE(ArrowArrayAppendInt(arr->children[0], i) == NANOARROW_OK); + REQUIRE(ArrowArrayFinishElement(arr.get()) == NANOARROW_OK); + } + REQUIRE(ArrowArrayFinishBuildingDefault(arr.get(), nullptr) == NANOARROW_OK); + + ArrowArray raw_arr = {}; + ArrowSchema raw_schema = {}; + ArrowArrayMove(arr.get(), &raw_arr); + ArrowSchemaMove(schema.get(), &raw_schema); + + Array vx = Array::from_arrow(&raw_arr, &raw_schema, false); + REQUIRE(vx.size() == 3); + REQUIRE(vx.has_dtype(DataTypeVariant::Struct)); + + Array a = vx.field(0); + REQUIRE(a.is_primitive(I32)); + auto view = a.values(session); + REQUIRE(view.values()[0] == 10); + REQUIRE(view.values()[2] == 30); +} + +TEST_CASE("Scan partition to ArrowArrayStream", "[arrow]") { + Session session; + TempPath path = write_sample(session); + DataSource ds = DataSource::open(session, {path.string()}); + + Scan scan = ds.scan(); + auto partition = scan.next_partition(); + REQUIRE(partition.has_value()); + + ArrowStream vx_stream = std::move(partition.value()).into_arrow_stream(); + nanoarrow::UniqueArrayStream owned; + ArrowArrayStreamMove(vx_stream.raw(), owned.get()); + + nanoarrow::UniqueSchema schema; + ArrowError err {}; + REQUIRE(ArrowArrayStreamGetSchema(owned.get(), schema.get(), &err) == NANOARROW_OK); + REQUIRE(schema->n_children == 2); + + size_t rows = 0; + while (true) { + nanoarrow::UniqueArray chunk; + int rc = owned->get_next(owned.get(), chunk.get()); + REQUIRE(rc == NANOARROW_OK); + if (chunk->release == nullptr) { + break; + } + rows += chunk->length; + } + REQUIRE(rows == vortex_test::SAMPLE_ROWS); +} +} // namespace diff --git a/lang/cpp/tests/common.hpp b/lang/cpp/tests/common.hpp new file mode 100644 index 00000000000..b1b3d982ffa --- /dev/null +++ b/lang/cpp/tests/common.hpp @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +#include "vortex/array.hpp" +#include +#include +#include +#include + +#include +#include + +namespace vortex_test { + +namespace fs = std::filesystem; +using namespace vortex; + +class TempPath { +public: + TempPath() = default; + explicit TempPath(fs::path p) : path_(std::move(p)) { + } + + TempPath(const TempPath &) = delete; + TempPath &operator=(const TempPath &) = delete; + + TempPath(TempPath &&other) noexcept : path_(std::move(other.path_)) { + other.path_.clear(); + } + TempPath &operator=(TempPath &&other) noexcept { + if (this != &other) { + reset(); + path_ = std::move(other.path_); + other.path_.clear(); + } + return *this; + } + + ~TempPath() { + reset(); + } + + const fs::path &path() const noexcept { + return path_; + } + std::string string() const { + return path_.string(); + } + + static TempPath unique() { + auto dir = fs::temp_directory_path() / "vortex_cxx_test"; + fs::create_directories(dir); + std::string name = std::to_string(std::random_device {}()) + ".vortex"; + return TempPath {dir / name}; + } + +private: + void reset() noexcept { + if (!path_.empty()) { + std::error_code ec; + fs::remove(path_, ec); + } + } + + fs::path path_; +}; + +inline DataType sample_dtype() { + return dtype::struct_({ + {"age", dtype::uint8()}, + {"height", dtype::uint16(dtype::Nullable)}, + }); +} + +constexpr size_t SAMPLE_ROWS = 100; + +inline std::vector sample_ages() { + std::vector buf(SAMPLE_ROWS); + for (size_t i = 0; i < SAMPLE_ROWS; ++i) { + buf[i] = static_cast(i); + } + return buf; +} + +inline std::vector sample_heights() { + std::vector buf(SAMPLE_ROWS); + for (size_t i = 0; i < SAMPLE_ROWS; ++i) { + buf[i] = static_cast((i + 1) % 200); + } + return buf; +} + +inline Array sample_array() { + auto ages = sample_ages(); + auto heights = sample_heights(); + return make_struct({ + {"age", Array::primitive(ages)}, + {"height", Array::primitive(heights, ValidityType::AllValid)}, + }); +} + +inline TempPath write_sample(const Session &session) { + TempPath path = TempPath::unique(); + Writer writer = Writer::open(session, path.string(), sample_dtype()); + writer.push(sample_array()); + writer.finish(); + return path; +} +} // namespace vortex_test diff --git a/lang/cpp/tests/dtype.cpp b/lang/cpp/tests/dtype.cpp new file mode 100644 index 00000000000..1b5e3573092 --- /dev/null +++ b/lang/cpp/tests/dtype.cpp @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include "vortex/dtype.hpp" +#include +#include + +using namespace vortex; + +namespace { +using enum vortex::PType; + +TEST_CASE("Null dtype", "[dtype]") { + auto d = dtype::null(); + REQUIRE(d.variant() == DataTypeVariant::Null); + REQUIRE(d.nullable()); +} + +TEST_CASE("Decimal dtype", "[dtype]") { + auto d = dtype::decimal(5, 2, false); + REQUIRE(d.variant() == DataTypeVariant::Decimal); + REQUIRE(d.decimal_precision() == 5); + REQUIRE(d.decimal_scale() == 2); + REQUIRE_FALSE(d.nullable()); + + REQUIRE_THROWS_AS(d.fields(), VortexException); + REQUIRE_THROWS_AS(d.list_element(), VortexException); +} + +TEST_CASE("copy dtype", "[dtype]") { + auto d = dtype::int32(true); + DataType d2 = d; + REQUIRE(d2.variant() == DataTypeVariant::Primitive); + REQUIRE(d2.primitive_type() == I32); + REQUIRE(d2.nullable()); + REQUIRE(d.variant() == DataTypeVariant::Primitive); +} + +TEST_CASE("list dtype", "[dtype]") { + auto d = dtype::list(dtype::float64(), dtype::Nullable); + REQUIRE(d.variant() == DataTypeVariant::List); + REQUIRE(d.nullable()); + REQUIRE(d.list_element().primitive_type() == F64); + + auto fsl = dtype::fixed_size_list(dtype::int16(), 4); + REQUIRE(fsl.variant() == DataTypeVariant::FixedSizeList); + REQUIRE(fsl.fixed_size_list_size() == 4); + REQUIRE(fsl.fixed_size_list_element().primitive_type() == I16); +} + +TEST_CASE("Struct DataType", "[dtype]") { + DataType d = dtype::struct_({ + {"col1", dtype::uint8()}, + {"col2", dtype::binary(dtype::Nullable)}, + }); + + REQUIRE(d.variant() == DataTypeVariant::Struct); + REQUIRE_FALSE(d.nullable()); + const std::vector fields = d.fields(); + REQUIRE(fields.size() == 2); + REQUIRE(fields[0].name == "col1"); + REQUIRE(fields[1].name == "col2"); + REQUIRE(fields[0].dtype.primitive_type() == U8); + REQUIRE(fields[1].dtype.variant() == DataTypeVariant::Binary); + REQUIRE(fields[1].dtype.nullable()); + + std::vector fields_vec; + fields_vec.emplace_back("col1", dtype::uint8()); + fields_vec.emplace_back("col2", dtype::utf8()); + d = dtype::struct_(fields_vec, dtype::Nullable); + + REQUIRE(d.variant() == DataTypeVariant::Struct); + REQUIRE(d.nullable()); + const std::vector built = d.fields(); + REQUIRE(built.size() == 2); + REQUIRE(built[1].name == "col2"); + + fields_vec = {}; + fields_vec.emplace_back("\xFF\xFE", dtype::uint8()); + REQUIRE_THROWS_AS(dtype::struct_(fields_vec), VortexException); +} +} // namespace diff --git a/lang/cpp/tests/enum_sizes.cpp b/lang/cpp/tests/enum_sizes.cpp new file mode 100644 index 00000000000..a100fa3b583 --- /dev/null +++ b/lang/cpp/tests/enum_sizes.cpp @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include + +#include +#include + +using magic_enum::enum_count; + +static_assert(enum_count() == enum_count()); +static_assert(enum_count() == enum_count()); +static_assert(enum_count() == enum_count()); +static_assert(enum_count() == enum_count()); +static_assert(enum_count() == enum_count()); +static_assert(enum_count() == enum_count()); diff --git a/lang/cpp/tests/expression.cpp b/lang/cpp/tests/expression.cpp new file mode 100644 index 00000000000..1a55364df02 --- /dev/null +++ b/lang/cpp/tests/expression.cpp @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include +#include + +#include +#include +#include + +using namespace vortex; + +namespace { +using enum vortex::PType; + +TEST_CASE("Expression copy", "[expr]") { + Expression clone = expr::root(); + { + Expression orig = expr::col("age"); + clone = orig; + } + Expression another = clone; + (void)another; + SUCCEED(); +} + +TEST_CASE("Apply root()", "[expr]") { + Session session; + std::vector data = {10, 20, 30, 40, 50}; + Array array = Array::primitive(data); + + Array applied = array.apply(expr::root()); + + REQUIRE(applied.size() == data.size()); + REQUIRE(applied.is_primitive(I32)); + auto values = applied.values(session); + for (size_t i = 0; i < data.size(); ++i) { + REQUIRE(values.values()[i] == data[i]); + } +} + +TEST_CASE("Apply projection", "[expr]") { + Session session; + std::vector ages = {10, 20, 30}; + std::vector heights = {150, 160, 170}; + + Array struct_arr = make_struct({ + {"age", Array::primitive(ages)}, + {"height", Array::primitive(heights)}, + }); + + Array projected = struct_arr.apply(expr::col("age")); + REQUIRE(projected.size() == ages.size()); + REQUIRE(projected.is_primitive(U8)); + auto values = projected.values(session); + for (size_t i = 0; i < ages.size(); ++i) { + REQUIRE(values.values()[i] == ages[i]); + } +} + +TEST_CASE("Apply arithmetic", "[expr]") { + Session session; + std::vector data = {1, 2, 3, 4, 5}; + Array array = Array::primitive(data); + + Expression e = expr::add(expr::root(), expr::lit(10)); + Array applied = array.apply(e); + + REQUIRE(applied.is_primitive(I32)); + auto values = applied.values(session); + for (size_t i = 0; i < data.size(); ++i) { + REQUIRE(values.values()[i] == data[i] + 10); + } +} + +TEST_CASE("Operator overloading", "[expr]") { + using namespace vortex::expr::ops; + Session session; + std::vector data = {1, 2, 2, 7}; + Array array = Array::primitive(data); + + Array applied = array.apply(expr::root() == expr::lit(2)); + auto bits = applied.bools(session); + REQUIRE(bits.size() == data.size()); + REQUIRE_FALSE(bits.value(0)); + REQUIRE(bits.value(1)); + REQUIRE(bits.value(2)); + REQUIRE_FALSE(bits.value(3)); +} + +TEST_CASE("Apply error", "[expr]") { + std::vector data = {1, 2, 3}; + Array array = Array::primitive(data); + + Expression bad = expr::add(expr::root(), expr::lit(1)); + REQUIRE_THROWS_AS(array.apply(bad), VortexException); +} + +TEST_CASE("Empty conjunction", "[expr]") { + REQUIRE_THROWS_AS(expr::and_all({}), VortexException); + REQUIRE_THROWS_AS(expr::or_all({}), VortexException); +} +} // namespace diff --git a/lang/cpp/tests/float16_t.cpp b/lang/cpp/tests/float16_t.cpp new file mode 100644 index 00000000000..7a08a92a05b --- /dev/null +++ b/lang/cpp/tests/float16_t.cpp @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "vortex/dtype.hpp" +#include +#include + +using namespace vortex; + +namespace { + +static_assert(float(1.0f16) == 1.0f); + +TEST_CASE("float16_t to_float", "[float]") { + REQUIRE(float(1.0f16) == 1.0F); + REQUIRE(float(-2.0f16) == -2.0F); + REQUIRE(float(0.0f) == 0.0F); + REQUIRE(float(-0.0f) == -0.0F); +} + +#if __STDCPP_FLOAT16_T__ == 1 +TEST_CASE("F16 scalar", "[scalar]") { + std::float16_t float16t = 1.0f16; + + Scalar scalar = scalar::of(float16t); + REQUIRE(scalar.dtype().variant() == DataTypeVariant::Primitive); + REQUIRE(scalar.dtype().primitive_type() == vortex::PType::F16); + + _Float16 float16t_alias = 1.0f16; + scalar = scalar::of(float16t_alias); + REQUIRE(scalar.dtype().variant() == DataTypeVariant::Primitive); + REQUIRE(scalar.dtype().primitive_type() == vortex::PType::F16); +} +#endif +} // namespace diff --git a/lang/cpp/tests/float16_t_compat.cpp b/lang/cpp/tests/float16_t_compat.cpp new file mode 100644 index 00000000000..94acb0df071 --- /dev/null +++ b/lang/cpp/tests/float16_t_compat.cpp @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include +#include +// This is UB but no other way to test compatibility macro, fine +// as we're running tests +#undef __STDCPP_FLOAT16_T__ +#include +#include + +using namespace vortex; + +namespace { + +constexpr float16_t one {0x3C00}; +static_assert(static_cast(one) == 1.0f); + +TEST_CASE("float16_t to_float (compatibility)", "[float]") { + REQUIRE(float(float16_t {0x3C00}) == 1.0F); + REQUIRE(float(float16_t {0xC000}) == -2.0F); + REQUIRE(float(float16_t {0}) == 0.0F); + REQUIRE(float(float16_t {0x8000}) == -0.0F); + REQUIRE(float(float16_t {0x7C00}) == std::numeric_limits::infinity()); + REQUIRE(float(float16_t {0xFC00}) == -std::numeric_limits::infinity()); + REQUIRE(std::fpclassify(float(float16_t {0x7E01})) == FP_NAN); + // Denormalized float16_t always gets to normal floats on conversion to float + REQUIRE(std::fpclassify(float(float16_t {0x0001})) == FP_NORMAL); + REQUIRE(std::fpclassify(float(float16_t {0x83FF})) == FP_NORMAL); +} + +TEST_CASE("F16 scalar (compatibility)", "[scalar]") { + const float16_t float16t {0x3C00}; + Scalar scalar = scalar::of(float16t); + REQUIRE(scalar.dtype().variant() == DataTypeVariant::Primitive); + REQUIRE(scalar.dtype().primitive_type() == vortex::PType::F16); +} +} // namespace diff --git a/lang/cpp/tests/scalar.cpp b/lang/cpp/tests/scalar.cpp new file mode 100644 index 00000000000..27ff73f656a --- /dev/null +++ b/lang/cpp/tests/scalar.cpp @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include "vortex/error.hpp" +#include +#include +#include + +using namespace vortex; +using namespace std::string_view_literals; + +namespace { +using enum vortex::PType; + +TEST_CASE("Boolean scalar", "[scalar]") { + Scalar s = scalar::of(true); + REQUIRE_FALSE(s.is_null()); + REQUIRE(s.dtype().variant() == DataTypeVariant::Bool); +} + +TEST_CASE("Integer scalars", "[scalar]") { + REQUIRE(scalar::of(42).dtype().primitive_type() == U8); + REQUIRE(scalar::of(42).dtype().primitive_type() == U16); + REQUIRE(scalar::of(42).dtype().primitive_type() == U32); + REQUIRE(scalar::of(42).dtype().primitive_type() == U64); + REQUIRE(scalar::of(-1).dtype().primitive_type() == I8); + REQUIRE(scalar::of(-1).dtype().primitive_type() == I16); + REQUIRE(scalar::of(-1).dtype().primitive_type() == I32); + REQUIRE(scalar::of(-1).dtype().primitive_type() == I64); +} + +TEST_CASE("Float scalars", "[scalar]") { + REQUIRE(scalar::of(1.5F).dtype().primitive_type() == F32); + REQUIRE(scalar::of(1.5).dtype().primitive_type() == F64); + REQUIRE(scalar::of(float16_t {0x3C00}).dtype().primitive_type() == F16); +} + +TEST_CASE("Nullable scalars", "[scalar]") { + Scalar s = scalar::of(0, true); + REQUIRE(s.dtype().nullable()); + REQUIRE_FALSE(s.is_null()); +} + +TEST_CASE("Null scalar", "[scalar]") { + Scalar s = scalar::null(dtype::int32(true)); + REQUIRE(s.is_null()); + REQUIRE(s.dtype().variant() == DataTypeVariant::Primitive); +} + +TEST_CASE("UTF-8 scalar", "[scalar]") { + Scalar s = scalar::of("hello"sv); + REQUIRE_FALSE(s.is_null()); + REQUIRE(s.dtype().variant() == DataTypeVariant::Utf8); + + REQUIRE_FALSE(scalar::of(""sv).is_null()); + s = scalar::of("Широкая строка"sv); + REQUIRE(s.dtype().variant() == DataTypeVariant::Utf8); + + REQUIRE_THROWS_AS(scalar::of("\xFF\xFE"sv), VortexException); +} + +TEST_CASE("Binary scalar", "[scalar]") { + const std::byte bytes[] = {std::byte {1}, std::byte {2}, std::byte {0}, std::byte {4}}; + Scalar s = scalar::of(std::span {bytes}); + REQUIRE_FALSE(s.is_null()); + REQUIRE(s.dtype().variant() == DataTypeVariant::Binary); +} + +TEST_CASE("Decimal scalars", "[scalar]") { + Scalar d8 = scalar::decimal(56, 5, 2); + REQUIRE(d8.dtype().variant() == DataTypeVariant::Decimal); + REQUIRE(d8.dtype().decimal_precision() == 5); + REQUIRE(d8.dtype().decimal_scale() == 2); + + Scalar d16 = scalar::decimal(int16_t(1234), 5, 2); + REQUIRE(d16.dtype().variant() == DataTypeVariant::Decimal); + REQUIRE(d16.dtype().decimal_precision() == 5); + REQUIRE(d16.dtype().decimal_scale() == 2); + + Scalar d32 = scalar::decimal(int32_t(1234), 5, 2); + REQUIRE(d32.dtype().variant() == DataTypeVariant::Decimal); + REQUIRE(d32.dtype().decimal_precision() == 5); + REQUIRE(d32.dtype().decimal_scale() == 2); + + Scalar d64 = scalar::decimal(99999, 12, 3); + REQUIRE(d64.dtype().variant() == DataTypeVariant::Decimal); + REQUIRE(d64.dtype().decimal_precision() == 12); + REQUIRE(d64.dtype().decimal_scale() == 3); +} + +TEST_CASE("Copy scalar", "[scalar]") { + Scalar a = scalar::of(42); + Scalar b = a; + REQUIRE(a.dtype().primitive_type() == I64); + REQUIRE(b.dtype().primitive_type() == I64); + + Scalar c = scalar::of(1); + c = b; + REQUIRE(c.dtype().primitive_type() == I64); +} +} // namespace diff --git a/lang/cpp/tests/scan.cpp b/lang/cpp/tests/scan.cpp new file mode 100644 index 00000000000..82854e11f2a --- /dev/null +++ b/lang/cpp/tests/scan.cpp @@ -0,0 +1,348 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include +#include +#include +#include +#include +#include + +#include "common.hpp" +#include "vortex/estimate.hpp" + +using namespace vortex; +using Catch::Matchers::ContainsSubstring; +using vortex_test::SAMPLE_ROWS; +using vortex_test::TempPath; +using vortex_test::write_sample; + +namespace { +using enum vortex::PType; + +TEST_CASE("Non-existent data source", "[datasource]") { + Session session; + + try { + DataSource::open(session, {"nonexistent"}); + FAIL("expected exception"); + } catch (const VortexException &e) { + REQUIRE_THAT(std::string(e.what()), ContainsSubstring("No files matched")); + } +} + +TEST_CASE("Read dtype and row count", "[datasource]") { + Session session; + TempPath path = write_sample(session); + + std::array paths = {path.string()}; + + DataSource other = DataSource::open(session, paths); + DataSource ds(other); + + Estimate row_count = ds.row_count(); + REQUIRE(row_count.type() == EstimateType::Exact); + REQUIRE(row_count.value() == SAMPLE_ROWS); + + DataType dt = ds.dtype(); + REQUIRE(dt.variant() == DataTypeVariant::Struct); + const std::vector fields = dt.fields(); + REQUIRE(fields.size() == 2); + REQUIRE(fields[0].name == "age"); + REQUIRE(fields[1].name == "height"); +} + +void verify_age_field(const Session &session, + const Array &age, + uint8_t first = 0, + size_t rows = SAMPLE_ROWS) { + REQUIRE(age.is_primitive(U8)); + REQUIRE(age.size() == rows); + auto view = age.values(session); + for (size_t i = 0; i < rows; ++i) { + REQUIRE(view.values()[i] == static_cast(first + i)); + } +} + +void verify_sample_array(const Session &session, const Array &array) { + REQUIRE(array.size() == SAMPLE_ROWS); + REQUIRE(array.has_dtype(DataTypeVariant::Struct)); + verify_age_field(session, array.field(0)); + + Array height = array.field(1); + REQUIRE(height.is_primitive(U16)); + auto view = height.values(session); + for (size_t i = 0; i < SAMPLE_ROWS; ++i) { + REQUIRE(view.values()[i] == (i + 1) % 200); + } + + REQUIRE_THROWS_AS(array.field(2), VortexException); +} + +TEST_CASE("Basic scan", "[scan]") { + Session session; + TempPath path = write_sample(session); + DataSource ds = DataSource::open(session, {path.string()}); + + Scan scan = ds.scan(); + REQUIRE(scan.partition_count().type() == EstimateType::Exact); + REQUIRE(scan.partition_count().value() == 1); + REQUIRE(scan.dtype().variant() == DataTypeVariant::Struct); + + auto partition = scan.next_partition(); + REQUIRE(partition.has_value()); + + Estimate rows = partition->row_count(); + REQUIRE(rows.type() == EstimateType::Exact); + REQUIRE(rows.value() == SAMPLE_ROWS); + + REQUIRE_FALSE(scan.next_partition().has_value()); + + auto array = partition->next(); + REQUIRE(array.has_value()); + REQUIRE_FALSE(partition->next().has_value()); + + verify_sample_array(session, *array); +} + +TEST_CASE("Range-for over partitions and batches", "[scan]") { + Session session; + TempPath path = write_sample(session); + DataSource ds = DataSource::open(session, {path.string()}); + + Scan scan = ds.scan(); + size_t partitions = 0; + size_t rows = 0; + for (Partition &partition : scan.partitions()) { + ++partitions; + for (Array &batch : partition.batches()) { + rows += batch.size(); + } + } + REQUIRE(partitions == 1); + REQUIRE(rows == SAMPLE_ROWS); +} + +TEST_CASE("Scan from buffer", "[scan]") { + Session session; + TempPath path = write_sample(session); + + std::ifstream file(path.path(), std::ios::binary | std::ios::ate); + const auto size = file.tellg(); + file.seekg(0, std::ios::beg); + std::vector buffer(static_cast(size)); + REQUIRE(file.read(reinterpret_cast(buffer.data()), size)); + + DataSource ds = DataSource::from_buffer(session, buffer); + Scan scan = ds.scan(); + auto partition = scan.next_partition(); + REQUIRE(partition.has_value()); + auto array = partition->next(); + REQUIRE(array.has_value()); + verify_sample_array(session, *array); +} + +TEST_CASE("Multiple paths", "[scan]") { + Session session; + TempPath first = write_sample(session); + TempPath second = write_sample(session); + + const std::vector paths = {first.string(), second.string()}; + DataSource ds = DataSource::open(session, paths); + REQUIRE(ds.row_count().value_or(0) >= SAMPLE_ROWS); +} + +TEST_CASE("Multithreaded scan", "[scan]") { + Session session; + + constexpr size_t NUM_FILES = 8; + std::vector paths; + std::vector path_strings; + for (size_t i = 0; i < NUM_FILES; ++i) { + paths.push_back(write_sample(session)); + path_strings.push_back(paths.back().string()); + } + + DataSource ds = DataSource::open(session, path_strings); + Scan scan = ds.scan(); + + std::vector threads; + threads.reserve(NUM_FILES); + std::vector rows(NUM_FILES, 0); + + for (size_t i = 0; i < NUM_FILES; ++i) { + threads.emplace_back([&, i] { + while (auto partition = scan.next_partition()) { + while (auto batch = partition->next()) { + rows[i] += batch->size(); + } + } + }); + } + for (auto &t : threads) { + t.join(); + } + + size_t total = 0; + for (size_t n : rows) { + total += n; + } + REQUIRE(total == NUM_FILES * SAMPLE_ROWS); +} + +TEST_CASE("Project field", "[projection]") { + Session session; + TempPath path = write_sample(session); + DataSource ds = DataSource::open(session, {path.string()}); + + Scan scan = ds.scan({.projection = expr::col("age")}); + auto partition = scan.next_partition(); + REQUIRE(partition.has_value()); + auto array = partition->next(); + REQUIRE(array.has_value()); + verify_age_field(session, *array); +} + +TEST_CASE("Project none fields", "[projection]") { + Session session; + TempPath path = write_sample(session); + DataSource ds = DataSource::open(session, {path.string()}); + + Scan scan = ds.scan({.projection = expr::root().select({})}); + auto partition = scan.next_partition(); + REQUIRE(partition.has_value()); + auto array = partition->next(); + REQUIRE(array.has_value()); + REQUIRE(array->dtype().fields().empty()); +} + +TEST_CASE("Project multiple fields", "[projection]") { + Session session; + TempPath path = write_sample(session); + DataSource ds = DataSource::open(session, {path.string()}); + + Scan scan = ds.scan({.projection = expr::root().select({"age", "height"})}); + auto partition = scan.next_partition(); + REQUIRE(partition.has_value()); + auto array = partition->next(); + REQUIRE(array.has_value()); + verify_age_field(session, array->field("age")); +} + +TEST_CASE("Filter age", "[filter]") { + Session session; + TempPath path = write_sample(session); + DataSource ds = DataSource::open(session, {path.string()}); + + constexpr uint8_t threshold = 50; + Scan scan = ds.scan({ + .filter = expr::gte(expr::col("age"), expr::lit(threshold)), + }); + auto partition = scan.next_partition(); + REQUIRE(partition.has_value()); + auto array = partition->next(); + REQUIRE(array.has_value()); + + REQUIRE(array->size() == SAMPLE_ROWS - threshold); + verify_age_field(session, array->field(0), threshold, SAMPLE_ROWS - threshold); +} + +TEST_CASE("Filter invalid values", "[filter]") { + Session session; + TempPath path = write_sample(session); + DataSource ds = DataSource::open(session, {path.string()}); + + Scan scan = ds.scan({.filter = expr::col("age").is_null()}); + auto partition = scan.next_partition(); + REQUIRE(!partition.has_value()); +} + +TEST_CASE("Filter with operators", "[filter]") { + using namespace vortex::expr::ops; + Session session; + TempPath path = write_sample(session); + const std::string path_string = path.string(); + std::array paths = {path_string}; + DataSource ds = DataSource::open(session, paths); + + Scan scan1 = ds.scan({ + .filter = expr::col("age") >= expr::lit(90) && expr::col("age") < expr::lit(95), + }); + + Scan scan2 = ds.scan({ + .filter = !(expr::col("age") < expr::lit(90)) && expr::col("age") < expr::lit(95), + }); + + Scan scans[2] = {std::move(scan1), std::move(scan2)}; + for (Scan &scan : scans) { + auto partition = scan.next_partition(); + REQUIRE(partition.has_value()); + auto array = partition->next(); + REQUIRE(array.has_value()); + REQUIRE(array->size() == 5); + } +} + +TEST_CASE("Type-mismatched filter", "[filter]") { + Session session; + TempPath path = write_sample(session); + DataSource ds = DataSource::open(session, {path.string()}); + + Scan scan = ds.scan({ + .filter = expr::eq(expr::col("age"), expr::lit(67)), + }); + auto partition = scan.next_partition(); + REQUIRE(partition.has_value()); + REQUIRE_THROWS_AS(partition->next(), VortexException); +} + +TEST_CASE("Row range and limit", "[scan]") { + Session session; + TempPath path = write_sample(session); + DataSource ds = DataSource::open(session, {path.string()}); + + Scan scan = ds.scan({ + .row_range = RowRange {10, 25}, + .limit = 5, + }); + auto partition = scan.next_partition(); + REQUIRE(partition.has_value()); + auto array = partition->next(); + REQUIRE(array.has_value()); + + REQUIRE(array->size() == 5); + verify_age_field(session, array->field(0), 10, 5); +} + +TEST_CASE("Selection", "[scan]") { + using enum Selection::Kind; + Session session; + TempPath path = write_sample(session); + DataSource ds = DataSource::open(session, {path.string()}); + + Scan scan = ds.scan({ + .selection = {{Include, {3, 5, 8}}}, + }); + auto partition = scan.next_partition(); + REQUIRE(partition.has_value()); + auto array = partition->next(); + REQUIRE(array.has_value()); + REQUIRE(array->size() == 3); + + auto ages = array->field(0).values(session); + REQUIRE(ages.values()[0] == 3); + REQUIRE(ages.values()[1] == 5); + REQUIRE(ages.values()[2] == 8); + + scan = ds.scan({ + .selection = {{Exclude, {0, 1, 2}}}, + }); + partition = scan.next_partition(); + REQUIRE(partition.has_value()); + array = partition->next(); + REQUIRE(array.has_value()); + REQUIRE(array->size() == 97); + + ages = array->field(0).values(session); + REQUIRE(ages.values()[0] == 3); +} +} // namespace diff --git a/lang/cpp/tests/session.cpp b/lang/cpp/tests/session.cpp new file mode 100644 index 00000000000..f99d7d3bb19 --- /dev/null +++ b/lang/cpp/tests/session.cpp @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include +#include + +using namespace vortex; + +namespace { +TEST_CASE("Session copy and move", "[session]") { + const Session s; + Session other; + other = s; + Session moved; + moved = std::move(other); + Session other2(std::move(moved)); +} +} // namespace diff --git a/lang/cpp/tests/string_binary.cpp b/lang/cpp/tests/string_binary.cpp new file mode 100644 index 00000000000..16ddb7eaa3c --- /dev/null +++ b/lang/cpp/tests/string_binary.cpp @@ -0,0 +1,133 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include + +#include +#include + +typedef struct ArrowSchema FFI_ArrowSchema; +typedef struct ArrowArray FFI_ArrowArray; +typedef struct ArrowArrayStream FFI_ArrowArrayStream; +#define USE_OWN_ARROW 1 + +#include +#include +#include +#include + +#include "common.hpp" + +using namespace vortex; +using namespace std::string_view_literals; +using vortex_test::TempPath; + +namespace { + +Array strings_from_arrow(std::span values, bool with_null) { + nanoarrow::UniqueSchema schema; + REQUIRE(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRING) == NANOARROW_OK); + nanoarrow::UniqueArray arr; + REQUIRE(ArrowArrayInitFromSchema(arr.get(), schema.get(), nullptr) == NANOARROW_OK); + REQUIRE(ArrowArrayStartAppending(arr.get()) == NANOARROW_OK); + for (const auto &value : values) { + const ArrowStringView view {value.data(), static_cast(value.size())}; + REQUIRE(ArrowArrayAppendString(arr.get(), view) == NANOARROW_OK); + } + if (with_null) { + REQUIRE(ArrowArrayAppendNull(arr.get(), 1) == NANOARROW_OK); + } + REQUIRE(ArrowArrayFinishBuildingDefault(arr.get(), nullptr) == NANOARROW_OK); + + ArrowArray raw_arr = {}; + ArrowSchema raw_schema = {}; + ArrowArrayMove(arr.get(), &raw_arr); + ArrowSchemaMove(schema.get(), &raw_schema); + return Array::from_arrow(&raw_arr, &raw_schema, true); +} + +Array bytes_from_arrow(std::span values) { + nanoarrow::UniqueSchema schema; + REQUIRE(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_BINARY) == NANOARROW_OK); + nanoarrow::UniqueArray arr; + REQUIRE(ArrowArrayInitFromSchema(arr.get(), schema.get(), nullptr) == NANOARROW_OK); + REQUIRE(ArrowArrayStartAppending(arr.get()) == NANOARROW_OK); + for (const auto &value : values) { + const ArrowBufferView view {{reinterpret_cast(value.data())}, + static_cast(value.size())}; + REQUIRE(ArrowArrayAppendBytes(arr.get(), view) == NANOARROW_OK); + } + REQUIRE(ArrowArrayFinishBuildingDefault(arr.get(), nullptr) == NANOARROW_OK); + + ArrowArray raw_arr = {}; + ArrowSchema raw_schema = {}; + ArrowArrayMove(arr.get(), &raw_arr); + ArrowSchemaMove(schema.get(), &raw_schema); + return Array::from_arrow(&raw_arr, &raw_schema, true); +} + +TEST_CASE("String view over utf8 array", "[strings]") { + Session session; + + const std::string long1(40, 'x'); + const std::vector values = {"short"sv, "Широкая строка"sv, long1, ""sv}; + Array array = strings_from_arrow(values, true); + + StringView view = array.strings(session); + REQUIRE(view.size() == values.size() + 1); + for (size_t i = 0; i < values.size(); ++i) { + REQUIRE_FALSE(view.is_null(i)); + REQUIRE(view[i] == values[i]); + } + REQUIRE(view.is_null(values.size())); + REQUIRE_THROWS_AS(view[view.size()], VortexException); +} + +TEST_CASE("Bytes view over binary array", "[strings]") { + Session session; + + // Includes a NUL byte + const std::vector values = {std::string_view {"abc\0def", 7}, "ffff"sv}; + Array array = bytes_from_arrow(values); + + BytesView view = array.bytes(session); + REQUIRE(view.size() == 2); + for (size_t i = 0; i < values.size(); ++i) { + const auto bytes = view[i]; + REQUIRE(std::string_view(reinterpret_cast(bytes.data()), bytes.size()) == values[i]); + } +} + +TEST_CASE("Strings roundtrip", "[strings]") { + Session session; + TempPath path = TempPath::unique(); + + const std::string long1(64, 'y'); + const std::vector values = {"inlined"sv, long1}; + Array strings = strings_from_arrow(values, false); + + Writer writer = + Writer::open(session, path.string(), dtype::struct_({{"s", dtype::utf8(dtype::Nullable)}})); + writer.push(make_struct({{"s", strings}})); + writer.finish(); + + DataSource ds = DataSource::open(session, {path.string()}); + Scan scan = ds.scan(); + auto partition = scan.next_partition(); + REQUIRE(partition.has_value()); + auto batch = partition->next(); + REQUIRE(batch.has_value()); + + StringView view = batch->field(0).strings(session); + REQUIRE(view.size() == 2); + REQUIRE(view[0] == "inlined"sv); + REQUIRE(view[1] == long1); +} + +TEST_CASE("strings() on a non-utf8 array", "[strings]") { + Session session; + std::vector data = {1}; + Array a = Array::primitive(data); + REQUIRE_THROWS_AS(a.strings(session), VortexException); + REQUIRE_THROWS_AS(a.bytes(session), VortexException); +} +} // namespace diff --git a/lang/cpp/tests/writer.cpp b/lang/cpp/tests/writer.cpp new file mode 100644 index 00000000000..f8dc2676fa5 --- /dev/null +++ b/lang/cpp/tests/writer.cpp @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include + +#include "common.hpp" +#include "vortex/dtype.hpp" +#include "vortex/error.hpp" +#include "vortex/session.hpp" + +using namespace vortex; +using vortex_test::TempPath; + +TEST_CASE("Dangling session writer", "[writer]") { + TempPath path = TempPath::unique(); + + { + Writer writer = Writer::open(Session {}, path.string(), vortex_test::sample_dtype()); + writer.push(vortex_test::sample_array()); + } + + REQUIRE_THROWS_AS(DataSource::open(Session {}, {path.string()}), VortexException); +} + +TEST_CASE("Unfinished writer", "[writer]") { + Session session; + TempPath path = TempPath::unique(); + + { + Writer writer = Writer::open(session, path.string(), vortex_test::sample_dtype()); + writer.push({vortex_test::sample_array()}); + } + + REQUIRE_THROWS_AS(DataSource::open(session, {path.string()}), VortexException); +} + +TEST_CASE("Write finish() called twice", "[writer]") { + Session session; + TempPath path = TempPath::unique(); + + Writer writer = Writer::open(session, path.string(), vortex_test::sample_dtype()); + writer.push(vortex_test::sample_array()); + writer.finish(); + REQUIRE_THROWS_AS(writer.finish(), VortexException); + REQUIRE_THROWS_AS(writer.push(vortex_test::sample_array()), VortexException); +} + +TEST_CASE("Writer push with invalid dtype", "[writer]") { + Session session; + TempPath path = TempPath::unique(); + Writer writer = Writer::open(session, path.string(), dtype::null()); + const std::vector arrays = {vortex_test::sample_array(), vortex_test::sample_array()}; + REQUIRE_THROWS_AS(writer.push(arrays), VortexException); +} diff --git a/policy.yml b/policy.yml new file mode 100644 index 00000000000..e669833e011 --- /dev/null +++ b/policy.yml @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +--- +policy: + approval: + - or: + - a vortex committer has approved + - an untouched renovate pull request has an allowed approval + - claude or codex authored pull requests have two committer approvals + disapproval: + options: + methods: + disapprove: + comments: + - ":-1:" + - "-1" + - "👎" + github_review: true + revoke: + comments: + - ":+1:" + - "+1" + - "👍" + github_review: true + requires: + teams: + - "vortex-data/committers" + +approval_defaults: + options: + # Allow explicit PR comments and GitHub reviews, but not PR body text. + methods: + comments: + - ":+1:" + - "+1" + - "👍" + comment_patterns: [] + github_review: true + github_review_comment_patterns: [] + body_patterns: [] + allow_author: false + allow_contributor: false + allow_non_author_contributor: true + invalidate_on_push: false + +approval_rules: + - name: a vortex committer has approved + description: "Requires approval from the Committers team." + requires: + count: 1 + teams: + - "vortex-data/committers" + + - name: an untouched renovate pull request has an allowed approval + description: "Allows approval by renovate-approve or a committer only when every commit in the PR is still Renovate-authored or Renovate-committed." + if: + has_author_in: + users: + - "renovate[bot]" + only_has_contributors_in: + users: + - "renovate[bot]" + options: + allow_author: true + allow_contributor: true + requires: + count: 1 + users: + - "renovate-approve[bot]" + teams: + - "vortex-data/committers" + + - name: claude or codex authored pull requests have two committer approvals + if: + has_author_in: + users: + - "vortex-claude[bot]" + - "claude-code[bot]" + - "codex[bot]" + - "openai-codex[bot]" + requires: + count: 2 + teams: + - "vortex-data/committers" diff --git a/scripts/_measurement_id.py b/scripts/_measurement_id.py index e8834f4c097..6e670c38dd0 100644 --- a/scripts/_measurement_id.py +++ b/scripts/_measurement_id.py @@ -1,20 +1,20 @@ # 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 `measurement_id` xxhash64 functions -- the primary-key hash of the benchmarks database. + +This began as a byte-for-byte port of the Rust implementation in 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`); with that server crate retired, this module +is now the only implementation. Re-ingesting an existing `(commit, dim-tuple)` upserts the +existing row via `ON CONFLICT (measurement_id) DO UPDATE` instead of inserting a duplicate, so +the hash is frozen forever: every row already in the production Postgres carries an id computed +this way. + +The freeze is NOT assumed -- it is pinned by `scripts/measurement_id_golden.json`, whose vectors +were generated by the original Rust source of truth and must never be regenerated. +`scripts/tests/test_measurement_id.py` asserts this module reproduces every vector; any drift +fails that check. ## The hash, precisely diff --git a/scripts/measurement_id_golden.json b/scripts/measurement_id_golden.json new file mode 100644 index 00000000000..5df8fe9775d --- /dev/null +++ b/scripts/measurement_id_golden.json @@ -0,0 +1,825 @@ +{ + "note": "Frozen golden vectors for the measurement_id_* hash, generated by the original Rust implementation (server/src/db.rs in the retired vortex-data/benchmarks-website server crate) that wrote the production rows. NEVER regenerate this file: the ids it pins are the primary keys already stored in the benchmarks RDS Postgres, and any drift would break ON CONFLICT (measurement_id) idempotency. scripts/tests/test_measurement_id.py asserts the Python port in scripts/_measurement_id.py reproduces every measurement_id here.", + "seed": 0, + "vectors": [ + { + "fields": { + "commit_sha": "0000000000000000000000000000000000000000", + "dataset": "dataset-0", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -20, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": 8941066430020068021, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000009e3779b97f4a7c15", + "dataset": "dataset-1", + "dataset_variant": "variant-1", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -19, + "scale_factor": "sf-1", + "storage": "s3" + }, + "measurement_id": 3149995011396621799, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000003c6ef372fe94f82a", + "dataset": "dataset-2", + "dataset_variant": "variant-2", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -18, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": -6215740072949592228, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000daa66d2c7ddf743f", + "dataset": "dataset-3", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -17, + "scale_factor": "sf-3", + "storage": "s3" + }, + "measurement_id": 8041805059478031959, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "00000000000000000000000078dde6e5fd29f054", + "dataset": "dataset-4", + "dataset_variant": "variant-4", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -16, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": 4772600626526399345, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000001715609f7c746c69", + "dataset": "dataset-5", + "dataset_variant": "variant-5", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -15, + "scale_factor": "sf-5", + "storage": "s3" + }, + "measurement_id": -9016235375010445132, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000b54cda58fbbee87e", + "dataset": "dataset-6", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -14, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": -6411293753624450917, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000538454127b096493", + "dataset": "dataset-7", + "dataset_variant": "variant-7", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -13, + "scale_factor": "sf-7", + "storage": "s3" + }, + "measurement_id": 8460757218662476151, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000f1bbcdcbfa53e0a8", + "dataset": "dataset-8", + "dataset_variant": "variant-8", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -12, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": 7670946045315319280, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000008ff34785799e5cbd", + "dataset": "dataset-9", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -11, + "scale_factor": "sf-9", + "storage": "s3" + }, + "measurement_id": 1978441881557979666, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000002e2ac13ef8e8d8d2", + "dataset": "dataset-10", + "dataset_variant": "variant-10", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -10, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": 8135912901467482164, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000cc623af8783354e7", + "dataset": "dataset-11", + "dataset_variant": "variant-11", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -9, + "scale_factor": "sf-11", + "storage": "s3" + }, + "measurement_id": 805733765350646006, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000006a99b4b1f77dd0fc", + "dataset": "dataset-12", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -8, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": -4091026328830870798, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "00000000000000000000000008d12e6b76c84d11", + "dataset": "dataset-13", + "dataset_variant": "variant-13", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -7, + "scale_factor": "sf-13", + "storage": "s3" + }, + "measurement_id": -6061385115965648155, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000a708a824f612c926", + "dataset": "dataset-14", + "dataset_variant": "variant-14", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -6, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": -6684883584448916863, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000454021de755d453b", + "dataset": "dataset-15", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -5, + "scale_factor": "sf-15", + "storage": "s3" + }, + "measurement_id": 8814896453978445645, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000e3779b97f4a7c150", + "dataset": "dataset-16", + "dataset_variant": "variant-16", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -4, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": 6686567714742361044, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "00000000000000000000000081af155173f23d65", + "dataset": "dataset-17", + "dataset_variant": "variant-17", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -3, + "scale_factor": "sf-17", + "storage": "s3" + }, + "measurement_id": 3557917168552872470, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000001fe68f0af33cb97a", + "dataset": "dataset-18", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -2, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": -6165762661706299592, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000be1e08c47287358f", + "dataset": "dataset-19", + "dataset_variant": "variant-19", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": -1, + "scale_factor": "sf-19", + "storage": "s3" + }, + "measurement_id": -535031789163192453, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000005c55827df1d1b1a4", + "dataset": "dataset-20", + "dataset_variant": "variant-20", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 0, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": 4634990895302954655, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000fa8cfc37711c2db9", + "dataset": "dataset-21", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 1, + "scale_factor": "sf-21", + "storage": "s3" + }, + "measurement_id": 5936866269340399925, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "00000000000000000000000098c475f0f066a9ce", + "dataset": "dataset-22", + "dataset_variant": "variant-22", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 2, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": 4471015533681449689, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "00000000000000000000000036fbefaa6fb125e3", + "dataset": "dataset-23", + "dataset_variant": "variant-23", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 3, + "scale_factor": "sf-23", + "storage": "s3" + }, + "measurement_id": 3270457410965634007, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000d5336963eefba1f8", + "dataset": "dataset-24", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 4, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": -6630843892626754319, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000736ae31d6e461e0d", + "dataset": "dataset-25", + "dataset_variant": "variant-25", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 5, + "scale_factor": "sf-25", + "storage": "s3" + }, + "measurement_id": 5238548701197127815, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "00000000000000000000000011a25cd6ed909a22", + "dataset": "dataset-26", + "dataset_variant": "variant-26", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 6, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": -869863853831808211, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000afd9d6906cdb1637", + "dataset": "dataset-27", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 7, + "scale_factor": "sf-27", + "storage": "s3" + }, + "measurement_id": -6175971040302506030, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000004e115049ec25924c", + "dataset": "dataset-28", + "dataset_variant": "variant-28", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 8, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": 4758474844754348116, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000ec48ca036b700e61", + "dataset": "dataset-29", + "dataset_variant": "variant-29", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 9, + "scale_factor": "sf-29", + "storage": "s3" + }, + "measurement_id": 4027371045171197062, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000008a8043bceaba8a76", + "dataset": "dataset-30", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 10, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": 6213625536713947544, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "00000000000000000000000028b7bd766a05068b", + "dataset": "dataset-31", + "dataset_variant": "variant-31", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 11, + "scale_factor": "sf-31", + "storage": "s3" + }, + "measurement_id": 3264730627889711995, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000c6ef372fe94f82a0", + "dataset": "dataset-32", + "dataset_variant": "variant-32", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 12, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": 6426500585598625566, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000006526b0e96899feb5", + "dataset": "dataset-33", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 13, + "scale_factor": "sf-33", + "storage": "s3" + }, + "measurement_id": -470855441762601549, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000035e2aa2e7e47aca", + "dataset": "dataset-34", + "dataset_variant": "variant-34", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 14, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": -8702944831802119265, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000a195a45c672ef6df", + "dataset": "dataset-35", + "dataset_variant": "variant-35", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 15, + "scale_factor": "sf-35", + "storage": "s3" + }, + "measurement_id": 6237767006190293538, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000003fcd1e15e67972f4", + "dataset": "dataset-36", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 16, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": 5322839569260259127, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "000000000000000000000000de0497cf65c3ef09", + "dataset": "dataset-37", + "dataset_variant": "variant-37", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 17, + "scale_factor": "sf-37", + "storage": "s3" + }, + "measurement_id": 1995931045924102080, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000007c3c1188e50e6b1e", + "dataset": "dataset-38", + "dataset_variant": "variant-38", + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 18, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": -7932518466629682266, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0000000000000000000000001a738b426458e733", + "dataset": "dataset-39", + "dataset_variant": null, + "engine": "vortex", + "format": "vortex-file-compressed", + "query_idx": 19, + "scale_factor": "sf-39", + "storage": "s3" + }, + "measurement_id": 7227626056758470989, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "tpch", + "dataset_variant": null, + "engine": "vortex", + "format": "parquet", + "query_idx": 7, + "scale_factor": "1", + "storage": "nvme" + }, + "measurement_id": 6778969912115220139, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "\ub370\uc774\ud130\uc14b-\u65e5\u672c\u8a9e", + "dataset_variant": "caf\u00e9-\u03a9", + "engine": "\u30f4", + "format": "format-\u2713", + "query_idx": 0, + "scale_factor": "\u03c3", + "storage": "s3" + }, + "measurement_id": 8038081617766937373, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "boundary", + "dataset_variant": null, + "engine": "e", + "format": "f", + "query_idx": -2147483648, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": -1997262439178828726, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "boundary", + "dataset_variant": null, + "engine": "e", + "format": "f", + "query_idx": 2147483647, + "scale_factor": null, + "storage": "nvme" + }, + "measurement_id": -1111737216055603595, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "", + "dataset_variant": "", + "engine": "", + "format": "", + "query_idx": 0, + "scale_factor": "", + "storage": "" + }, + "measurement_id": -4121125782605521286, + "table": "query_measurements" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "chimp", + "dataset_variant": null, + "format": "vortex", + "op": "encode" + }, + "measurement_id": -1351591591404870516, + "table": "compression_times" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "chimp", + "dataset_variant": null, + "format": "vortex", + "op": "decode" + }, + "measurement_id": 8221793816172450272, + "table": "compression_times" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "taxi", + "dataset_variant": "v2", + "format": "parquet", + "op": "encode" + }, + "measurement_id": -192124532873811017, + "table": "compression_times" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "\uc555\ucd95-\u30c6\u30b9\u30c8", + "dataset_variant": "\u5909\u7a2e", + "format": "vortex-\u2713", + "op": "decode" + }, + "measurement_id": -2942165613664654219, + "table": "compression_times" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "chimp", + "dataset_variant": null, + "format": "vortex" + }, + "measurement_id": 4543757287978490921, + "table": "compression_sizes" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "taxi", + "dataset_variant": "v2", + "format": "parquet" + }, + "measurement_id": -7313339766868556874, + "table": "compression_sizes" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "\ud06c\uae30-\u30b5\u30a4\u30ba", + "dataset_variant": "\u5909\u7a2e-\u03a9", + "format": "vortex-\u2713" + }, + "measurement_id": -8576081510595521142, + "table": "compression_sizes" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "chimp", + "format": "vortex" + }, + "measurement_id": 4734678989716641189, + "table": "random_access_times" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "taxi", + "format": "parquet" + }, + "measurement_id": 164255076009922704, + "table": "random_access_times" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "\u7121\u4f5c\u70ba", + "format": "\u2713" + }, + "measurement_id": 832111591737631247, + "table": "random_access_times" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "cohere-large-10m", + "flavor": "vortex", + "layout": "trained", + "threshold": 0.0 + }, + "measurement_id": 2427216480405807826, + "table": "vector_search_runs" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "cohere-large-10m", + "flavor": "vortex", + "layout": "trained", + "threshold": 0.5 + }, + "measurement_id": -1711405494360824574, + "table": "vector_search_runs" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "cohere-large-10m", + "flavor": "vortex", + "layout": "trained", + "threshold": 0.83 + }, + "measurement_id": 2194658981253663824, + "table": "vector_search_runs" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "cohere-large-10m", + "flavor": "vortex", + "layout": "trained", + "threshold": 1.0 + }, + "measurement_id": -5633397304139473003, + "table": "vector_search_runs" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "cohere-large-10m", + "flavor": "vortex", + "layout": "trained", + "threshold": -1.5 + }, + "measurement_id": -5505640980838817192, + "table": "vector_search_runs" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "cohere-large-10m", + "flavor": "vortex", + "layout": "trained", + "threshold": 1e-09 + }, + "measurement_id": -812426564273547024, + "table": "vector_search_runs" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "cohere-large-10m", + "flavor": "vortex", + "layout": "trained", + "threshold": 1000000000000.0 + }, + "measurement_id": 8044969450004818277, + "table": "vector_search_runs" + }, + { + "fields": { + "commit_sha": "0123456789abcdef0123456789abcdef01234567", + "dataset": "cohere-large-10m", + "flavor": "vortex", + "layout": "trained", + "threshold": 3.141592653589793 + }, + "measurement_id": 3687626387279485390, + "table": "vector_search_runs" + } + ] +} diff --git a/scripts/tests/test_measurement_id.py b/scripts/tests/test_measurement_id.py new file mode 100644 index 00000000000..c8baf0cfb8d --- /dev/null +++ b/scripts/tests/test_measurement_id.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Golden-vector pin for the `measurement_id` hash in `scripts/_measurement_id.py`. + +The `measurement_id` computed by `scripts/post-ingest.py --postgres` is the primary key of every +fact row in the benchmarks-website RDS Postgres. Re-ingesting an existing `(commit, dim-tuple)` +upserts via `ON CONFLICT (measurement_id)`; if the hash ever drifted, re-ingests would insert +duplicate rows next to millions of existing ones instead. The hash is therefore frozen forever. + +`scripts/measurement_id_golden.json` pins it: the vectors were generated by the original Rust +implementation (`server/src/db.rs` in the now-deleted `vortex-data/benchmarks-website` `server/` +crate, the source of truth at the time the production rows were written). With the Rust side +retired, the golden file is a frozen artifact -- it must NEVER be regenerated, because the ids it +pins are the ids already stored in production. This test asserts the Python port reproduces every +vector byte-for-byte. + +Run with: + + uv run --no-project --with pytest --with xxhash pytest scripts/tests/test_measurement_id.py +""" + +import importlib.util +import json +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +MODULE_PATH = REPO_ROOT / "scripts" / "_measurement_id.py" +GOLDEN_PATH = REPO_ROOT / "scripts" / "measurement_id_golden.json" + + +def load_measurement_id_module(): + spec = importlib.util.spec_from_file_location("_measurement_id", MODULE_PATH) + assert spec is not None + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_every_golden_vector_reproduces(): + module = load_measurement_id_module() + golden = json.loads(GOLDEN_PATH.read_text()) + assert golden["seed"] == 0 + + for vector in golden["vectors"]: + port_fn = module.MEASUREMENT_ID_BY_TABLE[vector["table"]] + computed = port_fn(**vector["fields"]) + assert computed == vector["measurement_id"], ( + f"measurement_id drift for table {vector['table']!r} fields {vector['fields']!r}: " + f"computed {computed}, golden {vector['measurement_id']}" + ) + + +def test_every_fact_table_is_covered(): + """Every port function must have at least one golden vector, so a new fact table cannot be + wired into `MEASUREMENT_ID_BY_TABLE` without also pinning its hash.""" + module = load_measurement_id_module() + golden = json.loads(GOLDEN_PATH.read_text()) + + covered_tables = {vector["table"] for vector in golden["vectors"]} + assert covered_tables == set(module.MEASUREMENT_ID_BY_TABLE) diff --git a/uv.lock b/uv.lock index 8748a91544e..645780c261a 100644 --- a/uv.lock +++ b/uv.lock @@ -1270,7 +1270,7 @@ wheels = [ [[package]] name = "ray" -version = "2.55.1" +version = "2.56.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -1283,20 +1283,20 @@ dependencies = [ { name = "requests" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/88/7d/48ba2f49b40a34b0071ee27c0144a2573d8836094eaca213d59cef12c271/ray-2.55.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:0053fd5b400f7ac56263aa1bbd3d68fb79341b08b8dc697c88782d5aca7b3ed4", size = 65835271, upload-time = "2026-04-22T20:09:34.984Z" }, - { url = "https://files.pythonhosted.org/packages/8f/a3/d6db3a428e4ea17cc72e79f747cfe11e90e63e36e1705bb8324e45f334b7/ray-2.55.1-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:0ea2f670a7725833ad2333a8c46ab69865ad06c8e5de9f65695e0f8f35331cec", size = 72879783, upload-time = "2026-04-22T20:09:40.986Z" }, - { url = "https://files.pythonhosted.org/packages/46/59/41da0e72a59cd3e8978480ccfeb86ef4235ae5ceb9b8928168a764fa930a/ray-2.55.1-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:d5382da181c03ee2f502ef46cf0ae4bbc30157b5bd9a67d7651f6a272528a85a", size = 73706515, upload-time = "2026-04-22T20:09:47.079Z" }, - { url = "https://files.pythonhosted.org/packages/65/52/c16bbdc3e31a5178f97be88966ab56db6f7e04882640c5cf2fee5b87757b/ray-2.55.1-cp311-cp311-win_amd64.whl", hash = "sha256:5e56d2e8f304cafe990c198a2b894f5b813de018998cd7212869201f6dc17cff", size = 27882093, upload-time = "2026-04-22T20:09:52.943Z" }, - { url = "https://files.pythonhosted.org/packages/ac/3a/4d34f471a68b958b7f94c974c19ad6836a61a2dc16393df4294169a2e4b0/ray-2.55.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:137f9006eee28caab8260803cca314f37bbda3fc94fdfa31c770b5d019626ad8", size = 65822379, upload-time = "2026-04-22T20:09:58.064Z" }, - { url = "https://files.pythonhosted.org/packages/f1/13/0db535102d0256b350ca116d8987588aca1a1f9ebb4638e1e1ff88bbcef8/ray-2.55.1-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:26541f69bb55607ef8335baac75b2ed12ff2ce02d56313219b29eda003039221", size = 72910802, upload-time = "2026-04-22T20:10:04.382Z" }, - { url = "https://files.pythonhosted.org/packages/4c/f8/fffadf3f4285eebd460e4d7f2ed1c0cd641ed89613c3f49eb881ee9fa7e2/ray-2.55.1-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:263705f6bab29e7622a94f82da25fd7f9cead76cdf89a07aab28f79cdf8f9d95", size = 73765203, upload-time = "2026-04-22T20:10:10.495Z" }, - { url = "https://files.pythonhosted.org/packages/10/f7/5acb86fc9625a0e6bbc40e1c7d42c60770e78585439a921c32738b6d675a/ray-2.55.1-cp312-cp312-win_amd64.whl", hash = "sha256:9ad56704c8bd7e92130162f9c58e4ef473609515637673d5a36e761f95335206", size = 27865547, upload-time = "2026-04-22T20:10:15.364Z" }, - { url = "https://files.pythonhosted.org/packages/d5/95/898699cc1a6a5f304ea95376d079843b5c05f4c8c1ec7e55a5cc7ffcea50/ray-2.55.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:f9844a9272ef2e6eb5771025866072cf4234cf4c7cc1a31e235b7de7111864be", size = 65766823, upload-time = "2026-04-22T20:10:20.786Z" }, - { url = "https://files.pythonhosted.org/packages/c9/13/87deecc090c672e45a0cf6f5eef511de448b93f37ef18fd10eb8e8557a0d/ray-2.55.1-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:b415d590e062f248907e0fe42994943f11726b7178fcf4b1cf5546721fb1a5f8", size = 72818676, upload-time = "2026-04-22T20:10:26.705Z" }, - { url = "https://files.pythonhosted.org/packages/71/d7/fc95d3b8824c62105c64aa1b59c59600b581f608d78a2af753e010936dc9/ray-2.55.1-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:1380e043eb57cde69b7e9199c6f2558ceeb8f0fc41c97d1d5e50ea042115f302", size = 73678908, upload-time = "2026-04-22T20:10:32.795Z" }, - { url = "https://files.pythonhosted.org/packages/a9/03/7e552325572e067b23a4584bda8dc6a67af8bd7e03c424d2610bfa93112d/ray-2.55.1-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:b062045c64c2bce39a51661624f7292c7bbf30f2a9d878627aae31d46da5712d", size = 65774106, upload-time = "2026-04-22T20:10:39.885Z" }, - { url = "https://files.pythonhosted.org/packages/94/62/607a8859520ce350861425f11f8e15d66c15ee33e6aac812f9e2889b5df4/ray-2.55.1-cp314-cp314-manylinux2014_aarch64.whl", hash = "sha256:4e618d61e1b14b6fde9a586151f3fd9d435b0b85048b997bcaa7f4a533747b2b", size = 72814044, upload-time = "2026-04-22T20:10:46.985Z" }, - { url = "https://files.pythonhosted.org/packages/04/5a/0699bef04a72d7dc54462960d07ef7a19cd8b1e09979880aba2b6d13cca2/ray-2.55.1-cp314-cp314-manylinux2014_x86_64.whl", hash = "sha256:156ed3e72ad95b645d2006cd71a8dddbcc89b56bfc00027f6225adf78bd9cb74", size = 73644244, upload-time = "2026-04-22T20:10:52.973Z" }, + { url = "https://files.pythonhosted.org/packages/7b/f2/2cbbbef7a67adc6555b141c7a67557bd7ce21cfb85b85772627c6a582e97/ray-2.56.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:a9ad4e26941eb2f8dbd494ad07f9f2227143164c6114132b26b23ad4f20b1c6f", size = 66354212, upload-time = "2026-06-29T20:50:37.148Z" }, + { url = "https://files.pythonhosted.org/packages/d6/69/09c8320519aa4d075bd6ecb1694c28d26b5b07a5b46050703d9ac2c5a9b6/ray-2.56.0-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:aea655831d25084cb343002a8e67a77b6aa552ddb776a65461d49f62884f096a", size = 73299499, upload-time = "2026-06-29T20:50:42.584Z" }, + { url = "https://files.pythonhosted.org/packages/18/d3/84df370c773a0b853abdba498a6258106cb928d452ff56eec7f86c2f80b4/ray-2.56.0-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:c5bf1a4384c0e2aa4420c95474b734d064cac354b0f00760be02d886afe96ca4", size = 74143858, upload-time = "2026-06-29T20:50:48.228Z" }, + { url = "https://files.pythonhosted.org/packages/5a/1b/0b4c96879ded375d0b35441c6b3633ccbea7c7c5448335589811fe5f275c/ray-2.56.0-cp311-cp311-win_amd64.whl", hash = "sha256:8e57781685bb4332edf8a7cfb1135ff53d50002b6a0504006e68365bcd26aa7c", size = 28390306, upload-time = "2026-06-29T20:50:53.841Z" }, + { url = "https://files.pythonhosted.org/packages/05/d2/83777f52c10e04654c162bfa093f4e34cf6decf7a0bb4311e2432bfa8ecb/ray-2.56.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:684a427c50745989e92a332343f0812c93b8506f71c768b95b1eefc113492699", size = 66344824, upload-time = "2026-06-29T20:50:58.367Z" }, + { url = "https://files.pythonhosted.org/packages/6e/b2/d610ffe2878dae8ab903cfdded883a054c102b0104670bca7b34b662db51/ray-2.56.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:e1fd03c6ecc5fe4c31466569e41ce0a4faf26fb930798c9d1f1eb1f405a687c8", size = 73317367, upload-time = "2026-06-29T20:51:03.893Z" }, + { url = "https://files.pythonhosted.org/packages/39/d1/8d442116f41e6bebec418d89f1556c67c77d593b55cf59e716503e7a4a17/ray-2.56.0-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:78ef34a71383c1fcf335e531e0e590867857fce9069f06ed351be6ce7a58fc50", size = 74192607, upload-time = "2026-06-29T20:51:09.528Z" }, + { url = "https://files.pythonhosted.org/packages/b9/9e/cac7454325c79eb32f36f233285d9179858c7073bc27ac91afb74854ff12/ray-2.56.0-cp312-cp312-win_amd64.whl", hash = "sha256:a8fc809dab6fc07cf05d45ea93a776c852c990512eb1fac0e3d15819fe6df10a", size = 28370978, upload-time = "2026-06-29T20:51:14.168Z" }, + { url = "https://files.pythonhosted.org/packages/65/86/29c9444c9f11f0bf3d4da75d3dc7a27c0fb86f8604beea6513968dc2df26/ray-2.56.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:992047f50473b5bfea74c8f528f999968e0b4bc735af23ad476a0f4e04741aea", size = 66289930, upload-time = "2026-06-29T20:51:20.79Z" }, + { url = "https://files.pythonhosted.org/packages/59/59/3c4d533a92e99b12b9d2e8690e41d7ace269d77deecb0aef753df7924b9e/ray-2.56.0-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:a0e9cfe92c88ab74abca23923c15a592f49bc7617ffdda4190daca7785a7c4f6", size = 73252460, upload-time = "2026-06-29T20:51:26.216Z" }, + { url = "https://files.pythonhosted.org/packages/59/ba/285d409253b332af91884daa4b0a3be3ce799682aa2c1e0f8dbbafaf1da4/ray-2.56.0-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:54d2725f8b65d9615c933fec5ec62e54a67b8f2e14286026fb014606253670ef", size = 74130685, upload-time = "2026-06-29T20:51:31.59Z" }, + { url = "https://files.pythonhosted.org/packages/82/a0/07a00091cf002662946ed2a16fe4b27c80be6fce7fc831e41c7e1471e2f7/ray-2.56.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:f38e03b77c53e3d94091aedb84b14efe7ee5b581d85b7d13925066bcd48c44a2", size = 66298834, upload-time = "2026-06-29T20:51:36.82Z" }, + { url = "https://files.pythonhosted.org/packages/30/3c/2819cc8b40bf4034ff871c18d04a0422b7b0a51b32260c6f03a06460aaa5/ray-2.56.0-cp314-cp314-manylinux2014_aarch64.whl", hash = "sha256:c3a16d43d75283a3d64fa1d904a3adaf3f526f3f508f447505b8bb8dc70bad6c", size = 73242922, upload-time = "2026-06-29T20:51:42.467Z" }, + { url = "https://files.pythonhosted.org/packages/d5/5b/33ac5616706d1bc0bd40e22baf12acccd687b505f85d107f549355379d6f/ray-2.56.0-cp314-cp314-manylinux2014_x86_64.whl", hash = "sha256:73edb6fb5fd05481b1f358ac2e8a4c7f6a031d89f9b823d25a587ddb7529070f", size = 74099117, upload-time = "2026-06-29T20:51:48.19Z" }, ] [[package]] @@ -1486,11 +1486,11 @@ wheels = [ [[package]] name = "setuptools" -version = "80.9.0" +version = "83.0.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/18/5d/3bf57dcd21979b887f014ea83c24ae194cfcd12b9e0fda66b957c69d1fca/setuptools-80.9.0.tar.gz", hash = "sha256:f36b47402ecde768dbfafc46e8e4207b4360c654f1f3bb84475f0a28628fb19c", size = 1319958, upload-time = "2025-05-27T00:56:51.443Z" } +sdist = { url = "https://files.pythonhosted.org/packages/34/26/f5d29e25ffdb535afef2d35cdb55b325298f96debd670da4c325e08d70f4/setuptools-83.0.0.tar.gz", hash = "sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef", size = 1154254, upload-time = "2026-07-04T15:31:22.699Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a3/dc/17031897dae0efacfea57dfd3a82fdd2a2aeb58e0ff71b77b87e44edc772/setuptools-80.9.0-py3-none-any.whl", hash = "sha256:062d34222ad13e0cc312a4c02d73f059e86a4acbfbdea8f8f76b28c99f306922", size = 1201486, upload-time = "2025-05-27T00:56:49.664Z" }, + { url = "https://files.pythonhosted.org/packages/5d/40/e1e72872c6354b306daef1703549e8e83b4d43cfea356311bf722a043752/setuptools-83.0.0-py3-none-any.whl", hash = "sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3", size = 1008090, upload-time = "2026-07-04T15:31:20.885Z" }, ] [[package]] @@ -1546,23 +1546,23 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "alabaster", marker = "python_full_version < '3.12'" }, - { name = "babel", marker = "python_full_version < '3.12'" }, - { name = "colorama", marker = "python_full_version < '3.12' and sys_platform == 'win32'" }, - { name = "docutils", marker = "python_full_version < '3.12'" }, - { name = "imagesize", marker = "python_full_version < '3.12'" }, - { name = "jinja2", marker = "python_full_version < '3.12'" }, - { name = "packaging", marker = "python_full_version < '3.12'" }, - { name = "pygments", marker = "python_full_version < '3.12'" }, - { name = "requests", marker = "python_full_version < '3.12'" }, - { name = "roman-numerals", marker = "python_full_version < '3.12'" }, - { name = "snowballstemmer", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-applehelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-devhelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-htmlhelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-jsmath", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-qthelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-serializinghtml", marker = "python_full_version < '3.12'" }, + { name = "alabaster" }, + { name = "babel" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "docutils" }, + { name = "imagesize" }, + { name = "jinja2" }, + { name = "packaging" }, + { name = "pygments" }, + { name = "requests" }, + { name = "roman-numerals" }, + { name = "snowballstemmer" }, + { name = "sphinxcontrib-applehelp" }, + { name = "sphinxcontrib-devhelp" }, + { name = "sphinxcontrib-htmlhelp" }, + { name = "sphinxcontrib-jsmath" }, + { name = "sphinxcontrib-qthelp" }, + { name = "sphinxcontrib-serializinghtml" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/50/a8c6ccc36d5eacdfd7913ddccd15a9cee03ecafc5ee2bc40e1f168d85022/sphinx-9.0.4.tar.gz", hash = "sha256:594ef59d042972abbc581d8baa577404abe4e6c3b04ef61bd7fc2acbd51f3fa3", size = 8710502, upload-time = "2025-12-04T07:45:27.343Z" } wheels = [ @@ -1577,23 +1577,23 @@ resolution-markers = [ "python_full_version >= '3.12'", ] dependencies = [ - { name = "alabaster", marker = "python_full_version >= '3.12'" }, - { name = "babel", marker = "python_full_version >= '3.12'" }, - { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, - { name = "docutils", marker = "python_full_version >= '3.12'" }, - { name = "imagesize", marker = "python_full_version >= '3.12'" }, - { name = "jinja2", marker = "python_full_version >= '3.12'" }, - { name = "packaging", marker = "python_full_version >= '3.12'" }, - { name = "pygments", marker = "python_full_version >= '3.12'" }, - { name = "requests", marker = "python_full_version >= '3.12'" }, - { name = "roman-numerals", marker = "python_full_version >= '3.12'" }, - { name = "snowballstemmer", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-applehelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-devhelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-htmlhelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-jsmath", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-qthelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-serializinghtml", marker = "python_full_version >= '3.12'" }, + { name = "alabaster" }, + { name = "babel" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "docutils" }, + { name = "imagesize" }, + { name = "jinja2" }, + { name = "packaging" }, + { name = "pygments" }, + { name = "requests" }, + { name = "roman-numerals" }, + { name = "snowballstemmer" }, + { name = "sphinxcontrib-applehelp" }, + { name = "sphinxcontrib-devhelp" }, + { name = "sphinxcontrib-htmlhelp" }, + { name = "sphinxcontrib-jsmath" }, + { name = "sphinxcontrib-qthelp" }, + { name = "sphinxcontrib-serializinghtml" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cd/bd/f08eb0f4eed5c83f1ba2a3bd18f7745a2b1525fad70660a1c00224ec468a/sphinx-9.1.0.tar.gz", hash = "sha256:7741722357dd75f8190766926071fed3bdc211c74dd2d7d4df5404da95930ddb", size = 8718324, upload-time = "2025-12-31T15:09:27.646Z" } wheels = [ diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 98de7caf95d..a1aa5b047f8 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -23,15 +23,7 @@ workspace = true arbitrary = { workspace = true, optional = true } arc-swap = { workspace = true } arcref = { workspace = true } -arrow-arith = { workspace = true } -arrow-array = { workspace = true, features = ["ffi"] } arrow-buffer = { workspace = true } -arrow-cast = { workspace = true } -arrow-data = { workspace = true } -arrow-ord = { workspace = true } -arrow-schema = { workspace = true, features = ["canonical_extension_types"] } -arrow-select = { workspace = true } -arrow-string = { workspace = true } async-lock = { workspace = true } bytes = { workspace = true } cfg-if = { workspace = true } @@ -45,6 +37,7 @@ humansize = { workspace = true } inventory = { workspace = true } itertools = { workspace = true } jiff = { workspace = true } +memchr = { workspace = true } num-traits = { workspace = true } num_enum = { workspace = true } parking_lot = { workspace = true } @@ -55,6 +48,8 @@ primitive-types = { workspace = true, optional = true, features = [ ] } prost = { workspace = true } rand = { workspace = true } +regex = { workspace = true } +regex-syntax = { workspace = true } rstest = { workspace = true, optional = true } rstest_reuse = { workspace = true, optional = true } rustc-hash = { workspace = true } @@ -69,7 +64,7 @@ termtree = { workspace = true } tracing = { workspace = true } uuid = { workspace = true } vortex-array-macros = { workspace = true } -vortex-buffer = { workspace = true, features = ["arrow"] } +vortex-buffer = { workspace = true } vortex-compute = { workspace = true } vortex-error = { workspace = true, features = ["flatbuffers"] } vortex-flatbuffers = { workspace = true, features = ["array", "dtype"] } @@ -87,10 +82,10 @@ _test-harness = ["dep:goldenfile", "dep:rstest", "dep:rstest_reuse"] serde = ["dep:serde", "vortex-buffer/serde", "vortex-mask/serde"] [dev-dependencies] -arrow-cast = { workspace = true } divan = { workspace = true } futures = { workspace = true, features = ["executor"] } insta = { workspace = true } +mimalloc = { workspace = true } rand_distr = { workspace = true } rstest = { workspace = true } serde_json = { workspace = true } @@ -138,6 +133,10 @@ harness = false name = "kleene_bool" harness = false +[[bench]] +name = "like" +harness = false + [[bench]] name = "interleave" harness = false @@ -200,6 +199,10 @@ required-features = ["_test-harness"] name = "dict_mask" harness = false +[[bench]] +name = "validity_is_valid" +harness = false + [[bench]] name = "dict_unreferenced_mask" harness = false @@ -216,6 +219,10 @@ harness = false name = "primitive_zip" harness = false +[[bench]] +name = "listview_zip" +harness = false + [[bench]] name = "take_primitive" harness = false @@ -236,22 +243,26 @@ harness = false name = "filter_bool" harness = false +[[bench]] +name = "filter_fixed_width" +harness = false + [[bench]] name = "list_length" harness = false [[bench]] -name = "listview_rebuild" +name = "list_sum" harness = false [[bench]] -name = "listview_builder_extend" +name = "listview_rebuild" harness = false [[bench]] -name = "slice_dict_primitive" +name = "listview_builder_extend" harness = false [[bench]] -name = "to_arrow" +name = "slice_dict_primitive" harness = false diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index fdc43dbcabd..11477e57503 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -28,6 +28,7 @@ use vortex_buffer::Buffer; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } diff --git a/vortex-array/benches/aggregate_max.rs b/vortex-array/benches/aggregate_max.rs index 4da7d1fc653..a47bdb28164 100644 --- a/vortex-array/benches/aggregate_max.rs +++ b/vortex-array/benches/aggregate_max.rs @@ -7,16 +7,18 @@ use divan::Bencher; use rand::prelude::*; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } const N: usize = 100_000; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(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 81402b24750..98e6845edb7 100644 --- a/vortex-array/benches/aggregate_sum.rs +++ b/vortex-array/benches/aggregate_sum.rs @@ -7,17 +7,19 @@ use divan::Bencher; use rand::prelude::*; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::expr::stats::Stat; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } const N: usize = 100_000; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(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 index 98c69413193..0740eac43f0 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -11,10 +11,12 @@ use std::sync::LazyLock; use divan::Bencher; use divan::counter::ItemsCount; +use mimalloc::MiMalloc; use vortex_array::ArrayRef; use vortex_array::Executable; 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::PrimitiveArray; @@ -22,11 +24,15 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_session::VortexSession; +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + fn main() { + LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); const LEN: usize = 65_536; diff --git a/vortex-array/benches/bool_zip.rs b/vortex-array/benches/bool_zip.rs index 40496fd1f0b..d1fc96c457e 100644 --- a/vortex-array/benches/bool_zip.rs +++ b/vortex-array/benches/bool_zip.rs @@ -10,16 +10,18 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::builtins::ArrayBuiltins; use vortex_mask::Mask; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); const LEN: usize = 65_536; diff --git a/vortex-array/benches/cast_primitive.rs b/vortex-array/benches/cast_primitive.rs index 3fd89270517..69a4c5d2cc5 100644 --- a/vortex-array/benches/cast_primitive.rs +++ b/vortex-array/benches/cast_primitive.rs @@ -10,6 +10,7 @@ use rand::prelude::*; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; @@ -19,6 +20,7 @@ use vortex_array::expr::stats::Stat; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } @@ -26,7 +28,7 @@ fn main() { // the kernel cost shows up clearly rather than being hidden by DRAM bandwidth. const SIZES: &[usize] = &[65_536]; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); #[divan::bench(args = SIZES)] fn cast_u16_to_u32(bencher: Bencher, n: usize) { diff --git a/vortex-array/benches/chunk_array_builder.rs b/vortex-array/benches/chunk_array_builder.rs index 6582071a6e1..c2dedc6ef46 100644 --- a/vortex-array/benches/chunk_array_builder.rs +++ b/vortex-array/benches/chunk_array_builder.rs @@ -4,6 +4,7 @@ use std::sync::LazyLock; use divan::Bencher; +use mimalloc::MiMalloc; use rand::RngExt; use rand::SeedableRng; use rand::prelude::StdRng; @@ -11,6 +12,7 @@ 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::ChunkedArray; use vortex_array::arrays::ConstantArray; @@ -21,7 +23,11 @@ use vortex_array::dtype::DType; use vortex_error::VortexExpect; use vortex_session::VortexSession; +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + fn main() { + LazyLock::force(&SESSION); divan::main(); } @@ -32,7 +38,7 @@ const BENCH_ARGS: &[(usize, usize)] = &[ (1000, 10), ]; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(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 74722235e54..de2f40dcbf5 100644 --- a/vortex-array/benches/chunked_dict_builder.rs +++ b/vortex-array/benches/chunked_dict_builder.rs @@ -15,6 +15,7 @@ use vortex_error::VortexExpect; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } diff --git a/vortex-array/benches/chunked_fsl_canonicalize.rs b/vortex-array/benches/chunked_fsl_canonicalize.rs index 4277b33382e..9be8433ce9d 100644 --- a/vortex-array/benches/chunked_fsl_canonicalize.rs +++ b/vortex-array/benches/chunked_fsl_canonicalize.rs @@ -16,6 +16,7 @@ use divan::Bencher; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::validity::Validity; @@ -23,10 +24,11 @@ use vortex_buffer::Buffer; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(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 95702c9f8f5..4a399760dc2 100644 --- a/vortex-array/benches/compare.rs +++ b/vortex-array/benches/compare.rs @@ -4,69 +4,221 @@ #![expect(clippy::unwrap_used)] use divan::Bencher; +use mimalloc::MiMalloc; 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::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DecimalDType; +use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + fn main() { divan::main(); } const ARRAY_SIZE: usize = 65_536; -#[divan::bench] -fn compare_bool(bencher: Bencher) { - let mut rng = StdRng::seed_from_u64(0); - let range = Uniform::new(0u8, 1).unwrap(); - - 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(); +fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { let session = vortex_array::array_session(); - bencher - .with_inputs(|| (&arr1, &arr2, session.create_execution_ctx())) + .with_inputs(|| (&lhs, &rhs, session.create_execution_ctx())) .bench_refs(|input| { input .0 .clone() - .binary(input.1.clone(), Operator::Gte) + .binary(input.1.clone(), op) .unwrap() .execute::(&mut input.2) }); } -#[divan::bench] -fn compare_int(bencher: Bencher) { - let mut rng = StdRng::seed_from_u64(0); - let range = Uniform::new(0i64, 100_000_000).unwrap(); +fn bool_array(rng: &mut StdRng) -> ArrayRef { + BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.random_bool(0.5))).into_array() +} - let arr1 = (0..ARRAY_SIZE) +fn bool_array_nullable(rng: &mut StdRng) -> ArrayRef { + BoolArray::new( + (0..ARRAY_SIZE).map(|_| rng.random_bool(0.5)).collect(), + Validity::from_iter((0..ARRAY_SIZE).map(|_| rng.random_bool(0.9))), + ) + .into_array() +} + +fn int_array(rng: &mut StdRng) -> ArrayRef { + let range = Uniform::new(0i64, 100_000_000).unwrap(); + (0..ARRAY_SIZE) .map(|_| rng.sample(range)) .collect::>() - .into_array(); + .into_array() +} - let arr2 = (0..ARRAY_SIZE) - .map(|_| rng.sample(range)) +fn int_array_nullable(rng: &mut StdRng) -> ArrayRef { + let range = Uniform::new(0i64, 100_000_000).unwrap(); + PrimitiveArray::new( + (0..ARRAY_SIZE) + .map(|_| rng.sample(range)) + .collect::>(), + Validity::from_iter((0..ARRAY_SIZE).map(|_| rng.random_bool(0.9))), + ) + .into_array() +} + +fn float_array(rng: &mut StdRng) -> ArrayRef { + (0..ARRAY_SIZE) + .map(|_| rng.random_range(0.0f64..1.0)) .collect::>() - .into_array(); - let session = vortex_array::array_session(); + .into_array() +} - bencher - .with_inputs(|| (&arr1, &arr2, session.create_execution_ctx())) - .bench_refs(|input| { - input - .0 - .clone() - .binary(input.1.clone(), Operator::Gte) - .unwrap() - .execute::(&mut input.2) - }); +fn string_array(rng: &mut StdRng) -> ArrayRef { + VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| { + let len = rng.random_range(1usize..24); + (0..len) + .map(|_| char::from(rng.random_range(b'a'..=b'z'))) + .collect::() + })) + .into_array() +} + +fn decimal_array(rng: &mut StdRng) -> ArrayRef { + DecimalArray::from_iter::( + (0..ARRAY_SIZE).map(|_| rng.random_range(0i128..100_000_000)), + DecimalDType::new(38, 2), + ) + .into_array() +} + +#[divan::bench] +fn compare_bool(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = bool_array(&mut rng); + let arr2 = bool_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_bool_nullable(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = bool_array_nullable(&mut rng); + let arr2 = bool_array_nullable(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_bool_constant(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr = bool_array(&mut rng); + let constant = ConstantArray::new(true, ARRAY_SIZE).into_array(); + bench_compare(bencher, arr, constant, Operator::Eq); +} + +#[divan::bench] +fn compare_int(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = int_array(&mut rng); + let arr2 = int_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_int_nullable(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = int_array_nullable(&mut rng); + let arr2 = int_array_nullable(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_int_constant(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr = int_array(&mut rng); + let constant = ConstantArray::new(50_000_000i64, ARRAY_SIZE).into_array(); + bench_compare(bencher, arr, constant, Operator::Gte); +} + +#[divan::bench] +fn compare_int_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = int_array(&mut rng); + let arr2 = int_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + +#[divan::bench] +fn compare_float(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = float_array(&mut rng); + let arr2 = float_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_decimal(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = decimal_array(&mut rng); + let arr2 = decimal_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_string_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = string_array(&mut rng); + let arr2 = string_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + +#[divan::bench] +fn compare_string_lt(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = string_array(&mut rng); + let arr2 = string_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Lt); +} + +#[divan::bench] +fn compare_string_constant(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr = string_array(&mut rng); + let constant = ConstantArray::new(Scalar::from("mmmmmmmmmmmm"), ARRAY_SIZE).into_array(); + bench_compare(bencher, arr, constant, Operator::Lt); +} + +fn struct_array(rng: &mut StdRng) -> ArrayRef { + StructArray::from_fields(&[("a", int_array(rng)), ("b", int_array(rng))]) + .unwrap() + .into_array() +} + +#[divan::bench] +fn compare_struct_lt(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = struct_array(&mut rng); + let arr2 = struct_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Lt); +} + +#[divan::bench] +fn compare_struct_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = struct_array(&mut rng); + let arr2 = struct_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); } diff --git a/vortex-array/benches/dict_compare.rs b/vortex-array/benches/dict_compare.rs index 446623939f1..5b31cf51831 100644 --- a/vortex-array/benches/dict_compare.rs +++ b/vortex-array/benches/dict_compare.rs @@ -9,6 +9,7 @@ use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; @@ -24,10 +25,11 @@ use vortex_array::scalar_fn::fns::operators::Operator; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); const LENGTH_AND_UNIQUE_VALUES: &[(usize, usize)] = &[ // length, unique_values diff --git a/vortex-array/benches/dict_compress.rs b/vortex-array/benches/dict_compress.rs index 2901d82c6b5..fa53dae8c43 100644 --- a/vortex-array/benches/dict_compress.rs +++ b/vortex-array/benches/dict_compress.rs @@ -6,11 +6,13 @@ use std::sync::LazyLock; use divan::Bencher; +use mimalloc::MiMalloc; use rand::distr::Distribution; use rand::distr::StandardUniform; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::dict_test::gen_primitive_for_dict; @@ -19,7 +21,11 @@ use vortex_array::builders::dict::dict_encode; use vortex_array::dtype::NativePType; use vortex_session::VortexSession; +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + fn main() { + LazyLock::force(&SESSION); divan::main(); } @@ -37,7 +43,7 @@ const BENCH_ARGS: &[(usize, usize)] = &[ (10_000, 512), ]; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); #[divan::bench(types = [u8, f32, i64], args = BENCH_ARGS)] fn encode_primitives(bencher: Bencher, (len, unique_values): (usize, usize)) diff --git a/vortex-array/benches/dict_unreferenced_mask.rs b/vortex-array/benches/dict_unreferenced_mask.rs index 1ec143682af..ec8653d6688 100644 --- a/vortex-array/benches/dict_unreferenced_mask.rs +++ b/vortex-array/benches/dict_unreferenced_mask.rs @@ -3,14 +3,21 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::DictArrayExt; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(array_session); fn main() { divan::main(); @@ -41,8 +48,8 @@ fn bench_many_codes_few_values(bencher: Bencher, num_values: i32) { let array = DictArray::try_new(codes, values).unwrap(); bencher - .with_inputs(|| &array) - .bench_refs(|array| array.compute_referenced_values_mask(false).unwrap()); + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| array.compute_referenced_values_mask(false, ctx).unwrap()); } /// Benchmark with many nulls in the codes array. @@ -73,8 +80,8 @@ fn bench_many_nulls(bencher: Bencher, fraction_valid: f64) { let array = DictArray::try_new(codes, values).unwrap(); bencher - .with_inputs(|| &array) - .bench_refs(|array| array.compute_referenced_values_mask(false).unwrap()); + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| array.compute_referenced_values_mask(false, ctx).unwrap()); } /// Benchmark with sparse code coverage (many unreferenced values). @@ -107,6 +114,6 @@ fn bench_sparse_coverage(bencher: Bencher, fraction_coverage: f64) { let array = DictArray::try_new(codes, values).unwrap(); bencher - .with_inputs(|| &array) - .bench_refs(|array| array.compute_referenced_values_mask(false).unwrap()); + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| array.compute_referenced_values_mask(false, ctx).unwrap()); } diff --git a/vortex-array/benches/expr/case_when_bench.rs b/vortex-array/benches/expr/case_when_bench.rs index 1bada8fa9e1..c7b0e96686d 100644 --- a/vortex-array/benches/expr/case_when_bench.rs +++ b/vortex-array/benches/expr/case_when_bench.rs @@ -11,6 +11,7 @@ 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::StructArray; use vortex_array::expr::case_when; @@ -25,9 +26,10 @@ use vortex_array::expr::root; use vortex_buffer::Buffer; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); fn main() { + LazyLock::force(&SESSION); divan::main(); } diff --git a/vortex-array/benches/filter_bool.rs b/vortex-array/benches/filter_bool.rs index cd31805bb77..20782a539b8 100644 --- a/vortex-array/benches/filter_bool.rs +++ b/vortex-array/benches/filter_bool.rs @@ -18,16 +18,18 @@ use rand::prelude::*; use rand_distr::Zipf; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_buffer::BitBuffer; use vortex_mask::Mask; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(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/filter_fixed_width.rs b/vortex-array/benches/filter_fixed_width.rs new file mode 100644 index 00000000000..01334108848 --- /dev/null +++ b/vortex-array/benches/filter_fixed_width.rs @@ -0,0 +1,158 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compact fixed-width filter benchmarks. +//! +//! The cases cover the dispatch dimensions without a large Cartesian product so CodSpeed's +//! instruction-count simulation remains inexpensive. + +#![expect( + clippy::cast_possible_truncation, + clippy::cast_precision_loss, + clippy::unwrap_used +)] + +use std::sync::LazyLock; + +use divan::Bencher; +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::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::i256; +use vortex_buffer::BitBuffer; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +// Keep each case small: the sweep has 24 cases and the targeted sections add only seven more. +const LEN: usize = 16_384; +const DENSITIES: &[f64] = &[0.01, 0.5, 0.8, 0.95]; +const CACHED_DENSITIES: &[f64] = &[0.01, 0.1]; + +#[derive(Clone, Copy, Debug)] +enum Pattern { + Random, + Runs, + Contiguous, +} + +const PATTERNS: &[Pattern] = &[Pattern::Random, Pattern::Runs, Pattern::Contiguous]; + +fn random_mask(density: f64) -> Mask { + let threshold = (density * u64::MAX as f64) as u64; + let mut state = 0x1234_5678_9abc_def0u64; + Mask::from_buffer(BitBuffer::from_iter((0..LEN).map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state <= threshold + }))) +} + +fn pattern_mask(pattern: Pattern) -> Mask { + match pattern { + Pattern::Random => random_mask(0.5), + Pattern::Runs => Mask::from_iter((0..LEN).map(|index| (index / 32).is_multiple_of(2))), + Pattern::Contiguous => Mask::from_slices(LEN, vec![(LEN / 4, LEN * 3 / 4)]), + } +} + +fn bench_filter( + bencher: Bencher, + array: ArrayRef, + make_mask: impl Fn() -> Mask + Sync, + cache_indices: bool, +) { + bencher + .with_inputs(|| { + let mask = make_mask(); + if cache_indices { + let _ = mask.values().unwrap().indices(); + } + (array.clone(), mask, SESSION.create_execution_ctx()) + }) + .bench_refs(|(array, mask, ctx)| { + divan::black_box( + array + .clone() + .filter(mask.clone()) + .unwrap() + .execute::(ctx) + .unwrap(), + ); + }); +} + +fn i8_array() -> ArrayRef { + PrimitiveArray::from_iter((0..LEN).map(|index| index as i8)).into_array() +} + +fn i16_array() -> ArrayRef { + PrimitiveArray::from_iter((0..LEN).map(|index| index as i16)).into_array() +} + +fn i32_array() -> ArrayRef { + PrimitiveArray::from_iter((0..LEN).map(|index| index as i32)).into_array() +} + +fn i64_array() -> ArrayRef { + PrimitiveArray::from_iter((0..LEN).map(|index| index as i64)).into_array() +} + +fn i128_array() -> ArrayRef { + DecimalArray::from_iter( + (0..LEN).map(|index| index as i128), + DecimalDType::new(19, 0), + ) + .into_array() +} + +fn i256_array() -> ArrayRef { + DecimalArray::from_iter( + (0..LEN).map(|index| i256::from_i128(index as i128)), + DecimalDType::new(39, 0), + ) + .into_array() +} + +macro_rules! random_density_benchmark { + ($name:ident, $array:ident) => { + #[divan::bench(args = DENSITIES)] + fn $name(bencher: Bencher, density: f64) { + bench_filter(bencher, $array(), || random_mask(density), false); + } + }; +} + +random_density_benchmark!(random_i8, i8_array); +random_density_benchmark!(random_i16, i16_array); +random_density_benchmark!(random_i32, i32_array); +random_density_benchmark!(random_i64, i64_array); +random_density_benchmark!(random_i128, i128_array); +random_density_benchmark!(random_i256, i256_array); + +#[divan::bench(args = PATTERNS)] +fn patterns_i128(bencher: Bencher, pattern: Pattern) { + bench_filter(bencher, i128_array(), || pattern_mask(pattern), false); +} + +#[divan::bench(args = CACHED_DENSITIES)] +fn cached_indices_i32(bencher: Bencher, density: f64) { + bench_filter(bencher, i32_array(), || random_mask(density), true); +} + +#[divan::bench(args = CACHED_DENSITIES)] +fn cached_indices_i128(bencher: Bencher, density: f64) { + bench_filter(bencher, i128_array(), || random_mask(density), true); +} diff --git a/vortex-array/benches/kleene_bool.rs b/vortex-array/benches/kleene_bool.rs index b23a0580a5a..60a0171a9b9 100644 --- a/vortex-array/benches/kleene_bool.rs +++ b/vortex-array/benches/kleene_bool.rs @@ -13,7 +13,6 @@ 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; @@ -31,7 +30,6 @@ static SESSION: LazyLock = LazyLock::new(|| { VortexSession::empty() .with::() .with::() - .with::() }); const LEN: usize = 65_536; diff --git a/vortex-array/benches/like.rs b/vortex-array/benches/like.rs new file mode 100644 index 00000000000..68219724717 --- /dev/null +++ b/vortex-array/benches/like.rs @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use divan::Bencher; +use rand::RngExt; +use rand::SeedableRng; +use rand::distr::Uniform; +use rand::prelude::StdRng; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::scalar_fn::fns::like::Like; +use vortex_array::scalar_fn::fns::like::LikeOptions; + +fn main() { + divan::main(); +} + +const ARRAY_SIZE: usize = 2_048; + +/// Random lowercase strings of 4..=24 bytes, some with a `hello` infix. +fn strings() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(0); + let len_dist = Uniform::new_inclusive(4usize, 24).unwrap(); + VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|i| { + let len = rng.sample(len_dist); + let mut s: String = (0..len) + .map(|_| char::from(rng.random_range(b'a'..=b'z'))) + .collect(); + if i % 7 == 0 { + s.insert_str(len / 2, "hello"); + } + s + })) + .into_array() +} + +fn bench_like(bencher: Bencher, pattern: &str, options: LikeOptions) { + let session = vortex_array::array_session(); + let array = strings(); + bencher + .with_inputs(|| { + ( + Like.try_new_array( + ARRAY_SIZE, + options, + [ + array.clone(), + ConstantArray::new(pattern, ARRAY_SIZE).into_array(), + ], + ) + .unwrap(), + session.create_execution_ctx(), + ) + }) + .bench_values(|(array, mut ctx)| array.execute::(&mut ctx).unwrap()); +} + +#[divan::bench] +fn like_exact(bencher: Bencher) { + bench_like(bencher, "hello", LikeOptions::default()); +} + +#[divan::bench] +fn like_prefix(bencher: Bencher) { + bench_like(bencher, "hello%", LikeOptions::default()); +} + +#[divan::bench] +fn like_suffix(bencher: Bencher) { + bench_like(bencher, "%hello", LikeOptions::default()); +} + +#[divan::bench] +fn like_contains(bencher: Bencher) { + bench_like(bencher, "%hello%", LikeOptions::default()); +} + +#[divan::bench] +fn like_regex(bencher: Bencher) { + bench_like(bencher, "h_llo%w%d", LikeOptions::default()); +} + +#[divan::bench] +fn like_per_row_patterns(bencher: Bencher) { + let session = vortex_array::array_session(); + let array = strings(); + // A non-constant pattern child takes the per-row path; repeated patterns hit the + // compile cache. + let patterns = VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| "hello%")).into_array(); + bencher + .with_inputs(|| { + ( + Like.try_new_array( + ARRAY_SIZE, + LikeOptions::default(), + [array.clone(), patterns.clone()], + ) + .unwrap(), + session.create_execution_ctx(), + ) + }) + .bench_values(|(array, mut ctx)| array.execute::(&mut ctx).unwrap()); +} + +#[divan::bench] +fn ilike_contains(bencher: Bencher) { + bench_like( + bencher, + "%HELLO%", + LikeOptions { + negated: false, + case_insensitive: true, + }, + ); +} diff --git a/vortex-array/benches/list_sum.rs b/vortex-array/benches/list_sum.rs new file mode 100644 index 00000000000..6092b73a807 --- /dev/null +++ b/vortex-array/benches/list_sum.rs @@ -0,0 +1,192 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks for the `list_sum` scalar function over `List`, `ListView`, and +//! `FixedSizeList` inputs. + +#![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::FixedSizeListArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::expr::list_sum; +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 * BASE_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()), + ) +} + +/// `0..total` as `i32` elements, with ~1/8 nulled at random positions when `nullable`. +fn make_elements(total: i32, nullable: bool) -> ArrayRef { + if !nullable { + return PrimitiveArray::from_iter(0..total).into_array(); + } + let mut rng = StdRng::seed_from_u64(total as u64); + let null_dist = Uniform::new(0u32, 8).unwrap(); + PrimitiveArray::from_option_iter((0..total).map(|v| (rng.sample(null_dist) != 0).then_some(v))) + .into_array() +} + +/// A canonical `List` of `num_lists` variable-length lists, ~1/8 of them null. +fn make_list(num_lists: usize, nullable_elements: bool) -> ArrayRef { + let (sizes, validity) = random_lists(num_lists); + let total: i32 = sizes.iter().sum(); + let elements = make_elements(total, nullable_elements); + 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 = make_elements(total, false); + 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() +} + +/// A `FixedSizeList` of `num_lists` lists, ~1/8 of them null. +fn make_fsl(num_lists: usize) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(num_lists as u64); + let total = (num_lists * BASE_LIST_SIZE) as i32; + let elements = make_elements(total, false); + let null_dist = Uniform::new(0u32, 8).unwrap(); + let valid = (0..num_lists).map(|_| rng.sample(null_dist) != 0); + let validity = Validity::Array(BoolArray::from_iter(valid).into_array()); + FixedSizeListArray::new(elements, BASE_LIST_SIZE as u32, validity, num_lists).into_array() +} + +/// Apply `list_sum(root())` and materialize the result. +fn run(bencher: Bencher, array: ArrayRef) { + let expr = list_sum(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_sum_small(bencher: Bencher) { + run(bencher, make_list(SMALL, false)); +} + +#[divan::bench] +fn list_sum_medium(bencher: Bencher) { + run(bencher, make_list(MEDIUM, false)); +} + +#[divan::bench] +fn list_sum_large(bencher: Bencher) { + run(bencher, make_list(LARGE, false)); +} + +#[divan::bench] +fn list_sum_nullable_elements_medium(bencher: Bencher) { + run(bencher, make_list(MEDIUM, true)); +} + +#[divan::bench] +fn list_sum_nullable_elements_large(bencher: Bencher) { + run(bencher, make_list(LARGE, true)); +} + +#[divan::bench] +fn listview_sum_small(bencher: Bencher) { + run(bencher, make_listview(SMALL)); +} + +#[divan::bench] +fn listview_sum_medium(bencher: Bencher) { + run(bencher, make_listview(MEDIUM)); +} + +#[divan::bench] +fn listview_sum_large(bencher: Bencher) { + run(bencher, make_listview(LARGE)); +} + +#[divan::bench] +fn fsl_sum_small(bencher: Bencher) { + run(bencher, make_fsl(SMALL)); +} + +#[divan::bench] +fn fsl_sum_medium(bencher: Bencher) { + run(bencher, make_fsl(MEDIUM)); +} + +#[divan::bench] +fn fsl_sum_large(bencher: Bencher) { + run(bencher, make_fsl(LARGE)); +} diff --git a/vortex-array/benches/listview_builder_extend.rs b/vortex-array/benches/listview_builder_extend.rs index 75487564ea8..99f35e9fe82 100644 --- a/vortex-array/benches/listview_builder_extend.rs +++ b/vortex-array/benches/listview_builder_extend.rs @@ -3,14 +3,16 @@ #![expect(clippy::cast_possible_truncation)] #![expect(clippy::cast_possible_wrap)] +#![expect(clippy::unwrap_used)] use std::sync::Arc; use divan::Bencher; use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::builders::ArrayBuilder; use vortex_array::builders::ListViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::NonNullable; @@ -63,13 +65,14 @@ fn extend_from_array_zctl(bencher: Bencher, (num_lists, list_size): (usize, usiz let source = source.into_array(); bencher.with_inputs(|| &source).bench_refs(|source| { - let mut builder = ListViewBuilder::::with_capacity( + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ListViewBuilder::::with_capacity( Arc::new(DType::Primitive(I32, NonNullable)), NonNullable, num_lists * list_size, num_lists, ); - builder.extend_from_array(source); + source.append_to_builder(&mut builder, &mut ctx).unwrap(); divan::black_box(builder.finish_into_listview()) }); } @@ -85,13 +88,14 @@ fn extend_from_array_non_zctl_overlapping( let source = source.into_array(); bencher.with_inputs(|| &source).bench_refs(|source| { - let mut builder = ListViewBuilder::::with_capacity( + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ListViewBuilder::::with_capacity( Arc::new(DType::Primitive(I32, NonNullable)), Nullable, num_lists * list_size, num_lists, ); - builder.extend_from_array(source); + source.append_to_builder(&mut builder, &mut ctx).unwrap(); divan::black_box(builder.finish_into_listview()) }); } diff --git a/vortex-array/benches/listview_rebuild.rs b/vortex-array/benches/listview_rebuild.rs index 81157d52b5a..2e61b969244 100644 --- a/vortex-array/benches/listview_rebuild.rs +++ b/vortex-array/benches/listview_rebuild.rs @@ -12,6 +12,7 @@ use divan::Bencher; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; @@ -26,11 +27,12 @@ use vortex_buffer::Buffer; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } /// A shared session for the `ListView` rebuild benchmarks, used to create execution contexts. -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(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/listview_zip.rs b/vortex-array/benches/listview_zip.rs new file mode 100644 index 00000000000..c50180a0e1a --- /dev/null +++ b/vortex-array/benches/listview_zip.rs @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::validity::Validity; +use vortex_buffer::BufferMut; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +// Smaller than the value-path benches: listview zip cost is dominated by element concatenation and +// per-list canonicalization. A few thousand lists already exercise the select while keeping each +// case well under a few hundred microseconds under CodSpeed's instruction-count simulation, which +// runs ~10x the local walltime. +const LEN: usize = 4_096; + +/// Fragmented (alternating) mask: the worst case for the per-element branch this kernel replaces. +/// The branchless chunked select is mask-shape-independent, so one shape suffices. +fn mask() -> Mask { + Mask::from_iter((0..LEN).map(|i| i.is_multiple_of(2))) +} + +#[divan::bench] +fn nonnull(bencher: Bencher) { + run(bencher, list_view(0, false), list_view(1_000_000, false)); +} + +#[divan::bench] +fn nullable(bencher: Bencher) { + run(bencher, list_view(0, true), list_view(1_000_000, true)); +} + +fn run(bencher: Bencher, if_true: ArrayRef, if_false: ArrayRef) { + let mask = mask(); + bencher + .with_inputs(|| { + ( + if_true.clone(), + if_false.clone(), + mask.clone().into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(t, f, m, ctx)| { + m.zip(t.clone(), f.clone()) + .unwrap() + .execute::(ctx) + .unwrap(); + }); +} + +/// `LEN` single-element lists: `list[i] = [base + i]`. When `nullable`, every 7th list is null +/// (list-level validity backed by a `BoolArray`), exercising the `zip_validity` path. +fn list_view(base: i64, nullable: bool) -> ArrayRef { + let mut elements = BufferMut::::with_capacity(LEN); + elements.extend((0..LEN as i64).map(|i| base + i)); + let offsets: BufferMut = (0..LEN as u64).collect(); + let sizes: BufferMut = std::iter::repeat_n(1u64, LEN).collect(); + + let validity = if nullable { + Validity::Array(BoolArray::from_iter((0..LEN).map(|i| !i.is_multiple_of(7))).into_array()) + } else { + Validity::NonNullable + }; + + ListViewArray::try_new( + elements.freeze().into_array(), + offsets.freeze().into_array(), + sizes.freeze().into_array(), + validity, + ) + .unwrap() + .into_array() +} diff --git a/vortex-array/benches/primitive_zip.rs b/vortex-array/benches/primitive_zip.rs index adcabf51d7d..93e83b372dc 100644 --- a/vortex-array/benches/primitive_zip.rs +++ b/vortex-array/benches/primitive_zip.rs @@ -10,20 +10,24 @@ use std::sync::LazyLock; use divan::Bencher; +use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; +use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_mask::Mask; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(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. @@ -49,7 +53,7 @@ fn nullable(bencher: Bencher) { run(bencher, if_true, if_false); } -fn run(bencher: Bencher, if_true: vortex_array::ArrayRef, if_false: vortex_array::ArrayRef) { +fn run(bencher: Bencher, if_true: ArrayRef, if_false: ArrayRef) { let mask = mask(); bencher .with_inputs(|| { @@ -71,10 +75,7 @@ fn run(bencher: Bencher, if_true: vortex_array::ArrayRef, if_false: vortex_array fn nonnull_array(base: i64) -> PrimitiveArray { let mut values = BufferMut::::with_capacity(LEN); values.extend((0..LEN as i64).map(|i| base + i)); - PrimitiveArray::new( - values.freeze(), - vortex_array::validity::Validity::NonNullable, - ) + PrimitiveArray::new(values.freeze(), Validity::NonNullable) } fn nullable_array(base: i64, null_every: usize) -> PrimitiveArray { diff --git a/vortex-array/benches/scalar_at_struct.rs b/vortex-array/benches/scalar_at_struct.rs index 3921f0a2366..8dca44e408a 100644 --- a/vortex-array/benches/scalar_at_struct.rs +++ b/vortex-array/benches/scalar_at_struct.rs @@ -13,6 +13,7 @@ use rand::rngs::StdRng; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::StructArray; use vortex_array::dtype::FieldNames; use vortex_array::validity::Validity; @@ -20,13 +21,14 @@ use vortex_buffer::Buffer; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } const ARRAY_SIZE: usize = 100_000; const NUM_ACCESSES: usize = 1000; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(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 ee9aba1055d..ef004d2e270 100644 --- a/vortex-array/benches/scalar_subtract.rs +++ b/vortex-array/benches/scalar_subtract.rs @@ -13,18 +13,21 @@ use rand::rngs::StdRng; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; use vortex_array::builtins::ArrayBuiltins; +use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_buffer::Buffer; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); #[divan::bench] fn scalar_subtract(bencher: Bencher) { @@ -50,11 +53,7 @@ fn scalar_subtract(bencher: Bencher) { chunked .clone() .binary( - ConstantArray::new( - vortex_array::scalar::Scalar::from(to_subtract), - chunked.len(), - ) - .into_array(), + ConstantArray::new(Scalar::from(to_subtract), chunked.len()).into_array(), Operator::Sub, ) .unwrap() diff --git a/vortex-array/benches/take_filter.rs b/vortex-array/benches/take_filter.rs index b6ec1b4040c..96852a91966 100644 --- a/vortex-array/benches/take_filter.rs +++ b/vortex-array/benches/take_filter.rs @@ -28,6 +28,7 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::validity::Validity; @@ -38,6 +39,7 @@ use vortex_mask::Mask; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } @@ -57,7 +59,7 @@ const INDEX_SEED: u64 = 43; const LIST_SIZE: usize = 4; const NULL_INDEX_INTERVAL: usize = 8; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(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 438faae4983..5dc491c28c5 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -19,16 +19,18 @@ use rand::rngs::StdRng; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::FixedSizeListArray; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(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 5dc0c04c325..d3b99e8ea1b 100644 --- a/vortex-array/benches/take_patches.rs +++ b/vortex-array/benches/take_patches.rs @@ -12,18 +12,20 @@ use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::ArrayRef; use vortex_array::IntoArray; -#[expect(deprecated)] -use vortex_array::ToCanonical as _; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::PrimitiveArray; use vortex_array::patches::Patches; use vortex_buffer::Buffer; +use vortex_error::VortexExpect; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); const BENCH_ARGS: &[(f64, f64)] = &[ // patches_sparsity, index_multiple @@ -54,8 +56,10 @@ fn take_search(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) bencher .with_inputs(|| (&patches, &indices, SESSION.create_execution_ctx())) .bench_refs(|(patches, indices, ctx)| { - #[expect(deprecated)] - let prim = indices.to_primitive(); + let prim = indices + .clone() + .execute::(ctx) + .vortex_expect("operation should succeed in benchmark"); patches.take_search(prim, false, ctx) }); } @@ -73,8 +77,10 @@ fn take_search_chunked(bencher: Bencher, (patches_sparsity, index_multiple): (f6 bencher .with_inputs(|| (&patches, &indices, SESSION.create_execution_ctx())) .bench_refs(|(patches, indices, ctx)| { - #[expect(deprecated)] - let prim = indices.to_primitive(); + let prim = indices + .clone() + .execute::(ctx) + .vortex_expect("operation should succeed in benchmark"); patches.take_search(prim, false, ctx) }); } @@ -92,8 +98,10 @@ fn take_map(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) { bencher .with_inputs(|| (&patches, &indices, SESSION.create_execution_ctx())) .bench_refs(|(patches, indices, ctx)| { - #[expect(deprecated)] - let prim = indices.to_primitive(); + let prim = indices + .clone() + .execute::(ctx) + .vortex_expect("operation should succeed in benchmark"); patches.take_map(prim, false, ctx) }); } @@ -105,15 +113,7 @@ fn fixture(len: usize, sparsity: f64, rng: &mut StdRng) -> Patches { .collect::>(); let sparse_len = indices.len(); let values = Buffer::from_iter((0..sparse_len).map(|x| x as u64)).into_array(); - Patches::new( - len, - 0, - indices.into_array(), - values, - // TODO(0ax1): handle chunk offsets - None, - ) - .unwrap() + Patches::new(len, 0, indices.into_array(), values, None).unwrap() } fn fixture_with_chunk_offsets(len: usize, sparsity: f64, rng: &mut StdRng) -> Patches { diff --git a/vortex-array/benches/take_primitive.rs b/vortex-array/benches/take_primitive.rs index 383b525dcc6..471363b99e4 100644 --- a/vortex-array/benches/take_primitive.rs +++ b/vortex-array/benches/take_primitive.rs @@ -17,11 +17,13 @@ use rand_distr::Zipf; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } @@ -31,7 +33,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(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); // --- DictArray canonicalization benchmarks --- diff --git a/vortex-array/benches/take_struct.rs b/vortex-array/benches/take_struct.rs index d49ce9e88a6..9f7e518cb3a 100644 --- a/vortex-array/benches/take_struct.rs +++ b/vortex-array/benches/take_struct.rs @@ -13,6 +13,7 @@ use rand::rngs::StdRng; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::StructArray; use vortex_array::dtype::FieldNames; use vortex_array::validity::Validity; @@ -20,10 +21,11 @@ use vortex_buffer::Buffer; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); const ARRAY_SIZE: usize = 100_000; const TAKE_SIZE: usize = 1000; diff --git a/vortex-array/benches/validity_is_valid.rs b/vortex-array/benches/validity_is_valid.rs new file mode 100644 index 00000000000..dc22fcb5478 --- /dev/null +++ b/vortex-array/benches/validity_is_valid.rs @@ -0,0 +1,71 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks the cost of checking array-backed [`Validity`] per element. +//! +//! For the `Validity::Array` variant, `Validity::execute_is_valid` runs a scalar lookup through +//! the compute stack on *every* call, and the (now deprecated) `is_valid` additionally spins up a +//! fresh `ExecutionCtx` per call. Either way, calling it in a `for i in 0..n` loop is +//! pathologically slow. The fix used by callers is to materialize the validity into a `Mask` once +//! (`execute_mask`) and then do cheap O(1) bit reads via `Mask::value`. This bench contrasts the two. +//! +//! Sizes are kept small because the per-element variant is intentionally the slow one we are +//! measuring. + +#![allow(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +const SIZES: &[usize] = &[256, 1024]; + +/// Build an array-backed validity (~10% nulls so it is `Validity::Array`). +fn array_validity(len: usize) -> Validity { + Validity::from(BitBuffer::from_iter( + (0..len).map(|i| !i.is_multiple_of(10)), + )) +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +/// Per-element validity check over array-backed validity (the antipattern). This mirrors the +/// deprecated `Validity::is_valid(i)`: a fresh `ExecutionCtx` plus a scalar lookup on every call. +#[divan::bench(args = SIZES)] +fn is_valid_per_element(bencher: Bencher, len: usize) { + let validity = array_validity(len); + bencher + .with_inputs(|| (&validity, SESSION.create_execution_ctx())) + .bench_refs(|(validity, ctx)| { + let mut count = 0usize; + for i in 0..len { + count += validity.execute_is_valid(i, ctx).unwrap() as usize; + } + count + }); +} + +/// Materialize the validity into a `Mask` once, then read bits (the fix). +#[divan::bench(args = SIZES)] +fn execute_mask_then_value(bencher: Bencher, len: usize) { + let validity = array_validity(len); + bencher + .with_inputs(|| (&validity, SESSION.create_execution_ctx())) + .bench_refs(|(validity, ctx)| { + let mask = validity.execute_mask(len, ctx).unwrap(); + let mut count = 0usize; + for i in 0..len { + count += mask.value(i) as usize; + } + count + }); +} diff --git a/vortex-array/benches/varbinview_compact.rs b/vortex-array/benches/varbinview_compact.rs index 99243ce3606..80f7e18ec30 100644 --- a/vortex-array/benches/varbinview_compact.rs +++ b/vortex-array/benches/varbinview_compact.rs @@ -1,20 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::LazyLock; + use divan::Bencher; +use mimalloc::MiMalloc; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::ArrayRef; use vortex_array::IntoArray; -#[expect(deprecated)] -use vortex_array::ToCanonical as _; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::VarBinViewArray; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_buffer::Buffer; use vortex_error::VortexExpect; +use vortex_session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; fn main() { divan::main(); @@ -28,6 +35,8 @@ const ARGS: &[(usize, usize)] = &[ (1 << 14, 90), ]; +static SESSION: LazyLock = LazyLock::new(array_session); + #[divan::bench(args = ARGS)] fn compact(bencher: Bencher, args: (usize, usize)) { compact_impl(bencher, args); @@ -46,31 +55,39 @@ fn compact_impl(bencher: Bencher, (output_size, utilization_pct): (usize, usize) .into_array() .take(indices) .vortex_expect("operation should succeed in benchmark"); - #[expect(deprecated)] - let array = taken.to_varbinview(); - - bencher.with_inputs(|| &array).bench_refs(|array| { - array - .compact_buffers() - .vortex_expect("operation should succeed in benchmark") - }) + let mut ctx = SESSION.create_execution_ctx(); + let array = taken + .execute::(&mut ctx) + .vortex_expect("operation should succeed in benchmark"); + + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| { + array + .compact_buffers(ctx) + .vortex_expect("operation should succeed in benchmark") + }) } fn compact_sliced_impl(bencher: Bencher, (output_size, utilization_pct): (usize, usize)) { + let mut ctx = SESSION.create_execution_ctx(); let base_size = (output_size * 100) / utilization_pct; let base_array = build_varbinview_fixture(base_size); let sliced = base_array .into_array() .slice(0..output_size) .vortex_expect("slice should succeed"); - #[expect(deprecated)] - let array = sliced.to_varbinview(); - - bencher.with_inputs(|| &array).bench_refs(|array| { - array - .compact_buffers() - .vortex_expect("operation should succeed in benchmark") - }) + let array = sliced + .execute::(&mut ctx) + .vortex_expect("operation should succeed in benchmark"); + + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| { + array + .compact_buffers(ctx) + .vortex_expect("operation should succeed in benchmark") + }) } /// Creates a base VarBinViewArray with mix of inlined and outlined strings. diff --git a/vortex-array/benches/varbinview_zip.rs b/vortex-array/benches/varbinview_zip.rs index e010291e9dc..b2d55b1b60d 100644 --- a/vortex-array/benches/varbinview_zip.rs +++ b/vortex-array/benches/varbinview_zip.rs @@ -9,6 +9,7 @@ use divan::Bencher; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::VarBinViewArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; @@ -17,10 +18,11 @@ use vortex_mask::Mask; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); /// Benchmarks zip on VarBinView arrays with a highly fragmented mask (worst case for per-slice lookup paths). #[divan::bench] 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 c3ee7f5d4e7..ada9e6eb2e5 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -247,7 +247,7 @@ impl AggregateFnVTable for IsSorted { | DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) - | DType::Union(_) + | DType::Union(..) | DType::Variant(..) | DType::Extension(_) => None, DType::Bool(_) @@ -264,7 +264,7 @@ impl AggregateFnVTable for IsSorted { | DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) - | DType::Union(_) + | DType::Union(..) | DType::Variant(..) | DType::Extension(_) => None, DType::Bool(_) @@ -668,16 +668,14 @@ mod tests { // Tests migrated from arrays/decimal/compute/is_sorted.rs #[test] fn test_decimal_is_sorted() -> VortexResult<()> { - use arrow_array::types::Decimal128Type; - use arrow_cast::parse::parse_decimal; - use crate::arrays::DecimalArray; use crate::dtype::DecimalDType; 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())?; + // "100.00" and "200.00" at scale 2. + let i100 = 10_000i128; + let i200 = 20_000i128; let sorted = buffer![i100, i200, i200]; let unsorted = buffer![i200, i100, i200]; @@ -693,17 +691,14 @@ mod tests { #[test] fn test_decimal_is_strict_sorted() -> VortexResult<()> { - use arrow_array::types::Decimal128Type; - use arrow_cast::parse::parse_decimal; - use crate::arrays::DecimalArray; use crate::dtype::DecimalDType; 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())?; - let i300 = parse_decimal::("300.00", dtype.precision(), dtype.scale())?; + // "100.00", "200.00" and "300.00" at scale 2. + let i100 = 10_000i128; + let i200 = 20_000i128; + let i300 = 30_000i128; let strict_sorted = buffer![i100, i200, i300]; let sorted = buffer![i100, i200, i200]; diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index b421d69a966..20b7f15834c 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -18,10 +18,16 @@ use crate::aggregate_fn::combined::CombinedOptions; use crate::aggregate_fn::combined::PairOptions; use crate::aggregate_fn::fns::count::Count; use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum::sum_decimal_dtype; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::MAX_PRECISION; +use crate::dtype::MAX_SCALE; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::dtype::i256; +use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::scalar_fn::fns::operators::Operator; @@ -45,10 +51,7 @@ pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { /// /// Implemented as `Sum / Count` via [`BinaryCombined`]. /// -/// Coercion / return type: -/// - Booleans and primitive numeric types are coerced to `f64` and the result -/// is a nullable `f64`. -/// - Decimals are kept as decimals but not implemented currently +/// Booleans and primitive numeric types are cast to f64. Decimals stay decimals. #[derive(Clone, Debug)] pub struct Mean; @@ -88,18 +91,18 @@ impl BinaryCombined for Mean { } fn finalize(&self, sum: ArrayRef, count: ArrayRef) -> VortexResult { - let target = match sum.dtype() { - DType::Decimal(..) => sum.dtype().with_nullability(Nullability::Nullable), - _ => DType::Primitive(PType::F64, Nullability::Nullable), - }; + if let DType::Decimal(..) = sum.dtype() { + vortex_bail!("grouped mean over decimals is not yet supported"); + } + let target = DType::Primitive(PType::F64, Nullability::Nullable); let sum_cast = sum.cast(target.clone())?; let count_cast = count.cast(target)?; sum_cast.binary(count_cast, Operator::Div) } fn finalize_scalar(&self, left_scalar: Scalar, right_scalar: Scalar) -> VortexResult { - if let DType::Decimal(..) = left_scalar.dtype() { - vortex_bail!("mean::finalize_scalar not yet implemented for decimal inputs"); + if let DType::Decimal(decimal_dtype, _) = *left_scalar.dtype() { + return finalize_decimal_scalar(&left_scalar, &right_scalar, decimal_dtype); } let target = DType::Primitive(PType::F64, Nullability::Nullable); @@ -140,13 +143,12 @@ impl BinaryCombined for Mean { /// - Bool stays as bool — `Sum` has a native bool path and bool → f64 isn't /// currently a direct cast in vortex. /// - Primitive numerics → `f64` so the sum and finalize work without overflow. +/// - Decimals stay as decimals fn coerced_input_dtype(input_dtype: &DType) -> Option { match input_dtype { DType::Bool(_) => Some(input_dtype.clone()), DType::Primitive(_, n) => Some(DType::Primitive(PType::F64, *n)), - DType::Decimal(..) => { - unimplemented!("mean is not implemented for decimals yet") - } + DType::Decimal(..) => Some(input_dtype.clone()), _ => None, } } @@ -156,13 +158,59 @@ fn mean_output_dtype(input_dtype: &DType) -> Option { DType::Bool(_) | DType::Primitive(..) => { Some(DType::Primitive(PType::F64, Nullability::Nullable)) } - DType::Decimal(..) => { - unimplemented!("mean for decimals is not yet implemented"); - } + DType::Decimal(decimal_dtype, _) => Some(DType::Decimal( + mean_decimal_dtype(&sum_decimal_dtype(decimal_dtype)), + Nullability::Nullable, + )), _ => None, } } +/// mean() output decimal type mimicking Spark/DataFusion/MySQL: decimal(p+4, s+4) +fn mean_decimal_dtype(sum: &DecimalDType) -> DecimalDType { + DecimalDType::new( + u8::min(MAX_PRECISION, sum.precision().saturating_sub(6)), + i8::min(MAX_SCALE, sum.scale() + 4), + ) +} + +fn finalize_decimal_scalar( + sum: &Scalar, + count: &Scalar, + sum_decimal: DecimalDType, +) -> VortexResult { + let target_decimal_dtype = mean_decimal_dtype(&sum_decimal); + let target_dtype = DType::Decimal(target_decimal_dtype, Nullability::Nullable); + + // overflow + let Some(sum_value) = sum.as_decimal().decimal_value() else { + return Ok(Scalar::null(target_dtype)); + }; + // empty input + let count = count.as_primitive().typed_value::().unwrap_or(0); + if count == 0 { + return Ok(Scalar::null(target_dtype)); + } + + let Ok(sum) = DecimalValue::rescale_i256( + sum_value.as_i256(), + sum_decimal.scale(), + target_decimal_dtype.scale(), + ) else { + return Ok(Scalar::null(target_dtype)); + }; + let mean = sum / i256::from_i128(i128::from(count)); + + let Ok(mean) = DecimalValue::try_from_i256(mean, target_decimal_dtype) else { + return Ok(Scalar::null(target_dtype)); + }; + Ok(Scalar::decimal( + mean, + target_decimal_dtype, + Nullability::Nullable, + )) +} + #[cfg(test)] mod tests { use vortex_buffer::buffer; @@ -175,7 +223,9 @@ mod tests { use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; + use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; + use crate::dtype::DecimalDType; use crate::validity::Validity; #[test] @@ -273,6 +323,73 @@ mod tests { Ok(()) } + #[test] + fn mean_decimal() -> VortexResult<()> { + let dtype = DecimalDType::new(6, 2); + let array = + DecimalArray::new(buffer![100i32, 200, 300], dtype, Validity::NonNullable).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let result = mean(&array, &mut ctx)?; + assert_eq!( + result.dtype(), + &DType::Decimal(DecimalDType::new(10, 6), Nullability::Nullable) + ); + // mean(1.00, 2.00, 3.00) = 2.000000 + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(2_000_000))) + ); + Ok(()) + } + + #[test] + fn mean_decimal_null() -> VortexResult<()> { + let dtype = DecimalDType::new(6, 2); + let validity = Validity::from_iter([true, false, true]); + let array = DecimalArray::new(buffer![150i32, 0, 450], dtype, validity).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let result = mean(&array, &mut ctx)?; + // mean(1.50, 4.50) = 3.000000 + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(3_000_000))) + ); + Ok(()) + } + + #[test] + fn mean_decimal_chunked() -> VortexResult<()> { + let dtype = DecimalDType::new(6, 2); + let validity = Validity::NonNullable; + let chunk1 = DecimalArray::new(buffer![100i32, 200], dtype, validity.clone()).into_array(); + let chunk2 = DecimalArray::new(buffer![300i32, 400, 500], dtype, validity).into_array(); + let dtype = chunk1.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk1, chunk2], dtype)?; + let mut ctx = array_session().create_execution_ctx(); + let result = mean(&chunked.into_array(), &mut ctx)?; + // mean(1.00, 2.00, 3.00, 4.00, 5.00) = 3.000000 + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(3_000_000))) + ); + Ok(()) + } + + #[test] + fn mean_decimal_33() -> VortexResult<()> { + let dtype = DecimalDType::new(6, 2); + let buf = buffer![100i32, 0, 0]; + let array = DecimalArray::new(buf, dtype, Validity::NonNullable).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let result = mean(&array, &mut ctx)?; + // mean(1.00, 0.00, 0.00) = 1/3 => 0.333333 + assert_eq!( + result.as_decimal().decimal_value(), + Some(DecimalValue::I256(i256::from_i128(333_333))) + ); + Ok(()) + } + #[test] fn mean_multi_batch() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 001b0f61657..010371fee4f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -77,6 +77,16 @@ pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { #[derive(Clone, Debug)] pub struct Sum; +// Both Spark and DataFusion use this heuristic. +// - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66 +// - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188 +pub(crate) fn sum_decimal_dtype(input: &DecimalDType) -> DecimalDType { + DecimalDType::new( + u8::min(MAX_PRECISION, input.precision() + 10), + input.scale(), + ) +} + impl AggregateFnVTable for Sum { type Options = NumericalAggregateOpts; type Partial = SumPartial; @@ -118,14 +128,7 @@ impl AggregateFnVTable for Sum { } }, DType::Decimal(decimal_dtype, _) => { - // Both Spark and DataFusion use this heuristic. - // - https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala#L66 - // - https://github.com/apache/datafusion/blob/4153adf2c0f6e317ef476febfdc834208bd46622/datafusion/functions-aggregate/src/sum.rs#L188 - let precision = u8::min(MAX_PRECISION, decimal_dtype.precision() + 10); - DType::Decimal( - DecimalDType::new(precision, decimal_dtype.scale()), - Nullable, - ) + DType::Decimal(sum_decimal_dtype(decimal_dtype), Nullable) } // Unsupported types _ => return None, 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 d73e670f50f..1d04a074d6f 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 @@ -289,7 +289,9 @@ fn supports_uncompressed_size_in_bytes(dtype: &DType) -> bool { DType::Struct(fields, _) => fields .fields() .all(|field| supports_uncompressed_size_in_bytes(&field)), - DType::Union(_) => todo!("TODO(connor)[Union]: unimplemented"), + DType::Union(variants, _) => variants + .variants() + .all(|variant| supports_uncompressed_size_in_bytes(&variant)), DType::Variant(_) => false, DType::Extension(ext_dtype) => { supports_uncompressed_size_in_bytes(ext_dtype.storage_dtype()) @@ -313,6 +315,7 @@ pub(crate) fn packed_bit_buffer_size_in_bytes(len: usize) -> VortexResult { #[cfg(test)] mod tests { use vortex_buffer::buffer; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -355,9 +358,12 @@ mod tests { fn materialized_uncompressed_size_in_bytes(array: &ArrayRef) -> u64 { let mut builder = builder_with_capacity(array.dtype(), array.len()); - unsafe { - builder.extend_from_array_unchecked(array); - } + array + .append_to_builder( + builder.as_mut(), + &mut array_session().create_execution_ctx(), + ) + .vortex_expect("appended"); builder.finish().nbytes() } diff --git a/vortex-array/src/aggregate_fn/plugin.rs b/vortex-array/src/aggregate_fn/plugin.rs index b7ff8b893ac..baf6f303466 100644 --- a/vortex-array/src/aggregate_fn/plugin.rs +++ b/vortex-array/src/aggregate_fn/plugin.rs @@ -10,6 +10,7 @@ use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; +use crate::dtype::DType; /// Reference-counted pointer to an aggregate function plugin. pub type AggregateFnPluginRef = Arc; @@ -28,6 +29,19 @@ pub trait AggregateFnPlugin: 'static + Send + Sync { /// Deserialize an aggregate function from serialized metadata. fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult; + + /// The default zone statistic (per-chunk) for a column of `input_dtype`, or `None` if the + /// dtype is not supported (or this aggregate is not a zone statistic at all). + /// + /// This is how a registered aggregate volunteers itself as a per-chunk statistic: when a + /// zoned layout writer opens a column, it collects every plugin's default via + /// [`AggregateFnSession::zone_stat_defaults`], so new statistics can be added without the + /// writer knowing about them. + /// + /// [`AggregateFnSession::zone_stat_defaults`]: crate::aggregate_fn::session::AggregateFnSession::zone_stat_defaults + fn zone_stat_default(&self, _input_dtype: &DType) -> Option { + None + } } impl std::fmt::Debug for dyn AggregateFnPlugin { @@ -51,4 +65,8 @@ impl AggregateFnPlugin for V { let options = AggregateFnVTable::deserialize(self, metadata, session)?; Ok(AggregateFn::new(self.clone(), options).erased()) } + + fn zone_stat_default(&self, input_dtype: &DType) -> Option { + AggregateFnVTable::zone_stat_default(self, input_dtype) + } } diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index 50c585233f0..92fac87892a 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -195,9 +195,8 @@ mod tests { #[test] fn unknown_aggregate_fn_id_allow_unknown() { - let session = VortexSession::empty() - .with::() - .allow_unknown(); + let session = VortexSession::empty().with::(); + session.allow_unknown(); let proto = pb::AggregateFn { id: "vortex.test.foreign_aggregate".to_string(), diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index c496f7a9bf8..72ef78a400c 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -10,6 +10,7 @@ use vortex_session::SessionVar; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnPluginRef; +use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::fns::all_nan::AllNan; use crate::aggregate_fn::fns::all_non_distinct::AllNonDistinct; @@ -44,6 +45,7 @@ use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate; use crate::arrays::dict::compute::is_constant::DictIsConstantKernel; use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel; use crate::arrays::dict::compute::min_max::DictMinMaxKernel; +use crate::dtype::DType; /// Session state for aggregate functions and encoding-specific aggregate kernels. /// @@ -134,6 +136,22 @@ impl AggregateFnSession { self.registry.insert(id, pluginref); } + /// The default per-chunk zone statistics for a column of `input_dtype`, collected from every + /// registered aggregate's `zone_stat_default`. + /// + /// Each call scans the whole plugin registry, so this is intended to be called once per + /// column when a zoned writer is opened, not per chunk or per row. + pub fn zone_stat_defaults(&self, input_dtype: &DType) -> Vec { + self.registry.read(|registry| { + let mut fns: Vec = registry + .values() + .filter_map(|plugin| plugin.zone_stat_default(input_dtype)) + .collect(); + fns.sort_by_key(|f| f.id()); + fns + }) + } + /// Returns the aggregate kernel registered for `array_id` and `agg_fn_id`, if any. /// /// Lookup first checks for a kernel registered for the exact aggregate function, then falls diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 49b28dd26d7..3f0a4cf567c 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -93,6 +93,12 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// Returns `None` if the aggregate function cannot be applied to the input dtype. fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option; + /// If this aggregate should be computed as a default zone statistic for `input_dtype`, return + /// the bound aggregate to store. Default: not a zone-map default. + fn zone_stat_default(&self, _input_dtype: &DType) -> Option { + None + } + /// DType of the intermediate partial accumulator state. /// /// Use a struct dtype when multiple fields are needed diff --git a/vortex-array/src/arc_swap_map.rs b/vortex-array/src/arc_swap_map.rs index 73108ae5d57..1788f704feb 100644 --- a/vortex-array/src/arc_swap_map.rs +++ b/vortex-array/src/arc_swap_map.rs @@ -28,7 +28,7 @@ use vortex_utils::aliases::hash_map::HashMap; /// 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 { +pub struct ArcSwapMap { inner: Arc>>, } @@ -56,7 +56,7 @@ impl Debug for ArcSwapMap { impl ArcSwapMap { /// Return the currently published map snapshot. - pub(crate) fn snapshot(&self) -> Arc> { + pub fn snapshot(&self) -> Arc> { self.inner.load_full() } @@ -64,7 +64,7 @@ impl ArcSwapMap { /// /// Every lookup inside `f` observes the same snapshot, which matters when a /// single logical read consults more than one key. - pub(crate) fn read(&self, f: impl FnOnce(&HashMap) -> R) -> R { + pub fn read(&self, f: impl FnOnce(&HashMap) -> R) -> R { f(&self.inner.load()) } @@ -87,7 +87,7 @@ impl ArcSwapMap { impl ArcSwapMap { /// Return a clone of the value stored under `key`, if present. - pub(crate) fn get(&self, key: &Q) -> Option + pub fn get(&self, key: &Q) -> Option where K: Borrow, Q: Eq + Hash + ?Sized, @@ -96,7 +96,7 @@ impl ArcSwapMap { } /// Insert `value` under `key`, replacing any existing value. - pub(crate) fn insert(&self, key: K, value: V) + pub fn insert(&self, key: K, value: V) where K: Clone, { @@ -111,7 +111,7 @@ impl ArcSwapMap> { /// /// Each key maps to an immutable `Arc<[T]>`; appending rebuilds that slice /// copy-on-write so existing readers keep their previous snapshot. - pub(crate) fn extend(&self, key: K, values: &[T]) { + pub fn extend(&self, key: K, values: &[T]) { self.modify(|map| { let merged: Arc<[T]> = match map.get(&key) { Some(existing) => existing.iter().chain(values).cloned().collect(), @@ -123,7 +123,7 @@ impl ArcSwapMap> { /// Append a single `value` to the list stored under `key`, creating it if /// absent. - pub(crate) fn push(&self, key: K, value: T) { + pub fn push(&self, key: K, value: T) { self.extend(key, &[value]); } } diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 7fb3d7baf01..c3dc2e0eed7 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -26,7 +26,6 @@ use crate::Canonical; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VTable; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::sum::sum; @@ -50,6 +49,7 @@ use crate::dtype::DType; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProviderExt; +use crate::legacy_session; use crate::matcher::Matcher; use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; @@ -269,8 +269,9 @@ impl ArrayRef { note = "Use `execute_scalar` instead, which allows passing an execution context for more \ efficient execution when fetching multiple scalars from the same array." )] + #[allow(clippy::disallowed_methods)] pub fn scalar_at(&self, index: usize) -> VortexResult { - self.execute_scalar(index, &mut LEGACY_SESSION.create_execution_ctx()) + self.execute_scalar(index, &mut legacy_session().create_execution_ctx()) } /// Execute the array to extract a scalar at the given index. @@ -360,8 +361,9 @@ impl ArrayRef { /// Returns the canonical representation of the array. #[deprecated(note = "use `array.execute::(ctx)` instead")] + #[allow(clippy::disallowed_methods)] pub fn into_canonical(self) -> VortexResult { - self.execute(&mut LEGACY_SESSION.create_execution_ctx()) + self.execute(&mut legacy_session().create_execution_ctx()) } /// Returns the canonical representation of the array. diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index 0d3bbe9cae3..3eea66c6a68 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -158,7 +158,7 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { /// Execute the array by taking a single encoding-specific execution step. /// /// This is the checked entry point. If the encoding reports - /// [`ExecutionStep::Done`](crate::ExecutionStep::Done), implementations must validate that the + /// [`ExecutionStep::Done`](ExecutionStep::Done), implementations must validate that the /// returned array preserves this array's logical `len` and `dtype`, and must transfer this /// array's statistics to the returned array. fn execute(&self, this: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; diff --git a/vortex-array/src/array/typed.rs b/vortex-array/src/array/typed.rs index b88f3792a84..4769073b82c 100644 --- a/vortex-array/src/array/typed.rs +++ b/vortex-array/src/array/typed.rs @@ -23,12 +23,12 @@ use crate::ArrayRef; use crate::ArraySlots; use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; use crate::dtype::DType; +use crate::legacy_session; use crate::stats::ArrayStats; use crate::stats::StatsSet; use crate::stats::StatsSetRef; @@ -382,9 +382,10 @@ impl Array { note = "Use `execute_scalar` instead, which allows passing an execution context for more \ efficient execution when fetching multiple scalars from the same array." )] + #[allow(clippy::disallowed_methods)] pub fn scalar_at(&self, index: usize) -> VortexResult { self.inner - .execute_scalar(index, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(index, &mut legacy_session().create_execution_ctx()) } /// Execute the array to extract a scalar at the given index. diff --git a/vortex-array/src/array/vtable/mod.rs b/vortex-array/src/array/vtable/mod.rs index 617b211a4c2..344aac64a88 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -193,8 +193,7 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { .clone() .execute::(ctx)? .into_array(); - builder.extend_from_array(&canonical); - Ok(()) + canonical.append_to_builder(builder, ctx) } /// Returns the name of the slot at the given index. diff --git a/vortex-array/src/arrays/arbitrary.rs b/vortex-array/src/arrays/arbitrary.rs index 035c23fd4db..5d5d9f60bf9 100644 --- a/vortex-array/src/arrays/arbitrary.rs +++ b/vortex-array/src/arrays/arbitrary.rs @@ -15,11 +15,10 @@ use vortex_error::VortexExpect; use crate::ArrayRef; use crate::IntoArray; -#[expect(deprecated)] -use crate::ToCanonical as _; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::NullArray; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinArray; @@ -130,9 +129,8 @@ fn random_array_chunk( PType::I32 => random_primitive::(u, *n, chunk_len), PType::I64 => random_primitive::(u, *n, chunk_len), PType::F16 => { - #[expect(deprecated)] let prim = random_primitive::(u, *n, chunk_len)? - .to_primitive() + .as_::() .reinterpret_cast(PType::F16) .into_array(); Ok(prim) diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index 8b83d0f8128..2dd08c55832 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -4,7 +4,6 @@ use std::fmt::Display; use std::fmt::Formatter; -use arrow_array::BooleanArray; use smallvec::smallvec; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMeta; @@ -338,14 +337,16 @@ impl FromIterator for BoolArray { impl FromIterator> for BoolArray { fn from_iter>>(iter: I) -> Self { - let (buffer, nulls) = BooleanArray::from_iter(iter).into_parts(); + let iter = iter.into_iter(); + let capacity = iter.size_hint().0; + let mut bits = BitBufferMut::with_capacity(capacity); + let mut validity = BitBufferMut::with_capacity(capacity); + for value in iter { + bits.append(value.unwrap_or_default()); + validity.append(value.is_some()); + } - BoolArray::new( - BitBuffer::from(buffer), - nulls - .map(|n| Validity::from(BitBuffer::from(n.into_inner()))) - .unwrap_or(Validity::AllValid), - ) + BoolArray::new(bits.freeze(), Validity::from(validity.freeze())) } } diff --git a/vortex-array/src/arrays/bool/compute/cast.rs b/vortex-array/src/arrays/bool/compute/cast.rs index 3b47ce62c8f..36849f55a18 100644 --- a/vortex-array/src/arrays/bool/compute/cast.rs +++ b/vortex-array/src/arrays/bool/compute/cast.rs @@ -123,7 +123,7 @@ mod tests { #[case(BoolArray::from_iter(vec![true]))] #[case(BoolArray::from_iter(vec![false, false]))] fn test_cast_bool_conformance(#[case] array: BoolArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[rstest] diff --git a/vortex-array/src/arrays/bool/compute/fill_null.rs b/vortex-array/src/arrays/bool/compute/fill_null.rs index 59b9ca0b717..5fa18299106 100644 --- a/vortex-array/src/arrays/bool/compute/fill_null.rs +++ b/vortex-array/src/arrays/bool/compute/fill_null.rs @@ -48,11 +48,11 @@ mod tests { use vortex_buffer::bitbuffer; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::builtins::ArrayBuiltins; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -62,16 +62,17 @@ mod tests { #[case(true, bitbuffer![true, true, false, true])] #[case(false, bitbuffer![true, false, false, false])] fn bool_fill_null(#[case] fill_value: bool, #[case] expected: BitBuffer) { + let mut ctx = array_session().create_execution_ctx(); let bool_array = BoolArray::new( BitBuffer::from_iter([true, true, false, false]), Validity::from_iter([true, false, true, false]), ); - #[expect(deprecated)] let non_null_array = bool_array .into_array() .fill_null(Scalar::from(fill_value)) .unwrap() - .to_bool(); + .execute::(&mut ctx) + .unwrap(); assert_eq!(non_null_array.to_bit_buffer(), expected); assert_eq!( non_null_array.dtype(), diff --git a/vortex-array/src/arrays/bool/compute/filter.rs b/vortex-array/src/arrays/bool/compute/filter.rs index f463f27f437..98f52ed11ee 100644 --- a/vortex-array/src/arrays/bool/compute/filter.rs +++ b/vortex-array/src/arrays/bool/compute/filter.rs @@ -3,6 +3,7 @@ use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; +use vortex_buffer::CpuKernel; use vortex_buffer::get_bit; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -90,14 +91,19 @@ pub fn filter_bitbuffer_by_mask( mask_buf: &BitBuffer, true_count: usize, ) -> BitBuffer { - #[cfg(target_arch = "x86_64")] - { - if std::arch::is_x86_feature_detected!("bmi2") { - // SAFETY: BMI2 confirmed available; the inner function is compiled with BMI2. - return unsafe { filter_pext_bmi2(src, mask_buf, true_count) }; + type FilterKernel = unsafe fn(&BitBuffer, &BitBuffer, usize) -> BitBuffer; + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + if std::arch::is_x86_feature_detected!("bmi2") { + return filter_pext_bmi2; + } } - } - filter_pext_fallback(src, mask_buf, true_count) + filter_pext_fallback + }); + // SAFETY: the selector only returns kernels that are safe or whose required CPU + // features were probed before selection. + unsafe { KERNEL.get()(src, mask_buf, true_count) } } /// BMI2-native filter: entire function compiled with BMI2+POPCNT enabled. @@ -366,7 +372,10 @@ mod tests { #[case(BoolArray::from_iter((0..100).map(|i| i % 2 == 0)))] #[case(BoolArray::from_iter((0..1024).map(|i| i % 3 != 0)))] fn test_filter_bool_conformance(#[case] array: BoolArray) { - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[cfg(target_arch = "x86_64")] diff --git a/vortex-array/src/arrays/bool/compute/mask.rs b/vortex-array/src/arrays/bool/compute/mask.rs index 6fd278f2d3d..783c933fe5a 100644 --- a/vortex-array/src/arrays/bool/compute/mask.rs +++ b/vortex-array/src/arrays/bool/compute/mask.rs @@ -29,6 +29,8 @@ mod test { use rstest::rstest; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::compute::conformance::mask::test_mask_conformance; @@ -39,6 +41,9 @@ mod test { #[case(BoolArray::from_iter([false, false]))] #[case(BoolArray::from_iter((0..100).map(|i| i % 2 == 0)))] fn test_mask_bool_conformance(#[case] array: BoolArray) { - test_mask_conformance(&array.into_array()); + test_mask_conformance( + &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 abff3527bfe..8469c7e25c2 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -88,8 +88,6 @@ mod test { use vortex_buffer::buffer; use crate::IntoArray as _; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; @@ -110,11 +108,11 @@ mod test { Some(false), ]); - #[expect(deprecated)] let b = reference .take(buffer![0, 3, 4].into_array()) .unwrap() - .to_bool(); + .execute::(&mut ctx) + .unwrap(); assert_eq!( b.to_bit_buffer(), BoolArray::from_iter([Some(false), None, Some(false)]).to_bit_buffer() @@ -190,6 +188,9 @@ mod test { #[case(BoolArray::from_iter([true, false]))] #[case(BoolArray::from_iter([true]))] fn test_take_bool_conformance(#[case] array: BoolArray) { - test_take_conformance(&array.into_array()); + test_take_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/bool/compute/zip.rs b/vortex-array/src/arrays/bool/compute/zip.rs index 80e7de80c82..ad311572395 100644 --- a/vortex-array/src/arrays/bool/compute/zip.rs +++ b/vortex-array/src/arrays/bool/compute/zip.rs @@ -34,7 +34,7 @@ impl ZipKernel for Bool { }; // Null mask entries select `if_false`, matching `Zip`'s SQL ELSE semantics. - let mask = mask.try_to_mask_fill_null_false(ctx)?; + let mask = mask.clone().null_as_false().execute(ctx)?; let mask_values = match &mask { // Defer trivial masks to the generic zip, which just casts the surviving side. Mask::AllTrue(_) | Mask::AllFalse(_) => return Ok(None), diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index 2420642d972..8a9cd961da8 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -28,7 +28,6 @@ use crate::builders::BoolBuilder; use crate::dtype::DType; use crate::serde::ArrayChildren; use crate::validity::Validity; -mod canonical; mod kernel; mod operations; mod validity; @@ -198,12 +197,10 @@ impl VTable for Bool { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> 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(()) + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for Bool requires a BoolBuilder"); + }; + builder.append_bool_array(&array.into_owned(), ctx) } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrays/bool/vtable/operations.rs b/vortex-array/src/arrays/bool/vtable/operations.rs index 9bae2b9e19c..c29ab20331b 100644 --- a/vortex-array/src/arrays/bool/vtable/operations.rs +++ b/vortex-array/src/arrays/bool/vtable/operations.rs @@ -28,8 +28,6 @@ mod tests { use std::iter; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; @@ -38,9 +36,14 @@ mod tests { #[test] fn test_slice_hundred_elements() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter(iter::repeat_n(Some(true), 100)); - #[expect(deprecated)] - let sliced_arr = arr.into_array().slice(8..16).unwrap().to_bool(); + let sliced_arr = arr + .into_array() + .slice(8..16) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_eq!(sliced_arr.len(), 8); assert_eq!(sliced_arr.to_bit_buffer().len(), 8); assert_eq!(sliced_arr.to_bit_buffer().offset(), 0); @@ -50,8 +53,12 @@ mod tests { 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(); + let sliced_arr = arr + .into_array() + .slice(1..4) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( sliced_arr, diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 54544ba1b4d..3963d6ba950 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -18,6 +18,8 @@ use vortex_error::vortex_bail; use crate::ArrayRef; use crate::ArraySlots; +use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray; use crate::array::Array; use crate::array::ArrayParts; @@ -194,7 +196,12 @@ impl Array { Ok(unsafe { Self::new_unchecked(chunks, dtype) }) } - pub fn rechunk(&self, target_bytesize: u64, target_rowsize: usize) -> VortexResult { + pub fn rechunk( + &self, + target_bytesize: u64, + target_rowsize: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { let mut new_chunks = Vec::new(); let mut chunks_to_combine = Vec::new(); let mut new_chunk_n_bytes = 0; @@ -207,12 +214,11 @@ impl Array { || new_chunk_n_elements + n_elements > target_rowsize) && !chunks_to_combine.is_empty() { - #[expect(deprecated)] let canonical = unsafe { Array::::new_unchecked(chunks_to_combine, self.dtype().clone()) } .into_array() - .to_canonical()? + .execute::(ctx)? .into_array(); new_chunks.push(canonical); @@ -231,11 +237,10 @@ impl Array { } if !chunks_to_combine.is_empty() { - #[expect(deprecated)] let canonical = unsafe { Array::::new_unchecked(chunks_to_combine, self.dtype().clone()) } .into_array() - .to_canonical()? + .execute::(ctx)? .into_array(); new_chunks.push(canonical); } @@ -298,7 +303,7 @@ mod test { ) .unwrap(); - let rechunked = chunked.rechunk(1 << 16, 1 << 16).unwrap(); + let rechunked = chunked.rechunk(1 << 16, 1 << 16, &mut ctx).unwrap(); assert_arrays_eq!(chunked, rechunked, &mut ctx); } @@ -312,7 +317,7 @@ mod test { ) .unwrap(); - let rechunked = chunked.rechunk(1 << 16, 1 << 16).unwrap(); + let rechunked = chunked.rechunk(1 << 16, 1 << 16, &mut ctx).unwrap(); assert_eq!(rechunked.nchunks(), 1); assert_arrays_eq!(chunked, rechunked, &mut ctx); @@ -330,7 +335,7 @@ mod test { ) .unwrap(); - let rechunked = chunked.rechunk(1 << 16, 5).unwrap(); + let rechunked = chunked.rechunk(1 << 16, 5, &mut ctx).unwrap(); assert_eq!(rechunked.nchunks(), 2); assert!(rechunked.iter_chunks().all(|c| c.len() < 5)); @@ -352,7 +357,7 @@ mod test { ) .unwrap(); - let rechunked = chunked.rechunk(1 << 16, 5).unwrap(); + let rechunked = chunked.rechunk(1 << 16, 5, &mut ctx).unwrap(); // greedy so should be: [0, 1, 2] [42, 42, 42, 42, 42, 42] [4, 5, 6, 7] [8, 9] assert_eq!(rechunked.nchunks(), 4); @@ -361,6 +366,7 @@ mod test { #[test] fn test_empty_chunks_all_valid() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Create chunks where some are empty but all non-empty chunks have all valid values let chunks = vec![ PrimitiveArray::new(buffer![1u64, 2, 3], Validity::AllValid).into_array(), @@ -373,18 +379,15 @@ 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 array_session().create_execution_ctx())?); - assert!( - !chunked - .into_array() - .all_invalid(&mut array_session().create_execution_ctx())? - ); + assert!(chunked.all_valid(&mut ctx)?); + assert!(!chunked.into_array().all_invalid(&mut ctx)?); Ok(()) } #[test] fn test_empty_chunks_all_invalid() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // 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(), @@ -397,18 +400,15 @@ mod test { 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 array_session().create_execution_ctx())?); - assert!( - chunked - .into_array() - .all_invalid(&mut array_session().create_execution_ctx())? - ); + assert!(!chunked.all_valid(&mut ctx)?); + assert!(chunked.into_array().all_invalid(&mut ctx)?); Ok(()) } #[test] fn test_empty_chunks_mixed_validity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Create chunks with mixed validity including empty chunks let chunks = vec![ PrimitiveArray::new(buffer![1u64, 2], Validity::AllValid).into_array(), @@ -421,12 +421,8 @@ mod test { ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?; // Should be neither all_valid nor all_invalid - assert!(!chunked.all_valid(&mut array_session().create_execution_ctx())?); - assert!( - !chunked - .into_array() - .all_invalid(&mut array_session().create_execution_ctx())? - ); + assert!(!chunked.all_valid(&mut ctx)?); + assert!(!chunked.into_array().all_invalid(&mut ctx)?); Ok(()) } diff --git a/vortex-array/src/arrays/chunked/compute/cast.rs b/vortex-array/src/arrays/chunked/compute/cast.rs index 80c6848b973..25e312de748 100644 --- a/vortex-array/src/arrays/chunked/compute/cast.rs +++ b/vortex-array/src/arrays/chunked/compute/cast.rs @@ -97,6 +97,6 @@ mod test { DType::Primitive(PType::U8, Nullability::NonNullable) ).unwrap().into_array())] fn test_cast_chunked_conformance(#[case] array: crate::ArrayRef) { - test_cast_conformance(&array); + test_cast_conformance(&array, &mut array_session().create_execution_ctx()); } } diff --git a/vortex-array/src/arrays/chunked/compute/filter.rs b/vortex-array/src/arrays/chunked/compute/filter.rs index 0e32543761d..5f50ce18387 100644 --- a/vortex-array/src/arrays/chunked/compute/filter.rs +++ b/vortex-array/src/arrays/chunked/compute/filter.rs @@ -203,6 +203,8 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::filter::test_filter_conformance; @@ -275,6 +277,9 @@ mod test { DType::Primitive(PType::I64, Nullability::NonNullable), ).unwrap())] fn test_filter_chunked_conformance(#[case] chunked: ChunkedArray) { - test_filter_conformance(&chunked.into_array()); + test_filter_conformance( + &chunked.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/chunked/compute/mask.rs b/vortex-array/src/arrays/chunked/compute/mask.rs index bec7b5e83d0..01170a232bb 100644 --- a/vortex-array/src/arrays/chunked/compute/mask.rs +++ b/vortex-array/src/arrays/chunked/compute/mask.rs @@ -45,6 +45,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::compute::conformance::mask::test_mask_conformance; @@ -80,6 +82,9 @@ mod test { DType::Primitive(PType::F32, Nullability::NonNullable), ).unwrap())] fn test_mask_chunked_conformance(#[case] chunked: ChunkedArray) { - test_mask_conformance(&chunked.into_array()); + test_mask_conformance( + &chunked.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 34f443f1ce1..1b867981236 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -126,8 +126,6 @@ mod test { use vortex_error::VortexResult; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; @@ -268,6 +266,7 @@ mod test { #[test] fn test_take_shuffled_large() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let nchunks: i32 = 100; let chunk_len: i32 = 1_000; let total = nchunks * chunk_len; @@ -297,8 +296,7 @@ mod test { let result = arr.take(indices_arr.into_array())?; // Verify every element. - #[expect(deprecated)] - let result = result.to_primitive(); + let result = result.execute::(&mut ctx)?; let result_vals = result.as_slice::(); for (pos, &idx) in indices.iter().enumerate() { assert_eq!( @@ -353,14 +351,20 @@ mod test { .clone(), ) .unwrap(); - test_take_conformance(&arr.into_array()); + test_take_conformance( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + ); // Test with nullable chunked array let a = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]); let b = PrimitiveArray::from_option_iter([Some(4i32), Some(5)]); let dtype = a.dtype().clone(); let arr = ChunkedArray::try_new(vec![a.into_array(), b.into_array()], dtype).unwrap(); - test_take_conformance(&arr.into_array()); + test_take_conformance( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + ); // Test with multiple identical chunks let chunk = buffer![10i32, 20, 30, 40, 50].into_array(); @@ -369,6 +373,9 @@ mod test { chunk.dtype().clone(), ) .unwrap(); - test_take_conformance(&arr.into_array()); + test_take_conformance( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/chunked/compute/zip.rs b/vortex-array/src/arrays/chunked/compute/zip.rs index 4037fd1dc49..dc95c194173 100644 --- a/vortex-array/src/arrays/chunked/compute/zip.rs +++ b/vortex-array/src/arrays/chunked/compute/zip.rs @@ -51,12 +51,11 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; + use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; @@ -65,6 +64,7 @@ mod tests { #[test] fn test_chunked_zip_aligns_across_boundaries() { + let mut ctx = array_session().create_execution_ctx(); let if_true = ChunkedArray::try_new( vec![ buffer![1i32, 2].into_array(), @@ -103,8 +103,7 @@ mod tests { assert_eq!(zipped.nchunks(), 4); let mut values: Vec = Vec::new(); for chunk in zipped.chunks() { - #[expect(deprecated)] - let primitive = chunk.to_primitive(); + let primitive = chunk.execute::(&mut ctx).unwrap(); values.extend_from_slice(primitive.as_slice::()); } assert_eq!(values, vec![1, 11, 3, 13, 5]); diff --git a/vortex-array/src/arrays/chunked/paired_chunks.rs b/vortex-array/src/arrays/chunked/paired_chunks.rs index 2145c88dbbd..c4a73a09518 100644 --- a/vortex-array/src/arrays/chunked/paired_chunks.rs +++ b/vortex-array/src/arrays/chunked/paired_chunks.rs @@ -121,9 +121,10 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; + use crate::arrays::PrimitiveArray; use crate::arrays::chunked::paired_chunks::PairedChunksExt; use crate::dtype::DType; use crate::dtype::Nullability; @@ -138,13 +139,20 @@ mod tests { left: &ChunkedArray, right: &ChunkedArray, ) -> VortexResult, Vec, std::ops::Range)>> { + let mut ctx = array_session().create_execution_ctx(); let mut result = Vec::new(); for pair in left.paired_chunks(right) { let pair = pair?; - #[expect(deprecated)] - let l: Vec = pair.left.to_primitive().as_slice::().to_vec(); - #[expect(deprecated)] - let r: Vec = pair.right.to_primitive().as_slice::().to_vec(); + let l: Vec = pair + .left + .execute::(&mut ctx)? + .as_slice::() + .to_vec(); + let r: Vec = pair + .right + .execute::(&mut ctx)? + .as_slice::() + .to_vec(); result.push((l, r, pair.pos)); } Ok(result) diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index 578b0a7c068..a32d03e207b 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -12,9 +12,11 @@ use vortex_session::VortexSession; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::ListArray; +use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; @@ -23,8 +25,6 @@ use crate::arrays::dict_test::gen_dict_primitive_chunks; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; use crate::builders::builder_with_capacity; -#[expect(deprecated)] -use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -32,7 +32,7 @@ use crate::dtype::PType::I32; use crate::executor::execute_into_builder; use crate::validity::Validity; -static SESSION: LazyLock = LazyLock::new(crate::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); fn chunked_array() -> ChunkedArray { ChunkedArray::try_new( @@ -375,33 +375,14 @@ pub fn pack_nested_structs() -> VortexResult<()> { dtype, )? .into_array(); - #[expect(deprecated)] - let canonical_struct = chunked.to_struct(); - #[expect(deprecated)] - let canonical_varbin = canonical_struct.unmasked_fields()[0].to_varbinview(); - #[expect(deprecated)] - let original_varbin = struct_array.unmasked_fields()[0].to_varbinview(); - 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); + let canonical_struct = chunked.execute::(&mut ctx)?; + let canonical_varbin = canonical_struct.unmasked_fields()[0] + .clone() + .execute::(&mut ctx)?; + let original_varbin = struct_array.unmasked_fields()[0] + .clone() + .execute::(&mut ctx)?; + assert_arrays_eq!(original_varbin, canonical_varbin, &mut ctx); Ok(()) } @@ -430,8 +411,12 @@ pub fn pack_nested_lists() { ), ); - #[expect(deprecated)] - let canon_values = chunked_list.unwrap().as_array().to_listview(); + let canon_values = chunked_list + .unwrap() + .as_array() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!( l1.execute_scalar(0, &mut ctx).unwrap(), diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 995d90bc0aa..12c3bbcd9e4 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -73,7 +73,7 @@ pub(super) fn _canonicalize( _ => { let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); array.array().append_to_builder(builder.as_mut(), ctx)?; - builder.finish_into_canonical() + builder.finish_into_canonical(ctx) } }) } diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 673855cfc0a..2eabd883742 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -22,14 +22,14 @@ use crate::EqMode; use crate::ExecutionCtx; use crate::ExecutionResult; use crate::IntoArray; -#[expect(deprecated)] -use crate::ToCanonical as _; +use crate::VortexSessionExecute; 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::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::chunked::ChunkedData; use crate::arrays::chunked::array::CHUNK_OFFSETS_SLOT; @@ -179,7 +179,7 @@ impl VTable for Chunked { metadata: &[u8], _buffers: &[BufferHandle], children: &dyn ArrayChildren, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult> { if !metadata.is_empty() { vortex_bail!( @@ -197,8 +197,11 @@ impl VTable for Chunked { &DType::Primitive(PType::U64, Nullability::NonNullable), nchunks + 1, )?; - #[expect(deprecated)] - let chunk_offsets_buf = chunk_offsets.to_primitive().to_buffer::(); + let mut ctx = session.create_execution_ctx(); + let chunk_offsets_buf = chunk_offsets + .clone() + .execute::(&mut ctx)? + .to_buffer::(); let chunk_offsets_usize = chunk_offsets_buf .iter() .copied() diff --git a/vortex-array/src/arrays/constant/arbitrary.rs b/vortex-array/src/arrays/constant/arbitrary.rs index c40bb5e82b7..37bde6414dd 100644 --- a/vortex-array/src/arrays/constant/arbitrary.rs +++ b/vortex-array/src/arrays/constant/arbitrary.rs @@ -16,15 +16,15 @@ pub struct ArbitraryConstantArray(pub ConstantArray); impl<'a> Arbitrary<'a> for ArbitraryConstantArray { fn arbitrary(u: &mut Unstructured<'a>) -> Result { let dtype: DType = u.arbitrary()?; - Self::with_dtype(u, &dtype, None) + Self::with_dtype(u, &dtype) } } impl ArbitraryConstantArray { /// Generate an arbitrary ConstantArray with the given dtype. - pub fn with_dtype(u: &mut Unstructured, dtype: &DType, len: Option) -> Result { + pub fn with_dtype(u: &mut Unstructured, dtype: &DType) -> Result { let scalar = random_scalar(u, dtype)?; - let len = len.unwrap_or(u.int_in_range(0..=100)?); + let len = u.int_in_range(0..=2048)?; Ok(ArbitraryConstantArray(ConstantArray::new(scalar, len))) } } diff --git a/vortex-array/src/arrays/constant/compute/cast.rs b/vortex-array/src/arrays/constant/compute/cast.rs index db3a41829f0..439bf8367b8 100644 --- a/vortex-array/src/arrays/constant/compute/cast.rs +++ b/vortex-array/src/arrays/constant/compute/cast.rs @@ -25,8 +25,8 @@ mod tests { use rstest::rstest; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ConstantArray; use crate::builtins::ArrayBuiltins; use crate::compute::conformance::cast::test_cast_conformance; @@ -44,10 +44,11 @@ mod tests { #[case(ConstantArray::new(Scalar::null_native::(), 4).into_array())] #[case(ConstantArray::new(Scalar::from(255u8), 1).into_array())] fn test_cast_constant_conformance(#[case] array: crate::ArrayRef) { - test_cast_conformance(&array); + test_cast_conformance(&array, &mut array_session().create_execution_ctx()); } #[test] + #[allow(clippy::disallowed_methods)] 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) @@ -57,7 +58,7 @@ mod tests { assert_eq!(casted.dtype(), &target_dtype); let scalar = casted - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut array_session().create_execution_ctx()) .unwrap(); assert_eq!( scalar.as_decimal().decimal_value(), diff --git a/vortex-array/src/arrays/constant/compute/fill_null.rs b/vortex-array/src/arrays/constant/compute/fill_null.rs index 1c9c2aaf903..67551a0592a 100644 --- a/vortex-array/src/arrays/constant/compute/fill_null.rs +++ b/vortex-array/src/arrays/constant/compute/fill_null.rs @@ -25,7 +25,7 @@ mod test { use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::ConstantArray; - use crate::arrow::ArrowSessionExt; + use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; @@ -43,21 +43,7 @@ mod test { assert!(!actual.dtype().is_nullable()); - let actual_arrow = array_session() - .arrow() - .execute_arrow(actual.clone(), None, &mut ctx) - .unwrap(); - let expected_arrow = array_session() - .arrow() - .execute_arrow(expected.clone(), None, &mut ctx) - .unwrap(); - assert_eq!( - &actual_arrow, - &expected_arrow, - "{}, {}", - actual.display_values(), - expected.display_values() - ); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] @@ -71,21 +57,7 @@ mod test { assert!(!actual.dtype().is_nullable()); - let actual_arrow = array_session() - .arrow() - .execute_arrow(actual.clone(), None, &mut ctx) - .unwrap(); - let expected_arrow = array_session() - .arrow() - .execute_arrow(expected.clone(), None, &mut ctx) - .unwrap(); - assert_eq!( - &actual_arrow, - &expected_arrow, - "{}, {}", - actual.display_values(), - expected.display_values() - ); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] @@ -98,26 +70,19 @@ mod test { Some(1.into()), )) .unwrap(); - let expected = ConstantArray::new(Scalar::from(1), 3).into_array(); + let expected = ConstantArray::new( + Scalar::new( + DType::Primitive(PType::I32, Nullability::Nullable), + Some(1.into()), + ), + 3, + ) + .into_array(); assert!(!Scalar::from(1).dtype().is_nullable()); assert!(actual.dtype().is_nullable()); - let actual_arrow = array_session() - .arrow() - .execute_arrow(actual.clone(), None, &mut ctx) - .unwrap(); - let expected_arrow = array_session() - .arrow() - .execute_arrow(expected.clone(), None, &mut ctx) - .unwrap(); - assert_eq!( - &actual_arrow, - &expected_arrow, - "{}, {}", - actual.display_values(), - expected.display_values() - ); + assert_arrays_eq!(actual, expected, &mut ctx); } } diff --git a/vortex-array/src/arrays/constant/compute/mod.rs b/vortex-array/src/arrays/constant/compute/mod.rs index 80fab02dc88..a996bf79a0f 100644 --- a/vortex-array/src/arrays/constant/compute/mod.rs +++ b/vortex-array/src/arrays/constant/compute/mod.rs @@ -29,21 +29,41 @@ mod test { #[test] fn test_mask_constant() { - test_mask_conformance(&ConstantArray::new(Scalar::null_native::(), 5).into_array()); - test_mask_conformance(&ConstantArray::new(Scalar::from(3u16), 5).into_array()); - test_mask_conformance(&ConstantArray::new(Scalar::from(1.0f32 / 0.0f32), 5).into_array()); + test_mask_conformance( + &ConstantArray::new(Scalar::null_native::(), 5).into_array(), + &mut array_session().create_execution_ctx(), + ); + test_mask_conformance( + &ConstantArray::new(Scalar::from(3u16), 5).into_array(), + &mut array_session().create_execution_ctx(), + ); + test_mask_conformance( + &ConstantArray::new(Scalar::from(1.0f32 / 0.0f32), 5).into_array(), + &mut array_session().create_execution_ctx(), + ); test_mask_conformance( &ConstantArray::new(Scalar::from(f16::from_f32(3.0f32)), 5).into_array(), + &mut array_session().create_execution_ctx(), ); } #[test] fn test_filter_constant() { - test_filter_conformance(&ConstantArray::new(Scalar::null_native::(), 5).into_array()); - test_filter_conformance(&ConstantArray::new(Scalar::from(3u16), 5).into_array()); - test_filter_conformance(&ConstantArray::new(Scalar::from(1.0f32 / 0.0f32), 5).into_array()); + test_filter_conformance( + &ConstantArray::new(Scalar::null_native::(), 5).into_array(), + &mut array_session().create_execution_ctx(), + ); + test_filter_conformance( + &ConstantArray::new(Scalar::from(3u16), 5).into_array(), + &mut array_session().create_execution_ctx(), + ); + test_filter_conformance( + &ConstantArray::new(Scalar::from(1.0f32 / 0.0f32), 5).into_array(), + &mut array_session().create_execution_ctx(), + ); test_filter_conformance( &ConstantArray::new(Scalar::from(f16::from_f32(3.0f32)), 5).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 9fb1373a5cf..ecd9db452ce 100644 --- a/vortex-array/src/arrays/constant/compute/take.rs +++ b/vortex-array/src/arrays/constant/compute/take.rs @@ -6,7 +6,6 @@ use vortex_mask::AllOr; use crate::ArrayRef; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array::ArrayView; use crate::arrays::Constant; @@ -14,13 +13,15 @@ use crate::arrays::ConstantArray; use crate::arrays::MaskedArray; use crate::arrays::dict::TakeReduce; use crate::arrays::dict::TakeReduceAdaptor; +use crate::legacy_session; use crate::optimizer::rules::ParentRuleSet; use crate::scalar::Scalar; use crate::validity::Validity; impl TakeReduce for Constant { + #[allow(clippy::disallowed_methods)] fn take(array: ArrayView<'_, Constant>, indices: &ArrayRef) -> VortexResult> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = legacy_session().create_execution_ctx(); let result = match indices .validity()? .execute_mask(indices.len(), &mut ctx)? @@ -73,8 +74,6 @@ mod tests { use vortex_mask::AllOr; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::ConstantArray; @@ -104,8 +103,7 @@ mod tests { taken.dtype() ); assert_arrays_eq!( - #[expect(deprecated)] - taken.to_primitive(), + taken.clone().execute::(&mut ctx).unwrap(), PrimitiveArray::new( buffer![42i32, 42, 42], Validity::from_iter([false, true, false]) @@ -116,7 +114,7 @@ mod tests { taken .validity() .unwrap() - .execute_mask(taken.len(), &mut array_session().create_execution_ctx()) + .execute_mask(taken.len(), &mut ctx) .unwrap() .indices(), AllOr::Some(valid_indices) @@ -135,8 +133,7 @@ mod tests { taken.dtype() ); assert_arrays_eq!( - #[expect(deprecated)] - taken.to_primitive(), + taken.clone().execute::(&mut ctx).unwrap(), PrimitiveArray::new(buffer![42i32, 42, 42], Validity::AllValid), &mut ctx ); @@ -144,7 +141,7 @@ mod tests { taken .validity() .unwrap() - .execute_mask(taken.len(), &mut array_session().create_execution_ctx()) + .execute_mask(taken.len(), &mut ctx) .unwrap() .indices(), AllOr::All @@ -158,6 +155,9 @@ mod tests { #[case(ConstantArray::new(Scalar::null_native::(), 5))] #[case(ConstantArray::new(true, 1))] fn test_take_constant_conformance(#[case] array: ConstantArray) { - test_take_conformance(&array.into_array()); + test_take_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index d3df3a14b36..a2057077335 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -246,7 +246,7 @@ impl VTable for Constant { .clone() .execute::(ctx)? .into_array(); - builder.extend_from_array(&canonical); + canonical.append_to_builder(builder, ctx)?; } } diff --git a/vortex-array/src/arrays/datetime/test.rs b/vortex-array/src/arrays/datetime/test.rs index 1ae114874d9..60081650f9f 100644 --- a/vortex-array/src/arrays/datetime/test.rs +++ b/vortex-array/src/arrays/datetime/test.rs @@ -8,8 +8,6 @@ use vortex_error::VortexResult; use crate::EqMode; use crate::IntoArray; -#[expect(deprecated)] -use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::PrimitiveArray; @@ -202,8 +200,11 @@ fn test_validity_preservation(#[case] validity: Validity) { let temporal_array = TemporalData::new_timestamp(milliseconds, TimeUnit::Milliseconds, Some("UTC".into())); - #[expect(deprecated)] - let prim = temporal_array.temporal_values().to_primitive(); + let prim = temporal_array + .temporal_values() + .clone() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap(); assert!( prim.validity() .vortex_expect("temporal validity should be derivable") diff --git a/vortex-array/src/arrays/decimal/compute/between.rs b/vortex-array/src/arrays/decimal/compute/between.rs index 75c2da923ea..50623b077ee 100644 --- a/vortex-array/src/arrays/decimal/compute/between.rs +++ b/vortex-array/src/arrays/decimal/compute/between.rs @@ -132,8 +132,10 @@ fn between_impl( ) -> ArrayRef { let buffer = arr.buffer::(); BoolArray::new( - BitBuffer::collect_bool(buffer.len(), |idx| { - let value = buffer[idx]; + BitBuffer::collect_bool_multiversioned(buffer.len(), |idx| { + // SAFETY: `collect_bool_multiversioned` invokes the predicate with indices + // `0..buffer.len()` only. + let value = unsafe { *buffer.get_unchecked(idx) }; lower.is_none_or(|l| lower_op(l, value)) & upper.is_none_or(|u| upper_op(value, u)) }), arr.validity() diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index 80fafde5059..20cac1d5106 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -1,7 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use num_traits::AsPrimitive; use num_traits::CheckedMul; +use num_traits::ToPrimitive as NumToPrimitive; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_compute::lane_kernels::IndexedSourceExt; @@ -19,11 +21,14 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Decimal; use crate::arrays::DecimalArray; +use crate::arrays::PrimitiveArray; use crate::dtype::BigCast; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; +use crate::dtype::Nullability; +use crate::dtype::PType; use crate::dtype::i256; use crate::match_each_decimal_value_type; use crate::scalar::DecimalValue; @@ -77,16 +82,19 @@ impl CastKernel for Decimal { dtype: &DType, ctx: &mut ExecutionCtx, ) -> VortexResult> { - // Early return if not casting to decimal - let DType::Decimal(to_decimal_dtype, to_nullability) = dtype else { - return Ok(None); - }; let DType::Decimal(from_decimal_dtype, _) = array.dtype() else { vortex_panic!( "DecimalArray must have decimal dtype, got {:?}", array.dtype() ); }; + if let DType::Primitive(PType::F64, nullability) = dtype { + let scale = from_decimal_dtype.scale(); + return cast_to_f64(array, scale, *nullability, ctx).map(Some); + } + let DType::Decimal(to_decimal_dtype, to_nullability) = dtype else { + return Ok(None); + }; // If the dtype is exactly the same, return self if array.dtype() == dtype { @@ -143,6 +151,57 @@ impl CastKernel for Decimal { } } +fn cast_to_f64( + array: ArrayView<'_, Decimal>, + scale: i8, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let source_validity = array.validity()?; + let n = array.len(); + let mask = source_validity.execute_mask(n, ctx)?; + let validity = source_validity.cast_nullability(nullability, n, ctx)?; + + let scale: i32 = scale.as_(); + let inv_factor: f64 = 10f64.powi(-scale); + + let buffer = match &mask { + Mask::AllFalse(_) => BufferMut::::zeroed(n), + Mask::AllTrue(_) => match_each_decimal_value_type!(array.values_type(), |F| { + let values = array.buffer::(); + let values = values.as_slice(); + let mut out = BufferMut::::with_capacity(n); + values.map_into(&mut out.spare_capacity_mut()[..n], |v: F| { + to_f64_lossy::(v) * inv_factor + }); + // SAFETY: map_into wrote every lane before returning. + unsafe { out.set_len(n) }; + out + }), + Mask::Values(mask_values) => match_each_decimal_value_type!(array.values_type(), |F| { + let values = array.buffer::(); + let values = values.as_slice(); + let mut out = BufferMut::::with_capacity(n); + let write_result = values.try_map_masked_into( + mask_values.bit_buffer(), + &mut out.spare_capacity_mut()[..n], + |v: F| Some(to_f64_lossy::(v) * inv_factor), + ); + debug_assert!(write_result.is_ok()); + // SAFETY: try_map_masked_into wrote every lane before returning Ok. + unsafe { out.set_len(n) }; + out + }), + }; + + Ok(PrimitiveArray::new(buffer, validity).into_array()) +} + +#[inline] +fn to_f64_lossy(value: F) -> f64 { + NumToPrimitive::to_f64(&value).unwrap_or(f64::NAN) +} + fn cast_decimal_values( array: ArrayView<'_, Decimal>, from_decimal_dtype: DecimalDType, @@ -386,22 +445,24 @@ mod tests { use vortex_buffer::buffer; use super::upcast_decimal_values; + use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::DecimalArray; use crate::builtins::ArrayBuiltins; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; use crate::compute::conformance::cast::test_cast_conformance; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::DecimalType; use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::scalar::Scalar; use crate::validity::Validity; #[test] fn cast_decimal_to_nullable() { + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(10, 2); let array = DecimalArray::new( buffer![100i32, 200, 300], @@ -411,12 +472,12 @@ mod tests { // Cast to nullable let nullable_dtype = DType::Decimal(decimal_dtype, Nullability::Nullable); - #[expect(deprecated)] let casted = array .into_array() .cast(nullable_dtype.clone()) .unwrap() - .to_decimal(); + .execute::(&mut ctx) + .unwrap(); assert_eq!(casted.dtype(), &nullable_dtype); assert!(matches!(casted.validity(), Ok(Validity::AllValid))); @@ -425,6 +486,7 @@ mod tests { #[test] fn cast_nullable_to_non_nullable() { + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(10, 2); // Create nullable array with no nulls @@ -432,12 +494,12 @@ mod tests { // Cast to non-nullable let non_nullable_dtype = DType::Decimal(decimal_dtype, Nullability::NonNullable); - #[expect(deprecated)] let casted = array .into_array() .cast(non_nullable_dtype.clone()) .unwrap() - .to_decimal(); + .execute::(&mut ctx) + .unwrap(); assert_eq!(casted.dtype(), &non_nullable_dtype); assert!(matches!(casted.validity(), Ok(Validity::NonNullable))); @@ -446,6 +508,7 @@ mod tests { #[test] #[should_panic(expected = "Cannot cast array with invalid values to non-nullable type")] fn cast_nullable_with_nulls_to_non_nullable_fails() { + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(10, 2); // Create nullable array with nulls @@ -453,16 +516,16 @@ mod tests { // Attempt to cast to non-nullable should fail let non_nullable_dtype = DType::Decimal(decimal_dtype, Nullability::NonNullable); - #[expect(deprecated)] - let result = array + array .into_array() .cast(non_nullable_dtype) - .and_then(|a| a.to_canonical().map(|c| c.into_array())); - result.unwrap(); + .and_then(|a| a.execute::(&mut ctx).map(|c| c.into_array())) + .unwrap(); } #[test] fn cast_different_scale_rescales() { + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::new( buffer![100i32], DecimalDType::new(10, 2), @@ -471,12 +534,12 @@ mod tests { // 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 casted = array .into_array() .cast(different_dtype) .unwrap() - .to_decimal(); + .execute::(&mut ctx) + .unwrap(); assert_eq!(casted.precision(), 15); assert_eq!(casted.scale(), 3); @@ -486,6 +549,7 @@ mod tests { #[test] fn cast_downcast_precision_succeeds_when_values_fit() { + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::new( buffer![100i64], DecimalDType::new(18, 2), @@ -494,8 +558,12 @@ mod tests { // 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(); + let casted = array + .into_array() + .cast(smaller_dtype) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_eq!(casted.precision(), 10); assert_eq!(casted.scale(), 2); @@ -504,6 +572,7 @@ mod tests { #[test] fn cast_downcast_precision_checks_values() { + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::new( buffer![1000i64], DecimalDType::new(18, 0), @@ -511,11 +580,10 @@ mod tests { ); 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())); + .and_then(|a| a.execute::(&mut ctx).map(|c| c.into_array())); assert!(result.is_err()); assert!( @@ -528,6 +596,7 @@ mod tests { #[test] fn cast_lower_scale_requires_exact_rescale() { + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::new( buffer![123456i64], DecimalDType::new(10, 4), @@ -535,11 +604,10 @@ mod tests { ); let lower_scale_dtype = DType::Decimal(DecimalDType::new(10, 2), Nullability::NonNullable); - #[expect(deprecated)] let result = array .into_array() .cast(lower_scale_dtype) - .and_then(|a| a.to_canonical().map(|c| c.into_array())); + .and_then(|a| a.execute::(&mut ctx).map(|c| c.into_array())); assert!(result.is_err()); assert!( @@ -552,6 +620,7 @@ mod tests { #[test] fn cast_lower_scale_ignores_null_lane_failures() { + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::new( buffer![100i64, 123456], DecimalDType::new(10, 4), @@ -559,12 +628,12 @@ mod tests { ); 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(); + .execute::(&mut ctx) + .unwrap(); let mask = casted .as_ref() @@ -582,6 +651,7 @@ mod tests { #[test] fn cast_upcast_precision_succeeds() { + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::new( buffer![100i32, 200, 300], DecimalDType::new(10, 2), @@ -590,8 +660,12 @@ mod tests { // Cast to higher precision with same scale - should succeed let wider_dtype = DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable); - #[expect(deprecated)] - let casted = array.into_array().cast(wider_dtype).unwrap().to_decimal(); + let casted = array + .into_array() + .cast(wider_dtype) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_eq!(casted.precision(), 38); assert_eq!(casted.scale(), 2); @@ -602,6 +676,7 @@ mod tests { #[test] fn cast_widening_same_physical_type_is_zero_copy() { + let mut ctx = array_session().create_execution_ctx(); // 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( @@ -612,8 +687,12 @@ mod tests { 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(); + let casted = array + .into_array() + .cast(wider_dtype) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_eq!(casted.precision(), 18); assert_eq!(casted.scale(), 2); @@ -629,6 +708,7 @@ mod tests { #[test] fn cast_to_non_decimal_returns_err() { + let mut ctx = array_session().create_execution_ctx(); let array = DecimalArray::new( buffer![100i32], DecimalDType::new(10, 2), @@ -636,11 +716,10 @@ mod tests { ); // Try to cast to non-decimal type - should fail since no kernel can handle it - #[expect(deprecated)] let result = array .into_array() .cast(DType::Utf8(Nullability::NonNullable)) - .and_then(|a| a.to_canonical().map(|c| c.into_array())); + .and_then(|a| a.execute::(&mut ctx).map(|c| c.into_array())); assert!(result.is_err()); assert!( @@ -657,7 +736,10 @@ mod tests { #[case(DecimalArray::from_option_iter([Some(100i32), None, Some(300)], DecimalDType::new(10, 2)))] #[case(DecimalArray::new(buffer![42i32], DecimalDType::new(5, 1), Validity::NonNullable))] fn test_cast_decimal_conformance(#[case] array: DecimalArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] @@ -769,4 +851,68 @@ mod tests { .contains("Cannot downcast decimal values") ); } + + #[test] + fn cast_decimal_f64() { + let dtype = DecimalDType::new(6, 2); + let array = DecimalArray::new(buffer![125i32, 250, -375], dtype, Validity::NonNullable); + let target = DType::Primitive(PType::F64, Nullability::NonNullable); + let casted = array.into_array().cast(target.clone()).unwrap(); + assert_eq!(casted.dtype(), &target); + + let mut ctx = array_session().create_execution_ctx(); + let primitive = casted + .execute::(&mut ctx) + .unwrap() + .into_primitive(); + assert_eq!(primitive.to_buffer::().as_ref(), &[1.25, 2.5, -3.75]); + } + + #[test] + fn cast_decimal_f64_null() { + let values = [Some(100i32), None, Some(-200)]; + let array = DecimalArray::from_option_iter(values, DecimalDType::new(6, 2)); + let target = DType::Primitive(PType::F64, Nullability::Nullable); + let casted = array.into_array().cast(target.clone()).unwrap(); + assert_eq!(casted.dtype(), &target); + + let mut ctx = array_session().create_execution_ctx(); + let primitive = casted + .execute::(&mut ctx) + .unwrap() + .into_primitive(); + + assert!(primitive.is_valid(0, &mut ctx).unwrap()); + assert!(!primitive.is_valid(1, &mut ctx).unwrap()); + assert!(primitive.is_valid(2, &mut ctx).unwrap()); + + assert_eq!( + primitive.execute_scalar(0, &mut ctx).unwrap(), + Scalar::from(1f64) + ); + assert_eq!( + primitive.execute_scalar(2, &mut ctx).unwrap(), + Scalar::from(-2f64) + ); + } + + #[test] + fn cast_decimal_f64_all_null() { + let dtype = DecimalDType::new(6, 2); + let buf = buffer![i32::MAX, i32::MIN, 12345]; + let array = DecimalArray::new(buf, dtype, Validity::AllInvalid); + let target = DType::Primitive(PType::F64, Nullability::Nullable); + let casted = array.into_array().cast(target.clone()).unwrap(); + assert_eq!(casted.dtype(), &target); + + let primitive = casted + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() + .into_primitive(); + let mut ctx = array_session().create_execution_ctx(); + for i in 0..2 { + assert!(!primitive.is_valid(i, &mut ctx).unwrap()); + } + assert_eq!(primitive.to_buffer::().as_ref(), &[0.0, 0.0, 0.0]); + } } diff --git a/vortex-array/src/arrays/decimal/compute/fill_null.rs b/vortex-array/src/arrays/decimal/compute/fill_null.rs index 1c2904b7b6b..c4de15de5cc 100644 --- a/vortex-array/src/arrays/decimal/compute/fill_null.rs +++ b/vortex-array/src/arrays/decimal/compute/fill_null.rs @@ -94,8 +94,6 @@ mod tests { use crate::arrays::DecimalArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; use crate::dtype::DecimalDType; use crate::dtype::Nullability; use crate::scalar::DecimalValue; @@ -110,7 +108,6 @@ mod tests { [None, Some(800i128), None, Some(1000i128), None], decimal_dtype, ); - #[expect(deprecated)] let p = arr .into_array() .fill_null(Scalar::decimal( @@ -119,7 +116,8 @@ mod tests { Nullability::NonNullable, )) .unwrap() - .to_decimal(); + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( p, DecimalArray::from_iter([4200, 800, 4200, 1000, 4200], decimal_dtype), @@ -152,7 +150,6 @@ mod tests { decimal_dtype, ); - #[expect(deprecated)] let p = arr .into_array() .fill_null(Scalar::decimal( @@ -161,7 +158,8 @@ mod tests { Nullability::NonNullable, )) .unwrap() - .to_decimal(); + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( p, DecimalArray::from_iter([25500, 25500, 25500, 25500, 25500], decimal_dtype), @@ -176,7 +174,6 @@ mod tests { 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. - #[expect(deprecated)] let result = arr .into_array() .fill_null(Scalar::decimal( @@ -185,7 +182,8 @@ mod tests { Nullability::NonNullable, )) .unwrap() - .to_decimal(); + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( result, DecimalArray::from_iter([200i16, 10, 200], decimal_dtype), @@ -203,7 +201,6 @@ mod tests { decimal_dtype, Validity::NonNullable, ); - #[expect(deprecated)] let p = arr .into_array() .fill_null(Scalar::decimal( @@ -212,7 +209,8 @@ mod tests { Nullability::NonNullable, )) .unwrap() - .to_decimal(); + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( p, DecimalArray::from_iter([800i128, 1000, 1200, 1400, 1600], decimal_dtype), diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 762ac3d7365..ce786f8e4ab 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -131,6 +131,9 @@ mod tests { ) })] fn test_take_decimal_conformance(#[case] array: DecimalArray) { - test_take_conformance(&array.into_array()); + test_take_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/decimal/mod.rs b/vortex-array/src/arrays/decimal/mod.rs index e740a75ea91..d8af8529e87 100644 --- a/vortex-array/src/arrays/decimal/mod.rs +++ b/vortex-array/src/arrays/decimal/mod.rs @@ -19,17 +19,3 @@ pub(crate) fn initialize(session: &vortex_session::VortexSession) { mod utils; pub use utils::*; - -#[cfg(test)] -mod tests { - use arrow_array::Decimal128Array; - - #[test] - fn test_decimal() { - // They pass it b/c the DType carries the information. No other way to carry a - // dtype except via the array. - let value = Decimal128Array::new_null(100); - let numeric = value.value(10); - assert_eq!(numeric, 0i128); - } -} diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index 674f8598256..1ed839b813e 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -20,6 +20,8 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::arrays::decimal::DecimalData; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::DecimalBuilder; use crate::dtype::DType; use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; @@ -207,6 +209,17 @@ impl VTable for Decimal { Ok(ExecutionResult::done(array)) } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for Decimal requires a DecimalBuilder"); + }; + builder.append_decimal_array(&array.into_owned(), ctx) + } + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/vortex-array/src/arrays/dict/arbitrary.rs b/vortex-array/src/arrays/dict/arbitrary.rs index 4eb938181b9..2a70dd2586c 100644 --- a/vortex-array/src/arrays/dict/arbitrary.rs +++ b/vortex-array/src/arrays/dict/arbitrary.rs @@ -37,7 +37,6 @@ impl ArbitraryDictArray { pub fn with_dtype(u: &mut Unstructured, dtype: &DType, len: Option) -> Result { // Generate the number of unique values (dictionary size) let values_len = u.int_in_range(1..=20)?; - // Generate values array with the given dtype let values = ArbitraryArray::arbitrary_with_config( u, &ArbitraryArrayConfig { diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index f73f4965439..cefc5a381fd 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -15,15 +15,13 @@ use vortex_mask::AllOr; use crate::ArrayRef; use crate::ArraySlots; -use crate::LEGACY_SESSION; -#[expect(deprecated)] -use crate::ToCanonical as _; -use crate::VortexSessionExecute; +use crate::ExecutionCtx; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; use crate::array_slots; use crate::arrays::Dict; +use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::PType; use crate::match_each_integer_ptype; @@ -129,13 +127,13 @@ pub trait DictArrayExt: TypedArrayRef + DictArraySlotsExt { self.all_values_referenced } - fn validate_all_values_referenced(&self) -> VortexResult<()> { + fn validate_all_values_referenced(&self, ctx: &mut ExecutionCtx) -> VortexResult<()> { if self.has_all_values_referenced() { if !self.codes().is_host() { return Ok(()); } - let referenced_mask = self.compute_referenced_values_mask(true)?; + let referenced_mask = self.compute_referenced_values_mask(true, ctx)?; let all_referenced = referenced_mask.true_count() == referenced_mask.len(); vortex_ensure!(all_referenced, "value in dict not referenced"); @@ -144,13 +142,14 @@ pub trait DictArrayExt: TypedArrayRef + DictArraySlotsExt { Ok(()) } - fn compute_referenced_values_mask(&self, referenced: bool) -> VortexResult { + fn compute_referenced_values_mask( + &self, + referenced: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult { let codes = self.codes(); - let codes_validity = codes - .validity()? - .execute_mask(codes.len(), &mut LEGACY_SESSION.create_execution_ctx())?; - #[expect(deprecated)] - let codes_primitive = self.codes().to_primitive(); + let codes_validity = codes.validity()?.execute_mask(codes.len(), ctx)?; + let codes_primitive = codes.clone().execute::(ctx)?; let values_len = self.values().len(); let init_value = !referenced; @@ -268,18 +267,9 @@ impl Array { self.into_data() .set_all_values_referenced(all_values_referenced) }; - let array = unsafe { + unsafe { Array::from_parts_unchecked(ArrayParts::new(Dict, dtype, len, data).with_slots(slots)) - }; - - #[cfg(debug_assertions)] - if all_values_referenced { - array - .validate_all_values_referenced() - .vortex_expect("validation should succeed when all values are referenced"); } - - array } } @@ -299,8 +289,6 @@ mod test { use crate::ArrayRef; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::ChunkedArray; @@ -468,9 +456,8 @@ mod test { &mut array_session().create_execution_ctx(), )?; - #[expect(deprecated)] - let into_prim = array.to_primitive(); - let prim_into = builder.finish_into_canonical().into_primitive(); + let into_prim = array.execute::(&mut ctx)?; + let prim_into = builder.finish_into_canonical(&mut ctx).into_primitive(); 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 56a433627a2..2597436f98f 100644 --- a/vortex-array/src/arrays/dict/compute/cast.rs +++ b/vortex-array/src/arrays/dict/compute/cast.rs @@ -184,7 +184,7 @@ mod tests { #[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: ArrayRef) { - test_cast_conformance(&array); + test_cast_conformance(&array, &mut SESSION.create_execution_ctx()); } #[test] diff --git a/vortex-array/src/arrays/dict/compute/is_constant.rs b/vortex-array/src/arrays/dict/compute/is_constant.rs index 82518963106..85438683442 100644 --- a/vortex-array/src/arrays/dict/compute/is_constant.rs +++ b/vortex-array/src/arrays/dict/compute/is_constant.rs @@ -47,7 +47,7 @@ impl DynAggregateKernel for DictIsConstantKernel { let result = if dict.has_all_values_referenced() { is_constant(dict.values(), ctx)? } else { - let referenced_mask = dict.compute_referenced_values_mask(true)?; + let referenced_mask = dict.compute_referenced_values_mask(true, ctx)?; let mask = Mask::from(referenced_mask); let filtered_values = dict.values().filter(mask)?; is_constant(&filtered_values, ctx)? diff --git a/vortex-array/src/arrays/dict/compute/min_max.rs b/vortex-array/src/arrays/dict/compute/min_max.rs index ee38c0392d4..0122761c7bf 100644 --- a/vortex-array/src/arrays/dict/compute/min_max.rs +++ b/vortex-array/src/arrays/dict/compute/min_max.rs @@ -45,7 +45,7 @@ impl DynAggregateKernel for DictMinMaxKernel { 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 referenced_mask = dict.compute_referenced_values_mask(true, ctx)?; let mask = Mask::from(referenced_mask); let filtered_values = dict.values().filter(mask)?; min_max(&filtered_values, ctx, *options)? diff --git a/vortex-array/src/arrays/dict/compute/mod.rs b/vortex-array/src/arrays/dict/compute/mod.rs index e2b0d07b97d..5c013b2d321 100644 --- a/vortex-array/src/arrays/dict/compute/mod.rs +++ b/vortex-array/src/arrays/dict/compute/mod.rs @@ -215,7 +215,7 @@ mod test { &mut SESSION.create_execution_ctx(), ) .unwrap(); - test_mask_conformance(&array.into_array()); + test_mask_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); let array = dict_encode( &PrimitiveArray::from_option_iter([Some(2), None, Some(2), Some(0), Some(10)]) @@ -223,7 +223,7 @@ mod test { &mut SESSION.create_execution_ctx(), ) .unwrap(); - test_mask_conformance(&array.into_array()); + test_mask_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); let array = dict_encode( &VarBinArray::from_iter( @@ -240,7 +240,7 @@ mod test { &mut SESSION.create_execution_ctx(), ) .unwrap(); - test_mask_conformance(&array.into_array()); + test_mask_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] @@ -250,7 +250,7 @@ mod test { &mut SESSION.create_execution_ctx(), ) .unwrap(); - test_filter_conformance(&array.into_array()); + test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); let array = dict_encode( &PrimitiveArray::from_option_iter([Some(2), None, Some(2), Some(0), Some(10)]) @@ -258,7 +258,7 @@ mod test { &mut SESSION.create_execution_ctx(), ) .unwrap(); - test_filter_conformance(&array.into_array()); + test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); let array = dict_encode( &VarBinArray::from_iter( @@ -275,7 +275,7 @@ mod test { &mut SESSION.create_execution_ctx(), ) .unwrap(); - test_filter_conformance(&array.into_array()); + test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] @@ -302,7 +302,7 @@ mod test { &mut SESSION.create_execution_ctx(), ) .unwrap(); - test_take_conformance(&array.into_array()); + test_take_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); let array = dict_encode( &PrimitiveArray::from_option_iter([Some(2), None, Some(2), Some(0), Some(10)]) @@ -310,7 +310,7 @@ mod test { &mut SESSION.create_execution_ctx(), ) .unwrap(); - test_take_conformance(&array.into_array()); + test_take_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); let array = dict_encode( &VarBinArray::from_iter( @@ -327,7 +327,7 @@ mod test { &mut SESSION.create_execution_ctx(), ) .unwrap(); - test_take_conformance(&array.into_array()); + test_take_conformance(&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 e4d102ee861..a15f2641eab 100644 --- a/vortex-array/src/arrays/dict/compute/rules.rs +++ b/vortex-array/src/arrays/dict/compute/rules.rs @@ -267,6 +267,7 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; @@ -284,6 +285,7 @@ mod tests { use crate::scalar_fn::fns::not::Not; #[test] + #[allow(clippy::disallowed_methods)] 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(); @@ -298,7 +300,7 @@ mod tests { assert!(ArrayRef::ptr_eq(dict.values(), &values)); assert_eq!(codes.nchunks(), 2); - let mut ctx = crate::LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_arrays_eq!( optimized, PrimitiveArray::from_iter([10u32, 20, 30, 10, 20]), @@ -309,6 +311,7 @@ mod tests { } #[test] + #[allow(clippy::disallowed_methods)] 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(); @@ -321,7 +324,7 @@ mod tests { let optimized = array.optimize()?; assert!(optimized.is::()); - let mut ctx = crate::LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_arrays_eq!( optimized, PrimitiveArray::from_iter([10u32, 20, 30, 10, 20]), diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index acf12b47070..d32fcfb6192 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -363,7 +363,7 @@ mod tests { #[case(create_timestamp_array(TimeUnit::Nanoseconds, false))] #[case(create_timestamp_array(TimeUnit::Seconds, true))] fn test_cast_extension_conformance(#[case] array: ExtensionArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } fn create_timestamp_array(time_unit: TimeUnit, nullable: bool) -> ExtensionArray { diff --git a/vortex-array/src/arrays/extension/compute/compare.rs b/vortex-array/src/arrays/extension/compute/compare.rs index ecd6bbd5614..b3f978ef503 100644 --- a/vortex-array/src/arrays/extension/compute/compare.rs +++ b/vortex-array/src/arrays/extension/compute/compare.rs @@ -22,6 +22,12 @@ impl CompareKernel for Extension { operator: CompareOperator, _ctx: &mut ExecutionCtx, ) -> VortexResult> { + // Storage values are only comparable when both sides share the same extension dtype + // (e.g. timestamps in different units must not compare their raw storage). + if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) { + return Ok(None); + } + // If the RHS is a constant, we can extract the storage scalar. if let Some(const_ext) = rhs.as_constant() { let storage_scalar = const_ext.as_extension().to_storage_scalar(); diff --git a/vortex-array/src/arrays/extension/compute/mod.rs b/vortex-array/src/arrays/extension/compute/mod.rs index 8309ae43be2..256f31ac49b 100644 --- a/vortex-array/src/arrays/extension/compute/mod.rs +++ b/vortex-array/src/arrays/extension/compute/mod.rs @@ -15,6 +15,8 @@ mod test { 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::filter::test_filter_conformance; @@ -31,14 +33,20 @@ mod test { // Create storage array let storage = buffer![1i32, 2, 3, 4, 5].into_array(); let array = ExtensionArray::new(ext_dtype.clone(), storage); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); // Test with nullable extension type let ext_dtype_nullable = ext_dtype.with_nullability(Nullability::Nullable); let storage = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None]) .into_array(); let array = ExtensionArray::new(ext_dtype_nullable, storage); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[rstest] @@ -77,7 +85,10 @@ mod test { ExtensionArray::new(ext_dtype_large, storage) })] fn test_take_extension_array_conformance(#[case] array: ExtensionArray) { - test_take_conformance(&array.into_array()); + test_take_conformance( + &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 396729b6347..2d02d1ae7a7 100644 --- a/vortex-array/src/arrays/extension/compute/rules.rs +++ b/vortex-array/src/arrays/extension/compute/rules.rs @@ -87,8 +87,7 @@ mod tests { use crate::EmptyMetadata; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::Extension; @@ -179,8 +178,11 @@ mod tests { assert_eq!(ext_result.ext_dtype(), &ext_dtype); // Check the storage values - #[expect(deprecated)] - let storage_prim = ext_result.storage_array().to_primitive(); + let storage_prim = ext_result + .storage_array() + .clone() + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); let storage_result: &[i64] = &storage_prim.to_buffer::(); assert_eq!(storage_result, &[1, 3, 5]); } @@ -207,8 +209,11 @@ mod tests { assert_eq!(ext_result.len(), 3); // Check values: should be [Some(1), None, None] - #[expect(deprecated)] - let canonical = ext_result.storage_array().to_primitive(); + let canonical = ext_result + .storage_array() + .clone() + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); assert_eq!(canonical.len(), 3); } diff --git a/vortex-array/src/arrays/extension/vtable/mod.rs b/vortex-array/src/arrays/extension/vtable/mod.rs index c27e70476dc..ef371e46afc 100644 --- a/vortex-array/src/arrays/extension/vtable/mod.rs +++ b/vortex-array/src/arrays/extension/vtable/mod.rs @@ -27,6 +27,8 @@ use crate::arrays::extension::array::STORAGE_SLOT; use crate::arrays::extension::compute::rules::PARENT_RULES; use crate::arrays::extension::compute::rules::RULES; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::ExtensionBuilder; use crate::dtype::DType; use crate::serde::ArrayChildren; @@ -188,6 +190,17 @@ impl VTable for Extension { Ok(ExecutionResult::done(array)) } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for Extension requires an ExtensionBuilder"); + }; + builder.append_extension_array(&array.into_owned(), ctx) + } + fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { RULES.evaluate(array) } diff --git a/vortex-array/src/arrays/filter/execute/bool.rs b/vortex-array/src/arrays/filter/execute/bool.rs index a22e55366fc..08364e3cac5 100644 --- a/vortex-array/src/arrays/filter/execute/bool.rs +++ b/vortex-array/src/arrays/filter/execute/bool.rs @@ -30,18 +30,22 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::filter::execute::bool::BoolArray; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; 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]); - #[expect(deprecated)] - let filtered = arr.filter(mask).unwrap().to_bool(); + let filtered = arr + .filter(mask) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_eq!(2, filtered.len()); assert_eq!( @@ -58,6 +62,9 @@ mod test { #[case(BoolArray::from_iter((0..100).map(|i| i % 2 == 0)))] #[case(BoolArray::from_iter((0..1024).map(|i| i % 3 != 0)))] fn test_filter_bool_conformance(#[case] array: BoolArray) { - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/filter/execute/decimal.rs b/vortex-array/src/arrays/filter/execute/decimal.rs index 869e1fabe0c..5ffb4a963b2 100644 --- a/vortex-array/src/arrays/filter/execute/decimal.rs +++ b/vortex-array/src/arrays/filter/execute/decimal.rs @@ -34,6 +34,8 @@ pub fn filter_decimal(array: &DecimalArray, mask: &Arc) -> DecimalAr #[cfg(test)] mod test { use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::filter::execute::decimal::DecimalArray; use crate::compute::conformance::filter::test_filter_conformance; use crate::dtype::DecimalDType; @@ -49,7 +51,10 @@ mod test { Some(99999), ]; let array = DecimalArray::from_option_iter(values, decimal_dtype); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] @@ -57,6 +62,9 @@ mod test { let decimal_dtype = DecimalDType::new(38, 4); let values = vec![Some(12345i128), None, Some(-12345), Some(0), None]; let array = DecimalArray::from_option_iter(values, decimal_dtype); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs index 392a6373f76..6facbe3e38d 100644 --- a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs +++ b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs @@ -137,7 +137,10 @@ mod test { fn test_filter_fixed_size_list_conformance() { let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6, 7, 8, 9]); let array = FixedSizeListArray::new(elements.into_array(), 3, Validity::NonNullable, 3); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] @@ -146,7 +149,10 @@ mod test { PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), Some(5), None]); let validity = Validity::from_iter([true, false, true]); let array = FixedSizeListArray::new(elements.into_array(), 2, validity, 3); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] diff --git a/vortex-array/src/arrays/filter/execute/listview.rs b/vortex-array/src/arrays/filter/execute/listview.rs index 0507452c89d..6d1521541bc 100644 --- a/vortex-array/src/arrays/filter/execute/listview.rs +++ b/vortex-array/src/arrays/filter/execute/listview.rs @@ -85,7 +85,7 @@ mod test { let sizes = buffer![2u32, 2, 2].into_array(); let array = ListViewArray::new(elements.into_array(), offsets, sizes, Validity::NonNullable); - test_filter_conformance(&array.into_array()); + test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] @@ -96,7 +96,7 @@ mod test { let sizes = buffer![2u32, 2, 2].into_array(); let validity = Validity::from_iter([true, false, true]); let array = ListViewArray::new(elements.into_array(), offsets, sizes, validity); - test_filter_conformance(&array.into_array()); + test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] @@ -109,7 +109,7 @@ mod test { ListViewArray::new_unchecked(elements, offsets, sizes, Validity::NonNullable) .with_zero_copy_to_list(true) }; - test_filter_conformance(&array.into_array()); + test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] @@ -120,7 +120,7 @@ mod test { let offsets = buffer![5u32, 2, 8, 0, 1].into_array(); let sizes = buffer![3u32, 2, 2, 2, 4].into_array(); let array = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); - test_filter_conformance(&array.into_array()); + test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] @@ -130,7 +130,7 @@ mod test { let offsets = buffer![0u32, 100, 200, 300, 400, 500, 600, 700, 800, 900].into_array(); let sizes = buffer![50u32, 50, 50, 50, 50, 50, 50, 50, 50, 50].into_array(); let array = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); - test_filter_conformance(&array.into_array()); + test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] diff --git a/vortex-array/src/arrays/filter/execute/primitive.rs b/vortex-array/src/arrays/filter/execute/primitive.rs index 276e08b5526..b7abd8efa4b 100644 --- a/vortex-array/src/arrays/filter/execute/primitive.rs +++ b/vortex-array/src/arrays/filter/execute/primitive.rs @@ -59,20 +59,24 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; use crate::compute::conformance::filter::LARGE_SIZE; use crate::compute::conformance::filter::MEDIUM_SIZE; use crate::compute::conformance::filter::test_filter_conformance; #[test] fn filter_run_variant_mixed_test() { + let mut ctx = array_session().create_execution_ctx(); let mask = [true, true, false, true, true, true, false, true]; let arr = PrimitiveArray::from_iter([1u32, 24, 54, 2, 3, 2, 3, 2]); - #[expect(deprecated)] - let filtered = arr.filter(Mask::from_iter(mask)).unwrap().to_primitive(); + let filtered = arr + .filter(Mask::from_iter(mask)) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_eq!( filtered.len(), mask.iter().filter(|x| **x).collect_vec().len() @@ -104,6 +108,9 @@ mod test { #[case(PrimitiveArray::from_iter([0.1f32, 0.2, 0.3, 0.4, 0.5]))] #[case(PrimitiveArray::from_option_iter([Some(1.1f64), None, Some(2.2), Some(3.3), None]))] fn test_filter_primitive_conformance(#[case] array: PrimitiveArray) { - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/filter/execute/struct_.rs b/vortex-array/src/arrays/filter/execute/struct_.rs index 8980c422ef3..e4a8157d2e2 100644 --- a/vortex-array/src/arrays/filter/execute/struct_.rs +++ b/vortex-array/src/arrays/filter/execute/struct_.rs @@ -70,7 +70,10 @@ mod test { ]; let array = StructArray::try_new(["a", "b"].into(), fields, 5, Validity::NonNullable).unwrap(); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] @@ -83,7 +86,10 @@ mod test { ]; let array = StructArray::try_new(["a", "b"].into(), fields, 5, Validity::NonNullable).unwrap(); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] @@ -150,6 +156,7 @@ mod test { &StructArray::try_new(FieldNames::empty(), vec![], 5, Validity::NonNullable) .unwrap() .into_array(), + &mut array_session().create_execution_ctx(), ); } @@ -184,6 +191,7 @@ mod test { ) .unwrap() .into_array(), + &mut array_session().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/filter/execute/varbinview.rs b/vortex-array/src/arrays/filter/execute/varbinview.rs index ded9c804cfd..58a348b41b0 100644 --- a/vortex-array/src/arrays/filter/execute/varbinview.rs +++ b/vortex-array/src/arrays/filter/execute/varbinview.rs @@ -3,48 +3,39 @@ use std::sync::Arc; -use arrow_array::BooleanArray; -use vortex_error::VortexExpect; -use vortex_mask::Mask; +use vortex_buffer::Buffer; use vortex_mask::MaskValues; -use crate::ArrayRef; -use crate::IntoArray; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; -use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; -use crate::arrow::ArrowSessionExt; -use crate::arrow::FromArrowArray; +use crate::arrays::filter::execute::buffer; +use crate::arrays::filter::execute::filter_validity; +use crate::arrays::varbinview::BinaryView; +use crate::arrays::varbinview::VarBinViewArrayExt; +use crate::buffer::BufferHandle; pub fn filter_varbinview(array: &VarBinViewArray, mask: &Arc) -> VarBinViewArray { - // Delegate to the Arrow implementation of filter over `VarBinView`. - arrow_filter_fn(&array.clone().into_array(), &Mask::Values(Arc::clone(mask))) - .vortex_expect("VarBinViewArray is Arrow-compatible and supports arrow_filter_fn") - .as_::() - .into_owned() -} - -fn arrow_filter_fn(array: &ArrayRef, mask: &Mask) -> vortex_error::VortexResult { - let values = match &mask { - Mask::Values(values) => values, - Mask::AllTrue(_) | Mask::AllFalse(_) => unreachable!("check in filter invoke"), - }; - - let array_ref = LEGACY_SESSION.arrow().execute_arrow( - array.clone(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; - let mask_array = BooleanArray::new(values.bit_buffer().clone().into(), None); - let filtered = arrow_select::filter::filter(array_ref.as_ref(), &mask_array)?; - - ArrayRef::from_arrow(filtered.as_ref(), array.dtype().is_nullable()) + let filtered_validity = filter_validity(array.varbinview_validity(), mask); + + let views = Buffer::::from_byte_buffer(array.views_handle().as_host().clone()); + let filtered_views = buffer::filter_buffer(views, mask.as_ref()); + + // SAFETY: the filtered views are a subset of the original views and reference the same data + // buffers, and the validity is filtered by the same mask so lengths stay aligned. + unsafe { + VarBinViewArray::new_handle_unchecked( + BufferHandle::new_host(filtered_views.into_byte_buffer()), + Arc::clone(array.data_buffers()), + array.dtype().clone(), + filtered_validity, + ) + } } #[cfg(test)] mod test { use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::VarBinViewArray; use crate::compute::conformance::filter::test_filter_conformance; @@ -52,6 +43,7 @@ mod test { fn test_filter_varbinview_conformance() { test_filter_conformance( &VarBinViewArray::from_iter_str(["one", "two", "three", "four", "five"]).into_array(), + &mut array_session().create_execution_ctx(), ); test_filter_conformance( @@ -63,6 +55,7 @@ mod test { Some("five"), ]) .into_array(), + &mut array_session().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/fixed_size_list/array.rs b/vortex-array/src/arrays/fixed_size_list/array.rs index f24001c4e78..b6328163727 100644 --- a/vortex-array/src/arrays/fixed_size_list/array.rs +++ b/vortex-array/src/arrays/fixed_size_list/array.rs @@ -12,8 +12,6 @@ use vortex_error::vortex_ensure; use crate::ArrayRef; use crate::ArraySlots; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; @@ -231,6 +229,7 @@ pub trait FixedSizeListArrayExt: TypedArrayRef { child_to_validity(self.as_ref().slots()[VALIDITY_SLOT].as_ref(), nullability) } + #[allow(clippy::disallowed_methods)] fn fixed_size_list_elements_at(&self, index: usize) -> VortexResult { debug_assert!( index < self.as_ref().len(), @@ -238,14 +237,6 @@ pub trait FixedSizeListArrayExt: TypedArrayRef { index, self.as_ref().len(), ); - #[expect(clippy::debug_assert_with_mut_call)] - { - debug_assert!( - self.fixed_size_list_validity() - .execute_is_valid(index, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap_or(false) - ); - } let start = self.list_size() as usize * index; let end = self.list_size() as usize * (index + 1); 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 6f609120901..69677400692 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs @@ -182,7 +182,7 @@ fn test_filter_nested_fixed_size_lists() { #[case(create_fsl_single_element())] #[case(create_fsl_empty())] fn test_filter_fsl_conformance(#[case] array: ArrayRef) { - test_filter_conformance(&array); + test_filter_conformance(&array, &mut array_session().create_execution_ctx()); } // Helper functions for creating test arrays. 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 44420d22cb3..40fd2bc066a 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs @@ -6,8 +6,6 @@ use std::sync::Arc; use vortex_buffer::buffer; use crate::IntoArray; -#[expect(deprecated)] -use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::FixedSizeListArray; @@ -85,9 +83,10 @@ fn test_fsl_of_fsl_basic() { let first_outer_list = outer_fsl.fixed_size_list_elements_at(0).unwrap(); // Check first inner list [1,2]. - #[expect(deprecated)] let inner_list_0 = first_outer_list - .to_fixed_size_list() + .clone() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() .fixed_size_list_elements_at(0) .unwrap(); assert_eq!( @@ -104,9 +103,10 @@ fn test_fsl_of_fsl_basic() { ); // Check second inner list [3,4]. - #[expect(deprecated)] let inner_list_1 = first_outer_list - .to_fixed_size_list() + .clone() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() .fixed_size_list_elements_at(1) .unwrap(); assert_eq!( @@ -123,9 +123,9 @@ fn test_fsl_of_fsl_basic() { ); // Check third inner list [5,6]. - #[expect(deprecated)] let inner_list_2 = first_outer_list - .to_fixed_size_list() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() .fixed_size_list_elements_at(2) .unwrap(); assert_eq!( @@ -145,9 +145,10 @@ fn test_fsl_of_fsl_basic() { let second_outer_list = outer_fsl.fixed_size_list_elements_at(1).unwrap(); // Check first inner list [7,8]. - #[expect(deprecated)] let inner_list_0 = second_outer_list - .to_fixed_size_list() + .clone() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() .fixed_size_list_elements_at(0) .unwrap(); assert_eq!( @@ -164,9 +165,10 @@ fn test_fsl_of_fsl_basic() { ); // Check second inner list [9,10]. - #[expect(deprecated)] let inner_list_1 = second_outer_list - .to_fixed_size_list() + .clone() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() .fixed_size_list_elements_at(1) .unwrap(); assert_eq!( @@ -183,9 +185,9 @@ fn test_fsl_of_fsl_basic() { ); // Check third inner list [11,12]. - #[expect(deprecated)] let inner_list_2 = second_outer_list - .to_fixed_size_list() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() .fixed_size_list_elements_at(2) .unwrap(); assert_eq!( @@ -307,21 +309,23 @@ fn test_deeply_nested_fsl() { // Check the actual deeply nested values. // Structure: [[[1,2],[3,4]],[[5,6],[7,8]]]. let top_level = level3.fixed_size_list_elements_at(0).unwrap(); - #[expect(deprecated)] let level2_0 = top_level - .to_fixed_size_list() + .clone() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() .fixed_size_list_elements_at(0) .unwrap(); - #[expect(deprecated)] let level2_1 = top_level - .to_fixed_size_list() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() .fixed_size_list_elements_at(1) .unwrap(); // First level-2 list: [[1,2],[3,4]]. - #[expect(deprecated)] let level1_0_0 = level2_0 - .to_fixed_size_list() + .clone() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() .fixed_size_list_elements_at(0) .unwrap(); assert_eq!( @@ -337,9 +341,9 @@ fn test_deeply_nested_fsl() { 2i32.into() ); - #[expect(deprecated)] let level1_0_1 = level2_0 - .to_fixed_size_list() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() .fixed_size_list_elements_at(1) .unwrap(); assert_eq!( @@ -356,9 +360,10 @@ fn test_deeply_nested_fsl() { ); // Second level-2 list: [[5,6],[7,8]]. - #[expect(deprecated)] let level1_1_0 = level2_1 - .to_fixed_size_list() + .clone() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() .fixed_size_list_elements_at(0) .unwrap(); assert_eq!( @@ -374,9 +379,9 @@ fn test_deeply_nested_fsl() { 6i32.into() ); - #[expect(deprecated)] let level1_1_1 = level2_1 - .to_fixed_size_list() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap() .fixed_size_list_elements_at(1) .unwrap(); assert_eq!( 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 e75139e5fc8..af8f94f2d5f 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -33,7 +33,10 @@ use crate::validity::Validity; #[case::single_element(create_single_element_fsl())] #[case::empty(create_empty_fsl())] fn test_take_fsl_conformance(#[case] fsl: FixedSizeListArray) { - test_take_conformance(&fsl.into_array()); + test_take_conformance( + &fsl.into_array(), + &mut array_session().create_execution_ctx(), + ); } // FSL-specific edge case tests that aren't covered by conformance. @@ -88,7 +91,10 @@ fn test_take_degenerate_lists( let elements = PrimitiveArray::empty::(Nullability::NonNullable); let fsl = FixedSizeListArray::new(elements.into_array(), 0, validity, 5); - test_take_conformance(&fsl.clone().into_array()); + test_take_conformance( + &fsl.clone().into_array(), + &mut array_session().create_execution_ctx(), + ); // Also test the specific behavior. let indices_array = PrimitiveArray::from_option_iter(indices); 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 8f6fef6a909..9393b8563d4 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs @@ -31,6 +31,8 @@ use crate::arrays::fixed_size_list::array::NUM_SLOTS; use crate::arrays::fixed_size_list::array::SLOT_NAMES; use crate::arrays::fixed_size_list::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::FixedSizeListBuilder; use crate::dtype::DType; use crate::serde::ArrayChildren; use crate::validity::Validity; @@ -201,4 +203,15 @@ impl VTable for FixedSizeList { fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done(array)) } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for FixedSizeList requires a FixedSizeListBuilder"); + }; + builder.append_fixed_size_list_array(&array.into_owned(), ctx) + } } diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index 945991b7e03..8350b1a3558 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -18,7 +18,6 @@ use crate::ArraySlots; use crate::Canonical; 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; @@ -29,10 +28,12 @@ use crate::array::child_to_validity; use crate::array::validity_to_child; use crate::arrays::ConstantArray; use crate::arrays::List; +use crate::arrays::ListArray; use crate::arrays::Primitive; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::NativePType; +use crate::legacy_session; use crate::match_each_integer_ptype; use crate::match_each_native_ptype; use crate::scalar_fn::fns::operators::Operator; @@ -182,6 +183,7 @@ impl ListData { /// Validates the components that would be used to create a `ListArray`. /// /// This function checks all the invariants required by `ListArray::new_unchecked`. + #[allow(clippy::disallowed_methods)] pub fn validate( elements: &ArrayRef, offsets: &ArrayRef, @@ -202,7 +204,7 @@ impl ListData { // We can safely unwrap the DType as primitive now let offsets_ptype = offsets.dtype().as_ptype(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = legacy_session().create_execution_ctx(); // Offsets must be sorted (but not strictly sorted, zero-length lists are allowed) if let Some(is_sorted) = offsets.statistics().compute_is_sorted(&mut ctx) { @@ -296,6 +298,7 @@ pub trait ListArrayExt: TypedArrayRef { ) } + #[allow(clippy::disallowed_methods)] fn offset_at(&self, index: usize) -> VortexResult { vortex_ensure!( index <= self.as_ref().len(), @@ -309,7 +312,7 @@ pub trait ListArrayExt: TypedArrayRef { })) } else { self.offsets() - .execute_scalar(index, &mut LEGACY_SESSION.create_execution_ctx())? + .execute_scalar(index, &mut legacy_session().create_execution_ctx())? .as_primitive() .as_::() .ok_or_else(|| vortex_error::vortex_err!("offset value does not fit in usize")) @@ -336,7 +339,6 @@ pub trait ListArrayExt: TypedArrayRef { let mut elements = self.sliced_elements()?; if recurse && elements.is_canonical() { let compacted = elements - .clone() .execute::(ctx)? .compact(ctx)? .into_array(); @@ -355,7 +357,8 @@ pub trait ListArrayExt: TypedArrayRef { Operator::Sub, )?; - Array::::try_new(elements, adjusted_offsets, self.list_validity()) + // SAFETY: By resetting the offsets we simply "shift" everything left and discard trailing garbage, so all invariants remain the same. + Ok(unsafe { ListArray::new_unchecked(elements, adjusted_offsets, self.list_validity()) }) } } impl> ListArrayExt for T {} diff --git a/vortex-array/src/arrays/list/compute/cast.rs b/vortex-array/src/arrays/list/compute/cast.rs index 084d9613e75..9fd27f731f1 100644 --- a/vortex-array/src/arrays/list/compute/cast.rs +++ b/vortex-array/src/arrays/list/compute/cast.rs @@ -181,7 +181,7 @@ mod tests { #[case(create_nested_list())] #[case(create_empty_lists())] fn test_cast_list_conformance(#[case] array: ListArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } fn create_simple_list() -> ListArray { diff --git a/vortex-array/src/arrays/list/compute/mod.rs b/vortex-array/src/arrays/list/compute/mod.rs index d8ddbd37583..ce6cd960f74 100644 --- a/vortex-array/src/arrays/list/compute/mod.rs +++ b/vortex-array/src/arrays/list/compute/mod.rs @@ -37,7 +37,10 @@ mod tests { let array = ListArray::try_new(elements.into_array(), offsets.into_array(), validity).unwrap(); - test_mask_conformance(&array.into_array()); + test_mask_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] @@ -48,7 +51,10 @@ mod tests { let array = ListArray::try_new(elements.into_array(), offsets.into_array(), validity).unwrap(); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[rstest] diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 2c6089e0c58..14e89d4c238 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -212,12 +212,11 @@ mod test { use vortex_buffer::buffer; use crate::IntoArray as _; - #[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; use crate::arrays::PrimitiveArray; use crate::compute::conformance::take::test_take_conformance; use crate::dtype::DType; @@ -228,6 +227,7 @@ mod test { #[test] fn nullable_take() { + let mut ctx = array_session().create_execution_ctx(); let list = ListArray::try_new( buffer![0i32, 5, 3, 4].into_array(), buffer![0, 2, 3, 4, 4].into_array(), @@ -249,8 +249,7 @@ mod test { ) ); - #[expect(deprecated)] - let result = result.to_listview(); + let result = result.execute::(&mut ctx).unwrap(); assert_eq!(result.len(), 4); @@ -332,6 +331,7 @@ mod test { #[test] fn non_nullable_take() { + let mut ctx = array_session().create_execution_ctx(); let list = ListArray::try_new( buffer![0i32, 5, 3, 4].into_array(), buffer![0, 2, 3, 3, 4].into_array(), @@ -352,8 +352,7 @@ mod test { ) ); - #[expect(deprecated)] - let result = result.to_listview(); + let result = result.execute::(&mut ctx).unwrap(); assert_eq!(result.len(), 3); @@ -466,11 +465,15 @@ mod test { Validity::NonNullable, ).unwrap())] fn test_take_list_conformance(#[case] list: ListArray) { - test_take_conformance(&list.into_array()); + test_take_conformance( + &list.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] fn test_u64_offset_accumulation_non_nullable() { + let mut ctx = array_session().create_execution_ctx(); let elements = buffer![0i32; 200].into_array(); let offsets = buffer![0u8, 200].into_array(); let list = ListArray::try_new(elements, offsets, Validity::NonNullable) @@ -483,8 +486,7 @@ mod test { assert_eq!(result.len(), 2); - #[expect(deprecated)] - let result_view = result.to_listview(); + let result_view = result.execute::(&mut ctx).unwrap(); assert_eq!(result_view.len(), 2); assert!( result_view @@ -500,6 +502,7 @@ mod test { #[test] fn test_u64_offset_accumulation_nullable() { + let mut ctx = array_session().create_execution_ctx(); let elements = buffer![0i32; 150].into_array(); let offsets = buffer![0u8, 150, 150].into_array(); let validity = BoolArray::from_iter(vec![true, false]).into_array(); @@ -513,8 +516,7 @@ mod test { assert_eq!(result.len(), 3); - #[expect(deprecated)] - let result_view = result.to_listview(); + let result_view = result.execute::(&mut ctx).unwrap(); assert_eq!(result_view.len(), 3); assert!( result_view diff --git a/vortex-array/src/arrays/list/test_harness.rs b/vortex-array/src/arrays/list/test_harness.rs index 42351484c19..ba2c9d335f1 100644 --- a/vortex-array/src/arrays/list/test_harness.rs +++ b/vortex-array/src/arrays/list/test_harness.rs @@ -6,7 +6,6 @@ use std::sync::Arc; use itertools::Itertools; use vortex_error::VortexResult; -use crate::ArrayRef; use crate::arrays::ListArray; use crate::builders::ArrayBuilder; use crate::builders::ListBuilder; @@ -21,7 +20,7 @@ impl ListArray { pub fn from_iter_slow( iter: I, dtype: Arc, - ) -> VortexResult + ) -> VortexResult where I::Item: IntoIterator, ::Item: Into, @@ -42,13 +41,13 @@ impl ListArray { ); builder.append_value(elem.as_list())? } - Ok(builder.finish()) + Ok(builder.finish_into_list()) } pub fn from_iter_opt_slow>, T>( iter: I, dtype: Arc, - ) -> VortexResult + ) -> VortexResult where T: IntoIterator, T::Item: Into, @@ -73,6 +72,6 @@ impl ListArray { builder.append_null() } } - Ok(builder.finish()) + Ok(builder.finish_into_list()) } } diff --git a/vortex-array/src/arrays/list/vtable/mod.rs b/vortex-array/src/arrays/list/vtable/mod.rs index 79316e76be7..55f7753c995 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -35,6 +35,7 @@ use crate::arrays::list::array::SLOT_NAMES; use crate::arrays::list::compute::rules::PARENT_RULES; use crate::arrays::listview::list_view_from_list; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -199,6 +200,14 @@ impl VTable for List { list_view_from_list(array, ctx)?.into_array(), )) } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + builder.append_list_array(array, ctx) + } } #[derive(Clone, Debug)] diff --git a/vortex-array/src/arrays/listview/array.rs b/vortex-array/src/arrays/listview/array.rs index 1d6297ba7c1..daf7100dc59 100644 --- a/vortex-array/src/arrays/listview/array.rs +++ b/vortex-array/src/arrays/listview/array.rs @@ -18,9 +18,6 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::ArraySlots; use crate::ExecutionCtx; -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; @@ -39,6 +36,7 @@ use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::PType; use crate::expr::stats::Stat; +use crate::legacy_session; use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; use crate::scalar_fn::fns::operators::Operator; @@ -259,10 +257,10 @@ impl ListViewData { // Skip host-only validation when offsets/sizes are not host-resident. if offsets.is_host() && sizes.is_host() { - #[expect(deprecated)] - let offsets_primitive = offsets.to_primitive(); - #[expect(deprecated)] - let sizes_primitive = sizes.to_primitive(); + #[allow(clippy::disallowed_methods)] + let mut ctx = legacy_session().create_execution_ctx(); + let offsets_primitive = offsets.clone().execute::(&mut ctx)?; + let sizes_primitive = sizes.clone().execute::(&mut ctx)?; // Offsets and sizes are non-negative; reinterpret to unsigned to dispatch over 4 widths // each (4x4 instead of 8x8). This is a read-only validation, so result types are moot. let offsets_primitive = @@ -382,6 +380,7 @@ pub trait ListViewArrayExt: TypedArrayRef { ) } + #[allow(clippy::disallowed_methods)] fn offset_at(&self, index: usize) -> usize { assert!( index < self.as_ref().len(), @@ -393,7 +392,7 @@ pub trait ListViewArrayExt: TypedArrayRef { .map(|p| match_each_integer_ptype!(p.ptype(), |P| { p.as_slice::

()[index].as_() })) .unwrap_or_else(|| { self.offsets() - .execute_scalar(index, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(index, &mut legacy_session().create_execution_ctx()) .vortex_expect("offsets must support execute_scalar") .as_primitive() .as_::() @@ -401,6 +400,7 @@ pub trait ListViewArrayExt: TypedArrayRef { }) } + #[allow(clippy::disallowed_methods)] fn size_at(&self, index: usize) -> usize { assert!( index < self.as_ref().len(), @@ -413,7 +413,7 @@ pub trait ListViewArrayExt: TypedArrayRef { .map(|p| match_each_integer_ptype!(p.ptype(), |P| { p.as_slice::

()[index].as_() })) .unwrap_or_else(|| { self.sizes() - .execute_scalar(index, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(index, &mut legacy_session().create_execution_ctx()) .vortex_expect("sizes must support execute_scalar") .as_primitive() .as_::() @@ -427,14 +427,6 @@ pub trait ListViewArrayExt: TypedArrayRef { self.elements().slice(offset..offset + size) } - fn verify_is_zero_copy_to_list(&self) -> bool { - #[expect(deprecated)] - let offsets_primitive = self.offsets().to_primitive(); - #[expect(deprecated)] - let sizes_primitive = self.sizes().to_primitive(); - validate_zctl(self.elements(), offsets_primitive, sizes_primitive).is_ok() - } - /// Returns a [`Mask`] of length `elements.len()` where each bit is set iff that /// position in `elements` is referenced by at least one view. Caller must ensure `elements` /// is non-empty. @@ -648,10 +640,18 @@ impl Array { /// See [`ListViewData::with_zero_copy_to_list`]. pub unsafe fn with_zero_copy_to_list(self, is_zctl: bool) -> Self { if cfg!(debug_assertions) && is_zctl { - #[expect(deprecated)] - let offsets_primitive = self.offsets().to_primitive(); - #[expect(deprecated)] - let sizes_primitive = self.sizes().to_primitive(); + #[allow(clippy::disallowed_methods)] + let mut ctx = legacy_session().create_execution_ctx(); + let offsets_primitive = self + .offsets() + .clone() + .execute::(&mut ctx) + .vortex_expect("offsets must canonicalize to primitive"); + let sizes_primitive = self + .sizes() + .clone() + .execute::(&mut ctx) + .vortex_expect("sizes must canonicalize to primitive"); validate_zctl(self.elements(), offsets_primitive, sizes_primitive) .vortex_expect("Failed to validate zero-copy to list flag"); } @@ -740,6 +740,7 @@ where /// Helper function to validate if the `ListViewArray` components are actually zero-copyable to /// [`ListArray`](crate::arrays::ListArray). +#[allow(clippy::disallowed_methods)] fn validate_zctl( elements: &ArrayRef, offsets_primitive: PrimitiveArray, @@ -747,7 +748,7 @@ fn validate_zctl( ) -> VortexResult<()> { // Offsets must be sorted (but not strictly sorted, zero-length lists are allowed), even // if there are null views. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = legacy_session().create_execution_ctx(); if let Some(is_sorted) = offsets_primitive.statistics().compute_is_sorted(&mut ctx) { vortex_ensure!(is_sorted, "offsets must be sorted"); } else { diff --git a/vortex-array/src/arrays/listview/compute/zip.rs b/vortex-array/src/arrays/listview/compute/zip.rs index 46537f1fafe..71640d4fb21 100644 --- a/vortex-array/src/arrays/listview/compute/zip.rs +++ b/vortex-array/src/arrays/listview/compute/zip.rs @@ -1,12 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; use std::ops::BitAnd; use std::ops::BitOr; use std::ops::Not; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -48,7 +50,7 @@ impl ZipKernel for ListView { }; // Null mask entries select `if_false`, matching `Zip`'s SQL ELSE semantics. - let mask = mask.try_to_mask_fill_null_false(ctx)?; + let mask = mask.clone().null_as_false().execute(ctx)?; match &mask { // Defer the trivial masks to the generic zip, which just casts one side. Mask::AllTrue(_) | Mask::AllFalse(_) => return Ok(None), @@ -85,24 +87,65 @@ impl ZipKernel for ListView { let mut offsets = BufferMut::::with_capacity(len); let mut sizes = BufferMut::::with_capacity(len); - for ((idx, (out_offsets, out_sizes)), selected) in offsets - .spare_capacity_mut() - .iter_mut() - .zip(sizes.spare_capacity_mut().iter_mut()) - .take(len) - .enumerate() - .zip(mask.iter()) { - if selected { - out_offsets.write(true_offsets[idx]); - out_sizes.write(true_sizes[idx]); - } else { - out_offsets.write(false_offsets[idx] + false_shift); - out_sizes.write(false_sizes[idx]); + let true_offsets = true_offsets.as_slice(); + let true_sizes = true_sizes.as_slice(); + let false_offsets = false_offsets.as_slice(); + let false_sizes = false_sizes.as_slice(); + + let offsets_out = offsets.spare_capacity_mut(); + let sizes_out = sizes.spare_capacity_mut(); + + // We matched `Mask::Values` above, so the bit buffer is materialized. `unaligned_chunks` + // iterates faster than `chunks`: it exposes the byte-aligned body as a plain `&[u64]` + // with no per-word reshifting, isolating any bit misalignment into a leading `prefix` + // and trailing `suffix` word. We blend both sides branchlessly per row so the compiler + // vectorizes the inner select instead of mispredicting a data-dependent branch. + let mask_bits = mask + .values() + .vortex_expect("mask is Mask::Values") + .bit_buffer(); + let unaligned = mask_bits.unaligned_chunks(); + // The prefix word's low `lead` bits are padding; shifting them out aligns row 0 to bit 0, + // after which every chunk and the suffix start cleanly on a row boundary. + let lead = unaligned.lead_padding(); + + let mut select_block = |word: u64, base: usize, n: usize| { + let end = base + n; + // `if_false` views address the second half of the concatenated elements, so shift + // their offsets by `false_shift`; sizes are taken verbatim from the chosen side. + select_column( + word, + &true_offsets[base..end], + &false_offsets[base..end], + false_shift, + &mut offsets_out[base..end], + ); + select_column( + word, + &true_sizes[base..end], + &false_sizes[base..end], + 0, + &mut sizes_out[base..end], + ); + }; + + let mut base = 0; + if let Some(prefix) = unaligned.prefix() { + let n = (64 - lead).min(len); + select_block(prefix >> lead, base, n); + base += n; + } + for &word in unaligned.chunks() { + select_block(word, base, 64); + base += 64; + } + if let Some(suffix) = unaligned.suffix() { + select_block(suffix, base, len - base); } } - // SAFETY: the loop above initialized exactly `len` slots in both buffers. + // SAFETY: `select_column` initialized exactly `len` slots in both buffers. unsafe { offsets.set_len(len); sizes.set_len(len); @@ -122,6 +165,30 @@ impl ZipKernel for ListView { } } +/// Branchlessly select one `u64` column per row from `if_true` or `if_false`. +/// +/// `word` holds the mask bits for this block, bit `j` (LSB-first) selecting row `j`: a set bit keeps +/// `true_vals[j]`, an unset bit keeps `false_vals[j] + false_add`. The bit is expanded to a +/// full-width lane mask and blended, so the inner loop is branch-free and auto-vectorizable. Inputs +/// are sliced to the output length up front so the compiler can elide bounds checks across the block. +#[inline] +fn select_column( + word: u64, + true_vals: &[u64], + false_vals: &[u64], + false_add: u64, + out: &mut [MaybeUninit], +) { + let n = out.len(); + let true_vals = &true_vals[..n]; + let false_vals = &false_vals[..n]; + for j in 0..n { + // 0 for an unset bit, `u64::MAX` for a set bit. + let lane = 0u64.wrapping_sub((word >> j) & 1); + out[j].write((true_vals[j] & lane) | ((false_vals[j] + false_add) & !lane)); + } +} + /// Appends `array`'s element chunks to `chunks`, flattening a top-level [`ChunkedArray`] so the /// concatenated elements never nest chunked arrays. fn push_element_chunks(array: ArrayRef, chunks: &mut Vec) { @@ -164,6 +231,12 @@ fn zip_validity( #[cfg(test)] mod tests { + #![allow( + clippy::cast_possible_truncation, + reason = "test fixtures use small indices that fit the target widths" + )] + + use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -200,6 +273,7 @@ mod tests { /// `zip` of two list views selects whole lists per the mask and keeps the list encoding. #[test] fn zip_selects_lists() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // [[1, 2], [3], [4, 5, 6]] let if_true = list_view( buffer![1i32, 2, 3, 4, 5, 6].into_array(), @@ -216,7 +290,6 @@ mod tests { ); let mask = Mask::from_iter([true, false, true]); - let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -239,6 +312,7 @@ mod tests { /// `zip` selects list-level validity from the chosen side and widens nullability. #[test] fn zip_selects_validity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // [[1], null, [2]] (list-level nulls) let if_true = list_view( buffer![1i32, 2].into_array(), @@ -256,7 +330,6 @@ mod tests { // true -> if_true, false -> if_false let mask = Mask::from_iter([false, true, true]); - let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -277,6 +350,7 @@ mod tests { /// is nullable. #[test] fn zip_out_of_order_offsets_and_widening() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // [[5, 6], [7], [8, 9]] expressed with out-of-order offsets. let if_true = list_view( buffer![7i32, 8, 9, 5, 6].into_array(), @@ -293,7 +367,6 @@ mod tests { ); let mask = Mask::from_iter([true, true, false]); - let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -311,10 +384,153 @@ mod tests { Ok(()) } + /// Zipping more rows than fit in a single 64-bit mask chunk exercises both the chunked select + /// loop and the trailing remainder, including the `false_shift` applied to `if_false` views. + #[test] + fn zip_spans_multiple_mask_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // 130 single-element lists per side: `if_true[i] = [i]`, `if_false[i] = [1000 + i]`. + let len = 130usize; + let true_elements: Vec = (0..len as i32).collect(); + let false_elements: Vec = (0..len as i32).map(|i| 1000 + i).collect(); + let offsets: Vec = (0..len as u64).collect(); + let sizes: Vec = vec![1; len]; + + let if_true = list_view( + true_elements + .iter() + .copied() + .collect::>() + .into_array(), + offsets + .iter() + .copied() + .collect::>() + .into_array(), + sizes.iter().copied().collect::>().into_array(), + Validity::NonNullable, + ); + let if_false = list_view( + false_elements + .iter() + .copied() + .collect::>() + .into_array(), + offsets + .iter() + .copied() + .collect::>() + .into_array(), + sizes.iter().copied().collect::>().into_array(), + Validity::NonNullable, + ); + + // A non-trivial pattern that straddles the chunk boundary (index 63/64) and the remainder. + let mask_bits: Vec = (0..len).map(|i| i.is_multiple_of(3) || i == 64).collect(); + let mask = Mask::from_iter(mask_bits.iter().copied()); + + let result = mask + .into_array() + .zip(if_true, if_false)? + .execute::(&mut ctx)?; + assert!(result.is::()); + + // Each row collapses to a single element: `i` when the mask is set, else `1000 + i`. + let expected_elements: Vec = (0..len) + .map(|i| { + if mask_bits[i] { + i as i32 + } else { + 1000 + i as i32 + } + }) + .collect(); + let expected = list_view( + expected_elements + .iter() + .copied() + .collect::>() + .into_array(), + offsets + .iter() + .copied() + .collect::>() + .into_array(), + sizes.iter().copied().collect::>().into_array(), + Validity::NonNullable, + ); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + /// A mask whose bit buffer starts at a non-byte-aligned offset (here from slicing a bool array) + /// has non-zero `unaligned_chunks` lead padding, exercising the prefix word alongside the + /// aligned chunk body and the suffix. + #[test] + fn zip_handles_offset_mask() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // 200 single-element lists per side: `if_true[i] = [i]`, `if_false[i] = [1000 + i]`. With a + // 3-bit lead offset the mask spans more than 16 bytes, so `unaligned_chunks` exposes a + // non-empty aligned `chunks` body between the prefix and suffix words. + let len = 200usize; + let true_elements: Vec = (0..len as i32).collect(); + let false_elements: Vec = (0..len as i32).map(|i| 1000 + i).collect(); + let offsets: Vec = (0..len as u64).collect(); + let sizes: Vec = vec![1; len]; + + let single_element_view = |elements: &[i32]| { + list_view( + elements + .iter() + .copied() + .collect::>() + .into_array(), + offsets + .iter() + .copied() + .collect::>() + .into_array(), + sizes.iter().copied().collect::>().into_array(), + Validity::NonNullable, + ) + }; + let if_true = single_element_view(&true_elements); + let if_false = single_element_view(&false_elements); + + // Slice off the first `offset` bits so the mask's bit buffer keeps a sub-byte offset while + // remaining `len` rows long. A non-trivial pattern straddles the prefix/body and chunk + // boundaries within the sliced window. + let offset = 3usize; + let mask_bits: Vec = (0..offset + len) + .map(|i| i.is_multiple_of(3) || i == offset + 64) + .collect(); + let mask = BoolArray::from_iter(mask_bits.iter().copied()) + .into_array() + .slice(offset..offset + len)?; + + let result = mask.zip(if_true, if_false)?.execute::(&mut ctx)?; + assert!(result.is::()); + + // Each row collapses to a single element: `i` when the sliced mask is set, else `1000 + i`. + let expected_elements: Vec = (0..len) + .map(|i| { + if mask_bits[offset + i] { + i as i32 + } else { + 1000 + i as i32 + } + }) + .collect(); + let expected = single_element_view(&expected_elements); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + /// When an input's `elements` is already a [`ChunkedArray`], its chunks are spliced in rather /// than nesting a chunked array inside the concatenated elements. #[test] fn zip_flattens_chunked_elements() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // elements [1, 2, 3] stored as two chunks; lists [[1, 2], [3]]. let chunked_elements = ChunkedArray::try_new( vec![buffer![1i32, 2].into_array(), buffer![3i32].into_array()], @@ -336,7 +552,6 @@ mod tests { ); let mask = Mask::from_iter([true, false]); - let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? diff --git a/vortex-array/src/arrays/listview/mod.rs b/vortex-array/src/arrays/listview/mod.rs index 80b2540a897..b240854c6cf 100644 --- a/vortex-array/src/arrays/listview/mod.rs +++ b/vortex-array/src/arrays/listview/mod.rs @@ -12,7 +12,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::ListView; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &VortexSession) { vtable::initialize(session); } @@ -25,6 +25,7 @@ mod rebuild; pub use rebuild::DEFAULT_REBUILD_DENSITY_THRESHOLD; pub use rebuild::DEFAULT_TRIM_ELEMENTS_THRESHOLD; pub use rebuild::ListViewRebuildMode; +use vortex_session::VortexSession; #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index bae2c4f3fb9..342c9eccf19 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -253,7 +253,7 @@ impl ListViewArray { } /// Rebuilds elements list-by-list: canonicalize elements upfront, then for each list `slice` - /// the relevant range and `extend_from_array` into a typed builder. + /// the relevant range and `append_to_builder` into a typed builder. fn rebuild_list_by_list( &self, ctx: &mut ExecutionCtx, @@ -317,7 +317,9 @@ impl ListViewArray { new_offsets.push(n_elements); new_sizes.push(size); - new_elements_builder.extend_from_array(&elements_canonical.slice(start..stop)?); + elements_canonical + .slice(start..stop)? + .append_to_builder(new_elements_builder.as_mut(), ctx)?; n_elements += num_traits::cast(size).vortex_expect("Cast failed"); } diff --git a/vortex-array/src/arrays/listview/tests/basic.rs b/vortex-array/src/arrays/listview/tests/basic.rs index 3d4d6db284d..f3082757ca0 100644 --- a/vortex-array/src/arrays/listview/tests/basic.rs +++ b/vortex-array/src/arrays/listview/tests/basic.rs @@ -401,9 +401,15 @@ fn test_verify_is_zero_copy_to_list() { let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); - // Should return true since offsets are sorted and no overlaps exist. - assert!(listview.verify_is_zero_copy_to_list()); + // Marking as zero-copy-to-list validates the components and should succeed since offsets + // are sorted and no overlaps exist. + let zctl = unsafe { listview.with_zero_copy_to_list(true) }; + assert!(zctl.is_zero_copy_to_list()); +} +#[test] +#[should_panic(expected = "Zero-copy-to-list requires views to be non-overlapping and ordered")] +fn test_verify_is_zero_copy_to_list_overlapping() { // Create a ListView that is NOT zero-copyable to List due to overlapping views. // Logical lists: [[1,2], [2,3,4], [3,4]] let elements = buffer![1i32, 2, 3, 4, 5].into_array(); @@ -412,8 +418,10 @@ fn test_verify_is_zero_copy_to_list() { let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); - // Should return false due to overlapping list views. - assert!(!listview.verify_is_zero_copy_to_list()); + // Validation should reject overlapping list views. + unsafe { + let _zctl = listview.with_zero_copy_to_list(true); + } } #[test] @@ -460,9 +468,6 @@ fn test_validate_monotonic_ends_correct_nulls() { let listview = ListViewArray::new(elements, offsets, sizes, validity); // Should be valid as zero-copy-to-list - this should NOT panic - let zctl_listview = unsafe { listview.clone().with_zero_copy_to_list(true) }; + let zctl_listview = unsafe { listview.with_zero_copy_to_list(true) }; assert!(zctl_listview.is_zero_copy_to_list()); - - // verify_is_zero_copy_to_list should also return true - assert!(listview.verify_is_zero_copy_to_list()); } diff --git a/vortex-array/src/arrays/listview/tests/filter.rs b/vortex-array/src/arrays/listview/tests/filter.rs index 25cc0de4aa8..e8cc7636b30 100644 --- a/vortex-array/src/arrays/listview/tests/filter.rs +++ b/vortex-array/src/arrays/listview/tests/filter.rs @@ -33,7 +33,7 @@ static SESSION: LazyLock = LazyLock::new(crate::array_session); #[case::overlapping(create_overlapping_listview())] #[case::large(create_large_listview())] fn test_filter_listview_conformance(#[case] listview: ListViewArray) { - test_filter_conformance(&listview.into_array()); + test_filter_conformance(&listview.into_array(), &mut SESSION.create_execution_ctx()); } #[test] diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index f49680a854e..51f03b1abeb 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -13,8 +13,6 @@ use super::common::create_large_listview; use super::common::create_nullable_listview; use crate::ArrayRef; use crate::IntoArray; -#[expect(deprecated)] -use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_constant::is_constant; use crate::array_session; @@ -258,8 +256,9 @@ fn test_cast_numeric_types(#[case] from_ptype: PType, #[case] to_ptype: PType) { let result = listview.cast(target_dtype.clone()).unwrap(); assert_eq!(result.dtype(), &target_dtype); - #[expect(deprecated)] - let result_list = result.to_listview(); + let result_list = result + .execute::(&mut array_session().create_execution_ctx()) + .unwrap(); assert!( result_list.len() == 3 || result_list.len() == 2, "Expected 2 or 3 lists" @@ -295,8 +294,9 @@ fn test_cast_with_nulls() { let result = listview.cast(target_dtype.clone()).unwrap(); assert_eq!(result.dtype(), &target_dtype); - #[expect(deprecated)] - let result_list = result.to_listview(); + let result_list = result + .execute::(&mut array_session().create_execution_ctx()) + .unwrap(); assert!( result_list .is_valid(0, &mut array_session().create_execution_ctx()) @@ -346,8 +346,9 @@ fn test_cast_special_patterns(#[case] expected_sizes: Vec, #[case] list_c }; let result = listview.cast(target_dtype).unwrap(); - #[expect(deprecated)] - let result_list = result.to_listview(); + let result_list = result + .execute::(&mut array_session().create_execution_ctx()) + .unwrap(); assert_eq!(result_list.len(), list_count); @@ -379,8 +380,9 @@ fn test_cast_large_dataset() { ); let result = listview.cast(target_dtype).unwrap(); - #[expect(deprecated)] - let result_list = result.to_listview(); + let result_list = result + .execute::(&mut array_session().create_execution_ctx()) + .unwrap(); assert_eq!(result_list.len(), 20); for i in 0..20 { @@ -622,7 +624,10 @@ fn test_constant_repeated_same_lists() { #[case::nullable(create_nullable_listview())] #[case::large(create_large_listview())] fn test_mask_listview_conformance(#[case] listview: ListViewArray) { - test_mask_conformance(&listview.into_array()); + test_mask_conformance( + &listview.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] @@ -644,8 +649,9 @@ fn test_mask_preserves_structure() { let result = listview.mask((!&selection).into_array()).unwrap(); assert_eq!(result.len(), 4); // Length is preserved. - #[expect(deprecated)] - let result_list = result.to_listview(); + let result_list = result + .execute::(&mut array_session().create_execution_ctx()) + .unwrap(); // Check validity: true in selection means null. assert!( @@ -698,8 +704,9 @@ fn test_mask_with_existing_nulls() { // Mask additional elements. let selection = Mask::from_iter([false, true, true]); let result = listview.mask((!&selection).into_array()).unwrap(); - #[expect(deprecated)] - let result_list = result.to_listview(); + let result_list = result + .execute::(&mut array_session().create_execution_ctx()) + .unwrap(); // Check combined validity: assert!( @@ -731,8 +738,9 @@ fn test_mask_with_gaps() { let selection = Mask::from_iter([true, false, false]); let result = listview.mask((!&selection).into_array()).unwrap(); - #[expect(deprecated)] - let result_list = result.to_listview(); + let result_list = result + .execute::(&mut array_session().create_execution_ctx()) + .unwrap(); assert_eq!(result_list.len(), 3); assert!( @@ -776,8 +784,9 @@ fn test_mask_constant_arrays() { let selection = Mask::from_iter([false, true, false]); let result = const_list.mask((!&selection).into_array()).unwrap(); - #[expect(deprecated)] - let result_list = result.to_listview(); + let result_list = result + .execute::(&mut array_session().create_execution_ctx()) + .unwrap(); assert_eq!(result_list.len(), 3); assert!( diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index 2a68f711d30..2af397e262f 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -32,7 +32,7 @@ static SESSION: LazyLock = LazyLock::new(crate::array_session); #[case::overlapping(create_overlapping_listview())] #[case::large(create_large_listview())] fn test_take_listview_conformance(#[case] listview: ListViewArray) { - test_take_conformance(&listview.into_array()); + test_take_conformance(&listview.into_array(), &mut SESSION.create_execution_ctx()); } // ListView-specific tests that aren't covered by conformance. diff --git a/vortex-array/src/arrays/listview/vtable/mod.rs b/vortex-array/src/arrays/listview/vtable/mod.rs index f7a7a5486b6..546e77898f8 100644 --- a/vortex-array/src/arrays/listview/vtable/mod.rs +++ b/vortex-array/src/arrays/listview/vtable/mod.rs @@ -35,6 +35,7 @@ use crate::arrays::listview::array::SIZES_SLOT; use crate::arrays::listview::array::SLOT_NAMES; use crate::arrays::listview::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -225,6 +226,14 @@ impl VTable for ListView { Ok(ExecutionResult::done(array)) } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + builder.append_listview_array(array, ctx) + } + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/vortex-array/src/arrays/masked/array.rs b/vortex-array/src/arrays/masked/array.rs index 5ba830cc0e5..077446b3f02 100644 --- a/vortex-array/src/arrays/masked/array.rs +++ b/vortex-array/src/arrays/masked/array.rs @@ -9,7 +9,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use crate::ArrayRef; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; @@ -18,6 +17,7 @@ use crate::array::child_to_validity; use crate::array::validity_to_child; use crate::array_slots; use crate::arrays::Masked; +use crate::legacy_session; use crate::validity::Validity; #[array_slots(Masked)] @@ -75,13 +75,14 @@ impl MaskedData { impl Array { /// Constructs a new `MaskedArray`. + #[allow(clippy::disallowed_methods)] pub fn try_new(child: ArrayRef, validity: Validity) -> VortexResult { let dtype = child.dtype().as_nullable(); let len = child.len(); let validity_slot = validity_to_child(&validity, len); let data = MaskedData::try_new( len, - child.all_valid(&mut LEGACY_SESSION.create_execution_ctx())?, + child.all_valid(&mut legacy_session().create_execution_ctx())?, validity, )?; Ok(unsafe { diff --git a/vortex-array/src/arrays/masked/compute/filter.rs b/vortex-array/src/arrays/masked/compute/filter.rs index 9ea6b168105..287a996a12b 100644 --- a/vortex-array/src/arrays/masked/compute/filter.rs +++ b/vortex-array/src/arrays/masked/compute/filter.rs @@ -33,6 +33,8 @@ mod tests { use rstest::rstest; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::filter::test_filter_conformance; @@ -58,6 +60,9 @@ mod tests { ).unwrap() )] fn test_filter_masked_conformance(#[case] array: MaskedArray) { - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/masked/compute/mask.rs b/vortex-array/src/arrays/masked/compute/mask.rs index 514d0aeaafc..e830f15bfa0 100644 --- a/vortex-array/src/arrays/masked/compute/mask.rs +++ b/vortex-array/src/arrays/masked/compute/mask.rs @@ -34,6 +34,8 @@ mod tests { use rstest::rstest; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::mask::test_mask_conformance; @@ -59,6 +61,9 @@ mod tests { ).unwrap() )] fn test_mask_masked_conformance(#[case] array: MaskedArray) { - test_mask_conformance(&array.into_array()); + test_mask_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/masked/compute/take.rs b/vortex-array/src/arrays/masked/compute/take.rs index 1eff686b4d9..49ad0fc9fe6 100644 --- a/vortex-array/src/arrays/masked/compute/take.rs +++ b/vortex-array/src/arrays/masked/compute/take.rs @@ -44,6 +44,8 @@ mod tests { use rstest::rstest; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::take::test_take_conformance; @@ -69,6 +71,9 @@ mod tests { ).unwrap() )] fn test_take_masked_conformance(#[case] array: MaskedArray) { - test_take_conformance(&array.into_array()); + test_take_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index 60b045d521d..842599bfa52 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -8,8 +8,6 @@ use vortex_error::VortexResult; use super::*; use crate::Canonical; use crate::IntoArray; -#[expect(deprecated)] -use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::PrimitiveArray; @@ -63,11 +61,13 @@ fn test_masked_child_with_validity() { let array = MaskedArray::try_new(child, Validity::from_iter([true, false, true, false, true])).unwrap(); - #[expect(deprecated)] - let prim = array.as_array().to_primitive(); - // Positions where validity is false should be null in masked_child. let mut ctx = array_session().create_execution_ctx(); + let prim = array + .as_array() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!(prim.valid_count(&mut ctx).unwrap(), 3); assert!( prim.is_valid(0, &mut array_session().create_execution_ctx()) diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index 64d39f63531..c7e32a7ee3c 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -23,7 +23,6 @@ use crate::ArrayRef; use crate::Canonical; use crate::EqMode; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayId; @@ -42,6 +41,7 @@ use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; +use crate::legacy_session; use crate::require_child; use crate::scalar::Scalar; use crate::serde::ArrayChildren; @@ -122,6 +122,7 @@ impl VTable for Masked { Ok(Some(vec![])) } + #[allow(clippy::disallowed_methods)] fn deserialize( &self, dtype: &DType, @@ -160,7 +161,7 @@ impl VTable for Masked { let validity_slot = validity_to_child(&validity, len); let data = MaskedData::try_new( len, - child.all_valid(&mut LEGACY_SESSION.create_execution_ctx())?, + child.all_valid(&mut legacy_session().create_execution_ctx())?, validity, )?; Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data) diff --git a/vortex-array/src/arrays/null/compute/cast.rs b/vortex-array/src/arrays/null/compute/cast.rs index a2527e5fb73..c1058156127 100644 --- a/vortex-array/src/arrays/null/compute/cast.rs +++ b/vortex-array/src/arrays/null/compute/cast.rs @@ -90,6 +90,9 @@ mod tests { #[case(NullArray::new(100))] #[case(NullArray::new(0))] fn test_cast_null_conformance(#[case] array: NullArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/null/compute/mod.rs b/vortex-array/src/arrays/null/compute/mod.rs index 5e756da27d0..73a23fe7a8d 100644 --- a/vortex-array/src/arrays/null/compute/mod.rs +++ b/vortex-array/src/arrays/null/compute/mod.rs @@ -15,8 +15,6 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::NullArray; @@ -29,8 +27,11 @@ mod test { #[test] fn test_slice_nulls() { let nulls = NullArray::new(10); - #[expect(deprecated)] - let sliced = nulls.slice(0..4).unwrap().to_null(); + let sliced = nulls + .slice(0..4) + .unwrap() + .execute::(&mut array_session().create_execution_ctx()) + .unwrap(); assert_eq!(sliced.len(), 4); let sliced_arr = sliced.as_array(); @@ -50,11 +51,11 @@ mod test { #[test] fn test_take_nulls() { let nulls = NullArray::new(10); - #[expect(deprecated)] let taken = nulls .take(buffer![0u64, 2, 4, 6, 8].into_array()) .unwrap() - .to_null(); + .execute::(&mut array_session().create_execution_ctx()) + .unwrap(); assert_eq!(taken.len(), 5); let taken_arr = taken.as_array(); @@ -81,21 +82,42 @@ mod test { #[test] fn test_filter_null_array() { - test_filter_conformance(&NullArray::new(5).into_array()); - test_filter_conformance(&NullArray::new(1).into_array()); - test_filter_conformance(&NullArray::new(10).into_array()); + test_filter_conformance( + &NullArray::new(5).into_array(), + &mut array_session().create_execution_ctx(), + ); + test_filter_conformance( + &NullArray::new(1).into_array(), + &mut array_session().create_execution_ctx(), + ); + test_filter_conformance( + &NullArray::new(10).into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] fn test_mask_null_array() { - test_mask_conformance(&NullArray::new(5).into_array()); + test_mask_conformance( + &NullArray::new(5).into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] fn test_take_null_array_conformance() { - test_take_conformance(&NullArray::new(5).into_array()); - test_take_conformance(&NullArray::new(1).into_array()); - test_take_conformance(&NullArray::new(10).into_array()); + test_take_conformance( + &NullArray::new(5).into_array(), + &mut array_session().create_execution_ctx(), + ); + test_take_conformance( + &NullArray::new(1).into_array(), + &mut array_session().create_execution_ctx(), + ); + test_take_conformance( + &NullArray::new(10).into_array(), + &mut array_session().create_execution_ctx(), + ); } #[rstest] diff --git a/vortex-array/src/arrays/null/compute/take.rs b/vortex-array/src/arrays/null/compute/take.rs index 5352ae67653..dc328614743 100644 --- a/vortex-array/src/arrays/null/compute/take.rs +++ b/vortex-array/src/arrays/null/compute/take.rs @@ -6,11 +6,11 @@ use vortex_error::vortex_bail; use crate::ArrayRef; use crate::IntoArray; -#[expect(deprecated)] -use crate::ToCanonical as _; +use crate::VortexSessionExecute; use crate::array::ArrayView; use crate::arrays::Null; use crate::arrays::NullArray; +use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeReduce; use crate::arrays::dict::TakeReduceAdaptor; use crate::match_each_integer_ptype; @@ -19,8 +19,9 @@ use crate::optimizer::rules::ParentRuleSet; impl TakeReduce for Null { #[expect(clippy::cast_possible_truncation)] fn take(array: ArrayView<'_, Null>, indices: &ArrayRef) -> VortexResult> { - #[expect(deprecated)] - let indices = indices.to_primitive(); + #[allow(clippy::disallowed_methods)] + let mut ctx = crate::legacy_session().create_execution_ctx(); + let indices = indices.clone().execute::(&mut ctx)?; // Enforce all indices are valid match_each_integer_ptype!(indices.ptype(), |T| { diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs index 9c8c8e41892..8479b7137cc 100644 --- a/vortex-array/src/arrays/null/mod.rs +++ b/vortex-array/src/arrays/null/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_session::VortexSession; @@ -21,6 +22,8 @@ use crate::array::ValidityVTable; use crate::array::with_empty_buffers; use crate::arrays::null::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::NullBuilder; use crate::dtype::DType; use crate::scalar::Scalar; use crate::serde::ArrayChildren; @@ -118,6 +121,18 @@ impl VTable for Null { fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done(array)) } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for Null requires a NullBuilder"); + }; + builder.append_nulls(array.len()); + Ok(()) + } } /// A array where all values are null. diff --git a/vortex-array/src/arrays/patched/array.rs b/vortex-array/src/arrays/patched/array.rs index b1e5367607b..5e80545665e 100644 --- a/vortex-array/src/arrays/patched/array.rs +++ b/vortex-array/src/arrays/patched/array.rs @@ -16,7 +16,6 @@ use crate::ArraySlots; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; @@ -31,6 +30,7 @@ use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::NativePType; use crate::dtype::PType; +use crate::legacy_session; use crate::match_each_native_ptype; use crate::match_each_unsigned_integer_ptype; use crate::patches::Patches; @@ -110,17 +110,18 @@ pub trait PatchedArrayExt: PatchedArraySlotsExt { } #[inline] + #[allow(clippy::disallowed_methods)] fn lane_range(&self, chunk: usize, lane: usize) -> VortexResult> { assert!(chunk * 1024 <= self.as_ref().len() + self.offset()); assert!(lane < self.n_lanes()); let start = self.lane_offsets().execute_scalar( chunk * self.n_lanes() + lane, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut legacy_session().create_execution_ctx(), )?; let stop = self.lane_offsets().execute_scalar( chunk * self.n_lanes() + lane + 1, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut legacy_session().create_execution_ctx(), )?; let start = start diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index bbe592785ec..e36d426c0be 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -210,8 +210,7 @@ impl VTable for Patched { .clone() .execute::(ctx)? .into_array(); - builder.extend_from_array(&canonical); - return Ok(()); + return canonical.append_to_builder(builder, ctx); } let ptype = dtype.as_ptype(); diff --git a/vortex-array/src/arrays/primitive/array/patch.rs b/vortex-array/src/arrays/primitive/array/patch.rs index d183552b66f..d4b4899077d 100644 --- a/vortex-array/src/arrays/primitive/array/patch.rs +++ b/vortex-array/src/arrays/primitive/array/patch.rs @@ -135,8 +135,6 @@ mod tests { use vortex_buffer::buffer; use super::*; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::assert_arrays_eq; @@ -178,8 +176,7 @@ mod tests { 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(); + let sliced_primitive = sliced.execute::(&mut ctx).unwrap(); assert_arrays_eq!( sliced_primitive, PrimitiveArray::new(buffer![2u32; 6], Validity::AllValid), diff --git a/vortex-array/src/arrays/primitive/array/top_value.rs b/vortex-array/src/arrays/primitive/array/top_value.rs index 7ae45c89405..a6233eb89ee 100644 --- a/vortex-array/src/arrays/primitive/array/top_value.rs +++ b/vortex-array/src/arrays/primitive/array/top_value.rs @@ -10,8 +10,7 @@ use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_utils::aliases::hash_map::HashMap; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; +use crate::ExecutionCtx; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::NativeValue; use crate::dtype::NativePType; @@ -20,7 +19,7 @@ use crate::scalar::PValue; impl PrimitiveArray { /// Compute most common present value of this array - pub fn top_value(&self) -> VortexResult> { + pub fn top_value(&self, ctx: &mut ExecutionCtx) -> VortexResult> { if self.is_empty() { return Ok(None); } @@ -32,10 +31,9 @@ impl PrimitiveArray { match_each_native_ptype!(self.ptype(), |P| { let (top, count) = typed_top_value( self.as_slice::

(), - self.as_ref().validity()?.execute_mask( - self.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), - )?, + self.as_ref() + .validity()? + .execute_mask(self.as_ref().len(), ctx)?, ); Ok(Some((top.into(), count))) }) diff --git a/vortex-array/src/arrays/primitive/compute/between.rs b/vortex-array/src/arrays/primitive/compute/between.rs index 9a6c7f60c1b..85f2b3d42e3 100644 --- a/vortex-array/src/arrays/primitive/compute/between.rs +++ b/vortex-array/src/arrays/primitive/compute/between.rs @@ -105,7 +105,7 @@ where { let slice = arr.as_slice::(); BoolArray::new( - BitBuffer::collect_bool(slice.len(), |idx| { + BitBuffer::collect_bool_multiversioned(slice.len(), |idx| { // We only iterate upto arr len and |arr| == |slice|. let i = unsafe { *slice.get_unchecked(idx) }; lower_fn(lower, i) & upper_fn(i, upper) diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index 36b95b4c554..09defa442b8 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -277,8 +277,6 @@ mod test { use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; use crate::compute::conformance::cast::test_cast_conformance; use crate::dtype::DType; use crate::dtype::Nullability; @@ -291,18 +289,21 @@ mod test { 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(); + let p = arr + .cast(PType::U8.into()) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200]), &mut ctx); assert!(matches!(p.validity(), Ok(Validity::NonNullable))); // to nullable - #[expect(deprecated)] let p = p .into_array() .cast(DType::Primitive(PType::U8, Nullability::Nullable)) .unwrap() - .to_primitive(); + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( p, PrimitiveArray::new(buffer![0u8, 10, 200], Validity::AllValid), @@ -311,22 +312,22 @@ mod test { assert!(matches!(p.validity(), Ok(Validity::AllValid))); // back to non-nullable - #[expect(deprecated)] let p = p .into_array() .cast(DType::Primitive(PType::U8, Nullability::NonNullable)) .unwrap() - .to_primitive(); + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200]), &mut ctx); assert!(matches!(p.validity(), Ok(Validity::NonNullable))); // to nullable u32 - #[expect(deprecated)] let p = p .into_array() .cast(DType::Primitive(PType::U32, Nullability::Nullable)) .unwrap() - .to_primitive(); + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( p, PrimitiveArray::new(buffer![0u32, 10, 200], Validity::AllValid), @@ -335,12 +336,12 @@ mod test { assert!(matches!(p.validity(), Ok(Validity::AllValid))); // to non-nullable u8 - #[expect(deprecated)] let p = p .into_array() .cast(DType::Primitive(PType::U8, Nullability::NonNullable)) .unwrap() - .to_primitive(); + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200]), &mut ctx); assert!(matches!(p.validity(), Ok(Validity::NonNullable))); } @@ -349,8 +350,11 @@ mod 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(); + let u8arr = arr + .cast(PType::F32.into()) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( u8arr, PrimitiveArray::from_iter([0.0f32, 10., 200.]), @@ -394,12 +398,12 @@ mod test { buffer![-1i32, 0, 10], Validity::from_iter([false, true, true]), ); - #[expect(deprecated)] let p = arr .into_array() .cast(DType::Primitive(PType::U32, Nullability::Nullable)) .unwrap() - .to_primitive(); + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( p, PrimitiveArray::from_option_iter([None, Some(0u32), Some(10)]), @@ -426,8 +430,10 @@ mod test { let src = PrimitiveArray::from_iter([0u32, 10, 100]); let src_ptr = src.as_slice::().as_ptr(); - #[expect(deprecated)] - let dst = src.into_array().cast(PType::I32.into())?.to_primitive(); + let dst = src + .into_array() + .cast(PType::I32.into())? + .execute::(&mut ctx)?; let dst_ptr = dst.as_slice::().as_ptr(); // Zero-copy: the data pointer should be identical. @@ -453,12 +459,12 @@ mod test { /// touching the buffer contents. #[test] fn cast_same_width_all_null() -> vortex_error::VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new(buffer![0xFFu8, 0xFF], Validity::AllInvalid); - #[expect(deprecated)] let casted = arr .into_array() .cast(DType::Primitive(PType::I8, Nullability::Nullable))? - .to_primitive(); + .execute::(&mut ctx)?; assert_eq!(casted.len(), 2); assert!(matches!(casted.validity(), Ok(Validity::AllInvalid))); Ok(()) @@ -475,11 +481,10 @@ mod test { buffer![u32::MAX, 0u32, 42u32], Validity::from_iter([false, true, true]), ); - #[expect(deprecated)] let casted = arr .into_array() .cast(DType::Primitive(PType::I32, Nullability::Nullable))? - .to_primitive(); + .execute::(&mut ctx)?; assert_arrays_eq!( casted, PrimitiveArray::from_option_iter([None, Some(0i32), Some(42)]), @@ -495,11 +500,10 @@ mod test { buffer![1000u32, 10u32, 42u32], Validity::from_iter([false, true, true]), ); - #[expect(deprecated)] let casted = arr .into_array() .cast(DType::Primitive(PType::U8, Nullability::Nullable))? - .to_primitive(); + .execute::(&mut ctx)?; assert_arrays_eq!( casted, PrimitiveArray::from_option_iter([None, Some(10u8), Some(42)]), @@ -525,6 +529,6 @@ mod test { #[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: ArrayRef) { - test_cast_conformance(&array); + test_cast_conformance(&array, &mut array_session().create_execution_ctx()); } } diff --git a/vortex-array/src/arrays/primitive/compute/fill_null.rs b/vortex-array/src/arrays/primitive/compute/fill_null.rs index d73442a150e..e8b0d48b322 100644 --- a/vortex-array/src/arrays/primitive/compute/fill_null.rs +++ b/vortex-array/src/arrays/primitive/compute/fill_null.rs @@ -57,8 +57,6 @@ mod test { use crate::arrays::primitive::compute::fill_null::BoolArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; use crate::scalar::Scalar; use crate::validity::Validity; @@ -66,12 +64,12 @@ mod 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 .into_array() .fill_null(Scalar::from(42u8)) .unwrap() - .to_primitive(); + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( p, PrimitiveArray::from_iter([42u8, 8, 42, 10, 42]), @@ -95,12 +93,12 @@ mod test { let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([Option::::None, None, None, None, None]); - #[expect(deprecated)] let p = arr .into_array() .fill_null(Scalar::from(255u8)) .unwrap() - .to_primitive(); + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( p, PrimitiveArray::from_iter([255u8, 255, 255, 255, 255]), @@ -126,12 +124,12 @@ mod test { buffer![8u8, 10, 12, 14, 16], Validity::Array(BoolArray::from_iter([true, true, true, true, true]).into_array()), ); - #[expect(deprecated)] let p = arr .into_array() .fill_null(Scalar::from(255u8)) .unwrap() - .to_primitive(); + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( p, PrimitiveArray::from_iter([8u8, 10, 12, 14, 16]), @@ -154,8 +152,11 @@ mod 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(); + let p = arr + .fill_null(Scalar::from(255u8)) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( p, PrimitiveArray::from_iter([8u8, 10, 12, 14, 16]), diff --git a/vortex-array/src/arrays/primitive/compute/mask.rs b/vortex-array/src/arrays/primitive/compute/mask.rs index 6768efd6fbf..023310b3807 100644 --- a/vortex-array/src/arrays/primitive/compute/mask.rs +++ b/vortex-array/src/arrays/primitive/compute/mask.rs @@ -30,6 +30,8 @@ mod test { use rstest::rstest; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::compute::conformance::mask::test_mask_conformance; @@ -42,6 +44,9 @@ mod test { #[case(PrimitiveArray::from_iter([0.1f32, 0.2, 0.3, 0.4, 0.5]))] #[case(PrimitiveArray::from_option_iter([Some(1.1f64), None, Some(2.2), Some(3.3), None]))] fn test_mask_primitive_conformance(#[case] array: PrimitiveArray) { - test_mask_conformance(&array.into_array()); + test_mask_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 9294659708e..dd562281608 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -170,6 +170,7 @@ mod test { #[test] fn test_take_with_null_indices() { + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::new( buffer![1i32, 2, 3, 4, 5], Validity::Array(BoolArray::from_iter([true, true, false, false, true]).into_array()), @@ -180,23 +181,17 @@ mod test { ); let actual = values.take(indices.into_array()).unwrap(); assert_eq!( - actual - .execute_scalar(0, &mut array_session().create_execution_ctx()) - .vortex_expect("no fail"), + actual.execute_scalar(0, &mut ctx).vortex_expect("no fail"), Scalar::from(Some(1)) ); // position 3 is null assert_eq!( - actual - .execute_scalar(1, &mut array_session().create_execution_ctx()) - .vortex_expect("no fail"), + actual.execute_scalar(1, &mut ctx).vortex_expect("no fail"), Scalar::null_native::() ); // the third index is null assert_eq!( - actual - .execute_scalar(2, &mut array_session().create_execution_ctx()) - .vortex_expect("no fail"), + actual.execute_scalar(2, &mut ctx).vortex_expect("no fail"), Scalar::null_native::() ); } @@ -213,7 +208,10 @@ mod test { ))] #[case(PrimitiveArray::from_option_iter([Some(1), None, Some(3), Some(4), None]))] fn test_take_primitive_conformance(#[case] array: PrimitiveArray) { - test_take_conformance(&array.into_array()); + test_take_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/primitive/compute/zip.rs b/vortex-array/src/arrays/primitive/compute/zip.rs index 83f03b10614..35e49831acc 100644 --- a/vortex-array/src/arrays/primitive/compute/zip.rs +++ b/vortex-array/src/arrays/primitive/compute/zip.rs @@ -46,7 +46,7 @@ impl ZipKernel for Primitive { } // Null mask entries select `if_false`, matching `Zip`'s SQL ELSE semantics. - let mask = mask.try_to_mask_fill_null_false(ctx)?; + let mask = mask.clone().null_as_false().execute(ctx)?; match &mask { // Defer trivial masks to the generic zip, which just casts the surviving side. Mask::AllTrue(_) | Mask::AllFalse(_) => return Ok(None), diff --git a/vortex-array/src/arrays/primitive/tests.rs b/vortex-array/src/arrays/primitive/tests.rs index f623479c7b0..15569b2cebb 100644 --- a/vortex-array/src/arrays/primitive/tests.rs +++ b/vortex-array/src/arrays/primitive/tests.rs @@ -1,10 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::LazyLock; + use vortex_buffer::buffer; +use vortex_session::VortexSession; use crate::ArrayRef; use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::filter::test_filter_conformance; @@ -12,12 +17,14 @@ use crate::compute::conformance::mask::test_mask_conformance; use crate::compute::conformance::search_sorted::rstest_reuse::apply; use crate::compute::conformance::search_sorted::search_sorted_conformance; use crate::compute::conformance::search_sorted::*; -use crate::scalar::PValue; use crate::search_sorted::SearchResult; use crate::search_sorted::SearchSorted; +use crate::search_sorted::SearchSortedPrimitiveArray; use crate::search_sorted::SearchSortedSide; use crate::validity::Validity; +static SESSION: LazyLock = LazyLock::new(array_session); + #[apply(search_sorted_conformance)] fn test_search_sorted_primitive( #[case] array: ArrayRef, @@ -25,9 +32,8 @@ fn test_search_sorted_primitive( #[case] side: SearchSortedSide, #[case] expected: SearchResult, ) -> vortex_error::VortexResult<()> { - let res = array - .as_primitive_typed() - .search_sorted(&Some(PValue::from(value)), side)?; + let res = SearchSortedPrimitiveArray::::new(&array, &mut SESSION.create_execution_ctx()) + .search_sorted(&value, side)?; assert_eq!(res, expected); Ok(()) } @@ -36,12 +42,15 @@ fn test_search_sorted_primitive( fn test_mask_primitive_array() { test_mask_conformance( &PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::NonNullable).into_array(), + &mut array_session().create_execution_ctx(), ); test_mask_conformance( &PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::AllValid).into_array(), + &mut array_session().create_execution_ctx(), ); test_mask_conformance( &PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::AllInvalid).into_array(), + &mut array_session().create_execution_ctx(), ); test_mask_conformance( &PrimitiveArray::new( @@ -49,6 +58,7 @@ fn test_mask_primitive_array() { Validity::Array(BoolArray::from_iter([true, false, true, false, true]).into_array()), ) .into_array(), + &mut array_session().create_execution_ctx(), ); } @@ -57,20 +67,25 @@ fn test_filter_primitive_array() { // Test various sizes test_filter_conformance( &PrimitiveArray::new(buffer![42i32], Validity::NonNullable).into_array(), + &mut array_session().create_execution_ctx(), ); test_filter_conformance( &PrimitiveArray::new(buffer![0, 1], Validity::NonNullable).into_array(), + &mut array_session().create_execution_ctx(), ); test_filter_conformance( &PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::NonNullable).into_array(), + &mut array_session().create_execution_ctx(), ); test_filter_conformance( &PrimitiveArray::new(buffer![0, 1, 2, 3, 4, 5, 6, 7], Validity::NonNullable).into_array(), + &mut array_session().create_execution_ctx(), ); // Test with validity test_filter_conformance( &PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::AllValid).into_array(), + &mut array_session().create_execution_ctx(), ); test_filter_conformance( &PrimitiveArray::new( @@ -80,5 +95,6 @@ fn test_filter_primitive_array() { ), ) .into_array(), + &mut array_session().create_execution_ctx(), ); } diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index a184253527e..50189bf0037 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -220,8 +220,7 @@ impl VTable for Primitive { } }); - builder.extend_from_array(array.as_ref()); - Ok(()) + vortex_bail!("append_to_builder for Primitive requires a matching PrimitiveBuilder"); } fn reduce_parent( diff --git a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs index 2ac376155e3..39ab8c2c466 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs @@ -2,19 +2,22 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_error::vortex_bail; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array::ArrayView; use crate::array::ValidityVTable; +use crate::arrays::ConstantArray; use crate::arrays::scalar_fn::ScalarFnArrayExt; use crate::arrays::scalar_fn::vtable::ArrayExpr; use crate::arrays::scalar_fn::vtable::FakeEq; use crate::arrays::scalar_fn::vtable::ScalarFn; use crate::expr::Expression; use crate::expr::lit; +use crate::legacy_session; use crate::scalar_fn::TypedScalarFnInstance; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::fns::literal::Literal; @@ -24,30 +27,32 @@ use crate::validity::Validity; /// Execute an expression tree recursively. /// /// This assumes all leaf expressions are either ArrayExpr (wrapping actual arrays) or Literals. -fn execute_expr(expr: &Expression, row_count: usize) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - +fn execute_expr( + expr: &Expression, + row_count: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { // Handle Root expression - this should not happen in validity expressions if expr.is::() { - vortex_error::vortex_bail!("Root expression cannot be executed in validity context"); + vortex_bail!("Root expression cannot be executed in validity context"); } // Handle Literal expression - create a constant array if expr.is::() { let scalar = expr.as_::(); - return Ok(crate::arrays::ConstantArray::new(scalar.clone(), row_count).into_array()); + return Ok(ConstantArray::new(scalar.clone(), row_count).into_array()); } // Recursively execute child expressions to get input arrays let inputs: Vec = expr .children() .iter() - .map(|child| execute_expr(child, row_count)) + .map(|child| execute_expr(child, row_count, ctx)) .collect::>()?; let args = VecExecutionArgs::new(inputs, row_count); - Ok(expr.scalar_fn().execute(&args, &mut ctx)?.into_array()) + Ok(expr.scalar_fn().execute(&args, ctx)?.into_array()) } impl ValidityVTable for ScalarFn { @@ -68,7 +73,13 @@ impl ValidityVTable for ScalarFn { let expr = Expression::try_new(array.scalar_fn().clone(), inputs)?; let validity_expr = array.scalar_fn().validity(&expr)?; + #[allow(clippy::disallowed_methods)] + let ctx = &mut legacy_session().create_execution_ctx(); // Execute the validity expression. All leaves are ArrayExpr nodes. - Ok(Validity::Array(execute_expr(&validity_expr, array.len())?)) + Ok(Validity::Array(execute_expr( + &validity_expr, + array.len(), + ctx, + )?)) } } diff --git a/vortex-array/src/arrays/struct_/compute/cast.rs b/vortex-array/src/arrays/struct_/compute/cast.rs index a7e1e760fd3..bfae815ceb5 100644 --- a/vortex-array/src/arrays/struct_/compute/cast.rs +++ b/vortex-array/src/arrays/struct_/compute/cast.rs @@ -188,7 +188,7 @@ mod tests { #[case(create_nested_struct())] #[case(create_simple_struct())] fn test_cast_struct_conformance(#[case] array: StructArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] diff --git a/vortex-array/src/arrays/struct_/compute/mod.rs b/vortex-array/src/arrays/struct_/compute/mod.rs index c1539a638bf..4b9a4e396fd 100644 --- a/vortex-array/src/arrays/struct_/compute/mod.rs +++ b/vortex-array/src/arrays/struct_/compute/mod.rs @@ -117,6 +117,7 @@ mod tests { &StructArray::try_new(FieldNames::empty(), vec![], 5, Validity::NonNullable) .unwrap() .into_array(), + &mut array_session().create_execution_ctx(), ); } @@ -151,6 +152,7 @@ mod tests { ) .unwrap() .into_array(), + &mut array_session().create_execution_ctx(), ); } @@ -270,6 +272,7 @@ mod tests { &StructArray::try_new(FieldNames::empty(), vec![], 5, Validity::NonNullable) .unwrap() .into_array(), + &mut array_session().create_execution_ctx(), ); } @@ -286,6 +289,7 @@ mod tests { &StructArray::try_new(["xs", "ys"].into(), vec![xs, ys], 5, Validity::NonNullable) .unwrap() .into_array(), + &mut array_session().create_execution_ctx(), ); } @@ -307,6 +311,7 @@ mod tests { ) .unwrap() .into_array(), + &mut array_session().create_execution_ctx(), ); } @@ -335,6 +340,7 @@ mod tests { ) .unwrap() .into_array(), + &mut array_session().create_execution_ctx(), ); } @@ -347,6 +353,7 @@ mod tests { &StructArray::try_new(["xs", "ys"].into(), vec![xs, ys], 1, Validity::NonNullable) .unwrap() .into_array(), + &mut array_session().create_execution_ctx(), ); } @@ -370,6 +377,7 @@ mod tests { ) .unwrap() .into_array(), + &mut array_session().create_execution_ctx(), ); } diff --git a/vortex-array/src/arrays/struct_/compute/zip.rs b/vortex-array/src/arrays/struct_/compute/zip.rs index cf7b7b91d6d..394cc5c8424 100644 --- a/vortex-array/src/arrays/struct_/compute/zip.rs +++ b/vortex-array/src/arrays/struct_/compute/zip.rs @@ -37,7 +37,7 @@ impl ZipKernel for Struct { let fields = if_true .iter_unmasked_fields() .zip(if_false.iter_unmasked_fields()) - .map(|(t, f)| ArrayBuiltins::zip(mask, t.clone(), f.clone())) + .map(|(t, f)| mask.zip(t.clone(), f.clone())) .collect::>>()?; let v1 = if_true.validity()?; @@ -48,11 +48,11 @@ impl ZipKernel for Struct { (Validity::AllInvalid, Validity::AllInvalid) => Validity::AllInvalid, (v1, v2) => { - let mask_mask = mask.try_to_mask_fill_null_false(ctx)?; + let mask_mask = mask.clone().null_as_false().execute(ctx)?; let v1m = v1.execute_mask(if_true.len(), ctx)?; let v2m = v2.execute_mask(if_false.len(), ctx)?; - let combined = (v1m.bitand(&mask_mask)).bitor(&v2m.bitand(&mask_mask.not())); + let combined = v1m.bitand(&mask_mask).bitor(&v2m.bitand(&mask_mask.not())); Validity::from_mask( combined, if_true.dtype().nullability() | if_false.dtype().nullability(), diff --git a/vortex-array/src/arrays/struct_/vtable/mod.rs b/vortex-array/src/arrays/struct_/vtable/mod.rs index 1b23ecf804c..0222a3c4712 100644 --- a/vortex-array/src/arrays/struct_/vtable/mod.rs +++ b/vortex-array/src/arrays/struct_/vtable/mod.rs @@ -23,6 +23,8 @@ use crate::arrays::struct_::array::VALIDITY_SLOT; use crate::arrays::struct_::array::make_struct_slots; use crate::arrays::struct_::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::StructBuilder; use crate::dtype::DType; use crate::serde::ArrayChildren; use crate::validity::Validity; @@ -197,6 +199,17 @@ impl VTable for Struct { Ok(ExecutionResult::done(array)) } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for Struct requires a StructBuilder"); + }; + builder.append_struct_array(&array.into_owned(), ctx) + } + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, diff --git a/vortex-array/src/arrays/varbin/array.rs b/vortex-array/src/arrays/varbin/array.rs index 220fe948a73..4cb74fe48b4 100644 --- a/vortex-array/src/arrays/varbin/array.rs +++ b/vortex-array/src/arrays/varbin/array.rs @@ -15,7 +15,6 @@ use vortex_error::vortex_err; use crate::ArrayRef; use crate::ArraySlots; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; @@ -28,6 +27,7 @@ use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::Nullability; +use crate::legacy_session; use crate::match_each_integer_ptype; use crate::validity::Validity; @@ -230,6 +230,7 @@ impl VarBinData { } /// Validates that every non-null value is valid UTF-8. + #[allow(clippy::disallowed_methods)] fn validate_utf8(offsets: &ArrayRef, bytes: &[u8], validity: &Validity) -> VortexResult<()> { let validate_at = |i: usize, start: usize, end: usize| -> VortexResult<()> { let string_bytes = &bytes[start..end]; @@ -242,7 +243,7 @@ impl VarBinData { Ok(()) }; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = legacy_session().create_execution_ctx(); // TODO(joe): update the created VarBin with this decompressed Array. let primitive_offsets = offsets.clone().execute::(&mut ctx)?; @@ -335,6 +336,7 @@ pub trait VarBinArrayExt: TypedArrayRef { ) } + #[allow(clippy::disallowed_methods)] fn offset_at(&self, index: usize) -> usize { assert!( index <= self.as_ref().len(), @@ -344,7 +346,7 @@ pub trait VarBinArrayExt: TypedArrayRef { (&self .offsets() - .execute_scalar(index, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(index, &mut legacy_session().create_execution_ctx()) .vortex_expect("offsets must support execute_scalar")) .try_into() .vortex_expect("Failed to convert offset to usize") diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 404f574af52..422139cc851 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -8,8 +8,6 @@ use vortex_error::vortex_panic; use crate::IntoArray; #[cfg(debug_assertions)] -use crate::LEGACY_SESSION; -#[cfg(debug_assertions)] use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; @@ -17,6 +15,8 @@ use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::expr::stats::Precision; use crate::expr::stats::Stat; +#[cfg(debug_assertions)] +use crate::legacy_session; use crate::validity::Validity; pub struct VarBinBuilder { @@ -94,6 +94,7 @@ impl VarBinBuilder { self.validity.append_n(true, num); } + #[allow(clippy::disallowed_methods)] pub fn finish(self, dtype: DType) -> VarBinArray { let offsets = PrimitiveArray::new(self.offsets.freeze(), Validity::NonNullable); let nulls = self.validity.freeze(); @@ -107,7 +108,7 @@ impl VarBinBuilder { { let offsets_are_sorted = offsets .statistics() - .compute_is_sorted(&mut LEGACY_SESSION.create_execution_ctx()) + .compute_is_sorted(&mut legacy_session().create_execution_ctx()) .unwrap_or(false); debug_assert!(offsets_are_sorted, "VarBinBuilder offsets must be sorted"); } diff --git a/vortex-array/src/arrays/varbin/compute/cast.rs b/vortex-array/src/arrays/varbin/compute/cast.rs index 5983886ea30..b0c4325c774 100644 --- a/vortex-array/src/arrays/varbin/compute/cast.rs +++ b/vortex-array/src/arrays/varbin/compute/cast.rs @@ -130,6 +130,6 @@ mod tests { #[case(VarBinArray::from_iter(vec![Some(b"test".as_slice()), None], DType::Binary(Nullability::Nullable)))] #[case(VarBinArray::from_iter(vec![Some("single")], DType::Utf8(Nullability::NonNullable)))] fn test_cast_varbin_conformance(#[case] array: VarBinArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index ede3c409d94..f6fff972c9b 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -1,17 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use arrow_array::BinaryArray; -use arrow_array::LargeBinaryArray; -use arrow_array::LargeStringArray; -use arrow_array::StringArray; -use arrow_ord::cmp; -use arrow_schema::DataType; +use std::cmp::Ordering; + use vortex_buffer::BitBuffer; use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_err; use crate::ArrayRef; use crate::ExecutionCtx; @@ -20,19 +15,15 @@ use crate::array::ArrayView; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; -use crate::arrays::VarBinViewArray; use crate::arrays::varbin::VarBinArrayExt; -use crate::arrow::Datum; -use crate::arrow::from_arrow_columnar; -use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::match_each_integer_ptype; use crate::scalar_fn::fns::binary::CompareKernel; use crate::scalar_fn::fns::operators::CompareOperator; -use crate::scalar_fn::fns::operators::Operator; -// This implementation exists so we can have custom translation of RHS to arrow that's not the same as IntoCanonical +// This implementation exists so we can compare against a constant in encoded space, without +// canonicalizing the VarBin array to VarBinView. impl CompareKernel for VarBin { fn compare( lhs: ArrayView<'_, VarBin>, @@ -40,106 +31,66 @@ impl CompareKernel for VarBin { operator: CompareOperator, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if let Some(rhs_const) = rhs.as_constant() { - let nullable = lhs.dtype().is_nullable() || rhs_const.dtype().is_nullable(); - let len = lhs.len(); - - let rhs_is_empty = match rhs_const.dtype() { - DType::Binary(_) => rhs_const - .as_binary() - .is_empty() - .vortex_expect("RHS should not be null"), - DType::Utf8(_) => rhs_const - .as_utf8() - .is_empty() - .vortex_expect("RHS should not be null"), - _ => vortex_bail!("VarBinArray can only have type of Binary or Utf8"), - }; - - if rhs_is_empty { - let buffer = match operator { - 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)?; - match_each_integer_ptype!(lhs_offsets.ptype(), |P| { - compare_offsets_to_empty::

(lhs_offsets, true) - }) - } - CompareOperator::NotEq | CompareOperator::Gt => { - let lhs_offsets = lhs.offsets().clone().execute::(ctx)?; - match_each_integer_ptype!(lhs_offsets.ptype(), |P| { - compare_offsets_to_empty::

(lhs_offsets, false) - }) - } - }; - - return Ok(Some( - BoolArray::new( - buffer, - lhs.validity()?.union_nullability(rhs.dtype().nullability()), - ) - .into_array(), - )); - } + let Some(rhs_const) = rhs.as_constant() else { + return Ok(None); + }; - let lhs = Datum::try_new(lhs.array(), ctx)?; + let len = lhs.len(); - // The RHS scalar must match the LHS Arrow data type. VarBin with i64 offsets is - // converted to LargeBinary/LargeUtf8 (see `preferred_arrow_type`), and Arrow refuses to - // compare LargeBinary with Binary (or LargeUtf8 with Utf8). - let arrow_rhs: &dyn arrow_array::Datum = match (rhs_const.dtype(), lhs.data_type()) { - (DType::Utf8(_), DataType::LargeUtf8) => &rhs_const - .as_utf8() - .value() - .map(LargeStringArray::new_scalar) - .unwrap_or_else(|| arrow_array::Scalar::new(LargeStringArray::new_null(1))), - (DType::Utf8(_), _) => &rhs_const - .as_utf8() - .value() - .map(StringArray::new_scalar) - .unwrap_or_else(|| arrow_array::Scalar::new(StringArray::new_null(1))), - (DType::Binary(_), DataType::LargeBinary) => &rhs_const - .as_binary() - .value() - .map(LargeBinaryArray::new_scalar) - .unwrap_or_else(|| arrow_array::Scalar::new(LargeBinaryArray::new_null(1))), - (DType::Binary(_), _) => &rhs_const - .as_binary() - .value() - .map(BinaryArray::new_scalar) - .unwrap_or_else(|| arrow_array::Scalar::new(BinaryArray::new_null(1))), - _ => vortex_bail!( - "VarBin array RHS can only be Utf8 or Binary, given {}", - rhs_const.dtype() - ), - }; + // The compare adaptor resolves null constants before dispatching to this kernel, so + // the scalar always carries a value. + let rhs_bytes: &[u8] = match rhs_const.dtype() { + DType::Binary(_) => rhs_const + .as_binary() + .value() + .vortex_expect("RHS should not be null") + .as_slice(), + DType::Utf8(_) => rhs_const + .as_utf8() + .value() + .vortex_expect("RHS should not be null") + .as_str() + .as_bytes(), + _ => vortex_bail!("VarBinArray can only have type of Binary or Utf8"), + }; - let array = match operator { - CompareOperator::Eq => cmp::eq(&lhs, arrow_rhs), - CompareOperator::NotEq => cmp::neq(&lhs, arrow_rhs), - CompareOperator::Gt => cmp::gt(&lhs, arrow_rhs), - CompareOperator::Gte => cmp::gt_eq(&lhs, arrow_rhs), - CompareOperator::Lt => cmp::lt(&lhs, arrow_rhs), - CompareOperator::Lte => cmp::lt_eq(&lhs, arrow_rhs), + let buffer = if rhs_bytes.is_empty() { + // Comparisons against "" only need the value lengths, i.e. the offset deltas. + match operator { + 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)?; + match_each_integer_ptype!(lhs_offsets.ptype(), |P| { + compare_offsets_to_empty::

(lhs_offsets, true) + }) + } + CompareOperator::NotEq | CompareOperator::Gt => { + let lhs_offsets = lhs.offsets().clone().execute::(ctx)?; + match_each_integer_ptype!(lhs_offsets.ptype(), |P| { + compare_offsets_to_empty::

(lhs_offsets, false) + }) + } } - .map_err(|err| vortex_err!("Failed to compare VarBin array: {}", err))?; - - Ok(Some(from_arrow_columnar(&array, len, nullable, ctx)?)) - } else if !rhs.is::() { - // NOTE: If the rhs is not a VarBin array it will be canonicalized to a VarBinView - // Arrow doesn't support comparing VarBin to VarBinView arrays, so we convert ourselves - // to VarBinView and re-invoke. - Ok(Some( - lhs.array() - .clone() - .execute::(ctx)? - .into_array() - .binary(rhs.clone(), Operator::from(operator))?, - )) } else { - Ok(None) - } + let lhs_offsets = lhs.offsets().clone().execute::(ctx)?; + match_each_integer_ptype!(lhs_offsets.ptype(), |P| { + compare_bytes_to_constant( + lhs_offsets.as_slice::

(), + lhs.bytes().as_slice(), + rhs_bytes, + operator, + ) + }) + }; + + Ok(Some( + BoolArray::new( + buffer, + lhs.validity()?.union_nullability(rhs.dtype().nullability()), + ) + .into_array(), + )) } } @@ -153,16 +104,80 @@ fn compare_offsets_to_empty(offsets: PrimitiveArray, eq: bool) }) } +/// Compare every value in a VarBin array against a constant, resolving values through the +/// offsets. Dispatches the operator outside the lane loop so each predicate inlines into its +/// own loop. +fn compare_bytes_to_constant( + offsets: &[P], + bytes: &[u8], + constant: &[u8], + operator: CompareOperator, +) -> BitBuffer { + match operator { + CompareOperator::Eq => { + collect_lane_bits(offsets, |start, end| value_eq(bytes, start, end, constant)) + } + CompareOperator::NotEq => { + collect_lane_bits(offsets, |start, end| !value_eq(bytes, start, end, constant)) + } + CompareOperator::Gt => collect_lane_bits(offsets, |start, end| { + value_cmp(bytes, start, end, constant).is_gt() + }), + CompareOperator::Gte => collect_lane_bits(offsets, |start, end| { + value_cmp(bytes, start, end, constant).is_ge() + }), + CompareOperator::Lt => collect_lane_bits(offsets, |start, end| { + value_cmp(bytes, start, end, constant).is_lt() + }), + CompareOperator::Lte => collect_lane_bits(offsets, |start, end| { + value_cmp(bytes, start, end, constant).is_le() + }), + } +} + +/// Bit-pack `predicate(offsets[i], offsets[i + 1])` over each lane of a VarBin array. +fn collect_lane_bits( + offsets: &[P], + predicate: impl Fn(usize, usize) -> bool, +) -> BitBuffer { + BitBuffer::collect_bool(offsets.len() - 1, |idx| { + // SAFETY: `collect_bool` yields idx < offsets.len() - 1. + let start = unsafe { offsets.get_unchecked(idx) }.as_(); + let end = unsafe { offsets.get_unchecked(idx + 1) }.as_(); + predicate(start, end) + }) +} + +/// Whether `bytes[start..end]` equals `constant`, comparing lengths first so lanes of a +/// different length never touch the value bytes. +/// +/// Offsets at null positions are not validated, so an out-of-bounds or inverted range is +/// possible there; such lanes answer `false`, and validity masks them out of the result anyway. +#[inline(always)] +fn value_eq(bytes: &[u8], start: usize, end: usize, constant: &[u8]) -> bool { + // A lane can only match when its length equals the constant's, so lanes of a different + // length answer without touching the value bytes. An inverted garbage range (start > end) + // wraps to a huge value that never equals `constant.len()`. + end.wrapping_sub(start) == constant.len() + && bytes.get(start..end).is_some_and(|value| value == constant) +} + +/// Order `bytes[start..end]` against `constant`, treating the unvalidated garbage ranges that +/// can appear at null positions as empty; validity masks those lanes out of the result anyway. +#[inline(always)] +fn value_cmp(bytes: &[u8], start: usize, end: usize, constant: &[u8]) -> Ordering { + bytes.get(start..end).unwrap_or_default().cmp(constant) +} + #[cfg(test)] mod test { use vortex_buffer::BitBuffer; use vortex_buffer::ByteBuffer; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; @@ -175,11 +190,11 @@ mod test { #[test] fn test_binary_compare() { + let mut ctx = array_session().create_execution_ctx(); let array = VarBinArray::from_iter( [Some(b"abc".to_vec()), None, Some(b"def".to_vec())], DType::Binary(Nullability::Nullable), ); - #[expect(deprecated)] let result = array .into_array() .binary( @@ -191,17 +206,15 @@ mod test { Operator::Eq, ) .unwrap() - .to_bool(); + .execute::(&mut ctx) + .unwrap(); assert_eq!( &result .as_ref() .validity() .unwrap() - .execute_mask( - result.as_ref().len(), - &mut array_session().create_execution_ctx() - ) + .execute_mask(result.as_ref().len(), &mut ctx) .unwrap() .to_bit_buffer(), &BitBuffer::from_iter([true, false, true]) @@ -214,6 +227,7 @@ mod test { #[test] fn varbinview_compare() { + let mut ctx = array_session().create_execution_ctx(); let array = VarBinArray::from_iter( [Some(b"abc".to_vec()), None, Some(b"def".to_vec())], DType::Binary(Nullability::Nullable), @@ -222,22 +236,19 @@ mod test { [None, None, Some(b"def".to_vec())], DType::Binary(Nullability::Nullable), ); - #[expect(deprecated)] let result = array .into_array() .binary(vbv.into_array(), Operator::Eq) .unwrap() - .to_bool(); + .execute::(&mut ctx) + .unwrap(); assert_eq!( result .as_ref() .validity() .unwrap() - .execute_mask( - result.as_ref().len(), - &mut array_session().create_execution_ctx() - ) + .execute_mask(result.as_ref().len(), &mut ctx) .unwrap() .to_bit_buffer(), BitBuffer::from_iter([false, false, true]) @@ -282,11 +293,9 @@ mod tests { ); } - /// Regression: a [`VarBinArray`] built with `i64` offsets is canonicalised to - /// Arrow `LargeUtf8` / `LargeBinary` by `preferred_arrow_type`. Without an explicit - /// branch in [`CompareKernel`], the constant RHS is wrapped in a `StringArray` / - /// `BinaryArray` and Arrow rejects the `LargeUtf8 == Utf8` mismatch. Triggering - /// this only requires `i64` offsets, not large data. + /// Regression: [`CompareKernel`] must handle every offset width; a `VarBinArray` built with + /// `i64` offsets once failed the constant comparison. Triggering this only requires `i64` + /// offsets, not large data. /// /// [`CompareKernel`]: super::CompareKernel #[test] diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index ab969cf3def..400beb9aa0a 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -360,12 +360,18 @@ mod test { vec!["hello", "world", "filter", "good", "bye"], DType::Utf8(NonNullable), ); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); let array = VarBinArray::from_iter( vec![Some("hello"), None, Some("filter"), Some("good"), None], DType::Utf8(Nullable), ); - test_filter_conformance(&array.into_array()); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/varbin/compute/mask.rs b/vortex-array/src/arrays/varbin/compute/mask.rs index 1c5567adea7..19236cee47b 100644 --- a/vortex-array/src/arrays/varbin/compute/mask.rs +++ b/vortex-array/src/arrays/varbin/compute/mask.rs @@ -29,6 +29,8 @@ impl MaskReduce for VarBin { #[cfg(test)] mod test { use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::VarBinArray; use crate::compute::conformance::mask::test_mask_conformance; use crate::dtype::DType; @@ -40,12 +42,18 @@ mod test { vec!["hello", "world", "filter", "good", "bye"], DType::Utf8(Nullability::NonNullable), ); - test_mask_conformance(&array.into_array()); + test_mask_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); let array = VarBinArray::from_iter( vec![Some("hello"), None, Some("filter"), Some("good"), None], DType::Utf8(Nullability::Nullable), ); - test_mask_conformance(&array.into_array()); + test_mask_conformance( + &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 74e0bd6af9d..6a6454d0603 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -303,7 +303,10 @@ mod tests { ))] #[case(VarBinArray::from_iter(["single"].map(Some), DType::Utf8(Nullability::NonNullable)))] fn test_take_varbin_conformance(#[case] array: VarBinArray) { - test_take_conformance(&array.into_array()); + test_take_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] diff --git a/vortex-array/src/arrays/varbin/vtable/canonical.rs b/vortex-array/src/arrays/varbin/vtable/canonical.rs index 4865df5e1a3..fca9671e3b4 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -53,8 +53,6 @@ mod tests { use crate::arrays::VarBinViewArray; use crate::arrays::varbin::builder::VarBinBuilder; use crate::assert_arrays_eq; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::Nullability; @@ -73,8 +71,8 @@ mod tests { let varbin = varbin.slice(1..4).unwrap(); - #[expect(deprecated)] - let canonical = varbin.to_varbinview(); + let mut ctx = array_session().create_execution_ctx(); + let canonical = varbin.execute::(&mut ctx).unwrap(); assert_eq!(canonical.dtype(), &dtype); assert!( @@ -98,8 +96,11 @@ mod tests { 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(); + let canonical = varbin + .as_array() + .clone() + .execute::(&mut ctx) + .unwrap(); let expected = match dtype { DType::Utf8(_) => VarBinViewArray::from_iter_str(["foo", "bar", "baz"]), _ => VarBinViewArray::from_iter_bin(["foo", "bar", "baz"]), @@ -112,8 +113,12 @@ mod tests { fn test_canonical_varbin_empty() { let varbin = VarBinArray::from_iter_nonnull([] as [&str; 0], DType::Utf8(Nullability::NonNullable)); - #[expect(deprecated)] - let canonical = varbin.as_array().to_varbinview(); + let mut ctx = array_session().create_execution_ctx(); + let canonical = varbin + .as_array() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!(canonical.len(), 0); } } diff --git a/vortex-array/src/arrays/varbinview/array.rs b/vortex-array/src/arrays/varbinview/array.rs index 3ba90409939..7fa66eaea7a 100644 --- a/vortex-array/src/arrays/varbinview/array.rs +++ b/vortex-array/src/arrays/varbinview/array.rs @@ -18,7 +18,6 @@ use vortex_error::vortex_err; use vortex_error::vortex_panic; use crate::ArraySlots; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::array::Array; use crate::array::ArrayParts; @@ -32,6 +31,7 @@ use crate::builders::ArrayBuilder; use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::dtype::Nullability; +use crate::legacy_session; use crate::validity::Validity; /// The validity bitmap indicating which elements are non-null. @@ -309,6 +309,7 @@ impl VarBinViewData { Ok(()) } + #[allow(clippy::disallowed_methods)] fn validate_views( views: &Buffer, buffers: &Arc<[ByteBuffer]>, @@ -370,7 +371,7 @@ impl VarBinViewData { // into a mask once and zip it with the views, validating only the valid (non-null) // entries. Validity::Array(_) => { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = legacy_session().create_execution_ctx(); let mask = validity.execute_mask(views.len(), &mut ctx)?; for ((idx, view), valid) in views.iter().enumerate().zip(mask.iter()) { if valid { diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index fc3f3477ffc..2aba21ae448 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -10,12 +10,9 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; -use crate::IntoArray; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; +use crate::ExecutionCtx; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::Ref; -use crate::builders::ArrayBuilder; use crate::builders::VarBinViewBuilder; const DEFAULT_COMPACTION_THRESHOLD: f64 = 0.5; @@ -30,16 +27,16 @@ impl VarBinViewArray { /// that are no longer visible. We detect when there is wasted space in any of the buffers, and if /// so, will aggressively compact all visible outlined string data into new buffers while keeping /// well-utilized buffers unchanged. - pub fn compact_buffers(&self) -> VortexResult { + pub fn compact_buffers(&self, ctx: &mut ExecutionCtx) -> VortexResult { // If there is nothing to be gained by compaction, return the original array untouched. - if !self.should_compact()? { + if !self.should_compact(ctx)? { return Ok(self.clone()); } - self.compact_with_threshold(DEFAULT_COMPACTION_THRESHOLD) + self.compact_with_threshold(DEFAULT_COMPACTION_THRESHOLD, ctx) } - fn should_compact(&self) -> VortexResult { + fn should_compact(&self, ctx: &mut ExecutionCtx) -> VortexResult { let nbuffers = self.data_buffers().len(); // If the array is entirely inlined strings, do not attempt to compact. @@ -62,21 +59,22 @@ impl VarBinViewArray { return Ok(false); } - let bytes_referenced: u64 = self.count_referenced_bytes()?; + let bytes_referenced: u64 = self.count_referenced_bytes(ctx)?; Ok((bytes_referenced as f64 / buffer_total_bytes as f64) < DEFAULT_COMPACTION_THRESHOLD) } /// Iterates over all valid, non-inlined views, calling the provided /// closure for each one. #[inline(always)] - fn iter_valid_views(&self, mut f: F) -> VortexResult<()> + fn iter_valid_views(&self, ctx: &mut ExecutionCtx, mut f: F) -> VortexResult<()> where F: FnMut(&Ref), { - match self.as_ref().validity()?.execute_mask( - self.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), - )? { + match self + .as_ref() + .validity()? + .execute_mask(self.as_ref().len(), ctx)? + { Mask::AllTrue(_) => { for &view in self.views().iter() { if !view.is_inlined() { @@ -98,13 +96,16 @@ impl VarBinViewArray { /// Count the number of bytes addressed by the views, not including null /// values or any inlined strings. - fn count_referenced_bytes(&self) -> VortexResult { + fn count_referenced_bytes(&self, ctx: &mut ExecutionCtx) -> VortexResult { let mut total = 0u64; - self.iter_valid_views(|view| total += view.size as u64)?; + self.iter_valid_views(ctx, |view| total += view.size as u64)?; Ok(total) } - pub(crate) fn buffer_utilizations(&self) -> VortexResult> { + pub(crate) fn buffer_utilizations( + &self, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { let mut utilizations: Vec = self .data_buffers() .iter() @@ -114,7 +115,7 @@ impl VarBinViewArray { }) .collect(); - self.iter_valid_views(|view| { + self.iter_valid_views(ctx, |view| { utilizations[view.buffer_index as usize].add(view.offset, view.size); })?; @@ -138,13 +139,14 @@ impl VarBinViewArray { pub fn compact_with_threshold( &self, buffer_utilization_threshold: f64, // [0, 1] + ctx: &mut ExecutionCtx, ) -> VortexResult { let mut builder = VarBinViewBuilder::with_compaction( self.dtype().clone(), self.len(), buffer_utilization_threshold, ); - builder.extend_from_array(&self.clone().into_array()); + builder.append_varbinview_array(self, ctx)?; Ok(builder.finish_into_varbinview()) } } @@ -227,14 +229,12 @@ mod tests { // Take only the first and last elements (indices 0 and 4) let indices = buffer![0u32, 4u32].into_array(); let taken = original.take(indices).unwrap(); - let taken = taken - .execute::(&mut array_session().create_execution_ctx()) - .unwrap(); + let taken = taken.execute::(&mut ctx).unwrap(); // The taken array should still have the same number of buffers assert_eq!(taken.data_buffers().len(), original_buffers); // Now optimize the taken array - let optimized_array = taken.compact_buffers().unwrap(); + let optimized_array = taken.compact_buffers(&mut ctx).unwrap(); // The optimized array should have compacted buffers // Since both remaining strings are short, they should be inlined @@ -272,7 +272,7 @@ mod tests { .execute::(&mut array_session().create_execution_ctx()) .unwrap(); - let optimized_array = taken_array.compact_with_threshold(1.0).unwrap(); + let optimized_array = taken_array.compact_with_threshold(1.0, &mut ctx).unwrap(); // The optimized array should have exactly 1 buffer (consolidated) assert_eq!(optimized_array.data_buffers().len(), 1); @@ -295,7 +295,7 @@ mod tests { assert_eq!(original.data_buffers().len(), 0); // Optimize should return the same array - let optimized_array = original.compact_buffers().unwrap(); + let optimized_array = original.compact_buffers(&mut ctx).unwrap(); assert_eq!(optimized_array.data_buffers().len(), 0); @@ -315,7 +315,7 @@ mod tests { assert_eq!(original.buffer(0).len(), str1.len() + str2.len()); // Optimize should return the same array (no change needed) - let optimized_array = original.compact_buffers().unwrap(); + let optimized_array = original.compact_buffers(&mut ctx).unwrap(); assert_eq!(optimized_array.data_buffers().len(), 1); @@ -341,7 +341,7 @@ mod tests { .execute::(&mut array_session().create_execution_ctx()) .unwrap(); // Compact with threshold=0 (should not compact) - let compacted = taken.compact_with_threshold(0.0).unwrap(); + let compacted = taken.compact_with_threshold(0.0, &mut ctx).unwrap(); // Should still have the same number of buffers as the taken array assert_eq!(compacted.data_buffers().len(), taken.data_buffers().len()); @@ -370,7 +370,7 @@ mod tests { let original_buffers = taken.data_buffers().len(); // Compact with threshold=1.0 (aggressive compaction) - let compacted = taken.compact_with_threshold(1.0).unwrap(); + let compacted = taken.compact_with_threshold(1.0, &mut ctx).unwrap(); // Should have compacted buffers assert!(compacted.data_buffers().len() <= original_buffers); @@ -393,7 +393,7 @@ mod tests { assert_eq!(original.data_buffers().len(), 1); // Compact with high threshold - let compacted = original.compact_with_threshold(0.8).unwrap(); + let compacted = original.compact_with_threshold(0.8, &mut ctx).unwrap(); // Well-utilized buffer should be preserved assert_eq!(compacted.data_buffers().len(), 1); @@ -425,7 +425,7 @@ mod tests { .unwrap(); // Compact with moderate threshold - let compacted = taken.compact_with_threshold(0.7).unwrap(); + let compacted = taken.compact_with_threshold(0.7, &mut ctx).unwrap(); let expected = VarBinViewArray::from_iter( [0, 2, 4, 6, 8].map(|i| Some(strings[i].as_str())), @@ -451,12 +451,12 @@ mod tests { .execute::(&mut array_session().create_execution_ctx()) .unwrap(); // Get buffer stats before compaction - let utils_before = taken.buffer_utilizations().unwrap(); + let utils_before = taken.buffer_utilizations(&mut ctx).unwrap(); let original_buffer_count = taken.data_buffers().len(); // Compact with a threshold that should trigger slicing // The range utilization should be high even if overall utilization is low - let compacted = taken.compact_with_threshold(0.8).unwrap(); + let compacted = taken.compact_with_threshold(0.8, &mut ctx).unwrap(); // After compaction, we should still have buffers (sliced, not rewritten) assert!( @@ -497,9 +497,13 @@ mod tests { #[case] expected_bytes: u64, #[case] expected_utils: &[f64], ) { - assert_eq!(arr.count_referenced_bytes().unwrap(), expected_bytes); + let mut ctx = array_session().create_execution_ctx(); + assert_eq!( + arr.count_referenced_bytes(&mut ctx).unwrap(), + expected_bytes + ); let utils: Vec = arr - .buffer_utilizations() + .buffer_utilizations(&mut ctx) .unwrap() .iter() .map(|u| u.overall_utilization()) diff --git a/vortex-array/src/arrays/varbinview/compute/cast.rs b/vortex-array/src/arrays/varbinview/compute/cast.rs index 476c74e7309..660a7842ed9 100644 --- a/vortex-array/src/arrays/varbinview/compute/cast.rs +++ b/vortex-array/src/arrays/varbinview/compute/cast.rs @@ -135,6 +135,6 @@ mod tests { #[case(VarBinViewArray::from_iter(vec![Some("single")], DType::Utf8(Nullability::NonNullable)))] #[case(VarBinViewArray::from_iter(vec![Some("very long string that exceeds the inline size to test view functionality with multiple buffers")], DType::Utf8(Nullability::NonNullable)))] fn test_cast_varbinview_conformance(#[case] array: VarBinViewArray) { - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/vortex-array/src/arrays/varbinview/compute/mask.rs b/vortex-array/src/arrays/varbinview/compute/mask.rs index 6172c42ab7d..e470dceba32 100644 --- a/vortex-array/src/arrays/varbinview/compute/mask.rs +++ b/vortex-array/src/arrays/varbinview/compute/mask.rs @@ -33,6 +33,8 @@ impl MaskReduce for VarBinView { #[cfg(test)] mod tests { use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::VarBinViewArray; use crate::compute::conformance::mask::test_mask_conformance; @@ -40,6 +42,7 @@ mod tests { fn take_mask_var_bin_view_array() { test_mask_conformance( &VarBinViewArray::from_iter_str(["one", "two", "three", "four", "five"]).into_array(), + &mut array_session().create_execution_ctx(), ); test_mask_conformance( @@ -51,6 +54,7 @@ mod tests { Some("five"), ]) .into_array(), + &mut array_session().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/varbinview/compute/mod.rs b/vortex-array/src/arrays/varbinview/compute/mod.rs index 0358ad01460..e25b6eafb53 100644 --- a/vortex-array/src/arrays/varbinview/compute/mod.rs +++ b/vortex-array/src/arrays/varbinview/compute/mod.rs @@ -15,8 +15,6 @@ mod tests { use crate::IntoArray; use crate::arrays::VarBinViewArray; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; #[test] fn take_nullable() -> VortexResult<()> { let arr = VarBinViewArray::from_iter_nullable_str([ @@ -32,8 +30,7 @@ mod tests { assert!(taken.dtype().is_nullable()); let mut ctx = array_session().create_execution_ctx(); - #[expect(deprecated)] - let taken = taken.to_varbinview(); + let taken = taken.execute::(&mut ctx)?; let mask = taken.validity()?.execute_mask(taken.len(), &mut ctx)?; let result = (0..taken.len()) .map(|i| { diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 1b7435729f1..ce10abda3d0 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -173,6 +173,9 @@ mod tests { ))] #[case(VarBinViewArray::from_iter(["single"].map(Some), DType::Utf8(NonNullable)))] fn test_take_varbinview_conformance(#[case] array: VarBinViewArray) { - test_take_conformance(&array.into_array()); + test_take_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } diff --git a/vortex-array/src/arrays/varbinview/compute/zip.rs b/vortex-array/src/arrays/varbinview/compute/zip.rs index 75e30808821..8966a327f9d 100644 --- a/vortex-array/src/arrays/varbinview/compute/zip.rs +++ b/vortex-array/src/arrays/varbinview/compute/zip.rs @@ -57,7 +57,7 @@ impl ZipKernel for VarBinView { let true_validity = if_true.varbinview_validity().execute_mask(len, ctx)?; let false_validity = if_false.varbinview_validity().execute_mask(len, ctx)?; - let mask = mask.try_to_mask_fill_null_false(ctx)?; + let mask = mask.clone().null_as_false().execute(ctx)?; let if_false_view = if_false; match mask.slices() { AllOr::All => push_range( @@ -217,8 +217,6 @@ mod tests { use crate::array_session; use crate::arrays::VarBinViewArray; use crate::builtins::ArrayBuiltins; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::Nullability; @@ -250,12 +248,12 @@ mod tests { let mask = Mask::from_iter([true, false, true, false, false, true]); - #[expect(deprecated)] + let mut ctx = array_session().create_execution_ctx(); let zipped = mask .clone() .into_array() .zip(a.into_array(), b.into_array())? - .to_varbinview(); + .execute::(&mut ctx)?; let mut ctx = array_session().create_execution_ctx(); let validity_mask = zipped.validity()?.execute_mask(zipped.len(), &mut ctx)?; diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index 9335ed092ca..f730282ddb4 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -245,12 +245,10 @@ impl VTable for VarBinView { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> 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(()) + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + vortex_bail!("append_to_builder for VarBinView requires a VarBinViewBuilder"); + }; + builder.append_varbinview_array(&array.into_owned(), ctx) } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrow/executor/map.rs b/vortex-array/src/arrow/executor/map.rs deleted file mode 100644 index 9a84281e60d..00000000000 --- a/vortex-array/src/arrow/executor/map.rs +++ /dev/null @@ -1,54 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::Arc; - -use arrow_array::ArrayRef as ArrowArrayRef; -use arrow_array::MapArray as ArrowMapArray; -use arrow_array::StructArray as ArrowStructArray; -use arrow_schema::FieldRef; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; - -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::arrow::executor::list::to_arrow_list; - -/// Convert a Vortex List> array into an Arrow MapArray. -pub(super) fn to_arrow_map( - array: ArrayRef, - entries_field: &FieldRef, - ordered: bool, - ctx: &mut ExecutionCtx, -) -> VortexResult { - // First, convert to Arrow ListArray since Map uses i32 offsets. - let list_array = to_arrow_list::(array, entries_field, ctx)?; - - // Downcast to GenericListArray to extract its components. - let Some(list_array) = list_array - .as_any() - .downcast_ref::>() - else { - vortex_bail!("to_arrow_list returned a non-ListArray when building a MapArray"); - }; - - // Extract components from the ListArray. - let (_list_field, offsets, entries, nulls) = list_array.clone().into_parts(); - - // The entries should be a StructArray. Downcast it. - let Some(entries_struct) = entries.as_any().downcast_ref::() else { - vortex_bail!("Map entries must be a StructArray"); - }; - let entries_struct = entries_struct.clone(); - - // Build the MapArray from the components. - let map_array = ArrowMapArray::try_new( - Arc::clone(entries_field), - offsets, - entries_struct, - nulls, - ordered, - )?; - - Ok(Arc::new(map_array)) -} diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index f26cbf04b57..73d2089e221 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -5,7 +5,6 @@ use std::any::Any; use std::mem; use vortex_buffer::BitBufferMut; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_mask::Mask; @@ -13,16 +12,12 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::builders::ArrayBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::LazyBitBufferBuilder; use crate::canonical::Canonical; -#[expect(deprecated)] -use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -127,13 +122,6 @@ impl ArrayBuilder for BoolBuilder { Ok(()) } - unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { - #[expect(deprecated)] - let bool_array = array.to_bool(); - 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) { self.inner.reserve(additional); self.nulls.reserve_exact(additional); @@ -147,7 +135,7 @@ impl ArrayBuilder for BoolBuilder { self.finish_into_bool().into_array() } - fn finish_into_canonical(&mut self) -> Canonical { + fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { Canonical::Bool(self.finish_into_bool()) } } @@ -170,8 +158,6 @@ mod tests { use crate::builders::BoolBuilder; use crate::builders::bool::BoolArray; use crate::builders::builder_with_capacity; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -205,10 +191,8 @@ mod tests { .clone() .append_to_builder(builder.as_mut(), &mut ctx)?; - #[expect(deprecated)] - let canon_into = builder.finish().to_bool(); - #[expect(deprecated)] - let into_canon = chunk.to_bool(); + let canon_into = builder.finish().execute::(&mut ctx)?; + let into_canon = chunk.clone().execute::(&mut ctx)?; assert!(canon_into.validity()?.mask_eq( &into_canon.validity()?, diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index aa777051e39..c08dfbc4d87 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -12,11 +12,8 @@ use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; -#[expect(deprecated)] -use crate::ToCanonical as _; -use crate::VortexSessionExecute; use crate::arrays::DecimalArray; use crate::builders::ArrayBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; @@ -129,6 +126,28 @@ impl DecimalBuilder { self.nulls.append_n_non_nulls(n); } + /// Appends the values of a canonical [`DecimalArray`] to the builder, coercing the physical + /// storage type to the builder's type as needed. + pub(crate) fn append_decimal_array( + &mut self, + array: &DecimalArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + match_each_decimal_value_type!(array.values_type(), |D| { + // Extends the values buffer from another buffer of type D where D can be coerced to the + // builder type. + self.values.extend(array.buffer::().iter().copied()); + }); + + self.nulls.append_validity_mask( + &array + .as_ref() + .validity()? + .execute_mask(array.as_ref().len(), ctx)?, + ); + Ok(()) + } + /// Finishes the builder directly into a [`DecimalArray`]. pub fn finish_into_decimal(&mut self) -> DecimalArray { let validity = self.nulls.finish_with_nullability(self.dtype.nullability()); @@ -195,30 +214,6 @@ impl ArrayBuilder for DecimalBuilder { Ok(()) } - unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { - #[expect(deprecated)] - let decimal_array = array.to_decimal(); - - match_each_decimal_value_type!(decimal_array.values_type(), |D| { - // Extends the values buffer from another buffer of type D where D can be coerced to the - // builder type. - self.values - .extend(decimal_array.buffer::().iter().copied()); - }); - - self.nulls.append_validity_mask( - &decimal_array - .as_ref() - .validity() - .vortex_expect("validity_mask") - .execute_mask( - decimal_array.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("Failed to compute validity mask"), - ); - } - fn reserve_exact(&mut self, additional: usize) { self.values.reserve(additional); self.nulls.reserve_exact(additional); @@ -232,7 +227,7 @@ impl ArrayBuilder for DecimalBuilder { self.finish_into_decimal().into_array() } - fn finish_into_canonical(&mut self) -> Canonical { + fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { Canonical::Decimal(self.finish_into_decimal()) } } @@ -328,7 +323,8 @@ mod tests { let i8s = i8s.finish(); let mut i128s = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); - i128s.extend_from_array(&i8s); + i8s.append_to_builder(&mut i128s, &mut array_session().create_execution_ctx()) + .unwrap(); let i128s = i128s.finish(); for i in 0..i8s.len() { diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index 61d5fe8d7b0..aa8e1b76aa2 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -8,6 +8,7 @@ use vortex_error::vortex_ensure; use vortex_mask::Mask; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ExtensionArray; use crate::arrays::extension::ExtensionArrayExt; @@ -15,8 +16,6 @@ use crate::builders::ArrayBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::builder_with_capacity; use crate::canonical::Canonical; -#[expect(deprecated)] -use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::extension::ExtDTypeRef; use crate::scalar::ExtScalar; @@ -47,6 +46,18 @@ impl ExtensionBuilder { self.storage.append_scalar(&value.to_storage_scalar()) } + /// Appends the values of a canonical [`ExtensionArray`] to the builder by appending its + /// storage array to the underlying storage builder. + pub(crate) fn append_extension_array( + &mut self, + array: &ExtensionArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + array + .storage_array() + .append_to_builder(self.storage.as_mut(), ctx) + } + /// Finishes the builder directly into a [`ExtensionArray`]. pub fn finish_into_extension(&mut self) -> ExtensionArray { let storage = self.storage.finish(); @@ -101,12 +112,6 @@ impl ArrayBuilder for ExtensionBuilder { self.append_value(scalar.as_extension()) } - unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { - #[expect(deprecated)] - let ext_array = array.to_extension(); - self.storage.extend_from_array(ext_array.storage_array()) - } - fn reserve_exact(&mut self, capacity: usize) { self.storage.reserve_exact(capacity) } @@ -119,7 +124,7 @@ impl ArrayBuilder for ExtensionBuilder { self.finish_into_extension().into_array() } - fn finish_into_canonical(&mut self) -> Canonical { + fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { Canonical::Extension(self.finish_into_extension()) } } diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 9129d8cb644..9150d37b50b 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -14,8 +14,6 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::builders::ArrayBuilder; @@ -23,8 +21,6 @@ use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::LazyBitBufferBuilder; use crate::builders::builder_with_capacity; use crate::canonical::Canonical; -#[expect(deprecated)] -use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::ListScalar; @@ -108,6 +104,25 @@ impl FixedSizeListBuilder { Ok(()) } + /// Appends the values of a canonical [`FixedSizeListArray`] to the builder, recursing into the + /// elements builder. + pub(crate) fn append_fixed_size_list_array( + &mut self, + array: &FixedSizeListArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if array.is_empty() { + return Ok(()); + } + + array + .elements() + .append_to_builder(self.elements_builder.as_mut(), ctx)?; + self.nulls + .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); + Ok(()) + } + /// Appends a fixed-size list `value` to the builder. /// /// Note that a [`ListScalar`] can represent both a [`ListArray`] scalar **and** a @@ -242,23 +257,6 @@ impl ArrayBuilder for FixedSizeListBuilder { /// This will increase the capacity if extending with this `array` would go past the original /// capacity. - unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { - #[expect(deprecated)] - let fsl = array.to_fixed_size_list(); - if fsl.is_empty() { - return; - } - - self.elements_builder.extend_from_array(fsl.elements()); - self.nulls.append_validity_mask( - &array - .validity() - .vortex_expect("validity_mask in extend_from_array_unchecked") - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) - .vortex_expect("Failed to compute validity mask"), - ); - } - fn reserve_exact(&mut self, additional: usize) { self.elements_builder .reserve_exact(additional * self.list_size() as usize); @@ -273,7 +271,7 @@ impl ArrayBuilder for FixedSizeListBuilder { self.finish_into_fixed_size_list().into_array() } - fn finish_into_canonical(&mut self) -> Canonical { + fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { Canonical::FixedSizeList(self.finish_into_fixed_size_list()) } } @@ -287,8 +285,6 @@ mod tests { use super::FixedSizeListBuilder; use crate::IntoArray as _; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::PrimitiveArray; @@ -342,8 +338,8 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 2); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let mut ctx = array_session().create_execution_ctx(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.elements().len(), 6); assert_eq!(fsl_array.list_size(), 3); } @@ -366,8 +362,8 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 100); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let mut ctx = array_session().create_execution_ctx(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.list_size(), 0); // The elements array should be empty since list_size is 0. assert_eq!(fsl_array.elements().len(), 0); @@ -396,8 +392,8 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 100); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let mut ctx = array_session().create_execution_ctx(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.list_size(), 0); assert_eq!(fsl_array.elements().len(), 0); } @@ -426,8 +422,8 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 5); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let mut ctx = array_session().create_execution_ctx(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.elements().len(), 10); } @@ -440,8 +436,8 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 0); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let mut ctx = array_session().create_execution_ctx(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.list_size(), 100000000); assert_eq!(fsl_array.elements().len(), 0); } @@ -474,8 +470,7 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 3); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert!( fsl_array .validity() @@ -538,8 +533,8 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 2); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let mut ctx = array_session().create_execution_ctx(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.elements().len(), 6); } @@ -553,14 +548,17 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 5); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let mut ctx = array_session().create_execution_ctx(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.list_size(), 3); assert_eq!(fsl_array.elements().len(), 15); // Check that all elements are zeros. - #[expect(deprecated)] - let elements_array = fsl_array.elements().to_primitive(); + let elements_array = fsl_array + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); let elements = elements_array.as_slice::(); assert!(elements.iter().all(|&x| x == 0)); } @@ -579,8 +577,7 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 3); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.list_size(), 2); // Check that all lists are null. @@ -611,8 +608,7 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 1); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.list_size(), 2); // Check that all lists are null. @@ -637,8 +633,8 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 1000); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let mut ctx = array_session().create_execution_ctx(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.list_size(), 0); assert_eq!(fsl_array.elements().len(), 0); } @@ -684,14 +680,17 @@ mod tests { let mut builder = FixedSizeListBuilder::with_capacity(dtype, 2, Nullable, 0); let source_array = source.into_array(); - builder.extend_from_array(&source_array); - builder.extend_from_array(&source_array); + source_array + .append_to_builder(&mut builder, &mut ctx) + .unwrap(); + source_array + .append_to_builder(&mut builder, &mut ctx) + .unwrap(); let fsl = builder.finish(); assert_eq!(fsl.len(), 6); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.elements().len(), 12); // Check validity pattern is repeated. @@ -761,14 +760,19 @@ mod tests { let mut builder = FixedSizeListBuilder::with_capacity(dtype, 0, Nullable, 0); - builder.extend_from_array(&source1.into_array()); - builder.extend_from_array(&source2.into_array()); + source1 + .into_array() + .append_to_builder(&mut builder, &mut ctx) + .unwrap(); + source2 + .into_array() + .append_to_builder(&mut builder, &mut ctx) + .unwrap(); let fsl = builder.finish(); assert_eq!(fsl.len(), 5); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.list_size(), 0); assert_eq!(fsl_array.elements().len(), 0); @@ -812,6 +816,7 @@ mod tests { #[test] fn test_extend_empty_array() { + let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); // Create an empty source array. @@ -838,7 +843,10 @@ mod tests { .unwrap(); // Extend with empty array (should be no-op). - builder.extend_from_array(&source.into_array()); + source + .into_array() + .append_to_builder(&mut builder, &mut ctx) + .unwrap(); let fsl = builder.finish(); assert_eq!(fsl.len(), 1); @@ -876,13 +884,15 @@ mod tests { Validity::AllValid, 1, ); - builder.extend_from_array(&source.into_array()); + source + .into_array() + .append_to_builder(&mut builder, &mut ctx) + .unwrap(); let fsl = builder.finish(); assert_eq!(fsl.len(), 6); - #[expect(deprecated)] - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.execute::(&mut ctx).unwrap(); assert_eq!(fsl_array.elements().len(), 12); // Check validity. @@ -927,7 +937,7 @@ mod tests { .vortex_expect("fixed-size-list validity should be derivable") .execute_is_valid(5, &mut ctx) .unwrap() - ); // extend_from_array + ); } #[test] @@ -1043,8 +1053,11 @@ mod tests { assert_eq!(fsl.list_size(), 3); // Verify elements array: [1, 2, 3, 10, 11, 12, 4, 5, 6, 20, 21, 22]. - #[expect(deprecated)] - let elements = fsl.elements().to_primitive(); + let elements = fsl + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!( elements.as_slice::(), &[1, 2, 3, 10, 11, 12, 4, 5, 6, 20, 21, 22] diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index 12c74ba53a4..130f37534cf 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -4,6 +4,7 @@ use std::any::Any; use std::sync::Arc; +use num_traits::AsPrimitive; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -15,17 +16,19 @@ use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; +use crate::array::ArrayView; +use crate::arrays::List; use crate::arrays::ListArray; +use crate::arrays::ListView; +use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::list::ListArrayExt; use crate::arrays::listview::ListViewArrayExt; use crate::builders::ArrayBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::LazyBitBufferBuilder; use crate::builders::PrimitiveBuilder; use crate::builders::builder_with_capacity; -#[expect(deprecated)] -use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::Nullability; @@ -108,6 +111,7 @@ impl ListBuilder { self.element_dtype() ); + self.elements_builder.reserve_exact(array.len()); array.append_to_builder(self.elements_builder.as_mut(), ctx)?; self.nulls.append_non_null(); self.offsets_builder.append_value( @@ -172,6 +176,55 @@ impl ListBuilder { } } +/// Appends `ListViewArray`-layout lists (`n` offsets and sizes) into a [`ListBuilder`], converting +/// into the `ListArray` (`n + 1` offsets) layout. +fn extend_from_listview( + builder: &mut ListBuilder, + new_elements: &ArrayRef, + new_offsets: &[OffsetType], + new_sizes: &[SizeType], + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + O: IntegerPType, + OffsetType: IntegerPType, + SizeType: IntegerPType, +{ + let num_lists = new_offsets.len(); + debug_assert_eq!(num_lists, new_sizes.len()); + + let total_elements: usize = new_sizes.iter().map(|size| size.as_()).sum(); + builder.elements_builder.reserve_exact(total_elements); + + let mut curr_offset = builder.elements_builder.len(); + builder.offsets_builder.reserve_exact(num_lists); + let mut offsets_range = builder.offsets_builder.uninit_range(num_lists); + + // We need to append each list individually, converting from `ListViewArray` format to + // the `ListArray` format that `ListBuilder` expects. + for i in 0..new_offsets.len() { + let offset: usize = new_offsets[i].as_(); + let size: usize = new_sizes[i].as_(); + + if size > 0 { + let list_elements = new_elements + .slice(offset..offset + size) + .vortex_expect("list builder slice"); + list_elements.append_to_builder(builder.elements_builder.as_mut(), ctx)?; + curr_offset += size; + } + + let new_offset = O::from_usize(curr_offset).vortex_expect("Failed to convert offset"); + + offsets_range.set_value(i, new_offset); + } + + // SAFETY: We have initialized all `num_lists` values, and since the `offsets` array is + // non-nullable, we are done. + unsafe { offsets_range.finish() }; + Ok(()) +} + impl ArrayBuilder for ListBuilder { fn as_any(&self) -> &dyn Any { self @@ -222,100 +275,104 @@ impl ArrayBuilder for ListBuilder { self.append_value(scalar.as_list()) } - unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { - #[expect(deprecated)] - let list = array.to_listview(); - if list.is_empty() { - return; - } + fn reserve_exact(&mut self, additional: usize) { + self.elements_builder.reserve_exact(additional); + self.offsets_builder.reserve_exact(additional); + self.nulls.reserve_exact(additional); + } - // Append validity information. - self.nulls.append_validity_mask( - &array - .validity() - .vortex_expect("validity_mask in extend_from_array_unchecked") - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) - .vortex_expect("Failed to compute validity mask"), - ); + unsafe fn set_validity_unchecked(&mut self, validity: Mask) { + self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); + } - // Note that `ListViewArray` has `n` offsets and sizes, not `n+1` offsets like `ListArray`. - let elements = list.elements(); - #[expect(deprecated)] - let offsets = list.offsets().to_primitive(); - #[expect(deprecated)] - let sizes = list.sizes().to_primitive(); - - fn extend_inner( - builder: &mut ListBuilder, - new_elements: &ArrayRef, - new_offsets: &[OffsetType], - new_sizes: &[SizeType], - ) where - O: IntegerPType, - OffsetType: IntegerPType, - SizeType: IntegerPType, - { - let num_lists = new_offsets.len(); - debug_assert_eq!(num_lists, new_sizes.len()); - - let mut curr_offset = builder.elements_builder.len(); - let mut offsets_range = builder.offsets_builder.uninit_range(num_lists); - - // We need to append each list individually, converting from `ListViewArray` format to - // the `ListArray` format that `ListBuilder` expects. - for i in 0..new_offsets.len() { - let offset: usize = new_offsets[i].as_(); - let size: usize = new_sizes[i].as_(); - - if size > 0 { - let list_elements = new_elements - .slice(offset..offset + size) - .vortex_expect("list builder slice"); - builder.elements_builder.extend_from_array(&list_elements); - curr_offset += size; - } + fn finish(&mut self) -> ArrayRef { + self.finish_into_list().into_array() + } - let new_offset = - O::from_usize(curr_offset).vortex_expect("Failed to convert offset"); + fn finish_into_canonical(&mut self, ctx: &mut ExecutionCtx) -> Canonical { + let listview = self + .finish() + .execute::(ctx) + .vortex_expect("list builder should canonicalize to listview"); + Canonical::List(listview) + } + + fn append_list_array( + &mut self, + array: ArrayView<'_, List>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if array.is_empty() { + return Ok(()); + } + + self.nulls + .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); - offsets_range.set_value(i, new_offset); + let num_lists = array.len(); + let offsets = array.offsets().clone().execute::(ctx)?; + match_each_integer_ptype!(offsets.ptype(), |OffsetType| { + let offsets = offsets.as_slice::(); + let first: usize = offsets[0].as_(); + let last: usize = offsets[num_lists].as_(); + + // Lists in a `ListArray` are contiguous, so the referenced elements can be appended + // in bulk and the offsets rebased onto this builder's elements. + let elements_base = self.elements_builder.len(); + if last > first { + self.elements_builder.reserve_exact(last - first); + array + .elements() + .slice(first..last)? + .append_to_builder(self.elements_builder.as_mut(), ctx)?; } + self.offsets_builder.reserve_exact(num_lists); + let mut offsets_range = self.offsets_builder.uninit_range(num_lists); + for i in 0..num_lists { + let end: usize = offsets[i + 1].as_(); + offsets_range.set_value( + i, + O::from_usize(end - first + elements_base) + .vortex_expect("Failed to convert offset"), + ); + } // SAFETY: We have initialized all `num_lists` values, and since the `offsets` array is // non-nullable, we are done. unsafe { offsets_range.finish() }; + }); + Ok(()) + } + + fn append_listview_array( + &mut self, + array: ArrayView<'_, ListView>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if array.is_empty() { + return Ok(()); } + self.nulls + .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); + + // Note that `ListViewArray` has `n` offsets and sizes, not `n+1` offsets like `ListArray`. + let elements = array.elements(); + let offsets = array.offsets().clone().execute::(ctx)?; + let sizes = array.sizes().clone().execute::(ctx)?; + match_each_integer_ptype!(offsets.ptype(), |OffsetType| { match_each_integer_ptype!(sizes.ptype(), |SizeType| { - extend_inner( + extend_from_listview( self, elements, offsets.as_slice::(), sizes.as_slice::(), - ) + ctx, + )? }) - }) - } - - fn reserve_exact(&mut self, additional: usize) { - self.elements_builder.reserve_exact(additional); - self.offsets_builder.reserve_exact(additional); - self.nulls.reserve_exact(additional); - } - - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - - fn finish(&mut self) -> ArrayRef { - self.finish_into_list().into_array() - } - - fn finish_into_canonical(&mut self) -> Canonical { - #[expect(deprecated)] - let listview = self.finish_into_list().into_array().to_listview(); - Canonical::List(listview) + }); + Ok(()) } } @@ -327,17 +384,19 @@ mod tests { use Nullability::Nullable; use vortex_buffer::buffer; use vortex_error::VortexExpect; + use vortex_error::VortexResult; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::array_session; use crate::arrays::ChunkedArray; + use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::list::ListArrayExt; use crate::arrays::listview::ListViewArrayExt; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; + use crate::builders::ListViewBuilder; + use crate::builders::builder_with_capacity; use crate::builders::list::ListArray; use crate::builders::list::ListBuilder; use crate::dtype::DType; @@ -387,8 +446,8 @@ mod tests { let list = builder.finish(); assert_eq!(list.len(), 2); - #[expect(deprecated)] - let list_array = list.to_listview(); + let mut ctx = array_session().create_execution_ctx(); + let list_array = list.execute::(&mut ctx).unwrap(); assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3); assert_eq!(list_array.list_elements_at(1).unwrap().len(), 3); @@ -440,8 +499,8 @@ mod tests { let list = builder.finish(); assert_eq!(list.len(), 3); - #[expect(deprecated)] - let list_array = list.to_listview(); + let mut ctx = array_session().create_execution_ctx(); + let list_array = list.execute::(&mut ctx).unwrap(); assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3); assert_eq!(list_array.list_elements_at(1).unwrap().len(), 0); @@ -453,18 +512,24 @@ mod tests { [Some(vec![0, 1, 2]), None, Some(vec![4, 5])], Arc::new(I32.into()), ) - .unwrap(); + .unwrap() + .into_array(); assert_eq!(list.len(), 3); 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); - builder.extend_from_array(&list); - builder.extend_from_array(&list.slice(0..0).unwrap()); - builder.extend_from_array(&list.slice(1..3).unwrap()); + list.append_to_builder(&mut builder, &mut ctx).unwrap(); + list.append_to_builder(&mut builder, &mut ctx).unwrap(); + list.slice(0..0) + .unwrap() + .append_to_builder(&mut builder, &mut ctx) + .unwrap(); + list.slice(1..3) + .unwrap() + .append_to_builder(&mut builder, &mut ctx) + .unwrap(); - #[expect(deprecated)] let expected = ListArray::from_iter_opt_slow::( [ Some(vec![0, 1, 2]), @@ -479,9 +544,11 @@ mod tests { Arc::new(DType::Primitive(I32, NonNullable)), ) .unwrap() - .to_listview(); + .into_array() + .execute::(&mut ctx) + .unwrap(); - let actual = builder.finish_into_canonical().into_listview(); + let actual = builder.finish_into_canonical(&mut ctx).into_listview(); assert_arrays_eq!(actual.elements(), expected.elements(), &mut ctx); @@ -502,6 +569,90 @@ mod tests { ); } + /// `append_to_builder` must handle any list builder kind without assuming the offset/size + /// integer types produced by `builder_with_capacity`. It appends a `List`-encoded array and a + /// `ListView`-encoded array into `ListViewBuilder`s and `ListBuilder`s with assorted (and + /// non-`u64`) offset/size types. + #[test] + fn test_append_to_builder_any_list_builder() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let list = ListArray::from_iter_opt_slow::( + [Some(vec![0, 1, 2]), None, Some(vec![4, 5])], + Arc::new(I32.into()), + )? + .into_array(); + let listview = list + .clone() + .execute::(&mut ctx)? + .into_array(); + let elem_dtype = || Arc::new(I32.into()); + + // `builder_with_capacity` produces a `ListViewBuilder` for `DType::List`; appending the + // `List`-encoded array must dispatch into it instead of bailing. + let mut listview_builder = builder_with_capacity(list.dtype(), list.len()); + list.append_to_builder(listview_builder.as_mut(), &mut ctx)?; + assert_arrays_eq!(listview_builder.finish(), list, &mut ctx); + + // A `ListViewBuilder` with non-`u64` (including signed) offset and size types must work + // for both source encodings. + let mut lv_u32_u8 = ListViewBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); + list.append_to_builder(&mut lv_u32_u8, &mut ctx)?; + assert_arrays_eq!(lv_u32_u8.finish(), list, &mut ctx); + + let mut lv_i32_i16 = + ListViewBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); + list.append_to_builder(&mut lv_i32_i16, &mut ctx)?; + assert_arrays_eq!(lv_i32_i16.finish(), list, &mut ctx); + + let mut lv_u16_u16 = + ListViewBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); + listview.append_to_builder(&mut lv_u16_u16, &mut ctx)?; + assert_arrays_eq!(lv_u16_u16.finish(), list, &mut ctx); + + // Both source encodings appended into `ListBuilder`s with non-`u64` (including signed) + // offset types. + let mut list_builder = ListBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); + list.append_to_builder(&mut list_builder, &mut ctx)?; + assert_arrays_eq!(list_builder.finish(), list, &mut ctx); + + let mut list_builder_i16 = ListBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); + listview.append_to_builder(&mut list_builder_i16, &mut ctx)?; + assert_arrays_eq!(list_builder_i16.finish(), list, &mut ctx); + + Ok(()) + } + + #[test] + fn test_append_list_arrays_grow_builder() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype: Arc = Arc::new(I32.into()); + + // Enough lists to exceed the offsets capacity of a zero-capacity builder, so appending + // must grow the builder rather than panic in `uninit_range`. + let lists: Vec>> = + (0..100).map(|i| (i % 10 != 0).then(|| vec![i])).collect(); + let source = ListArray::from_iter_opt_slow::(lists.clone(), Arc::clone(&dtype))?; + let expected = ListArray::from_iter_opt_slow::( + lists.iter().cloned().chain(lists.iter().cloned()), + Arc::clone(&dtype), + )?; + + // Appending twice checks growth from a non-empty builder and offset rebasing. + let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + builder.append_list_array(source.as_view(), &mut ctx)?; + builder.append_list_array(source.as_view(), &mut ctx)?; + assert_arrays_eq!(builder.finish(), expected, &mut ctx); + + let source_listview = source.into_array().execute::(&mut ctx)?; + let mut builder = ListBuilder::::with_capacity(dtype, Nullable, 0, 0); + builder.append_listview_array(source_listview.as_view(), &mut ctx)?; + builder.append_listview_array(source_listview.as_view(), &mut ctx)?; + assert_arrays_eq!(builder.finish(), expected, &mut ctx); + + Ok(()) + } + #[test] fn test_extend_builder() { test_extend_builder_gen::(); @@ -539,8 +690,13 @@ mod tests { DType::List(Arc::new(DType::Primitive(I32, NonNullable)), NonNullable), ); - #[expect(deprecated)] - let canon_values = chunked_list.unwrap().as_array().to_listview(); + let mut ctx = array_session().create_execution_ctx(); + let canon_values = chunked_list + .unwrap() + .as_array() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!( one_trailing_unused_element diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index cb131048339..4e6523bd7d3 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -21,11 +21,13 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; +use crate::array::ArrayView; use crate::array::IntoArray; +use crate::arrays::List; +use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::list::ListArrayExt; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::listview::ListViewRebuildMode; use crate::builders::ArrayBuilder; @@ -140,6 +142,7 @@ impl ListViewBuilder { "appending this list would cause an offset overflow" ); + self.elements_builder.reserve_exact(num_elements); array.append_to_builder(self.elements_builder.as_mut(), ctx)?; self.nulls.append_non_null(); @@ -208,8 +211,9 @@ impl ListViewBuilder { // - The offsets, sizes, and validity have the same length since we always appended the same // amount. // - We checked on construction that the sizes type fits into the offsets. - // - In every method that adds values to this builder (`append_value`, `append_scalar`, and - // `extend_from_array_unchecked`), we checked that `offset + size` does not overflow. + // - In every method that adds values to this builder (`append_value`, `append_scalar`, + // `append_list_array`, and `append_listview_array`), we checked that `offset + size` + // does not overflow. // - We constructed everything in a way that builds the `ListViewArray` similar to the shape // of a `ListArray`, so we know the resulting array is zero-copyable to a `ListArray`. unsafe { @@ -296,41 +300,75 @@ impl ArrayBuilder for ListViewBuilder { self.append_value(list_scalar) } - unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { - // TODO: The `ArrayBuilder` trait does not thread an `ExecutionCtx` through its extend - // methods, so we are forced to mint a fresh `LEGACY_SESSION` context here on every call - // (which for chunked input means once per chunk). Once the trait carries a `&mut - // ExecutionCtx`, the caller's session should be reused instead. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + fn reserve_exact(&mut self, capacity: usize) { + self.elements_builder.reserve_exact(capacity * 2); + self.offsets_builder.reserve_exact(capacity); + self.sizes_builder.reserve_exact(capacity); + self.nulls.reserve_exact(capacity); + } - let listview = array - .clone() - .execute::(&mut ctx) - .vortex_expect("failed to execute array into ListViewArray in extend_from_array"); - if listview.is_empty() { - return; + unsafe fn set_validity_unchecked(&mut self, validity: Mask) { + self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); + } + + fn finish(&mut self) -> ArrayRef { + self.finish_into_listview().into_array() + } + + fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { + Canonical::List(self.finish_into_listview()) + } + + fn append_list_array( + &mut self, + array: ArrayView<'_, List>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if array.is_empty() { + return Ok(()); + } + + self.nulls + .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); + + let offsets = array.offsets().clone().execute::(ctx)?; + match_each_integer_ptype!(offsets.ptype(), |OffsetType| { + extend_from_list( + self, + array.elements(), + offsets.as_slice::(), + ctx, + )? + }); + Ok(()) + } + + fn append_listview_array( + &mut self, + array: ArrayView<'_, ListView>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if array.is_empty() { + return Ok(()); } // Normalize to an exact zero-copy-to-list layout and then bulk append. This avoids the // very expensive scalar_at-per-list path for overlapping / out-of-order list views. - let listview = listview - .rebuild(ListViewRebuildMode::MakeExact, &mut ctx) - .vortex_expect("ListViewArray::rebuild(MakeExact) failed in extend_from_array"); + let listview = array + .into_owned() + .rebuild(ListViewRebuildMode::MakeExact, ctx)?; debug_assert!(listview.is_zero_copy_to_list()); - self.nulls.append_validity_mask( - &array - .validity() - .vortex_expect("validity_mask in extend_from_array_unchecked") - .execute_mask(array.len(), &mut ctx) - .vortex_expect("Failed to compute validity mask"), - ); + self.nulls + .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); // Bulk append the new elements (which should have no gaps or overlaps). let old_elements_len = self.elements_builder.len(); self.elements_builder .reserve_exact(listview.elements().len()); - self.elements_builder.extend_from_array(listview.elements()); + listview + .elements() + .append_to_builder(self.elements_builder.as_mut(), ctx)?; let new_elements_len = self.elements_builder.len(); // Reserve enough space for the new views. @@ -342,23 +380,15 @@ impl ArrayBuilder for ListViewBuilder { let cast_sizes = listview .sizes() .clone() - .cast(self.sizes_builder.dtype().clone()) - .vortex_expect( - "was somehow unable to cast the new sizes to the type of the builder sizes", - ); - self.sizes_builder.extend_from_array(&cast_sizes); + .cast(self.sizes_builder.dtype().clone())?; + cast_sizes.append_to_builder(&mut self.sizes_builder, ctx)?; // Now we need to adjust all of the offsets by adding the current number of elements in the // builder. - let uninit_range = self.offsets_builder.uninit_range(extend_length); // This should be cheap because we didn't compress after rebuilding. - let new_offsets = listview - .offsets() - .clone() - .execute::(&mut ctx) - .vortex_expect("failed to execute list view offsets into a PrimitiveArray"); + let new_offsets = listview.offsets().clone().execute::(ctx)?; match_each_integer_ptype!(new_offsets.ptype(), |A| { adjust_and_extend_offsets::( @@ -367,33 +397,75 @@ impl ArrayBuilder for ListViewBuilder { old_elements_len, new_elements_len, ); - }) - } - - fn reserve_exact(&mut self, capacity: usize) { - self.elements_builder.reserve_exact(capacity * 2); - self.offsets_builder.reserve_exact(capacity); - self.sizes_builder.reserve_exact(capacity); - self.nulls.reserve_exact(capacity); - } - - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); + }); + Ok(()) } +} - fn finish(&mut self) -> ArrayRef { - self.finish_into_listview().into_array() +/// Appends `ListArray`-layout lists (`n + 1` cumulative offsets) into a [`ListViewBuilder`]. +/// +/// Lists in a `ListArray` are contiguous, so the referenced elements can be appended in bulk, +/// with the offsets rebased onto the builder's elements and the sizes taken from consecutive +/// offset differences. +fn extend_from_list( + builder: &mut ListViewBuilder, + elements: &ArrayRef, + offsets: &[OffsetType], + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + O: IntegerPType, + S: IntegerPType, + OffsetType: IntegerPType, +{ + let num_lists = offsets.len() - 1; + let first: usize = offsets[0].as_(); + let last: usize = offsets[num_lists].as_(); + + let elements_base = builder.elements_builder.len(); + + // We must assert this even in release mode to ensure that the safety comment in + // `finish_into_listview` is correct. + assert!( + ((elements_base + (last - first)) as u64) < O::max_value_as_u64(), + "appending this list would cause an offset overflow" + ); + + if last > first { + builder.elements_builder.reserve_exact(last - first); + elements + .slice(first..last)? + .append_to_builder(builder.elements_builder.as_mut(), ctx)?; } - fn finish_into_canonical(&mut self) -> Canonical { - Canonical::List(self.finish_into_listview()) + builder.offsets_builder.reserve_exact(num_lists); + builder.sizes_builder.reserve_exact(num_lists); + let mut offsets_range = builder.offsets_builder.uninit_range(num_lists); + let mut sizes_range = builder.sizes_builder.uninit_range(num_lists); + for i in 0..num_lists { + let start: usize = offsets[i].as_(); + let end: usize = offsets[i + 1].as_(); + offsets_range.set_value( + i, + O::from_usize(start - first + elements_base) + .vortex_expect("Failed to convert from usize to `O`"), + ); + sizes_range.set_value( + i, + S::from_usize(end - start).vortex_expect("Failed to convert from usize to `S`"), + ); } + // SAFETY: We have initialized all `num_lists` values in both ranges, and both the `offsets` + // and the `sizes` builders are non-nullable. + unsafe { offsets_range.finish() }; + unsafe { sizes_range.finish() }; + Ok(()) } /// Given new offsets, adds them to the `UninitRange` after adding the `old_elements_len` to each /// offset. -fn adjust_and_extend_offsets<'a, O: IntegerPType, A: IntegerPType>( - mut uninit_range: UninitRange<'a, O>, +fn adjust_and_extend_offsets( + mut uninit_range: UninitRange, new_offsets: PrimitiveArray, old_elements_len: usize, new_elements_len: usize, @@ -436,6 +508,7 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexExpect; + use vortex_error::VortexResult; use super::ListViewBuilder; use crate::IntoArray; @@ -674,7 +747,13 @@ mod tests { .unwrap(); // Extend from the ListArray. - builder.extend_from_array(&source.into_array()); + let source = source + .into_array() + .execute::(&mut ctx) + .unwrap(); + builder + .append_listview_array(source.as_view(), &mut ctx) + .unwrap(); // Extend from empty array (should be no-op). let empty_source = ListArray::from_iter_opt_slow::>( @@ -682,7 +761,13 @@ mod tests { Arc::new(I32.into()), ) .unwrap(); - builder.extend_from_array(&empty_source.into_array()); + let empty_source = empty_source + .into_array() + .execute::(&mut ctx) + .unwrap(); + builder + .append_listview_array(empty_source.as_view(), &mut ctx) + .unwrap(); let listview = builder.finish_into_listview(); assert_eq!(listview.len(), 4); @@ -719,6 +804,35 @@ mod tests { ); } + #[test] + fn test_append_list_array_grows_builder() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype: Arc = Arc::new(I32.into()); + + // Enough lists to exceed the offsets/sizes capacity of a zero-capacity builder, so + // appending must grow the builder rather than panic in `uninit_range`. + let lists: Vec>> = + (0..100).map(|i| (i % 10 != 0).then(|| vec![i])).collect(); + let source = ListArray::from_iter_opt_slow::(lists.clone(), Arc::clone(&dtype))?; + + let mut builder = + ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + builder.append_list_array(source.as_view(), &mut ctx)?; + // Append a second time to check growth from a non-empty builder and offset rebasing. + builder.append_list_array(source.as_view(), &mut ctx)?; + + let listview = builder.finish_into_listview(); + assert!(listview.is_zero_copy_to_list()); + + let expected = ListArray::from_iter_opt_slow::( + lists.iter().cloned().chain(lists.iter().cloned()), + dtype, + )?; + assert_arrays_eq!(listview, expected, &mut ctx); + + Ok(()) + } + #[test] fn test_extend_from_array_overlapping_listview() { let mut ctx = array_session().create_execution_ctx(); @@ -740,7 +854,9 @@ mod tests { let mut builder = ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); - builder.extend_from_array(&source.into_array()); + builder + .append_listview_array(source.as_view(), &mut ctx) + .unwrap(); let listview = builder.finish_into_listview(); assert_eq!(listview.len(), 3); diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index 48a347aa3d5..561efc2711e 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -34,10 +34,14 @@ use std::any::Any; use std::sync::Arc; use vortex_error::VortexResult; -use vortex_error::vortex_panic; +use vortex_error::vortex_bail; use vortex_mask::Mask; use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::array::ArrayView; +use crate::arrays::List; +use crate::arrays::ListView; use crate::canonical::Canonical; use crate::dtype::DType; use crate::match_each_decimal_value_type; @@ -153,30 +157,6 @@ pub trait ArrayBuilder: Send { /// A generic function to append a scalar to the builder. fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()>; - /// The inner part of `extend_from_array`. - /// - /// # Safety - /// - /// The array that must have an equal [`DType`] to the array builder's `DType` (with nullability - /// superset semantics). - unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef); - - /// Extends the array with the provided array, canonicalizing if necessary. - /// - /// Implementors must validate that the passed in [`ArrayRef`] has the correct [`DType`]. - fn extend_from_array(&mut self, array: &ArrayRef) { - if !self.dtype().eq_with_nullability_superset(array.dtype()) { - vortex_panic!( - "tried to extend a builder with `DType` {} with an array with `DType {}", - self.dtype(), - array.dtype() - ); - } - - // SAFETY: We checked that the array had a valid `DType` above. - unsafe { self.extend_from_array_unchecked(array) } - } - /// Allocate space for extra `additional` items fn reserve_exact(&mut self, additional: usize); @@ -214,7 +194,40 @@ pub trait ArrayBuilder: Send { /// This method provides a default implementation that creates an [`ArrayRef`] via `finish` and /// then converts it to canonical form. Specific builders can override this with optimized /// implementations that avoid the intermediate [`ArrayRef`] creation. - fn finish_into_canonical(&mut self) -> Canonical; + fn finish_into_canonical(&mut self, ctx: &mut ExecutionCtx) -> Canonical; + + /// Appends the values of a [`List`]-encoded `array` to this builder. + /// + /// Only list-typed builders support this; the default implementation returns an error. List + /// encodings dispatch through this method because the concrete list builders are generic over + /// their offset/size integer types, which cannot be named through a `dyn ArrayBuilder`. + fn append_list_array( + &mut self, + array: ArrayView<'_, List>, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_bail!( + "cannot append a List array of dtype {} to a {} builder", + array.dtype(), + self.dtype() + ) + } + + /// Appends the values of a [`ListView`]-encoded `array` to this builder. + /// + /// See [`append_list_array`](Self::append_list_array); this is the same hook for the canonical + /// [`ListViewArray`](crate::arrays::ListViewArray) encoding. + fn append_listview_array( + &mut self, + array: ArrayView<'_, ListView>, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_bail!( + "cannot append a ListView array of dtype {} to a {} builder", + array.dtype(), + self.dtype() + ) + } } /// Construct a new canonical builder for the given [`DType`]. diff --git a/vortex-array/src/builders/null.rs b/vortex-array/src/builders/null.rs index 45b6e7ed974..28eab14dd7a 100644 --- a/vortex-array/src/builders/null.rs +++ b/vortex-array/src/builders/null.rs @@ -8,6 +8,7 @@ use vortex_error::vortex_ensure; use vortex_mask::Mask; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::NullArray; use crate::builders::ArrayBuilder; @@ -69,10 +70,6 @@ impl ArrayBuilder for NullBuilder { Ok(()) } - unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { - self.append_nulls(array.len()); - } - fn reserve_exact(&mut self, _additional: usize) {} unsafe fn set_validity_unchecked(&mut self, _validity: Mask) {} @@ -81,7 +78,7 @@ impl ArrayBuilder for NullBuilder { NullArray::new(self.length).into_array() } - fn finish_into_canonical(&mut self) -> Canonical { + fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { Canonical::Null(NullArray::new(self.length)) } } diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index e7e8d3d7512..4ce1aafe1f7 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -13,15 +13,11 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::builders::ArrayBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::LazyBitBufferBuilder; use crate::canonical::Canonical; -#[expect(deprecated)] -use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; @@ -200,14 +196,6 @@ impl ArrayBuilder for PrimitiveBuilder { Ok(()) } - unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { - #[expect(deprecated)] - let array = array.to_primitive(); - - 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) { self.values.reserve(additional); self.nulls.reserve_exact(additional); @@ -221,7 +209,7 @@ impl ArrayBuilder for PrimitiveBuilder { self.finish_into_primitive().into_array() } - fn finish_into_canonical(&mut self) -> Canonical { + fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { Canonical::Primitive(self.finish_into_primitive()) } } diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index 016d42508ee..da08e262760 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -12,9 +12,8 @@ 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; use crate::arrays::StructArray; use crate::arrays::struct_::StructArrayExt; use crate::builders::ArrayBuilder; @@ -22,8 +21,6 @@ use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::LazyBitBufferBuilder; use crate::builders::builder_with_capacity; use crate::canonical::Canonical; -#[expect(deprecated)] -use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::StructFields; @@ -122,6 +119,25 @@ impl StructBuilder { struct_fields } + + /// Appends the values of a canonical [`StructArray`] to the builder, recursing into each + /// field's builder. + pub(crate) fn append_struct_array( + &mut self, + array: &StructArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + for (field, builder) in array + .iter_unmasked_fields() + .zip_eq(self.builders.iter_mut()) + { + field.append_to_builder(builder.as_mut(), ctx)?; + } + + self.nulls + .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); + Ok(()) + } } impl ArrayBuilder for StructBuilder { @@ -168,26 +184,6 @@ impl ArrayBuilder for StructBuilder { self.append_value(scalar.as_struct()) } - unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { - #[expect(deprecated)] - let array = array.to_struct(); - - for (a, builder) in array - .iter_unmasked_fields() - .zip_eq(self.builders.iter_mut()) - { - builder.extend_from_array(a); - } - - self.nulls.append_validity_mask( - &array - .validity() - .vortex_expect("validity_mask") - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) - .vortex_expect("Failed to compute validity mask"), - ); - } - fn reserve_exact(&mut self, capacity: usize) { self.builders.iter_mut().for_each(|builder| { builder.reserve_exact(capacity); @@ -203,7 +199,7 @@ impl ArrayBuilder for StructBuilder { self.finish_into_struct().into_array() } - fn finish_into_canonical(&mut self) -> Canonical { + fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { Canonical::Struct(self.finish_into_struct()) } } diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index 02babab4a7d..138e6941def 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -8,7 +8,10 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; +use crate::Canonical; +use crate::ExecutionCtx; use crate::VortexSessionExecute; +use crate::array::IntoArray; use crate::array_session; use crate::builders::ArrayBuilder; use crate::builders::builder_with_capacity; @@ -223,12 +226,10 @@ fn test_append_defaults_behavior(#[case] dtype: DType, #[case] should_be_null: b /// Helper function that fills two builders with the same values and compares the results /// of `to_canonical()` vs `finish().to_canonical()`. -fn compare_to_canonical_methods(dtype: &DType, mut fill_builder: F) +fn compare_to_canonical_methods(dtype: &DType, ctx: &mut ExecutionCtx, mut fill_builder: F) where F: FnMut(&mut dyn ArrayBuilder), { - use crate::IntoArray; - // Create two identical builders. let mut builder1 = builder_with_capacity(dtype, 10); let mut builder2 = builder_with_capacity(dtype, 10); @@ -238,11 +239,10 @@ where fill_builder(builder2.as_mut()); // Get canonical arrays using both methods. - let canonical_direct = builder1.finish_into_canonical(); - #[expect(deprecated)] + let canonical_direct = builder1.finish_into_canonical(ctx); let canonical_indirect = builder2 .finish() - .to_canonical() + .execute::(ctx) .vortex_expect("to_canonical failed"); // Convert both to arrays for comparison. @@ -254,12 +254,8 @@ where // Compare each element. for i in 0..array_direct.len() { - let scalar_direct = array_direct - .execute_scalar(i, &mut array_session().create_execution_ctx()) - .unwrap(); - let scalar_indirect = array_indirect - .execute_scalar(i, &mut array_session().create_execution_ctx()) - .unwrap(); + let scalar_direct = array_direct.execute_scalar(i, ctx).unwrap(); + let scalar_indirect = array_indirect.execute_scalar(i, ctx).unwrap(); assert_eq!( scalar_direct, scalar_indirect, @@ -271,8 +267,9 @@ where #[test] fn test_to_canonical_bool() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Bool(Nullability::NonNullable); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { for i in 0..5 { let value = Scalar::bool(i % 2 == 0, Nullability::NonNullable); builder.append_scalar(&value).unwrap(); @@ -282,8 +279,9 @@ fn test_to_canonical_bool() { #[test] fn test_to_canonical_bool_nullable() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Bool(Nullability::Nullable); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { for i in 0..5 { let value = Scalar::bool(i % 2 == 0, Nullability::Nullable); builder.append_scalar(&value).unwrap(); @@ -294,8 +292,9 @@ fn test_to_canonical_bool_nullable() { #[test] fn test_to_canonical_i32() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { for i in 0..5 { let value = Scalar::primitive(i, Nullability::NonNullable); builder.append_scalar(&value).unwrap(); @@ -305,8 +304,9 @@ fn test_to_canonical_i32() { #[test] fn test_to_canonical_i32_nullable() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { for i in 0..5 { let value = Scalar::primitive(i, Nullability::Nullable); builder.append_scalar(&value).unwrap(); @@ -317,8 +317,9 @@ fn test_to_canonical_i32_nullable() { #[test] fn test_to_canonical_f64() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { for i in 0..5 { let value = Scalar::primitive(i as f64, Nullability::NonNullable); builder.append_scalar(&value).unwrap(); @@ -328,8 +329,9 @@ fn test_to_canonical_f64() { #[test] fn test_to_canonical_utf8() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Utf8(Nullability::NonNullable); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { let values = ["hello", "world", "test", "data", "vortex"]; for value in &values { let scalar = Scalar::utf8(*value, Nullability::NonNullable); @@ -340,8 +342,9 @@ fn test_to_canonical_utf8() { #[test] fn test_to_canonical_utf8_nullable() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Utf8(Nullability::Nullable); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { let values = ["hello", "world", "test"]; for value in &values { let scalar = Scalar::utf8(*value, Nullability::Nullable); @@ -353,8 +356,9 @@ fn test_to_canonical_utf8_nullable() { #[test] fn test_to_canonical_binary() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Binary(Nullability::NonNullable); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { let values = [b"hello", b"world", b"vortx", b"bytes", b"tests"]; for value in &values { let scalar = Scalar::binary(value.to_vec(), Nullability::NonNullable); @@ -365,6 +369,7 @@ fn test_to_canonical_binary() { #[test] fn test_to_canonical_struct() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Struct( StructFields::from_iter([ ("a", DType::Primitive(PType::I32, Nullability::NonNullable)), @@ -372,7 +377,7 @@ fn test_to_canonical_struct() { ]), Nullability::NonNullable, ); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { for _ in 0..3 { let value = Scalar::default_value(&dtype); builder.append_scalar(&value).unwrap(); @@ -382,9 +387,10 @@ fn test_to_canonical_struct() { #[test] fn test_to_canonical_extension() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Extension(Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased()); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { let ext_dtype = match &dtype { DType::Extension(ext) => ext.clone(), _ => unreachable!(), @@ -399,16 +405,18 @@ fn test_to_canonical_extension() { #[test] fn test_to_canonical_null() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Null; - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { builder.append_nulls(5); }); } #[test] fn test_to_canonical_decimal() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Decimal(DecimalDType::new(10, 2), Nullability::NonNullable); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { for _ in 0..5 { let value = Scalar::default_value(&dtype); builder.append_scalar(&value).unwrap(); @@ -418,8 +426,9 @@ fn test_to_canonical_decimal() { #[test] fn test_to_canonical_i8() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I8, Nullability::NonNullable); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { for i in 0..5i8 { let value = Scalar::primitive(i, Nullability::NonNullable); builder.append_scalar(&value).unwrap(); @@ -429,8 +438,9 @@ fn test_to_canonical_i8() { #[test] fn test_to_canonical_u64() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::U64, Nullability::NonNullable); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { for i in 0..5 { let value = Scalar::primitive(i as u64, Nullability::NonNullable); builder.append_scalar(&value).unwrap(); @@ -440,8 +450,9 @@ fn test_to_canonical_u64() { #[test] fn test_to_canonical_f32() { + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::NonNullable); - compare_to_canonical_methods(&dtype, |builder| { + compare_to_canonical_methods(&dtype, &mut ctx, |builder| { for i in 0..5 { let value = Scalar::primitive(i as f32, Nullability::NonNullable); builder.append_scalar(&value).unwrap(); diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 3aabfd87c31..1d530a384ef 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -22,8 +22,6 @@ 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; @@ -31,8 +29,6 @@ use crate::arrays::varbinview::compact::BufferUtilization; use crate::builders::ArrayBuilder; use crate::builders::LazyBitBufferBuilder; use crate::canonical::Canonical; -#[expect(deprecated)] -use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::scalar::Scalar; @@ -280,6 +276,7 @@ impl VarBinViewBuilder { .extend_from_compaction(BuffersWithOffsets::from_array( array, self.compaction_threshold, + ctx, )); match view_adjustment { @@ -372,13 +369,6 @@ impl ArrayBuilder for VarBinViewBuilder { Ok(()) } - unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { - #[expect(deprecated)] - let array = array.to_varbinview(); - 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) { self.views_builder.reserve(additional); self.nulls.reserve_exact(additional); @@ -392,7 +382,7 @@ impl ArrayBuilder for VarBinViewBuilder { self.finish_into_varbinview().into_array() } - fn finish_into_canonical(&mut self) -> Canonical { + fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { Canonical::VarBinView(self.finish_into_varbinview()) } } @@ -636,7 +626,11 @@ enum BuffersWithOffsets { } impl BuffersWithOffsets { - pub fn from_array(array: &VarBinViewArray, compaction_threshold: f64) -> Self { + pub fn from_array( + array: &VarBinViewArray, + compaction_threshold: f64, + ctx: &mut ExecutionCtx, + ) -> Self { if compaction_threshold == 0.0 { return Self::AllKept { buffers: Arc::from( @@ -652,7 +646,7 @@ impl BuffersWithOffsets { } let buffer_utilizations = array - .buffer_utilizations() + .buffer_utilizations(ctx) .vortex_expect("buffer_utilizations in BuffersWithOffsets::from_array"); let mut has_rewrite = false; let mut has_nonzero_offset = false; @@ -923,11 +917,11 @@ mod tests { let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); builder.append_value("Hello1"); - builder.extend_from_array(&array); + array.append_to_builder(&mut builder, &mut ctx).unwrap(); builder.append_nulls(2); builder.append_value("Hello3"); - let actual = builder.finish_into_canonical(); + let actual = builder.finish_into_canonical(&mut ctx); let expected = >::from_iter([ Some("Hello1"), None, diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 02029fc8e3e..71da52e1300 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -85,14 +85,14 @@ use crate::validity::Validity; /// /// The full list of canonical types and their equivalent Arrow array types are: /// -/// * `NullArray`: [`arrow_array::NullArray`] -/// * `BoolArray`: [`arrow_array::BooleanArray`] -/// * `PrimitiveArray`: [`arrow_array::PrimitiveArray`] -/// * `DecimalArray`: [`arrow_array::Decimal128Array`] and [`arrow_array::Decimal256Array`] -/// * `VarBinViewArray`: [`arrow_array::GenericByteViewArray`] -/// * `ListViewArray`: [`arrow_array::ListViewArray`] -/// * `FixedSizeListArray`: [`arrow_array::FixedSizeListArray`] -/// * `StructArray`: [`arrow_array::StructArray`] +/// * `NullArray`: `arrow_array::NullArray` +/// * `BoolArray`: `arrow_array::BooleanArray` +/// * `PrimitiveArray`: `arrow_array::PrimitiveArray` +/// * `DecimalArray`: `arrow_array::Decimal128Array` and `arrow_array::Decimal256Array` +/// * `VarBinViewArray`: `arrow_array::GenericByteViewArray` +/// * `ListViewArray`: `arrow_array::ListViewArray` +/// * `FixedSizeListArray`: `arrow_array::FixedSizeListArray` +/// * `StructArray`: `arrow_array::StructArray` /// /// Vortex uses a logical type system, unlike Arrow which uses physical encodings for its types. /// As an example, there are at least six valid physical encodings for a `Utf8` array. This can @@ -102,7 +102,7 @@ use crate::validity::Validity; /// variants to hold the data. /// /// To disambiguate, we choose a canonical physical encoding for every Vortex [`DType`], which -/// will correspond to an arrow-rs [`arrow_schema::DataType`]. +/// will correspond to an arrow-rs `arrow_schema::DataType`. /// /// # Views support /// @@ -111,7 +111,7 @@ use crate::validity::Validity; /// fully supported by the Datafusion query engine. We use them as our canonical string encoding /// for all `Utf8` and `Binary` typed arrays in Vortex. They provide considerably faster filter /// execution than the core `StringArray` and `BinaryArray` types, at the expense of potentially -/// needing [garbage collection][arrow_array::GenericByteViewArray::gc] to clear unreferenced items +/// needing garbage collection (`arrow_array::GenericByteViewArray::gc`) to clear unreferenced items /// from memory. /// /// # For Developers @@ -262,7 +262,7 @@ impl Canonical { /// and copy operations. pub fn compact(&self, ctx: &mut ExecutionCtx) -> VortexResult { match self { - Canonical::VarBinView(array) => Ok(Canonical::VarBinView(array.compact_buffers()?)), + Canonical::VarBinView(array) => Ok(Canonical::VarBinView(array.compact_buffers(ctx)?)), Canonical::List(array) => Ok(Canonical::List( array.rebuild(ListViewRebuildMode::TrimElements, ctx)?, )), @@ -1118,25 +1118,8 @@ impl Matcher for AnyCanonical { #[cfg(test)] mod test { - use std::sync::Arc; use std::sync::LazyLock; - use arrow_array::Array as ArrowArray; - use arrow_array::ArrayRef as ArrowArrayRef; - use arrow_array::ListArray as ArrowListArray; - use arrow_array::PrimitiveArray as ArrowPrimitiveArray; - use arrow_array::StringArray; - use arrow_array::StringViewArray; - use arrow_array::StructArray as ArrowStructArray; - use arrow_array::cast::AsArray; - use arrow_array::types::Int32Type; - use arrow_array::types::Int64Type; - use arrow_array::types::UInt64Type; - use arrow_buffer::NullBufferBuilder; - use arrow_buffer::OffsetBuffer; - use arrow_schema::DataType; - use arrow_schema::Field; - use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::VortexSession; @@ -1154,8 +1137,6 @@ mod test { use crate::arrays::VariantArray; use crate::arrays::struct_::StructArrayExt; use crate::arrays::variant::VariantArrayExt; - use crate::arrow::ArrowSessionExt; - use crate::arrow::FromArrowArray; use crate::canonical::StructArray; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -1207,134 +1188,4 @@ mod test { Ok(()) } - - #[test] - fn test_canonicalize_nested_struct() { - let mut ctx = SESSION.create_execution_ctx(); - // Create a struct array with multiple internal components. - let nested_struct_array = StructArray::from_fields(&[ - ("a", buffer![1u64].into_array()), - ( - "b", - StructArray::from_fields(&[( - "inner_a", - // The nested struct contains a ConstantArray representing the primitive array - // [100i64] - // ConstantArray is not a canonical type, so converting `into_arrow()` should - // map this to the nearest canonical type (PrimitiveArray). - ConstantArray::new(100i64, 1).into_array(), - )]) - .unwrap() - .into_array(), - ), - ]) - .unwrap(); - - let arrow_struct = SESSION - .arrow() - .execute_arrow(nested_struct_array.into_array(), None, &mut ctx) - .unwrap() - .as_any() - .downcast_ref::() - .cloned() - .unwrap(); - - assert!( - arrow_struct - .column(0) - .as_any() - .downcast_ref::>() - .is_some() - ); - - let inner_struct = Arc::clone(arrow_struct.column(1)) - .as_any() - .downcast_ref::() - .cloned() - .unwrap(); - - let inner_a = inner_struct - .column(0) - .as_any() - .downcast_ref::>(); - assert!(inner_a.is_some()); - - assert_eq!( - inner_a.cloned().unwrap(), - ArrowPrimitiveArray::from_iter([100i64]) - ); - } - - #[test] - fn roundtrip_struct() { - let mut ctx = SESSION.create_execution_ctx(); - let mut nulls = NullBufferBuilder::new(6); - nulls.append_n_non_nulls(4); - nulls.append_null(); - nulls.append_non_null(); - let names = Arc::new(StringViewArray::from_iter(vec![ - Some("Joseph"), - None, - Some("Angela"), - Some("Mikhail"), - None, - None, - ])); - let ages = Arc::new(ArrowPrimitiveArray::::from(vec![ - Some(25), - Some(31), - None, - Some(57), - None, - None, - ])); - - let arrow_struct = ArrowStructArray::new( - vec![ - Arc::new(Field::new("name", DataType::Utf8View, true)), - Arc::new(Field::new("age", DataType::Int32, true)), - ] - .into(), - vec![names, ages], - nulls.finish(), - ); - - let vortex_struct = ArrayRef::from_arrow(&arrow_struct, true).unwrap(); - let vortex_struct = SESSION - .arrow() - .execute_arrow(vortex_struct, None, &mut ctx) - .unwrap(); - assert_eq!(&arrow_struct, vortex_struct.as_struct()); - } - - #[test] - fn roundtrip_list() { - let mut ctx = SESSION.create_execution_ctx(); - let names = Arc::new(StringArray::from_iter(vec![ - Some("Joseph"), - Some("Angela"), - Some("Mikhail"), - ])); - - let arrow_list = ArrowListArray::new( - Arc::new(Field::new_list_field(DataType::Utf8, true)), - OffsetBuffer::from_lengths(vec![0, 2, 1]), - names, - None, - ); - let list_data_type = arrow_list.data_type(); - let list_field = Field::new(String::new(), list_data_type.clone(), true); - - let vortex_list = ArrayRef::from_arrow(&arrow_list, true).unwrap(); - - let rt_arrow_list = SESSION - .arrow() - .execute_arrow(vortex_list, Some(&list_field), &mut ctx) - .unwrap(); - - assert_eq!( - (Arc::new(arrow_list.clone()) as ArrowArrayRef).as_ref(), - rt_arrow_list.as_ref() - ); - } } diff --git a/vortex-array/src/compute/conformance/cast.rs b/vortex-array/src/compute/conformance/cast.rs index c823b2a778a..abe23e6e2dc 100644 --- a/vortex-array/src/compute/conformance/cast.rs +++ b/vortex-array/src/compute/conformance/cast.rs @@ -6,13 +6,12 @@ use vortex_error::VortexResult; use vortex_error::vortex_panic; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; 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; @@ -20,10 +19,14 @@ use crate::dtype::PType; use crate::scalar::Scalar; /// Cast and force execution via `to_canonical`, returning the canonical array. -fn cast_and_execute(array: &ArrayRef, dtype: DType) -> VortexResult { +fn cast_and_execute( + array: &ArrayRef, + dtype: DType, + ctx: &mut ExecutionCtx, +) -> VortexResult { Ok(array .cast(dtype)? - .execute::(&mut array_session().create_execution_ctx())? + .execute::(ctx)? .0 .into_array()) } @@ -37,18 +40,18 @@ fn cast_and_execute(array: &ArrayRef, dtype: DType) -> VortexResult { /// - Casting with nullability changes /// - Casting between string types (Utf8/Binary) /// - Edge cases like overflow behavior -pub fn test_cast_conformance(array: &ArrayRef) { +pub fn test_cast_conformance(array: &ArrayRef, ctx: &mut ExecutionCtx) { let dtype = array.dtype(); // Always test identity cast and nullability changes - test_cast_identity(array); + test_cast_identity(array, ctx); - test_cast_to_non_nullable(array); - test_cast_to_nullable(array); + test_cast_to_non_nullable(array, ctx); + test_cast_to_nullable(array, ctx); // Test based on the specific DType match dtype { - DType::Null => test_cast_from_null(array), + DType::Null => test_cast_from_null(array, ctx), DType::Primitive(ptype, ..) => match ptype { PType::U8 | PType::U16 @@ -57,16 +60,16 @@ pub fn test_cast_conformance(array: &ArrayRef) { | PType::I8 | PType::I16 | PType::I32 - | PType::I64 => test_cast_to_integral_types(array), - PType::F16 | PType::F32 | PType::F64 => test_cast_from_floating_point_types(array), + | PType::I64 => test_cast_to_integral_types(array, ctx), + PType::F16 | PType::F32 | PType::F64 => test_cast_from_floating_point_types(array, ctx), }, _ => {} } } -fn test_cast_identity(array: &ArrayRef) { +fn test_cast_identity(array: &ArrayRef, ctx: &mut ExecutionCtx) { // Casting to the same type should be a no-op - let result = cast_and_execute(&array.clone(), array.dtype().clone()) + let result = cast_and_execute(&array.clone(), array.dtype().clone(), ctx) .vortex_expect("cast should succeed in conformance test"); assert_eq!(result.len(), array.len()); assert_eq!(result.dtype(), array.dtype()); @@ -75,18 +78,18 @@ fn test_cast_identity(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } } -fn test_cast_from_null(array: &ArrayRef) { +fn test_cast_from_null(array: &ArrayRef, ctx: &mut ExecutionCtx) { // Null can be cast to itself - let result = cast_and_execute(&array.clone(), DType::Null) + let result = cast_and_execute(&array.clone(), DType::Null, ctx) .vortex_expect("cast should succeed in conformance test"); assert_eq!(result.len(), array.len()); assert_eq!(result.dtype(), &DType::Null); @@ -101,7 +104,7 @@ fn test_cast_from_null(array: &ArrayRef) { ]; for dtype in nullable_types { - let result = cast_and_execute(&array.clone(), dtype.clone()) + let result = cast_and_execute(&array.clone(), dtype.clone(), ctx) .vortex_expect("cast should succeed in conformance test"); assert_eq!(result.len(), array.len()); assert_eq!(result.dtype(), &dtype); @@ -110,7 +113,7 @@ fn test_cast_from_null(array: &ArrayRef) { for i in 0..array.len().min(10) { assert!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") .is_null() ); @@ -124,17 +127,17 @@ fn test_cast_from_null(array: &ArrayRef) { ]; for dtype in non_nullable_types { - assert!(cast_and_execute(&array.clone(), dtype.clone()).is_err()); + assert!(cast_and_execute(&array.clone(), dtype.clone(), ctx).is_err()); } } -fn test_cast_to_non_nullable(array: &ArrayRef) { +fn test_cast_to_non_nullable(array: &ArrayRef, ctx: &mut ExecutionCtx) { if array - .invalid_count(&mut array_session().create_execution_ctx()) + .invalid_count(ctx) .vortex_expect("invalid_count should succeed in conformance test") == 0 { - let non_nullable = cast_and_execute(&array.clone(), array.dtype().as_nonnullable()) + let non_nullable = cast_and_execute(&array.clone(), array.dtype().as_nonnullable(), ctx) .vortex_expect("arrays without nulls can cast to non-nullable"); assert_eq!(non_nullable.dtype(), &array.dtype().as_nonnullable()); assert_eq!(non_nullable.len(), array.len()); @@ -142,15 +145,15 @@ fn test_cast_to_non_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), non_nullable - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } - let back_to_nullable = cast_and_execute(&non_nullable, array.dtype().clone()) + let back_to_nullable = cast_and_execute(&non_nullable, array.dtype().clone(), ctx) .vortex_expect("non-nullable arrays can cast to nullable"); assert_eq!(back_to_nullable.dtype(), array.dtype()); assert_eq!(back_to_nullable.len(), array.len()); @@ -158,10 +161,10 @@ fn test_cast_to_non_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), back_to_nullable - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -171,7 +174,7 @@ fn test_cast_to_non_nullable(array: &ArrayRef) { // array can be casted to DType::Null. return; } - cast_and_execute(&array.clone(), array.dtype().as_nonnullable()) + cast_and_execute(&array.clone(), array.dtype().as_nonnullable(), ctx) .err() .unwrap_or_else(|| { vortex_panic!( @@ -182,8 +185,8 @@ fn test_cast_to_non_nullable(array: &ArrayRef) { } } -fn test_cast_to_nullable(array: &ArrayRef) { - let nullable = cast_and_execute(&array.clone(), array.dtype().as_nullable()) +fn test_cast_to_nullable(array: &ArrayRef, ctx: &mut ExecutionCtx) { + let nullable = cast_and_execute(&array.clone(), array.dtype().as_nullable(), ctx) .vortex_expect("arrays without nulls can cast to nullable"); assert_eq!(nullable.dtype(), &array.dtype().as_nullable()); assert_eq!(nullable.len(), array.len()); @@ -191,15 +194,15 @@ fn test_cast_to_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), nullable - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } - let back = cast_and_execute(&nullable, array.dtype().clone()) + let back = cast_and_execute(&nullable, array.dtype().clone(), ctx) .vortex_expect("casting to nullable and back should be a no-op"); assert_eq!(back.dtype(), array.dtype()); assert_eq!(back.len(), array.len()); @@ -207,38 +210,43 @@ fn test_cast_to_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), - back.execute_scalar(i, &mut array_session().create_execution_ctx()) + back.execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } } -fn test_cast_from_floating_point_types(array: &ArrayRef) { - let ptype = array.as_primitive_typed().ptype(); - test_cast_to_primitive(array, PType::I8, false); - test_cast_to_primitive(array, PType::U8, false); - test_cast_to_primitive(array, PType::I16, false); - test_cast_to_primitive(array, PType::U16, false); - test_cast_to_primitive(array, PType::I32, false); - test_cast_to_primitive(array, PType::U32, false); - test_cast_to_primitive(array, PType::I64, false); - test_cast_to_primitive(array, PType::U64, false); - test_cast_to_primitive(array, PType::F16, matches!(ptype, PType::F16)); - test_cast_to_primitive(array, PType::F32, matches!(ptype, PType::F16 | PType::F32)); - test_cast_to_primitive(array, PType::F64, true); +fn test_cast_from_floating_point_types(array: &ArrayRef, ctx: &mut ExecutionCtx) { + let ptype = array.dtype().as_ptype(); + test_cast_to_primitive(array, PType::I8, false, ctx); + test_cast_to_primitive(array, PType::U8, false, ctx); + test_cast_to_primitive(array, PType::I16, false, ctx); + test_cast_to_primitive(array, PType::U16, false, ctx); + test_cast_to_primitive(array, PType::I32, false, ctx); + test_cast_to_primitive(array, PType::U32, false, ctx); + test_cast_to_primitive(array, PType::I64, false, ctx); + test_cast_to_primitive(array, PType::U64, false, ctx); + test_cast_to_primitive(array, PType::F16, matches!(ptype, PType::F16), ctx); + test_cast_to_primitive( + array, + PType::F32, + matches!(ptype, PType::F16 | PType::F32), + ctx, + ); + test_cast_to_primitive(array, PType::F64, true, ctx); } -fn test_cast_to_integral_types(array: &ArrayRef) { - test_cast_to_primitive(array, PType::I8, true); - test_cast_to_primitive(array, PType::U8, true); - test_cast_to_primitive(array, PType::I16, true); - test_cast_to_primitive(array, PType::U16, true); - test_cast_to_primitive(array, PType::I32, true); - test_cast_to_primitive(array, PType::U32, true); - test_cast_to_primitive(array, PType::I64, true); - test_cast_to_primitive(array, PType::U64, true); +fn test_cast_to_integral_types(array: &ArrayRef, ctx: &mut ExecutionCtx) { + test_cast_to_primitive(array, PType::I8, true, ctx); + test_cast_to_primitive(array, PType::U8, true, ctx); + test_cast_to_primitive(array, PType::I16, true, ctx); + test_cast_to_primitive(array, PType::U16, true, ctx); + test_cast_to_primitive(array, PType::I32, true, ctx); + test_cast_to_primitive(array, PType::U32, true, ctx); + test_cast_to_primitive(array, PType::I64, true, ctx); + test_cast_to_primitive(array, PType::U64, true, ctx); } /// Does this scalar fit in this type? @@ -247,9 +255,13 @@ fn fits(value: &Scalar, ptype: PType) -> bool { value.cast(&dtype).is_ok() } -fn test_cast_to_primitive(array: &ArrayRef, target_ptype: PType, test_round_trip: bool) { - let mut ctx = array_session().create_execution_ctx(); - let maybe_min_max = min_max(array, &mut ctx, NumericalAggregateOpts::default()) +fn test_cast_to_primitive( + array: &ArrayRef, + target_ptype: PType, + test_round_trip: bool, + ctx: &mut ExecutionCtx, +) { + let maybe_min_max = min_max(array, ctx, NumericalAggregateOpts::default()) .vortex_expect("cast should succeed in conformance test"); if let Some(MinMaxResult { min, max }) = maybe_min_max @@ -258,6 +270,7 @@ fn test_cast_to_primitive(array: &ArrayRef, target_ptype: PType, test_round_trip cast_and_execute( &array.clone(), DType::Primitive(target_ptype, array.dtype().nullability()), + ctx, ) .err() .unwrap_or_else(|| { @@ -277,6 +290,7 @@ fn test_cast_to_primitive(array: &ArrayRef, target_ptype: PType, test_round_trip let casted = cast_and_execute( &array.clone(), DType::Primitive(target_ptype, array.dtype().nullability()), + ctx, ) .unwrap_or_else(|e| { vortex_panic!( @@ -289,20 +303,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 array_session().create_execution_ctx()) + .execute_mask(array.len(), ctx) .vortex_expect("Failed to compute validity mask"), casted .validity() .vortex_expect("validity_mask should succeed in conformance test") - .execute_mask(casted.len(), &mut array_session().create_execution_ctx()) + .execute_mask(casted.len(), ctx) .vortex_expect("Failed to compute validity mask") ); for i in 0..array.len().min(10) { let original = array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let casted = casted - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( original @@ -325,10 +339,15 @@ fn test_cast_to_primitive(array: &ArrayRef, target_ptype: PType, test_round_trip #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_buffer::buffer; + use vortex_session::VortexSession; use super::*; use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::NullArray; @@ -339,40 +358,42 @@ mod tests { use crate::dtype::FieldNames; use crate::dtype::Nullability; + static SESSION: LazyLock = LazyLock::new(array_session); + #[test] fn test_cast_conformance_u32() { let array = buffer![0u32, 100, 200, 65535, 1000000].into_array(); - test_cast_conformance(&array); + test_cast_conformance(&array, &mut SESSION.create_execution_ctx()); } #[test] fn test_cast_conformance_i32() { let array = buffer![-100i32, -1, 0, 1, 100].into_array(); - test_cast_conformance(&array); + test_cast_conformance(&array, &mut SESSION.create_execution_ctx()); } #[test] fn test_cast_conformance_f32() { let array = buffer![0.0f32, 1.5, -2.5, 100.0, 1e6].into_array(); - test_cast_conformance(&array); + test_cast_conformance(&array, &mut SESSION.create_execution_ctx()); } #[test] fn test_cast_conformance_nullable() { let array = PrimitiveArray::from_option_iter([Some(1u8), None, Some(255), Some(0), None]); - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] fn test_cast_conformance_bool() { let array = BoolArray::from_iter(vec![true, false, true, false]); - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] fn test_cast_conformance_null() { let array = NullArray::new(5); - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] @@ -381,7 +402,7 @@ mod tests { vec![Some("hello"), None, Some("world")], DType::Utf8(Nullability::Nullable), ); - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] @@ -390,7 +411,7 @@ mod tests { vec![Some(b"data".as_slice()), None, Some(b"bytes".as_slice())], DType::Binary(Nullability::Nullable), ); - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] @@ -407,7 +428,7 @@ mod tests { let array = StructArray::try_new(names, vec![a, b], 3, crate::validity::Validity::NonNullable) .unwrap(); - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[test] @@ -417,6 +438,6 @@ mod tests { let array = ListArray::try_new(data, offsets, crate::validity::Validity::NonNullable).unwrap(); - test_cast_conformance(&array.into_array()); + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); } } diff --git a/vortex-array/src/compute/conformance/filter.rs b/vortex-array/src/compute/conformance/filter.rs index 01bfb3bb2fa..5c78e0d70ad 100644 --- a/vortex-array/src/compute/conformance/filter.rs +++ b/vortex-array/src/compute/conformance/filter.rs @@ -5,9 +5,9 @@ use vortex_error::VortexExpect; use vortex_mask::Mask; use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray; -use crate::VortexSessionExecute; -use crate::array_session; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -18,16 +18,16 @@ pub const LARGE_SIZE: usize = 1024; /// Test filter compute function with various array sizes and patterns. /// The input array can be of any length. -pub fn test_filter_conformance(array: &ArrayRef) { +pub fn test_filter_conformance(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); // Test with arrays of any size if len > 0 { - test_all_filter(array); + test_all_filter(array, ctx); test_none_filter(array); - test_selective_filter(array); - test_single_element_filter(array); - test_alternating_pattern_filter(array); + test_selective_filter(array, ctx); + test_single_element_filter(array, ctx); + test_alternating_pattern_filter(array, ctx); test_runs_pattern_filter(array); test_sparse_true_filter(array); test_sparse_false_filter(array); @@ -61,14 +61,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(); +fn test_all_filter(array: &ArrayRef, ctx: &mut ExecutionCtx) { 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, &mut ctx); + assert_arrays_eq!(filtered, array, ctx); } /// Tests that filtering with an all-false mask returns an empty array with the same dtype @@ -82,7 +81,7 @@ fn test_none_filter(array: &ArrayRef) { assert_eq!(filtered.dtype(), array.dtype()); } -fn test_selective_filter(array: &ArrayRef) { +fn test_selective_filter(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len < 2 { return; // Skip for very small arrays @@ -101,10 +100,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 array_session().create_execution_ctx()) + .execute_scalar(filtered_idx, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -121,24 +120,24 @@ fn test_selective_filter(array: &ArrayRef) { assert_eq!(filtered.len(), 2); assert_eq!( filtered - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); assert_eq!( filtered - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar(1, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(len - 1, &mut array_session().create_execution_ctx()) + .execute_scalar(len - 1, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } } -fn test_single_element_filter(array: &ArrayRef) { +fn test_single_element_filter(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len == 0 { return; @@ -154,10 +153,10 @@ fn test_single_element_filter(array: &ArrayRef) { assert_eq!(filtered.len(), 1); assert_eq!( filtered - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); @@ -172,18 +171,16 @@ fn test_single_element_filter(array: &ArrayRef) { assert_eq!(filtered.len(), 1); assert_eq!( filtered - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(len - 1, &mut array_session().create_execution_ctx()) + .execute_scalar(len - 1, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } } fn test_empty_array_filter(dtype: &DType) { - use crate::Canonical; - let empty_array = Canonical::empty(dtype).into_array(); let empty_mask = Mask::new_false(0); let filtered = empty_array @@ -221,7 +218,7 @@ fn test_mismatched_lengths(array: &ArrayRef) { } /// Tests filtering with alternating true/false pattern -fn test_alternating_pattern_filter(array: &ArrayRef) { +fn test_alternating_pattern_filter(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); let pattern = create_alternating_pattern(len); let expected_count = pattern.iter().filter(|&&v| v).count(); @@ -238,10 +235,10 @@ fn test_alternating_pattern_filter(array: &ArrayRef) { if keep { assert_eq!( filtered - .execute_scalar(filtered_idx, &mut array_session().create_execution_ctx()) + .execute_scalar(filtered_idx, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, 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 7a002662d38..50930ba2876 100644 --- a/vortex-array/src/compute/conformance/mask.rs +++ b/vortex-array/src/compute/conformance/mask.rs @@ -5,38 +5,38 @@ use vortex_error::VortexExpect; use vortex_mask::Mask; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; -use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::builtins::ArrayBuiltins; +use crate::validity::Validity; /// Test mask compute function with various array sizes and patterns. /// The mask operation sets elements to null where the mask is true. -pub fn test_mask_conformance(array: &ArrayRef) { +pub fn test_mask_conformance(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len > 0 { - test_heterogenous_mask(array); - test_empty_mask(array); - test_full_mask(array); - test_alternating_mask(array); - test_sparse_mask(array); - test_single_element_mask(array); + test_heterogenous_mask(array, ctx); + test_empty_mask(array, ctx); + test_full_mask(array, ctx); + test_alternating_mask(array, ctx); + test_sparse_mask(array, ctx); + test_single_element_mask(array, ctx); } if len >= 5 { - test_double_mask(array); + test_double_mask(array, ctx); } if len > 0 { - test_nullable_mask_input(array); + test_nullable_mask_input(array, ctx); } } /// Tests masking with a heterogeneous pattern -fn test_heterogenous_mask(array: &ArrayRef) { +fn test_heterogenous_mask(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); // Create a pattern where roughly half the values are masked @@ -54,16 +54,16 @@ fn test_heterogenous_mask(array: &ArrayRef) { if masked_out { assert!( !masked - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid(i, ctx) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( masked - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -72,7 +72,7 @@ fn test_heterogenous_mask(array: &ArrayRef) { } /// Tests that an empty mask (all false) preserves all elements -fn test_empty_mask(array: &ArrayRef) { +fn test_empty_mask(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); let all_unmasked = vec![false; len]; let mask_array = Mask::from_iter(all_unmasked); @@ -87,10 +87,10 @@ fn test_empty_mask(array: &ArrayRef) { for i in 0..len { assert_eq!( masked - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -98,7 +98,7 @@ fn test_empty_mask(array: &ArrayRef) { } /// Tests that a full mask (all true) makes all elements null -fn test_full_mask(array: &ArrayRef) { +fn test_full_mask(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); let all_masked = vec![true; len]; let mask_array = Mask::from_iter(all_masked); @@ -113,14 +113,14 @@ fn test_full_mask(array: &ArrayRef) { for i in 0..len { assert!( !masked - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid(i, ctx) .vortex_expect("is_valid should succeed in conformance test") ); } } /// Tests alternating mask pattern -fn test_alternating_mask(array: &ArrayRef) { +fn test_alternating_mask(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); let pattern: Vec = (0..len).map(|i| i % 2 == 0).collect(); let mask_array = Mask::from_iter(pattern); @@ -135,16 +135,16 @@ fn test_alternating_mask(array: &ArrayRef) { if i % 2 == 0 { assert!( !masked - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid(i, ctx) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( masked - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -153,7 +153,7 @@ fn test_alternating_mask(array: &ArrayRef) { } /// Tests sparse mask (only a few elements masked) -fn test_sparse_mask(array: &ArrayRef) { +fn test_sparse_mask(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len < 10 { return; // Skip for small arrays @@ -173,7 +173,7 @@ fn test_sparse_mask(array: &ArrayRef) { let valid_count = (0..len) .filter(|&i| { masked - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid(i, 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 array_session().create_execution_ctx()) + .is_valid(i, ctx) .vortex_expect("is_valid should succeed in conformance test") }) .count(); @@ -194,7 +194,7 @@ fn test_sparse_mask(array: &ArrayRef) { } /// Tests masking a single element -fn test_single_element_mask(array: &ArrayRef) { +fn test_single_element_mask(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); // Mask only the first element @@ -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 array_session().create_execution_ctx()) + .is_valid(0, ctx) .vortex_expect("is_valid should succeed in conformance test") ); for i in 1..len { assert_eq!( masked - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -226,7 +226,7 @@ fn test_single_element_mask(array: &ArrayRef) { } /// Tests double masking operations -fn test_double_mask(array: &ArrayRef) { +fn test_double_mask(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); // Create two different mask patterns @@ -249,16 +249,16 @@ fn test_double_mask(array: &ArrayRef) { if mask1_pattern[i] || mask2_pattern[i] { assert!( !double_masked - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid(i, ctx) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( double_masked - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -267,7 +267,7 @@ fn test_double_mask(array: &ArrayRef) { } /// Tests masking with nullable mask (nulls treated as false) -fn test_nullable_mask_input(array: &ArrayRef) { +fn test_nullable_mask_input(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len < 3 { return; // Skip for very small arrays @@ -278,11 +278,10 @@ fn test_nullable_mask_input(array: &ArrayRef) { let validity_values: Vec = (0..len).map(|i| i % 3 != 0).collect(); let bool_array = BoolArray::from_iter(bool_values.clone()); - let validity = crate::validity::Validity::from_iter(validity_values.clone()); + let validity = Validity::from_iter(validity_values.clone()); let nullable_mask = BoolArray::new(bool_array.to_bit_buffer(), validity); - let mask_array = - nullable_mask.to_mask_fill_null_false(&mut array_session().create_execution_ctx()); + let mask_array = nullable_mask.to_mask_fill_null_false(ctx); let masked = array .clone() .mask((!&mask_array).into_array()) @@ -293,16 +292,16 @@ fn test_nullable_mask_input(array: &ArrayRef) { if bool_values[i] && validity_values[i] { assert!( !masked - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid(i, ctx) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( masked - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); diff --git a/vortex-array/src/compute/conformance/take.rs b/vortex-array/src/compute/conformance/take.rs index babaf392eb9..47bc8f1bd82 100644 --- a/vortex-array/src/compute/conformance/take.rs +++ b/vortex-array/src/compute/conformance/take.rs @@ -6,9 +6,8 @@ use vortex_error::VortexExpect; use crate::ArrayRef; use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray as _; -use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::Nullability; @@ -21,39 +20,39 @@ use crate::dtype::Nullability; /// - Taking with out-of-bounds indices (should panic) /// - Taking with nullable indices /// - Edge cases like empty arrays -pub fn test_take_conformance(array: &ArrayRef) { +pub fn test_take_conformance(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); if len > 0 { - test_take_all(array); + test_take_all(array, ctx); test_take_none(array); - test_take_selective(array); - test_take_first_and_last(array); - test_take_with_nullable_indices(array); - test_take_repeated_indices(array); + test_take_selective(array, ctx); + test_take_first_and_last(array, ctx); + test_take_with_nullable_indices(array, ctx); + test_take_repeated_indices(array, ctx); } test_empty_indices(array); // Additional edge cases for non-empty arrays if len > 0 { - test_take_reverse(array); - test_take_single_middle(array); + test_take_reverse(array, ctx); + test_take_single_middle(array, ctx); } if len > 3 { - test_take_random_unsorted(array); - test_take_contiguous_range(array); - test_take_mixed_repeated(array); + test_take_random_unsorted(array, ctx); + test_take_contiguous_range(array, ctx); + test_take_mixed_repeated(array, ctx); } // Test for larger arrays if len >= 1024 { - test_take_large_indices(array); + test_take_large_indices(array, ctx); } } -fn test_take_all(array: &ArrayRef) { +fn test_take_all(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); let indices = PrimitiveArray::from_iter(0..len as u64); let result = array @@ -64,13 +63,13 @@ fn test_take_all(array: &ArrayRef) { assert_eq!(result.dtype(), array.dtype()); // Verify elements match - #[expect(deprecated)] let array_canonical = array - .to_canonical() + .clone() + .execute::(ctx) .vortex_expect("to_canonical failed on array"); - #[expect(deprecated)] let result_canonical = result - .to_canonical() + .clone() + .execute::(ctx) .vortex_expect("to_canonical failed on result"); match (array_canonical, result_canonical) { (Canonical::Primitive(orig_prim), Canonical::Primitive(result_prim)) => { @@ -84,10 +83,10 @@ fn test_take_all(array: &ArrayRef) { for i in 0..len { assert_eq!( array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -106,7 +105,7 @@ fn test_take_none(array: &ArrayRef) { } #[expect(clippy::cast_possible_truncation)] -fn test_take_selective(array: &ArrayRef) { +fn test_take_selective(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); // Take every other element @@ -123,19 +122,16 @@ fn test_take_selective(array: &ArrayRef) { for (result_idx, &original_idx) in indices.iter().enumerate() { assert_eq!( array - .execute_scalar( - original_idx as usize, - &mut array_session().create_execution_ctx() - ) + .execute_scalar(original_idx as usize, ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(result_idx, &mut array_session().create_execution_ctx()) + .execute_scalar(result_idx, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } } -fn test_take_first_and_last(array: &ArrayRef) { +fn test_take_first_and_last(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); let indices = PrimitiveArray::from_iter([0u64, (len - 1) as u64]); let result = array @@ -145,24 +141,24 @@ fn test_take_first_and_last(array: &ArrayRef) { assert_eq!(result.len(), 2); assert_eq!( array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); assert_eq!( array - .execute_scalar(len - 1, &mut array_session().create_execution_ctx()) + .execute_scalar(len - 1, ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar(1, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } #[expect(clippy::cast_possible_truncation)] -fn test_take_with_nullable_indices(array: &ArrayRef) { +fn test_take_with_nullable_indices(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); // Create indices with some null values @@ -190,17 +186,17 @@ fn test_take_with_nullable_indices(array: &ArrayRef) { match idx_opt { Some(idx) => { let expected = array - .execute_scalar(*idx as usize, &mut array_session().create_execution_ctx()) + .execute_scalar(*idx as usize, ctx) .vortex_expect("scalar_at should succeed in conformance test"); let actual = result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!(expected, actual); } None => { assert!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") .is_null() ); @@ -209,7 +205,7 @@ fn test_take_with_nullable_indices(array: &ArrayRef) { } } -fn test_take_repeated_indices(array: &ArrayRef) { +fn test_take_repeated_indices(array: &ArrayRef, ctx: &mut ExecutionCtx) { if array.is_empty() { return; } @@ -222,12 +218,12 @@ fn test_take_repeated_indices(array: &ArrayRef) { assert_eq!(result.len(), 3); let first_elem = array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test"); for i in 0..3 { assert_eq!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), first_elem ); @@ -244,7 +240,7 @@ fn test_empty_indices(array: &ArrayRef) { assert_eq!(result.dtype(), array.dtype()); } -fn test_take_reverse(array: &ArrayRef) { +fn test_take_reverse(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); // Take elements in reverse order let indices = PrimitiveArray::from_iter((0..len as u64).rev()); @@ -258,16 +254,16 @@ fn test_take_reverse(array: &ArrayRef) { for i in 0..len { assert_eq!( array - .execute_scalar(len - 1 - i, &mut array_session().create_execution_ctx()) + .execute_scalar(len - 1 - i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } } -fn test_take_single_middle(array: &ArrayRef) { +fn test_take_single_middle(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); let middle_idx = len / 2; @@ -279,20 +275,20 @@ fn test_take_single_middle(array: &ArrayRef) { assert_eq!(result.len(), 1); assert_eq!( array - .execute_scalar(middle_idx, &mut array_session().create_execution_ctx()) + .execute_scalar(middle_idx, ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar(0, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } #[expect(clippy::cast_possible_truncation)] -fn test_take_random_unsorted(array: &ArrayRef) { +fn test_take_random_unsorted(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); // Create a pseudo-random but deterministic pattern - let mut indices = Vec::new(); + let mut indices = Vec::with_capacity(len.min(10)); let mut idx = 1u64; for _ in 0..len.min(10) { indices.push((idx * 7 + 3) % len as u64); @@ -310,16 +306,16 @@ fn test_take_random_unsorted(array: &ArrayRef) { for (i, &idx) in indices.iter().enumerate() { assert_eq!( array - .execute_scalar(idx as usize, &mut array_session().create_execution_ctx()) + .execute_scalar(idx as usize, ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } } -fn test_take_contiguous_range(array: &ArrayRef) { +fn test_take_contiguous_range(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); let start = len / 4; let end = len / 2; @@ -336,17 +332,17 @@ fn test_take_contiguous_range(array: &ArrayRef) { for i in 0..(end - start) { assert_eq!( array - .execute_scalar(start + i, &mut array_session().create_execution_ctx()) + .execute_scalar(start + i, ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } } #[expect(clippy::cast_possible_truncation)] -fn test_take_mixed_repeated(array: &ArrayRef) { +fn test_take_mixed_repeated(array: &ArrayRef, ctx: &mut ExecutionCtx) { let len = array.len(); // Create pattern with some repeated indices @@ -372,17 +368,17 @@ fn test_take_mixed_repeated(array: &ArrayRef) { for (i, &idx) in indices.iter().enumerate() { assert_eq!( array - .execute_scalar(idx as usize, &mut array_session().create_execution_ctx()) + .execute_scalar(idx as usize, ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } } #[expect(clippy::cast_possible_truncation)] -fn test_take_large_indices(array: &ArrayRef) { +fn test_take_large_indices(array: &ArrayRef, ctx: &mut ExecutionCtx) { // Test with a large number of indices to stress test performance let len = array.len(); let num_indices = 10000.min(len * 3); @@ -404,10 +400,10 @@ fn test_take_large_indices(array: &ArrayRef) { let expected_idx = indices[i] as usize; assert_eq!( array - .execute_scalar(expected_idx, &mut array_session().create_execution_ctx()) + .execute_scalar(expected_idx, ctx) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar(i, ctx) .vortex_expect("scalar_at should succeed in conformance test") ); } diff --git a/vortex-array/src/display/mod.rs b/vortex-array/src/display/mod.rs index 0134a28558c..1c12afcf8d2 100644 --- a/vortex-array/src/display/mod.rs +++ b/vortex-array/src/display/mod.rs @@ -6,6 +6,7 @@ mod extractors; mod tree_display; use std::fmt::Display; +use std::iter::repeat_n; pub use extractor::IndentedFormatter; pub use extractor::TreeContext; @@ -19,8 +20,8 @@ use itertools::Itertools as _; pub use tree_display::TreeDisplay; use crate::ArrayRef; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::legacy_session; /// Describe how to convert an array to a string. /// @@ -521,6 +522,7 @@ impl ArrayRef { DisplayArrayAs(self, DisplayOptions::TableDisplay) } + #[allow(clippy::disallowed_methods)] fn fmt_as(&self, f: &mut std::fmt::Formatter, options: &DisplayOptions) -> std::fmt::Result { match options { DisplayOptions::MetadataOnly => EncodingSummaryExtractor::write(self, f), @@ -536,7 +538,7 @@ impl ArrayRef { let is_truncated = self.len() > limit; let fmt_scalar = |i| { - self.execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx()) + self.execute_scalar(i, &mut legacy_session().create_execution_ctx()) .map_or_else(|e| format!(""), |s| s.to_string()) }; write!( @@ -544,10 +546,7 @@ impl ArrayRef { "{opening_brace}{}{closing_brace}", (0..limit.saturating_sub(3)) .map(fmt_scalar) - .chain(std::iter::repeat_n( - "...".to_string(), - is_truncated as usize - )) + .chain(repeat_n("...".to_string(), is_truncated as usize)) .chain((self.len().saturating_sub(3)..self.len()).map(fmt_scalar)) .format(sep) ) @@ -587,7 +586,7 @@ impl ArrayRef { let mut builder = tabled::builder::Builder::default(); // Reuse a single execution context across all per-row accesses below. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = legacy_session().create_execution_ctx(); // Special logic for struct arrays. let DType::Struct(sf, _) = self.dtype() else { diff --git a/vortex-array/src/dtype/coercion.rs b/vortex-array/src/dtype/coercion.rs index d57c0b9c1f4..4d30a80c216 100644 --- a/vortex-array/src/dtype/coercion.rs +++ b/vortex-array/src/dtype/coercion.rs @@ -79,6 +79,10 @@ impl DType { /// The core primitive — what type can hold both `self` and `other`? /// Returns `None` if no common supertype exists. pub fn least_supertype(&self, other: &DType) -> Option { + if let (DType::Union(lhs, lhs_null), DType::Union(rhs, rhs_null)) = (self, other) { + return (lhs == rhs).then(|| DType::Union(lhs.clone(), *lhs_null | *rhs_null)); + } + let union_null = self.nullability() | other.nullability(); if let ( @@ -308,6 +312,7 @@ mod tests { use crate::dtype::DType; use crate::dtype::PType; + use crate::dtype::UnionVariants; use crate::dtype::decimal::DecimalDType; use crate::dtype::nullability::Nullability::NonNullable; use crate::dtype::nullability::Nullability::Nullable; @@ -349,6 +354,52 @@ mod tests { ); } + #[test] + fn least_supertype_union_requires_identical_variants() { + let nonnullable = DType::Union( + UnionVariants::new( + ["value"].into(), + vec![DType::Primitive(PType::I32, NonNullable)], + ) + .unwrap(), + NonNullable, + ); + let nullable = DType::Union( + UnionVariants::new( + ["value"].into(), + vec![DType::Primitive(PType::I32, Nullable)], + ) + .unwrap(), + NonNullable, + ); + + assert_eq!( + nonnullable.least_supertype(&nonnullable), + Some(nonnullable.clone()) + ); + assert!(nonnullable.least_supertype(&nullable).is_none()); + assert!(nullable.least_supertype(&nonnullable).is_none()); + } + + #[test] + fn least_supertype_null_makes_union_outer_nullable() { + let nonnullable = DType::Union( + UnionVariants::new( + ["value"].into(), + vec![DType::Primitive(PType::I32, NonNullable)], + ) + .unwrap(), + NonNullable, + ); + let nullable = nonnullable.as_nullable(); + + assert_eq!( + DType::Null.least_supertype(&nonnullable), + Some(nullable.clone()) + ); + assert_eq!(nonnullable.least_supertype(&DType::Null), Some(nullable)); + } + #[test] fn least_supertype_unsigned_widening() { let u8_nn = DType::Primitive(PType::U8, NonNullable); diff --git a/vortex-array/src/dtype/dtype_impl.rs b/vortex-array/src/dtype/dtype_impl.rs index 87fff761d81..20ff3d6a182 100644 --- a/vortex-array/src/dtype/dtype_impl.rs +++ b/vortex-array/src/dtype/dtype_impl.rs @@ -15,6 +15,7 @@ use crate::dtype::FieldDType; use crate::dtype::FieldName; use crate::dtype::PType; use crate::dtype::StructFields; +use crate::dtype::UnionVariants; use crate::dtype::decimal::DecimalDType; use crate::dtype::decimal::DecimalType; use crate::dtype::extension::ExtDTypeRef; @@ -63,7 +64,7 @@ impl DType { | List(_, null) | FixedSizeList(_, _, null) | Struct(_, null) - | Union(null) + | Union(_, null) | Variant(null) => matches!(null, Nullability::Nullable), Extension(ext_dtype) => ext_dtype.storage_dtype().is_nullable(), } @@ -79,7 +80,9 @@ impl DType { self.with_nullability(Nullability::Nullable) } - /// Get a new DType with the given nullability (but otherwise the same as `self`) + /// Get a new DType with the given nullability (but otherwise the same as `self`). + /// + /// [`DType::Null`] has intrinsic nullability and is returned unchanged. pub fn with_nullability(&self, nullability: Nullability) -> Self { match self { Null => Null, @@ -91,7 +94,7 @@ impl DType { List(edt, _) => List(Arc::clone(edt), nullability), FixedSizeList(edt, size, _) => FixedSizeList(Arc::clone(edt), *size, nullability), Struct(sf, _) => Struct(sf.clone(), nullability), - Union(_) => Union(nullability), + Union(vs, _) => Union(vs.clone(), nullability), Variant(_) => Variant(nullability), Extension(ext) => Extension(ext.with_nullability(nullability)), } @@ -123,7 +126,15 @@ impl DType { .zip_eq(rhs_dtype.fields()) .all(|(l, r)| l.eq_ignore_nullability(&r))) } - (Union(_), Union(_)) => true, + (Union(lhs, _), Union(rhs, _)) => { + // Equal `names` implies equal length by FieldNames equality. + lhs.names() == rhs.names() + && lhs.type_ids() == rhs.type_ids() + && lhs + .variants() + .zip_eq(rhs.variants()) + .all(|(l, r)| l.eq_ignore_nullability(&r)) + } (Variant(_), Variant(_)) => true, (Extension(lhs_extdtype), Extension(rhs_extdtype)) => { lhs_extdtype.eq_ignore_nullability(rhs_extdtype) @@ -425,6 +436,24 @@ impl DType { } } + /// Get the [`UnionVariants`] if `self` is a [`DType::Union`], otherwise `None`. + pub fn as_union_variants_opt(&self) -> Option<&UnionVariants> { + if let Union(uv, _) = self { + Some(uv) + } else { + None + } + } + + /// Owned version of [Self::as_union_variants_opt]. + pub fn into_union_variants_opt(self) -> Option { + if let Union(uv, _) = self { + Some(uv) + } else { + None + } + } + /// Downcast a `DType` to an `ExtDType` pub fn as_extension(&self) -> &ExtDTypeRef { let Extension(ext) = self else { @@ -476,7 +505,7 @@ impl Display for DType { .map(|(field_null, dt)| format!("{field_null}={dt}")) .join(", "), ), - Union(null) => write!(f, "union(){null}"), + Union(uv, null) => write!(f, "union({uv}){null}"), Variant(null) => write!(f, "variant{null}"), Extension(ext) => write!(f, "{}", ext), } diff --git a/vortex-array/src/dtype/mod.rs b/vortex-array/src/dtype/mod.rs index 009c2610c73..8b579c70d39 100644 --- a/vortex-array/src/dtype/mod.rs +++ b/vortex-array/src/dtype/mod.rs @@ -15,7 +15,6 @@ #[cfg(feature = "arbitrary")] mod arbitrary; -pub mod arrow; mod bigint; mod coercion; mod decimal; @@ -31,6 +30,7 @@ mod ptype; pub mod serde; pub mod session; mod struct_; +mod union; use std::sync::Arc; @@ -110,11 +110,14 @@ pub enum DType { /// 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), + /// A `Union` is composed of one or more **variants**, each with a name and a `DType`. A per-row + /// `u8` tag selects which variant is "live" at that row. + /// + /// A union has independent outer nullability in addition to the nullability of each variant. + /// This distinguishes a null union from a non-null union whose selected child is null. + /// + /// See [`UnionVariants`] for the type-tag conventions and accessors. + Union(UnionVariants, Nullability), /// Dynamically typed values stored as Vortex scalars. /// @@ -146,7 +149,8 @@ impl PartialEq for DType { } // StructFields handles its own Arc::ptr_eq in its PartialEq impl. (Self::Struct(a, na), Self::Struct(b, nb)) => na == nb && a == b, - (Self::Union(a), Self::Union(b)) => a == b, + // UnionVariants handles its own Arc::ptr_eq in its PartialEq impl. + (Self::Union(a, na), Self::Union(b, nb)) => na == nb && a == b, (Self::Variant(a), Self::Variant(b)) => a == b, (Self::Extension(a), Self::Extension(b)) => a == b, // Every variant is listed in the first position so that adding a new @@ -178,6 +182,7 @@ pub use half; pub use nullability::*; pub use ptype::*; pub use struct_::*; +pub use union::*; use crate::dtype::extension::ExtDTypeRef; diff --git a/vortex-array/src/dtype/serde/flatbuffers.rs b/vortex-array/src/dtype/serde/flatbuffers.rs index a61fc8f5527..a796490f58a 100644 --- a/vortex-array/src/dtype/serde/flatbuffers.rs +++ b/vortex-array/src/dtype/serde/flatbuffers.rs @@ -23,6 +23,7 @@ use crate::dtype::DecimalDType; use crate::dtype::FieldDType; use crate::dtype::PType; use crate::dtype::StructFields; +use crate::dtype::UnionVariants; use crate::dtype::extension::ExtId; use crate::dtype::extension::ForeignExtDType; use crate::dtype::flatbuffers as fb; @@ -90,6 +91,43 @@ impl StructFields { } } +impl UnionVariants { + /// Creates a new instance from a flatbuffer-defined object and its underlying buffer. + fn from_fb( + fb_union: fbd::Union<'_>, + buffer: FlatBuffer, + session: VortexSession, + ) -> VortexResult { + let names = fb_union + .names() + .ok_or_else(|| vortex_err!("failed to parse union names from flatbuffer"))? + .iter() + .collect(); + + let dtypes = fb_union + .dtypes() + .ok_or_else(|| vortex_err!("failed to parse union dtypes from flatbuffer"))? + .iter() + .map(|dt| { + FieldDType::from(ViewedDType::from_fb_loc( + dt._tab.loc(), + buffer.clone(), + session.clone(), + )) + }) + .collect::>(); + + let type_ids: Vec = fb_union + .type_ids() + .ok_or_else(|| vortex_err!("failed to parse union type_ids from flatbuffer"))? + .iter() + .map(i8::cast_unsigned) + .collect(); + + UnionVariants::try_from_fields(names, dtypes, type_ids) + } +} + impl DType { /// Create a [`DType`] from a flatbuffer buffer. pub fn from_flatbuffer(buffer: FlatBuffer, session: &VortexSession) -> VortexResult { @@ -195,7 +233,9 @@ impl TryFrom for DType { let fb_union = fb .type__as_union() .ok_or_else(|| vortex_err!("failed to parse union from flatbuffer"))?; - Ok(Self::Union(fb_union.nullable().into())) + let variants = + UnionVariants::from_fb(fb_union, vfdt.buffer().clone(), vfdt.session.clone())?; + Ok(Self::Union(variants, fb_union.nullable().into())) } fb::Type::Variant => { let fb_variant = fb @@ -341,13 +381,41 @@ impl WriteFlatBuffer for DType { ) .as_union_value() } - Self::Union(n) => fb::Union::create( - fbb, - &fb::UnionArgs { - nullable: (*n).into(), - }, - ) - .as_union_value(), + Self::Union(uv, n) => { + let names = uv + .names() + .iter() + .map(|name| fbb.create_string(name.as_ref())) + .collect_vec(); + let names = Some(fbb.create_vector(&names)); + + let dtypes = uv + .variants() + .map(|dtype| dtype.write_flatbuffer(fbb)) + .collect::>>()?; + let dtypes = Some(fbb.create_vector(&dtypes)); + + // The FlatBuffers schema retains its original signed byte wire type for backwards + // compatibility. Reinterpreting the bits preserves the full u8 tag range. + let type_ids = uv + .type_ids() + .iter() + .copied() + .map(u8::cast_signed) + .collect_vec(); + let type_ids = Some(fbb.create_vector(&type_ids)); + + fb::Union::create( + fbb, + &fb::UnionArgs { + names, + dtypes, + type_ids, + nullable: (*n).into(), + }, + ) + .as_union_value() + } Self::Variant(n) => fb::Variant::create( fbb, &fb::VariantArgs { @@ -446,6 +514,7 @@ mod test { use crate::dtype::DType; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::dtype::UnionVariants; use crate::dtype::flatbuffers as fb; use crate::dtype::nullability::Nullability; use crate::dtype::serde::flatbuffers::ViewedDType; @@ -502,4 +571,165 @@ mod test { )); roundtrip_dtype(DType::Variant(Nullability::Nullable)); } + + fn make_basic_union() -> DType { + DType::Union( + UnionVariants::new( + ["int", "str"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + ], + ) + .unwrap(), + Nullability::NonNullable, + ) + } + + #[test] + fn test_union_round_trip_flatbuffer() { + roundtrip_dtype(make_basic_union()); + } + + #[test] + fn test_union_round_trip_flatbuffer_with_nullable_variant() { + let dtype = DType::Union( + UnionVariants::new( + ["null_variant", "str"].into(), + vec![DType::Null, DType::Utf8(Nullability::NonNullable)], + ) + .unwrap(), + Nullability::NonNullable, + ); + roundtrip_dtype(dtype); + } + + #[test] + fn test_union_round_trip_flatbuffer_with_type_id_indirection() { + let dtype = DType::Union( + UnionVariants::try_new( + ["a", "b", "c"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + vec![0, 5, u8::MAX], + ) + .unwrap(), + Nullability::Nullable, + ); + + let bytes = dtype.write_flatbuffer_bytes().unwrap(); + let root_fb = root::(&bytes).unwrap(); + let view = ViewedDType::from_fb_loc( + root_fb._tab.loc(), + FlatBuffer::from(bytes.clone()), + SESSION.clone(), + ); + + let deserialized = DType::try_from(view).unwrap(); + assert_eq!(dtype, deserialized); + let DType::Union(uv, _) = &deserialized else { + panic!("Expected Union"); + }; + assert_eq!(uv.type_ids(), &[0, 5, u8::MAX]); + } + + #[test] + fn test_nested_union_round_trip_flatbuffer() { + let inner_union = make_basic_union(); + let struct_with_union = DType::Struct( + StructFields::from_iter([ + ("id", DType::Primitive(PType::I64, Nullability::NonNullable)), + ("inner", inner_union), + ]), + Nullability::NonNullable, + ); + + let outer_union = DType::Union( + UnionVariants::new( + ["plain", "nested"].into(), + vec![DType::Utf8(Nullability::NonNullable), struct_with_union], + ) + .unwrap(), + Nullability::Nullable, + ); + + roundtrip_dtype(outer_union); + } + + /// A `UnionVariants` parsed lazily from a flatbuffer must compare equal to its eager + /// counterpart. + #[test] + fn test_union_viewed_dtype_equality() { + let eager = make_basic_union(); + + let bytes = eager.write_flatbuffer_bytes().unwrap(); + let root_fb = root::(&bytes).unwrap(); + let view = ViewedDType::from_fb_loc( + root_fb._tab.loc(), + FlatBuffer::from(bytes.clone()), + SESSION.clone(), + ); + + let viewed = DType::try_from(view).unwrap(); + assert_eq!(eager, viewed); + assert_eq!(viewed, eager); + } + + /// A malformed flatbuffer (here, `dtypes.len() != type_ids.len()`) must round-trip + /// to `Err`, not panic. + #[test] + fn test_union_malformed_flatbuffer_errors() { + use flatbuffers::FlatBufferBuilder; + use vortex_buffer::ByteBuffer; + use vortex_flatbuffers::WriteFlatBuffer; + + let mut fbb = FlatBufferBuilder::new(); + + // Build a Union with 2 names + 2 dtypes but only 1 type_id. + let name_a = fbb.create_string("a"); + let name_b = fbb.create_string("b"); + let names = fbb.create_vector(&[name_a, name_b]); + + let dtype_a = DType::Primitive(PType::I32, Nullability::NonNullable) + .write_flatbuffer(&mut fbb) + .unwrap(); + let dtype_b = DType::Utf8(Nullability::NonNullable) + .write_flatbuffer(&mut fbb) + .unwrap(); + let dtypes = fbb.create_vector(&[dtype_a, dtype_b]); + + let type_ids = fbb.create_vector::(&[0]); + + let union_table = fb::Union::create( + &mut fbb, + &fb::UnionArgs { + names: Some(names), + dtypes: Some(dtypes), + type_ids: Some(type_ids), + nullable: false, + }, + ); + + let dtype = fb::DType::create( + &mut fbb, + &fb::DTypeArgs { + type_type: fb::Type::Union, + type_: Some(union_table.as_union_value()), + }, + ); + fbb.finish_minimal(dtype); + let (vec, start) = fbb.collapse(); + let end = vec.len(); + let buffer = FlatBuffer::align_from(ByteBuffer::from(vec).slice(start..end)); + + let root_fb = root::(&buffer).unwrap(); + let view = ViewedDType::from_fb_loc(root_fb._tab.loc(), buffer, SESSION.clone()); + + let result = DType::try_from(view); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("length mismatch")); + } } diff --git a/vortex-array/src/dtype/serde/mod.rs b/vortex-array/src/dtype/serde/mod.rs index d0f43f0c4ac..c6837f4ab2d 100644 --- a/vortex-array/src/dtype/serde/mod.rs +++ b/vortex-array/src/dtype/serde/mod.rs @@ -24,6 +24,7 @@ mod test { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::dtype::UnionVariants; use crate::dtype::serde::DTypeSerde; use crate::dtype::test::SESSION; @@ -181,4 +182,25 @@ mod test { .unwrap(); assert_eq!(DType::Variant(Nullability::Nullable), deserialized); } + + #[test] + fn test_serde_union_dtype_json_roundtrip() { + let dtype = DType::Union( + UnionVariants::new(["value"].into(), vec![DType::Utf8(Nullability::Nullable)]).unwrap(), + Nullability::Nullable, + ); + + let json = serde_json::to_string(&dtype).unwrap(); + assert_eq!( + json, + r#"{"Union":[{"names":["value"],"dtypes":[{"Utf8":true}],"type_ids":[0]},true]}"# + ); + let mut deserializer = serde_json::Deserializer::from_str(&json); + let deserialized: DType = DTypeSerde::::new(&SESSION) + .deserialize(&mut deserializer) + .unwrap(); + + assert_eq!(deserialized, dtype); + assert_eq!(deserialized.nullability(), Nullability::Nullable); + } } diff --git a/vortex-array/src/dtype/serde/proto.rs b/vortex-array/src/dtype/serde/proto.rs index 2f992dce591..8389a90200a 100644 --- a/vortex-array/src/dtype/serde/proto.rs +++ b/vortex-array/src/dtype/serde/proto.rs @@ -12,6 +12,7 @@ use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::PType; use crate::dtype::StructFields; +use crate::dtype::UnionVariants; use crate::dtype::extension::ExtId; use crate::dtype::extension::ForeignExtDType; use crate::dtype::field::Field; @@ -86,7 +87,25 @@ impl DType { ), s.nullable.into(), )), - DtypeType::Union(u) => Ok(Self::Union(u.nullable.into())), + DtypeType::Union(u) => { + let names = u.names.iter().map(|s| s.as_str()).collect(); + let dtypes = u + .dtypes + .iter() + .map(|dt| DType::from_proto(dt, session)) + .collect::>>()?; + let type_ids = u + .type_ids + .iter() + .map(|t| { + u8::try_from(*t).map_err(|_| { + vortex_err!("Union type_id {t} somehow does not fit in u8") + }) + }) + .collect::>>()?; + let variants = UnionVariants::try_new(names, dtypes, type_ids)?; + Ok(Self::Union(variants, u.nullable.into())) + } DtypeType::Variant(v) => Ok(Self::Variant(v.nullable.into())), DtypeType::Extension(e) => { #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] @@ -154,7 +173,13 @@ impl TryFrom<&DType> for pb::DType { .collect::>>()?, nullable: (*null).into(), }), - DType::Union(null) => DtypeType::Union(pb::Union { + DType::Union(uv, null) => DtypeType::Union(pb::Union { + names: uv.names().iter().map(|n| n.as_ref().to_string()).collect(), + dtypes: uv + .variants() + .map(|d| Self::try_from(&d)) + .collect::>>()?, + type_ids: uv.type_ids().iter().map(|t| *t as i32).collect(), nullable: (*null).into(), }), DType::Variant(null) => DtypeType::Variant(pb::Variant { @@ -231,6 +256,7 @@ mod tests { use std::sync::Arc; use super::*; + use crate::array_session; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::Field; @@ -238,6 +264,7 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::dtype::UnionVariants; use crate::dtype::proto::dtype::d_type::DtypeType; use crate::dtype::proto::dtype::field::FieldType; use crate::dtype::test::SESSION; @@ -380,6 +407,118 @@ mod tests { assert_eq!(DType::Variant(Nullability::Nullable), converted); } + #[test] + fn test_union_round_trip_proto() { + let variants = UnionVariants::new( + ["int", "str"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + ], + ) + .unwrap(); + let dtype = DType::Union(variants, Nullability::NonNullable); + let converted = round_trip_dtype(&dtype); + assert_eq!(dtype, converted); + } + + #[test] + fn test_union_round_trip_proto_with_inner_and_outer_nullability() { + let variants = UnionVariants::new( + ["null_variant", "str"].into(), + vec![DType::Null, DType::Utf8(Nullability::NonNullable)], + ) + .unwrap(); + + let dtype = DType::Union(variants, Nullability::Nullable); + let converted = round_trip_dtype(&dtype); + assert_eq!(dtype, converted); + } + + #[test] + fn test_union_round_trip_proto_with_type_id_indirection() { + let variants = UnionVariants::try_new( + ["a", "b", "c"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + vec![0, 5, 7], + ) + .unwrap(); + + let dtype = DType::Union(variants, Nullability::Nullable); + let converted = round_trip_dtype(&dtype); + assert_eq!(dtype, converted); + + let DType::Union(uv, _) = &converted else { + panic!("Expected Union"); + }; + assert_eq!(uv.type_ids(), &[0, 5, 7]); + assert!(!uv.is_consecutive()); + } + + #[test] + fn test_union_proto_rejects_out_of_range_type_id() { + let proto = pb::DType { + dtype_type: Some(DtypeType::Union(pb::Union { + names: vec!["a".to_string()], + dtypes: vec![ + pb::DType::try_from(&DType::Primitive(PType::I32, Nullability::NonNullable)) + .unwrap(), + ], + // 256 does not fit in u8. + type_ids: vec![256], + nullable: false, + })), + }; + + let result = DType::from_proto(&proto, &SESSION); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("does not fit in u8") + ); + } + + #[test] + fn test_nested_union_round_trip_proto() { + let inner_union = DType::Union( + UnionVariants::new( + ["a", "b"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + ], + ) + .unwrap(), + Nullability::NonNullable, + ); + + let struct_with_union = DType::Struct( + StructFields::from_iter([ + ("id", DType::Primitive(PType::I64, Nullability::NonNullable)), + ("inner", inner_union), + ]), + Nullability::NonNullable, + ); + + let outer_union = DType::Union( + UnionVariants::new( + ["plain", "nested"].into(), + vec![DType::Utf8(Nullability::NonNullable), struct_with_union], + ) + .unwrap(), + Nullability::Nullable, + ); + + let converted = round_trip_dtype(&outer_union); + assert_eq!(outer_union, converted); + } + #[test] fn test_field_path_round_trip() { let test_paths = vec![ @@ -502,7 +641,8 @@ mod tests { #[test] fn test_unknown_extension_allow_unknown() { - let session = crate::array_session().allow_unknown(); + let session = array_session(); + 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 789fd17306b..b3870b8c03e 100644 --- a/vortex-array/src/dtype/serde/serde.rs +++ b/vortex-array/src/dtype/serde/serde.rs @@ -27,6 +27,7 @@ use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; +use crate::dtype::UnionVariants; use crate::dtype::decimal::DecimalDType; use crate::dtype::extension::ExtDTypeRef; use crate::dtype::extension::ExtId; @@ -115,7 +116,12 @@ impl Serialize for DType { state.serialize_field(n)?; state.end() } - DType::Union(n) => serializer.serialize_newtype_variant("DType", 11, "Union", n), + DType::Union(uv, n) => { + let mut state = serializer.serialize_tuple_variant("DType", 11, "Union", 2)?; + state.serialize_field(uv)?; + state.serialize_field(n)?; + state.end() + } DType::Variant(n) => serializer.serialize_newtype_variant("DType", 10, "Variant", n), DType::Extension(ext) => { serializer.serialize_newtype_variant("DType", 9, "Extension", ext) @@ -137,6 +143,20 @@ impl Serialize for StructFields { } } +impl Serialize for UnionVariants { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let mut state = serializer.serialize_struct("UnionVariants", 3)?; + state.serialize_field("names", self.names())?; + let dtypes: Vec = self.variants().collect(); + state.serialize_field("dtypes", &dtypes)?; + state.serialize_field("type_ids", self.type_ids())?; + state.end() + } +} + // ============================================================================ // DType Deserialization with session context (DeserializeSeed) // ============================================================================ @@ -217,10 +237,9 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, DType> { "Struct" => access.newtype_variant_seed(StructFieldsSeed { session: self.session, }), - "Union" => { - let n = access.newtype_variant()?; - Ok(DType::Union(n)) - } + "Union" => access.newtype_variant_seed(UnionFieldsSeed { + session: self.session, + }), "Variant" => { let n = access.newtype_variant()?; Ok(DType::Variant(n)) @@ -391,6 +410,126 @@ impl<'de> DeserializeSeed<'de> for StructFieldsSeed<'_> { } } +struct UnionFieldsSeed<'a> { + session: &'a VortexSession, +} + +impl<'de> DeserializeSeed<'de> for UnionFieldsSeed<'_> { + type Value = DType; + + fn deserialize(self, deserializer: D) -> Result + where + D: Deserializer<'de>, + { + struct UnionVisitor<'a> { + session: &'a VortexSession, + } + + impl<'de> Visitor<'de> for UnionVisitor<'_> { + type Value = DType; + + fn expecting(&self, f: &mut Formatter) -> fmt::Result { + f.write_str("Union tuple (variants, nullability)") + } + + fn visit_seq(self, mut seq: A) -> Result + where + A: SeqAccess<'de>, + { + let variants = seq + .next_element_seed(DTypeSerde::::new(self.session))? + .ok_or_else(|| de::Error::invalid_length(0, &self))?; + let nullability = seq + .next_element()? + .ok_or_else(|| de::Error::invalid_length(1, &self))?; + Ok(DType::Union(variants, nullability)) + } + } + + deserializer.deserialize_tuple( + 2, + UnionVisitor { + session: self.session, + }, + ) + } +} + +impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, UnionVariants> { + type Value = UnionVariants; + + fn deserialize(self, deserializer: D) -> Result + where + D: Deserializer<'de>, + { + const FIELDS: &[&str] = &["names", "dtypes", "type_ids"]; + + struct UnionVariantsInnerVisitor<'a> { + session: &'a VortexSession, + } + + impl<'de> Visitor<'de> for UnionVariantsInnerVisitor<'_> { + type Value = UnionVariants; + + fn expecting(&self, f: &mut Formatter) -> fmt::Result { + f.write_str("struct UnionVariants") + } + + fn visit_map(self, mut map: V) -> Result + where + V: MapAccess<'de>, + { + let mut names: Option = None; + let mut dtypes: Option> = None; + let mut type_ids: Option> = None; + + while let Some(key) = map.next_key::>()? { + match key.as_ref() { + "names" => { + if names.is_some() { + return Err(de::Error::duplicate_field("names")); + } + names = Some(map.next_value()?); + } + "dtypes" => { + if dtypes.is_some() { + return Err(de::Error::duplicate_field("dtypes")); + } + dtypes = Some( + map.next_value_seed(DTypeSerde::>::new(self.session))?, + ); + } + "type_ids" => { + if type_ids.is_some() { + return Err(de::Error::duplicate_field("type_ids")); + } + type_ids = Some(map.next_value()?); + } + _ => { + let _ = map.next_value::()?; + } + } + } + + let names = names.ok_or_else(|| de::Error::missing_field("names"))?; + let dtypes = dtypes.ok_or_else(|| de::Error::missing_field("dtypes"))?; + let type_ids = type_ids.ok_or_else(|| de::Error::missing_field("type_ids"))?; + + UnionVariants::try_new(names, dtypes, type_ids) + .map_err(|e| de::Error::custom(e.to_string())) + } + } + + deserializer.deserialize_struct( + "UnionVariants", + FIELDS, + UnionVariantsInnerVisitor { + session: self.session, + }, + ) + } +} + impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, StructFields> { type Value = StructFields; diff --git a/vortex-array/src/dtype/union.rs b/vortex-array/src/dtype/union.rs new file mode 100644 index 00000000000..7ce53155adc --- /dev/null +++ b/vortex-array/src/dtype/union.rs @@ -0,0 +1,586 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::hash::Hash; +use std::sync::Arc; + +use itertools::Itertools; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; + +use crate::dtype::DType; +use crate::dtype::FieldDType; +use crate::dtype::FieldNames; + +/// Type information for a union array. +/// +/// A `UnionVariants` describes the possible alternative types for each row of a union, along with a +/// per-variant `u8` type tag. We use the term **variants** (rather than "fields") because a union +/// is a sum type: each row chooses exactly one alternative. +/// +/// By default, tag `i` selects the child at offset `i` (`type_ids = [0, 1, ..., N-1]`). Vortex uses +/// unsigned tags and supports up to 256 variants. +/// +/// Schemas may also use non-consecutive tags (e.g. `[0, 5, 7]`), in which case the value of +/// `type_ids[i]` is the tag used in the data to select the child at offset `i`. Supporting +/// non-consecutive tags lets the schema remove children without renumbering the remaining tags. +/// +/// Variant names must be distinct. Unlike [`StructFields`](crate::dtype::StructFields), which +/// permits duplicate field names for Arrow/Parquet round-trip fidelity, duplicates have no +/// meaningful semantics in a sum type and are rejected at construction. +/// +/// Note that the Arrow spec does not require distinct names (children are selected by the type-id +/// buffer, never by name), but DuckDB requires distinct member names for its `UNION` type and +/// resolves members by name (e.g. in `union_extract`), so permitting duplicates would break DuckDB +/// round-trips while enabling no useful semantics. +/// +/// ``` +/// use vortex_array::dtype::{DType, Nullability, PType, UnionVariants}; +/// +/// let variants = UnionVariants::new( +/// ["a", "b"].into(), +/// vec![ +/// DType::Primitive(PType::I32, Nullability::NonNullable), +/// DType::Utf8(Nullability::NonNullable), +/// ], +/// ) +/// .unwrap(); +/// +/// assert_eq!(variants.len(), 2); +/// assert_eq!(variants.type_ids(), &[0, 1]); +/// ``` +#[allow( + clippy::derived_hash_with_manual_eq, + reason = "manual PartialEq adds Arc::ptr_eq fast path only" +)] +#[derive(Clone, Eq, Hash)] +pub struct UnionVariants(Arc); + +impl PartialEq for UnionVariants { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) || self.0 == other.0 + } +} + +impl fmt::Debug for UnionVariants { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnionVariants") + .field("names", &self.0.names) + .field("dtypes", &self.0.dtypes) + .field("type_ids", &self.0.type_ids) + .finish() + } +} + +impl fmt::Display for UnionVariants { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // Surface non-consecutive type tags so they aren't hidden by the format. Consecutive + // `[0, 1, ..., N-1]` is the common case and is left implicit. + let show_tags = !self.is_consecutive(); + write!( + f, + "{}", + self.names() + .iter() + .zip(self.variants()) + .zip(self.type_ids().iter()) + .map(|((name, dt), tag)| { + if show_tags { + format!("{name}@{tag}={dt}") + } else { + format!("{name}={dt}") + } + }) + .join(", ") + ) + } +} + +struct UnionVariantsInner { + /// The names of the variants. This is called `FieldNames` because it is shared with the + /// [`StructFields`](crate::dtype::StructFields) implementation. + names: FieldNames, + + /// The types of each of the variants. + dtypes: Arc<[FieldDType]>, + + /// One tag per variant, in variant order. The common case where children are referenced by + /// consecutive offsets is `[0, 1, ..., N-1]`. + /// + /// For schemas with explicit `typeIds` indirection (e.g. `[0, 5, 7]`), this stores those tags. + type_ids: Arc<[u8]>, +} + +impl UnionVariantsInner { + fn from_fields(names: FieldNames, dtypes: Arc<[FieldDType]>, type_ids: Arc<[u8]>) -> Self { + Self { + names, + dtypes, + type_ids, + } + } +} + +impl PartialEq for UnionVariantsInner { + fn eq(&self, other: &Self) -> bool { + self.names == other.names && self.dtypes == other.dtypes && self.type_ids == other.type_ids + } +} + +impl Eq for UnionVariantsInner {} + +impl Hash for UnionVariantsInner { + fn hash(&self, state: &mut H) { + self.names.hash(state); + self.dtypes.hash(state); + self.type_ids.hash(state); + } +} + +impl Default for UnionVariants { + fn default() -> Self { + Self::empty() + } +} + +impl UnionVariants { + /// The variants of the empty union. + pub fn empty() -> Self { + Self(Arc::new(UnionVariantsInner::from_fields( + FieldNames::default(), + Arc::from([]), + Arc::from([]), + ))) + } + + /// Validate that `names`, `dtypes`, and `type_ids` are mutually consistent. + fn validate_shape(names: &FieldNames, n_dtypes: usize, type_ids: &[u8]) -> VortexResult<()> { + vortex_ensure_eq!( + names.len(), + n_dtypes, + "length mismatch between names ({}) and dtypes ({})", + names.len(), + n_dtypes + ); + vortex_ensure_eq!( + names.len(), + type_ids.len(), + "length mismatch between names ({}) and type_ids ({})", + names.len(), + type_ids.len() + ); + + vortex_ensure!( + type_ids.iter().all_unique(), + "type_ids must be distinct, got {:?}", + type_ids + ); + vortex_ensure!( + type_ids.len() <= u8::MAX as usize + 1, + "union supports at most {} variants, got {}", + u8::MAX as usize + 1, + type_ids.len() + ); + vortex_ensure!( + names.iter().all_unique(), + "union variant names must be distinct, got {:?}", + names + ); + + Ok(()) + } + + /// Create a new [`UnionVariants`] with explicit `type_ids`. + /// + /// # Errors + /// + /// Returns an error if names, dtypes, or type IDs do not all have the same length, or if there + /// are any duplicate names or type IDs, or if there are more than 256 variants. + pub fn try_new(names: FieldNames, dtypes: Vec, type_ids: Vec) -> VortexResult { + Self::validate_shape(&names, dtypes.len(), &type_ids)?; + + let dtypes: Arc<[FieldDType]> = dtypes.into_iter().map(FieldDType::from).collect(); + let type_ids: Arc<[u8]> = Arc::from(type_ids); + + Ok(Self(Arc::new(UnionVariantsInner::from_fields( + names, dtypes, type_ids, + )))) + } + + /// Create a new [`UnionVariants`] with consecutive `type_ids = [0, 1, ..., N-1]`. + /// + /// # Errors + /// + /// `names` and `dtypes` must have the same length, and `names.len()` cannot be more than + /// `u8::MAX as usize + 1` (256). + pub fn new(names: FieldNames, dtypes: Vec) -> VortexResult { + const MAX_VARIANTS: usize = u8::MAX as usize + 1; + vortex_ensure!( + names.len() <= MAX_VARIANTS, + "union supports at most {} consecutive variants, got {}", + MAX_VARIANTS, + names.len() + ); + + #[expect( + clippy::cast_possible_truncation, + reason = "the MAX_VARIANTS bound above guarantees `i as u8` is in range" + )] + let type_ids: Vec = (0..names.len()).map(|i| i as u8).collect(); + + Self::try_new(names, dtypes, type_ids) + } + + /// Create a new [`UnionVariants`] from pre-constructed [`FieldDType`]s, which may be owned or + /// backed by a flatbuffer view. + /// + /// Used by deserialization paths where the children may be lazily backed by a flatbuffer. + /// + /// # Errors + /// + /// Returns an error if names, dtypes, or type IDs do not all have the same length, or if there + /// are any duplicate names or type IDs, or if there are more than 256 variants. + pub(crate) fn try_from_fields( + names: FieldNames, + dtypes: Vec, + type_ids: Vec, + ) -> VortexResult { + Self::validate_shape(&names, dtypes.len(), &type_ids)?; + + Ok(Self(Arc::new(UnionVariantsInner::from_fields( + names, + dtypes.into(), + Arc::from(type_ids), + )))) + } + + /// Get the names of the variants in the union. + pub fn names(&self) -> &FieldNames { + &self.0.names + } + + /// Returns the number of variants in the union. + pub fn len(&self) -> usize { + self.0.names.len() + } + + /// Returns true if the union has no variants. + pub fn is_empty(&self) -> bool { + self.0.names.is_empty() + } + + /// Returns the per-variant type tag vector. Entry `i` is the tag that the data uses to + /// select the variant at offset `i`. + pub fn type_ids(&self) -> &[u8] { + &self.0.type_ids + } + + /// Returns `true` if the type tags are the consecutive sequence `[0, 1, ..., N-1]`. + pub fn is_consecutive(&self) -> bool { + self.0 + .type_ids + .iter() + .enumerate() + .all(|(i, &tag)| u8::try_from(i).is_ok_and(|i| i == tag)) + } + + /// Find the offset of a variant by name. Returns `None` if no variant has the name. + /// + /// This is a linear scan over [`Self::names`]. A union has at most 256 variants (and usually + /// far fewer), so a linear scan beats building and probing a `HashMap` in practice. + pub fn find(&self, name: impl AsRef) -> Option { + let name = name.as_ref(); + self.0.names.iter().position(|n| n.as_ref() == name) + } + + /// Get the [`DType`] of a variant by name. Returns `None` if no variant has the name. + pub fn variant(&self, name: impl AsRef) -> Option { + let index = self.find(name)?; + Some( + self.0.dtypes[index] + .value() + .vortex_expect("variant DType must be valid"), + ) + } + + /// Get the [`DType`] of a variant by offset. + pub fn variant_by_index(&self, index: usize) -> Option { + Some( + self.0 + .dtypes + .get(index)? + .value() + .vortex_expect("variant DType must be valid"), + ) + } + + /// Returns an ordered iterator over the variants. + pub fn variants(&self) -> impl ExactSizeIterator + '_ { + self.0 + .dtypes + .iter() + .map(|dt| dt.value().vortex_expect("variant DType must be valid")) + } + + /// Convert a data-level type tag to a child offset. Returns `None` if the tag is not in + /// `type_ids`. + /// + /// This is a linear scan over [`Self::type_ids`]. The number of variants is bounded by + /// 256, so a linear scan is faster than a `HashMap` lookup in practice for the + /// small unions typically encountered. + pub fn tag_to_child_index(&self, tag: u8) -> Option { + self.0.type_ids.iter().position(|&t| t == tag) + } + + /// Convert a child offset to its data-level type tag. + /// + /// # Panics + /// + /// Panics if `index >= self.len()`. + pub fn child_index_to_tag(&self, index: usize) -> u8 { + self.0.type_ids[index] + } +} + +#[cfg(test)] +mod tests { + use crate::dtype::DType; + use crate::dtype::FieldNames; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::dtype::UnionVariants; + + fn i32_variants() -> UnionVariants { + UnionVariants::new( + ["int", "str"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + ], + ) + .unwrap() + } + + #[test] + fn test_consecutive_type_ids() { + let variants = i32_variants(); + assert_eq!(variants.type_ids(), &[0, 1]); + assert!(variants.is_consecutive()); + assert_eq!(variants.tag_to_child_index(0), Some(0)); + assert_eq!(variants.tag_to_child_index(1), Some(1)); + assert_eq!(variants.tag_to_child_index(2), None); + assert_eq!(variants.child_index_to_tag(0), 0); + assert_eq!(variants.child_index_to_tag(1), 1); + } + + #[test] + fn test_type_id_indirection() { + let variants = UnionVariants::try_new( + ["a", "b", "c"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + vec![0, 5, 7], + ) + .unwrap(); + + assert_eq!(variants.type_ids(), &[0, 5, 7]); + assert!(!variants.is_consecutive()); + assert_eq!(variants.tag_to_child_index(0), Some(0)); + assert_eq!(variants.tag_to_child_index(5), Some(1)); + assert_eq!(variants.tag_to_child_index(7), Some(2)); + assert_eq!(variants.tag_to_child_index(1), None); + } + + #[test] + fn test_find() { + let variants = i32_variants(); + assert_eq!(variants.find("int"), Some(0)); + assert_eq!(variants.find("str"), Some(1)); + assert!(variants.find("missing").is_none()); + + let value = variants.variant("int").unwrap(); + assert_eq!( + value, + DType::Primitive(PType::I32, Nullability::NonNullable) + ); + } + + #[test] + fn test_duplicate_names_rejected() { + let result = UnionVariants::new( + ["dup", "dup"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Primitive(PType::I64, Nullability::NonNullable), + ], + ); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("distinct")); + } + + #[test] + fn test_high_type_id() { + let variants = UnionVariants::try_new( + ["high"].into(), + vec![DType::Primitive(PType::I32, Nullability::NonNullable)], + vec![u8::MAX], + ) + .unwrap(); + assert_eq!(variants.type_ids(), &[u8::MAX]); + } + + #[test] + fn test_length_mismatch_rejected() { + let result = UnionVariants::try_new( + ["a", "b"].into(), + vec![DType::Primitive(PType::I32, Nullability::NonNullable)], + vec![0, 1], + ); + assert!(result.is_err()); + + let result = UnionVariants::try_new( + ["a"].into(), + vec![DType::Primitive(PType::I32, Nullability::NonNullable)], + vec![0, 1], + ); + assert!(result.is_err()); + } + + #[test] + fn test_duplicate_type_ids_rejected() { + let result = UnionVariants::try_new( + ["a", "b"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + ], + vec![3, 3], + ); + assert!(result.is_err()); + } + + #[test] + fn test_nested_union_nullability_is_independent() { + let inner = DType::Union( + UnionVariants::new(["value"].into(), vec![DType::Utf8(Nullability::Nullable)]).unwrap(), + Nullability::NonNullable, + ); + let outer = DType::Union( + UnionVariants::new( + ["nested", "number"].into(), + vec![ + inner, + DType::Primitive(PType::I32, Nullability::NonNullable), + ], + ) + .unwrap(), + Nullability::Nullable, + ); + + assert_eq!(outer.nullability(), Nullability::Nullable); + let variants = outer.as_union_variants_opt().unwrap(); + assert_eq!( + variants.variant_by_index(0).unwrap().nullability(), + Nullability::NonNullable + ); + } + + #[test] + fn test_with_nullability_changes_only_outer_union() { + let nonnullable = DType::Union(i32_variants(), Nullability::NonNullable); + assert_eq!(nonnullable.nullability(), Nullability::NonNullable); + + let nullable = nonnullable.as_nullable(); + assert_eq!(nullable.nullability(), Nullability::Nullable); + assert_eq!( + nullable.as_union_variants_opt(), + nonnullable.as_union_variants_opt() + ); + assert_eq!(nullable.as_nonnullable(), nonnullable); + } + + #[test] + fn test_display() { + let variants = i32_variants(); + let dtype = DType::Union(variants, Nullability::NonNullable); + assert_eq!(dtype.to_string(), "union(int=i32, str=utf8)"); + + let nullable = DType::Union( + UnionVariants::new( + ["int", "maybe_str"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + ], + ) + .unwrap(), + Nullability::Nullable, + ); + assert_eq!(nullable.to_string(), "union(int=i32, maybe_str=utf8?)?"); + assert_eq!( + nullable.as_nonnullable().to_string(), + "union(int=i32, maybe_str=utf8?)" + ); + } + + #[test] + fn test_display_with_type_id_indirection() { + let variants = UnionVariants::try_new( + ["a", "b", "c"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + DType::Bool(Nullability::NonNullable), + ], + vec![0, 5, 7], + ) + .unwrap(); + let dtype = DType::Union(variants, Nullability::NonNullable); + assert_eq!(dtype.to_string(), "union(a@0=i32, b@5=utf8, c@7=bool)"); + } + + #[test] + fn test_new_max_size() { + let names: Vec = (0..256).map(|i| format!("v{i}")).collect(); + let dtypes: Vec = (0..256) + .map(|_| DType::Primitive(PType::I32, Nullability::NonNullable)) + .collect(); + let names: FieldNames = names.into_iter().collect(); + let variants = UnionVariants::new(names, dtypes).unwrap(); + assert_eq!(variants.len(), 256); + assert_eq!(variants.type_ids()[255], 255); + assert!(variants.is_consecutive()); + } + + #[test] + fn test_new_too_large_rejected() { + let names: Vec = (0..257).map(|i| format!("v{i}")).collect(); + let dtypes: Vec = (0..257) + .map(|_| DType::Primitive(PType::I32, Nullability::NonNullable)) + .collect(); + let names: FieldNames = names.into_iter().collect(); + let result = UnionVariants::new(names, dtypes); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("at most 256 consecutive variants") + ); + } + + #[test] + fn test_empty() { + let v = UnionVariants::empty(); + assert!(v.is_empty()); + assert_eq!(v.len(), 0); + assert_eq!(v.type_ids(), &[] as &[u8]); + assert!(v.is_consecutive()); + } +} diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 4f41302c91f..a49f98b4d55 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -495,9 +495,9 @@ impl Executable for ArrayRef { /// Execute `array` into the given `builder`. /// /// This uses the encoding's [`crate::array::VTable::append_to_builder`] implementation. Most -/// encodings use the default path of `execute::` followed by `builder.extend_from_array`, -/// while encodings like `Chunked` can override that to append child-by-child without materializing -/// the entire parent. +/// encodings use the default path of `execute::` followed by re-dispatching +/// `append_to_builder` on the canonical array, while encodings like `Chunked` can override that to +/// append child-by-child without materializing the entire parent. /// /// The builder must have a [`DType`] that is a nullability-superset of `array.dtype()`. pub fn execute_into_builder( diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index ec12a321552..da70427a0f7 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -10,6 +10,7 @@ use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_utils::iter::ReduceBalancedIterExt; +use crate::aggregate_fn::NumericalAggregateOpts; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; @@ -38,6 +39,7 @@ 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::list_sum::ListSum; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::merge::DuplicateHandling; @@ -201,15 +203,17 @@ pub fn nested_case_when( /// ``` /// # use vortex_array::arrays::{BoolArray, PrimitiveArray}; /// # use vortex_array::arrays::bool::BoolArrayExt; -/// # use vortex_array::{IntoArray, ToCanonical}; +/// # use vortex_array::IntoArray; +/// # use vortex_array::{VortexSessionExecute, array_session}; /// # use vortex_array::validity::Validity; /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{eq, root, lit}; /// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); /// let result = xs.into_array().apply(&eq(root(), lit(3))).unwrap(); +/// let mut ctx = array_session().create_execution_ctx(); /// /// assert_eq!( -/// result.to_bool().to_bit_buffer(), +/// result.execute::(&mut ctx).unwrap().to_bit_buffer(), /// BoolArray::from_iter(vec![false, false, true]).to_bit_buffer(), /// ); /// ``` @@ -226,15 +230,17 @@ pub fn eq(lhs: Expression, rhs: Expression) -> Expression { /// ``` /// # use vortex_array::arrays::{BoolArray, PrimitiveArray}; /// # use vortex_array::arrays::bool::BoolArrayExt; -/// # use vortex_array::{ IntoArray, ToCanonical}; +/// # use vortex_array::IntoArray; +/// # use vortex_array::{VortexSessionExecute, array_session}; /// # use vortex_array::validity::Validity; /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{root, lit, not_eq}; /// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); /// let result = xs.into_array().apply(¬_eq(root(), lit(3))).unwrap(); +/// let mut ctx = array_session().create_execution_ctx(); /// /// assert_eq!( -/// result.to_bool().to_bit_buffer(), +/// result.execute::(&mut ctx).unwrap().to_bit_buffer(), /// BoolArray::from_iter(vec![true, true, false]).to_bit_buffer(), /// ); /// ``` @@ -251,15 +257,17 @@ pub fn not_eq(lhs: Expression, rhs: Expression) -> Expression { /// ``` /// # use vortex_array::arrays::{BoolArray, PrimitiveArray }; /// # use vortex_array::arrays::bool::BoolArrayExt; -/// # use vortex_array::{IntoArray, ToCanonical}; +/// # use vortex_array::IntoArray; +/// # use vortex_array::{VortexSessionExecute, array_session}; /// # use vortex_array::validity::Validity; /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{gt_eq, root, lit}; /// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); /// let result = xs.into_array().apply(>_eq(root(), lit(3))).unwrap(); +/// let mut ctx = array_session().create_execution_ctx(); /// /// assert_eq!( -/// result.to_bool().to_bit_buffer(), +/// result.execute::(&mut ctx).unwrap().to_bit_buffer(), /// BoolArray::from_iter(vec![false, false, true]).to_bit_buffer(), /// ); /// ``` @@ -276,15 +284,17 @@ pub fn gt_eq(lhs: Expression, rhs: Expression) -> Expression { /// ``` /// # use vortex_array::arrays::{BoolArray, PrimitiveArray }; /// # use vortex_array::arrays::bool::BoolArrayExt; -/// # use vortex_array::{IntoArray, ToCanonical}; +/// # use vortex_array::IntoArray; +/// # use vortex_array::{VortexSessionExecute, array_session}; /// # use vortex_array::validity::Validity; /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{gt, root, lit}; /// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); /// let result = xs.into_array().apply(>(root(), lit(2))).unwrap(); +/// let mut ctx = array_session().create_execution_ctx(); /// /// assert_eq!( -/// result.to_bool().to_bit_buffer(), +/// result.execute::(&mut ctx).unwrap().to_bit_buffer(), /// BoolArray::from_iter(vec![false, false, true]).to_bit_buffer(), /// ); /// ``` @@ -301,15 +311,17 @@ pub fn gt(lhs: Expression, rhs: Expression) -> Expression { /// ``` /// # use vortex_array::arrays::{BoolArray, PrimitiveArray }; /// # use vortex_array::arrays::bool::BoolArrayExt; -/// # use vortex_array::{IntoArray, ToCanonical}; +/// # use vortex_array::IntoArray; +/// # use vortex_array::{VortexSessionExecute, array_session}; /// # use vortex_array::validity::Validity; /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{root, lit, lt_eq}; /// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); /// let result = xs.into_array().apply(<_eq(root(), lit(2))).unwrap(); +/// let mut ctx = array_session().create_execution_ctx(); /// /// assert_eq!( -/// result.to_bool().to_bit_buffer(), +/// result.execute::(&mut ctx).unwrap().to_bit_buffer(), /// BoolArray::from_iter(vec![true, true, false]).to_bit_buffer(), /// ); /// ``` @@ -326,15 +338,17 @@ pub fn lt_eq(lhs: Expression, rhs: Expression) -> Expression { /// ``` /// # use vortex_array::arrays::{BoolArray, PrimitiveArray }; /// # use vortex_array::arrays::bool::BoolArrayExt; -/// # use vortex_array::{IntoArray, ToCanonical}; +/// # use vortex_array::IntoArray; +/// # use vortex_array::{VortexSessionExecute, array_session}; /// # use vortex_array::validity::Validity; /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{root, lit, lt}; /// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); /// let result = xs.into_array().apply(<(root(), lit(3))).unwrap(); +/// let mut ctx = array_session().create_execution_ctx(); /// /// assert_eq!( -/// result.to_bool().to_bit_buffer(), +/// result.execute::(&mut ctx).unwrap().to_bit_buffer(), /// BoolArray::from_iter(vec![true, true, false]).to_bit_buffer(), /// ); /// ``` @@ -351,13 +365,15 @@ pub fn lt(lhs: Expression, rhs: Expression) -> Expression { /// ``` /// # use vortex_array::arrays::BoolArray; /// # use vortex_array::arrays::bool::BoolArrayExt; -/// # use vortex_array::{IntoArray, ToCanonical}; +/// # use vortex_array::IntoArray; +/// # use vortex_array::{VortexSessionExecute, array_session}; /// # use vortex_array::expr::{root, lit, or}; /// let xs = BoolArray::from_iter(vec![true, false, true]); /// let result = xs.into_array().apply(&or(root(), lit(false))).unwrap(); +/// let mut ctx = array_session().create_execution_ctx(); /// /// assert_eq!( -/// result.to_bool().to_bit_buffer(), +/// result.execute::(&mut ctx).unwrap().to_bit_buffer(), /// BoolArray::from_iter(vec![true, false, true]).to_bit_buffer(), /// ); /// ``` @@ -387,13 +403,15 @@ where /// ``` /// # use vortex_array::arrays::BoolArray; /// # use vortex_array::arrays::bool::BoolArrayExt; -/// # use vortex_array::{IntoArray, ToCanonical}; +/// # use vortex_array::IntoArray; +/// # use vortex_array::{VortexSessionExecute, array_session}; /// # use vortex_array::expr::{and, root, lit}; /// let xs = BoolArray::from_iter(vec![true, false, true]).into_array(); /// let result = xs.apply(&and(root(), lit(true))).unwrap(); +/// let mut ctx = array_session().create_execution_ctx(); /// /// assert_eq!( -/// result.to_bool().to_bit_buffer(), +/// result.execute::(&mut ctx).unwrap().to_bit_buffer(), /// BoolArray::from_iter(vec![true, false, true]).to_bit_buffer(), /// ); /// ``` @@ -422,7 +440,8 @@ where /// /// ``` /// # use vortex_array::IntoArray; -/// # use vortex_array::arrow::ArrowArrayExecutor; +/// # use vortex_array::arrays::PrimitiveArray; +/// # use vortex_array::builtins::ArrayBuiltins; /// # use vortex_array::{VortexSessionExecute, array_session}; /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{checked_add, lit, root}; @@ -430,13 +449,8 @@ where /// let result = xs.apply(&checked_add(root(), lit(5))).unwrap(); /// /// let mut ctx = array_session().create_execution_ctx(); -/// assert_eq!( -/// &result.execute_arrow(None, &mut ctx).unwrap(), -/// &buffer![6, 7, 8] -/// .into_array() -/// .execute_arrow(None, &mut ctx) -/// .unwrap() -/// ); +/// let result = result.execute::(&mut ctx).unwrap(); +/// assert_eq!(result.as_slice::(), [6, 7, 8]); /// ``` pub fn checked_add(lhs: Expression, rhs: Expression) -> Expression { Binary @@ -765,3 +779,27 @@ pub fn ext_storage(input: Expression) -> Expression { pub fn list_length(input: Expression) -> Expression { ListLength.new_expr(EmptyOptions, [input]) } + +// ---- ListSum ---- + +/// Creates an expression that sums the elements of each list for `List` and +/// `FixedSizeList` inputs, akin to DuckDB's `list_sum()`. +/// +/// Follows SQL `SUM` semantics per list: null lists, empty lists, and lists whose elements are +/// all null yield null; null elements are skipped; integer and decimal overflow yields a null +/// value. The result dtype follows `sum`'s widening rules and is always nullable. NaN float +/// elements are skipped by default; see [`list_sum_opts`] for the NaN-including variant. +/// +/// ```rust +/// # use vortex_array::expr::{list_sum, root}; +/// let expr = list_sum(root()); +/// ``` +pub fn list_sum(input: Expression) -> Expression { + ListSum.new_expr(NumericalAggregateOpts::default(), [input]) +} + +/// Creates a [`list_sum`] expression with explicit [`NumericalAggregateOpts`], controlling +/// whether NaN float elements are skipped (the default) or poison the list's sum to NaN. +pub fn list_sum_opts(input: Expression, options: NumericalAggregateOpts) -> Expression { + ListSum.new_expr(options, [input]) +} diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index 9f2b81bde84..0a8d71f2954 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -116,9 +116,8 @@ mod tests { #[test] fn unknown_expression_id_allow_unknown() { - let session = VortexSession::empty() - .with::() - .allow_unknown(); + let session = VortexSession::empty().with::(); + session.allow_unknown(); let expr_proto = pb::Expr { id: "vortex.test.foreign_scalar_fn".to_string(), diff --git a/vortex-array/src/extension/uuid/mod.rs b/vortex-array/src/extension/uuid/mod.rs index 8a178035d2f..e4347c2513c 100644 --- a/vortex-array/src/extension/uuid/mod.rs +++ b/vortex-array/src/extension/uuid/mod.rs @@ -13,7 +13,6 @@ mod metadata; pub use metadata::UuidMetadata; -mod arrow; pub(crate) mod vtable; /// The VTable for the UUID extension type. diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 7496c527172..9d3a59eb7a0 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -95,7 +95,6 @@ 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; @@ -106,10 +105,9 @@ use crate::stats::session::StatsSession; pub mod aggregate_fn; #[doc(hidden)] pub mod aliases; -mod arc_swap_map; +pub mod arc_swap_map; mod array; pub mod arrays; -pub mod arrow; pub mod buffer; pub mod builders; pub mod builtins; @@ -144,7 +142,6 @@ pub mod stream; #[cfg(any(test, feature = "_test-harness"))] pub mod test_harness; pub mod validity; -pub mod variants; pub mod flatbuffers { //! Re-exported autogenerated code from the core Vortex flatbuffer definitions. @@ -164,11 +161,13 @@ pub fn initialize(session: &VortexSession) { /// 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. +/// 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. +/// register file, layout, or runtime state — those live in higher-level crates. Arrow +/// interoperability lives in the `vortex-arrow` crate, whose session variables are registered +/// lazily on first use (or eagerly by consumers). pub fn array_session() -> VortexSession { VortexSession::empty() .with::() @@ -177,13 +176,18 @@ pub fn array_session() -> VortexSession { .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(array_session); +/// Returns the hidden global [`VortexSession`] used as a back-compat shim for code paths that +/// cannot yet thread an explicit session through. +#[inline] +pub fn legacy_session() -> &'static VortexSession { + static LEGACY_SESSION: LazyLock = LazyLock::new(array_session); + &LEGACY_SESSION +} pub type ArrayContext = Context; diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 2df53a8e7f9..641efd61404 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -23,7 +23,6 @@ use crate::ArrayRef; use crate::ArraySlots; use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; @@ -36,12 +35,12 @@ use crate::dtype::Nullability; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType; use crate::dtype::UnsignedPType; -use crate::match_each_integer_ptype; +use crate::legacy_session; use crate::match_each_unsigned_integer_ptype; -use crate::scalar::PValue; use crate::scalar::Scalar; use crate::search_sorted::SearchResult; use crate::search_sorted::SearchSorted; +use crate::search_sorted::SearchSortedPrimitiveArray; use crate::search_sorted::SearchSortedSide; use crate::validity::Validity; @@ -238,6 +237,7 @@ pub struct Patches { } impl Patches { + #[allow(clippy::disallowed_methods)] pub fn new( array_len: usize, offset: usize, @@ -266,7 +266,7 @@ impl Patches { if indices.is_host() && values.is_host() { let max = usize::try_from(&indices.execute_scalar( indices.len() - 1, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut legacy_session().create_execution_ctx(), )?) .map_err(|_| vortex_err!("indices must be a number"))?; vortex_ensure!( @@ -276,9 +276,8 @@ impl Patches { #[cfg(debug_assertions)] { - use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_sorted::is_sorted; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = legacy_session().create_execution_ctx(); assert!( is_sorted(&indices, &mut ctx).unwrap_or(false), "Patch indices must be sorted" @@ -381,13 +380,14 @@ impl Patches { } #[inline] + #[allow(clippy::disallowed_methods)] pub fn chunk_offset_at(&self, idx: usize) -> VortexResult { let Some(chunk_offsets) = &self.chunk_offsets else { vortex_bail!("chunk_offsets must be set to retrieve offset at index") }; chunk_offsets - .execute_scalar(idx, &mut LEGACY_SESSION.create_execution_ctx())? + .execute_scalar(idx, &mut legacy_session().create_execution_ctx())? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("chunk offset does not fit in usize")) @@ -441,12 +441,13 @@ impl Patches { } /// Get the patched value at a given index if it exists. + #[allow(clippy::disallowed_methods)] pub fn get_patched(&self, index: usize) -> VortexResult> { self.search_index(index)? .to_found() .map(|patch_idx| { self.values() - .execute_scalar(patch_idx, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(patch_idx, &mut legacy_session().create_execution_ctx()) }) .transpose() } @@ -473,32 +474,7 @@ impl Patches { return self.search_index_chunked(index); } - Self::search_index_binary_search(&self.indices, index + self.offset) - } - - /// Binary searches for `needle` in the indices array. - /// - /// # Returns - /// [`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 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. - // - // The needle is a non-negative integer (a usize); therefore, it must be larger - // than all values in this array. - return Ok(SearchResult::NotFound(primitive.len())); - }; - return primitive - .as_slice::() - .search_sorted(&needle, SearchSortedSide::Left); - }); - } - indices - .as_primitive_typed() - .search_sorted(&PValue::U64(needle as u64), SearchSortedSide::Left) + search_index_binary_search(&self.indices, index + self.offset) } /// Constant time searches for `index` in the indices array. @@ -546,7 +522,7 @@ impl Patches { }; let chunk_indices = self.indices.slice(patches_start_idx..patches_end_idx)?; - let result = Self::search_index_binary_search(&chunk_indices, index + self.offset)?; + let result = search_index_binary_search(&chunk_indices, index + self.offset)?; Ok(match result { SearchResult::Found(idx) => SearchResult::Found(patches_start_idx + idx), @@ -625,10 +601,11 @@ impl Patches { } /// Returns the minimum patch index + #[allow(clippy::disallowed_methods)] pub fn min_index(&self) -> VortexResult { let first = self .indices - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())? + .execute_scalar(0, &mut legacy_session().create_execution_ctx())? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("index does not fit in usize"))?; @@ -636,12 +613,13 @@ impl Patches { } /// Returns the maximum patch index + #[allow(clippy::disallowed_methods)] pub fn max_index(&self) -> VortexResult { let last = self .indices .execute_scalar( self.indices.len() - 1, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut legacy_session().create_execution_ctx(), )? .as_primitive() .as_::() @@ -743,6 +721,7 @@ impl Patches { } /// Slice the patches by a range of the patched array. + #[allow(clippy::disallowed_methods)] pub fn slice(&self, range: Range) -> VortexResult> { let slice_start_idx = self.search_index(range.start)?.to_index(); let slice_end_idx = self.search_index(range.end)?.to_index(); @@ -769,7 +748,7 @@ impl Patches { .as_ref() .map(|new_chunk_offsets| -> VortexResult { let new_chunk_base = new_chunk_offsets - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())? + .execute_scalar(0, &mut legacy_session().create_execution_ctx())? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("chunk offset does not fit in usize"))?; @@ -1004,6 +983,41 @@ impl Patches { } } +/// Binary searches for `needle` in the indices array. +/// +/// # Returns +/// [`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 let Some(primitive) = indices.as_opt::() { + match_each_unsigned_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. + // + // The needle is a non-negative integer (a usize); therefore, it must be larger + // than all values in this array. + return Ok(SearchResult::NotFound(primitive.len())); + }; + return primitive + .as_slice::() + .search_sorted(&needle, SearchSortedSide::Left); + }); + } + + search_index_binary_search_scalar(indices, needle) +} + +#[allow(clippy::disallowed_methods)] +fn search_index_binary_search_scalar( + indices: &ArrayRef, + needle: usize, +) -> VortexResult { + match_each_unsigned_integer_ptype!(indices.dtype().as_ptype(), |T| { + SearchSortedPrimitiveArray::::new(indices, &mut legacy_session().create_execution_ctx()) + .search_sorted(&needle, SearchSortedSide::Left) + }) +} + #[expect(clippy::too_many_arguments)] // private function, can clean up one day fn take_map, T: NativePType>( indices: &[I], @@ -1214,8 +1228,6 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::assert_arrays_eq; @@ -1273,10 +1285,16 @@ mod test { ) .unwrap() .unwrap(); - #[expect(deprecated)] - let primitive_values = taken.values().to_primitive(); - #[expect(deprecated)] - let primitive_indices = taken.indices().to_primitive(); + let primitive_values = taken + .values() + .clone() + .execute::(&mut ctx) + .unwrap(); + let primitive_indices = taken + .indices() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!(taken.array_len(), 2); assert_arrays_eq!( primitive_values, @@ -1320,8 +1338,11 @@ mod test { .unwrap() .unwrap(); - #[expect(deprecated)] - let primitive_values = taken.values().to_primitive(); + let primitive_values = taken + .values() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!(taken.array_len(), 2); assert_arrays_eq!( primitive_values, @@ -1634,8 +1655,11 @@ mod test { ); // Values should be the null and 300 - #[expect(deprecated)] - let masked_values = masked.values().to_primitive(); + let masked_values = masked + .values() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!(masked_values.len(), 2); 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 @@ -1830,8 +1854,11 @@ mod test { ) .unwrap(); - #[expect(deprecated)] - let values = patches.values().to_primitive(); + let values = patches + .values() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!( i32::try_from(&values.execute_scalar(0, &mut ctx).unwrap()).unwrap(), 100i32 diff --git a/vortex-array/src/scalar/mod.rs b/vortex-array/src/scalar/mod.rs index ff8d2c6134b..dbc1e17d04c 100644 --- a/vortex-array/src/scalar/mod.rs +++ b/vortex-array/src/scalar/mod.rs @@ -15,7 +15,6 @@ #[cfg(feature = "arbitrary")] pub mod arbitrary; -mod arrow; mod cast; mod constructor; diff --git a/vortex-array/src/scalar/proto.rs b/vortex-array/src/scalar/proto.rs index ebb15809abc..f172d5491f3 100644 --- a/vortex-array/src/scalar/proto.rs +++ b/vortex-array/src/scalar/proto.rs @@ -441,9 +441,13 @@ fn list_from_proto( dtype: &DType, session: &VortexSession, ) -> VortexResult { - let element_dtype = dtype - .as_list_element_opt() - .ok_or_else(|| vortex_err!(Serde: "expected List dtype for ListValue, got {dtype}"))?; + let element_dtype = match dtype { + DType::List(edt, _) => edt, + DType::FixedSizeList(edt, ..) => edt, + _ => { + vortex_bail!(Serde: "expected List or FixedSizeList dtype for ListValue, got {dtype}") + } + }; let mut values = Vec::with_capacity(v.values.len()); for elem in v.values.iter() { diff --git a/vortex-array/src/scalar/tests/round_trip.rs b/vortex-array/src/scalar/tests/round_trip.rs index 7d434a95796..ab315483d72 100644 --- a/vortex-array/src/scalar/tests/round_trip.rs +++ b/vortex-array/src/scalar/tests/round_trip.rs @@ -254,6 +254,25 @@ mod tests { assert_eq!(extracted.as_slice(), &large_binary); } + // Test that fixed-size list scalars round-trip through protobuf. This is + // the path taken by a ConstantArray of a fixed-size list dtype (e.g. a + // UUID column stored as fixed_size_list(u8)[16]), whose scalar metadata + // is serialized as a protobuf ListValue. + #[test] + fn test_protobuf_fixed_size_list_round_trip() { + let fsl = Scalar::fixed_size_list( + Arc::new(DType::Primitive(PType::U8, Nullability::NonNullable)), + (0u8..16) + .map(|i| Scalar::primitive(i, Nullability::NonNullable)) + .collect(), + Nullability::NonNullable, + ); + + let pb_fsl = pb::Scalar::from(&fsl); + let round_tripped = Scalar::from_proto(&pb_fsl, &SESSION).unwrap(); + assert_eq!(fsl, round_tripped); + } + // Test that nullable and non-nullable types are preserved #[test] fn test_nullability_preservation() { diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index 2151a310dde..01c4fe06d31 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -153,7 +153,7 @@ fn between_canonical( upper.clone(), Operator::from(options.upper_strict.to_compare_operator()), )?; - execute_boolean(&lower_cmp, &upper_cmp, Operator::And, ctx) + execute_boolean(lower_cmp, upper_cmp, Operator::And, ctx) } /// An optimized scalar expression to compute whether values fall between two bounds. diff --git a/vortex-array/src/scalar_fn/fns/binary/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/boolean.rs index b7f838c6d13..550b6c2c701 100644 --- a/vortex-array/src/scalar_fn/fns/binary/boolean.rs +++ b/vortex-array/src/scalar_fn/fns/binary/boolean.rs @@ -3,7 +3,6 @@ use std::iter::repeat_n; -use arrow_array::cast::AsArray; use vortex_buffer::BitBuffer; use vortex_buffer::BufferMut; use vortex_buffer::read_u64_le; @@ -27,8 +26,6 @@ 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::dtype::Nullability; @@ -116,67 +113,32 @@ pub fn or_kleene(lhs: &ArrayRef, rhs: &ArrayRef) -> VortexResult { /// This is the entry point for boolean operations from the binary expression. /// Handles constants and canonical boolean arrays directly, otherwise falls back to Arrow. pub(crate) fn execute_boolean( - lhs: &ArrayRef, - rhs: &ArrayRef, + lhs: ArrayRef, + rhs: ArrayRef, op: Operator, ctx: &mut ExecutionCtx, ) -> VortexResult { - let nullable = boolean_nullability(lhs, rhs); + 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)? { + 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)? - { + let lhs = lhs.execute::(ctx)?; + if let Some(result) = ::boolean(lhs.as_view(), &rhs, op, ctx)? { return Ok(result); } - arrow_execute_boolean(lhs.clone(), rhs.clone(), op, ctx) -} - -/// Arrow implementation for Kleene boolean operations using [`Operator`]. -fn arrow_execute_boolean( - lhs: ArrayRef, - rhs: ArrayRef, - op: Operator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let nullable = boolean_nullability(&lhs, &rhs); - let session = ctx.session().clone(); - - let lhs = session - .arrow() - .execute_arrow(lhs, None, ctx)? - .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_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 => vortex_bail!("Not a boolean operator: {other}"), + let rhs = rhs.execute::(ctx)?; + let Some(result) = ::boolean(rhs.as_view(), &lhs.into_array(), op, ctx)? + else { + vortex_bail!("No boolean kernel for two BoolArrays"); }; - - ArrayRef::from_arrow(&array, nullable == Nullability::Nullable) + Ok(result) } /// Handles boolean operations where at least one operand is a constant array. @@ -763,8 +725,6 @@ mod tests { 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; @@ -864,9 +824,9 @@ mod tests { BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(), )] fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) { + let mut ctx = array_session().create_execution_ctx(); let r = lhs.binary(rhs, Operator::Or).unwrap(); - #[expect(deprecated)] - let r = r.to_bool().into_array(); + let r = r.execute::(&mut ctx).unwrap().into_array(); let v0 = r .execute_scalar(0, &mut array_session().create_execution_ctx()) @@ -905,11 +865,12 @@ mod tests { BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(), )] fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) { - #[expect(deprecated)] + let mut ctx = array_session().create_execution_ctx(); let r = lhs .binary(rhs, Operator::And) .unwrap() - .to_bool() + .execute::(&mut ctx) + .unwrap() .into_array(); let v0 = r diff --git a/vortex-array/src/scalar_fn/fns/binary/compare.rs b/vortex-array/src/scalar_fn/fns/binary/compare.rs deleted file mode 100644 index b7a53c21d62..00000000000 --- a/vortex-array/src/scalar_fn/fns/binary/compare.rs +++ /dev/null @@ -1,590 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::cmp::Ordering; - -use arrow_array::BooleanArray; -use arrow_buffer::NullBuffer; -use arrow_ord::cmp; -use arrow_ord::ord::make_comparator; -use arrow_schema::Field; -use arrow_schema::SortOptions; -use vortex_error::VortexResult; -use vortex_error::vortex_err; - -use crate::ArrayRef; -use crate::Canonical; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::array::VTable; -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::Datum; -use crate::arrow::from_arrow_columnar; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::kernel::ExecuteParentKernel; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::operators::CompareOperator; - -/// Trait for encoding-specific comparison kernels that operate in encoded space. -/// -/// Implementations can compare an encoded array against another array (typically a constant) -/// without first decompressing. The adaptor normalizes operand order so `array` is always -/// the left-hand side, swapping the operator when necessary. -pub trait CompareKernel: VTable { - fn compare( - lhs: ArrayView<'_, Self>, - rhs: &ArrayRef, - operator: CompareOperator, - ctx: &mut ExecutionCtx, - ) -> VortexResult>; -} - -/// Adaptor that bridges [`CompareKernel`] implementations to [`ExecuteParentKernel`]. -/// -/// When a `ScalarFnArray(Binary, cmp_op)` wraps a child that implements `CompareKernel`, -/// this adaptor extracts the comparison operator and other operand, normalizes operand order -/// (swapping the operator if the encoded array is on the RHS), and delegates to the kernel. -#[derive(Default, Debug)] -pub struct CompareExecuteAdaptor(pub V); - -impl ExecuteParentKernel for CompareExecuteAdaptor -where - V: CompareKernel, -{ - type Parent = ExactScalarFn; - - fn execute_parent( - &self, - array: ArrayView<'_, V>, - parent: ScalarFnArrayView<'_, Binary>, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - // Only handle comparison operators - let Ok(cmp_op) = CompareOperator::try_from(*parent.options) else { - return Ok(None); - }; - - // Get the ScalarFnArray to access children - let Some(scalar_fn_array) = parent.as_opt::() else { - return Ok(None); - }; - // Normalize so `array` is always LHS, swapping the operator if needed - // TODO(joe): should be go this here or in the Rule/Kernel - let (cmp_op, other) = match child_idx { - 0 => (cmp_op, scalar_fn_array.get_child(1)), - 1 => (cmp_op.swap(), scalar_fn_array.get_child(0)), - _ => return Ok(None), - }; - - let len = array.len(); - let nullable = array.dtype().is_nullable() || other.dtype().is_nullable(); - - // Empty array → empty bool result - if len == 0 { - return Ok(Some( - Canonical::empty(&DType::Bool(nullable.into())).into_array(), - )); - } - - // Null constant on either side → all-null bool result - if other.as_constant().is_some_and(|s| s.is_null()) { - return Ok(Some( - ConstantArray::new(Scalar::null(DType::Bool(nullable.into())), len).into_array(), - )); - } - - V::compare(array, other, cmp_op, ctx) - } -} - -/// Execute a compare operation between two arrays. -/// -/// This is the entry point for compare operations from the binary expression. -/// Handles empty, constant-null, and constant-constant directly, otherwise falls back to Arrow. -pub(crate) fn execute_compare( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: CompareOperator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable(); - - if lhs.is_empty() { - return Ok(Canonical::empty(&DType::Bool(nullable.into())).into_array()); - } - - let left_constant_null = lhs.as_constant().map(|l| l.is_null()).unwrap_or(false); - let right_constant_null = rhs.as_constant().map(|r| r.is_null()).unwrap_or(false); - if left_constant_null || right_constant_null { - return Ok( - ConstantArray::new(Scalar::null(DType::Bool(nullable.into())), lhs.len()).into_array(), - ); - } - - // Constant-constant fast path - if let (Some(lhs_const), Some(rhs_const)) = (lhs.as_opt::(), rhs.as_opt::()) - { - let result = scalar_cmp(lhs_const.scalar(), rhs_const.scalar(), op)?; - return Ok(ConstantArray::new(result, lhs.len()).into_array()); - } - - arrow_compare_arrays(lhs, rhs, op, ctx) -} - -/// Fall back to Arrow for comparison. -fn arrow_compare_arrays( - left: &ArrayRef, - right: &ArrayRef, - operator: CompareOperator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - assert_eq!(left.len(), right.len()); - - let nullable = left.dtype().is_nullable() || right.dtype().is_nullable(); - - // Arrow's vectorized comparison kernels don't support nested types. - // For nested types, fall back to `make_comparator` which does element-wise comparison. - let arrow_array: BooleanArray = if left.dtype().is_nested() || right.dtype().is_nested() { - let session = ctx.session().clone(); - let lhs = session.arrow().execute_arrow(left.clone(), None, ctx)?; - let target_field = Field::new("", lhs.data_type().clone(), right.dtype().is_nullable()); - let rhs = session - .arrow() - .execute_arrow(right.clone(), Some(&target_field), ctx)?; - - compare_nested_arrow_arrays(lhs.as_ref(), rhs.as_ref(), operator)? - } else { - // Fast path: use vectorized kernels for primitive types. - let lhs = Datum::try_new(left, ctx)?; - let rhs = Datum::try_new_with_target_datatype(right, lhs.data_type(), ctx)?; - - match operator { - CompareOperator::Eq => cmp::eq(&lhs, &rhs)?, - CompareOperator::NotEq => cmp::neq(&lhs, &rhs)?, - CompareOperator::Gt => cmp::gt(&lhs, &rhs)?, - CompareOperator::Gte => cmp::gt_eq(&lhs, &rhs)?, - CompareOperator::Lt => cmp::lt(&lhs, &rhs)?, - CompareOperator::Lte => cmp::lt_eq(&lhs, &rhs)?, - } - }; - - from_arrow_columnar(&arrow_array, left.len(), nullable, ctx) -} - -pub fn scalar_cmp(lhs: &Scalar, rhs: &Scalar, operator: CompareOperator) -> VortexResult { - if lhs.is_null() | rhs.is_null() { - return Ok(Scalar::null(DType::Bool(Nullability::Nullable))); - } - - let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); - - // We use `partial_cmp` to ensure we do not lose a type mismatch error. - let ordering = lhs.partial_cmp(rhs).ok_or_else(|| { - vortex_err!( - "Cannot compare scalars with incompatible types: {} and {}", - lhs.dtype(), - rhs.dtype() - ) - })?; - - let b = match operator { - CompareOperator::Eq => ordering.is_eq(), - CompareOperator::NotEq => ordering.is_ne(), - CompareOperator::Gt => ordering.is_gt(), - CompareOperator::Gte => ordering.is_ge(), - CompareOperator::Lt => ordering.is_lt(), - CompareOperator::Lte => ordering.is_le(), - }; - - Ok(Scalar::bool(b, nullability)) -} - -/// Compare two Arrow arrays element-wise using [`make_comparator`]. -/// -/// This function is required for nested types (Struct, List, FixedSizeList) because Arrow's -/// vectorized comparison kernels ([`cmp::eq`], [`cmp::neq`], etc.) do not support them. -/// -/// The vectorized kernels are faster but only work on primitive types, so for non-nested types, -/// prefer using the vectorized kernels directly for better performance. -pub fn compare_nested_arrow_arrays( - lhs: &dyn arrow_array::Array, - rhs: &dyn arrow_array::Array, - operator: CompareOperator, -) -> VortexResult { - let compare_arrays_at = make_comparator(lhs, rhs, SortOptions::default())?; - - let cmp_fn = match operator { - CompareOperator::Eq => Ordering::is_eq, - CompareOperator::NotEq => Ordering::is_ne, - CompareOperator::Gt => Ordering::is_gt, - CompareOperator::Gte => Ordering::is_ge, - CompareOperator::Lt => Ordering::is_lt, - CompareOperator::Lte => Ordering::is_le, - }; - - let values = (0..lhs.len()) - .map(|i| cmp_fn(compare_arrays_at(i, i))) - .collect(); - let nulls = NullBuffer::union(lhs.nulls(), rhs.nulls()); - - Ok(BooleanArray::new(values, nulls)) -} - -#[cfg(test)] -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::VortexSessionExecute; - use crate::array_session; - use crate::arrays::BoolArray; - use crate::arrays::ListArray; - use crate::arrays::ListViewArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::StructArray; - use crate::arrays::VarBinArray; - use crate::arrays::VarBinViewArray; - use crate::assert_arrays_eq; - use crate::builtins::ArrayBuiltins; - use crate::dtype::DType; - use crate::dtype::FieldName; - use crate::dtype::FieldNames; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::extension::datetime::TimeUnit; - use crate::extension::datetime::Timestamp; - use crate::extension::datetime::TimestampOptions; - use crate::scalar::Scalar; - use crate::scalar_fn::fns::binary::compare::ConstantArray; - use crate::scalar_fn::fns::binary::scalar_cmp; - use crate::scalar_fn::fns::operators::CompareOperator; - use crate::scalar_fn::fns::operators::Operator; - use crate::test_harness::to_int_indices; - use crate::validity::Validity; - - #[test] - fn test_bool_basic_comparisons() { - 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]), - ); - - let matches = arr - .clone() - .into_array() - .binary(arr.clone().into_array(), Operator::Eq) - .unwrap() - .execute::(ctx) - .vortex_expect("must be a bool array"); - assert_eq!(to_int_indices(matches, ctx).unwrap(), [1u64, 2, 3, 4]); - - let matches = arr - .clone() - .into_array() - .binary(arr.clone().into_array(), Operator::NotEq) - .unwrap() - .execute::(ctx) - .vortex_expect("must be a bool array"); - let empty: [u64; 0] = []; - 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]), - ); - - let matches = arr - .clone() - .into_array() - .binary(other.clone().into_array(), Operator::Lte) - .unwrap() - .execute::(ctx) - .vortex_expect("must be a bool array"); - assert_eq!(to_int_indices(matches, ctx).unwrap(), [2u64, 3, 4]); - - let matches = arr - .clone() - .into_array() - .binary(other.clone().into_array(), Operator::Lt) - .unwrap() - .execute::(ctx) - .vortex_expect("must be a bool array"); - assert_eq!(to_int_indices(matches, ctx).unwrap(), [4u64]); - - let matches = other - .clone() - .into_array() - .binary(arr.clone().into_array(), Operator::Gte) - .unwrap() - .execute::(ctx) - .vortex_expect("must be a bool array"); - assert_eq!(to_int_indices(matches, ctx).unwrap(), [2u64, 3, 4]); - - let matches = other - .into_array() - .binary(arr.into_array(), Operator::Gt) - .unwrap() - .execute::(ctx) - .vortex_expect("must be a bool array"); - assert_eq!(to_int_indices(matches, ctx).unwrap(), [4u64]); - } - - #[test] - fn constant_compare() { - let left = ConstantArray::new(Scalar::from(2u32), 10); - let right = ConstantArray::new(Scalar::from(10u32), 10); - - let result = left - .into_array() - .binary(right.into_array(), Operator::Gt) - .unwrap(); - assert_eq!(result.len(), 10); - let scalar = result - .execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(); - assert_eq!(scalar.as_bool().value(), Some(false)); - } - - #[rstest] - #[case(VarBinArray::from(vec!["a", "b"]).into_array(), VarBinViewArray::from_iter_str(["a", "b"]).into_array())] - #[case(VarBinViewArray::from_iter_str(["a", "b"]).into_array(), VarBinArray::from(vec!["a", "b"]).into_array())] - #[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, &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( - values1.into_array(), - offsets1.into_array(), - Validity::NonNullable, - ) - .unwrap(); - - let values2 = PrimitiveArray::from_iter([1i32, 2, 3, 4, 7, 8]); - let offsets2 = PrimitiveArray::from_iter([0i32, 2, 4, 6]); - let list2 = ListArray::try_new( - values2.into_array(), - offsets2.into_array(), - Validity::NonNullable, - ) - .unwrap(); - - let result = list1 - .clone() - .into_array() - .binary(list2.clone().into_array(), Operator::Eq) - .unwrap(); - let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected, &mut ctx); - - let result = list1 - .clone() - .into_array() - .binary(list2.clone().into_array(), Operator::NotEq) - .unwrap(); - let expected = BoolArray::from_iter([false, false, true]); - 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, &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( - values.into_array(), - offsets.into_array(), - Validity::NonNullable, - ) - .unwrap(); - - let list_scalar = Scalar::list( - Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), - vec![3i32.into(), 4i32.into()], - Nullability::NonNullable, - ); - let constant = ConstantArray::new(list_scalar, 3); - - let result = list - .into_array() - .binary(constant.into_array(), Operator::Eq) - .unwrap(); - let expected = BoolArray::from_iter([false, true, false]); - 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]); - - let bool_field2 = BoolArray::from_iter([Some(true), Some(false), Some(false)]); - let int_field2 = PrimitiveArray::from_iter([1i32, 2, 4]); - - let struct1 = StructArray::from_fields(&[ - ("bool_col", bool_field1.into_array()), - ("int_col", int_field1.into_array()), - ]) - .unwrap(); - - let struct2 = StructArray::from_fields(&[ - ("bool_col", bool_field2.into_array()), - ("int_col", int_field2.into_array()), - ]) - .unwrap(); - - let result = struct1 - .clone() - .into_array() - .binary(struct2.clone().into_array(), Operator::Eq) - .unwrap(); - let expected = BoolArray::from_iter([true, true, false]); - 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, &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(), - 5, - Validity::NonNullable, - ) - .unwrap(); - - let empty2 = StructArray::try_new( - FieldNames::from(Vec::::new()), - Vec::new(), - 5, - Validity::NonNullable, - ) - .unwrap(); - - let result = empty1 - .into_array() - .binary(empty2.into_array(), Operator::Eq) - .unwrap(); - let expected = BoolArray::from_iter([true, true, true, true, true]); - 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(), - "banana".as_bytes(), - "cherry".as_bytes(), - ]); - let struct1 = StructArray::from_fields(&[("data", bin_field1.into_array())]).unwrap(); - - // RHS: struct with a VarBinViewArray (view-based) binary field — same logical DType - let bin_field2 = VarBinViewArray::from_iter_bin([ - "apple".as_bytes(), - "banana".as_bytes(), - "durian".as_bytes(), - ]); - let struct2 = StructArray::from_fields(&[("data", bin_field2.into_array())]).unwrap(); - - let result = struct1 - .into_array() - .binary(struct2.into_array(), Operator::Eq) - .unwrap(); - let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected, &mut ctx); - } - - /// Regression test: `scalar_cmp` must error when comparing scalars with incompatible - /// extension types (e.g., timestamps with different time units) rather than silently - /// returning a wrong result. - #[test] - fn scalar_cmp_incompatible_extension_types_errors() { - let ms_scalar = Scalar::extension::( - TimestampOptions { - unit: TimeUnit::Milliseconds, - tz: None, - }, - Scalar::from(1704067200000i64), - ); - let s_scalar = Scalar::extension::( - TimestampOptions { - unit: TimeUnit::Seconds, - tz: None, - }, - Scalar::from(1704067200i64), - ); - - // Ordering comparisons must error on incompatible types. - assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Gt).is_err()); - assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Lt).is_err()); - assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Gte).is_err()); - assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Lte).is_err()); - assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Eq).is_err()); - assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::NotEq).is_err()); - } - - #[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(), - buffer![0i32, 0i32, 0i32].into_array(), - Validity::AllValid, - ); - - let result = list - .clone() - .into_array() - .binary(list.into_array(), Operator::Eq) - .unwrap(); - 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/compare/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/compare/boolean.rs new file mode 100644 index 00000000000..15bd486d341 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/boolean.rs @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native comparison of boolean arrays using word-wise bit operations. + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Constant; +use crate::dtype::Nullability; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::operators::CompareOperator; +use crate::validity::Validity; + +enum BoolOperand { + Array { bits: BitBuffer, validity: Validity }, + Constant { value: bool, validity: Validity }, +} + +impl BoolOperand { + fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(constant) = array.as_opt::() { + let value = constant + .scalar() + .as_bool_opt() + .ok_or_else(|| vortex_err!("expected boolean scalar"))? + .value() + .ok_or_else(|| vortex_err!("null constant handled by execute_compare"))?; + return Ok(Self::Constant { + value, + validity: if constant.scalar().dtype().is_nullable() { + Validity::AllValid + } else { + Validity::NonNullable + }, + }); + } + + let array = array.clone().execute::(ctx)?; + let validity = array.validity()?; + Ok(Self::Array { + bits: array.into_bit_buffer(), + validity, + }) + } + + fn validity(&self) -> Validity { + match self { + Self::Array { validity, .. } | Self::Constant { validity, .. } => validity.clone(), + } + } +} + +/// Compare two boolean arrays. +/// +/// Values compare as `false < true`; every operator reduces to at most two word-wise passes over +/// the value bit buffers. +pub(super) fn compare_bool( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let lhs = BoolOperand::try_new(lhs, ctx)?; + let rhs = BoolOperand::try_new(rhs, ctx)?; + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + + let bits = match (lhs, rhs) { + (BoolOperand::Array { bits: l, .. }, BoolOperand::Array { bits: r, .. }) => { + compare_bits(l, r, op) + } + (BoolOperand::Array { bits, .. }, BoolOperand::Constant { value, .. }) => { + compare_bits_constant(bits, value, op) + } + (BoolOperand::Constant { value, .. }, BoolOperand::Array { bits, .. }) => { + compare_bits_constant(bits, value, op.swap()) + } + (BoolOperand::Constant { value: l, .. }, BoolOperand::Constant { value: r, .. }) => { + // Unreachable through `execute_compare` (constant-constant is folded there), but + // cheap to answer anyway. + let result = super::ordering_predicate(op)(l.cmp(&r)); + BitBuffer::full(result, len) + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +fn compare_bits(lhs: BitBuffer, rhs: BitBuffer, op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => !(lhs ^ &rhs), + CompareOperator::NotEq => lhs ^ &rhs, + // a < b ⟺ !a & b + CompareOperator::Lt => rhs.into_bitand_not(&lhs), + // a <= b ⟺ !(a & !b) + CompareOperator::Lte => !lhs.into_bitand_not(&rhs), + // a > b ⟺ a & !b + CompareOperator::Gt => lhs.into_bitand_not(&rhs), + // a >= b ⟺ !(!a & b) + CompareOperator::Gte => !rhs.into_bitand_not(&lhs), + } +} + +/// Compare array bits against a non-null constant: `bits value`. +fn compare_bits_constant(bits: BitBuffer, value: bool, op: CompareOperator) -> BitBuffer { + let len = bits.len(); + match (op, value) { + (CompareOperator::Eq, true) + | (CompareOperator::NotEq, false) + | (CompareOperator::Gt, false) + | (CompareOperator::Gte, true) => bits, + (CompareOperator::Eq, false) + | (CompareOperator::NotEq, true) + | (CompareOperator::Lt, true) + | (CompareOperator::Lte, false) => !bits, + (CompareOperator::Lt, false) | (CompareOperator::Gt, true) => BitBuffer::new_unset(len), + (CompareOperator::Lte, true) | (CompareOperator::Gte, false) => BitBuffer::new_set(len), + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs b/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs new file mode 100644 index 00000000000..70c56411d08 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs @@ -0,0 +1,400 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native comparison of UTF-8 and binary arrays over canonical [`VarBinViewArray`]s. +//! +//! Equality first compares the leading 8 bytes of each view (length plus 4-byte prefix), which +//! answers most lanes without touching the data buffers. Ordering compares the inline 4-byte +//! prefixes first and only dereferences the full value on a prefix tie. UTF-8 values compare by +//! their byte representation, which matches code-point order. + +use std::cmp::Ordering; + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Constant; +use crate::arrays::VarBinViewArray; +use crate::arrays::varbinview::BinaryView; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::binary::compare::ordering_predicate; +use crate::scalar_fn::fns::operators::CompareOperator; +use crate::validity::Validity; + +enum BytesOperand { + Array { + values: VarBinViewArray, + validity: Validity, + }, + Constant { + value: Vec, + validity: Validity, + }, +} + +impl BytesOperand { + fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(constant) = array.as_opt::() { + let value = constant_bytes(constant.scalar())?; + return Ok(Self::Constant { + value, + validity: if constant.scalar().dtype().is_nullable() { + Validity::AllValid + } else { + Validity::NonNullable + }, + }); + } + + let values = array.clone().execute::(ctx)?; + let validity = values.validity()?; + Ok(Self::Array { values, validity }) + } + + fn validity(&self) -> Validity { + match self { + Self::Array { validity, .. } | Self::Constant { validity, .. } => validity.clone(), + } + } +} + +fn constant_bytes(scalar: &Scalar) -> VortexResult> { + let value = match scalar.dtype() { + DType::Utf8(_) => scalar + .as_utf8() + .value() + .map(|s| s.as_str().as_bytes().to_vec()), + DType::Binary(_) => scalar.as_binary().value().map(|b| b.to_vec()), + _ => vortex_bail!("expected utf8 or binary scalar, got {}", scalar.dtype()), + }; + value.ok_or_else(|| vortex_err!("null constant handled by execute_compare")) +} + +/// A resolved view over a canonical [`VarBinViewArray`]: the view structs plus borrowed slices of +/// every data buffer, supporting cheap per-lane byte access. +struct ViewsSide<'a> { + views: &'a [BinaryView], + buffers: Vec<&'a [u8]>, +} + +impl<'a> ViewsSide<'a> { + fn new(array: &'a VarBinViewArray) -> Self { + Self { + views: array.views(), + buffers: (0..array.data_buffers().len()) + .map(|idx| array.buffer(idx).as_slice()) + .collect(), + } + } + + fn len(&self) -> usize { + self.views.len() + } + + /// The view at `index` without a bounds check. + /// + /// # Safety + /// + /// `index` must be strictly less than `self.len()`. + #[inline] + unsafe fn view_unchecked(&self, index: usize) -> &'a BinaryView { + // SAFETY: caller guarantees index < self.views.len(). + unsafe { self.views.get_unchecked(index) } + } + + /// The full bytes of `view`, which must belong to this side. + #[inline] + fn view_bytes(&self, view: &'a BinaryView) -> &'a [u8] { + if view.is_inlined() { + view.as_inlined().value() + } else { + let view = view.as_view(); + &self.buffers[view.buffer_index as usize][view.as_range()] + } + } +} + +/// The leading 8 bytes of a view: the `u32` length plus the first 4 bytes of the value +/// (zero-padded for values shorter than 4 bytes). +#[inline] +#[expect(clippy::cast_possible_truncation, reason = "intentional bit slicing")] +fn view_head(view: &BinaryView) -> u64 { + view.as_u128() as u64 +} + +/// The first 4 value bytes of a view as a big-endian `u32` (zero-padded for values shorter than +/// 4 bytes), so that numeric comparison matches lexicographic byte order. +/// +/// Zero padding preserves lexicographic order: a padded position differs from a real byte only +/// when one value is a strict prefix of the other within the first 4 bytes, and the shorter +/// value orders first exactly as `0` orders before any later real byte. A padded tie falls +/// through to the tail or full byte comparison. +#[inline] +#[expect(clippy::cast_possible_truncation, reason = "intentional bit slicing")] +fn view_prefix(view: &BinaryView) -> u32 { + ((view.as_u128() >> 32) as u32).swap_bytes() +} + +/// Value bytes 4..12 of an *inlined* view as a big-endian `u64` (zero-padded past the value's +/// length). Only meaningful for inlined views: reference views store the buffer index and offset +/// in these bytes. +#[inline] +fn view_tail(view: &BinaryView) -> u64 { + ((view.as_u128() >> 64) as u64).swap_bytes() +} + +/// Compare `lhs_view` from `lhs` against `rhs_view` from `rhs` for equality. +#[inline] +fn view_eq( + lhs: &ViewsSide<'_>, + lhs_view: &BinaryView, + rhs: &ViewsSide<'_>, + rhs_view: &BinaryView, +) -> bool { + if view_head(lhs_view) != view_head(rhs_view) { + return false; + } + if lhs_view.is_inlined() { + // Lengths are equal and at most 12: the whole value lives in the view. + return lhs_view.as_u128() == rhs_view.as_u128(); + } + // Equal lengths above 12 and equal prefixes: compare the out-of-line suffixes. + lhs.view_bytes(lhs_view)[4..] == rhs.view_bytes(rhs_view)[4..] +} + +/// Compare `lhs_view` from `lhs` against `rhs_view` from `rhs`. +#[inline] +fn view_cmp( + lhs: &ViewsSide<'_>, + lhs_view: &BinaryView, + rhs: &ViewsSide<'_>, + rhs_view: &BinaryView, +) -> Ordering { + let lhs_prefix = view_prefix(lhs_view); + let rhs_prefix = view_prefix(rhs_view); + if lhs_prefix != rhs_prefix { + return lhs_prefix.cmp(&rhs_prefix); + } + if lhs_view.is_inlined() && rhs_view.is_inlined() { + // Both values live entirely in their views: compare the remaining 8 (zero-padded) + // value bytes, then lengths. A tie on padded windows means the shorter value is a + // prefix of the longer one. + let lhs_tail = view_tail(lhs_view); + let rhs_tail = view_tail(rhs_view); + if lhs_tail != rhs_tail { + return lhs_tail.cmp(&rhs_tail); + } + return lhs_view.len().cmp(&rhs_view.len()); + } + lhs.view_bytes(lhs_view).cmp(rhs.view_bytes(rhs_view)) +} + +/// Compare two UTF-8 or binary arrays. +pub(super) fn compare_bytes( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let lhs = BytesOperand::try_new(lhs, ctx)?; + let rhs = BytesOperand::try_new(rhs, ctx)?; + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + + let bits = match (&lhs, &rhs) { + (BytesOperand::Array { values: l, .. }, BytesOperand::Array { values: r, .. }) => { + compare_views(&ViewsSide::new(l), &ViewsSide::new(r), op) + } + (BytesOperand::Array { values, .. }, BytesOperand::Constant { value, .. }) => { + compare_views_constant(&ViewsSide::new(values), value, op) + } + (BytesOperand::Constant { value, .. }, BytesOperand::Array { values, .. }) => { + compare_views_constant(&ViewsSide::new(values), value, op.swap()) + } + (BytesOperand::Constant { value: l, .. }, BytesOperand::Constant { value: r, .. }) => { + // Unreachable through `execute_compare` (constant-constant is folded there), but + // cheap to answer anyway. + BitBuffer::full(ordering_predicate(op)(l.as_slice().cmp(r.as_slice())), len) + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +fn compare_views(lhs: &ViewsSide<'_>, rhs: &ViewsSide<'_>, op: CompareOperator) -> BitBuffer { + let len = lhs.len(); + // The unchecked view accesses below index both sides with `i < len`, so this must hold even + // in release builds. + assert_eq!(len, rhs.len(), "compared views must have equal lengths"); + // Dispatch the operator outside the lane loop so each predicate inlines into its own loop; + // a shared `fn(Ordering) -> bool` pointer would cost an indirect call per lane. + match op { + CompareOperator::Eq => BitBuffer::collect_bool(len, |i| { + // SAFETY: `collect_bool` yields i < len == views.len() for both sides. + unsafe { view_eq(lhs, lhs.view_unchecked(i), rhs, rhs.view_unchecked(i)) } + }), + CompareOperator::NotEq => BitBuffer::collect_bool(len, |i| { + // SAFETY: `collect_bool` yields i < len == views.len() for both sides. + unsafe { !view_eq(lhs, lhs.view_unchecked(i), rhs, rhs.view_unchecked(i)) } + }), + CompareOperator::Gt => collect_ordering_bits(lhs, rhs, Ordering::is_gt), + CompareOperator::Gte => collect_ordering_bits(lhs, rhs, Ordering::is_ge), + CompareOperator::Lt => collect_ordering_bits(lhs, rhs, Ordering::is_lt), + CompareOperator::Lte => collect_ordering_bits(lhs, rhs, Ordering::is_le), + } +} + +/// Bit-pack `predicate(view_cmp(lhs[i], rhs[i]))` over two equal-length view sides. +fn collect_ordering_bits( + lhs: &ViewsSide<'_>, + rhs: &ViewsSide<'_>, + predicate: impl Fn(Ordering) -> bool, +) -> BitBuffer { + let len = lhs.len(); + assert_eq!(len, rhs.len(), "compared views must have equal lengths"); + BitBuffer::collect_bool(len, |i| { + // SAFETY: `collect_bool` yields i < len == views.len() for both sides. + predicate(unsafe { view_cmp(lhs, lhs.view_unchecked(i), rhs, rhs.view_unchecked(i)) }) + }) +} + +fn compare_views_constant(lhs: &ViewsSide<'_>, constant: &[u8], op: CompareOperator) -> BitBuffer { + let len = lhs.len(); + // The same head/prefix/tail words a view stores, precomputed once for the constant. + let mut prefix_bytes = [0u8; 4]; + let prefix_len = constant.len().min(4); + prefix_bytes[..prefix_len].copy_from_slice(&constant[..prefix_len]); + let mut tail_bytes = [0u8; 8]; + if constant.len() > 4 { + let tail_len = (constant.len() - 4).min(8); + tail_bytes[..tail_len].copy_from_slice(&constant[4..4 + tail_len]); + } + let constant_head = + (constant.len() as u64) | (u64::from(u32::from_le_bytes(prefix_bytes)) << 32); + let constant_prefix = u32::from_be_bytes(prefix_bytes); + let constant_tail = u64::from_be_bytes(tail_bytes); + // The full 16-byte word an inlined view holding `constant` would carry; only meaningful when + // the constant is short enough to inline, and only reached in that case (a longer constant + // never head-matches an inlined view). + let constant_inlined = + u128::from(constant_head) | (u128::from(u64::from_le_bytes(tail_bytes)) << 64); + + match op { + CompareOperator::Eq => BitBuffer::collect_bool(len, |i| { + // SAFETY: `collect_bool` yields i < len == views.len(). + let view = unsafe { lhs.view_unchecked(i) }; + constant_eq(lhs, view, constant, constant_head, constant_inlined) + }), + CompareOperator::NotEq => BitBuffer::collect_bool(len, |i| { + // SAFETY: `collect_bool` yields i < len == views.len(). + let view = unsafe { lhs.view_unchecked(i) }; + !constant_eq(lhs, view, constant, constant_head, constant_inlined) + }), + CompareOperator::Gt => collect_constant_ordering_bits( + lhs, + constant, + constant_prefix, + constant_tail, + Ordering::is_gt, + ), + CompareOperator::Gte => collect_constant_ordering_bits( + lhs, + constant, + constant_prefix, + constant_tail, + Ordering::is_ge, + ), + CompareOperator::Lt => collect_constant_ordering_bits( + lhs, + constant, + constant_prefix, + constant_tail, + Ordering::is_lt, + ), + CompareOperator::Lte => collect_constant_ordering_bits( + lhs, + constant, + constant_prefix, + constant_tail, + Ordering::is_le, + ), + } +} + +/// Bit-pack `predicate(constant_cmp(lhs[i], constant))` over one view side. +fn collect_constant_ordering_bits( + lhs: &ViewsSide<'_>, + constant: &[u8], + constant_prefix: u32, + constant_tail: u64, + predicate: impl Fn(Ordering) -> bool, +) -> BitBuffer { + BitBuffer::collect_bool(lhs.len(), |i| { + // SAFETY: `collect_bool` yields i < len == views.len(). + let view = unsafe { lhs.view_unchecked(i) }; + predicate(constant_cmp( + lhs, + view, + constant, + constant_prefix, + constant_tail, + )) + }) +} + +/// Compare a view against a constant for equality using the constant's precomputed head and +/// inline words. +#[inline] +fn constant_eq( + side: &ViewsSide<'_>, + view: &BinaryView, + constant: &[u8], + constant_head: u64, + constant_inlined: u128, +) -> bool { + if view_head(view) != constant_head { + return false; + } + if view.is_inlined() { + // An equal head implies equal lengths, so the constant is also at most 12 bytes and + // `constant_inlined` holds its exact inlined representation. + return view.as_u128() == constant_inlined; + } + side.view_bytes(view)[4..] == constant[4..] +} + +/// Compare a view against a constant using the constant's precomputed prefix and tail words. +#[inline] +fn constant_cmp( + side: &ViewsSide<'_>, + view: &BinaryView, + constant: &[u8], + constant_prefix: u32, + constant_tail: u64, +) -> Ordering { + let prefix = view_prefix(view); + if prefix != constant_prefix { + return prefix.cmp(&constant_prefix); + } + if view.is_inlined() { + // The view's whole value lives in its first 12 (zero-padded) bytes, and the constant's + // first 12 (zero-padded) bytes are precomputed: compare tails, then lengths. A padded + // tie means one value is a prefix of the other within the first 12 bytes. + let tail = view_tail(view); + if tail != constant_tail { + return tail.cmp(&constant_tail); + } + return (view.len() as usize).cmp(&constant.len()); + } + side.view_bytes(view).cmp(constant) +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs b/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs new file mode 100644 index 00000000000..0825d0674af --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs @@ -0,0 +1,188 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native comparison of decimal arrays. +//! +//! Both operands share a logical [`DecimalDType`] (equal precision and scale), so comparing the +//! unscaled integer values is sufficient. The physical storage width may differ per operand; when +//! it does, both sides are widened once to the larger storage type before the lane loop. +//! +//! [`DecimalDType`]: crate::dtype::DecimalDType + +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Constant; +use crate::arrays::DecimalArray; +use crate::dtype::NativeDecimalType; +use crate::dtype::Nullability; +use crate::dtype::i256; +use crate::match_each_decimal_value_type; +use crate::scalar::DecimalValue; +use crate::scalar_fn::fns::binary::compare::collect_bits; +use crate::scalar_fn::fns::binary::compare::collect_zip_bits; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::operators::CompareOperator; +use crate::validity::Validity; + +enum DecimalOperand { + Array { + values: DecimalArray, + validity: Validity, + }, + Constant { + value: DecimalValue, + validity: Validity, + }, +} + +impl DecimalOperand { + fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(constant) = array.as_opt::() { + let value = constant + .scalar() + .as_decimal() + .decimal_value() + .ok_or_else(|| vortex_err!("null constant handled by execute_compare"))?; + return Ok(Self::Constant { + value, + validity: if constant.scalar().dtype().is_nullable() { + Validity::AllValid + } else { + Validity::NonNullable + }, + }); + } + + let values = array.clone().execute::(ctx)?; + let validity = values.validity()?; + Ok(Self::Array { values, validity }) + } + + fn validity(&self) -> Validity { + match self { + Self::Array { validity, .. } | Self::Constant { validity, .. } => validity.clone(), + } + } +} + +/// Compare two decimal arrays with the same logical decimal dtype. +pub(super) fn compare_decimal( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let lhs = DecimalOperand::try_new(lhs, ctx)?; + let rhs = DecimalOperand::try_new(rhs, ctx)?; + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + + let bits = match (lhs, rhs) { + (DecimalOperand::Array { values: l, .. }, DecimalOperand::Array { values: r, .. }) => { + compare_decimal_values(&l, &r, op) + } + (DecimalOperand::Array { values, .. }, DecimalOperand::Constant { value, .. }) => { + compare_decimal_constant(&values, value, op) + } + (DecimalOperand::Constant { value, .. }, DecimalOperand::Array { values, .. }) => { + compare_decimal_constant(&values, value, op.swap()) + } + (DecimalOperand::Constant { value: l, .. }, DecimalOperand::Constant { value: r, .. }) => { + // Unreachable through `execute_compare` (constant-constant is folded there), but + // cheap to answer anyway. + let ordering = l.as_i256().cmp(&r.as_i256()); + BitBuffer::full(super::ordering_predicate(op)(ordering), len) + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +fn compare_decimal_values( + lhs: &DecimalArray, + rhs: &DecimalArray, + op: CompareOperator, +) -> BitBuffer { + let common = lhs.values_type().max(rhs.values_type()); + match_each_decimal_value_type!(common, |W| { + let lhs = widened_buffer::(lhs); + let rhs = widened_buffer::(rhs); + compare_slices::(&lhs, &rhs, op) + }) +} + +/// Return the array's unscaled values widened to `W`, which must be at least as wide as the +/// array's storage type. +pub(super) fn widened_buffer(array: &DecimalArray) -> Buffer { + if array.values_type() == W::DECIMAL_TYPE { + return array.buffer::(); + } + match_each_decimal_value_type!(array.values_type(), |T| { + array + .buffer::() + .iter() + .map(|v| W::from(*v).vortex_expect("widening decimal cast must succeed")) + .collect() + }) +} + +fn compare_decimal_constant( + array: &DecimalArray, + constant: DecimalValue, + op: CompareOperator, +) -> BitBuffer { + match_each_decimal_value_type!(array.values_type(), |T| { + match constant.cast::() { + Some(value) => compare_slice_constant::(&array.buffer::(), value, op), + None => { + // The constant does not fit the array's storage type, so it is either greater + // than every possible array value or less than every possible array value; the + // sign tells us which. + let constant_greater = constant.as_i256() > i256::ZERO; + let result = match op { + CompareOperator::Eq => false, + CompareOperator::NotEq => true, + // array constant + CompareOperator::Lt | CompareOperator::Lte => constant_greater, + CompareOperator::Gt | CompareOperator::Gte => !constant_greater, + }; + BitBuffer::full(result, array.len()) + } + } + }) +} + +fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a == b), + CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| a != b), + CompareOperator::Gt => collect_zip_bits(lhs, rhs, |a: T, b: T| a > b), + CompareOperator::Gte => collect_zip_bits(lhs, rhs, |a: T, b: T| a >= b), + CompareOperator::Lt => collect_zip_bits(lhs, rhs, |a: T, b: T| a < b), + CompareOperator::Lte => collect_zip_bits(lhs, rhs, |a: T, b: T| a <= b), + } +} + +fn compare_slice_constant( + values: &[T], + constant: T, + op: CompareOperator, +) -> BitBuffer { + match op { + CompareOperator::Eq => collect_bits(values, |a: T| a == constant), + CompareOperator::NotEq => collect_bits(values, |a: T| a != constant), + CompareOperator::Gt => collect_bits(values, |a: T| a > constant), + CompareOperator::Gte => collect_bits(values, |a: T| a >= constant), + CompareOperator::Lt => collect_bits(values, |a: T| a < constant), + CompareOperator::Lte => collect_bits(values, |a: T| a <= constant), + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs new file mode 100644 index 00000000000..70e98639b54 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -0,0 +1,320 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native comparison kernels. +//! +//! [`execute_compare`] dispatches on the logical [`DType`] of its operands and evaluates every +//! comparison directly over Vortex canonical arrays — bit buffers for booleans, lane kernels from +//! `vortex-compute` for primitives and decimals, binary views for strings/bytes, and a row-wise +//! comparator for nested types. There is no Arrow fallback. +//! +//! Floating point values compare with Vortex's total ordering (`NaN` is the largest value, +//! `-0.0 < +0.0`, and equality is bitwise), matching [`Scalar`] comparison semantics. + +use std::cmp::Ordering; + +use vortex_buffer::BitBuffer; +use vortex_buffer::BufferMut; +use vortex_compute::lane_kernels::IndexedSourceExt; +use vortex_compute::lane_kernels::LaneZip; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::array::VTable; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::ExtensionArray; +use crate::arrays::ScalarFn; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::scalar_fn::ExactScalarFn; +use crate::arrays::scalar_fn::ScalarFnArrayExt; +use crate::arrays::scalar_fn::ScalarFnArrayView; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::kernel::ExecuteParentKernel; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::fns::operators::CompareOperator; +use crate::validity::Validity; + +mod boolean; +mod bytes; +mod decimal; +mod nested; +mod primitive; +#[cfg(test)] +mod tests; + +/// Trait for encoding-specific comparison kernels that operate in encoded space. +/// +/// Implementations can compare an encoded array against another array (typically a constant) +/// without first decompressing. The adaptor normalizes operand order so `array` is always +/// the left-hand side, swapping the operator when necessary. +pub trait CompareKernel: VTable { + fn compare( + lhs: ArrayView<'_, Self>, + rhs: &ArrayRef, + operator: CompareOperator, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; +} + +/// Adaptor that bridges [`CompareKernel`] implementations to [`ExecuteParentKernel`]. +/// +/// When a `ScalarFnArray(Binary, cmp_op)` wraps a child that implements `CompareKernel`, +/// this adaptor extracts the comparison operator and other operand, normalizes operand order +/// (swapping the operator if the encoded array is on the RHS), and delegates to the kernel. +#[derive(Default, Debug)] +pub struct CompareExecuteAdaptor(pub V); + +impl ExecuteParentKernel for CompareExecuteAdaptor +where + V: CompareKernel, +{ + type Parent = ExactScalarFn; + + fn execute_parent( + &self, + array: ArrayView<'_, V>, + parent: ScalarFnArrayView<'_, Binary>, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + // Only handle comparison operators + let Ok(cmp_op) = CompareOperator::try_from(*parent.options) else { + return Ok(None); + }; + + // Get the ScalarFnArray to access children + let Some(scalar_fn_array) = parent.as_opt::() else { + return Ok(None); + }; + // Normalize so `array` is always LHS, swapping the operator if needed + // TODO(joe): should be go this here or in the Rule/Kernel + let (cmp_op, other) = match child_idx { + 0 => (cmp_op, scalar_fn_array.get_child(1)), + 1 => (cmp_op.swap(), scalar_fn_array.get_child(0)), + _ => return Ok(None), + }; + + let len = array.len(); + let nullable = array.dtype().is_nullable() || other.dtype().is_nullable(); + + // Empty array → empty bool result + if len == 0 { + return Ok(Some( + Canonical::empty(&DType::Bool(nullable.into())).into_array(), + )); + } + + // Null constant on either side → all-null bool result + if other.as_constant().is_some_and(|s| s.is_null()) { + return Ok(Some( + ConstantArray::new(Scalar::null(DType::Bool(nullable.into())), len).into_array(), + )); + } + + V::compare(array, other, cmp_op, ctx) + } +} + +/// Execute a compare operation between two arrays. +/// +/// This is the entry point for compare operations from the binary expression. +/// Handles empty, constant-null, and constant-constant directly, otherwise dispatches to a +/// native per-dtype kernel. +pub(crate) fn execute_compare( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable(); + + if lhs.len() != rhs.len() { + vortex_bail!( + "compare operator requires equal lengths, got {} and {}", + lhs.len(), + rhs.len() + ); + } + + if lhs.is_empty() { + return Ok(Canonical::empty(&DType::Bool(nullable.into())).into_array()); + } + + let left_constant_null = lhs.as_constant().map(|l| l.is_null()).unwrap_or(false); + let right_constant_null = rhs.as_constant().map(|r| r.is_null()).unwrap_or(false); + if left_constant_null || right_constant_null { + return Ok( + ConstantArray::new(Scalar::null(DType::Bool(nullable.into())), lhs.len()).into_array(), + ); + } + + // Constant-constant fast path + if let (Some(lhs_const), Some(rhs_const)) = (lhs.as_opt::(), rhs.as_opt::()) + { + let result = scalar_cmp(lhs_const.scalar(), rhs_const.scalar(), op)?; + return Ok(ConstantArray::new(result, lhs.len()).into_array()); + } + + compare_arrays(lhs, rhs, op, ctx) +} + +/// Dispatch a comparison to the native kernel for the operands' logical dtype. +fn compare_arrays( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let nullability = Nullability::from(lhs.dtype().is_nullable() || rhs.dtype().is_nullable()); + + // Extension arrays compare through their storage. When both sides are extensions they must + // agree on the full extension dtype (a timestamp in milliseconds must not compare its raw + // storage against one in seconds); a lone extension side may compare against its raw storage. + if lhs.dtype().is_extension() || rhs.dtype().is_extension() { + if lhs.dtype().is_extension() + && rhs.dtype().is_extension() + && !lhs.dtype().eq_ignore_nullability(rhs.dtype()) + { + vortex_bail!( + "Cannot compare extension dtypes {} and {}", + lhs.dtype(), + rhs.dtype() + ); + } + let lhs = extension_storage(lhs, ctx)?; + let rhs = extension_storage(rhs, ctx)?; + return compare_arrays(&lhs, &rhs, op, ctx); + } + + if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) { + vortex_bail!( + "Cannot compare different DTypes {} and {}", + lhs.dtype(), + rhs.dtype() + ); + } + + match lhs.dtype() { + // Every value of the null dtype is null, so every comparison result is null. + DType::Null => Ok(ConstantArray::new( + Scalar::null(DType::Bool(Nullability::Nullable)), + lhs.len(), + ) + .into_array()), + DType::Bool(_) => boolean::compare_bool(lhs, rhs, op, nullability, ctx), + DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, nullability, ctx), + DType::Decimal(..) => decimal::compare_decimal(lhs, rhs, op, nullability, ctx), + DType::Utf8(_) | DType::Binary(_) => bytes::compare_bytes(lhs, rhs, op, nullability, ctx), + DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => { + nested::compare_nested(lhs, rhs, op, nullability, ctx) + } + DType::Union(..) | DType::Variant(_) | DType::Extension(_) => { + vortex_bail!("compare is not supported for dtype {}", lhs.dtype()) + } + } +} + +/// Replace an extension array (or extension constant) with its storage. Non-extension arrays are +/// returned unchanged. +fn extension_storage(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if !array.dtype().is_extension() { + return Ok(array.clone()); + } + if let Some(constant) = array.as_opt::() { + return Ok(ConstantArray::new( + constant.scalar().as_extension().to_storage_scalar(), + constant.len(), + ) + .into_array()); + } + let ext = array.clone().execute::(ctx)?; + Ok(ext.storage_array().clone()) +} + +pub fn scalar_cmp(lhs: &Scalar, rhs: &Scalar, operator: CompareOperator) -> VortexResult { + if lhs.is_null() | rhs.is_null() { + return Ok(Scalar::null(DType::Bool(Nullability::Nullable))); + } + + let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); + + // We use `partial_cmp` to ensure we do not lose a type mismatch error. + let ordering = lhs.partial_cmp(rhs).ok_or_else(|| { + vortex_err!( + "Cannot compare scalars with incompatible types: {} and {}", + lhs.dtype(), + rhs.dtype() + ) + })?; + + let b = match operator { + CompareOperator::Eq => ordering.is_eq(), + CompareOperator::NotEq => ordering.is_ne(), + CompareOperator::Gt => ordering.is_gt(), + CompareOperator::Gte => ordering.is_ge(), + CompareOperator::Lt => ordering.is_lt(), + CompareOperator::Lte => ordering.is_le(), + }; + + Ok(Scalar::bool(b, nullability)) +} + +/// Returns the predicate `Ordering -> bool` for a comparison operator. +#[inline] +pub(super) fn ordering_predicate(op: CompareOperator) -> fn(Ordering) -> bool { + match op { + CompareOperator::Eq => Ordering::is_eq, + CompareOperator::NotEq => Ordering::is_ne, + CompareOperator::Gt => Ordering::is_gt, + CompareOperator::Gte => Ordering::is_ge, + CompareOperator::Lt => Ordering::is_lt, + CompareOperator::Lte => Ordering::is_le, + } +} + +/// Freeze `len` bits packed into `words` (LSB-first, 64 lanes per word) into a [`BitBuffer`]. +pub(super) fn bit_buffer_from_words(words: BufferMut, len: usize) -> BitBuffer { + debug_assert!(words.len() * 64 >= len); + let mut bytes = words.into_byte_buffer(); + bytes.truncate(len.div_ceil(8)); + BitBuffer::new(bytes.freeze(), len) +} + +/// Bit-pack the predicate `f(lhs[i], rhs[i])` over two equal-length slices into a [`BitBuffer`]. +pub(super) fn collect_zip_bits( + lhs: &[T], + rhs: &[T], + mut f: impl FnMut(T, T) -> bool, +) -> BitBuffer { + let len = lhs.len(); + let mut words = BufferMut::::zeroed(len.div_ceil(64)); + LaneZip::new(lhs, rhs).map_bits_into(words.as_mut_slice(), |(a, b)| f(a, b)); + bit_buffer_from_words(words, len) +} + +/// Bit-pack the predicate `f(values[i])` over a slice into a [`BitBuffer`]. +pub(super) fn collect_bits(values: &[T], f: impl FnMut(T) -> bool) -> BitBuffer { + let len = values.len(); + let mut words = BufferMut::::zeroed(len.div_ceil(64)); + values.map_bits_into(words.as_mut_slice(), f); + bit_buffer_from_words(words, len) +} + +/// Combine operand validities into the validity of a comparison result with the given result +/// nullability. +pub(super) fn compare_validity( + lhs: Validity, + rhs: Validity, + nullability: Nullability, +) -> VortexResult { + Ok(lhs.and(rhs)?.union_nullability(nullability)) +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs new file mode 100644 index 00000000000..56d8cbf0496 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs @@ -0,0 +1,256 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Row-wise comparison of nested (struct, list, fixed-size list) arrays. +//! +//! Nested comparisons canonicalize both operands recursively and build a tree of row +//! comparators, one per nested level. The ordering semantics match [`Scalar`] comparison: +//! structs compare field-by-field in declaration order, lists compare element-by-element and +//! then by length, and a null value (at any nesting depth) orders before every non-null value. +//! Only top-level nulls make the comparison result null. +//! +//! [`Scalar`]: crate::scalar::Scalar + +use std::cmp::Ordering; + +use num_traits::AsPrimitive; +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::RecursiveCanonical; +use crate::arrays::BoolArray; +use crate::arrays::DecimalArray; +use crate::arrays::ExtensionArray; +use crate::arrays::FixedSizeListArray; +use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::arrays::VarBinViewArray; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::struct_::StructArrayExt; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::match_each_integer_ptype; +use crate::match_each_native_ptype; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::operators::CompareOperator; + +/// A row comparator: compares row `i` of the left operand against row `j` of the right operand. +type RowComparator = Box Ordering>; + +/// Compare two nested arrays row by row. +pub(super) fn compare_nested( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let lhs = lhs + .clone() + .execute::(ctx)? + .0 + .into_array(); + let rhs = rhs + .clone() + .execute::(ctx)? + .0 + .into_array(); + + let validity = compare_validity(lhs.validity()?, rhs.validity()?, nullability)?; + let comparator = build_comparator(&lhs, &rhs, ctx)?; + // Dispatch the operator outside the row loop so the predicate inlines into each loop; the + // comparator call itself stays virtual. + let bits = match op { + CompareOperator::Eq => collect_ordering_bits(len, &comparator, Ordering::is_eq), + CompareOperator::NotEq => collect_ordering_bits(len, &comparator, Ordering::is_ne), + CompareOperator::Gt => collect_ordering_bits(len, &comparator, Ordering::is_gt), + CompareOperator::Gte => collect_ordering_bits(len, &comparator, Ordering::is_ge), + CompareOperator::Lt => collect_ordering_bits(len, &comparator, Ordering::is_lt), + CompareOperator::Lte => collect_ordering_bits(len, &comparator, Ordering::is_le), + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +/// Bit-pack `predicate(comparator(i, i))` over `len` rows. +fn collect_ordering_bits( + len: usize, + comparator: &RowComparator, + predicate: impl Fn(Ordering) -> bool, +) -> BitBuffer { + BitBuffer::collect_bool(len, |i| predicate(comparator(i, i))) +} + +/// The validity mask of a recursively canonical array. +fn validity_mask(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + array.validity()?.execute_mask(array.len(), ctx) +} + +/// Build a row comparator over two recursively canonical arrays of the same logical dtype +/// (ignoring nullability). Null values order before all non-null values at every level. +fn build_comparator( + lhs: &ArrayRef, + rhs: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let lhs_mask = validity_mask(lhs, ctx)?; + let rhs_mask = validity_mask(rhs, ctx)?; + let values = build_values_comparator(lhs, rhs, ctx)?; + + if lhs_mask.all_true() && rhs_mask.all_true() { + return Ok(values); + } + + Ok(Box::new(move |i, j| { + match (lhs_mask.value(i), rhs_mask.value(j)) { + (true, true) => values(i, j), + (true, false) => Ordering::Greater, + (false, true) => Ordering::Less, + (false, false) => Ordering::Equal, + } + })) +} + +/// Build a comparator over the non-null values of two recursively canonical arrays. +fn build_values_comparator( + lhs: &ArrayRef, + rhs: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) { + vortex_bail!( + "Cannot compare different DTypes {} and {}", + lhs.dtype(), + rhs.dtype() + ); + } + + Ok(match lhs.dtype() { + // Null values compare through the validity wrapper; the value comparator is trivial. + DType::Null => Box::new(|_, _| Ordering::Equal), + DType::Bool(_) => { + let lhs = lhs.clone().execute::(ctx)?.into_bit_buffer(); + let rhs = rhs.clone().execute::(ctx)?.into_bit_buffer(); + Box::new(move |i, j| lhs.value(i).cmp(&rhs.value(j))) + } + DType::Primitive(ptype, _) => { + match_each_native_ptype!(*ptype, |T| { primitive_comparator::(lhs, rhs, ctx)? }) + } + DType::Decimal(..) => { + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + let common = lhs.values_type().max(rhs.values_type()); + crate::match_each_decimal_value_type!(common, |W| { + let lhs = super::decimal::widened_buffer::(&lhs); + let rhs = super::decimal::widened_buffer::(&rhs); + Box::new(move |i: usize, j: usize| lhs[i].cmp(&rhs[j])) as RowComparator + }) + } + DType::Utf8(_) | DType::Binary(_) => { + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + Box::new(move |i, j| view_bytes(&lhs, i).cmp(view_bytes(&rhs, j))) + } + DType::Struct(..) => { + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + let fields = lhs + .iter_unmasked_fields() + .zip(rhs.iter_unmasked_fields()) + .map(|(lhs_field, rhs_field)| build_comparator(lhs_field, rhs_field, ctx)) + .collect::>>()?; + Box::new(move |i, j| { + fields + .iter() + .map(|cmp| cmp(i, j)) + .find(|ordering| ordering.is_ne()) + .unwrap_or(Ordering::Equal) + }) + } + DType::List(..) => { + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + let elements = build_comparator(lhs.elements(), rhs.elements(), ctx)?; + // Materialize offsets and sizes once instead of paying `offset_at`/`size_at`'s + // per-call encoding dispatch inside the row loop. + let lhs_offsets = usize_values(lhs.offsets(), ctx)?; + let lhs_sizes = usize_values(lhs.sizes(), ctx)?; + let rhs_offsets = usize_values(rhs.offsets(), ctx)?; + let rhs_sizes = usize_values(rhs.sizes(), ctx)?; + Box::new(move |i, j| { + let lhs_len = lhs_sizes[i]; + let rhs_len = rhs_sizes[j]; + (0..lhs_len.min(rhs_len)) + .map(|el| elements(lhs_offsets[i] + el, rhs_offsets[j] + el)) + .find(|ordering| ordering.is_ne()) + .unwrap_or_else(|| lhs_len.cmp(&rhs_len)) + }) + } + DType::FixedSizeList(_, size, _) => { + let size = *size as usize; + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + let elements = build_comparator(lhs.elements(), rhs.elements(), ctx)?; + Box::new(move |i, j| { + (0..size) + .map(|el| elements(i * size + el, j * size + el)) + .find(|ordering| ordering.is_ne()) + .unwrap_or(Ordering::Equal) + }) + } + DType::Extension(_) => { + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + build_comparator(lhs.storage_array(), rhs.storage_array(), ctx)? + } + DType::Union(..) | DType::Variant(_) => { + vortex_bail!("compare is not supported for dtype {}", lhs.dtype()) + } + }) +} + +fn primitive_comparator( + lhs: &ArrayRef, + rhs: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let lhs = lhs + .clone() + .execute::(ctx)? + .into_buffer::(); + let rhs = rhs + .clone() + .execute::(ctx)? + .into_buffer::(); + Ok(Box::new(move |i, j| lhs[i].total_compare(rhs[j]))) +} + +/// Materialize an integer array as `usize` values for `O(1)` row access. +fn usize_values(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + let primitive = array.clone().execute::(ctx)?; + match_each_integer_ptype!(primitive.ptype(), |P| { + Ok(primitive.as_slice::

().iter().map(|v| v.as_()).collect()) + }) +} + +/// The bytes of row `index`, resolved directly from the view without cloning buffer handles the +/// way `VarBinViewArray::bytes_at` does. +fn view_bytes(array: &VarBinViewArray, index: usize) -> &[u8] { + let view = &array.views()[index]; + if view.is_inlined() { + view.as_inlined().value() + } else { + let view = view.as_view(); + &array.buffer(view.buffer_index as usize).as_slice()[view.as_range()] + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs new file mode 100644 index 00000000000..3bfb11a266e --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native comparison of primitive arrays via bit-packing lane kernels. + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::match_each_native_ptype; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::PrimitiveOperand; +use crate::scalar_fn::fns::binary::compare::collect_bits; +use crate::scalar_fn::fns::binary::compare::collect_zip_bits; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::operators::CompareOperator; + +/// Compare two primitive arrays of the same [`PType`]. +/// +/// Floats compare with Vortex's total ordering: `NaN` is the largest value, `-0.0 < +0.0`, and +/// equality is bitwise. +pub(super) fn compare_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let ptype = PType::try_from(lhs.dtype())?; + match_each_native_ptype!(ptype, |T| { + compare_primitive_typed::(lhs, rhs, op, nullability, ctx) + }) +} + +fn compare_primitive_typed( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; + let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; + if lhs.len() != rhs.len() { + vortex_bail!( + "compare operator requires equal lengths, got {} and {}", + lhs.len(), + rhs.len() + ); + } + + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + + let bits = match (&lhs, &rhs) { + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slices(lhs, rhs, op), + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => compare_slice_constant(lhs, *rhs, op), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slice_constant(rhs, *lhs, op.swap()), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => { + // Unreachable through `execute_compare` (constant-constant is folded there), but + // cheap to answer anyway. + BitBuffer::full(apply_op(*lhs, *rhs, op), len) + } + (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { + return Ok( + ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) + .into_array(), + ); + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +#[inline(always)] +fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { + match op { + CompareOperator::Eq => lhs.is_eq(rhs), + CompareOperator::NotEq => !lhs.is_eq(rhs), + CompareOperator::Gt => lhs.is_gt(rhs), + CompareOperator::Gte => lhs.is_ge(rhs), + CompareOperator::Lt => lhs.is_lt(rhs), + CompareOperator::Lte => lhs.is_le(rhs), + } +} + +fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { + // Dispatch the operator outside the lane loop so each instantiation vectorizes a single + // branch-free predicate. + match op { + CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a.is_eq(b)), + CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| !a.is_eq(b)), + CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), + CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), + CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), + CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), + } +} + +fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_bits(lhs, |a: T| a.is_eq(rhs)), + CompareOperator::NotEq => collect_bits(lhs, |a: T| !a.is_eq(rhs)), + CompareOperator::Gt => collect_bits(lhs, |a: T| a.is_gt(rhs)), + CompareOperator::Gte => collect_bits(lhs, |a: T| a.is_ge(rhs)), + CompareOperator::Lt => collect_bits(lhs, |a: T| a.is_lt(rhs)), + CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs)), + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs new file mode 100644 index 00000000000..c2af83561a4 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -0,0 +1,673 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::ExtensionArray; +use crate::arrays::FixedSizeListArray; +use crate::arrays::ListArray; +use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::arrays::VarBinArray; +use crate::arrays::VarBinViewArray; +use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::extension::datetime::TimeUnit; +use crate::extension::datetime::Timestamp; +use crate::extension::datetime::TimestampOptions; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::scalar_cmp; +use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::fns::operators::Operator; +use crate::test_harness::to_int_indices; +use crate::validity::Validity; + +#[test] +fn test_bool_basic_comparisons() { + 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]), + ); + + let matches = arr + .clone() + .into_array() + .binary(arr.clone().into_array(), Operator::Eq) + .unwrap() + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [1u64, 2, 3, 4]); + + let matches = arr + .clone() + .into_array() + .binary(arr.clone().into_array(), Operator::NotEq) + .unwrap() + .execute::(ctx) + .vortex_expect("must be a bool array"); + let empty: [u64; 0] = []; + 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]), + ); + + let matches = arr + .clone() + .into_array() + .binary(other.clone().into_array(), Operator::Lte) + .unwrap() + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [2u64, 3, 4]); + + let matches = arr + .clone() + .into_array() + .binary(other.clone().into_array(), Operator::Lt) + .unwrap() + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [4u64]); + + let matches = other + .clone() + .into_array() + .binary(arr.clone().into_array(), Operator::Gte) + .unwrap() + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [2u64, 3, 4]); + + let matches = other + .into_array() + .binary(arr.into_array(), Operator::Gt) + .unwrap() + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [4u64]); +} + +#[test] +fn constant_compare() { + let left = ConstantArray::new(Scalar::from(2u32), 10); + let right = ConstantArray::new(Scalar::from(10u32), 10); + + let result = left + .into_array() + .binary(right.into_array(), Operator::Gt) + .unwrap(); + assert_eq!(result.len(), 10); + let scalar = result + .execute_scalar(0, &mut array_session().create_execution_ctx()) + .unwrap(); + assert_eq!(scalar.as_bool().value(), Some(false)); +} + +#[rstest] +#[case(VarBinArray::from(vec!["a", "b"]).into_array(), VarBinViewArray::from_iter_str(["a", "b"]).into_array())] +#[case(VarBinViewArray::from_iter_str(["a", "b"]).into_array(), VarBinArray::from(vec!["a", "b"]).into_array())] +#[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 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, &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( + values1.into_array(), + offsets1.into_array(), + Validity::NonNullable, + ) + .unwrap(); + + let values2 = PrimitiveArray::from_iter([1i32, 2, 3, 4, 7, 8]); + let offsets2 = PrimitiveArray::from_iter([0i32, 2, 4, 6]); + let list2 = ListArray::try_new( + values2.into_array(), + offsets2.into_array(), + Validity::NonNullable, + ) + .unwrap(); + + let result = list1 + .clone() + .into_array() + .binary(list2.clone().into_array(), Operator::Eq) + .unwrap(); + let expected = BoolArray::from_iter([true, true, false]); + assert_arrays_eq!(result, expected, &mut ctx); + + let result = list1 + .clone() + .into_array() + .binary(list2.clone().into_array(), Operator::NotEq) + .unwrap(); + let expected = BoolArray::from_iter([false, false, true]); + 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, &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( + values.into_array(), + offsets.into_array(), + Validity::NonNullable, + ) + .unwrap(); + + let list_scalar = Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + vec![3i32.into(), 4i32.into()], + Nullability::NonNullable, + ); + let constant = ConstantArray::new(list_scalar, 3); + + let result = list + .into_array() + .binary(constant.into_array(), Operator::Eq) + .unwrap(); + let expected = BoolArray::from_iter([false, true, false]); + 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]); + + let bool_field2 = BoolArray::from_iter([Some(true), Some(false), Some(false)]); + let int_field2 = PrimitiveArray::from_iter([1i32, 2, 4]); + + let struct1 = StructArray::from_fields(&[ + ("bool_col", bool_field1.into_array()), + ("int_col", int_field1.into_array()), + ]) + .unwrap(); + + let struct2 = StructArray::from_fields(&[ + ("bool_col", bool_field2.into_array()), + ("int_col", int_field2.into_array()), + ]) + .unwrap(); + + let result = struct1 + .clone() + .into_array() + .binary(struct2.clone().into_array(), Operator::Eq) + .unwrap(); + let expected = BoolArray::from_iter([true, true, false]); + 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, &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(), + 5, + Validity::NonNullable, + ) + .unwrap(); + + let empty2 = StructArray::try_new( + FieldNames::from(Vec::::new()), + Vec::new(), + 5, + Validity::NonNullable, + ) + .unwrap(); + + let result = empty1 + .into_array() + .binary(empty2.into_array(), Operator::Eq) + .unwrap(); + let expected = BoolArray::from_iter([true, true, true, true, true]); + 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(), + "banana".as_bytes(), + "cherry".as_bytes(), + ]); + let struct1 = StructArray::from_fields(&[("data", bin_field1.into_array())]).unwrap(); + + // RHS: struct with a VarBinViewArray (view-based) binary field — same logical DType + let bin_field2 = VarBinViewArray::from_iter_bin([ + "apple".as_bytes(), + "banana".as_bytes(), + "durian".as_bytes(), + ]); + let struct2 = StructArray::from_fields(&[("data", bin_field2.into_array())]).unwrap(); + + let result = struct1 + .into_array() + .binary(struct2.into_array(), Operator::Eq) + .unwrap(); + let expected = BoolArray::from_iter([true, true, false]); + assert_arrays_eq!(result, expected, &mut ctx); +} + +/// Regression test: `scalar_cmp` must error when comparing scalars with incompatible +/// extension types (e.g., timestamps with different time units) rather than silently +/// returning a wrong result. +#[test] +fn scalar_cmp_incompatible_extension_types_errors() { + let ms_scalar = Scalar::extension::( + TimestampOptions { + unit: TimeUnit::Milliseconds, + tz: None, + }, + Scalar::from(1704067200000i64), + ); + let s_scalar = Scalar::extension::( + TimestampOptions { + unit: TimeUnit::Seconds, + tz: None, + }, + Scalar::from(1704067200i64), + ); + + // Ordering comparisons must error on incompatible types. + assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Gt).is_err()); + assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Lt).is_err()); + assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Gte).is_err()); + assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Lte).is_err()); + assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Eq).is_err()); + assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::NotEq).is_err()); +} + +#[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(), + buffer![0i32, 0i32, 0i32].into_array(), + Validity::AllValid, + ); + + let result = list + .clone() + .into_array() + .binary(list.into_array(), Operator::Eq) + .unwrap(); + 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()); +} + +fn execute_compare_test(lhs: ArrayRef, rhs: ArrayRef, op: Operator) -> ArrayRef { + lhs.binary(rhs, op).unwrap() +} + +#[rstest] +#[case(Operator::Eq, [false, true, false, false])] +#[case(Operator::NotEq, [true, false, true, true])] +#[case(Operator::Lt, [true, false, false, true])] +#[case(Operator::Lte, [true, true, false, true])] +#[case(Operator::Gt, [false, false, true, false])] +#[case(Operator::Gte, [false, true, true, false])] +fn int_all_operators(#[case] op: Operator, #[case] expected: [bool; 4]) { + let mut ctx = array_session().create_execution_ctx(); + let lhs = buffer![1i32, 5, 9, 2].into_array(); + let rhs = buffer![3i32, 5, 7, 4].into_array(); + let result = execute_compare_test(lhs, rhs, op); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); +} + +#[rstest] +#[case(Operator::Lt, [Some(true), None, Some(false), None])] +#[case(Operator::Eq, [Some(false), None, Some(false), None])] +fn int_nullable(#[case] op: Operator, #[case] expected: [Option; 4]) { + let mut ctx = array_session().create_execution_ctx(); + let lhs = PrimitiveArray::from_option_iter([Some(1i64), None, Some(9), Some(2)]).into_array(); + let rhs = PrimitiveArray::from_option_iter([Some(3i64), Some(5), Some(7), None]).into_array(); + let result = execute_compare_test(lhs, rhs, op); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); +} + +#[rstest] +#[case(Operator::Gt, [false, false, true, false])] +#[case(Operator::Lte, [true, true, false, true])] +fn int_constant_lhs_and_rhs(#[case] op: Operator, #[case] expected: [bool; 4]) { + let mut ctx = array_session().create_execution_ctx(); + let array = buffer![1i32, 5, 9, 2].into_array(); + let constant = ConstantArray::new(5i32, 4).into_array(); + + let result = execute_compare_test(array.clone(), constant.clone(), op); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); + + // The swapped form must produce the swapped result. + let swapped = execute_compare_test(constant, array, op); + let expected_swapped: Vec = match op { + Operator::Gt => vec![true, false, false, true], + Operator::Lte => vec![false, true, true, false], + _ => unreachable!(), + }; + assert_arrays_eq!(swapped, BoolArray::from_iter(expected_swapped), &mut ctx); +} + +/// Floats compare with Vortex's total ordering: NaN is the largest value, equality is bitwise, +/// and -0.0 < +0.0. This matches `Scalar` comparison semantics. +#[test] +fn float_total_order() { + let mut ctx = array_session().create_execution_ctx(); + let lhs = buffer![f64::NAN, f64::NAN, -0.0f64, 1.0].into_array(); + let rhs = buffer![f64::NAN, f64::INFINITY, 0.0f64, f64::NAN].into_array(); + + let result = execute_compare_test(lhs.clone(), rhs.clone(), Operator::Eq); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, false, false, false]), + &mut ctx + ); + + let result = execute_compare_test(lhs, rhs, Operator::Lt); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, false, true, true]), + &mut ctx + ); +} + +#[rstest] +#[case(Operator::Eq, [true, false, true, true])] +#[case(Operator::Lt, [false, true, false, false])] +#[case(Operator::Gte, [true, false, true, true])] +fn bool_constant(#[case] op: Operator, #[case] expected: [bool; 4]) { + let mut ctx = array_session().create_execution_ctx(); + let array = BoolArray::from_iter([true, false, true, true]).into_array(); + let constant = ConstantArray::new(true, 4).into_array(); + let result = execute_compare_test(array, constant, op); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); +} + +#[rstest] +// Inlined vs inlined, decided by prefix. +#[case("bad", "bat", Operator::Lt, true)] +// Inlined prefix tie decided by the tail bytes. +#[case("abcdefgh", "abcdefgi", Operator::Lt, true)] +// Inlined prefix and tail tie decided by length. +#[case("abc", "abcd", Operator::Lt, true)] +// Out-of-line values with equal prefixes. +#[case("aaaaaaaaaaaaaaaaaaaab", "aaaaaaaaaaaaaaaaaaaac", Operator::Lt, true)] +// Inlined vs out-of-line where one is a prefix of the other. +#[case("aaaa", "aaaaaaaaaaaaaaaaaaaa", Operator::Lt, true)] +// Equality across the inlined/out-of-line boundary. +#[case("aaaaaaaaaaaaaaaaaaaa", "aaaaaaaaaaaaaaaaaaaa", Operator::Eq, true)] +#[case("aaaaaaaaaaaaaaaaaaaa", "aaaaaaaaaaaaaaaaaaab", Operator::Eq, false)] +// Embedded NUL: "a\0" > "a" even though the padded prefixes tie. +#[case("a\0", "a", Operator::Gt, true)] +fn string_compare_cases( + #[case] lhs: &str, + #[case] rhs: &str, + #[case] op: Operator, + #[case] expected: bool, +) { + let mut ctx = array_session().create_execution_ctx(); + let lhs = VarBinViewArray::from_iter_str([lhs]).into_array(); + let rhs = VarBinViewArray::from_iter_str([rhs]).into_array(); + let result = execute_compare_test(lhs, rhs, op); + assert_arrays_eq!(result, BoolArray::from_iter([expected]), &mut ctx); +} + +#[test] +fn string_constant_compare() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter_str([ + "apple", + "banana", + "banan", + "bananarama-bananarama", + "cherry", + ]) + .into_array(); + let constant = ConstantArray::new(Scalar::from("banana"), 5).into_array(); + + let result = execute_compare_test(array.clone(), constant.clone(), Operator::Eq); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, true, false, false, false]), + &mut ctx + ); + + let result = execute_compare_test(array.clone(), constant.clone(), Operator::Lt); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, false, true, false, false]), + &mut ctx + ); + + let result = execute_compare_test(constant, array, Operator::Lt); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, false, false, true, true]), + &mut ctx + ); +} + +#[test] +fn string_constant_longer_than_inline() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter_str(["short", "averyveryverylongstring", "averyvery"]) + .into_array(); + let constant = ConstantArray::new(Scalar::from("averyveryverylongstring"), 3).into_array(); + + let result = execute_compare_test(array.clone(), constant.clone(), Operator::Eq); + assert_arrays_eq!(result, BoolArray::from_iter([false, true, false]), &mut ctx); + + let result = execute_compare_test(array, constant, Operator::Lte); + assert_arrays_eq!(result, BoolArray::from_iter([false, true, true]), &mut ctx); +} + +#[test] +fn decimal_compare() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let lhs = DecimalArray::from_iter::([100, 250, 300], dtype).into_array(); + let rhs = DecimalArray::from_iter::([200, 250, 100], dtype).into_array(); + + let result = execute_compare_test(lhs, rhs, Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); +} + +/// Two decimal arrays with the same logical dtype but different storage widths compare through +/// the widened common storage type. +#[test] +fn decimal_compare_mixed_storage_widths() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let lhs = DecimalArray::from_iter::([100, 250, 300], dtype).into_array(); + let rhs = DecimalArray::from_iter::([200, 250, 100], dtype).into_array(); + + let result = execute_compare_test(lhs, rhs, Operator::Lte); + assert_arrays_eq!(result, BoolArray::from_iter([true, true, false]), &mut ctx); +} + +/// A decimal constant that does not fit the array's narrow storage type still compares +/// correctly: it is greater than every representable array value. +#[test] +fn decimal_constant_out_of_storage_range() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(20, 2); + let array = DecimalArray::from_iter::([1, 50, 100], dtype).into_array(); + let constant = ConstantArray::new( + Scalar::decimal( + DecimalValue::from(10_000_000i64), + dtype, + Nullability::NonNullable, + ), + 3, + ) + .into_array(); + + let result = execute_compare_test(array.clone(), constant.clone(), Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([true, true, true]), &mut ctx); + + let result = execute_compare_test(array, constant, Operator::Gte); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, false, false]), + &mut ctx + ); +} + +/// Extension arrays compare through their storage values. +#[test] +fn extension_timestamp_compare() { + let mut ctx = array_session().create_execution_ctx(); + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); + let lhs = ExtensionArray::new(ext_dtype.clone(), buffer![1000i64, 2000, 3000].into_array()) + .into_array(); + let rhs = + ExtensionArray::new(ext_dtype, buffer![1500i64, 2000, 2500].into_array()).into_array(); + + let result = execute_compare_test(lhs.clone(), rhs.clone(), Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); + + let result = execute_compare_test(lhs, rhs, Operator::Eq); + assert_arrays_eq!(result, BoolArray::from_iter([false, true, false]), &mut ctx); +} + +/// Comparing extension arrays with different extension dtypes (e.g. timestamps in different +/// units) must error rather than silently comparing raw storage values. +#[test] +fn extension_mismatched_units_errors() { + let mut ctx = array_session().create_execution_ctx(); + let ms = ExtensionArray::new( + Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(), + buffer![1000i64, 2000].into_array(), + ) + .into_array(); + let secs = ExtensionArray::new( + Timestamp::new(TimeUnit::Seconds, Nullability::NonNullable).erased(), + buffer![1i64, 3].into_array(), + ) + .into_array(); + + let result = ms + .binary(secs, Operator::Eq) + .and_then(|a| a.execute::(&mut ctx)); + assert!(result.is_err()); +} + +/// Struct fields containing nulls order null-first, matching `Scalar` comparison semantics; +/// only top-level nulls make the result null. +#[test] +fn struct_compare_null_fields_order_first() { + let mut ctx = array_session().create_execution_ctx(); + let lhs = StructArray::from_fields(&[( + "a", + PrimitiveArray::from_option_iter([None, Some(5i32), None]).into_array(), + )]) + .unwrap() + .into_array(); + let rhs = StructArray::from_fields(&[( + "a", + PrimitiveArray::from_option_iter([Some(1i32), Some(5), None]).into_array(), + )]) + .unwrap() + .into_array(); + + // null field < non-null field; null == null. + let result = execute_compare_test(lhs.clone(), rhs.clone(), Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); + + let result = execute_compare_test(lhs, rhs, Operator::Eq); + assert_arrays_eq!(result, BoolArray::from_iter([false, true, true]), &mut ctx); +} + +#[test] +fn fixed_size_list_compare() { + let mut ctx = array_session().create_execution_ctx(); + let lhs = FixedSizeListArray::new( + buffer![1i32, 2, 3, 4, 5, 6].into_array(), + 2, + Validity::NonNullable, + 3, + ) + .into_array(); + let rhs = FixedSizeListArray::new( + buffer![1i32, 2, 3, 5, 4, 6].into_array(), + 2, + Validity::NonNullable, + 3, + ) + .into_array(); + + let result = execute_compare_test(lhs.clone(), rhs.clone(), Operator::Eq); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); + + let result = execute_compare_test(lhs, rhs, Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([false, true, false]), &mut ctx); +} + +/// Binary (non-utf8) comparison over VarBinView arrays. +#[test] +fn binary_compare() { + let mut ctx = array_session().create_execution_ctx(); + let lhs = VarBinViewArray::from_iter_bin([b"bad".as_slice(), b"\xff\x00", b""]).into_array(); + let rhs = VarBinViewArray::from_iter_bin([b"bat".as_slice(), b"\xff", b""]).into_array(); + + let result = execute_compare_test(lhs.clone(), rhs.clone(), Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); + + let result = execute_compare_test(lhs, rhs, Operator::Eq); + assert_arrays_eq!(result, BoolArray::from_iter([false, false, true]), &mut ctx); +} diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 156540a2820..361c4fab240 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -155,8 +155,8 @@ impl ScalarFnVTable for Binary { Operator::Lte => execute_compare(&lhs, &rhs, CompareOperator::Lte, ctx), Operator::Gt => execute_compare(&lhs, &rhs, CompareOperator::Gt, ctx), Operator::Gte => execute_compare(&lhs, &rhs, CompareOperator::Gte, ctx), - Operator::And => execute_boolean(&lhs, &rhs, Operator::And, ctx), - Operator::Or => execute_boolean(&lhs, &rhs, Operator::Or, ctx), + Operator::And => execute_boolean(lhs, rhs, Operator::And, ctx), + Operator::Or => execute_boolean(lhs, rhs, Operator::Or, ctx), Operator::Add => execute_numeric(&lhs, &rhs, NumericOperator::Add, ctx), Operator::Sub => execute_numeric(&lhs, &rhs, NumericOperator::Sub, ctx), Operator::Mul => execute_numeric(&lhs, &rhs, NumericOperator::Mul, ctx), diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index e95627155e9..ac76d2e4934 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric.rs @@ -212,7 +212,9 @@ where } } -enum PrimitiveOperand { +/// A primitive binary-operator operand: a materialized buffer, a non-null constant, or an +/// all-null constant. +pub(super) enum PrimitiveOperand { Array { values: Buffer, validity: Validity, @@ -226,7 +228,7 @@ enum PrimitiveOperand { } impl PrimitiveOperand { - fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + pub(super) 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::()? { @@ -250,14 +252,14 @@ impl PrimitiveOperand { Ok(Self::Array { values, validity }) } - fn len(&self) -> usize { + pub(super) fn len(&self) -> usize { match self { Self::Array { values, .. } => values.len(), Self::Constant { len, .. } | Self::Null(len) => *len, } } - fn validity(&self) -> Validity { + pub(super) fn validity(&self) -> Validity { match self { Self::Array { validity, .. } => validity.clone(), Self::Constant { validity, .. } => validity.clone(), diff --git a/vortex-array/src/scalar_fn/fns/like/mod.rs b/vortex-array/src/scalar_fn/fns/like/mod.rs index 63ced07c8e7..bdedce7112b 100644 --- a/vortex-array/src/scalar_fn/fns/like/mod.rs +++ b/vortex-array/src/scalar_fn/fns/like/mod.rs @@ -2,26 +2,36 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod kernel; +mod pattern; use std::borrow::Cow; use std::fmt::Display; use std::fmt::Formatter; pub use kernel::*; +use pattern::LikePattern; use prost::Message; +use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_err; use vortex_proto::expr as pb; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::ArrayRef; +use crate::Canonical; use crate::ExecutionCtx; -use crate::arrow::Datum; -use crate::arrow::from_arrow_columnar; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::VarBinViewArray; +use crate::arrays::varbinview::BinaryView; use crate::dtype::DType; +use crate::dtype::Nullability; use crate::expr::Expression; use crate::expr::and; +use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; @@ -140,7 +150,7 @@ impl ScalarFnVTable for Like { let child = args.get(0)?; let pattern = args.get(1)?; - arrow_like(&child, &pattern, *options, ctx) + execute_like(&child, &pattern, *options, ctx) } fn validity( @@ -163,34 +173,261 @@ impl ScalarFnVTable for Like { } } -/// Implementation of LIKE using the Arrow crate. -pub(crate) fn arrow_like( +/// Native implementation of LIKE over canonical [`VarBinViewArray`]s. +/// +/// Patterns are compiled once per distinct pattern (see [`LikePattern`]) and evaluated per +/// value directly on the view bytes. +pub(crate) fn execute_like( array: &ArrayRef, pattern: &ArrayRef, options: LikeOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - let nullable = array.dtype().is_nullable() | pattern.dtype().is_nullable(); - let len = array.len(); assert_eq!( array.len(), pattern.len(), - "Arrow Like: length mismatch for {}", + "LIKE: length mismatch for {}", array.encoding_id() ); + let len = array.len(); + let nullability = + Nullability::from(array.dtype().is_nullable() || pattern.dtype().is_nullable()); - // convert the pattern to the preferred array datatype - let lhs = Datum::try_new(array, ctx)?; - let rhs = Datum::try_new_with_target_datatype(pattern, lhs.data_type(), ctx)?; + if len == 0 { + return Ok(Canonical::empty(&DType::Bool(nullability)).into_array()); + } - let result = match (options.negated, options.case_insensitive) { - (false, false) => arrow_string::like::like(&lhs, &rhs)?, - (true, false) => arrow_string::like::nlike(&lhs, &rhs)?, - (false, true) => arrow_string::like::ilike(&lhs, &rhs)?, - (true, true) => arrow_string::like::nilike(&lhs, &rhs)?, - }; + if let Some(pattern_const) = pattern.as_constant() { + let Some(pattern_str) = pattern_const.as_utf8().value() else { + // A null pattern makes every row null. + return Ok( + ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) + .into_array(), + ); + }; + let values = array.clone().execute::(ctx)?; + let haystack = ResolvedViews::new(&values); + // The ASCII case-insensitive fast paths are only sound when the haystack is pure + // ASCII; see `LikePattern::compile`. + let ascii_haystack = + options.case_insensitive && pattern_str.as_str().is_ascii() && haystack.is_ascii(); + let compiled = LikePattern::compile( + pattern_str.as_str(), + options.case_insensitive, + ascii_haystack, + )?; + let bits = eval_pattern(&haystack, &compiled, options.negated); + let validity = values.validity()?.union_nullability(nullability); + return Ok(BoolArray::new(bits, validity).into_array()); + } - from_arrow_columnar(&result, len, nullable, ctx) + // Per-row patterns: compile each pattern individually. + let values = array.clone().execute::(ctx)?; + let patterns = pattern.clone().execute::(ctx)?; + let haystack = ResolvedViews::new(&values); + let pattern_views = ResolvedViews::new(&patterns); + let ascii_haystack = options.case_insensitive && haystack.is_ascii(); + + let mut bits = Vec::with_capacity(len); + // Reuse the previous row's compiled pattern while the pattern bytes repeat, so runs of + // identical patterns (the common case for non-constant pattern children) compile once. + let mut cached: Option<(&[u8], LikePattern)> = None; + for i in 0..len { + let pattern_bytes = pattern_views.bytes(i); + let compiled = match &cached { + Some((bytes, compiled)) if *bytes == pattern_bytes => compiled, + _ => { + let pattern_str = std::str::from_utf8(pattern_bytes) + .map_err(|e| vortex_err!("LIKE pattern is not valid UTF-8: {e}"))?; + let compiled = LikePattern::compile( + pattern_str, + options.case_insensitive, + ascii_haystack && pattern_str.is_ascii(), + )?; + &cached.insert((pattern_bytes, compiled)).1 + } + }; + bits.push(compiled.matches(haystack.bytes(i)) != options.negated); + } + let validity = values + .validity()? + .and(patterns.validity()?)? + .union_nullability(nullability); + Ok(BoolArray::new(BitBuffer::from_iter(bits), validity).into_array()) +} + +/// Resolved views over a canonical [`VarBinViewArray`]: the view structs plus borrowed slices +/// of every data buffer, supporting cheap per-element byte access. +struct ResolvedViews<'a> { + views: &'a [BinaryView], + buffers: Vec<&'a [u8]>, +} + +impl<'a> ResolvedViews<'a> { + fn new(array: &'a VarBinViewArray) -> Self { + Self { + views: array.views(), + buffers: (0..array.data_buffers().len()) + .map(|idx| array.buffer(idx).as_slice()) + .collect(), + } + } + + #[inline] + fn bytes(&self, index: usize) -> &'a [u8] { + let view = &self.views[index]; + if view.is_inlined() { + view.as_inlined().value() + } else { + let view = view.as_view(); + &self.buffers[view.buffer_index as usize][view.as_range()] + } + } + + /// Whether every value (including values under a null) is pure ASCII. + fn is_ascii(&self) -> bool { + (0..self.views.len()).all(|i| self.bytes(i).is_ascii()) + } + + /// The last `suffix_len` bytes of `view`, which must belong to this array. + /// + /// # Safety + /// + /// `suffix_len` must be at most `view.len()`. Buffer bounds are guaranteed by + /// [`VarBinViewArray::validate`], which checks every view against its data buffer at + /// construction. + #[inline] + unsafe fn suffix_bytes_unchecked(&self, view: &'a BinaryView, suffix_len: usize) -> &'a [u8] { + let len = view.len() as usize; + if view.is_inlined() { + // SAFETY: inlined values hold `len <= 12` value bytes, and the caller + // guarantees `suffix_len <= len`. + unsafe { view.as_inlined().value().get_unchecked(len - suffix_len..) } + } else { + let view = view.as_view(); + let end = view.offset as usize + len; + // SAFETY: validated views reference `buffer_index < buffers.len()` and bytes + // `offset..offset + len` within that buffer. + unsafe { + self.buffers + .get_unchecked(view.buffer_index as usize) + .get_unchecked(end - suffix_len..end) + } + } + } +} + +/// Evaluate `pattern` against every element of `haystack`. +/// +/// The equality, prefix, and suffix patterns exploit the view layout: a view stores the value +/// length and its first four bytes inline, which settles most elements without touching the +/// data buffers (values of up to 12 bytes are stored entirely inline). +fn eval_pattern(haystack: &ResolvedViews<'_>, pattern: &LikePattern, negated: bool) -> BitBuffer { + let len = haystack.views.len(); + match pattern { + LikePattern::Eq(needle) if needle.len() <= BinaryView::MAX_INLINED_SIZE => { + // The needle fits in a view, so equality is a single 16-byte comparison: a view + // of a different length or prefix can never share the same bit pattern. + let needle_view = BinaryView::new_inlined(needle).as_u128(); + BitBuffer::collect_bool(len, |i| { + (haystack.views[i].as_u128() == needle_view) != negated + }) + } + LikePattern::Eq(needle) => { + // Compare the view head (length plus 4-byte prefix) first; only views that agree + // on both dereference their data buffer for the remaining bytes. + let needle_head = needle_head(needle); + BitBuffer::collect_bool(len, |i| { + let view = &haystack.views[i]; + let matched = + view_head(view) == needle_head && haystack.bytes(i)[4..] == needle[4..]; + matched != negated + }) + } + LikePattern::StartsWith(needle) => { + // A branch-free masked comparison of the view's inline 4-byte prefix rejects + // almost every element without touching the data buffers; only elements whose + // prefix matches compare the remaining needle bytes. + let needle_len = needle.len(); + let prefix_len = needle_len.min(4); + let needle_prefix = u32::from_le_bytes({ + let mut padded = [0u8; 4]; + padded[..prefix_len].copy_from_slice(&needle[..prefix_len]); + padded + }); + let prefix_mask = if prefix_len == 4 { + u32::MAX + } else { + (1u32 << (8 * prefix_len)) - 1 + }; + BitBuffer::collect_bool(len, |i| { + let view = &haystack.views[i]; + let matched = view.len() as usize >= needle_len + && (view_prefix(view) & prefix_mask) == needle_prefix + && (needle_len <= 4 || haystack.bytes(i)[4..needle_len] == needle[4..]); + matched != negated + }) + } + LikePattern::EndsWith(needle) => { + // Inlined values compare their suffix inside the view struct without touching the + // data buffers; reference views slice exactly the suffix out of their buffer. + let needle_len = needle.len(); + BitBuffer::collect_bool(len, |i| { + // SAFETY: `i` is below the array length, and the suffix length is only read + // once the view is known to be at least `needle_len` long. + let matched = unsafe { + let view = haystack.views.get_unchecked(i); + view.len() as usize >= needle_len + && bytes_eq(haystack.suffix_bytes_unchecked(view, needle_len), needle) + }; + matched != negated + }) + } + LikePattern::IEqAscii(needle) => BitBuffer::collect_bool(len, |i| { + let view = &haystack.views[i]; + let matched = view.len() as usize == needle.len() + && haystack.bytes(i).eq_ignore_ascii_case(needle); + matched != negated + }), + LikePattern::Contains(finder, needle_len) => BitBuffer::collect_bool(len, |i| { + let view = &haystack.views[i]; + let matched = + view.len() as usize >= *needle_len && finder.find(haystack.bytes(i)).is_some(); + matched != negated + }), + _ => BitBuffer::collect_bool(len, |i| pattern.matches(haystack.bytes(i)) != negated), + } +} + +/// Byte equality as an inlined loop with early exit. +/// +/// Faster than slice `==` for the short, unpredictable-length comparisons in this module, +/// which otherwise lower to a `memcmp` libcall per element. +#[inline] +fn bytes_eq(lhs: &[u8], rhs: &[u8]) -> bool { + lhs.len() == rhs.len() && std::iter::zip(lhs, rhs).all(|(l, r)| l == r) +} + +/// The leading 8 bytes of a view: the `u32` length plus the first 4 bytes of the value +/// (zero-padded for values shorter than 4 bytes). +#[inline] +#[expect(clippy::cast_possible_truncation, reason = "intentional bit slicing")] +fn view_head(view: &BinaryView) -> u64 { + view.as_u128() as u64 +} + +/// The view head a needle of more than 4 bytes would have: its length plus first 4 bytes. +fn needle_head(needle: &[u8]) -> u64 { + let prefix: [u8; 4] = [needle[0], needle[1], needle[2], needle[3]]; + (needle.len() as u64) | (u64::from(u32::from_le_bytes(prefix)) << 32) +} + +/// The first 4 value bytes stored inline in any view (zero-padded for values shorter than +/// 4 bytes), as a raw little-endian `u32` in memory order. +#[inline] +#[expect(clippy::cast_possible_truncation, reason = "intentional bit slicing")] +fn view_prefix(view: &BinaryView) -> u32 { + (view.as_u128() >> 32) as u32 } /// Variants of the LIKE filter that we know how to turn into a stats pruning predicate. @@ -243,10 +480,16 @@ impl<'a> LikeVariant<'a> { mod tests { use std::borrow::Cow; + use rstest::rstest; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; + use crate::arrays::ConstantArray; + use crate::arrays::VarBinArray; + use crate::arrays::VarBinViewArray; + use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::Nullability; @@ -256,8 +499,205 @@ mod tests { use crate::expr::not; use crate::expr::not_ilike; use crate::expr::root; + use crate::scalar::Scalar; + use crate::scalar_fn::fns::like::Like; + use crate::scalar_fn::fns::like::LikeOptions; use crate::scalar_fn::fns::like::LikeVariant; + fn run_like( + array: crate::ArrayRef, + pattern: crate::ArrayRef, + options: LikeOptions, + ) -> crate::ArrayRef { + let len = array.len(); + Like.try_new_array(len, options, [array, pattern]).unwrap() + } + + #[rstest] + // Exact, prefix, suffix, contains fast paths. + #[case("hello", [true, false, false, false])] + #[case("he%", [true, false, true, false])] + #[case("%llo", [true, false, false, true])] + #[case("%ell%", [true, false, false, true])] + // Wildcards that require the regex path. + #[case("h_llo", [true, false, false, false])] + #[case("h%o", [true, false, false, false])] + #[case("%", [true, true, true, true])] + #[case("_____", [true, true, false, true])] + fn test_like_patterns(#[case] pattern: &str, #[case] expected: [bool; 4]) { + let mut ctx = array_session().create_execution_ctx(); + let array = + VarBinViewArray::from_iter_str(["hello", "world", "help", "jello"]).into_array(); + let result = run_like( + array, + ConstantArray::new(pattern, 4).into_array(), + LikeOptions::default(), + ); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); + } + + #[test] + fn test_like_escapes() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter_str(["100%", "100x", "a_b", "axb"]).into_array(); + + let result = run_like( + array.clone(), + ConstantArray::new(r"100\%", 4).into_array(), + LikeOptions::default(), + ); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, false, false, false]), + &mut ctx + ); + + let result = run_like( + array, + ConstantArray::new(r"a\_b", 4).into_array(), + LikeOptions::default(), + ); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, false, true, false]), + &mut ctx + ); + } + + #[test] + fn test_like_regex_meta_characters_are_literal() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter_str(["a.c", "abc", "a$c"]).into_array(); + let result = run_like( + array, + ConstantArray::new("a.%", 3).into_array(), + LikeOptions::default(), + ); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); + } + + #[test] + fn test_like_unicode() { + let mut ctx = array_session().create_execution_ctx(); + let array = + VarBinViewArray::from_iter_str(["h\u{00a3}llo", "hxllo", "h\u{00a3}xllo"]).into_array(); + // `_` matches exactly one character of any byte width. + let result = run_like( + array, + ConstantArray::new("h_llo", 3).into_array(), + LikeOptions::default(), + ); + assert_arrays_eq!(result, BoolArray::from_iter([true, true, false]), &mut ctx); + } + + #[test] + fn test_nlike() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter_str(["hello", "world"]).into_array(); + let result = run_like( + array, + ConstantArray::new("he%", 2).into_array(), + LikeOptions { + negated: true, + case_insensitive: false, + }, + ); + assert_arrays_eq!(result, BoolArray::from_iter([false, true]), &mut ctx); + } + + #[test] + fn test_ilike() { + let mut ctx = array_session().create_execution_ctx(); + let ilike = LikeOptions { + negated: false, + case_insensitive: true, + }; + + // Pure-ASCII data takes the ASCII fast paths. + let array = VarBinViewArray::from_iter_str(["HELLO", "world", "Help"]).into_array(); + let result = run_like(array, ConstantArray::new("he%", 3).into_array(), ilike); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, true]), &mut ctx); + + // Non-ASCII data requires full case folding: `k` matches U+212A KELVIN SIGN. + let array = VarBinViewArray::from_iter_str(["\u{212a}", "k", "x"]).into_array(); + let result = run_like(array, ConstantArray::new("k", 3).into_array(), ilike); + assert_arrays_eq!(result, BoolArray::from_iter([true, true, false]), &mut ctx); + } + + #[test] + fn test_nilike() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter_str(["HELLO", "world"]).into_array(); + let result = run_like( + array, + ConstantArray::new("he%", 2).into_array(), + LikeOptions { + negated: true, + case_insensitive: true, + }, + ); + assert_arrays_eq!(result, BoolArray::from_iter([false, true]), &mut ctx); + } + + #[test] + fn test_like_nullable_input() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter_nullable_str([Some("hello"), None, Some("help")]) + .into_array(); + let result = run_like( + array, + ConstantArray::new("he%", 3).into_array(), + LikeOptions::default(), + ); + assert_arrays_eq!( + result, + BoolArray::from_iter([Some(true), None, Some(true)]), + &mut ctx + ); + } + + #[test] + fn test_like_null_pattern() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter_str(["hello", "world"]).into_array(); + let result = run_like( + array, + ConstantArray::new(Scalar::null(DType::Utf8(Nullability::Nullable)), 2).into_array(), + LikeOptions::default(), + ); + assert_arrays_eq!(result, BoolArray::from_iter([None, None]), &mut ctx); + } + + #[test] + fn test_like_per_row_patterns() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter_str(["hello", "hello", "hello"]).into_array(); + let patterns = VarBinViewArray::from_iter_str(["he%", "%world", "h_llo"]).into_array(); + let result = run_like(array, patterns, LikeOptions::default()); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, true]), &mut ctx); + } + + #[test] + fn test_like_non_canonical_input() { + let mut ctx = array_session().create_execution_ctx(); + // VarBin (not the canonical VarBinView) is canonicalized before evaluation. + let array = VarBinArray::from_iter( + [Some("hello"), Some("world")], + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let result = run_like( + array, + ConstantArray::new("he%", 2).into_array(), + LikeOptions::default(), + ); + assert_arrays_eq!( + result, + BoolArray::from_iter([Some(true), Some(false)]), + &mut ctx + ); + } + #[test] fn invert_booleans() { let not_expr = not(root()); diff --git a/vortex-array/src/scalar_fn/fns/like/pattern.rs b/vortex-array/src/scalar_fn/fns/like/pattern.rs new file mode 100644 index 00000000000..8a667159ecb --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/like/pattern.rs @@ -0,0 +1,279 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compiled SQL LIKE patterns evaluated over UTF-8 byte slices. +//! +//! A [`LikePattern`] is compiled once per pattern and then evaluated per value. Patterns that +//! reduce to plain equality, prefix, suffix, or substring searches use direct byte comparisons +//! (with `memchr`'s SIMD substring search for the contains case); everything else falls back to +//! a regex translated from the LIKE pattern. +//! +//! The pattern grammar and its regex translation mirror the `arrow-string` LIKE kernels so that +//! results are identical to the previous Arrow-backed implementation: `%` matches any sequence +//! of characters (including across newlines), `_` matches exactly one character, and `\` escapes +//! the next character. + +use memchr::memchr3; +use memchr::memmem::Finder; +use regex::bytes::Regex; +use regex::bytes::RegexBuilder; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +/// A LIKE pattern compiled for repeated evaluation against UTF-8 haystacks. +pub(crate) enum LikePattern { + /// The pattern has no wildcards: plain byte equality. + Eq(Vec), + /// `prefix%`: byte prefix comparison. + StartsWith(Vec), + /// `%suffix`: byte suffix comparison. + EndsWith(Vec), + /// `%needle%`: SIMD substring search, along with the needle length. + /// + /// The searcher is boxed to keep this variant close in size to the others. + Contains(Box>, usize), + /// Equality ignoring ASCII case (ILIKE over ASCII-only data). + IEqAscii(Vec), + /// Prefix comparison ignoring ASCII case (ILIKE over ASCII-only data). + IStartsWithAscii(Vec), + /// Suffix comparison ignoring ASCII case (ILIKE over ASCII-only data). + IEndsWithAscii(Vec), + /// Everything else: the LIKE pattern translated to an anchored regex. + Regex(Regex), +} + +impl LikePattern { + /// Compile `pattern`, with `case_insensitive` selecting ILIKE semantics. + /// + /// `ascii_haystack` enables the ASCII-only case-insensitive fast paths and must only be + /// `true` when every haystack this pattern will be evaluated against is pure ASCII: + /// full case folding can match ASCII pattern characters against non-ASCII haystack + /// characters (e.g. `k` folds to U+212A KELVIN SIGN), which the ASCII fast paths would + /// miss. + pub(crate) fn compile( + pattern: &str, + case_insensitive: bool, + ascii_haystack: bool, + ) -> VortexResult { + if case_insensitive { + if ascii_haystack && pattern.is_ascii() { + if !contains_like_pattern(pattern) { + return Ok(Self::IEqAscii(pattern.as_bytes().to_vec())); + } else if pattern.ends_with('%') + && !contains_like_pattern(&pattern[..pattern.len() - 1]) + { + return Ok(Self::IStartsWithAscii( + pattern.as_bytes()[..pattern.len() - 1].to_vec(), + )); + } else if pattern.starts_with('%') && !contains_like_pattern(&pattern[1..]) { + return Ok(Self::IEndsWithAscii(pattern.as_bytes()[1..].to_vec())); + } + } + return Ok(Self::Regex(regex_like(pattern, true)?)); + } + + Ok(if !contains_like_pattern(pattern) { + Self::Eq(pattern.as_bytes().to_vec()) + } else if pattern.ends_with('%') && !contains_like_pattern(&pattern[..pattern.len() - 1]) { + Self::StartsWith(pattern.as_bytes()[..pattern.len() - 1].to_vec()) + } else if pattern.starts_with('%') && !contains_like_pattern(&pattern[1..]) { + Self::EndsWith(pattern.as_bytes()[1..].to_vec()) + } else if pattern.starts_with('%') + && pattern.ends_with('%') + && !contains_like_pattern(&pattern[1..pattern.len() - 1]) + { + Self::Contains( + Box::new(Finder::new(&pattern.as_bytes()[1..pattern.len() - 1]).into_owned()), + pattern.len() - 2, + ) + } else { + Self::Regex(regex_like(pattern, false)?) + }) + } + + /// Evaluate this pattern against `haystack`, which must be valid UTF-8. + #[inline] + pub(crate) fn matches(&self, haystack: &[u8]) -> bool { + match self { + Self::Eq(needle) => needle == haystack, + Self::StartsWith(needle) => { + needle.len() <= haystack.len() && &haystack[..needle.len()] == needle.as_slice() + } + Self::EndsWith(needle) => { + needle.len() <= haystack.len() + && &haystack[haystack.len() - needle.len()..] == needle.as_slice() + } + Self::Contains(finder, needle_len) => { + haystack.len() >= *needle_len && finder.find(haystack).is_some() + } + Self::IEqAscii(needle) => haystack.eq_ignore_ascii_case(needle), + Self::IStartsWithAscii(needle) => { + needle.len() <= haystack.len() + && haystack[..needle.len()].eq_ignore_ascii_case(needle) + } + Self::IEndsWithAscii(needle) => { + needle.len() <= haystack.len() + && haystack[haystack.len() - needle.len()..].eq_ignore_ascii_case(needle) + } + Self::Regex(regex) => regex.is_match(haystack), + } + } +} + +/// Returns whether `pattern` contains a wildcard (`%`, `_`) or an escape (`\`). +fn contains_like_pattern(pattern: &str) -> bool { + memchr3(b'%', b'_', b'\\', pattern.as_bytes()).is_some() +} + +/// Transforms a LIKE `pattern` into an equivalent anchored regex. +/// +/// This is a port of the `arrow-string` translation so that pattern semantics are unchanged: +/// +/// 1. `%` => `.*` (a leading/trailing `%` truncates the anchor instead, e.g. `%foo%` => `foo` +/// rather than `^.*foo.*$`); +/// 2. `_` => `.`; +/// 3. regex meta characters are escaped so they match literally; +/// 4. `\x` matches `x` literally (a trailing `\` matches a literal backslash). +/// +/// The regex runs over the haystack bytes directly (`regex::bytes`) in Unicode mode, which +/// matches identically to a `&str` regex on valid UTF-8 input. +fn regex_like(pattern: &str, case_insensitive: bool) -> VortexResult { + let mut result = String::with_capacity(pattern.len() * 2); + let mut chars_iter = pattern.chars().peekable(); + match chars_iter.peek() { + // If the pattern starts with `%`, avoid starting the regex with a slow but + // meaningless `^.*`. + Some('%') => { + chars_iter.next(); + } + _ => result.push('^'), + }; + + while let Some(c) = chars_iter.next() { + match c { + '\\' => { + match chars_iter.peek() { + Some(&next) => { + if regex_syntax::is_meta_character(next) { + result.push('\\'); + } + result.push(next); + // Skip the next char as it is already appended. + chars_iter.next(); + } + None => { + // A trailing backslash matches a literal backslash. + result.push('\\'); + result.push('\\'); + } + } + } + '%' => result.push_str(".*"), + '_' => result.push('.'), + c => { + if regex_syntax::is_meta_character(c) { + result.push('\\'); + } + result.push(c); + } + } + } + // Instead of ending the regex with `.*$` and making it needlessly slow, just end it. + if result.ends_with(".*") { + result.pop(); + result.pop(); + } else { + result.push('$'); + } + RegexBuilder::new(&result) + .case_insensitive(case_insensitive) + .dot_matches_new_line(true) + .build() + .map_err(|e| vortex_err!("Unable to build regex from LIKE pattern: {e}")) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn like(pattern: &str) -> LikePattern { + LikePattern::compile(pattern, false, false).unwrap() + } + + #[test] + fn compile_fast_paths() { + assert!(matches!(like("foo"), LikePattern::Eq(_))); + assert!(matches!(like("foo%"), LikePattern::StartsWith(_))); + assert!(matches!(like("%foo"), LikePattern::EndsWith(_))); + assert!(matches!(like("%foo%"), LikePattern::Contains(..))); + assert!(matches!(like("%"), LikePattern::StartsWith(_))); + assert!(matches!(like("%%"), LikePattern::Contains(..))); + assert!(matches!(like("f%o"), LikePattern::Regex(_))); + assert!(matches!(like("f_o"), LikePattern::Regex(_))); + assert!(matches!(like(r"foo\%"), LikePattern::Regex(_))); + } + + #[test] + fn regex_translation() { + // (pattern, expected regex) pairs from the arrow-string translation. + let cases = [ + (r"%foobar%", r"foobar"), + (r"foo%bar", r"^foo.*bar$"), + (r"foo_bar", r"^foo.bar$"), + (r"\%\_", r"^%_$"), + (r"\a", r"^a$"), + (r"\\%", r"^\\"), + (r"\\a", r"^\\a$"), + (r".", r"^\.$"), + (r"$", r"^\$$"), + (r"\\", r"^\\$"), + ]; + for (pattern, expected) in cases { + assert_eq!(regex_like(pattern, false).unwrap().to_string(), expected); + } + } + + #[test] + fn matches_wildcards() { + assert!(like("h_llo w%d").matches(b"hello world")); + assert!(like("h_llo w%d").matches("h\u{00a3}llo wd".as_bytes())); + assert!(!like("h_llo w%d").matches(b"hxxllo world")); + // `%` crosses newlines. + assert!(like("hello%world").matches(b"hello\nworld")); + // `_` matches exactly one character, of any width. + assert!(like("_").matches("\u{00a3}".as_bytes())); + assert!(!like("_").matches("\u{00a3}x".as_bytes())); + } + + #[test] + fn matches_escapes() { + assert!(like(r"100\%").matches(b"100%")); + assert!(!like(r"100\%").matches(b"100x")); + assert!(like(r"a\_b").matches(b"a_b")); + assert!(!like(r"a\_b").matches(b"axb")); + // A trailing backslash matches a literal backslash. + assert!(like("trailing\\").matches(b"trailing\\")); + } + + #[test] + fn matches_case_insensitive() { + let ilike = |pattern: &str| LikePattern::compile(pattern, true, false).unwrap(); + assert!(ilike("hello%").matches(b"HELLO WORLD")); + assert!(ilike("%WORLD").matches(b"hello world")); + // Full case folding: the ASCII pattern `k` matches U+212A KELVIN SIGN. + assert!(ilike("k").matches("\u{212a}".as_bytes())); + // Greek sigma case folds across all three forms. + assert!(ilike("\u{03c3}").matches("\u{03a3}".as_bytes())); + assert!(ilike("\u{03c3}").matches("\u{03c2}".as_bytes())); + + // The ASCII fast paths agree with the regex on ASCII haystacks. + let ascii = |pattern: &str| LikePattern::compile(pattern, true, true).unwrap(); + assert!(matches!(ascii("abc"), LikePattern::IEqAscii(_))); + assert!(matches!(ascii("abc%"), LikePattern::IStartsWithAscii(_))); + assert!(matches!(ascii("%abc"), LikePattern::IEndsWithAscii(_))); + assert!(ascii("abc").matches(b"AbC")); + assert!(ascii("abc%").matches(b"ABCDEF")); + assert!(ascii("%def").matches(b"ABCDEF")); + assert!(!ascii("abc").matches(b"AbCd")); + } +} 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 5185e513e70..edf530cb75f 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -400,6 +400,7 @@ mod tests { use rstest::rstest; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -410,8 +411,6 @@ mod tests { 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::Nullability; use crate::dtype::PType::I32; @@ -666,15 +665,14 @@ mod tests { // -- Tests migrated from compute/list_contains.rs -- fn nonnull_strings(values: Vec>) -> ArrayRef { - #[expect(deprecated)] - let result = ListArray::from_iter_slow::( - values, - Arc::new(DType::Utf8(Nullability::NonNullable)), - ) - .unwrap() - .to_listview() - .into_array(); - result + let mut ctx = array_session().create_execution_ctx(); + + ListArray::from_iter_slow::(values, Arc::new(DType::Utf8(Nullability::NonNullable))) + .unwrap() + .into_array() + .execute::(&mut ctx) + .vortex_expect("failed to convert to listview") + .into_array() } fn null_strings(values: Vec>>) -> ArrayRef { @@ -693,13 +691,15 @@ mod tests { let elements = VarBinArray::from_iter(elements, DType::Utf8(Nullability::Nullable)).into_array(); - #[expect(deprecated)] - let result = ListArray::try_new(elements, offsets, Validity::NonNullable) + let mut ctx = array_session().create_execution_ctx(); + + ListArray::try_new(elements, offsets, Validity::NonNullable) .unwrap() .as_array() - .to_listview() - .into_array(); - result + .clone() + .execute::(&mut ctx) + .vortex_expect("failed to convert to listview") + .into_array() } fn bool_array(values: Vec, validity: Validity) -> BoolArray { diff --git a/vortex-array/src/scalar_fn/fns/list_sum.rs b/vortex-array/src/scalar_fn/fns/list_sum.rs new file mode 100644 index 00000000000..a5f8d219db6 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_sum.rs @@ -0,0 +1,585 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::Canonical; +use crate::Columnar; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::aggregate_fn::AggregateFnVTable; +use crate::aggregate_fn::DynGroupedAccumulator; +use crate::aggregate_fn::GroupRanges; +use crate::aggregate_fn::GroupedAccumulator; +use crate::aggregate_fn::GroupedArray; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::sum::Sum; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::FixedSizeList; +use crate::arrays::ListView; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +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::validity::Validity; + +/// Sum of the elements in each list of a `List` or `FixedSizeList` typed array. +/// +/// Follows SQL `SUM` semantics per list, matching DuckDB's `list_sum`: null lists, empty +/// lists, and lists whose elements are all null yield a null sum; null elements are skipped. +/// Integer and decimal overflow yields a null sum value, matching [`Sum`]. The result dtype +/// follows [`Sum`]'s widening rules and is always nullable. +/// +/// NaN handling for float elements is controlled by [`NumericalAggregateOpts`]: with +/// `skip_nans` (the default) NaN values contribute nothing, otherwise any NaN poisons the +/// list's sum to NaN. +#[derive(Clone)] +pub struct ListSum; + +impl ScalarFnVTable for ListSum { + type Options = NumericalAggregateOpts; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.list.sum"); + *ID + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.serialize())) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + NumericalAggregateOpts::deserialize(metadata) + } + + 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 list_sum()"), + } + } + + fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + let elem_dtype = match &arg_dtypes[0] { + DType::List(elem, _) | DType::FixedSizeList(elem, ..) => elem.as_ref(), + other => vortex_bail!("list_sum() requires List or FixedSizeList, got {other}"), + }; + Sum.return_dtype(options, elem_dtype) + .ok_or_else(|| vortex_err!("list_sum() cannot sum elements of type {elem_dtype}")) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + + let elem_dtype = match input.dtype() { + DType::List(elem, _) | DType::FixedSizeList(elem, ..) => elem.as_ref().clone(), + other => vortex_bail!("list_sum() requires List or FixedSizeList, got {other}"), + }; + + // `mask_empty_lists` needs access to list elements validity and sizes + let columnar = input.execute::(ctx)?; + + match columnar { + Columnar::Constant(constant) => { + // Canonicalize one row of the constant and broadcast its sum. + let one_row = ConstantArray::new(constant.scalar().clone(), 1) + .into_array() + .execute::(ctx)? + .into_array(); + let sum = + list_sum_impl(one_row, elem_dtype, options, ctx)?.execute_scalar(0, ctx)?; + Ok(ConstantArray::new(sum, constant.len()).into_array()) + } + Columnar::Canonical(canonical) => { + list_sum_impl(canonical.into_array(), elem_dtype, options, ctx) + } + } + } + + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { + false + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +/// Sum each list of a canonical `array` into one value per list. +/// +/// Note that we need to nullify sums produced by empty or all-null lists, +/// since grouped sum kernels default to 0 for these. +fn list_sum_impl( + canonical: ArrayRef, + elem_dtype: DType, + options: &NumericalAggregateOpts, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut acc = GroupedAccumulator::try_new(Sum, *options, elem_dtype)?; + acc.accumulate_list(&canonical, ctx)?; + let sums = acc.finish()?; + + let grouped: GroupedArray = if let Some(fsl) = canonical.as_opt::() { + fsl.into_owned().into() + } else if let Some(lv) = canonical.as_opt::() { + lv.into_owned().into() + } else { + let dtype = canonical.dtype(); + vortex_bail!("list_sum() requires List or FixedSizeList but got {dtype}") + }; + + mask_empty_lists(grouped, sums, ctx) +} + +/// Applies a mask to `sums` that nullifies entries produced by lists without at least +/// one valid element. This is necessary because the grouped `Sum` aggregate only produces +/// nulls for null lists and sums that overflow. +fn mask_empty_lists( + grouped: GroupedArray, + sums: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let elements = grouped.elements(); + let elem_mask = elements.validity()?.execute_mask(elements.len(), ctx)?; + let ranges = grouped.group_ranges(ctx)?; + + let has_valid_element: BitBuffer = match elem_mask.bit_buffer() { + AllOr::All => match &ranges { + // fixed-size lists of non-zero width cannot have empty lists. + GroupRanges::FixedSizeList { size, .. } if *size > 0 => return Ok(sums), + GroupRanges::FixedSizeList { len, .. } => BitBuffer::full(false, *len), + GroupRanges::ListView { ranges } => ranges.iter().map(|&(_, size)| size > 0).collect(), + }, + AllOr::None => BitBuffer::full(false, ranges.len()), + AllOr::Some(bits) => ranges + .iter() + .map(|(offset, size)| size > 0 && bits.count_range(offset, offset + size) > 0) + .collect(), + }; + if has_valid_element.true_count() == has_valid_element.len() { + return Ok(sums); + } + + sums.mask(BoolArray::new(has_valid_element, Validity::NonNullable).into_array()) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use prost::Message; + use rstest::rstest; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_proto::expr as pb; + + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::aggregate_fn::NumericalAggregateOpts; + 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::Expression; + use crate::expr::list_sum; + use crate::expr::list_sum_opts; + use crate::expr::proto::ExprSerializeProtoExt; + use crate::expr::root; + use crate::scalar::Scalar; + use crate::scalar_fn::ScalarFnVTable; + use crate::scalar_fn::fns::list_sum::ListSum; + 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_sum(#[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_sum(root()))?; + let mut ctx = array_session().create_execution_ctx(); + // [1, 2] = 3; [3, 4, 5] = 12; [] = null; [6, null] = 6 (nulls skipped). + let expected = + PrimitiveArray::from_option_iter::([Some(3), Some(12), None, Some(6)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_nullable_list_sum() -> VortexResult<()> { + let elements = create_list_elements(); + let list = ListArray::try_new( + elements, + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::Array(BoolArray::from_iter([true, false, true, false]).into_array()), + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + // Row 1 and 3 are null lists; row 2 is a valid but empty list. + let expected = PrimitiveArray::from_option_iter::([Some(3), None, None, None]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_all_null_elements_sum_to_null() -> VortexResult<()> { + let elements = PrimitiveArray::from_option_iter::([None, None, Some(1)]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 2, 3].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([None, Some(1)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_overflow_sums_to_null() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter([i64::MAX, 1, 1, 2]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([None, Some(3)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_unsigned_widening() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter([1u8, 2, 3]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 3].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(6)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_bool_elements() -> VortexResult<()> { + let elements = BoolArray::from_iter([true, true, false, true]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 3, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(2), Some(1)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_nan_skipped_by_default() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter([1.0f64, f64::NAN, 2.0]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 3].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(3.0)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_all_nan_list_sums_to_zero() -> VortexResult<()> { + // NaN elements are *valid*: with the default `skip_nans` they contribute nothing, but + // the list still has valid elements, so the sum is `0.0` rather than null (unlike an + // all-null list). + let elements = PrimitiveArray::from_iter([f64::NAN, f64::NAN]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 2].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(0.0)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_nan_poisons_with_include_nans() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter([1.0f64, f64::NAN, 2.0]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32, 3].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum_opts( + root(), + NumericalAggregateOpts::include_nans(), + ))?; + + let mut ctx = array_session().create_execution_ctx(); + assert!(result.is_valid(0, &mut ctx)?); + let result = result.execute::(&mut ctx)?; + assert!(result.as_slice::()[0].is_nan()); + Ok(()) + } + + #[test] + fn test_listview_sum() -> VortexResult<()> { + let elements = create_list_elements(); + // Overlapping and out-of-order views over the shared 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_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + // [6, null] = 6; [1, 2, 3] = 6; [] = null; [2, 3] = 5. + let expected = + PrimitiveArray::from_option_iter::([Some(6), Some(6), None, Some(5)]); + assert_arrays_eq!(result, expected, &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_sum() -> VortexResult<()> { + let fsl = create_fixed_size_list(Validity::NonNullable); + let result = fsl.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = + PrimitiveArray::from_option_iter::([Some(3), Some(7), Some(11), Some(15)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_fixed_size_list_sum_nullable() -> VortexResult<()> { + let fsl = create_fixed_size_list(Validity::Array( + BoolArray::from_iter([true, false, true, false]).into_array(), + )); + let result = fsl.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(3), None, Some(11), None]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_list_sum_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_sum(root()))?; + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(6), Some(3), None]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_list_sum_slice() -> 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 sliced = list.slice(1..4)?; + + let result = sliced.apply(&list_sum(root()))?; + let mut ctx = array_session().create_execution_ctx(); + let expected = PrimitiveArray::from_option_iter::([Some(12), None, Some(6)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_empty_array() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter([0i32; 0]); + let list = ListArray::try_new( + elements.into_array(), + buffer![0u32].into_array(), + Validity::NonNullable, + )? + .into_array(); + let result = list.apply(&list_sum(root()))?; + + let mut ctx = array_session().create_execution_ctx(); + let result = result.execute::(&mut ctx)?; + assert_eq!(result.len(), 0); + Ok(()) + } + + #[test] + fn test_constant_list_sum() -> 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 mut ctx = array_session().create_execution_ctx(); + let scalar = list.execute_scalar(1, &mut ctx)?; + + let constant = ConstantArray::new(scalar, 3).into_array(); + let result = constant.apply(&list_sum(root()))?; + let expected = PrimitiveArray::from_option_iter::([Some(12), Some(12), Some(12)]); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_null_scalar_list_sum() -> 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_sum(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_unsupported_dtypes() { + let opts = NumericalAggregateOpts::default(); + // Non-list inputs are rejected. + assert!( + ListSum + .return_dtype( + &opts, + &[DType::Primitive(PType::I32, Nullability::NonNullable)] + ) + .is_err() + ); + // Non-numeric element types are rejected. + assert!( + ListSum + .return_dtype( + &opts, + &[DType::List( + Arc::new(DType::Utf8(Nullability::NonNullable)), + Nullability::NonNullable, + )], + ) + .is_err() + ); + } + + #[test] + fn test_display() { + assert_eq!(list_sum(root()).to_string(), "vortex.list.sum($)"); + assert_eq!( + list_sum_opts(root(), NumericalAggregateOpts::include_nans()).to_string(), + "vortex.list.sum($, opts=skip_nans=false)" + ); + } + + #[test] + fn test_proto_round_trip() -> VortexResult<()> { + for expr in [ + list_sum(root()), + list_sum_opts(root(), NumericalAggregateOpts::include_nans()), + ] { + let proto = expr.serialize_proto()?; + let buf = proto.encode_to_vec(); + let decoded = pb::Expr::decode(buf.as_slice())?; + let deser = Expression::from_proto(&decoded, &array_session())?; + assert_eq!(expr, deser); + } + Ok(()) + } +} diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index c9f5783de27..7b9d57fcdca 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -277,8 +277,6 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::PrimitiveArray; @@ -300,21 +298,27 @@ mod tests { use crate::scalar_fn::fns::pack::Pack; fn primitive_field(array: &ArrayRef, field_path: &[&str]) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); let mut field_path = field_path.iter(); let Some(field) = field_path.next() else { vortex_bail!("empty field path"); }; - #[expect(deprecated)] - let mut array = array.to_struct().unmasked_field_by_name(field)?.clone(); + let mut array = array + .clone() + .execute::(&mut ctx)? + .unmasked_field_by_name(field)? + .clone(); for field in field_path { - #[expect(deprecated)] - let next = array.to_struct().unmasked_field_by_name(field)?.clone(); + let next = array + .clone() + .execute::(&mut ctx)? + .unmasked_field_by_name(field)? + .clone(); array = next; } - #[expect(deprecated)] - let result = array.to_primitive(); + let result = array.execute::(&mut ctx)?; Ok(result) } @@ -364,7 +368,7 @@ mod tests { let actual_array = test_array.apply(&expr).unwrap(); assert_eq!( - actual_array.as_struct_typed().names(), + actual_array.dtype().as_struct_fields().names(), ["a", "b", "c", "d", "e"] ); @@ -450,7 +454,7 @@ mod tests { .into_array(); let actual_array = test_array.clone().apply(&expr).unwrap(); assert_eq!(actual_array.len(), test_array.len()); - assert_eq!(actual_array.as_struct_typed().nfields(), 0); + assert_eq!(actual_array.nchildren(), 0); } #[test] @@ -491,14 +495,19 @@ mod tests { ]) .unwrap() .into_array(); - #[expect(deprecated)] - let actual_array = test_array.apply(&expr).unwrap().to_struct(); + let mut ctx = array_session().create_execution_ctx(); + let actual_array = test_array + .apply(&expr) + .unwrap() + .execute::(&mut ctx) + .unwrap(); - #[expect(deprecated)] let inner_struct = actual_array .unmasked_field_by_name("a") .unwrap() - .to_struct(); + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!( inner_struct .names() @@ -511,6 +520,7 @@ mod tests { #[test] pub fn test_merge_order() { + let mut ctx = array_session().create_execution_ctx(); let expr = merge(vec![get_item("0", root()), get_item("1", root())]); let test_array = StructArray::from_fields(&[ @@ -535,8 +545,11 @@ mod tests { ]) .unwrap() .into_array(); - #[expect(deprecated)] - let actual_array = test_array.apply(&expr).unwrap().to_struct(); + let actual_array = test_array + .apply(&expr) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_eq!(actual_array.names(), ["a", "c", "b", "d"]); } diff --git a/vortex-array/src/scalar_fn/fns/mod.rs b/vortex-array/src/scalar_fn/fns/mod.rs index f7e98e676b0..f4618154329 100644 --- a/vortex-array/src/scalar_fn/fns/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mod.rs @@ -15,6 +15,7 @@ pub mod is_null; pub mod like; pub mod list_contains; pub mod list_length; +pub mod list_sum; pub mod literal; pub mod mask; pub mod merge; diff --git a/vortex-array/src/scalar_fn/fns/not/mod.rs b/vortex-array/src/scalar_fn/fns/not/mod.rs index 156a6568df7..d0f9e400bde 100644 --- a/vortex-array/src/scalar_fn/fns/not/mod.rs +++ b/vortex-array/src/scalar_fn/fns/not/mod.rs @@ -110,8 +110,8 @@ impl ScalarFnVTable for Not { #[cfg(test)] mod tests { use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::bool::BoolArrayExt; use crate::dtype::DType; use crate::dtype::Nullability; @@ -124,10 +124,15 @@ mod tests { #[test] fn invert_booleans() { + let mut ctx = array_session().create_execution_ctx(); let not_expr = not(root()); let bools = BoolArray::from_iter([false, true, false, false, true, true]); - #[expect(deprecated)] - let result = bools.into_array().apply(¬_expr).unwrap().to_bool(); + let result = bools + .into_array() + .apply(¬_expr) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_eq!( result.to_bit_buffer().iter().collect::>(), vec![true, false, true, true, false, false] diff --git a/vortex-array/src/scalar_fn/fns/pack.rs b/vortex-array/src/scalar_fn/fns/pack.rs index 81019e22d6f..45475623d32 100644 --- a/vortex-array/src/scalar_fn/fns/pack.rs +++ b/vortex-array/src/scalar_fn/fns/pack.rs @@ -170,8 +170,6 @@ mod tests { use super::PackOptions; use crate::ArrayRef; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::PrimitiveArray; @@ -194,26 +192,33 @@ mod tests { } fn primitive_field(array: &ArrayRef, field_path: &[&str]) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); let mut field_path = field_path.iter(); let Some(field) = field_path.next() else { vortex_bail!("empty field path"); }; - #[expect(deprecated)] - let mut array = array.to_struct().unmasked_field_by_name(field)?.clone(); + let mut array = array + .clone() + .execute::(&mut ctx)? + .unmasked_field_by_name(field)? + .clone(); for field in field_path { - #[expect(deprecated)] - let next = array.to_struct().unmasked_field_by_name(field)?.clone(); + let next = array + .clone() + .execute::(&mut ctx)? + .unmasked_field_by_name(field)? + .clone(); array = next; } - #[expect(deprecated)] - let result = array.to_primitive(); + let result = array.execute::(&mut ctx)?; Ok(result) } #[test] pub fn test_empty_pack() { + let mut ctx = array_session().create_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: Default::default(), @@ -225,8 +230,11 @@ mod tests { let test_array = test_array(); let actual_array = test_array.clone().apply(&expr).unwrap(); assert_eq!(actual_array.len(), test_array.len()); - #[expect(deprecated)] - let nfields = actual_array.to_struct().struct_fields().nfields(); + let nfields = actual_array + .execute::(&mut ctx) + .unwrap() + .struct_fields() + .nfields(); assert_eq!(nfields, 0); } @@ -241,8 +249,11 @@ mod tests { [col("a"), col("b"), col("a")], ); - #[expect(deprecated)] - let actual_array = test_array().apply(&expr).unwrap().to_struct(); + let actual_array = test_array() + .apply(&expr) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_eq!(actual_array.names(), ["one", "two", "three"]); assert!(matches!(actual_array.validity(), Ok(Validity::NonNullable))); @@ -285,8 +296,11 @@ mod tests { ], ); - #[expect(deprecated)] - let actual_array = test_array().apply(&expr).unwrap().to_struct(); + let actual_array = test_array() + .apply(&expr) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_eq!(actual_array.names(), ["one", "two", "three"]); @@ -314,6 +328,7 @@ mod tests { #[test] pub fn test_pack_nullable() { + let mut ctx = array_session().create_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), @@ -322,8 +337,11 @@ mod tests { [col("a"), col("b"), col("a")], ); - #[expect(deprecated)] - let actual_array = test_array().apply(&expr).unwrap().to_struct(); + let actual_array = test_array() + .apply(&expr) + .unwrap() + .execute::(&mut ctx) + .unwrap(); assert_eq!(actual_array.names(), ["one", "two", "three"]); assert!(matches!(actual_array.validity(), Ok(Validity::AllValid))); diff --git a/vortex-array/src/scalar_fn/fns/select.rs b/vortex-array/src/scalar_fn/fns/select.rs index 54b63a00f89..2fbe4593503 100644 --- a/vortex-array/src/scalar_fn/fns/select.rs +++ b/vortex-array/src/scalar_fn/fns/select.rs @@ -309,8 +309,8 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; - #[expect(deprecated)] - use crate::ToCanonical as _; + use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; use crate::dtype::FieldName; @@ -336,20 +336,30 @@ mod tests { #[test] pub fn include_columns() { + let mut ctx = array_session().create_execution_ctx(); let st = test_array(); let select = select(vec![FieldName::from("a")], root()); - #[expect(deprecated)] - let selected = st.into_array().apply(&select).unwrap().to_struct(); + let selected = st + .into_array() + .apply(&select) + .unwrap() + .execute::(&mut ctx) + .unwrap(); let selected_names = selected.names().clone(); assert_eq!(selected_names.as_ref(), &["a"]); } #[test] pub fn exclude_columns() { + let mut ctx = array_session().create_execution_ctx(); let st = test_array(); let select = select_exclude(vec![FieldName::from("a")], root()); - #[expect(deprecated)] - let selected = st.into_array().apply(&select).unwrap().to_struct(); + let selected = st + .into_array() + .apply(&select) + .unwrap() + .execute::(&mut ctx) + .unwrap(); let selected_names = selected.names().clone(); assert_eq!(selected_names.as_ref(), &["b"]); } 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 1a4854a6a4a..2f7b563202c 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs @@ -138,7 +138,7 @@ impl ScalarFnVTable for VariantGet { builder.append_scalar(&output)?; } - return Ok(builder.finish_into_canonical().into_array()); + return Ok(builder.finish()); } // TODO(variant): replace this with a Variant builder once one exists. diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 48449bc8428..40fedbe85eb 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -267,8 +267,6 @@ pub(crate) fn zip_validity( #[cfg(test)] mod tests { - use arrow_array::cast::AsArray; - use arrow_select::zip::zip as arrow_zip; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -283,7 +281,7 @@ mod tests { use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::VarBinView; - use crate::arrow::ArrowSessionExt; + use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::BufferGrowthStrategy; @@ -489,35 +487,34 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); let zipped = mask_array - .zip(if_true.clone(), if_false.clone()) + .zip(if_true, if_false) .unwrap() .execute::(&mut ctx) .unwrap(); let zipped = zipped.as_opt::().unwrap(); assert_eq!(zipped.data_buffers().len(), 2); - let mut arrow_ctx = array_session().create_execution_ctx(); - let expected = arrow_zip( - array_session() - .arrow() - .execute_arrow(mask.into_array(), None, &mut arrow_ctx) - .unwrap() - .as_boolean(), - &array_session() - .arrow() - .execute_arrow(if_true, None, &mut arrow_ctx) - .unwrap(), - &array_session() - .arrow() - .execute_arrow(if_false, None, &mut arrow_ctx) - .unwrap(), - ) - .unwrap(); - - let actual = array_session() - .arrow() - .execute_arrow(zipped.array().clone(), None, &mut arrow_ctx) - .unwrap(); - assert_eq!(actual.as_ref(), expected.as_ref()); + let true_value = |i: usize| { + if i.is_multiple_of(2) { + "Hello" + } else { + "Hello this is a long string that won't be inlined." + } + }; + let false_value = |i: usize| { + if i.is_multiple_of(2) { + "Hello2" + } else { + "Hello2 this is a long string that won't be inlined." + } + }; + let expected = VarBinViewArray::from_iter_str((0..200).map(|i| { + if mask.value(i) { + true_value(i) + } else { + false_value(i) + } + })); + assert_arrays_eq!(zipped.array().clone(), expected, &mut ctx); } } diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index 477d5758990..e937ac8d316 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -22,6 +22,7 @@ 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::list_sum::ListSum; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::merge::Merge; use crate::scalar_fn::fns::not::Not; @@ -70,6 +71,7 @@ impl Default for ScalarFnSession { this.register(Like); this.register(ListContains); this.register(ListLength); + this.register(ListSum); this.register(Literal); this.register(Merge); this.register(Not); diff --git a/vortex-array/src/search_sorted.rs b/vortex-array/src/search_sorted/mod.rs similarity index 98% rename from vortex-array/src/search_sorted.rs rename to vortex-array/src/search_sorted/mod.rs index 3db5215e049..cda27cfc511 100644 --- a/vortex-array/src/search_sorted.rs +++ b/vortex-array/src/search_sorted/mod.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod primitive; + use std::cmp::Ordering; use std::cmp::Ordering::Equal; use std::cmp::Ordering::Greater; @@ -10,11 +12,12 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hint; +pub use primitive::*; use vortex_error::VortexResult; use crate::ArrayRef; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::legacy_session; use crate::scalar::Scalar; #[derive(Debug, Copy, Clone, Eq, PartialEq)] @@ -270,8 +273,9 @@ fn search_sorted_side_idx VortexResult>( } impl IndexOrd for ArrayRef { + #[allow(clippy::disallowed_methods)] fn index_cmp(&self, idx: usize, elem: &Scalar) -> VortexResult> { - let scalar_a = self.execute_scalar(idx, &mut LEGACY_SESSION.create_execution_ctx())?; + let scalar_a = self.execute_scalar(idx, &mut legacy_session().create_execution_ctx())?; Ok(scalar_a.partial_cmp(elem)) } diff --git a/vortex-array/src/search_sorted/primitive.rs b/vortex-array/src/search_sorted/primitive.rs new file mode 100644 index 00000000000..3da28c7be2b --- /dev/null +++ b/vortex-array/src/search_sorted/primitive.rs @@ -0,0 +1,137 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::cell::RefCell; +use std::cmp::Ordering; +use std::marker::PhantomData; + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::NativePType; +use crate::search_sorted::IndexOrd; + +/// A [`SearchSorted`](crate::search_sorted::SearchSorted) adapter over a sorted primitive-typed +/// array, comparing elements as the native type `T`. +/// +/// Values can be searched as `T`, `Option`, or `usize`. Searching as `T` or `usize` treats +/// null elements as `T::zero()`; use `Option` when the array may contain nulls, in which case +/// nulls sort before all non-null values. +pub struct SearchSortedPrimitiveArray<'a, T>( + &'a ArrayRef, + RefCell<&'a mut ExecutionCtx>, + PhantomData, +); + +impl<'a, T: NativePType> SearchSortedPrimitiveArray<'a, T> { + /// Wraps `array` for searching, panicking if the array's [`PType`](crate::dtype::PType) is + /// not `T::PTYPE`. + pub fn new(array: &'a ArrayRef, ctx: &'a mut ExecutionCtx) -> Self { + assert_eq!( + array.dtype().as_ptype(), + T::PTYPE, + "Array PType must match primitive type" + ); + Self(array, RefCell::new(ctx), PhantomData) + } + + /// Returns the value at `idx`, with nulls mapped to `T::zero()`. + fn value(&self, idx: usize) -> VortexResult { + Ok(self + .0 + .execute_scalar(idx, &mut self.1.borrow_mut())? + .as_primitive() + .typed_value::() + .unwrap_or_else(|| T::zero())) + } +} + +impl IndexOrd for SearchSortedPrimitiveArray<'_, T> { + fn index_cmp(&self, idx: usize, elem: &T) -> VortexResult> { + let value = self.value(idx)?; + Ok(Some(value.total_compare(*elem))) + } + + fn index_len(&self) -> usize { + self.0.len() + } +} + +impl IndexOrd> for SearchSortedPrimitiveArray<'_, T> { + fn index_cmp(&self, idx: usize, elem: &Option) -> VortexResult> { + // The borrow must end before `self.value` re-borrows the ctx. + let valid = self.0.is_valid(idx, &mut self.1.borrow_mut())?; + let value = valid.then(|| self.value(idx)).transpose()?; + + Ok(match (value, elem.as_ref()) { + (Some(l), Some(r)) => Some(l.total_compare(*r)), + (Some(_), None) => Some(Ordering::Greater), + (None, Some(_)) => Some(Ordering::Less), + (None, None) => Some(Ordering::Equal), + }) + } + + fn index_len(&self) -> usize { + self.0.len() + } +} + +impl IndexOrd for SearchSortedPrimitiveArray<'_, T> { + fn index_cmp(&self, idx: usize, elem: &usize) -> VortexResult> { + let value = self.value(idx)?; + + let Some(elem_t) = T::from_usize(*elem) else { + return Ok(Some(Ordering::Less)); + }; + + Ok(Some(value.total_compare(elem_t))) + } + + fn index_len(&self) -> usize { + self.0.len() + } +} + +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::array_session; + use crate::arrays::PrimitiveArray; + use crate::executor::VortexSessionExecute; + use crate::search_sorted::SearchResult; + use crate::search_sorted::SearchSorted; + use crate::search_sorted::SearchSortedPrimitiveArray; + use crate::search_sorted::SearchSortedSide; + use crate::validity::Validity; + + #[test] + fn search_sorted_optional_value() -> VortexResult<()> { + let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::AllValid).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let res = SearchSortedPrimitiveArray::::new(&array, &mut ctx) + .search_sorted(&Some(3i32), SearchSortedSide::Left)?; + assert_eq!(res, SearchResult::Found(2)); + Ok(()) + } + + #[test] + fn search_sorted_optional_value_with_nulls() -> VortexResult<()> { + let array = + PrimitiveArray::from_option_iter([None, None, Some(2i32), Some(3)]).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let searcher = SearchSortedPrimitiveArray::::new(&array, &mut ctx); + assert_eq!( + searcher.search_sorted(&Some(2i32), SearchSortedSide::Left)?, + SearchResult::Found(2) + ); + assert_eq!( + searcher.search_sorted(&None, SearchSortedSide::Right)?, + SearchResult::Found(2) + ); + Ok(()) + } +} diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index 637b57324c0..9fdf72f39fb 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -329,7 +329,7 @@ impl SerializedArray { if session.allows_unknown() { return self.decode_foreign(encoding_id, dtype, len, ctx); } - return Err(vortex_err!("Unknown encoding: {}", encoding_id)); + vortex_bail!("Unknown encoding: {}", encoding_id); }; let children = SerializedArrayChildren { diff --git a/vortex-array/src/stats/stats_set.rs b/vortex-array/src/stats/stats_set.rs index 6025afed2c6..cb4af0b2c7f 100644 --- a/vortex-array/src/stats/stats_set.rs +++ b/vortex-array/src/stats/stats_set.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::type_name; use std::fmt::Debug; use enum_iterator::all; @@ -133,12 +134,7 @@ impl StatsSet { .vortex_expect("failed to construct a scalar statistic"), ) .unwrap_or_else(|err| { - vortex_panic!( - err, - "Failed to get stat {} as {}", - stat, - std::any::type_name::() - ) + vortex_panic!(err, "Failed to get stat {} as {}", stat, type_name::()) }) }) } diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 9c80e34267e..3bb47eda0ae 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -24,7 +24,6 @@ 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::ChunkedArray; @@ -33,6 +32,7 @@ use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; +use crate::legacy_session; use crate::optimizer::ArrayOptimizer; use crate::patches::Patches; use crate::scalar::Scalar; @@ -186,15 +186,17 @@ impl Validity { /// Returns whether the `index` item is valid. #[deprecated(note = "use `execute_is_valid` with an explicit `ExecutionCtx`")] #[inline] + #[allow(clippy::disallowed_methods)] pub fn is_valid(&self, index: usize) -> VortexResult { - self.execute_is_valid(index, &mut LEGACY_SESSION.create_execution_ctx()) + self.execute_is_valid(index, &mut legacy_session().create_execution_ctx()) } /// Returns whether the `index` item is null. #[deprecated(note = "use `execute_is_null` with an explicit `ExecutionCtx`")] #[inline] + #[allow(clippy::disallowed_methods)] pub fn is_null(&self, index: usize) -> VortexResult { - self.execute_is_null(index, &mut LEGACY_SESSION.create_execution_ctx()) + self.execute_is_null(index, &mut legacy_session().create_execution_ctx()) } #[inline] diff --git a/vortex-array/src/variants.rs b/vortex-array/src/variants.rs deleted file mode 100644 index 2bc1f7fecfd..00000000000 --- a/vortex-array/src/variants.rs +++ /dev/null @@ -1,237 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! This module defines extension functionality specific to each Vortex DType. -use std::cmp::Ordering; - -use vortex_error::VortexExpect; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_panic; -use vortex_mask::Mask; - -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; -use crate::aggregate_fn::fns::sum::sum; -use crate::arrays::BoolArray; -use crate::arrays::bool::BoolArrayExt; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::FieldNames; -use crate::dtype::PType; -use crate::dtype::extension::ExtDTypeRef; -use crate::scalar::PValue; -use crate::scalar::Scalar; -use crate::search_sorted::IndexOrd; - -impl ArrayRef { - /// Downcasts the array for null-specific behavior. - pub fn as_null_typed(&self) -> NullTyped<'_> { - matches!(self.dtype(), DType::Null) - .then(|| NullTyped(self)) - .vortex_expect("Array does not have DType::Null") - } - - /// Downcasts the array for bool-specific behavior. - pub fn as_bool_typed(&self) -> BoolTyped<'_> { - matches!(self.dtype(), DType::Bool(..)) - .then(|| BoolTyped(self)) - .vortex_expect("Array does not have DType::Bool") - } - - /// Downcasts the array for primitive-specific behavior. - pub fn as_primitive_typed(&self) -> PrimitiveTyped<'_> { - matches!(self.dtype(), DType::Primitive(..)) - .then(|| PrimitiveTyped(self)) - .vortex_expect("Array does not have DType::Primitive") - } - - /// Downcasts the array for decimal-specific behavior. - pub fn as_decimal_typed(&self) -> DecimalTyped<'_> { - matches!(self.dtype(), DType::Decimal(..)) - .then(|| DecimalTyped(self)) - .vortex_expect("Array does not have DType::Decimal") - } - - /// Downcasts the array for utf8-specific behavior. - pub fn as_utf8_typed(&self) -> Utf8Typed<'_> { - matches!(self.dtype(), DType::Utf8(..)) - .then(|| Utf8Typed(self)) - .vortex_expect("Array does not have DType::Utf8") - } - - /// Downcasts the array for binary-specific behavior. - pub fn as_binary_typed(&self) -> BinaryTyped<'_> { - matches!(self.dtype(), DType::Binary(..)) - .then(|| BinaryTyped(self)) - .vortex_expect("Array does not have DType::Binary") - } - - /// Downcasts the array for struct-specific behavior. - pub fn as_struct_typed(&self) -> StructTyped<'_> { - matches!(self.dtype(), DType::Struct(..)) - .then(|| StructTyped(self)) - .vortex_expect("Array does not have DType::Struct") - } - - /// Downcasts the array for list-specific behavior. - pub fn as_list_typed(&self) -> ListTyped<'_> { - matches!(self.dtype(), DType::List(..)) - .then(|| ListTyped(self)) - .vortex_expect("Array does not have DType::List") - } - - /// Downcasts the array for extension-specific behavior. - pub fn as_extension_typed(&self) -> ExtensionTyped<'_> { - matches!(self.dtype(), DType::Extension(..)) - .then(|| ExtensionTyped(self)) - .vortex_expect("Array does not have DType::Extension") - } - - pub fn try_to_mask_fill_null_false(&self, ctx: &mut ExecutionCtx) -> VortexResult { - if !matches!(self.dtype(), DType::Bool(_)) { - vortex_bail!("mask must be bool array, has dtype {}", self.dtype()); - } - - // Convert nulls to false first in case this can be done cheaply by the encoding. - let array = self - .clone() - .fill_null(Scalar::bool(false, self.dtype().nullability()))?; - - Ok(array - .execute::(ctx)? - .to_mask_fill_null_false(ctx)) - } -} - -#[expect(dead_code)] -pub struct NullTyped<'a>(&'a ArrayRef); - -pub struct BoolTyped<'a>(&'a ArrayRef); - -impl BoolTyped<'_> { - #[deprecated( - note = "Relies on the hidden global `LEGACY_SESSION`; use `sum(array, ctx)` with an explicit `ExecutionCtx` instead" - )] - pub fn true_count(&self) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let true_count = sum(self.0, &mut ctx)?; - Ok(true_count - .as_primitive() - .as_::() - .vortex_expect("true count should never be null")) - } -} - -pub struct PrimitiveTyped<'a>(&'a ArrayRef); - -impl PrimitiveTyped<'_> { - pub fn ptype(&self) -> PType { - let DType::Primitive(ptype, _) = self.0.dtype() else { - vortex_panic!("Expected Primitive DType") - }; - *ptype - } - - /// Return the primitive value at the given index. - #[deprecated( - note = "Relies on the hidden global `LEGACY_SESSION`; use `is_valid`/`execute_scalar` with an explicit `ExecutionCtx` instead" - )] - #[allow(deprecated)] - pub fn value(&self, idx: usize) -> VortexResult> { - self.0 - .is_valid(idx, &mut LEGACY_SESSION.create_execution_ctx())? - .then(|| self.value_unchecked(idx)) - .transpose() - } - - /// Return the primitive value at the given index, ignoring nullability. - #[deprecated( - note = "Relies on the hidden global `LEGACY_SESSION`; use `execute_scalar` with an explicit `ExecutionCtx` instead" - )] - pub fn value_unchecked(&self, idx: usize) -> VortexResult { - Ok(self - .0 - .execute_scalar(idx, &mut LEGACY_SESSION.create_execution_ctx())? - .as_primitive() - .pvalue() - .unwrap_or_else(|| PValue::zero(&self.ptype()))) - } -} - -impl IndexOrd> for PrimitiveTyped<'_> { - #[allow(deprecated)] - fn index_cmp(&self, idx: usize, elem: &Option) -> VortexResult> { - let value = self.value(idx)?; - Ok(value.partial_cmp(elem)) - } - - fn index_len(&self) -> usize { - self.0.len() - } -} - -// TODO(ngates): add generics to the `value` function and implement this over T. -impl IndexOrd for PrimitiveTyped<'_> { - #[allow(deprecated)] - fn index_cmp(&self, idx: usize, elem: &PValue) -> VortexResult> { - assert!( - self.0 - .all_valid(&mut LEGACY_SESSION.create_execution_ctx())? - ); - let value = self.value_unchecked(idx)?; - Ok(value.partial_cmp(elem)) - } - - fn index_len(&self) -> usize { - self.0.len() - } -} - -#[expect(dead_code)] -pub struct Utf8Typed<'a>(&'a ArrayRef); - -#[expect(dead_code)] -pub struct BinaryTyped<'a>(&'a ArrayRef); - -#[expect(dead_code)] -pub struct DecimalTyped<'a>(&'a ArrayRef); - -pub struct StructTyped<'a>(&'a ArrayRef); - -impl StructTyped<'_> { - pub fn names(&self) -> &FieldNames { - let DType::Struct(st, _) = self.0.dtype() else { - unreachable!() - }; - st.names() - } - - pub fn dtypes(&self) -> Vec { - let DType::Struct(st, _) = self.0.dtype() else { - unreachable!() - }; - st.fields().collect() - } - - pub fn nfields(&self) -> usize { - self.names().len() - } -} - -#[expect(dead_code)] -pub struct ListTyped<'a>(&'a ArrayRef); - -pub struct ExtensionTyped<'a>(&'a ArrayRef); - -impl ExtensionTyped<'_> { - /// Returns the extension logical [`DType`]. - pub fn ext_dtype(&self) -> &ExtDTypeRef { - let DType::Extension(ext_dtype) = self.0.dtype() else { - vortex_panic!("Expected ExtDType") - }; - ext_dtype - } -} diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml new file mode 100644 index 00000000000..9d69a06fd7d --- /dev/null +++ b/vortex-arrow/Cargo.toml @@ -0,0 +1,46 @@ +[package] +name = "vortex-arrow" +authors = { workspace = true } +categories = { workspace = true } +description = "Apache Arrow interoperability for Vortex arrays" +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 + +[lints] +workspace = true + +[dependencies] +arrow-array = { workspace = true, features = ["ffi"] } +arrow-buffer = { workspace = true } +arrow-cast = { workspace = true } +arrow-data = { workspace = true } +arrow-schema = { workspace = true, features = ["canonical_extension_types"] } +arrow-select = { workspace = true } +itertools = { workspace = true } +num-traits = { workspace = true } +tracing = { workspace = true } +vortex-array = { workspace = true } +vortex-buffer = { workspace = true, features = ["arrow"] } +vortex-error = { workspace = true } +vortex-mask = { workspace = true } +vortex-runend = { workspace = true } +vortex-session = { workspace = true } + +[dev-dependencies] +divan = { workspace = true } +rstest = { workspace = true } +vortex-array = { workspace = true, features = ["_test-harness"] } + +[[bench]] +name = "to_arrow" +harness = false diff --git a/vortex-array/benches/to_arrow.rs b/vortex-arrow/benches/to_arrow.rs similarity index 87% rename from vortex-array/benches/to_arrow.rs rename to vortex-arrow/benches/to_arrow.rs index 48e5e917476..1cec274c28e 100644 --- a/vortex-array/benches/to_arrow.rs +++ b/vortex-arrow/benches/to_arrow.rs @@ -10,28 +10,32 @@ use divan::Bencher; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; -#[expect( - deprecated, - reason = "benchmark comparing deprecated method with new one" -)] -use vortex_array::arrow::ArrowArrayExecutor; -use vortex_array::arrow::ArrowSessionExt; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; +#[expect( + deprecated, + reason = "benchmark comparing deprecated method with new one" +)] +use vortex_arrow::ArrowArrayExecutor; +use vortex_arrow::ArrowSessionExt; +#[allow(deprecated)] +use vortex_arrow::dtype::ToArrowType as _; use vortex_session::VortexSession; fn main() { + LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); fn schema() -> DType { let fields = StructFields::from_iter([ @@ -89,9 +93,6 @@ fn to_arrow_dtype(bencher: Bencher) { #[allow(non_snake_case)] #[divan::bench] fn ArrowExportVTable_to_arrow_field(bencher: Bencher) { - // Warm the ArrowSession - drop(SESSION.arrow().to_arrow_field("", &schema()).unwrap()); - bencher .with_inputs(schema) .bench_values(|dtype| SESSION.arrow().to_arrow_field("", &dtype).unwrap()) @@ -110,13 +111,6 @@ fn to_arrow_array(bencher: Bencher) { #[allow(non_snake_case)] #[divan::bench] fn ArrowExportVTable_execute_arrow(bencher: Bencher) { - // Warm the ArrowSession - drop( - SESSION - .arrow() - .execute_arrow(array(), None, &mut SESSION.create_execution_ctx()), - ); - bencher .with_inputs(|| (array(), SESSION.create_execution_ctx())) .bench_values(|(array, mut ctx)| { diff --git a/vortex-array/src/arrow/convert.rs b/vortex-arrow/src/convert.rs similarity index 92% rename from vortex-array/src/arrow/convert.rs rename to vortex-arrow/src/convert.rs index e62b3a15d22..0bab0937a8a 100644 --- a/vortex-array/src/arrow/convert.rs +++ b/vortex-arrow/src/convert.rs @@ -13,7 +13,6 @@ use arrow_array::GenericByteArray; use arrow_array::GenericByteViewArray; use arrow_array::GenericListArray; use arrow_array::GenericListViewArray; -use arrow_array::MapArray as ArrowMapArray; use arrow_array::NullArray as ArrowNullArray; use arrow_array::OffsetSizeTrait; use arrow_array::PrimitiveArray as ArrowPrimitiveArray; @@ -58,6 +57,28 @@ use arrow_buffer::buffer::NullBuffer; use arrow_buffer::buffer::OffsetBuffer; use arrow_schema::DataType; use arrow_schema::TimeUnit as ArrowTimeUnit; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::NullArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::TemporalArray; +use vortex_array::arrays::VarBinArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::dtype::i256; +use vortex_array::extension::datetime::TimeUnit; +use vortex_array::validity::Validity; use vortex_buffer::Alignment; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; @@ -69,31 +90,19 @@ use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use crate::ArrayRef; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::DecimalArray; -use crate::arrays::DictArray; -use crate::arrays::FixedSizeListArray; -use crate::arrays::ListArray; -use crate::arrays::ListViewArray; -use crate::arrays::NullArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::StructArray; -use crate::arrays::TemporalArray; -use crate::arrays::VarBinArray; -use crate::arrays::VarBinViewArray; -use crate::arrow::FromArrowArray; -use crate::dtype::DType; -use crate::dtype::DecimalDType; -use crate::dtype::IntegerPType; -use crate::dtype::NativePType; -use crate::dtype::PType; -use crate::dtype::i256; -use crate::extension::datetime::TimeUnit; -use crate::validity::Validity; - -impl IntoArray for ArrowBuffer { +use crate::FromArrowArray; +use crate::dtype::FromArrowType; + +/// Zero-copy conversion of Arrow buffers into non-nullable Vortex arrays. +/// +/// This mirrors [`IntoArray`] for Arrow buffer types; a separate trait is required because both +/// [`IntoArray`] and the Arrow buffer types are foreign to this crate. +pub trait IntoVortexArray { + /// Convert this Arrow buffer into a non-nullable Vortex array without copying. + fn into_array(self) -> ArrayRef; +} + +impl IntoVortexArray for ArrowBuffer { fn into_array(self) -> ArrayRef { PrimitiveArray::from_byte_buffer( ByteBuffer::from_arrow_buffer(self, Alignment::of::()), @@ -104,13 +113,13 @@ impl IntoArray for ArrowBuffer { } } -impl IntoArray for BooleanBuffer { +impl IntoVortexArray for BooleanBuffer { fn into_array(self) -> ArrayRef { BoolArray::new(self.into(), Validity::NonNullable).into_array() } } -impl IntoArray for ScalarBuffer +impl IntoVortexArray for ScalarBuffer where T: ArrowNativeType + NativePType, { @@ -123,7 +132,7 @@ where } } -impl IntoArray for OffsetBuffer +impl IntoVortexArray for OffsetBuffer where O: IntegerPType + OffsetSizeTrait, { @@ -258,10 +267,14 @@ where Ok(match value.data_type() { DataType::Timestamp(time_unit, tz) => { - TemporalArray::new_timestamp(arr, time_unit.into(), tz.clone()).into() + TemporalArray::new_timestamp(arr, TimeUnit::from_arrow(time_unit), tz.clone()).into() + } + DataType::Time32(time_unit) => { + TemporalArray::new_time(arr, TimeUnit::from_arrow(time_unit)).into() + } + DataType::Time64(time_unit) => { + TemporalArray::new_time(arr, TimeUnit::from_arrow(time_unit)).into() } - DataType::Time32(time_unit) => TemporalArray::new_time(arr, time_unit.into()).into(), - DataType::Time64(time_unit) => TemporalArray::new_time(arr, time_unit.into()).into(), DataType::Date32 => TemporalArray::new_date(arr, TimeUnit::Days).into(), DataType::Date64 => TemporalArray::new_date(arr, TimeUnit::Milliseconds).into(), DataType::Duration(_) => unimplemented!(), @@ -464,18 +477,6 @@ impl FromArrowArray<&ArrowFixedSizeListArray> for ArrayRef { } } -impl FromArrowArray<&ArrowMapArray> for ArrayRef { - fn from_arrow(value: &ArrowMapArray, nullable: bool) -> VortexResult { - // Arrow Map is logically List> with i32 offsets. - // We convert it to a ListArray of structs. - let entries = Self::from_arrow(value.entries() as &dyn ArrowArray, false)?; - let offsets = value.offsets().clone().into_array(); - let nulls = nulls(value.nulls(), nullable)?; - - Ok(ListArray::try_new(entries, offsets, nulls)?.into_array()) - } -} - impl FromArrowArray<&ArrowNullArray> for ArrayRef { fn from_arrow(value: &ArrowNullArray, nullable: bool) -> VortexResult { vortex_ensure!( @@ -545,7 +546,6 @@ impl FromArrowArray<&dyn ArrowArray> for ArrayRef { DataType::ListView(_) => Self::from_arrow(array.as_list_view::(), nullable), DataType::LargeListView(_) => Self::from_arrow(array.as_list_view::(), nullable), DataType::FixedSizeList(..) => Self::from_arrow(array.as_fixed_size_list(), nullable), - DataType::Map(..) => Self::from_arrow(array.as_map(), nullable), DataType::Null => Self::from_arrow(as_null_array(array), nullable), DataType::Timestamp(u, _) => match u { ArrowTimeUnit::Second => { @@ -707,28 +707,26 @@ mod tests { use arrow_schema::Fields; use arrow_schema::Schema; use rstest::rstest; - - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::Decimal; - use crate::arrays::FixedSizeList; - use crate::arrays::List; - use crate::arrays::ListView; - use crate::arrays::Primitive; - use crate::arrays::Struct; - use crate::arrays::VarBinView; - use crate::arrays::fixed_size_list::FixedSizeListArrayExt; - 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::dtype::DType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::extension::datetime::TimeUnit; - use crate::extension::datetime::Timestamp; + use vortex_array::ArrayRef; + use vortex_array::arrays::Decimal; + use vortex_array::arrays::FixedSizeList; + use vortex_array::arrays::List; + use vortex_array::arrays::ListView; + use vortex_array::arrays::Primitive; + use vortex_array::arrays::Struct; + use vortex_array::arrays::VarBinView; + use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; + use vortex_array::arrays::list::ListArrayExt; + use vortex_array::arrays::listview::ListViewArrayExt; + use vortex_array::arrays::struct_::StructArrayExt; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::extension::datetime::TimeUnit; + use vortex_array::extension::datetime::Timestamp; + + use crate::FromArrowArray as _; + use crate::IntoVortexArray as _; #[rstest] #[case::i8( @@ -1571,58 +1569,4 @@ mod tests { ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).is_err() ); } - - #[test] - fn test_map_array_conversion() { - use arrow_array::MapArray; - use arrow_array::builder::MapBuilder; - use arrow_array::builder::StringBuilder; - - // Build a MapArray: map - let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); - // First map entry: {"a": 1, "b": 2} - builder.keys().append_value("a"); - builder.values().append_value(1); - builder.keys().append_value("b"); - builder.values().append_value(2); - builder.append(true).unwrap(); - - // Second map entry: null - builder.append(false).unwrap(); - - // Third map entry: {"c": 3} - builder.keys().append_value("c"); - builder.values().append_value(3); - builder.append(true).unwrap(); - - let arrow_map = builder.finish(); - assert_eq!(arrow_map.len(), 3); - - // Convert Arrow MapArray → Vortex ListArray - let vortex_array = ArrayRef::from_arrow(&arrow_map, true).unwrap(); - assert_eq!(vortex_array.len(), 3); - - // Verify it's stored as List> - let list_array = vortex_array.as_::(); - assert_eq!(list_array.elements().len(), 3); // 3 total key-value pairs - 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 via an ArrowSession. - let map_dtype = arrow_map.data_type().clone(); - 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() - .downcast_ref::() - .expect("Should be a MapArray"); - assert_eq!(map_back.len(), 3); - assert_eq!(map_back.entries().len(), 3); - assert!(map_back.is_null(1)); // Second entry was null - } } diff --git a/vortex-array/src/arrow/datum.rs b/vortex-arrow/src/datum.rs similarity index 89% rename from vortex-array/src/arrow/datum.rs rename to vortex-arrow/src/datum.rs index b2f047b5948..5019642c452 100644 --- a/vortex-array/src/arrow/datum.rs +++ b/vortex-arrow/src/datum.rs @@ -6,19 +6,19 @@ use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::Datum as ArrowDatum; use arrow_schema::DataType; use arrow_schema::Field; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::legacy_session; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; -use crate::ArrayRef; -use crate::IntoArray; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; -use crate::arrays::Constant; -use crate::arrays::ConstantArray; -use crate::arrow::ArrowSessionExt; -use crate::arrow::FromArrowArray; -use crate::executor::ExecutionCtx; +use crate::ArrowSessionExt; +use crate::FromArrowArray; /// A wrapper around a generic Arrow array that can be used as a Datum in Arrow compute. #[derive(Debug)] @@ -106,8 +106,9 @@ impl ArrowDatum for Datum { /// /// The provided array must have length #[deprecated( - note = "Relies on the hidden global `LEGACY_SESSION`; use `from_arrow_columnar` with an explicit `ExecutionCtx` instead" + note = "Relies on the hidden global `legacy_session()`; use `from_arrow_columnar` with an explicit `ExecutionCtx` instead" )] +#[allow(clippy::disallowed_methods)] pub fn from_arrow_array_with_len(array: A, len: usize, nullable: bool) -> VortexResult where ArrayRef: FromArrowArray, @@ -128,7 +129,7 @@ where Ok(ConstantArray::new( array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut legacy_session().create_execution_ctx()) .vortex_expect("array of length 1 must support execute_scalar(0)"), len, ) diff --git a/vortex-array/src/dtype/arrow.rs b/vortex-arrow/src/dtype.rs similarity index 89% rename from vortex-array/src/dtype/arrow.rs rename to vortex-arrow/src/dtype.rs index 41d8082a131..01ddc898f49 100644 --- a/vortex-array/src/dtype/arrow.rs +++ b/vortex-arrow/src/dtype.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Convert between Vortex [`crate::dtype::DType`] and Apache Arrow [`arrow_schema::DataType`]. +//! Convert between Vortex [`vortex_array::dtype::DType`] and Apache Arrow [`arrow_schema::DataType`]. //! //! Apache Arrow's type system includes physical information, which could lead to ambiguities as //! Vortex treats encodings as separate from logical types. @@ -23,26 +23,24 @@ use arrow_schema::Schema; use arrow_schema::SchemaBuilder; use arrow_schema::SchemaRef; use arrow_schema::TimeUnit as ArrowTimeUnit; -use vortex_error::VortexError; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::extension::datetime::AnyTemporal; +use vortex_array::extension::datetime::Date; +use vortex_array::extension::datetime::TemporalMetadata; +use vortex_array::extension::datetime::Time; +use vortex_array::extension::datetime::TimeUnit; +use vortex_array::extension::datetime::Timestamp; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use crate::dtype::DType; -use crate::dtype::DecimalDType; -use crate::dtype::FieldName; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::dtype::StructFields; -use crate::extension::datetime::AnyTemporal; -use crate::extension::datetime::Date; -use crate::extension::datetime::TemporalMetadata; -use crate::extension::datetime::Time; -use crate::extension::datetime::TimeUnit; -use crate::extension::datetime::Timestamp; - /// Trait for converting Arrow types to Vortex types. pub trait FromArrowType: Sized { /// Convert the Arrow type to a Vortex type. @@ -93,14 +91,14 @@ impl TryFromArrowType<&DataType> for DecimalDType { } } -impl From<&ArrowTimeUnit> for TimeUnit { - fn from(value: &ArrowTimeUnit) -> Self { - (*value).into() +impl FromArrowType<&ArrowTimeUnit> for TimeUnit { + fn from_arrow(value: &ArrowTimeUnit) -> Self { + Self::from_arrow(*value) } } -impl From for TimeUnit { - fn from(value: ArrowTimeUnit) -> Self { +impl FromArrowType for TimeUnit { + fn from_arrow(value: ArrowTimeUnit) -> Self { match value { ArrowTimeUnit::Second => Self::Seconds, ArrowTimeUnit::Millisecond => Self::Milliseconds, @@ -110,18 +108,19 @@ impl From for TimeUnit { } } -impl TryFrom for ArrowTimeUnit { - type Error = VortexError; - - fn try_from(value: TimeUnit) -> VortexResult { - Ok(match value { - TimeUnit::Seconds => Self::Second, - TimeUnit::Milliseconds => Self::Millisecond, - TimeUnit::Microseconds => Self::Microsecond, - TimeUnit::Nanoseconds => Self::Nanosecond, - _ => vortex_bail!("Cannot convert {value} to Arrow TimeUnit"), - }) - } +/// Convert a Vortex [`TimeUnit`] to an Arrow [`ArrowTimeUnit`]. +/// +/// # Errors +/// +/// Returns an error for units with no Arrow equivalent (e.g. [`TimeUnit::Days`]). +pub fn to_arrow_time_unit(value: TimeUnit) -> VortexResult { + Ok(match value { + TimeUnit::Seconds => ArrowTimeUnit::Second, + TimeUnit::Milliseconds => ArrowTimeUnit::Millisecond, + TimeUnit::Microseconds => ArrowTimeUnit::Microsecond, + TimeUnit::Nanoseconds => ArrowTimeUnit::Nanosecond, + _ => vortex_bail!("Cannot convert {value} to Arrow TimeUnit"), + }) } impl FromArrowType for DType { @@ -204,13 +203,14 @@ impl TryFromArrowType<(&DataType, Nullability)> for DType { DType::Extension(Date::new(TimeUnit::Milliseconds, nullability).erased()) } DataType::Time32(unit) => { - DType::Extension(Time::new(unit.into(), nullability).erased()) + DType::Extension(Time::new(TimeUnit::from_arrow(unit), nullability).erased()) } DataType::Time64(unit) => { - DType::Extension(Time::new(unit.into(), nullability).erased()) + DType::Extension(Time::new(TimeUnit::from_arrow(unit), nullability).erased()) } DataType::Timestamp(unit, tz) => DType::Extension( - Timestamp::new_with_tz(unit.into(), tz.clone(), nullability).erased(), + Timestamp::new_with_tz(TimeUnit::from_arrow(unit), tz.clone(), nullability) + .erased(), ), DataType::List(e) | DataType::LargeList(e) @@ -263,14 +263,32 @@ impl TryFromArrowType<&Field> for DType { } } -impl DType { +/// Extension trait converting Vortex [`DType`]s into Arrow schemas and data types. +/// +/// This mirrors inherent methods that lived on [`DType`] before Arrow interoperability moved +/// into this crate. +pub trait ToArrowType { /// Convert a Vortex [`DType`] into an Arrow [`Schema`]. /// /// **Prefer `ArrowSession::to_arrow_schema`.** This method is not plugin-aware and /// strips any `ARROW:extension:name` metadata for non-builtin extensions (only /// `arrow.parquet.variant` is special-cased here). Use the session method when you /// need round-trippable extension metadata. - pub fn to_arrow_schema(&self) -> VortexResult { + fn to_arrow_schema(&self) -> VortexResult; + + /// Returns the Arrow [`DataType`] that best corresponds to this Vortex [`DType`]. + /// + /// **Prefer `ArrowSession::to_arrow_datatype` (or `to_arrow_field`).** This method + /// has no awareness of registered Arrow extension plugins, so any [`DType::Extension`] + /// outside the builtin temporal set will fail or silently lose its + /// `ARROW:extension:name` metadata. The session methods recurse through containers + /// and dispatch plugins at every extension node. + #[deprecated(note = "Use `ArrowSession::to_arrow_datatype` instead")] + fn to_arrow_dtype(&self) -> VortexResult; +} + +impl ToArrowType for DType { + fn to_arrow_schema(&self) -> VortexResult { let DType::Struct(struct_dtype, nullable) = self else { vortex_bail!("only DType::Struct can be converted to arrow schema"); }; @@ -303,15 +321,7 @@ impl DType { Ok(builder.finish()) } - /// Returns the Arrow [`DataType`] that best corresponds to this Vortex [`DType`]. - /// - /// **Prefer `ArrowSession::to_arrow_data_type` (or `to_arrow_field`).** This method - /// has no awareness of registered Arrow extension plugins, so any [`DType::Extension`] - /// outside the builtin temporal set will fail or silently lose its - /// `ARROW:extension:name` metadata. The session methods recurse through containers - /// and dispatch plugins at every extension node. - #[deprecated(note = "Use `ArrowSession::to_arrow_datatype` instead")] - pub fn to_arrow_dtype(&self) -> VortexResult { + fn to_arrow_dtype(&self) -> VortexResult { to_data_type_naive(self) } } @@ -390,7 +400,7 @@ pub(crate) fn to_data_type_naive(dtype: &DType) -> VortexResult { if let Some(temporal) = ext_dtype.metadata_opt::() { return Ok(match temporal { TemporalMetadata::Timestamp(unit, tz) => { - DataType::Timestamp(ArrowTimeUnit::try_from(*unit)?, tz.clone()) + DataType::Timestamp(to_arrow_time_unit(*unit)?, tz.clone()) } TemporalMetadata::Date(unit) => match unit { TimeUnit::Days => DataType::Date32, @@ -433,14 +443,14 @@ mod test { use arrow_schema::Schema; use rstest::fixture; use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::FieldName; + use vortex_array::dtype::FieldNames; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::StructFields; use super::*; - use crate::dtype::DType; - use crate::dtype::FieldName; - use crate::dtype::FieldNames; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::StructFields; #[test] fn test_dtype_conversion_success() { @@ -497,7 +507,7 @@ mod test { #[case(39, DataType::Decimal256(39, 0))] #[case(76, DataType::Decimal256(76, 0))] fn test_decimal_dtype_to_arrow(#[case] precision: u8, #[case] expected: DataType) { - use crate::dtype::DecimalDType; + use vortex_array::dtype::DecimalDType; let dtype = DType::Decimal(DecimalDType::new(precision, 0), Nullability::NonNullable); assert_eq!(dtype.to_arrow_dtype().unwrap(), expected); diff --git a/vortex-array/src/arrow/executor/bool.rs b/vortex-arrow/src/executor/bool.rs similarity index 83% rename from vortex-array/src/arrow/executor/bool.rs rename to vortex-arrow/src/executor/bool.rs index 68f2451cd77..64bffa804f2 100644 --- a/vortex-array/src/arrow/executor/bool.rs +++ b/vortex-arrow/src/executor/bool.rs @@ -5,13 +5,13 @@ use std::sync::Arc; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::BooleanArray as ArrowBooleanArray; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::bool::BoolArrayExt; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::arrays::BoolArray; -use crate::arrays::bool::BoolArrayExt; -use crate::arrow::null_buffer::to_null_buffer; +use crate::null_buffer::to_null_buffer; /// Convert a canonical BoolArray directly to Arrow. pub fn canonical_bool_to_arrow( diff --git a/vortex-array/src/arrow/executor/byte.rs b/vortex-arrow/src/executor/byte.rs similarity index 81% rename from vortex-array/src/arrow/executor/byte.rs rename to vortex-arrow/src/executor/byte.rs index 2db9d3e6494..ae40ee01c26 100644 --- a/vortex-array/src/arrow/executor/byte.rs +++ b/vortex-arrow/src/executor/byte.rs @@ -8,22 +8,22 @@ use arrow_array::GenericByteArray; use arrow_array::types::BinaryViewType; use arrow_array::types::ByteArrayType; use arrow_array::types::StringViewType; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::VarBin; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; use vortex_error::VortexError; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::Canonical; -use crate::ExecutionCtx; -use crate::array::ArrayView; -use crate::arrays::VarBin; -use crate::arrays::VarBinViewArray; -use crate::arrays::varbin::VarBinArrayExt; -use crate::arrow::byte_view::execute_varbinview_to_arrow; -use crate::arrow::executor::validity::to_arrow_null_buffer; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::NativePType; -use crate::dtype::Nullability; +use crate::byte_view::execute_varbinview_to_arrow; +use crate::executor::validity::to_arrow_null_buffer; /// Convert a Vortex array into an Arrow GenericBinaryArray. pub(super) fn to_arrow_byte_array( @@ -79,18 +79,18 @@ mod tests { use arrow_array::Array; use arrow_array::cast::AsArray; use arrow_schema::DataType; + use arrow_schema::Field; use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; 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; - use crate::dtype::Nullability; + use crate::ArrowSessionExt; + use crate::executor::byte::VarBinViewArray; fn make_utf8_array() -> VarBinViewArray { VarBinViewArray::from_iter_str(["hello", "world", "this is a longer string for testing"]) @@ -122,10 +122,12 @@ mod tests { #[case] vortex_array: VarBinViewArray, #[case] target_dtype: DataType, ) { - let mut ctx = array_session().create_execution_ctx(); - let arrow = vortex_array - .into_array() - .execute_arrow(Some(&target_dtype), &mut ctx) + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("test_field", target_dtype.clone(), true); + let arrow = session + .arrow() + .execute_arrow(vortex_array.into_array(), Some(&field), &mut ctx) .unwrap(); assert_eq!(arrow.data_type(), &target_dtype); @@ -169,10 +171,12 @@ mod tests { vortex_dtype, ); - let mut ctx = array_session().create_execution_ctx(); - let arrow = vortex_array - .into_array() - .execute_arrow(Some(&target_dtype), &mut ctx) + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("test_field", target_dtype.clone(), true); + let arrow = session + .arrow() + .execute_arrow(vortex_array.into_array(), Some(&field), &mut ctx) .unwrap(); assert_eq!(arrow.data_type(), &target_dtype); @@ -192,10 +196,11 @@ mod tests { .into_array() .filter(Mask::from_iter([true, false, false]))?; - let arrow = filtered.execute_arrow( - Some(&DataType::Utf8View), - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let arrow = session + .arrow() + .execute_arrow(filtered.into_array(), None, &mut ctx)?; assert_eq!(arrow.as_string_view().value(0), "selected"); assert!( diff --git a/vortex-array/src/arrow/executor/byte_view.rs b/vortex-arrow/src/executor/byte_view.rs similarity index 85% rename from vortex-array/src/arrow/executor/byte_view.rs rename to vortex-arrow/src/executor/byte_view.rs index 89f37681220..ffddc49a4dd 100644 --- a/vortex-array/src/arrow/executor/byte_view.rs +++ b/vortex-arrow/src/executor/byte_view.rs @@ -7,16 +7,16 @@ use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::GenericByteViewArray; use arrow_array::types::ByteViewType; use arrow_buffer::ScalarBuffer; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::arrays::VarBinViewArray; -use crate::arrow::null_buffer::to_null_buffer; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::arrow::FromArrowType; +use crate::dtype::FromArrowType; +use crate::null_buffer::to_null_buffer; /// Convert a canonical VarBinViewArray directly to Arrow. pub fn canonical_varbinview_to_arrow( @@ -47,7 +47,7 @@ pub fn execute_varbinview_to_arrow( array: &VarBinViewArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - let compacted = array.compact_buffers()?; + let compacted = array.compact_buffers(ctx)?; canonical_varbinview_to_arrow::(&compacted, ctx) } diff --git a/vortex-array/src/arrow/executor/decimal.rs b/vortex-arrow/src/executor/decimal.rs similarity index 92% rename from vortex-array/src/arrow/executor/decimal.rs rename to vortex-arrow/src/executor/decimal.rs index a9aa9fdd2f5..d5eabc77145 100644 --- a/vortex-array/src/arrow/executor/decimal.rs +++ b/vortex-arrow/src/executor/decimal.rs @@ -13,15 +13,15 @@ use arrow_schema::DataType; use itertools::Itertools; use num_traits::AsPrimitive; use num_traits::ToPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::DecimalArray; +use vortex_array::dtype::DecimalType; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::arrays::DecimalArray; -use crate::arrow::null_buffer::to_null_buffer; -use crate::dtype::DecimalType; +use crate::null_buffer::to_null_buffer; pub(super) fn to_arrow_decimal( array: ArrayRef, @@ -72,7 +72,7 @@ fn to_arrow_decimal32(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResu }) .process_results(|iter| Buffer::from_trusted_len_iter(iter))?, DecimalType::I256 => array - .buffer::() + .buffer::() .into_iter() .map(|x| { x.to_i32() @@ -116,7 +116,7 @@ fn to_arrow_decimal64(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResu }) .process_results(|iter| Buffer::from_trusted_len_iter(iter))?, DecimalType::I256 => array - .buffer::() + .buffer::() .into_iter() .map(|x| { x.to_i64() @@ -155,7 +155,7 @@ fn to_arrow_decimal128(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexRes } DecimalType::I128 => array.buffer::(), DecimalType::I256 => array - .buffer::() + .buffer::() .into_iter() .map(|x| { x.to_i128() @@ -196,7 +196,7 @@ fn to_arrow_decimal256(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexRes array .buffer::() .into_iter() - .map(|x| crate::dtype::i256::from_i128(x).into()), + .map(|x| vortex_array::dtype::i256::from_i128(x).into()), ), DecimalType::I256 => { Buffer::::from_byte_buffer(array.buffer_handle().clone().into_host_sync()) @@ -219,19 +219,19 @@ mod tests { use arrow_buffer::i256; use arrow_schema::DataType; use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::builders::ArrayBuilder; + use vortex_array::builders::DecimalBuilder; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::NativeDecimalType; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; - 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; - use crate::builders::DecimalBuilder; - use crate::dtype::DecimalDType; - use crate::dtype::NativeDecimalType; - use crate::validity::Validity; + use crate::ArrowArrayExecutor; + use crate::executor::decimal::DecimalArray; #[test] fn decimal_to_arrow() -> VortexResult<()> { @@ -260,7 +260,7 @@ mod tests { #[case(0i32)] #[case(0i64)] #[case(0i128)] - #[case(crate::dtype::i256::ZERO)] + #[case(vortex_array::dtype::i256::ZERO)] fn test_to_arrow_decimal128( #[case] _decimal_type: T, ) -> VortexResult<()> { @@ -288,7 +288,7 @@ mod tests { #[case(0i32)] #[case(0i64)] #[case(0i128)] - #[case(crate::dtype::i256::ZERO)] + #[case(vortex_array::dtype::i256::ZERO)] fn test_to_arrow_decimal32(#[case] _decimal_type: T) -> VortexResult<()> { use arrow_array::Decimal32Array; @@ -316,7 +316,7 @@ mod tests { #[case(0i32)] #[case(0i64)] #[case(0i128)] - #[case(crate::dtype::i256::ZERO)] + #[case(vortex_array::dtype::i256::ZERO)] fn test_to_arrow_decimal64(#[case] _decimal_type: T) -> VortexResult<()> { use arrow_array::Decimal64Array; @@ -344,7 +344,7 @@ mod tests { #[case(0i32)] #[case(0i64)] #[case(0i128)] - #[case(crate::dtype::i256::ZERO)] + #[case(vortex_array::dtype::i256::ZERO)] fn test_to_arrow_decimal256( #[case] _decimal_type: T, ) -> VortexResult<()> { diff --git a/vortex-array/src/arrow/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs similarity index 87% rename from vortex-array/src/arrow/executor/dictionary.rs rename to vortex-arrow/src/executor/dictionary.rs index 97d25c3f7cc..f5be1076a4b 100644 --- a/vortex-array/src/arrow/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -10,19 +10,19 @@ use arrow_array::cast::AsArray; use arrow_array::new_null_array; use arrow_array::types::*; use arrow_schema::DataType; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Dict; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::Constant; -use crate::arrays::ConstantArray; -use crate::arrays::Dict; -use crate::arrays::DictArray; -use crate::arrays::dict::DictArraySlotsExt; -use crate::arrow::ArrowArrayExecutor; +use crate::ArrowArrayExecutor; pub(super) fn to_arrow_dictionary( array: ArrayRef, @@ -145,30 +145,33 @@ mod tests { use arrow_array::types::UInt32Type; use arrow_schema::DataType; use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexResult; - use crate::IntoArray; - use crate::array_session; - use crate::arrays::PrimitiveArray; - use crate::arrays::VarBinViewArray; - use crate::arrow::ArrowArrayExecutor; - use crate::arrow::executor::dictionary::ConstantArray; - use crate::arrow::executor::dictionary::DictArray; - use crate::dtype::DType; - use crate::dtype::Nullability::Nullable; - use crate::executor::VortexSessionExecute; - use crate::scalar::Scalar; + use crate::ArrowArrayExecutor; + use crate::executor::dictionary::ConstantArray; + use crate::executor::dictionary::DictArray; fn dict_type(codes: DataType, values: DataType) -> DataType { DataType::Dictionary(Box::new(codes), Box::new(values)) } - fn execute(array: crate::ArrayRef, dt: &DataType) -> VortexResult { + fn execute( + array: vortex_array::ArrayRef, + dt: &DataType, + ) -> VortexResult { array.execute_arrow(Some(dt), &mut array_session().create_execution_ctx()) } - fn dict_basic_input() -> crate::ArrayRef { + fn dict_basic_input() -> vortex_array::ArrayRef { DictArray::try_new( buffer![0u8, 1, 0].into_array(), VarBinViewArray::from_iter_str(["a", "b"]).into_array(), @@ -177,7 +180,7 @@ mod tests { .into_array() } - fn dict_with_null_codes_input() -> crate::ArrayRef { + fn dict_with_null_codes_input() -> vortex_array::ArrayRef { DictArray::try_new( PrimitiveArray::from_option_iter(vec![Some(0u8), None, Some(1)]).into_array(), VarBinViewArray::from_iter_str(["a", "b"]).into_array(), @@ -213,7 +216,7 @@ mod tests { Arc::new(vec![Some("a"), None, Some("a"), Some("b"), Some("a")].into_iter().collect::>()) as arrow_array::ArrayRef, )] fn to_arrow_dictionary( - #[case] input: crate::ArrayRef, + #[case] input: vortex_array::ArrayRef, #[case] target_type: DataType, #[case] expected: arrow_array::ArrayRef, ) -> VortexResult<()> { diff --git a/vortex-array/src/arrow/executor/fixed_size_list.rs b/vortex-arrow/src/executor/fixed_size_list.rs similarity index 86% rename from vortex-array/src/arrow/executor/fixed_size_list.rs rename to vortex-arrow/src/executor/fixed_size_list.rs index caeb5804da2..72ce7c57ff0 100644 --- a/vortex-array/src/arrow/executor/fixed_size_list.rs +++ b/vortex-arrow/src/executor/fixed_size_list.rs @@ -4,16 +4,16 @@ use std::sync::Arc; use arrow_schema::FieldRef; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::FixedSizeList; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::arrays::FixedSizeList; -use crate::arrays::FixedSizeListArray; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; -use crate::arrow::executor::validity::to_arrow_null_buffer; -use crate::arrow::session::ArrowSessionExt; +use crate::executor::validity::to_arrow_null_buffer; +use crate::session::ArrowSessionExt; pub(super) fn to_arrow_fixed_list( array: ArrayRef, diff --git a/vortex-array/src/arrow/executor/list.rs b/vortex-arrow/src/executor/list.rs similarity index 90% rename from vortex-array/src/arrow/executor/list.rs rename to vortex-arrow/src/executor/list.rs index 7bf7a16f496..0b0d7f625b5 100644 --- a/vortex-array/src/arrow/executor/list.rs +++ b/vortex-arrow/src/executor/list.rs @@ -8,30 +8,30 @@ use arrow_array::GenericListArray; use arrow_array::OffsetSizeTrait; use arrow_buffer::OffsetBuffer; use arrow_schema::FieldRef; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::List; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListView; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::chunked::ChunkedArrayExt; +use vortex_array::arrays::list::ListArrayExt; +use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::listview::ListViewDataParts; +use vortex_array::arrays::listview::ListViewRebuildMode; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use crate::ArrayRef; -use crate::Canonical; -use crate::ExecutionCtx; -use crate::arrays::Chunked; -use crate::arrays::List; -use crate::arrays::ListArray; -use crate::arrays::ListView; -use crate::arrays::ListViewArray; -use crate::arrays::chunked::ChunkedArrayExt; -use crate::arrays::list::ListArrayExt; -use crate::arrays::listview::ListViewArrayExt; -use crate::arrays::listview::ListViewDataParts; -use crate::arrays::listview::ListViewRebuildMode; -use crate::arrow::executor::validity::to_arrow_null_buffer; -use crate::arrow::session::ArrowSessionExt; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::NativePType; -use crate::dtype::Nullability; +use crate::executor::validity::to_arrow_null_buffer; +use crate::session::ArrowSessionExt; #[allow(rustdoc::broken_intra_doc_links)] /// Convert a Vortex VarBinArray into an Arrow [`GenericListArray`](arrow_array:array::GenericListArray). @@ -210,22 +210,22 @@ mod tests { use arrow_array::Int32Array; use arrow_schema::DataType; use arrow_schema::Field; + use vortex_array::Canonical; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::Canonical; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrays::PrimitiveArray; - use crate::arrow::ArrowArrayExecutor; - use crate::arrow::executor::list::ListViewArray; - use crate::dtype::DType; - use crate::dtype::Nullability::NonNullable; - use crate::validity::Validity; + use crate::ArrowArrayExecutor; + use crate::executor::list::ListViewArray; /// A shared session for these list-executor tests, used to create execution contexts. - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_to_arrow_list_i32() -> VortexResult<()> { @@ -368,7 +368,10 @@ mod tests { fn test_to_arrow_list_empty_zctl() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let dtype = DType::List( - Arc::new(DType::Primitive(crate::dtype::PType::I32, NonNullable)), + Arc::new(DType::Primitive( + vortex_array::dtype::PType::I32, + NonNullable, + )), NonNullable, ); let list_array = unsafe { diff --git a/vortex-array/src/arrow/executor/list_view.rs b/vortex-arrow/src/executor/list_view.rs similarity index 89% rename from vortex-array/src/arrow/executor/list_view.rs rename to vortex-arrow/src/executor/list_view.rs index 7d5666b3762..da242eb9719 100644 --- a/vortex-array/src/arrow/executor/list_view.rs +++ b/vortex-arrow/src/executor/list_view.rs @@ -6,24 +6,24 @@ use std::sync::Arc; use arrow_array::GenericListViewArray; use arrow_array::OffsetSizeTrait; use arrow_schema::FieldRef; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::listview::DEFAULT_REBUILD_DENSITY_THRESHOLD; +use vortex_array::arrays::listview::DEFAULT_TRIM_ELEMENTS_THRESHOLD; +use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::listview::ListViewDataParts; +use vortex_array::arrays::listview::ListViewRebuildMode; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::Nullability::NonNullable; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::arrays::ListViewArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::listview::DEFAULT_REBUILD_DENSITY_THRESHOLD; -use crate::arrays::listview::DEFAULT_TRIM_ELEMENTS_THRESHOLD; -use crate::arrays::listview::ListViewArrayExt; -use crate::arrays::listview::ListViewDataParts; -use crate::arrays::listview::ListViewRebuildMode; -use crate::arrow::executor::validity::to_arrow_null_buffer; -use crate::arrow::session::ArrowSessionExt; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::IntegerPType; -use crate::dtype::Nullability::NonNullable; +use crate::executor::validity::to_arrow_null_buffer; +use crate::session::ArrowSessionExt; pub(super) fn to_arrow_list_view( array: ArrayRef, @@ -115,19 +115,19 @@ mod tests { use arrow_array::GenericListViewArray; use arrow_schema::DataType; use arrow_schema::Field; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::arrow::ArrowArrayExecutor; - use crate::arrow::executor::list_view::ListViewArray; - use crate::arrow::executor::list_view::PrimitiveArray; - use crate::validity::Validity; + use crate::ArrowArrayExecutor; + use crate::executor::list_view::ListViewArray; + use crate::executor::list_view::PrimitiveArray; /// A shared session for these list-view-executor tests, used to create execution contexts. - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn trims_zero_copy_with_significant_trailing_waste() -> VortexResult<()> { diff --git a/vortex-array/src/arrow/executor/mod.rs b/vortex-arrow/src/executor/mod.rs similarity index 87% rename from vortex-array/src/arrow/executor/mod.rs rename to vortex-arrow/src/executor/mod.rs index 504511b3ae6..875577788b4 100644 --- a/vortex-array/src/arrow/executor/mod.rs +++ b/vortex-arrow/src/executor/mod.rs @@ -14,7 +14,6 @@ mod dictionary; mod fixed_size_list; mod list; mod list_view; -mod map; pub mod null; pub mod primitive; mod run_end; @@ -31,36 +30,35 @@ use arrow_schema::Field; use arrow_schema::FieldRef; use arrow_schema::Schema; use itertools::Itertools; -use vortex_array::dtype::arrow::to_data_type_naive; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::List; +use vortex_array::arrays::VarBin; +use vortex_array::arrays::list::ListArrayExt; +use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use crate::ArrayRef; -use crate::arrays::List; -use crate::arrays::VarBin; -use crate::arrays::list::ListArrayExt; -use crate::arrays::varbin::VarBinArrayExt; -use crate::arrow::executor::bool::to_arrow_bool; -use crate::arrow::executor::byte::to_arrow_byte_array; -use crate::arrow::executor::byte_view::to_arrow_byte_view; -use crate::arrow::executor::decimal::to_arrow_decimal; -use crate::arrow::executor::dictionary::to_arrow_dictionary; -use crate::arrow::executor::fixed_size_list::to_arrow_fixed_list; -use crate::arrow::executor::list::to_arrow_list; -use crate::arrow::executor::list_view::to_arrow_list_view; -use crate::arrow::executor::map::to_arrow_map; -use crate::arrow::executor::null::to_arrow_null; -use crate::arrow::executor::primitive::to_arrow_primitive; -use crate::arrow::executor::run_end::to_arrow_run_end; -use crate::arrow::executor::struct_::to_arrow_struct; -use crate::arrow::executor::temporal::to_arrow_date; -use crate::arrow::executor::temporal::to_arrow_time; -use crate::arrow::executor::temporal::to_arrow_timestamp; -use crate::arrow::session::ArrowSessionExt; -use crate::dtype::DType; -use crate::dtype::PType; -use crate::executor::ExecutionCtx; +use crate::dtype::to_data_type_naive; +use crate::executor::bool::to_arrow_bool; +use crate::executor::byte::to_arrow_byte_array; +use crate::executor::byte_view::to_arrow_byte_view; +use crate::executor::decimal::to_arrow_decimal; +use crate::executor::dictionary::to_arrow_dictionary; +use crate::executor::fixed_size_list::to_arrow_fixed_list; +use crate::executor::list::to_arrow_list; +use crate::executor::list_view::to_arrow_list_view; +use crate::executor::null::to_arrow_null; +use crate::executor::primitive::to_arrow_primitive; +use crate::executor::run_end::to_arrow_run_end; +use crate::executor::struct_::to_arrow_struct; +use crate::executor::temporal::to_arrow_date; +use crate::executor::temporal::to_arrow_time; +use crate::executor::temporal::to_arrow_timestamp; +use crate::session::ArrowSessionExt; /// Trait for executing a Vortex array to produce an Arrow array. #[deprecated(note = "Use an `ArrowSession` to perform conversions to/from Arrow arrays")] @@ -190,11 +188,11 @@ pub(crate) fn execute_arrow_naive( DataType::RunEndEncoded(ends_type, values_type) => { to_arrow_run_end(array, ends_type.data_type(), values_type, ctx) } - DataType::Map(entries_field, ordered) => to_arrow_map(array, entries_field, *ordered, ctx), dt @ (DataType::Date32 | DataType::Date64) => to_arrow_date(array, dt, ctx), dt @ (DataType::Time32(_) | DataType::Time64(_)) => to_arrow_time(array, dt, ctx), dt @ DataType::Timestamp(..) => to_arrow_timestamp(array, dt, ctx), DataType::FixedSizeBinary(_) + | DataType::Map(..) | DataType::Duration(_) | DataType::Interval(_) | DataType::Union(..) => { diff --git a/vortex-array/src/arrow/executor/null.rs b/vortex-arrow/src/executor/null.rs similarity index 86% rename from vortex-array/src/arrow/executor/null.rs rename to vortex-arrow/src/executor/null.rs index 9d2b5ef6fc4..c5ffc991006 100644 --- a/vortex-array/src/arrow/executor/null.rs +++ b/vortex-arrow/src/executor/null.rs @@ -5,12 +5,11 @@ use std::sync::Arc; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::NullArray as ArrowNullArray; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::NullArray; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::arrays::NullArray; - /// Convert a canonical NullArray directly to Arrow. pub fn canonical_null_to_arrow(array: &NullArray) -> ArrowArrayRef { Arc::new(ArrowNullArray::new(array.len())) diff --git a/vortex-array/src/arrow/executor/primitive.rs b/vortex-arrow/src/executor/primitive.rs similarity index 81% rename from vortex-array/src/arrow/executor/primitive.rs rename to vortex-arrow/src/executor/primitive.rs index 0d41176972b..6cea08c0a72 100644 --- a/vortex-array/src/arrow/executor/primitive.rs +++ b/vortex-arrow/src/executor/primitive.rs @@ -6,16 +6,16 @@ use std::sync::Arc; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::ArrowPrimitiveType; use arrow_array::PrimitiveArray as ArrowPrimitiveArray; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::arrays::PrimitiveArray; -use crate::arrow::null_buffer::to_null_buffer; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::NativePType; -use crate::dtype::Nullability; +use crate::null_buffer::to_null_buffer; /// Convert a canonical PrimitiveArray directly to Arrow. pub fn canonical_primitive_to_arrow( diff --git a/vortex-array/src/arrow/executor/run_end.rs b/vortex-arrow/src/executor/run_end.rs similarity index 76% rename from vortex-array/src/arrow/executor/run_end.rs rename to vortex-arrow/src/executor/run_end.rs index 0632f871f29..b6913c1ea1c 100644 --- a/vortex-array/src/arrow/executor/run_end.rs +++ b/vortex-arrow/src/executor/run_end.rs @@ -11,35 +11,20 @@ use arrow_array::types::*; use arrow_buffer::ArrowNativeType; use arrow_schema::DataType; use arrow_schema::Field; -use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_runend::RunEnd; +use vortex_runend::RunEndArray; +use vortex_runend::RunEndArrayExt; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::Constant; -use crate::arrays::ConstantArray; -use crate::arrow::ArrowArrayExecutor; -use crate::session::ArraySessionExt; - -/// The encoding ID used by `vortex-runend`. We match on this string to avoid a crate dependency. -const VORTEX_RUNEND_ID: &str = "vortex.runend"; - -/// Mirror of `RunEndMetadata` from the `vortex-runend` crate. Same prost field tags so we can -/// decode the serialized metadata without depending on that crate. -#[derive(Clone, prost::Message)] -struct RunEndMetadata { - #[prost(int32, tag = "1")] - pub ends_ptype: i32, - #[prost(uint64, tag = "2")] - pub num_runs: u64, - #[prost(uint64, tag = "3")] - pub offset: u64, -} +use crate::ArrowArrayExecutor; pub(super) fn to_arrow_run_end( array: ArrayRef, @@ -57,47 +42,32 @@ pub(super) fn to_arrow_run_end( // Execute to unwrap any wrapper VTables (Slice, Filter, etc.) which may // reveal a RunEndArray. let array = array.execute::(ctx)?; - if array.encoding_id().as_ref() == VORTEX_RUNEND_ID { - // NOTE(ngates): while this module still lives in vortex-array, we cannot depend on the - // vortex-runend crate. Therefore, we match on the encoding ID string and extract the children - // and metadata directly. - return run_end_to_arrow(array, ends_type, values_type, ctx); + match array.try_downcast::() { + Ok(run_end) => run_end_to_arrow(run_end, ends_type, values_type, ctx), + Err(array) => { + // Fallback: canonicalize to flat Arrow, then cast to REE. + let flat = array.execute_arrow(Some(values_type.data_type()), ctx)?; + let ree_type = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", ends_type.clone(), false)), + Arc::new(values_type.clone()), + ); + arrow_cast::cast(&flat, &ree_type).map_err(VortexError::from) + } } - - // Fallback: canonicalize to flat Arrow, then cast to REE. - let flat = array.execute_arrow(Some(values_type.data_type()), ctx)?; - let ree_type = DataType::RunEndEncoded( - Arc::new(Field::new("run_ends", ends_type.clone(), false)), - Arc::new(values_type.clone()), - ); - arrow_cast::cast(&flat, &ree_type).map_err(VortexError::from) } fn run_end_to_arrow( - array: ArrayRef, + array: RunEndArray, ends_type: &DataType, values_type: &Field, ctx: &mut ExecutionCtx, ) -> VortexResult { let length = array.len(); - let metadata_bytes = ctx - .session() - .array_serialize(&array)? - .ok_or_else(|| vortex_err!("RunEndArray missing metadata"))?; - let metadata = RunEndMetadata::decode(&*metadata_bytes) - .map_err(|e| vortex_err!("Failed to decode RunEndMetadata: {e}"))?; - let offset = usize::try_from(metadata.offset) - .map_err(|_| vortex_err!("RunEndArray offset {} overflows usize", metadata.offset))?; - - let children = array.children(); - vortex_ensure!( - children.len() == 2, - "Expected RunEndArray to have 2 children, got {}", - children.len() - ); - - let arrow_ends = children[0].clone().execute_arrow(Some(ends_type), ctx)?; - let arrow_values = children[1] + let offset = array.offset(); + + let arrow_ends = array.ends().clone().execute_arrow(Some(ends_type), ctx)?; + let arrow_values = array + .values() .clone() .execute_arrow(Some(values_type.data_type()), ctx)?; @@ -204,21 +174,21 @@ mod tests { use arrow_schema::DataType; use arrow_schema::Field; use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::VortexSession; - use crate::IntoArray; - use crate::arrays::PrimitiveArray; - use crate::arrow::ArrowArrayExecutor; - use crate::arrow::executor::run_end::ConstantArray; - use crate::dtype::DType; - use crate::dtype::Nullability::Nullable; - use crate::dtype::PType; - use crate::executor::VortexSessionExecute; - use crate::scalar::Scalar; + use crate::ArrowArrayExecutor; + use crate::executor::run_end::ConstantArray; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn ree_type(ends: DataType, values_dtype: DataType) -> DataType { DataType::RunEndEncoded( @@ -227,7 +197,10 @@ mod tests { ) } - fn execute(array: crate::ArrayRef, dt: &DataType) -> VortexResult { + fn execute( + array: vortex_array::ArrayRef, + dt: &DataType, + ) -> VortexResult { array.execute_arrow(Some(dt), &mut SESSION.create_execution_ctx()) } @@ -279,7 +252,7 @@ mod tests { ), )] fn constant_to_ree( - #[case] input: crate::ArrayRef, + #[case] input: vortex_array::ArrayRef, #[case] target_type: DataType, #[case] expected: arrow_array::ArrayRef, ) -> VortexResult<()> { diff --git a/vortex-array/src/arrow/executor/struct_.rs b/vortex-arrow/src/executor/struct_.rs similarity index 89% rename from vortex-array/src/arrow/executor/struct_.rs rename to vortex-arrow/src/executor/struct_.rs index d3b7bcda2dc..e2ce2700d4a 100644 --- a/vortex-array/src/arrow/executor/struct_.rs +++ b/vortex-arrow/src/executor/struct_.rs @@ -9,27 +9,27 @@ use arrow_buffer::NullBuffer; use arrow_schema::Field; use arrow_schema::Fields; use itertools::Itertools; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::ScalarFn; +use vortex_array::arrays::Struct; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; +use vortex_array::arrays::struct_::StructDataParts; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::StructFields; +use vortex_array::scalar_fn::fns::pack::Pack; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::Chunked; -use crate::arrays::ScalarFn; -use crate::arrays::Struct; -use crate::arrays::StructArray; -use crate::arrays::scalar_fn::ScalarFnArrayExt; -use crate::arrays::struct_::StructDataParts; -use crate::arrow::ArrowArrayExecutor; -use crate::arrow::executor::validity::to_arrow_null_buffer; -use crate::arrow::session::ArrowSessionExt; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::FieldNames; -use crate::dtype::StructFields; -use crate::dtype::arrow::FromArrowType; -use crate::scalar_fn::fns::pack::Pack; +use crate::ArrowArrayExecutor; +use crate::dtype::FromArrowType; +use crate::executor::validity::to_arrow_null_buffer; +use crate::session::ArrowSessionExt; pub(super) fn to_arrow_struct( array: ArrayRef, @@ -91,7 +91,7 @@ pub(super) fn to_arrow_struct( // We apply a cast to ensure we push down casting where possible into the struct fields. array.cast(DType::Struct( vx_fields, - crate::dtype::Nullability::Nullable, + vortex_array::dtype::Nullability::Nullable, ))? } else { array @@ -204,20 +204,21 @@ mod tests { use arrow_buffer::NullBuffer; use arrow_schema::DataType; use arrow_schema::Field; + use vortex_array as array; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::dtype::FieldNames; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array; - use crate::array_session; - use crate::arrays; - use crate::arrays::PrimitiveArray; - use crate::arrays::StructArray; - use crate::arrow::ArrowArrayExecutor; - use crate::arrow::FromArrowArray; - use crate::dtype::FieldNames; - use crate::validity::Validity; + use crate::ArrowArrayExecutor; + use crate::FromArrowArray; + use crate::dtype::to_data_type_naive; #[test] fn struct_nullable_non_null_to_arrow() -> VortexResult<()> { @@ -328,7 +329,7 @@ mod tests { ), ])?); - let arrow_dtype = array.dtype().to_arrow_dtype()?; + let arrow_dtype = to_data_type_naive(array.dtype())?; assert_eq!( &array .into_array() diff --git a/vortex-array/src/arrow/executor/temporal.rs b/vortex-arrow/src/executor/temporal.rs similarity index 92% rename from vortex-array/src/arrow/executor/temporal.rs rename to vortex-arrow/src/executor/temporal.rs index bed70f8baec..8af9aa0b2c4 100644 --- a/vortex-array/src/arrow/executor/temporal.rs +++ b/vortex-arrow/src/executor/temporal.rs @@ -26,20 +26,20 @@ use arrow_array::types::TimestampNanosecondType; use arrow_array::types::TimestampSecondType; use arrow_schema::DataType; use arrow_schema::TimeUnit as ArrowTimeUnit; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::extension::datetime::AnyTemporal; +use vortex_array::extension::datetime::TemporalMetadata; +use vortex_array::extension::datetime::TimeUnit; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::arrays::PrimitiveArray; -use crate::arrow::null_buffer::to_null_buffer; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::NativePType; -use crate::dtype::Nullability; -use crate::extension::datetime::AnyTemporal; -use crate::extension::datetime::TemporalMetadata; -use crate::extension::datetime::TimeUnit; +use crate::null_buffer::to_null_buffer; pub(super) fn to_arrow_date( array: ArrayRef, @@ -168,7 +168,7 @@ fn validate_temporal_extension(array: &ArrayRef, target: &DataType) -> VortexRes Ok(()) } (TemporalMetadata::Timestamp(unit, src_tz), DataType::Timestamp(arrow_unit, tgt_tz)) => { - let src_arrow_unit = ArrowTimeUnit::try_from(*unit)?; + let src_arrow_unit = crate::dtype::to_arrow_time_unit(*unit)?; if src_arrow_unit != *arrow_unit { vortex_bail!( "Cannot convert Timestamp({unit}) to Arrow Timestamp({arrow_unit:?}): unit mismatch" diff --git a/vortex-array/src/arrays/bool/vtable/canonical.rs b/vortex-arrow/src/executor/validity.rs similarity index 63% rename from vortex-array/src/arrays/bool/vtable/canonical.rs rename to vortex-arrow/src/executor/validity.rs index 0d735177e5d..ab8a88b5ee3 100644 --- a/vortex-array/src/arrays/bool/vtable/canonical.rs +++ b/vortex-arrow/src/executor/validity.rs @@ -1,2 +1,4 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors + +pub(super) use crate::null_buffer::to_arrow_null_buffer; diff --git a/vortex-array/src/arrow/iter.rs b/vortex-arrow/src/iter.rs similarity index 86% rename from vortex-array/src/arrow/iter.rs rename to vortex-arrow/src/iter.rs index f05bb0a4e7b..0d7ebe5c82a 100644 --- a/vortex-array/src/arrow/iter.rs +++ b/vortex-arrow/src/iter.rs @@ -2,14 +2,14 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use arrow_array::ffi_stream; +use vortex_array::ArrayRef; +use vortex_array::dtype::DType; +use vortex_array::iter::ArrayIterator; use vortex_error::VortexError; use vortex_error::VortexResult; -use crate::ArrayRef; -use crate::arrow::FromArrowArray; -use crate::dtype::DType; -use crate::dtype::arrow::FromArrowType; -use crate::iter::ArrayIterator; +use crate::FromArrowArray; +use crate::dtype::FromArrowType; /// An adapter for converting an `ArrowArrayStreamReader` into a Vortex `ArrayStream`. pub struct ArrowArrayStreamAdapter { diff --git a/vortex-array/src/arrow/mod.rs b/vortex-arrow/src/lib.rs similarity index 57% rename from vortex-array/src/arrow/mod.rs rename to vortex-arrow/src/lib.rs index 5259005dffc..73fa4c724ce 100644 --- a/vortex-array/src/arrow/mod.rs +++ b/vortex-arrow/src/lib.rs @@ -1,38 +1,62 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Utilities to work with `Arrow` data and types. +//! Apache Arrow interoperability for Vortex. +//! +//! This crate owns every conversion between Vortex and Arrow: importing Arrow arrays, record +//! batches, schemas, and fields into Vortex ([`FromArrowArray`], [`FromArrowType`], +//! [`TryFromArrowType`]), and executing Vortex arrays into Arrow ([`ArrowSession::execute_arrow`] +//! and the [`ArrowArrayExecutor`] convenience trait). +//! +//! The [`ArrowSession`] session variable is created lazily on first use of +//! [`ArrowSessionExt::arrow`], so no explicit registration is required for plain conversions. use arrow_array::ArrayRef as ArrowArrayRef; use arrow_schema::DataType; +use vortex_array::ArrayRef; +use vortex_array::VortexSessionExecute; +use vortex_array::legacy_session; use vortex_error::VortexResult; +use vortex_session::VortexSession; mod convert; mod datum; +pub mod dtype; mod executor; mod iter; mod null_buffer; +mod run_end_import; +mod scalar; mod session; +mod uuid; +pub use convert::IntoVortexArray; pub(crate) use convert::nulls; pub use datum::*; +pub use dtype::FromArrowType; +pub use dtype::ToArrowType; +pub use dtype::TryFromArrowType; pub use executor::*; pub use iter::*; pub use null_buffer::to_arrow_null_buffer; pub use null_buffer::to_null_buffer; +pub use scalar::ToArrowDatum; pub use session::*; -use crate::ArrayRef; -use crate::LEGACY_SESSION; -use crate::VortexSessionExecute; +/// Register vortex-arrow's session extensions into `session`. +/// +/// This eagerly registers the [`ArrowSession`], which is otherwise created lazily on first use. +pub fn initialize(session: &VortexSession) { + session.register(ArrowSession::default()); +} /// Construct a Vortex array from an Arrow array (or other Arrow container) of type `A`. /// /// Implementations reuse the underlying Arrow buffers without copying wherever the Arrow and /// Vortex memory layouts allow it. pub trait FromArrowArray { - /// Convert `array` into a Vortex array whose [`DType`](crate::dtype::DType) has the requested - /// `nullable` [`Nullability`](crate::dtype::Nullability). + /// Convert `array` into a Vortex array whose [`DType`](vortex_array::dtype::DType) has the requested + /// `nullable` [`Nullability`](vortex_array::dtype::Nullability). /// /// An Arrow array can carry a validity (null) buffer regardless of whether its schema declares /// the field nullable, so the desired nullability is supplied separately by the caller @@ -64,13 +88,18 @@ pub trait IntoArrowArray { #[allow(deprecated)] impl IntoArrowArray for ArrayRef { - /// Convert this [`crate::ArrayRef`] into an Arrow [`crate::ArrayRef`] by using the array's + /// Convert this [`vortex_array::ArrayRef`] into an Arrow [`ArrowArrayRef`] by using the array's /// preferred (cheapest) Arrow [`DataType`]. + #[allow(clippy::disallowed_methods)] fn into_arrow_preferred(self) -> VortexResult { - self.execute_arrow(None, &mut LEGACY_SESSION.create_execution_ctx()) + self.execute_arrow(None, &mut legacy_session().create_execution_ctx()) } + #[allow(clippy::disallowed_methods)] fn into_arrow(self, data_type: &DataType) -> VortexResult { - self.execute_arrow(Some(data_type), &mut LEGACY_SESSION.create_execution_ctx()) + self.execute_arrow( + Some(data_type), + &mut legacy_session().create_execution_ctx(), + ) } } diff --git a/vortex-array/src/arrow/null_buffer.rs b/vortex-arrow/src/null_buffer.rs similarity index 93% rename from vortex-array/src/arrow/null_buffer.rs rename to vortex-arrow/src/null_buffer.rs index 484dd5e55bc..b50adb30cd2 100644 --- a/vortex-array/src/arrow/null_buffer.rs +++ b/vortex-arrow/src/null_buffer.rs @@ -3,12 +3,11 @@ use arrow_buffer::BooleanBuffer; use arrow_buffer::NullBuffer; +use vortex_array::ExecutionCtx; +use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_mask::Mask; -use crate::ExecutionCtx; -use crate::validity::Validity; - /// Converts a [`Validity`] to an Arrow [`NullBuffer`], executing the validity array if needed. pub fn to_arrow_null_buffer( validity: Validity, @@ -38,7 +37,7 @@ mod tests { use vortex_buffer::BitBuffer; use vortex_mask::Mask; - use crate::arrow::null_buffer::to_null_buffer; + use crate::null_buffer::to_null_buffer; #[test] fn test_mask_to_null_buffer() { diff --git a/encodings/runend/src/arrow.rs b/vortex-arrow/src/run_end_import.rs similarity index 89% rename from encodings/runend/src/arrow.rs rename to vortex-arrow/src/run_end_import.rs index 193937122ef..69deba71fdf 100644 --- a/encodings/runend/src/arrow.rs +++ b/vortex-arrow/src/run_end_import.rs @@ -5,26 +5,25 @@ use arrow_array::RunArray; use arrow_array::types::RunEndIndexType; 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::primitive::PrimitiveArrayExt; -use vortex_array::arrow::FromArrowArray; use vortex_array::dtype::NativePType; -use vortex_array::scalar::PValue; -use vortex_array::search_sorted::SearchSorted; -use vortex_array::search_sorted::SearchSortedSide; +use vortex_array::legacy_session; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; +use vortex_runend::RunEndData; +use vortex_runend::ops::find_physical_index; +use vortex_runend::ops::find_slice_end_index; -use crate::RunEndData; -use crate::ops::find_slice_end_index; +use crate::FromArrowArray; impl FromArrowArray<&RunArray> for RunEndData where R::Native: NativePType, { + #[allow(clippy::disallowed_methods)] fn from_arrow(array: &RunArray, nullable: bool) -> VortexResult { let offset = array.run_ends().offset(); let len = array.run_ends().len(); @@ -40,14 +39,13 @@ where ends.validity()?, ) .into_array(); + let mut ctx = legacy_session().create_execution_ctx(); + let (ends_slice, values_slice) = if offset == 0 && len == array.run_ends().max_value() { (ends_array, values) } else { - let slice_begin = ends_array - .as_primitive_typed() - .search_sorted(&PValue::from(offset), SearchSortedSide::Right)? - .to_ends_index(ends_array.len()); - let slice_end = find_slice_end_index(&ends_array, offset + len)?; + let slice_begin = find_physical_index(&ends_array, offset, &mut ctx)?; + let slice_end = find_slice_end_index(&ends_array, offset + len, &mut ctx)?; ( ends_array.slice(slice_begin..slice_end)?, @@ -56,8 +54,6 @@ where }; // SAFETY: arrow-rs enforces the RunEndArray invariants, we inherit their guarantees. - // TODO(ctx): trait fixes - FromArrowArray::from_arrow has a fixed signature. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); RunEndData::validate_parts(&ends_slice, &values_slice, offset, len, &mut ctx)?; Ok(unsafe { RunEndData::new_unchecked(offset) }) } @@ -84,35 +80,34 @@ mod tests { use vortex_array::VortexSessionExecute as _; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; - use vortex_array::arrow::ArrowSessionExt; - use vortex_array::arrow::FromArrowArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::scalar::PValue; - use vortex_array::search_sorted::SearchSorted; - use vortex_array::search_sorted::SearchSortedSide; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_runend::RunEnd; + use vortex_runend::RunEndArray; + use vortex_runend::ops::find_physical_index; + use vortex_runend::ops::find_slice_end_index; use vortex_session::VortexSession; - use crate::RunEnd; - use crate::ops::find_slice_end_index; + use crate::ArrowSessionExt; + use crate::FromArrowArray; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); - crate::initialize(&session); + vortex_runend::initialize(&session); session }); fn decode_run_array( array: &RunArray, nullable: bool, - ) -> VortexResult + ) -> VortexResult where R::Native: NativePType, { @@ -130,14 +125,12 @@ mod tests { ends.validity()?, ) .into_array(); + let mut ctx = SESSION.create_execution_ctx(); let (ends_slice, values_slice) = if offset == 0 && len == array.run_ends().max_value() { (ends_array, values) } else { - let slice_begin = ends_array - .as_primitive_typed() - .search_sorted(&PValue::from(offset), SearchSortedSide::Right)? - .to_ends_index(ends_array.len()); - let slice_end = find_slice_end_index(&ends_array, offset + len)?; + let slice_begin = find_physical_index(&ends_array, offset, &mut ctx)?; + let slice_end = find_slice_end_index(&ends_array, offset + len, &mut ctx)?; ( ends_array.slice(slice_begin..slice_end)?, @@ -145,13 +138,7 @@ mod tests { ) }; - RunEnd::try_new_offset_length( - ends_slice, - values_slice, - offset, - array.len(), - &mut SESSION.create_execution_ctx(), - ) + RunEnd::try_new_offset_length(ends_slice, values_slice, offset, array.len(), &mut ctx) } #[test] diff --git a/vortex-array/src/scalar/arrow.rs b/vortex-arrow/src/scalar.rs similarity index 83% rename from vortex-array/src/scalar/arrow.rs rename to vortex-arrow/src/scalar.rs index cca6749cb19..aeac56f43d3 100644 --- a/vortex-array/src/scalar/arrow.rs +++ b/vortex-arrow/src/scalar.rs @@ -7,24 +7,23 @@ use std::sync::Arc; use arrow_array::Scalar as ArrowScalar; use arrow_array::*; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::extension::datetime::AnyTemporal; +use vortex_array::extension::datetime::TemporalMetadata; +use vortex_array::extension::datetime::TimeUnit; +use vortex_array::scalar::BinaryScalar; +use vortex_array::scalar::BoolScalar; +use vortex_array::scalar::DecimalScalar; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::ExtScalar; +use vortex_array::scalar::PrimitiveScalar; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::Utf8Scalar; use vortex_error::VortexError; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use crate::dtype::DType; -use crate::dtype::PType; -use crate::extension::datetime::AnyTemporal; -use crate::extension::datetime::TemporalMetadata; -use crate::extension::datetime::TimeUnit; -use crate::scalar::BinaryScalar; -use crate::scalar::BoolScalar; -use crate::scalar::DecimalScalar; -use crate::scalar::DecimalValue; -use crate::scalar::ExtScalar; -use crate::scalar::PrimitiveScalar; -use crate::scalar::Scalar; -use crate::scalar::Utf8Scalar; - /// Arrow represents scalars as single-element arrays. This constant is the length of those arrays. const SCALAR_ARRAY_LEN: usize = 1; @@ -50,10 +49,18 @@ macro_rules! timestamp_to_arrow_scalar { }}; } -impl TryFrom<&Scalar> for Arc { - type Error = VortexError; +/// Convert a Vortex [`Scalar`] into an Arrow [`Datum`] (a single-element Arrow array). +/// +/// This mirrors a `TryFrom<&Scalar> for Arc` conversion; a separate trait is +/// required because both `Scalar` and `Datum` are foreign to this crate. +pub trait ToArrowDatum { + /// Convert this scalar to an Arrow [`Datum`]. + fn to_arrow_datum(&self) -> Result, VortexError>; +} - fn try_from(value: &Scalar) -> Result, Self::Error> { +impl ToArrowDatum for Scalar { + fn to_arrow_datum(&self) -> Result, VortexError> { + let value = self; match value.dtype() { DType::Null => Ok(Arc::new(NullArray::new(SCALAR_ARRAY_LEN))), DType::Bool(_) => bool_to_arrow(value.as_bool()), @@ -192,149 +199,149 @@ fn extension_to_arrow(scalar: ExtScalar<'_>) -> Result, VortexErr mod tests { use std::sync::Arc; - use arrow_array::Datum; use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::FieldDType; + use vortex_array::dtype::NativeDType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::StructFields; + use vortex_array::dtype::extension::ExtDType; + use vortex_array::dtype::extension::ExtId; + use vortex_array::dtype::extension::ExtVTable; + use vortex_array::dtype::i256; + use vortex_array::extension::datetime::Date; + use vortex_array::extension::datetime::Time; + use vortex_array::extension::datetime::TimeUnit; + use vortex_array::extension::datetime::Timestamp; + use vortex_array::extension::datetime::TimestampOptions; + use vortex_array::scalar::DecimalValue; + use vortex_array::scalar::Scalar; + use vortex_array::scalar::ScalarValue; use vortex_error::VortexResult; use vortex_error::vortex_bail; - use crate::dtype::DType; - use crate::dtype::DecimalDType; - use crate::dtype::FieldDType; - use crate::dtype::NativeDType; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::StructFields; - use crate::dtype::extension::ExtDType; - use crate::dtype::extension::ExtId; - use crate::dtype::extension::ExtVTable; - use crate::dtype::i256; - use crate::extension::datetime::Date; - use crate::extension::datetime::Time; - use crate::extension::datetime::TimeUnit; - use crate::extension::datetime::Timestamp; - use crate::extension::datetime::TimestampOptions; - use crate::scalar::DecimalValue; - use crate::scalar::Scalar; - use crate::scalar::ScalarValue; + use super::ToArrowDatum; #[test] fn test_null_scalar_to_arrow() { let scalar = Scalar::null(DType::Null); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_bool_scalar_to_arrow() { let scalar = Scalar::bool(true, Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_null_bool_scalar_to_arrow() { let scalar = Scalar::null(bool::dtype().as_nullable()); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_primitive_u8_to_arrow() { let scalar = Scalar::primitive(42u8, Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_primitive_u16_to_arrow() { let scalar = Scalar::primitive(1000u16, Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_primitive_u32_to_arrow() { let scalar = Scalar::primitive(100000u32, Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_primitive_u64_to_arrow() { let scalar = Scalar::primitive(10000000000u64, Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_primitive_i8_to_arrow() { let scalar = Scalar::primitive(-42i8, Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_primitive_i16_to_arrow() { let scalar = Scalar::primitive(-1000i16, Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_primitive_i32_to_arrow() { let scalar = Scalar::primitive(-100000i32, Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_primitive_i64_to_arrow() { let scalar = Scalar::primitive(-10000000000i64, Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_primitive_f16_to_arrow() { - use crate::dtype::half::f16; + use vortex_array::dtype::half::f16; let scalar = Scalar::primitive(f16::from_f32(1.234), Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_primitive_f32_to_arrow() { let scalar = Scalar::primitive(1.234f32, Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_primitive_f64_to_arrow() { let scalar = Scalar::primitive(1.234567890123f64, Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_null_primitive_to_arrow() { let scalar = Scalar::null(i32::dtype().as_nullable()); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_utf8_scalar_to_arrow() { let scalar = Scalar::utf8("hello world".to_string(), Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_null_utf8_scalar_to_arrow() { let scalar = Scalar::null(String::dtype().as_nullable()); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } @@ -342,14 +349,14 @@ mod tests { fn test_binary_scalar_to_arrow() { let data = vec![1u8, 2, 3, 4, 5]; let scalar = Scalar::binary(data, Nullability::NonNullable); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } #[test] fn test_null_binary_scalar_to_arrow() { let scalar = Scalar::null(DType::Binary(Nullability::Nullable)); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } @@ -363,35 +370,35 @@ mod tests { decimal_dtype, Nullability::NonNullable, ); - assert!(Arc::::try_from(&scalar_i8).is_ok()); + assert!(scalar_i8.to_arrow_datum().is_ok()); let scalar_i16 = Scalar::decimal( DecimalValue::I16(10000), decimal_dtype, Nullability::NonNullable, ); - assert!(Arc::::try_from(&scalar_i16).is_ok()); + assert!(scalar_i16.to_arrow_datum().is_ok()); let scalar_i32 = Scalar::decimal( DecimalValue::I32(99999), decimal_dtype, Nullability::NonNullable, ); - assert!(Arc::::try_from(&scalar_i32).is_ok()); + assert!(scalar_i32.to_arrow_datum().is_ok()); let scalar_i64 = Scalar::decimal( DecimalValue::I64(99999), decimal_dtype, Nullability::NonNullable, ); - assert!(Arc::::try_from(&scalar_i64).is_ok()); + assert!(scalar_i64.to_arrow_datum().is_ok()); let scalar_i128 = Scalar::decimal( DecimalValue::I128(99999), decimal_dtype, Nullability::NonNullable, ); - assert!(Arc::::try_from(&scalar_i128).is_ok()); + assert!(scalar_i128.to_arrow_datum().is_ok()); // Test i256 @@ -401,14 +408,14 @@ mod tests { decimal_dtype, Nullability::NonNullable, ); - assert!(Arc::::try_from(&scalar_i256).is_ok()); + assert!(scalar_i256.to_arrow_datum().is_ok()); } #[test] fn test_null_decimal_to_arrow() { let decimal_dtype = DecimalDType::new(10, 2); let scalar = Scalar::null(DType::Decimal(decimal_dtype, Nullability::Nullable)); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } @@ -427,7 +434,7 @@ mod tests { struct_dtype, vec![Scalar::primitive(42i32, Nullability::NonNullable)], ); - Arc::::try_from(&struct_scalar).unwrap(); + struct_scalar.to_arrow_datum().unwrap(); } #[test] @@ -443,7 +450,7 @@ mod tests { Nullability::NonNullable, ); - Arc::::try_from(&list_scalar).unwrap(); + list_scalar.to_arrow_datum().unwrap(); } #[test] @@ -485,7 +492,7 @@ mod tests { Scalar::primitive(42i32, Nullability::NonNullable), ); - Arc::::try_from(&scalar).unwrap(); + scalar.to_arrow_datum().unwrap(); } #[rstest] @@ -510,7 +517,7 @@ mod tests { }, ); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } @@ -534,7 +541,7 @@ mod tests { }, ); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } @@ -552,7 +559,7 @@ mod tests { Scalar::primitive(value, Nullability::NonNullable), ); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } @@ -576,7 +583,7 @@ mod tests { Scalar::primitive(value, Nullability::NonNullable), ); - let result = Arc::::try_from(&scalar); + let result = scalar.to_arrow_datum(); assert!(result.is_ok()); } @@ -587,7 +594,7 @@ mod tests { Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), ); - let _result = Arc::::try_from(&scalar).unwrap(); + let _result = scalar.to_arrow_datum().unwrap(); } #[test] diff --git a/vortex-array/src/arrow/session.rs b/vortex-arrow/src/session.rs similarity index 93% rename from vortex-array/src/arrow/session.rs rename to vortex-arrow/src/session.rs index d3f43443d5d..745e8afeb47 100644 --- a/vortex-array/src/arrow/session.rs +++ b/vortex-arrow/src/session.rs @@ -3,8 +3,8 @@ //! Plugin layer for moving Arrow extension types in and out of Vortex. //! -//! Vortex's canonical Arrow conversion (see [`crate::dtype::arrow`] and the executor in -//! [`crate::arrow::executor`]) handles every non-extension Arrow type and the builtin temporal +//! Vortex's canonical Arrow conversion (see [`crate::dtype`] and the executor in +//! [`crate::executor`]) handles every non-extension Arrow type and the builtin temporal //! extensions. The plugins registered here cover the remaining case: **Arrow extension types**. //! //! * An [`ArrowExportVTable`] is dispatched purely by the **target Arrow extension Id** — @@ -36,6 +36,23 @@ use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; use arrow_schema::extension::ExtensionType; use tracing::debug; use tracing::trace; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arc_swap_map::ArcSwapMap; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::StructArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::dtype::extension::ExtId; +use vortex_array::extension::datetime::AnyTemporal; +use vortex_array::extension::uuid::Uuid; +use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -44,26 +61,13 @@ use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_session::registry::Id; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arc_swap_map::ArcSwapMap; -use crate::arrays::StructArray; -use crate::arrow::FromArrowArray; -use crate::arrow::convert::nulls; -use crate::arrow::convert::remove_nulls; -use crate::arrow::executor::execute_arrow_naive; -use crate::dtype::DType; -use crate::dtype::FieldName; -use crate::dtype::FieldNames; -use crate::dtype::Nullability; -use crate::dtype::StructFields; -use crate::dtype::arrow::TryFromArrowType; -use crate::dtype::arrow::to_data_type_naive; -use crate::dtype::extension::ExtId; -use crate::extension::datetime::AnyTemporal; -use crate::extension::uuid::Uuid; -use crate::validity::Validity; +use crate::FromArrowArray; +use crate::IntoVortexArray; +use crate::convert::nulls; +use crate::convert::remove_nulls; +use crate::dtype::TryFromArrowType; +use crate::dtype::to_data_type_naive; +use crate::executor::execute_arrow_naive; /// Outcome of a successful call to [`ArrowExportVTable::execute_arrow`]. /// @@ -296,9 +300,7 @@ impl ArrowSession { /// extensions are preserved via [`Self::to_arrow_field`]. pub fn to_arrow_schema(&self, dtype: &DType) -> VortexResult { let DType::Struct(struct_dtype, _) = dtype else { - vortex_error::vortex_bail!( - "to_arrow_schema requires a top-level struct dtype, got {dtype}" - ); + vortex_bail!("to_arrow_schema requires a top-level struct dtype, got {dtype}"); }; let mut fields = Vec::with_capacity(struct_dtype.names().len()); for (name, field_dtype) in struct_dtype.names().iter().zip(struct_dtype.fields()) { @@ -537,7 +539,7 @@ impl ArrowSession { .from_arrow_array(ArrowArrayRef::clone(list.values()), elem_field.as_ref())?; let offsets = list.offsets().clone().into_array(); let validity = nulls(list.nulls(), field.is_nullable())?; - Ok(crate::arrays::ListArray::try_new(elements, offsets, validity)?.into_array()) + Ok(ListArray::try_new(elements, offsets, validity)?.into_array()) } DataType::LargeList(elem_field) => { let list = array.as_list::(); @@ -545,20 +547,17 @@ impl ArrowSession { .from_arrow_array(ArrowArrayRef::clone(list.values()), elem_field.as_ref())?; let offsets = list.offsets().clone().into_array(); let validity = nulls(list.nulls(), field.is_nullable())?; - Ok(crate::arrays::ListArray::try_new(elements, offsets, validity)?.into_array()) + Ok(ListArray::try_new(elements, offsets, validity)?.into_array()) } DataType::FixedSizeList(elem_field, list_size) => { let fsl = array.as_fixed_size_list(); let elements = self.from_arrow_array(ArrowArrayRef::clone(fsl.values()), elem_field.as_ref())?; let validity = nulls(fsl.nulls(), field.is_nullable())?; - Ok(crate::arrays::FixedSizeListArray::try_new( - elements, - *list_size as u32, - validity, - fsl.len(), - )? - .into_array()) + Ok( + FixedSizeListArray::try_new(elements, *list_size as u32, validity, fsl.len())? + .into_array(), + ) } DataType::ListView(elem_field) => { let list = array.as_list_view::(); @@ -567,10 +566,7 @@ impl ArrowSession { let offsets = list.offsets().clone().into_array(); let sizes = list.sizes().clone().into_array(); let validity = nulls(list.nulls(), field.is_nullable())?; - Ok( - crate::arrays::ListViewArray::try_new(elements, offsets, sizes, validity)? - .into_array(), - ) + Ok(ListViewArray::try_new(elements, offsets, sizes, validity)?.into_array()) } DataType::LargeListView(elem_field) => { let list = array.as_list_view::(); @@ -579,10 +575,7 @@ impl ArrowSession { let offsets = list.offsets().clone().into_array(); let sizes = list.sizes().clone().into_array(); let validity = nulls(list.nulls(), field.is_nullable())?; - Ok( - crate::arrays::ListViewArray::try_new(elements, offsets, sizes, validity)? - .into_array(), - ) + Ok(ListViewArray::try_new(elements, offsets, sizes, validity)?.into_array()) } _ => ArrayRef::from_arrow(array.as_ref(), field.is_nullable()), } @@ -631,20 +624,20 @@ mod tests { use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::extension::Uuid as ArrowUuid; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::dtype::DType; + use vortex_array::dtype::FieldName; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::StructFields; + use vortex_array::dtype::extension::ExtDType; + use vortex_array::dtype::extension::ExtVTable; + use vortex_array::extension::uuid::Uuid; + use vortex_array::extension::uuid::UuidMetadata; use vortex_error::VortexResult; use super::*; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::dtype::DType; - use crate::dtype::FieldName; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::dtype::StructFields; - use crate::dtype::extension::ExtDType; - use crate::dtype::extension::ExtVTable; - use crate::extension::uuid::Uuid; - use crate::extension::uuid::UuidMetadata; fn uuid_dtype(nullable: bool) -> DType { let storage = DType::FixedSizeList( diff --git a/vortex-array/src/extension/uuid/arrow.rs b/vortex-arrow/src/uuid.rs similarity index 85% rename from vortex-array/src/extension/uuid/arrow.rs rename to vortex-arrow/src/uuid.rs index cd4b97ea60a..5901f56389b 100644 --- a/vortex-array/src/extension/uuid/arrow.rs +++ b/vortex-arrow/src/uuid.rs @@ -18,8 +18,22 @@ use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::extension::ExtensionType; use arrow_schema::extension::Uuid as ArrowUuid; -use vortex_array::arrow::ArrowSession; -use vortex_array::arrow::has_valid_extension_type; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::extension::uuid::Uuid; +use vortex_array::extension::uuid::UuidMetadata; +use vortex_array::validity::Validity; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_error::VortexExpect; @@ -27,28 +41,14 @@ use vortex_error::VortexResult; use vortex_session::registry::CachedId; use vortex_session::registry::Id; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ExtensionArray; -use crate::arrays::FixedSizeListArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::extension::ExtensionArrayExt; -use crate::arrow::ArrowExport; -use crate::arrow::ArrowExportVTable; -use crate::arrow::ArrowImport; -use crate::arrow::ArrowImportVTable; -use crate::arrow::ArrowSessionExt; -use crate::arrow::nulls; -use crate::buffer::BufferHandle; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::dtype::extension::ExtDType; -use crate::dtype::extension::ExtVTable; -use crate::extension::uuid::Uuid; -use crate::extension::uuid::UuidMetadata; -use crate::validity::Validity; +use crate::ArrowExport; +use crate::ArrowExportVTable; +use crate::ArrowImport; +use crate::ArrowImportVTable; +use crate::ArrowSession; +use crate::ArrowSessionExt; +use crate::nulls; +use crate::session::has_valid_extension_type; const UUID_BYTE_LEN: i32 = 16; diff --git a/vortex-arrow/tests/canonical.rs b/vortex-arrow/tests/canonical.rs new file mode 100644 index 00000000000..f67e0b73115 --- /dev/null +++ b/vortex-arrow/tests/canonical.rs @@ -0,0 +1,168 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::tests_outside_test_module)] + +//! Round-trip tests between canonical Vortex arrays and Arrow, moved from +//! `vortex-array/src/canonical.rs` when Arrow interoperability moved into this crate. + +use std::sync::Arc; +use std::sync::LazyLock; + +use arrow_array::Array as ArrowArray; +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::ListArray as ArrowListArray; +use arrow_array::PrimitiveArray as ArrowPrimitiveArray; +use arrow_array::StringArray; +use arrow_array::StringViewArray; +use arrow_array::StructArray as ArrowStructArray; +use arrow_array::cast::AsArray; +use arrow_array::types::Int32Type; +use arrow_array::types::Int64Type; +use arrow_array::types::UInt64Type; +use arrow_buffer::NullBufferBuilder; +use arrow_buffer::OffsetBuffer; +use arrow_schema::DataType; +use arrow_schema::Field; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::StructArray; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; +use vortex_buffer::buffer; +use vortex_session::VortexSession; + +/// A shared session for these canonical tests, used to create execution contexts. +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +#[test] +fn test_canonicalize_nested_struct() { + let mut ctx = SESSION.create_execution_ctx(); + // Create a struct array with multiple internal components. + let nested_struct_array = StructArray::from_fields(&[ + ("a", buffer![1u64].into_array()), + ( + "b", + StructArray::from_fields(&[( + "inner_a", + // The nested struct contains a ConstantArray representing the primitive array + // [100i64] + // ConstantArray is not a canonical type, so converting `into_arrow()` should + // map this to the nearest canonical type (PrimitiveArray). + ConstantArray::new(100i64, 1).into_array(), + )]) + .unwrap() + .into_array(), + ), + ]) + .unwrap(); + + let arrow_struct = SESSION + .arrow() + .execute_arrow(nested_struct_array.into_array(), None, &mut ctx) + .unwrap() + .as_any() + .downcast_ref::() + .cloned() + .unwrap(); + + assert!( + arrow_struct + .column(0) + .as_any() + .downcast_ref::>() + .is_some() + ); + + let inner_struct = Arc::clone(arrow_struct.column(1)) + .as_any() + .downcast_ref::() + .cloned() + .unwrap(); + + let inner_a = inner_struct + .column(0) + .as_any() + .downcast_ref::>(); + assert!(inner_a.is_some()); + + assert_eq!( + inner_a.cloned().unwrap(), + ArrowPrimitiveArray::from_iter([100i64]) + ); +} + +#[test] +fn roundtrip_struct() { + let mut ctx = SESSION.create_execution_ctx(); + let mut nulls = NullBufferBuilder::new(6); + nulls.append_n_non_nulls(4); + nulls.append_null(); + nulls.append_non_null(); + let names = Arc::new(StringViewArray::from_iter(vec![ + Some("Joseph"), + None, + Some("Angela"), + Some("Mikhail"), + None, + None, + ])); + let ages = Arc::new(ArrowPrimitiveArray::::from(vec![ + Some(25), + Some(31), + None, + Some(57), + None, + None, + ])); + + let arrow_struct = ArrowStructArray::new( + vec![ + Arc::new(Field::new("name", DataType::Utf8View, true)), + Arc::new(Field::new("age", DataType::Int32, true)), + ] + .into(), + vec![names, ages], + nulls.finish(), + ); + + let vortex_struct = ArrayRef::from_arrow(&arrow_struct, true).unwrap(); + let vortex_struct = SESSION + .arrow() + .execute_arrow(vortex_struct, None, &mut ctx) + .unwrap(); + assert_eq!(&arrow_struct, vortex_struct.as_struct()); +} + +#[test] +fn roundtrip_list() { + let mut ctx = SESSION.create_execution_ctx(); + let names = Arc::new(StringArray::from_iter(vec![ + Some("Joseph"), + Some("Angela"), + Some("Mikhail"), + ])); + + let arrow_list = ArrowListArray::new( + Arc::new(Field::new_list_field(DataType::Utf8, true)), + OffsetBuffer::from_lengths(vec![0, 2, 1]), + names, + None, + ); + let list_data_type = arrow_list.data_type(); + let list_field = Field::new(String::new(), list_data_type.clone(), true); + + let vortex_list = ArrayRef::from_arrow(&arrow_list, true).unwrap(); + + let rt_arrow_list = SESSION + .arrow() + .execute_arrow(vortex_list, Some(&list_field), &mut ctx) + .unwrap(); + + assert_eq!( + (Arc::new(arrow_list.clone()) as ArrowArrayRef).as_ref(), + rt_arrow_list.as_ref() + ); +} diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index 9cd7fc92d30..f228e1571d3 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -23,6 +23,7 @@ vortex = { workspace = true, features = [ "tokio", "zstd", ] } +vortex-arrow = { workspace = true } vortex-geo = { workspace = true } vortex-tensor = { workspace = true } # TODO(connor): In the future, this might be inside vortex. @@ -34,6 +35,10 @@ async-trait = { workspace = true } bzip2 = { workspace = true } clap = { workspace = true, features = ["derive"] } futures = { workspace = true } +geo = { workspace = true } +geo-traits = { workspace = true } +geoarrow = { workspace = true } +geoarrow-cast = { workspace = true } get_dir = { workspace = true } glob = { workspace = true } humansize = { workspace = true } @@ -69,6 +74,7 @@ tracing-subscriber = { workspace = true, features = [ ] } url = { workspace = true } uuid = { workspace = true, features = ["v4"] } +wkb = { workspace = true } [dev-dependencies] insta = { workspace = true } diff --git a/vortex-bench/appian/q1.sql b/vortex-bench/sql/appian/q1.sql similarity index 100% rename from vortex-bench/appian/q1.sql rename to vortex-bench/sql/appian/q1.sql diff --git a/vortex-bench/appian/q2.sql b/vortex-bench/sql/appian/q2.sql similarity index 100% rename from vortex-bench/appian/q2.sql rename to vortex-bench/sql/appian/q2.sql diff --git a/vortex-bench/appian/q3.sql b/vortex-bench/sql/appian/q3.sql similarity index 100% rename from vortex-bench/appian/q3.sql rename to vortex-bench/sql/appian/q3.sql diff --git a/vortex-bench/appian/q4.sql b/vortex-bench/sql/appian/q4.sql similarity index 100% rename from vortex-bench/appian/q4.sql rename to vortex-bench/sql/appian/q4.sql diff --git a/vortex-bench/appian/q5.sql b/vortex-bench/sql/appian/q5.sql similarity index 100% rename from vortex-bench/appian/q5.sql rename to vortex-bench/sql/appian/q5.sql diff --git a/vortex-bench/appian/q6.sql b/vortex-bench/sql/appian/q6.sql similarity index 100% rename from vortex-bench/appian/q6.sql rename to vortex-bench/sql/appian/q6.sql diff --git a/vortex-bench/appian/q7.sql b/vortex-bench/sql/appian/q7.sql similarity index 100% rename from vortex-bench/appian/q7.sql rename to vortex-bench/sql/appian/q7.sql diff --git a/vortex-bench/appian/q8.sql b/vortex-bench/sql/appian/q8.sql similarity index 100% rename from vortex-bench/appian/q8.sql rename to vortex-bench/sql/appian/q8.sql diff --git a/vortex-bench/clickbench_queries.sql b/vortex-bench/sql/clickbench_queries.sql similarity index 100% rename from vortex-bench/clickbench_queries.sql rename to vortex-bench/sql/clickbench_queries.sql diff --git a/vortex-bench/sql/fineweb.sql b/vortex-bench/sql/fineweb.sql new file mode 100644 index 00000000000..7f97d303236 --- /dev/null +++ b/vortex-bench/sql/fineweb.sql @@ -0,0 +1,30 @@ +-- Some basic string-focused queries against the HuggingFace FineWeb dataset, numbered from +-- Q0 in file order. The harness splits the file on semicolons, so a comment must never +-- contain one. + +-- Q0: simple summary. +SELECT count(DISTINCT dump) FROM fineweb; + +-- Q1: selective string equality filter. +SELECT * FROM fineweb WHERE dump = 'CC-MAIN-2016-30'; + +-- Q2: LIKE with prefix filter. +SELECT * FROM fineweb WHERE date LIKE '2020-10-%'; + +-- Q3: LIKE with simple containment filter. +SELECT * FROM fineweb WHERE url LIKE '%google%' AND text LIKE '%Google%'; + +-- Q4: LIKE with larger containment filter. +SELECT * FROM fineweb WHERE url LIKE '%.google.%' OR text LIKE '% Google %'; + +-- Q5: LIKE with rare containment filter. +SELECT * FROM fineweb WHERE text LIKE '% vortex %'; + +-- Q6: more LIKE filters. +SELECT * FROM fineweb WHERE url LIKE '%espn%' AND language = 'en' AND language_score > 0.92; + +-- Q7: more LIKE filters. +SELECT * FROM fineweb WHERE url LIKE '%espn%' OR url LIKE '%www.espn.go.com%' OR url LIKE '%espn.go.com%'; + +-- Q8: no results, stats cannot prune but tokenized bloom filters could. +SELECT * FROM fineweb WHERE file_path LIKE '%/CC-MAIN-2014-%'; diff --git a/vortex-bench/sql/gharchive.sql b/vortex-bench/sql/gharchive.sql new file mode 100644 index 00000000000..a3280bd25d4 --- /dev/null +++ b/vortex-bench/sql/gharchive.sql @@ -0,0 +1,17 @@ +-- GitHub Archive queries over nested event data, numbered from Q0 in file order. The +-- harness splits the file on semicolons, so a comment must never contain one. + +-- Q0. +select count(*) from events where payload.ref = 'refs/heads/main'; + +-- Q1. +select distinct repo.name from events where repo.name like 'spiraldb/%'; + +-- Q2. +select distinct org.id as org_id from events order by org_id limit 100; + +-- Q3. +select actor.login, count() as freq from events group by actor.login order by freq desc limit 10; + +-- Q4. +select actor.avatar_url from events where actor.login = 'renovate[bot]'; diff --git a/vortex-bench/polarsignals.sql b/vortex-bench/sql/polarsignals.sql similarity index 100% rename from vortex-bench/polarsignals.sql rename to vortex-bench/sql/polarsignals.sql diff --git a/vortex-bench/spatialbench.sql b/vortex-bench/sql/spatialbench.sql similarity index 100% rename from vortex-bench/spatialbench.sql rename to vortex-bench/sql/spatialbench.sql diff --git a/vortex-bench/statpopgen.sql b/vortex-bench/sql/statpopgen.sql similarity index 100% rename from vortex-bench/statpopgen.sql rename to vortex-bench/sql/statpopgen.sql diff --git a/vortex-bench/tpcds/01.sql b/vortex-bench/sql/tpcds/01.sql similarity index 100% rename from vortex-bench/tpcds/01.sql rename to vortex-bench/sql/tpcds/01.sql diff --git a/vortex-bench/tpcds/02.sql b/vortex-bench/sql/tpcds/02.sql similarity index 100% rename from vortex-bench/tpcds/02.sql rename to vortex-bench/sql/tpcds/02.sql diff --git a/vortex-bench/tpcds/03.sql b/vortex-bench/sql/tpcds/03.sql similarity index 100% rename from vortex-bench/tpcds/03.sql rename to vortex-bench/sql/tpcds/03.sql diff --git a/vortex-bench/tpcds/04.sql b/vortex-bench/sql/tpcds/04.sql similarity index 100% rename from vortex-bench/tpcds/04.sql rename to vortex-bench/sql/tpcds/04.sql diff --git a/vortex-bench/tpcds/05.sql b/vortex-bench/sql/tpcds/05.sql similarity index 100% rename from vortex-bench/tpcds/05.sql rename to vortex-bench/sql/tpcds/05.sql diff --git a/vortex-bench/tpcds/06.sql b/vortex-bench/sql/tpcds/06.sql similarity index 100% rename from vortex-bench/tpcds/06.sql rename to vortex-bench/sql/tpcds/06.sql diff --git a/vortex-bench/tpcds/07.sql b/vortex-bench/sql/tpcds/07.sql similarity index 100% rename from vortex-bench/tpcds/07.sql rename to vortex-bench/sql/tpcds/07.sql diff --git a/vortex-bench/tpcds/08.sql b/vortex-bench/sql/tpcds/08.sql similarity index 100% rename from vortex-bench/tpcds/08.sql rename to vortex-bench/sql/tpcds/08.sql diff --git a/vortex-bench/tpcds/09.sql b/vortex-bench/sql/tpcds/09.sql similarity index 100% rename from vortex-bench/tpcds/09.sql rename to vortex-bench/sql/tpcds/09.sql diff --git a/vortex-bench/tpcds/10.sql b/vortex-bench/sql/tpcds/10.sql similarity index 100% rename from vortex-bench/tpcds/10.sql rename to vortex-bench/sql/tpcds/10.sql diff --git a/vortex-bench/tpcds/11.sql b/vortex-bench/sql/tpcds/11.sql similarity index 100% rename from vortex-bench/tpcds/11.sql rename to vortex-bench/sql/tpcds/11.sql diff --git a/vortex-bench/tpcds/12.sql b/vortex-bench/sql/tpcds/12.sql similarity index 100% rename from vortex-bench/tpcds/12.sql rename to vortex-bench/sql/tpcds/12.sql diff --git a/vortex-bench/tpcds/13.sql b/vortex-bench/sql/tpcds/13.sql similarity index 100% rename from vortex-bench/tpcds/13.sql rename to vortex-bench/sql/tpcds/13.sql diff --git a/vortex-bench/tpcds/14.sql b/vortex-bench/sql/tpcds/14.sql similarity index 100% rename from vortex-bench/tpcds/14.sql rename to vortex-bench/sql/tpcds/14.sql diff --git a/vortex-bench/tpcds/15.sql b/vortex-bench/sql/tpcds/15.sql similarity index 100% rename from vortex-bench/tpcds/15.sql rename to vortex-bench/sql/tpcds/15.sql diff --git a/vortex-bench/tpcds/16.sql b/vortex-bench/sql/tpcds/16.sql similarity index 100% rename from vortex-bench/tpcds/16.sql rename to vortex-bench/sql/tpcds/16.sql diff --git a/vortex-bench/tpcds/17.sql b/vortex-bench/sql/tpcds/17.sql similarity index 100% rename from vortex-bench/tpcds/17.sql rename to vortex-bench/sql/tpcds/17.sql diff --git a/vortex-bench/tpcds/18.sql b/vortex-bench/sql/tpcds/18.sql similarity index 100% rename from vortex-bench/tpcds/18.sql rename to vortex-bench/sql/tpcds/18.sql diff --git a/vortex-bench/tpcds/19.sql b/vortex-bench/sql/tpcds/19.sql similarity index 100% rename from vortex-bench/tpcds/19.sql rename to vortex-bench/sql/tpcds/19.sql diff --git a/vortex-bench/tpcds/20.sql b/vortex-bench/sql/tpcds/20.sql similarity index 100% rename from vortex-bench/tpcds/20.sql rename to vortex-bench/sql/tpcds/20.sql diff --git a/vortex-bench/tpcds/21.sql b/vortex-bench/sql/tpcds/21.sql similarity index 100% rename from vortex-bench/tpcds/21.sql rename to vortex-bench/sql/tpcds/21.sql diff --git a/vortex-bench/tpcds/22.sql b/vortex-bench/sql/tpcds/22.sql similarity index 100% rename from vortex-bench/tpcds/22.sql rename to vortex-bench/sql/tpcds/22.sql diff --git a/vortex-bench/tpcds/23.sql b/vortex-bench/sql/tpcds/23.sql similarity index 100% rename from vortex-bench/tpcds/23.sql rename to vortex-bench/sql/tpcds/23.sql diff --git a/vortex-bench/tpcds/24.sql b/vortex-bench/sql/tpcds/24.sql similarity index 100% rename from vortex-bench/tpcds/24.sql rename to vortex-bench/sql/tpcds/24.sql diff --git a/vortex-bench/tpcds/25.sql b/vortex-bench/sql/tpcds/25.sql similarity index 100% rename from vortex-bench/tpcds/25.sql rename to vortex-bench/sql/tpcds/25.sql diff --git a/vortex-bench/tpcds/26.sql b/vortex-bench/sql/tpcds/26.sql similarity index 100% rename from vortex-bench/tpcds/26.sql rename to vortex-bench/sql/tpcds/26.sql diff --git a/vortex-bench/tpcds/27.sql b/vortex-bench/sql/tpcds/27.sql similarity index 100% rename from vortex-bench/tpcds/27.sql rename to vortex-bench/sql/tpcds/27.sql diff --git a/vortex-bench/tpcds/28.sql b/vortex-bench/sql/tpcds/28.sql similarity index 100% rename from vortex-bench/tpcds/28.sql rename to vortex-bench/sql/tpcds/28.sql diff --git a/vortex-bench/tpcds/29.sql b/vortex-bench/sql/tpcds/29.sql similarity index 100% rename from vortex-bench/tpcds/29.sql rename to vortex-bench/sql/tpcds/29.sql diff --git a/vortex-bench/tpcds/30.sql b/vortex-bench/sql/tpcds/30.sql similarity index 100% rename from vortex-bench/tpcds/30.sql rename to vortex-bench/sql/tpcds/30.sql diff --git a/vortex-bench/tpcds/31.sql b/vortex-bench/sql/tpcds/31.sql similarity index 100% rename from vortex-bench/tpcds/31.sql rename to vortex-bench/sql/tpcds/31.sql diff --git a/vortex-bench/tpcds/32.sql b/vortex-bench/sql/tpcds/32.sql similarity index 100% rename from vortex-bench/tpcds/32.sql rename to vortex-bench/sql/tpcds/32.sql diff --git a/vortex-bench/tpcds/33.sql b/vortex-bench/sql/tpcds/33.sql similarity index 100% rename from vortex-bench/tpcds/33.sql rename to vortex-bench/sql/tpcds/33.sql diff --git a/vortex-bench/tpcds/34.sql b/vortex-bench/sql/tpcds/34.sql similarity index 100% rename from vortex-bench/tpcds/34.sql rename to vortex-bench/sql/tpcds/34.sql diff --git a/vortex-bench/tpcds/35.sql b/vortex-bench/sql/tpcds/35.sql similarity index 100% rename from vortex-bench/tpcds/35.sql rename to vortex-bench/sql/tpcds/35.sql diff --git a/vortex-bench/tpcds/36.sql b/vortex-bench/sql/tpcds/36.sql similarity index 100% rename from vortex-bench/tpcds/36.sql rename to vortex-bench/sql/tpcds/36.sql diff --git a/vortex-bench/tpcds/37.sql b/vortex-bench/sql/tpcds/37.sql similarity index 100% rename from vortex-bench/tpcds/37.sql rename to vortex-bench/sql/tpcds/37.sql diff --git a/vortex-bench/tpcds/38.sql b/vortex-bench/sql/tpcds/38.sql similarity index 100% rename from vortex-bench/tpcds/38.sql rename to vortex-bench/sql/tpcds/38.sql diff --git a/vortex-bench/tpcds/39.sql b/vortex-bench/sql/tpcds/39.sql similarity index 100% rename from vortex-bench/tpcds/39.sql rename to vortex-bench/sql/tpcds/39.sql diff --git a/vortex-bench/tpcds/40.sql b/vortex-bench/sql/tpcds/40.sql similarity index 100% rename from vortex-bench/tpcds/40.sql rename to vortex-bench/sql/tpcds/40.sql diff --git a/vortex-bench/tpcds/41.sql b/vortex-bench/sql/tpcds/41.sql similarity index 100% rename from vortex-bench/tpcds/41.sql rename to vortex-bench/sql/tpcds/41.sql diff --git a/vortex-bench/tpcds/42.sql b/vortex-bench/sql/tpcds/42.sql similarity index 100% rename from vortex-bench/tpcds/42.sql rename to vortex-bench/sql/tpcds/42.sql diff --git a/vortex-bench/tpcds/43.sql b/vortex-bench/sql/tpcds/43.sql similarity index 100% rename from vortex-bench/tpcds/43.sql rename to vortex-bench/sql/tpcds/43.sql diff --git a/vortex-bench/tpcds/44.sql b/vortex-bench/sql/tpcds/44.sql similarity index 100% rename from vortex-bench/tpcds/44.sql rename to vortex-bench/sql/tpcds/44.sql diff --git a/vortex-bench/tpcds/45.sql b/vortex-bench/sql/tpcds/45.sql similarity index 100% rename from vortex-bench/tpcds/45.sql rename to vortex-bench/sql/tpcds/45.sql diff --git a/vortex-bench/tpcds/46.sql b/vortex-bench/sql/tpcds/46.sql similarity index 100% rename from vortex-bench/tpcds/46.sql rename to vortex-bench/sql/tpcds/46.sql diff --git a/vortex-bench/tpcds/47.sql b/vortex-bench/sql/tpcds/47.sql similarity index 100% rename from vortex-bench/tpcds/47.sql rename to vortex-bench/sql/tpcds/47.sql diff --git a/vortex-bench/tpcds/48.sql b/vortex-bench/sql/tpcds/48.sql similarity index 100% rename from vortex-bench/tpcds/48.sql rename to vortex-bench/sql/tpcds/48.sql diff --git a/vortex-bench/tpcds/49.sql b/vortex-bench/sql/tpcds/49.sql similarity index 100% rename from vortex-bench/tpcds/49.sql rename to vortex-bench/sql/tpcds/49.sql diff --git a/vortex-bench/tpcds/50.sql b/vortex-bench/sql/tpcds/50.sql similarity index 100% rename from vortex-bench/tpcds/50.sql rename to vortex-bench/sql/tpcds/50.sql diff --git a/vortex-bench/tpcds/51.sql b/vortex-bench/sql/tpcds/51.sql similarity index 100% rename from vortex-bench/tpcds/51.sql rename to vortex-bench/sql/tpcds/51.sql diff --git a/vortex-bench/tpcds/52.sql b/vortex-bench/sql/tpcds/52.sql similarity index 100% rename from vortex-bench/tpcds/52.sql rename to vortex-bench/sql/tpcds/52.sql diff --git a/vortex-bench/tpcds/53.sql b/vortex-bench/sql/tpcds/53.sql similarity index 100% rename from vortex-bench/tpcds/53.sql rename to vortex-bench/sql/tpcds/53.sql diff --git a/vortex-bench/tpcds/54.sql b/vortex-bench/sql/tpcds/54.sql similarity index 100% rename from vortex-bench/tpcds/54.sql rename to vortex-bench/sql/tpcds/54.sql diff --git a/vortex-bench/tpcds/55.sql b/vortex-bench/sql/tpcds/55.sql similarity index 100% rename from vortex-bench/tpcds/55.sql rename to vortex-bench/sql/tpcds/55.sql diff --git a/vortex-bench/tpcds/56.sql b/vortex-bench/sql/tpcds/56.sql similarity index 100% rename from vortex-bench/tpcds/56.sql rename to vortex-bench/sql/tpcds/56.sql diff --git a/vortex-bench/tpcds/57.sql b/vortex-bench/sql/tpcds/57.sql similarity index 100% rename from vortex-bench/tpcds/57.sql rename to vortex-bench/sql/tpcds/57.sql diff --git a/vortex-bench/tpcds/58.sql b/vortex-bench/sql/tpcds/58.sql similarity index 100% rename from vortex-bench/tpcds/58.sql rename to vortex-bench/sql/tpcds/58.sql diff --git a/vortex-bench/tpcds/59.sql b/vortex-bench/sql/tpcds/59.sql similarity index 100% rename from vortex-bench/tpcds/59.sql rename to vortex-bench/sql/tpcds/59.sql diff --git a/vortex-bench/tpcds/60.sql b/vortex-bench/sql/tpcds/60.sql similarity index 100% rename from vortex-bench/tpcds/60.sql rename to vortex-bench/sql/tpcds/60.sql diff --git a/vortex-bench/tpcds/61.sql b/vortex-bench/sql/tpcds/61.sql similarity index 100% rename from vortex-bench/tpcds/61.sql rename to vortex-bench/sql/tpcds/61.sql diff --git a/vortex-bench/tpcds/62.sql b/vortex-bench/sql/tpcds/62.sql similarity index 100% rename from vortex-bench/tpcds/62.sql rename to vortex-bench/sql/tpcds/62.sql diff --git a/vortex-bench/tpcds/63.sql b/vortex-bench/sql/tpcds/63.sql similarity index 100% rename from vortex-bench/tpcds/63.sql rename to vortex-bench/sql/tpcds/63.sql diff --git a/vortex-bench/tpcds/64.sql b/vortex-bench/sql/tpcds/64.sql similarity index 100% rename from vortex-bench/tpcds/64.sql rename to vortex-bench/sql/tpcds/64.sql diff --git a/vortex-bench/tpcds/65.sql b/vortex-bench/sql/tpcds/65.sql similarity index 100% rename from vortex-bench/tpcds/65.sql rename to vortex-bench/sql/tpcds/65.sql diff --git a/vortex-bench/tpcds/66.sql b/vortex-bench/sql/tpcds/66.sql similarity index 100% rename from vortex-bench/tpcds/66.sql rename to vortex-bench/sql/tpcds/66.sql diff --git a/vortex-bench/tpcds/67.sql b/vortex-bench/sql/tpcds/67.sql similarity index 100% rename from vortex-bench/tpcds/67.sql rename to vortex-bench/sql/tpcds/67.sql diff --git a/vortex-bench/tpcds/68.sql b/vortex-bench/sql/tpcds/68.sql similarity index 100% rename from vortex-bench/tpcds/68.sql rename to vortex-bench/sql/tpcds/68.sql diff --git a/vortex-bench/tpcds/69.sql b/vortex-bench/sql/tpcds/69.sql similarity index 100% rename from vortex-bench/tpcds/69.sql rename to vortex-bench/sql/tpcds/69.sql diff --git a/vortex-bench/tpcds/70.sql b/vortex-bench/sql/tpcds/70.sql similarity index 100% rename from vortex-bench/tpcds/70.sql rename to vortex-bench/sql/tpcds/70.sql diff --git a/vortex-bench/tpcds/71.sql b/vortex-bench/sql/tpcds/71.sql similarity index 100% rename from vortex-bench/tpcds/71.sql rename to vortex-bench/sql/tpcds/71.sql diff --git a/vortex-bench/tpcds/72.sql b/vortex-bench/sql/tpcds/72.sql similarity index 100% rename from vortex-bench/tpcds/72.sql rename to vortex-bench/sql/tpcds/72.sql diff --git a/vortex-bench/tpcds/73.sql b/vortex-bench/sql/tpcds/73.sql similarity index 100% rename from vortex-bench/tpcds/73.sql rename to vortex-bench/sql/tpcds/73.sql diff --git a/vortex-bench/tpcds/74.sql b/vortex-bench/sql/tpcds/74.sql similarity index 100% rename from vortex-bench/tpcds/74.sql rename to vortex-bench/sql/tpcds/74.sql diff --git a/vortex-bench/tpcds/75.sql b/vortex-bench/sql/tpcds/75.sql similarity index 100% rename from vortex-bench/tpcds/75.sql rename to vortex-bench/sql/tpcds/75.sql diff --git a/vortex-bench/tpcds/76.sql b/vortex-bench/sql/tpcds/76.sql similarity index 100% rename from vortex-bench/tpcds/76.sql rename to vortex-bench/sql/tpcds/76.sql diff --git a/vortex-bench/tpcds/77.sql b/vortex-bench/sql/tpcds/77.sql similarity index 100% rename from vortex-bench/tpcds/77.sql rename to vortex-bench/sql/tpcds/77.sql diff --git a/vortex-bench/tpcds/78.sql b/vortex-bench/sql/tpcds/78.sql similarity index 100% rename from vortex-bench/tpcds/78.sql rename to vortex-bench/sql/tpcds/78.sql diff --git a/vortex-bench/tpcds/79.sql b/vortex-bench/sql/tpcds/79.sql similarity index 100% rename from vortex-bench/tpcds/79.sql rename to vortex-bench/sql/tpcds/79.sql diff --git a/vortex-bench/tpcds/80.sql b/vortex-bench/sql/tpcds/80.sql similarity index 100% rename from vortex-bench/tpcds/80.sql rename to vortex-bench/sql/tpcds/80.sql diff --git a/vortex-bench/tpcds/81.sql b/vortex-bench/sql/tpcds/81.sql similarity index 100% rename from vortex-bench/tpcds/81.sql rename to vortex-bench/sql/tpcds/81.sql diff --git a/vortex-bench/tpcds/82.sql b/vortex-bench/sql/tpcds/82.sql similarity index 100% rename from vortex-bench/tpcds/82.sql rename to vortex-bench/sql/tpcds/82.sql diff --git a/vortex-bench/tpcds/83.sql b/vortex-bench/sql/tpcds/83.sql similarity index 100% rename from vortex-bench/tpcds/83.sql rename to vortex-bench/sql/tpcds/83.sql diff --git a/vortex-bench/tpcds/84.sql b/vortex-bench/sql/tpcds/84.sql similarity index 100% rename from vortex-bench/tpcds/84.sql rename to vortex-bench/sql/tpcds/84.sql diff --git a/vortex-bench/tpcds/85.sql b/vortex-bench/sql/tpcds/85.sql similarity index 100% rename from vortex-bench/tpcds/85.sql rename to vortex-bench/sql/tpcds/85.sql diff --git a/vortex-bench/tpcds/86.sql b/vortex-bench/sql/tpcds/86.sql similarity index 100% rename from vortex-bench/tpcds/86.sql rename to vortex-bench/sql/tpcds/86.sql diff --git a/vortex-bench/tpcds/87.sql b/vortex-bench/sql/tpcds/87.sql similarity index 100% rename from vortex-bench/tpcds/87.sql rename to vortex-bench/sql/tpcds/87.sql diff --git a/vortex-bench/tpcds/88.sql b/vortex-bench/sql/tpcds/88.sql similarity index 100% rename from vortex-bench/tpcds/88.sql rename to vortex-bench/sql/tpcds/88.sql diff --git a/vortex-bench/tpcds/89.sql b/vortex-bench/sql/tpcds/89.sql similarity index 100% rename from vortex-bench/tpcds/89.sql rename to vortex-bench/sql/tpcds/89.sql diff --git a/vortex-bench/tpcds/90.sql b/vortex-bench/sql/tpcds/90.sql similarity index 100% rename from vortex-bench/tpcds/90.sql rename to vortex-bench/sql/tpcds/90.sql diff --git a/vortex-bench/tpcds/91.sql b/vortex-bench/sql/tpcds/91.sql similarity index 100% rename from vortex-bench/tpcds/91.sql rename to vortex-bench/sql/tpcds/91.sql diff --git a/vortex-bench/tpcds/92.sql b/vortex-bench/sql/tpcds/92.sql similarity index 100% rename from vortex-bench/tpcds/92.sql rename to vortex-bench/sql/tpcds/92.sql diff --git a/vortex-bench/tpcds/93.sql b/vortex-bench/sql/tpcds/93.sql similarity index 100% rename from vortex-bench/tpcds/93.sql rename to vortex-bench/sql/tpcds/93.sql diff --git a/vortex-bench/tpcds/94.sql b/vortex-bench/sql/tpcds/94.sql similarity index 100% rename from vortex-bench/tpcds/94.sql rename to vortex-bench/sql/tpcds/94.sql diff --git a/vortex-bench/tpcds/95.sql b/vortex-bench/sql/tpcds/95.sql similarity index 100% rename from vortex-bench/tpcds/95.sql rename to vortex-bench/sql/tpcds/95.sql diff --git a/vortex-bench/tpcds/96.sql b/vortex-bench/sql/tpcds/96.sql similarity index 100% rename from vortex-bench/tpcds/96.sql rename to vortex-bench/sql/tpcds/96.sql diff --git a/vortex-bench/tpcds/97.sql b/vortex-bench/sql/tpcds/97.sql similarity index 100% rename from vortex-bench/tpcds/97.sql rename to vortex-bench/sql/tpcds/97.sql diff --git a/vortex-bench/tpcds/98.sql b/vortex-bench/sql/tpcds/98.sql similarity index 100% rename from vortex-bench/tpcds/98.sql rename to vortex-bench/sql/tpcds/98.sql diff --git a/vortex-bench/tpcds/99.sql b/vortex-bench/sql/tpcds/99.sql similarity index 100% rename from vortex-bench/tpcds/99.sql rename to vortex-bench/sql/tpcds/99.sql diff --git a/vortex-bench/tpch/q1.sql b/vortex-bench/sql/tpch/q1.sql similarity index 100% rename from vortex-bench/tpch/q1.sql rename to vortex-bench/sql/tpch/q1.sql diff --git a/vortex-bench/tpch/q10.sql b/vortex-bench/sql/tpch/q10.sql similarity index 100% rename from vortex-bench/tpch/q10.sql rename to vortex-bench/sql/tpch/q10.sql diff --git a/vortex-bench/tpch/q11.sql b/vortex-bench/sql/tpch/q11.sql similarity index 100% rename from vortex-bench/tpch/q11.sql rename to vortex-bench/sql/tpch/q11.sql diff --git a/vortex-bench/tpch/q12.sql b/vortex-bench/sql/tpch/q12.sql similarity index 100% rename from vortex-bench/tpch/q12.sql rename to vortex-bench/sql/tpch/q12.sql diff --git a/vortex-bench/tpch/q13.sql b/vortex-bench/sql/tpch/q13.sql similarity index 100% rename from vortex-bench/tpch/q13.sql rename to vortex-bench/sql/tpch/q13.sql diff --git a/vortex-bench/tpch/q14.sql b/vortex-bench/sql/tpch/q14.sql similarity index 100% rename from vortex-bench/tpch/q14.sql rename to vortex-bench/sql/tpch/q14.sql diff --git a/vortex-bench/tpch/q15.sql b/vortex-bench/sql/tpch/q15.sql similarity index 100% rename from vortex-bench/tpch/q15.sql rename to vortex-bench/sql/tpch/q15.sql diff --git a/vortex-bench/tpch/q16.sql b/vortex-bench/sql/tpch/q16.sql similarity index 100% rename from vortex-bench/tpch/q16.sql rename to vortex-bench/sql/tpch/q16.sql diff --git a/vortex-bench/tpch/q17.sql b/vortex-bench/sql/tpch/q17.sql similarity index 100% rename from vortex-bench/tpch/q17.sql rename to vortex-bench/sql/tpch/q17.sql diff --git a/vortex-bench/tpch/q18.sql b/vortex-bench/sql/tpch/q18.sql similarity index 100% rename from vortex-bench/tpch/q18.sql rename to vortex-bench/sql/tpch/q18.sql diff --git a/vortex-bench/tpch/q19.sql b/vortex-bench/sql/tpch/q19.sql similarity index 100% rename from vortex-bench/tpch/q19.sql rename to vortex-bench/sql/tpch/q19.sql diff --git a/vortex-bench/tpch/q2.sql b/vortex-bench/sql/tpch/q2.sql similarity index 100% rename from vortex-bench/tpch/q2.sql rename to vortex-bench/sql/tpch/q2.sql diff --git a/vortex-bench/tpch/q20.sql b/vortex-bench/sql/tpch/q20.sql similarity index 100% rename from vortex-bench/tpch/q20.sql rename to vortex-bench/sql/tpch/q20.sql diff --git a/vortex-bench/tpch/q21.sql b/vortex-bench/sql/tpch/q21.sql similarity index 100% rename from vortex-bench/tpch/q21.sql rename to vortex-bench/sql/tpch/q21.sql diff --git a/vortex-bench/tpch/q22.sql b/vortex-bench/sql/tpch/q22.sql similarity index 100% rename from vortex-bench/tpch/q22.sql rename to vortex-bench/sql/tpch/q22.sql diff --git a/vortex-bench/tpch/q3.sql b/vortex-bench/sql/tpch/q3.sql similarity index 100% rename from vortex-bench/tpch/q3.sql rename to vortex-bench/sql/tpch/q3.sql diff --git a/vortex-bench/tpch/q4.sql b/vortex-bench/sql/tpch/q4.sql similarity index 100% rename from vortex-bench/tpch/q4.sql rename to vortex-bench/sql/tpch/q4.sql diff --git a/vortex-bench/tpch/q5.sql b/vortex-bench/sql/tpch/q5.sql similarity index 100% rename from vortex-bench/tpch/q5.sql rename to vortex-bench/sql/tpch/q5.sql diff --git a/vortex-bench/tpch/q6.sql b/vortex-bench/sql/tpch/q6.sql similarity index 100% rename from vortex-bench/tpch/q6.sql rename to vortex-bench/sql/tpch/q6.sql diff --git a/vortex-bench/tpch/q7.sql b/vortex-bench/sql/tpch/q7.sql similarity index 100% rename from vortex-bench/tpch/q7.sql rename to vortex-bench/sql/tpch/q7.sql diff --git a/vortex-bench/tpch/q8.sql b/vortex-bench/sql/tpch/q8.sql similarity index 100% rename from vortex-bench/tpch/q8.sql rename to vortex-bench/sql/tpch/q8.sql diff --git a/vortex-bench/tpch/q9.sql b/vortex-bench/sql/tpch/q9.sql similarity index 100% rename from vortex-bench/tpch/q9.sql rename to vortex-bench/sql/tpch/q9.sql diff --git a/vortex-bench/sql/vortex/0_sum-with-filter.sql b/vortex-bench/sql/vortex/0_sum-with-filter.sql new file mode 100644 index 00000000000..355f9569a92 --- /dev/null +++ b/vortex-bench/sql/vortex/0_sum-with-filter.sql @@ -0,0 +1,4 @@ +-- As this aggregation has a filter, Vortex has to use a linear scan. Once stats +-- are propagated to arrays, this should use zone maps to aggregate instead of +-- decoding and processing each row. +SELECT sum(col) FROM test WHERE col2 > 0 AND col2 < 1000; diff --git a/vortex-bench/sql/vortex/1_sum.sql b/vortex-bench/sql/vortex/1_sum.sql new file mode 100644 index 00000000000..7640c0d3394 --- /dev/null +++ b/vortex-bench/sql/vortex/1_sum.sql @@ -0,0 +1,3 @@ +-- When Footer changes land, vortex-duckdb should populate statistics from +-- Footer without loading and decoding the data. +SELECT sum(col) FROM test; diff --git a/vortex-bench/sql/vortex/init.sql b/vortex-bench/sql/vortex/init.sql new file mode 100644 index 00000000000..90002573995 --- /dev/null +++ b/vortex-bench/sql/vortex/init.sql @@ -0,0 +1,8 @@ +-- Script that prepares Parquet data for our SQL microbenchmarks. + +COPY ( + SELECT + i % 1000 AS col, + (i * 2654435761) % 100000 AS col2 + FROM range(25000000) t(i) +) TO 'test.parquet' (FORMAT parquet); diff --git a/vortex-bench/src/appian/mod.rs b/vortex-bench/src/appian/mod.rs index 3f22bf9c07f..cdc82ac200a 100644 --- a/vortex-bench/src/appian/mod.rs +++ b/vortex-bench/src/appian/mod.rs @@ -74,14 +74,15 @@ const TABLES: &[&str] = &[ ]; /// Eight join-heavy queries from `duckdb/duckdb:benchmark/appian_benchmarks/queries/`, -/// stored byte-identically under `vortex-bench/appian/q{1..8}.sql` (sibling of the TPC-H -/// `tpch/q*.sql` layout). Upstream refreshes are a pure copy into that directory. +/// stored byte-identically under `vortex-bench/sql/appian/q{1..8}.sql` (sibling of the TPC-H +/// `sql/tpch/q*.sql` layout). Upstream refreshes are a pure copy into that directory. pub fn appian_queries() -> impl Iterator { (1..=8).map(|q| (q, appian_query(q))) } fn appian_query(query_idx: usize) -> String { let path = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("sql") .join("appian") .join(format!("q{query_idx}")) .with_extension("sql"); diff --git a/vortex-bench/src/benchmark.rs b/vortex-bench/src/benchmark.rs index fe16df5cd3d..d59708886a6 100644 --- a/vortex-bench/src/benchmark.rs +++ b/vortex-bench/src/benchmark.rs @@ -3,6 +3,8 @@ //! Core benchmark trait and types. +use std::path::Path; + use arrow_schema::Schema; use glob::Pattern; use url::Url; @@ -47,6 +49,13 @@ pub trait Benchmark: Send + Sync { /// call this method to ensure base data exists, then perform their own format conversion. async fn generate_base_data(&self) -> anyhow::Result<()>; + /// Prepare benchmark- and format-specific data beyond the Parquet base that + /// [`Benchmark::generate_base_data`] produced. Called once per requested format, after the base + /// data exists. Default: nothing. + async fn prepare_format(&self, _format: Format, _base_path: &Path) -> anyhow::Result<()> { + Ok(()) + } + /// Get expected row counts for validation (optional) /// If None, no validation will be performed fn expected_row_counts(&self) -> Option> { diff --git a/vortex-bench/src/clickbench/benchmark.rs b/vortex-bench/src/clickbench/benchmark.rs index c036d1c084b..346496dbda0 100644 --- a/vortex-bench/src/clickbench/benchmark.rs +++ b/vortex-bench/src/clickbench/benchmark.rs @@ -60,7 +60,9 @@ impl ClickBenchSortedBenchmark { 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"), + None => Path::new(env!("CARGO_MANIFEST_DIR")) + .join("sql") + .join("clickbench_queries.sql"), }; Ok(fs::read_to_string(queries_filepath)? diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index ad84ef61b1d..92064c25785 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -4,6 +4,7 @@ use std::fs; use std::path::Path; use std::path::PathBuf; +use std::sync::Arc; use futures::StreamExt; use futures::TryStreamExt; @@ -22,29 +23,43 @@ use tracing::info; use tracing::trace; use vortex::VortexSessionDefault; use vortex::array::ArrayRef; +use vortex::array::ExecutionCtx; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ChunkedArray; use vortex::array::arrays::ExtensionArray; use vortex::array::arrays::Struct; use vortex::array::arrays::StructArray; +use vortex::array::arrays::VarBinViewArray; 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::compressor::BtrBlocksCompressorBuilder; use vortex::dtype::DType; +use vortex::dtype::FieldPath; 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::VortexWriteOptions; use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; +use vortex::layout::LayoutStrategy; +use vortex::layout::layouts::chunked::writer::ChunkedLayoutStrategy; +use vortex::layout::layouts::compressed::CompressingStrategy; +use vortex::layout::layouts::flat::writer::FlatLayoutStrategy; use vortex::session::VortexSession; use vortex::utils::aliases::hash_set::HashSet; +use vortex_arrow::FromArrowArray; +use vortex_arrow::FromArrowType; use vortex_geo::extension::GeoMetadata; use vortex_geo::extension::WellKnownBinary; +use wkb::Endianness; +use wkb::reader::read_wkb; +use wkb::writer::WriteOptions; +use wkb::writer::write_geometry; use crate::CompactionStrategy; use crate::Format; @@ -142,8 +157,7 @@ pub async fn convert_parquet_file_to_vortex( .open(output_path) .await?; - compaction - .apply_options(SESSION.write_options()) + write_options_for(compaction, &dtype, is_spatialbench(parquet_path)) .write( &mut output_file, ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype, stream)), @@ -153,6 +167,54 @@ pub async fn convert_parquet_file_to_vortex( Ok(()) } +/// Whether `path` points at SpatialBench data. +fn is_spatialbench(path: &Path) -> bool { + path.components() + .any(|component| component.as_os_str() == "spatialbench") +} + +/// Vortex write options for converting `dtype`-shaped data. +/// +/// For SpatialBench (`skip_binary_dict`), the geometry blobs are large and +/// unique, so the dictionary builder balloons memory (tens of GB) for zero gain. +fn write_options_for( + compaction: CompactionStrategy, + dtype: &DType, + skip_binary_dict: bool, +) -> VortexWriteOptions { + let binary_fields: Vec<_> = match dtype { + DType::Struct(fields, _) if skip_binary_dict => fields + .names() + .iter() + .zip(fields.fields()) + .filter(|(_, field)| matches!(field, DType::Binary(_))) + .map(|(name, _)| name.clone()) + .collect(), + _ => Vec::new(), + }; + if binary_fields.is_empty() { + return compaction.apply_options(SESSION.write_options()); + } + + let mut builder = WriteStrategyBuilder::default(); + if matches!(compaction, CompactionStrategy::Compact) { + builder = + builder.with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()); + } + for name in binary_fields { + builder = builder.with_field_writer(FieldPath::from_name(name), no_dict_layout()); + } + SESSION.write_options().with_strategy(builder.build()) +} + +/// A chunked + compressed layout that skips dictionary encoding for opaque `Binary` blobs. +fn no_dict_layout() -> Arc { + Arc::new(CompressingStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + BtrBlocksCompressorBuilder::default().build(), + )) +} + /// Convert all Parquet files in a directory to Vortex format. /// /// This function reads Parquet files from `{input_path}/parquet/` and writes Vortex files to @@ -251,7 +313,7 @@ pub async fn add_geoparquet_metadata(parquet_path: &Path, geo_json: &str) -> any return Ok(()); } - let schema = std::sync::Arc::clone(builder.schema()); + let schema = Arc::clone(builder.schema()); let mut reader = builder.build()?; let tmp_path = parquet_path.with_extension("parquet.tmp"); @@ -314,7 +376,8 @@ fn tag_geo_dtype(dtype: DType, geo_columns: &HashSet) -> VortexResult

) -> VortexResult { if geo_columns.is_empty() { return Ok(chunk); @@ -325,17 +388,51 @@ fn tag_geo_array(chunk: ArrayRef, geo_columns: &HashSet) -> VortexResult 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()) + let mut ctx = SESSION.create_execution_ctx(); + let mut tagged = Vec::with_capacity(names.len()); + for (name, field) in names.iter().zip(struct_array.iter_unmasked_fields()) { + if geo_columns.contains(name.as_ref()) && field.dtype().is_binary() { + let ext = wkb_ext_dtype(field.dtype())?; + let little_endian = wkb_field_to_little_endian(field, &mut ctx)?; + tagged.push(ExtensionArray::try_new(ext, little_endian)?.into_array()); + } else { + tagged.push(field.clone()); + } + } + Ok(StructArray::try_new(names, tagged, len, validity)?.into_array()) +} + +/// Re-encode a binary `field`'s WKB values as little-endian (NDR), the only byte order DuckDB's +/// `GEOMETRY` accepts. Byte order is uniform within a column, so the array is returned untouched (no +/// copy) unless its first non-null value is big-endian; only then is each value re-encoded. +fn wkb_field_to_little_endian(field: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let values = field.clone().execute::(ctx)?; + let len = values.len(); + let validity = values.validity()?.execute_mask(len, ctx)?; + let is_big_endian = (0..len) + .find(|&i| validity.value(i)) + .is_some_and(|i| values.bytes_at(i).as_slice().first() == Some(&0x00)); + if !is_big_endian { + return Ok(field.clone()); + } + let little_endian: Vec>> = (0..len) + .map(|i| { + if !validity.value(i) { + return Ok(None); } + let wkb = values.bytes_at(i); + let geometry = read_wkb(wkb.as_slice()).map_err(|e| vortex_err!("invalid WKB: {e}"))?; + let mut encoded = Vec::with_capacity(wkb.len()); + write_geometry( + &mut encoded, + &geometry, + &WriteOptions { + endianness: Endianness::LittleEndian, + }, + ) + .map_err(|e| vortex_err!("re-encoding WKB as little-endian: {e}"))?; + Ok(Some(encoded)) }) - .collect::>>()?; - Ok(StructArray::try_new(names, tagged, len, validity)?.into_array()) + .collect::>()?; + Ok(VarBinViewArray::from_iter_nullable_bin(little_endian).into_array()) } diff --git a/vortex-bench/src/datasets/mod.rs b/vortex-bench/src/datasets/mod.rs index a550a712eba..7eb341edbb7 100644 --- a/vortex-bench/src/datasets/mod.rs +++ b/vortex-bench/src/datasets/mod.rs @@ -81,6 +81,8 @@ pub enum BenchmarkDataset { Fineweb, #[serde(rename = "gharchive")] GhArchive, + #[serde(rename = "vortex")] + VortexQueries, } impl BenchmarkDataset { @@ -97,6 +99,7 @@ impl BenchmarkDataset { BenchmarkDataset::PolarSignals { .. } => "polarsignals", BenchmarkDataset::Fineweb => "fineweb", BenchmarkDataset::GhArchive => "gharchive", + BenchmarkDataset::VortexQueries => "vortex", } } } @@ -122,6 +125,7 @@ impl Display for BenchmarkDataset { } BenchmarkDataset::Fineweb => write!(f, "fineweb"), BenchmarkDataset::GhArchive => write!(f, "gharchive"), + BenchmarkDataset::VortexQueries => write!(f, "vortex"), } } } @@ -179,6 +183,8 @@ impl BenchmarkDataset { BenchmarkDataset::PolarSignals { .. } => &["stacktraces"], BenchmarkDataset::Fineweb => &["fineweb"], BenchmarkDataset::GhArchive => &["events"], + // See VortexBenchmark::table_specs + BenchmarkDataset::VortexQueries => &[], } } } diff --git a/vortex-bench/src/datasets/struct_list_of_ints.rs b/vortex-bench/src/datasets/struct_list_of_ints.rs index fe9aa991681..164ad4534fc 100644 --- a/vortex-bench/src/datasets/struct_list_of_ints.rs +++ b/vortex-bench/src/datasets/struct_list_of_ints.rs @@ -24,9 +24,10 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::chunked::ChunkedArrayExt; use vortex::array::arrays::listview::recursive_list_from_list_view; -use vortex::array::arrow::ArrowSessionExt; use vortex::array::validity::Validity; use vortex::dtype::FieldNames; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ToArrowType; use crate::IdempotentPath; use crate::SESSION; diff --git a/vortex-bench/src/fineweb/mod.rs b/vortex-bench/src/fineweb/mod.rs index 93151cc6fca..b8c86925905 100644 --- a/vortex-bench/src/fineweb/mod.rs +++ b/vortex-bench/src/fineweb/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::fs; use std::path::PathBuf; use tracing::info; @@ -11,30 +12,11 @@ use crate::BenchmarkDataset; use crate::TableSpec; use crate::datasets::data_downloads::download_data; use crate::utils::file::resolve_data_url; +use crate::workspace_root; /// URL to the sample file const SAMPLE_URL: &str = "https://huggingface.co/datasets/HuggingFaceFW/fineweb/resolve/v1.4.0/sample/10BT/001_00000.parquet"; -/// Some basic string-focused queries. -const QUERIES: &[&str] = &[ - // simple summary - "SELECT count(DISTINCT dump) FROM fineweb", - // selective string equality filter - "SELECT * FROM fineweb WHERE dump = 'CC-MAIN-2016-30'", - // LIKE with prefix filter - "SELECT * FROM fineweb WHERE date LIKE '2020-10-%'", - // LIKE with simple containment filter - "SELECT * FROM fineweb WHERE url LIKE '%google%' AND text LIKE '%Google%'", - // LIKE with larger containment filter - "SELECT * FROM fineweb WHERE url LIKE '%.google.%' OR text LIKE '% Google %'", - "SELECT * FROM fineweb WHERE text LIKE '% vortex %'", - // More LIKE filters - "SELECT * FROM fineweb WHERE url LIKE '%espn%' AND language = 'en' AND language_score > 0.92", - "SELECT * FROM fineweb WHERE url LIKE '%espn%' OR url LIKE '%www.espn.go.com%' OR url LIKE '%espn.go.com%'", - // no results, stats cannot prune but tokenized bloom filters could - "SELECT * FROM fineweb WHERE file_path LIKE '%/CC-MAIN-2014-%'", -]; - /// A benchmark using the HuggingFace FineWeb dataset. /// /// This is a very string-heavy dataset, and exercises dictionary and FSST encoding heavily. @@ -71,8 +53,22 @@ impl FinewebBenchmark { #[async_trait::async_trait] impl Benchmark for FinewebBenchmark { + /// Some basic string-focused queries, numbered from Q0 in `sql/fineweb.sql` file order. fn queries(&self) -> anyhow::Result> { - Ok(QUERIES.iter().map(|s| s.to_string()).enumerate().collect()) + // `;`-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("sql") + .join("fineweb") + .with_extension("sql"); + let contents = fs::read_to_string(queries_file)?; + Ok(contents + .split_terminator(';') + .map(str::trim) + .filter(|stmt| !stmt.is_empty()) + .map(str::to_string) + .enumerate() + .collect()) } async fn generate_base_data(&self) -> anyhow::Result<()> { diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index daef12c7e78..abf4e1ccea3 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -36,6 +36,7 @@ use vortex::file::WriteStrategyBuilder; use vortex::utils::aliases::hash_map::HashMap; use crate::spatialbench::SpatialBenchBenchmark; +use crate::vortex_queries::VortexBenchmark; pub mod appian; pub mod benchmark; @@ -61,6 +62,7 @@ pub mod tpch; pub mod utils; pub mod v3; pub mod vector_dataset; +pub mod vortex_queries; pub use benchmark::Benchmark; pub use benchmark::TableSpec; @@ -76,8 +78,11 @@ use vortex::session::VortexSession; #[global_allocator] static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; -pub static SESSION: LazyLock = - LazyLock::new(|| VortexSession::default().with_tokio()); +pub static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::default().with_tokio(); + vortex_geo::initialize(&session); + session +}); #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] pub struct Target { @@ -146,6 +151,9 @@ pub enum Format { #[clap(name = "vortex-compact")] #[serde(rename = "vortex-compact")] VortexCompact, + #[clap(name = "vortex-geo-native")] + #[serde(rename = "vortex-geo-native")] + VortexNative, #[clap(name = "duckdb")] #[serde(rename = "duckdb")] OnDiskDuckDB, @@ -185,6 +193,7 @@ impl Format { Format::Parquet => "parquet", Format::OnDiskVortex => "vortex-file-compressed", Format::VortexCompact => "vortex-compact", + Format::VortexNative => "vortex-geo-native", Format::OnDiskDuckDB => "duckdb", Format::Lance => "lance", } @@ -197,6 +206,7 @@ impl Format { Format::Parquet => "parquet", Format::OnDiskVortex => "vortex", Format::VortexCompact => "vortex", + Format::VortexNative => "vortex", Format::OnDiskDuckDB => "duckdb", Format::Lance => "lance", } @@ -273,6 +283,8 @@ pub enum BenchmarkArg { PublicBi, #[clap(name = "spatialbench")] SpatialBench, + #[clap(name = "vortex")] + VortexQueries, } /// Default scale factor for TPC-related benchmarks @@ -345,6 +357,13 @@ pub fn create_benchmark(b: BenchmarkArg, opts: &Opts) -> anyhow::Result { + let mut benchmark = VortexBenchmark::new()?; + if let Some(query) = opts.get("query") { + benchmark = benchmark.with_query(query)?; + } + Ok(Box::new(benchmark) as _) + } } } diff --git a/vortex-bench/src/polarsignals/benchmark.rs b/vortex-bench/src/polarsignals/benchmark.rs index 0ef7b4246e6..fce692ceb68 100644 --- a/vortex-bench/src/polarsignals/benchmark.rs +++ b/vortex-bench/src/polarsignals/benchmark.rs @@ -62,6 +62,7 @@ impl Benchmark for PolarSignalsBenchmark { fn queries(&self) -> Result> { let queries_file = workspace_root() .join("vortex-bench") + .join("sql") .join("polarsignals") .with_extension("sql"); let contents = fs::read_to_string(queries_file)?; diff --git a/vortex-bench/src/realnest/gharchive.rs b/vortex-bench/src/realnest/gharchive.rs index 3ec6a67f340..a84395f4a10 100644 --- a/vortex-bench/src/realnest/gharchive.rs +++ b/vortex-bench/src/realnest/gharchive.rs @@ -5,6 +5,7 @@ //! //! This dataset applies a bunch of events this way +use std::fs; use std::path::PathBuf; use std::process::Command; @@ -18,6 +19,7 @@ use crate::TableSpec; use crate::idempotent; use crate::idempotent_async; use crate::utils::file::resolve_data_url; +use crate::workspace_root; /// Template URL for raw JSON dataset fn raw_json_url(hour: usize) -> String { @@ -25,14 +27,6 @@ fn raw_json_url(hour: usize) -> String { format!("https://data.gharchive.org/2024-10-01-{hour}.json.gz") } -const QUERIES: &[&str] = &[ - "select count(*) from events where payload.ref = 'refs/heads/main'", - "select distinct repo.name from events where repo.name like 'spiraldb/%'", - "select distinct org.id as org_id from events order by org_id limit 100", - "select actor.login, count() as freq from events group by actor.login order by freq desc limit 10", - "select actor.avatar_url from events where actor.login = 'renovate[bot]'", -]; - pub struct GithubArchiveBenchmark { data_url: Url, } @@ -72,8 +66,22 @@ impl GithubArchiveBenchmark { #[async_trait::async_trait] impl Benchmark for GithubArchiveBenchmark { + /// GitHub Archive queries, numbered from Q0 in `sql/gharchive.sql` file order. fn queries(&self) -> anyhow::Result> { - Ok(QUERIES.iter().map(|s| s.to_string()).enumerate().collect()) + // `;`-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("sql") + .join("gharchive") + .with_extension("sql"); + let contents = fs::read_to_string(queries_file)?; + Ok(contents + .split_terminator(';') + .map(str::trim) + .filter(|stmt| !stmt.is_empty()) + .map(str::to_string) + .enumerate() + .collect()) } async fn generate_base_data(&self) -> anyhow::Result<()> { diff --git a/vortex-bench/src/spatialbench/benchmark.rs b/vortex-bench/src/spatialbench/benchmark.rs index 1f5ae37fc53..48fa84f8ec1 100644 --- a/vortex-bench/src/spatialbench/benchmark.rs +++ b/vortex-bench/src/spatialbench/benchmark.rs @@ -4,6 +4,7 @@ //! SpatialBench benchmark implementation use std::fs; +use std::path::Path; use url::Url; @@ -17,6 +18,9 @@ use crate::spatialbench::datagen::Table; use crate::utils::file::resolve_data_url; use crate::workspace_root; +/// Data-dir subfolder for the native-geometry Vortex files (the `vortex-geo-native` lane). +pub const NATIVE_DIR: &str = "vortex-geo-native"; + /// 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 . @@ -35,6 +39,21 @@ impl SpatialBenchBenchmark { scale_factor, }) } + + /// Tables to materialize and register: the in-process base tables (`trip`, `building`, + /// `customer`) plus the externally-sourced `zone` when its parquet is present. Shared by native + /// data-gen and table registration so both lanes cover the same set. + fn base_tables(&self) -> Vec

{ + let mut tables = vec![Table::Trip, Table::Building, Table::Customer]; + let zone_present = match self.data_url.to_file_path() { + Ok(base) => zone_parquet_present(&base.join(Format::Parquet.name())), + Err(()) => true, + }; + if zone_present { + tables.push(Table::Zone); + } + tables + } } #[async_trait::async_trait] @@ -44,6 +63,7 @@ impl Benchmark for SpatialBenchBenchmark { // `;`-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("sql") .join("spatialbench") .with_extension("sql"); let contents = fs::read_to_string(queries_file)?; @@ -76,6 +96,29 @@ impl Benchmark for SpatialBenchBenchmark { crate::conversions::add_geoparquet_metadata(&zone_file, &geo).await?; } } + + // Cluster the source parquet along a spatial curve so all lanes read the same layout and + // the geometry zone-map prune can skip chunks. + let derived_dirs = [ + base_data_dir.join(Format::OnDiskVortex.name()), + base_data_dir.join(Format::VortexCompact.name()), + base_data_dir.join(NATIVE_DIR), + ]; + datagen::spatially_sort_tables(&base_data_dir.join(Format::Parquet.name()), &derived_dirs) + .await?; + Ok(()) + } + + /// The `vortex-geo-native` lane decodes each table's WKB geometry to native GeoArrow once, into the + /// `vortex-geo-native` dir, so its queries read DuckDB `GEOMETRY` directly. Idempotent. + async fn prepare_format(&self, format: Format, base_path: &Path) -> anyhow::Result<()> { + if format == Format::VortexNative { + let parquet_dir = base_path.join(Format::Parquet.name()); + let native_dir = base_path.join(NATIVE_DIR); + for table in self.base_tables() { + datagen::write_native_vortex(table, &parquet_dir, &native_dir).await?; + } + } Ok(()) } @@ -83,6 +126,16 @@ impl Benchmark for SpatialBenchBenchmark { &self.data_url } + /// The `vortex-geo-native` lane reads the native-geometry Vortex dir; every other format reads its + /// own `{format}` subfolder. + fn format_path(&self, format: Format, base_url: &Url) -> anyhow::Result { + let dir = match format { + Format::VortexNative => NATIVE_DIR, + other => other.name(), + }; + Ok(base_url.join(&format!("{dir}/"))?) + } + 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 @@ -110,16 +163,11 @@ impl Benchmark for SpatialBenchBenchmark { format!("spatialbench(sf={})", self.scale_factor) } + /// Both lanes register the same tables (WKB reads `parquet`/`vortex`, native reads + /// `vortex-geo-native`); `zone` is externally sourced and optional, registered only when present. 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) + self.base_tables() + .iter() .map(|table| TableSpec::new(table.name(), None)) .collect() } @@ -146,7 +194,7 @@ impl Benchmark for SpatialBenchBenchmark { /// 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 { +fn zone_parquet_present(parquet_dir: &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 index 8ebd1b35d86..4c0d4bd2f81 100644 --- a/vortex-bench/src/spatialbench/datagen/mod.rs +++ b/vortex-bench/src/spatialbench/datagen/mod.rs @@ -1,11 +1,16 @@ // 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. +//! SpatialBench data preparation. [`wkb`] generates the canonical WKB base tables; [`native`] derives +//! native-geometry Vortex files from them for `points=native`. The [`table`] catalog is the single +//! source of truth for the base tables both stages share. +pub mod native; +pub mod spatial_sort; pub mod table; pub mod wkb; +pub use native::write_native_vortex; +pub use spatial_sort::spatially_sort_tables; pub use table::Table; pub use wkb::generate_tables; diff --git a/vortex-bench/src/spatialbench/datagen/native.rs b/vortex-bench/src/spatialbench/datagen/native.rs new file mode 100644 index 00000000000..2e03b61b39f --- /dev/null +++ b/vortex-bench/src/spatialbench/datagen/native.rs @@ -0,0 +1,165 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native-geometry prep for `points=native`: decode a table's WKB geometry to native +//! `vortex.geo.{point,polygon,multipolygon}` via `geoarrow_cast` (so Vortex never decodes WKB), then +//! write a Vortex file. A one-time cost; queries then see DuckDB `GEOMETRY` directly. + +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; + +use anyhow::Context; +use arrow_array::RecordBatch; +use arrow_schema::DataType; +use arrow_schema::Schema; +use futures::TryStreamExt; +use geoarrow::array::GenericWkbArray; +use geoarrow::array::GeoArrowArray; +use geoarrow::array::WkbViewArray; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Dimension; +use geoarrow::datatypes::GeoArrowType; +use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::MultiPolygonType; +use geoarrow::datatypes::PointType; +use geoarrow::datatypes::PolygonType; +use geoarrow::datatypes::WkbType; +use geoarrow_cast::cast::cast; +use parquet::arrow::ParquetRecordBatchStreamBuilder; +use tokio::fs::File as TokioFile; +use vortex::array::ArrayRef; +use vortex::array::IntoArray; +use vortex::array::arrays::ChunkedArray; +use vortex::file::WriteOptionsSessionExt; +use vortex_arrow::ArrowSessionExt; + +use super::table::GeometryKind; +use super::table::Table; +use crate::SESSION; +use crate::utils::file::idempotent_async; + +fn geo_metadata() -> Arc { + Arc::new(Metadata::new(Crs::default(), None)) +} + +/// Write `{native_dir}/{table}_0.vortex` with native geometry columns from the WKB parquet. Idempotent. +pub async fn write_native_vortex( + table: Table, + parquet_dir: &Path, + native_dir: &Path, +) -> anyhow::Result { + idempotent_async( + native_dir.join(format!("{}_0.vortex", table.name())), + |path| async move { + let chunks = map_source_batches(parquet_dir, table, |b| native_chunk(b, table)).await?; + + let dtype = chunks[0].dtype().clone(); + let chunked = ChunkedArray::try_new(chunks, dtype)?.into_array(); + let mut file = TokioFile::create(&path).await?; + SESSION + .write_options() + .write(&mut file, chunked.to_array_stream()) + .await?; + tracing::info!(path = %path.display(), table = table.name(), "wrote native geometry table"); + Ok(()) + }, + ) + .await +} + +/// Apply `f` to every batch of `table`'s base WKB parquet parts. All columns are kept; only the +/// geometry columns are rewritten to native types. +async fn map_source_batches( + parquet_dir: &Path, + table: Table, + mut f: impl FnMut(RecordBatch) -> anyhow::Result, +) -> anyhow::Result> { + let pattern = parquet_dir.join(format!("{}_*.parquet", table.name())); + let mut files: Vec = + glob::glob(&pattern.to_string_lossy())?.collect::>()?; + files.sort(); + anyhow::ensure!(!files.is_empty(), "no parquet matching {pattern:?}"); + + let mut out = Vec::new(); + for file in files { + let builder = ParquetRecordBatchStreamBuilder::new(TokioFile::open(&file).await?).await?; + let mut stream = builder.build()?; + while let Some(batch) = stream.try_next().await? { + out.push(f(batch)?); + } + } + Ok(out) +} + +/// Rewrite each of `table`'s WKB geometry columns to its native-lane type, tagging the field with +/// the matching `geoarrow.*` extension. +fn native_record_batch(batch: RecordBatch, table: Table) -> anyhow::Result { + let schema = batch.schema(); + let mut fields = schema.fields().to_vec(); + let mut columns = batch.columns().to_vec(); + + for geom in table.geometry_columns() { + let idx = schema.index_of(geom.name)?; + let column = batch.column(idx).as_ref(); + let wkb_type = WkbType::new(geo_metadata()); + + // Wrap the source WKB. SpatialBench tables emit `Binary`; the external `zone` parquet uses + // `BinaryView`. + let wkb: Box = match column.data_type() { + DataType::Binary => Box::new(GenericWkbArray::::try_from((column, wkb_type))?), + DataType::LargeBinary => { + Box::new(GenericWkbArray::::try_from((column, wkb_type))?) + } + DataType::BinaryView => Box::new(WkbViewArray::try_from((column, wkb_type))?), + other => anyhow::bail!("{}: unsupported WKB column type {other}", geom.name), + }; + + // Decode to a native, separated-XY GeoArrow type. The columnar round-trip also normalizes + // WKB endianness (Overture ships big-endian; native types carry none). + let native: Arc = match geom.kind { + GeometryKind::Point => cast( + wkb.as_ref(), + &GeoArrowType::Point( + PointType::new(Dimension::XY, geo_metadata()) + .with_coord_type(CoordType::Separated), + ), + )?, + GeometryKind::Polygon => cast( + wkb.as_ref(), + &GeoArrowType::Polygon( + PolygonType::new(Dimension::XY, geo_metadata()) + .with_coord_type(CoordType::Separated), + ), + )?, + // Polygon promotes to a one-element multipolygon, so this also covers the mixed + // `Polygon`/`MultiPolygon` zone boundaries. + GeometryKind::MultiPolygon => cast( + wkb.as_ref(), + &GeoArrowType::MultiPolygon( + MultiPolygonType::new(Dimension::XY, geo_metadata()) + .with_coord_type(CoordType::Separated), + ), + )?, + }; + + columns[idx] = native.to_array_ref(); + fields[idx] = Arc::new(native.data_type().to_field(geom.name, false)); + } + + Ok(RecordBatch::try_new( + Arc::new(Schema::new(fields)), + columns, + )?) +} + +/// Convert a WKB batch to a Vortex struct chunk with `table`'s geometry columns as native types. +fn native_chunk(batch: RecordBatch, table: Table) -> anyhow::Result { + let native_batch = native_record_batch(batch, table)?; + let native_schema = native_batch.schema(); + SESSION + .arrow() + .from_arrow_record_batch(native_batch, &native_schema) + .context("importing native batch") +} diff --git a/vortex-bench/src/spatialbench/datagen/spatial_sort.rs b/vortex-bench/src/spatialbench/datagen/spatial_sort.rs new file mode 100644 index 00000000000..d77aea15846 --- /dev/null +++ b/vortex-bench/src/spatialbench/datagen/spatial_sort.rs @@ -0,0 +1,255 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Spatial clustering of the source parquet, in place, so every downstream lane (parquet, +//! vortex-WKB, `vortex-geo-native`) reads the same layout. +//! +//! The geometry zone-map prune skips a chunk only when its bounding box is disjoint from the query +//! region. In generation order every chunk's box spans the whole map, so nothing prunes; ordering +//! rows on a Z-order (Morton) curve of each geometry's bounding-box center gives each chunk a +//! compact box instead. The center works uniformly for every geometry type. +//! +//! Every table with a geometry column is sorted by its first one, each part independently — global +//! order is unnecessary for per-chunk pruning. A parquet marker makes it idempotent; derived vortex +//! files from pre-sort parquet are deleted so the existence-keyed conversions regenerate. + +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; + +use anyhow::Context; +use arrow_array::Array; +use arrow_array::ArrayRef; +use arrow_array::RecordBatch; +use arrow_array::UInt64Array; +use arrow_schema::DataType; +use arrow_select::concat::concat_batches; +use arrow_select::take::take; +use futures::TryStreamExt; +use geo::BoundingRect; +use geo::Geometry; +use geo_traits::to_geo::ToGeoGeometry; +use geoarrow::array::GenericWkbArray; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::WkbViewArray; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::WkbType; +use parquet::arrow::AsyncArrowWriter; +use parquet::arrow::ParquetRecordBatchStreamBuilder; +use parquet::basic::Compression; +use parquet::file::metadata::KeyValue; +use parquet::file::properties::WriterProperties; +use tokio::fs::File as TokioFile; +use tracing::info; + +use super::table::Table; +use super::wkb::geo_parquet_metadata; + +/// Parquet metadata marker: this file is already spatially sorted (makes the step idempotent). +const SORTED_KEY: &str = "vortex_spatial_sorted"; + +fn geo_metadata() -> Arc { + Arc::new(Metadata::new(Crs::default(), None)) +} + +/// Spatially sort every table with a geometry column, in place. +pub async fn spatially_sort_tables( + parquet_dir: &Path, + derived_dirs: &[PathBuf], +) -> anyhow::Result<()> { + for table in Table::ALL.into_iter().filter(|table| table.is_generated()) { + sort_table_parquet(table, parquet_dir, derived_dirs).await?; + } + Ok(()) +} + +/// Sort every unsorted `{table}_*.parquet` part by its first geometry column's bounding-box center. +async fn sort_table_parquet( + table: Table, + parquet_dir: &Path, + derived_dirs: &[PathBuf], +) -> anyhow::Result<()> { + let Some(geom) = table.geometry_columns().first() else { + return Ok(()); + }; + + let pattern = parquet_dir.join(format!("{}_*.parquet", table.name())); + let mut files: Vec = + glob::glob(&pattern.to_string_lossy())?.collect::>()?; + files.sort(); + let mut pending = Vec::new(); + for file in files { + if !is_sorted(&file).await? { + pending.push(file); + } + } + if pending.is_empty() { + return Ok(()); + } + + // Delete before rewriting: an interrupted run leaves unsorted parts unmarked and re-deletes. + for dir in derived_dirs { + let stale = dir.join(format!("{}_*.vortex", table.name())); + for file in glob::glob(&stale.to_string_lossy())?.flatten() { + tokio::fs::remove_file(&file).await?; + info!(path = %file.display(), "removed stale derived file from pre-sort parquet"); + } + } + + for file in pending { + sort_part(&file, table, geom.name).await?; + } + Ok(()) +} + +/// Whether `path` already carries the [`SORTED_KEY`] marker. +async fn is_sorted(path: &Path) -> anyhow::Result { + let builder = ParquetRecordBatchStreamBuilder::new(TokioFile::open(path).await?).await?; + Ok(builder + .metadata() + .file_metadata() + .key_value_metadata() + .is_some_and(|kvs| kvs.iter().any(|kv| kv.key == SORTED_KEY))) +} + +/// Rewrite one parquet part with its rows in Z-order of `geom_col`. +async fn sort_part(path: &Path, table: Table, geom_col: &str) -> anyhow::Result<()> { + let builder = ParquetRecordBatchStreamBuilder::new(TokioFile::open(path).await?).await?; + let schema = Arc::clone(builder.schema()); + let mut reader = builder.build()?; + let mut batches = Vec::new(); + while let Some(batch) = reader.try_next().await? { + batches.push(batch); + } + if batches.iter().all(|batch| batch.num_rows() == 0) { + return Ok(()); + } + let batch = concat_batches(&schema, &batches)?; + drop(batches); + + let geom_idx = schema.index_of(geom_col)?; + let (xs, ys) = wkb_centers(batch.column(geom_idx).as_ref()) + .with_context(|| format!("decoding geometry column {geom_col} of {}", path.display()))?; + let indices = morton_sort_indices(&xs, &ys); + + let columns: Vec = batch + .columns() + .iter() + .map(|column| take(column.as_ref(), &indices, None)) + .collect::>()?; + let sorted = RecordBatch::try_new(Arc::clone(&schema), columns)?; + let num_rows = sorted.num_rows(); + + let tmp_path = path.with_extension("parquet.sorttmp"); + let props = WriterProperties::builder() + .set_compression(Compression::SNAPPY) + .build(); + let mut writer = + AsyncArrowWriter::try_new(TokioFile::create(&tmp_path).await?, schema, Some(props))?; + writer.write(&sorted).await?; + // A fresh write drops metadata: re-tag geo so DuckDB reads `GEOMETRY`, and add the sorted marker. + if let Some(geo) = geo_parquet_metadata(table) { + writer.append_key_value_metadata(KeyValue::new("geo".to_string(), Some(geo))); + } + writer.append_key_value_metadata(KeyValue::new( + SORTED_KEY.to_string(), + Some(format!("morton:{geom_col}")), + )); + writer.close().await?; + tokio::fs::rename(&tmp_path, path).await?; + + info!( + path = %path.display(), + rows = num_rows, + column = geom_col, + "spatially sorted parquet (morton z-order)" + ); + Ok(()) +} + +/// The bounding-box center `(x, y)` of every geometry in a WKB column, whatever its geometry type. +fn wkb_centers(column: &dyn Array) -> anyhow::Result<(Vec, Vec)> { + let wkb_type = WkbType::new(geo_metadata()); + // Expanded per concrete WKB array type. + macro_rules! centers { + ($array:expr) => {{ + let mut xy = Vec::new(); + for item in $array.iter() { + xy.push(match item { + Some(geometry) => bbox_center(&geometry?.to_geometry()), + None => (f64::NAN, f64::NAN), + }); + } + Ok(xy.into_iter().unzip()) + }}; + } + match column.data_type() { + DataType::Binary => centers!(GenericWkbArray::::try_from((column, wkb_type))?), + DataType::LargeBinary => centers!(GenericWkbArray::::try_from((column, wkb_type))?), + DataType::BinaryView => centers!(WkbViewArray::try_from((column, wkb_type))?), + other => anyhow::bail!("unsupported WKB column type {other}"), + } +} + +/// The center of a geometry's bounding box, or `NaN` for an empty geometry. +fn bbox_center(geometry: &Geometry) -> (f64, f64) { + geometry + .bounding_rect() + .map(|r| ((r.min().x + r.max().x) / 2.0, (r.min().y + r.max().y) / 2.0)) + .unwrap_or((f64::NAN, f64::NAN)) +} + +/// Row indices ordering `(x, y)` along a 32-bit-per-axis Z-order curve, with each axis +/// quantized over the data's own extent. +fn morton_sort_indices(xs: &[f64], ys: &[f64]) -> UInt64Array { + let (mut xmin, mut xmax) = (f64::INFINITY, f64::NEG_INFINITY); + let (mut ymin, mut ymax) = (f64::INFINITY, f64::NEG_INFINITY); + for (&x, &y) in xs.iter().zip(ys) { + if !x.is_nan() && !y.is_nan() { + xmin = xmin.min(x); + xmax = xmax.max(x); + ymin = ymin.min(y); + ymax = ymax.max(y); + } + } + + let keys: Vec = xs + .iter() + .zip(ys) + .map(|(&x, &y)| { + if x.is_nan() || y.is_nan() { + 0 + } else { + interleave(quantize(x, xmin, xmax)) | (interleave(quantize(y, ymin, ymax)) << 1) + } + }) + .collect(); + + let mut order: Vec = (0..xs.len()).collect(); + order.sort_unstable_by_key(|&i| keys[i]); + UInt64Array::from_iter_values(order.into_iter().map(|i| i as u64)) +} + +/// Map `value` within `[min, max]` to the full `u32` range. +// `t` is clamped to [0, 1], so the product never exceeds `u32::MAX`. +#[expect(clippy::cast_possible_truncation)] +fn quantize(value: f64, min: f64, max: f64) -> u32 { + let range = max - min; + if range <= 0.0 { + return 0; + } + let t = ((value - min) / range).clamp(0.0, 1.0); + (t * f64::from(u32::MAX)) as u32 +} + +/// Spread a 32-bit value's bits into the even positions of a `u64`. +fn interleave(value: u32) -> u64 { + let mut n = u64::from(value); + n = (n | (n << 16)) & 0x0000_FFFF_0000_FFFF; + n = (n | (n << 8)) & 0x00FF_00FF_00FF_00FF; + n = (n | (n << 4)) & 0x0F0F_0F0F_0F0F_0F0F; + n = (n | (n << 2)) & 0x3333_3333_3333_3333; + n = (n | (n << 1)) & 0x5555_5555_5555_5555; + n +} diff --git a/vortex-bench/src/spatialbench/datagen/table.rs b/vortex-bench/src/spatialbench/datagen/table.rs index c84de192fdf..d45624eec10 100644 --- a/vortex-bench/src/spatialbench/datagen/table.rs +++ b/vortex-bench/src/spatialbench/datagen/table.rs @@ -1,8 +1,8 @@ // 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. +//! The shared SpatialBench table catalog: one source of truth for the base tables, used by both the +//! WKB generation ([`super::wkb`]) and the native geometry conversion ([`super::native`]). /// A SpatialBench table. #[derive(Clone, Copy)] @@ -13,6 +13,20 @@ pub enum Table { Zone, } +/// A geometry column and the geometry type its WKB bytes decode to. +pub(crate) struct GeometryColumn { + pub(crate) name: &'static str, + pub(crate) kind: GeometryKind, +} + +/// Geometry types a column can decode to on the native lane. +#[derive(Clone, Copy, Debug)] +pub(crate) enum GeometryKind { + Point, + Polygon, + MultiPolygon, +} + impl Table { /// Every SpatialBench table, in registration order. pub(crate) const ALL: [Table; 4] = [Table::Trip, Table::Building, Table::Customer, Table::Zone]; @@ -32,13 +46,29 @@ impl Table { !matches!(self, Table::Zone) } - /// Geometry columns SpatialBench tables. - pub(crate) fn geometry_columns(self) -> &'static [&'static str] { + /// Geometry columns to decode from WKB to native, with their geometry type. Empty for tables with + /// no geometry (e.g. `customer`). + pub(crate) fn geometry_columns(self) -> &'static [GeometryColumn] { match self { - Table::Trip => &["t_pickuploc", "t_dropoffloc"], - Table::Building => &["b_boundary"], - Table::Zone => &["z_boundary"], + Table::Trip => &[ + GeometryColumn { + name: "t_pickuploc", + kind: GeometryKind::Point, + }, + GeometryColumn { + name: "t_dropoffloc", + kind: GeometryKind::Point, + }, + ], + Table::Building => &[GeometryColumn { + name: "b_boundary", + kind: GeometryKind::Polygon, + }], Table::Customer => &[], + Table::Zone => &[GeometryColumn { + name: "z_boundary", + kind: GeometryKind::MultiPolygon, + }], } } } diff --git a/vortex-bench/src/spatialbench/datagen/wkb.rs b/vortex-bench/src/spatialbench/datagen/wkb.rs index 422505a3623..51b34bf3016 100644 --- a/vortex-bench/src/spatialbench/datagen/wkb.rs +++ b/vortex-bench/src/spatialbench/datagen/wkb.rs @@ -5,9 +5,11 @@ //! Geometry is emitted as WKB, which DuckDB reads directly as `GEOMETRY` via `ST_GeomFromWKB`. use std::fs; +use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use anyhow::Context; 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. @@ -23,6 +25,7 @@ use spatialbench_parquet::basic::Compression; use spatialbench_parquet::file::properties::WriterProperties; use spatialbench_parquet::format::KeyValue; use tokio::fs::File as TokioFile; +use tokio::process::Command; use tracing::info; use super::table::Table; @@ -109,6 +112,52 @@ pub async fn generate_tables(scale_factor: &str, output_dir: PathBuf) -> Result< } } + // `zone` isn't generated in-process (`Table::is_generated` is false); it comes from the upstream + // `spatialbench-cli` (install it, or set `SPATIALBENCH_CLI` to its binary). + generate_zone(scale_factor, &parquet_dir).await?; + + Ok(()) +} + +/// Generate the externally-sourced `zone` table by shelling out to the upstream +/// `spatialbench-cli`. +async fn generate_zone(scale_factor: f64, parquet_dir: &Path) -> Result<()> { + let cli = std::env::var("SPATIALBENCH_CLI").unwrap_or_else(|_| "spatialbench-cli".to_string()); + idempotent_async(parquet_dir.join("zone_0.parquet"), |zone_path| async move { + info!( + scale_factor, + cli, "Generating SpatialBench zone table via spatialbench-cli" + ); + // The CLI writes a fixed `zone.parquet` name into an output directory, so + // generate into a scratch dir and move the produced parquet onto `zone_path`. + let scratch = zone_path.with_extension("zone-scratch"); + fs::create_dir_all(&scratch)?; + let status = Command::new(&cli) + .arg("-s") + .arg(scale_factor.to_string()) + .args(["-T", "zone", "-f", "parquet", "-o"]) + .arg(&scratch) + .status() + .await + .with_context(|| format!("failed to spawn `{cli}` (is it installed / on PATH?)"))?; + anyhow::ensure!( + status.success(), + "`{cli}` exited with {status} while generating zone" + ); + let produced = glob::glob(&scratch.join("**/*.parquet").to_string_lossy())? + .flatten() + .next() + .with_context(|| { + format!( + "`{cli}` produced no zone parquet under {}", + scratch.display() + ) + })?; + fs::rename(&produced, &zone_path)?; + fs::remove_dir_all(&scratch).ok(); + Ok(()) + }) + .await?; Ok(()) } @@ -118,9 +167,9 @@ pub(crate) fn geo_parquet_metadata(table: Table) -> Option { let primary = geometry_columns.first()?; let columns: serde_json::Map = geometry_columns .iter() - .map(|&column| { + .map(|column| { ( - column.to_string(), + column.name.to_string(), serde_json::json!({ "encoding": "WKB", "geometry_types": [] }), ) }) @@ -128,7 +177,7 @@ pub(crate) fn geo_parquet_metadata(table: Table) -> Option { Some( serde_json::json!({ "version": "1.0.0", - "primary_column": primary, + "primary_column": primary.name, "columns": columns, }) .to_string(), diff --git a/vortex-bench/src/statpopgen/statpopgen_benchmark.rs b/vortex-bench/src/statpopgen/statpopgen_benchmark.rs index 58c4b6aed51..0861f7bfc88 100644 --- a/vortex-bench/src/statpopgen/statpopgen_benchmark.rs +++ b/vortex-bench/src/statpopgen/statpopgen_benchmark.rs @@ -106,6 +106,7 @@ impl Benchmark for StatPopGenBenchmark { fn queries(&self) -> Result> { let queries_file = workspace_root() .join("vortex-bench") + .join("sql") .join("statpopgen") .with_extension("sql"); let contents = fs::read_to_string(queries_file)?; diff --git a/vortex-bench/src/tpcds/mod.rs b/vortex-bench/src/tpcds/mod.rs index 3ece61f5941..b8eda522e48 100644 --- a/vortex-bench/src/tpcds/mod.rs +++ b/vortex-bench/src/tpcds/mod.rs @@ -17,6 +17,7 @@ pub fn tpcds_queries() -> impl Iterator { // A few tpcds queries have multiple statements, this handles that fn tpcds_query(query_idx: usize) -> String { let manifest_dir = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("sql") .join("tpcds") .join(format!("{query_idx:02}")) .with_extension("sql"); diff --git a/vortex-bench/src/tpch/mod.rs b/vortex-bench/src/tpch/mod.rs index 7c36ffb6950..5a2bba97c84 100644 --- a/vortex-bench/src/tpch/mod.rs +++ b/vortex-bench/src/tpch/mod.rs @@ -29,6 +29,7 @@ pub fn tpch_queries() -> impl Iterator { // A few tpch queries have multiple statements, this handles that fn tpch_query(query_idx: usize) -> String { let manifest_dir = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("sql") .join("tpch") .join(format!("q{query_idx}")) .with_extension("sql"); diff --git a/vortex-bench/src/tpch/tpchgen.rs b/vortex-bench/src/tpch/tpchgen.rs index 5dbeb380a7e..79b91200e48 100644 --- a/vortex-bench/src/tpch/tpchgen.rs +++ b/vortex-bench/src/tpch/tpchgen.rs @@ -32,12 +32,12 @@ use tpchgen::generators::SupplierGenerator; use tpchgen_arrow::RecordBatchIterator; use tracing::info; use vortex::array::ArrayRef; -use vortex::array::arrow::FromArrowArray; use vortex::array::stream::ArrayStreamAdapter; use vortex::dtype::DType; -use vortex::dtype::arrow::FromArrowType; use vortex::error::VortexExpect; use vortex::file::WriteOptionsSessionExt; +use vortex_arrow::FromArrowArray; +use vortex_arrow::FromArrowType; use crate::CompactionStrategy; use crate::Format; diff --git a/vortex-bench/src/utils/file.rs b/vortex-bench/src/utils/file.rs index f05113963f7..4af68c029ab 100644 --- a/vortex-bench/src/utils/file.rs +++ b/vortex-bench/src/utils/file.rs @@ -56,8 +56,12 @@ pub trait IdempotentPath { fn to_data_path(&self) -> PathBuf; } +pub fn bench_dir() -> PathBuf { + workspace_root().join("vortex-bench") +} + pub fn data_dir() -> PathBuf { - workspace_root().join("vortex-bench").join("data") + bench_dir().join("data") } /// Find the workspace's root by looking for Cargo's lock file diff --git a/vortex-bench/src/v3.rs b/vortex-bench/src/v3.rs index a7529866f80..4bdd67ed18f 100644 --- a/vortex-bench/src/v3.rs +++ b/vortex-bench/src/v3.rs @@ -296,6 +296,7 @@ fn canonical_tpc_scale_factor(scale_factor: &str) -> String { /// | `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. | +/// | `VortexQueries` | `vortex` | `None` | `None` | Own microbenchmarks | pub fn benchmark_dataset_dims(d: &BenchmarkDataset) -> (String, Option, Option) { match d { BenchmarkDataset::TpcH { scale_factor } => ( @@ -331,6 +332,7 @@ pub fn benchmark_dataset_dims(d: &BenchmarkDataset) -> (String, Option, BenchmarkDataset::Fineweb => ("fineweb".to_string(), None, None), BenchmarkDataset::GhArchive => ("gharchive".to_string(), None, None), BenchmarkDataset::Appian => ("appian".to_string(), None, None), + BenchmarkDataset::VortexQueries => ("vortex".to_string(), None, None), } } diff --git a/vortex-bench/src/vortex_queries.rs b/vortex-bench/src/vortex_queries.rs new file mode 100644 index 00000000000..f611238bb6b --- /dev/null +++ b/vortex-bench/src/vortex_queries.rs @@ -0,0 +1,179 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmark that runs queries which aren't part of any existing benchmark +//! suite but which performance we want to track. + +use std::fs; +use std::path::PathBuf; +use std::process::Command; + +use anyhow::Context; +use anyhow::Result; +use anyhow::anyhow; +use anyhow::bail; +use glob::Pattern; +use tracing::debug; +use url::Url; +use vortex::error::VortexExpect; + +use crate::Benchmark; +use crate::BenchmarkDataset; +use crate::Format; +use crate::TableSpec; +use crate::bench_dir; +use crate::resolve_data_url; + +// Path to script that creates Parquet data +const INIT_SQL: &str = "init.sql"; + +pub struct VortexBenchmark { + data_url: Url, + queries_dir: PathBuf, + query: Option, +} + +impl VortexBenchmark { + pub fn new() -> Result { + let queries_dir = bench_dir().join("sql").join("vortex"); + let data_url = resolve_data_url(None, "vortex")?; + Ok(Self { + data_url, + queries_dir, + query: None, + }) + } + + // Same as new(), but run only for "query" SQL file + pub fn with_query(mut self, query: &str) -> Result { + let as_path = PathBuf::from(query); + let path = if as_path.is_file() { + as_path + } else if query.ends_with(".sql") { + self.queries_dir.join(query) + } else { + self.queries_dir.join(format!("{query}.sql")) + }; + if !path.is_file() { + bail!("{query} file not found in {}", self.queries_dir.display()); + } + self.query = Some(path); + Ok(self) + } + + fn query_files(&self) -> Result> { + if let Some(query) = &self.query { + return Ok(vec![query.clone()]); + } + + let entries = fs::read_dir(&self.queries_dir) + .with_context(|| format!("cannot list queries in {}", self.queries_dir.display()))? + .collect::>>()?; + + let mut files: Vec = entries + .into_iter() + .map(|entry| entry.path()) + .filter(|path| { + path.extension().is_some_and(|ext| ext == "sql") + && path.file_name().is_some_and(|name| name != INIT_SQL) + }) + .collect(); + files.sort(); + + if files.is_empty() { + bail!("no query files found in {}", self.queries_dir.display()); + } + Ok(files) + } +} + +#[async_trait::async_trait] +impl Benchmark for VortexBenchmark { + fn queries(&self) -> Result> { + self.query_files()? + .iter() + .map(|path| { + let idx = path + .file_name() + .vortex_expect("no file name") + .to_str() + .vortex_expect("not utf-8") + .split_once("_") + .vortex_expect("query without a number") + .0 + .parse::()?; + let query = fs::read_to_string(path) + .with_context(|| format!("cannot read query {}", path.display()))?; + debug!(idx, file = %path.display(), "Loaded vortex query"); + Ok((idx, query)) + }) + .collect() + } + + async fn generate_base_data(&self) -> Result<()> { + let data_dir = self + .data_url + .to_file_path() + .map_err(|_| anyhow!("Invalid file URL: {}", self.data_url.as_str()))?; + let parquet_dir = data_dir.join(Format::Parquet.name()); + fs::create_dir_all(&parquet_dir)?; + let parquet_file = parquet_dir.join("test.parquet"); + + if parquet_file.exists() { + debug!("Parquet data present in {}", parquet_dir.display()); + return Ok(()); + } + + let init_path = self.queries_dir.join(INIT_SQL); + let script = fs::read_to_string(&init_path) + .with_context(|| format!("cannot read {}", init_path.display()))?; + + let output = Command::new("duckdb") + .current_dir(&parquet_dir) + .arg("-c") + .arg(&script) + .output() + .context("cannot run duckdb")?; + if !output.status.success() { + bail!( + "duckdb {INIT_SQL} failed: stdout={:?} stderr={:?}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr), + ); + } + + if !parquet_file.exists() { + bail!("{INIT_SQL} did not create Parquet files"); + } + + debug!("Parquet data generated in {}", parquet_dir.display()); + Ok(()) + } + + fn dataset(&self) -> BenchmarkDataset { + BenchmarkDataset::VortexQueries + } + + fn dataset_name(&self) -> &str { + "vortex" + } + + fn dataset_display(&self) -> String { + "vortex".to_owned() + } + + fn data_url(&self) -> &Url { + &self.data_url + } + + fn table_specs(&self) -> Vec { + vec![TableSpec::new("test", None)] + } + + fn pattern(&self, table_name: &str, format: Format) -> Option { + Some( + Pattern::new(&format!("{table_name}.{}", format.ext())) + .expect("table name is a valid identifier"), + ) + } +} diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 5ee2bad5402..255d2fbd075 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -18,7 +18,6 @@ all-features = true [dependencies] itertools = { workspace = true } -num-traits = { workspace = true } pco = { workspace = true, optional = true } rand = { workspace = true } vortex-alp = { workspace = true } @@ -41,6 +40,7 @@ vortex-zstd = { workspace = true, optional = true } [dev-dependencies] divan = { workspace = true } +insta = { workspace = true } rstest = { workspace = true } test-with = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } diff --git a/vortex-btrblocks/REUSE.toml b/vortex-btrblocks/REUSE.toml new file mode 100644 index 00000000000..6bee2b6904f --- /dev/null +++ b/vortex-btrblocks/REUSE.toml @@ -0,0 +1,7 @@ +version = 1 + +# `insta` snapshot files do not allow leading comment lines. +[[annotations]] +path = "tests/snapshots/**.snap" +SPDX-FileCopyrightText = "Copyright the Vortex contributors" +SPDX-License-Identifier = "Apache-2.0" diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 6489f2d4e52..e51fac9a9d5 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -11,7 +11,6 @@ use crate::Scheme; use crate::SchemeExt; use crate::SchemeId; use crate::schemes::binary; -use crate::schemes::bool; use crate::schemes::decimal; use crate::schemes::float; use crate::schemes::integer; @@ -23,14 +22,9 @@ use crate::schemes::temporal; /// This list is order-sensitive: the builder preserves this order when constructing /// the final scheme list, so that tie-breaking is deterministic. pub const ALL_SCHEMES: &[&dyn Scheme] = &[ - //////////////////////////////////////////////////////////////////////////////////////////////// - // Bool schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - &bool::BoolConstantScheme, //////////////////////////////////////////////////////////////////////////////////////////////// // Integer schemes. //////////////////////////////////////////////////////////////////////////////////////////////// - &integer::IntConstantScheme, // NOTE: FoR must precede BitPacking to avoid unnecessary patches. &integer::FoRScheme, // NOTE: ZigZag should precede BitPacking because we don't want negative numbers. @@ -47,7 +41,6 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ //////////////////////////////////////////////////////////////////////////////////////////////// // Float schemes. //////////////////////////////////////////////////////////////////////////////////////////////// - &float::FloatConstantScheme, &float::ALPScheme, &float::ALPRDScheme, &float::FloatDictScheme, @@ -62,13 +55,11 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ &string::FSSTScheme, #[cfg(feature = "unstable_encodings")] &string::OnPairScheme, - &string::StringConstantScheme, &string::NullDominatedSparseScheme, //////////////////////////////////////////////////////////////////////////////////////////////// // Binary schemes. //////////////////////////////////////////////////////////////////////////////////////////////// &binary::BinaryDictScheme, - &binary::BinaryConstantScheme, // Decimal schemes. &decimal::DecimalScheme, // Temporal schemes. @@ -174,10 +165,14 @@ impl BtrBlocksCompressorBuilder { // dictionary expansion at decode time, which is incompatible with // pure-GPU decompression paths. Strip whichever string-fragment // scheme is enabled by feature. - #[cfg_attr(not(feature = "unstable_encodings"), allow(unused_mut))] + #[cfg_attr( + not(any(feature = "pco", feature = "unstable_encodings")), + allow(unused_mut) + )] let mut excluded: Vec = vec![ integer::SparseScheme.id(), integer::IntRLEScheme.id(), + float::ALPRDScheme.id(), float::FloatRLEScheme.id(), float::NullDominatedSparseScheme.id(), string::StringDictScheme.id(), @@ -190,6 +185,8 @@ impl BtrBlocksCompressorBuilder { // is incompatible with pure-GPU decompression paths. #[cfg(feature = "unstable_encodings")] excluded.push(integer::DeltaScheme::default().id()); + #[cfg(feature = "pco")] + excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]); let builder = self.exclude_schemes(excluded); #[cfg(all(feature = "zstd", feature = "unstable_encodings"))] @@ -232,4 +229,27 @@ mod tests { let builder = BtrBlocksCompressorBuilder::default(); assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); } + + #[test] + fn cuda_compatible_excludes_alprd() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + assert!( + !builder + .schemes + .iter() + .any(|s| s.id() == float::ALPRDScheme.id()) + ); + } + + #[test] + #[cfg(feature = "pco")] + fn cuda_compatible_excludes_pco() { + let builder = BtrBlocksCompressorBuilder::default() + .with_new_scheme(&integer::PcoScheme) + .with_new_scheme(&float::PcoScheme) + .only_cuda_compatible(); + for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] { + assert!(!builder.schemes.iter().any(|s| s.id() == scheme)); + } + } } diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 69ddb59738c..37915b427b6 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -78,8 +78,8 @@ pub use builder::BtrBlocksCompressorBuilder; pub use canonical_compressor::BtrBlocksCompressor; pub use schemes::patches::compress_patches; pub use vortex_compressor::CascadingCompressor; -pub use vortex_compressor::ctx::CompressorContext; -pub use vortex_compressor::ctx::MAX_CASCADE; +pub use vortex_compressor::scheme::CompressorContext; +pub use vortex_compressor::scheme::MAX_CASCADE; pub use vortex_compressor::scheme::Scheme; pub use vortex_compressor::scheme::SchemeExt; pub use vortex_compressor::scheme::SchemeId; diff --git a/vortex-btrblocks/src/schemes/binary/mod.rs b/vortex-btrblocks/src/schemes/binary/mod.rs index 44f09f2daa3..af66d345cc6 100644 --- a/vortex-btrblocks/src/schemes/binary/mod.rs +++ b/vortex-btrblocks/src/schemes/binary/mod.rs @@ -9,7 +9,6 @@ mod zstd; mod zstd_buffers; // Re-export builtin schemes from vortex-compressor. -pub use vortex_compressor::builtins::BinaryConstantScheme; pub use vortex_compressor::builtins::BinaryDictScheme; #[cfg(feature = "zstd")] pub use zstd::ZstdScheme; diff --git a/vortex-btrblocks/src/schemes/binary/zstd.rs b/vortex-btrblocks/src/schemes/binary/zstd.rs index 56be76cdcb8..83c09d24e94 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd.rs @@ -7,8 +7,8 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -45,7 +45,10 @@ impl Scheme for ZstdScheme { _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let compacted = data.array_as_varbinview().into_owned().compact_buffers()?; + let compacted = data + .array_as_varbinview() + .into_owned() + .compact_buffers(exec_ctx)?; Ok( vortex_zstd::Zstd::from_var_bin_view_without_dict(&compacted, 3, 8192, exec_ctx)? .into_array(), diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index 269665866a2..a84672ef860 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -7,8 +7,8 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; use crate::ArrayAndStats; diff --git a/vortex-btrblocks/src/schemes/bool.rs b/vortex-btrblocks/src/schemes/bool.rs deleted file mode 100644 index c27251a8599..00000000000 --- a/vortex-btrblocks/src/schemes/bool.rs +++ /dev/null @@ -1,7 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Bool compression schemes. - -pub use vortex_compressor::builtins::BoolConstantScheme; -pub use vortex_compressor::stats::BoolStats; diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index d60e5c6b175..20c5cd80a3c 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -11,8 +11,8 @@ use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::decimal::narrowed_decimal; use vortex_array::dtype::DecimalType; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::EstimateVerdict; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::EstimateVerdict; use vortex_decimal_byte_parts::DecimalByteParts; use vortex_error::VortexResult; diff --git a/vortex-btrblocks/src/schemes/float/alp.rs b/vortex-btrblocks/src/schemes/float/alp.rs index ba8d97f676c..b8a26abf348 100644 --- a/vortex-btrblocks/src/schemes/float/alp.rs +++ b/vortex-btrblocks/src/schemes/float/alp.rs @@ -15,9 +15,9 @@ use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::PType; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; -use vortex_compressor::estimate::EstimateVerdict; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use crate::ArrayAndStats; diff --git a/vortex-btrblocks/src/schemes/float/alprd.rs b/vortex-btrblocks/src/schemes/float/alprd.rs index 969afabad49..662a47a1d9c 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/vortex-btrblocks/src/schemes/float/alprd.rs @@ -12,9 +12,9 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::PType; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; -use vortex_compressor::estimate::EstimateVerdict; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_error::vortex_panic; diff --git a/vortex-btrblocks/src/schemes/float/mod.rs b/vortex-btrblocks/src/schemes/float/mod.rs index 3713fa2eb63..1301184ac0c 100644 --- a/vortex-btrblocks/src/schemes/float/mod.rs +++ b/vortex-btrblocks/src/schemes/float/mod.rs @@ -18,7 +18,6 @@ pub use pco::PcoScheme; pub use rle::FloatRLEScheme; pub use sparse::NullDominatedSparseScheme; // Re-export builtin schemes from vortex-compressor. -pub use vortex_compressor::builtins::FloatConstantScheme; pub use vortex_compressor::builtins::FloatDictScheme; pub use vortex_compressor::stats::FloatStats; diff --git a/vortex-btrblocks/src/schemes/float/pco.rs b/vortex-btrblocks/src/schemes/float/pco.rs index 34147f3dfdc..dc9a96133d5 100644 --- a/vortex-btrblocks/src/schemes/float/pco.rs +++ b/vortex-btrblocks/src/schemes/float/pco.rs @@ -7,8 +7,8 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; use crate::ArrayAndStats; diff --git a/vortex-btrblocks/src/schemes/float/rle.rs b/vortex-btrblocks/src/schemes/float/rle.rs index 54406fcb9fd..c2683ed131c 100644 --- a/vortex-btrblocks/src/schemes/float/rle.rs +++ b/vortex-btrblocks/src/schemes/float/rle.rs @@ -6,11 +6,11 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; -use vortex_compressor::estimate::EstimateVerdict; use vortex_compressor::scheme::AncestorExclusion; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use crate::ArrayAndStats; diff --git a/vortex-btrblocks/src/schemes/float/sparse.rs b/vortex-btrblocks/src/schemes/float/sparse.rs index 8f41143c971..42d09e323c5 100644 --- a/vortex-btrblocks/src/schemes/float/sparse.rs +++ b/vortex-btrblocks/src/schemes/float/sparse.rs @@ -9,10 +9,10 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::EstimateVerdict; use vortex_compressor::scheme::ChildSelection; +use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_sparse::Sparse; use vortex_sparse::SparseExt as _; diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index f742f8deac1..2424402649e 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -10,9 +10,9 @@ use vortex_array::IntoArray; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; -use vortex_compressor::estimate::EstimateVerdict; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_fastlanes::BitPacked; use vortex_fastlanes::bitpack_compress::bit_width_histogram; diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index f25bc6ce466..ed087f5d845 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -12,13 +12,13 @@ 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::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; +use vortex_compressor::scheme::EstimateScore; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_fastlanes::Delta; diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/vortex-btrblocks/src/schemes/integer/for_.rs index ec23f1114db..832ed4819fd 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -12,10 +12,10 @@ 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::EstimateVerdict; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_fastlanes::FoR; diff --git a/vortex-btrblocks/src/schemes/integer/mod.rs b/vortex-btrblocks/src/schemes/integer/mod.rs index abe5868f5c8..3aae2ae5601 100644 --- a/vortex-btrblocks/src/schemes/integer/mod.rs +++ b/vortex-btrblocks/src/schemes/integer/mod.rs @@ -30,7 +30,6 @@ pub use runend::RunEndScheme; pub use sequence::SequenceScheme; pub use sparse::SparseScheme; // Re-export builtin schemes from vortex-compressor. -pub use vortex_compressor::builtins::IntConstantScheme; pub use vortex_compressor::builtins::IntDictScheme; pub use vortex_compressor::stats::IntegerStats; pub use zigzag::ZigZagScheme; diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/vortex-btrblocks/src/schemes/integer/pco.rs index 63e5917a996..d7f182f588e 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/vortex-btrblocks/src/schemes/integer/pco.rs @@ -7,9 +7,9 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; -use vortex_compressor::estimate::EstimateVerdict; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use crate::ArrayAndStats; diff --git a/vortex-btrblocks/src/schemes/integer/rle.rs b/vortex-btrblocks/src/schemes/integer/rle.rs index df0f0183780..5bc4e5296f1 100644 --- a/vortex-btrblocks/src/schemes/integer/rle.rs +++ b/vortex-btrblocks/src/schemes/integer/rle.rs @@ -9,11 +9,11 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; -use vortex_compressor::estimate::EstimateVerdict; use vortex_compressor::scheme::AncestorExclusion; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; +use vortex_compressor::scheme::EstimateVerdict; #[cfg(feature = "unstable_encodings")] use vortex_compressor::scheme::SchemeId; use vortex_error::VortexResult; diff --git a/vortex-btrblocks/src/schemes/integer/runend.rs b/vortex-btrblocks/src/schemes/integer/runend.rs index b68567d789c..175f33e7406 100644 --- a/vortex-btrblocks/src/schemes/integer/runend.rs +++ b/vortex-btrblocks/src/schemes/integer/runend.rs @@ -12,12 +12,12 @@ 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::EstimateVerdict; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_runend::RunEnd; use vortex_runend::compress::runend_encode; diff --git a/vortex-btrblocks/src/schemes/integer/sequence.rs b/vortex-btrblocks/src/schemes/integer/sequence.rs index 604d7567d00..fda9995b510 100644 --- a/vortex-btrblocks/src/schemes/integer/sequence.rs +++ b/vortex-btrblocks/src/schemes/integer/sequence.rs @@ -10,12 +10,12 @@ 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::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::EstimateScore; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; diff --git a/vortex-btrblocks/src/schemes/integer/sparse.rs b/vortex-btrblocks/src/schemes/integer/sparse.rs index 609a258d495..e843d9d7999 100644 --- a/vortex-btrblocks/src/schemes/integer/sparse.rs +++ b/vortex-btrblocks/src/schemes/integer/sparse.rs @@ -12,10 +12,10 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::scalar::Scalar; use vortex_compressor::builtins::IntDictScheme; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::EstimateVerdict; use vortex_compressor::scheme::ChildSelection; +use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_sparse::Sparse; diff --git a/vortex-btrblocks/src/schemes/integer/zigzag.rs b/vortex-btrblocks/src/schemes/integer/zigzag.rs index 87fb38e02b8..959889dabac 100644 --- a/vortex-btrblocks/src/schemes/integer/zigzag.rs +++ b/vortex-btrblocks/src/schemes/integer/zigzag.rs @@ -12,12 +12,12 @@ 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::EstimateVerdict; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_zigzag::ZigZag; use vortex_zigzag::ZigZagArrayExt; diff --git a/vortex-btrblocks/src/schemes/mod.rs b/vortex-btrblocks/src/schemes/mod.rs index 16123429e86..a0e9b042a66 100644 --- a/vortex-btrblocks/src/schemes/mod.rs +++ b/vortex-btrblocks/src/schemes/mod.rs @@ -4,7 +4,6 @@ //! Compression scheme implementations. pub mod binary; -pub mod bool; pub mod float; pub mod integer; pub mod string; diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index 5e6baf73933..bd5bd010396 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -11,8 +11,8 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::varbin::VarBinArrayExt; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; use vortex_fsst::FSST; use vortex_fsst::FSSTArrayExt; diff --git a/vortex-btrblocks/src/schemes/string/mod.rs b/vortex-btrblocks/src/schemes/string/mod.rs index acb16ef4323..ac8e5b4b8df 100644 --- a/vortex-btrblocks/src/schemes/string/mod.rs +++ b/vortex-btrblocks/src/schemes/string/mod.rs @@ -19,7 +19,6 @@ pub use fsst::FSSTScheme; pub use onpair::OnPairScheme; pub use sparse::NullDominatedSparseScheme; // Re-export builtin schemes from vortex-compressor. -pub use vortex_compressor::builtins::StringConstantScheme; pub use vortex_compressor::builtins::StringDictScheme; pub use vortex_compressor::stats::StringStats; #[cfg(feature = "zstd")] diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index 7bbe480d059..f0a8966089c 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -9,8 +9,8 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::SchemeId; use vortex_error::VortexResult; use vortex_onpair::DEFAULT_DICT12_CONFIG; diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/vortex-btrblocks/src/schemes/string/sparse.rs index 37f6c4b05cf..750dd978030 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -9,10 +9,10 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::EstimateVerdict; use vortex_compressor::scheme::ChildSelection; +use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; +use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_sparse::Sparse; use vortex_sparse::SparseExt as _; diff --git a/vortex-btrblocks/src/schemes/string/zstd.rs b/vortex-btrblocks/src/schemes/string/zstd.rs index 6e99cc27aa3..177acfa1543 100644 --- a/vortex-btrblocks/src/schemes/string/zstd.rs +++ b/vortex-btrblocks/src/schemes/string/zstd.rs @@ -7,8 +7,8 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -45,7 +45,10 @@ impl Scheme for ZstdScheme { _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let compacted = data.array_as_varbinview().into_owned().compact_buffers()?; + let compacted = data + .array_as_varbinview() + .into_owned() + .compact_buffers(exec_ctx)?; Ok( vortex_zstd::Zstd::from_var_bin_view_without_dict(&compacted, 3, 8192, exec_ctx)? .into_array(), diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index 0456bbc4204..a8d16cc837f 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -7,8 +7,8 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::DeferredEstimate; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; use crate::ArrayAndStats; diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/vortex-btrblocks/src/schemes/temporal.rs index c5c0abc46aa..6989b6a2aba 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -7,8 +7,6 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::aggregate_fn::fns::is_constant::is_constant; -use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; @@ -17,8 +15,8 @@ use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::extension::Matcher; use vortex_array::extension::datetime::AnyTemporal; use vortex_array::extension::datetime::TemporalMetadata; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::EstimateVerdict; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::EstimateVerdict; use vortex_datetime_parts::DateTimeParts; use vortex_datetime_parts::TemporalParts; use vortex_datetime_parts::split_temporal; @@ -79,17 +77,7 @@ impl Scheme for TemporalScheme { ) -> VortexResult { let array = data.array().clone(); let ext_array = array.execute::(exec_ctx)?; - let temporal_array = TemporalArray::try_from(ext_array.clone().into_array())?; - - // Check for constant array and return early if so. - let is_constant = is_constant(&ext_array.clone().into_array(), exec_ctx)?; - - if is_constant { - return Ok( - ConstantArray::new(ext_array.execute_scalar(0, exec_ctx)?, ext_array.len()) - .into_array(), - ); - } + let temporal_array = TemporalArray::try_from(ext_array.into_array())?; let dtype = temporal_array.dtype().clone(); let TemporalParts { diff --git a/vortex-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs new file mode 100644 index 00000000000..c60752b064a --- /dev/null +++ b/vortex-btrblocks/tests/golden.rs @@ -0,0 +1,415 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Golden-corpus determinism tests for the default compressor. +//! +//! Compresses a fixed, seed-generated corpus and snapshots each entry's full encoding tree +//! and exact byte counts. These snapshots pin the default compressor's *decisions*: any +//! refactor of scheme selection (see the compressor cost-model track) must leave every +//! snapshot untouched, so snapshot churn in a later change is the reviewable signal of a +//! behavior change. +//! +//! Three variants cover the feature matrix: +//! +//! - `default`: the default feature set and [`BtrBlocksCompressor::default`]. +//! - `unstable`: `unstable_encodings` enabled, default builder — pins Delta / OnPair +//! selection (compiled out of `ALL_SCHEMES` otherwise). +//! - `compact`: `unstable_encodings` + `zstd` + `pco`, with +//! [`BtrBlocksCompressorBuilder::with_compact`] — pins Zstd / Pco selection. +//! +//! Every corpus entry is longer than 1024 values so the sampling-based estimation path is +//! exercised, and each entry is compressed twice per run to assert determinism directly. + +#![allow(clippy::cast_possible_truncation, clippy::tests_outside_test_module)] + +use std::fmt; +use std::sync::Arc; +use std::sync::LazyLock; + +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::TemporalArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::display::EncodingSummaryExtractor; +use vortex_array::display::MetadataExtractor; +use vortex_array::display::TreeContext; +use vortex_array::display::TreeExtractor; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::Nullability; +use vortex_array::extension::datetime::TimeUnit; +use vortex_array::validity::Validity; +use vortex_btrblocks::BtrBlocksCompressor; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +/// Number of values in each numeric corpus entry: comfortably above the 1024-value sampling +/// threshold so scheme selection runs on sampled estimates, as it does for real file chunks. +const N: usize = 16_384; + +/// Header extractor printing exact byte counts (the built-in [`NbytesExtractor`] rounds +/// through `humansize`, which could mask small size regressions). +/// +/// [`NbytesExtractor`]: vortex_array::display::NbytesExtractor +struct ExactNbytesExtractor; + +impl TreeExtractor for ExactNbytesExtractor { + fn write_header( + &self, + array: &ArrayRef, + _ctx: &TreeContext, + f: &mut fmt::Formatter<'_>, + ) -> fmt::Result { + write!(f, " nbytes={}", array.nbytes()) + } +} + +/// Renders the full snapshot content for one compressed corpus entry. +fn render(input: &ArrayRef, compressed: &ArrayRef) -> String { + format!( + "input: {}, len={}, nbytes={}\n{}", + input.dtype(), + input.len(), + input.nbytes(), + compressed + .tree_display_builder() + .with(EncodingSummaryExtractor) + .with(ExactNbytesExtractor) + .with(MetadataExtractor) + ) +} + +/// Compresses every corpus entry twice (direct determinism check) and snapshots the result +/// under `{variant}__{entry}`. +fn golden_corpus_snapshots(variant: &str, compressor: &BtrBlocksCompressor) -> VortexResult<()> { + for (name, array) in corpus()? { + let rendered = { + let mut exec_ctx = SESSION.create_execution_ctx(); + render(&array, &compressor.compress(&array, &mut exec_ctx)?) + }; + let rendered_again = { + let mut exec_ctx = SESSION.create_execution_ctx(); + render(&array, &compressor.compress(&array, &mut exec_ctx)?) + }; + assert_eq!( + rendered, rendered_again, + "compressing corpus entry {name} twice produced different results" + ); + + insta::assert_snapshot!(format!("{variant}__{name}"), rendered); + } + Ok(()) +} + +/// The fixed corpus: deterministic synthetic arrays covering each scheme's habitat. +fn corpus() -> VortexResult> { + Ok(vec![ + ("int_monotone_jitter", int_monotone_jitter()), + ("int_arithmetic_sequence", int_arithmetic_sequence()), + ("int_low_cardinality", int_low_cardinality()), + ("int_runs", int_runs()), + ("int_sparse_outliers", int_sparse_outliers()), + ("int_mostly_null", int_mostly_null()), + ("int_negatives", int_negatives()), + ("int_wide_random", int_wide_random()), + ("float_alp_prices", float_alp_prices()), + ("float_low_cardinality", float_low_cardinality()), + ("float_full_precision", float_full_precision()), + ("float_mostly_null", float_mostly_null()), + ("string_fsst_structured", string_fsst_structured()), + ("string_low_cardinality", string_low_cardinality()), + ("binary_low_cardinality", binary_low_cardinality()), + ("decimal_prices", decimal_prices()), + ("temporal_timestamp_micros", temporal_timestamp_micros()), + ("bool_random", bool_random()), + ("struct_mixed", struct_mixed()?), + ("list_of_int_runs", list_of_int_runs()?), + ]) +} + +/// Near-monotone u64 (timestamp-like): FoR/BitPacking habitat, Delta habitat when enabled. +fn int_monotone_jitter() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(101); + let mut value = 1_700_000_000_000u64; + let values: Buffer = (0..N) + .map(|_| { + value += 900 + rng.random_range(0..200); + value + }) + .collect(); + PrimitiveArray::new(values, Validity::NonNullable).into_array() +} + +/// Exact arithmetic sequence: Sequence habitat (distinct == len, no nulls). +fn int_arithmetic_sequence() -> ArrayRef { + let values: Buffer = (0..N as i64).map(|i| 10_000 + 7 * i).collect(); + PrimitiveArray::new(values, Validity::NonNullable).into_array() +} + +/// A handful of widely-spaced distinct values: IntDict habitat. +fn int_low_cardinality() -> ArrayRef { + const DISTINCT: [i64; 6] = [0, 123_400, 617_000, 1_234_000, 12_340_000, 37_020_000]; + let mut rng = StdRng::seed_from_u64(102); + let values: Buffer = (0..N).map(|_| DISTINCT[rng.random_range(0..6)]).collect(); + PrimitiveArray::new(values, Validity::NonNullable).into_array() +} + +/// Long runs over a moderate value set: RunEnd/RLE habitat. +fn int_runs() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(103); + let mut values: Vec = Vec::with_capacity(N); + while values.len() < N { + let value = rng.random_range(-50_000..50_000i32); + let run = rng.random_range(8..25); + values.extend(std::iter::repeat_n(value, run)); + } + values.truncate(N); + PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable).into_array() +} + +/// One dominant value with rare large outliers: Sparse habitat. +fn int_sparse_outliers() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(104); + let values: Buffer = (0..N) + .map(|_| { + if rng.random_range(0..100) < 5 { + rng.random_range(1_000_000_000..2_000_000_000i64) + } else { + 1_000_000 + } + }) + .collect(); + PrimitiveArray::new(values, Validity::NonNullable).into_array() +} + +/// 95% nulls over small values: null-dominated integer habitat. +fn int_mostly_null() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(105); + let mut validity: Vec = Vec::with_capacity(N); + let values: Buffer = (0..N) + .map(|_| { + let valid = rng.random_range(0..100) < 5; + validity.push(valid); + if valid { rng.random_range(0..1000) } else { 0 } + }) + .collect(); + PrimitiveArray::new( + values, + Validity::Array(BoolArray::from_iter(validity).into_array()), + ) + .into_array() +} + +/// Small values of mixed sign: ZigZag habitat. +fn int_negatives() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(106); + let values: Buffer = (0..N).map(|_| rng.random_range(-128..128i64)).collect(); + PrimitiveArray::new(values, Validity::NonNullable).into_array() +} + +/// Full-width random u64: essentially incompressible; pins the "no scheme wins" path. +fn int_wide_random() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(107); + let values: Buffer = (0..N).map(|_| rng.random::()).collect(); + PrimitiveArray::new(values, Validity::NonNullable).into_array() +} + +/// Two-decimal-digit "prices": ALP habitat. +fn float_alp_prices() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(108); + let values: Buffer = (0..N) + .map(|_| rng.random_range(0..10_000_000i64) as f64 / 100.0) + .collect(); + PrimitiveArray::new(values, Validity::NonNullable).into_array() +} + +/// A handful of distinct floats: FloatDict habitat. +fn float_low_cardinality() -> ArrayRef { + const DISTINCT: [f64; 8] = [0.0, 0.5, 1.25, 2.75, 3.5, 10.125, 100.0625, 1000.03125]; + let mut rng = StdRng::seed_from_u64(109); + let values: Buffer = (0..N).map(|_| DISTINCT[rng.random_range(0..8)]).collect(); + PrimitiveArray::new(values, Validity::NonNullable).into_array() +} + +/// Full-precision uniform floats: ALP-RD habitat. +fn float_full_precision() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(110); + let values: Buffer = (0..N).map(|_| rng.random::()).collect(); + PrimitiveArray::new(values, Validity::NonNullable).into_array() +} + +/// 95% nulls over floats: null-dominated sparse float habitat. +fn float_mostly_null() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(111); + let mut validity: Vec = Vec::with_capacity(N); + let values: Buffer = (0..N) + .map(|_| { + let valid = rng.random_range(0..100) < 5; + validity.push(valid); + if valid { + rng.random::() * 100.0 + } else { + 0.0 + } + }) + .collect(); + PrimitiveArray::new( + values, + Validity::Array(BoolArray::from_iter(validity).into_array()), + ) + .into_array() +} + +/// High-cardinality strings with shared substructure (emails): FSST habitat. +fn string_fsst_structured() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(112); + let strings: Vec = (0..N) + .map(|_| { + format!( + "user{:06}@example{}.com", + rng.random_range(0..1_000_000), + rng.random_range(0..100) + ) + }) + .collect(); + VarBinViewArray::from_iter_str(strings.iter().map(String::as_str)).into_array() +} + +/// A dozen distinct strings: StringDict habitat. +fn string_low_cardinality() -> ArrayRef { + const DISTINCT: [&str; 12] = [ + "alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel", "india", + "juliett", "kilo", "lima", + ]; + let mut rng = StdRng::seed_from_u64(113); + let strings: Vec> = (0..N) + .map(|_| Some(DISTINCT[rng.random_range(0..12)])) + .collect(); + VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)).into_array() +} + +/// Low-cardinality binary blobs: BinaryDict habitat (Zstd habitat under `compact`). +fn binary_low_cardinality() -> ArrayRef { + const DISTINCT: [&[u8]; 5] = [ + &[0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08], + &[0xDE, 0xAD, 0xBE, 0xEF], + &[0x00; 16], + &[0xFF, 0x00, 0xFF, 0x00, 0xFF, 0x00], + &[0x42; 12], + ]; + let mut rng = StdRng::seed_from_u64(114); + let blobs: Vec> = (0..N) + .map(|_| Some(DISTINCT[rng.random_range(0..5)])) + .collect(); + VarBinViewArray::from_iter(blobs, DType::Binary(Nullability::NonNullable)).into_array() +} + +/// Two-decimal-place decimals: DecimalScheme (byte-parts) habitat. +fn decimal_prices() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(115); + let values: Buffer = (0..N).map(|_| rng.random_range(0..10_000_000i64)).collect(); + DecimalArray::new(values, DecimalDType::new(12, 2), Validity::NonNullable).into_array() +} + +/// Near-monotone microsecond timestamps: TemporalScheme (datetime-parts) habitat. +fn temporal_timestamp_micros() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(116); + let mut value = 1_700_000_000_000_000i64; + let values: Buffer = (0..N) + .map(|_| { + value += rng.random_range(1_000..1_000_000); + value + }) + .collect(); + TemporalArray::new_timestamp( + PrimitiveArray::new(values, Validity::NonNullable).into_array(), + TimeUnit::Microseconds, + Some(Arc::from("UTC")), + ) + .into_array() +} + +/// Random booleans: no bool scheme is registered, pinning the "stays canonical" path. +fn bool_random() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(117); + BoolArray::from_iter((0..N).map(|_| rng.random::())).into_array() +} + +/// Struct of int/string/float fields: pins the structural recursion path. +fn struct_mixed() -> VortexResult { + Ok(StructArray::from_fields(&[ + ("id", int_arithmetic_sequence()), + ("category", string_low_cardinality()), + ("value", float_alp_prices()), + ])? + .into_array()) +} + +/// Variable-length lists of run-heavy ints: pins the list offsets/elements path. +fn list_of_int_runs() -> VortexResult { + let mut rng = StdRng::seed_from_u64(118); + let elements = int_runs(); + let mut offsets: Vec = Vec::with_capacity(N / 4 + 1); + let mut offset = 0i32; + offsets.push(offset); + while (offset as usize) < N { + offset = (offset + rng.random_range(1..8)).min(N as i32); + offsets.push(offset); + } + let offsets = PrimitiveArray::new(Buffer::copy_from(&offsets), Validity::NonNullable); + Ok(ListArray::try_new(elements, offsets.into_array(), Validity::NonNullable)?.into_array()) +} + +/// Excludes OnPair from the golden compressors: its dictionary training (upstream `onpair` +/// crate) iterates randomly-seeded `hashbrown` maps, so its compressed output — and therefore +/// its sampled estimate — differs run-to-run. A nondeterministic scheme cannot serve as a +/// golden baseline; excluding it keeps the remaining unstable schemes pinned. +#[cfg(feature = "unstable_encodings")] +fn without_onpair( + builder: vortex_btrblocks::BtrBlocksCompressorBuilder, +) -> vortex_btrblocks::BtrBlocksCompressorBuilder { + use vortex_btrblocks::SchemeExt; + use vortex_btrblocks::schemes::string::OnPairScheme; + + builder.exclude_schemes([OnPairScheme.id()]) +} + +#[cfg(not(feature = "unstable_encodings"))] +#[test] +fn golden_default() -> VortexResult<()> { + golden_corpus_snapshots("default", &BtrBlocksCompressor::default()) +} + +#[cfg(feature = "unstable_encodings")] +#[test] +fn golden_unstable() -> VortexResult<()> { + use vortex_btrblocks::BtrBlocksCompressorBuilder; + + golden_corpus_snapshots( + "unstable", + &without_onpair(BtrBlocksCompressorBuilder::default()).build(), + ) +} + +#[cfg(all(feature = "unstable_encodings", feature = "zstd", feature = "pco"))] +#[test] +fn golden_compact() -> VortexResult<()> { + use vortex_btrblocks::BtrBlocksCompressorBuilder; + + golden_corpus_snapshots( + "compact", + &without_onpair(BtrBlocksCompressorBuilder::default().with_compact()).build(), + ) +} diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__binary_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__compact__binary_low_cardinality.snap new file mode 100644 index 00000000000..d68474ce61c --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__binary_low_cardinality.snap @@ -0,0 +1,11 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: binary, len=16384, nbytes=315856 +root: vortex.dict(binary, len=16384) nbytes=6199 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 + metadata: bit_width: 3, offset: 0 + values: vortex.zstd(binary, len=5) nbytes=55 + metadata: nrows: 5, slice: 0..5 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__bool_random.snap b/vortex-btrblocks/tests/snapshots/golden__compact__bool_random.snap new file mode 100644 index 00000000000..d80e2fa4464 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__bool_random.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: bool, len=16384, nbytes=2048 +root: vortex.bool(bool, len=16384) nbytes=2048 + metadata: offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__decimal_prices.snap b/vortex-btrblocks/tests/snapshots/golden__compact__decimal_prices.snap new file mode 100644 index 00000000000..780ef30a6b0 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__decimal_prices.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: decimal(12,2), len=16384, nbytes=131072 +root: vortex.decimal_byte_parts(decimal(12,2), len=16384) nbytes=47666 + metadata: + msp: vortex.pco(i32, len=16384) nbytes=47666 + metadata: ptype: i32, nrows: 16384, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__float_alp_prices.snap b/vortex-btrblocks/tests/snapshots/golden__compact__float_alp_prices.snap new file mode 100644 index 00000000000..b1a3b41c4e7 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__float_alp_prices.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64, len=16384, nbytes=131072 +root: vortex.alp(f64, len=16384) nbytes=47699 + metadata: exponents: e: 14, f: 12 + encoded: vortex.pco(i64, len=16384) nbytes=47699 + metadata: ptype: i64, nrows: 16384, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__float_full_precision.snap b/vortex-btrblocks/tests/snapshots/golden__compact__float_full_precision.snap new file mode 100644 index 00000000000..20fecf31e22 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__float_full_precision.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64, len=16384, nbytes=131072 +root: vortex.pco(f64, len=16384) nbytes=110692 + metadata: ptype: f64, nrows: 16384, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__float_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__compact__float_low_cardinality.snap new file mode 100644 index 00000000000..006f5a93b69 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__float_low_cardinality.snap @@ -0,0 +1,13 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64, len=16384, nbytes=131072 +root: vortex.dict(f64, len=16384) nbytes=6184 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 + metadata: bit_width: 3, offset: 0 + values: vortex.alp(f64, len=8) nbytes=40 + metadata: exponents: e: 16, f: 11 + encoded: vortex.pco(i64, len=8) nbytes=40 + metadata: ptype: i64, nrows: 8, slice: 0..8 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__float_mostly_null.snap b/vortex-btrblocks/tests/snapshots/golden__compact__float_mostly_null.snap new file mode 100644 index 00000000000..134acccdc1c --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__float_mostly_null.snap @@ -0,0 +1,13 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64?, len=16384, nbytes=133120 +root: vortex.sparse(f64?, len=16384) nbytes=7545 + metadata: fill_value: null + patch_indices: vortex.pco(u16, len=850) nbytes=636 + metadata: ptype: u16, nrows: 850, slice: 0..850 + patch_values: vortex.primitive(f64?, len=850) nbytes=6907 + metadata: ptype: f64 + validity: vortex.bool(bool, len=850) nbytes=107 + metadata: offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__int_arithmetic_sequence.snap b/vortex-btrblocks/tests/snapshots/golden__compact__int_arithmetic_sequence.snap new file mode 100644 index 00000000000..873f3655a1f --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__int_arithmetic_sequence.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i64, len=16384, nbytes=131072 +root: vortex.sequence(i64, len=16384) nbytes=0 + metadata: base: 10000i64, multiplier: 7i64 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__int_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__compact__int_low_cardinality.snap new file mode 100644 index 00000000000..0cbe2e06814 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__int_low_cardinality.snap @@ -0,0 +1,11 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i64, len=16384, nbytes=131072 +root: vortex.dict(i64, len=16384) nbytes=6177 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 + metadata: bit_width: 3, offset: 0 + values: vortex.pco(i64, len=6) nbytes=33 + metadata: ptype: i64, nrows: 6, slice: 0..6 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__int_monotone_jitter.snap b/vortex-btrblocks/tests/snapshots/golden__compact__int_monotone_jitter.snap new file mode 100644 index 00000000000..4a63aa0da71 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__int_monotone_jitter.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: u64, len=16384, nbytes=131072 +root: vortex.pco(u64, len=16384) nbytes=15715 + metadata: ptype: u64, nrows: 16384, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__int_mostly_null.snap b/vortex-btrblocks/tests/snapshots/golden__compact__int_mostly_null.snap new file mode 100644 index 00000000000..71e1c671a03 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__int_mostly_null.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i32?, len=16384, nbytes=67584 +root: vortex.pco(i32?, len=16384) nbytes=3086 + metadata: ptype: i32, nrows: 16384, slice: 0..16384 + validity: vortex.bool(bool, len=16384) nbytes=2048 + metadata: offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__int_negatives.snap b/vortex-btrblocks/tests/snapshots/golden__compact__int_negatives.snap new file mode 100644 index 00000000000..24949893190 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__int_negatives.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i64, len=16384, nbytes=131072 +root: fastlanes.for(i64, len=16384) nbytes=16384 + metadata: reference: -128i64 + encoded: fastlanes.bitpacked(i64, len=16384) nbytes=16384 + metadata: bit_width: 8, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__compact__int_runs.snap new file mode 100644 index 00000000000..4ab2e46cead --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__int_runs.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i32, len=16384, nbytes=65536 +root: vortex.zigzag(i32, len=16384) nbytes=2969 + metadata: + encoded: vortex.pco(u32, len=16384) nbytes=2969 + metadata: ptype: u32, nrows: 16384, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__int_sparse_outliers.snap b/vortex-btrblocks/tests/snapshots/golden__compact__int_sparse_outliers.snap new file mode 100644 index 00000000000..0303e0e4ef1 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__int_sparse_outliers.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i64, len=16384, nbytes=131072 +root: vortex.pco(i64, len=16384) nbytes=3817 + metadata: ptype: i64, nrows: 16384, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__int_wide_random.snap b/vortex-btrblocks/tests/snapshots/golden__compact__int_wide_random.snap new file mode 100644 index 00000000000..90e1894f910 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__int_wide_random.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: u64, len=16384, nbytes=131072 +root: vortex.primitive(u64, len=16384) nbytes=131072 + metadata: ptype: u64 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__list_of_int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__compact__list_of_int_runs.snap new file mode 100644 index 00000000000..88274daf4b1 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__list_of_int_runs.snap @@ -0,0 +1,13 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: list(i32), len=4066, nbytes=81804 +root: vortex.list(list(i32), len=4066) nbytes=4434 + metadata: + elements: vortex.zigzag(i32, len=16384) nbytes=2969 + metadata: + encoded: vortex.pco(u32, len=16384) nbytes=2969 + metadata: ptype: u32, nrows: 16384, slice: 0..16384 + offsets: vortex.pco(u16, len=4067) nbytes=1465 + metadata: ptype: u16, nrows: 4067, slice: 0..4067 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__string_fsst_structured.snap b/vortex-btrblocks/tests/snapshots/golden__compact__string_fsst_structured.snap new file mode 100644 index 00000000000..49ad248087b --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__string_fsst_structured.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: utf8, len=16384, nbytes=653785 +root: vortex.zstd(utf8, len=16384) nbytes=109119 + metadata: nrows: 16384, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__string_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__compact__string_low_cardinality.snap new file mode 100644 index 00000000000..a65054ac178 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__string_low_cardinality.snap @@ -0,0 +1,11 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: utf8, len=16384, nbytes=262144 +root: vortex.dict(utf8, len=16384) nbytes=8287 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=8192 + metadata: bit_width: 4, offset: 0 + values: vortex.zstd(utf8, len=12) nbytes=95 + metadata: nrows: 12, slice: 0..12 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__struct_mixed.snap b/vortex-btrblocks/tests/snapshots/golden__compact__struct_mixed.snap new file mode 100644 index 00000000000..83e34c583dc --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__struct_mixed.snap @@ -0,0 +1,19 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: {id=i64, category=utf8, value=f64}, len=16384, nbytes=524288 +root: vortex.struct({id=i64, category=utf8, value=f64}, len=16384) nbytes=55986 + metadata: + id: vortex.sequence(i64, len=16384) nbytes=0 + metadata: base: 10000i64, multiplier: 7i64 + category: vortex.dict(utf8, len=16384) nbytes=8287 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=8192 + metadata: bit_width: 4, offset: 0 + values: vortex.zstd(utf8, len=12) nbytes=95 + metadata: nrows: 12, slice: 0..12 + value: vortex.alp(f64, len=16384) nbytes=47699 + metadata: exponents: e: 14, f: 12 + encoded: vortex.pco(i64, len=16384) nbytes=47699 + metadata: ptype: i64, nrows: 16384, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__temporal_timestamp_micros.snap b/vortex-btrblocks/tests/snapshots/golden__compact__temporal_timestamp_micros.snap new file mode 100644 index 00000000000..0422808ce29 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__compact__temporal_timestamp_micros.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: vortex.timestamp[µs, tz=UTC](i64), len=16384, nbytes=131072 +root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=40985 + metadata: + storage: vortex.pco(i64, len=16384) nbytes=40985 + metadata: ptype: i64, nrows: 16384, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__binary_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__default__binary_low_cardinality.snap new file mode 100644 index 00000000000..56918c0d096 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__binary_low_cardinality.snap @@ -0,0 +1,11 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: binary, len=16384, nbytes=315856 +root: vortex.dict(binary, len=16384) nbytes=6240 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 + metadata: bit_width: 3, offset: 0 + values: vortex.varbinview(binary, len=5) nbytes=96 + metadata: diff --git a/vortex-btrblocks/tests/snapshots/golden__default__bool_random.snap b/vortex-btrblocks/tests/snapshots/golden__default__bool_random.snap new file mode 100644 index 00000000000..d80e2fa4464 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__bool_random.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: bool, len=16384, nbytes=2048 +root: vortex.bool(bool, len=16384) nbytes=2048 + metadata: offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__decimal_prices.snap b/vortex-btrblocks/tests/snapshots/golden__default__decimal_prices.snap new file mode 100644 index 00000000000..f669755e4b1 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__decimal_prices.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: decimal(12,2), len=16384, nbytes=131072 +root: vortex.decimal_byte_parts(decimal(12,2), len=16384) nbytes=49152 + metadata: + msp: fastlanes.bitpacked(i32, len=16384) nbytes=49152 + metadata: bit_width: 24, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__float_alp_prices.snap b/vortex-btrblocks/tests/snapshots/golden__default__float_alp_prices.snap new file mode 100644 index 00000000000..f70d90a301d --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__float_alp_prices.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64, len=16384, nbytes=131072 +root: vortex.alp(f64, len=16384) nbytes=49152 + metadata: exponents: e: 14, f: 12 + encoded: fastlanes.bitpacked(i64, len=16384) nbytes=49152 + metadata: bit_width: 24, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__float_full_precision.snap b/vortex-btrblocks/tests/snapshots/golden__default__float_full_precision.snap new file mode 100644 index 00000000000..487e7c7900b --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__float_full_precision.snap @@ -0,0 +1,15 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64, len=16384, nbytes=131072 +root: vortex.alprd(f64, len=16384) nbytes=112880 + metadata: right_bit_width: 52, patch_offset: 0 + left_parts: fastlanes.bitpacked(u16, len=16384) nbytes=6144 + metadata: bit_width: 3, offset: 0 + right_parts: fastlanes.bitpacked(u64, len=16384) nbytes=106496 + metadata: bit_width: 52, offset: 0 + patch_indices: vortex.primitive(u16, len=60) nbytes=120 + metadata: ptype: u16 + patch_values: vortex.primitive(u16, len=60) nbytes=120 + metadata: ptype: u16 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__float_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__default__float_low_cardinality.snap new file mode 100644 index 00000000000..38d952262d6 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__float_low_cardinality.snap @@ -0,0 +1,19 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64, len=16384, nbytes=131072 +root: vortex.alp(f64, len=16384) nbytes=10233 + metadata: exponents: e: 16, f: 12, patch_offset: 0 + encoded: vortex.dict(i64, len=16384) nbytes=6200 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 + metadata: bit_width: 3, offset: 0 + values: vortex.primitive(i64, len=7) nbytes=56 + metadata: ptype: i64 + patch_indices: vortex.primitive(u16, len=1996) nbytes=3992 + metadata: ptype: u16 + patch_values: vortex.constant(f64, len=1996) nbytes=9 + metadata: scalar: 1000.03125f64 + patch_chunk_offsets: vortex.primitive(u16, len=16) nbytes=32 + metadata: ptype: u16 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__float_mostly_null.snap b/vortex-btrblocks/tests/snapshots/golden__default__float_mostly_null.snap new file mode 100644 index 00000000000..7b88d819472 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__float_mostly_null.snap @@ -0,0 +1,13 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64?, len=16384, nbytes=133120 +root: vortex.sparse(f64?, len=16384) nbytes=8609 + metadata: fill_value: null + patch_indices: vortex.primitive(u16, len=850) nbytes=1700 + metadata: ptype: u16 + patch_values: vortex.primitive(f64?, len=850) nbytes=6907 + metadata: ptype: f64 + validity: vortex.bool(bool, len=850) nbytes=107 + metadata: offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_arithmetic_sequence.snap b/vortex-btrblocks/tests/snapshots/golden__default__int_arithmetic_sequence.snap new file mode 100644 index 00000000000..873f3655a1f --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__int_arithmetic_sequence.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i64, len=16384, nbytes=131072 +root: vortex.sequence(i64, len=16384) nbytes=0 + metadata: base: 10000i64, multiplier: 7i64 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__default__int_low_cardinality.snap new file mode 100644 index 00000000000..8648200a138 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__int_low_cardinality.snap @@ -0,0 +1,11 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i64, len=16384, nbytes=131072 +root: vortex.dict(i64, len=16384) nbytes=6192 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 + metadata: bit_width: 3, offset: 0 + values: vortex.primitive(i64, len=6) nbytes=48 + metadata: ptype: i64 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap b/vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap new file mode 100644 index 00000000000..4800aae289d --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: u64, len=16384, nbytes=131072 +root: fastlanes.for(u64, len=16384) nbytes=49152 + metadata: reference: 1700000001036u64 + encoded: fastlanes.bitpacked(u64, len=16384) nbytes=49152 + metadata: bit_width: 24, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_mostly_null.snap b/vortex-btrblocks/tests/snapshots/golden__default__int_mostly_null.snap new file mode 100644 index 00000000000..73566369279 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__int_mostly_null.snap @@ -0,0 +1,13 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i32?, len=16384, nbytes=67584 +root: vortex.sparse(i32?, len=16384) nbytes=3031 + metadata: fill_value: null + patch_indices: vortex.primitive(u16, len=823) nbytes=1646 + metadata: ptype: u16 + patch_values: fastlanes.bitpacked(i32?, len=823) nbytes=1383 + metadata: bit_width: 10, offset: 0 + validity_child: vortex.bool(bool, len=823) nbytes=103 + metadata: offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_negatives.snap b/vortex-btrblocks/tests/snapshots/golden__default__int_negatives.snap new file mode 100644 index 00000000000..24949893190 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__int_negatives.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i64, len=16384, nbytes=131072 +root: fastlanes.for(i64, len=16384) nbytes=16384 + metadata: reference: -128i64 + encoded: fastlanes.bitpacked(i64, len=16384) nbytes=16384 + metadata: bit_width: 8, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__default__int_runs.snap new file mode 100644 index 00000000000..271de8faf93 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__int_runs.snap @@ -0,0 +1,15 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i32, len=16384, nbytes=65536 +root: vortex.runend(i32, len=16384) nbytes=3968 + metadata: offset: 0 + ends: fastlanes.for(u16, len=1020) nbytes=1792 + metadata: reference: 13u16 + encoded: fastlanes.bitpacked(u16, len=1020) nbytes=1792 + metadata: bit_width: 14, offset: 0 + values: fastlanes.for(i32, len=1020) nbytes=2176 + metadata: reference: -49931i32 + encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176 + metadata: bit_width: 17, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_sparse_outliers.snap b/vortex-btrblocks/tests/snapshots/golden__default__int_sparse_outliers.snap new file mode 100644 index 00000000000..53a5d1e023c --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__int_sparse_outliers.snap @@ -0,0 +1,13 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i64, len=16384, nbytes=131072 +root: vortex.sparse(i64, len=16384) nbytes=5540 + metadata: fill_value: 1000000i64 + patch_indices: vortex.primitive(u16, len=848) nbytes=1696 + metadata: ptype: u16 + patch_values: fastlanes.for(i64, len=848) nbytes=3840 + metadata: reference: 1000830099i64 + encoded: fastlanes.bitpacked(i64, len=848) nbytes=3840 + metadata: bit_width: 30, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_wide_random.snap b/vortex-btrblocks/tests/snapshots/golden__default__int_wide_random.snap new file mode 100644 index 00000000000..90e1894f910 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__int_wide_random.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: u64, len=16384, nbytes=131072 +root: vortex.primitive(u64, len=16384) nbytes=131072 + metadata: ptype: u64 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap new file mode 100644 index 00000000000..c9554add05a --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap @@ -0,0 +1,25 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: list(i32), len=4066, nbytes=81804 +root: vortex.list(list(i32), len=4066) nbytes=11146 + metadata: + elements: vortex.runend(i32, len=16384) nbytes=3968 + metadata: offset: 0 + ends: fastlanes.for(u16, len=1020) nbytes=1792 + metadata: reference: 13u16 + encoded: fastlanes.bitpacked(u16, len=1020) nbytes=1792 + metadata: bit_width: 14, offset: 0 + values: fastlanes.for(i32, len=1020) nbytes=2176 + metadata: reference: -49931i32 + encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176 + metadata: bit_width: 17, offset: 0 + offsets: fastlanes.bitpacked(u16, len=4067) nbytes=7178 + metadata: bit_width: 14, offset: 0 + patch_indices: vortex.primitive(u16, len=1) nbytes=2 + metadata: ptype: u16 + patch_values: vortex.constant(u16, len=1) nbytes=4 + metadata: scalar: 16384u16 + patch_chunk_offsets: vortex.primitive(u8, len=4) nbytes=4 + metadata: ptype: u8 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap b/vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap new file mode 100644 index 00000000000..327f050b0d7 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap @@ -0,0 +1,15 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: utf8, len=16384, nbytes=653785 +root: vortex.fsst(utf8, len=16384) nbytes=151382 + metadata: len: 16384, nsymbols: 223 + uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=3154 + metadata: fill_value: 24u8 + patch_indices: vortex.primitive(u16, len=1575) nbytes=3150 + metadata: ptype: u16 + patch_values: vortex.constant(u8, len=1575) nbytes=2 + metadata: scalar: 23u8 + codes_offsets: fastlanes.bitpacked(u32, len=16385) nbytes=36992 + metadata: bit_width: 17, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__string_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__default__string_low_cardinality.snap new file mode 100644 index 00000000000..ec79989d46a --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__string_low_cardinality.snap @@ -0,0 +1,15 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: utf8, len=16384, nbytes=262144 +root: vortex.dict(utf8, len=16384) nbytes=8373 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=8192 + metadata: bit_width: 4, offset: 0 + values: vortex.fsst(utf8, len=12) nbytes=181 + metadata: len: 12, nsymbols: 9 + uncompressed_lengths: vortex.primitive(u8, len=12) nbytes=12 + metadata: ptype: u8 + codes_offsets: vortex.primitive(u8, len=13) nbytes=13 + metadata: ptype: u8 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__struct_mixed.snap b/vortex-btrblocks/tests/snapshots/golden__default__struct_mixed.snap new file mode 100644 index 00000000000..be8f48b779b --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__struct_mixed.snap @@ -0,0 +1,23 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: {id=i64, category=utf8, value=f64}, len=16384, nbytes=524288 +root: vortex.struct({id=i64, category=utf8, value=f64}, len=16384) nbytes=57525 + metadata: + id: vortex.sequence(i64, len=16384) nbytes=0 + metadata: base: 10000i64, multiplier: 7i64 + category: vortex.dict(utf8, len=16384) nbytes=8373 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=8192 + metadata: bit_width: 4, offset: 0 + values: vortex.fsst(utf8, len=12) nbytes=181 + metadata: len: 12, nsymbols: 9 + uncompressed_lengths: vortex.primitive(u8, len=12) nbytes=12 + metadata: ptype: u8 + codes_offsets: vortex.primitive(u8, len=13) nbytes=13 + metadata: ptype: u8 + value: vortex.alp(f64, len=16384) nbytes=49152 + metadata: exponents: e: 14, f: 12 + encoded: fastlanes.bitpacked(i64, len=16384) nbytes=49152 + metadata: bit_width: 24, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap b/vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap new file mode 100644 index 00000000000..6f38e8e6f5d --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap @@ -0,0 +1,11 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: vortex.timestamp[µs, tz=UTC](i64), len=16384, nbytes=131072 +root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=67584 + metadata: + storage: fastlanes.for(i64, len=16384) nbytes=67584 + metadata: reference: 1700000000891673i64 + encoded: fastlanes.bitpacked(i64, len=16384) nbytes=67584 + metadata: bit_width: 33, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__binary_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__binary_low_cardinality.snap new file mode 100644 index 00000000000..56918c0d096 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__binary_low_cardinality.snap @@ -0,0 +1,11 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: binary, len=16384, nbytes=315856 +root: vortex.dict(binary, len=16384) nbytes=6240 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 + metadata: bit_width: 3, offset: 0 + values: vortex.varbinview(binary, len=5) nbytes=96 + metadata: diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__bool_random.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__bool_random.snap new file mode 100644 index 00000000000..d80e2fa4464 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__bool_random.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: bool, len=16384, nbytes=2048 +root: vortex.bool(bool, len=16384) nbytes=2048 + metadata: offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__decimal_prices.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__decimal_prices.snap new file mode 100644 index 00000000000..f669755e4b1 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__decimal_prices.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: decimal(12,2), len=16384, nbytes=131072 +root: vortex.decimal_byte_parts(decimal(12,2), len=16384) nbytes=49152 + metadata: + msp: fastlanes.bitpacked(i32, len=16384) nbytes=49152 + metadata: bit_width: 24, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__float_alp_prices.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__float_alp_prices.snap new file mode 100644 index 00000000000..f70d90a301d --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__float_alp_prices.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64, len=16384, nbytes=131072 +root: vortex.alp(f64, len=16384) nbytes=49152 + metadata: exponents: e: 14, f: 12 + encoded: fastlanes.bitpacked(i64, len=16384) nbytes=49152 + metadata: bit_width: 24, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__float_full_precision.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__float_full_precision.snap new file mode 100644 index 00000000000..487e7c7900b --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__float_full_precision.snap @@ -0,0 +1,15 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64, len=16384, nbytes=131072 +root: vortex.alprd(f64, len=16384) nbytes=112880 + metadata: right_bit_width: 52, patch_offset: 0 + left_parts: fastlanes.bitpacked(u16, len=16384) nbytes=6144 + metadata: bit_width: 3, offset: 0 + right_parts: fastlanes.bitpacked(u64, len=16384) nbytes=106496 + metadata: bit_width: 52, offset: 0 + patch_indices: vortex.primitive(u16, len=60) nbytes=120 + metadata: ptype: u16 + patch_values: vortex.primitive(u16, len=60) nbytes=120 + metadata: ptype: u16 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__float_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__float_low_cardinality.snap new file mode 100644 index 00000000000..38d952262d6 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__float_low_cardinality.snap @@ -0,0 +1,19 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64, len=16384, nbytes=131072 +root: vortex.alp(f64, len=16384) nbytes=10233 + metadata: exponents: e: 16, f: 12, patch_offset: 0 + encoded: vortex.dict(i64, len=16384) nbytes=6200 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 + metadata: bit_width: 3, offset: 0 + values: vortex.primitive(i64, len=7) nbytes=56 + metadata: ptype: i64 + patch_indices: vortex.primitive(u16, len=1996) nbytes=3992 + metadata: ptype: u16 + patch_values: vortex.constant(f64, len=1996) nbytes=9 + metadata: scalar: 1000.03125f64 + patch_chunk_offsets: vortex.primitive(u16, len=16) nbytes=32 + metadata: ptype: u16 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__float_mostly_null.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__float_mostly_null.snap new file mode 100644 index 00000000000..7b88d819472 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__float_mostly_null.snap @@ -0,0 +1,13 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: f64?, len=16384, nbytes=133120 +root: vortex.sparse(f64?, len=16384) nbytes=8609 + metadata: fill_value: null + patch_indices: vortex.primitive(u16, len=850) nbytes=1700 + metadata: ptype: u16 + patch_values: vortex.primitive(f64?, len=850) nbytes=6907 + metadata: ptype: f64 + validity: vortex.bool(bool, len=850) nbytes=107 + metadata: offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_arithmetic_sequence.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_arithmetic_sequence.snap new file mode 100644 index 00000000000..873f3655a1f --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__int_arithmetic_sequence.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i64, len=16384, nbytes=131072 +root: vortex.sequence(i64, len=16384) nbytes=0 + metadata: base: 10000i64, multiplier: 7i64 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_low_cardinality.snap new file mode 100644 index 00000000000..8648200a138 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__int_low_cardinality.snap @@ -0,0 +1,11 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i64, len=16384, nbytes=131072 +root: vortex.dict(i64, len=16384) nbytes=6192 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 + metadata: bit_width: 3, offset: 0 + values: vortex.primitive(i64, len=6) nbytes=48 + metadata: ptype: i64 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_monotone_jitter.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_monotone_jitter.snap new file mode 100644 index 00000000000..f1ded4c3e59 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__int_monotone_jitter.snap @@ -0,0 +1,15 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: u64, len=16384, nbytes=131072 +root: fastlanes.delta(u64, len=16384) nbytes=19840 + metadata: offset: 0 + bases: vortex.primitive(u64, len=256) nbytes=2048 + metadata: ptype: u64 + deltas: vortex.dict(u64, len=16384) nbytes=17792 + metadata: all_values_referenced: true + codes: vortex.primitive(u8, len=16384) nbytes=16384 + metadata: ptype: u8 + values: fastlanes.bitpacked(u64, len=201) nbytes=1408 + metadata: bit_width: 11, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_mostly_null.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_mostly_null.snap new file mode 100644 index 00000000000..73566369279 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__int_mostly_null.snap @@ -0,0 +1,13 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i32?, len=16384, nbytes=67584 +root: vortex.sparse(i32?, len=16384) nbytes=3031 + metadata: fill_value: null + patch_indices: vortex.primitive(u16, len=823) nbytes=1646 + metadata: ptype: u16 + patch_values: fastlanes.bitpacked(i32?, len=823) nbytes=1383 + metadata: bit_width: 10, offset: 0 + validity_child: vortex.bool(bool, len=823) nbytes=103 + metadata: offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_negatives.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_negatives.snap new file mode 100644 index 00000000000..24949893190 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__int_negatives.snap @@ -0,0 +1,9 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i64, len=16384, nbytes=131072 +root: fastlanes.for(i64, len=16384) nbytes=16384 + metadata: reference: -128i64 + encoded: fastlanes.bitpacked(i64, len=16384) nbytes=16384 + metadata: bit_width: 8, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_runs.snap new file mode 100644 index 00000000000..271de8faf93 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__int_runs.snap @@ -0,0 +1,15 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i32, len=16384, nbytes=65536 +root: vortex.runend(i32, len=16384) nbytes=3968 + metadata: offset: 0 + ends: fastlanes.for(u16, len=1020) nbytes=1792 + metadata: reference: 13u16 + encoded: fastlanes.bitpacked(u16, len=1020) nbytes=1792 + metadata: bit_width: 14, offset: 0 + values: fastlanes.for(i32, len=1020) nbytes=2176 + metadata: reference: -49931i32 + encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176 + metadata: bit_width: 17, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_sparse_outliers.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_sparse_outliers.snap new file mode 100644 index 00000000000..53a5d1e023c --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__int_sparse_outliers.snap @@ -0,0 +1,13 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: i64, len=16384, nbytes=131072 +root: vortex.sparse(i64, len=16384) nbytes=5540 + metadata: fill_value: 1000000i64 + patch_indices: vortex.primitive(u16, len=848) nbytes=1696 + metadata: ptype: u16 + patch_values: fastlanes.for(i64, len=848) nbytes=3840 + metadata: reference: 1000830099i64 + encoded: fastlanes.bitpacked(i64, len=848) nbytes=3840 + metadata: bit_width: 30, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__int_wide_random.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__int_wide_random.snap new file mode 100644 index 00000000000..90e1894f910 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__int_wide_random.snap @@ -0,0 +1,7 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: u64, len=16384, nbytes=131072 +root: vortex.primitive(u64, len=16384) nbytes=131072 + metadata: ptype: u64 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap new file mode 100644 index 00000000000..c9554add05a --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap @@ -0,0 +1,25 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: list(i32), len=4066, nbytes=81804 +root: vortex.list(list(i32), len=4066) nbytes=11146 + metadata: + elements: vortex.runend(i32, len=16384) nbytes=3968 + metadata: offset: 0 + ends: fastlanes.for(u16, len=1020) nbytes=1792 + metadata: reference: 13u16 + encoded: fastlanes.bitpacked(u16, len=1020) nbytes=1792 + metadata: bit_width: 14, offset: 0 + values: fastlanes.for(i32, len=1020) nbytes=2176 + metadata: reference: -49931i32 + encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176 + metadata: bit_width: 17, offset: 0 + offsets: fastlanes.bitpacked(u16, len=4067) nbytes=7178 + metadata: bit_width: 14, offset: 0 + patch_indices: vortex.primitive(u16, len=1) nbytes=2 + metadata: ptype: u16 + patch_values: vortex.constant(u16, len=1) nbytes=4 + metadata: scalar: 16384u16 + patch_chunk_offsets: vortex.primitive(u8, len=4) nbytes=4 + metadata: ptype: u8 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap new file mode 100644 index 00000000000..327f050b0d7 --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap @@ -0,0 +1,15 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: utf8, len=16384, nbytes=653785 +root: vortex.fsst(utf8, len=16384) nbytes=151382 + metadata: len: 16384, nsymbols: 223 + uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=3154 + metadata: fill_value: 24u8 + patch_indices: vortex.primitive(u16, len=1575) nbytes=3150 + metadata: ptype: u16 + patch_values: vortex.constant(u8, len=1575) nbytes=2 + metadata: scalar: 23u8 + codes_offsets: fastlanes.bitpacked(u32, len=16385) nbytes=36992 + metadata: bit_width: 17, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__string_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__string_low_cardinality.snap new file mode 100644 index 00000000000..ec79989d46a --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__string_low_cardinality.snap @@ -0,0 +1,15 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: utf8, len=16384, nbytes=262144 +root: vortex.dict(utf8, len=16384) nbytes=8373 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=8192 + metadata: bit_width: 4, offset: 0 + values: vortex.fsst(utf8, len=12) nbytes=181 + metadata: len: 12, nsymbols: 9 + uncompressed_lengths: vortex.primitive(u8, len=12) nbytes=12 + metadata: ptype: u8 + codes_offsets: vortex.primitive(u8, len=13) nbytes=13 + metadata: ptype: u8 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__struct_mixed.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__struct_mixed.snap new file mode 100644 index 00000000000..be8f48b779b --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__struct_mixed.snap @@ -0,0 +1,23 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: {id=i64, category=utf8, value=f64}, len=16384, nbytes=524288 +root: vortex.struct({id=i64, category=utf8, value=f64}, len=16384) nbytes=57525 + metadata: + id: vortex.sequence(i64, len=16384) nbytes=0 + metadata: base: 10000i64, multiplier: 7i64 + category: vortex.dict(utf8, len=16384) nbytes=8373 + metadata: all_values_referenced: true + codes: fastlanes.bitpacked(u8, len=16384) nbytes=8192 + metadata: bit_width: 4, offset: 0 + values: vortex.fsst(utf8, len=12) nbytes=181 + metadata: len: 12, nsymbols: 9 + uncompressed_lengths: vortex.primitive(u8, len=12) nbytes=12 + metadata: ptype: u8 + codes_offsets: vortex.primitive(u8, len=13) nbytes=13 + metadata: ptype: u8 + value: vortex.alp(f64, len=16384) nbytes=49152 + metadata: exponents: e: 14, f: 12 + encoded: fastlanes.bitpacked(i64, len=16384) nbytes=49152 + metadata: bit_width: 24, offset: 0 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__temporal_timestamp_micros.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__temporal_timestamp_micros.snap new file mode 100644 index 00000000000..a763fa28f4b --- /dev/null +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__temporal_timestamp_micros.snap @@ -0,0 +1,13 @@ +--- +source: vortex-btrblocks/tests/golden.rs +expression: rendered +--- +input: vortex.timestamp[µs, tz=UTC](i64), len=16384, nbytes=131072 +root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=43008 + metadata: + storage: fastlanes.delta(i64, len=16384) nbytes=43008 + metadata: offset: 0 + bases: vortex.primitive(i64, len=256) nbytes=2048 + metadata: ptype: i64 + deltas: fastlanes.bitpacked(i64, len=16384) nbytes=40960 + metadata: bit_width: 20, offset: 0 diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index e0397bacc91..ed9df797535 100644 --- a/vortex-buffer/Cargo.toml +++ b/vortex-buffer/Cargo.toml @@ -39,7 +39,6 @@ workspace = true [dev-dependencies] divan = { workspace = true } num-traits = { workspace = true } -rand = { workspace = true } rstest = { workspace = true } [[bench]] @@ -49,3 +48,11 @@ harness = false [[bench]] name = "vortex_bitbuffer" harness = false + +[[bench]] +name = "cpu_dispatch" +harness = false + +[[bench]] +name = "collect_bool" +harness = false diff --git a/vortex-buffer/benches/collect_bool.rs b/vortex-buffer/benches/collect_bool.rs new file mode 100644 index 00000000000..3bfe9a55b69 --- /dev/null +++ b/vortex-buffer/benches/collect_bool.rs @@ -0,0 +1,209 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks for `collect_bool` bit packing. +//! +//! Three groups, each with a scalar baseline that replicates the previous bit-at-a-time +//! implementation. Under CodSpeed only the shipped entry points are tracked +//! (`from_bool_slice`, `collect_bool_u32_gt`); the baselines, the dispatch loop, +//! per-SIMD-level loops, and the portable SWAR kernel compile out and are for local A/B runs: +//! +//! - `pack_words_*`: the portable SWAR kernel vs the scalar loop, word by word. The SIMD +//! kernels are covered by `words_gather_*` instead, where they can inline. +//! - `words_gather_*`: each multiversioned word loop with a bounds-check-free bool gather, +//! measuring the fused fill + pack pipeline per SIMD level. +//! - `collect_bool_*` / `from_bool_slice`: the public entry points end to end, with a +//! boolean-gather predicate and a `u32` comparison predicate. + +use divan::Bencher; +use vortex_buffer::BitBuffer; +#[cfg(not(codspeed))] +use vortex_buffer::collect_bool_word_scalar; +#[cfg(not(codspeed))] +use vortex_buffer::pack_bool_word_swar; + +fn main() { + // Pre-warm CPUID feature detection so the one-time probe cost is never + // included in any benchmark iteration. + #[cfg(target_arch = "x86_64")] + { + let _ = is_x86_feature_detected!("avx2"); + let _ = is_x86_feature_detected!("avx512f"); + let _ = is_x86_feature_detected!("avx512bw"); + } + + divan::main(); +} + +const INPUT_SIZE: &[usize] = &[1024, 65_536]; + +/// Deterministic pseudo-random words (LCG), the source for all benchmark inputs. +fn make_words(len: usize) -> impl Iterator { + let mut state = 0x9E37_79B9_7F4A_7C15u64; + (0..len).map(move |_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + state + }) +} + +fn make_bools(len: usize) -> Vec { + make_words(len).map(|w| (w >> 33) & 1 == 1).collect() +} + +fn make_u32s(len: usize) -> Vec { + make_words(len).map(|w| (w >> 32) as u32).collect() +} + +/// Benchmark a per-word packing kernel over `len` pre-materialized bools. +#[cfg(not(codspeed))] +fn bench_pack_words(bencher: Bencher, len: usize, pack: impl Fn(&[bool; 64]) -> u64 + Sync) { + let bools = make_bools(len); + bencher + .with_inputs(|| vec![0u64; len / 64]) + .bench_refs(|out| { + let (chunks, _) = bools.as_chunks::<64>(); + for (word, chunk) in out.iter_mut().zip(chunks) { + *word = pack(chunk); + } + }); +} + +#[cfg(not(codspeed))] +#[divan::bench(args = INPUT_SIZE)] +fn pack_words_scalar(bencher: Bencher, len: usize) { + bench_pack_words(bencher, len, |chunk| { + collect_bool_word_scalar(64, |i| chunk[i]) + }); +} + +#[cfg(not(codspeed))] +#[divan::bench(args = INPUT_SIZE)] +fn pack_words_swar(bencher: Bencher, len: usize) { + bench_pack_words(bencher, len, pack_bool_word_swar); +} + +/// Benchmark a full multiversioned word loop with a bounds-check-free bool gather, measuring +/// the packing pipeline itself (fill + pack fully inlined) rather than predicate evaluation. +#[cfg(not(codspeed))] +fn bench_words_gather( + bencher: Bencher, + len: usize, + collect: impl Fn(&mut [u64], usize, &[bool]) + Sync, +) { + let bools = make_bools(len); + bencher + .with_inputs(|| vec![0u64; len.div_ceil(64)]) + .bench_refs(|words| collect(words, len, &bools)); +} + +#[cfg(not(codspeed))] +#[divan::bench(args = INPUT_SIZE)] +fn words_gather_dispatch(bencher: Bencher, len: usize) { + bench_words_gather(bencher, len, |words, len, bools| { + // SAFETY: `collect_bool_words` invokes the predicate with indices `0..len` only. + vortex_buffer::collect_bool_words(words, len, |i| unsafe { *bools.get_unchecked(i) }) + }); +} + +#[cfg(not(codspeed))] +#[divan::bench(args = INPUT_SIZE)] +fn words_gather_scalar(bencher: Bencher, len: usize) { + bench_words_gather(bencher, len, |words, len, bools| { + // SAFETY: `collect_bool_words_old` invokes the predicate with indices `0..len` only. + collect_bool_words_old(words, len, |i| unsafe { *bools.get_unchecked(i) }) + }); +} + +#[cfg(target_arch = "x86_64")] +#[cfg(not(codspeed))] +#[divan::bench(args = INPUT_SIZE)] +fn words_gather_sse2(bencher: Bencher, len: usize) { + bench_words_gather(bencher, len, |words, len, bools| { + // SAFETY: SSE2 is part of the x86-64 baseline; indices passed are `0..len`. + unsafe { vortex_buffer::collect_bool_words_sse2(words, len, |i| *bools.get_unchecked(i)) } + }); +} + +#[cfg(target_arch = "x86_64")] +#[cfg(not(codspeed))] +#[divan::bench(args = INPUT_SIZE)] +fn words_gather_avx2(bencher: Bencher, len: usize) { + if !is_x86_feature_detected!("avx2") { + return; + } + bench_words_gather(bencher, len, |words, len, bools| { + // SAFETY: runtime detection guarantees AVX2; indices passed are `0..len`. + unsafe { vortex_buffer::collect_bool_words_avx2(words, len, |i| *bools.get_unchecked(i)) } + }); +} + +#[cfg(target_arch = "x86_64")] +#[cfg(not(codspeed))] +#[divan::bench(args = INPUT_SIZE)] +fn words_gather_avx512(bencher: Bencher, len: usize) { + if !(is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw")) { + return; + } + bench_words_gather(bencher, len, |words, len, bools| { + // SAFETY: runtime detection guarantees AVX-512F/BW; indices passed are `0..len`. + unsafe { vortex_buffer::collect_bool_words_avx512(words, len, |i| *bools.get_unchecked(i)) } + }); +} + +#[cfg(target_arch = "aarch64")] +#[cfg(not(codspeed))] +#[divan::bench(args = INPUT_SIZE)] +fn words_gather_neon(bencher: Bencher, len: usize) { + bench_words_gather(bencher, len, |words, len, bools| { + // SAFETY: NEON is part of the aarch64 baseline; indices passed are `0..len`. + unsafe { vortex_buffer::collect_bool_words_neon(words, len, |i| *bools.get_unchecked(i)) } + }); +} + +/// Faithful copy of the previous scalar-only `collect_bool_words` word loop, used as the +/// baseline for the end-to-end comparison. +#[cfg(not(codspeed))] +fn collect_bool_words_old(words: &mut [u64], len: usize, mut f: impl FnMut(usize) -> bool) { + let full = len / 64; + let remainder = len % 64; + for word_idx in 0..full { + let offset = word_idx * 64; + words[word_idx] = collect_bool_word_scalar(64, |bit_idx| f(offset + bit_idx)); + } + if remainder != 0 { + let offset = full * 64; + words[full] = collect_bool_word_scalar(remainder, |bit_idx| f(offset + bit_idx)); + } +} + +#[cfg(not(codspeed))] +#[divan::bench(args = INPUT_SIZE)] +fn from_bool_slice_old_scalar(bencher: Bencher, len: usize) { + let bools = make_bools(len); + bencher + .with_inputs(|| vec![0u64; len.div_ceil(64)]) + .bench_refs(|words| collect_bool_words_old(words, len, |i| bools[i])); +} + +#[divan::bench(args = INPUT_SIZE)] +fn from_bool_slice(bencher: Bencher, len: usize) { + let bools = make_bools(len); + bencher.bench(|| vortex_buffer::BitBufferMut::from(bools.as_slice())); +} + +#[cfg(not(codspeed))] +#[divan::bench(args = INPUT_SIZE)] +fn collect_bool_u32_gt_old_scalar(bencher: Bencher, len: usize) { + let values = make_u32s(len); + bencher + .with_inputs(|| vec![0u64; len.div_ceil(64)]) + .bench_refs(|words| collect_bool_words_old(words, len, |i| values[i] > u32::MAX / 2)); +} + +#[divan::bench(args = INPUT_SIZE)] +fn collect_bool_u32_gt(bencher: Bencher, len: usize) { + let values = make_u32s(len); + bencher.bench(|| BitBuffer::collect_bool(len, |i| values[i] > u32::MAX / 2)); +} diff --git a/vortex-buffer/benches/cpu_dispatch.rs b/vortex-buffer/benches/cpu_dispatch.rs new file mode 100644 index 00000000000..6891dd22119 --- /dev/null +++ b/vortex-buffer/benches/cpu_dispatch.rs @@ -0,0 +1,160 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Overhead comparison of one-time CPU-feature dispatch mechanisms. +//! +//! Every variant ends in a call to the same `#[inline(never)]` kernel, so the +//! difference between rows is pure dispatch overhead: +//! +//! * `direct`: plain call — the floor. +//! * `cpu_kernel`: [`CpuKernel`] — relaxed load + null check + indirect call. +//! * `cpu_kernel_unconditional`: a selector with no probes (the aarch64 shape) — +//! same steady state as `cpu_kernel`. +//! * `lazy_lock`: `std::sync::LazyLock` — Once-state check + value load + +//! indirect call. +//! * `feature_detect_each_call_1/_3`: the pre-`CpuKernel` pattern — 1 or 3 +//! `is_x86_feature_detected!` cached lookups per call, then a direct call. + +use std::sync::LazyLock; + +use divan::Bencher; +use divan::black_box; +use divan::counter::ItemsCount; +use vortex_buffer::CpuKernel; + +fn main() { + // Resolve every dispatcher and warm the std feature-detection cache so no + // benchmark iteration pays a one-time probe. + let seed = black_box(7); + let _ = via_direct(1, seed) + + via_cpu_kernel(1, seed) + + via_cpu_kernel_unconditional(1, seed) + + via_lazy_lock(1, seed); + #[cfg(target_arch = "x86_64")] + let _ = via_feature_detect_1(1, seed) + via_feature_detect_3(1, seed); + + divan::main(); +} + +type Kernel = fn(u64, u64) -> u64; + +#[inline(never)] +fn kernel_a(x: u64, seed: u64) -> u64 { + x.wrapping_mul(31).wrapping_add(seed) +} + +#[inline(never)] +fn kernel_b(x: u64, seed: u64) -> u64 { + x.wrapping_mul(33).wrapping_add(seed) +} + +fn has_bmi2() -> bool { + #[cfg(target_arch = "x86_64")] + { + std::arch::is_x86_feature_detected!("bmi2") + } + #[cfg(not(target_arch = "x86_64"))] + { + false + } +} + +// ── dispatch variants ─────────────────────────────────────────────────── + +#[inline(never)] +fn via_direct(x: u64, seed: u64) -> u64 { + kernel_a(x, seed) +} + +#[cfg(target_arch = "x86_64")] +#[inline(never)] +fn via_feature_detect_1(x: u64, seed: u64) -> u64 { + if std::arch::is_x86_feature_detected!("bmi2") { + kernel_a(x, seed) + } else { + kernel_b(x, seed) + } +} + +#[cfg(target_arch = "x86_64")] +#[inline(never)] +fn via_feature_detect_3(x: u64, seed: u64) -> u64 { + // Non-baseline features that are present on any recent x86_64 part, so all + // three cached lookups actually execute (mirrors the old select_in_chunk). + if std::arch::is_x86_feature_detected!("avx") + && std::arch::is_x86_feature_detected!("avx2") + && std::arch::is_x86_feature_detected!("bmi2") + { + kernel_a(x, seed) + } else { + kernel_b(x, seed) + } +} + +static LAZY: LazyLock = LazyLock::new(|| if has_bmi2() { kernel_a } else { kernel_b }); + +#[inline(never)] +fn via_lazy_lock(x: u64, seed: u64) -> u64 { + (*LAZY)(x, seed) +} + +static CPU_KERNEL: CpuKernel = + CpuKernel::new(|| if has_bmi2() { kernel_a } else { kernel_b }); + +#[inline(never)] +fn via_cpu_kernel(x: u64, seed: u64) -> u64 { + CPU_KERNEL.get()(x, seed) +} + +static CPU_KERNEL_UNCONDITIONAL: CpuKernel = CpuKernel::new(|| kernel_a); + +#[inline(never)] +fn via_cpu_kernel_unconditional(x: u64, seed: u64) -> u64 { + CPU_KERNEL_UNCONDITIONAL.get()(x, seed) +} + +// ── benches ───────────────────────────────────────────────────────────── + +const CALLS: u64 = 1024; + +fn bench_via(bencher: Bencher, via: fn(u64, u64) -> u64) { + bencher.counter(ItemsCount::new(CALLS)).bench(|| { + let mut acc = 0u64; + for i in 0..CALLS { + acc = acc.wrapping_add(via(black_box(i), black_box(7))); + } + acc + }) +} + +#[divan::bench] +fn direct(bencher: Bencher) { + bench_via(bencher, via_direct); +} + +#[divan::bench] +fn cpu_kernel_unconditional(bencher: Bencher) { + bench_via(bencher, via_cpu_kernel_unconditional); +} + +#[divan::bench] +fn cpu_kernel(bencher: Bencher) { + bench_via(bencher, via_cpu_kernel); +} + +#[divan::bench] +fn lazy_lock(bencher: Bencher) { + bench_via(bencher, via_lazy_lock); +} + +#[cfg(target_arch = "x86_64")] +#[divan::bench] +fn feature_detect_each_call_1(bencher: Bencher) { + bench_via(bencher, via_feature_detect_1); +} + +#[cfg(target_arch = "x86_64")] +#[divan::bench] +fn feature_detect_each_call_3(bencher: Bencher) { + bench_via(bencher, via_feature_detect_3); +} diff --git a/vortex-buffer/benches/vortex_bitbuffer.rs b/vortex-buffer/benches/vortex_bitbuffer.rs index 7efc6c6a114..b67b3f1bf39 100644 --- a/vortex-buffer/benches/vortex_bitbuffer.rs +++ b/vortex-buffer/benches/vortex_bitbuffer.rs @@ -19,6 +19,12 @@ fn main() { let _ = is_x86_feature_detected!("avx512vpopcntdq"); } + // Pre-resolve the one-time CpuKernel selections for the count/select paths so + // no benchmark iteration pays the first-call cost. + let warm = BitBuffer::from_iter((0..512).map(|i| i % 3 == 0)); + let _ = warm.true_count(); + let _ = warm.select(1); + divan::main(); } @@ -246,7 +252,12 @@ fn bitwise_not_vortex_buffer(bencher: Bencher, length: usize) { .bench_values(|buffer| !&buffer); } -#[divan::bench(args = INPUT_SIZE)] +/// The in-place NOT on an owned `BitBufferMut` performs no allocation and no copy, so below a +/// few thousand bits the measurement is fixed harness overhead and binary code layout rather +/// than the loop itself. Only the sizes where the loop dominates are worth measuring. +const NOT_MUT_INPUT_SIZE: &[usize] = &[16_384, 65_536]; + +#[divan::bench(args = NOT_MUT_INPUT_SIZE)] fn bitwise_not_vortex_buffer_mut(bencher: Bencher, length: usize) { bencher .with_inputs(|| BitBufferMut::from_iter((0..length).map(|i| i % 2 == 0))) diff --git a/vortex-buffer/benches/vortex_buffer.rs b/vortex-buffer/benches/vortex_buffer.rs index 74a402887f0..e1e6683e505 100644 --- a/vortex-buffer/benches/vortex_buffer.rs +++ b/vortex-buffer/benches/vortex_buffer.rs @@ -170,10 +170,17 @@ fn slice_tight_loop_arrow(bencher: Bencher, len: usize) { }); } +/// Loops like `slice_tight_loop_vortex` above: a single empty slice is a few nanoseconds, far +/// below the measurement noise floor, so the loop amortizes fixed overhead until the slice +/// itself dominates. #[divan::bench] -fn slice_empty_vortex(bencher: Bencher) { +fn slice_empty_tight_loop_vortex(bencher: Bencher) { let buf = Buffer::::from_iter((0..1024).map(|i| i % i32::MAX)); - bencher.bench(|| divan::black_box(buf.slice(8..8))); + bencher.bench(|| { + for _ in 0..1024 { + divan::black_box(buf.slice(8..8)); + } + }); } #[divan::bench(args = INPUT_SIZE)] diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index a8bd9e1929c..8dc5e38a42c 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -185,11 +185,43 @@ impl BitBuffer { } /// Invokes `f` with indexes `0..len` collecting the boolean results into a new [`BitBuffer`]. + /// + /// `f` is invoked exactly once per index, in ascending order, and the results are packed + /// with the baseline SIMD byte→bit instruction of the target. + /// + /// # Performance + /// + /// The packing is a few instructions per 64 bits, so evaluating `f` is usually the + /// bottleneck. In particular, a bounds-checked slice access in `f` (`|i| values[i] > x`) + /// blocks vectorization of the gather and can cost ~10x the packing itself. Since `f` only + /// ever sees indices `0..len`, callers reading from a slice with `len <= values.len()` may + /// soundly use `|i| unsafe { *values.get_unchecked(i) }`. + /// + /// Prefer this entry point for every predicate. Only switch to + /// [`Self::collect_bool_multiversioned`] after carefully checking that your specific `f` + /// meets its contract (a trivially cheap, bounds-check-free gather or comparison) — + /// ideally with a benchmark. #[inline] pub fn collect_bool bool>(len: usize, f: F) -> Self { BitBufferMut::collect_bool(len, f).freeze() } + /// Like [`Self::collect_bool`], but compiles the packing loop — with `f` inside it — once + /// per CPU feature level (AVX-512BW/AVX2/baseline) and selects a clone by runtime feature + /// detection. + /// + /// Calling this asserts that `f` is small and simple enough (e.g. a bounds-check-free slice + /// gather or comparison) that duplicating it per feature level and paying a + /// `#[target_feature]` call boundary beats inlining it once into your function. For any + /// non-trivial `f` that assertion is false — the boundary deoptimizes the predicate — so + /// unless you have carefully checked (ideally benchmarked) that your specific `f` + /// qualifies, use [`Self::collect_bool`]. See + /// [`collect_bool_words_multiversioned`](crate::bit::collect_bool_words_multiversioned). + #[inline] + pub fn collect_bool_multiversioned bool>(len: usize, f: F) -> Self { + BitBufferMut::collect_bool_multiversioned(len, f).freeze() + } + /// Maps over each bit in this buffer, calling `f(index, bit_value)` and collecting results. /// /// This is more efficient than `collect_bool` when you need to read the current bit value, @@ -379,6 +411,20 @@ impl BitBuffer { count_ones(self.buffer.as_slice(), self.offset, self.len) } + /// Get the number of set bits in the bit range `[start, end)`. + /// + /// Unlike `self.slice(start..end).true_count()`, this counts directly over the + /// existing backing buffer without allocating or cloning a new [`BitBuffer`], + /// making it cheap to call repeatedly over many small ranges. + /// + /// Panics if `start > end` or `end > len`. + #[inline] + pub fn count_range(&self, start: usize, end: usize) -> usize { + assert!(start <= end, "start {start} exceeds end {end}"); + assert!(end <= self.len, "end {end} exceeds len {}", self.len); + count_ones(self.buffer.as_slice(), self.offset + start, end - start) + } + /// Returns the position of the `nth` set bit (0-indexed). /// /// This is the "select" operation on a bitmap: given a rank `nth`, find @@ -410,6 +456,33 @@ impl BitBuffer { BitSliceIterator::new(self.buffer.as_slice(), self.offset, self.len) } + /// Invoke `f(index)` for every set bit, in ascending order, processing a `u64` + /// word at a time. + /// + /// This is the fast way to "do something for each set bit": it skips all-zero + /// words, fast-paths all-one words, and walks the remaining bits with + /// `trailing_zeros`. Prefer it over `for i in 0..len { if buf.value(i) { f(i) } }` + /// (which pays a branch per element) and over collecting [`Self::set_indices`] + /// (whose per-`next` iterator state does not inline as well). + #[inline] + pub fn for_each_set_index(&self, mut f: F) { + let mut base = 0usize; + for word in self.chunks().iter_padded() { + if word == u64::MAX { + for k in 0..64 { + f(base + k); + } + } else { + let mut w = word; + while w != 0 { + f(base + w.trailing_zeros() as usize); + w &= w - 1; + } + } + base += 64; + } + } + /// Created a new BitBuffer with offset reset to 0 pub fn sliced(&self) -> Self { if self.offset.is_multiple_of(8) { @@ -857,6 +930,75 @@ mod tests { } } + #[rstest] + #[case(0, 0)] + #[case(0, 64)] + #[case(5, 70)] + #[case(64, 130)] + #[case(0, 200)] + fn test_count_range(#[case] start: usize, #[case] end: usize) { + let len = 200; + let buf = BitBuffer::collect_bool(len, |i| i % 3 == 0); + let expected = (start..end).filter(|i| i % 3 == 0).count(); + assert_eq!(buf.count_range(start, end), expected); + // Must agree with slicing then counting. + assert_eq!( + buf.count_range(start, end), + buf.slice(start..end).true_count() + ); + } + + #[rstest] + #[case(3)] + #[case(7)] + fn test_count_range_with_offset(#[case] offset: usize) { + let len = 150; + let buf = BitBuffer::collect_bool(offset + len, |i| i % 2 == 0); + let view = BitBuffer::new_with_offset(buf.inner().clone(), len, offset); + for (start, end) in [(0, len), (10, 100), (1, 2), (63, 129)] { + let expected = (offset + start..offset + end) + .filter(|i| i % 2 == 0) + .count(); + assert_eq!(view.count_range(start, end), expected, "[{start}, {end})"); + } + } + + #[rstest] + #[case(0)] + #[case(1)] + #[case(63)] + #[case(64)] + #[case(65)] + #[case(200)] + #[case(1000)] + fn test_for_each_set_index_matches_set_indices(#[case] len: usize) { + let buf = BitBuffer::collect_bool(len, |i| i % 5 == 0 || i % 7 == 0); + let expected: Vec = buf.set_indices().collect(); + let mut got = Vec::new(); + buf.for_each_set_index(|i| got.push(i)); + assert_eq!(got, expected); + } + + #[rstest] + #[case(3, 200)] + #[case(7, 130)] + fn test_for_each_set_index_with_offset(#[case] offset: usize, #[case] len: usize) { + let base = BitBuffer::collect_bool(offset + len, |i| i % 3 == 0); + let view = BitBuffer::new_with_offset(base.inner().clone(), len, offset); + let expected: Vec = view.set_indices().collect(); + let mut got = Vec::new(); + view.for_each_set_index(|i| got.push(i)); + assert_eq!(got, expected); + } + + #[test] + fn test_for_each_set_index_all_set() { + let buf = BitBuffer::new_set(130); + let mut got = Vec::new(); + buf.for_each_set_index(|i| got.push(i)); + assert_eq!(got, (0..130).collect::>()); + } + #[test] fn test_map_cmp_conditional() { // map_cmp with conditional logic based on index and bit value diff --git a/vortex-buffer/src/bit/buf_mut.rs b/vortex-buffer/src/bit/buf_mut.rs index df38b0cd434..b3e7207a640 100644 --- a/vortex-buffer/src/bit/buf_mut.rs +++ b/vortex-buffer/src/bit/buf_mut.rs @@ -11,6 +11,7 @@ use crate::ByteBufferMut; use crate::bit::collect_bool_words; use crate::bit::get_bit_unchecked; use crate::bit::ops; +use crate::bit::pack::collect_bool_words_multiversioned; use crate::bit::set_bit_unchecked; use crate::bit::unset_bit_unchecked; use crate::buffer_mut; @@ -184,15 +185,56 @@ impl BitBufferMut { } /// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BitBufferMut` + /// + /// `f` is invoked exactly once per index, in ascending order, and the results are packed + /// with the baseline SIMD byte→bit instruction of the target. + /// + /// # Performance + /// + /// The packing is a few instructions per 64 bits, so evaluating `f` is usually the + /// bottleneck. In particular, a bounds-checked slice access in `f` (`|i| values[i] > x`) + /// blocks vectorization of the gather and can cost ~10x the packing itself. Since `f` only + /// ever sees indices `0..len`, callers reading from a slice with `len <= values.len()` may + /// soundly use `|i| unsafe { *values.get_unchecked(i) }`. + /// + /// Prefer this entry point for every predicate. Only switch to + /// [`Self::collect_bool_multiversioned`] after carefully checking that your specific `f` + /// meets its contract (a trivially cheap, bounds-check-free gather or comparison) — + /// ideally with a benchmark. #[inline] pub fn collect_bool bool>(len: usize, f: F) -> Self { + Self::collect_words(len, |words| collect_bool_words(words, len, f)) + } + + /// Like [`Self::collect_bool`], but compiles the packing loop — with `f` inside it — once + /// per CPU feature level (AVX-512BW/AVX2/baseline) and selects a clone by runtime feature + /// detection. + /// + /// Calling this asserts that `f` is small and simple enough (e.g. a bounds-check-free slice + /// gather or comparison) that duplicating it per feature level and paying a + /// `#[target_feature]` call boundary beats inlining it once into your function. For any + /// non-trivial `f` that assertion is false — the boundary deoptimizes the predicate — so + /// unless you have carefully checked (ideally benchmarked) that your specific `f` + /// qualifies, use [`Self::collect_bool`]. See + /// [`collect_bool_words_multiversioned`]. + #[inline] + pub fn collect_bool_multiversioned bool>(len: usize, f: F) -> Self { + Self::collect_words(len, |words| { + collect_bool_words_multiversioned(words, len, f) + }) + } + + /// Allocate a zero-copy word buffer for `len` bits, let `fill` populate it, and wrap it as a + /// `BitBufferMut`. + #[inline] + fn collect_words(len: usize, fill: impl FnOnce(&mut [u64])) -> Self { let num_words = len.div_ceil(64); let mut buffer: BufferMut = BufferMut::with_capacity(num_words); - // SAFETY: `collect_bool_words` writes every word in `0..num_words` below - // before any read; `u64` has no invalid bit patterns and the assignments - // inside `collect_bool_words` are pure writes. + // SAFETY: `fill` (a `collect_bool_words` variant) writes every word in `0..num_words` + // below before any read; `u64` has no invalid bit patterns and the assignments inside + // `collect_bool_words` are pure writes. unsafe { buffer.set_len(num_words) }; - collect_bool_words(buffer.as_mut_slice(), len, f); + fill(buffer.as_mut_slice()); let mut bytes = buffer.into_byte_buffer(); bytes.truncate(len.div_ceil(8)); @@ -378,6 +420,10 @@ impl BitBufferMut { return; } + assert!( + self.offset <= usize::MAX - len, + "Truncate on BitBufferMut overflowed" + ); let end_bit = self.offset + len; let new_len_bytes = end_bit.div_ceil(8); self.buffer.truncate(new_len_bytes); @@ -444,6 +490,13 @@ impl BitBufferMut { return; } + assert!( + self.offset + .checked_add(self.len) + .and_then(|v| v.checked_add(n)) + .is_some(), + "Append on BitBufferMut overflowed" + ); let end_bit_pos = self.offset + self.len + n; let required_bytes = end_bit_pos.div_ceil(8); @@ -590,14 +643,23 @@ impl Not for BitBufferMut { impl From<&[bool]> for BitBufferMut { fn from(value: &[bool]) -> Self { - BitBufferMut::collect_bool(value.len(), |i| value[i]) + // SAFETY: the predicate is invoked with indices `0..value.len()` only. + // Skipping the bounds check lets the gather loop vectorize. + BitBufferMut::collect_bool_multiversioned(value.len(), |i| unsafe { + *value.get_unchecked(i) + }) } } // allow building a buffer from a set of truthy byte values. impl From<&[u8]> for BitBufferMut { fn from(value: &[u8]) -> Self { - BitBufferMut::collect_bool(value.len(), |i| value[i] > 0) + // SAFETY: the predicate is invoked with indices `0..value.len()` only. + // Skipping the bounds check lets the gather loop vectorize. + BitBufferMut::collect_bool_multiversioned( + value.len(), + |i| unsafe { *value.get_unchecked(i) } > 0, + ) } } diff --git a/vortex-buffer/src/bit/count_ones.rs b/vortex-buffer/src/bit/count_ones.rs index df5844a2914..b977600cbe5 100644 --- a/vortex-buffer/src/bit/count_ones.rs +++ b/vortex-buffer/src/bit/count_ones.rs @@ -4,6 +4,8 @@ #[cfg(target_arch = "x86_64")] use vortex_error::VortexExpect; +use crate::dispatch::CpuKernel; + #[inline] pub fn count_ones(bytes: &[u8], offset: usize, len: usize) -> usize { if bytes.is_empty() { @@ -70,24 +72,66 @@ fn mask_byte(byte: u8, bit_offset: usize, bit_len: usize) -> u8 { shifted & mask } +type CountOnes = unsafe fn(&[u8]) -> usize; + #[inline] fn count_ones_aligned(bytes: &[u8]) -> usize { - #[cfg(target_arch = "x86_64")] - { - if bytes.len() >= 64 - && is_x86_feature_detected!("avx512f") - && is_x86_feature_detected!("avx512vpopcntdq") + // SIMD kernels only pay off from 32 bytes. Below that, call the scalar kernel + // directly: it stays inlinable and skips the dispatch indirection, which would + // otherwise dominate the couple of word popcounts. + if bytes.len() < 32 { + return count_ones_aligned_scalar(bytes); + } + + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] { - // SAFETY: Runtime detection guarantees the required target features. - return unsafe { count_ones_aligned_avx512(bytes) }; + if is_x86_feature_detected!("avx512f") + && is_x86_feature_detected!("avx512vpopcntdq") + && is_x86_feature_detected!("avx2") + { + return count_ones_aligned_avx512_any_len; + } + if is_x86_feature_detected!("avx2") { + return count_ones_aligned_avx2_any_len; + } } + count_ones_aligned_scalar + }); + // SAFETY: the selector only returns kernels that are safe or whose required CPU + // features were probed before selection. + unsafe { KERNEL.get()(bytes) } +} - if bytes.len() >= 32 && is_x86_feature_detected!("avx2") { - // SAFETY: Runtime detection guarantees the required target features. - return unsafe { count_ones_aligned_avx2(bytes) }; - } +/// Length-aware AVX-512 entry: SIMD only pays off from 64 (AVX-512) / 32 (AVX2) bytes. +/// +/// # Safety +/// Requires AVX-512F, AVX-512VPOPCNTDQ, and AVX2. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx512f,avx512vpopcntdq,avx2")] +unsafe fn count_ones_aligned_avx512_any_len(bytes: &[u8]) -> usize { + if bytes.len() >= 64 { + // SAFETY: the caller guarantees the required target features. + return unsafe { count_ones_aligned_avx512(bytes) }; + } + if bytes.len() >= 32 { + // SAFETY: every AVX-512F CPU also supports AVX2. + return unsafe { count_ones_aligned_avx2(bytes) }; } + count_ones_aligned_scalar(bytes) +} +/// Length-aware AVX2 entry: SIMD only pays off from 32 bytes. +/// +/// # Safety +/// Requires AVX2. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +unsafe fn count_ones_aligned_avx2_any_len(bytes: &[u8]) -> usize { + if bytes.len() >= 32 { + // SAFETY: the caller guarantees AVX2. + return unsafe { count_ones_aligned_avx2(bytes) }; + } count_ones_aligned_scalar(bytes) } diff --git a/vortex-buffer/src/bit/mod.rs b/vortex-buffer/src/bit/mod.rs index a39c93e3f25..46d1bee89be 100644 --- a/vortex-buffer/src/bit/mod.rs +++ b/vortex-buffer/src/bit/mod.rs @@ -14,6 +14,7 @@ mod count_ones; mod macros; mod meta; mod ops; +mod pack; mod select; mod view; @@ -27,29 +28,45 @@ pub use arrow_buffer::bit_iterator::BitSliceIterator; pub use buf::*; pub use buf_mut::*; pub use meta::*; +pub use pack::*; pub use view::*; /// Packs up to 64 boolean values into a little-endian `u64` word. +/// +/// This is [`collect_bool_words`] for a single word: a full 64-bit word is materialized as a +/// `[bool; 64]` and packed with the baseline SIMD byte→bit instruction of the target; shorter +/// lengths fall back to the bit-at-a-time [`collect_bool_word_scalar`] loop. #[inline] -pub fn collect_bool_word(len: usize, mut f: F) -> u64 +pub fn collect_bool_word(len: usize, f: F) -> u64 where F: FnMut(usize) -> bool, { assert!(len <= 64, "cannot pack {len} bits into a u64 word"); - let mut packed = 0; - for bit_idx in 0..len { - packed |= (f(bit_idx) as u64) << bit_idx; - } - packed + let mut word = [0u64; 1]; + collect_bool_words_inline(&mut word, len, f); + word[0] } /// Pack `len` boolean values returned by `f` into the prefix of `words`, LSB-first, /// 64 bits per `u64`. `words` must have capacity for at least `len.div_ceil(64)` entries. /// +/// `f` is invoked exactly once per index, in ascending order `0..len`. +/// /// Writes via `=` (not `|=`), so the destination need not be zero-initialised. +/// +/// The word loop packs with the baseline SIMD kernel of the target (SSE2 on x86-64, NEON on +/// aarch64), which inlines fully into the caller together with the predicate and the +/// `[bool; 64]` materialization — wider kernels would sit behind a non-inlinable +/// `#[target_feature]` boundary that deoptimizes expensive predicates. See +/// [`BitBuffer::collect_bool`] for the performance note on +/// avoiding bounds checks in `f`. +/// +/// Prefer this entry point for every predicate; only switch to +/// [`collect_bool_words_multiversioned`] after carefully checking that your specific `f` +/// meets its contract. #[inline] -pub fn collect_bool_words(words: &mut [u64], len: usize, mut f: F) +pub fn collect_bool_words(words: &mut [u64], len: usize, f: F) where F: FnMut(usize) -> bool, { @@ -60,18 +77,7 @@ where words.len(), ); - let full = len / 64; - let remainder = len % 64; - - for word_idx in 0..full { - let offset = word_idx * 64; - words[word_idx] = collect_bool_word(64, |bit_idx| f(offset + bit_idx)); - } - - if remainder != 0 { - let offset = full * 64; - words[full] = collect_bool_word(remainder, |bit_idx| f(offset + bit_idx)); - } + collect_bool_words_inline(words, len, f) } /// Read up to 8 bytes as a little-endian `u64`, zero-padding the high bytes when fewer than 8 diff --git a/vortex-buffer/src/bit/pack.rs b/vortex-buffer/src/bit/pack.rs new file mode 100644 index 00000000000..565ce1e8345 --- /dev/null +++ b/vortex-buffer/src/bit/pack.rs @@ -0,0 +1,521 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Kernels for packing boolean values into bitmap words. +//! +//! `collect_bool` materializes each full 64-bit chunk as a `[bool; 64]` (a loop the +//! auto-vectorizer turns into vector stores for simple predicates) and then packs the 64 bytes +//! into a single `u64` with a byte→bit kernel: +//! +//! - x86-64 AVX-512BW: one `vptestmb` produces the full 64-bit mask. +//! - x86-64 AVX2: two `vpmovmskb` halves. +//! - x86-64 SSE2 (baseline): four `pmovmskb` quarters. +//! - aarch64 NEON (baseline): per-lane `ushl` by the bit position, then an `addp` reduction tree. +//! - elsewhere (and under Miri): a branch-free SWAR multiply. +//! +//! There are two tiers. The default ([`collect_bool_words_inline`]) compiles the loop once +//! with the widest *statically-enabled* kernel (SSE2 / NEON / SWAR on stock targets) and +//! inlines fully into the caller — safe for arbitrary predicates. The opt-in tier +//! ([`collect_bool_words_multiversioned`]) compiles the loop — with `f` inside — once per CPU +//! feature level and selects a clone by runtime detection; only for predicates small and +//! simple enough that the per-level duplication and its `#[target_feature]` call boundary pay +//! off. +//! +//! The bit-at-a-time loop lives on as [`collect_bool_word_scalar`], used for tail chunks and as +//! the reference implementation for tests and benchmarks. + +/// Packs up to 64 boolean values into a little-endian `u64` word one bit at a time. +/// +/// This is the scalar reference implementation behind +/// [`collect_bool_word`](crate::bit::collect_bool_word); prefer calling that entry point, which +/// takes the SIMD fast path for full 64-bit words. +#[inline] +pub fn collect_bool_word_scalar(len: usize, mut f: F) -> u64 +where + F: FnMut(usize) -> bool, +{ + assert!(len <= 64, "cannot pack {len} bits into a u64 word"); + + let mut packed = 0; + for bit_idx in 0..len { + packed |= (f(bit_idx) as u64) << bit_idx; + } + packed +} + +/// Body of [`collect_bool_words`](crate::bit::collect_bool_words) (and, via a one-word slice, +/// of [`collect_bool_word`](crate::bit::collect_bool_word)): the word loop with the widest +/// pack kernel enabled *at compile time* — SSE2 on stock x86-64 (AVX2/AVX-512BW when built +/// with e.g. `-C target-cpu=native`), NEON on aarch64, SWAR elsewhere and under Miri. +/// +/// Statically-enabled kernels are part of every function's feature set, so this loop +/// (predicate, the `[bool; 64]` materialization, and the pack) inlines fully into the caller +/// with no `#[target_feature]` boundary. That boundary is why *runtime*-detected wider kernels +/// are not used here: hiding an expensive, non-vectorizable predicate (e.g. FSST's per-string +/// DFA scan) behind a non-inlinable AVX-512 loop copy costs far more (~30% end to end) than +/// the wider pack saves — and an indirect call per word is worse still (~4x on cheap +/// predicates), since an opaque call target blocks fill/pack fusion regardless of how cheap +/// the kernel *selection* is. For provably cheap predicates, use [`collect_bool_words_multiversioned`]. +#[inline(always)] +pub(crate) fn collect_bool_words_inline(words: &mut [u64], len: usize, f: F) +where + F: FnMut(usize) -> bool, +{ + #[cfg(all( + target_arch = "x86_64", + target_feature = "avx512f", + target_feature = "avx512bw", + not(miri) + ))] + { + // SAFETY: AVX-512F/BW are statically enabled for this build (e.g. -C + // target-cpu=native), so they are in every function's feature set and the kernel + // inlines here like any other function. + collect_bool_words_with(words, len, f, |bools| unsafe { + pack_bool_word_avx512(bools) + }) + } + #[cfg(all( + target_arch = "x86_64", + target_feature = "avx2", + not(all(target_feature = "avx512f", target_feature = "avx512bw")), + not(miri) + ))] + { + // SAFETY: AVX2 is statically enabled for this build. + collect_bool_words_with(words, len, f, |bools| unsafe { pack_bool_word_avx2(bools) }) + } + #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2"), not(miri)))] + { + // SAFETY: SSE2 is part of the x86-64 baseline instruction set. + collect_bool_words_with(words, len, f, |bools| unsafe { pack_bool_word_sse2(bools) }) + } + #[cfg(all(target_arch = "aarch64", not(miri)))] + { + // SAFETY: NEON is part of the aarch64 baseline instruction set. + collect_bool_words_with(words, len, f, |bools| unsafe { pack_bool_word_neon(bools) }) + } + #[cfg(any(not(any(target_arch = "x86_64", target_arch = "aarch64")), miri))] + collect_bool_words_with(words, len, f, pack_bool_word_swar) +} + +/// Word loop with the *widest* pack kernel the CPU offers (AVX-512BW, then AVX2, then the +/// baseline), for predicates known to be cheap. +/// +/// The wide loop copies live behind a `#[target_feature]` function boundary that cannot inline +/// into the caller, which deoptimizes expensive predicates (see the module docs and +/// `collect_bool_words_inline`). Only route a predicate here when its evaluation is trivial +/// — e.g. the bounds-check-free slice gathers in the `From<&[bool]>` / `From<&[u8]>` +/// conversions, or unchecked slice comparisons like the `between` kernels — where the fused +/// AVX-512 loop is worth another ~2x over the baseline kernel. +/// +/// `words` must hold at least `len.div_ceil(64)` entries and `f` is invoked with `0..len`, +/// exactly once per index in ascending order. +/// +/// Panics if `words` is too short. +#[inline] +pub fn collect_bool_words_multiversioned(words: &mut [u64], len: usize, f: F) +where + F: FnMut(usize) -> bool, +{ + let num_words = len.div_ceil(64); + assert!( + words.len() >= num_words, + "words slice has {} entries, need at least {num_words}", + words.len(), + ); + + // Without a full 64-bit word only the scalar tail would run; skip feature detection and + // the `#[target_feature]` call boundary entirely. + if len < 64 { + return collect_bool_words_inline(words, len, f); + } + + #[cfg(all(target_arch = "x86_64", not(miri)))] + { + if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") { + // SAFETY: runtime detection guarantees the required target features. + return unsafe { collect_bool_words_avx512(words, len, f) }; + } + if is_x86_feature_detected!("avx2") { + // SAFETY: runtime detection guarantees the required target features. + return unsafe { collect_bool_words_avx2(words, len, f) }; + } + } + collect_bool_words_inline(words, len, f) +} + +/// Shared word loop: materialize each full 64-bit chunk as a `[bool; 64]` and pack it with +/// `pack`; the tail chunk goes through the scalar loop. +/// +/// Marked `#[inline(always)]` so each `#[target_feature]` wrapper gets its own fully-inlined +/// copy compiled with that feature set. +#[inline(always)] +fn collect_bool_words_with(words: &mut [u64], len: usize, mut f: F, pack: P) +where + F: FnMut(usize) -> bool, + P: Fn(&[bool; 64]) -> u64, +{ + let full = len / 64; + let remainder = len % 64; + + for (word_idx, word) in words[..full].iter_mut().enumerate() { + let offset = word_idx * 64; + let mut bools = [false; 64]; + for (bit_idx, b) in bools.iter_mut().enumerate() { + *b = f(offset + bit_idx); + } + *word = pack(&bools); + } + + if remainder != 0 { + let offset = full * 64; + words[full] = collect_bool_word_scalar(remainder, |bit_idx| f(offset + bit_idx)); + } +} + +/// SSE2 copy of the [`collect_bool_words`](crate::bit::collect_bool_words) word loop. +/// +/// `words` must hold at least `len.div_ceil(64)` entries and `f` is invoked with `0..len`. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports SSE2 (always true on x86-64). +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "sse2")] +pub unsafe fn collect_bool_words_sse2 bool>( + words: &mut [u64], + len: usize, + f: F, +) { + // SAFETY: the caller guarantees SSE2 support. + collect_bool_words_with(words, len, f, |bools| unsafe { pack_bool_word_sse2(bools) }) +} + +/// AVX2 copy of the [`collect_bool_words`](crate::bit::collect_bool_words) word loop. +/// +/// `words` must hold at least `len.div_ceil(64)` entries and `f` is invoked with `0..len`. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports AVX2. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +pub unsafe fn collect_bool_words_avx2 bool>( + words: &mut [u64], + len: usize, + f: F, +) { + // SAFETY: the caller guarantees AVX2 support. + collect_bool_words_with(words, len, f, |bools| unsafe { pack_bool_word_avx2(bools) }) +} + +/// AVX-512BW copy of the [`collect_bool_words`](crate::bit::collect_bool_words) word loop. +/// +/// `words` must hold at least `len.div_ceil(64)` entries and `f` is invoked with `0..len`. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports AVX-512F and AVX-512BW. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx512f,avx512bw")] +pub unsafe fn collect_bool_words_avx512 bool>( + words: &mut [u64], + len: usize, + f: F, +) { + // SAFETY: the caller guarantees AVX-512F and AVX-512BW support. + collect_bool_words_with(words, len, f, |bools| unsafe { + pack_bool_word_avx512(bools) + }) +} + +/// NEON copy of the [`collect_bool_words`](crate::bit::collect_bool_words) word loop. +/// +/// `words` must hold at least `len.div_ceil(64)` entries and `f` is invoked with `0..len`. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports NEON (always true on aarch64). +#[cfg(target_arch = "aarch64")] +#[target_feature(enable = "neon")] +pub unsafe fn collect_bool_words_neon bool>( + words: &mut [u64], + len: usize, + f: F, +) { + // SAFETY: the caller guarantees NEON support. + collect_bool_words_with(words, len, f, |bools| unsafe { pack_bool_word_neon(bools) }) +} + +/// Portable branch-free byte→bit pack kernel, used when no SIMD kernel is available. +/// +/// Reads the bools eight at a time as a `u64` and gathers the eight `0x00`/`0x01` bytes into +/// eight contiguous bits with a single multiply: byte `i` contributes `2^(8i)`, and the magic +/// constant `0x0102_0408_1020_4080 = Σ 2^(56-7i)` shifts each contribution to bit `56 + i` +/// without any cross-term collisions, so the mask falls out of the top byte of the product. +#[inline] +pub fn pack_bool_word_swar(bools: &[bool; 64]) -> u64 { + const MAGIC: u64 = 0x0102_0408_1020_4080; + + let (chunks, rest) = bools.as_chunks::<8>(); + debug_assert!(rest.is_empty()); + + let mut packed = 0u64; + for (chunk_idx, chunk) in chunks.iter().enumerate() { + let word = u64::from_le_bytes(chunk.map(|b| b as u8)); + packed |= (word.wrapping_mul(MAGIC) >> 56) << (8 * chunk_idx); + } + packed +} + +/// SSE2 byte→bit pack kernel: four 16-byte `pcmpeqb`-against-zero + `pmovmskb` rounds. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports SSE2 (always true on x86-64). +#[cfg(target_arch = "x86_64")] +#[inline] +#[target_feature(enable = "sse2")] +pub unsafe fn pack_bool_word_sse2(bools: &[bool; 64]) -> u64 { + use std::arch::x86_64::__m128i; + use std::arch::x86_64::_mm_cmpeq_epi8; + use std::arch::x86_64::_mm_loadu_si128; + use std::arch::x86_64::_mm_movemask_epi8; + use std::arch::x86_64::_mm_setzero_si128; + + let ptr = bools.as_ptr().cast::(); + let zero = _mm_setzero_si128(); + + let mut packed = 0u64; + for lane in 0..4 { + // SAFETY: `lane * 16 + 16 <= 64`, so the 16-byte load is in bounds. + let chunk = unsafe { _mm_loadu_si128(ptr.add(lane * 16).cast::<__m128i>()) }; + // `cmpeq` against zero sets 0xFF for *false* bytes; invert to get the truthy mask. + let zero_mask = _mm_movemask_epi8(_mm_cmpeq_epi8(chunk, zero)) as u32 as u64; + packed |= (!zero_mask & 0xFFFF) << (16 * lane); + } + packed +} + +/// AVX2 byte→bit pack kernel: two 32-byte `vpcmpeqb`-against-zero + `vpmovmskb` rounds. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports AVX2. +#[cfg(target_arch = "x86_64")] +#[inline] +#[target_feature(enable = "avx2")] +pub unsafe fn pack_bool_word_avx2(bools: &[bool; 64]) -> u64 { + use std::arch::x86_64::__m256i; + use std::arch::x86_64::_mm256_cmpeq_epi8; + use std::arch::x86_64::_mm256_loadu_si256; + use std::arch::x86_64::_mm256_movemask_epi8; + use std::arch::x86_64::_mm256_setzero_si256; + + let ptr = bools.as_ptr().cast::(); + let zero = _mm256_setzero_si256(); + + // SAFETY: both 32-byte loads are within the 64-byte array. + let lo = unsafe { _mm256_loadu_si256(ptr.cast::<__m256i>()) }; + // SAFETY: see above. + let hi = unsafe { _mm256_loadu_si256(ptr.add(32).cast::<__m256i>()) }; + + // `cmpeq` against zero sets 0xFF for *false* bytes; invert to get the truthy mask. + let lo_mask = !(_mm256_movemask_epi8(_mm256_cmpeq_epi8(lo, zero)) as u32) as u64; + let hi_mask = !(_mm256_movemask_epi8(_mm256_cmpeq_epi8(hi, zero)) as u32) as u64; + lo_mask | (hi_mask << 32) +} + +/// AVX-512BW byte→bit pack kernel: a single 64-byte `vptestmb` produces the whole word. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports AVX-512F and AVX-512BW. +#[cfg(target_arch = "x86_64")] +#[inline] +#[target_feature(enable = "avx512f,avx512bw")] +pub unsafe fn pack_bool_word_avx512(bools: &[bool; 64]) -> u64 { + use std::arch::x86_64::__m512i; + use std::arch::x86_64::_mm512_loadu_si512; + use std::arch::x86_64::_mm512_test_epi8_mask; + + // SAFETY: the 64-byte load covers exactly the `[bool; 64]` array. + let chunk = unsafe { _mm512_loadu_si512(bools.as_ptr().cast::<__m512i>()) }; + // Mask bit `i` is set iff byte `i` AND byte `i` is nonzero, i.e. iff `bools[i]`. + _mm512_test_epi8_mask(chunk, chunk) +} + +/// NEON byte→bit pack kernel: shift each `0x00`/`0x01` byte left by its bit position +/// (`ushl`), then fold the four vectors into one `u64` with a pairwise-add (`addp`) tree. +/// +/// # Safety +/// +/// The caller must ensure the CPU supports NEON (always true on aarch64). +#[cfg(target_arch = "aarch64")] +#[inline] +#[target_feature(enable = "neon")] +pub unsafe fn pack_bool_word_neon(bools: &[bool; 64]) -> u64 { + use std::arch::aarch64::vgetq_lane_u64; + use std::arch::aarch64::vld1q_s8; + use std::arch::aarch64::vld1q_u8; + use std::arch::aarch64::vpaddq_u8; + use std::arch::aarch64::vreinterpretq_u64_u8; + use std::arch::aarch64::vshlq_u8; + + const BIT_SHIFTS: [i8; 16] = [0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7]; + + let ptr = bools.as_ptr().cast::(); + // SAFETY: loading 16 constant bytes from `BIT_SHIFTS`; the four 16-byte data loads below are + // all within the 64-byte array. + unsafe { + let shifts = vld1q_s8(BIT_SHIFTS.as_ptr()); + + // Byte j of each vector becomes `bools[16v + j] << (j % 8)`. + let m0 = vshlq_u8(vld1q_u8(ptr), shifts); + let m1 = vshlq_u8(vld1q_u8(ptr.add(16)), shifts); + let m2 = vshlq_u8(vld1q_u8(ptr.add(32)), shifts); + let m3 = vshlq_u8(vld1q_u8(ptr.add(48)), shifts); + + // Three rounds of pairwise adds sum each group of 8 weighted bytes into one mask byte, + // yielding the 8 mask bytes in order in the low half of the final vector. + let sum01 = vpaddq_u8(m0, m1); + let sum23 = vpaddq_u8(m2, m3); + let sum = vpaddq_u8(sum01, sum23); + let sum = vpaddq_u8(sum, sum); + vgetq_lane_u64::<0>(vreinterpretq_u64_u8(sum)) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::collect_bool_word_scalar; + use super::pack_bool_word_swar; + + fn patterns() -> Vec<[bool; 64]> { + let mut patterns = vec![ + [false; 64], + [true; 64], + std::array::from_fn(|i| i % 2 == 0), + std::array::from_fn(|i| i % 3 == 0), + std::array::from_fn(|i| i < 32), + std::array::from_fn(|i| i == 0 || i == 63), + ]; + // A few deterministic pseudo-random patterns. + let mut state = 0x9E37_79B9_7F4A_7C15u64; + for _ in 0..8 { + patterns.push(std::array::from_fn(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (state >> 33) & 1 == 1 + })); + } + patterns + } + + fn reference(bools: &[bool; 64]) -> u64 { + collect_bool_word_scalar(64, |i| bools[i]) + } + + #[test] + fn swar_matches_scalar() { + for bools in patterns() { + assert_eq!(pack_bool_word_swar(&bools), reference(&bools)); + } + } + + #[test] + fn dispatch_matches_scalar() { + for bools in patterns() { + assert_eq!( + crate::bit::collect_bool_word(64, |i| bools[i]), + reference(&bools) + ); + } + } + + #[cfg(all(target_arch = "x86_64", not(miri)))] + #[test] + fn sse2_matches_scalar() { + for bools in patterns() { + // SAFETY: SSE2 is part of the x86-64 baseline instruction set. + assert_eq!( + unsafe { super::pack_bool_word_sse2(&bools) }, + reference(&bools) + ); + } + } + + #[cfg(all(target_arch = "x86_64", not(miri)))] + #[test] + fn avx2_matches_scalar() { + if !is_x86_feature_detected!("avx2") { + return; + } + for bools in patterns() { + // SAFETY: runtime detection guarantees AVX2. + assert_eq!( + unsafe { super::pack_bool_word_avx2(&bools) }, + reference(&bools) + ); + } + } + + #[cfg(all(target_arch = "x86_64", not(miri)))] + #[test] + fn avx512_matches_scalar() { + if !(is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw")) { + return; + } + for bools in patterns() { + // SAFETY: runtime detection guarantees AVX-512F and AVX-512BW. + assert_eq!( + unsafe { super::pack_bool_word_avx512(&bools) }, + reference(&bools) + ); + } + } + + #[cfg(all(target_arch = "aarch64", not(miri)))] + #[test] + fn neon_matches_scalar() { + for bools in patterns() { + // SAFETY: NEON is part of the aarch64 baseline instruction set. + assert_eq!( + unsafe { super::pack_bool_word_neon(&bools) }, + reference(&bools) + ); + } + } + + #[rstest] + #[case(0)] + #[case(1)] + #[case(63)] + #[case(64)] + #[case(65)] + #[case(200)] + fn multiversioned_matches_inline(#[case] len: usize) { + let pattern = |i: usize| i.is_multiple_of(3) || i.is_multiple_of(7); + let num_words = len.div_ceil(64); + let mut multiversioned = vec![0u64; num_words]; + super::collect_bool_words_multiversioned(&mut multiversioned, len, pattern); + let mut inline = vec![0u64; num_words]; + super::collect_bool_words_inline(&mut inline, len, pattern); + assert_eq!(multiversioned, inline); + } + + #[rstest] + #[case(0)] + #[case(1)] + #[case(5)] + #[case(63)] + #[case(64)] + fn collect_bool_word_partial_lens_match(#[case] len: usize) { + let expected = collect_bool_word_scalar(len, |i| i % 3 == 0); + assert_eq!(crate::bit::collect_bool_word(len, |i| i % 3 == 0), expected); + } +} diff --git a/vortex-buffer/src/bit/select.rs b/vortex-buffer/src/bit/select.rs index 53eded23dd4..deae368ddfb 100644 --- a/vortex-buffer/src/bit/select.rs +++ b/vortex-buffer/src/bit/select.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use super::count_ones::align_offset_len; +use crate::dispatch::CpuKernel; /// Returns the position of the `nth` set bit (0-indexed) within the logical range /// `[offset, offset + len)` of the given byte slice. @@ -82,32 +83,47 @@ pub fn bit_select(bytes: &[u8], offset: usize, len: usize, nth: usize) -> Option /// /// If `chunk_index < chunks.len()`, the target bit is inside that chunk and `remaining` /// is the rank *within* that chunk. Otherwise all chunks were consumed. -#[inline] -fn scan_chunks(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usize, usize) { - scan_chunks_impl(chunks, remaining, pos) -} +type ScanChunks = unsafe fn(&[[u8; 64]], usize, usize) -> (usize, usize, usize); -#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))] -#[inline] -fn scan_chunks_impl(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usize, usize) { - scan_chunks_scalar(chunks, remaining, pos) -} - -#[cfg(target_arch = "x86_64")] #[inline] -fn scan_chunks_impl(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usize, usize) { - if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512vpopcntdq") { - // SAFETY: runtime detection guarantees the required target features. - return unsafe { scan_chunks_avx512_vpopcnt(chunks, remaining, pos) }; +fn scan_chunks(chunks: &[[u8; 64]], remaining: usize, pos: usize) -> (usize, usize, usize) { + // Scans of a couple of chunks don't amortize the dispatch indirection: call the + // per-architecture unconditional kernel directly so it stays inlinable (see the + // size-gating note in the CpuKernel docs). + if chunks.len() <= 2 { + #[cfg(target_arch = "aarch64")] + return scan_chunks_neon(chunks, remaining, pos); + #[allow(unreachable_code)] + { + return scan_chunks_scalar(chunks, remaining, pos); + } } - scan_chunks_scalar(chunks, remaining, pos) + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512vpopcntdq") { + return scan_chunks_avx512_vpopcnt; + } + } + #[cfg(target_arch = "aarch64")] + return scan_chunks_neon; + // The aarch64 arm above returns unconditionally (NEON needs no probe), making + // this portable default unreachable there. + #[allow(unreachable_code)] + { + scan_chunks_scalar + } + }); + // SAFETY: the selector only returns kernels that are safe or whose required CPU + // features were probed before selection. + unsafe { KERNEL.get()(chunks, remaining, pos) } } #[cfg(target_arch = "aarch64")] #[allow(clippy::cast_possible_truncation)] // u64 → usize is lossless on aarch64 (64-bit) #[inline] -fn scan_chunks_impl( +fn scan_chunks_neon( chunks: &[[u8; 64]], mut remaining: usize, mut pos: usize, @@ -185,7 +201,6 @@ unsafe fn scan_chunks_avx512_vpopcnt( (remaining, pos, chunks.len()) } -#[cfg(not(target_arch = "aarch64"))] #[inline] fn scan_chunks_scalar( chunks: &[[u8; 64]], @@ -280,21 +295,26 @@ fn scan_words_scalar( // ── In-chunk select ───────────────────────────────────────────────────── +type SelectInChunk = unsafe fn(&[u8; 64], usize) -> usize; + /// Position of the `nth` set bit inside a 64-byte chunk (0-indexed). #[inline] fn select_in_chunk(chunk: &[u8; 64], nth: usize) -> usize { - #[cfg(target_arch = "x86_64")] - { - if is_x86_feature_detected!("avx512f") - && is_x86_feature_detected!("avx512vpopcntdq") - && is_x86_feature_detected!("avx512vbmi2") + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] { - // SAFETY: runtime detection guarantees the required target features. - return unsafe { select_in_chunk_vbmi2(chunk, nth) }; + if is_x86_feature_detected!("avx512f") + && is_x86_feature_detected!("avx512vpopcntdq") + && is_x86_feature_detected!("avx512vbmi2") + { + return select_in_chunk_vbmi2; + } } - } - - select_in_chunk_scalar(chunk, nth) + select_in_chunk_scalar + }); + // SAFETY: the selector only returns kernels that are safe or whose required CPU + // features were probed before selection. + unsafe { KERNEL.get()(chunk, nth) } } #[cfg(target_arch = "x86_64")] @@ -343,7 +363,6 @@ fn select_in_chunk_scalar(chunk: &[u8; 64], mut nth: usize) -> usize { unreachable!("select_in_chunk: nth exceeds popcount") } -#[cfg(not(target_arch = "aarch64"))] #[inline] fn count_ones_chunk(chunk: &[u8; 64]) -> usize { let words = chunk.as_chunks::<8>().0; @@ -359,17 +378,23 @@ fn count_ones_chunk(chunk: &[u8; 64]) -> usize { // ── In-word select ────────────────────────────────────────────────────── +type SelectInWord = unsafe fn(u64, usize) -> usize; + /// Position of the `nth` set bit inside a u64 (0-indexed, little-endian bit order). #[inline] fn select_in_word(word: u64, nth: usize) -> usize { - #[cfg(target_arch = "x86_64")] - { - if is_x86_feature_detected!("bmi2") { - // SAFETY: runtime detection guarantees the required target feature. - return unsafe { select_in_word_bmi2(word, nth) }; + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("bmi2") { + return select_in_word_bmi2; + } } - } - select_in_word_scalar(word, nth) + select_in_word_scalar + }); + // SAFETY: the selector only returns kernels that are safe or whose required CPU + // features were probed before selection. + unsafe { KERNEL.get()(word, nth) } } /// BMI2: deposit a single bit at the nth set-bit position, then count trailing zeros. diff --git a/vortex-buffer/src/dispatch.rs b/vortex-buffer/src/dispatch.rs new file mode 100644 index 00000000000..6040bf0c885 --- /dev/null +++ b/vortex-buffer/src/dispatch.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! One-time CPU-feature-based function dispatch. +//! +//! [`CpuKernel`] holds a function pointer chosen by a *selector* exactly once, on the +//! first call. Every later call is a relaxed atomic load, a never-taken predicted +//! branch, and an indirect call — measured indistinguishable from a direct call (see +//! `benches/cpu_dispatch.rs`). +//! +//! The selector is ordinary code that models both dispatch dimensions: +//! +//! * **Compile time** (x86_64 vs aarch64): one `#[cfg(target_arch = ...)]` block per +//! architecture, each *returning early* with its chosen kernel. +//! * **Runtime** (AVX-512 vs BMI2 vs ...): an if-chain of feature probes inside that +//! block. +//! * The **portable default** is the plain tail expression. Because the architecture +//! arms return early, no `#[cfg(not(any(...)))]` negation is ever needed. An arm +//! that returns unconditionally (e.g. NEON, architecturally guaranteed on aarch64) +//! makes the tail unreachable on that architecture — wrap just the tail in an +//! `#[allow(unreachable_code)]` block. +//! +//! When kernels are `unsafe` `#[target_feature]` functions, make `F` an +//! `unsafe fn(...)` pointer type: the bare kernel names then coerce directly (no +//! closure wrappers), and the one dispatched call is wrapped in `unsafe` with a +//! SAFETY comment stating that the selector probed the required features. +//! +//! Races are benign: the slot only ever holds valid function pointers of the same +//! type, and every candidate must compute the same result. +//! +//! # When NOT to use it +//! +//! Do not put the dispatched call inside a per-element hot loop: the indirect call +//! blocks inlining. Hoist the decision to the outermost per-buffer entry point and +//! monomorphize the loop instead, like the `Bmi2`/`Portable` type-parameter pattern in +//! `vortex-mask::intersect_by_rank`. +//! +//! For the same reason, gate on input size *before* [`get`](CpuKernel::get) when tiny +//! inputs are common: call the portable kernel directly below the size where SIMD pays +//! off, so those calls stay inlinable and skip the dispatch entirely (see +//! `count_ones_aligned`). + +use core::mem::transmute_copy; +use core::ptr; +use core::sync::atomic::AtomicPtr; +use core::sync::atomic::Ordering; + +/// A function pointer selected by CPU-feature detection once, on first use. +/// +/// `F` must be a plain function-pointer type (`fn(...) -> ...`). The selector passed +/// to [`new`](Self::new) runs at most once per process (once per racing thread in the +/// worst case), on the first [`get`](Self::get), and its result is cached. +/// Non-capturing closures coerce to function pointers, so the selector can be written +/// inline in the `static`, like `LazyLock`. +/// +/// # Example +/// +/// ``` +/// use vortex_buffer::CpuKernel; +/// +/// /// Sums a slice, using the best kernel for the current CPU. +/// fn sum(values: &[u64]) -> u64 { +/// static KERNEL: CpuKernel u64> = CpuKernel::new(|| { +/// // Compile-time arm per architecture; runtime probes inside it. +/// #[cfg(target_arch = "x86_64")] +/// { +/// if std::arch::is_x86_feature_detected!("avx2") { +/// // return |values| unsafe { sum_avx2(values) }; +/// } +/// } +/// // Portable default: plain tail, no cfg(not(...)) needed. +/// |values| values.iter().sum() +/// }); +/// KERNEL.get()(values) +/// } +/// +/// assert_eq!(sum(&[1, 2, 3]), 6); +/// ``` +pub struct CpuKernel { + selected: AtomicPtr<()>, + select: fn() -> F, +} + +impl CpuKernel { + /// Create a kernel slot whose kernel is chosen by `select` on the first + /// [`get`](Self::get). + pub const fn new(select: fn() -> F) -> Self { + assert!( + size_of::() == size_of::<*mut ()>(), + "CpuKernel requires a function-pointer type" + ); + Self { + selected: AtomicPtr::new(ptr::null_mut()), + select, + } + } + + /// Return the selected kernel, running the selector on the first call. + /// + /// Steady state is a relaxed load plus a never-taken predicted branch. + #[inline] + pub fn get(&self) -> F { + let fn_ptr = self.selected.load(Ordering::Relaxed); + if fn_ptr.is_null() { + return self.select_slow(); + } + // SAFETY: non-null values in `selected` are always the bits of an `F` stored + // by `select_slow`, and `F` is pointer-sized (asserted in `new`). Function + // pointers are never null, so the null sentinel stays unambiguous. + unsafe { transmute_copy::<*mut (), F>(&fn_ptr) } + } + + #[cold] + fn select_slow(&self) -> F { + let kernel = (self.select)(); + // SAFETY: `F` is pointer-sized (asserted in `new`); a bitwise copy into a raw + // pointer preserves the function pointer for the transmute back in `get`. + let fn_ptr = unsafe { transmute_copy::(&kernel) }; + self.selected.store(fn_ptr, Ordering::Relaxed); + kernel + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use super::CpuKernel; + + static SELECT_CALLS: AtomicUsize = AtomicUsize::new(0); + + fn add_one(x: u64) -> u64 { + static KERNEL: CpuKernel u64> = CpuKernel::new(|| { + SELECT_CALLS.fetch_add(1, Ordering::Relaxed); + |x| x + 1 + }); + KERNEL.get()(x) + } + + #[test] + fn selects_once_then_dispatches() { + assert_eq!(add_one(41), 42); + assert_eq!(add_one(1), 2); + assert_eq!(add_one(2), 3); + assert_eq!(SELECT_CALLS.load(Ordering::Relaxed), 1); + } + + #[test] + fn selector_early_returns_skip_the_default() { + static KERNEL: CpuKernel u64> = CpuKernel::new(|| { + if 1 + 1 == 2 { + return |x| x + 2; + } + |x| x + 100 + }); + assert_eq!(KERNEL.get()(0), 2); + } + + type XorFn = fn(&[u8; 4], &mut [u8; 4]); + + #[test] + fn supports_reference_arguments() { + static KERNEL: CpuKernel = CpuKernel::new(|| { + |input, output| { + for (o, i) in output.iter_mut().zip(input) { + *o ^= *i; + } + } + }); + let selected = KERNEL.get(); + let input = [1, 2, 3, 4]; + let mut output = [0, 0, 0, 0]; + selected(&input, &mut output); + assert_eq!(output, input); + } +} diff --git a/vortex-buffer/src/lib.rs b/vortex-buffer/src/lib.rs index 8319fffa387..ee113481353 100644 --- a/vortex-buffer/src/lib.rs +++ b/vortex-buffer/src/lib.rs @@ -52,6 +52,7 @@ pub use buffer::*; pub use buffer_mut::*; pub use bytes::*; pub use r#const::*; +pub use dispatch::*; pub use string::*; mod alignment; #[cfg(feature = "arrow")] @@ -62,6 +63,7 @@ mod buffer_mut; mod bytes; mod r#const; mod debug; +mod dispatch; mod macros; #[cfg(feature = "memmap2")] mod memmap2; diff --git a/vortex-compressor/Cargo.toml b/vortex-compressor/Cargo.toml index 2d28d86a5e6..24a2d8e40d5 100644 --- a/vortex-compressor/Cargo.toml +++ b/vortex-compressor/Cargo.toml @@ -28,8 +28,6 @@ vortex-utils = { workspace = true } [dev-dependencies] divan = { workspace = true } -rstest = { workspace = true } -tracing-subscriber = { workspace = true, features = ["env-filter"] } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-session = { workspace = true } diff --git a/vortex-compressor/benches/dict_encode.rs b/vortex-compressor/benches/dict_encode.rs index 2c4e24108a7..ec6343b6f40 100644 --- a/vortex-compressor/benches/dict_encode.rs +++ b/vortex-compressor/benches/dict_encode.rs @@ -8,6 +8,7 @@ use std::sync::LazyLock; use divan::Bencher; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::builders::dict::dict_encode; @@ -15,7 +16,7 @@ use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); fn make_array() -> PrimitiveArray { let values: BufferMut = (0..50).cycle().take(64_000).collect(); @@ -41,5 +42,6 @@ fn encode_generic(bencher: Bencher) { } fn main() { + LazyLock::force(&SESSION); divan::main() } diff --git a/vortex-compressor/src/builtins/constant/binary.rs b/vortex-compressor/src/builtins/constant/binary.rs deleted file mode 100644 index 4ad24fe57b5..00000000000 --- a/vortex-compressor/src/builtins/constant/binary.rs +++ /dev/null @@ -1,84 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Constant encoding for binary arrays. - -use vortex_array::ArrayRef; -use vortex_array::Canonical; -use vortex_array::ExecutionCtx; -use vortex_array::aggregate_fn::fns::is_constant::is_constant; -use vortex_error::VortexResult; - -use crate::CascadingCompressor; -use crate::builtins::constant::compress_constant_array_with_validity; -use crate::ctx::CompressorContext; -use crate::estimate::CompressionEstimate; -use crate::estimate::DeferredEstimate; -use crate::estimate::EstimateVerdict; -use crate::scheme::Scheme; -use crate::stats::ArrayAndStats; - -/// Constant encoding for binary arrays with a single distinct value. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct BinaryConstantScheme; - -impl Scheme for BinaryConstantScheme { - fn scheme_name(&self) -> &'static str { - "vortex.binary.constant" - } - - fn matches(&self, canonical: &Canonical) -> bool { - canonical.dtype().is_binary() - } - - fn expected_compression_ratio( - &self, - data: &ArrayAndStats, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - // Constant detection on a sample is a false positive, since the sample being constant does - // not mean the full array is constant. - if compress_ctx.is_sample() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - let array_len = data.array().len(); - let stats = data.varbinview_stats(exec_ctx); - - // We want to use `Constant` if there are only nulls in the array. - if stats.value_count() == 0 { - debug_assert_eq!(stats.null_count() as usize, array_len); - return CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse); - } - - // Since the estimated distinct count is always going to be less than or equal to the actual - // distinct count, if this is not equal to 1 the actual is definitely not equal to 1. - if stats.estimated_distinct_count().is_some_and(|c| c > 1) { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - // Otherwise our best bet is to actually check if the array is constant. - // This is an expensive check, but the alternative of not compressing a constant array is - // far less preferable. - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, data, _best_so_far, _ctx, exec_ctx| { - if is_constant(data.array(), exec_ctx)? { - Ok(EstimateVerdict::AlwaysUse) - } else { - Ok(EstimateVerdict::Skip) - } - }, - ))) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - data: &ArrayAndStats, - _compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - compress_constant_array_with_validity(data.array(), exec_ctx) - } -} diff --git a/vortex-compressor/src/builtins/constant/bool.rs b/vortex-compressor/src/builtins/constant/bool.rs deleted file mode 100644 index a3bdcb0216e..00000000000 --- a/vortex-compressor/src/builtins/constant/bool.rs +++ /dev/null @@ -1,69 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Constant encoding for bool arrays. - -use vortex_array::ArrayRef; -use vortex_array::Canonical; -use vortex_array::ExecutionCtx; -use vortex_error::VortexResult; - -use crate::CascadingCompressor; -use crate::builtins::constant::compress_constant_array_with_validity; -use crate::ctx::CompressorContext; -use crate::estimate::CompressionEstimate; -use crate::estimate::EstimateVerdict; -use crate::scheme::Scheme; -use crate::stats::ArrayAndStats; - -/// Constant encoding for bool arrays where all valid values are the same. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct BoolConstantScheme; - -impl Scheme for BoolConstantScheme { - fn scheme_name(&self) -> &'static str { - "vortex.bool.constant" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches!(canonical, Canonical::Bool(_)) - } - - fn expected_compression_ratio( - &self, - data: &ArrayAndStats, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - // Constant detection on a sample is a false positive, since the sample being constant does - // not mean the full array is constant. - if compress_ctx.is_sample() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - let array_len = data.array().len(); - let stats = data.bool_stats(exec_ctx); - - // We want to use `Constant` if there are only nulls in the array. - if stats.value_count() == 0 { - debug_assert_eq!(stats.null_count() as usize, array_len); - return CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse); - } - - if stats.is_constant() { - return CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse); - } - - CompressionEstimate::Verdict(EstimateVerdict::Skip) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - data: &ArrayAndStats, - _compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - compress_constant_array_with_validity(data.array(), exec_ctx) - } -} diff --git a/vortex-compressor/src/builtins/constant/float.rs b/vortex-compressor/src/builtins/constant/float.rs deleted file mode 100644 index 0480a1b7a53..00000000000 --- a/vortex-compressor/src/builtins/constant/float.rs +++ /dev/null @@ -1,90 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Constant encoding for float arrays. - -use vortex_array::ArrayRef; -use vortex_array::Canonical; -use vortex_array::ExecutionCtx; -use vortex_array::aggregate_fn::fns::is_constant::is_constant; -use vortex_error::VortexResult; - -use crate::CascadingCompressor; -use crate::builtins::constant::compress_constant_array_with_validity; -use crate::ctx::CompressorContext; -use crate::estimate::CompressionEstimate; -use crate::estimate::DeferredEstimate; -use crate::estimate::EstimateVerdict; -use crate::scheme::Scheme; -use crate::stats::ArrayAndStats; - -/// Constant encoding for float arrays with a single distinct value. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct FloatConstantScheme; - -impl Scheme for FloatConstantScheme { - fn scheme_name(&self) -> &'static str { - "vortex.float.constant" - } - - fn matches(&self, canonical: &Canonical) -> bool { - canonical.dtype().is_float() - } - - fn expected_compression_ratio( - &self, - data: &ArrayAndStats, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - // Constant detection on a sample is a false positive, since the sample being constant does - // not mean the full array is constant. - if compress_ctx.is_sample() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - let array_len = data.array().len(); - let stats = data.float_stats(exec_ctx); - - // Note that we only compute distinct counts if other schemes have requested it. - if let Some(distinct_count) = stats.distinct_count() { - if distinct_count > 1 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } else { - debug_assert_eq!(distinct_count, 1); - return CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse); - } - } - - // We want to use `Constant` if there are only nulls in the array. - if stats.value_count() == 0 { - debug_assert_eq!(stats.null_count() as usize, array_len); - return CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse); - } - - // TODO(connor): Can we be smart here with the max and min like with integers? - - // Otherwise our best bet is to actually check if the array is constant. - // This is an expensive check, but in practice the distinct count is known because we often - // include dictionary encoding in our set of schemes, so we rarely call this. - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, data, _best_so_far, _ctx, exec_ctx| { - if is_constant(data.array(), exec_ctx)? { - Ok(EstimateVerdict::AlwaysUse) - } else { - Ok(EstimateVerdict::Skip) - } - }, - ))) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - data: &ArrayAndStats, - _compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - compress_constant_array_with_validity(data.array(), exec_ctx) - } -} diff --git a/vortex-compressor/src/builtins/constant/integer.rs b/vortex-compressor/src/builtins/constant/integer.rs deleted file mode 100644 index 3f324c36e17..00000000000 --- a/vortex-compressor/src/builtins/constant/integer.rs +++ /dev/null @@ -1,80 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Constant encoding for integer arrays. - -use vortex_array::ArrayRef; -use vortex_array::Canonical; -use vortex_array::ExecutionCtx; -use vortex_error::VortexResult; - -use crate::CascadingCompressor; -use crate::builtins::constant::compress_constant_array_with_validity; -use crate::ctx::CompressorContext; -use crate::estimate::CompressionEstimate; -use crate::estimate::EstimateVerdict; -use crate::scheme::Scheme; -use crate::stats::ArrayAndStats; - -/// Constant encoding for integer arrays with a single distinct value. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct IntConstantScheme; - -impl Scheme for IntConstantScheme { - fn scheme_name(&self) -> &'static str { - "vortex.int.constant" - } - - fn matches(&self, canonical: &Canonical) -> bool { - canonical.dtype().is_int() - } - - fn expected_compression_ratio( - &self, - data: &ArrayAndStats, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - // Constant detection on a sample is a false positive, since the sample being constant does - // not mean the full array is constant. - if compress_ctx.is_sample() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - let array_len = data.array().len(); - let stats = data.integer_stats(exec_ctx); - - // Note that we only compute distinct counts if other schemes have requested it. - if let Some(distinct_count) = stats.distinct_count() { - if distinct_count > 1 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } else { - debug_assert_eq!(distinct_count, 1); - return CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse); - } - } - - // We want to use `Constant` if there are only nulls in the array. - if stats.value_count() == 0 { - debug_assert_eq!(stats.null_count() as usize, array_len); - return CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse); - } - - // Otherwise, use the max and min to determine if there is a single value. - match stats.erased().max_minus_min().checked_ilog2() { - Some(_) => CompressionEstimate::Verdict(EstimateVerdict::Skip), - // If max-min == 0, then we know that there is only 1 value. - None => CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse), - } - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - data: &ArrayAndStats, - _compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - compress_constant_array_with_validity(data.array(), exec_ctx) - } -} diff --git a/vortex-compressor/src/builtins/constant/mod.rs b/vortex-compressor/src/builtins/constant/mod.rs deleted file mode 100644 index d5927366fdf..00000000000 --- a/vortex-compressor/src/builtins/constant/mod.rs +++ /dev/null @@ -1,55 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Constant encoding schemes for binary, bool, float, integer, and string arrays. - -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::MaskedArray; -use vortex_array::scalar::Scalar; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; - -mod binary; -mod bool; -mod float; -mod integer; -mod string; - -pub use binary::BinaryConstantScheme; -pub use bool::BoolConstantScheme; -pub use float::FloatConstantScheme; -pub use integer::IntConstantScheme; -pub use string::StringConstantScheme; - -/// Shared helper for compressing a constant array (binary, bool, int, float, string) into a -/// [`ConstantArray`]. -/// -/// Assumes that the source array has constant valid scalars. -/// -/// If the array has any nulls, returns a [`MaskedArray`] with a [`ConstantArray`] child.` -fn compress_constant_array_with_validity( - source: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult { - if source.all_invalid(ctx)? { - return Ok( - ConstantArray::new(Scalar::null(source.dtype().clone()), source.len()).into_array(), - ); - } - - let scalar_idx = (0..source.len()) - .position(|idx| source.is_valid(idx, ctx).unwrap_or(false)) - .vortex_expect("We checked that there exists a scalar that is not invalid"); - - let scalar = source.execute_scalar(scalar_idx, ctx)?; - let const_arr = ConstantArray::new(scalar, source.len()).into_array(); - - if !source.all_valid(ctx)? { - Ok(MaskedArray::try_new(const_arr, source.validity()?)?.into_array()) - } else { - Ok(const_arr) - } -} diff --git a/vortex-compressor/src/builtins/constant/string.rs b/vortex-compressor/src/builtins/constant/string.rs deleted file mode 100644 index f55c1661660..00000000000 --- a/vortex-compressor/src/builtins/constant/string.rs +++ /dev/null @@ -1,84 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Constant encoding for string arrays. - -use vortex_array::ArrayRef; -use vortex_array::Canonical; -use vortex_array::ExecutionCtx; -use vortex_array::aggregate_fn::fns::is_constant::is_constant; -use vortex_error::VortexResult; - -use crate::CascadingCompressor; -use crate::builtins::constant::compress_constant_array_with_validity; -use crate::ctx::CompressorContext; -use crate::estimate::CompressionEstimate; -use crate::estimate::DeferredEstimate; -use crate::estimate::EstimateVerdict; -use crate::scheme::Scheme; -use crate::stats::ArrayAndStats; - -/// Constant encoding for string arrays with a single distinct value. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct StringConstantScheme; - -impl Scheme for StringConstantScheme { - fn scheme_name(&self) -> &'static str { - "vortex.string.constant" - } - - fn matches(&self, canonical: &Canonical) -> bool { - canonical.dtype().is_utf8() - } - - fn expected_compression_ratio( - &self, - data: &ArrayAndStats, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - // Constant detection on a sample is a false positive, since the sample being constant does - // not mean the full array is constant. - if compress_ctx.is_sample() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - let array_len = data.array().len(); - let stats = data.varbinview_stats(exec_ctx); - - // We want to use `Constant` if there are only nulls in the array. - if stats.value_count() == 0 { - debug_assert_eq!(stats.null_count() as usize, array_len); - return CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse); - } - - // Since the estimated distinct count is always going to be less than or equal to the actual - // distinct count, if this is not equal to 1 the actual is definitely not equal to 1. - if stats.estimated_distinct_count().is_some_and(|c| c > 1) { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - // Otherwise our best bet is to actually check if the array is constant. - // This is an expensive check, but the alternative of not compressing a constant array is - // far less preferable. - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, data, _best_so_far, _ctx, exec_ctx| { - if is_constant(data.array(), exec_ctx)? { - Ok(EstimateVerdict::AlwaysUse) - } else { - Ok(EstimateVerdict::Skip) - } - }, - ))) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - data: &ArrayAndStats, - _compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - compress_constant_array_with_validity(data.array(), exec_ctx) - } -} diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index bd194aa1121..72b5f01141a 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -21,12 +21,12 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; -use crate::ctx::CompressorContext; -use crate::estimate::CompressionEstimate; -use crate::estimate::DeferredEstimate; -use crate::estimate::EstimateVerdict; use crate::scheme::ChildSelection; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::DeferredEstimate; use crate::scheme::DescendantExclusion; +use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index 5288637a609..074d3ce5e07 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -25,12 +25,12 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; -use crate::ctx::CompressorContext; -use crate::estimate::CompressionEstimate; -use crate::estimate::DeferredEstimate; -use crate::estimate::EstimateVerdict; use crate::scheme::ChildSelection; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::DeferredEstimate; use crate::scheme::DescendantExclusion; +use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 8ec3bf53345..4e91bace3ac 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -23,9 +23,9 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::CascadingCompressor; -use crate::ctx::CompressorContext; -use crate::estimate::CompressionEstimate; -use crate::estimate::EstimateVerdict; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index 72db72c2cc5..64ed469dde9 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -21,12 +21,12 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; -use crate::ctx::CompressorContext; -use crate::estimate::CompressionEstimate; -use crate::estimate::DeferredEstimate; -use crate::estimate::EstimateVerdict; use crate::scheme::ChildSelection; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::DeferredEstimate; use crate::scheme::DescendantExclusion; +use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; diff --git a/vortex-compressor/src/builtins/mod.rs b/vortex-compressor/src/builtins/mod.rs index d014f1a6cf3..11059fa6ed4 100644 --- a/vortex-compressor/src/builtins/mod.rs +++ b/vortex-compressor/src/builtins/mod.rs @@ -3,12 +3,10 @@ //! Built-in compression schemes that use only `vortex-array` encodings. //! -//! These schemes produce arrays using types already in `vortex-array` ([`ConstantArray`], -//! [`DictArray`], [`MaskedArray`], etc.) and have no external encoding crate dependencies. +//! These schemes produce arrays using types already in `vortex-array` ([`DictArray`], etc.) and +//! have no external encoding crate dependencies. //! -//! [`ConstantArray`]: vortex_array::arrays::ConstantArray //! [`DictArray`]: vortex_array::arrays::DictArray -//! [`MaskedArray`]: vortex_array::arrays::MaskedArray mod dict; @@ -18,11 +16,3 @@ pub use dict::IntDictScheme; pub use dict::StringDictScheme; pub use dict::float_dictionary_encode; pub use dict::integer_dictionary_encode; - -mod constant; - -pub use constant::BinaryConstantScheme; -pub use constant::BoolConstantScheme; -pub use constant::FloatConstantScheme; -pub use constant::IntConstantScheme; -pub use constant::StringConstantScheme; diff --git a/vortex-compressor/src/compressor.rs b/vortex-compressor/src/compressor.rs deleted file mode 100644 index 965c719bf4b..00000000000 --- a/vortex-compressor/src/compressor.rs +++ /dev/null @@ -1,1328 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Cascading array compression implementation. - -use vortex_array::ArrayRef; -use vortex_array::ArraySlots; -use vortex_array::Canonical; -use vortex_array::CanonicalValidity; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::FixedSizeListArray; -use vortex_array::arrays::ListArray; -use vortex_array::arrays::ListViewArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::StructArray; -use vortex_array::arrays::Variant; -use vortex_array::arrays::VariantArray; -use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; -use vortex_array::arrays::list::ListArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; -use vortex_array::arrays::listview::list_from_list_view; -use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_array::arrays::scalar_fn::AnyScalarFn; -use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::arrays::variant::VariantArrayExt; -use vortex_array::scalar::Scalar; -use vortex_error::VortexResult; - -use crate::builtins::IntDictScheme; -use crate::ctx::CompressorContext; -use crate::estimate::CompressionEstimate; -use crate::estimate::DeferredEstimate; -use crate::estimate::EstimateScore; -use crate::estimate::EstimateVerdict; -use crate::estimate::WinnerEstimate; -use crate::estimate::estimate_compression_ratio_with_sampling; -use crate::estimate::is_better_score; -use crate::scheme::ChildSelection; -use crate::scheme::DescendantExclusion; -use crate::scheme::Scheme; -use crate::scheme::SchemeExt; -use crate::scheme::SchemeId; -use crate::stats::ArrayAndStats; -use crate::stats::GenerateStatsOptions; -use crate::trace; - -/// Synthetic scheme ID used for the compressor's own root-level cascading. -pub(crate) const ROOT_SCHEME_ID: SchemeId = SchemeId { - name: "vortex.compressor.root", -}; - -/// Child indices for the compressor's list/listview compression. -mod root_list_children { - /// List/ListView offsets child. - pub const OFFSETS: usize = 1; - /// ListView sizes child. - pub const SIZES: usize = 2; -} - -/// The main compressor type implementing cascading adaptive compression. -/// -/// This compressor applies adaptive compression [`Scheme`]s to arrays based on their data types and -/// characteristics. It recursively compresses nested structures like structs and lists, and chooses -/// optimal compression schemes for leaf types. -/// -/// The compressor works by: -/// 1. Canonicalizing input arrays to a standard representation. -/// 2. Pre-filtering schemes by [`Scheme::matches`] and exclusion rules. -/// 3. Evaluating each matching scheme's compression estimate and resolving deferred work. -/// 4. Compressing with the best scheme and verifying the result is smaller. -/// -/// 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. - schemes: Vec<&'static dyn Scheme>, - - /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from - /// list offsets). - root_exclusions: Vec, -} - -impl CascadingCompressor { - /// Creates a new compressor with the given schemes. - /// - /// Root-level exclusion rules (e.g. excluding Dict from list offsets) are built - /// automatically. - pub fn new(schemes: Vec<&'static dyn Scheme>) -> Self { - // Root exclusion: exclude IntDict from list/listview offsets (monotonically - // increasing data where dictionary encoding is wasteful). - let root_exclusions = vec![DescendantExclusion { - excluded: IntDictScheme.id(), - children: ChildSelection::One(root_list_children::OFFSETS), - }]; - Self { - schemes, - root_exclusions, - } - } - - /// Compresses an array using cascading adaptive compression. - /// - /// First canonicalizes and compacts the array, then applies optimal compression schemes. - /// - /// # Errors - /// - /// Returns an error if canonicalization or compression fails. - pub fn compress( - &self, - array: &ArrayRef, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - let before_nbytes = array.nbytes(); - let span = trace::compress_span(array.len(), array.dtype(), before_nbytes); - let _enter = span.enter(); - - let canonical = array.clone().execute::(exec_ctx)?.0; - let compact = canonical.compact(exec_ctx)?; - let compressed = self.compress_canonical(compact, CompressorContext::new(), exec_ctx)?; - - trace::record_compress_outcome(&span, before_nbytes, compressed.nbytes()); - - Ok(compressed) - } - - /// Compresses a child array produced by a cascading scheme. - /// - /// If the cascade budget is exhausted, the canonical array is returned as-is. Otherwise, the - /// child context is created by descending and recording the parent scheme + child index, and - /// compression proceeds normally. - /// - /// # Errors - /// - /// Returns an error if compression fails. - pub fn compress_child( - &self, - child: &ArrayRef, - parent_ctx: &CompressorContext, - parent_id: SchemeId, - child_index: usize, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - if parent_ctx.finished_cascading() { - trace::cascade_exhausted(parent_id, child_index); - return Ok(child.clone()); - } - - let canonical = child.clone().execute::(exec_ctx)?.0; - let compact = canonical.compact(exec_ctx)?; - - let child_ctx = parent_ctx - .clone() - .descend_with_scheme(parent_id, child_index); - self.compress_canonical(compact, child_ctx, exec_ctx) - } - - /// Compresses a canonical array by dispatching to type-specific logic. - /// - /// # Errors - /// - /// Returns an error if compression of any sub-array fails. - fn compress_canonical( - &self, - array: Canonical, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - match array { - Canonical::Null(null_array) => Ok(null_array.into_array()), - Canonical::Bool(bool_array) => { - self.choose_and_compress(Canonical::Bool(bool_array), compress_ctx, exec_ctx) - } - Canonical::Primitive(primitive) => { - self.choose_and_compress(Canonical::Primitive(primitive), compress_ctx, exec_ctx) - } - Canonical::Decimal(decimal) => { - self.choose_and_compress(Canonical::Decimal(decimal), compress_ctx, exec_ctx) - } - Canonical::Struct(struct_array) => { - let fields = struct_array - .iter_unmasked_fields() - .map(|field| self.compress(field, exec_ctx)) - .collect::, _>>()?; - - Ok(StructArray::try_new( - struct_array.names().clone(), - fields, - struct_array.len(), - struct_array.validity()?, - )? - .into_array()) - } - Canonical::List(list_view_array) => { - if list_view_array.is_zero_copy_to_list() || list_view_array.elements().is_empty() { - let list_array = list_from_list_view(list_view_array, exec_ctx)?; - self.compress_list_array(list_array, compress_ctx, exec_ctx) - } else { - self.compress_list_view_array(list_view_array, compress_ctx, exec_ctx) - } - } - Canonical::FixedSizeList(fsl_array) => { - let compressed_elems = self.compress(fsl_array.elements(), exec_ctx)?; - - Ok(FixedSizeListArray::try_new( - compressed_elems, - fsl_array.list_size(), - fsl_array.validity()?, - fsl_array.len(), - )? - .into_array()) - } - Canonical::VarBinView(varbinview) => { - self.choose_and_compress(Canonical::VarBinView(varbinview), compress_ctx, exec_ctx) - } - Canonical::Extension(ext_array) => { - // Try scheme-based compression first. - let scheme_compressed = self.choose_and_compress( - Canonical::Extension(ext_array.clone()), - compress_ctx, - exec_ctx, - )?; - // TODO(connor): HACK TO SUPPORT L2 DENORMALIZATION!!! - if scheme_compressed.is::() { - return Ok(scheme_compressed); - } - - // Also compress the underlying storage array. Some extension schemes can beat the - // extension storage but still lose to ordinary storage compression. - let compressed_storage = self.compress(ext_array.storage_array(), exec_ctx)?; - let storage_compressed = - ExtensionArray::new(ext_array.ext_dtype().clone(), compressed_storage) - .into_array(); - - if scheme_compressed.nbytes() < storage_compressed.nbytes() { - Ok(scheme_compressed) - } else { - Ok(storage_compressed) - } - } - Canonical::Variant(variant_array) => { - let core_storage = - self.compress_physical_slots(variant_array.core_storage(), exec_ctx)?; - let shredded = variant_array - .shredded() - .map(|arr| { - // Avoid stack-overflow for variant shredded values - if arr.is::() { - self.compress_physical_slots(arr, exec_ctx) - } else { - self.compress(arr, exec_ctx) - } - }) - .transpose()?; - - Ok(VariantArray::try_new(core_storage, shredded)?.into_array()) - } - } - } - - /// The main scheme-selection entry point for a single leaf array. - /// - /// Filters allowed schemes by [`matches`] and exclusion rules, merges their [`stats_options`] - /// into a single [`GenerateStatsOptions`], and picks the winner by estimated compression - /// ratio. - /// - /// If a winner is found and its compressed output is actually smaller, that output is - /// returned. Otherwise, the original array is returned unchanged. - /// - /// Empty and all-null arrays are short-circuited before any scheme evaluation. - /// - /// [`matches`]: Scheme::matches - /// [`stats_options`]: Scheme::stats_options - fn choose_and_compress( - &self, - canonical: Canonical, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - let eligible_schemes: Vec<&'static dyn Scheme> = self - .schemes - .iter() - .copied() - .filter(|s| s.matches(&canonical) && !self.is_excluded(*s, &compress_ctx)) - .collect(); - - let array: ArrayRef = canonical.into(); - - if eligible_schemes.is_empty() || array.is_empty() { - return Ok(array); - } - - if array.all_invalid(exec_ctx)? { - return Ok( - ConstantArray::new(Scalar::null(array.dtype().clone()), array.len()).into_array(), - ); - } - - let before_nbytes = array.nbytes(); - - let merged_opts = eligible_schemes - .iter() - .fold(GenerateStatsOptions::default(), |acc, s| { - acc.merge(s.stats_options()) - }); - let compress_ctx = compress_ctx.with_merged_stats_options(merged_opts); - - let data = ArrayAndStats::new(array, merged_opts); - - let Some((winner, winner_estimate)) = - self.choose_best_scheme(&eligible_schemes, &data, compress_ctx.clone(), exec_ctx)? - else { - return Ok(data.into_array()); - }; - - // Run the winning scheme's `compress`. On failure, emit an ERROR event carrying the - // scheme name and cascade history before propagating. - let error_ctx = trace::enabled_error_context(&compress_ctx); - let _winner_span = trace::winner_compress_span(winner.id(), before_nbytes).entered(); - let compressed = winner - .compress(self, &data, compress_ctx, exec_ctx) - .inspect_err(|err| { - // NB: this is the only way we can tell which scheme panicked / bailed on their - // data, especially for third-party schemes where the error site may not carry any - // compressor context. - trace::scheme_compress_failed(winner.id(), before_nbytes, error_ctx.as_ref(), err); - })?; - - let after_nbytes = compressed.nbytes(); - let actual_ratio = (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64); - - // TODO(connor): HACK TO SUPPORT L2 DENORMALIZATION!!! - let accepted = after_nbytes < before_nbytes || compressed.is::(); - - trace::record_winner_compress_result( - after_nbytes, - winner_estimate.trace_ratio(), - actual_ratio, - accepted, - ); - - if accepted { - Ok(compressed) - } else { - Ok(data.into_array()) - } - } - - /// Calls [`expected_compression_ratio`] on each candidate and returns the winning scheme along - /// with its resolved winner estimate, or `None` if no scheme beats the canonical encoding. - /// - /// Selection runs in two passes. Pass 1 evaluates every immediate - /// [`CompressionEstimate::Verdict`] and tracks the running best. [`Scheme`]s returning - /// [`CompressionEstimate::Deferred`] are stashed for pass 2 so that we do not make any - /// expensive computations if we don't have to. - /// - /// Pass 2 evaluates the deferred work and, for each [`DeferredEstimate::Callback`], passes the - /// current best [`EstimateScore`] as an early-exit hint so the callback can return - /// [`EstimateVerdict::Skip`] without doing expensive work when it cannot beat the threshold. - /// - /// Ties are broken by registration order within each pass. - /// - /// [`expected_compression_ratio`]: Scheme::expected_compression_ratio - fn choose_best_scheme( - &self, - schemes: &[&'static dyn Scheme], - data: &ArrayAndStats, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let mut best: Option<(&'static dyn Scheme, EstimateScore)> = None; - let mut deferred: Vec<(&'static dyn Scheme, DeferredEstimate)> = Vec::new(); - - // Pass 1: evaluate every immediate verdict. Stash deferred work for pass 2. - { - let _verdict_pass = trace::verdict_pass_span().entered(); - for &scheme in schemes { - match scheme.expected_compression_ratio(data, compress_ctx.clone(), exec_ctx) { - CompressionEstimate::Verdict(EstimateVerdict::Skip) => {} - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) => { - return Ok(Some((scheme, WinnerEstimate::AlwaysUse))); - } - CompressionEstimate::Verdict(EstimateVerdict::Ratio(ratio)) => { - let score = EstimateScore::FiniteCompression(ratio); - - if is_better_score(score, best.as_ref()) { - best = Some((scheme, score)); - } - } - CompressionEstimate::Deferred(deferred_estimate) => { - deferred.push((scheme, deferred_estimate)); - } - } - } - } - - // Pass 2: run deferred work. Callbacks receive the current best as a threshold so they can - // short-circuit with `Skip` when they cannot beat it. - for (scheme, deferred_estimate) in deferred { - let _span = trace::scheme_eval_span(scheme.id()).entered(); - let threshold: Option = best.map(|(_, score)| score); - match deferred_estimate { - DeferredEstimate::Sample => { - let score = estimate_compression_ratio_with_sampling( - self, - scheme, - data.array(), - compress_ctx.clone(), - exec_ctx, - )?; - - if is_better_score(score, best.as_ref()) { - best = Some((scheme, score)); - } - } - DeferredEstimate::Callback(callback) => { - match callback(self, data, threshold, compress_ctx.clone(), exec_ctx)? { - EstimateVerdict::Skip => {} - EstimateVerdict::AlwaysUse => { - return Ok(Some((scheme, WinnerEstimate::AlwaysUse))); - } - EstimateVerdict::Ratio(ratio) => { - let score = EstimateScore::FiniteCompression(ratio); - - if is_better_score(score, best.as_ref()) { - best = Some((scheme, score)); - } - } - } - } - } - } - - Ok(best.map(|(scheme, score)| (scheme, WinnerEstimate::Score(score)))) - } - - // TODO(connor): Lots of room for optimization here. - /// Returns `true` if the candidate scheme should be excluded based on the cascade history and - /// exclusion rules. - fn is_excluded(&self, candidate: &dyn Scheme, ctx: &CompressorContext) -> bool { - let id = candidate.id(); - let history = ctx.cascade_history(); - - // Self-exclusion: no scheme appears twice in any chain. - if history.iter().any(|&(sid, _)| sid == id) { - return true; - } - - let mut iter = history.iter().copied().peekable(); - - // The root entry is always first in the history (if present). Check if the root has - // excluded us. - if let Some((_, child_idx)) = iter.next_if(|&(sid, _)| sid == ROOT_SCHEME_ID) - && self - .root_exclusions - .iter() - .any(|rule| rule.excluded == id && rule.children.contains(child_idx)) - { - return true; - } - - // Push rules: Check if any of our ancestors have excluded us. - for (ancestor_id, child_idx) in iter { - if let Some(ancestor) = self.schemes.iter().find(|s| s.id() == ancestor_id) - && ancestor - .descendant_exclusions() - .iter() - .any(|rule| rule.excluded == id && rule.children.contains(child_idx)) - { - return true; - } - } - - // Pull rules: Check if we have excluded ourselves because of our ancestors. - for rule in candidate.ancestor_exclusions() { - if history - .iter() - .any(|(sid, cidx)| *sid == rule.ancestor && rule.children.contains(*cidx)) - { - return true; - } - } - - false - } - - /// Compresses a [`ListArray`] by narrowing offsets and recursively compressing elements. - fn compress_list_array( - &self, - list_array: ListArray, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - let list_array = list_array.reset_offsets(true, exec_ctx)?; - - let compressed_elems = self.compress(list_array.elements(), exec_ctx)?; - - // Record the root scheme with the offsets child index so root exclusion rules apply. - let offset_ctx = - compress_ctx.descend_with_scheme(ROOT_SCHEME_ID, root_list_children::OFFSETS); - let list_offsets_primitive = list_array - .offsets() - .clone() - .execute::(exec_ctx)? - .narrow(exec_ctx)?; - let compressed_offsets = self.compress_canonical( - Canonical::Primitive(list_offsets_primitive), - offset_ctx, - exec_ctx, - )?; - - Ok( - ListArray::try_new(compressed_elems, compressed_offsets, list_array.validity()?)? - .into_array(), - ) - } - - /// Compresses a [`ListViewArray`] by narrowing offsets/sizes and recursively compressing - /// elements. - fn compress_list_view_array( - &self, - list_view: ListViewArray, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - let compressed_elems = self.compress(list_view.elements(), exec_ctx)?; - - let offset_ctx = compress_ctx - .clone() - .descend_with_scheme(ROOT_SCHEME_ID, root_list_children::OFFSETS); - let list_view_offsets_primitive = list_view - .offsets() - .clone() - .execute::(exec_ctx)? - .narrow(exec_ctx)?; - let compressed_offsets = self.compress_canonical( - Canonical::Primitive(list_view_offsets_primitive), - offset_ctx, - exec_ctx, - )?; - - let sizes_ctx = compress_ctx.descend_with_scheme(ROOT_SCHEME_ID, root_list_children::SIZES); - let list_view_sizes_primitive = list_view - .sizes() - .clone() - .execute::(exec_ctx)? - .narrow(exec_ctx)?; - let compressed_sizes = self.compress_canonical( - Canonical::Primitive(list_view_sizes_primitive), - sizes_ctx, - exec_ctx, - )?; - - Ok(ListViewArray::try_new( - compressed_elems, - compressed_offsets, - compressed_sizes, - list_view.validity()?, - )? - .into_array()) - } - - /// Compress very child slot of the array, then re-build it from them. - fn compress_physical_slots( - &self, - array: &ArrayRef, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - let slots = array - .slots() - .iter() - .map(|slot| { - slot.as_ref() - .map(|child| self.compress(child, exec_ctx)) - .transpose() - }) - .collect::>()?; - - // 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) } - } -} - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use parking_lot::Mutex; - use vortex_array::ArrayRef; - use vortex_array::Canonical; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::Constant; - use vortex_array::arrays::NullArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::validity::Validity; - use vortex_buffer::buffer; - use vortex_session::VortexSession; - - use super::*; - use crate::builtins::FloatDictScheme; - use crate::builtins::IntDictScheme; - use crate::builtins::StringDictScheme; - use crate::ctx::CompressorContext; - use crate::estimate::CompressionEstimate; - use crate::estimate::DeferredEstimate; - use crate::estimate::EstimateScore; - use crate::estimate::EstimateVerdict; - use crate::estimate::WinnerEstimate; - use crate::scheme::SchemeExt; - - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); - - fn compressor() -> CascadingCompressor { - CascadingCompressor::new(vec![&IntDictScheme, &FloatDictScheme, &StringDictScheme]) - } - - fn estimate_test_data() -> ArrayAndStats { - let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); - ArrayAndStats::new(array, GenerateStatsOptions::default()) - } - - fn matches_integer_primitive(canonical: &Canonical) -> bool { - matches!(canonical, Canonical::Primitive(primitive) if primitive.ptype().is_int()) - } - - #[derive(Debug)] - struct DirectRatioScheme; - - impl Scheme for DirectRatioScheme { - fn scheme_name(&self) -> &'static str { - "test.direct_ratio" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::Ratio(2.0)) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - unreachable!("test helper should never be selected for compression") - } - } - - #[derive(Debug)] - struct ImmediateAlwaysUseScheme; - - impl Scheme for ImmediateAlwaysUseScheme { - fn scheme_name(&self) -> &'static str { - "test.immediate_always_use" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - unreachable!("test helper should never be selected for compression") - } - } - - #[derive(Debug)] - struct CallbackAlwaysUseScheme; - - impl Scheme for CallbackAlwaysUseScheme { - fn scheme_name(&self) -> &'static str { - "test.callback_always_use" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, _data, _ctx, _exec_ctx, _best_so_far| Ok(EstimateVerdict::AlwaysUse), - ))) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - unreachable!("test helper should never be selected for compression") - } - } - - #[derive(Debug)] - struct CallbackSkipScheme; - - impl Scheme for CallbackSkipScheme { - fn scheme_name(&self) -> &'static str { - "test.callback_skip" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, _data, _ctx, _exec_ctx, _best_so_far| Ok(EstimateVerdict::Skip), - ))) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - unreachable!("test helper should never be selected for compression") - } - } - - #[derive(Debug)] - struct CallbackRatioScheme; - - impl Scheme for CallbackRatioScheme { - fn scheme_name(&self) -> &'static str { - "test.callback_ratio" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, _data, _ctx, _exec_ctx, _best_so_far| Ok(EstimateVerdict::Ratio(3.0)), - ))) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - unreachable!("test helper should never be selected for compression") - } - } - - #[derive(Debug)] - struct HugeRatioScheme; - - impl Scheme for HugeRatioScheme { - fn scheme_name(&self) -> &'static str { - "test.huge_ratio" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::Ratio(100.0)) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - unreachable!("test helper should never be selected for compression") - } - } - - #[derive(Debug)] - struct ZeroBytesSamplingScheme; - - impl Scheme for ZeroBytesSamplingScheme { - fn scheme_name(&self) -> &'static str { - "test.zero_bytes_sampling" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Sample) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - Ok(NullArray::new(data.array().len()).into_array()) - } - } - - #[test] - fn test_self_exclusion() { - let c = compressor(); - let ctx = CompressorContext::default().descend_with_scheme(IntDictScheme.id(), 0); - - // IntDictScheme is in the history, so it should be excluded. - assert!(c.is_excluded(&IntDictScheme, &ctx)); - } - - #[test] - fn test_root_exclusion_list_offsets() { - let c = compressor(); - let ctx = CompressorContext::default() - .descend_with_scheme(ROOT_SCHEME_ID, root_list_children::OFFSETS); - - // IntDict should be excluded for list offsets. - assert!(c.is_excluded(&IntDictScheme, &ctx)); - } - - #[test] - fn test_push_rule_float_dict_excludes_int_dict_from_codes() { - let c = compressor(); - // FloatDict cascading through codes (child 1). - let ctx = CompressorContext::default().descend_with_scheme(FloatDictScheme.id(), 1); - - // IntDict should be excluded from FloatDict's codes child. - assert!(c.is_excluded(&IntDictScheme, &ctx)); - } - - #[test] - fn test_push_rule_float_dict_excludes_int_dict_from_values() { - let c = compressor(); - // FloatDict cascading through values (child 0). - let ctx = CompressorContext::default().descend_with_scheme(FloatDictScheme.id(), 0); - - // IntDict should also be excluded from FloatDict's values child (ALP propagation - // replacement). - assert!(c.is_excluded(&IntDictScheme, &ctx)); - } - - #[test] - fn test_no_exclusion_without_history() { - let c = compressor(); - let ctx = CompressorContext::default(); - - // No history means no exclusions. - assert!(!c.is_excluded(&IntDictScheme, &ctx)); - } - - #[test] - fn immediate_always_use_wins_immediately() -> VortexResult<()> { - let compressor = - CascadingCompressor::new(vec![&DirectRatioScheme, &ImmediateAlwaysUseScheme]); - let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &ImmediateAlwaysUseScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(), - &mut exec_ctx, - )?; - - assert!(matches!( - winner, - Some((scheme, WinnerEstimate::AlwaysUse)) - if scheme.id() == ImmediateAlwaysUseScheme.id() - )); - Ok(()) - } - - #[test] - fn callback_always_use_wins_immediately() -> VortexResult<()> { - let compressor = - CascadingCompressor::new(vec![&DirectRatioScheme, &CallbackAlwaysUseScheme]); - let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &CallbackAlwaysUseScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(), - &mut exec_ctx, - )?; - - assert!(matches!( - winner, - Some((scheme, WinnerEstimate::AlwaysUse)) - if scheme.id() == CallbackAlwaysUseScheme.id() - )); - Ok(()) - } - - #[test] - fn callback_skip_is_ignored() -> VortexResult<()> { - let compressor = CascadingCompressor::new(vec![&CallbackSkipScheme, &DirectRatioScheme]); - let schemes: [&'static dyn Scheme; 2] = [&CallbackSkipScheme, &DirectRatioScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(), - &mut exec_ctx, - )?; - - assert!(matches!( - winner, - Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(2.0)))) - if scheme.id() == DirectRatioScheme.id() - )); - Ok(()) - } - - #[test] - fn callback_ratio_competes_numerically() -> VortexResult<()> { - let compressor = CascadingCompressor::new(vec![&DirectRatioScheme, &CallbackRatioScheme]); - let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &CallbackRatioScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(), - &mut exec_ctx, - )?; - - assert!(matches!( - winner, - Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(3.0)))) - if scheme.id() == CallbackRatioScheme.id() - )); - Ok(()) - } - - #[test] - fn zero_byte_sample_loses_to_finite_ratio() -> VortexResult<()> { - let compressor = CascadingCompressor::new(vec![&HugeRatioScheme, &ZeroBytesSamplingScheme]); - let schemes: [&'static dyn Scheme; 2] = [&HugeRatioScheme, &ZeroBytesSamplingScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(), - &mut exec_ctx, - )?; - - assert!(matches!( - winner, - Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(100.0)))) - if scheme.id() == HugeRatioScheme.id() - )); - Ok(()) - } - - #[test] - fn finite_ratio_displaces_zero_byte_sample() -> VortexResult<()> { - let compressor = CascadingCompressor::new(vec![&ZeroBytesSamplingScheme, &HugeRatioScheme]); - let schemes: [&'static dyn Scheme; 2] = [&ZeroBytesSamplingScheme, &HugeRatioScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(), - &mut exec_ctx, - )?; - - assert!(matches!( - winner, - Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(100.0)))) - if scheme.id() == HugeRatioScheme.id() - )); - Ok(()) - } - - #[test] - fn zero_byte_sample_alone_selects_no_scheme() -> VortexResult<()> { - let compressor = CascadingCompressor::new(vec![&ZeroBytesSamplingScheme]); - let schemes: [&'static dyn Scheme; 1] = [&ZeroBytesSamplingScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(), - &mut exec_ctx, - )?; - - assert!(winner.is_none()); - Ok(()) - } - - // Observer helper used by threshold-related tests. Captures the `best_so_far` value the - // compressor passes to its deferred callback. `OBSERVER_LOCK` serializes tests that share - // `OBSERVED_THRESHOLD` so they do not race. - static OBSERVER_LOCK: Mutex<()> = Mutex::new(()); - static OBSERVED_THRESHOLD: Mutex>> = Mutex::new(None); - - #[derive(Debug)] - struct ThresholdObservingScheme; - - impl Scheme for ThresholdObservingScheme { - fn scheme_name(&self) -> &'static str { - "test.threshold_observing" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, _data, best_so_far, _ctx, _exec_ctx| { - *OBSERVED_THRESHOLD.lock() = Some(best_so_far); - Ok(EstimateVerdict::Skip) - }, - ))) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - unreachable!("test helper should never be selected for compression") - } - } - - #[derive(Debug)] - struct CallbackMatchingRatioScheme; - - impl Scheme for CallbackMatchingRatioScheme { - fn scheme_name(&self) -> &'static str { - "test.callback_matching_ratio" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, _data, _ctx, _exec_ctx, _best_so_far| Ok(EstimateVerdict::Ratio(2.0)), - ))) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - unreachable!("test helper should never be selected for compression") - } - } - - #[test] - fn callback_always_use_overrides_pass_one_best() -> VortexResult<()> { - // `HugeRatioScheme` returns an immediate `Ratio(100.0)` in pass 1; - // `CallbackAlwaysUseScheme` returns `AlwaysUse` from its deferred callback in pass 2. - // The deferred `AlwaysUse` must still win. - let compressor = CascadingCompressor::new(vec![&HugeRatioScheme, &CallbackAlwaysUseScheme]); - let schemes: [&'static dyn Scheme; 2] = [&HugeRatioScheme, &CallbackAlwaysUseScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(), - &mut exec_ctx, - )?; - - assert!(matches!( - winner, - Some((scheme, WinnerEstimate::AlwaysUse)) - if scheme.id() == CallbackAlwaysUseScheme.id() - )); - Ok(()) - } - - #[test] - fn threshold_reflects_pass_one_best() -> VortexResult<()> { - let _guard = OBSERVER_LOCK.lock(); - *OBSERVED_THRESHOLD.lock() = None; - - let compressor = - CascadingCompressor::new(vec![&DirectRatioScheme, &ThresholdObservingScheme]); - let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &ThresholdObservingScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; - - let observed = *OBSERVED_THRESHOLD.lock(); - assert!(matches!( - observed, - Some(Some(EstimateScore::FiniteCompression(r))) if r == 2.0 - )); - Ok(()) - } - - #[test] - fn threshold_is_none_when_only_prior_is_zero_bytes() -> VortexResult<()> { - let _guard = OBSERVER_LOCK.lock(); - *OBSERVED_THRESHOLD.lock() = None; - - let compressor = - CascadingCompressor::new(vec![&ZeroBytesSamplingScheme, &ThresholdObservingScheme]); - let schemes: [&'static dyn Scheme; 2] = - [&ZeroBytesSamplingScheme, &ThresholdObservingScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; - - // The observing callback was invoked (outer `Some`) and `best_so_far` was `None` (inner - // `None`) because the zero-byte sample is never stored as the best. - let observed = *OBSERVED_THRESHOLD.lock(); - assert_eq!(observed, Some(None)); - Ok(()) - } - - #[test] - fn threshold_is_none_when_no_prior_scheme() -> VortexResult<()> { - let _guard = OBSERVER_LOCK.lock(); - *OBSERVED_THRESHOLD.lock() = None; - - let compressor = CascadingCompressor::new(vec![&ThresholdObservingScheme]); - let schemes: [&'static dyn Scheme; 1] = [&ThresholdObservingScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; - - let observed = *OBSERVED_THRESHOLD.lock(); - assert_eq!(observed, Some(None)); - Ok(()) - } - - #[test] - fn threshold_updates_from_earlier_deferred_callback() -> VortexResult<()> { - let _guard = OBSERVER_LOCK.lock(); - *OBSERVED_THRESHOLD.lock() = None; - - // Both schemes are deferred. The first callback registers `Ratio(3.0)`; the second - // callback must observe it as its threshold. - let compressor = - CascadingCompressor::new(vec![&CallbackRatioScheme, &ThresholdObservingScheme]); - let schemes: [&'static dyn Scheme; 2] = [&CallbackRatioScheme, &ThresholdObservingScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; - - let observed = *OBSERVED_THRESHOLD.lock(); - assert!(matches!( - observed, - Some(Some(EstimateScore::FiniteCompression(r))) if r == 3.0 - )); - Ok(()) - } - - #[test] - fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<()> { - // Both schemes produce the same `Ratio(2.0)`, one from pass 1 (immediate) and one from - // pass 2 (deferred callback). Pass 1 locks in first, and strict `>` tie-breaking means - // the deferred callback's equal ratio cannot displace it. - let compressor = - CascadingCompressor::new(vec![&CallbackMatchingRatioScheme, &DirectRatioScheme]); - let schemes: [&'static dyn Scheme; 2] = [&CallbackMatchingRatioScheme, &DirectRatioScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(), - &mut exec_ctx, - )?; - - assert!(matches!( - winner, - Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(r)))) - if scheme.id() == DirectRatioScheme.id() && r == 2.0 - )); - Ok(()) - } - - #[test] - fn all_null_array_compresses_to_constant() -> VortexResult<()> { - let array = PrimitiveArray::new( - buffer![0i32, 0, 0, 0, 0], - Validity::Array(BoolArray::from_iter([false, false, false, false, false]).into_array()), - ) - .into_array(); - - // The compressor should produce a `ConstantArray` for an all-null array regardless of - // which schemes are registered. - let compressor = CascadingCompressor::new(vec![&IntDictScheme]); - let mut exec_ctx = SESSION.create_execution_ctx(); - let compressed = compressor.compress(&array, &mut exec_ctx)?; - assert!(compressed.is::()); - Ok(()) - } - - /// Regression test for . - /// - /// `estimate_compression_ratio_with_sampling` must use the *scheme's* stats options - /// (which request distinct-value counting) rather than the context's stats options - /// (which may not). With the old code this panicked inside `dictionary_encode` because - /// distinct values were never computed for the sample. - #[test] - fn sampling_uses_scheme_stats_options() -> VortexResult<()> { - // Low-cardinality float array so FloatDictScheme considers it compressible. - let array = PrimitiveArray::new( - buffer![1.0f32, 2.0, 1.0, 2.0, 1.0, 2.0, 1.0, 2.0], - Validity::NonNullable, - ) - .into_array(); - - let compressor = CascadingCompressor::new(vec![&FloatDictScheme]); - - // A context with default stats_options (count_distinct_values = false) and - // marked as a sample so the function skips the sampling step and compresses - // the array directly. - let ctx = CompressorContext::new().with_sampling(); - - // Before the fix this panicked with: - // "this must be present since `DictScheme` declared that we need distinct values" - let mut exec_ctx = SESSION.create_execution_ctx(); - let score = estimate_compression_ratio_with_sampling( - &compressor, - &FloatDictScheme, - &array, - ctx, - &mut exec_ctx, - )?; - assert!(matches!(score, EstimateScore::FiniteCompression(ratio) if ratio.is_finite())); - Ok(()) - } -} diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs new file mode 100644 index 00000000000..9dbbc611448 --- /dev/null +++ b/vortex-compressor/src/compressor/cascade.rs @@ -0,0 +1,324 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Core cascading compression flow. + +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::CanonicalValidity; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::Masked; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::Variant; +use vortex_array::arrays::VariantArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::listview::list_from_list_view; +use vortex_array::arrays::masked::MaskedArraySlotsExt; +use vortex_array::arrays::scalar_fn::AnyScalarFn; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::arrays::variant::VariantArrayExt; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; + +use super::CascadingCompressor; +use super::constant; +use crate::scheme::CompressorContext; +use crate::scheme::Scheme; +use crate::scheme::SchemeExt; +use crate::scheme::SchemeId; +use crate::stats::ArrayAndStats; +use crate::stats::GenerateStatsOptions; +use crate::trace; + +impl CascadingCompressor { + /// Compresses an array using cascading adaptive compression. + /// + /// First canonicalizes and compacts the array, then applies optimal compression schemes. + /// + /// # Errors + /// + /// Returns an error if canonicalization or compression fails. + pub fn compress( + &self, + array: &ArrayRef, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let before_nbytes = array.nbytes(); + let span = trace::compress_span(array.len(), array.dtype(), before_nbytes); + let _enter = span.enter(); + + let canonical = array.clone().execute::(exec_ctx)?.0; + let compact = canonical.compact(exec_ctx)?; + let compressed = self.compress_canonical(compact, CompressorContext::new(), exec_ctx)?; + + trace::record_compress_outcome(&span, before_nbytes, compressed.nbytes()); + + Ok(compressed) + } + + /// Compresses a child array produced by a cascading scheme. + /// + /// If the cascade budget is exhausted, the canonical array is returned as-is. Otherwise, the + /// child context is created by descending and recording the parent scheme + child index, and + /// compression proceeds normally. + /// + /// # Errors + /// + /// Returns an error if compression fails. + pub fn compress_child( + &self, + child: &ArrayRef, + parent_ctx: &CompressorContext, + parent_id: SchemeId, + child_index: usize, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + if parent_ctx.finished_cascading() { + trace::cascade_exhausted(parent_id, child_index); + return Ok(child.clone()); + } + + let canonical = child.clone().execute::(exec_ctx)?.0; + let compact = canonical.compact(exec_ctx)?; + + let child_ctx = parent_ctx + .clone() + .descend_with_scheme(parent_id, child_index); + self.compress_canonical(compact, child_ctx, exec_ctx) + } + + /// Compresses a canonical array by dispatching to type-specific logic. + /// + /// # Errors + /// + /// Returns an error if compression of any sub-array fails. + pub(super) fn compress_canonical( + &self, + array: Canonical, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + match array { + Canonical::Null(null_array) => Ok(null_array.into_array()), + Canonical::Bool(bool_array) => { + self.choose_and_compress(Canonical::Bool(bool_array), compress_ctx, exec_ctx) + } + Canonical::Primitive(primitive) => { + self.choose_and_compress(Canonical::Primitive(primitive), compress_ctx, exec_ctx) + } + Canonical::Decimal(decimal) => { + self.choose_and_compress(Canonical::Decimal(decimal), compress_ctx, exec_ctx) + } + Canonical::Struct(struct_array) => { + let fields = struct_array + .iter_unmasked_fields() + .map(|field| self.compress(field, exec_ctx)) + .collect::, _>>()?; + + Ok(StructArray::try_new( + struct_array.names().clone(), + fields, + struct_array.len(), + struct_array.validity()?, + )? + .into_array()) + } + Canonical::List(list_view_array) => { + if list_view_array.is_zero_copy_to_list() || list_view_array.elements().is_empty() { + let list_array = list_from_list_view(list_view_array, exec_ctx)?; + self.compress_list_array(list_array, compress_ctx, exec_ctx) + } else { + self.compress_list_view_array(list_view_array, compress_ctx, exec_ctx) + } + } + Canonical::FixedSizeList(fsl_array) => { + let compressed_elems = self.compress(fsl_array.elements(), exec_ctx)?; + + Ok(FixedSizeListArray::try_new( + compressed_elems, + fsl_array.list_size(), + fsl_array.validity()?, + fsl_array.len(), + )? + .into_array()) + } + Canonical::VarBinView(varbinview) => { + self.choose_and_compress(Canonical::VarBinView(varbinview), compress_ctx, exec_ctx) + } + Canonical::Extension(ext_array) => { + // Try scheme-based compression first. + let scheme_compressed = self.choose_and_compress( + Canonical::Extension(ext_array.clone()), + compress_ctx, + exec_ctx, + )?; + // TODO(connor): HACK TO SUPPORT L2 DENORMALIZATION!!! + if scheme_compressed.is::() { + return Ok(scheme_compressed); + } + + // A constant extension array (that might be masked) is already in its terminal + // representation, and compressing the storage separately cannot do better. + if scheme_compressed.is::() { + return Ok(scheme_compressed); + } + if let Some(masked) = scheme_compressed.as_opt::() + && masked.child().is::() + { + return Ok(scheme_compressed); + } + + // Also compress the underlying storage array. Some extension schemes can beat the + // extension storage but still lose to ordinary storage compression. + let compressed_storage = self.compress(ext_array.storage_array(), exec_ctx)?; + let storage_compressed = + ExtensionArray::new(ext_array.ext_dtype().clone(), compressed_storage) + .into_array(); + + if scheme_compressed.nbytes() < storage_compressed.nbytes() { + Ok(scheme_compressed) + } else { + Ok(storage_compressed) + } + } + Canonical::Variant(variant_array) => { + let core_storage = + self.compress_physical_slots(variant_array.core_storage(), exec_ctx)?; + let shredded = variant_array + .shredded() + .map(|arr| { + // Avoid stack-overflow for variant shredded values + if arr.is::() { + self.compress_physical_slots(arr, exec_ctx) + } else { + self.compress(arr, exec_ctx) + } + }) + .transpose()?; + + Ok(VariantArray::try_new(core_storage, shredded)?.into_array()) + } + } + } + + /// The main scheme-selection entry point for a single leaf array. + /// + /// Filters allowed schemes by [`matches`] and exclusion rules, merges their [`stats_options`] + /// into a single [`GenerateStatsOptions`], and picks the winner by estimated compression + /// ratio. + /// + /// If a winner is found and its compressed output is actually smaller, that output is + /// returned. Otherwise, the original array is returned unchanged. + /// + /// Empty, all-null, and constant arrays are handled by the compressor itself before any + /// scheme evaluation (constant detection is skipped while compressing samples). + /// + /// [`matches`]: Scheme::matches + /// [`stats_options`]: Scheme::stats_options + fn choose_and_compress( + &self, + canonical: Canonical, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let eligible_schemes: Vec<&'static dyn Scheme> = self + .schemes + .iter() + .copied() + .filter(|s| s.matches(&canonical) && !self.is_excluded(*s, &compress_ctx)) + .collect(); + + let array: ArrayRef = canonical.into(); + + if array.is_empty() { + return Ok(array); + } + + if array.all_invalid(exec_ctx)? { + return Ok( + ConstantArray::new(Scalar::null(array.dtype().clone()), array.len()).into_array(), + ); + } + + let before_nbytes = array.nbytes(); + + let merged_opts = eligible_schemes + .iter() + .fold(GenerateStatsOptions::default(), |acc, s| { + acc.merge(s.stats_options()) + }); + let compress_ctx = compress_ctx.with_merged_stats_options(merged_opts); + + let data = ArrayAndStats::new(array, merged_opts); + + // Constant detection is built into the compressor: a constant leaf always short-circuits + // scheme selection. Samples are exempt because a constant sample does not imply that the + // full array is constant. + if !compress_ctx.is_sample() && constant::is_constant_for_compression(&data, exec_ctx)? { + let _winner_span = + trace::winner_compress_span(constant::CONSTANT_SCHEME_ID, before_nbytes).entered(); + let compressed = constant::compress_constant(data.array(), exec_ctx)?; + + let after_nbytes = compressed.nbytes(); + let actual_ratio = + (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64); + let accepted = after_nbytes < before_nbytes; + trace::record_winner_compress_result(after_nbytes, None, actual_ratio, accepted); + + return if accepted { + Ok(compressed) + } else { + Ok(data.into_array()) + }; + } + + if eligible_schemes.is_empty() { + return Ok(data.into_array()); + } + + let Some((winner, winner_estimate)) = + self.choose_best_scheme(&eligible_schemes, &data, compress_ctx.clone(), exec_ctx)? + else { + return Ok(data.into_array()); + }; + + // Run the winning scheme's `compress`. On failure, emit an ERROR event carrying the + // scheme name and cascade history before propagating. + let error_ctx = trace::enabled_error_context(&compress_ctx); + let _winner_span = trace::winner_compress_span(winner.id(), before_nbytes).entered(); + let compressed = winner + .compress(self, &data, compress_ctx, exec_ctx) + .inspect_err(|err| { + // NB: this is the only way we can tell which scheme panicked / bailed on their + // data, especially for third-party schemes where the error site may not carry any + // compressor context. + trace::scheme_compress_failed(winner.id(), before_nbytes, error_ctx.as_ref(), err); + })?; + + let after_nbytes = compressed.nbytes(); + let actual_ratio = (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64); + + // TODO(connor): HACK TO SUPPORT L2 DENORMALIZATION!!! + let accepted = after_nbytes < before_nbytes || compressed.is::(); + + trace::record_winner_compress_result( + after_nbytes, + winner_estimate.trace_ratio(), + actual_ratio, + accepted, + ); + + if accepted { + Ok(compressed) + } else { + Ok(data.into_array()) + } + } +} diff --git a/vortex-compressor/src/compressor/constant.rs b/vortex-compressor/src/compressor/constant.rs new file mode 100644 index 00000000000..7e8edc73f2f --- /dev/null +++ b/vortex-compressor/src/compressor/constant.rs @@ -0,0 +1,231 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Built-in constant detection and encoding. +//! +//! Constant arrays are not compressed through a pluggable [`Scheme`]: the compressor always +//! detects constant leaf arrays itself, before evaluating any registered scheme. Detection is +//! skipped while compressing samples, since a constant sample does not imply that the full array +//! is constant. +//! +//! [`Scheme`]: crate::scheme::Scheme + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::fns::is_constant::is_constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::MaskedArray; +use vortex_array::dtype::DType; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; + +use crate::scheme::SchemeId; +use crate::stats::ArrayAndStats; + +/// Synthetic scheme ID reported in traces when the compressor's built-in constant encoding wins. +pub(crate) const CONSTANT_SCHEME_ID: SchemeId = SchemeId { + name: "vortex.compressor.constant", +}; + +/// Returns `true` if all valid values of the canonical array are equal, meaning the array can be +/// encoded by [`compress_constant`]. +/// +/// The caller must have already handled empty and all-null arrays. +/// +/// Uses the cheapest available evidence per type: distinct counts when another scheme already +/// requested them, `O(1)` conclusions from type stats where possible, and otherwise a vectorized +/// equality scan via [`is_constant`]. +/// +/// Note that for types where the check falls through to [`is_constant`] (floats without distinct +/// counts, strings, binary, decimals, and extension types), arrays that contain any nulls are +/// reported as not constant, while stats-based checks detect constant valid values under nulls. +/// This mirrors the behavior of the per-type constant schemes this module replaced. +pub(crate) fn is_constant_for_compression( + data: &ArrayAndStats, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let dtype = data.array().dtype(); + + if matches!(dtype, DType::Bool(_)) { + return Ok(data.bool_stats(exec_ctx).is_constant()); + } + + if dtype.is_int() { + let stats = data.integer_stats(exec_ctx); + + // Distinct counts are only computed when a registered scheme requested them. + if let Some(distinct_count) = stats.distinct_count() { + return Ok(distinct_count == 1); + } + + // If max - min == 0 over the valid values, there is only one distinct value. + return Ok(stats.erased().max_minus_min() == 0); + } + + if dtype.is_float() { + let stats = data.float_stats(exec_ctx); + + if let Some(distinct_count) = stats.distinct_count() { + return Ok(distinct_count == 1); + } + + return is_constant(data.array(), exec_ctx); + } + + if dtype.is_utf8() || dtype.is_binary() { + let stats = data.varbinview_stats(exec_ctx); + + // The estimated distinct count is a lower bound on the actual distinct count, so a value + // above 1 proves the array is not constant without scanning it. + if stats.estimated_distinct_count().is_some_and(|c| c > 1) { + return Ok(false); + } + + return is_constant(data.array(), exec_ctx); + } + + // Decimal, extension, and any other leaf type: fall back to the generic constant check. + is_constant(data.array(), exec_ctx) +} + +/// Encodes an array whose valid values are all equal. +/// +/// Returns a [`ConstantArray`], wrapped in a [`MaskedArray`] when the array has some nulls, or a +/// null [`ConstantArray`] when the array is all-null. +/// +/// # Errors +/// +/// Returns an error if computing validity or extracting the constant scalar fails. +pub(crate) fn compress_constant( + source: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = source.validity()?; + let mask = validity.execute_mask(source.len(), ctx)?; + + let Some(first_valid) = mask.first() else { + return Ok( + ConstantArray::new(Scalar::null(source.dtype().clone()), source.len()).into_array(), + ); + }; + + let scalar = source.execute_scalar(first_valid, ctx)?; + let const_arr = ConstantArray::new(scalar, source.len()).into_array(); + + if mask.all_true() { + Ok(const_arr) + } else { + Ok(MaskedArray::try_new(const_arr, validity)?.into_array()) + } +} + +#[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::Constant; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::Masked; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::TemporalArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::DecimalDType; + use vortex_array::extension::datetime::TimeUnit; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::CascadingCompressor; + + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + + /// Constant detection is built into the compressor, so it must work with no schemes at all. + fn empty_compressor() -> CascadingCompressor { + CascadingCompressor::new(Vec::new()) + } + + #[test] + fn constant_int_compresses_without_schemes() -> VortexResult<()> { + let array = PrimitiveArray::new(buffer![7i64; 100], Validity::NonNullable).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let compressed = empty_compressor().compress(&array, &mut ctx)?; + assert!(compressed.is::()); + Ok(()) + } + + #[test] + fn constant_int_with_nulls_compresses_to_masked_constant() -> VortexResult<()> { + let validity = + Validity::Array(BoolArray::from_iter((0..100).map(|i| i % 10 != 0)).into_array()); + let array = PrimitiveArray::new(buffer![7i64; 100], validity).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let compressed = empty_compressor().compress(&array, &mut ctx)?; + assert!(compressed.is::()); + Ok(()) + } + + #[test] + fn constant_string_compresses_without_schemes() -> VortexResult<()> { + let array = VarBinViewArray::from_iter_str(std::iter::repeat_n("hello", 100)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let compressed = empty_compressor().compress(&array, &mut ctx)?; + assert!(compressed.is::()); + Ok(()) + } + + #[test] + fn constant_bool_compresses_without_schemes() -> VortexResult<()> { + let array = BoolArray::from_iter(std::iter::repeat_n(true, 100)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let compressed = empty_compressor().compress(&array, &mut ctx)?; + assert!(compressed.is::()); + Ok(()) + } + + #[test] + fn constant_decimal_compresses_without_schemes() -> VortexResult<()> { + let array = DecimalArray::new( + buffer![123_456i128; 100], + DecimalDType::new(20, 2), + Validity::NonNullable, + ) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let compressed = empty_compressor().compress(&array, &mut ctx)?; + assert!(compressed.is::()); + Ok(()) + } + + #[test] + fn constant_timestamp_compresses_without_schemes() -> VortexResult<()> { + let ts = PrimitiveArray::from_iter(std::iter::repeat_n(1_704_067_200_000i64, 100)); + let array = TemporalArray::new_timestamp(ts.into_array(), TimeUnit::Milliseconds, None) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let compressed = empty_compressor().compress(&array, &mut ctx)?; + assert!(compressed.is::()); + Ok(()) + } + + #[test] + fn non_constant_int_is_left_canonical_without_schemes() -> VortexResult<()> { + let array = PrimitiveArray::from_iter(0..100i64).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let compressed = empty_compressor().compress(&array, &mut ctx)?; + assert!(!compressed.is::()); + assert_eq!(compressed.dtype(), array.dtype()); + Ok(()) + } +} diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs new file mode 100644 index 00000000000..a661970950c --- /dev/null +++ b/vortex-compressor/src/compressor/mod.rs @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Cascading array compression implementation. + +mod cascade; +mod constant; +mod sample; +mod select; +mod structural; + +use crate::builtins::IntDictScheme; +use crate::scheme::ChildSelection; +use crate::scheme::DescendantExclusion; +use crate::scheme::Scheme; +use crate::scheme::SchemeExt; +use crate::scheme::SchemeId; + +/// Synthetic scheme ID used for the compressor's own root-level cascading. +pub(crate) const ROOT_SCHEME_ID: SchemeId = SchemeId { + name: "vortex.compressor.root", +}; + +/// The main compressor type implementing cascading adaptive compression. +/// +/// This compressor applies adaptive compression [`Scheme`]s to arrays based on their data types and +/// characteristics. It recursively compresses nested structures like structs and lists, and chooses +/// optimal compression schemes for leaf types. +/// +/// The compressor works by: +/// 1. Canonicalizing input arrays to a standard representation. +/// 2. Pre-filtering schemes by [`Scheme::matches`] and exclusion rules. +/// 3. Evaluating each matching scheme's compression estimate and resolving deferred work. +/// 4. Compressing with the best scheme and verifying the result is smaller. +/// +/// 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. + schemes: Vec<&'static dyn Scheme>, + + /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from + /// list offsets). + root_exclusions: Vec, +} + +impl CascadingCompressor { + /// Creates a new compressor with the given schemes. + /// + /// Root-level exclusion rules (e.g. excluding Dict from list offsets) are built automatically. + pub fn new(schemes: Vec<&'static dyn Scheme>) -> Self { + // Root exclusion: exclude IntDict from list/listview offsets (monotonically + // increasing data where dictionary encoding is wasteful). + let root_exclusions = vec![DescendantExclusion { + excluded: IntDictScheme.id(), + children: ChildSelection::One(structural::root_list_children::OFFSETS), + }]; + + Self { + schemes, + root_exclusions, + } + } +} + +// NB: Cascading compression logic is located in `vortex-compressor/src/compressor/cascade.rs`. + +#[cfg(test)] +mod tests; diff --git a/vortex-compressor/src/sample.rs b/vortex-compressor/src/compressor/sample.rs similarity index 74% rename from vortex-compressor/src/sample.rs rename to vortex-compressor/src/compressor/sample.rs index deb5bbe6f37..ba1271d4a6d 100644 --- a/vortex-compressor/src/sample.rs +++ b/vortex-compressor/src/compressor/sample.rs @@ -7,9 +7,20 @@ use rand::RngExt; use rand::SeedableRng; use rand::prelude::StdRng; use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::CascadingCompressor; +use crate::scheme::CompressorContext; +use crate::scheme::EstimateScore; +use crate::scheme::Scheme; +use crate::scheme::SchemeExt; +use crate::stats::ArrayAndStats; +use crate::trace; /// The size of each sampled run. pub const SAMPLE_SIZE: u32 = 64; @@ -130,6 +141,54 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> .collect() } +/// Estimates compression ratio by compressing a ~1% sample of the data. +/// +/// Creates a new [`ArrayAndStats`] for the sample so that stats are generated from the sample, not +/// the full array. +/// +/// # Errors +/// +/// Returns an error if sample compression fails. +pub(super) fn estimate_compression_ratio_with_sampling( + compressor: &CascadingCompressor, + scheme: &S, + array: &ArrayRef, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let sample_array = if compress_ctx.is_sample() { + array.clone() + } else { + let sample_count = sample_count_approx_one_percent(array.len()); + // `ArrayAndStats` expects a canonical array (so that it can easily compute lazy stats). + let canonical: Canonical = sample(array, SAMPLE_SIZE, sample_count).execute(exec_ctx)?; + canonical.into_array() + }; + + let sample_data = ArrayAndStats::new(sample_array, scheme.stats_options()); + let error_ctx = trace::enabled_error_context(&compress_ctx); + let sample_ctx = compress_ctx.with_sampling(); + + let compressed = match scheme.compress(compressor, &sample_data, sample_ctx, exec_ctx) { + Ok(compressed) => compressed, + Err(err) => { + trace::sample_compress_failed(scheme.id(), error_ctx.as_ref(), &err); + return Err(err); + } + }; + + let after = compressed.nbytes(); + let before = sample_data.array().nbytes(); + + let score = EstimateScore::from_sample_sizes(before, after); + + if matches!(score, EstimateScore::ZeroBytes) { + trace::zero_byte_sample_result(scheme.id(), before); + } + + Ok(score) +} + #[cfg(test)] mod tests { use vortex_array::IntoArray; diff --git a/vortex-compressor/src/compressor/select.rs b/vortex-compressor/src/compressor/select.rs new file mode 100644 index 00000000000..3c73d2d4cdb --- /dev/null +++ b/vortex-compressor/src/compressor/select.rs @@ -0,0 +1,187 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scheme selection: estimating each eligible scheme and choosing the winner. + +use vortex_array::ExecutionCtx; +use vortex_error::VortexResult; + +use super::ROOT_SCHEME_ID; +use super::sample::estimate_compression_ratio_with_sampling; +use crate::CascadingCompressor; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::DeferredEstimate; +use crate::scheme::EstimateScore; +use crate::scheme::EstimateVerdict; +use crate::scheme::Scheme; +use crate::scheme::SchemeExt; +use crate::stats::ArrayAndStats; +use crate::trace; + +/// Winner estimate carried from scheme selection into result tracing. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(super) enum WinnerEstimate { + /// The scheme must be used immediately. + AlwaysUse, + /// The scheme won by a ranked estimate. + Score(EstimateScore), +} + +impl WinnerEstimate { + /// Returns the traceable numeric ratio for the winning estimate. + pub(super) fn trace_ratio(self) -> Option { + match self { + Self::AlwaysUse => None, + Self::Score(score) => score.finite_ratio(), + } + } +} + +/// Returns `true` if `score` beats the current best estimate. +fn is_better_score( + score: EstimateScore, + best: Option<&(&'static dyn Scheme, EstimateScore)>, +) -> bool { + score.is_valid() && best.is_none_or(|(_, best_score)| score.beats(*best_score)) +} + +impl CascadingCompressor { + /// Calls [`expected_compression_ratio`] on each candidate and returns the winning scheme along + /// with its resolved winner estimate, or `None` if no scheme beats the canonical encoding. + /// + /// Selection runs in two passes. Pass 1 evaluates every immediate + /// [`CompressionEstimate::Verdict`] and tracks the running best. [`Scheme`]s returning + /// [`CompressionEstimate::Deferred`] are stashed for pass 2 so that we do not make any + /// expensive computations if we don't have to. + /// + /// Pass 2 evaluates the deferred work and, for each [`DeferredEstimate::Callback`], passes the + /// current best [`EstimateScore`] as an early-exit hint so the callback can return + /// [`EstimateVerdict::Skip`] without doing expensive work when it cannot beat the threshold. + /// + /// Ties are broken by registration order within each pass. + /// + /// [`expected_compression_ratio`]: Scheme::expected_compression_ratio + pub(super) fn choose_best_scheme( + &self, + schemes: &[&'static dyn Scheme], + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let mut best: Option<(&'static dyn Scheme, EstimateScore)> = None; + let mut deferred: Vec<(&'static dyn Scheme, DeferredEstimate)> = Vec::new(); + + // Pass 1: evaluate every immediate verdict. Stash deferred work for pass 2. + { + let _verdict_pass = trace::verdict_pass_span().entered(); + for &scheme in schemes { + match scheme.expected_compression_ratio(data, compress_ctx.clone(), exec_ctx) { + CompressionEstimate::Verdict(EstimateVerdict::Skip) => {} + CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) => { + return Ok(Some((scheme, WinnerEstimate::AlwaysUse))); + } + CompressionEstimate::Verdict(EstimateVerdict::Ratio(ratio)) => { + let score = EstimateScore::FiniteCompression(ratio); + + if is_better_score(score, best.as_ref()) { + best = Some((scheme, score)); + } + } + CompressionEstimate::Deferred(deferred_estimate) => { + deferred.push((scheme, deferred_estimate)); + } + } + } + } + + // Pass 2: run deferred work. Callbacks receive the current best as a threshold so they can + // short-circuit with `Skip` when they cannot beat it. + for (scheme, deferred_estimate) in deferred { + let _span = trace::scheme_eval_span(scheme.id()).entered(); + let threshold: Option = best.map(|(_, score)| score); + match deferred_estimate { + DeferredEstimate::Sample => { + let score = estimate_compression_ratio_with_sampling( + self, + scheme, + data.array(), + compress_ctx.clone(), + exec_ctx, + )?; + + if is_better_score(score, best.as_ref()) { + best = Some((scheme, score)); + } + } + DeferredEstimate::Callback(callback) => { + match callback(self, data, threshold, compress_ctx.clone(), exec_ctx)? { + EstimateVerdict::Skip => {} + EstimateVerdict::AlwaysUse => { + return Ok(Some((scheme, WinnerEstimate::AlwaysUse))); + } + EstimateVerdict::Ratio(ratio) => { + let score = EstimateScore::FiniteCompression(ratio); + + if is_better_score(score, best.as_ref()) { + best = Some((scheme, score)); + } + } + } + } + } + } + + Ok(best.map(|(scheme, score)| (scheme, WinnerEstimate::Score(score)))) + } + + // TODO(connor): Lots of room for optimization here. + /// Returns `true` if the candidate scheme should be excluded based on the cascade history and + /// exclusion rules. + pub(super) fn is_excluded(&self, candidate: &dyn Scheme, ctx: &CompressorContext) -> bool { + let id = candidate.id(); + let history = ctx.cascade_history(); + + // Self-exclusion: no scheme appears twice in any chain. + if history.iter().any(|&(sid, _)| sid == id) { + return true; + } + + let mut iter = history.iter().copied().peekable(); + + // The root entry is always first in the history (if present). Check if the root has + // excluded us. + if let Some((_, child_idx)) = iter.next_if(|&(sid, _)| sid == ROOT_SCHEME_ID) + && self + .root_exclusions + .iter() + .any(|rule| rule.excluded == id && rule.children.contains(child_idx)) + { + return true; + } + + // Push rules: Check if any of our ancestors have excluded us. + for (ancestor_id, child_idx) in iter { + if let Some(ancestor) = self.schemes.iter().find(|s| s.id() == ancestor_id) + && ancestor + .descendant_exclusions() + .iter() + .any(|rule| rule.excluded == id && rule.children.contains(child_idx)) + { + return true; + } + } + + // Pull rules: Check if we have excluded ourselves because of our ancestors. + for rule in candidate.ancestor_exclusions() { + if history + .iter() + .any(|(sid, cidx)| *sid == rule.ancestor && rule.children.contains(*cidx)) + { + return true; + } + } + + false + } +} diff --git a/vortex-compressor/src/compressor/structural.rs b/vortex-compressor/src/compressor/structural.rs new file mode 100644 index 00000000000..ec867c1c1dc --- /dev/null +++ b/vortex-compressor/src/compressor/structural.rs @@ -0,0 +1,128 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Recursive compression of structural arrays: lists, list views, and physical slots. + +use vortex_array::ArrayRef; +use vortex_array::ArraySlots; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::list::ListArrayExt; +use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_error::VortexResult; + +use super::ROOT_SCHEME_ID; +use crate::CascadingCompressor; +use crate::scheme::CompressorContext; + +/// Child indices for the compressor's list/listview compression. +pub(super) mod root_list_children { + /// List/ListView offsets child. + pub const OFFSETS: usize = 1; + /// ListView sizes child. + pub const SIZES: usize = 2; +} + +impl CascadingCompressor { + /// Compresses a [`ListArray`] by narrowing offsets and recursively compressing elements. + pub(super) fn compress_list_array( + &self, + list_array: ListArray, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let list_array = list_array.reset_offsets(true, exec_ctx)?; + + let compressed_elems = self.compress(list_array.elements(), exec_ctx)?; + + // Record the root scheme with the offsets child index so root exclusion rules apply. + let offset_ctx = + compress_ctx.descend_with_scheme(ROOT_SCHEME_ID, root_list_children::OFFSETS); + let list_offsets_primitive = list_array + .offsets() + .clone() + .execute::(exec_ctx)? + .narrow(exec_ctx)?; + let compressed_offsets = self.compress_canonical( + Canonical::Primitive(list_offsets_primitive), + offset_ctx, + exec_ctx, + )?; + + Ok( + ListArray::try_new(compressed_elems, compressed_offsets, list_array.validity()?)? + .into_array(), + ) + } + + /// Compresses a [`ListViewArray`] by narrowing offsets/sizes and recursively compressing + /// elements. + pub(super) fn compress_list_view_array( + &self, + list_view: ListViewArray, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let compressed_elems = self.compress(list_view.elements(), exec_ctx)?; + + let offset_ctx = compress_ctx + .clone() + .descend_with_scheme(ROOT_SCHEME_ID, root_list_children::OFFSETS); + let list_view_offsets_primitive = list_view + .offsets() + .clone() + .execute::(exec_ctx)? + .narrow(exec_ctx)?; + let compressed_offsets = self.compress_canonical( + Canonical::Primitive(list_view_offsets_primitive), + offset_ctx, + exec_ctx, + )?; + + let sizes_ctx = compress_ctx.descend_with_scheme(ROOT_SCHEME_ID, root_list_children::SIZES); + let list_view_sizes_primitive = list_view + .sizes() + .clone() + .execute::(exec_ctx)? + .narrow(exec_ctx)?; + let compressed_sizes = self.compress_canonical( + Canonical::Primitive(list_view_sizes_primitive), + sizes_ctx, + exec_ctx, + )?; + + Ok(ListViewArray::try_new( + compressed_elems, + compressed_offsets, + compressed_sizes, + list_view.validity()?, + )? + .into_array()) + } + + /// Compress very child slot of the array, then re-build it from them. + pub(super) fn compress_physical_slots( + &self, + array: &ArrayRef, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let slots = array + .slots() + .iter() + .map(|slot| { + slot.as_ref() + .map(|child| self.compress(child, exec_ctx)) + .transpose() + }) + .collect::>()?; + + // 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) } + } +} diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs new file mode 100644 index 00000000000..08eafa369de --- /dev/null +++ b/vortex-compressor/src/compressor/tests.rs @@ -0,0 +1,706 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use parking_lot::Mutex; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::NullArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::validity::Validity; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use super::CascadingCompressor; +use super::ROOT_SCHEME_ID; +use super::sample::estimate_compression_ratio_with_sampling; +use super::select::WinnerEstimate; +use super::structural; +use crate::builtins::FloatDictScheme; +use crate::builtins::IntDictScheme; +use crate::builtins::StringDictScheme; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::DeferredEstimate; +use crate::scheme::EstimateScore; +use crate::scheme::EstimateVerdict; +use crate::scheme::Scheme; +use crate::scheme::SchemeExt; +use crate::stats::ArrayAndStats; +use crate::stats::GenerateStatsOptions; + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +fn compressor() -> CascadingCompressor { + CascadingCompressor::new(vec![&IntDictScheme, &FloatDictScheme, &StringDictScheme]) +} + +fn estimate_test_data() -> ArrayAndStats { + let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + ArrayAndStats::new(array, GenerateStatsOptions::default()) +} + +fn matches_integer_primitive(canonical: &Canonical) -> bool { + matches!(canonical, Canonical::Primitive(primitive) if primitive.ptype().is_int()) +} + +#[derive(Debug)] +struct DirectRatioScheme; + +impl Scheme for DirectRatioScheme { + fn scheme_name(&self) -> &'static str { + "test.direct_ratio" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::Ratio(2.0)) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper should never be selected for compression") + } +} + +#[derive(Debug)] +struct ImmediateAlwaysUseScheme; + +impl Scheme for ImmediateAlwaysUseScheme { + fn scheme_name(&self) -> &'static str { + "test.immediate_always_use" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper should never be selected for compression") + } +} + +#[derive(Debug)] +struct CallbackAlwaysUseScheme; + +impl Scheme for CallbackAlwaysUseScheme { + fn scheme_name(&self) -> &'static str { + "test.callback_always_use" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + |_compressor, _data, _ctx, _exec_ctx, _best_so_far| Ok(EstimateVerdict::AlwaysUse), + ))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper should never be selected for compression") + } +} + +#[derive(Debug)] +struct CallbackSkipScheme; + +impl Scheme for CallbackSkipScheme { + fn scheme_name(&self) -> &'static str { + "test.callback_skip" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + |_compressor, _data, _ctx, _exec_ctx, _best_so_far| Ok(EstimateVerdict::Skip), + ))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper should never be selected for compression") + } +} + +#[derive(Debug)] +struct CallbackRatioScheme; + +impl Scheme for CallbackRatioScheme { + fn scheme_name(&self) -> &'static str { + "test.callback_ratio" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + |_compressor, _data, _ctx, _exec_ctx, _best_so_far| Ok(EstimateVerdict::Ratio(3.0)), + ))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper should never be selected for compression") + } +} + +#[derive(Debug)] +struct HugeRatioScheme; + +impl Scheme for HugeRatioScheme { + fn scheme_name(&self) -> &'static str { + "test.huge_ratio" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::Ratio(100.0)) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper should never be selected for compression") + } +} + +#[derive(Debug)] +struct ZeroBytesSamplingScheme; + +impl Scheme for ZeroBytesSamplingScheme { + fn scheme_name(&self) -> &'static str { + "test.zero_bytes_sampling" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Sample) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(NullArray::new(data.array().len()).into_array()) + } +} + +#[test] +fn test_self_exclusion() { + let c = compressor(); + let ctx = CompressorContext::default().descend_with_scheme(IntDictScheme.id(), 0); + + // IntDictScheme is in the history, so it should be excluded. + assert!(c.is_excluded(&IntDictScheme, &ctx)); +} + +#[test] +fn test_root_exclusion_list_offsets() { + let c = compressor(); + let ctx = CompressorContext::default() + .descend_with_scheme(ROOT_SCHEME_ID, structural::root_list_children::OFFSETS); + + // IntDict should be excluded for list offsets. + assert!(c.is_excluded(&IntDictScheme, &ctx)); +} + +#[test] +fn test_push_rule_float_dict_excludes_int_dict_from_codes() { + let c = compressor(); + // FloatDict cascading through codes (child 1). + let ctx = CompressorContext::default().descend_with_scheme(FloatDictScheme.id(), 1); + + // IntDict should be excluded from FloatDict's codes child. + assert!(c.is_excluded(&IntDictScheme, &ctx)); +} + +#[test] +fn test_push_rule_float_dict_excludes_int_dict_from_values() { + let c = compressor(); + // FloatDict cascading through values (child 0). + let ctx = CompressorContext::default().descend_with_scheme(FloatDictScheme.id(), 0); + + // IntDict should also be excluded from FloatDict's values child (ALP propagation + // replacement). + assert!(c.is_excluded(&IntDictScheme, &ctx)); +} + +#[test] +fn test_no_exclusion_without_history() { + let c = compressor(); + let ctx = CompressorContext::default(); + + // No history means no exclusions. + assert!(!c.is_excluded(&IntDictScheme, &ctx)); +} + +#[test] +fn immediate_always_use_wins_immediately() -> VortexResult<()> { + let compressor = CascadingCompressor::new(vec![&DirectRatioScheme, &ImmediateAlwaysUseScheme]); + let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &ImmediateAlwaysUseScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(matches!( + winner, + Some((scheme, WinnerEstimate::AlwaysUse)) + if scheme.id() == ImmediateAlwaysUseScheme.id() + )); + Ok(()) +} + +#[test] +fn callback_always_use_wins_immediately() -> VortexResult<()> { + let compressor = CascadingCompressor::new(vec![&DirectRatioScheme, &CallbackAlwaysUseScheme]); + let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &CallbackAlwaysUseScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(matches!( + winner, + Some((scheme, WinnerEstimate::AlwaysUse)) + if scheme.id() == CallbackAlwaysUseScheme.id() + )); + Ok(()) +} + +#[test] +fn callback_skip_is_ignored() -> VortexResult<()> { + let compressor = CascadingCompressor::new(vec![&CallbackSkipScheme, &DirectRatioScheme]); + let schemes: [&'static dyn Scheme; 2] = [&CallbackSkipScheme, &DirectRatioScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(matches!( + winner, + Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(2.0)))) + if scheme.id() == DirectRatioScheme.id() + )); + Ok(()) +} + +#[test] +fn callback_ratio_competes_numerically() -> VortexResult<()> { + let compressor = CascadingCompressor::new(vec![&DirectRatioScheme, &CallbackRatioScheme]); + let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &CallbackRatioScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(matches!( + winner, + Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(3.0)))) + if scheme.id() == CallbackRatioScheme.id() + )); + Ok(()) +} + +#[test] +fn zero_byte_sample_loses_to_finite_ratio() -> VortexResult<()> { + let compressor = CascadingCompressor::new(vec![&HugeRatioScheme, &ZeroBytesSamplingScheme]); + let schemes: [&'static dyn Scheme; 2] = [&HugeRatioScheme, &ZeroBytesSamplingScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(matches!( + winner, + Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(100.0)))) + if scheme.id() == HugeRatioScheme.id() + )); + Ok(()) +} + +#[test] +fn finite_ratio_displaces_zero_byte_sample() -> VortexResult<()> { + let compressor = CascadingCompressor::new(vec![&ZeroBytesSamplingScheme, &HugeRatioScheme]); + let schemes: [&'static dyn Scheme; 2] = [&ZeroBytesSamplingScheme, &HugeRatioScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(matches!( + winner, + Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(100.0)))) + if scheme.id() == HugeRatioScheme.id() + )); + Ok(()) +} + +#[test] +fn zero_byte_sample_alone_selects_no_scheme() -> VortexResult<()> { + let compressor = CascadingCompressor::new(vec![&ZeroBytesSamplingScheme]); + let schemes: [&'static dyn Scheme; 1] = [&ZeroBytesSamplingScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(winner.is_none()); + Ok(()) +} + +// Observer helper used by threshold-related tests. Captures the `best_so_far` value the +// compressor passes to its deferred callback. `OBSERVER_LOCK` serializes tests that share +// `OBSERVED_THRESHOLD` so they do not race. +static OBSERVER_LOCK: Mutex<()> = Mutex::new(()); +static OBSERVED_THRESHOLD: Mutex>> = Mutex::new(None); + +#[derive(Debug)] +struct ThresholdObservingScheme; + +impl Scheme for ThresholdObservingScheme { + fn scheme_name(&self) -> &'static str { + "test.threshold_observing" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + |_compressor, _data, best_so_far, _ctx, _exec_ctx| { + *OBSERVED_THRESHOLD.lock() = Some(best_so_far); + Ok(EstimateVerdict::Skip) + }, + ))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper should never be selected for compression") + } +} + +#[derive(Debug)] +struct CallbackMatchingRatioScheme; + +impl Scheme for CallbackMatchingRatioScheme { + fn scheme_name(&self) -> &'static str { + "test.callback_matching_ratio" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + |_compressor, _data, _ctx, _exec_ctx, _best_so_far| Ok(EstimateVerdict::Ratio(2.0)), + ))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper should never be selected for compression") + } +} + +#[test] +fn callback_always_use_overrides_pass_one_best() -> VortexResult<()> { + // `HugeRatioScheme` returns an immediate `Ratio(100.0)` in pass 1; + // `CallbackAlwaysUseScheme` returns `AlwaysUse` from its deferred callback in pass 2. + // The deferred `AlwaysUse` must still win. + let compressor = CascadingCompressor::new(vec![&HugeRatioScheme, &CallbackAlwaysUseScheme]); + let schemes: [&'static dyn Scheme; 2] = [&HugeRatioScheme, &CallbackAlwaysUseScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(matches!( + winner, + Some((scheme, WinnerEstimate::AlwaysUse)) + if scheme.id() == CallbackAlwaysUseScheme.id() + )); + Ok(()) +} + +#[test] +fn threshold_reflects_pass_one_best() -> VortexResult<()> { + let _guard = OBSERVER_LOCK.lock(); + *OBSERVED_THRESHOLD.lock() = None; + + let compressor = CascadingCompressor::new(vec![&DirectRatioScheme, &ThresholdObservingScheme]); + let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &ThresholdObservingScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + let observed = *OBSERVED_THRESHOLD.lock(); + assert!(matches!( + observed, + Some(Some(EstimateScore::FiniteCompression(r))) if r == 2.0 + )); + Ok(()) +} + +#[test] +fn threshold_is_none_when_only_prior_is_zero_bytes() -> VortexResult<()> { + let _guard = OBSERVER_LOCK.lock(); + *OBSERVED_THRESHOLD.lock() = None; + + let compressor = + CascadingCompressor::new(vec![&ZeroBytesSamplingScheme, &ThresholdObservingScheme]); + let schemes: [&'static dyn Scheme; 2] = [&ZeroBytesSamplingScheme, &ThresholdObservingScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + // The observing callback was invoked (outer `Some`) and `best_so_far` was `None` (inner + // `None`) because the zero-byte sample is never stored as the best. + let observed = *OBSERVED_THRESHOLD.lock(); + assert_eq!(observed, Some(None)); + Ok(()) +} + +#[test] +fn threshold_is_none_when_no_prior_scheme() -> VortexResult<()> { + let _guard = OBSERVER_LOCK.lock(); + *OBSERVED_THRESHOLD.lock() = None; + + let compressor = CascadingCompressor::new(vec![&ThresholdObservingScheme]); + let schemes: [&'static dyn Scheme; 1] = [&ThresholdObservingScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + let observed = *OBSERVED_THRESHOLD.lock(); + assert_eq!(observed, Some(None)); + Ok(()) +} + +#[test] +fn threshold_updates_from_earlier_deferred_callback() -> VortexResult<()> { + let _guard = OBSERVER_LOCK.lock(); + *OBSERVED_THRESHOLD.lock() = None; + + // Both schemes are deferred. The first callback registers `Ratio(3.0)`; the second + // callback must observe it as its threshold. + let compressor = + CascadingCompressor::new(vec![&CallbackRatioScheme, &ThresholdObservingScheme]); + let schemes: [&'static dyn Scheme; 2] = [&CallbackRatioScheme, &ThresholdObservingScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + let observed = *OBSERVED_THRESHOLD.lock(); + assert!(matches!( + observed, + Some(Some(EstimateScore::FiniteCompression(r))) if r == 3.0 + )); + Ok(()) +} + +#[test] +fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<()> { + // Both schemes produce the same `Ratio(2.0)`, one from pass 1 (immediate) and one from + // pass 2 (deferred callback). Pass 1 locks in first, and strict `>` tie-breaking means + // the deferred callback's equal ratio cannot displace it. + let compressor = + CascadingCompressor::new(vec![&CallbackMatchingRatioScheme, &DirectRatioScheme]); + let schemes: [&'static dyn Scheme; 2] = [&CallbackMatchingRatioScheme, &DirectRatioScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(matches!( + winner, + Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(r)))) + if scheme.id() == DirectRatioScheme.id() && r == 2.0 + )); + Ok(()) +} + +#[test] +fn all_null_array_compresses_to_constant() -> VortexResult<()> { + let array = PrimitiveArray::new( + buffer![0i32, 0, 0, 0, 0], + Validity::Array(BoolArray::from_iter([false, false, false, false, false]).into_array()), + ) + .into_array(); + + // The compressor should produce a `ConstantArray` for an all-null array regardless of + // which schemes are registered. + let compressor = CascadingCompressor::new(vec![&IntDictScheme]); + let mut exec_ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array, &mut exec_ctx)?; + assert!(compressed.is::()); + Ok(()) +} + +/// Regression test for . +/// +/// `estimate_compression_ratio_with_sampling` must use the *scheme's* stats options +/// (which request distinct-value counting) rather than the context's stats options +/// (which may not). With the old code this panicked inside `dictionary_encode` because +/// distinct values were never computed for the sample. +#[test] +fn sampling_uses_scheme_stats_options() -> VortexResult<()> { + // Low-cardinality float array so FloatDictScheme considers it compressible. + let array = PrimitiveArray::new( + buffer![1.0f32, 2.0, 1.0, 2.0, 1.0, 2.0, 1.0, 2.0], + Validity::NonNullable, + ) + .into_array(); + + let compressor = CascadingCompressor::new(vec![&FloatDictScheme]); + + // A context with default stats_options (count_distinct_values = false) and + // marked as a sample so the function skips the sampling step and compresses + // the array directly. + let ctx = CompressorContext::new().with_sampling(); + + // Before the fix this panicked with: + // "this must be present since `DictScheme` declared that we need distinct values" + let mut exec_ctx = SESSION.create_execution_ctx(); + let score = estimate_compression_ratio_with_sampling( + &compressor, + &FloatDictScheme, + &array, + ctx, + &mut exec_ctx, + )?; + assert!(matches!(score, EstimateScore::FiniteCompression(ratio) if ratio.is_finite())); + Ok(()) +} diff --git a/vortex-compressor/src/lib.rs b/vortex-compressor/src/lib.rs index 7e6854eacbe..55bb9b188f6 100644 --- a/vortex-compressor/src/lib.rs +++ b/vortex-compressor/src/lib.rs @@ -19,7 +19,8 @@ //! # 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 +//! still canonicalizes supported inputs, recursively handles nested structure, and encodes +//! constant leaves (constant detection is built into the compressor), but no other leaf //! compression is selected. //! //! ```rust @@ -62,13 +63,10 @@ //! with a short `jq` query. pub mod builtins; -pub mod ctx; -pub mod estimate; pub mod scheme; pub mod stats; -mod sample; - mod compressor; -mod trace; pub use compressor::CascadingCompressor; + +mod trace; diff --git a/vortex-compressor/src/ctx.rs b/vortex-compressor/src/scheme/ctx.rs similarity index 95% rename from vortex-compressor/src/ctx.rs rename to vortex-compressor/src/scheme/ctx.rs index 4e0619ff8ee..4eed7538daa 100644 --- a/vortex-compressor/src/ctx.rs +++ b/vortex-compressor/src/scheme/ctx.rs @@ -44,7 +44,7 @@ impl CompressorContext { /// Creates a new `CompressorContext`. /// /// This should **only** be created by the compressor. - pub(super) fn new() -> Self { + pub(crate) fn new() -> Self { Self { is_sample: false, allowed_cascading: MAX_CASCADE, @@ -102,13 +102,13 @@ impl CompressorContext { } /// Returns a context with the given stats options. - pub(super) fn with_merged_stats_options(mut self, opts: GenerateStatsOptions) -> Self { + pub(crate) fn with_merged_stats_options(mut self, opts: GenerateStatsOptions) -> Self { self.merged_stats_options = opts; self } /// Returns a context marked as sample compression. - pub(super) fn with_sampling(mut self) -> Self { + pub(crate) fn with_sampling(mut self) -> Self { self.is_sample = true; self } @@ -118,7 +118,7 @@ impl CompressorContext { /// /// The `child_index` identifies which child of the scheme is being compressed (e.g. for /// Dict: values=0, codes=1). - pub(super) fn descend_with_scheme(mut self, id: SchemeId, child_index: usize) -> Self { + pub(crate) fn descend_with_scheme(mut self, id: SchemeId, child_index: usize) -> Self { self.allowed_cascading = self .allowed_cascading .checked_sub(1) diff --git a/vortex-compressor/src/estimate.rs b/vortex-compressor/src/scheme/estimate.rs similarity index 63% rename from vortex-compressor/src/estimate.rs rename to vortex-compressor/src/scheme/estimate.rs index 70dd75d13b0..20b9dc77e9b 100644 --- a/vortex-compressor/src/estimate.rs +++ b/vortex-compressor/src/scheme/estimate.rs @@ -1,25 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Compression ratio estimation types and sampling-based estimation. +//! Compression ratio estimation types returned by schemes. use std::fmt; -use vortex_array::ArrayRef; -use vortex_array::Canonical; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; use vortex_error::VortexResult; use crate::CascadingCompressor; -use crate::ctx::CompressorContext; -use crate::sample::SAMPLE_SIZE; -use crate::sample::sample; -use crate::sample::sample_count_approx_one_percent; -use crate::scheme::Scheme; -use crate::scheme::SchemeExt; +use crate::scheme::CompressorContext; use crate::stats::ArrayAndStats; -use crate::trace; /// Closure type for [`DeferredEstimate::Callback`]. /// @@ -44,10 +35,12 @@ pub type EstimateFn = dyn FnOnce( + Send + Sync; -/// The result of a [`Scheme`]'s compression ratio estimation. +/// The result of a [`Scheme`](crate::scheme::Scheme)'s compression ratio estimation. /// -/// This type is returned by [`Scheme::expected_compression_ratio`] to tell the compressor how -/// promising this scheme is for a given array without performing any expensive work. +/// This type is returned by +/// [`Scheme::expected_compression_ratio`](crate::scheme::Scheme::expected_compression_ratio) to +/// tell the compressor how promising this scheme is for a given array without performing any +/// expensive work. /// /// [`CompressionEstimate::Verdict`] means the scheme already knows the terminal answer. /// [`CompressionEstimate::Deferred`] means the compressor must do extra work before the scheme can @@ -69,7 +62,7 @@ pub enum EstimateVerdict { /// Always use this scheme, as it is definitively the best choice. /// - /// Some examples include constant detection, decimal byte parts, and temporal decomposition. + /// Some examples include decimal byte parts and temporal decomposition. /// /// The compressor will select this scheme immediately without evaluating further candidates. /// Schemes that return `AlwaysUse` must be mutually exclusive per canonical type (enforced by @@ -119,7 +112,7 @@ pub enum EstimateScore { impl EstimateScore { /// Converts measured sample sizes into a ranked estimate. - pub(super) fn from_sample_sizes(before_nbytes: u64, after_nbytes: u64) -> Self { + pub(crate) fn from_sample_sizes(before_nbytes: u64, after_nbytes: u64) -> Self { if after_nbytes == 0 { Self::ZeroBytes } else { @@ -139,7 +132,7 @@ impl EstimateScore { } /// Returns whether this estimate is eligible to compete. - fn is_valid(self) -> bool { + pub(crate) fn is_valid(self) -> bool { match self { Self::FiniteCompression(ratio) => { ratio.is_finite() && !ratio.is_subnormal() && ratio > 1.0 @@ -149,7 +142,7 @@ impl EstimateScore { } /// Returns whether this estimate beats another valid estimate. - fn beats(self, other: Self) -> bool { + pub(crate) fn beats(self, other: Self) -> bool { match (self, other) { (Self::ZeroBytes, _) => false, (Self::FiniteCompression(_), Self::ZeroBytes) => true, @@ -160,81 +153,6 @@ impl EstimateScore { } } -/// Winner estimate carried from scheme selection into result tracing. -#[derive(Debug, Clone, Copy, PartialEq)] -pub(super) enum WinnerEstimate { - /// The scheme must be used immediately. - AlwaysUse, - /// The scheme won by a ranked estimate. - Score(EstimateScore), -} - -impl WinnerEstimate { - /// Returns the traceable numeric ratio for the winning estimate. - pub(super) fn trace_ratio(self) -> Option { - match self { - Self::AlwaysUse => None, - Self::Score(score) => score.finite_ratio(), - } - } -} - -/// Returns `true` if `score` beats the current best estimate. -pub(super) fn is_better_score( - score: EstimateScore, - best: Option<&(&'static dyn Scheme, EstimateScore)>, -) -> bool { - score.is_valid() && best.is_none_or(|(_, best_score)| score.beats(*best_score)) -} - -/// Estimates compression ratio by compressing a ~1% sample of the data. -/// -/// Creates a new [`ArrayAndStats`] for the sample so that stats are generated from the sample, not -/// the full array. -/// -/// # Errors -/// -/// Returns an error if sample compression fails. -pub(super) fn estimate_compression_ratio_with_sampling( - compressor: &CascadingCompressor, - scheme: &S, - array: &ArrayRef, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, -) -> VortexResult { - let sample_array = if compress_ctx.is_sample() { - array.clone() - } else { - let sample_count = sample_count_approx_one_percent(array.len()); - // `ArrayAndStats` expects a canonical array (so that it can easily compute lazy stats). - let canonical: Canonical = sample(array, SAMPLE_SIZE, sample_count).execute(exec_ctx)?; - canonical.into_array() - }; - - let sample_data = ArrayAndStats::new(sample_array, scheme.stats_options()); - let error_ctx = trace::enabled_error_context(&compress_ctx); - let sample_ctx = compress_ctx.with_sampling(); - - let compressed = match scheme.compress(compressor, &sample_data, sample_ctx, exec_ctx) { - Ok(compressed) => compressed, - Err(err) => { - trace::sample_compress_failed(scheme.id(), error_ctx.as_ref(), &err); - return Err(err); - } - }; - - let after = compressed.nbytes(); - let before = sample_data.array().nbytes(); - - let score = EstimateScore::from_sample_sizes(before, after); - - if matches!(score, EstimateScore::ZeroBytes) { - trace::zero_byte_sample_result(scheme.id(), before); - } - - Ok(score) -} - impl fmt::Debug for DeferredEstimate { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { diff --git a/vortex-compressor/src/scheme/exclusion.rs b/vortex-compressor/src/scheme/exclusion.rs new file mode 100644 index 00000000000..2dba6b85046 --- /dev/null +++ b/vortex-compressor/src/scheme/exclusion.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Exclusion rules that keep incompatible schemes out of a cascade chain. + +use crate::scheme::SchemeId; + +/// Selects which children of a cascading scheme a rule applies to. +#[derive(Debug, Clone, Copy)] +pub enum ChildSelection { + /// Rule applies to all children. + All, + /// Rule applies to a single child. + One(usize), + /// Rule applies to multiple specific children. + Many(&'static [usize]), +} + +impl ChildSelection { + /// Returns `true` if this selection includes the given child index. + pub fn contains(&self, child_index: usize) -> bool { + match self { + ChildSelection::All => true, + ChildSelection::One(idx) => *idx == child_index, + ChildSelection::Many(indices) => indices.contains(&child_index), + } + } +} + +/// Push rule: declared by a cascading scheme to exclude another scheme from the subtree +/// rooted at the specified children. +/// +/// Use this when the declaring scheme (the ancestor) knows about the excluded scheme. For example, +/// `ZigZag` excludes `Dict` from all its children. +#[derive(Debug, Clone, Copy)] +pub struct DescendantExclusion { + /// The scheme to exclude from descendants. + pub excluded: SchemeId, + /// Which children of the declaring scheme this rule applies to. + pub children: ChildSelection, +} + +/// Pull rule: declared by a scheme to exclude itself when the specified ancestor is in the +/// cascade chain. +/// +/// Use this when the excluded scheme (the descendant) knows about the ancestor. For example, +/// `Sequence` excludes itself when `IntDict` is an ancestor on its codes child. +#[derive(Debug, Clone, Copy)] +pub struct AncestorExclusion { + /// The ancestor scheme that makes the declaring scheme ineligible. + pub ancestor: SchemeId, + /// Which children of the ancestor this rule applies to. + pub children: ChildSelection, +} diff --git a/vortex-compressor/src/scheme.rs b/vortex-compressor/src/scheme/mod.rs similarity index 77% rename from vortex-compressor/src/scheme.rs rename to vortex-compressor/src/scheme/mod.rs index 6e6899a3ea1..341acc8acbf 100644 --- a/vortex-compressor/src/scheme.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -1,21 +1,34 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Unified compression scheme trait and exclusion rules. +//! Everything a scheme author implements or receives: the [`Scheme`] trait, exclusion rules, +//! compression estimates, and the compression context. +mod ctx; +pub use ctx::CompressorContext; +pub use ctx::MAX_CASCADE; + +pub(crate) mod estimate; +mod exclusion; use std::fmt; use std::fmt::Debug; use std::hash::Hash; use std::hash::Hasher; +pub use estimate::CompressionEstimate; +pub use estimate::DeferredEstimate; +pub use estimate::EstimateFn; +pub use estimate::EstimateScore; +pub use estimate::EstimateVerdict; +pub use exclusion::AncestorExclusion; +pub use exclusion::ChildSelection; +pub use exclusion::DescendantExclusion; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_error::VortexResult; use crate::CascadingCompressor; -use crate::ctx::CompressorContext; -use crate::estimate::CompressionEstimate; use crate::stats::ArrayAndStats; use crate::stats::GenerateStatsOptions; @@ -39,54 +52,6 @@ impl fmt::Display for SchemeId { } } -/// Selects which children of a cascading scheme a rule applies to. -#[derive(Debug, Clone, Copy)] -pub enum ChildSelection { - /// Rule applies to all children. - All, - /// Rule applies to a single child. - One(usize), - /// Rule applies to multiple specific children. - Many(&'static [usize]), -} - -impl ChildSelection { - /// Returns `true` if this selection includes the given child index. - pub fn contains(&self, child_index: usize) -> bool { - match self { - ChildSelection::All => true, - ChildSelection::One(idx) => *idx == child_index, - ChildSelection::Many(indices) => indices.contains(&child_index), - } - } -} - -/// Push rule: declared by a cascading scheme to exclude another scheme from the subtree -/// rooted at the specified children. -/// -/// Use this when the declaring scheme (the ancestor) knows about the excluded scheme. For example, -/// `ZigZag` excludes `Dict` from all its children. -#[derive(Debug, Clone, Copy)] -pub struct DescendantExclusion { - /// The scheme to exclude from descendants. - pub excluded: SchemeId, - /// Which children of the declaring scheme this rule applies to. - pub children: ChildSelection, -} - -/// Pull rule: declared by a scheme to exclude itself when the specified ancestor is in the -/// cascade chain. -/// -/// Use this when the excluded scheme (the descendant) knows about the ancestor. For example, -/// `Sequence` excludes itself when `IntDict` is an ancestor on its codes child. -#[derive(Debug, Clone, Copy)] -pub struct AncestorExclusion { - /// The ancestor scheme that makes the declaring scheme ineligible. - pub ancestor: SchemeId, - /// Which children of the ancestor this rule applies to. - pub children: ChildSelection, -} - // TODO(connor): Remove all default implemented methods. /// A single compression encoding that the [`CascadingCompressor`] can select from. /// @@ -199,20 +164,20 @@ pub trait Scheme: Debug + Send + Sync { /// entire array. /// /// [`CompressionEstimate::Verdict`] means the scheme already knows the terminal - /// [`crate::estimate::EstimateVerdict`]. `CompressionEstimate::Deferred(DeferredEstimate::Sample)` + /// [`crate::scheme::EstimateVerdict`]. `CompressionEstimate::Deferred(DeferredEstimate::Sample)` /// asks the compressor to sample. `CompressionEstimate::Deferred(DeferredEstimate::Callback(...))` /// asks the compressor to run custom deferred work. Deferred callbacks must return a - /// [`crate::estimate::EstimateVerdict`] directly, never another deferred request. + /// [`crate::scheme::EstimateVerdict`] directly, never another deferred request. /// /// Note that the compressor will also use this method when compressing samples, so some /// statistics that might hold for the samples may not hold for the entire array (e.g., - /// `Constant`). Implementations should check `ctx.is_sample` to make sure that they are + /// constancy). Implementations should check `ctx.is_sample` to make sure that they are /// returning the correct information. /// /// The compressor guarantees that empty and all-null arrays are handled before this method is - /// called. Implementations may assume the array has at least one valid element. However, a - /// constant scheme should still be registered with the compressor to detect single-value arrays - /// that are not all-null. + /// called, so implementations may assume the array has at least one valid element. Outside of + /// sample compression, the compressor also encodes constant arrays itself before evaluating + /// schemes, so implementations only see constant arrays when `ctx.is_sample()` is `true`. fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-compressor/src/trace.rs b/vortex-compressor/src/trace.rs index 84027272b31..ee594621199 100644 --- a/vortex-compressor/src/trace.rs +++ b/vortex-compressor/src/trace.rs @@ -5,7 +5,7 @@ use std::fmt; -use crate::ctx::CompressorContext; +use crate::scheme::CompressorContext; use crate::scheme::SchemeId; /// Shared tracing target for compressor decisions and coarse cascade structure. diff --git a/vortex-compute/src/lane_kernels/map_into.rs b/vortex-compute/src/lane_kernels/map_into.rs index d20b261c274..258913e9fc7 100644 --- a/vortex-compute/src/lane_kernels/map_into.rs +++ b/vortex-compute/src/lane_kernels/map_into.rs @@ -156,6 +156,67 @@ pub trait IndexedSourceExt: IndexedSource + Sized { } } + /// Apply the predicate `f(value)` lane-by-lane and bit-pack the results into + /// `words`, LSB-first, 64 lanes per `u64`. + /// + /// This is the kernel shape behind comparison operators: each lane read is an + /// independent indexed load (drive two columns via [`LaneZip`]) and the 64 + /// per-lane booleans of a chunk reduce into a single word with `OR + shift`, + /// which the autovectorizer lowers to a vector compare plus movemask. + /// + /// Words are written with `=` (not `|=`), so `words` need not be + /// zero-initialised. Bits at positions `>= self.len()` in the last word are + /// written as zero. + /// + /// Like [`map_into`], this kernel has no validity awareness; pair the packed + /// bits with a separately computed validity mask. + /// + /// [`LaneZip`]: crate::lane_kernels::LaneZip + /// [`map_into`]: IndexedSourceExt::map_into + /// + /// # Panics + /// + /// Panics if `words.len() < self.len().div_ceil(64)`. + #[inline] + fn map_bits_into(self, words: &mut [u64], mut f: F) + where + F: FnMut(Self::Item) -> bool, + { + #[inline(always)] + fn chunk(values: &S, f: &mut F, base: usize, count: usize) -> u64 + where + S: IndexedSource, + F: FnMut(S::Item) -> bool, + { + let mut packed: u64 = 0; + for bit_idx in 0..count { + // SAFETY: caller guarantees base + count <= len. + let val = unsafe { values.get_unchecked(base + bit_idx) }; + packed |= (f(val) as u64) << bit_idx; + } + packed + } + + let values = self; + let len = values.len(); + let num_words = len.div_ceil(64); + assert!( + words.len() >= num_words, + "words slice has {} entries, need at least {num_words}", + words.len(), + ); + + let full = len / 64; + let remainder = len % 64; + + for word_idx in 0..full { + words[word_idx] = chunk(&values, &mut f, word_idx * 64, 64); + } + if remainder != 0 { + words[full] = chunk(&values, &mut f, full * 64, remainder); + } + } + /// Fallible map with **no validity awareness at all** — every `None` returned /// by the closure is treated as a failure, even at null lanes. /// @@ -485,6 +546,43 @@ mod tests { assert!(res.is_ok(), "null lane should bypass the range check"); } + #[test] + fn map_bits_into_packs_full_and_remainder_words() { + let values: Vec = (0..130).collect(); + let mut words = vec![u64::MAX; 3]; + values.as_slice().map_bits_into(&mut words, |v| v % 2 == 0); + + for idx in 0..130 { + let bit = (words[idx / 64] >> (idx % 64)) & 1; + assert_eq!(bit == 1, idx % 2 == 0, "lane {idx}"); + } + // Bits past `len` in the remainder word must be written as zero. + assert_eq!(words[2] >> 2, 0); + } + + #[test] + fn map_bits_into_lane_zip_compare() { + use crate::lane_kernels::LaneZip; + + let lhs: Vec = (0..100).collect(); + let rhs: Vec = (0..100).rev().collect(); + let mut words = vec![0u64; 2]; + LaneZip::new(lhs.as_slice(), rhs.as_slice()).map_bits_into(&mut words, |(a, b)| a >= b); + + for idx in 0..100 { + let bit = (words[idx / 64] >> (idx % 64)) & 1; + assert_eq!(bit == 1, lhs[idx] >= rhs[idx], "lane {idx}"); + } + } + + #[test] + #[should_panic(expected = "words slice has 1 entries")] + fn map_bits_into_words_too_short_panics() { + let values: Vec = (0..65).collect(); + let mut words = vec![0u64; 1]; + values.as_slice().map_bits_into(&mut words, |v| v > 0); + } + #[test] fn try_map_masked_into_overflow_in_remainder() { let mut values: Vec = (0..130).collect(); diff --git a/vortex-cuda/Cargo.toml b/vortex-cuda/Cargo.toml index cc34bee74fb..be6faa77e7b 100644 --- a/vortex-cuda/Cargo.toml +++ b/vortex-cuda/Cargo.toml @@ -26,7 +26,6 @@ unstable_encodings = ["vortex/unstable_encodings"] arc-swap = { workspace = true } arrow-schema = { workspace = true, features = ["ffi"] } async-trait = { workspace = true } -bytes = { workspace = true } cudarc = { workspace = true, features = ["f16"] } futures = { workspace = true, features = ["executor"] } itertools = { workspace = true } @@ -39,6 +38,7 @@ tokio = { workspace = true, features = ["fs"] } tracing = { workspace = true, features = ["std", "attributes"] } vortex = { workspace = true, features = ["zstd"] } vortex-array = { workspace = true, features = ["cudarc"] } +vortex-arrow = { workspace = true } vortex-cub = { path = "cub" } vortex-cuda-macros = { workspace = true } vortex-error = { workspace = true, features = ["object_store"] } diff --git a/vortex-cuda/benches/dynamic_dispatch_cuda.rs b/vortex-cuda/benches/dynamic_dispatch_cuda.rs index 81ee3a45167..e4726d7c46d 100644 --- a/vortex-cuda/benches/dynamic_dispatch_cuda.rs +++ b/vortex-cuda/benches/dynamic_dispatch_cuda.rs @@ -7,6 +7,7 @@ mod bench_config; +use std::f64::consts::PI; use std::marker::PhantomData; use std::mem::size_of; use std::os::raw::c_void; @@ -39,6 +40,7 @@ use vortex::dtype::NativePType; use vortex::dtype::PType; use vortex::encodings::alp::ALP; use vortex::encodings::alp::ALPArrayExt; +use vortex::encodings::alp::ALPArrayOwnedExt; use vortex::encodings::alp::ALPArraySlotsExt; use vortex::encodings::alp::ALPFloat; use vortex::encodings::alp::Exponents; @@ -745,40 +747,50 @@ fn bench_alp_for_bitpacked_f64(c: &mut Criterion) { for (len, len_str) in BENCH_SIZES { group.throughput(Throughput::Bytes((len * size_of::()) as u64)); - // Generate f64 values that ALP-encode without patches. - let floats: Vec = (0..*len) - .map(|i| ::decode_single(10 + (i as i64 % 64), exponents)) - .collect(); - let float_prim = PrimitiveArray::new(Buffer::from(floats), NonNullable); - - // Encode: ALP → FoR → BitPacked - let alp = - alp_encode(float_prim.as_view(), Some(exponents), &mut ctx).vortex_expect("alp_encode"); - assert!(alp.patches().is_none()); - let for_arr = FoRData::encode( - alp.encoded() - .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) - .vortex_expect("bitpack encode"); - - let tree = ALP::new( - FoR::try_new(bp.into_array(), for_arr.reference_scalar().clone()) - .vortex_expect("for_new") - .into_array(), - exponents, - None, - ); - let array = tree.into_array(); + for (patch_interval, benchmark_name) in [ + (None, "cuda/alp_for_bp_6bw_f64/dispatch_f64"), + ( + Some(100), + "cuda/alp_for_bp_6bw_f64_1pct_patches/dispatch_f64", + ), + ] { + let floats: Vec = (0..*len) + .map(|i| { + if patch_interval.is_some_and(|interval| i % interval == 0) { + PI + } else { + ::decode_single(10 + (i as i64 % 64), exponents) + } + }) + .collect(); + let float_prim = PrimitiveArray::new(Buffer::from(floats), NonNullable); + + // Encode: ALP → FoR → BitPacked, preserving ALP's exception patches. + let alp = alp_encode(float_prim.as_view(), Some(exponents), &mut ctx) + .vortex_expect("alp_encode"); + assert_eq!(alp.patches().is_some(), patch_interval.is_some()); + let (encoded, alp_exponents, patches) = alp.into_parts(); + let for_arr = FoRData::encode( + encoded + .execute::(&mut ctx) + .vortex_expect("to primitive"), + &mut ctx, + ) + .vortex_expect("for encode"); + let bp = BitPackedData::encode(for_arr.encoded(), bit_width, &mut ctx) + .vortex_expect("bitpack encode"); + assert!(bp.patches().is_none(), "expected only ALP patches"); + + let tree = ALP::new( + FoR::try_new(bp.into_array(), for_arr.reference_scalar().clone()) + .vortex_expect("for_new") + .into_array(), + alp_exponents, + patches, + ); + let array = tree.into_array(); - group.bench_with_input( - BenchmarkId::new("cuda/alp_for_bp_6bw_f64/dispatch_f64", len_str), - len, - |b, &n| { + group.bench_with_input(BenchmarkId::new(benchmark_name, len_str), len, |b, &n| { let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx"); @@ -791,8 +803,8 @@ fn bench_alp_for_bitpacked_f64(c: &mut Criterion) { } total_time }); - }, - ); + }); + } } group.finish(); diff --git a/vortex-cuda/benches/fsst_cuda.rs b/vortex-cuda/benches/fsst_cuda.rs index 21db8744919..49930c9a30a 100644 --- a/vortex-cuda/benches/fsst_cuda.rs +++ b/vortex-cuda/benches/fsst_cuda.rs @@ -4,7 +4,6 @@ //! CUDA benchmarks for FSST decompression. #![expect(clippy::unwrap_used)] -#![expect(clippy::cast_possible_truncation)] #[allow(dead_code)] mod bench_config; @@ -18,12 +17,20 @@ use criterion::BenchmarkId; use criterion::Criterion; use criterion::Throughput; use futures::executor::block_on; +use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::match_each_integer_ptype; +use vortex::dtype::DType; +use vortex::dtype::Nullability; +use vortex::encodings::fsst::FSST; use vortex::encodings::fsst::FSSTArrayExt; use vortex::error::VortexExpect; +use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; +use vortex_cuda::VarBinExportLayout; +use vortex_cuda::arrow::DeviceArrayExt; +use vortex_cuda::arrow::release_device_array; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda_macros::cuda_available; use vortex_cuda_macros::cuda_not_available; @@ -36,30 +43,63 @@ use crate::timed_launch_strategy::TimedLaunchStrategy; // other kernels benchmark. const BENCH_SIZES: &[(usize, &str)] = &[(10_000_000, "10M")]; +fn session_with_varbin_layout(layout: VarBinExportLayout) -> vortex::session::VortexSession { + vortex::array::array_session().with_some( + CudaSession::try_default() + .vortex_expect("failed to create CUDA session") + .with_varbin_export_layout(layout), + ) +} + +struct FSSTBenchFixture { + utf8: ArrayRef, + binary: ArrayRef, + uncompressed_size: u64, +} + +fn make_fixture(n: usize) -> FSSTBenchFixture { + 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()); + + let lens = fsst + .uncompressed_lengths() + .clone() + .execute::(setup_ctx.execution_ctx()) + .vortex_expect("canonicalize uncompressed_lengths"); + #[allow(clippy::unnecessary_cast)] + let uncompressed_size = match_each_integer_ptype!(lens.ptype(), |P| { + lens.as_slice::

().iter().map(|x| *x as u64).sum() + }); + + let binary = FSST::try_new( + DType::Binary(Nullability::NonNullable), + fsst.symbols().clone(), + fsst.symbol_lengths().clone(), + fsst.codes(), + fsst.uncompressed_lengths().clone(), + setup_ctx.execution_ctx(), + ) + .vortex_expect("rebuild FSST fixture with Binary dtype") + .into_array(); + + FSSTBenchFixture { + utf8: fsst.into_array(), + binary, + uncompressed_size, + } +} + 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(&vortex_cuda::cuda_session()) - .vortex_expect("failed to create execution context"); - let fsst = make_fsst_clickbench_urls(n, setup_ctx.execution_ctx()); - - let lens = fsst - .uncompressed_lengths() - .clone() - .execute::(setup_ctx.execution_ctx()) - .vortex_expect("canonicalize uncompressed_lengths"); - let total_size: usize = match_each_integer_ptype!(lens.ptype(), |P| { - lens.as_slice::

().iter().map(|x| *x as usize).sum() - }); - let uncompressed_size = total_size as u64; - - let fsst_array = fsst.into_array(); - - group.throughput(Throughput::Bytes(uncompressed_size)); + let fixture = make_fixture(n); + + group.throughput(Throughput::Bytes(fixture.uncompressed_size)); group.bench_with_input( - BenchmarkId::new("cuda/fsst/decompress", len_str), - &fsst_array, + BenchmarkId::new("cuda/fsst/decompress_to_varbinview", len_str), + &fixture.utf8, |b, fsst_array| { b.iter_custom(|iters| { let timed = TimedLaunchStrategy::default(); @@ -68,6 +108,7 @@ fn benchmark_fsst_cuda_decompress(c: &mut Criterion) { 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 { @@ -77,6 +118,51 @@ fn benchmark_fsst_cuda_decompress(c: &mut Criterion) { }); }, ); + + for (name, array, layout) in [ + ( + "cuda/fsst/decompress_to_varbin", + &fixture.utf8, + VarBinExportLayout::VarBin, + ), + ( + "cuda/fsst/export_binary", + &fixture.binary, + VarBinExportLayout::VarBin, + ), + ( + "cuda/fsst/export_utf8_view", + &fixture.utf8, + VarBinExportLayout::VarBinView, + ), + ( + "cuda/fsst/export_binary_view", + &fixture.binary, + VarBinExportLayout::VarBinView, + ), + ] { + group.bench_with_input(BenchmarkId::new(name, len_str), array, |b, array| { + b.iter_custom(|iters| { + let timed = TimedLaunchStrategy::default(); + let timer = timed.timer(); + + let session = session_with_varbin_layout(layout); + let mut cuda_ctx = CudaSession::create_execution_ctx(&session) + .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) + .with_launch_strategy(Arc::new(timed)); + + for _ in 0..iters { + let mut exported = + block_on((*array).clone().export_device_array(&mut cuda_ctx)) + .vortex_expect("export FSST device array"); + release_device_array(&mut exported); + } + + Duration::from_nanos(timer.load(Ordering::Relaxed)) + }); + }); + } } group.finish(); diff --git a/vortex-cuda/build.rs b/vortex-cuda/build.rs index e201345cd73..0621a2884b3 100644 --- a/vortex-cuda/build.rs +++ b/vortex-cuda/build.rs @@ -35,6 +35,13 @@ fn main() { println!("cargo:rerun-if-env-changed=PROFILE"); + // nvcc availability depends on PATH (e.g. entering/leaving a nix shell that + // provides cuda_nvcc). Without this, cargo caches a stale build.rs result + // from an environment without nvcc, and switching to one with nvcc does not + // trigger PTX recompilation — producing a binary with an empty embedded_ptx + // table that silently falls back to CPU at runtime. + println!("cargo:rerun-if-env-changed=PATH"); + // Regenerate bit_unpack kernels only when the generator changes println!( "cargo:rerun-if-changed={}", diff --git a/vortex-cuda/cub/Cargo.toml b/vortex-cuda/cub/Cargo.toml index ed166318569..c37fd18ec83 100644 --- a/vortex-cuda/cub/Cargo.toml +++ b/vortex-cuda/cub/Cargo.toml @@ -17,6 +17,9 @@ rust-version = { workspace = true } version = { workspace = true } publish = false +[package.metadata.cargo-shear] +ignored = ["libloading"] + [lints] workspace = true diff --git a/vortex-cuda/ffi/Cargo.toml b/vortex-cuda/ffi/Cargo.toml index a8ca51ad954..c1e1efdb23c 100644 --- a/vortex-cuda/ffi/Cargo.toml +++ b/vortex-cuda/ffi/Cargo.toml @@ -22,7 +22,6 @@ vortex-cuda = { path = ".." } vortex-ffi = { path = "../../vortex-ffi" } [dev-dependencies] -vortex-array = { workspace = true, features = ["_test-harness"] } vortex-cuda-macros = { workspace = true } [lib] diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index 3c030095374..1b07e12f0df 100644 --- a/vortex-cuda/ffi/README.md +++ b/vortex-cuda/ffi/README.md @@ -16,3 +16,10 @@ Use this crate as the CUDA-enabled FFI artifact. Include both headers: and link the CUDA FFI library (`vortex_cuda_ffi`). Do not pass Vortex handles between independently linked Rust FFI libraries. Use `vx_cuda_session_new` to initialize CUDA once and reuse it across exports. + +Use `vx_cuda_array_sink_open_file` to open a standard Vortex file sink configured to produce CUDA-readable files. +Push host-resident arrays and close the sink using the standard `vx_array_sink_*` APIs. + +Use `vx_cuda_scan_path_arrow_device_stream` to read such a local file through pinned host buffers +and receive an Arrow C Device stream. Reuse the same CUDA session across scans so the pinned buffer +pool and CUDA state are reused as well. diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 0b06a660afa..971ad9872f4 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -63,6 +63,33 @@ struct ArrowDeviceArrayStream { */ vx_session *vx_cuda_session_new(vx_error **error_out); +/** + * Open a Vortex file sink configured to produce CUDA-readable files. + * + * Push host-resident arrays and close or abort the returned sink with the standard + * `vx_array_sink_*` functions. This API configures the on-disk encodings and layout; it does not + * move arrays to the GPU during the write. + */ +vx_array_sink *vx_cuda_array_sink_open_file(const vx_session *session, + vx_view path, + const vx_dtype *dtype, + vx_error **error_out); + +/** + * Scan a local CUDA-compatible Vortex file as an Arrow C Device stream. + * + * Files written by `vx_cuda_array_sink_open_file` are compatible with this path. Reusing the same + * CUDA session across calls also reuses the pinned host buffers used to stage file reads. + * + * On success returns 0 and writes an owned `ArrowDeviceArrayStream` to `out_stream`. The caller + * must release the stream and each produced `ArrowDeviceArray` through their embedded Arrow + * release callbacks. On error returns 1 and writes a `vx_error` to `error_out` when non-NULL. + */ +int vx_cuda_scan_path_arrow_device_stream(const vx_session *session, + vx_view path, + struct ArrowDeviceArrayStream *out_stream, + vx_error **error_out); + /** * Export a borrowed Vortex array for cuDF's Arrow Device import path. * diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 611f644e535..12bc584dad6 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -10,27 +10,42 @@ use std::os::raw::c_int; use std::ptr; +use std::sync::Arc; use arrow_schema::ffi::FFI_ArrowSchema; +use vortex::array::stream::ArrayStreamExt; +use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::error::VortexResult; use vortex::error::vortex_ensure; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; +use vortex::io::VortexReadAt; +use vortex::io::runtime::BlockingRuntime; +use vortex::io::session::RuntimeSessionExt; use vortex::session::SessionExt; use vortex::session::VortexSession; use vortex_cuda::CudaSession; +use vortex_cuda::PooledFileReadAt; use vortex_cuda::arrow::ArrowDeviceArray; use vortex_cuda::arrow::ArrowDeviceArrayStream; use vortex_cuda::arrow::DeviceArrayExt; use vortex_cuda::arrow::DeviceArrayStreamExt; +use vortex_cuda::layout::CudaFlatLayoutStrategy; +use vortex_cuda::layout::register_cuda_layout; 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_array_sink; +use vortex_ffi::vx_array_sink_open_file_with_strategy; +use vortex_ffi::vx_dtype; 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; +use vortex_ffi::vx_view; const VX_CUDA_OK: c_int = 0; const VX_CUDA_ERR: c_int = 1; @@ -41,6 +56,7 @@ const VX_CUDA_ERR: c_int = 1; /// returns a new session cloned from `session` with a default [`CudaSession`] attached. fn session_with_cuda(session: &VortexSession) -> VortexResult { session.get::(); + register_cuda_layout(session); Ok(session.clone()) } @@ -59,11 +75,93 @@ pub unsafe extern "C-unwind" fn vx_cuda_session_new( try_or(error_out, ptr::null_mut(), || { let cuda_session = CudaSession::try_default()?; Ok(vx_session_new_with(|session| { - session.with_some(cuda_session) + let session = session.with_some(cuda_session); + register_cuda_layout(&session); + session })) }) } +/// Open a Vortex file sink configured to produce CUDA-readable files. +/// +/// Push host-resident arrays and close or abort the returned sink with the standard +/// `vx_array_sink_*` functions. This function configures the on-disk encodings and layout; it does +/// not move arrays to the GPU during the write. +/// +/// # Safety +/// +/// `session`, `path`, and `dtype` must satisfy the same requirements as +/// `vx_array_sink_open_file`. If `error_out` is non-null, it must be valid for writing one error +/// pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file( + session: *const vx_session, + path: vx_view, + dtype: *const vx_dtype, + error_out: *mut *mut vx_error, +) -> *mut vx_array_sink { + try_or(error_out, ptr::null_mut(), || { + session_with_cuda(unsafe { vx_session_ref(session) }?)?; + let strategy = WriteStrategyBuilder::default() + .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().only_cuda_compatible()) + .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())) + .build(); + unsafe { vx_array_sink_open_file_with_strategy(session, path, dtype, strategy) } + }) +} + +/// Scan a local Vortex file through pinned host buffers and export an Arrow C Device stream. +/// +/// The file must use encodings and layouts supported by the CUDA execution path, such as files +/// written by [`vx_cuda_array_sink_open_file`]. Pinned staging buffers are reused across scans made +/// with the same CUDA session. +/// +/// On success returns `0` and writes an owned [`ArrowDeviceArrayStream`] to `out_stream`. The +/// caller must release the stream and each array produced by it through their embedded Arrow +/// release callbacks. +/// +/// 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`. `path` must be valid for the +/// duration of this call and contain UTF-8. `out_stream` must be a valid writable pointer. If +/// `error_out` is non-null, it must be valid for writing one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( + session: *const vx_session, + path: vx_view, + out_stream: *mut ArrowDeviceArrayStream, + error_out: *mut *mut vx_error, +) -> c_int { + try_or(error_out, VX_CUDA_ERR, || { + vortex_ensure!(!out_stream.is_null(), "null ArrowDeviceArrayStream output"); + + let path = unsafe { path.as_str() }?.to_owned(); + let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?; + let cuda_session = session.get::(); + let stream = cuda_session.stream()?; + let pool = Arc::clone(cuda_session.pinned_buffer_pool()); + drop(cuda_session); + + let reader: Arc = Arc::new(PooledFileReadAt::open( + path, + session.handle(), + pool, + stream, + )?); + let array_stream = ffi_runtime().block_on(async { + let file = session.open_options().open(reader).await?; + Ok::<_, vortex::error::VortexError>(file.scan()?.into_array_stream()?.boxed()) + })?; + let device_stream = array_stream.export_device_array_stream(&session, ffi_runtime())?; + + unsafe { ptr::write(out_stream, device_stream) }; + Ok(VX_CUDA_OK) + }) +} + /// Export a borrowed Vortex array for cuDF's Arrow Device import path. /// /// On success returns `0` and writes independently releasable `out_schema` and `out_array`; the diff --git a/vortex-cuda/gpu-scan-cli/src/main.rs b/vortex-cuda/gpu-scan-cli/src/main.rs index 9078e30c691..00f6cd1ac1c 100644 --- a/vortex-cuda/gpu-scan-cli/src/main.rs +++ b/vortex-cuda/gpu-scan-cli/src/main.rs @@ -29,13 +29,12 @@ use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; use vortex::io::session::RuntimeSessionExt; +use vortex::session::SessionExt; use vortex::session::VortexSession; use vortex_cuda::CudaSession; -use vortex_cuda::PinnedByteBufferPool; use vortex_cuda::PooledByteBufferReadAt; use vortex_cuda::PooledFileReadAt; use vortex_cuda::TracingLaunchStrategy; -use vortex_cuda::VortexCudaStreamPool; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda::layout::CudaFlatLayoutStrategy; use vortex_cuda::layout::register_cuda_layout; @@ -155,11 +154,10 @@ async fn cmd_scan(path: PathBuf, gpu_file: bool, json_output: bool) -> VortexRes let mut cuda_ctx = CudaSession::create_execution_ctx(&session)? .with_launch_strategy(Arc::new(TracingLaunchStrategy)); - let pool = Arc::new(PinnedByteBufferPool::new(Arc::clone( - cuda_ctx.stream().context(), - ))); - let cuda_stream = - VortexCudaStreamPool::new(Arc::clone(cuda_ctx.stream().context()), 1).stream()?; + let cuda_session = session.get::(); + let pool = Arc::clone(cuda_session.pinned_buffer_pool()); + let cuda_stream = cuda_session.stream()?; + drop(cuda_session); let handle = session.handle(); let gpu_file_handle = if gpu_file { diff --git a/vortex-cuda/kernels/src/dict.cu b/vortex-cuda/kernels/src/dict.cu index 3422980c65e..a456b476821 100644 --- a/vortex-cuda/kernels/src/dict.cu +++ b/vortex-cuda/kernels/src/dict.cu @@ -19,12 +19,51 @@ __device__ void dict_kernel(const IndexT *const __restrict codes, (block_start + elements_per_block < codes_len) ? (block_start + elements_per_block) : codes_len; for (uint64_t idx = block_start + threadIdx.x; idx < block_end; idx += blockDim.x) { - IndexT code = codes[idx]; + const IndexT code = codes[idx]; output[idx] = values[code]; } } -// Macro to generate dict kernels for all value/index type combinations +// Nullable dictionary values require gathering validity through the row codes: +// +// row_validity[i] = values_validity[codes[i]] +// +// Vortex represents this gather lazily as `Dict`; null codes are carried by the resulting +// boolean array's own validity. Consequently, this kernel may materialize validity for a +// dictionary whose logical values are strings or another non-boolean type. +// +// Vortex booleans are bit-packed, so fixed-width `output[i] = values[codes[i]]` semantics do not +// apply. Assigning each complete output byte to one thread avoids races between individual bit +// writes. +template +__device__ void dict_bool_kernel(const IndexT *const __restrict codes, + uint64_t codes_len, + const uint8_t *const __restrict values, + uint64_t values_bit_offset, + uint8_t *const __restrict output) { + const uint64_t output_len = (codes_len + 7) / 8; + const uint32_t elements_per_block = blockDim.x * ELEMENTS_PER_THREAD; + const uint64_t block_start = static_cast(blockIdx.x) * elements_per_block; + const uint64_t block_end = + (block_start + elements_per_block < output_len) ? (block_start + elements_per_block) : output_len; + + for (uint64_t output_idx = block_start + threadIdx.x; output_idx < block_end; output_idx += blockDim.x) { + const uint64_t row_start = output_idx * 8; + uint8_t packed = 0; +#pragma unroll + for (uint32_t bit = 0; bit < 8; ++bit) { + const uint64_t row = row_start + bit; + if (row < codes_len) { + const uint64_t value_idx = values_bit_offset + static_cast(codes[row]); + const uint8_t value = (values[value_idx / 8] >> (value_idx % 8)) & 1; + packed |= static_cast(value << bit); + } + } + output[output_idx] = packed; + } +} + +// Macro to generate dict kernels for all fixed-width value/index type combinations. #define GENERATE_DICT_KERNEL(value_suffix, ValueType, index_suffix, IndexType) \ extern "C" __global__ void dict_##value_suffix##_##index_suffix( \ const IndexType *const __restrict codes, \ @@ -41,5 +80,21 @@ __device__ void dict_kernel(const IndexT *const __restrict codes, GENERATE_DICT_KERNEL(value_suffix, ValueType, u32, uint32_t) \ GENERATE_DICT_KERNEL(value_suffix, ValueType, u64, uint64_t) -// Generate for all native ptypes & decimal values +#define GENERATE_DICT_BOOL_KERNEL(index_suffix, IndexType) \ + extern "C" __global__ void dict_bool_##index_suffix(const IndexType *const __restrict codes, \ + uint64_t codes_len, \ + const uint8_t *const __restrict values, \ + uint64_t values_bit_offset, \ + uint8_t *const __restrict output) { \ + dict_bool_kernel(codes, codes_len, values, values_bit_offset, output); \ + } + +// Generate fixed-width kernels for all native ptypes and decimal values. FOR_EACH_NUMERIC(GENERATE_DICT_FOR_ALL_INDICES) + +// Boolean values use a different physical layout and launch unit, but still dispatch over the +// same unsigned dictionary-code types. +GENERATE_DICT_BOOL_KERNEL(u8, uint8_t) +GENERATE_DICT_BOOL_KERNEL(u16, uint16_t) +GENERATE_DICT_BOOL_KERNEL(u32, uint32_t) +GENERATE_DICT_BOOL_KERNEL(u64, uint64_t) diff --git a/vortex-cuda/kernels/src/dynamic_dispatch.cu b/vortex-cuda/kernels/src/dynamic_dispatch.cu index fa19b47fe1a..3a7f28e1576 100644 --- a/vortex-cuda/kernels/src/dynamic_dispatch.cu +++ b/vortex-cuda/kernels/src/dynamic_dispatch.cu @@ -117,12 +117,8 @@ __shared__ uint64_t runend_cursors[BLOCK_SIZE]; // ═══════════════════════════════════════════════════════════════════════════ /// Apply one scalar operation to N values in registers. -/// -/// `abs_pos` is the absolute output position of the first value to process. -/// It is used by scalar operations that apply patches, e.g. ALP. template -__device__ inline void -scalar_op(T *values, const struct ScalarOp &op, char *__restrict smem, uint64_t abs_pos = 0) { +__device__ inline void scalar_op(T *values, const struct ScalarOp &op, char *__restrict smem) { switch (op.op_code) { case ScalarOp::FOR: { const T ref = static_cast(op.params.frame_of_ref.reference); @@ -165,30 +161,9 @@ scalar_op(T *values, const struct ScalarOp &op, char *__restrict smem, uint64_t values[i] = static_cast(__double_as_longlong(r)); } } - // Apply ALP patches: override positions whose float value couldn't - // be reconstructed through the ALP encode/decode cycle. - // Per-value cursor — with a slice offset, a tile's N values can - // straddle two FL chunks, so each value needs its own lookup. - if (op.params.alp.patches_ptr != 0) { - const auto &patches = *reinterpret_cast(op.params.alp.patches_ptr); - const uint32_t chunk_start = patches.offset / FL_CHUNK; -#pragma unroll - for (uint32_t i = 0; i < N; ++i) { - uint64_t my_pos = (N > 1) ? abs_pos + i * blockDim.x + threadIdx.x : abs_pos; - uint64_t orig = my_pos + patches.offset; - uint32_t chunk = static_cast(orig / FL_CHUNK) - chunk_start; - uint32_t within = static_cast(orig % FL_CHUNK); - PatchesCursor cursor(patches, chunk, 0, 1); - auto patch = cursor.next(); - while (patch.index != FL_CHUNK) { - if (patch.index == within) { - values[i] = patch.value; - break; - } - patch = cursor.next(); - } - } - } + // ALP patches are scattered cooperatively after the stage has decoded. + // Looking up patches here would restart a chunk cursor for every value, + // turning sparse exception handling into O(values * patches). break; } case ScalarOp::DICT: { @@ -225,6 +200,58 @@ scatter_patches_chunk(const GPUPatches &patches, T *__restrict out, uint32_t chu } } +/// Scatter patches overlapping a logical output range. +/// +/// `logical_start` and `range_len` are relative to the sliced array represented by +/// `patches`. Patch indices are stored in the coordinate space of the original array, +/// so add `patches.offset` while locating chunks and subtract it when addressing `out`. +/// Every thread cooperates on each overlapping FastLanes chunk. +template +__device__ inline void scatter_patches_range(const GPUPatches &patches, + T *__restrict out, + uint64_t logical_start, + uint32_t range_len) { + if (range_len == 0) { + return; + } + + const uint64_t original_start = logical_start + patches.offset; + const uint64_t original_end = original_start + range_len; + const uint32_t patch_chunk_start = patches.offset / FL_CHUNK; + const uint32_t first_chunk = static_cast(original_start / FL_CHUNK); + const uint32_t last_chunk = static_cast((original_end - 1) / FL_CHUNK); + + for (uint32_t original_chunk = first_chunk; original_chunk <= last_chunk; ++original_chunk) { + const uint32_t chunk = original_chunk - patch_chunk_start; + PatchesCursor cursor(patches, chunk, threadIdx.x, blockDim.x); + auto patch = cursor.next(); + while (patch.index != FL_CHUNK) { + const uint64_t original_pos = static_cast(original_chunk) * FL_CHUNK + patch.index; + if (original_pos >= original_start && original_pos < original_end) { + out[original_pos - patches.offset] = patch.value; + } + patch = cursor.next(); + } + } +} + +/// Apply patch payloads attached to scalar operations after their stage has decoded. +template +__device__ inline void +scatter_scalar_patches(const Stage &stage, T *__restrict out, uint64_t logical_start, uint32_t range_len) { + for (uint8_t op_idx = 0; op_idx < stage.num_scalar_ops; ++op_idx) { + const auto &op = stage.scalar_ops[op_idx]; + if (op.op_code == ScalarOp::ALP && op.params.alp.patches_ptr != 0) { + const auto &patches = *reinterpret_cast(op.params.alp.patches_ptr); + // All ordinary writes must complete before exception values overwrite them. + __syncthreads(); + scatter_patches_range(patches, out, logical_start, range_len); + // A later stage may consume this patched shared-memory output. + __syncthreads(); + } + } +} + // ═══════════════════════════════════════════════════════════════════════════ // Source ops // ═══════════════════════════════════════════════════════════════════════════ @@ -491,7 +518,7 @@ __device__ void execute_output_stage(T *__restrict output, smem); for (uint8_t op = 0; op < stage.num_scalar_ops; ++op) { - scalar_op(values, stage.scalar_ops[op], smem, tile_start); + scalar_op(values, stage.scalar_ops[op], smem); } #pragma unroll @@ -513,7 +540,7 @@ __device__ void execute_output_stage(T *__restrict output, source_op(&val, src, raw_input, ptype, smem_src, i, gpos, smem); for (uint8_t op = 0; op < stage.num_scalar_ops; ++op) { - scalar_op(&val, stage.scalar_ops[op], smem, gpos); + scalar_op(&val, stage.scalar_ops[op], smem); } __stcs(&output[gpos], val); } @@ -525,6 +552,8 @@ __device__ void execute_output_stage(T *__restrict output, } elem_idx += chunk_len; } + + scatter_scalar_patches(stage, output, block_start, block_len); } // ═══════════════════════════════════════════════════════════════════════════ @@ -559,7 +588,7 @@ __device__ void execute_input_stage(const Stage &stage, char *__restrict smem) { for (uint32_t elem_idx = threadIdx.x; elem_idx < stage.len; elem_idx += blockDim.x) { T val = smem_out[elem_idx]; for (uint8_t op = 0; op < stage.num_scalar_ops; ++op) { - scalar_op(&val, stage.scalar_ops[op], smem, elem_idx); + scalar_op(&val, stage.scalar_ops[op], smem); } smem_out[elem_idx] = val; } @@ -583,7 +612,7 @@ __device__ void execute_input_stage(const Stage &stage, char *__restrict smem) { T val; source_op(&val, src, raw_input, stage.source_ptype, nullptr, 0, elem_idx, smem); for (uint8_t op = 0; op < stage.num_scalar_ops; ++op) { - scalar_op(&val, stage.scalar_ops[op], smem, elem_idx); + scalar_op(&val, stage.scalar_ops[op], smem); } smem_out[elem_idx] = val; } @@ -591,6 +620,8 @@ __device__ void execute_input_stage(const Stage &stage, char *__restrict smem) { // stages to read. __syncthreads(); } + + scatter_scalar_patches(stage, smem_out, 0, stage.len); } // ═══════════════════════════════════════════════════════════════════════════ diff --git a/vortex-cuda/kernels/src/fsst.cu b/vortex-cuda/kernels/src/fsst.cu index 938f262c3da..22e8790a479 100644 --- a/vortex-cuda/kernels/src/fsst.cu +++ b/vortex-cuda/kernels/src/fsst.cu @@ -53,7 +53,9 @@ // the next add (≤ 8 bytes) within the 24-byte capacity. // // `codes_offsets` is templated over the four unsigned integer widths -// (u8/u16/u32/u64). `output_offsets` is uint64_t. +// (u8/u16/u32/u64). `output_offsets` is uint64_t for the view kernels +// (`fsst_*`, which also take an optional views pointer) and int32_t Arrow +// varbin offsets for the bytes-only varbin kernels (`fsst_varbin_*`). // 24-byte scratch buffer split across three u64 lanes. `cursor` is the // number of bytes currently buffered and the next-push offset. @@ -123,13 +125,18 @@ struct Scratch { } }; -template +// Arrow/Vortex variable-length view records are 16 bytes. Values up to 12 +// bytes are stored inline after the u32 length. Longer values store their +// first four bytes, backing-buffer index, and byte offset. +constexpr uint32_t MAX_INLINED_SIZE = 12; + +template struct FSSTArgs { // Compressed FSST code stream, contiguous across all strings. String // `sid`'s codes live in `[codes_offsets[sid], codes_offsets[sid + 1])`. const uint8_t *__restrict codes_bytes; // Per-string offsets into `codes_bytes`, length `num_strings + 1`. - const OffT *__restrict codes_offsets; + const CodeOffsetT *__restrict codes_offsets; // FSST symbol table. const uint64_t *__restrict symbols; // Length in bytes (1..=8) of each entry in `symbols`. The remaining bits @@ -138,19 +145,57 @@ struct FSSTArgs { // Buffer to write decoded data into. uint8_t *__restrict output_bytes; // Per-string offsets into `output_bytes`, length `num_strings + 1`. - const uint64_t *__restrict output_offsets; + const OutputOffsetT *__restrict output_offsets; // Validity of each string. const uint8_t *__restrict validity_bits; + // Bit offset of string zero within the first validity byte (0..7). + uint64_t validity_bit_offset; + // Optional output views, one 16-byte uint4 per string. A null pointer + // requests bytes-only decoding for heaps that need the host rollover path. + uint4 *__restrict output_views; }; -template -__device__ inline void fsst_decode_string(const FSSTArgs &args, uint64_t sid) { - if (((args.validity_bits[sid >> 3] >> (sid & 7u)) & 1u) == 0u) { +// Build one BinaryView from the bytes this thread just decoded. The Rust +// caller only provides output_views when every offset fits in the view's u32 +// fields and the decoded heap is exposed as backing buffer zero. +template +__device__ inline void fsst_write_view(const FSSTArgs &args, uint64_t sid) { + if (args.output_views == nullptr) { return; } - OffT in_pos = args.codes_offsets[sid]; - const OffT in_end = args.codes_offsets[sid + 1]; + const uint64_t start = args.output_offsets[sid]; + const uint32_t len = (uint32_t)(args.output_offsets[sid + 1] - start); + if (len <= MAX_INLINED_SIZE) { + uint32_t words[3] = {0, 0, 0}; +#pragma unroll + for (uint32_t i = 0; i < MAX_INLINED_SIZE; i++) { + if (i < len) { + words[i >> 2] |= (uint32_t)args.output_bytes[start + i] << (8u * (i & 3u)); + } + } + args.output_views[sid] = make_uint4(len, words[0], words[1], words[2]); + return; + } + + const uint32_t prefix = + (uint32_t)args.output_bytes[start] | ((uint32_t)args.output_bytes[start + 1] << 8u) | + ((uint32_t)args.output_bytes[start + 2] << 16u) | ((uint32_t)args.output_bytes[start + 3] << 24u); + args.output_views[sid] = make_uint4(len, prefix, 0, (uint32_t)start); +} + +template +__device__ inline void fsst_decode_string(const FSSTArgs &args, uint64_t sid) { + const uint64_t validity_index = sid + args.validity_bit_offset; + if (((args.validity_bits[validity_index >> 3] >> (validity_index & 7u)) & 1u) == 0u) { + if (args.output_views != nullptr) { + args.output_views[sid] = make_uint4(0, 0, 0, 0); + } + return; + } + + CodeOffsetT in_pos = args.codes_offsets[sid]; + const CodeOffsetT in_end = args.codes_offsets[sid + 1]; uint64_t out_pos = args.output_offsets[sid]; const uint64_t out_end = args.output_offsets[sid + 1]; @@ -181,25 +226,38 @@ __device__ inline void fsst_decode_string(const FSSTArgs &args, uint64_t s sym &= mask; scratch.push(sym, len); - in_pos += (OffT)consumed; + in_pos += (CodeOffsetT)consumed; } // Epilogue: drain everything that's left. while (scratch.cursor > 0) { scratch.drain(args.output_bytes, out_pos, out_end); } + + fsst_write_view(args, sid); } -#define GENERATE_FSST_KERNEL(suffix, OffT) \ +#define FSST_GRID_STRIDE_LOOP(CodeOffsetT, OutputOffsetT, args) \ + const uint64_t elements_per_block = (uint64_t)blockDim.x * ELEMENTS_PER_THREAD; \ + const uint64_t block_start = (uint64_t)blockIdx.x * elements_per_block; \ + const uint64_t block_end = \ + (block_start + elements_per_block < num_strings) ? (block_start + elements_per_block) : num_strings; \ + for (uint64_t sid = block_start + threadIdx.x; sid < block_end; sid += blockDim.x) { \ + fsst_decode_string(args, sid); \ + } + +#define GENERATE_FSST_VIEW_KERNEL(suffix, CodeOffsetT) \ extern "C" __global__ void fsst_##suffix(const uint8_t *__restrict codes_bytes, \ - const OffT *__restrict codes_offsets, \ + const CodeOffsetT *__restrict codes_offsets, \ const uint64_t *__restrict symbols, \ const uint8_t *__restrict symbol_lengths, \ const uint64_t *__restrict output_offsets, \ const uint8_t *__restrict validity_bits, \ + uint64_t validity_bit_offset, \ uint8_t *__restrict output_bytes, \ + uint4 *__restrict output_views, \ uint64_t num_strings) { \ - const FSSTArgs args = { \ + const FSSTArgs args = { \ codes_bytes, \ codes_offsets, \ symbols, \ @@ -207,20 +265,42 @@ __device__ inline void fsst_decode_string(const FSSTArgs &args, uint64_t s output_bytes, \ output_offsets, \ validity_bits, \ + validity_bit_offset, \ + output_views, \ }; \ - \ - const uint64_t elements_per_block = (uint64_t)blockDim.x * ELEMENTS_PER_THREAD; \ - const uint64_t block_start = (uint64_t)blockIdx.x * elements_per_block; \ - const uint64_t block_end = (block_start + elements_per_block < num_strings) \ - ? (block_start + elements_per_block) \ - : num_strings; \ - \ - for (uint64_t sid = block_start + threadIdx.x; sid < block_end; sid += blockDim.x) { \ - fsst_decode_string(args, sid); \ - } \ + FSST_GRID_STRIDE_LOOP(CodeOffsetT, uint64_t, args) \ } -GENERATE_FSST_KERNEL(u8, uint8_t) -GENERATE_FSST_KERNEL(u16, uint16_t) -GENERATE_FSST_KERNEL(u32, uint32_t) -GENERATE_FSST_KERNEL(u64, uint64_t) +#define GENERATE_FSST_VARBIN_KERNEL(suffix, CodeOffsetT) \ + extern "C" __global__ void fsst_varbin_##suffix(const uint8_t *__restrict codes_bytes, \ + const CodeOffsetT *__restrict codes_offsets, \ + const uint64_t *__restrict symbols, \ + const uint8_t *__restrict symbol_lengths, \ + const int32_t *__restrict output_offsets, \ + const uint8_t *__restrict validity_bits, \ + uint64_t validity_bit_offset, \ + uint8_t *__restrict output_bytes, \ + uint64_t num_strings) { \ + const FSSTArgs args = { \ + codes_bytes, \ + codes_offsets, \ + symbols, \ + symbol_lengths, \ + output_bytes, \ + output_offsets, \ + validity_bits, \ + validity_bit_offset, \ + nullptr, \ + }; \ + FSST_GRID_STRIDE_LOOP(CodeOffsetT, int32_t, args) \ + } + +GENERATE_FSST_VIEW_KERNEL(u8, uint8_t) +GENERATE_FSST_VIEW_KERNEL(u16, uint16_t) +GENERATE_FSST_VIEW_KERNEL(u32, uint32_t) +GENERATE_FSST_VIEW_KERNEL(u64, uint64_t) + +GENERATE_FSST_VARBIN_KERNEL(u8, uint8_t) +GENERATE_FSST_VARBIN_KERNEL(u16, uint16_t) +GENERATE_FSST_VARBIN_KERNEL(u32, uint32_t) +GENERATE_FSST_VARBIN_KERNEL(u64, uint64_t) diff --git a/vortex-cuda/nvcomp/Cargo.toml b/vortex-cuda/nvcomp/Cargo.toml index 5fdd88d4369..4caeef8c7a2 100644 --- a/vortex-cuda/nvcomp/Cargo.toml +++ b/vortex-cuda/nvcomp/Cargo.toml @@ -14,6 +14,9 @@ rust-version = { workspace = true } version = { workspace = true } publish = false +[package.metadata.cargo-shear] +ignored = ["libloading"] + [lints] workspace = true diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index 109e343e424..66c50a0413c 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -54,6 +54,8 @@ use vortex::dtype::NativeDecimalType; use vortex::dtype::Nullability; use vortex::dtype::PType; use vortex::dtype::i256; +use vortex::encodings::fsst::FSST; +use vortex::encodings::fsst::FSSTArray; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_ensure; @@ -63,6 +65,7 @@ use vortex::extension::datetime::AnyTemporal; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; use crate::CudaExecutionCtx; +use crate::VarBinExportLayout; use crate::arrow::ARROW_DEVICE_CUDA; use crate::arrow::ArrowArray; use crate::arrow::ArrowDeviceArray; @@ -75,8 +78,11 @@ use crate::arrow::arrow_schema_for_array; use crate::arrow::cuda_decimal_value_type; use crate::arrow::list_view::export_device_list_view; use crate::cub::exclusive_sum_i32; +use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; use crate::executor::CudaArrayExt; use crate::executor::execute_validity_cuda; +use crate::kernel::FSSTVarBin; +use crate::kernel::decode_fsst_varbin; /// An implementation of `ExportDeviceArray` that exports Vortex arrays to `ArrowDeviceArray` by /// first decoding the array on the GPU and then converting the canonical type to the nearest @@ -219,6 +225,15 @@ fn export_array( Ok(list_view) => return export_list_view(list_view, ctx).await, Err(array) => array, }; + // The offset-based FSST export always uses the standalone varbin kernel; + // `CudaDispatchMode` only governs `execute_cuda`'s fused-vs-standalone planning. + let array = match array.try_downcast::() { + Ok(fsst) if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin => { + return export_fsst_varbin(fsst, ctx).await; + } + Ok(fsst) => fsst.into_array(), + Err(array) => array, + }; let cuda_array = array.execute_cuda(ctx).await?; export_canonical(cuda_array, ctx).await @@ -302,60 +317,10 @@ fn export_canonical( export_fixed_size_list(fixed_size_list, ctx).await } Canonical::VarBinView(varbinview) => { - if matches!(varbinview.dtype(), DType::Binary(_)) { - return export_binary(varbinview, ctx).await; - } - - let len = varbinview.len(); - let VarBinViewDataParts { - views, - buffers: data_buffers, - validity, - .. - } = varbinview.into_data_parts(); - - let (validity_buffer, null_count) = - export_arrow_validity_buffer(validity, len, 0, ctx).await?; - - let views = ctx.ensure_on_device(views).await?; - let mut buffers = Vec::with_capacity(data_buffers.len() + 3); - buffers.push(validity_buffer); - buffers.push(Some(views)); - for buffer in data_buffers.iter() { - buffers.push(Some(ctx.ensure_on_device(buffer.clone()).await?)); + if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin { + return export_varbin(varbinview, ctx).await; } - // Nanoarrow's Utf8View/BinaryView C layout stores the variadic data buffer sizes - // as the final buffer slot, after the null bitmap, views, and data buffers. - let variadic_buffer_sizes = data_buffers - .iter() - .map(|buffer| i64::try_from(buffer.len())) - .collect::, _>>()?; - buffers.push(Some( - ctx.ensure_on_device(BufferHandle::new_host( - Buffer::from(variadic_buffer_sizes).into_byte_buffer(), - )) - .await?, - )); - - let n_buffers = i64::try_from(buffers.len())?; - let mut private_data = PrivateData::new(buffers, vec![], ctx)?; - let sync_event = private_data.sync_event(); - let arrow_array = ArrowArray { - length: len as i64, - null_count, - offset: 0, - // Arrow Utf8View/BinaryView layout: optional null bitmap, views, data buffers, - // and trailing variadic buffer sizes. - n_buffers, - buffers: private_data.buffer_ptrs.as_mut_ptr(), - n_children: 0, - children: ptr::null_mut(), - release: Some(release_array), - dictionary: ptr::null_mut(), - private_data: Box::into_raw(private_data).cast(), - }; - - Ok((arrow_array, sync_event)) + export_varbinview(varbinview, ctx).await } c => vortex_bail!("unsupported Arrow Device export for {} array", c.dtype()), } @@ -370,12 +335,14 @@ async fn export_dict( ctx: &mut CudaExecutionCtx, ) -> VortexResult<(ArrowArray, SyncEvent)> { let len = array.len(); + let validity = array.validity()?; let parts = array.into_parts(); - let PrimitiveDataParts { - buffer, validity, .. - } = export_dictionary_codes(parts.codes, ctx).await?; + let PrimitiveDataParts { buffer, .. } = export_dictionary_codes(parts.codes, ctx).await?; let (validity_buffer, null_count) = export_arrow_validity_buffer(validity, len, 0, ctx).await?; let codes_buffer = ctx.ensure_on_device(buffer).await?; + // Arrow permits null dictionary values, so preserve the child's validity bitmap. The outer + // bitmap independently marks each row whose code selects a null dictionary value, ensuring + // consumers that require non-null dictionary keys do not lose the logical nulls. let (dictionary, _) = export_array(parts.values, ctx).await?; let mut private_data = PrivateData::new_with_dictionary( @@ -538,8 +505,64 @@ where Ok(BufferHandle::new_device(Arc::new(output_device))) } -/// Export Vortex binary views as an Arrow Device array with standard `Binary` layout. -async fn export_binary( +/// Export Vortex binary views as an Arrow Device array with `Utf8View`/`BinaryView` layout. +async fn export_varbinview( + varbinview: VarBinViewArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<(ArrowArray, SyncEvent)> { + let len = varbinview.len(); + let VarBinViewDataParts { + views, + buffers: data_buffers, + validity, + .. + } = varbinview.into_data_parts(); + + let (validity_buffer, null_count) = export_arrow_validity_buffer(validity, len, 0, ctx).await?; + + let views = ctx.ensure_on_device(views).await?; + let mut buffers = Vec::with_capacity(data_buffers.len() + 3); + buffers.push(validity_buffer); + buffers.push(Some(views)); + for buffer in data_buffers.iter() { + buffers.push(Some(ctx.ensure_on_device(buffer.clone()).await?)); + } + // Nanoarrow's Utf8View/BinaryView C layout stores the variadic data buffer sizes + // as the final buffer slot, after the null bitmap, views, and data buffers. + let variadic_buffer_sizes = data_buffers + .iter() + .map(|buffer| i64::try_from(buffer.len())) + .collect::, _>>()?; + buffers.push(Some( + ctx.ensure_on_device(BufferHandle::new_host( + Buffer::from(variadic_buffer_sizes).into_byte_buffer(), + )) + .await?, + )); + + let n_buffers = i64::try_from(buffers.len())?; + let mut private_data = PrivateData::new(buffers, vec![], ctx)?; + let sync_event = private_data.sync_event(); + let arrow_array = ArrowArray { + length: len as i64, + null_count, + offset: 0, + // Arrow Utf8View/BinaryView layout: optional null bitmap, views, data buffers, + // and trailing variadic buffer sizes. + n_buffers, + buffers: private_data.buffer_ptrs.as_mut_ptr(), + n_children: 0, + children: ptr::null_mut(), + release: Some(release_array), + dictionary: ptr::null_mut(), + private_data: Box::into_raw(private_data).cast(), + }; + + Ok((arrow_array, sync_event)) +} + +/// Export Vortex binary views as an Arrow Device array with standard `Utf8`/`Binary` layout. +async fn export_varbin( varbinview: VarBinViewArray, ctx: &mut CudaExecutionCtx, ) -> VortexResult<(ArrowArray, SyncEvent)> { @@ -555,16 +578,44 @@ async fn export_binary( let views = ctx.ensure_on_device(views).await?; let (offsets, values) = export_binary_buffers(&views, &data_buffers, validity_buffer.as_ref(), len, ctx).await?; + export_varbin_buffers(len, validity_buffer, null_count, offsets, values, ctx) +} - let buffers = vec![validity_buffer, Some(offsets), Some(values)]; +async fn export_fsst_varbin( + fsst: FSSTArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<(ArrowArray, SyncEvent)> { + let FSSTVarBin { + dtype, + len, + offsets, + values, + validity, + } = decode_fsst_varbin(fsst, ctx).await?; + vortex_ensure!( + matches!(dtype, DType::Utf8(_) | DType::Binary(_)), + "FSST produced invalid variable-length dtype {dtype}" + ); + let (validity_buffer, null_count) = export_arrow_validity_buffer(validity, len, 0, ctx).await?; + export_varbin_buffers(len, validity_buffer, null_count, offsets, values, ctx) +} +fn export_varbin_buffers( + len: usize, + validity_buffer: Option, + null_count: i64, + offsets: BufferHandle, + values: BufferHandle, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<(ArrowArray, SyncEvent)> { + let buffers = vec![validity_buffer, Some(offsets), Some(values)]; let mut private_data = PrivateData::new(buffers, vec![], ctx)?; let sync_event = private_data.sync_event(); let arrow_array = ArrowArray { length: len as i64, null_count, offset: 0, - // Arrow Binary layout: optional null bitmap, i32 offsets, contiguous bytes. + // Arrow Utf8/Binary layout: optional null bitmap, i32 offsets, contiguous bytes. n_buffers: 3, buffers: private_data.buffer_ptrs.as_mut_ptr(), n_children: 0, @@ -764,7 +815,7 @@ fn gather_binary_values( /// /// Returns `None` for the buffer when Arrow can omit validity because all rows are valid. /// -/// Returned buffers use zeroed 4-byte padding so cuDF's word-sized mask reads stay in bounds. +/// Returned buffers use zeroed cuDF-sized padding so mask reads stay in bounds. /// Bits at positions `>= len + arrow_offset` within the final data byte are unspecified, as /// Arrow permits. pub(super) async fn export_arrow_validity_buffer( @@ -773,6 +824,11 @@ pub(super) async fn export_arrow_validity_buffer( arrow_offset: usize, ctx: &mut CudaExecutionCtx, ) -> VortexResult<(Option, i64)> { + // Empty arrays do not need a validity buffer; avoid zero-sized CUDA allocations. + if len == 0 { + return Ok((None, 0)); + } + // 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?; @@ -796,16 +852,18 @@ pub(super) async fn export_arrow_validity_buffer( })?; 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 { - repack_arrow_validity_buffer(&bitmap, meta.offset(), len, arrow_offset, ctx)? - }; + let bitmap = + match export_arrow_validity_bitmap(&bitmap, meta.offset(), len, arrow_offset, ctx)? + { + Some(bitmap) => bitmap, + None => 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)) @@ -826,12 +884,78 @@ fn device_zeroed_byte_buffer( byte_len: usize, ctx: &mut CudaExecutionCtx, ) -> VortexResult { - let allocation_len = byte_len.next_multiple_of(size_of::()).max(1); + vortex_ensure!( + byte_len > 0, + "zero-length validity buffers should be omitted" + ); + let allocation_len = byte_len.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING); let mut buffer = ctx.device_alloc::(allocation_len)?; ctx.stream() .memset_zeros(&mut buffer) .map_err(|err| vortex_err!("Failed to zero Arrow validity buffer: {err}"))?; - Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(buffer))).slice(0..byte_len)) + // The memset above zeroed the whole allocation, including cuDF tail padding. + Ok( + BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new_with_zeroed_tail(buffer, 0)?)) + .slice(0..byte_len), + ) +} + +/// Exports a matching-offset bitmap by reusing it or copying it into zero-padded storage. +fn export_arrow_validity_bitmap( + bitmap: &BufferHandle, + input_offset: usize, + len: usize, + arrow_offset: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult> { + if input_offset != arrow_offset { + return Ok(None); + } + + let output_bytes = validity_bitmap_byte_len(len, arrow_offset)?; + let allocation_bytes = output_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING); + if bitmap.has_zeroed_tail_padding(output_bytes, allocation_bytes)? { + return Ok(Some(bitmap.slice(0..output_bytes))); + } + + copy_arrow_validity_buffer(bitmap, output_bytes, ctx).map(Some) +} + +/// Copies a validity bitmap into a new cuDF-padded buffer without shifting bits. +fn copy_arrow_validity_buffer( + input_buffer: &BufferHandle, + output_bytes: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + vortex_ensure!( + output_bytes > 0, + "zero-length validity buffers should be omitted" + ); + vortex_ensure!( + input_buffer.len() >= output_bytes, + "Arrow validity bitmap has {} bytes, expected at least {output_bytes}", + input_buffer.len() + ); + + let allocation_bytes = output_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING); + let mut output = ctx.device_alloc::(allocation_bytes)?; + ctx.stream() + .memset_zeros(&mut output) + .map_err(|err| vortex_err!("Failed to zero Arrow validity buffer padding: {err}"))?; + + let input_view = input_buffer.cuda_view::()?.slice(0..output_bytes); + let mut output_view = output.slice_mut(0..output_bytes); + ctx.stream() + .memcpy_dtod(&input_view, &mut output_view) + .map_err(|err| vortex_err!("Failed to copy Arrow validity buffer: {err}"))?; + + Ok( + BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new_with_zeroed_tail( + output, + output_bytes, + )?)) + .slice(0..output_bytes), + ) } pub fn count_arrow_validity_nulls( @@ -894,8 +1018,8 @@ pub fn count_arrow_validity_nulls( /// /// 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). +/// bit-level offset. The output handle keeps Arrow's logical byte length, while the backing +/// allocation is zero-padded to cuDF's mask allocation size for consumers that read full masks. pub fn repack_arrow_validity_buffer( input_buffer: &BufferHandle, input_offset: usize, @@ -904,7 +1028,18 @@ pub fn repack_arrow_validity_buffer( ctx: &mut CudaExecutionCtx, ) -> VortexResult { let output_bytes = validity_bitmap_byte_len(len, arrow_offset)?; + vortex_ensure!( + output_bytes > 0, + "zero-length validity buffers should be omitted" + ); + // The CUDA kernel writes the bitmap as u64 words, so round the logical byte length up to the + // number of words that cover the exported Arrow bytes. let output_words = output_bytes.div_ceil(size_of::()); + // `device_alloc::` takes a word count, while the padding policy is expressed in bytes. + // Round up so the padded byte allocation is fully represented by whole u64 words. + let allocation_words = output_bytes + .next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING) + .div_ceil(size_of::()); // The kernel loads the input bitmap as 64-bit words. if !input_buffer @@ -914,8 +1049,14 @@ pub fn repack_arrow_validity_buffer( vortex_bail!("Arrow validity repack requires an 8-byte aligned device buffer"); } - let output = ctx.device_alloc::(output_words.max(1))?; - let output_device = CudaDeviceBuffer::new(output); + let mut output = ctx.device_alloc::(allocation_words.max(1))?; + // The repack kernel writes only the logical bitmap words. Zero the whole backing allocation so + // cuDF's padded mask reads see invalid rows, not uninitialized CUDA memory. + ctx.stream() + .memset_zeros(&mut output) + .map_err(|err| vortex_err!("Failed to zero Arrow validity buffer padding: {err}"))?; + // The memset above zeroed all allocation bytes after the logical output. + let output_device = CudaDeviceBuffer::new_with_zeroed_tail(output, output_bytes)?; if output_words > 0 { let input_view = input_buffer.cuda_view::()?; @@ -1300,6 +1441,7 @@ mod tests { use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::arrays::BoolArray; + use vortex::array::arrays::ChunkedArray; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::DictArray; use vortex::array::arrays::FixedSizeListArray; @@ -1309,6 +1451,7 @@ mod tests { use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::TemporalArray; + use vortex::array::arrays::VarBinArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::primitive::PrimitiveArrayExt; use vortex::array::arrays::varbinview::BinaryView; @@ -1326,21 +1469,29 @@ mod tests { use vortex::dtype::PType; use vortex::dtype::half::f16; use vortex::dtype::i256; + use vortex::encodings::fsst::FSST; + use vortex::encodings::fsst::FSSTArrayExt; + use vortex::encodings::fsst::fsst_compress; + use vortex::encodings::fsst::fsst_train_compressor; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::extension::datetime::TimeUnit; + use crate::CudaBufferExt; use crate::CudaExecutionCtx; use crate::arrow::ARROW_DEVICE_CUDA; use crate::arrow::ArrowArray; use crate::arrow::ArrowDeviceArray; use crate::arrow::DeviceArrayExt; use crate::arrow::PrivateData; + use crate::arrow::arrow_schema_for_array; use crate::arrow::canonical::export_arrow_validity_buffer; use crate::arrow::canonical::repack_arrow_validity_buffer; + use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; use crate::device_buffer::cuda_backing_allocation; use crate::session::CudaSession; + use crate::session::VarBinExportLayout; unsafe fn release_exported_array(array: *mut ArrowArray) { unsafe { @@ -1350,6 +1501,58 @@ mod tests { } } + fn cuda_ctx_with_varbin_layout(layout: VarBinExportLayout) -> VortexResult { + let session = vortex::array::array_session() + .with_some(CudaSession::try_default()?.with_varbin_export_layout(layout)); + CudaSession::create_execution_ctx(&session) + } + + // Compress values into a concrete FSST array so exports take the direct FSST path. + fn fsst_array_from( + values: &[Option<&'static [u8]>], + dtype: DType, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let varbin = VarBinArray::from_iter(values.iter().copied(), dtype).into_array(); + let compressor = fsst_train_compressor(&varbin, ctx.execution_ctx())?; + Ok(fsst_compress(&varbin, &compressor, ctx.execution_ctx())?.into_array()) + } + + // Assert an exported varbin array against the logical values: offsets are the prefix sum + // of lengths (nulls contribute zero), values are the non-null bytes concatenated, and the + // Arrow null bitmap marks exactly the null slots. + fn assert_varbin_contents( + array: &ArrowArray, + values: &[Option<&'static [u8]>], + ) -> VortexResult<()> { + let mut expected_offsets = vec![0i32]; + let mut expected_values = Vec::new(); + for value in values { + if let Some(value) = value { + expected_values.extend_from_slice(value); + } + expected_offsets.push(i32::try_from(expected_values.len())?); + } + let null_count = values.iter().filter(|value| value.is_none()).count(); + + assert_binary_layout( + array, + i64::try_from(values.len())?, + i64::try_from(null_count)?, + &expected_offsets, + &expected_values, + )?; + + if null_count > 0 { + let bitmap = private_data_buffer_bytes(array, 0)?; + for (index, value) in values.iter().enumerate() { + let bit = (bitmap.as_ref()[index / 8] >> (index % 8)) & 1; + assert_eq!(bit == 1, value.is_some(), "validity bit {index}"); + } + } + Ok(()) + } + // Assert Arrow Device metadata that consumers use before reading buffers. fn assert_device_metadata( device_array: &ArrowDeviceArray, @@ -1717,7 +1920,7 @@ mod tests { field, Field::new( "", - DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8View)), + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), true, ) ); @@ -1738,7 +1941,42 @@ mod tests { ); let dictionary = unsafe { &*exported.array.array.dictionary }; - assert_varbinview_layout(dictionary, 2, 0, &[out_of_line.len()])?; + assert_binary_layout( + dictionary, + 2, + 0, + &[0, 5, i32::try_from(5 + out_of_line.len())?], + ["alpha", out_of_line].concat().as_bytes(), + )?; + + unsafe { release_exported_array(&raw mut exported.array.array) }; + Ok(()) + } + + #[crate::test] + async fn test_export_dictionary_propagates_value_nulls_to_codes() -> VortexResult<()> { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let array = DictArray::try_new( + PrimitiveArray::from_iter([0u8, 1, 2, 1]).into_array(), + VarBinViewArray::from_iter_nullable_str([ + Some("alpha"), + None, + Some("a dictionary value stored out-of-line"), + ]) + .into_array(), + )? + .into_array(); + let mut exported = array.export_device_array_with_schema(&mut ctx).await?; + + assert_eq!(exported.array.array.null_count, 2); + assert_eq!( + private_data_buffer_bytes(&exported.array.array, 0)?.as_ref(), + &[0b0000_0101] + ); + let dictionary = unsafe { &*exported.array.array.dictionary }; + assert_eq!(dictionary.null_count, 1); unsafe { release_exported_array(&raw mut exported.array.array) }; Ok(()) @@ -1770,7 +2008,7 @@ mod tests { "", DataType::Struct(Fields::from(vec![Field::new( "dict", - DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8View)), + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), true, )])), false, @@ -1807,7 +2045,11 @@ mod tests { true, ) ); - assert_eq!(exported.array.array.null_count, 0); + assert_eq!(exported.array.array.null_count, 1); + assert_eq!( + private_data_buffer_bytes(&exported.array.array, 0)?.as_ref(), + &[0b0000_0101] + ); assert_eq!( private_data_buffer_i16_values(&exported.array.array, 1)?, [0, 1, 0] @@ -2054,9 +2296,8 @@ mod tests { } #[crate::test] - async fn test_export_varbinview() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) - .vortex_expect("failed to create execution context"); + async fn test_export_varbinview_opt_in() -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBinView)?; let out_of_line = "this is a longer string for out-of-line storage"; let array = VarBinViewArray::from_iter_str(["hello", "world", out_of_line]).into_array(); @@ -2217,6 +2458,230 @@ mod tests { Ok(()) } + // Regression test: when the export schema is derived from the dtype alone (the top-level + // array is not one of the concretely-handled encodings), list element fields must still + // reflect the session's varbin export layout, or consumers would read the offset-based + // element data through a view-typed schema. + #[rstest] + #[case::varbin(VarBinExportLayout::VarBin, DataType::Utf8, 3)] + #[case::varbin_view(VarBinExportLayout::VarBinView, DataType::Utf8View, 4)] + #[crate::test] + async fn test_export_chunked_list_utf8_element_matches_schema( + #[case] layout: VarBinExportLayout, + #[case] expected_element_type: DataType, + #[case] expected_element_n_buffers: i64, + ) -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(layout)?; + + let elements = VarBinViewArray::from_iter_str([ + "hello", + "world", + "this is a longer string for out-of-line storage", + ]) + .into_array(); + let chunk = ListArray::try_new( + elements, + PrimitiveArray::from_iter([0i32, 2, 3]).into_array(), + Validity::NonNullable, + )? + .into_array(); + let dtype = chunk.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk], dtype)?.into_array(); + + let mut exported = chunked.export_device_array_with_schema(&mut ctx).await?; + + let field = Field::try_from(&exported.schema)?; + let DataType::List(element_field) = field.data_type() else { + vortex_bail!("expected List schema, got {:?}", field.data_type()); + }; + assert_eq!(element_field.data_type(), &expected_element_type); + + assert_eq!(exported.array.array.n_children, 1); + let children = unsafe { std::slice::from_raw_parts(exported.array.array.children, 1) }; + let element_array = unsafe { &*children[0] }; + assert_eq!(element_array.n_buffers, expected_element_n_buffers); + + unsafe { release_exported_array(&raw mut exported.array.array) }; + Ok(()) + } + + // Standard Arrow Utf8/Binary uses i32 offsets. Oversized FSST exports keep that stable schema + // and return a clear error rather than changing layout based on batch contents. + #[crate::test] + async fn test_oversized_fsst_varbin_export_errors() -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + + let varbin = VarBinArray::from_iter( + [Some(&b"short"[..]), Some(&b"another value"[..])], + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let compressor = fsst_train_compressor(&varbin, ctx.execution_ctx())?; + let fsst = fsst_compress(&varbin, &compressor, ctx.execution_ctx())?; + + // Same codes, but uncompressed lengths whose sum exceeds i32::MAX. + let oversized = FSST::try_new( + DType::Utf8(Nullability::NonNullable), + fsst.symbols().clone(), + fsst.symbol_lengths().clone(), + fsst.codes(), + PrimitiveArray::from_iter([i32::MAX, i32::MAX]).into_array(), + ctx.execution_ctx(), + )? + .into_array(); + let schema = arrow_schema_for_array(&oversized, &mut ctx)?; + assert_eq!( + Field::try_from(&schema)?, + Field::new("", DataType::Utf8, false) + ); + + let error = oversized + .export_device_array(&mut ctx) + .await + .err() + .vortex_expect("oversized FSST varbin export must fail"); + assert!( + error + .to_string() + .contains("FSST decoded size exceeds Arrow i32 offset range") + ); + Ok(()) + } + + // Content coverage for the direct FSST varbin export: offsets, values, and null bitmap, + // not just schema shape. + #[rstest] + #[case::utf8_inline_and_outlined( + vec![Some(&b""[..]), + Some(&b"short"[..]), + Some(&b"this value is stored out of line in the heap"[..])], + DType::Utf8(Nullability::NonNullable), + )] + #[case::partial_nulls( + vec![Some(&b"alpha"[..]), None, Some(&b"gamma"[..]), None, Some(&b"epsilon"[..])], + DType::Utf8(Nullability::Nullable), + )] + #[case::all_nulls( + vec![None, None, None, None, None], + DType::Binary(Nullability::Nullable), + )] + #[case::all_empty( + vec![Some(&b""[..]), Some(&b""[..]), Some(&b""[..])], + DType::Binary(Nullability::NonNullable), + )] + #[case::empty(vec![], DType::Utf8(Nullability::NonNullable))] + #[crate::test] + async fn test_export_fsst_varbin_contents( + #[case] values: Vec>, + #[case] dtype: DType, + ) -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let fsst = fsst_array_from(&values, dtype.clone(), &mut ctx)?; + + let mut exported = fsst.export_device_array_with_schema(&mut ctx).await?; + let expected_data_type = if matches!(dtype, DType::Utf8(_)) { + DataType::Utf8 + } else { + DataType::Binary + }; + assert_eq!( + Field::try_from(&exported.schema)?, + Field::new("", expected_data_type, dtype.is_nullable()) + ); + assert_varbin_contents(&exported.array.array, &values)?; + + unsafe { release_exported_array(&raw mut exported.array.array) }; + Ok(()) + } + + // A sliced FSST array keeps its encoding, so the direct varbin export must respect the + // slice's codes offsets and validity. + #[crate::test] + async fn test_export_sliced_fsst_varbin() -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let values: &[Option<&'static [u8]>] = &[ + Some(&b"alpha"[..]), + Some(&b"this value is stored out of line in the heap"[..]), + None, + Some(&b"delta"[..]), + Some(&b"echo"[..]), + ]; + let fsst = fsst_array_from(values, DType::Utf8(Nullability::Nullable), &mut ctx)?; + let sliced = fsst.slice(1..4)?; + assert!(sliced.as_opt::().is_some()); + + let mut exported = sliced.export_device_array_with_schema(&mut ctx).await?; + assert_eq!( + Field::try_from(&exported.schema)?, + Field::new("", DataType::Utf8, true) + ); + assert_varbin_contents(&exported.array.array, &values[1..4])?; + + unsafe { release_exported_array(&raw mut exported.array.array) }; + Ok(()) + } + + // FSST fields inside a struct take the direct varbin export path through the child + // recursion, and the struct schema reflects the offset-based layout. + #[crate::test] + async fn test_export_struct_with_fsst_field() -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let values: &[Option<&'static [u8]>] = &[ + Some(&b"short"[..]), + Some(&b"this value is stored out of line in the heap"[..]), + ]; + let fsst = fsst_array_from(values, DType::Utf8(Nullability::NonNullable), &mut ctx)?; + let array = StructArray::new( + FieldNames::from_iter(["s"]), + vec![fsst], + values.len(), + Validity::NonNullable, + ) + .into_array(); + + let mut exported = array.export_device_array_with_schema(&mut ctx).await?; + assert_eq!( + Schema::try_from(&exported.schema)?, + Schema::new(vec![Field::new("s", DataType::Utf8, false)]) + ); + assert_eq!(exported.array.array.n_children, 1); + let children = unsafe { std::slice::from_raw_parts(exported.array.array.children, 1) }; + assert_varbin_contents(unsafe { &*children[0] }, values)?; + + unsafe { release_exported_array(&raw mut exported.array.array) }; + Ok(()) + } + + // FSST dictionary values take the direct varbin export path, and the dictionary schema + // reflects the offset-based layout. + #[crate::test] + async fn test_export_dict_with_fsst_values() -> VortexResult<()> { + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let values: &[Option<&'static [u8]>] = &[ + Some(&b"alpha"[..]), + Some(&b"this dictionary value is stored out of line"[..]), + ]; + let fsst = fsst_array_from(values, DType::Utf8(Nullability::NonNullable), &mut ctx)?; + let array = DictArray::try_new(PrimitiveArray::from_iter([0u8, 1, 0]).into_array(), fsst)? + .into_array(); + + let mut exported = array.export_device_array_with_schema(&mut ctx).await?; + assert_eq!( + Field::try_from(&exported.schema)?, + Field::new( + "", + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), + false, + ) + ); + assert!(!exported.array.array.dictionary.is_null()); + let dictionary = unsafe { &*exported.array.array.dictionary }; + assert_varbin_contents(dictionary, values)?; + + unsafe { release_exported_array(&raw mut exported.array.array) }; + Ok(()) + } + #[crate::test] async fn test_export_host_contiguous_list_view() -> VortexResult<()> { let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) @@ -2433,27 +2898,38 @@ mod tests { #[rstest] #[case::utf8( multi_buffer_varbinview(DType::Utf8(Nullability::NonNullable)), - DataType::Utf8View + VarBinExportLayout::VarBin, + DataType::Utf8 )] #[case::binary( multi_buffer_varbinview(DType::Binary(Nullability::NonNullable)), + VarBinExportLayout::VarBin, DataType::Binary )] + #[case::utf8_view( + multi_buffer_varbinview(DType::Utf8(Nullability::NonNullable)), + VarBinExportLayout::VarBinView, + DataType::Utf8View + )] + #[case::binary_view( + multi_buffer_varbinview(DType::Binary(Nullability::NonNullable)), + VarBinExportLayout::VarBinView, + DataType::BinaryView + )] #[crate::test] async fn test_export_varbinview_multiple_variadic_buffers( #[case] fixture: (ArrayRef, [usize; 2]), + #[case] layout: VarBinExportLayout, #[case] expected_data_type: DataType, ) -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) - .vortex_expect("failed to create execution context"); + let mut ctx = cuda_ctx_with_varbin_layout(layout)?; let (array, expected_data_buffer_lengths) = fixture; let mut exported = array.export_device_array_with_schema(&mut ctx).await?; let field = Field::try_from(&exported.schema)?; - let is_binary = expected_data_type == DataType::Binary; assert_eq!(field, Field::new("", expected_data_type, false)); - if is_binary { + if layout == VarBinExportLayout::VarBin { assert_binary_layout( &exported.array.array, 3, @@ -2598,7 +3074,11 @@ mod tests { let children = unsafe { std::slice::from_raw_parts(exported.array.array.children, 1) }; let elements = unsafe { &*children[0] }; - assert_eq!(elements.null_count, 0); + assert_eq!(elements.null_count, 2); + assert_eq!( + private_data_buffer_bytes(elements, 0)?.as_ref(), + &[0b0001_0101] + ); assert_eq!( private_data_buffer_i16_values(elements, 1)?, [0, 1, 0, 1, 2] @@ -2849,8 +3329,15 @@ mod tests { .slice(1..4)?; let mut exported = utf8.export_device_array_with_schema(&mut ctx).await?; let field = Field::try_from(&exported.schema)?; - assert_eq!(field, Field::new("", DataType::Utf8View, true)); - assert_varbinview_shape(&exported.array.array, 3, 1)?; + assert_eq!(field, Field::new("", DataType::Utf8, true)); + let sliced_out_of_line = "this out-of-line value remains in the slice"; + assert_binary_layout( + &exported.array.array, + 3, + 1, + &[0, 5, 5, i32::try_from(5 + sliced_out_of_line.len())?], + ["hello", sliced_out_of_line].concat().as_bytes(), + )?; assert_eq!(exported.array.device_type, ARROW_DEVICE_CUDA); let private_data = unsafe { &*exported.array.array.private_data.cast::() }; @@ -2955,13 +3442,152 @@ mod tests { let backing_bytes = backing.to_host_sync(); assert_eq!( backing_bytes.len(), - output_bytes.next_multiple_of(size_of::()) + output_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING) + ); + assert!(backing_bytes[output_bytes..].iter().all(|byte| *byte == 0)); + + Ok(()) + } + + #[crate::test] + async fn test_export_validity_buffer_pads_matching_offset() -> VortexResult<()> { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let len = 3; + let arrow_offset = 0; + let (buffer, null_count) = export_arrow_validity_buffer( + Validity::from(BitBuffer::from_iter([true, false, true])), + len, + arrow_offset, + &mut ctx, + ) + .await?; + ctx.synchronize_stream()?; + + assert_eq!(null_count, 1); + let buffer = buffer.vortex_expect("nullable validity should export a null buffer"); + let output_bytes = (len + arrow_offset).div_ceil(8); + assert_eq!(buffer.len(), output_bytes); + let actual = BitBuffer::new(buffer.to_host_sync(), len + arrow_offset) + .iter() + .collect::>(); + assert_eq!(actual, [true, false, true]); + + let backing = cuda_backing_allocation(&buffer)?; + let backing_bytes = backing.to_host_sync(); + assert_eq!( + backing_bytes.len(), + output_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING) ); assert!(backing_bytes[output_bytes..].iter().all(|byte| *byte == 0)); Ok(()) } + #[crate::test] + async fn test_export_validity_buffer_reuses_matching_padded_device_bitmap() -> VortexResult<()> + { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let len = 3; + let source = BitBuffer::from_iter([true, false, true]); + let (input_offset, _, input_buffer) = source.into_inner(); + let input_buffer = ctx + .ensure_on_device(BufferHandle::new_host(input_buffer)) + .await?; + let input_ptr = input_buffer.cuda_device_ptr()?; + let validity = BoolArray::new_handle( + input_buffer.clone(), + input_offset, + len, + Validity::NonNullable, + ) + .into_array(); + + let (buffer, null_count) = + export_arrow_validity_buffer(Validity::Array(validity), len, input_offset, &mut ctx) + .await?; + ctx.synchronize_stream()?; + + assert_eq!(null_count, 1); + let buffer = buffer.vortex_expect("nullable validity should export a null buffer"); + assert_eq!(buffer.cuda_device_ptr()?, input_ptr); + assert_eq!(buffer.len(), (len + input_offset).div_ceil(8)); + let actual = BitBuffer::new(buffer.to_host_sync(), len + input_offset) + .iter() + .collect::>(); + let expected = std::iter::repeat_n(false, input_offset) + .chain([true, false, true]) + .collect::>(); + assert_eq!(actual, expected); + + Ok(()) + } + + #[crate::test] + async fn test_export_validity_buffer_repacks_matching_offset_without_tail_padding() + -> VortexResult<()> { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let len = 3; + let source = BitBuffer::from_iter((0..80).map(|idx| idx % 3 != 1)); + let (input_offset, _, input_buffer) = source.into_inner(); + let input_buffer = ctx + .ensure_on_device(BufferHandle::new_host(input_buffer)) + .await?; + let input_ptr = input_buffer.cuda_device_ptr()?; + let validity = BoolArray::new_handle( + input_buffer.clone(), + input_offset, + len, + Validity::NonNullable, + ) + .into_array(); + + let (buffer, null_count) = + export_arrow_validity_buffer(Validity::Array(validity), len, input_offset, &mut ctx) + .await?; + ctx.synchronize_stream()?; + + assert_eq!(null_count, 1); + let buffer = buffer.vortex_expect("nullable validity should export a null buffer"); + assert_ne!(buffer.cuda_device_ptr()?, input_ptr); + let output_bytes = (len + input_offset).div_ceil(8); + assert_eq!(buffer.len(), output_bytes); + let actual = BitBuffer::new(buffer.to_host_sync(), len + input_offset) + .iter() + .collect::>(); + let expected = std::iter::repeat_n(false, input_offset) + .chain([true, false, true]) + .collect::>(); + assert_eq!(actual, expected); + + let backing = cuda_backing_allocation(&buffer)?; + assert_eq!( + backing.len(), + output_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING) + ); + + Ok(()) + } + + #[crate::test] + async fn test_export_empty_validity_buffer_is_omitted() -> VortexResult<()> { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let (buffer, null_count) = + export_arrow_validity_buffer(Validity::AllInvalid, 0, 0, &mut ctx).await?; + + assert_eq!(null_count, 0); + assert!(buffer.is_none()); + + Ok(()) + } + #[crate::test] async fn test_export_all_false_validity_buffer_is_zeroed_on_device() -> VortexResult<()> { let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) @@ -2983,6 +3609,11 @@ mod tests { let bytes = buffer.to_host_sync(); assert_eq!(bytes.len(), (len + arrow_offset).div_ceil(8)); assert!(bytes.iter().all(|byte| *byte == 0)); + let backing = cuda_backing_allocation(&buffer)?; + assert_eq!( + backing.len(), + bytes.len().next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING) + ); Ok(()) } @@ -3113,7 +3744,7 @@ mod tests { Ok(()) } - // Check nullable string-view exports include Arrow null bitmaps. + // Check nullable variable-length exports include Arrow null bitmaps. #[crate::test] async fn test_export_nullable_varbinview() -> VortexResult<()> { let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) @@ -3126,7 +3757,7 @@ mod tests { Some("this is a longer string for out-of-line storage"), ]) .into_array(), - 4, + 3, 1, &mut ctx, ) @@ -3227,7 +3858,7 @@ mod tests { assert_eq!(nested_primitive_child.n_children, 0); let string_child = unsafe { &*nested_children[1] }; - assert_eq!(string_child.n_buffers, 4); + assert_eq!(string_child.n_buffers, 3); assert_eq!(string_child.n_children, 0); let string_buffers = unsafe { std::slice::from_raw_parts( @@ -3238,7 +3869,6 @@ mod tests { assert!(string_buffers[0].is_null()); assert!(!string_buffers[1].is_null()); assert!(!string_buffers[2].is_null()); - assert!(!string_buffers[3].is_null()); unsafe { release_exported_array(&raw mut device_array.array) }; Ok(()) @@ -3347,7 +3977,7 @@ mod tests { Schema::new(vec![ Field::new("a", DataType::UInt32, false), Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Utf8View, false), + Field::new("c", DataType::Utf8, false), ]) ); assert_eq!(exported.array.array.length, 5); @@ -3378,7 +4008,7 @@ mod tests { "nested", DataType::Struct(Fields::from(vec![ Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Utf8View, false), + Field::new("c", DataType::Utf8, false), ])), false, ), @@ -3473,23 +4103,35 @@ mod tests { } #[crate::test] - async fn test_export_varbinview_with_schema_uses_utf8_view_layout() -> VortexResult<()> { + async fn test_export_varbinview_with_schema_uses_utf8_layout() -> VortexResult<()> { let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let japanese = "こんにちは"; let long_emoji = "🦀 and 🚀 make this string out-of-line"; - let array = VarBinViewArray::from_iter_str(["", "hello", "é", "🦀", japanese, long_emoji]) - .into_array(); + let values = ["", "hello", "é", "🦀", japanese, long_emoji]; + let array = VarBinViewArray::from_iter_str(values).into_array(); let mut exported = array.export_device_array_with_schema(&mut ctx).await?; let field = Field::try_from(&exported.schema)?; - assert_eq!(field, Field::new("", DataType::Utf8View, false)); - assert_varbinview_layout( + assert_eq!(field, Field::new("", DataType::Utf8, false)); + let mut expected_offsets = Vec::with_capacity(values.len() + 1); + expected_offsets.push(0i32); + for value in values { + expected_offsets.push( + expected_offsets + .last() + .copied() + .vortex_expect("offsets is non-empty") + + i32::try_from(value.len())?, + ); + } + assert_binary_layout( &exported.array.array, 6, 0, - &[japanese.len() + long_emoji.len()], + &expected_offsets, + values.concat().as_bytes(), )?; assert_eq!(exported.array.device_type, ARROW_DEVICE_CUDA); diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index 9545ba6abee..071e22085c4 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -41,7 +41,6 @@ use vortex::array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex::array::arrays::list::ListArrayExt; 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; @@ -55,9 +54,11 @@ use vortex::error::vortex_err; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::current::CurrentThreadRuntime; use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; use crate::CudaBufferExt; use crate::CudaExecutionCtx; +use crate::VarBinExportLayout; mod arrow_c_abi { #![allow(dead_code)] @@ -814,14 +815,27 @@ fn arrow_device_export_field( .arrow() .to_arrow_field(name.as_ref(), dtype)?; - let data_type = match dtype { - DType::Binary(_) => DataType::Binary, - DType::Decimal(decimal_dtype, _) => arrow_device_export_decimal_data_type(*decimal_dtype), - DType::Struct(struct_dtype, _) => { - DataType::Struct(arrow_device_export_struct_fields(struct_dtype, ctx)?.into()) - } - _ => return Ok(field), - }; + let data_type = + match (ctx.cuda_session().varbin_export_layout(), dtype) { + (VarBinExportLayout::VarBin, DType::Utf8(_)) => DataType::Utf8, + (VarBinExportLayout::VarBin, DType::Binary(_)) => DataType::Binary, + (VarBinExportLayout::VarBinView, DType::Utf8(_)) => DataType::Utf8View, + (VarBinExportLayout::VarBinView, DType::Binary(_)) => DataType::BinaryView, + (_, DType::Decimal(decimal_dtype, _)) => { + arrow_device_export_decimal_data_type(*decimal_dtype) + } + (_, DType::Struct(struct_dtype, _)) => { + DataType::Struct(arrow_device_export_struct_fields(struct_dtype, ctx)?.into()) + } + // List elements are exported through the same layout-dependent paths as top-level + // arrays, so their fields must be rebuilt recursively. Without this, an element such + // as Utf8 would keep `to_arrow_field`'s Utf8View mapping while the data is exported + // with the offset-based layout. + (_, DType::List(element_dtype, _)) => DataType::List(Arc::new( + arrow_device_export_field(Field::LIST_FIELD_DEFAULT_NAME, element_dtype, ctx)?, + )), + _ => return Ok(field), + }; Ok( Field::new(field.name().clone(), data_type, field.is_nullable()) diff --git a/vortex-cuda/src/canonical.rs b/vortex-cuda/src/canonical.rs index 4f3d6fd37e3..9f1fce7e68e 100644 --- a/vortex-cuda/src/canonical.rs +++ b/vortex-cuda/src/canonical.rs @@ -7,7 +7,6 @@ use async_trait::async_trait; use futures::future::try_join_all; 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; @@ -23,6 +22,7 @@ use vortex::array::arrays::struct_::StructDataParts; use vortex::array::arrays::varbinview::BinaryView; use vortex::array::arrays::varbinview::VarBinViewDataParts; use vortex::array::buffer::BufferHandle; +use vortex::array::legacy_session; use vortex::buffer::BitBuffer; use vortex::buffer::Buffer; use vortex::buffer::ByteBuffer; @@ -38,6 +38,7 @@ pub trait CanonicalCudaExt { #[async_trait] impl CanonicalCudaExt for Canonical { + #[allow(clippy::disallowed_methods)] async fn into_host(self) -> VortexResult { match self { Canonical::Struct(struct_array) => { @@ -55,7 +56,7 @@ impl CanonicalCudaExt for Canonical { host_fields.push( field .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut legacy_session().create_execution_ctx())? .into_host() .await? .into_array(), @@ -144,7 +145,7 @@ impl CanonicalCudaExt for Canonical { let host_storage = ext .storage_array() .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut legacy_session().create_execution_ctx())? .into_host() .await? .into_array(); diff --git a/vortex-cuda/src/device_buffer.rs b/vortex-cuda/src/device_buffer.rs index bc07cd5daaa..954731fab44 100644 --- a/vortex-cuda/src/device_buffer.rs +++ b/vortex-cuda/src/device_buffer.rs @@ -24,11 +24,16 @@ use vortex::buffer::ByteBuffer; use vortex::buffer::ByteBufferMut; use vortex::error::VortexExpect; use vortex::error::VortexResult; +use vortex::error::vortex_ensure; use vortex::error::vortex_err; use vortex::error::vortex_panic; use crate::stream::await_stream_callback; +/// cuDF may read Arrow validity masks through a 64-byte padded extent, so keep +/// exported masks logically sliced but zero-pad their backing allocation. +pub(crate) const CUDF_VALIDITY_BUFFER_PADDING: usize = 64; + /// A [`DeviceBuffer`] wrapping a CUDA GPU allocation. /// /// Like the host `BufferHandle` variant, all slicing/referencing works in terms of byte units. @@ -43,6 +48,12 @@ pub struct CudaDeviceBuffer { device_ptr: u64, /// Minimum required alignment of the buffer. alignment: Alignment, + /// Allocation-relative byte offset where zeroed tail padding begins. + /// + /// `None` means the tail contents are not tracked. Arrow validity export uses this to decide + /// whether a sliced logical bitmap can safely reuse this allocation for cuDF's padded mask + /// reads without copying or repacking. + zeroed_tail_start: Option, } mod private { @@ -101,9 +112,25 @@ impl CudaDeviceBuffer { len, device_ptr, alignment: Alignment::of::(), + zeroed_tail_start: None, } } + /// Wraps a CUDA allocation and records that bytes from `zeroed_tail_start` are already zeroed. + pub(crate) fn new_with_zeroed_tail( + cuda_slice: CudaSlice, + zeroed_tail_start: usize, + ) -> VortexResult { + let mut buffer = Self::new(cuda_slice); + vortex_ensure!( + zeroed_tail_start <= buffer.len, + "zeroed tail start {zeroed_tail_start} exceeds CUDA allocation length {}", + buffer.len + ); + buffer.zeroed_tail_start = Some(zeroed_tail_start); + Ok(buffer) + } + /// Returns the byte offset within the allocated buffer. pub fn offset(&self) -> usize { self.offset @@ -149,6 +176,7 @@ pub(crate) fn cuda_backing_allocation(handle: &BufferHandle) -> VortexResult VortexResult; + + /// Returns whether this buffer has zeroed tail padding for the requested extent. + /// + /// # Errors + /// + /// Returns an error if the buffer is not a CUDA buffer. + fn has_zeroed_tail_padding(&self, logical_len: usize, padded_len: usize) -> VortexResult; } impl CudaBufferExt for BufferHandle { @@ -197,6 +232,40 @@ impl CudaBufferExt for BufferHandle { Ok(ptr) } + + fn has_zeroed_tail_padding(&self, logical_len: usize, padded_len: usize) -> VortexResult { + let device_buffer = self + .as_device_opt() + .ok_or_else(|| vortex_err!("Buffer is not on device"))?; + + let cuda_buf = device_buffer + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("expected CudaDeviceBuffer, was {device_buffer:?}"))?; + + if logical_len > padded_len || logical_len > cuda_buf.len { + return Ok(false); + } + + let Some(required_end) = cuda_buf.offset.checked_add(padded_len) else { + return Ok(false); + }; + if required_end > cuda_buf.allocation.as_bytes_view().len() { + return Ok(false); + } + + if logical_len == padded_len { + return Ok(true); + } + + let Some(zeroed_tail_start) = cuda_buf.zeroed_tail_start else { + return Ok(false); + }; + let Some(logical_end) = cuda_buf.offset.checked_add(logical_len) else { + return Ok(false); + }; + Ok(zeroed_tail_start <= logical_end) + } } impl Debug for CudaDeviceBuffer { @@ -206,6 +275,7 @@ impl Debug for CudaDeviceBuffer { .field("device_ptr", &self.device_ptr) .field("offset", &self.offset) .field("len", &self.len) + .field("zeroed_tail_start", &self.zeroed_tail_start) .finish() } } @@ -345,6 +415,7 @@ impl DeviceBuffer for CudaDeviceBuffer { len: new_len, device_ptr: self.device_ptr, alignment: self.alignment, + zeroed_tail_start: self.zeroed_tail_start, }) } @@ -361,6 +432,7 @@ impl DeviceBuffer for CudaDeviceBuffer { len: self.len, device_ptr: self.device_ptr, alignment, + zeroed_tail_start: self.zeroed_tail_start, })) } else if alignment > Alignment::new(256) { vortex_panic!("we do not support alignment greater than 256") diff --git a/vortex-cuda/src/kernel/arrays/dict.rs b/vortex-cuda/src/kernel/arrays/dict.rs index 394de699f49..840a3cdbbac 100644 --- a/vortex-cuda/src/kernel/arrays/dict.rs +++ b/vortex-cuda/src/kernel/arrays/dict.rs @@ -10,11 +10,13 @@ use tracing::instrument; use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::IntoArray; +use vortex::array::arrays::BoolArray; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::Dict; use vortex::array::arrays::DictArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::VarBinViewArray; +use vortex::array::arrays::bool::BoolDataParts; use vortex::array::arrays::decimal::DecimalDataParts; use vortex::array::arrays::dict::DictArraySlotsExt; use vortex::array::arrays::primitive::PrimitiveDataParts; @@ -29,6 +31,7 @@ use vortex::dtype::NativePType; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; @@ -55,6 +58,8 @@ impl CudaExecute for DictExecutor { let values_dtype = dict_array.values().dtype().clone(); match &values_dtype { + // Nullable dictionary values expose their logical validity as a lazy `Dict`. + DType::Bool(..) => execute_dict_bool(dict_array, ctx).await, DType::Decimal(..) => execute_dict_decimal(dict_array, ctx).await, DType::Primitive(..) => execute_dict_prim(dict_array, ctx).await, DType::Utf8(..) | DType::Binary(..) => execute_dict_varbinview(dict_array, ctx).await, @@ -86,6 +91,80 @@ async fn execute_dict_prim(dict: DictArray, ctx: &mut CudaExecutionCtx) -> Vorte }) } +/// Gather bit-packed boolean values through dictionary codes. +/// +/// This path is especially important for dictionary validity. When dictionary values are +/// nullable, `Dict::validity()` represents `take(values_validity, codes)` lazily as a `Dict`. +/// Materializing that validity on the GPU therefore requires a boolean dictionary gather even +/// when the user-visible array contains strings or another non-boolean type. +async fn execute_dict_bool(dict: DictArray, ctx: &mut CudaExecutionCtx) -> VortexResult { + let values = dict.values().clone().execute_cuda(ctx).await?.into_bool(); + let codes = dict + .codes() + .clone() + .execute_cuda(ctx) + .await? + .into_primitive(); + let codes_ptype = codes.ptype(); + + match_each_integer_ptype!(codes_ptype, |I| { + execute_dict_bool_typed::(values, codes, ctx).await + }) +} + +async fn execute_dict_bool_typed( + values: BoolArray, + codes: PrimitiveArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + vortex_ensure!(!codes.is_empty(), "cannot CUDA-decode an empty dictionary"); + let codes_len = codes.len(); + + let values_len = values.len(); + let values_validity = values.validity()?; + let BoolDataParts { + bits: values_buffer, + meta: values_meta, + } = values.into_data().into_parts(values_len); + let output_validity = values_validity.take(&codes.clone().into_array())?; + let PrimitiveDataParts { + buffer: codes_buffer, + .. + } = codes.into_data_parts(); + + let values_device = ctx.ensure_on_device(values_buffer).await?; + let codes_device = ctx.ensure_on_device(codes_buffer).await?; + + // Each CUDA thread owns complete output bytes, avoiding races between threads writing + // different bits in the same byte. The kernel handles the final partial byte explicitly. + let output_bytes = codes_len.div_ceil(8); + let output_slice = ctx.device_alloc::(output_bytes)?; + let output_device = CudaDeviceBuffer::new(output_slice); + + let values_view = values_device.cuda_view::()?; + let codes_view = codes_device.cuda_view::()?; + let output_view = output_device.as_view::(); + let codes_len_u64 = codes_len as u64; + let values_offset_u64 = values_meta.offset() as u64; + + let codes_ptype = I::PTYPE.to_string(); + let kernel_function = ctx.load_function_with_suffixes("dict", &["bool", &codes_ptype])?; + ctx.launch_kernel(&kernel_function, output_bytes, |args| { + args.arg(&codes_view) + .arg(&codes_len_u64) + .arg(&values_view) + .arg(&values_offset_u64) + .arg(&output_view); + })?; + + Ok(Canonical::Bool(BoolArray::new_handle( + BufferHandle::new_device(Arc::new(output_device)), + 0, + codes_len, + output_validity, + ))) +} + async fn execute_dict_prim_typed( values: PrimitiveArray, codes: PrimitiveArray, @@ -298,6 +377,7 @@ async fn execute_dict_varbinview( #[cfg(test)] mod tests { use vortex::array::IntoArray; + use vortex::array::arrays::BoolArray; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::DictArray; use vortex::array::arrays::PrimitiveArray; @@ -323,6 +403,43 @@ mod tests { )) } + #[crate::test] + async fn test_cuda_dict_bool_gathers_packed_validity_bits() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + // Slicing leaves the dictionary values at a non-zero bit offset. Thirteen codes also + // exercise a final partial output byte. + let values = BoolArray::from_iter([ + false, true, false, true, false, true, true, false, true, false, + ]) + .into_array() + .slice(3..8)?; + let codes = PrimitiveArray::new( + Buffer::from(vec![0u8, 1, 2, 3, 4, 3, 2, 1, 0, 4, 1, 3, 0]), + NonNullable, + ); + let expected = DictArray::try_new(codes.clone().into_array(), values.clone())?.into_array(); + + let codes_handle = cuda_ctx + .ensure_on_device(codes.buffer_handle().clone()) + .await?; + let device_codes = + PrimitiveArray::from_buffer_handle(codes_handle, codes.ptype(), codes.validity()?); + let dict = DictArray::try_new(device_codes.into_array(), values)?.into_array(); + + let actual = DictExecutor + .execute(dict, &mut cuda_ctx) + .await? + .into_host() + .await? + .into_bool(); + + assert_arrays_eq!(actual.into_array(), expected, &mut ctx); + Ok(()) + } + #[crate::test] async fn test_cuda_dict_u32_values_u8_codes() -> VortexResult<()> { let mut ctx = vortex_array::array_session().create_execution_ctx(); diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index eef0db049b8..de46d0b8fa1 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -17,14 +17,16 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::arrays::varbin::VarBinArrayExt; -use vortex::array::arrays::varbinview::BinaryView; use vortex::array::arrays::varbinview::build_views::MAX_BUFFER_LEN; use vortex::array::arrays::varbinview::build_views::build_views; +use vortex::array::buffer::BufferHandle; use vortex::array::buffer::DeviceBuffer; use vortex::array::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::buffer::ByteBuffer; +use vortex::dtype::DType; use vortex::dtype::NativePType; use vortex::encodings::fsst::FSST; use vortex::encodings::fsst::FSSTArray; @@ -38,6 +40,34 @@ use crate::CudaDeviceBuffer; use crate::executor::CudaExecute; use crate::executor::CudaExecutionCtx; +/// Device-resident offset-based result of FSST decompression. +pub(crate) struct FSSTVarBin { + pub(crate) dtype: DType, + pub(crate) len: usize, + pub(crate) offsets: BufferHandle, + pub(crate) values: BufferHandle, + pub(crate) validity: Validity, +} + +/// Returns validity backing bytes and the bit offset of the first row for the FSST kernels. +/// +/// `BitBuffer` slices normalize whole-byte offsets but can retain a sub-byte offset. Passing that +/// offset to CUDA lets the kernels address sliced validity directly without repacking it on host. +fn cuda_validity( + validity: &Validity, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<(u64, ByteBuffer)> { + let bits = validity + .clone() + .execute_mask(len, ctx.execution_ctx())? + .into_bit_buffer(); + let (bit_offset, bit_len, bytes) = bits.into_inner(); + debug_assert_eq!(bit_len, len); + let byte_len = (bit_offset + bit_len).div_ceil(8); + Ok((bit_offset as u64, bytes.slice(0..byte_len))) +} + /// CUDA decoder for FSST. #[derive(Debug)] pub(crate) struct FSSTExecutor; @@ -61,16 +91,15 @@ impl CudaExecute for FSSTExecutor { let dtype = fsst.dtype().clone(); let validity = fsst.codes().validity()?; - if fsst.is_empty() || validity.definitely_all_null() { - let empty = unsafe { - VarBinViewArray::new_unchecked( - Buffer::::zeroed(fsst.len()), - Arc::from([]), - dtype, - validity, - ) - }; - return Ok(Canonical::VarBinView(empty)); + if fsst.is_empty() { + return Ok(Canonical::empty(&dtype)); + } + + if validity.definitely_all_null() { + let views = ctx.copy_to_device(vec![0i128; fsst.len()])?.await?; + return Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([]), dtype, validity) + })); } let lens = fsst @@ -105,6 +134,138 @@ impl CudaExecute for FSSTExecutor { } } +/// Decode FSST directly into Arrow-compatible i32 offsets and contiguous values on device. +pub(crate) async fn decode_fsst_varbin( + fsst: FSSTArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let dtype = fsst.dtype().clone(); + let validity = fsst.codes().validity()?; + let len = fsst.len(); + let lens = fsst + .uncompressed_lengths() + .clone() + .execute::(ctx.execution_ctx())?; + let codes_offsets = fsst + .codes() + .offsets() + .clone() + .execute::(ctx.execution_ctx())?; + + let output_offsets = match_each_integer_ptype!(lens.ptype(), |P| { + let mut offsets = Vec::with_capacity(lens.len() + 1); + let mut acc = 0u64; + offsets.push(0i32); + #[allow(clippy::unnecessary_cast)] + for &length in lens.as_slice::

() { + let length = u64::try_from(length as i128) + .map_err(|_| vortex_err!("FSST uncompressed length cannot be negative"))?; + acc = acc + .checked_add(length) + .ok_or_else(|| vortex_err!("FSST decoded size overflow"))?; + offsets + .push(i32::try_from(acc).map_err(|_| { + vortex_err!("FSST decoded size exceeds Arrow i32 offset range") + })?); + } + VortexResult::Ok(offsets) + })?; + let total_size = usize::try_from( + *output_offsets + .last() + .vortex_expect("output_offsets has at least one entry"), + )?; + + if total_size == 0 { + let offsets = ctx.copy_to_device(output_offsets)?.await?; + let allocation = CudaDeviceBuffer::new(ctx.device_alloc::(1)?); + let values = BufferHandle::new_device(allocation.slice(0..0)); + return Ok(FSSTVarBin { + dtype, + len, + offsets, + values, + validity, + }); + } + + match_each_unsigned_integer_ptype!(codes_offsets.ptype().to_unsigned(), |U| { + decode_fsst_varbin_typed::(fsst, codes_offsets, output_offsets, total_size, ctx).await + }) +} + +async fn decode_fsst_varbin_typed( + fsst: FSSTArray, + codes_offsets: PrimitiveArray, + output_offsets: Vec, + total_size: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + U: NativePType + DeviceRepr + Send + Sync + 'static, +{ + let dtype = fsst.dtype().clone(); + let validity = fsst.codes().validity()?; + let len = fsst.len(); + let len_u64 = len as u64; + let symbols_u64: Vec = fsst.symbols().iter().map(|s| s.to_u64()).collect(); + let symbol_lengths = fsst.symbol_lengths().clone(); + let codes_bytes_handle = fsst.codes_bytes_handle().clone(); + let PrimitiveDataParts { + buffer: codes_offsets_buffer, + .. + } = codes_offsets.into_data_parts(); + let (validity_bit_offset, validity_bits) = cuda_validity(&validity, len, ctx)?; + + let (symbols, symbol_lengths, output_offsets, validity_device, codes_bytes, codes_offsets) = futures::try_join!( + ctx.copy_to_device(symbols_u64)?, + ctx.copy_to_device(symbol_lengths)?, + ctx.copy_to_device(output_offsets)?, + ctx.copy_to_device(validity_bits)?, + ctx.ensure_on_device(codes_bytes_handle), + ctx.ensure_on_device(codes_offsets_buffer), + )?; + + // The kernel gates store widths on `out_pos % N` relative to the base, so the base must + // satisfy the widest store (u128 → 16). + let output = ctx.device_alloc::(total_size)?; + let (output_base_ptr, _) = output.device_ptr(ctx.stream()); + assert_eq!( + output_base_ptr % 16, + 0, + "output base not 16-aligned: {output_base_ptr:#x}", + ); + + let codes_bytes_view = codes_bytes.cuda_view::()?; + let codes_offsets_view = codes_offsets.cuda_view::()?; + let symbols_view = symbols.cuda_view::()?; + let symbol_lengths_view = symbol_lengths.cuda_view::()?; + let output_offsets_view = output_offsets.cuda_view::()?; + let validity_view = validity_device.cuda_view::()?; + let ptype = U::PTYPE.to_string(); + let cuda_function = ctx.load_function_with_suffixes("fsst", &["varbin", &ptype])?; + + ctx.launch_kernel(&cuda_function, len, |args| { + args.arg(&codes_bytes_view) + .arg(&codes_offsets_view) + .arg(&symbols_view) + .arg(&symbol_lengths_view) + .arg(&output_offsets_view) + .arg(&validity_view) + .arg(&validity_bit_offset) + .arg(&output) + .arg(&len_u64); + })?; + + Ok(FSSTVarBin { + dtype, + len, + offsets: output_offsets, + values: BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(output))), + validity, + }) +} + async fn decode_fsst( fsst: FSSTArray, codes_offsets: PrimitiveArray, @@ -126,6 +287,13 @@ where ) .vortex_expect("total_size fits in usize"); + if total_size == 0 { + let views = ctx.copy_to_device(vec![0i128; num_strings])?.await?; + return Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([]), dtype, validity) + })); + } + let symbols_u64: Vec = fsst.symbols().iter().map(|s| s.to_u64()).collect(); let symbol_lengths = fsst.symbol_lengths().clone(); let codes_bytes_handle = fsst.codes_bytes_handle().clone(); @@ -134,18 +302,13 @@ where .. } = codes_offsets.into_data_parts(); - let (.., validity_bits) = validity - .clone() - .execute_mask(num_strings, ctx.execution_ctx())? - .into_bit_buffer() - .sliced() - .into_inner(); + let (validity_bit_offset, validity_bits) = cuda_validity(&validity, num_strings, ctx)?; let (symbols, symbol_lengths, output_offsets, validity_device, codes_bytes, codes_offsets) = futures::try_join!( ctx.copy_to_device(symbols_u64)?, ctx.copy_to_device(symbol_lengths)?, ctx.copy_to_device(output_offsets)?, - ctx.copy_to_device(validity_bits.to_vec())?, + ctx.copy_to_device(validity_bits)?, ctx.ensure_on_device(codes_bytes_handle), ctx.ensure_on_device(codes_offsets_buffer), )?; @@ -153,6 +316,9 @@ where // The kernel checks store alignment relative to the base via // `out_pos % N`, so the base must satisfy the widest store (u128 → 16). let device_output = ctx.device_alloc::(total_size)?; + let device_views = (total_size <= MAX_BUFFER_LEN) + .then(|| ctx.device_alloc::(num_strings)) + .transpose()?; let (output_base_ptr, _) = device_output.device_ptr(ctx.stream()); assert_eq!( output_base_ptr % 16, @@ -168,6 +334,7 @@ where let validity_view = validity_device.cuda_view::()?; let cuda_function = ctx.load_function("fsst", &[U::PTYPE])?; + let null_views = 0u64; ctx.launch_kernel(&cuda_function, num_strings, |args| { args.arg(&codes_bytes_view) .arg(&codes_offsets_view) @@ -175,10 +342,29 @@ where .arg(&symbol_lengths_view) .arg(&output_offsets_view) .arg(&validity_view) - .arg(&device_output) - .arg(&num_strings_u64); + .arg(&validity_bit_offset) + .arg(&device_output); + if let Some(device_views) = device_views.as_ref() { + args.arg(device_views); + } else { + args.arg(&null_views); + } + args.arg(&num_strings_u64); })?; + // Fast path: the decoded heap fits in one BinaryView backing buffer, so the kernel wrote + // views directly and both buffers can remain on-device. Larger heaps use the host rollover + // path below to split decoded bytes across multiple backing buffers. + if let Some(device_views) = device_views { + let views = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(device_views))); + let bytes = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(device_output))); + return Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([bytes]), dtype, validity) + })); + } + + // BinaryView offsets are u32. Retain the host rollover path for decoded heaps + // that need multiple backing buffers; ordinary batches stay entirely on-device. let host_bytes = CudaDeviceBuffer::new(device_output) .copy_to_host(Alignment::new(1))? .await?; @@ -200,10 +386,13 @@ where #[cfg(test)] mod tests { + use arrow_schema::DataType; + use arrow_schema::Field; use rstest::rstest; use vortex::array::IntoArray; use vortex::array::arrays::VarBinArray; use vortex::array::assert_arrays_eq; + use vortex::buffer::Buffer; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::encodings::fsst::fsst_compress; @@ -213,34 +402,73 @@ mod tests { use super::*; use crate::CanonicalCudaExt; + use crate::arrow::DeviceArrayExt; + use crate::arrow::release_device_array; + use crate::arrow::release_schema; use crate::session::CudaSession; + use crate::session::VarBinExportLayout; + + fn cuda_ctx_with_varbin_layout(layout: VarBinExportLayout) -> VortexResult { + let session = vortex::array::array_session() + .with_some(CudaSession::try_default()?.with_varbin_export_layout(layout)); + CudaSession::create_execution_ctx(&session) + } + + fn assert_device_resident(canonical: &Canonical) { + let varbinview = canonical.as_varbinview(); + assert!(varbinview.views_handle().is_on_device()); + assert!( + varbinview + .data_buffers() + .iter() + .all(BufferHandle::is_on_device) + ); + } #[rstest] - #[case::non_null( + #[case::binary_non_null( vec![Some(&b"the quick brown fox"[..]), Some(&b"jumps over the lazy dog"[..]), Some(&b"hello world"[..]), Some(&b"vortex fsst test string"[..])], - Nullability::NonNullable, + DType::Binary(Nullability::NonNullable), )] - #[case::partial_nulls( + #[case::utf8_non_null( + vec![Some(&b"the quick brown fox"[..]), + Some(&b"jumps over the lazy dog"[..]), + Some(&b"hello world"[..]), + Some(&b"vortex fsst test string"[..])], + DType::Utf8(Nullability::NonNullable), + )] + #[case::utf8_inline_boundary( + vec![Some(&b""[..]), + Some(&b"123456789012"[..]), + Some(&b"1234567890123"[..]), + Some(&b"this is another outlined value"[..])], + DType::Utf8(Nullability::NonNullable), + )] + #[case::utf8_partial_nulls( vec![Some(&b"alpha"[..]), None, Some(&b"gamma"[..]), None, Some(&b"epsilon"[..])], - Nullability::Nullable, + DType::Utf8(Nullability::Nullable), + )] + #[case::binary_all_empty( + vec![Some(&b""[..]), Some(&b""[..]), Some(&b""[..])], + DType::Binary(Nullability::NonNullable), )] - #[case::all_nulls( + #[case::binary_all_nulls( vec![None, None, None, None, None], - Nullability::Nullable, + DType::Binary(Nullability::Nullable), )] #[crate::test] async fn test_cuda_fsst_decompression_roundtrip( #[case] strings: Vec>, - #[case] nullability: Nullability, + #[case] dtype: DType, ) -> VortexResult<()> { let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); - let varbin = VarBinArray::from_iter(strings, DType::Binary(nullability)).into_array(); + let varbin = VarBinArray::from_iter(strings, dtype.clone()).into_array(); let compressor = fsst_train_compressor(&varbin, cuda_ctx.execution_ctx())?; let fsst_array = fsst_compress(&varbin, &compressor, cuda_ctx.execution_ctx())?.into_array(); @@ -248,12 +476,128 @@ mod tests { let gpu_result = FSSTExecutor .execute(fsst_array.clone(), &mut cuda_ctx) .await - .vortex_expect("GPU decompression failed") - .into_host() - .await? - .into_array(); + .vortex_expect("GPU decompression failed"); + assert_eq!(gpu_result.dtype(), &dtype); + assert_device_resident(&gpu_result); + + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(fsst_array, host_result, &mut ctx); + Ok(()) + } + + /// Verifies that the view kernel applies a sliced validity bitmap's nonzero bit offset. + #[crate::test] + async fn test_cuda_fsst_decompression_sliced_validity() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let values = [ + Some(&b"before"[..]), + None, + Some(&b"gamma"[..]), + None, + Some(&b"after"[..]), + ]; + let varbin = + VarBinArray::from_iter(values, DType::Utf8(Nullability::Nullable)).into_array(); + let compressor = fsst_train_compressor(&varbin, cuda_ctx.execution_ctx())?; + let fsst = fsst_compress(&varbin, &compressor, cuda_ctx.execution_ctx())?.into_array(); + let sliced = fsst.slice(1..4)?; + + let gpu_result = FSSTExecutor.execute(sliced.clone(), &mut cuda_ctx).await?; + assert_device_resident(&gpu_result); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(sliced, host_result, &mut ctx); + Ok(()) + } + + #[crate::test] + async fn test_cuda_fsst_direct_varbin_output() -> VortexResult<()> { + let mut cuda_ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let values: [&[u8]; 3] = [ + b"", + b"short", + b"this value is stored directly in the values buffer", + ]; + let varbin = VarBinArray::from_iter( + values.into_iter().map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let compressor = fsst_train_compressor(&varbin, cuda_ctx.execution_ctx())?; + let fsst = fsst_compress(&varbin, &compressor, cuda_ctx.execution_ctx())?; + + let output = decode_fsst_varbin(fsst, &mut cuda_ctx).await?; + assert_eq!(output.dtype, DType::Utf8(Nullability::NonNullable)); + assert_eq!(output.len, values.len()); + assert!(output.offsets.is_on_device()); + assert!(output.values.is_on_device()); + + let offsets = Buffer::::from_byte_buffer(output.offsets.try_to_host()?.await?); + assert_eq!( + offsets.as_slice(), + &[0, 0, 5, i32::try_from(5 + values[2].len())?,] + ); + assert_eq!( + output.values.try_to_host()?.await?.as_ref(), + values.concat() + ); + Ok(()) + } + + #[rstest] + #[case::binary( + DType::Binary(Nullability::NonNullable), + VarBinExportLayout::VarBin, + DataType::Binary, + 3 + )] + #[case::utf8( + DType::Utf8(Nullability::NonNullable), + VarBinExportLayout::VarBin, + DataType::Utf8, + 3 + )] + #[case::binary_view( + DType::Binary(Nullability::NonNullable), + VarBinExportLayout::VarBinView, + DataType::BinaryView, + 4 + )] + #[case::utf8_view( + DType::Utf8(Nullability::NonNullable), + VarBinExportLayout::VarBinView, + DataType::Utf8View, + 4 + )] + #[crate::test] + async fn test_cuda_fsst_arrow_export_uses_dtype_layout( + #[case] dtype: DType, + #[case] layout: VarBinExportLayout, + #[case] expected_data_type: DataType, + #[case] expected_n_buffers: i64, + ) -> VortexResult<()> { + let mut cuda_ctx = cuda_ctx_with_varbin_layout(layout)?; + let values = [ + Some(&b"short"[..]), + Some(&b"this value is stored out of line"[..]), + ]; + let varbin = VarBinArray::from_iter(values, dtype).into_array(); + let compressor = fsst_train_compressor(&varbin, cuda_ctx.execution_ctx())?; + let fsst_array = + fsst_compress(&varbin, &compressor, cuda_ctx.execution_ctx())?.into_array(); - assert_arrays_eq!(fsst_array, gpu_result, &mut ctx); + let mut exported = fsst_array + .export_device_array_with_schema(&mut cuda_ctx) + .await?; + assert_eq!( + Field::try_from(&exported.schema)?, + Field::new("", expected_data_type, false) + ); + assert_eq!(exported.array.array.n_buffers, expected_n_buffers); + + release_device_array(&mut exported.array); + release_schema(&mut exported.schema); Ok(()) } @@ -271,12 +615,11 @@ mod tests { let gpu_result = FSSTExecutor .execute(fsst_array.clone(), &mut cuda_ctx) .await - .vortex_expect("GPU decompression failed") - .into_host() - .await? - .into_array(); + .vortex_expect("GPU decompression failed"); + assert_device_resident(&gpu_result); - assert_arrays_eq!(fsst_array, gpu_result, &mut ctx); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(fsst_array, host_result, &mut ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/mod.rs b/vortex-cuda/src/kernel/encodings/mod.rs index b97cba13412..6ad37e7ead1 100644 --- a/vortex-cuda/src/kernel/encodings/mod.rs +++ b/vortex-cuda/src/kernel/encodings/mod.rs @@ -21,6 +21,8 @@ pub(crate) use date_time_parts::DateTimePartsExecutor; pub(crate) use decimal_byte_parts::DecimalBytePartsExecutor; pub(crate) use for_::FoRExecutor; pub(crate) use fsst::FSSTExecutor; +pub(crate) use fsst::FSSTVarBin; +pub(crate) use fsst::decode_fsst_varbin; pub(crate) use runend::RunEndExecutor; pub(crate) use sequence::SequenceExecutor; pub(crate) use zigzag::ZigZagExecutor; diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 07cfe40b06f..5817e15436b 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -62,6 +62,7 @@ pub use pooled_read_at::PooledFileReadAt; pub use pooled_read_at::PooledObjectStoreReadAt; pub use session::CudaSession; pub use session::CudaSessionExt; +pub use session::VarBinExportLayout; pub use stream::VortexCudaStream; pub use stream_pool::VortexCudaStreamPool; use vortex::array::ArrayVTable; diff --git a/vortex-cuda/src/pinned.rs b/vortex-cuda/src/pinned.rs index cd97db63164..52b1aa46e2a 100644 --- a/vortex-cuda/src/pinned.rs +++ b/vortex-cuda/src/pinned.rs @@ -113,6 +113,15 @@ pub struct PinnedByteBufferPool { puts: AtomicU64, } +impl std::fmt::Debug for PinnedByteBufferPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PinnedByteBufferPool") + .field("max_keep_per_size", &self.max_keep_per_size) + .field("stats", &self.stats()) + .finish_non_exhaustive() + } +} + struct InflightPinnedBuffer { event: Arc, buffer: PinnedByteBuffer, diff --git a/vortex-cuda/src/session.rs b/vortex-cuda/src/session.rs index 5bc4e55de69..a5410db0c99 100644 --- a/vortex-cuda/src/session.rs +++ b/vortex-cuda/src/session.rs @@ -23,12 +23,23 @@ use crate::executor::CudaExecute; pub use crate::executor::CudaExecutionCtx; use crate::initialize_cuda; use crate::kernel::KernelLoader; +use crate::pinned::PinnedByteBufferPool; use crate::stream::VortexCudaStream; use crate::stream_pool::VortexCudaStreamPool; /// Default maximum number of streams in the pool. const DEFAULT_STREAM_POOL_CAPACITY: usize = 4; +/// Arrow Device layout used when exporting variable-length UTF-8 and binary arrays. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum VarBinExportLayout { + /// Offset-based Arrow `Utf8`/`Binary` with one contiguous values buffer. + #[default] + VarBin, + /// Arrow `Utf8View`/`BinaryView` with 16-byte views and variadic data buffers. + VarBinView, +} + /// CUDA session for GPU accelerated execution. /// /// Maintains a registry of CUDA kernel implementations for array encodings. @@ -38,8 +49,10 @@ pub struct CudaSession { context: Arc, kernels: Arc>, export_device_array: Arc, + varbin_export_layout: VarBinExportLayout, kernel_loader: Arc, stream_pool: Arc, + pinned_buffer_pool: Arc, } impl CudaSession { @@ -57,15 +70,29 @@ impl CudaSession { Arc::clone(&context), stream_pool_capacity, )); + let pinned_buffer_pool = Arc::new(PinnedByteBufferPool::new(Arc::clone(&context))); Self { context, kernels: Arc::new(DashMap::default()), kernel_loader: Arc::new(KernelLoader::new()), export_device_array: Arc::new(CanonicalDeviceArrayExport), + varbin_export_layout: VarBinExportLayout::default(), stream_pool, + pinned_buffer_pool, } } + /// Selects the Arrow Device layout for variable-length UTF-8 and binary exports. + pub fn with_varbin_export_layout(mut self, layout: VarBinExportLayout) -> Self { + self.varbin_export_layout = layout; + self + } + + /// Returns the Arrow Device layout used for variable-length UTF-8 and binary exports. + pub fn varbin_export_layout(&self) -> VarBinExportLayout { + self.varbin_export_layout + } + /// Creates a default CUDA session using device 0, with all GPU array kernels preloaded. /// /// Unlike [`Default::default`], this returns an error instead of panicking when CUDA cannot be @@ -105,6 +132,11 @@ impl CudaSession { self.stream_pool.stream() } + /// Returns the session-scoped pool used for staging file reads in pinned host memory. + pub fn pinned_buffer_pool(&self) -> &Arc { + &self.pinned_buffer_pool + } + /// Registers CUDA support for an array encoding. /// /// # Arguments diff --git a/vortex-cuda/src/stream.rs b/vortex-cuda/src/stream.rs index 6342c39ab28..4d31a7937f9 100644 --- a/vortex-cuda/src/stream.rs +++ b/vortex-cuda/src/stream.rs @@ -23,6 +23,11 @@ use vortex::error::vortex_ensure; use vortex::error::vortex_err; use crate::CudaDeviceBuffer; +use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; + +// cuDF imports Arrow validity masks into padded buffers and kernels may read through that +// padded extent. Keep copied device buffers padded and zero-tailed so Arrow validity exports +// can safely reuse matching bitmaps without repacking. #[derive(Clone)] pub struct VortexCudaStream(pub(crate) Arc); @@ -68,8 +73,8 @@ impl VortexCudaStream { /// (guaranteed by the returned future capturing it). /// /// The returned [`BufferHandle`] keeps the source byte length, while its - /// CUDA allocation may include zeroed tail padding. This is needed for - /// Arrow validity buffers passed to cuDF, which reads masks as 32-bit words. + /// CUDA allocation may include zeroed tail padding for consumers such as cuDF + /// that read validity masks through padded extents. pub(crate) fn copy_to_device( &self, data: D, @@ -90,7 +95,8 @@ impl VortexCudaStream { zero_padding(self, &mut cuda_slice, host_slice.len())?; - let cuda_buf = CudaDeviceBuffer::new(cuda_slice); + // `zero_padding` zeroed all allocation bytes after `byte_count`. + let cuda_buf = CudaDeviceBuffer::new_with_zeroed_tail(cuda_slice, byte_count)?; let buffer = BufferHandle::new_device(Arc::new(cuda_buf)).slice(0..byte_count); let stream = Arc::clone(&self.0); @@ -131,29 +137,30 @@ impl VortexCudaStream { zero_padding(self, &mut cuda_slice, data.len())?; - let cuda_buf = CudaDeviceBuffer::new(cuda_slice); + // `zero_padding` zeroed all allocation bytes after `byte_count`. + let cuda_buf = CudaDeviceBuffer::new_with_zeroed_tail(cuda_slice, byte_count)?; Ok(BufferHandle::new_device(Arc::new(cuda_buf)).slice(0..byte_count)) } } /// Returns the typed CUDA allocation length for `byte_count`. /// -/// The backing allocation is padded for cuDF's 32-bit validity mask reads. -/// The returned length is in `T` elements. +/// The backing allocation is padded for consumers such as cuDF that read validity masks +/// through padded extents. The returned length is in `T` elements. fn padded_device_allocation_len(byte_count: usize) -> VortexResult { let element_size = size_of::(); vortex_ensure!( element_size != 0, "cannot copy zero-sized values to CUDA device" ); - let min_allocation_bytes = byte_count.next_multiple_of(size_of::()); + let min_allocation_bytes = byte_count.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING); Ok(min_allocation_bytes.div_ceil(element_size)) } /// Zeroes the allocation tail after the copied values. /// /// Returned handles are sliced to the copied byte count; the trailing padding -/// exists so a final 32-bit mask read stays within the backing allocation. +/// exists so padded mask reads stay within the backing allocation. fn zero_padding( stream: &VortexCudaStream, cuda_slice: &mut CudaSlice, @@ -250,19 +257,38 @@ fn register_stream_callback(stream: &CudaStream) -> VortexResult VortexResult<()> { assert_eq!(padded_device_allocation_len::(0)?, 0); - assert_eq!(padded_device_allocation_len::(1)?, 4); - assert_eq!(padded_device_allocation_len::(4)?, 4); - assert_eq!(padded_device_allocation_len::(5)?, 8); - assert_eq!(padded_device_allocation_len::(1)?, 1); - assert_eq!(padded_device_allocation_len::(5)?, 2); + assert_eq!( + padded_device_allocation_len::(1)?, + CUDF_VALIDITY_BUFFER_PADDING + ); + assert_eq!( + padded_device_allocation_len::(4)?, + CUDF_VALIDITY_BUFFER_PADDING + ); + assert_eq!( + padded_device_allocation_len::(5)?, + CUDF_VALIDITY_BUFFER_PADDING + ); + assert_eq!( + padded_device_allocation_len::(1)?, + CUDF_VALIDITY_BUFFER_PADDING / size_of::() + ); + assert_eq!( + padded_device_allocation_len::(5)?, + CUDF_VALIDITY_BUFFER_PADDING / size_of::() + ); Ok(()) } @@ -275,6 +301,12 @@ mod tests { let host = handle.try_to_host()?.await?; assert_eq!(host.as_slice(), &[0xab]); + let backing = cuda_backing_allocation(&handle)?; + assert_eq!(backing.len(), CUDF_VALIDITY_BUFFER_PADDING); + let backing_host = backing.try_to_host()?.await?; + assert_eq!(backing_host[0], 0xab); + assert!(backing_host[1..].iter().all(|byte| *byte == 0)); + Ok(()) } @@ -287,6 +319,12 @@ mod tests { let host = handle.try_to_host()?.await?; assert_eq!(host.as_slice(), &[1, 2, 3, 4, 5]); + let backing = cuda_backing_allocation(&handle)?; + assert_eq!(backing.len(), CUDF_VALIDITY_BUFFER_PADDING); + let backing_host = backing.try_to_host()?.await?; + assert_eq!(&backing_host[..5], &[1, 2, 3, 4, 5]); + assert!(backing_host[5..].iter().all(|byte| *byte == 0)); + Ok(()) } } diff --git a/vortex-cxx/.clang-tidy b/vortex-cxx/.clang-tidy deleted file mode 100644 index 80c94310cbc..00000000000 --- a/vortex-cxx/.clang-tidy +++ /dev/null @@ -1,93 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright the Vortex contributors - -Checks: | - -*, - clang-diagnostic-*, - bugprone-*, - performance-*, - google-explicit-constructor, - google-build-using-namespace, - google-runtime-int, - misc-definitions-in-headers, - modernize-use-nullptr, - modernize-use-override, - -bugprone-macro-parentheses, - readability-braces-around-statements, - -bugprone-branch-clone, - readability-identifier-naming, - hicpp-exception-baseclass, - misc-throw-by-value-catch-by-reference, - -bugprone-signed-char-misuse, - -bugprone-misplaced-widening-cast, - -bugprone-sizeof-expression, - -bugprone-easily-swappable-parameters, - google-global-names-in-headers, - llvm-header-guard, - misc-definitions-in-headers, - modernize-use-emplace, - modernize-use-bool-literals, - -performance-inefficient-string-concatenation, - -performance-no-int-to-ptr, - readability-container-size-empty, - cppcoreguidelines-pro-type-cstyle-cast, - -llvm-header-guard, - -performance-enum-size, - cppcoreguidelines-pro-type-const-cast, - cppcoreguidelines-interfaces-global-init, - cppcoreguidelines-slicing, - cppcoreguidelines-rvalue-reference-param-not-moved, - cppcoreguidelines-virtual-class-destructor, - -readability-identifier-naming, - -bugprone-exception-escape, - -bugprone-unused-local-non-trivial-variable, - -bugprone-empty-catch, -WarningsAsErrors: '*' -HeaderFilterRegex: '(cpp|examples)/.*\.(cpp|hpp)$' -FormatStyle: none -CheckOptions: - - key: readability-identifier-naming.ClassCase - value: CamelCase - - key: readability-identifier-naming.EnumCase - value: CamelCase - - key: readability-identifier-naming.TypedefCase - value: lower_case - - key: readability-identifier-naming.TypedefSuffix - value: _t - - key: readability-identifier-naming.FunctionCase - value: CamelCase - - key: readability-identifier-naming.MemberCase - value: lower_case - - key: readability-identifier-naming.ParameterCase - value: lower_case - - key: readability-identifier-naming.ConstantCase - value: aNy_CasE - - key: readability-identifier-naming.ConstantParameterCase - value: lower_case - - key: readability-identifier-naming.NamespaceCase - value: lower_case - - key: readability-identifier-naming.MacroDefinitionCase - value: UPPER_CASE - - key: readability-identifier-naming.StaticConstantCase - value: UPPER_CASE - - key: readability-identifier-naming.ConstantMemberCase - value: aNy_CasE - - key: readability-identifier-naming.StaticVariableCase - value: UPPER_CASE - - key: readability-identifier-naming.ClassConstantCase - value: UPPER_CASE - - key: readability-identifier-naming.EnumConstantCase - value: UPPER_CASE - - key: readability-identifier-naming.ConstexprVariableCase - value: aNy_CasE - - key: readability-identifier-naming.StaticConstantCase - value: UPPER_CASE - - key: readability-identifier-naming.TemplateTemplateParameterCase - value: UPPER_CASE - - key: readability-identifier-naming.TypeTemplateParameterCase - value: UPPER_CASE - - key: readability-identifier-naming.VariableCase - value: lower_case - - key: cppcoreguidelines-rvalue-reference-param-not-moved.IgnoreUnnamedParams - value: true - diff --git a/vortex-cxx/CMakeLists.txt b/vortex-cxx/CMakeLists.txt deleted file mode 100644 index ab5474d0617..00000000000 --- a/vortex-cxx/CMakeLists.txt +++ /dev/null @@ -1,120 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright the Vortex contributors - -cmake_minimum_required(VERSION 3.22) - -# FetchContent from URL timestamp handling -cmake_policy(SET CMP0135 NEW) - -project(vortex) - -set(CMAKE_CXX_STANDARD 20) -set(CMAKE_CXX_STANDARD_REQUIRED ON) -set(CMAKE_EXPORT_COMPILE_COMMANDS ON) - -find_program(SCCACHE_PROGRAM sccache) -if (SCCACHE_PROGRAM) - set(CMAKE_C_COMPILER_LAUNCHER "${SCCACHE_PROGRAM}") - set(CMAKE_CXX_COMPILER_LAUNCHER "${SCCACHE_PROGRAM}") - message(STATUS "Sccache found: ${SCCACHE_PROGRAM}") -else () - message(STATUS "Sccache not found") -endif () - -if (NOT CMAKE_BUILD_TYPE) - set(CMAKE_BUILD_TYPE Debug) -endif() - -set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -Wpedantic -Wno-dollar-in-identifier-extension") - -option(VORTEX_ENABLE_TESTING "Enable building test binary for vortex-cxx" OFF) -option(VORTEX_ENABLE_ASAN "Enable address sanitizer" OFF) - -include(FetchContent) -FetchContent_Declare( - Corrosion - GIT_REPOSITORY https://github.com/corrosion-rs/corrosion.git - GIT_TAG v0.5 -) -FetchContent_MakeAvailable(Corrosion) - -set(RUST_SOURCE_FILE lib.rs) - -# Import Rust crate using Corrosion -corrosion_import_crate( - MANIFEST_PATH ${CMAKE_CURRENT_SOURCE_DIR}/Cargo.toml - FEATURES ${CORROSION_FEATURES} - CRATES vortex-cxx -) - -# Enable CXX bridge for the Rust crate -corrosion_add_cxxbridge(vortex_cxx_bridge CRATE vortex_cxx FILES ${RUST_SOURCE_FILE}) - -FetchContent_Declare( - nanoarrow - GIT_REPOSITORY https://github.com/apache/arrow-nanoarrow.git - GIT_TAG a579fbf5d192e85b6249935e117de7d02a6dc4e9 # v0.8.0 -) -FetchContent_MakeAvailable(nanoarrow) - -file(GLOB_RECURSE CPP_SOURCE_FILE CONFIGURE_DEPENDS - "${CMAKE_CURRENT_SOURCE_DIR}/cpp/src/*.cpp" -) - -# Public headers -set(CPP_INCLUDE_DIRS - ${CMAKE_CURRENT_SOURCE_DIR}/cpp/include -) - -# Private headers -set(CPP_PRIVATE_INCLUDE_DIRS - ${CMAKE_CURRENT_SOURCE_DIR}/cpp/src -) - -# Create the main library combining C++ and Rust code -add_library(vortex STATIC ${CPP_SOURCE_FILE}) -target_include_directories(vortex PUBLIC ${CPP_INCLUDE_DIRS} - ${CMAKE_CURRENT_BINARY_DIR}/corrosion_generated/cxxbridge/vortex_cxx_bridge/include) -target_include_directories(vortex PRIVATE - ${CPP_PRIVATE_INCLUDE_DIRS} -) -target_link_libraries(vortex - PUBLIC nanoarrow_static - PRIVATE vortex_cxx_bridge -) - -if (VORTEX_ENABLE_ASAN) - target_compile_options(vortex PRIVATE -fsanitize=leak,address,undefined -fno-omit-frame-pointer -fno-common -O1) - target_link_options(vortex PRIVATE -fsanitize=leak,address,undefined) -endif() - -# Tests -if (VORTEX_ENABLE_TESTING) - FetchContent_Declare( - googletest - GIT_REPOSITORY https://github.com/google/googletest.git - GIT_TAG v1.17.0 - GIT_SHALLOW TRUE - ) - FetchContent_MakeAvailable(googletest) - - enable_testing() - add_executable(vortex_cxx_test cpp/tests/basic_test.cpp cpp/tests/test_data_generator.cpp) - target_include_directories(vortex_cxx_test PUBLIC ${CPP_INCLUDE_DIRS}) - target_include_directories(vortex_cxx_test PRIVATE cpp/tests) - target_link_libraries(vortex_cxx_test PRIVATE gtest_main vortex nanoarrow_static) - target_include_directories(vortex_cxx_test PRIVATE - ${CMAKE_CURRENT_BINARY_DIR}/corrosion_generated/cxxbridge/vortex_cxx_bridge/include - ) - # Platform-specific configuration - if(APPLE) - set(APPLE_LINK_FLAGS "-framework CoreFoundation -framework Security") - endif() - target_link_libraries(vortex_cxx_test PRIVATE vortex_cxx_bridge ${APPLE_LINK_FLAGS}) - if (VORTEX_ENABLE_ASAN) - target_compile_options(vortex_cxx_test PRIVATE -fsanitize=leak,address,undefined -fno-omit-frame-pointer -fno-common -O1) - target_link_options(vortex_cxx_test PRIVATE -fsanitize=leak,address,undefined) - endif() - include(GoogleTest) - gtest_discover_tests(vortex_cxx_test) -endif() diff --git a/vortex-cxx/Cargo.toml b/vortex-cxx/Cargo.toml deleted file mode 100644 index d75f7b74cee..00000000000 --- a/vortex-cxx/Cargo.toml +++ /dev/null @@ -1,32 +0,0 @@ -[package] -name = "vortex-cxx" -authors = { workspace = true } -categories = { workspace = true } -description = "C++ bindings for Vortex generated from Rust using cxx" -edition = { workspace = true } -homepage = { workspace = true } -include = { workspace = true } -keywords = { workspace = true } -license = { workspace = true } -publish = false -readme = { workspace = true } -repository = { workspace = true } -rust-version = { workspace = true } -version = { workspace = true } - -[lints] -workspace = true - -[lib] -crate-type = ["staticlib"] - -[dependencies] -anyhow = { workspace = true } -arrow-array = { workspace = true, features = ["ffi"] } -arrow-schema = { workspace = true } -async-fs = { workspace = true } -cxx = "1.0" -futures = { workspace = true } -paste = { workspace = true } -take_mut = { workspace = true } -vortex = { workspace = true } diff --git a/vortex-cxx/README.md b/vortex-cxx/README.md deleted file mode 100644 index aa761f98b1a..00000000000 --- a/vortex-cxx/README.md +++ /dev/null @@ -1,35 +0,0 @@ -# Vortex C++ Bindings - -This directory contains C++ bindings for Vortex using the [cxx](https://cxx.rs/) crate. The bindings provide a C++ interface to Vortex file operations, including roundtripping with Arrow Array stream with advanced pushdown support. - -## Building - -### Requirements - -- CMake 3.22 or higher -- C++20 compatible compiler -- Rust toolchain (for building the Rust components) -- (optional) Ninja (`ninja-build`) - -### Build Steps - -```bash -mkdir build -cmake -Bbuild -GNinja -cmake --build build -j -``` - -### Running Tests - -```bash -# Enable tests in CMake -cmake -Bbuild -DVORTEX_ENABLE_TESTING=ON -GNinja -cmake --build build -j -./vortex_cxx_test -``` - -## C++ Coding Convention - -We use `.clang-tidy` and `.clang-format` to setup the coding convention. Both are borrowed from DuckDB. - -`cppcoreguidelines-avoid-non-const-global-variables` is removed from `.clang-tidy` because GTest violates it. diff --git a/vortex-cxx/cpp/include/vortex.hpp b/vortex-cxx/cpp/include/vortex.hpp deleted file mode 100644 index 1f2dd1625c1..00000000000 --- a/vortex-cxx/cpp/include/vortex.hpp +++ /dev/null @@ -1,11 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#pragma once - -#include "vortex/file.hpp" -#include "vortex/scan.hpp" -#include "vortex/write_options.hpp" -#include "vortex/exception.hpp" -#include "vortex/expr.hpp" -#include "vortex/scalar.hpp" \ No newline at end of file diff --git a/vortex-cxx/cpp/include/vortex/dtype.hpp b/vortex-cxx/cpp/include/vortex/dtype.hpp deleted file mode 100644 index 9bf4a16d9bc..00000000000 --- a/vortex-cxx/cpp/include/vortex/dtype.hpp +++ /dev/null @@ -1,71 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#pragma once - -#include -#include -#include "vortex_cxx_bridge/lib.h" - -namespace vortex { - -enum class PType : uint8_t { - U8 = 0, - U16, - U32, - U64, - I8, - I16, - I32, - I64, - F16, - F32, - F64, -}; - -namespace dtype { - class DType { - public: - DType() = delete; - explicit DType(rust::Box impl) : impl_(std::move(impl)) { - } - DType(DType &&other) noexcept = default; - DType &operator=(DType &&other) = default; - ~DType() = default; - - DType(const DType &) = delete; - DType &operator=(const DType &) = delete; - - std::string ToString() const; - - const rust::Box &GetImpl() { - return impl_; - } - - private: - rust::Box impl_; - }; - - // Factory functions - DType null(); - DType bool_(bool nullable = false); - DType primitive(PType ptype, bool nullable = false); - DType int8(bool nullable = false); - DType int16(bool nullable = false); - DType int32(bool nullable = false); - DType int64(bool nullable = false); - DType uint8(bool nullable = false); - DType uint16(bool nullable = false); - DType uint32(bool nullable = false); - DType uint64(bool nullable = false); - DType float16(bool nullable = false); - DType float32(bool nullable = false); - DType float64(bool nullable = false); - DType decimal(uint8_t precision = 10, int8_t scale = 0, bool nullable = false); - DType utf8(bool nullable = false); - DType binary(bool nullable = false); - /// TODO: Other DTypes are only supported by creating from Arrow for now. - DType from_arrow(struct ArrowSchema &schema, bool non_nullable = false); -} // namespace dtype - -} // namespace vortex \ No newline at end of file diff --git a/vortex-cxx/cpp/include/vortex/exception.hpp b/vortex-cxx/cpp/include/vortex/exception.hpp deleted file mode 100644 index a7fe44fb91c..00000000000 --- a/vortex-cxx/cpp/include/vortex/exception.hpp +++ /dev/null @@ -1,18 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#pragma once - -#include -#include - -namespace vortex { - -/// TODO(xinyu): better error handling -class VortexException : public std::runtime_error { -public: - explicit VortexException(const std::string &message) : std::runtime_error(message) { - } -}; - -} // namespace vortex \ No newline at end of file diff --git a/vortex-cxx/cpp/include/vortex/expr.hpp b/vortex-cxx/cpp/include/vortex/expr.hpp deleted file mode 100644 index 3060a8e3536..00000000000 --- a/vortex-cxx/cpp/include/vortex/expr.hpp +++ /dev/null @@ -1,51 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#pragma once - -#include "vortex/scalar.hpp" -#include "vortex_cxx_bridge/lib.h" -#include - -namespace vortex::expr { -class Expr { -public: - Expr() = delete; - explicit Expr(rust::Box impl) : impl_(std::move(impl)) { - } - Expr(Expr &&other) noexcept = default; - Expr &operator=(Expr &&other) noexcept = default; - ~Expr() = default; - - Expr(const Expr &) = delete; - Expr &operator=(const Expr &) = delete; - - rust::Box IntoImpl() && { - return std::move(impl_); - } - - const ffi::Expr &Impl() const & { - return *impl_; - } - -private: - rust::Box impl_; -}; - -Expr literal(scalar::Scalar scalar); -Expr root(); -Expr column(std::string_view name); -Expr get_item(std::string_view field, Expr expr); -Expr not_(Expr expr); -Expr is_null(Expr expr); -Expr eq(Expr lhs, Expr rhs); -Expr not_eq_(Expr lhs, Expr rhs); -Expr gt(Expr lhs, Expr rhs); -Expr gt_eq(Expr lhs, Expr rhs); -Expr lt(Expr lhs, Expr rhs); -Expr lt_eq(Expr lhs, Expr rhs); -Expr and_(Expr lhs, Expr rhs); -Expr or_(Expr lhs, Expr rhs); -Expr checked_add(Expr lhs, Expr rhs); -Expr select(const std::vector &fields, Expr child); -} // namespace vortex::expr \ No newline at end of file diff --git a/vortex-cxx/cpp/include/vortex/file.hpp b/vortex-cxx/cpp/include/vortex/file.hpp deleted file mode 100644 index c7b93e326d8..00000000000 --- a/vortex-cxx/cpp/include/vortex/file.hpp +++ /dev/null @@ -1,38 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#pragma once - -#include -#include "vortex_cxx_bridge/lib.h" - -namespace vortex { -class ScanBuilder; - -class VortexFile { -public: - static VortexFile Open(const std::string &path); - static VortexFile Open(const uint8_t *data, size_t length); - - VortexFile(VortexFile &&other) noexcept = default; - VortexFile &operator=(VortexFile &&other) noexcept = default; - ~VortexFile() = default; - - VortexFile(const VortexFile &) = delete; - VortexFile &operator=(const VortexFile &) = delete; - - /// Get the number of rows in the file. - uint64_t RowCount() const; - - /// Create a scan builder for the file. - /// The scan builder can be used to scan the file. - ScanBuilder CreateScanBuilder() const; - -private: - explicit VortexFile(rust::Box impl) : impl_(std::move(impl)) { - } - - rust::Box impl_; -}; - -} // namespace vortex \ No newline at end of file diff --git a/vortex-cxx/cpp/include/vortex/scalar.hpp b/vortex-cxx/cpp/include/vortex/scalar.hpp deleted file mode 100644 index 77706aef98a..00000000000 --- a/vortex-cxx/cpp/include/vortex/scalar.hpp +++ /dev/null @@ -1,47 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#pragma once - -#include -#include "dtype.hpp" -#include "vortex_cxx_bridge/lib.h" - -namespace vortex::scalar { -class Scalar { -public: - Scalar() = delete; - explicit Scalar(rust::Box impl) : impl_(std::move(impl)) { - } - Scalar(Scalar &&other) noexcept = default; - Scalar &operator=(Scalar &&other) noexcept = default; - ~Scalar() = default; - - Scalar(const Scalar &) = delete; - Scalar &operator=(const Scalar &) = delete; - - rust::Box IntoImpl() && { - return std::move(impl_); - } - -private: - rust::Box impl_; -}; - -// Factory functions for creating scalar values -Scalar bool_(bool value); -Scalar int8(int8_t value); -Scalar int16(int16_t value); -Scalar int32(int32_t value); -Scalar int64(int64_t value); -Scalar uint8(uint8_t value); -Scalar uint16(uint16_t value); -Scalar uint32(uint32_t value); -Scalar uint64(uint64_t value); -Scalar float32(float value); -Scalar float64(double value); -Scalar string(std::string_view value); -Scalar binary(const uint8_t *data, size_t length); -/// TODO: Other Scalars are only supported by casting for now. -Scalar cast(Scalar scalar, dtype::DType dtype); -} // namespace vortex::scalar diff --git a/vortex-cxx/cpp/include/vortex/scan.hpp b/vortex-cxx/cpp/include/vortex/scan.hpp deleted file mode 100644 index 7debb9ef6aa..00000000000 --- a/vortex-cxx/cpp/include/vortex/scan.hpp +++ /dev/null @@ -1,108 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#pragma once - -#include "vortex/expr.hpp" -#include -#include "vortex_cxx_bridge/lib.h" - -#include -#include - -namespace vortex { -/// The StreamDriver internally holds a `RecordBatchIteratorAdapter` from the Rust side, which is thread-safe -/// and cloneable. The `RecordBatchIteratorAdapter` internally holds a `WorkStealingArrayIterator`. -class StreamDriver { -public: - StreamDriver(StreamDriver &&other) noexcept = default; - StreamDriver &operator=(StreamDriver &&other) noexcept = default; - ~StreamDriver() = default; - - StreamDriver(const StreamDriver &) = delete; - StreamDriver &operator=(const StreamDriver &) = delete; - - /// Create a stream of record batches. - /// - /// This function is thread-safe and can be called from multiple threads to create one stream per - /// thread to make progress on the same StreamDriver that is built from a ScanBuilder concurrently. - /// - /// Within each thread, the record batches will be emitted in the original order they are within - /// the scan. Between threads, the order is not guaranteed. - /// - /// Example: If the scan contains batches [b0, b1, b2, b3, b4, b5] and two threads call this - /// function respectively to make progress on their own stream, Thread 1 might receive [b0, - /// b2, b4] and Thread 2 might receive [b1, b3, b5]. Each thread maintains order within its - /// subset, but overall ordering between threads is not guaranteed (e.g., Thread 2 could emit b1 - /// before Thread 1 emits b0). - ArrowArrayStream CreateArrayStream() const; - -private: - friend class ScanBuilder; - - explicit StreamDriver(rust::Box impl) : impl_(std::move(impl)) { - } - - rust::Box impl_; -}; - -class ScanBuilder { -public: - ScanBuilder(ScanBuilder &&other) noexcept = default; - ScanBuilder &operator=(ScanBuilder &&other) noexcept { - if (this != &other) { - impl_ = std::move(other.impl_); - } - return *this; - } - ~ScanBuilder() = default; - - ScanBuilder(const ScanBuilder &) = delete; - ScanBuilder &operator=(const ScanBuilder &) = delete; - - /// Only include rows that match the filter expressions. - ScanBuilder &WithFilter(expr::Expr &&expr) &; - ScanBuilder &WithFilter(const expr::Expr &expr) &; - ScanBuilder &&WithFilter(expr::Expr &&expr) &&; - ScanBuilder &&WithFilter(const expr::Expr &expr) &&; - - /// Only include columns that match the projection expressions. - ScanBuilder &WithProjection(expr::Expr &&expr) &; - ScanBuilder &WithProjection(const expr::Expr &expr) &; - ScanBuilder &&WithProjection(expr::Expr &&expr) &&; - ScanBuilder &&WithProjection(const expr::Expr &expr) &&; - - /// Only include rows in the range [row_range_start, row_range_end). - ScanBuilder &WithRowRange(uint64_t row_range_start, uint64_t row_range_end) &; - ScanBuilder &&WithRowRange(uint64_t row_range_start, uint64_t row_range_end) &&; - - /// Only include rows with the given indices. - ScanBuilder &WithIncludeByIndex(const uint64_t *indices, std::size_t size) &; - ScanBuilder &&WithIncludeByIndex(const uint64_t *indices, std::size_t size) &&; - - /// Set the limit on the number of rows to scan out. - ScanBuilder &WithLimit(uint64_t limit) &; - ScanBuilder &&WithLimit(uint64_t limit) &&; - - /// Set the output schema on the scan builder. - /// TODO: currently if pass in this option, the schema needs to be the schema after adding projection. - ScanBuilder &WithOutputSchema(ArrowSchema &output_schema) &; - ScanBuilder &&WithOutputSchema(ArrowSchema &output_schema) &&; - - /// Take ownership and consume the scan builder to a stream of record batches. - ArrowArrayStream IntoStream() &&; - - /// Take ownership and consume the scan builder to a stream driver. - /// Under the hood, this function calls `ScanBuilder::into_record_batch_reader` and holds a - /// `WorkStealingArrayIterator` in StreamDriver. - StreamDriver IntoStreamDriver() &&; - -private: - friend class VortexFile; - - explicit ScanBuilder(rust::Box impl) : impl_(std::move(impl)) { - } - - rust::Box impl_; -}; -} // namespace vortex \ No newline at end of file diff --git a/vortex-cxx/cpp/include/vortex/write_options.hpp b/vortex-cxx/cpp/include/vortex/write_options.hpp deleted file mode 100644 index 08b903bf2fb..00000000000 --- a/vortex-cxx/cpp/include/vortex/write_options.hpp +++ /dev/null @@ -1,29 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#pragma once - -#include -#include "vortex_cxx_bridge/lib.h" - -namespace vortex { - -class VortexWriteOptions { -public: - VortexWriteOptions() : impl_(ffi::write_options_new()) { - } - VortexWriteOptions(VortexWriteOptions &&other) noexcept = default; - VortexWriteOptions &operator=(VortexWriteOptions &&other) noexcept = default; - ~VortexWriteOptions() = default; - - VortexWriteOptions(const VortexWriteOptions &) = delete; - VortexWriteOptions &operator=(const VortexWriteOptions &) = delete; - - /// Write an ArrowArrayStream to a Vortex file - void WriteArrayStream(ArrowArrayStream &stream, const std::string &path); - -private: - rust::Box impl_; -}; - -} // namespace vortex \ No newline at end of file diff --git a/vortex-cxx/cpp/src/dtype.cpp b/vortex-cxx/cpp/src/dtype.cpp deleted file mode 100644 index 5b85f36835e..00000000000 --- a/vortex-cxx/cpp/src/dtype.cpp +++ /dev/null @@ -1,90 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#include "vortex/dtype.hpp" -#include "vortex/exception.hpp" - -#include "rust/cxx.h" - -namespace vortex::dtype { -DType null() { - return DType(ffi::dtype_null()); -} - -DType bool_(bool nullable) { - return DType(ffi::dtype_bool(nullable)); -} - -DType primitive(PType ptype, bool nullable) { - return DType(ffi::dtype_primitive(static_cast(ptype), nullable)); -} - -DType int8(bool nullable) { - return primitive(PType::I8, nullable); -} - -DType int16(bool nullable) { - return primitive(PType::I16, nullable); -} - -DType int32(bool nullable) { - return primitive(PType::I32, nullable); -} - -DType int64(bool nullable) { - return primitive(PType::I64, nullable); -} - -DType uint8(bool nullable) { - return primitive(PType::U8, nullable); -} - -DType uint16(bool nullable) { - return primitive(PType::U16, nullable); -} - -DType uint32(bool nullable) { - return primitive(PType::U32, nullable); -} - -DType uint64(bool nullable) { - return primitive(PType::U64, nullable); -} - -DType float16(bool nullable) { - return primitive(PType::F16, nullable); -} - -DType float32(bool nullable) { - return primitive(PType::F32, nullable); -} - -DType float64(bool nullable) { - return primitive(PType::F64, nullable); -} - -DType decimal(uint8_t precision, int8_t scale, bool nullable) { - return DType(ffi::dtype_decimal(precision, scale, nullable)); -} - -DType utf8(bool nullable) { - return DType(ffi::dtype_utf8(nullable)); -} - -DType binary(bool nullable) { - return DType(ffi::dtype_binary(nullable)); -} - -DType from_arrow(struct ArrowSchema &schema, bool non_nullable) { - try { - return DType(ffi::from_arrow(reinterpret_cast(&schema), non_nullable)); - } catch (const rust::cxxbridge1::Error &e) { - throw VortexException(e.what()); - } -} -// Methods -std::string DType::ToString() const { - auto rust_str = impl_->to_string(); - return std::string(rust_str.data(), rust_str.length()); -} -} // namespace vortex::dtype diff --git a/vortex-cxx/cpp/src/expr.cpp b/vortex-cxx/cpp/src/expr.cpp deleted file mode 100644 index 0e2d67395c1..00000000000 --- a/vortex-cxx/cpp/src/expr.cpp +++ /dev/null @@ -1,60 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#include "vortex/expr.hpp" -#include "vortex_cxx_bridge/lib.h" - -namespace vortex::expr { - -Expr literal(scalar::Scalar scalar) { - return Expr(ffi::literal(std::move(scalar).IntoImpl())); -} - -Expr root() { - return Expr(ffi::root()); -} - -Expr column(std::string_view name) { - return Expr(ffi::column(rust::String(name.data(), name.length()))); -} - -Expr get_item(std::string_view field, Expr child) { - return Expr(ffi::get_item(rust::String(field.data(), field.length()), std::move(child).IntoImpl())); -} - -Expr not_(Expr child) { - return Expr(ffi::not_(std::move(child).IntoImpl())); -} - -Expr is_null(Expr child) { - return Expr(ffi::is_null(std::move(child).IntoImpl())); -} - -// Macro to define binary operator functions -#define DEFINE_BINARY_OP(name) \ - Expr name(Expr lhs, Expr rhs) { \ - return Expr(ffi::name(std::move(lhs).IntoImpl(), std::move(rhs).IntoImpl())); \ - } - -DEFINE_BINARY_OP(eq) -DEFINE_BINARY_OP(not_eq_) -DEFINE_BINARY_OP(gt) -DEFINE_BINARY_OP(gt_eq) -DEFINE_BINARY_OP(lt) -DEFINE_BINARY_OP(lt_eq) -DEFINE_BINARY_OP(and_) -DEFINE_BINARY_OP(or_) -DEFINE_BINARY_OP(checked_add) - -#undef DEFINE_BINARY_OP - -Expr select(const std::vector &fields, Expr child) { - ::rust::Vec<::rust::String> rs_fields; - rs_fields.reserve(fields.size()); - for (std::string_view f : fields) { - rs_fields.emplace_back(f.data(), f.length()); - } - return Expr(ffi::select(rs_fields, std::move(child).IntoImpl())); -} - -} // namespace vortex::expr \ No newline at end of file diff --git a/vortex-cxx/cpp/src/file.cpp b/vortex-cxx/cpp/src/file.cpp deleted file mode 100644 index 078c305639d..00000000000 --- a/vortex-cxx/cpp/src/file.cpp +++ /dev/null @@ -1,36 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#include "vortex/file.hpp" -#include "vortex/scan.hpp" -#include "vortex/exception.hpp" -#include "rust/cxx.h" - -namespace vortex { - -VortexFile VortexFile::Open(const uint8_t *data, size_t length) { - try { - rust::Slice slice(data, length); - return VortexFile(ffi::open_file_from_buffer(slice)); - } catch (const rust::cxxbridge1::Error &e) { - throw VortexException(e.what()); - } -} - -VortexFile VortexFile::Open(const std::string &path) { - try { - return VortexFile(ffi::open_file(path)); - } catch (const rust::cxxbridge1::Error &e) { - throw VortexException(e.what()); - } -} - -uint64_t VortexFile::RowCount() const { - return impl_->row_count(); -} - -ScanBuilder VortexFile::CreateScanBuilder() const { - return ScanBuilder(impl_->scan_builder()); -} - -} // namespace vortex \ No newline at end of file diff --git a/vortex-cxx/cpp/src/scalar.cpp b/vortex-cxx/cpp/src/scalar.cpp deleted file mode 100644 index 87502ee2f00..00000000000 --- a/vortex-cxx/cpp/src/scalar.cpp +++ /dev/null @@ -1,66 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#include "vortex/scalar.hpp" - -#include "vortex_cxx_bridge/lib.h" - -namespace vortex::scalar { - -Scalar bool_(bool value) { - return Scalar(ffi::bool_scalar_new(value)); -} - -Scalar int8(int8_t value) { - return Scalar(ffi::i8_scalar_new(value)); -} - -Scalar int16(int16_t value) { - return Scalar(ffi::i16_scalar_new(value)); -} - -Scalar int32(int32_t value) { - return Scalar(ffi::i32_scalar_new(value)); -} - -Scalar int64(int64_t value) { - return Scalar(ffi::i64_scalar_new(value)); -} - -Scalar uint8(uint8_t value) { - return Scalar(ffi::u8_scalar_new(value)); -} - -Scalar uint16(uint16_t value) { - return Scalar(ffi::u16_scalar_new(value)); -} - -Scalar uint32(uint32_t value) { - return Scalar(ffi::u32_scalar_new(value)); -} - -Scalar uint64(uint64_t value) { - return Scalar(ffi::u64_scalar_new(value)); -} - -Scalar float32(float value) { - return Scalar(ffi::f32_scalar_new(value)); -} - -Scalar float64(double value) { - return Scalar(ffi::f64_scalar_new(value)); -} - -Scalar string(std::string_view value) { - return Scalar(ffi::string_scalar_new(rust::Str(value.data(), value.length()))); -} - -Scalar binary(const uint8_t *data, size_t length) { - return Scalar(ffi::binary_scalar_new(rust::Slice(data, length))); -} - -Scalar cast(Scalar scalar, dtype::DType dtype) { - return Scalar(std::move(scalar).IntoImpl()->cast_scalar(*std::move(dtype).GetImpl())); -} - -} // namespace vortex::scalar \ No newline at end of file diff --git a/vortex-cxx/cpp/src/scan.cpp b/vortex-cxx/cpp/src/scan.cpp deleted file mode 100644 index 9b40ec00e0d..00000000000 --- a/vortex-cxx/cpp/src/scan.cpp +++ /dev/null @@ -1,116 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#include "vortex/scan.hpp" -#include "vortex/exception.hpp" -#include "rust/cxx.h" -#include "vortex/expr.hpp" - -namespace vortex { -ScanBuilder &ScanBuilder::WithFilter(expr::Expr &&expr) & { - impl_->with_filter(std::move(expr).IntoImpl()); - return *this; -} -ScanBuilder &ScanBuilder::WithFilter(const expr::Expr &expr) & { - impl_->with_filter_ref(expr.Impl()); - return *this; -} -ScanBuilder &&ScanBuilder::WithFilter(expr::Expr &&expr) && { - impl_->with_filter(std::move(expr).IntoImpl()); - return std::move(*this); -} -ScanBuilder &&ScanBuilder::WithFilter(const expr::Expr &expr) && { - impl_->with_filter_ref(expr.Impl()); - return std::move(*this); -} - -ScanBuilder &ScanBuilder::WithProjection(expr::Expr &&expr) & { - impl_->with_projection(std::move(expr).IntoImpl()); - return *this; -} -ScanBuilder &ScanBuilder::WithProjection(const expr::Expr &expr) & { - impl_->with_projection_ref(expr.Impl()); - return *this; -} -ScanBuilder &&ScanBuilder::WithProjection(expr::Expr &&expr) && { - impl_->with_projection(std::move(expr).IntoImpl()); - return std::move(*this); -} -ScanBuilder &&ScanBuilder::WithProjection(const expr::Expr &expr) && { - impl_->with_projection_ref(expr.Impl()); - return std::move(*this); -} - -ScanBuilder &ScanBuilder::WithRowRange(uint64_t row_range_start, uint64_t row_range_end) & { - impl_->with_row_range(row_range_start, row_range_end); - return *this; -} -ScanBuilder &&ScanBuilder::WithRowRange(uint64_t row_range_start, uint64_t row_range_end) && { - impl_->with_row_range(row_range_start, row_range_end); - return std::move(*this); -} - -ScanBuilder &ScanBuilder::WithLimit(uint64_t limit) & { - impl_->with_limit(limit); - return *this; -} - -ScanBuilder &&ScanBuilder::WithLimit(uint64_t limit) && { - impl_->with_limit(limit); - return std::move(*this); -} - -ScanBuilder &ScanBuilder::WithIncludeByIndex(const uint64_t *indices, std::size_t size) & { - impl_->with_include_by_index(rust::Slice(indices, size)); - return *this; -} - -ScanBuilder &&ScanBuilder::WithIncludeByIndex(const uint64_t *indices, std::size_t size) && { - impl_->with_include_by_index(rust::Slice(indices, size)); - return std::move(*this); -} - -ScanBuilder &ScanBuilder::WithOutputSchema(ArrowSchema &output_schema) & { - try { - impl_->with_output_schema(reinterpret_cast(&output_schema)); - } catch (const rust::cxxbridge1::Error &e) { - throw VortexException(e.what()); - } - return *this; -} - -ScanBuilder &&ScanBuilder::WithOutputSchema(ArrowSchema &output_schema) && { - try { - impl_->with_output_schema(reinterpret_cast(&output_schema)); - } catch (const rust::cxxbridge1::Error &e) { - throw VortexException(e.what()); - } - return std::move(*this); -} - -ArrowArrayStream ScanBuilder::IntoStream() && { - try { - ArrowArrayStream stream; - ffi::scan_builder_into_stream(std::move(impl_), reinterpret_cast(&stream)); - return stream; - } catch (const rust::cxxbridge1::Error &e) { - throw VortexException(e.what()); - } -} - -StreamDriver ScanBuilder::IntoStreamDriver() && { - try { - rust::Box reader = - ffi::scan_builder_into_threadsafe_cloneable_reader(std::move(impl_)); - return StreamDriver(std::move(reader)); - } catch (const rust::cxxbridge1::Error &e) { - throw VortexException(e.what()); - } -} - -ArrowArrayStream StreamDriver::CreateArrayStream() const { - ArrowArrayStream stream; - impl_->clone_a_stream(reinterpret_cast(&stream)); - return stream; -} -} // namespace vortex \ No newline at end of file diff --git a/vortex-cxx/cpp/src/write_options.cpp b/vortex-cxx/cpp/src/write_options.cpp deleted file mode 100644 index 712a6933208..00000000000 --- a/vortex-cxx/cpp/src/write_options.cpp +++ /dev/null @@ -1,18 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#include "vortex/write_options.hpp" -#include "vortex/exception.hpp" - -#include "rust/cxx.h" - -namespace vortex { -void VortexWriteOptions::WriteArrayStream(ArrowArrayStream &stream, const std::string &path) { - try { - ffi::write_array_stream(std::move(impl_), reinterpret_cast(&stream), path); - } catch (const rust::cxxbridge1::Error &e) { - throw VortexException(e.what()); - } -} - -} // namespace vortex \ No newline at end of file diff --git a/vortex-cxx/cpp/tests/basic_test.cpp b/vortex-cxx/cpp/tests/basic_test.cpp deleted file mode 100644 index eb44715285a..00000000000 --- a/vortex-cxx/cpp/tests/basic_test.cpp +++ /dev/null @@ -1,510 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "vortex/file.hpp" -#include "vortex/scan.hpp" -#include "vortex/write_options.hpp" -#include "vortex/scalar.hpp" -#include "test_data_generator.hpp" -#include "vortex_cxx_bridge/lib.h" - -#include -#include - -class VortexTest : public ::testing::Test { -protected: - // Helper function to create unique temporary files for each test - static std::string GetUniqueTempFile(const std::string &suffix = "vortex") { - std::filesystem::path temp_dir = std::filesystem::temp_directory_path(); - std::filesystem::path vortex_test_dir = temp_dir / "vortex_test"; - - if (!std::filesystem::exists(vortex_test_dir)) { - std::filesystem::create_directories(vortex_test_dir); - } - - // Use a unique random filename to prevent races between parallel test runs - std::string unique_name = "test_" + std::to_string(std::random_device {}()) + "_" + suffix; - return (vortex_test_dir / unique_name).string(); - } - - // Write test data to a unique temporary file and return the path - static std::string WriteTestData(const std::string &suffix = "test_data.vortex") { - std::string path = GetUniqueTempFile(suffix); - auto stream = vortex::testing::CreateTestDataStream(); - auto write_options = vortex::ffi::write_options_new(); - vortex::ffi::write_array_stream(std::move(write_options), - reinterpret_cast(&stream), - path.c_str()); - return path; - } - - // Helper function to create and initialize array view - nanoarrow::UniqueArrayView CreateArrayView(const nanoarrow::UniqueArray &array, - const nanoarrow::UniqueSchema &schema) { - nanoarrow::UniqueArrayView array_view; - ArrowError error; - ArrowErrorCode init_result = ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), &error); - if (init_result != NANOARROW_OK) { - std::cerr << "Error: " << error.message << '\n'; - std::abort(); - } - ArrowErrorCode set_result = ArrowArrayViewSetArray(array_view.get(), array.get(), nullptr); - if (set_result != NANOARROW_OK) { - std::cerr << "Error: " << error.message << '\n'; - std::abort(); - } - return array_view; - } - - std::pair - StreamToUniqueStreamSchema(ArrowArrayStream &stream) { - nanoarrow::UniqueArrayStream array_stream; - ArrowArrayStreamMove(&stream, array_stream.get()); - ArrowError error; - nanoarrow::UniqueSchema schema; - ArrowErrorCode set_result = ArrowArrayStreamGetSchema(array_stream.get(), schema.get(), &error); - if (set_result != NANOARROW_OK) { - std::cerr << "Error: " << error.message << '\n'; - std::abort(); - } - return {std::move(array_stream), std::move(schema)}; - } - - nanoarrow::UniqueArray ReadFirstArrayFromUniqueStream(nanoarrow::UniqueArrayStream &array_stream) { - nanoarrow::UniqueArray array; - int get_next_result = array_stream->get_next(array_stream.get(), array.get()); - EXPECT_EQ(get_next_result, 0); - return array; - } - - std::pair - ReadFirstArrayFromStream(ArrowArrayStream reference_stream) { - auto [ref_array_stream, ref_schema] = StreamToUniqueStreamSchema(reference_stream); - auto ref_array = ReadFirstArrayFromUniqueStream(ref_array_stream); - return {std::move(ref_array), std::move(ref_schema)}; - } - - /// Both array are struct of int64 - void ValidateArray(const nanoarrow::UniqueArray &actual_array, - const nanoarrow::UniqueSchema &actual_schema, - const nanoarrow::UniqueArray &ref_array, - const nanoarrow::UniqueSchema &ref_schema) { - // Basic properties validation - ASSERT_EQ(actual_schema->n_children, ref_schema->n_children); - - auto actual_view = CreateArrayView(actual_array, actual_schema); - auto ref_view = CreateArrayView(ref_array, ref_schema); - - ASSERT_EQ(actual_array->length, ref_array->length); - - // Compare all fields - for (int64_t field_idx = 0; field_idx < actual_schema->n_children; ++field_idx) { - auto actual_field = actual_view->children[field_idx]; - auto expected_field = ref_view->children[field_idx]; - - ASSERT_EQ(actual_field->array->length, expected_field->array->length); - - for (int64_t i = 0; i < actual_field->array->length; ++i) { - int64_t actual_value = ArrowArrayViewGetIntUnsafe(actual_field, i); - int64_t expected_value = ArrowArrayViewGetIntUnsafe(expected_field, i); - - ASSERT_EQ(actual_value, expected_value); - } - } - } - - /// Both array are struct of int64 - void ValidateArrayWithSelection(const nanoarrow::UniqueArray &actual_array, - const nanoarrow::UniqueSchema &actual_schema, - const nanoarrow::UniqueArray &ref_array, - const nanoarrow::UniqueSchema &ref_schema, - const std::vector &row_indices) { - // Basic properties validation - ASSERT_EQ(actual_schema->n_children, ref_schema->n_children); - if (row_indices.empty()) { - ASSERT_EQ(actual_array->length, 0); - return; - } - auto actual_view = CreateArrayView(actual_array, actual_schema); - auto ref_view = CreateArrayView(ref_array, ref_schema); - - ASSERT_EQ(actual_array->length, row_indices.size()); - - // Selective row comparison using indices - ASSERT_EQ(actual_array->length, static_cast(row_indices.size())); - - for (int64_t i = 0; i < static_cast(row_indices.size()); ++i) { - int64_t ref_idx = row_indices[i]; - - for (int64_t field_idx = 0; field_idx < actual_schema->n_children; ++field_idx) { - int64_t actual_val = ArrowArrayViewGetIntUnsafe(actual_view->children[field_idx], i); - int64_t expected_val = ArrowArrayViewGetIntUnsafe(ref_view->children[field_idx], ref_idx); - - ASSERT_EQ(actual_val, expected_val); - } - } - } - - // Helper to execute scan builder and get array+schema - std::pair ScanFirstArrayFromTestData( - const std::function &configureScanBuilder) { - auto test_data_path = WriteTestData(); - auto file = vortex::VortexFile::Open(test_data_path); - auto scan_builder = file.CreateScanBuilder(); - auto stream = configureScanBuilder(scan_builder); - - auto [array_stream, schema] = StreamToUniqueStreamSchema(stream); - auto array = ReadFirstArrayFromUniqueStream(array_stream); - - return {std::move(array), std::move(schema)}; - } - - /// Validate array with projection - only checks specified field indices - void ValidateArrayWithProjection(const nanoarrow::UniqueArray &actual_array, - const nanoarrow::UniqueSchema &actual_schema, - const nanoarrow::UniqueArray &ref_array, - const nanoarrow::UniqueSchema &ref_schema, - const std::vector &field_idxs) { - ASSERT_EQ(actual_schema->n_children, field_idxs.size()); - - auto actual_view = CreateArrayView(actual_array, actual_schema); - auto ref_view = CreateArrayView(ref_array, ref_schema); - - ASSERT_EQ(actual_array->length, ref_array->length); - - // Compare only the specified fields - for (int64_t i = 0; i < actual_array->length; ++i) { - for (size_t proj_idx = 0; proj_idx < field_idxs.size(); ++proj_idx) { - int64_t ref_field_idx = field_idxs[proj_idx]; - int64_t actual_val = ArrowArrayViewGetIntUnsafe(actual_view->children[proj_idx], i); - int64_t expected_val = ArrowArrayViewGetIntUnsafe(ref_view->children[ref_field_idx], i); - ASSERT_EQ(actual_val, expected_val); - } - } - } - - // Top-level test helper that all tests can use - void - RunScanBuilderTest(const std::function &configureScanBuilder, - ArrowArrayStream expected_stream, - const std::vector &expected_row_indices = {}, - bool selection = false) { - - auto [array, schema] = ScanFirstArrayFromTestData(configureScanBuilder); - auto [ref_array, ref_schema] = ReadFirstArrayFromStream(expected_stream); - selection == false - ? ValidateArray(array, schema, ref_array, ref_schema) - : ValidateArrayWithSelection(array, schema, ref_array, ref_schema, expected_row_indices); - } - - // New helper for projection tests - void RunScanBuilderProjectionTest( - const std::function &configureScanBuilder, - ArrowArrayStream expected_stream, - const std::vector &field_idxs) { - - auto [array, schema] = ScanFirstArrayFromTestData(configureScanBuilder); - auto [ref_array, ref_schema] = ReadFirstArrayFromStream(expected_stream); - - ValidateArrayWithProjection(array, schema, ref_array, ref_schema, field_idxs); - } -}; - -TEST_F(VortexTest, ScanToStream) { - RunScanBuilderTest([](vortex::ScanBuilder &builder) { return std::move(builder).IntoStream(); }, - vortex::testing::CreateTestDataStream()); -} - -TEST_F(VortexTest, ScanBuilderWithLimitWithRowRange) { - // Test field "a" and "b" - should contain values from rows 1-2 from original data (indices 1 and - // 2) - RunScanBuilderTest( - [](vortex::ScanBuilder &scan_builder) { - return std::move(scan_builder.WithLimit(2).WithRowRange(1, 4)).IntoStream(); - }, - vortex::testing::CreateTestDataStream(), - {1, 2}, - true); -} - -TEST_F(VortexTest, ScanBuilderWithIncludeByIndex) { - std::vector include_by_index = {1, 3}; - - RunScanBuilderTest( - [&include_by_index](vortex::ScanBuilder &scan_builder) { - return std::move( - scan_builder.WithIncludeByIndex(include_by_index.data(), include_by_index.size())) - .IntoStream(); - }, - vortex::testing::CreateTestDataStream(), - {1, 3}, - true); -} - -TEST_F(VortexTest, ScanBuilderWithRowRangeWithIncludeByIndex) { - std::vector include_by_index = {1, 3, 4}; - - RunScanBuilderTest( - [&include_by_index](vortex::ScanBuilder &scan_builder) { - return std::move(scan_builder.WithRowRange(2, 5).WithIncludeByIndex(include_by_index.data(), - include_by_index.size())) - .IntoStream(); - }, - vortex::testing::CreateTestDataStream(), - {3, 4}, - true); -} - -TEST_F(VortexTest, WriteArrayStream) { - auto test_data_path = WriteTestData(); - auto file = vortex::VortexFile::Open(test_data_path); - auto stream = file.CreateScanBuilder().IntoStream(); - - // Write the stream to a new Vortex file - std::string test_output_path = GetUniqueTempFile("write_output.vortex"); - vortex::VortexWriteOptions write_options; - ASSERT_NO_THROW(write_options.WriteArrayStream(stream, test_output_path)); - - // Verify the written file - auto written_file = vortex::VortexFile::Open(test_output_path); - ASSERT_EQ(written_file.RowCount(), 5); - - // Verify data integrity by reading from the written file - auto out_stream = written_file.CreateScanBuilder().IntoStream(); - auto [array_stream, schema] = StreamToUniqueStreamSchema(out_stream); - auto array = ReadFirstArrayFromUniqueStream(array_stream); - auto [ref_array, ref_schema] = ReadFirstArrayFromStream(vortex::testing::CreateTestDataStream()); - ValidateArray(array, schema, ref_array, ref_schema); -} - -TEST_F(VortexTest, ConcurrentMultiStreamRead) { - std::string test_data_path_1m = GetUniqueTempFile("concurrent_1m.vortex"); - auto stream_1m = vortex::testing::CreateTestData1MStream(); - auto write_options = vortex::ffi::write_options_new(); - vortex::ffi::write_array_stream(std::move(write_options), - reinterpret_cast(&stream_1m), - test_data_path_1m.c_str()); - - auto file = vortex::VortexFile::Open(test_data_path_1m); - auto stream_driver = file.CreateScanBuilder().IntoStreamDriver(); - - // Structure to hold batch data with first ID and nanoarrow array - struct BatchData { - int64_t first_id; - nanoarrow::UniqueArray array; - - BatchData(int64_t first_id, nanoarrow::UniqueArray array) - : first_id(first_id), array(std::move(array)) { - } - }; - - std::vector thread1_batches; - std::vector thread2_batches; - - // Helper function to read from a stream and collect batches - auto read_stream = [&](std::vector &batches) { - // Each thread creates its own stream - auto stream = stream_driver.CreateArrayStream(); - auto [array_stream, schema] = StreamToUniqueStreamSchema(stream); - - std::vector local_batches; - - while (true) { - nanoarrow::UniqueArray array; - int get_next_result = array_stream->get_next(array_stream.get(), array.get()); - - if (get_next_result != 0) { - std::cerr << "Error: " << array_stream->get_last_error(array_stream.get()) << '\n'; - std::abort(); - } - - if (array->length == 0) { - break; // Empty array indicates end - } - - auto array_view = CreateArrayView(array, schema); - - int64_t first_id = ArrowArrayViewGetIntUnsafe(array_view->children[0], 0); - - local_batches.emplace_back(first_id, std::move(array)); - } - batches = std::move(local_batches); - }; - - // Launch two threads - std::thread thread1(read_stream, std::ref(thread1_batches)); - std::thread thread2(read_stream, std::ref(thread2_batches)); - - // Wait for both threads to complete - thread1.join(); - thread2.join(); - - // Combine all batches from both threads - std::vector all_batches; - all_batches.insert(all_batches.end(), - std::make_move_iterator(thread1_batches.begin()), - std::make_move_iterator(thread1_batches.end())); - all_batches.insert(all_batches.end(), - std::make_move_iterator(thread2_batches.begin()), - std::make_move_iterator(thread2_batches.end())); - - // Sort batches by first ID to ensure proper validation order - std::sort(all_batches.begin(), all_batches.end(), [](const BatchData &a, const BatchData &b) { - return a.first_id < b.first_id; - }); - - // Create reference data for validation - auto [ref_array, ref_schema] = ReadFirstArrayFromStream(vortex::testing::CreateTestData1MStream()); - auto ref_array_view = CreateArrayView(ref_array, ref_schema); - - // Validate all data against reference - constexpr size_t EXPECTED_ROWS = static_cast(1024) * 1024; - size_t total_rows_read = 0; - int64_t reference_offset = 0; - - auto stream_for_schema = stream_driver.CreateArrayStream(); - auto [_, schema] = StreamToUniqueStreamSchema(stream_for_schema); - for (const auto &batch : all_batches) { - auto array_view = CreateArrayView(batch.array, schema); - - for (int64_t i = 0; i < batch.array->length; ++i) { - - int64_t actual_id = ArrowArrayViewGetIntUnsafe(array_view->children[0], i); - int32_t actual_value = - static_cast(ArrowArrayViewGetIntUnsafe(array_view->children[1], i)); - - int64_t expected_id = ArrowArrayViewGetIntUnsafe(ref_array_view->children[0], reference_offset); - int32_t expected_value = static_cast( - ArrowArrayViewGetIntUnsafe(ref_array_view->children[1], reference_offset)); - - ASSERT_EQ(actual_id, expected_id); - ASSERT_EQ(actual_value, expected_value); - reference_offset++; - } - total_rows_read += batch.array->length; - } - - // Verify we read all expected data - ASSERT_EQ(total_rows_read, EXPECTED_ROWS) - << "Expected to read " << EXPECTED_ROWS << " rows, but read " << total_rows_read << " rows"; - ASSERT_EQ(reference_offset, EXPECTED_ROWS) << "Reference validation didn't cover all rows"; - - ASSERT_GT(all_batches.size(), 1) << "Expected multiple batches, but got " << all_batches.size(); -} - -namespace ve = vortex::expr; -namespace vs = vortex::scalar; - -TEST_F(VortexTest, ScanBuilderWithFilter) { - // Test filtering with eq(column("a"), val) - should return only rows where column "a" equals 30 - RunScanBuilderTest( - [](vortex::ScanBuilder &scan_builder) { - auto filter = ve::eq(ve::column("a"), ve::literal(vs::int32(30))); - return std::move(scan_builder.WithFilter(std::move(filter))).IntoStream(); - }, - vortex::testing::CreateTestDataStream(), - {2}, - true); // Row index 2 corresponds to value 30 -} - -TEST_F(VortexTest, ScanBuilderWithFilterLvalueref) { - // Test filtering with eq(column("a"), val) - should return only rows where column "a" equals 30 - RunScanBuilderTest( - [](vortex::ScanBuilder &scan_builder) { - const auto filter = ve::eq(ve::column("a"), ve::literal(vs::int32(30))); - return std::move(scan_builder.WithFilter(filter)).IntoStream(); - }, - vortex::testing::CreateTestDataStream(), - {2}, - true); // Row index 2 corresponds to value 30 -} - -TEST_F(VortexTest, ScanBuilderWithFilterNoMatches) { - // Test filtering with eq(column("a"), val) where no rows match - should return empty result - RunScanBuilderTest( - [](vortex::ScanBuilder &scan_builder) { - auto filter = ve::eq(ve::column("a"), - ve::literal(vs::int32(999)) // Value that doesn't exist in test data - ); - return std::move(scan_builder.WithFilter(std::move(filter))).IntoStream(); - }, - vortex::testing::CreateTestDataStream(), - {}, - true); // No matching rows -} - -TEST_F(VortexTest, ScanBuilderWithFilterUsingDTypeFromArrowAndScalarCast) { - // Test filtering using DType::from_arrow and Scalar::cast functionality - // This test creates a filter expression by casting a scalar to match the column type - - // Test DType::from_arrow with int32 schema - nanoarrow::UniqueSchema int32_schema; - ArrowSchemaInit(int32_schema.get()); - ArrowErrorCode result = ArrowSchemaSetType(int32_schema.get(), NANOARROW_TYPE_INT32); - EXPECT_EQ(result, NANOARROW_OK); - result = ArrowSchemaSetName(int32_schema.get(), "test_field"); - EXPECT_EQ(result, NANOARROW_OK); - - auto dtype = vortex::dtype::from_arrow(*int32_schema.get()); - - // Use the casted scalar in filter expression - create a new scalar for lambda - RunScanBuilderTest( - [&](vortex::ScanBuilder &scan_builder) { - auto test_scalar = vs::cast(vs::int64(30), std::move(dtype)); - auto filter = ve::eq(ve::column("a"), ve::literal(std::move(test_scalar))); - return std::move(scan_builder.WithFilter(std::move(filter))).IntoStream(); - }, - vortex::testing::CreateTestDataStream(), - {2}, - true); // Row index 2 corresponds to value 30 -} - -TEST_F(VortexTest, ScanBuilderWithProjectionSingleColumn) { - // Test projection selecting only column "a" (field index 0) - RunScanBuilderProjectionTest( - [](vortex::ScanBuilder &scan_builder) { - auto projection = ve::select({"a"}, ve::root()); - return std::move(scan_builder.WithProjection(std::move(projection))).IntoStream(); - }, - vortex::testing::CreateTestDataStream(), - {0}); -} - -TEST_F(VortexTest, OpenFromBuffer) { - std::string test_file_path = GetUniqueTempFile("buffer.vortex"); - auto stream = vortex::testing::CreateTestDataStream(); - auto write_options = vortex::ffi::write_options_new(); - vortex::ffi::write_array_stream(std::move(write_options), - reinterpret_cast(&stream), - test_file_path.c_str()); - - std::ifstream file(test_file_path, std::ios::binary | std::ios::ate); - ASSERT_TRUE(file.is_open()) << "Failed to open file: " << test_file_path; - - std::streamsize file_size = file.tellg(); - file.seekg(0, std::ios::beg); - - std::vector buffer(file_size); - ASSERT_TRUE(file.read(reinterpret_cast(buffer.data()), file_size)) - << "Failed to read file into buffer"; - file.close(); - - auto vortex_file = vortex::VortexFile::Open(buffer.data(), buffer.size()); - ASSERT_EQ(vortex_file.RowCount(), 5); - - auto scan_stream = vortex_file.CreateScanBuilder().IntoStream(); - auto [array_stream, schema] = StreamToUniqueStreamSchema(scan_stream); - auto array = ReadFirstArrayFromUniqueStream(array_stream); - - auto [ref_array, ref_schema] = ReadFirstArrayFromStream(vortex::testing::CreateTestDataStream()); - ValidateArray(array, schema, ref_array, ref_schema); -} diff --git a/vortex-cxx/cpp/tests/test_data_generator.cpp b/vortex-cxx/cpp/tests/test_data_generator.cpp deleted file mode 100644 index 67dcb290509..00000000000 --- a/vortex-cxx/cpp/tests/test_data_generator.cpp +++ /dev/null @@ -1,125 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#include "test_data_generator.hpp" -#include -#include -#include -#include -#include - -namespace vortex { -namespace testing { - - ArrowArrayStream CreateTestDataStream() { - // Create a simple two-column struct with int32 data - // Schema: struct{a: int32, b: int32} - nanoarrow::UniqueSchema schema; - ArrowSchemaInit(schema.get()); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetType(schema.get(), NANOARROW_TYPE_STRUCT)); - NANOARROW_THROW_NOT_OK(ArrowSchemaAllocateChildren(schema.get(), 2)); - ArrowSchemaInit(schema->children[0]); - ArrowSchemaInit(schema->children[1]); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[0], "a")); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetType(schema->children[0], NANOARROW_TYPE_INT32)); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[1], "b")); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetType(schema->children[1], NANOARROW_TYPE_INT32)); - - // Create arrays for each field - nanoarrow::UniqueArray field_a, field_b; - NANOARROW_THROW_NOT_OK(ArrowArrayInitFromType(field_a.get(), NANOARROW_TYPE_INT32)); - NANOARROW_THROW_NOT_OK(ArrowArrayInitFromType(field_b.get(), NANOARROW_TYPE_INT32)); - - // Reserve for 5 elements - NANOARROW_THROW_NOT_OK(ArrowArrayStartAppending(field_a.get())); - NANOARROW_THROW_NOT_OK(ArrowArrayStartAppending(field_b.get())); - - // Add data: [10, 20, 30, 40, 50] - std::vector data = {10, 20, 30, 40, 50}; - for (int32_t value : data) { - NANOARROW_THROW_NOT_OK(ArrowArrayAppendInt(field_a.get(), value)); - NANOARROW_THROW_NOT_OK(ArrowArrayAppendInt(field_b.get(), value)); - } - - NANOARROW_THROW_NOT_OK( - ArrowArrayFinishBuilding(field_a.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); - NANOARROW_THROW_NOT_OK( - ArrowArrayFinishBuilding(field_b.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); - - // Create struct array - nanoarrow::UniqueArray struct_array; - NANOARROW_THROW_NOT_OK(ArrowArrayInitFromType(struct_array.get(), NANOARROW_TYPE_STRUCT)); - NANOARROW_THROW_NOT_OK(ArrowArrayAllocateChildren(struct_array.get(), 2)); - struct_array->length = 5; - ArrowArrayMove(field_a.get(), struct_array->children[0]); - ArrowArrayMove(field_b.get(), struct_array->children[1]); - - // Create vector and move array into it - std::vector arrays; - arrays.push_back(std::move(struct_array)); - - // Create stream - ArrowArrayStream stream; - nanoarrow::VectorArrayStream vector_stream(schema.get(), std::move(arrays)); - vector_stream.ToArrayStream(&stream); - - return stream; - } - - ArrowArrayStream CreateTestData1MStream() { - constexpr size_t NUM_ROWS = 1024UL * 1024; - - // Create schema: struct{id: int64, value: int32} - nanoarrow::UniqueSchema schema; - ArrowSchemaInit(schema.get()); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetType(schema.get(), NANOARROW_TYPE_STRUCT)); - NANOARROW_THROW_NOT_OK(ArrowSchemaAllocateChildren(schema.get(), 2)); - ArrowSchemaInit(schema->children[0]); - ArrowSchemaInit(schema->children[1]); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[0], "id")); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetType(schema->children[0], NANOARROW_TYPE_INT64)); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[1], "value")); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetType(schema->children[1], NANOARROW_TYPE_INT32)); - - // Create arrays for each field - nanoarrow::UniqueArray id_field, value_field; - NANOARROW_THROW_NOT_OK(ArrowArrayInitFromType(id_field.get(), NANOARROW_TYPE_INT64)); - NANOARROW_THROW_NOT_OK(ArrowArrayInitFromType(value_field.get(), NANOARROW_TYPE_INT32)); - - // Reserve space - NANOARROW_THROW_NOT_OK(ArrowArrayStartAppending(id_field.get())); - NANOARROW_THROW_NOT_OK(ArrowArrayStartAppending(value_field.get())); - - // Add data - for (size_t i = 0; i < NUM_ROWS; ++i) { - NANOARROW_THROW_NOT_OK(ArrowArrayAppendInt(id_field.get(), static_cast(i))); - NANOARROW_THROW_NOT_OK(ArrowArrayAppendInt(value_field.get(), static_cast(i * 2))); - } - - NANOARROW_THROW_NOT_OK( - ArrowArrayFinishBuilding(id_field.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); - NANOARROW_THROW_NOT_OK( - ArrowArrayFinishBuilding(value_field.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); - - // Create struct array - nanoarrow::UniqueArray struct_array; - NANOARROW_THROW_NOT_OK(ArrowArrayInitFromType(struct_array.get(), NANOARROW_TYPE_STRUCT)); - NANOARROW_THROW_NOT_OK(ArrowArrayAllocateChildren(struct_array.get(), 2)); - struct_array->length = NUM_ROWS; - ArrowArrayMove(id_field.get(), struct_array->children[0]); - ArrowArrayMove(value_field.get(), struct_array->children[1]); - - // Create vector and move array into it - std::vector arrays; - arrays.push_back(std::move(struct_array)); - - // Create stream - ArrowArrayStream stream; - nanoarrow::VectorArrayStream vector_stream(schema.get(), std::move(arrays)); - vector_stream.ToArrayStream(&stream); - - return stream; - } - -} // namespace testing -} // namespace vortex \ No newline at end of file diff --git a/vortex-cxx/cpp/tests/test_data_generator.hpp b/vortex-cxx/cpp/tests/test_data_generator.hpp deleted file mode 100644 index ec5dba55d50..00000000000 --- a/vortex-cxx/cpp/tests/test_data_generator.hpp +++ /dev/null @@ -1,20 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#pragma once - -#include - -namespace vortex { -namespace testing { - - /// Create test data with structure {a: [10, 20, 30, 40, 50], b: [10, 20, 30, 40, 50]} - /// This stream only has one Array - ArrowArrayStream CreateTestDataStream(); - - /// Create 1M rows of test data with structure {id: [0..1M], value: [0, 2, 4, ..., 2M]} - /// This stream only has one Array - ArrowArrayStream CreateTestData1MStream(); - -} // namespace testing -} // namespace vortex \ No newline at end of file diff --git a/vortex-cxx/examples/.gitignore b/vortex-cxx/examples/.gitignore deleted file mode 100644 index 065a36090f7..00000000000 --- a/vortex-cxx/examples/.gitignore +++ /dev/null @@ -1 +0,0 @@ -!goldenfiles/example.parquet \ No newline at end of file diff --git a/vortex-cxx/examples/CMakeLists.txt b/vortex-cxx/examples/CMakeLists.txt deleted file mode 100644 index 348d80e1547..00000000000 --- a/vortex-cxx/examples/CMakeLists.txt +++ /dev/null @@ -1,34 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright the Vortex contributors - -cmake_minimum_required(VERSION 3.22) - -project(vortex-examples) - -set(CMAKE_CXX_STANDARD 20) -set(CMAKE_CXX_STANDARD_REQUIRED ON) - -find_program(SCCACHE_PROGRAM sccache) -if (SCCACHE_PROGRAM) - set(CMAKE_C_COMPILER_LAUNCHER "${SCCACHE_PROGRAM}") - set(CMAKE_CXX_COMPILER_LAUNCHER "${SCCACHE_PROGRAM}") - message(STATUS "Sccache found: ${SCCACHE_PROGRAM}") -else () - message(STATUS "Sccache not found") -endif () - -include(FetchContent) -FetchContent_Declare( - vortex - SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../../ - SOURCE_SUBDIR vortex-cxx -) -FetchContent_MakeAvailable(vortex) - -add_executable(hello-vortex hello-vortex.cpp) - -if(APPLE) - set(APPLE_LINK_FLAGS "-framework CoreFoundation -framework Security") -endif() - -target_link_libraries(hello-vortex PRIVATE vortex ${APPLE_LINK_FLAGS}) diff --git a/vortex-cxx/examples/README.md b/vortex-cxx/examples/README.md deleted file mode 100644 index 40e7708fa33..00000000000 --- a/vortex-cxx/examples/README.md +++ /dev/null @@ -1,11 +0,0 @@ -# C++ examples - -This example shows how to interface with the C++ API using CMake. - -```bash -mkdir -p build -cd build -cmake .. -make -j$(nproc) -./hello-vortex -``` \ No newline at end of file diff --git a/vortex-cxx/examples/hello-vortex.cpp b/vortex-cxx/examples/hello-vortex.cpp deleted file mode 100644 index c6a7fec3b9a..00000000000 --- a/vortex-cxx/examples/hello-vortex.cpp +++ /dev/null @@ -1,129 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#include "nanoarrow/common/inline_types.h" -#include "nanoarrow/hpp/unique.hpp" -#include "nanoarrow/nanoarrow.hpp" -#include "vortex/file.hpp" -#include "vortex/scan.hpp" -#include "vortex/write_options.hpp" -#include -#include -#include -#include -#include - -/// Create test data with structure {a: [10, 20, 30, 40, 50], b: [100, 200, 300, 400, 500]} -ArrowArrayStream CreateTestDataStream() { - // Create a simple two-column struct with int32 data - // Schema: struct{a: int32, b: int32} - nanoarrow::UniqueSchema schema; - ArrowSchemaInit(schema.get()); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetType(schema.get(), NANOARROW_TYPE_STRUCT)); - NANOARROW_THROW_NOT_OK(ArrowSchemaAllocateChildren(schema.get(), 2)); - ArrowSchemaInit(schema->children[0]); - ArrowSchemaInit(schema->children[1]); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[0], "a")); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetType(schema->children[0], NANOARROW_TYPE_INT32)); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[1], "b")); - NANOARROW_THROW_NOT_OK(ArrowSchemaSetType(schema->children[1], NANOARROW_TYPE_INT32)); - - // Create arrays for each field - nanoarrow::UniqueArray field_a, field_b; - NANOARROW_THROW_NOT_OK(ArrowArrayInitFromType(field_a.get(), NANOARROW_TYPE_INT32)); - NANOARROW_THROW_NOT_OK(ArrowArrayInitFromType(field_b.get(), NANOARROW_TYPE_INT32)); - - // Reserve for 5 elements - NANOARROW_THROW_NOT_OK(ArrowArrayStartAppending(field_a.get())); - NANOARROW_THROW_NOT_OK(ArrowArrayStartAppending(field_b.get())); - - // Add data: a=[10, 20, 30, 40, 50], b=[100, 200, 300, 400, 500] - std::vector data_a = {10, 20, 30, 40, 50}; - std::vector data_b = {100, 200, 300, 400, 500}; - for (size_t i = 0; i < data_a.size(); ++i) { - NANOARROW_THROW_NOT_OK(ArrowArrayAppendInt(field_a.get(), data_a[i])); - NANOARROW_THROW_NOT_OK(ArrowArrayAppendInt(field_b.get(), data_b[i])); - } - - NANOARROW_THROW_NOT_OK(ArrowArrayFinishBuilding(field_a.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); - NANOARROW_THROW_NOT_OK(ArrowArrayFinishBuilding(field_b.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); - - // Create struct array - nanoarrow::UniqueArray struct_array; - NANOARROW_THROW_NOT_OK(ArrowArrayInitFromType(struct_array.get(), NANOARROW_TYPE_STRUCT)); - NANOARROW_THROW_NOT_OK(ArrowArrayAllocateChildren(struct_array.get(), 2)); - struct_array->length = 5; - ArrowArrayMove(field_a.get(), struct_array->children[0]); - ArrowArrayMove(field_b.get(), struct_array->children[1]); - - // Create vector and move array into it - std::vector arrays; - arrays.push_back(std::move(struct_array)); - - // Create stream - ArrowArrayStream stream; - nanoarrow::VectorArrayStream vector_stream(schema.get(), std::move(arrays)); - vector_stream.ToArrayStream(&stream); - - return stream; -} - -int main() { - // Create a temporary file path - std::filesystem::path temp_dir = std::filesystem::temp_directory_path(); - std::string vortex_file = (temp_dir / "hello_vortex_example.vortex").string(); - - std::cout << "=== Vortex C++ Example ===" << '\n'; - std::cout << "Writing to: " << vortex_file << '\n'; - - // Write test data to a Vortex file - { - auto stream = CreateTestDataStream(); - vortex::VortexWriteOptions write_options; - write_options.WriteArrayStream(stream, vortex_file); - std::cout << "Wrote test data to file" << '\n'; - } - - auto check_stream = [](ArrowArrayStream &stream) { - nanoarrow::UniqueArray array; - int get_next_result = stream.get_next(&stream, array.get()); - assert(get_next_result == 0); - std::cout << "Number of rows: " << array->length << '\n'; - std::cout << "Number of columns in schema: " << array->n_children << '\n'; - }; - - // 1. Classic C++ builder pattern - std::cout << "\n1. Classic C++ builder pattern:" << '\n'; - { - auto builder = vortex::VortexFile::Open(vortex_file).CreateScanBuilder(); - builder.WithLimit(3); - auto stream = std::move(builder).IntoStream(); - check_stream(stream); - } - // 2. One-line Rusty function chain - std::cout << "\n2. One-line Rusty function chain:" << '\n'; - { - auto stream = vortex::VortexFile::Open(vortex_file).CreateScanBuilder().WithLimit(3).IntoStream(); - check_stream(stream); - } - // 3. Conditionally set the builder - std::cout << "\n3. Conditionally set the builder:" << '\n'; - { - auto limit = 1; - auto builder = vortex::VortexFile::Open(vortex_file).CreateScanBuilder(); - if (limit > 0) { - // prefer C++ way - builder.WithLimit(1); - // Rusty way is Ok, but you have to move the builder to an rvalue. - // builder = std::move(builder).WithLimit(3); - } - auto stream = std::move(builder).IntoStream(); - check_stream(stream); - } - - // Clean up - std::filesystem::remove(vortex_file); - std::cout << "\nCleaned up temporary file" << '\n'; - - return 0; -} \ No newline at end of file diff --git a/vortex-cxx/src/dtype.rs b/vortex-cxx/src/dtype.rs deleted file mode 100644 index 0c125bcf822..00000000000 --- a/vortex-cxx/src/dtype.rs +++ /dev/null @@ -1,94 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::fmt::Display; -use std::fmt::Formatter; - -use anyhow::Result; -use arrow_array::ffi::FFI_ArrowSchema; -use arrow_schema::Field; -use vortex::dtype::DType as RustDType; -use vortex::dtype::DecimalDType; -use vortex::dtype::Nullability; -use vortex::dtype::PType as RustPType; -use vortex::dtype::arrow::FromArrowType; - -use crate::ffi; -pub(crate) struct DType { - pub(crate) inner: RustDType, -} - -pub(crate) fn dtype_null() -> Box { - Box::new(DType { - inner: RustDType::Null, - }) -} - -pub(crate) fn dtype_bool(nullable: bool) -> Box { - Box::new(DType { - inner: RustDType::Bool(nullability_from_bool(nullable)), - }) -} - -pub(crate) fn dtype_primitive(ptype: ffi::PType, nullable: bool) -> Box { - let vortex_ptype = match ptype { - ffi::PType::U8 => RustPType::U8, - ffi::PType::U16 => RustPType::U16, - ffi::PType::U32 => RustPType::U32, - ffi::PType::U64 => RustPType::U64, - ffi::PType::I8 => RustPType::I8, - ffi::PType::I16 => RustPType::I16, - ffi::PType::I32 => RustPType::I32, - ffi::PType::I64 => RustPType::I64, - ffi::PType::F16 => RustPType::F16, - ffi::PType::F32 => RustPType::F32, - ffi::PType::F64 => RustPType::F64, - _ => unreachable!(), - }; - Box::new(DType { - inner: RustDType::Primitive(vortex_ptype, nullability_from_bool(nullable)), - }) -} - -pub(crate) fn dtype_decimal(precision: u8, scale: i8, nullable: bool) -> Box { - Box::new(DType { - inner: RustDType::Decimal( - DecimalDType::new(precision, scale), - nullability_from_bool(nullable), - ), - }) -} - -pub(crate) fn dtype_utf8(nullable: bool) -> Box { - Box::new(DType { - inner: RustDType::Utf8(nullability_from_bool(nullable)), - }) -} - -pub(crate) fn dtype_binary(nullable: bool) -> Box { - Box::new(DType { - inner: RustDType::Binary(nullability_from_bool(nullable)), - }) -} - -impl Display for DType { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "{0}", self.inner) - } -} - -pub(crate) fn nullability_from_bool(nullable: bool) -> Nullability { - if nullable { - Nullability::Nullable - } else { - Nullability::NonNullable - } -} - -pub(crate) unsafe fn from_arrow(ffi_schema: *mut u8, non_nullable: bool) -> Result> { - let arrow_schema = unsafe { FFI_ArrowSchema::from_raw(ffi_schema as *mut FFI_ArrowSchema) }; - let arrow_dtype = arrow_schema::DataType::try_from(&arrow_schema)?; - Ok(Box::new(DType { - inner: RustDType::from_arrow(&Field::new("_", arrow_dtype, !non_nullable)), - })) -} diff --git a/vortex-cxx/src/expr.rs b/vortex-cxx/src/expr.rs deleted file mode 100644 index b980db7111b..00000000000 --- a/vortex-cxx/src/expr.rs +++ /dev/null @@ -1,81 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex::dtype::FieldName; -use vortex::expr::Expression; - -use crate::scalar::Scalar; - -pub(crate) struct Expr { - pub(crate) inner: Expression, -} - -pub(crate) fn literal(scalar: Box) -> Box { - Box::new(Expr { - inner: vortex::expr::lit(scalar.inner), - }) -} - -pub(crate) fn root() -> Box { - Box::new(Expr { - inner: vortex::expr::root(), - }) -} - -pub(crate) fn column(name: String) -> Box { - Box::new(Expr { - inner: vortex::expr::get_item(name, vortex::expr::root()), - }) -} - -pub(crate) fn get_item(field: String, child: Box) -> Box { - Box::new(Expr { - inner: vortex::expr::get_item(field, child.inner), - }) -} - -pub(crate) fn not_(child: Box) -> Box { - Box::new(Expr { - inner: vortex::expr::not(child.inner), - }) -} - -pub(crate) fn is_null(child: Box) -> Box { - Box::new(Expr { - inner: vortex::expr::is_null(child.inner), - }) -} - -macro_rules! binary_op { - ($fn_name:ident $(, $suffix:tt)?) => { - paste::paste! { - pub(crate) fn [<$fn_name $($suffix)?>]( - lhs: Box, - rhs: Box, - ) -> Box { - Box::new(Expr { - inner: vortex::expr::$fn_name(lhs.inner, rhs.inner), - }) - } - } - }; -} - -binary_op!(eq); -binary_op!(not_eq, _); -binary_op!(gt); -binary_op!(gt_eq); -binary_op!(lt); -binary_op!(lt_eq); -binary_op!(and, _); -binary_op!(or, _); -binary_op!(checked_add); - -pub(crate) fn select(fields: Vec, child: Box) -> Box { - Box::new(Expr { - inner: vortex::expr::select( - fields.into_iter().map(FieldName::from).collect::>(), - child.inner, - ), - }) -} diff --git a/vortex-cxx/src/lib.rs b/vortex-cxx/src/lib.rs deleted file mode 100644 index b6d002513f6..00000000000 --- a/vortex-cxx/src/lib.rs +++ /dev/null @@ -1,139 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#![allow(clippy::boxed_local)] - -mod dtype; -mod expr; -mod read; -mod scalar; -mod write; - -use std::sync::LazyLock; - -use dtype::*; -use expr::*; -use read::*; -use scalar::*; -use vortex::VortexSessionDefault; -use vortex::io::runtime::BlockingRuntime; -use vortex::io::runtime::current::CurrentThreadRuntime; -use vortex::io::session::RuntimeSessionExt; -use vortex::session::VortexSession; -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 = - LazyLock::new(CurrentThreadRuntime::new); -pub(crate) static SESSION: LazyLock = - LazyLock::new(|| VortexSession::default().with_handle(RUNTIME.handle())); - -#[cxx::bridge(namespace = "vortex::ffi")] -#[allow(let_underscore_drop)] -#[allow(clippy::absolute_paths)] -mod ffi { - extern "Rust" { - type DType; - // Factory functions for creating DType - fn dtype_null() -> Box; - fn dtype_bool(nullable: bool) -> Box; - fn dtype_primitive(ptype: PType, nullable: bool) -> Box; - fn dtype_decimal(precision: u8, scale: i8, nullable: bool) -> Box; - fn dtype_utf8(nullable: bool) -> Box; - fn dtype_binary(nullable: bool) -> Box; - unsafe fn from_arrow(ffi_schema: *mut u8, non_nullable: bool) -> Result>; - // Methods for DType - fn to_string(self: &DType) -> String; - - type Scalar; - fn bool_scalar_new(value: bool) -> Box; - fn i8_scalar_new(value: i8) -> Box; - fn i16_scalar_new(value: i16) -> Box; - fn i32_scalar_new(value: i32) -> Box; - fn i64_scalar_new(value: i64) -> Box; - fn u8_scalar_new(value: u8) -> Box; - fn u16_scalar_new(value: u16) -> Box; - fn u32_scalar_new(value: u32) -> Box; - fn u64_scalar_new(value: u64) -> Box; - fn f32_scalar_new(value: f32) -> Box; - fn f64_scalar_new(value: f64) -> Box; - fn string_scalar_new(value: &str) -> Box; - fn binary_scalar_new(value: &[u8]) -> Box; - fn cast_scalar(self: &Scalar, dtype: &DType) -> Result>; - - type Expr; - fn literal(scalar: Box) -> Box; - fn root() -> Box; - fn column(name: String) -> Box; - fn get_item(field: String, child: Box) -> Box; - fn not_(child: Box) -> Box; - fn is_null(child: Box) -> Box; - // binary op - fn eq(lhs: Box, rhs: Box) -> Box; - fn not_eq_(lhs: Box, rhs: Box) -> Box; - fn gt(lhs: Box, rhs: Box) -> Box; - fn gt_eq(lhs: Box, rhs: Box) -> Box; - fn lt(lhs: Box, rhs: Box) -> Box; - fn lt_eq(lhs: Box, rhs: Box) -> Box; - fn and_(lhs: Box, rhs: Box) -> Box; - fn or_(lhs: Box, rhs: Box) -> Box; - fn checked_add(lhs: Box, rhs: Box) -> Box; - fn select(fields: Vec, child: Box) -> Box; - - type VortexFile; - fn row_count(self: &VortexFile) -> u64; - fn scan_builder(self: &VortexFile) -> Result>; - fn open_file(path: &str) -> Result>; - fn open_file_from_buffer(data: &[u8]) -> Result>; - - type VortexScanBuilder; - fn with_filter(self: &mut VortexScanBuilder, filter: Box); - fn with_filter_ref(self: &mut VortexScanBuilder, filter: &Expr); - fn with_projection(self: &mut VortexScanBuilder, projection: Box); - fn with_projection_ref(self: &mut VortexScanBuilder, projection: &Expr); - fn with_row_range(self: &mut VortexScanBuilder, row_range_start: u64, row_range_end: u64); - fn with_include_by_index(self: &mut VortexScanBuilder, include_by_index: &[u64]); - fn with_limit(self: &mut VortexScanBuilder, limit: usize); - unsafe fn with_output_schema( - self: &mut VortexScanBuilder, - output_schema: *mut u8, - ) -> Result<()>; - unsafe fn scan_builder_into_stream( - builder: Box, - out_stream: *mut u8, - ) -> Result<()>; - fn scan_builder_into_threadsafe_cloneable_reader( - builder: Box, - ) -> Result>; - - type ThreadsafeCloneableReader; - unsafe fn clone_a_stream(self: &ThreadsafeCloneableReader, out_stream: *mut u8); - - type VortexWriteOptions; - fn write_options_new() -> Box; - unsafe fn write_array_stream( - options: Box, - input_stream: *mut u8, - path: &str, - ) -> Result<()>; - } - - #[repr(u8)] - #[derive(Debug, Clone, Copy)] - enum PType { - U8, - U16, - U32, - U64, - I8, - I16, - I32, - I64, - F16, - F32, - F64, - } -} diff --git a/vortex-cxx/src/read.rs b/vortex-cxx/src/read.rs deleted file mode 100644 index 4ce229b7a2b..00000000000 --- a/vortex-cxx/src/read.rs +++ /dev/null @@ -1,193 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::Arc; - -use anyhow::Result; -use arrow_array::RecordBatch; -use arrow_array::RecordBatchReader; -use arrow_array::cast::AsArray; -use arrow_array::ffi::FFI_ArrowSchema; -use arrow_array::ffi_stream::FFI_ArrowArrayStream; -use arrow_schema::ArrowError; -use arrow_schema::Field; -use arrow_schema::Schema; -use arrow_schema::SchemaRef; -use futures::stream::TryStreamExt; -use vortex::array::ArrayRef; -use vortex::array::LEGACY_SESSION; -use vortex::array::VortexSessionExecute; -use vortex::array::arrow::ArrowSessionExt; -use vortex::buffer::Buffer; -use vortex::file::OpenOptionsSessionExt; -use vortex::io::runtime::BlockingRuntime; -use vortex::layout::scan::arrow::RecordBatchIteratorAdapter; -use vortex::layout::scan::scan_builder::ScanBuilder; -use vortex::scan::selection::Selection; - -use crate::RUNTIME; -use crate::SESSION; -use crate::expr::Expr; - -pub(crate) struct VortexFile { - inner: vortex::file::VortexFile, -} - -impl VortexFile { - pub(crate) fn row_count(&self) -> u64 { - self.inner.row_count() - } - - pub(crate) fn scan_builder(&self) -> Result> { - Ok(Box::new(VortexScanBuilder { - inner: self.inner.scan()?, - output_schema: None, - })) - } -} - -/// File operations - using blocking operations for simplicity -/// TODO(xinyu): object store (see vortex-ffi) -pub(crate) fn open_file(path: &str) -> Result> { - let file = RUNTIME.block_on(SESSION.open_options().open_path(std::path::Path::new(path)))?; - Ok(Box::new(VortexFile { inner: file })) -} - -pub(crate) fn open_file_from_buffer(data: &[u8]) -> Result> { - let buffer = Buffer::from(data.to_vec()); - let file = SESSION.open_options().open_buffer(buffer)?; - Ok(Box::new(VortexFile { inner: file })) -} - -pub(crate) struct VortexScanBuilder { - inner: ScanBuilder, - output_schema: Option, -} - -impl VortexScanBuilder { - pub(crate) fn with_filter(&mut self, filter: Box) { - take_mut::take(&mut self.inner, |inner| inner.with_filter(filter.inner)); - } - - pub(crate) fn with_filter_ref(&mut self, filter: &Expr) { - take_mut::take(&mut self.inner, |inner| { - inner.with_filter(filter.inner.clone()) - }); - } - - pub(crate) fn with_projection(&mut self, filter: Box) { - take_mut::take(&mut self.inner, |inner| inner.with_projection(filter.inner)); - } - - pub(crate) fn with_projection_ref(&mut self, filter: &Expr) { - take_mut::take(&mut self.inner, |inner| { - inner.with_projection(filter.inner.clone()) - }); - } - - pub(crate) fn with_row_range(&mut self, row_range_start: u64, row_range_end: u64) { - take_mut::take(&mut self.inner, |inner| { - inner.with_row_range(row_range_start..row_range_end) - }); - } - - pub(crate) fn with_include_by_index(&mut self, include_by_index: &[u64]) { - let selection = Selection::IncludeByIndex(Buffer::copy_from(include_by_index)); - take_mut::take(&mut self.inner, |inner| inner.with_selection(selection)); - } - - pub(crate) fn with_limit(&mut self, limit: usize) { - take_mut::take(&mut self.inner, |inner| inner.with_limit(limit as u64)); - } - - pub(crate) unsafe fn with_output_schema(&mut self, output_schema: *mut u8) -> Result<()> { - let ffi_schema = - unsafe { FFI_ArrowSchema::from_raw(output_schema as *mut FFI_ArrowSchema) }; - self.output_schema = Some(Arc::new(Schema::try_from(&ffi_schema)?)); - Ok(()) - } -} - -/// # Safety -/// -/// out_stream should be properly aligned according to the Arrow C stream interface and valid for write. -pub(crate) unsafe fn scan_builder_into_stream( - builder: Box, - out_stream: *mut u8, -) -> Result<()> { - let schema = match builder.output_schema { - Some(schema) => schema, - None => { - let dtype = builder.inner.dtype()?; - let arrow_schema = dtype.to_arrow_schema()?; - Arc::new(arrow_schema) - } - }; - let reader = builder.inner.into_record_batch_reader(schema, &*RUNTIME)?; - let stream = FFI_ArrowArrayStream::new(Box::new(reader)); - let out_stream = out_stream as *mut FFI_ArrowArrayStream; - // # Safety - // Arrow C stream interface - unsafe { std::ptr::write(out_stream, stream) }; - Ok(()) -} - -trait ThreadsafeCloneableReaderTrait: RecordBatchReader + Send + 'static { - fn clone_boxed(&self) -> Box; -} - -impl ThreadsafeCloneableReaderTrait for T -where - T: RecordBatchReader + Send + Clone + 'static, -{ - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -pub(crate) struct ThreadsafeCloneableReader { - inner: Box, -} - -pub(crate) fn scan_builder_into_threadsafe_cloneable_reader( - builder: Box, -) -> Result, Box> { - let schema = match builder.output_schema { - Some(schema) => schema, - None => { - let dtype = builder.inner.dtype()?; - let arrow_schema = dtype.to_arrow_schema()?; - Arc::new(arrow_schema) - } - }; - let target = Field::new_struct("", schema.fields.clone(), false); - - let stream = builder - .inner - .map(move |b| { - SESSION - .arrow() - .execute_arrow(b, Some(&target), &mut LEGACY_SESSION.create_execution_ctx()) - .map(|struct_array| RecordBatch::from(struct_array.as_struct())) - }) - .into_stream()? - .map_err(|e| ArrowError::ExternalError(Box::new(e))); - - let iter = RUNTIME.block_on_stream_thread_safe(|_h| stream); - let rbr = RecordBatchIteratorAdapter::new(iter, schema); - - Ok(Box::new(ThreadsafeCloneableReader { - inner: Box::new(rbr), - })) -} - -impl ThreadsafeCloneableReader { - pub(crate) fn clone_a_stream(&self, out_stream: *mut u8) { - let cloned_reader = self.inner.clone_boxed(); - let stream = FFI_ArrowArrayStream::new(cloned_reader); - let out_stream = out_stream as *mut FFI_ArrowArrayStream; - // # Safety - // Arrow C stream interface - unsafe { std::ptr::write(out_stream, stream) }; - } -} diff --git a/vortex-cxx/src/scalar.rs b/vortex-cxx/src/scalar.rs deleted file mode 100644 index 24cd38be0f1..00000000000 --- a/vortex-cxx/src/scalar.rs +++ /dev/null @@ -1,52 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use anyhow::Result; - -use crate::dtype::DType; - -pub(crate) struct Scalar { - pub(crate) inner: vortex::scalar::Scalar, -} - -macro_rules! primitive_scalar_new { - ($name:ident, $type:ty) => { - pub(crate) fn $name(value: $type) -> Box { - Box::new(Scalar { - inner: vortex::scalar::Scalar::from(value), - }) - } - }; -} - -primitive_scalar_new!(bool_scalar_new, bool); // bool is not primitive but reuse the macro here -primitive_scalar_new!(i8_scalar_new, i8); -primitive_scalar_new!(i16_scalar_new, i16); -primitive_scalar_new!(i32_scalar_new, i32); -primitive_scalar_new!(i64_scalar_new, i64); -primitive_scalar_new!(u8_scalar_new, u8); -primitive_scalar_new!(u16_scalar_new, u16); -primitive_scalar_new!(u32_scalar_new, u32); -primitive_scalar_new!(u64_scalar_new, u64); -primitive_scalar_new!(f32_scalar_new, f32); -primitive_scalar_new!(f64_scalar_new, f64); - -pub(crate) fn string_scalar_new(value: &str) -> Box { - Box::new(Scalar { - inner: vortex::scalar::Scalar::from(value), - }) -} - -pub(crate) fn binary_scalar_new(value: &[u8]) -> Box { - Box::new(Scalar { - inner: vortex::scalar::Scalar::from(value), - }) -} - -impl Scalar { - pub(crate) fn cast_scalar(&self, dtype: &DType) -> Result> { - Ok(Box::new(Scalar { - inner: self.inner.cast(&dtype.inner)?, - })) - } -} diff --git a/vortex-cxx/src/write.rs b/vortex-cxx/src/write.rs deleted file mode 100644 index ff77c5a3dc5..00000000000 --- a/vortex-cxx/src/write.rs +++ /dev/null @@ -1,70 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use anyhow::Result; -use arrow_array::RecordBatchReader; -use arrow_array::ffi_stream::ArrowArrayStreamReader; -use arrow_array::ffi_stream::FFI_ArrowArrayStream; -use vortex::array::ArrayRef; -use vortex::array::arrow::FromArrowArray; -use vortex::array::iter::ArrayIteratorAdapter; -use vortex::array::iter::ArrayIteratorExt; -use vortex::array::stream::ArrayStream; -use vortex::dtype::DType; -use vortex::dtype::arrow::FromArrowType; -use vortex::error::VortexError; -use vortex::file::VortexWriteOptions as WriteOptions; -use vortex::file::WriteOptionsSessionExt; -use vortex::io::VortexWrite; -use vortex::io::runtime::BlockingRuntime; - -use crate::RUNTIME; -use crate::SESSION; - -pub(crate) struct VortexWriteOptions { - inner: WriteOptions, -} - -pub(crate) fn write_options_new() -> Box { - Box::new(VortexWriteOptions { - inner: SESSION.write_options(), - }) -} - -/// Convert an ArrowArrayStreamReader to a Vortex ArrayStream -fn arrow_stream_to_vortex_stream(reader: ArrowArrayStreamReader) -> Result { - let array_iter = ArrayIteratorAdapter::new( - DType::from_arrow(reader.schema()), - reader.map(|result| { - result - .map_err(VortexError::from) - .and_then(|record_batch| ArrayRef::from_arrow(record_batch, false)) - }), - ); - - Ok(array_iter.into_array_stream()) -} - -/// # Safety -/// -/// input_stream should be valid FFI_ArrowArrayStream. -/// See [`FFI_ArrowArrayStream::from_raw`] -pub(crate) unsafe fn write_array_stream( - options: Box, - input_stream: *mut u8, - path: &str, -) -> Result<()> { - let path = path.to_string(); - - let stream_reader = - unsafe { ArrowArrayStreamReader::from_raw(input_stream as *mut FFI_ArrowArrayStream) }?; - - let vortex_stream = arrow_stream_to_vortex_stream(stream_reader)?; - - RUNTIME.block_on(async { - let mut file = async_fs::File::create(path).await?; - options.inner.write(&mut file, vortex_stream).await?; - file.shutdown().await?; - Ok(()) - }) -} diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index 86f31d2dfce..0f204f71fdf 100644 --- a/vortex-datafusion/Cargo.toml +++ b/vortex-datafusion/Cargo.toml @@ -38,6 +38,7 @@ tokio-stream = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } uuid = { workspace = true, features = ["v7"] } vortex = { workspace = true, features = ["object_store", "tokio", "files"] } +vortex-arrow = { workspace = true } vortex-utils = { workspace = true, features = ["dashmap"] } [dev-dependencies] diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 8f3ff81af31..701b4c995fa 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -24,7 +24,6 @@ use datafusion_physical_expr_common::physical_expr::is_dynamic_physical_expr; use datafusion_physical_plan::expressions as df_expr; use itertools::Itertools; use vortex::VortexSessionDefault; -use vortex::array::arrow::ArrowSessionExt; use vortex::dtype::Nullability; use vortex::expr::Expression; use vortex::expr::and_collect; @@ -47,6 +46,7 @@ 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 vortex_arrow::ArrowSessionExt; use crate::convert::FromDataFusion; @@ -1243,8 +1243,8 @@ mod tests { use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::VortexSessionExecute as _; - use vortex::array::arrow::FromArrowArray; use vortex::session::VortexSession; + use vortex_arrow::FromArrowArray; // Create test data let values = Arc::new(Int32Array::from(vec![1, 5, 10, 15, 20])); diff --git a/vortex-datafusion/src/convert/scalars.rs b/vortex-datafusion/src/convert/scalars.rs index 86c1309df57..ec262b479d9 100644 --- a/vortex-datafusion/src/convert/scalars.rs +++ b/vortex-datafusion/src/convert/scalars.rs @@ -14,7 +14,6 @@ use vortex::dtype::DecimalDType; use vortex::dtype::NativeDecimalType; use vortex::dtype::Nullability; use vortex::dtype::PType; -use vortex::dtype::arrow::FromArrowType; use vortex::dtype::half::f16; use vortex::dtype::i256; use vortex::error::VortexExpect; @@ -27,6 +26,7 @@ use vortex::extension::datetime::TemporalMetadata; use vortex::extension::datetime::TimeUnit; use vortex::scalar::DecimalValue; use vortex::scalar::Scalar; +use vortex_arrow::FromArrowType; use crate::convert::FromDataFusion; use crate::convert::TryToDataFusion; @@ -382,6 +382,7 @@ mod tests { use vortex::dtype::Nullability; use vortex::dtype::PType; use vortex::dtype::StructFields; + use vortex::dtype::UnionVariants; use vortex::dtype::i256; use vortex::scalar::DecimalValue; use vortex::scalar::Scalar; @@ -826,7 +827,14 @@ mod tests { 2, Nullability::Nullable )))] - #[case::union(Scalar::null(DType::Union(Nullability::Nullable)))] + #[case::union(Scalar::null(DType::Union( + UnionVariants::new( + ["a"].into(), + vec![DType::Primitive(PType::I32, Nullability::Nullable)], + ) + .unwrap(), + Nullability::Nullable, + )))] fn unsupported_vortex_scalars_return_errors(#[case] scalar: Scalar) { let err = scalar.try_to_df().unwrap_err(); diff --git a/vortex-datafusion/src/convert/schema.rs b/vortex-datafusion/src/convert/schema.rs index ecfaa72d01e..2ec17758593 100644 --- a/vortex-datafusion/src/convert/schema.rs +++ b/vortex-datafusion/src/convert/schema.rs @@ -6,8 +6,8 @@ use arrow_schema::Field; use arrow_schema::Schema; use datafusion_common::Result as DFResult; use datafusion_common::exec_datafusion_err; -use vortex::array::arrow::ArrowSession; use vortex::dtype::DType; +use vortex_arrow::ArrowSession; /// Calculate the physical Arrow schema for a Vortex file given its DType and the expected logical schema. /// diff --git a/vortex-datafusion/src/lib.rs b/vortex-datafusion/src/lib.rs index 2858cbe908e..65bb021d258 100644 --- a/vortex-datafusion/src/lib.rs +++ b/vortex-datafusion/src/lib.rs @@ -143,11 +143,11 @@ mod common_tests { use url::Url; use vortex::VortexSessionDefault; use vortex::array::ArrayRef; - use vortex::array::arrow::FromArrowArray; use vortex::file::WriteOptionsSessionExt; use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::session::VortexSession; + use vortex_arrow::FromArrowArray; use crate::VortexFormatFactory; use crate::VortexTableOptions; diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 3e8f932b272..d936f27d366 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -46,7 +46,6 @@ use futures::stream; use object_store::ObjectMeta; use object_store::ObjectStore; use vortex::VortexSessionDefault; -use vortex::array::arrow::ArrowSessionExt; use vortex::array::memory::MemorySessionExt; use vortex::dtype::DType; use vortex::dtype::Nullability; @@ -65,6 +64,7 @@ use vortex::io::session::RuntimeSessionExt; use vortex::scalar::Scalar; use vortex::scalar::ScalarValue as VortexScalarValue; use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; use super::cache::CachedVortexMetadata; use super::sink::VortexSink; diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 9f062b9bdfe..823a47af1ff 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -37,7 +37,6 @@ use itertools::Itertools; use object_store::path::Path; use tracing::Instrument; use vortex::array::VortexSessionExecute; -use vortex::array::arrow::ArrowSessionExt; use vortex::dtype::FieldMask; use vortex::error::VortexError; use vortex::error::VortexExpect; @@ -49,6 +48,7 @@ use vortex::layout::scan::split_by::SplitBy; use vortex::metrics::Label; use vortex::metrics::MetricsRegistry; use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; use vortex_utils::aliases::dash_map::DashMap; use vortex_utils::aliases::dash_map::Entry; @@ -596,7 +596,6 @@ mod tests { use rstest::rstest; use vortex::VortexSessionDefault; use vortex::array::ArrayRef; - use vortex::array::arrow::FromArrowArray; use vortex::buffer::Buffer; use vortex::file::WriteOptionsSessionExt; use vortex::io::VortexWrite; @@ -604,6 +603,7 @@ mod tests { use vortex::metrics::DefaultMetricsRegistry; use vortex::scan::selection::Selection; use vortex::session::VortexSession; + use vortex_arrow::FromArrowArray; use super::*; use crate::VortexAccessPlan; diff --git a/vortex-datafusion/src/persistent/sink.rs b/vortex-datafusion/src/persistent/sink.rs index 6d8d9edbe76..32be0c6873b 100644 --- a/vortex-datafusion/src/persistent/sink.rs +++ b/vortex-datafusion/src/persistent/sink.rs @@ -3,17 +3,13 @@ use std::sync::Arc; -use arrow_schema::Schema; use arrow_schema::SchemaRef; use async_trait::async_trait; use datafusion_common::DataFusionError; use datafusion_common::Result as DFResult; -use datafusion_common::arrow::array::RecordBatch; -use datafusion_common::arrow::array::RecordBatchOptions; use datafusion_common::exec_datafusion_err; use datafusion_common_runtime::JoinSet; use datafusion_common_runtime::SpawnedTask; -use datafusion_datasource::ListingTableUrl; use datafusion_datasource::display::FileGroupDisplay; use datafusion_datasource::file_sink_config::FileSink; use datafusion_datasource::file_sink_config::FileSinkConfig; @@ -25,132 +21,34 @@ use datafusion_execution::TaskContext; use datafusion_physical_plan::DisplayAs; use datafusion_physical_plan::DisplayFormatType; use datafusion_physical_plan::metrics::MetricsSet; -use futures::SinkExt; -use futures::Stream; use futures::StreamExt; -use futures::channel::mpsc; use object_store::ObjectStore; -use object_store::ObjectStoreExt; use object_store::path::Path; -use tokio::task::JoinHandle; use tokio_stream::wrappers::ReceiverStream; -use uuid::Uuid; -use vortex::array::arrow::ArrowSessionExt; use vortex::array::stream::ArrayStreamAdapter; -use vortex::dtype::DType; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteSummary; use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::session::VortexSession; -use vortex_utils::aliases::hash_set::HashSet; - -struct WriteOutputOptions<'a> { - base_output_path: &'a ListingTableUrl, - target_file_size: Option, - extension: &'a str, - write_id: &'a str, - partition_column_names: &'a [String], - keep_partition_by_columns: bool, - single_file_output: bool, -} - -#[derive(Clone, Copy)] -struct CompressionEstimate { - prev_compressed_bytes: u64, - prev_uncompressed_bytes: u64, -} - -impl CompressionEstimate { - fn identity() -> Self { - Self { - prev_compressed_bytes: 1, - prev_uncompressed_bytes: 1, - } - } - - fn from_file_sizes(compressed_bytes: u64, uncompressed_bytes: u64) -> DFResult { - if uncompressed_bytes == 0 { - return Err(exec_datafusion_err!( - "Cannot derive compression estimate from zero uncompressed bytes" - )); - } - - Ok(Self { - prev_compressed_bytes: compressed_bytes, - prev_uncompressed_bytes: uncompressed_bytes, - }) - } - - fn estimate_compressed_size(self, uncompressed_bytes: u64) -> DFResult { - if self.prev_uncompressed_bytes == 0 { - return Err(exec_datafusion_err!( - "Compression estimate denominator must be non-zero" - )); - } - - let estimated = u128::from(uncompressed_bytes) - .checked_mul(u128::from(self.prev_compressed_bytes)) - .ok_or_else(|| { - exec_datafusion_err!( - "Compressed size estimate overflow for {} * {}", - uncompressed_bytes, - self.prev_compressed_bytes - ) - })? - / u128::from(self.prev_uncompressed_bytes); - - u64::try_from(estimated).map_err(|_| { - exec_datafusion_err!("Compressed size estimate does not fit in u64: {estimated}") - }) - } -} - -struct ActiveFileWriter { - path: Path, - sender: mpsc::Sender, - task: JoinHandle>, -} +use vortex_arrow::ArrowSessionExt; /// Implements [`DataSink`] for writing Vortex files. pub struct VortexSink { config: FileSinkConfig, schema: SchemaRef, session: VortexSession, - target_file_size: Option, } impl VortexSink { /// Creates a new [`VortexSink`] instance. - pub fn new( - config: FileSinkConfig, - schema: SchemaRef, - session: VortexSession, - target_file_size: Option, - ) -> Self { + pub fn new(config: FileSinkConfig, schema: SchemaRef, session: VortexSession) -> Self { Self { config, schema, session, - target_file_size, } } - - fn base_output_path(&self) -> DFResult<&ListingTableUrl> { - self.config - .table_paths - .first() - .ok_or_else(|| exec_datafusion_err!("Vortex sink requires at least one table path")) - } - - fn writer_dtype(&self, writer_schema: &SchemaRef) -> DFResult { - self.session - .arrow() - .from_arrow_schema(writer_schema) - .map_err(|e| { - exec_datafusion_err!("Failed to derive Vortex DType from writer schema: {e}") - }) - } } impl std::fmt::Debug for VortexSink { @@ -190,46 +88,7 @@ impl DataSink for VortexSink { data: SendableRecordBatchStream, context: &Arc, ) -> DFResult { - if !self.config.table_partition_cols.is_empty() { - return FileSink::write_all(self, data, context).await; - } - - let object_store = context - .runtime_env() - .object_store(&self.config.object_store_url)?; - let writer_schema = get_writer_schema(&self.config); - let dtype = self.writer_dtype(&writer_schema)?; - let write_id = Uuid::now_v7().simple().to_string(); - let base_output_path = self.base_output_path()?; - let partition_column_names = self - .config - .table_partition_cols - .iter() - .map(|(name, _)| name.clone()) - .collect::>(); - - let summaries = write_record_batch_stream_to_files( - self.session.clone(), - object_store, - dtype, - data, - Arc::clone(&writer_schema), - &WriteOutputOptions { - base_output_path, - target_file_size: self.target_file_size, - extension: &self.config.file_extension, - write_id: &write_id, - partition_column_names: &partition_column_names, - keep_partition_by_columns: self.config.keep_partition_by_columns, - single_file_output: self - .config - .file_output_mode - .single_file_output(base_output_path), - }, - ) - .await?; - - aggregate_write_summaries(summaries) + FileSink::write_all(self, data, context).await } } @@ -246,50 +105,65 @@ impl FileSink for VortexSink { mut file_stream_rx: DemuxedStreamReceiver, object_store: Arc, ) -> DFResult { + let mut file_write_tasks: JoinSet> = JoinSet::new(); let writer_schema = get_writer_schema(&self.config); - let dtype = self.writer_dtype(&writer_schema)?; - let mut file_write_tasks: JoinSet>> = JoinSet::new(); + let dtype = self + .session + .arrow() + .from_arrow_schema(&writer_schema) + .map_err(|e| { + exec_datafusion_err!("Failed to derive Vortex DType from writer schema: {e}") + })?; + // TODO(adamg): + // 1. We can probably be better at signaling how much memory we're consuming (potentially when reading too), see ParquetSink::spawn_writer_tasks_and_join. while let Some((path, rx)) = file_stream_rx.recv().await { let session = self.session.clone(); let object_store = Arc::clone(&object_store); - let dtype = dtype.clone(); - let import_schema = Arc::clone(&writer_schema); - let target_file_size = self.target_file_size; - let extension = self.config.file_extension.clone(); + // We need to spawn work because there's a dependency between the different files. If one file has too many batches buffered, + // the demux task might deadlock itself. + let arrow_session = session.clone(); + let import_schema = Arc::clone(&writer_schema); + let dtype = dtype.clone(); file_write_tasks.spawn(async move { - let stream = ReceiverStream::new(rx).map(Ok); - write_record_batch_stream_to_paths( - session, - object_store, - dtype, - stream, - import_schema, - target_file_size, - |file_index| { - if file_index == 0 { - path.clone() - } else { - numbered_path(&path, file_index, &extension) - } - }, - ) - .await + let stream = ReceiverStream::new(rx).map(move |rb| { + arrow_session + .arrow() + .from_arrow_record_batch(rb, &import_schema) + }); + + let stream_adapter = ArrayStreamAdapter::new(dtype, stream); + + let mut object_writer = ObjectStoreWrite::new(object_store, &path) + .await + .map_err(|e| exec_datafusion_err!("Failed to create ObjectStoreWrite: {e}"))?; + + let summary = session + .write_options() + .write(&mut object_writer, stream_adapter) + .await + .map_err(|e| exec_datafusion_err!("Failed to write Vortex file: {e}"))?; + + object_writer + .shutdown() + .await + .map_err(|e| exec_datafusion_err!("Failed to shutdown Vortex writer: {e}"))?; + + Ok((path, summary)) }); } - let mut row_count = 0_u64; + let mut row_count = 0; + while let Some(result) = file_write_tasks.join_next().await { match result { - Ok(summaries) => { - row_count = row_count - .checked_add(aggregate_write_summaries(summaries?)?) - .ok_or_else(|| { - exec_datafusion_err!( - "Row count overflow while aggregating file writer tasks" - ) - })?; + Ok(r) => { + let (path, summary) = r?; + + row_count += summary.row_count(); + + tracing::info!(path = %path, "Successfully written file"); } Err(e) => { if e.is_panic() { @@ -310,328 +184,6 @@ impl FileSink for VortexSink { } } -fn aggregate_write_summaries(summaries: Vec<(Path, WriteSummary)>) -> DFResult { - let mut row_count = 0_u64; - for (path, summary) in summaries { - row_count = row_count.checked_add(summary.row_count()).ok_or_else(|| { - exec_datafusion_err!( - "Row count overflow while aggregating sink summaries (current={}, file={})", - row_count, - summary.row_count() - ) - })?; - tracing::debug!(path = %path, "Successfully written file"); - } - - Ok(row_count) -} - -/// Write batches from a single input stream to one or more output files. -/// -/// For collection paths, files are emitted using the `{write_id}_{file_index:05}.{extension}` -/// naming scheme as produced by the underlying writer implementation. -/// For a single-file path, the original target path is used unless rolling is needed, -/// in which case additional files follow the same naming scheme. -async fn write_record_batch_stream_to_files( - session: VortexSession, - object_store: Arc, - dtype: DType, - data: SendableRecordBatchStream, - import_schema: SchemaRef, - output_options: &WriteOutputOptions<'_>, -) -> DFResult> { - let stream = data.map(|batch| { - let batch = batch?; - if output_options.keep_partition_by_columns - || output_options.partition_column_names.is_empty() - { - Ok(batch) - } else { - remove_partition_columns(&batch, output_options.partition_column_names) - } - }); - - write_record_batch_stream_to_paths( - session, - object_store, - dtype, - stream, - import_schema, - output_options.target_file_size, - |file_index| { - output_file_path( - output_options.base_output_path, - file_index, - output_options.extension, - output_options.single_file_output, - output_options.write_id, - ) - }, - ) - .await -} - -async fn write_record_batch_stream_to_paths( - session: VortexSession, - object_store: Arc, - dtype: DType, - data: S, - import_schema: SchemaRef, - target_file_size: Option, - mut output_path: F, -) -> DFResult> -where - S: Stream>, - F: FnMut(usize) -> Path, -{ - let target = target_file_size.map(|t| t.max(1)); - let mut results: Vec<(Path, WriteSummary)> = Vec::new(); - let mut active_writer: Option = None; - let mut uncompressed_bytes_in_file = 0_u64; - let mut file_index = 0_usize; - let mut compression_estimate = CompressionEstimate::identity(); - - futures::pin_mut!(data); - - let write_result: DFResult<()> = async { - while let Some(batch) = data.next().await.transpose()? { - if active_writer.is_none() { - active_writer = Some(start_file_writer( - &session, - Arc::clone(&object_store), - output_path(file_index), - dtype.clone(), - Arc::clone(&import_schema), - )); - } - - let batch_bytes = batch_uncompressed_bytes(&batch)?; - let writer = active_writer.as_mut().ok_or_else(|| { - exec_datafusion_err!("Missing active file writer for sink output") - })?; - send_batch_to_active_writer(writer, batch).await?; - uncompressed_bytes_in_file = uncompressed_bytes_in_file - .checked_add(batch_bytes) - .ok_or_else(|| { - exec_datafusion_err!( - "Uncompressed byte counter overflow for output file {}", - writer.path - ) - })?; - - if let Some(target) = target { - let estimated_compressed = - compression_estimate.estimate_compressed_size(uncompressed_bytes_in_file)?; - if estimated_compressed >= target { - let writer = active_writer.take().ok_or_else(|| { - exec_datafusion_err!( - "Missing active file writer while finalizing rotated output file" - ) - })?; - let file_path = writer.path.clone(); - let summary = finish_file_writer(writer).await?; - if uncompressed_bytes_in_file > 0 { - compression_estimate = CompressionEstimate::from_file_sizes( - summary.size(), - uncompressed_bytes_in_file, - )?; - } - - results.push((file_path, summary)); - uncompressed_bytes_in_file = 0; - file_index += 1; - } - } - } - - if let Some(writer) = active_writer.take() { - let file_path = writer.path.clone(); - let summary = finish_file_writer(writer).await?; - results.push((file_path, summary)); - } - - Ok(()) - } - .await; - - if let Err(err) = write_result { - cleanup_failed_write( - object_store, - active_writer.into_iter().collect(), - results.iter().map(|(path, _)| path.clone()).collect(), - ) - .await; - return Err(err); - } - - Ok(results) -} - -/// Generate a numbered file path from an existing path for size-based splitting. -/// -/// Given `base/file.vortex`, produces `base/file_00000.vortex`. -/// If the path has no recognized extension, appends `_00000.{extension}`. -fn numbered_path(original: &Path, index: usize, extension: &str) -> Path { - let s = original.to_string(); - let suffix = format!(".{extension}"); - if let Some(stem) = s.strip_suffix(&suffix) { - Path::from(format!("{stem}_{index:05}{suffix}")) - } else { - Path::from(format!("{s}_{index:05}.{extension}")) - } -} - -fn start_file_writer( - session: &VortexSession, - object_store: Arc, - path: Path, - dtype: DType, - import_schema: SchemaRef, -) -> ActiveFileWriter { - let (sender, receiver) = mpsc::channel::(1); - let session = session.clone(); - let path_for_task = path.clone(); - - let task = tokio::spawn(async move { - let mut object_writer = ObjectStoreWrite::new(object_store, &path_for_task) - .await - .map_err(|e| { - exec_datafusion_err!( - "Failed to create ObjectStoreWrite for '{}': {e}", - path_for_task - ) - })?; - - let arrow_session = session.clone(); - let stream = receiver.map(move |rb| { - arrow_session - .arrow() - .from_arrow_record_batch(rb, &import_schema) - }); - let stream_adapter = ArrayStreamAdapter::new(dtype, stream); - - let summary = session - .write_options() - .write(&mut object_writer, stream_adapter) - .await - .map_err(|e| { - exec_datafusion_err!("Failed to write Vortex file '{}': {e}", path_for_task) - })?; - - object_writer.shutdown().await.map_err(|e| { - exec_datafusion_err!("Failed to shutdown Vortex writer '{}': {e}", path_for_task) - })?; - - Ok(summary) - }); - - ActiveFileWriter { path, sender, task } -} - -async fn cleanup_failed_write( - object_store: Arc, - active_writers: Vec, - finished_paths: Vec, -) { - let mut cleanup_paths = HashSet::new(); - - for writer in active_writers { - cleanup_paths.insert(writer.path.clone()); - writer.task.abort(); - drop(writer.task.await); - } - - for path in finished_paths { - cleanup_paths.insert(path); - } - - for path in cleanup_paths { - if let Err(e) = object_store.delete(&path).await { - tracing::warn!(path = %path, error = %e, "Failed to delete sink output during error cleanup"); - } - } -} - -async fn send_batch_to_active_writer( - writer: &mut ActiveFileWriter, - batch: RecordBatch, -) -> DFResult<()> { - writer.sender.send(batch).await.map_err(|e| { - exec_datafusion_err!( - "Failed to send batch to active writer for '{}': {e}", - writer.path - ) - }) -} - -async fn finish_file_writer(mut writer: ActiveFileWriter) -> DFResult { - writer.sender.close_channel(); - match writer.task.await { - Ok(result) => result, - Err(e) => Err(exec_datafusion_err!( - "Vortex writer task for '{}' failed to join: {e}", - writer.path - )), - } -} - -fn batch_uncompressed_bytes(batch: &RecordBatch) -> DFResult { - u64::try_from(batch.get_array_memory_size()).map_err(|_| { - exec_datafusion_err!( - "RecordBatch memory size does not fit in u64: {}", - batch.get_array_memory_size() - ) - }) -} - -fn remove_partition_columns( - batch: &RecordBatch, - partition_column_names: &[String], -) -> DFResult { - let partition_names: Vec<_> = partition_column_names.iter().map(String::as_str).collect(); - let (columns, fields): (Vec<_>, Vec<_>) = batch - .columns() - .iter() - .zip(batch.schema().fields()) - .filter(|(_, field)| !partition_names.contains(&field.name().as_str())) - .map(|(array, field)| (Arc::clone(array), (**field).clone())) - .unzip(); - - let schema = Schema::new(fields); - if columns.is_empty() { - let options = RecordBatchOptions::default().with_row_count(Some(batch.num_rows())); - return Ok(RecordBatch::try_new_with_options( - Arc::new(schema), - columns, - &options, - )?); - } - - Ok(RecordBatch::try_new(Arc::new(schema), columns)?) -} - -/// Build the output path for a rolling write. -fn output_file_path( - base_output_path: &ListingTableUrl, - file_index: usize, - extension: &str, - single_file_output: bool, - write_id: &str, -) -> Path { - if single_file_output { - if file_index == 0 { - return base_output_path.prefix().clone(); - } - return numbered_path(base_output_path.prefix(), file_index, extension); - } - - let mut base = base_output_path.prefix().to_string(); - if !base.ends_with('/') { - base.push('/'); - } - Path::from(format!("{base}{write_id}_{file_index:05}.{extension}")) -} - #[cfg(test)] mod tests { use std::sync::Arc; @@ -650,8 +202,6 @@ mod tests { use datafusion::logical_expr::Values; use datafusion::logical_expr::dml::InsertOp; use datafusion_common::ScalarValue; - use datafusion_common::exec_datafusion_err; - use datafusion_datasource::ListingTableUrl; use datafusion_datasource::TableSchema; use datafusion_datasource::file_format::format_as_file_type; use datafusion_datasource::file_groups::FileGroup; @@ -664,26 +214,12 @@ mod tests { use datafusion_physical_plan::empty::EmptyExec; use futures::TryStreamExt; use rstest::rstest; - use tokio::sync::oneshot; - use tokio::time::Duration; use vortex::file::VORTEX_FILE_EXTENSION; use vortex::session::VortexSession; use super::*; use crate::common_tests::TestSessionContext; use crate::persistent::VortexFormatFactory; - use crate::persistent::VortexTableOptions; - use crate::persistent::sink::ActiveFileWriter; - use crate::persistent::sink::finish_file_writer; - - fn split_path(base_path: &Path, file_index: usize, extension: &str) -> Path { - let mut base = base_path.to_string(); - if !base.ends_with('/') { - base.push('/'); - } - let filename = format!("part-{file_index:05}.{extension}"); - Path::from(format!("{base}{filename}")) - } #[tokio::test] async fn test_insert_into_sql() -> anyhow::Result<()> { @@ -783,11 +319,10 @@ mod tests { } /// Reproduction by . - /// Uses a 1MB target file size to exercise file splitting behavior. #[rstest] - #[case(1_000, 1)] - #[case(5_000_000, 6)] - #[case(10_000_000, 10)] + #[case(1000, 1)] + #[case(40_961, 4)] + #[case(1_000_000, 4)] #[tokio::test] async fn test_write_large_batch( #[case] entries: usize, @@ -795,26 +330,15 @@ mod tests { ) -> anyhow::Result<()> { let ctx = TestSessionContext::default(); - let opts = VortexTableOptions { - target_file_size_mb: 1, - ..Default::default() - }; - - let factory = VortexFormatFactory::new().with_options(opts); - - let values: Vec = (0..entries) - .map(|i| i8::try_from(i % 127)) - .collect::>()?; - let data = ctx.session.read_batch(RecordBatch::try_new( Arc::new(Schema::new(vec![Field::new("a", DataType::Int8, false)])), - vec![Arc::new(Int8Array::from(values))], + vec![Arc::new(Int8Array::from(vec![0i8; entries]))], )?)?; let logical_plan = LogicalPlanBuilder::copy_to( data.logical_plan().clone(), "/table/".to_string(), - format_as_file_type(Arc::new(factory)), + format_as_file_type(Arc::new(VortexFormatFactory::new())), Default::default(), vec![], )? @@ -850,57 +374,50 @@ mod tests { entries, count_value ); - let file_metas = ctx - .store - .list(Some(&"table".into())) - .try_collect::>() - .await?; - - assert!( - file_metas.len() >= expected_files, - "Expected at least {expected_files} files for {entries} values, got {}", - file_metas.len() - ); - - Ok(()) - } - - #[tokio::test] - async fn test_write_large_batch_default_target_is_128mb() -> anyhow::Result<()> { - let ctx = TestSessionContext::default(); - - let entries = 1_000_000; - let values: Vec = (0..entries) - .map(|i| i8::try_from(i % 127)) - .collect::>()?; - - let data = ctx.session.read_batch(RecordBatch::try_new( - Arc::new(Schema::new(vec![Field::new("a", DataType::Int8, false)])), - vec![Arc::new(Int8Array::from(values))], - )?)?; - - let logical_plan = LogicalPlanBuilder::copy_to( - data.logical_plan().clone(), - "/table/".to_string(), - format_as_file_type(Arc::new(VortexFormatFactory::new())), - Default::default(), - vec![], - )? - .build()?; - - ctx.session - .execute_logical_plan(logical_plan) + let all_data = ctx + .session + .sql("SELECT a FROM '/table/'") .await? .collect() .await?; + let mut total_rows = 0; + for batch in all_data { + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + + for i in 0..batch.num_rows() { + assert_eq!( + col.value(i), + 0i8, + "Expected value 0 at row {}, but found {}", + total_rows + i, + col.value(i) + ); + } + total_rows += batch.num_rows(); + } + + assert_eq!( + total_rows, entries, + "Total rows read ({}) doesn't match expected entries ({})", + total_rows, entries + ); + let file_metas = ctx .store .list(Some(&"/table".into())) .try_collect::>() .await?; - assert_eq!(file_metas.len(), 1); + assert_eq!( + file_metas.len(), + expected_files, + "Expected {expected_files} files for {entries} values" + ); Ok(()) } @@ -947,1023 +464,6 @@ mod tests { Ok(()) } - #[test] - fn test_split_path_basic() { - let path = Path::from("data/output"); - assert_eq!( - split_path(&path, 0, "vortex").to_string(), - "data/output/part-00000.vortex" - ); - assert_eq!( - split_path(&path, 12, "vortex").to_string(), - "data/output/part-00012.vortex" - ); - } - - #[test] - fn test_split_path_preserves_trailing_slash() { - let path = Path::from("nested/path/"); - assert_eq!( - split_path(&path, 3, "vx").to_string(), - "nested/path/part-00003.vx" - ); - } - - #[test] - fn test_numbered_path() { - use super::numbered_path; - - let path = Path::from("table/c1=alpha/abc123.vortex"); - assert_eq!( - numbered_path(&path, 0, "vortex").to_string(), - "table/c1=alpha/abc123_00000.vortex" - ); - assert_eq!( - numbered_path(&path, 5, "vortex").to_string(), - "table/c1=alpha/abc123_00005.vortex" - ); - } - - #[test] - fn test_numbered_path_no_extension() { - use super::numbered_path; - - let path = Path::from("table/output"); - assert_eq!( - numbered_path(&path, 0, "vortex").to_string(), - "table/output_00000.vortex" - ); - } - - #[test] - fn test_output_file_path_single_file_and_collection() { - use super::output_file_path; - - let single = ListingTableUrl::parse("file:///tmp/output.vortex").unwrap(); - assert_eq!( - output_file_path(&single, 0, "vortex", true, "wid").to_string(), - "tmp/output.vortex" - ); - assert_eq!( - output_file_path(&single, 2, "vortex", true, "wid").to_string(), - "tmp/output_00002.vortex" - ); - - let collection = ListingTableUrl::parse("file:///tmp/table/").unwrap(); - assert_eq!( - output_file_path(&collection, 3, "vortex", false, "wid").to_string(), - "tmp/table/wid_00003.vortex" - ); - } - - #[tokio::test] - async fn test_finish_file_writer_waits_for_task_completion() -> anyhow::Result<()> { - let (sender, receiver) = mpsc::channel::(1); - drop(receiver); - - let (gate_tx, gate_rx) = oneshot::channel::<()>(); - - let writer = ActiveFileWriter { - path: Path::from("table/pending.vortex"), - sender, - task: tokio::spawn(async move { - let _ = gate_rx.await; - Err(exec_datafusion_err!( - "synthetic writer failure after completion gate" - )) - }), - }; - - let mut finish_fut = Box::pin(finish_file_writer(writer)); - assert!( - tokio::time::timeout(Duration::from_millis(50), &mut finish_fut) - .await - .is_err(), - "finish_file_writer returned before writer task completed" - ); - - gate_tx - .send(()) - .map_err(|_| anyhow::anyhow!("failed to release writer completion gate in test"))?; - - let err = match finish_fut.await { - Ok(_) => { - return Err(anyhow::anyhow!( - "finish_file_writer unexpectedly succeeded after gate release" - )); - } - Err(err) => err, - }; - assert!( - err.to_string().contains("synthetic writer failure"), - "unexpected error: {err}" - ); - - Ok(()) - } - - #[test] - fn test_remove_partition_columns() -> anyhow::Result<()> { - use datafusion::arrow::array::StringArray; - - use super::remove_partition_columns; - - let batch = RecordBatch::try_new( - Arc::new(Schema::new(vec![ - Field::new("part", DataType::Utf8, false), - Field::new("val", DataType::Int64, false), - ])), - vec![ - Arc::new(StringArray::from(vec!["x", "y"])), - Arc::new(Int64Array::from(vec![1, 2])), - ], - )?; - - let out = remove_partition_columns(&batch, &["part".to_string()])?; - assert_eq!(out.num_columns(), 1); - assert_eq!(out.schema().field(0).name(), "val"); - assert_eq!(out.num_rows(), 2); - - Ok(()) - } - - #[test] - fn test_remove_partition_columns_all_columns_partitioned() -> anyhow::Result<()> { - use datafusion::arrow::array::StringArray; - - use super::remove_partition_columns; - - let batch = RecordBatch::try_new( - Arc::new(Schema::new(vec![Field::new("part", DataType::Utf8, false)])), - vec![Arc::new(StringArray::from(vec!["x", "y", "z"]))], - )?; - - let out = remove_partition_columns(&batch, &["part".to_string()])?; - assert_eq!(out.num_columns(), 0); - assert_eq!(out.num_rows(), 3); - - Ok(()) - } - - /// Generate `count` pseudo-random i64 values using a simple LCG. - /// These values resist compression (unlike sequential or modular data), - /// giving more realistic compressed file sizes. - fn pseudo_random_i64s(count: usize, seed: i64) -> Vec { - let mut v = seed; - (0..count) - .map(|_| { - v = v - .wrapping_mul(6_364_136_223_846_793_005) - .wrapping_add(1_442_695_040_888_963_407); - v - }) - .collect() - } - - /// Tests file splitting through the full DataFusion pipeline. - /// - /// Writes ~62MB of pseudo-random Int64 data (near 1:1 compression ratio) - /// via COPY TO with a 16MB target file size. Verifies that exactly 4 files - /// are produced and each file's compressed size is approximately 16MB. - /// - /// This exercises the complete COPY TO write path through DataFusion and - /// VortexSink, unlike a direct `write_stream_to_files` call. - #[tokio::test] - async fn test_file_splitting_62mb_into_4_files() -> anyhow::Result<()> { - use datafusion::datasource::MemTable; - use datafusion_datasource::file_format::format_as_file_type; - - let ctx = TestSessionContext::default(); - - let target_mb = 16_usize; - let opts = VortexTableOptions { - target_file_size_mb: target_mb, - ..Default::default() - }; - let factory = VortexFormatFactory::new().with_options(opts); - - let batch_rows = 8192_usize; - let total_elements = 62 * 1024 * 1024 / 8; // ~8,126,464 i64 values ≈ 62MB Arrow memory - let num_batches = total_elements / batch_rows; - let expected_total_rows = (num_batches * batch_rows) as i64; - - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut batches = Vec::new(); - for i in 0..num_batches { - let values = pseudo_random_i64s(batch_rows, (i * batch_rows) as i64); - batches.push(RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?); - } - - let table = MemTable::try_new(Arc::clone(&schema), vec![batches])?; - ctx.session.register_table("source", Arc::new(table))?; - - let source = ctx.session.table("source").await?; - let logical_plan = LogicalPlanBuilder::copy_to( - source.logical_plan().clone(), - "/table/".to_string(), - format_as_file_type(Arc::new(factory)), - Default::default(), - vec![], - )? - .build()?; - - ctx.session - .execute_logical_plan(logical_plan) - .await? - .collect() - .await?; - - let file_metas = ctx - .store - .list(Some(&"/table".into())) - .try_collect::>() - .await?; - - assert_eq!( - file_metas.len(), - 4, - "Expected 4 files for ~62MB data with {target_mb}MB target, got {} (sizes: {:?})", - file_metas.len(), - file_metas.iter().map(|m| m.size).collect::>() - ); - - let target_bytes = u64::try_from(target_mb * 1024 * 1024)?; - for meta in &file_metas { - assert!( - meta.size > target_bytes / 2, - "File {} is {}B, expected at least {}B (target/2)", - meta.location, - meta.size, - target_bytes / 2 - ); - } - - // Verify total row count. - let result = ctx - .session - .sql("SELECT COUNT(*) as cnt FROM '/table/'") - .await? - .collect() - .await?; - - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - /// Tests file splitting with compressible data through the full pipeline. - /// - /// Uses low-entropy Int64 values (repeating 0..255) which compress ~8:1 in - /// Vortex. With the current code that compares Arrow memory size against - /// `target_file_size`, files are split far too early, producing many tiny - /// compressed files instead of files that are close to the target. - /// - /// For ~32MB of Arrow data (~4MB compressed at 8:1) with a 1MB target: - /// - **Correct**: 4 files of ~1MB compressed each - /// - **Bug**: 32 files of ~0.125MB compressed each - #[tokio::test] - async fn test_file_splitting_compressible_data() -> anyhow::Result<()> { - use datafusion::datasource::MemTable; - use datafusion_datasource::file_format::format_as_file_type; - - let ctx = TestSessionContext::default(); - - let target_mb = 1_usize; - let opts = VortexTableOptions { - target_file_size_mb: target_mb, - ..Default::default() - }; - let factory = VortexFormatFactory::new().with_options(opts); - - // Generate low-entropy Int64 values: repeating 0..255. - // Arrow memory: 4M × 8 bytes = 32MB. - // Vortex compressed: each value only needs ~1 byte → ~4MB total. - let total_elements = 4_000_000_usize; - let batch_rows = 8192_usize; - let num_batches = total_elements / batch_rows; - - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut batches = Vec::new(); - for i in 0..num_batches { - let values: Vec = (0..batch_rows) - .map(|j| ((i * batch_rows + j) % 256) as i64) - .collect(); - batches.push(RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?); - } - - let table = MemTable::try_new(Arc::clone(&schema), vec![batches])?; - ctx.session.register_table("source", Arc::new(table))?; - - let source = ctx.session.table("source").await?; - let logical_plan = LogicalPlanBuilder::copy_to( - source.logical_plan().clone(), - "/table/".to_string(), - format_as_file_type(Arc::new(factory)), - Default::default(), - vec![], - )? - .build()?; - - ctx.session - .execute_logical_plan(logical_plan) - .await? - .collect() - .await?; - - let file_metas = ctx - .store - .list(Some(&"/table".into())) - .try_collect::>() - .await?; - - // With compressible data, there should be few files (not > 10). - // The buggy code produces many tiny files because it splits on Arrow - // memory (32MB / 1MB = 32 files) instead of compressed size (~4MB / 1MB = 4 files). - let total_compressed: u64 = file_metas.iter().map(|m| m.size).sum(); - let target_bytes = u64::try_from(target_mb * 1024 * 1024)?; - - // We should have at most ~(total_compressed / target) + 1 files, not - // ~(arrow_memory / target) files. - let max_expected = usize::try_from(total_compressed / target_bytes + 2)?; - assert!( - file_metas.len() <= max_expected, - "Too many files: got {} but total compressed is {}B with {}B target \ - (expected at most {max_expected}). Files are being split on Arrow memory \ - instead of compressed size. Sizes: {:?}", - file_metas.len(), - total_compressed, - target_bytes, - file_metas.iter().map(|m| m.size).collect::>() - ); - - // Every file except the first should be reasonably sized. The first - // file may be smaller because the compression ratio is unknown until - // the first write completes. - for meta in file_metas.iter().skip(1) { - assert!( - meta.size > target_bytes / 4, - "File {} is {}B, far below target {}B — splitting on Arrow memory, not compressed size", - meta.location, - meta.size, - target_bytes - ); - } - - Ok(()) - } - - #[tokio::test] - async fn test_write_large_batch_target_file_size_disabled() -> anyhow::Result<()> { - use datafusion::datasource::MemTable; - use datafusion_datasource::file_format::format_as_file_type; - - let ctx = TestSessionContext::default(); - - let opts = VortexTableOptions { - // Disable sink-side rolling/splitting. - target_file_size_mb: 0, - ..Default::default() - }; - let factory = VortexFormatFactory::new().with_options(opts); - - let rows_per_partition = 300_000_usize; - let num_partitions = 8_usize; - let expected_total_rows = (rows_per_partition * num_partitions) as i64; - - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut partitions: Vec> = Vec::new(); - for p in 0..num_partitions { - let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - partitions.push(vec![RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?]); - } - - let table = MemTable::try_new(schema, partitions)?; - ctx.session.register_table("source", Arc::new(table))?; - - let source = ctx.session.table("source").await?; - let logical_plan = LogicalPlanBuilder::copy_to( - source.logical_plan().clone(), - "/table/".to_string(), - format_as_file_type(Arc::new(factory)), - Default::default(), - vec![], - )? - .build()?; - - ctx.session - .execute_logical_plan(logical_plan) - .await? - .collect() - .await?; - - let file_metas = ctx - .store - .list(Some(&"/table".into())) - .try_collect::>() - .await?; - - let unique_write_ids: HashSet<_> = file_metas - .iter() - .filter_map(|m| { - m.location - .filename() - .and_then(|name| name.split_once('_')) - .map(|(prefix, _)| prefix.to_string()) - }) - .collect(); - - assert_eq!( - unique_write_ids.len(), - 1, - "Expected one write_id with target size disabled; got {:?} from files: {:?}", - unique_write_ids, - file_metas - .iter() - .map(|m| format!("{}: {}B", m.location, m.size)) - .collect::>() - ); - - assert_eq!( - file_metas.len(), - 1, - "Expected exactly one output file when target size is disabled, got {}", - file_metas.len() - ); - - let result = ctx - .session - .sql("SELECT COUNT(*) as cnt FROM '/table/'") - .await? - .collect() - .await?; - - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - #[tokio::test] - async fn test_target_file_size_uses_single_sink_input_partition() -> anyhow::Result<()> { - use datafusion::datasource::MemTable; - use datafusion_datasource::file_format::format_as_file_type; - - let ctx = TestSessionContext::default(); - - let opts = VortexTableOptions { - // Enable sink-side sizing, but make the threshold large enough - // that all input data should fit in a single file. - target_file_size_mb: 512, - ..Default::default() - }; - let factory = VortexFormatFactory::new().with_options(opts); - - let rows_per_partition = 300_000_usize; - let num_partitions = 8_usize; - let expected_total_rows = (rows_per_partition * num_partitions) as i64; - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - // Build a MemTable with multiple physical input partitions to mimic - // DataFusion's parallel writer inputs. - let mut partitions: Vec> = Vec::new(); - for p in 0..num_partitions { - let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - partitions.push(vec![RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?]); - } - - let table = MemTable::try_new(schema, partitions)?; - ctx.session.register_table("source", Arc::new(table))?; - - let source = ctx.session.table("source").await?; - let logical_plan = LogicalPlanBuilder::copy_to( - source.logical_plan().clone(), - "/table/".to_string(), - format_as_file_type(Arc::new(factory)), - Default::default(), - vec![], - )? - .build()?; - - ctx.session - .execute_logical_plan(logical_plan) - .await? - .collect() - .await?; - - let file_metas = ctx - .store - .list(Some(&"/table".into())) - .try_collect::>() - .await?; - - let unique_write_ids: HashSet<_> = file_metas - .iter() - .filter_map(|m| { - m.location - .filename() - .and_then(|name| name.split_once('_')) - .map(|(prefix, _)| prefix.to_string()) - }) - .collect(); - - assert_eq!( - unique_write_ids.len(), - 1, - "Expected one write_id (single sink stream), got {:?} from files: {:?}", - unique_write_ids, - file_metas - .iter() - .map(|m| format!("{}: {}B", m.location, m.size)) - .collect::>() - ); - - assert!( - file_metas.len() < num_partitions, - "Expected fewer output files than input partitions after coalescing; got {} files for {num_partitions} input partitions", - file_metas.len() - ); - - let result = ctx - .session - .sql("SELECT COUNT(*) AS cnt FROM '/table/'") - .await? - .collect() - .await?; - - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - #[tokio::test] - async fn test_insert_sql_target_size_multi_partition_source_single_write_id() - -> anyhow::Result<()> { - use datafusion::datasource::MemTable; - - let ctx = TestSessionContext::default(); - - ctx.session - .sql( - "CREATE EXTERNAL TABLE my_tbl \ - (a BIGINT NOT NULL) \ - STORED AS vortex \ - LOCATION 'table/' \ - OPTIONS(target_file_size_mb '64');", - ) - .await?; - - let rows_per_partition = 300_000_usize; - let num_partitions = 8_usize; - let expected_total_rows = (rows_per_partition * num_partitions) as i64; - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut partitions: Vec> = Vec::new(); - for p in 0..num_partitions { - let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - partitions.push(vec![RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?]); - } - - let source = MemTable::try_new(schema, partitions)?; - ctx.session.register_table("source", Arc::new(source))?; - - ctx.session - .sql("INSERT INTO my_tbl SELECT a FROM source") - .await? - .collect() - .await?; - - let all_files = ctx.store.list(None).try_collect::>().await?; - - let unique_write_ids: HashSet<_> = all_files - .iter() - .filter_map(|m| { - m.location - .filename() - .and_then(|name| name.split_once('_')) - .map(|(prefix, _)| prefix.to_string()) - }) - .collect(); - - assert_eq!( - unique_write_ids.len(), - 1, - "Expected one write_id, got {:?} from files: {:?}", - unique_write_ids, - all_files - .iter() - .map(|m| format!("{}: {}B", m.location, m.size)) - .collect::>() - ); - assert!( - all_files.len() < num_partitions, - "Expected fewer files than input partitions; got {} files for {num_partitions} input partitions", - all_files.len() - ); - - let result = ctx - .session - .sql("SELECT COUNT(*) AS cnt FROM my_tbl") - .await? - .collect() - .await?; - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - #[tokio::test] - async fn test_insert_sql_streaming_source_single_write_id() -> anyhow::Result<()> { - use arrow_schema::SchemaRef; - use datafusion::catalog::streaming::StreamingTable; - use datafusion::physical_plan::stream::RecordBatchStreamAdapter; - use datafusion::physical_plan::streaming::PartitionStream; - use futures::stream; - - #[derive(Debug)] - struct StaticPartitionStream { - schema: SchemaRef, - batch: RecordBatch, - } - - impl PartitionStream for StaticPartitionStream { - fn schema(&self) -> &SchemaRef { - &self.schema - } - - fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { - let schema = Arc::clone(&self.schema); - let batch = self.batch.clone(); - Box::pin(RecordBatchStreamAdapter::new( - schema, - stream::iter(vec![Ok(batch)]), - )) - } - } - - let ctx = TestSessionContext::default(); - - ctx.session - .sql( - "CREATE EXTERNAL TABLE my_tbl \ - (a BIGINT NOT NULL) \ - STORED AS vortex \ - LOCATION 'table/' \ - OPTIONS(target_file_size_mb '64');", - ) - .await?; - - let rows_per_partition = 300_000_usize; - let num_partitions = 8_usize; - let expected_total_rows = (rows_per_partition * num_partitions) as i64; - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut partitions: Vec> = Vec::new(); - for p in 0..num_partitions { - let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?; - - partitions.push(Arc::new(StaticPartitionStream { - schema: Arc::clone(&schema), - batch, - })); - } - - let source = StreamingTable::try_new(schema, partitions)?; - ctx.session - .register_table("source_stream", Arc::new(source))?; - - ctx.session - .sql("INSERT INTO my_tbl SELECT a FROM source_stream") - .await? - .collect() - .await?; - - let all_files = ctx.store.list(None).try_collect::>().await?; - - let unique_write_ids: HashSet<_> = all_files - .iter() - .filter_map(|m| { - m.location - .filename() - .and_then(|name| name.split_once('_')) - .map(|(prefix, _)| prefix.to_string()) - }) - .collect(); - - assert_eq!( - unique_write_ids.len(), - 1, - "Expected one write_id for streaming source insert, got {:?} from files: {:?}", - unique_write_ids, - all_files - .iter() - .map(|m| format!("{}: {}B", m.location, m.size)) - .collect::>() - ); - - let result = ctx - .session - .sql("SELECT COUNT(*) AS cnt FROM my_tbl") - .await? - .collect() - .await?; - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - #[tokio::test] - async fn test_listing_table_direct_insert_into_streaming_exec_single_write_id() - -> anyhow::Result<()> { - use arrow_schema::SchemaRef; - use datafusion::physical_plan::ExecutionPlan; - use datafusion::physical_plan::collect; - use datafusion::physical_plan::stream::RecordBatchStreamAdapter; - use datafusion::physical_plan::streaming::PartitionStream; - use datafusion::physical_plan::streaming::StreamingTableExec; - use datafusion_expr::dml::InsertOp; - use futures::stream; - - #[derive(Debug)] - struct StaticPartitionStream { - schema: SchemaRef, - batch: RecordBatch, - } - - impl PartitionStream for StaticPartitionStream { - fn schema(&self) -> &SchemaRef { - &self.schema - } - - fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { - let schema = Arc::clone(&self.schema); - let batch = self.batch.clone(); - Box::pin(RecordBatchStreamAdapter::new( - schema, - stream::iter(vec![Ok(batch)]), - )) - } - } - - let ctx = TestSessionContext::default(); - - ctx.session - .sql( - "CREATE EXTERNAL TABLE my_tbl \ - (a BIGINT NOT NULL) \ - STORED AS vortex \ - LOCATION 'table/' \ - OPTIONS(target_file_size_mb '64');", - ) - .await?; - - let table_provider = ctx.session.table_provider("my_tbl").await?; - - let rows_per_partition = 300_000_usize; - let num_partitions = 8_usize; - let expected_total_rows = (rows_per_partition * num_partitions) as i64; - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut partitions: Vec> = Vec::new(); - for p in 0..num_partitions { - let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?; - - partitions.push(Arc::new(StaticPartitionStream { - schema: Arc::clone(&schema), - batch, - })); - } - - let input = Arc::new(StreamingTableExec::try_new( - schema, - partitions, - None, - Vec::new(), - false, - None, - )?) as Arc; - - let plan = table_provider - .insert_into(&ctx.session.state(), input, InsertOp::Append) - .await?; - let _count_batches = collect(plan, ctx.session.task_ctx()).await?; - - let all_files = ctx.store.list(None).try_collect::>().await?; - - let unique_write_ids: HashSet<_> = all_files - .iter() - .filter_map(|m| { - m.location - .filename() - .and_then(|name| name.split_once('_')) - .map(|(prefix, _)| prefix.to_string()) - }) - .collect(); - - assert_eq!( - unique_write_ids.len(), - 1, - "Expected one write_id for direct insert_into streaming exec, got {:?} from files: {:?}", - unique_write_ids, - all_files - .iter() - .map(|m| format!("{}: {}B", m.location, m.size)) - .collect::>() - ); - - let result = ctx - .session - .sql("SELECT COUNT(*) AS cnt FROM my_tbl") - .await? - .collect() - .await?; - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - - #[tokio::test] - async fn test_listing_table_direct_insert_into_unbounded_streaming_exec_single_write_id() - -> anyhow::Result<()> { - use arrow_schema::SchemaRef; - use datafusion::physical_plan::ExecutionPlan; - use datafusion::physical_plan::collect; - use datafusion::physical_plan::stream::RecordBatchStreamAdapter; - use datafusion::physical_plan::streaming::PartitionStream; - use datafusion::physical_plan::streaming::StreamingTableExec; - use datafusion_expr::dml::InsertOp; - use futures::stream; - - #[derive(Debug)] - struct StaticPartitionStream { - schema: SchemaRef, - batch: RecordBatch, - } - - impl PartitionStream for StaticPartitionStream { - fn schema(&self) -> &SchemaRef { - &self.schema - } - - fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { - let schema = Arc::clone(&self.schema); - let batch = self.batch.clone(); - Box::pin(RecordBatchStreamAdapter::new( - schema, - stream::iter(vec![Ok(batch)]), - )) - } - } - - let ctx = TestSessionContext::default(); - - ctx.session - .sql( - "CREATE EXTERNAL TABLE my_tbl \ - (a BIGINT NOT NULL) \ - STORED AS vortex \ - LOCATION 'table/' \ - OPTIONS(target_file_size_mb '64');", - ) - .await?; - - let table_provider = ctx.session.table_provider("my_tbl").await?; - - let rows_per_partition = 100_000_usize; - let num_partitions = 8_usize; - let expected_total_rows = (rows_per_partition * num_partitions) as i64; - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let mut partitions: Vec> = Vec::new(); - for p in 0..num_partitions { - let values = pseudo_random_i64s(rows_per_partition, (p * rows_per_partition) as i64); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int64Array::from(values))], - )?; - - partitions.push(Arc::new(StaticPartitionStream { - schema: Arc::clone(&schema), - batch, - })); - } - - let input = Arc::new(StreamingTableExec::try_new( - schema, - partitions, - None, - Vec::new(), - true, - None, - )?) as Arc; - - let plan = table_provider - .insert_into(&ctx.session.state(), input, InsertOp::Append) - .await?; - let _count_batches = collect(plan, ctx.session.task_ctx()).await?; - - let all_files = ctx.store.list(None).try_collect::>().await?; - - let unique_write_ids: HashSet<_> = all_files - .iter() - .filter_map(|m| { - m.location - .filename() - .and_then(|name| name.split_once('_')) - .map(|(prefix, _)| prefix.to_string()) - }) - .collect(); - - assert_eq!( - unique_write_ids.len(), - 1, - "Expected one write_id for unbounded streaming insert, got {:?} from files: {:?}", - unique_write_ids, - all_files - .iter() - .map(|m| format!("{}: {}B", m.location, m.size)) - .collect::>() - ); - - let result = ctx - .session - .sql("SELECT COUNT(*) AS cnt FROM my_tbl") - .await? - .collect() - .await?; - let count = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, expected_total_rows, "Total row count mismatch"); - - Ok(()) - } - #[test] fn test_display_as() { let session = VortexSession::empty(); @@ -1986,7 +486,6 @@ mod tests { config: config.clone(), schema: Arc::clone(table_schema.file_schema()), session: session.clone(), - target_file_size: None, }; insta::assert_snapshot!(DefaultDisplay(get_sink()).to_string(), @"VortexSink(file_groups=[])"); diff --git a/vortex-datafusion/src/persistent/tests.rs b/vortex-datafusion/src/persistent/tests.rs index 220a1477b13..41d8141165f 100644 --- a/vortex-datafusion/src/persistent/tests.rs +++ b/vortex-datafusion/src/persistent/tests.rs @@ -494,7 +494,7 @@ async fn arrow_uuid_extension_roundtrip() -> anyhow::Result<()> { use datafusion::arrow::array::FixedSizeBinaryArray; use datafusion::arrow::array::RecordBatch; use datafusion::assert_batches_sorted_eq; - use vortex::array::arrow::ArrowSessionExt; + use vortex_arrow::ArrowSessionExt; let ctx = TestSessionContext::default(); // Default vortex session has importer/exporter for Arrow UUID @@ -561,7 +561,7 @@ async fn arrow_uuid_extension_roundtrip_nested_struct() -> anyhow::Result<()> { use datafusion::arrow::array::RecordBatch; use datafusion::arrow::array::StructArray as ArrowStructArray; use datafusion::assert_batches_sorted_eq; - use vortex::array::arrow::ArrowSessionExt; + use vortex_arrow::ArrowSessionExt; let ctx = TestSessionContext::default(); let session = VortexSession::default(); diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index 82fae9bf56d..325f8eae92f 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -98,7 +98,6 @@ use futures::StreamExt; use futures::TryStreamExt; use futures::future::try_join_all; use vortex::array::VortexSessionExecute; -use vortex::array::arrow::ArrowSessionExt; use vortex::dtype::DType; use vortex::dtype::FieldPath; use vortex::dtype::Nullability; @@ -115,6 +114,7 @@ use vortex::io::session::RuntimeSessionExt; use vortex::scan::DataSourceRef; use vortex::scan::ScanRequest; use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; use vortex_utils::parallelism::get_available_parallelism; use crate::convert::exprs::DefaultExpressionConvertor; diff --git a/vortex-duckdb/Cargo.toml b/vortex-duckdb/Cargo.toml index b6896910a89..59dbe8d4003 100644 --- a/vortex-duckdb/Cargo.toml +++ b/vortex-duckdb/Cargo.toml @@ -25,7 +25,6 @@ crate-type = ["rlib"] [dependencies] async-fs = { workspace = true } -async-trait = { workspace = true } bitvec = { workspace = true } custom-labels = { workspace = true } futures = { workspace = true } @@ -45,6 +44,7 @@ vortex-utils = { workspace = true, features = ["dashmap"] } [dev-dependencies] anyhow = { workspace = true } +divan = { workspace = true } geo-types = { workspace = true } jiff = { workspace = true } rstest = { workspace = true } @@ -54,6 +54,10 @@ vortex-runend = { workspace = true } vortex-sequence = { workspace = true } wkb = { workspace = true } +[[bench]] +name = "bool_export" +harness = false + [lints] workspace = true diff --git a/vortex-duckdb/benches/bool_export.rs b/vortex-duckdb/benches/bool_export.rs new file mode 100644 index 00000000000..f8094b554ae --- /dev/null +++ b/vortex-duckdb/benches/bool_export.rs @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmark for the bit -> byte-bool unpack performed by `BoolExporter::export`. +//! +//! DuckDB stores booleans as one byte per value (`&mut [bool]`), while Vortex stores them +//! bit-packed in a [`BitBuffer`]. On every boolean column export we unpack a slice of the +//! bit buffer into the destination byte slice. +//! +//! This bench isolates that pure unpacking logic so it can run without a live DuckDB +//! vector: a reused `Vec` stands in for the DuckDB byte-bool destination. It compares +//! the previous implementation (`.iter().collect::>()` followed by +//! `copy_from_slice`) against the new allocation-free direct zip-write. + +use divan::Bencher; +use vortex::buffer::BitBuffer; + +fn main() { + divan::main(); +} + +const SIZES: &[usize] = &[1024, 16_384]; + +/// Bit densities (true ratio), expressed as `(numerator, denominator)`. +const DENSITIES: &[(usize, usize)] = &[(1, 2), (1, 10), (9, 10)]; + +fn make_buffer(len: usize, density: (usize, usize)) -> BitBuffer { + let (num, den) = density; + // `+ 8` so we can slice off a non-byte-aligned offset, matching real exports. + BitBuffer::from_iter((0..len + 8).map(|i| (i % den) < num)) +} + +/// Previous implementation: allocate a throwaway `Vec` then copy into the destination. +#[divan::bench(args = SIZES, consts = [0usize, 1usize, 2usize])] +fn old_collect_copy(bencher: Bencher, len: usize) { + let buffer = make_buffer(len, DENSITIES[DENSITY_IDX]); + bencher + .with_inputs(|| vec![false; len]) + .bench_refs(|dst: &mut Vec| { + // Offset by 1 to exercise a non-byte-aligned slice like the export hot path. + dst.copy_from_slice(&buffer.slice(1..(1 + len)).iter().collect::>()); + divan::black_box(&dst); + }); +} + +/// New implementation: zip the sliced bit iterator directly into the destination slice. +#[divan::bench(args = SIZES, consts = [0usize, 1usize, 2usize])] +fn new_zip_write(bencher: Bencher, len: usize) { + let buffer = make_buffer(len, DENSITIES[DENSITY_IDX]); + bencher + .with_inputs(|| vec![false; len]) + .bench_refs(|dst: &mut Vec| { + for (slot, bit) in dst.iter_mut().zip(buffer.slice(1..(1 + len)).iter()) { + *slot = bit; + } + divan::black_box(&dst); + }); +} diff --git a/vortex-duckdb/build.rs b/vortex-duckdb/build.rs index 7647011e366..caae2a8158a 100644 --- a/vortex-duckdb/build.rs +++ b/vortex-duckdb/build.rs @@ -23,15 +23,19 @@ const DUCKDB_RELEASES_URL: &str = "https://ci-builds.vortex.dev"; const DUCKDB_SOURCE_RELEASE_URL: &str = "https://github.com/duckdb/duckdb/archive/refs/tags"; const DUCKDB_SOURCE_COMMIT_URL: &str = "https://github.com/duckdb/duckdb/archive"; -const DEFAULT_DUCKDB_VERSION: &str = "1.5.3"; +const DEFAULT_DUCKDB_VERSION: &str = "1.5.4"; const BUILD_ARTIFACTS: [&str; 3] = ["libduckdb.dylib", "libduckdb.so", "libduckdb_static.a"]; -const SOURCE_FILES: [&str; 7] = [ +const SOURCE_FILES: [&str; 11] = [ "cpp/vortex_duckdb.cpp", "cpp/copy_function.cpp", "cpp/expr.cpp", + "cpp/optimizer.cpp", "cpp/scalar_fn_pushdown.cpp", + "cpp/spatial_overrides.cpp", + "cpp/cast_pushdown.cpp", + "cpp/aggregate_fn_pushdown.cpp", "cpp/table_filter.cpp", "cpp/table_function.cpp", "cpp/vector.cpp", @@ -175,9 +179,10 @@ const DUCKDB_C_API_FUNCTIONS: [&str; 133] = [ "duckdb_vector_size", ]; -const DUCKDB_C_API_HEADERS: [&str; 6] = [ +const DUCKDB_C_API_HEADERS: [&str; 7] = [ "cpp/include/vortex_duckdb.h", "cpp/include/expr.h", + "cpp/include/spatial_overrides.h", "cpp/include/table_filter.h", "cpp/include/vector.h", "cpp/include/copy_function.h", diff --git a/vortex-duckdb/cpp/aggregate_fn_pushdown.cpp b/vortex-duckdb/cpp/aggregate_fn_pushdown.cpp new file mode 100644 index 00000000000..2caa3463d53 --- /dev/null +++ b/vortex-duckdb/cpp/aggregate_fn_pushdown.cpp @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include "aggregate_fn_pushdown.hpp" +#include "duckdb/planner/expression/bound_aggregate_expression.hpp" +#include "duckdb/planner/expression/bound_columnref_expression.hpp" +#include "duckdb/planner/operator/logical_aggregate.hpp" +#include "optimizer.hpp" +#include "table_function.hpp" + +using enum LogicalOperatorType; + +LogicalOperatorPtr TryPushdownAggregateFunctions(ClientContext &context, LogicalOperatorPtr plan) { + Analyses analyses; + Projections projections; + FindGetsAndProjections(*plan, analyses, projections); + if (analyses.empty()) { + return plan; + } + return RewriteAggregates(context, std::move(plan), analyses, projections); +} + +LogicalOperatorPtr RewriteAggregates(ClientContext &context, + LogicalOperatorPtr op, + Analyses &analyses, + const Projections &projections) { + for (auto &child : op->children) { + child = RewriteAggregates(context, std::move(child), analyses, projections); + } + if (op->type == LOGICAL_AGGREGATE_AND_GROUP_BY) { + return TryReplaceAggregate(context, std::move(op), analyses, projections); + } + return op; +} + +static bool IsUngrouped(const LogicalAggregate &agg) { + return agg.groups.empty() && agg.grouping_sets.empty() && agg.grouping_functions.empty() && + !agg.expressions.empty(); +} + +constexpr inline idx_t COUNT_STAR_PROJ_IDX = std::numeric_limits::max(); + +LogicalOperatorPtr TryReplaceAggregate(ClientContext &context, + LogicalOperatorPtr op, + Analyses &analyses, + const Projections &projections) { + LogicalAggregate &agg = op->Cast(); + if (!IsUngrouped(agg)) { + return op; + } + + LogicalGet *const get = GetChildGet(agg); + if (get == nullptr) { + return op; + } + + vector> input; + const idx_t aggregates_len = agg.expressions.size(); + input.reserve(aggregates_len); + + for (const auto &expr : agg.expressions) { + if (expr->GetExpressionClass() != ExpressionClass::BOUND_AGGREGATE) { + return op; + } + const auto &bound_aggr = expr->Cast(); + if (bound_aggr.IsDistinct() || bound_aggr.filter != nullptr || bound_aggr.order_bys != nullptr) { + return op; + } + + if (bound_aggr.function.name == "count_star") { + input.emplace_back(COUNT_STAR_PROJ_IDX, *expr); + continue; + } + + if (bound_aggr.children.size() != 1 || + bound_aggr.children[0]->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { + return op; + } + const auto &bound_col = bound_aggr.children[0]->Cast(); + const auto binding = Resolve(bound_col.binding, analyses, projections); + if (!binding || &binding->analysis.get != get) { + return op; + } + input.emplace_back(binding->column_index, *expr); + } + + if (!aggregate_pushdown(context, {*get, input})) { + return op; + } + + // GET now returns one column per aggregate. Expand existing columns + auto &column_ids = get->GetMutableColumnIds(); + get->types.resize(aggregates_len); + get->returned_types.resize(aggregates_len); + column_ids.resize(aggregates_len); + + vector names(aggregates_len); // need a copy because we reference original names + + for (idx_t i = 0; i < aggregates_len; i++) { + const auto &[column_index, expr] = input[i]; + if (column_index == COUNT_STAR_PROJ_IDX) { + names[i] = "count_star()"; + } else { + const TableColumnStorageIndex storage_index = get->GetColumnIds()[column_index].GetPrimaryIndex(); + names[i] = get->names[storage_index]; + } + get->types[i] = expr.return_type; + get->returned_types[i] = expr.return_type; + column_ids[i] = ColumnIndex {i}; + } + get->names = std::move(names); + get->projection_ids.clear(); + get->table_index = agg.aggregate_index; + + unique_ptr &child = agg.children[0]; + if (child->type == LOGICAL_GET) { + return std::move(child); + } + D_ASSERT(child->type == LOGICAL_PROJECTION); + D_ASSERT(child->children.size() == 1); + D_ASSERT(child->children[0]->type == LOGICAL_GET); + return std::move(child->children[0]); +} + +LogicalGet *GetChildGet(const LogicalAggregate &agg) { + if (agg.children.size() != 1) { + return nullptr; + } + LogicalOperator &child = *agg.children[0]; + LogicalOperator *op; + if (child.type == LOGICAL_GET) { + op = &child; + } else if (child.type == LOGICAL_PROJECTION && child.children.size() == 1 && + child.children[0]->type == LOGICAL_GET) { + op = child.children[0].get(); + } else { + return nullptr; + } + LogicalGet &get = op->Cast(); + return get.function.bind == duckdb_vx_table_function_bind ? &get : nullptr; +} diff --git a/vortex-duckdb/cpp/cast_pushdown.cpp b/vortex-duckdb/cpp/cast_pushdown.cpp new file mode 100644 index 00000000000..3acd5b04161 --- /dev/null +++ b/vortex-duckdb/cpp/cast_pushdown.cpp @@ -0,0 +1,178 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include "cast_pushdown.hpp" +#include "table_function.hpp" + +#include "duckdb/planner/operator/logical_get.hpp" +#include "duckdb/planner/operator/logical_projection.hpp" +#include "duckdb/planner/expression/bound_cast_expression.hpp" +#include "duckdb/planner/expression/bound_columnref_expression.hpp" + +/** + * A GET reachable through a single-child chain of projections. A join or any + * other multi-child operator breaks the chain. + * + * LOGICAL_FILTER also breaks the chain. In SELECT CAST/TRY_CAST(x) WHERE g(x) + * we can't push down cast if g(x) isn't pushed. If g(x) filters some rows + * on which cast would fail, the query passed by default but fails if cast is pushed + * and runs before g(x). + * + * See test/sql/copy/csv/test_insert_into_types.test in duckdb (cast not pushed past a join) + */ +static bool ReachesPushdownGet(const LogicalOperator &op) { + const LogicalOperator *cur = &op; + while (cur->children.size() == 1) { + cur = cur->children[0].get(); + switch (cur->type) { + case LogicalOperatorType::LOGICAL_GET: + return cur->Cast().function.bind == duckdb_vx_table_function_bind; + case LogicalOperatorType::LOGICAL_PROJECTION: + continue; + default: + return false; + } + } + return false; +} + +void CastCollect::VisitOperator(LogicalOperator &op) { + /* + * Logical projection expressions are columns which reference underlying + * GETs. Don't process them, as they would add conflicts for every column + * used in projection. Example: PROJECTION(col) -> GET(col). We don't want + * to visit BoundColumnRefExpression in PROJECTION to avoid registering a + * non-existent conflict. + * + * However, CastReplace will visit them because we need to update their + * types if pushdown succeeded. + */ + if (op.type != LogicalOperatorType::LOGICAL_PROJECTION) { + return LogicalOperatorVisitor::VisitOperator(op); + } + auto &projection = op.Cast(); + + // Only push casts from a projection that forwards just column refs and + // casts and reaches a GET without a join in between. A constant or other + // expression makes the projection ineligible. + // See test/sql/copy/csv/test_csv_error_message_type.test (top-level cast + // to VARCHAR must still push) and test_large_integer_detection.test (a + // nested cast to VARCHAR must not) in duckdb. + bool clean = ReachesPushdownGet(projection); + for (const auto &e : projection.expressions) { + switch (e->GetExpressionClass()) { + case ExpressionClass::BOUND_COLUMN_REF: + case ExpressionClass::BOUND_CAST: + continue; + default: + clean = false; + break; + } + } + if (clean) { + for (const auto &e : projection.expressions) { + if (e->GetExpressionClass() == ExpressionClass::BOUND_CAST) { + top_level_casts.insert(e.get()); + } + } + } + if (projections.count(projection.table_index)) { + VisitOperatorChildren(op); + return; + } + + LogicalOperatorVisitor::VisitOperator(op); +} + +ExpressionPtr CastCollect::VisitReplace(BoundColumnRefExpression &expr, ExpressionPtr *ptr) { + if (const auto binding = Resolve(expr.binding, analyses, projections)) { + // Column is used without cast applied to it, register a conflict. + // Not emplace() as we need to update the value if it was present + binding->analysis.col_to_expr[binding->column_index] = nullptr; + } + return std::move(*ptr); +} + +ExpressionPtr CastCollect::VisitReplace(BoundCastExpression &expr, ExpressionPtr *ptr) { + if (expr.child->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { + // Descend into children so e.g. fn(col, other) still sees "col" and + // registers a conflict + return nullptr; + } + const auto &bound_col = expr.child->Cast(); + const auto binding = Resolve(bound_col.binding, analyses, projections); + if (!binding) { + return nullptr; + } + auto &col_to_expr = binding->analysis.col_to_expr; + + if (auto it = col_to_expr.find(binding->column_index); it == col_to_expr.end()) { + // Only a top-level projection cast starts a candidate. + if (top_level_casts.count(&expr)) { + col_to_expr.emplace(binding->column_index, &expr); + } + } else if (it->second == nullptr || + it->second->Cast().return_type != expr.return_type || + // TODO(myrrc) this line needs upstreaming + it->second->Cast().try_cast != expr.try_cast) { + // Different target type, or already a conflict. + it->second = nullptr; + } + + return std::move(*ptr); +} + +ExpressionPtr CastReplace::VisitReplace(BoundColumnRefExpression &expr, ExpressionPtr *ptr) { + const auto binding = Resolve(expr.binding, analyses, projections); + if (!binding) { + return std::move(*ptr); + } + + const auto &[analysis, column_index, projection] = *binding; + if (CanPushdownColumn(analysis, column_index)) { + const idx_t storage_index = analysis.get.GetColumnIds()[column_index].GetPrimaryIndex(); + const LogicalType return_type = analysis.get.returned_types[storage_index]; + expr.return_type = return_type; + // LogicalProjection types are resolved by calling + // LogicalProjection::ResolveTypes, so we need to check whether types in + // projection have been resolved, and updated them only if needed. + if (projection != nullptr && !projection->types.empty()) { + projection->types[column_index] = return_type; + } + } + + return std::move(*ptr); +} + +ExpressionPtr CastReplace::VisitReplace(BoundCastExpression &expr, ExpressionPtr *ptr) { + if (expr.child->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { + return nullptr; // Same as in ScalarFnCollect::VisitReplace + } + auto &bound_col_base = expr.child; + const auto &bound_col = bound_col_base->Cast(); + const auto binding = Resolve(bound_col.binding, analyses, projections); + if (!binding) { + return nullptr; + } + + const auto &[analysis, column_index, projection] = *binding; + if (!CanPushdownColumn(analysis, column_index)) { + return std::move(*ptr); + } + + const idx_t storage_index = analysis.get.GetColumnIds()[column_index].GetPrimaryIndex(); + const LogicalType return_type = analysis.get.returned_types[storage_index]; + bound_col_base->return_type = return_type; + // Same as in CastReplace::VisitReplace(BoundColumnRefExpression) + if (projection != nullptr && !projection->types.empty()) { + projection->types[column_index] = return_type; + } + return std::move(bound_col_base); +} + +CastCollect::CastCollect(Analyses &analyses, const Projections &projections) + : analyses(analyses), projections(projections) { +} + +CastReplace::CastReplace(Analyses &analyses, const Projections &projections) + : analyses(analyses), projections(projections) { +} diff --git a/vortex-duckdb/cpp/expr.cpp b/vortex-duckdb/cpp/expr.cpp index afe2573adc2..54ef37cd959 100644 --- a/vortex-duckdb/cpp/expr.cpp +++ b/vortex-duckdb/cpp/expr.cpp @@ -3,10 +3,13 @@ #include "expr.h" #include "duckdb/function/scalar_function.hpp" +#include "duckdb/function/aggregate_function.hpp" #include "duckdb/planner/expression/bound_between_expression.hpp" +#include "duckdb/planner/expression/bound_cast_expression.hpp" #include "duckdb/planner/expression/bound_columnref_expression.hpp" #include "duckdb/planner/expression/bound_comparison_expression.hpp" #include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "duckdb/planner/expression/bound_aggregate_expression.hpp" #include "duckdb/planner/expression/bound_function_expression.hpp" #include "duckdb/planner/expression/bound_operator_expression.hpp" #include "duckdb/planner/expression/bound_conjunction_expression.hpp" @@ -21,6 +24,11 @@ extern "C" const char *duckdb_vx_sfunc_name(duckdb_vx_sfunc ffi_func) { return func->name.c_str(); } +extern "C" const char *duckdb_vx_agg_func_name(duckdb_vx_agg_func ffi) { + D_ASSERT(ffi); + return reinterpret_cast(ffi)->name.c_str(); +} + extern "C" const char *duckdb_vx_expr_to_string(duckdb_vx_expr ffi_expr) { if (!ffi_expr) { return nullptr; @@ -32,7 +40,6 @@ extern "C" const char *duckdb_vx_expr_to_string(duckdb_vx_expr ffi_expr) { return result; } -//! Create a DuckDB vortex error. extern "C" void duckdb_vx_destroy_expr(duckdb_vx_expr *ffi_expr) { auto expr = reinterpret_cast(ffi_expr); delete expr; @@ -129,3 +136,21 @@ extern "C" void duckdb_vx_expr_get_bound_function(duckdb_vx_expr ffi_expr, out->scalar_function = reinterpret_cast(&expr.function); out->bind_info = expr.bind_info.get(); } + +extern "C" duckdb_vx_expr duckdb_vx_expr_get_bound_cast_child(duckdb_vx_expr ffi_expr) { + D_ASSERT(ffi_expr); + auto &expr = reinterpret_cast(ffi_expr)->Cast(); + return reinterpret_cast(expr.child.get()); +} + +extern "C" bool duckdb_vx_expr_get_bound_cast_is_try(duckdb_vx_expr ffi_expr) { + D_ASSERT(ffi_expr); + auto &expr = reinterpret_cast(ffi_expr)->Cast(); + return expr.try_cast; +} + +extern "C" duckdb_vx_agg_func duckdb_vx_expr_get_bound_aggregate_function(duckdb_vx_expr ffi_expr) { + D_ASSERT(ffi_expr); + auto &expr = reinterpret_cast(ffi_expr)->Cast(); + return reinterpret_cast(&expr.function); +} diff --git a/vortex-duckdb/cpp/include/aggregate_fn_pushdown.hpp b/vortex-duckdb/cpp/include/aggregate_fn_pushdown.hpp new file mode 100644 index 00000000000..4deb716b7a5 --- /dev/null +++ b/vortex-duckdb/cpp/include/aggregate_fn_pushdown.hpp @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once +#include "optimizer.hpp" +#include "duckdb/optimizer/optimizer_extension.hpp" + +using namespace duckdb; + +// Push UNGROUPED_AGGREGATE's of form agg(T) and count_star() into GET. +LogicalOperatorPtr TryPushdownAggregateFunctions(ClientContext &context, LogicalOperatorPtr plan); + +LogicalOperatorPtr RewriteAggregates(ClientContext &context, + LogicalOperatorPtr op, + Analyses &analyses, + const Projections &projections); + +LogicalOperatorPtr TryReplaceAggregate(ClientContext &context, + LogicalOperatorPtr op, + Analyses &analyses, + const Projections &projections); + +// return GET for UNGROUPED_AGGREGATE -> [GET] or for UNGROUPED_AGGREGATE -> +// PROJECTION -> [GET], nullptr if not found. +LogicalGet *GetChildGet(const LogicalAggregate &agg); diff --git a/vortex-duckdb/cpp/include/cast_pushdown.hpp b/vortex-duckdb/cpp/include/cast_pushdown.hpp new file mode 100644 index 00000000000..f64e9da8b00 --- /dev/null +++ b/vortex-duckdb/cpp/include/cast_pushdown.hpp @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once +#include "optimizer.hpp" + +#include "duckdb/common/unordered_set.hpp" +#include "duckdb/main/client_context.hpp" +#include "duckdb/planner/expression.hpp" +#include "duckdb/planner/logical_operator.hpp" + +using namespace duckdb; + +/** + * Collect CAST(col) expressions. If "col" is used without CAST in "plan", + * record in "analyses.col_to_expr" + */ +struct CastCollect final : LogicalOperatorVisitor { + Analyses &analyses; + const Projections &projections; + // Casts that are direct outputs of a clean projection over a GET. Only these + // start a pushdown candidate; a nested cast may push down a different value. + // See test/sql/copy/csv/auto/test_large_integer_detection.test in duckdb + unordered_set top_level_casts; + + CastCollect(Analyses &analyses, const Projections &projections); + void VisitOperator(LogicalOperator &op) override; + ExpressionPtr VisitReplace(BoundColumnRefExpression &expr, ExpressionPtr *ptr) override; + ExpressionPtr VisitReplace(BoundCastExpression &expr, ExpressionPtr *ptr) override; +}; + +/* + * For "col" in columns collected by CastCollect, replace CAST(col) to "col" + * if "col" doesn't have conflicting usage. Update return types for bound + * columns and logical projections referencing this column. + */ +struct CastReplace final : LogicalOperatorVisitor { + Analyses &analyses; + const Projections &projections; + + CastReplace(Analyses &analyses, const Projections &projections); + ExpressionPtr VisitReplace(BoundColumnRefExpression &expr, ExpressionPtr *ptr) override; + ExpressionPtr VisitReplace(BoundCastExpression &expr, ExpressionPtr *ptr) override; +}; diff --git a/vortex-duckdb/cpp/include/expr.h b/vortex-duckdb/cpp/include/expr.h index 5b7997596d6..704c8406c90 100644 --- a/vortex-duckdb/cpp/include/expr.h +++ b/vortex-duckdb/cpp/include/expr.h @@ -10,11 +10,14 @@ extern "C" { #endif typedef struct duckdb_vx_sfunc_ *duckdb_vx_sfunc; +typedef struct duckdb_vx_agg_func_ *duckdb_vx_agg_func; const char *duckdb_vx_sfunc_name(duckdb_vx_sfunc ffi_func); typedef struct duckdb_vx_expr_ *duckdb_vx_expr; +const char *duckdb_vx_agg_func_name(duckdb_vx_agg_func func); + /// Return the string representation of the expression. Must be freed with `duckdb_free`. const char *duckdb_vx_expr_to_string(duckdb_vx_expr expr); @@ -264,6 +267,12 @@ typedef struct { void duckdb_vx_expr_get_bound_function(duckdb_vx_expr expr, duckdb_vx_expr_bound_function *out); +duckdb_vx_expr duckdb_vx_expr_get_bound_cast_child(duckdb_vx_expr expr); + +bool duckdb_vx_expr_get_bound_cast_is_try(duckdb_vx_expr expr); + +duckdb_vx_agg_func duckdb_vx_expr_get_bound_aggregate_function(duckdb_vx_expr expr); + #ifdef __cplusplus /* End C ABI */ } #endif diff --git a/vortex-duckdb/cpp/include/optimizer.hpp b/vortex-duckdb/cpp/include/optimizer.hpp new file mode 100644 index 00000000000..acb797928f6 --- /dev/null +++ b/vortex-duckdb/cpp/include/optimizer.hpp @@ -0,0 +1,137 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once +#include "table_function.hpp" + +#include "duckdb/planner/expression.hpp" +#include "duckdb/planner/operator/logical_get.hpp" + +#include + +// Aliases here are for ease of migration to duckdb 2.0 where these are +// separate types + +using namespace duckdb; + +/** + * Column index in requested scan. Example: + * + * CREATE TABLE t (a1 INTEGER, a2 INTEGER, a3 INTEGER); + * SELECT a2, a3 FROM t; + * + * a2's TableColumnScanIndex is 0, a3's TableColumnScanIndex is 1, + * index is index in SELECT clause. + */ +using TableColumnScanIndex = idx_t; +using ProjectionIndex = TableColumnScanIndex; + +/** + * Column index in table's storage. Example: + * + * CREATE TABLE t (a1 INTEGER, a2 INTEGER, a3 INTEGER); + * SELECT a2, a3 FROM t; + * + * a2's TableColumnStorageIndex is 1, a3's TableColumnStorageIndex is 2, + * index is index of column in table storage. + * + * for i: TableColumnScanIndex, column_ids[i].GetPrimaryIndex() is + * TableColumnStorageIndex + */ +using TableColumnStorageIndex = idx_t; + +using TableIndex = idx_t; + +using ExpressionPtr = unique_ptr; +using LogicalOperatorPtr = unique_ptr; + +struct GetAnalysis { + LogicalGet &get; + /** + * for fn(col), mapping of "col scan index" -> "expression applied to function". + * "expression" is nullptr iff column is used with a different expression + * or without expression application in the query plan (i.e. SELECT col). + */ + unordered_map col_to_expr; + + TableColumnStorageIndex StorageIndex(TableColumnScanIndex idx) const; +}; + +using Analyses = unordered_map; + +/* + * Using scalar function pushdown as a specific example, + * SELECT fn(col) FROM '*.vortex' yields a PROJECTION fn(col) -> GET (vortex) + * plan. PROJECTION's "col" table_index is 1, vortex GET's table_index is 0. + * So we want to track original table_index for GET in case column is found + * in filter we failed to push down (i.e. WHERE prefix(col, 'h')) as well as + * projection's table_index. + * + * So we keep a mapping of + * + * "projection table index" to "projection operator". + * + * to resolve this. + * For simplicity, current implementation is limited to one level i.e. + * PROJECTION -> GET (i.e. read from VIEW) is pushed down but VIEW->VIEW->GET + * or VIEW->CTE->GET is not. + * + * Storing a reference is fine because the plan outlives the optimizer pass. + */ +using Projections = unordered_map; + +void FindGetsAndProjections(LogicalOperator &op, Analyses &analyses, Projections &aliases); + +struct GetBinding { + GetAnalysis &analysis; + TableColumnScanIndex column_index; + // If column binding was part of a projection, this is non-nullptr + LogicalProjection *projection; +}; + +/* + * Given a column binding, resolve it to a GET and a GET's column scan index. + * Returns nullopt for virtual columns and columns which are neither part of + * GET nor part of PROJECTION wrapping a GET. + */ +std::optional Resolve(ColumnBinding binding, Analyses &analyses, const Projections &projections); + +// A passthrough projection only forwards its child columns, e.g. a VIEW's +// "SELECT col". +bool IsPassthrough(const LogicalProjection &projection); + +// There are no conflicting column usages in the plan +bool CanPushdownColumn(const GetAnalysis &analysis, TableColumnScanIndex idx); + +template +LogicalOperatorPtr TryPushdown(ClientContext &context, LogicalOperatorPtr plan) { + Analyses analyses; + Projections projections; + FindGetsAndProjections(*plan, analyses, projections); + if (analyses.empty()) { + return plan; + } + Collect(analyses, projections).VisitOperator(*plan); + + bool any_pushed = false; + for (auto &[_, analysis] : analyses) { + for (auto &[column_index, expr] : analysis.col_to_expr) { + if (expr == nullptr) { // Conflict for column + continue; + } + const TableColumnStorageIndex storage_index = analysis.StorageIndex(column_index); + TableFunctionProjectionExpressionInput input {analysis.get, *expr, storage_index}; + if (projection_expression_pushdown(context, input)) { + // LOGICAL_GET doesn't initialize .types of LogicalOperator + analysis.get.returned_types[storage_index] = expr->return_type; + any_pushed = true; + } else { // failed to push down expression, can't replace it + expr = nullptr; + } + } + } + + if (any_pushed) { + Replace(analyses, projections).VisitOperator(*plan); + } + return plan; +} diff --git a/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp b/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp index ef590c96dcc..19bad361cb0 100644 --- a/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp +++ b/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp @@ -1,82 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors #pragma once -#include "duckdb.h" - -#include "duckdb/optimizer/optimizer_extension.hpp" +#include "optimizer.hpp" #include "duckdb/planner/expression/bound_columnref_expression.hpp" #include "duckdb/planner/expression/bound_function_expression.hpp" -#include "duckdb/planner/operator/logical_get.hpp" -#include using namespace duckdb; -using ExpressionPtr = unique_ptr; -using LogicalOperatorPtr = unique_ptr; - -/** - * Column index in requested scan. Example: - * - * CREATE TABLE t (a1 INTEGER, a2 INTEGER, a3 INTEGER); - * SELECT a2, a3 FROM t; - * - * a2's TableColumnScanIndex is 0, a3's TableColumnScanIndex is 1, - * index is index in SELECT clause. - */ -using TableColumnScanIndex = idx_t; - -/** - * Column index in table's storage. Example: - * - * CREATE TABLE t (a1 INTEGER, a2 INTEGER, a3 INTEGER); - * SELECT a2, a3 FROM t; - * - * a2's TableColumnStorageIndex is 1, a3's TableColumnScanIndex is 2, - * index is index of column in table storage. - * - * for i: TableColumnScanIndex, column_ids[i].GetPrimaryIndex() is - * TableColumnStorageIndex - */ -using TableColumnStorageIndex = idx_t; - -using TableIndex = idx_t; - -struct GetAnalysis { - LogicalGet &get; - /** - * for fn(col), mapping of "col scan index" -> "fn expression". - * "fn expression" is nullptr iff column is used with a different function - * or without function application in the query plan. - */ - unordered_map col_to_fn; - - TableColumnStorageIndex StorageIndex(TableColumnScanIndex idx) const; -}; - -using Analyses = unordered_map; - -/* - * SELECT fn(col) FROM '*.vortex' yields a PROJECTION fn(col) -> GET (vortex) - * plan. PROJECTION's "col" table_index is 1, vortex GET's table_index is 0. - * So we want to track original table_index for GET in case column is found - * in filter we failed to push down (i.e. WHERE prefix(col, 'h')) as well as - * projection's table_index. - * - * So we keep a mapping of - * - * "projection table index" to "projection operator". - * - * to resolve this. - * For simplicity, current implementation is limited to one level i.e. - * PROJECTION -> GET (i.e. read from VIEW) is pushed down but VIEW->VIEW->GET - * or VIEW->CTE->GET is not. - * - * Storing a reference is fine because the plan outlives the optimizer pass. - */ -using Projections = unordered_map; - -LogicalOperatorPtr TryPushdownScalarFunctions(ClientContext &context, LogicalOperatorPtr plan); - /** * Collect fn(col) expressions i.e. expressions where a single function (not * a function chain) wraps a single bound column. If "col" is used without @@ -105,19 +35,3 @@ struct ScalarFnReplace final : LogicalOperatorVisitor { ExpressionPtr VisitReplace(BoundColumnRefExpression &expr, ExpressionPtr *ptr) override; ExpressionPtr VisitReplace(BoundFunctionExpression &expr, ExpressionPtr *ptr) override; }; - -void FindGetsAndProjections(LogicalOperator &op, Analyses &analyses, Projections &aliases); - -struct GetBinding { - GetAnalysis &analysis; - TableColumnScanIndex column_index; - // If column binding was part of a projection, this is non-nullptr - LogicalProjection *projection; -}; - -/* - * Given a column binding, resolve it to a GET and a GET's column scan index. - * Returns nullopt for virtual columns and columns which are neither part of - * GET nor part of PROJECTION wrapping a GET. - */ -std::optional Resolve(ColumnBinding binding, Analyses &analyses, const Projections &projections); diff --git a/vortex-duckdb/cpp/include/spatial_overrides.h b/vortex-duckdb/cpp/include/spatial_overrides.h new file mode 100644 index 00000000000..e5bbe4dde04 --- /dev/null +++ b/vortex-duckdb/cpp/include/spatial_overrides.h @@ -0,0 +1,19 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#pragma once + +#include "duckdb.h" + +#ifdef __cplusplus /* If compiled as C++, use C ABI */ +extern "C" { +#endif + +/// Shadow the spatial functions that block filter pushdown with pushable copies; see +/// `SPATIAL_OVERRIDES` in spatial_overrides.cpp for the list. Call after `LOAD spatial`; +/// does nothing when spatial is not loaded. +duckdb_state duckdb_vx_register_spatial_overrides(duckdb_database ffi_db); + +#ifdef __cplusplus +} +#endif diff --git a/vortex-duckdb/cpp/include/spatial_overrides.hpp b/vortex-duckdb/cpp/include/spatial_overrides.hpp new file mode 100644 index 00000000000..ce49d82fef8 --- /dev/null +++ b/vortex-duckdb/cpp/include/spatial_overrides.hpp @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once +#include "duckdb/main/client_context.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/planner/logical_operator_visitor.hpp" + +using namespace duckdb; + +/* + * Vortex shadows some spatial (`ST_*`) scalar functions with tweaked copies, because as + * spatial registers them their filters can never push into Vortex scans. + * + * `duckdb_vx_register_spatial_overrides` (spatial_overrides.h) installs the copies after + * `LOAD spatial`. `RestoreSpatialOverrides` below hands join conditions back to spatial's + * originals: Vortex never pushes join conditions, and spatial's join machinery expects its + * own functions. + */ + +// Rebinds overridden spatial calls in join conditions back to spatial's original, so spatial's +// own machinery handles joins. Filters are left untouched: they keep the override and push to +// Vortex scans. +struct SpatialOverrideRestore final : LogicalOperatorVisitor { + ClientContext &context; + + explicit SpatialOverrideRestore(ClientContext &context); + void VisitOperator(LogicalOperator &op) override; + unique_ptr VisitReplace(BoundFunctionExpression &expr, unique_ptr *ptr) override; +}; + +/// Rebind overridden spatial calls in join conditions to spatial's originals. +void RestoreSpatialOverrides(ClientContext &context, LogicalOperator &plan); diff --git a/vortex-duckdb/cpp/include/table_function.h b/vortex-duckdb/cpp/include/table_function.h index f59c0330300..65a1a3f1dd3 100644 --- a/vortex-duckdb/cpp/include/table_function.h +++ b/vortex-duckdb/cpp/include/table_function.h @@ -88,6 +88,10 @@ typedef struct { duckdb_state duckdb_vx_register_table_functions(duckdb_database ffi_db); +typedef struct duckdb_vx_agg_input_ *duckdb_vx_agg_input; +idx_t duckdb_vx_aggregate_len(duckdb_vx_agg_input ffi); +duckdb_vx_expr duckdb_vx_aggregate_at(duckdb_vx_agg_input ffi, idx_t index, idx_t *proj_idx); + #ifdef __cplusplus } #endif diff --git a/vortex-duckdb/cpp/include/table_function.hpp b/vortex-duckdb/cpp/include/table_function.hpp index 8cdb813f4a7..07d035917be 100644 --- a/vortex-duckdb/cpp/include/table_function.hpp +++ b/vortex-duckdb/cpp/include/table_function.hpp @@ -25,3 +25,11 @@ struct TableFunctionProjectionExpressionInput { // true if we can push down the expression, false otherwise bool projection_expression_pushdown(duckdb::ClientContext &context, const TableFunctionProjectionExpressionInput &input); + +struct TableFunctionUngroupedAggregateInput { + const duckdb::LogicalGet &get; + // Column scan index -> aggregate expression + const duckdb::vector> &projections; +}; + +bool aggregate_pushdown(duckdb::ClientContext &context, const TableFunctionUngroupedAggregateInput &input); diff --git a/vortex-duckdb/cpp/optimizer.cpp b/vortex-duckdb/cpp/optimizer.cpp new file mode 100644 index 00000000000..87082d57790 --- /dev/null +++ b/vortex-duckdb/cpp/optimizer.cpp @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include "optimizer.hpp" +#include "table_function.hpp" + +#include "duckdb/planner/expression/bound_columnref_expression.hpp" +#include "duckdb/planner/operator/logical_get.hpp" +#include "duckdb/planner/operator/logical_projection.hpp" + +void FindGetsAndProjections(LogicalOperator &op, Analyses &analyses, Projections &projections) { + using enum LogicalOperatorType; + switch (op.type) { + case LOGICAL_GET: { + if (auto &get = op.Cast(); get.function.bind == duckdb_vx_table_function_bind) { + analyses.emplace(get.table_index, GetAnalysis {get, {}}); + } + break; + } + case LOGICAL_PROJECTION: { + LogicalProjection &projection = op.Cast(); + D_ASSERT(projection.children.size() == 1); + auto &child = *projection.children[0]; + if (!IsPassthrough(projection) || child.type != LOGICAL_GET) { + break; + } + // The GET itself is recorded when recursion reaches it below. Only + // passthrough projections wrapping a vortex GET act as aliases. + if (auto &get = child.Cast(); get.function.bind == duckdb_vx_table_function_bind) { + projections.emplace(projection.table_index, projection); + } + break; + } + default: + break; + } + + for (auto &child : op.children) { + FindGetsAndProjections(*child, analyses, projections); + } +} + +TableColumnStorageIndex GetAnalysis::StorageIndex(TableColumnScanIndex idx) const { + return get.GetColumnIds()[idx].GetPrimaryIndex(); +} + +static bool IsVirtualColumn(const GetAnalysis &analysis, TableColumnScanIndex idx) { + return analysis.get.GetColumnIds()[idx].IsVirtualColumn(); +} + +std::optional Resolve(ColumnBinding binding, Analyses &analyses, const Projections &projections) { + if (const auto it = analyses.find(binding.table_index); it != analyses.end()) { + if (IsVirtualColumn(it->second, binding.column_index)) { + return std::nullopt; + } + return {{it->second, binding.column_index, nullptr}}; + } + + const auto projection_it = projections.find(binding.table_index); + if (projection_it == projections.end()) { + return std::nullopt; + } + + LogicalProjection &projection = projection_it->second; + const ExpressionPtr &inner = projection.expressions[binding.column_index]; + if (inner->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { + return std::nullopt; + } + const ColumnBinding &get_binding = inner->Cast().binding; + if (const auto it = analyses.find(get_binding.table_index); it != analyses.end()) { + if (IsVirtualColumn(it->second, get_binding.column_index)) { + return std::nullopt; + } + return {{it->second, get_binding.column_index, &projection}}; + } + return std::nullopt; +} + +bool CanPushdownColumn(const GetAnalysis &analysis, TableColumnScanIndex idx) { + const auto it = analysis.col_to_expr.find(idx); + return it != analysis.col_to_expr.end() && it->second != nullptr; +} + +bool IsPassthrough(const LogicalProjection &projection) { + if (projection.expressions.empty()) { + return false; // don't register empty projections in Projections + } + for (const auto &e : projection.expressions) { + if (e->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { + return false; + } + } + return true; +} diff --git a/vortex-duckdb/cpp/scalar_fn_pushdown.cpp b/vortex-duckdb/cpp/scalar_fn_pushdown.cpp index 057920f79f7..bcf23496443 100644 --- a/vortex-duckdb/cpp/scalar_fn_pushdown.cpp +++ b/vortex-duckdb/cpp/scalar_fn_pushdown.cpp @@ -1,9 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -#include "duckdb/catalog/catalog.hpp" -#include "duckdb/planner/operator/logical_projection.hpp" #include "scalar_fn_pushdown.hpp" -#include "table_function.hpp" + +#include "duckdb/planner/operator/logical_projection.hpp" +#include "duckdb/planner/operator/logical_get.hpp" + #include /** @@ -12,114 +13,6 @@ * If there are any functions left, this means they were not pushed down and * may produce conflicts (e.g. WHERE prefix("str", 'h')). */ - -LogicalOperatorPtr TryPushdownScalarFunctions(ClientContext &context, LogicalOperatorPtr plan) { - Analyses analyses; - Projections projections; - FindGetsAndProjections(*plan, analyses, projections); - if (analyses.empty()) { - return plan; - } - ScalarFnCollect(analyses, projections).VisitOperator(*plan); - - bool any_pushed = false; - for (auto &[_, analysis] : analyses) { - for (auto &[column_index, expr] : analysis.col_to_fn) { - if (expr == nullptr) { // Conflict for column - continue; - } - const TableColumnStorageIndex storage_index = analysis.StorageIndex(column_index); - TableFunctionProjectionExpressionInput input {analysis.get, *expr, storage_index}; - if (projection_expression_pushdown(context, input)) { - analysis.get.types[column_index] = expr->return_type; - analysis.get.returned_types[storage_index] = expr->return_type; - any_pushed = true; - } else { // failed to push down expression, can't replace it - expr = nullptr; - } - } - } - - if (any_pushed) { - ScalarFnReplace(analyses, projections).VisitOperator(*plan); - } - return plan; -} - -// A passthrough projection only forwards its child columns, e.g. a VIEW's -// "SELECT col". -static bool is_passthrough(const LogicalProjection &projection) { - if (projection.expressions.empty()) { - return false; // don't register empty projections in Projections - } - for (const auto &e : projection.expressions) { - if (e->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { - return false; - } - } - return true; -} - -void FindGetsAndProjections(LogicalOperator &op, Analyses &analyses, Projections &projections) { - using enum LogicalOperatorType; - switch (op.type) { - case LOGICAL_GET: { - if (auto &get = op.Cast(); get.function.bind == duckdb_vx_table_function_bind) { - analyses.emplace(get.table_index, GetAnalysis {get, {}}); - } - break; - } - case LOGICAL_PROJECTION: { - LogicalProjection &projection = op.Cast(); - D_ASSERT(projection.children.size() == 1); - auto &child = *projection.children[0]; - if (!is_passthrough(projection) || child.type != LOGICAL_GET) { - break; - } - // The GET itself is recorded when recursion reaches it below. Only - // passthrough projections wrapping a vortex GET act as aliases. - if (auto &get = child.Cast(); get.function.bind == duckdb_vx_table_function_bind) { - projections.emplace(projection.table_index, projection); - } - break; - } - default: - break; - } - - for (auto &child : op.children) { - FindGetsAndProjections(*child, analyses, projections); - } -} - -std::optional Resolve(ColumnBinding binding, Analyses &analyses, const Projections &projections) { - if (IsVirtualColumn(binding.column_index)) { - return std::nullopt; - } - if (const auto it = analyses.find(binding.table_index); it != analyses.end()) { - return {{it->second, binding.column_index, nullptr}}; - } - - const auto projection_it = projections.find(binding.table_index); - if (projection_it == projections.end()) { - return std::nullopt; - } - - LogicalProjection &projection = projection_it->second; - const ExpressionPtr &inner = projection.expressions[binding.column_index]; - if (inner->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) { - return std::nullopt; - } - const ColumnBinding &get_binding = inner->Cast().binding; - if (IsVirtualColumn(get_binding.column_index)) { - return std::nullopt; - } - if (const auto it = analyses.find(get_binding.table_index); it != analyses.end()) { - return {{it->second, get_binding.column_index, &projection}}; - } - return std::nullopt; -} - void ScalarFnCollect::VisitOperator(LogicalOperator &op) { /* * Logical projection expressions are columns which reference underlying @@ -143,7 +36,7 @@ ExpressionPtr ScalarFnCollect::VisitReplace(BoundColumnRefExpression &expr, Expr if (const auto binding = Resolve(expr.binding, analyses, projections)) { // Column is used without function applied to it, register a conflict. // Not emplace() as we need to update the value if it was present - binding->analysis.col_to_fn[binding->column_index] = nullptr; + binding->analysis.col_to_expr[binding->column_index] = nullptr; } return std::move(*ptr); } @@ -160,11 +53,11 @@ ExpressionPtr ScalarFnCollect::VisitReplace(BoundFunctionExpression &expr, Expre if (!binding) { return nullptr; } - auto &col_to_fn = binding->analysis.col_to_fn; + auto &col_to_expr = binding->analysis.col_to_expr; - if (auto it = col_to_fn.find(binding->column_index); it == col_to_fn.end()) { + if (auto it = col_to_expr.find(binding->column_index); it == col_to_expr.end()) { // This is the first time we see the column used by a single function. - col_to_fn.emplace(binding->column_index, &expr); + col_to_expr.emplace(binding->column_index, &expr); } else if (it->second == nullptr || !it->second->Equals(expr)) { // Either column is used with different function in "expr" or // there already is a conflict. @@ -174,11 +67,6 @@ ExpressionPtr ScalarFnCollect::VisitReplace(BoundFunctionExpression &expr, Expre return std::move(*ptr); } -static bool can_pushdown_column(const GetAnalysis &analysis, TableColumnScanIndex idx) { - const auto it = analysis.col_to_fn.find(idx); - return it != analysis.col_to_fn.end() && it->second != nullptr; -} - ExpressionPtr ScalarFnReplace::VisitReplace(BoundColumnRefExpression &expr, ExpressionPtr *ptr) { const auto binding = Resolve(expr.binding, analyses, projections); if (!binding) { @@ -186,9 +74,11 @@ ExpressionPtr ScalarFnReplace::VisitReplace(BoundColumnRefExpression &expr, Expr } const auto &[analysis, column_index, projection] = *binding; - if (can_pushdown_column(analysis, column_index)) { - expr.return_type = analysis.get.types[column_index]; - if (projection != nullptr) { + if (CanPushdownColumn(analysis, column_index)) { + const idx_t storage_index = analysis.get.GetColumnIds()[column_index].GetPrimaryIndex(); + const LogicalType return_type = analysis.get.returned_types[storage_index]; + expr.return_type = return_type; + if (projection != nullptr && !projection->types.empty()) { projection->types[column_index] = expr.return_type; } } @@ -209,13 +99,15 @@ ExpressionPtr ScalarFnReplace::VisitReplace(BoundFunctionExpression &expr, Expre } const auto &[analysis, column_index, projection] = *binding; - if (!can_pushdown_column(analysis, column_index)) { + if (!CanPushdownColumn(analysis, column_index)) { return std::move(*ptr); } - bound_col_base->return_type = analysis.get.types[column_index]; - if (projection != nullptr) { - projection->types[column_index] = bound_col_base->return_type; + const idx_t storage_index = analysis.get.GetColumnIds()[column_index].GetPrimaryIndex(); + const LogicalType return_type = analysis.get.returned_types[storage_index]; + bound_col_base->return_type = return_type; + if (projection != nullptr && !projection->types.empty()) { + projection->types[column_index] = return_type; } return std::move(bound_col_base); } @@ -227,7 +119,3 @@ ScalarFnCollect::ScalarFnCollect(Analyses &analyses, const Projections &projecti ScalarFnReplace::ScalarFnReplace(Analyses &analyses, const Projections &projections) : analyses(analyses), projections(projections) { } - -TableColumnStorageIndex GetAnalysis::StorageIndex(TableColumnScanIndex idx) const { - return get.GetColumnIds()[idx].GetPrimaryIndex(); -} diff --git a/vortex-duckdb/cpp/spatial_overrides.cpp b/vortex-duckdb/cpp/spatial_overrides.cpp new file mode 100644 index 00000000000..28f095a7b89 --- /dev/null +++ b/vortex-duckdb/cpp/spatial_overrides.cpp @@ -0,0 +1,157 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "spatial_overrides.hpp" + +#include "spatial_overrides.h" + +#include "duckdb/catalog/catalog.hpp" +#include "duckdb/catalog/catalog_entry/scalar_function_catalog_entry.hpp" +#include "duckdb/common/error_data.hpp" +#include "duckdb/function/function_binder.hpp" +#include "duckdb/function/scalar_function.hpp" +#include "duckdb/logging/logger.hpp" +#include "duckdb/main/capi/capi_internal.hpp" +#include "duckdb/main/connection.hpp" +#include "duckdb/main/database_manager.hpp" +#include "duckdb/parser/parsed_data/create_scalar_function_info.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/planner/logical_operator_visitor.hpp" +#include "duckdb/transaction/meta_transaction.hpp" + +#include +#include + +/// A spatial function Vortex shadows so that its filters can push into Vortex scans. +struct SpatialOverride { + const char *name; + idx_t arity; + void (*tweak)(ScalarFunction &); +}; + +/// Drop spatial's bind so the filter keeps the radius visible as `children[2]`. +static void DropBind(ScalarFunction &fn) { + fn.bind = nullptr; +} + +/// Clear the error mode so the filter pushes through view projections. +static void ClearErrorMode(ScalarFunction &fn) { + fn.SetErrorMode(FunctionErrors::CANNOT_ERROR); +} + +static constexpr SpatialOverride SPATIAL_OVERRIDES[] = { + {"st_dwithin", 3, DropBind}, + {"st_intersects", 2, ClearErrorMode}, + {"st_contains", 2, ClearErrorMode}, + {"st_within", 2, ClearErrorMode}, +}; + +/// Apply one override, later calls to the function bind to the pushable copy instead of +/// spatial's original. +static void RegisterSpatialOverride(ClientContext &context, const SpatialOverride &fn_override) { + auto &system = Catalog::GetSystemCatalog(context); + auto entry = system.GetEntry(context, + DEFAULT_SCHEMA, + fn_override.name, + OnEntryNotFound::RETURN_NULL); + if (!entry) { + return; + } + ScalarFunctionSet set(fn_override.name); + for (const auto &overload : entry->functions.functions) { + ScalarFunction copy = overload; + fn_override.tweak(copy); + set.AddFunction(copy); + } + CreateScalarFunctionInfo info(std::move(set)); + info.on_conflict = OnCreateConflict::REPLACE_ON_CONFLICT; + + info.internal = false; + // Register in the default database's catalog: unqualified calls resolve there before the + // system catalog, so the copy shadows spatial's entry. + auto &catalog = Catalog::GetCatalog(context, DatabaseManager::GetDefaultDatabase(context)); + // A file-backed catalog rejects writes unless the database is marked modified; the mark is + // harmless here because function entries are never persisted to disk. + MetaTransaction::Get(context).ModifyDatabase(catalog.GetAttached(), DatabaseModificationType()); + catalog.CreateFunction(context, info); +} + +/// Apply every override in `SPATIAL_OVERRIDES` in one transaction. +extern "C" duckdb_state duckdb_vx_register_spatial_overrides(duckdb_database ffi_db) { + if (!ffi_db) { + return DuckDBError; + } + const DatabaseWrapper &wrapper = *reinterpret_cast(ffi_db); + DatabaseInstance &db = *wrapper.database->instance; + try { + Connection conn(db); + ClientContext &context = *conn.context; + context.RunFunctionInTransaction([&]() { + for (const auto &fn_override : SPATIAL_OVERRIDES) { + RegisterSpatialOverride(context, fn_override); + } + }); + } catch (const std::exception &e) { + ErrorData data(e); + DUCKDB_LOG_ERROR(db, "Failed to register the spatial overrides:\t" + data.Message()); + return DuckDBError; + } + return DuckDBSuccess; +} + +SpatialOverrideRestore::SpatialOverrideRestore(ClientContext &context) : context(context) { +} + +void SpatialOverrideRestore::VisitOperator(LogicalOperator &op) { + using enum LogicalOperatorType; + switch (op.type) { + case LOGICAL_COMPARISON_JOIN: + case LOGICAL_ANY_JOIN: + case LOGICAL_DELIM_JOIN: + case LOGICAL_ASOF_JOIN: + VisitOperatorExpressions(op); + break; + default: + break; + } + VisitOperatorChildren(op); +} + +unique_ptr SpatialOverrideRestore::VisitReplace(BoundFunctionExpression &expr, + unique_ptr *) { + const bool overridden = + std::any_of(std::begin(SPATIAL_OVERRIDES), + std::end(SPATIAL_OVERRIDES), + [&](const SpatialOverride &o) { + return expr.function.name == o.name && expr.children.size() == o.arity; + }); + if (!overridden) { + return nullptr; // Not an overridden call: leave it as is. + } + // Spatial's original lives in the system catalog, where the override cannot shadow it. + auto original = + Catalog::GetSystemCatalog(context).GetEntry(context, + DEFAULT_SCHEMA, + expr.function.name, + OnEntryNotFound::RETURN_NULL); + if (!original) { + return nullptr; + } + // Rebind a copy of the call's arguments through the original function. + vector> children; + children.reserve(expr.children.size()); + for (const auto &child : expr.children) { + children.push_back(child->Copy()); + } + ErrorData error; + FunctionBinder binder(context); + auto bound = binder.BindScalarFunction(*original, std::move(children), error); + if (!bound) { + return nullptr; // No matching overload: the override call still executes, keep it. + } + return bound; +} + +void RestoreSpatialOverrides(ClientContext &context, LogicalOperator &plan) { + SpatialOverrideRestore(context).VisitOperator(plan); +} diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 2b351593450..42c4285c383 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -4,6 +4,7 @@ #include "data.hpp" #include "error.hpp" #include "table_function.hpp" +#include "expr.h" #include "vortex_duckdb.h" #include "table_function.h" #include "vortex.h" @@ -171,12 +172,16 @@ struct CTableBindResult { vector &names; }; +// This is a flaw of Duckdb API which doesn't allow passing non-const +// expressions. We never modify the value on Rust side. +static duckdb_vx_expr get_ffi_expr(const Expression &expr) { + return reinterpret_cast(const_cast(&expr)); +} + bool projection_expression_pushdown(ClientContext &, const TableFunctionProjectionExpressionInput &input) { const auto &bind = input.get.bind_data->Cast(); - // This is a flaw of Duckdb API which doesn't allow passing non-const - // expressions. We never modify the value on Rust side. - auto ffi_expr = reinterpret_cast(const_cast(&input.expression)); + duckdb_vx_expr ffi_expr = get_ffi_expr(input.expression); void *const ffi_bind = bind.ffi_data->DataPtr(); duckdb_vx_error error_out = nullptr; @@ -191,6 +196,33 @@ bool projection_expression_pushdown(ClientContext &, const TableFunctionProjecti return ret; } +extern "C" { +idx_t duckdb_vx_aggregate_len(duckdb_vx_agg_input ffi_input) { + return reinterpret_cast(ffi_input)->projections.size(); +} + +duckdb_vx_expr duckdb_vx_aggregate_at(duckdb_vx_agg_input ffi_input, idx_t i, idx_t *proj_idx) { + const auto &input = *reinterpret_cast(ffi_input); + const auto &[scan_index, expr] = input.projections[i]; + *proj_idx = scan_index == COUNT_STAR_PROJ_IDX ? scan_index + : input.get.GetColumnIds()[scan_index].GetPrimaryIndex(); + return get_ffi_expr(expr); +} +} + +bool aggregate_pushdown(ClientContext &, const TableFunctionUngroupedAggregateInput &input) { + const auto &bind = input.get.bind_data->Cast(); + void *const ffi_bind = bind.ffi_data->DataPtr(); + duckdb_vx_error error_out = nullptr; + const auto ffi_input = + reinterpret_cast(const_cast(&input)); + const bool res = duckdb_table_function_pushdown_projection_aggregates(ffi_bind, ffi_input, &error_out); + if (error_out) { + throw BinderException(IntoErrString(error_out)); + } + return res; +} + /** * Called for every new query. For example, if there is a VIEW over *.vortex, * and after a query another file is added matching the glob, for second query @@ -238,10 +270,11 @@ unique_ptr c_init_global(ClientContext &context, Table } unique_ptr -init_local(ExecutionContext &, TableFunctionInitInput &, GlobalTableFunctionState *global_state) { +init_local(ExecutionContext &, TableFunctionInitInput &input, GlobalTableFunctionState *global_state) { + const void *const ffi_bind = input.bind_data->Cast().ffi_data->DataPtr(); void *const ffi_global = global_state->Cast().ffi_data->DataPtr(); - duckdb_vx_data ffi_local_data = duckdb_table_function_init_local(ffi_global); + duckdb_vx_data ffi_local_data = duckdb_table_function_init_local(ffi_bind, ffi_global); auto cdata = unique_ptr(reinterpret_cast(ffi_local_data)); return make_uniq(std::move(cdata)); } diff --git a/vortex-duckdb/cpp/vortex_duckdb.cpp b/vortex-duckdb/cpp/vortex_duckdb.cpp index 091f98a1703..05ee9098583 100644 --- a/vortex-duckdb/cpp/vortex_duckdb.cpp +++ b/vortex-duckdb/cpp/vortex_duckdb.cpp @@ -1,9 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#include "aggregate_fn_pushdown.hpp" #include "data.hpp" #include "error.hpp" #include "scalar_fn_pushdown.hpp" +#include "spatial_overrides.hpp" +#include "cast_pushdown.hpp" #include "vortex_duckdb.h" #include "duckdb/catalog/catalog.hpp" @@ -268,11 +271,18 @@ extern "C" duckdb_blob duckdb_vx_value_get_geometry(duckdb_value value) { } static void VortexOptimizeFunction(OptimizerExtensionInput &input, unique_ptr &plan) { - plan = TryPushdownScalarFunctions(input.context, std::move(plan)); + plan = TryPushdown(input.context, std::move(plan)); + plan = TryPushdown(input.context, std::move(plan)); + plan = TryPushdownAggregateFunctions(input.context, std::move(plan)); +} + +static void VortexPreOptimizeFunction(OptimizerExtensionInput &input, unique_ptr &plan) { + RestoreSpatialOverrides(input.context, *plan); } struct VortexOptimizerExtension final : OptimizerExtension { - inline VortexOptimizerExtension() : OptimizerExtension(VortexOptimizeFunction, nullptr, {}) { + inline VortexOptimizerExtension() + : OptimizerExtension(VortexOptimizeFunction, VortexPreOptimizeFunction, {}) { } }; diff --git a/vortex-duckdb/include/vortex.h b/vortex-duckdb/include/vortex.h index c63b20cbbb0..e8e4550e108 100644 --- a/vortex-duckdb/include/vortex.h +++ b/vortex-duckdb/include/vortex.h @@ -11,6 +11,8 @@ #pragma once +#define COUNT_STAR_PROJ_IDX UINT64_MAX + #ifdef __cplusplus extern "C" { #endif // __cplusplus @@ -40,6 +42,11 @@ bool duckdb_table_function_pushdown_projection_expression(void *bind_data, size_t column_id, duckdb_vx_error *error_out); +extern +bool duckdb_table_function_pushdown_projection_aggregates(void *bind_data, + duckdb_vx_agg_input input, + duckdb_vx_error *error_out); + extern void duckdb_table_function_scan(void *global_init_data, void *local_init_data, @@ -56,7 +63,9 @@ extern duckdb_vx_data duckdb_table_function_init_global(const duckdb_vx_tfunc_init_input *init_input, duckdb_vx_error *error_out); -extern duckdb_vx_data duckdb_table_function_init_local(void *global_init_data); +extern +duckdb_vx_data duckdb_table_function_init_local(const void *bind_data, + void *global_init_data); extern duckdb_vx_data duckdb_table_function_bind(duckdb_vx_tfunc_bind_input bind_input, diff --git a/vortex-duckdb/src/convert/dtype.rs b/vortex-duckdb/src/convert/dtype.rs index 03aed0631aa..138a6f1bc69 100644 --- a/vortex-duckdb/src/convert/dtype.rs +++ b/vortex-duckdb/src/convert/dtype.rs @@ -58,6 +58,12 @@ use vortex::extension::datetime::Time; use vortex::extension::datetime::TimeUnit; use vortex::extension::datetime::Timestamp; use vortex_geo::extension::GeoMetadata; +use vortex_geo::extension::LineString; +use vortex_geo::extension::MultiLineString; +use vortex_geo::extension::MultiPoint; +use vortex_geo::extension::MultiPolygon; +use vortex_geo::extension::Point; +use vortex_geo::extension::Polygon; use vortex_geo::extension::WellKnownBinary; use crate::cpp::DUCKDB_TYPE; @@ -88,8 +94,8 @@ impl FromLogicalType for DType { DUCKDB_TYPE::DUCKDB_TYPE_USMALLINT => DType::Primitive(U16, nullability), DUCKDB_TYPE::DUCKDB_TYPE_UINTEGER => DType::Primitive(U32, nullability), DUCKDB_TYPE::DUCKDB_TYPE_UBIGINT => DType::Primitive(U64, nullability), - DUCKDB_TYPE::DUCKDB_TYPE_HUGEINT => todo!(), - DUCKDB_TYPE::DUCKDB_TYPE_UHUGEINT => todo!(), + DUCKDB_TYPE::DUCKDB_TYPE_HUGEINT => vortex_bail!("I128 is not in Vortex type system"), + DUCKDB_TYPE::DUCKDB_TYPE_UHUGEINT => vortex_bail!("U128 is not in Vortex type system"), DUCKDB_TYPE::DUCKDB_TYPE_FLOAT => DType::Primitive(F32, nullability), DUCKDB_TYPE::DUCKDB_TYPE_DOUBLE => DType::Primitive(F64, nullability), DUCKDB_TYPE::DUCKDB_TYPE_VARCHAR => DType::Utf8(nullability), @@ -173,17 +179,19 @@ impl FromLogicalType for DType { ) } DUCKDB_TYPE::DUCKDB_TYPE_VARIANT => DType::Variant(nullability), - DUCKDB_TYPE::DUCKDB_TYPE_TIME_TZ => todo!(), - DUCKDB_TYPE::DUCKDB_TYPE_INTERVAL => todo!(), - DUCKDB_TYPE::DUCKDB_TYPE_ENUM => todo!(), - DUCKDB_TYPE::DUCKDB_TYPE_MAP => todo!(), - DUCKDB_TYPE::DUCKDB_TYPE_UUID => todo!(), - DUCKDB_TYPE::DUCKDB_TYPE_UNION => todo!(), - DUCKDB_TYPE::DUCKDB_TYPE_BIT => todo!(), - DUCKDB_TYPE::DUCKDB_TYPE_ANY => todo!(), - DUCKDB_TYPE::DUCKDB_TYPE_BIGNUM => todo!(), - DUCKDB_TYPE::DUCKDB_TYPE_STRING_LITERAL => todo!(), - DUCKDB_TYPE::DUCKDB_TYPE_INTEGER_LITERAL => todo!(), + other @ (DUCKDB_TYPE::DUCKDB_TYPE_TIME_TZ + | DUCKDB_TYPE::DUCKDB_TYPE_INTERVAL + | DUCKDB_TYPE::DUCKDB_TYPE_ENUM + | DUCKDB_TYPE::DUCKDB_TYPE_MAP + | DUCKDB_TYPE::DUCKDB_TYPE_UUID + | DUCKDB_TYPE::DUCKDB_TYPE_UNION + | DUCKDB_TYPE::DUCKDB_TYPE_BIT + | DUCKDB_TYPE::DUCKDB_TYPE_ANY + | DUCKDB_TYPE::DUCKDB_TYPE_BIGNUM + | DUCKDB_TYPE::DUCKDB_TYPE_STRING_LITERAL + | DUCKDB_TYPE::DUCKDB_TYPE_INTEGER_LITERAL) => { + vortex_bail!("{other:?} -> DType conversion is not supported") + } }) } } @@ -235,19 +243,26 @@ impl TryFrom<&DType> for LogicalType { DType::Struct(struct_type, _) => { return LogicalType::try_from(struct_type); } - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), - DType::Variant(_) => { - vortex_bail!("Vortex Variant array aren't supported in DuckDB") - } + // TODO(connor): Union + DType::Union(..) => vortex_bail!("Vortex Union isn't supported"), + DType::Variant(_) => vortex_bail!("Vortex Variant array aren't supported"), DType::Extension(ext_dtype) => { // Handle first-party extension types that have DuckDB equivalents. if let Some(temporal) = ext_dtype.metadata_opt::() { return temporal_to_duckdb(temporal); } - if let Some(wkb) = ext_dtype.metadata_opt::() { - let crs = wkb.crs.as_ref(); - return LogicalType::geometry_type(crs.map(|crs| crs.as_str())); + // Native geometry types and WKB all surface to DuckDB as GEOMETRY so `ST_*` bind. + if let Some(geo) = ext_dtype + .metadata_opt::() + .or_else(|| ext_dtype.metadata_opt::()) + .or_else(|| ext_dtype.metadata_opt::()) + .or_else(|| ext_dtype.metadata_opt::()) + .or_else(|| ext_dtype.metadata_opt::()) + .or_else(|| ext_dtype.metadata_opt::()) + .or_else(|| ext_dtype.metadata_opt::()) + { + return LogicalType::geometry_type(geo.crs.as_deref()); } vortex_bail!("Unsupported extension type \"{}\"", ext_dtype.id()); diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 387b644fe30..aee065a2a20 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -1,9 +1,23 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::fmt::Display; +use std::fmt::Formatter; +use std::fmt::Result; use std::sync::Arc; use tracing::debug; +use vortex::aggregate_fn::Accumulator; +use vortex::aggregate_fn::DynAccumulator; +use vortex::aggregate_fn::EmptyOptions as AggregateEmptyOptions; +use vortex::aggregate_fn::NumericalAggregateOpts; +use vortex::aggregate_fn::combined::PairOptions; +use vortex::aggregate_fn::fns::count::Count; +use vortex::aggregate_fn::fns::first::First; +use vortex::aggregate_fn::fns::max::Max; +use vortex::aggregate_fn::fns::mean::Mean; +use vortex::aggregate_fn::fns::min::Min; +use vortex::aggregate_fn::fns::sum::Sum; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; @@ -28,6 +42,7 @@ use vortex::expr::not; use vortex::expr::or_collect; use vortex::expr::root; use vortex::scalar::Scalar; +use vortex::scalar_fn::EmptyOptions as ScalarEmptyOptions; use vortex::scalar_fn::ScalarFnVTableExt; use vortex::scalar_fn::fns::between::Between; use vortex::scalar_fn::fns::between::BetweenOptions; @@ -37,7 +52,19 @@ use vortex::scalar_fn::fns::like::Like; use vortex::scalar_fn::fns::like::LikeOptions; use vortex::scalar_fn::fns::literal::Literal; use vortex::scalar_fn::fns::operators::Operator; +use vortex_geo::extension::LineString; +use vortex_geo::extension::MultiLineString; +use vortex_geo::extension::MultiPoint; +use vortex_geo::extension::MultiPolygon; +use vortex_geo::extension::Point; +use vortex_geo::extension::Polygon; +use vortex_geo::extension::WellKnownBinary; +use vortex_geo::extension::native_geometry_scalar_from_wkb; +use vortex_geo::scalar_fn::contains::GeoContains; +use vortex_geo::scalar_fn::distance::GeoDistance; +use vortex_geo::scalar_fn::intersects::GeoIntersects; +use crate::convert::dtype::FromLogicalType; use crate::cpp::DUCKDB_TYPE; use crate::cpp::DUCKDB_VX_EXPR_TYPE; use crate::duckdb; @@ -45,6 +72,7 @@ use crate::duckdb::BoundFunction; use crate::duckdb::BoundOperator; use crate::duckdb::ExpressionClass; use crate::duckdb::ExpressionClass::BoundBetween; +use crate::duckdb::ExpressionClass::BoundCast; use crate::duckdb::ExpressionClass::BoundColumnRef; use crate::duckdb::ExpressionClass::BoundComparison; use crate::duckdb::ExpressionClass::BoundConjunction; @@ -73,29 +101,169 @@ fn build_list_length(expr: Expression, nullability: Nullability) -> Expression { cast(list_length(expr), DType::Primitive(PType::I64, nullability)) } +/// Read an `f64` from a constant expression (the `ST_DWithin` radius); `None` for non-constants. +fn from_bound_f64(value: &duckdb::ExpressionRef) -> VortexResult> { + match value.as_class().vortex_expect("unknown class") { + BoundConstant(constant) => Ok(Some(f64::try_from(&Scalar::try_from(constant.value)?)?)), + _ => Ok(None), + } +} + +/// Context threaded through expression conversion. +#[derive(Clone, Copy)] +struct ConvertCtx<'a> { + /// Substituted for `BoundRef` references when converting scan-scoped table filters. + col_sub: Option<&'a Expression>, + /// The scan's fields, when known. + fields: Option<&'a [DuckdbField]>, +} + +/// Whether `name` is a non-nullable native geometry column of the scan. The pushed geo kernels +/// reject nullable operands and cannot evaluate `vortex.geo.wkb` columns, which also surface to +/// DuckDB as `GEOMETRY`. +fn is_native_geo_column(fields: Option<&[DuckdbField]>, name: &str) -> bool { + fields + .into_iter() + .flatten() + .filter(|field| field.name == name && !field.dtype.is_nullable()) + .any(|field| match field.dtype.as_extension_opt() { + Some(ext) => { + ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + } + None => false, + }) +} + +/// Lower a geo operand: a `GEOMETRY` literal arrives as WKB, decoded once to its native type so the +/// pushed `GeoDistance` stays native; a column must be native geometry. `None` skips the push. +fn geo_operand( + value: &duckdb::ExpressionRef, + ctx: ConvertCtx<'_>, +) -> VortexResult> { + match value.as_class() { + Some(BoundConstant(constant)) => { + let scalar = Scalar::try_from(constant.value)?; + let DType::Extension(ext_dtype) = scalar.dtype() else { + return Ok(None); + }; + if !ext_dtype.is::() { + return Ok(None); + } + let storage = scalar.as_extension().to_storage_scalar(); + let Some(buf) = storage.as_binary_opt().and_then(|b| b.value()) else { + return Ok(None); + }; + Ok(native_geometry_scalar_from_wkb(buf.as_slice())?.map(lit)) + } + Some(BoundColumnRef(col_ref)) + if is_native_geo_column(ctx.fields, col_ref.name.as_ref()) => + { + try_from_expression_inner(value, ctx) + } + _ => Ok(None), + } +} + +/// Lower all geometry operands of a geo function. Returns `None`, skipping the push, when any +/// operand is neither a constant geometry nor a native geometry column. +fn geo_operands( + children: &[&duckdb::ExpressionRef], + ctx: ConvertCtx<'_>, +) -> VortexResult>> { + children + .iter() + .map(|child| geo_operand(child, ctx)) + .collect() +} + +/// Lower geo UDFs to native Vortex geo ops so the work runs in the scan. `None` otherwise. +fn try_from_geo_function( + name: &str, + func: &BoundFunction, + ctx: ConvertCtx<'_>, +) -> VortexResult> { + let children: Vec<_> = func.children().collect(); + let expr = match name.to_ascii_lowercase().as_str() { + // Spatial's own st_dwithin folds the radius into bind data; the override + // (cpp/spatial_overrides.cpp) keeps it visible here as `children[2]`. + "st_dwithin" => { + if children.len() != 3 { + return Ok(None); + } + let Some(operands) = geo_operands(&children[..2], ctx)? else { + return Ok(None); + }; + // A non-constant radius is left for DuckDB to evaluate. + let Some(distance) = from_bound_f64(children[2])? else { + return Ok(None); + }; + let geo_distance = GeoDistance.new_expr(ScalarEmptyOptions, operands); + Binary.new_expr(Operator::Lte, [geo_distance, lit(distance)]) + } + "st_distance" => { + if children.len() != 2 { + return Ok(None); + } + let Some(operands) = geo_operands(&children, ctx)? else { + return Ok(None); + }; + GeoDistance.new_expr(ScalarEmptyOptions, operands) + } + "st_intersects" => { + if children.len() != 2 { + return Ok(None); + } + let Some(operands) = geo_operands(&children, ctx)? else { + return Ok(None); + }; + GeoIntersects.new_expr(ScalarEmptyOptions, operands) + } + containment @ ("st_contains" | "st_within") => { + if children.len() != 2 { + return Ok(None); + } + let Some(mut operands) = geo_operands(&children, ctx)? else { + return Ok(None); + }; + // `st_within(a, b)` is `st_contains(b, a)`; both lower to the contains kernel. + if containment == "st_within" { + operands.swap(0, 1); + } + GeoContains.new_expr(ScalarEmptyOptions, operands) + } + _ => return Ok(None), + }; + + Ok(Some(expr)) +} + fn try_from_bound_function( func: &BoundFunction, - col_sub: Option<&Expression>, + ctx: ConvertCtx<'_>, ) -> VortexResult> { let expr = match func.scalar_function.name() { "strlen" => { let children: Vec<_> = func.children().collect(); vortex_ensure!(children.len() == 1); - let Some(col) = try_from_expression_inner(children[0], col_sub)? else { + let Some(col) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; let col = byte_length(col); // byte_length returns u64, strlen expects i64. // At this point we don't know column's dtype so we ultimately - // set it to be nullable. For non-nullable column the nullability - // will be AllValid so it's a marginal cost. + // set it to be nullable. let dtype = DType::Primitive(PType::I64, Nullability::Nullable); cast(col, dtype) } "struct_extract" => { let children: Vec<_> = func.children().collect(); vortex_ensure!(children.len() == 2); - let Some(child) = try_from_expression_inner(children[0], col_sub)? else { + let Some(child) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; let field = from_bound_str(children[1])?; @@ -104,10 +272,10 @@ fn try_from_bound_function( like @ ("~~" | "!~~") => { let children: Vec<_> = func.children().collect(); vortex_ensure!(children.len() == 2); - let Some(string) = try_from_expression_inner(children[0], col_sub)? else { + let Some(string) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; - let Some(target) = try_from_expression_inner(children[1], col_sub)? else { + let Some(target) = try_from_expression_inner(children[1], ctx)? else { return Ok(None); }; let opts = LikeOptions { @@ -119,7 +287,7 @@ fn try_from_bound_function( matchers @ ("contains" | "prefix" | "suffix") => { let children: Vec<_> = func.children().collect(); vortex_ensure!(children.len() == 2); - let Some(value) = try_from_expression_inner(children[0], col_sub)? else { + let Some(value) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; let pattern = from_bound_str(children[1])?; @@ -137,11 +305,11 @@ fn try_from_bound_function( if children.len() != 1 { return Ok(None); } - let Some(col) = try_from_expression_inner(children[0], col_sub)? else { + let Some(col) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; - // We don't know the column's nullability here, so we set it to nullable. + // We don't know the column's nullability here build_list_length(col, Nullability::Nullable) } // len/length semantics depend on the return type of underlying expr. @@ -151,21 +319,19 @@ fn try_from_bound_function( let child = children[0]; if returns_a_list(child) { - let Some(col) = try_from_expression_inner(child, col_sub)? else { + let Some(col) = try_from_expression_inner(child, ctx)? else { return Ok(None); }; - // Same nullability rationale as in "array_length" branch. + // We don't know the column's nullability here let list_len_expr = build_list_length(col, Nullability::Nullable); return Ok(Some(list_len_expr)); } else { return Ok(None); } } - _ => { - debug!("bound function {}", func.scalar_function.name()); - return Ok(None); - } + // Geo UDFs are handled here; non-geo names return `None` inside. + name => return try_from_geo_function(name, func, ctx), }; Ok(Some(expr)) @@ -173,15 +339,30 @@ fn try_from_bound_function( pub fn try_from_bound_expression( value: &duckdb::ExpressionRef, + fields: &[DuckdbField], ) -> VortexResult> { - try_from_expression_inner(value, None) + try_from_expression_inner( + value, + ConvertCtx { + col_sub: None, + fields: Some(fields), + }, + ) } pub(super) fn try_from_bound_expression_with_col_sub( value: &duckdb::ExpressionRef, col_sub: &Expression, ) -> VortexResult> { - try_from_expression_inner(value, Some(col_sub)) + // No fields: scan-time table filters never carry geo functions, because + // `can_push_expression` refuses them. + try_from_expression_inner( + value, + ConvertCtx { + col_sub: Some(col_sub), + fields: None, + }, + ) } fn is_supported_length_alias(func: &BoundFunction) -> bool { @@ -189,6 +370,19 @@ fn is_supported_length_alias(func: &BoundFunction) -> bool { children.len() == 1 && returns_a_list(children[0]) } +// We limit casting to Primitive types, because some conversions yield an error +// like vortex.date[days](i32) -> vortex.timestamp[µs](i64?). However, when we +// push down the cast, we don't have access to column's dtype, so we need to +// be overly restrictive. +// TODO(myrrc) change after https://github.com/vortex-data/vortex/issues/8570 +// is resolved +// +// We also don't push floats and doubles because Vortex truncates to zero and +// Duckdb rounds the result +fn can_push_cast(cast: &duckdb::BoundCast<'_>, target: &duckdb::LogicalTypeRef) -> bool { + !cast.is_try && target.is_primitive_integer() && cast.child.return_type().is_primitive_integer() +} + // Called before pushdown_complex_filter or a table filter expression call. // As we support complex filter pushdown, Duckdb pushes expressions to Vortex. // However, it doesn't know what type of expressions we can handle. Here we list @@ -199,15 +393,19 @@ fn is_supported_length_alias(func: &BoundFunction) -> bool { // // Example: we don't support substr() expression so we tell Duckdb we can't // push it. +// Example: we support CAST but not TRY_CAST. // Example: optional filters may fail to parse on our side (we return // Ok(None)), so we don't allow pushing these. pub fn can_push_expression(value: &duckdb::ExpressionRef) -> bool { - let Some(value) = value.as_class() else { + let Some(class) = value.as_class() else { return false; }; - match value { + match class { BoundColumnRef(_) => true, BoundConstant(_) => true, + BoundCast(cast) => { + can_push_cast(&cast, value.return_type()) && can_push_expression(cast.child) + } BoundRef => true, BoundComparison(comp) => can_push_expression(comp.left) && can_push_expression(comp.right), BoundBetween(between) => { @@ -227,6 +425,9 @@ pub fn can_push_expression(value: &duckdb::ExpressionRef) -> bool { || name == "strlen" || name == "array_length" || (matches!(name, "len" | "length") && is_supported_length_alias(&func)) + // Geo functions are absent on purpose: they push only via + // `pushdown_complex_filter`, which has the scan's fields to verify the geometry + // columns are native. } ExpressionClass::BoundOperator(op) => { if !matches!( @@ -241,6 +442,7 @@ pub fn can_push_expression(value: &duckdb::ExpressionRef) -> bool { } op.children().all(can_push_expression) } + ExpressionClass::BoundAggregate(_) => false, } } @@ -255,47 +457,126 @@ pub fn try_from_projection_expression( value: &duckdb::ExpressionRef, field: &DuckdbField, ) -> VortexResult> { - let Some(value) = value.as_class() else { + let Some(class) = value.as_class() else { return Ok(None); }; - let ExpressionClass::BoundFunction(func) = value else { - return Ok(None); - }; - Ok(match func.scalar_function.name() { - "strlen" => { - let col = byte_length(get_item(field.name.as_str(), root())); - // byte_length returns u64, strlen expects i64 - let dtype = DType::Primitive(PType::I64, field.dtype.nullability()); - let col = cast(col, dtype); - Some(col) + Ok(match class { + ExpressionClass::BoundFunction(func) => { + match func.scalar_function.name() { + "strlen" => { + let col = byte_length(get_item(field.name.as_str(), root())); + // byte_length returns u64, strlen expects i64 + let dtype = DType::Primitive(PType::I64, field.dtype.nullability()); + let col = cast(col, dtype); + Some(col) + } + "array_length" => { + // Only accept array_length(expr) rather than array_length(expr, dim). + (func.children().count() == 1).then(|| list_length_on_field(field)) + } + // len/length have different semantics depending on field dtype. + "len" | "length" => { + matches!(field.dtype, DType::List(..) | DType::FixedSizeList(..)) + .then(|| list_length_on_field(field)) + } + _ => None, + } } - "array_length" => { - // Only accept array_length(expr) rather than array_length(expr, dim). - (func.children().count() == 1).then(|| list_length_on_field(field)) + BoundCast(c) => { + let target = value.return_type(); + if !can_push_cast(&c, target) { + None + } else { + let dtype = DType::from_logical_type(target, field.dtype.nullability())?; + let col = get_item(field.name.as_str(), root()); + Some(cast(col, dtype)) + } } - // len/length have different semantics depending on field dtype. - "len" | "length" => matches!(field.dtype, DType::List(..) | DType::FixedSizeList(..)) - .then(|| list_length_on_field(field)), _ => None, }) } +/// Aggregations we have pushed down in Vortex +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PushedAggregate { + Min, + Max, + Sum, + Mean, + // Also used for ANY_VALUE() which is allowed by definition + First, + // Valid values in column + Count, +} + +impl Display for PushedAggregate { + fn fmt(&self, f: &mut Formatter<'_>) -> Result { + match self { + PushedAggregate::Min => f.write_str("min"), + PushedAggregate::Max => f.write_str("max"), + PushedAggregate::Sum => f.write_str("sum"), + PushedAggregate::Mean => f.write_str("mean"), + PushedAggregate::First => f.write_str("first"), + PushedAggregate::Count => f.write_str("count"), + } + } +} + +impl PushedAggregate { + pub fn build(self, dtype: DType) -> VortexResult> { + let opts = NumericalAggregateOpts::default(); + Ok(match self { + Self::Min => Box::new(Accumulator::try_new(Min, opts, dtype)?), + Self::Max => Box::new(Accumulator::try_new(Max, opts, dtype)?), + Self::Sum => Box::new(Accumulator::try_new(Sum, opts, dtype)?), + Self::Mean => Box::new(Accumulator::try_new( + Mean::combined(), + PairOptions(opts, opts), + dtype, + )?), + Self::First => Box::new(Accumulator::try_new(First, AggregateEmptyOptions, dtype)?), + Self::Count => Box::new(Accumulator::try_new(Count, opts, dtype)?), + }) + } +} + +/// Check if this is an aggregate function we can handle in Vortex +pub fn try_from_projection_aggregate( + expr: &duckdb::ExpressionRef, +) -> VortexResult> { + let Some(expr) = expr.as_class() else { + return Ok(None); + }; + let ExpressionClass::BoundAggregate(agg) = expr else { + return Ok(None); + }; + Ok(Some(match agg.aggregate_function.name() { + "min" => PushedAggregate::Min, + "max" => PushedAggregate::Max, + "sum" | "sum_no_overflow" => PushedAggregate::Sum, + "avg" | "mean" => PushedAggregate::Mean, + "first" | "any_value" => PushedAggregate::First, + "count" => PushedAggregate::Count, + _ => return Ok(None), + })) +} + // If you want to add support for other expressions, also change // can_push_expression fn try_from_expression_inner( value: &duckdb::ExpressionRef, - col_sub: Option<&Expression>, + ctx: ConvertCtx<'_>, ) -> VortexResult> { - let Some(value) = value.as_class() else { + let Some(class) = value.as_class() else { debug!( class_id = ?value.as_class_id(), "unknown expression class id" ); return Ok(None); }; - Ok(Some(match value { + Ok(Some(match class { BoundRef => { - let Some(col) = col_sub else { + let Some(col) = ctx.col_sub else { vortex_bail!("BoundRef requested but no column supplied"); }; col.clone() @@ -305,23 +586,23 @@ fn try_from_expression_inner( BoundComparison(compare) => { let operator: Operator = compare.op.try_into()?; - let Some(left) = try_from_expression_inner(compare.left, col_sub)? else { + let Some(left) = try_from_expression_inner(compare.left, ctx)? else { return Ok(None); }; - let Some(right) = try_from_expression_inner(compare.right, col_sub)? else { + let Some(right) = try_from_expression_inner(compare.right, ctx)? else { return Ok(None); }; Binary.new_expr(operator, [left, right]) } BoundBetween(between) => { - let Some(array) = try_from_expression_inner(between.input, col_sub)? else { + let Some(array) = try_from_expression_inner(between.input, ctx)? else { return Ok(None); }; - let Some(lower) = try_from_expression_inner(between.lower, col_sub)? else { + let Some(lower) = try_from_expression_inner(between.lower, ctx)? else { return Ok(None); }; - let Some(upper) = try_from_expression_inner(between.upper, col_sub)? else { + let Some(upper) = try_from_expression_inner(between.upper, ctx)? else { return Ok(None); }; Between.new_expr( @@ -346,7 +627,7 @@ fn try_from_expression_inner( | DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_OPERATOR_IS_NOT_NULL => { let children: Vec<_> = operator.children().collect(); vortex_ensure!(children.len() == 1); - let Some(child) = try_from_expression_inner(children[0], col_sub)? else { + let Some(child) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; match operator.op { @@ -359,10 +640,10 @@ fn try_from_expression_inner( } } DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_IN => { - return try_from_compare_in(operator, col_sub, false); + return try_from_compare_in(operator, ctx, false); } DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_NOT_IN => { - return try_from_compare_in(operator, col_sub, true); + return try_from_compare_in(operator, ctx, true); } _ => { debug!(op=?operator.op, "cannot push down operator"); @@ -370,12 +651,24 @@ fn try_from_expression_inner( } }, ExpressionClass::BoundFunction(func) => { - return try_from_bound_function(&func, col_sub); + return try_from_bound_function(&func, ctx); + } + BoundCast(cast_inner) => { + let target = value.return_type(); + if !can_push_cast(&cast_inner, target) { + return Ok(None); + } + let Some(child) = try_from_expression_inner(cast_inner.child, ctx)? else { + return Ok(None); + }; + // We don't know the column's nullability here + let dtype = DType::from_logical_type(target, Nullability::Nullable)?; + cast(child, dtype) } BoundConjunction(conj) => { let Some(children) = conj .children() - .map(|c| try_from_expression_inner(c, col_sub)) + .map(|c| try_from_expression_inner(c, ctx)) .collect::>>>()? else { return Ok(None); @@ -390,18 +683,19 @@ fn try_from_expression_inner( _ => vortex_bail!("unexpected operator {:?} in bound conjunction", conj.op), } } + ExpressionClass::BoundAggregate(_) => return Ok(None), })) } fn try_from_compare_in( operator: BoundOperator, - col_sub: Option<&Expression>, + ctx: ConvertCtx<'_>, not_in: bool, ) -> VortexResult> { // First child is element, rest form the list. let children: Vec<_> = operator.children().collect(); assert!(children.len() >= 2); - let Some(element) = try_from_expression_inner(children[0], col_sub)? else { + let Some(element) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; @@ -409,7 +703,7 @@ fn try_from_compare_in( .iter() .skip(1) .map(|c| { - let Some(value) = try_from_expression_inner(c, col_sub)? else { + let Some(value) = try_from_expression_inner(c, ctx)? else { return Ok(None); }; Ok(Some( @@ -438,14 +732,14 @@ impl TryFrom for Operator { fn try_from(value: DUCKDB_VX_EXPR_TYPE) -> VortexResult { Ok(match value { - DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_INVALID => vortex_bail!("invalid expr"), + DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_INVALID => vortex_bail!("invalid expression"), DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_EQUAL => Operator::Eq, DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_NOTEQUAL => Operator::NotEq, DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_LESSTHAN => Operator::Lt, DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_GREATERTHAN => Operator::Gt, DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_LESSTHANOREQUALTO => Operator::Lte, DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_GREATERTHANOREQUALTO => Operator::Gte, - _ => todo!("cannot convert {:?}", value), + _ => vortex_bail!("cannot convert {:?}", value), }) } } diff --git a/vortex-duckdb/src/convert/mod.rs b/vortex-duckdb/src/convert/mod.rs index 0d641a73457..f1b5ba5bb10 100644 --- a/vortex-duckdb/src/convert/mod.rs +++ b/vortex-duckdb/src/convert/mod.rs @@ -8,8 +8,10 @@ mod table_filter; mod vector; pub use dtype::FromLogicalType; +pub use expr::PushedAggregate; pub use expr::can_push_expression; pub use expr::try_from_bound_expression; +pub use expr::try_from_projection_aggregate; pub use expr::try_from_projection_expression; pub use scalar::*; pub use table_filter::try_from_table_filter; diff --git a/vortex-duckdb/src/convert/scalar.rs b/vortex-duckdb/src/convert/scalar.rs index 2591786d3d9..7ad12e89ca3 100644 --- a/vortex-duckdb/src/convert/scalar.rs +++ b/vortex-duckdb/src/convert/scalar.rs @@ -78,11 +78,14 @@ impl ToDuckDBScalar for Scalar { DType::Decimal(..) => self.as_decimal().try_to_duckdb_scalar(), DType::Utf8(_) => self.as_utf8().try_to_duckdb_scalar(), DType::Binary(_) => self.as_binary().try_to_duckdb_scalar(), - DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) => todo!(), - DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), - DType::Variant(_) => { - vortex_bail!("Vortex Variant scalars aren't supported in DuckDB") + DType::List(..) => vortex_bail!("Vortex List scalars aren't supported"), + DType::FixedSizeList(..) => { + vortex_bail!("Vortex FixedSizeList scalars aren't supported") } + DType::Variant(_) => vortex_bail!("Vortex Variant scalars aren't supported"), + DType::Struct(..) => vortex_bail!("Vortex Struct scalars aren't supported"), + // TODO(connor): Union + DType::Union(..) => vortex_bail!("Vortex Union scalars aren't supported"), DType::Extension(..) => self.as_extension().try_to_duckdb_scalar(), } } diff --git a/vortex-duckdb/src/convert/vector.rs b/vortex-duckdb/src/convert/vector.rs index 1d12dab9dbe..a3fe098ffbb 100644 --- a/vortex-duckdb/src/convert/vector.rs +++ b/vortex-duckdb/src/convert/vector.rs @@ -361,7 +361,7 @@ pub fn flat_vector_to_vortex(vector: &VectorRef, len: usize) -> VortexResult unimplemented!("missing impl for {type_id:?}"), + type_id => vortex_bail!("{type_id:?} flat Vector to Vortex array not supported"), } } diff --git a/vortex-duckdb/src/duckdb/aggregate_pushdown.rs b/vortex-duckdb/src/duckdb/aggregate_pushdown.rs new file mode 100644 index 00000000000..2b245b2a614 --- /dev/null +++ b/vortex-duckdb/src/duckdb/aggregate_pushdown.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::AsPrimitive; + +use crate::cpp; +use crate::duckdb::Expression; +use crate::duckdb::ExpressionRef; +use crate::lifetime_wrapper; + +lifetime_wrapper!( + /// Aggregates we want to push at Vortex at once + AggregatePushdownInput, + cpp::duckdb_vx_agg_input, |_| {}); + +pub struct AggregateExpression<'a> { + pub expr: &'a ExpressionRef, + /// Output column projection id after the pass has expanded columns with + /// multiple aggregations per column + pub projection_id: u64, +} + +impl AggregatePushdownInputRef { + pub fn len(&self) -> usize { + unsafe { cpp::duckdb_vx_aggregate_len(self.as_ptr()) }.as_() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn get(&'_ self, index: usize) -> AggregateExpression<'_> { + let mut projection_id = 0u64; + let expr = unsafe { + cpp::duckdb_vx_aggregate_at(self.as_ptr(), index as u64, &raw mut projection_id) + }; + let expr = unsafe { Expression::borrow(expr) }; + AggregateExpression { + expr, + projection_id, + } + } +} diff --git a/vortex-duckdb/src/duckdb/database.rs b/vortex-duckdb/src/duckdb/database.rs index ab86503b291..1d258f05751 100644 --- a/vortex-duckdb/src/duckdb/database.rs +++ b/vortex-duckdb/src/duckdb/database.rs @@ -90,4 +90,15 @@ impl DatabaseRef { ); Ok(()) } + + /// Shadow the spatial functions that block filter pushdown with pushable copies; see + /// cpp/spatial_overrides.cpp for the list. Call after `LOAD spatial`; does nothing when + /// spatial is not loaded. + pub fn register_spatial_overrides(&self) -> VortexResult<()> { + duckdb_try!( + unsafe { cpp::duckdb_vx_register_spatial_overrides(self.as_ptr()) }, + "Failed to register the spatial function overrides" + ); + Ok(()) + } } diff --git a/vortex-duckdb/src/duckdb/expr.rs b/vortex-duckdb/src/duckdb/expr.rs index 48f255f744e..61dbbab5dcf 100644 --- a/vortex-duckdb/src/duckdb/expr.rs +++ b/vortex-duckdb/src/duckdb/expr.rs @@ -2,13 +2,14 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::ffi::CStr; -use std::ffi::c_void; use std::fmt::Display; use std::fmt::Formatter; use std::ptr; use crate::cpp; use crate::cpp::duckdb_vx_expr_class; +use crate::duckdb::AggregateFunction; +use crate::duckdb::AggregateFunctionRef; use crate::duckdb::DDBString; use crate::duckdb::LogicalType; use crate::duckdb::LogicalTypeRef; @@ -44,6 +45,14 @@ impl ExpressionRef { pub fn as_class(&self) -> Option> { Some( match unsafe { cpp::duckdb_vx_expr_get_class(self.as_ptr()) } { + cpp::DUCKDB_VX_EXPR_CLASS::DUCKDB_VX_EXPR_CLASS_BOUND_CAST => { + let child = unsafe { + Expression::borrow(cpp::duckdb_vx_expr_get_bound_cast_child(self.as_ptr())) + }; + let is_try = + unsafe { cpp::duckdb_vx_expr_get_bound_cast_is_try(self.as_ptr()) }; + ExpressionClass::BoundCast(BoundCast { child, is_try }) + } cpp::DUCKDB_VX_EXPR_CLASS::DUCKDB_VX_EXPR_CLASS_BOUND_COLUMN_REF => { let name = unsafe { let ptr = cpp::duckdb_vx_expr_get_bound_column_ref_get_name(self.as_ptr()); @@ -143,9 +152,15 @@ impl ExpressionRef { ExpressionClass::BoundFunction(BoundFunction { children, scalar_function: unsafe { ScalarFunction::borrow(out.scalar_function) }, - bind_info: out.bind_info, }) } + cpp::DUCKDB_VX_EXPR_CLASS::DUCKDB_VX_EXPR_CLASS_BOUND_AGGREGATE => { + let aggregate_function = unsafe { + let ptr = cpp::duckdb_vx_expr_get_bound_aggregate_function(self.as_ptr()); + AggregateFunction::borrow(ptr) + }; + ExpressionClass::BoundAggregate(BoundAggregate { aggregate_function }) + } cpp::DUCKDB_VX_EXPR_CLASS::DUCKDB_VX_EXPR_CLASS_BOUND_REF => { ExpressionClass::BoundRef } @@ -165,14 +180,25 @@ pub enum ExpressionClass<'a> { BoundBetween(BoundBetween<'a>), BoundOperator(BoundOperator<'a>), BoundFunction(BoundFunction<'a>), + BoundCast(BoundCast<'a>), + BoundAggregate(BoundAggregate<'a>), /// Column inside ExpressionFilter for expression pushed down to Vortex. BoundRef, } +pub struct BoundCast<'a> { + pub child: &'a ExpressionRef, + pub is_try: bool, +} + pub struct BoundColumnRef { pub name: DDBString, } +pub struct BoundAggregate<'a> { + pub aggregate_function: &'a AggregateFunctionRef, +} + pub struct BoundConstant<'a> { pub value: &'a ValueRef, } @@ -222,7 +248,6 @@ impl<'a> BoundOperator<'a> { pub struct BoundFunction<'a> { children: &'a [cpp::duckdb_vx_expr], pub scalar_function: &'a ScalarFunctionRef, - pub bind_info: *const c_void, } impl<'a> BoundFunction<'a> { diff --git a/vortex-duckdb/src/duckdb/scalar_function.rs b/vortex-duckdb/src/duckdb/function.rs similarity index 58% rename from vortex-duckdb/src/duckdb/scalar_function.rs rename to vortex-duckdb/src/duckdb/function.rs index 175b0e70d5c..0441fe43548 100644 --- a/vortex-duckdb/src/duckdb/scalar_function.rs +++ b/vortex-duckdb/src/duckdb/function.rs @@ -8,6 +8,7 @@ use crate::cpp; use crate::lifetime_wrapper; lifetime_wrapper!(ScalarFunction, cpp::duckdb_vx_sfunc, |_| {}); +lifetime_wrapper!(AggregateFunction, cpp::duckdb_vx_agg_func, |_| {}); impl ScalarFunctionRef { pub fn name(&self) -> &str { @@ -20,3 +21,15 @@ impl ScalarFunctionRef { } } } + +impl AggregateFunctionRef { + pub fn name(&self) -> &str { + unsafe { + let name_ptr = cpp::duckdb_vx_agg_func_name(self.as_ptr()); + std::ffi::CStr::from_ptr(name_ptr) + .to_str() + .map_err(|e| vortex_err!("invalid utf-8: {e}")) + .vortex_expect("aggregate function name should be valid UTF-8") + } + } +} diff --git a/vortex-duckdb/src/duckdb/logical_type.rs b/vortex-duckdb/src/duckdb/logical_type.rs index 7a4627e6953..28c17cbcf02 100644 --- a/vortex-duckdb/src/duckdb/logical_type.rs +++ b/vortex-duckdb/src/duckdb/logical_type.rs @@ -153,6 +153,14 @@ impl LogicalType { Self::new(DUCKDB_TYPE::DUCKDB_TYPE_BLOB) } + pub fn uint8() -> Self { + Self::new(DUCKDB_TYPE::DUCKDB_TYPE_UTINYINT) + } + + pub fn uint16() -> Self { + Self::new(DUCKDB_TYPE::DUCKDB_TYPE_USMALLINT) + } + pub fn uint32() -> Self { Self::new(DUCKDB_TYPE::DUCKDB_TYPE_UINTEGER) } @@ -165,6 +173,14 @@ impl LogicalType { Self::new(DUCKDB_TYPE::DUCKDB_TYPE_UHUGEINT) } + pub fn int8() -> Self { + Self::new(DUCKDB_TYPE::DUCKDB_TYPE_TINYINT) + } + + pub fn int16() -> Self { + Self::new(DUCKDB_TYPE::DUCKDB_TYPE_SMALLINT) + } + pub fn int32() -> Self { Self::new(DUCKDB_TYPE::DUCKDB_TYPE_INTEGER) } @@ -234,6 +250,21 @@ impl LogicalTypeRef { matches!(self.as_type_id(), DUCKDB_TYPE::DUCKDB_TYPE_DECIMAL) } + /// True if this type maps to a Vortex Primitive and isn't a floating point + pub fn is_primitive_integer(&self) -> bool { + matches!( + self.as_type_id(), + DUCKDB_TYPE::DUCKDB_TYPE_TINYINT + | DUCKDB_TYPE::DUCKDB_TYPE_SMALLINT + | DUCKDB_TYPE::DUCKDB_TYPE_INTEGER + | DUCKDB_TYPE::DUCKDB_TYPE_BIGINT + | DUCKDB_TYPE::DUCKDB_TYPE_UTINYINT + | DUCKDB_TYPE::DUCKDB_TYPE_USMALLINT + | DUCKDB_TYPE::DUCKDB_TYPE_UINTEGER + | DUCKDB_TYPE::DUCKDB_TYPE_UBIGINT + ) + } + pub fn geometry_crs(&self) -> Option { unsafe { let c_string = duckdb_geometry_type_get_crs(self.as_ptr()); diff --git a/vortex-duckdb/src/duckdb/mod.rs b/vortex-duckdb/src/duckdb/mod.rs index 27ea994b1e7..95dc8adc186 100644 --- a/vortex-duckdb/src/duckdb/mod.rs +++ b/vortex-duckdb/src/duckdb/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod aggregate_pushdown; mod bind_input; mod connection; mod data; @@ -8,11 +9,11 @@ mod data_chunk; mod database; mod ddb_string; mod expr; +mod function; mod logical_type; mod macro_; mod query_result; mod reusable_dict; -mod scalar_function; mod selection_vector; mod string_map; mod table_filter; @@ -24,6 +25,7 @@ mod vector_buffer; use std::ffi::c_void; use std::ptr; +pub use aggregate_pushdown::*; pub use bind_input::*; pub use connection::*; pub use data::*; @@ -31,10 +33,10 @@ pub use data_chunk::*; pub use database::*; pub use ddb_string::*; pub use expr::*; +pub use function::*; pub use logical_type::*; pub use query_result::*; pub use reusable_dict::*; -pub use scalar_function::*; pub use selection_vector::*; pub use string_map::*; pub use table_filter::*; diff --git a/vortex-duckdb/src/e2e_test/geo_pushdown_test.rs b/vortex-duckdb/src/e2e_test/geo_pushdown_test.rs new file mode 100644 index 00000000000..f058de20d47 --- /dev/null +++ b/vortex-duckdb/src/e2e_test/geo_pushdown_test.rs @@ -0,0 +1,160 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Pushdown tests for the geo scalar functions: every lowered filter must reach the Vortex +//! scan on a direct file scan and through a view. + +use num_traits::AsPrimitive; +use rstest::rstest; +use tempfile::NamedTempFile; +use vortex::array::IntoArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::StructArray; +use vortex::file::WriteOptionsSessionExt; +use vortex::io::runtime::BlockingRuntime; +use vortex_array::arrays::ExtensionArray; +use vortex_array::dtype::extension::ExtDType; +use vortex_geo::extension::GeoMetadata; +use vortex_geo::extension::Point; + +use crate::RUNTIME; +use crate::SESSION; +use crate::cpp::duckdb_string_t; +use crate::duckdb::Connection; +use crate::duckdb::Database; + +/// The test points; the filters below state how many of them they match. +const POINTS: [(f64, f64); 5] = [ + (1.0, 1.0), + (4.0, 4.0), + (10.0, 10.0), + (-1.0, 5.0), + (2.0, 3.0), +]; + +/// Matches three of [`POINTS`]: two inside the polygon and one on its boundary. +const ST_INTERSECTS_FILTER: &str = + "ST_Intersects(geometry, ST_GeomFromText('POLYGON((0 0, 4 0, 4 4, 0 4, 0 0))'))"; + +/// Matches two of [`POINTS`]: those closer than `3` to `(1, 1)`. +const ST_DISTANCE_FILTER: &str = "ST_Distance(geometry, ST_GeomFromText('POINT (1 1)')) < 3.0"; + +/// Matches the same two points as [`ST_DISTANCE_FILTER`], phrased as `ST_DWithin`. +const ST_DWITHIN_FILTER: &str = "ST_DWithin(geometry, ST_GeomFromText('POINT (1 1)'), 3.0)"; + +/// Matches two of [`POINTS`]: the boundary point of [`ST_INTERSECTS_FILTER`]'s polygon is not +/// contained (OGC contains excludes the boundary). +const ST_CONTAINS_FILTER: &str = + "ST_Contains(ST_GeomFromText('POLYGON((0 0, 4 0, 4 4, 0 4, 0 0))'), geometry)"; + +/// Matches the same two points as [`ST_CONTAINS_FILTER`], phrased as `ST_Within`. +const ST_WITHIN_FILTER: &str = + "ST_Within(geometry, ST_GeomFromText('POLYGON((0 0, 4 0, 4 4, 0 4, 0 0))'))"; + +/// A vortex file whose single column `geometry` holds [`POINTS`] as native `Point`s. +fn native_point_file() -> NamedTempFile { + RUNTIME.block_on(async { + let xs = PrimitiveArray::from_iter(POINTS.map(|(x, _)| x)).into_array(); + let ys = PrimitiveArray::from_iter(POINTS.map(|(_, y)| y)).into_array(); + let storage = StructArray::from_fields(&[("x", xs), ("y", ys)]) + .unwrap() + .into_array(); + let dtype = + ExtDType::::try_new(GeoMetadata { crs: None }, storage.dtype().clone()).unwrap(); + let points = ExtensionArray::new(dtype.erased(), storage).into_array(); + + let file = NamedTempFile::with_suffix(".vortex").unwrap(); + let table = StructArray::from_fields(&[("geometry", points)]).unwrap(); + let mut writer = async_fs::File::create(&file).await.unwrap(); + SESSION + .write_options() + .write(&mut writer, table.into_array().to_array_stream()) + .await + .unwrap(); + file + }) +} + +/// An in-memory database with the Vortex extension initialized, `spatial` loaded, and the +/// spatial overrides registered. +fn spatial_database() -> (Database, Connection) { + let db = Database::open_in_memory().unwrap(); + db.register_vortex_scan_replacement().unwrap(); + crate::initialize(&db).unwrap(); + let conn = db.connect().unwrap(); + conn.query("INSTALL spatial; LOAD spatial;").unwrap(); + // Must follow `LOAD spatial`: the overrides copy spatial's catalog entries. + db.register_spatial_overrides().unwrap(); + (db, conn) +} + +/// Read back the single `i64` of a one-row, one-column query. +fn query_i64(conn: &Connection, query: &str) -> i64 { + let result = conn.query(query).unwrap(); + let chunk = result.into_iter().next().unwrap(); + chunk + .get_vector(0) + .as_slice_with_len::(chunk.len().as_())[0] +} + +/// The `EXPLAIN` physical plan of `query` as one string. +fn explain_plan(conn: &Connection, query: &str) -> String { + let explain = conn.query(&format!("EXPLAIN {query}")).unwrap(); + let mut plan = String::new(); + for mut chunk in explain { + let len = chunk.len().as_(); + let vec = chunk.get_vector_mut(1); + for value in unsafe { vec.as_slice_mut::(len) } { + let slice: &[u8] = unsafe { + std::slice::from_raw_parts( + crate::cpp::duckdb_string_t_data(&raw mut *value) as _, + crate::cpp::duckdb_string_t_length(*value) as usize, + ) + }; + plan.push_str(&String::from_utf8_lossy(slice)); + } + } + plan +} + +/// Assert that filtering `table` with `filter` counts `expected` points and leaves no `FILTER` +/// operator in the plan, i.e. the predicate ran inside the scan. +fn assert_pushed(conn: &Connection, table: &str, filter: &str, expected: i64) { + let query = format!("SELECT count(*) FROM {table} WHERE {filter}"); + assert_eq!(query_i64(conn, &query), expected); + let plan = explain_plan(conn, &query); + assert!(!plan.contains("FILTER"), "filter was not pushed:\n{plan}"); +} + +/// Every lowered geo filter pushes on a direct file scan. +#[rstest] +#[case::st_intersects(ST_INTERSECTS_FILTER, 3)] +#[case::st_distance(ST_DISTANCE_FILTER, 2)] +#[case::st_dwithin(ST_DWITHIN_FILTER, 2)] +#[case::st_contains(ST_CONTAINS_FILTER, 2)] +#[case::st_within(ST_WITHIN_FILTER, 2)] +fn geo_filter_pushes_on_file_scan(#[case] filter: &str, #[case] expected: i64) { + let file = native_point_file(); + let (_db, conn) = spatial_database(); + let table = format!("'{}'", file.path().to_string_lossy()); + assert_pushed(&conn, &table, filter, expected); +} + +/// Every lowered geo filter pushes through a view; without the overrides, DuckDB would keep +/// `ST_Intersects` (can-throw) above the view's projection and hide `ST_DWithin`'s radius. +#[rstest] +#[case::st_intersects(ST_INTERSECTS_FILTER, 3)] +#[case::st_distance(ST_DISTANCE_FILTER, 2)] +#[case::st_dwithin(ST_DWITHIN_FILTER, 2)] +#[case::st_contains(ST_CONTAINS_FILTER, 2)] +#[case::st_within(ST_WITHIN_FILTER, 2)] +fn geo_filter_pushes_through_view(#[case] filter: &str, #[case] expected: i64) { + let file = native_point_file(); + let (_db, conn) = spatial_database(); + conn.query(&format!( + "CREATE VIEW points_v AS SELECT * FROM read_vortex('{}')", + file.path().to_string_lossy() + )) + .unwrap(); + assert_pushed(&conn, "points_v", filter, expected); +} diff --git a/vortex-duckdb/src/e2e_test/mod.rs b/vortex-duckdb/src/e2e_test/mod.rs index 34182bd133d..5d41df51220 100644 --- a/vortex-duckdb/src/e2e_test/mod.rs +++ b/vortex-duckdb/src/e2e_test/mod.rs @@ -1,5 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#[cfg(test)] +mod geo_pushdown_test; #[cfg(test)] mod vortex_scan_test; diff --git a/vortex-duckdb/src/exporter/bool.rs b/vortex-duckdb/src/exporter/bool.rs index 4d0a2b29883..fc1dd3e10c3 100644 --- a/vortex-duckdb/src/exporter/bool.rs +++ b/vortex-duckdb/src/exporter/bool.rs @@ -46,13 +46,16 @@ impl ColumnExporter for BoolExporter { ) -> VortexResult<()> { // DuckDB uses byte bools, not bit bools. // maybe we can convert into these from a compressed array sometimes?. - unsafe { vector.as_slice_mut(len) }.copy_from_slice( - &self - .bit_buffer - .slice(offset..(offset + len)) - .iter() - .collect::>(), - ); + // Unpack the bits directly into the destination byte slice, avoiding the + // throwaway `Vec` allocation and the extra copy pass. The sliced + // iterator already accounts for the buffer's bit offset. + let dst = unsafe { vector.as_slice_mut(len) }; + for (slot, bit) in dst + .iter_mut() + .zip(self.bit_buffer.slice(offset..(offset + len)).iter()) + { + *slot = bit; + } Ok(()) } diff --git a/vortex-duckdb/src/exporter/extension.rs b/vortex-duckdb/src/exporter/extension.rs index 221dc92a85f..c07dc25a886 100644 --- a/vortex-duckdb/src/exporter/extension.rs +++ b/vortex-duckdb/src/exporter/extension.rs @@ -8,6 +8,18 @@ use vortex::array::arrays::extension::ExtensionArrayExt; use vortex::array::extension::datetime::AnyTemporal; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex_geo::extension::LineString; +use vortex_geo::extension::LineStringData; +use vortex_geo::extension::MultiLineString; +use vortex_geo::extension::MultiLineStringData; +use vortex_geo::extension::MultiPoint; +use vortex_geo::extension::MultiPointData; +use vortex_geo::extension::MultiPolygon; +use vortex_geo::extension::MultiPolygonData; +use vortex_geo::extension::Point; +use vortex_geo::extension::PointData; +use vortex_geo::extension::Polygon; +use vortex_geo::extension::PolygonData; use vortex_geo::extension::WellKnownBinary; use vortex_geo::extension::WellKnownBinaryData; @@ -27,5 +39,29 @@ pub(crate) fn new_exporter( return geo::new_wkb_exporter(WellKnownBinaryData::try_from(ext)?, ctx); } + if ext.ext_dtype().is::() { + return geo::new_point_exporter(PointData::try_from(ext)?, ctx); + } + + if ext.ext_dtype().is::() { + return geo::new_linestring_exporter(LineStringData::try_from(ext)?, ctx); + } + + if ext.ext_dtype().is::() { + return geo::new_multipoint_exporter(MultiPointData::try_from(ext)?, ctx); + } + + if ext.ext_dtype().is::() { + return geo::new_polygon_exporter(PolygonData::try_from(ext)?, ctx); + } + + if ext.ext_dtype().is::() { + return geo::new_multilinestring_exporter(MultiLineStringData::try_from(ext)?, ctx); + } + + if ext.ext_dtype().is::() { + return geo::new_multipolygon_exporter(MultiPolygonData::try_from(ext)?, ctx); + } + vortex_bail!("no non-temporal extension exporter") } diff --git a/vortex-duckdb/src/exporter/geo.rs b/vortex-duckdb/src/exporter/geo.rs index 1287ed019e2..ccc1893ea18 100644 --- a/vortex-duckdb/src/exporter/geo.rs +++ b/vortex-duckdb/src/exporter/geo.rs @@ -4,6 +4,12 @@ use vortex::array::ExecutionCtx; use vortex::array::arrays::VarBinViewArray; use vortex::error::VortexResult; +use vortex_geo::extension::LineStringData; +use vortex_geo::extension::MultiLineStringData; +use vortex_geo::extension::MultiPointData; +use vortex_geo::extension::MultiPolygonData; +use vortex_geo::extension::PointData; +use vortex_geo::extension::PolygonData; use vortex_geo::extension::WellKnownBinaryData; use crate::exporter::ColumnExporter; @@ -17,3 +23,66 @@ pub(crate) fn new_wkb_exporter( let values = array.wkb_values().clone().execute::(ctx)?; new_exporter(values, ctx) } + +/// Create an exporter for a native `Point` column. DuckDB `GEOMETRY` vectors carry WKB, so the +/// points are serialized to WKB via [`PointData::to_wkb`] (only for rows DuckDB materializes — +/// with predicate pushdown that's just the survivors). +pub(crate) fn new_point_exporter( + point: PointData, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let values = point.to_wkb(ctx)?.execute::(ctx)?; + new_exporter(values, ctx) +} + +/// Create an exporter for a native `Polygon` column. Like [`new_point_exporter`], DuckDB `GEOMETRY` +/// vectors carry WKB, so the polygons are serialized to WKB via [`PolygonData::to_wkb`]. +pub(crate) fn new_polygon_exporter( + polygon: PolygonData, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let values = polygon.to_wkb(ctx)?.execute::(ctx)?; + new_exporter(values, ctx) +} + +/// Create an exporter for a native `MultiPolygon` column, serialized to WKB via +/// [`MultiPolygonData::to_wkb`] (see [`new_point_exporter`]). +pub(crate) fn new_multipolygon_exporter( + multipolygon: MultiPolygonData, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let values = multipolygon.to_wkb(ctx)?.execute::(ctx)?; + new_exporter(values, ctx) +} + +/// Create an exporter for a native `LineString` column, serialized to WKB via +/// [`LineStringData::to_wkb`] (see [`new_point_exporter`]). +pub(crate) fn new_linestring_exporter( + linestring: LineStringData, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let values = linestring.to_wkb(ctx)?.execute::(ctx)?; + new_exporter(values, ctx) +} + +/// Create an exporter for a native `MultiPoint` column, serialized to WKB via +/// [`MultiPointData::to_wkb`] (see [`new_point_exporter`]). +pub(crate) fn new_multipoint_exporter( + multipoint: MultiPointData, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let values = multipoint.to_wkb(ctx)?.execute::(ctx)?; + new_exporter(values, ctx) +} + +/// Create an exporter for a native `MultiLineString` column, serialized to WKB via +/// [`MultiLineStringData::to_wkb`] (see [`new_point_exporter`]). +pub(crate) fn new_multilinestring_exporter( + multilinestring: MultiLineStringData, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let values = multilinestring + .to_wkb(ctx)? + .execute::(ctx)?; + new_exporter(values, ctx) +} diff --git a/vortex-duckdb/src/ffi.rs b/vortex-duckdb/src/ffi.rs index 880c3f7f01c..1ed78c23494 100644 --- a/vortex-duckdb/src/ffi.rs +++ b/vortex-duckdb/src/ffi.rs @@ -17,6 +17,7 @@ use crate::copy::copy_to_finalize; use crate::copy::copy_to_initialize_global; use crate::copy::copy_to_sink; use crate::cpp; +use crate::duckdb::AggregatePushdownInput; use crate::duckdb::BindInput; use crate::duckdb::BindResult; use crate::duckdb::Data; @@ -38,6 +39,7 @@ use crate::table_function::get_partition_data; use crate::table_function::init_global; use crate::table_function::init_local; use crate::table_function::pushdown_complex_filter; +use crate::table_function::pushdown_projection_aggregates; use crate::table_function::pushdown_projection_expression; use crate::table_function::scan; use crate::table_function::statistics; @@ -126,6 +128,20 @@ unsafe extern "C-unwind" fn duckdb_table_function_pushdown_projection_expression }) } +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_pushdown_projection_aggregates( + bind_data: *mut c_void, + input: cpp::duckdb_vx_agg_input, + error_out: *mut cpp::duckdb_vx_error, +) -> bool { + let bind_data = unsafe { bind_data.cast::().as_mut() } + .vortex_expect("bind_data null pointer"); + let input = unsafe { AggregatePushdownInput::borrow(input) }; + try_or(error_out, || { + pushdown_projection_aggregates(bind_data, input) + }) +} + #[unsafe(no_mangle)] unsafe extern "C-unwind" fn duckdb_table_function_scan( global_init_data: *mut c_void, @@ -207,12 +223,15 @@ pub unsafe extern "C-unwind" fn duckdb_table_function_init_global( #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn duckdb_table_function_init_local( + bind_data: *const c_void, global_init_data: *mut c_void, ) -> cpp::duckdb_vx_data { + let bind_data = unsafe { bind_data.cast::().as_ref() } + .vortex_expect("bind_data null pointer"); let global_init_data = unsafe { global_init_data.cast::().as_ref() } .vortex_expect("global_init_data null pointer"); - let init_data = init_local(global_init_data); + let init_data = init_local(bind_data, global_init_data); Data::from(Box::new(init_data)).as_ptr() } diff --git a/vortex-duckdb/src/lib.rs b/vortex-duckdb/src/lib.rs index 55039b88737..c68924868b3 100644 --- a/vortex-duckdb/src/lib.rs +++ b/vortex-duckdb/src/lib.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #![expect(clippy::missing_safety_doc)] +#![forbid(clippy::todo)] +#![forbid(clippy::unimplemented)] use std::ffi::c_char; use std::ffi::c_void; diff --git a/vortex-duckdb/src/multi_file.rs b/vortex-duckdb/src/multi_file.rs index bb9e015af5c..247b6a19909 100644 --- a/vortex-duckdb/src/multi_file.rs +++ b/vortex-duckdb/src/multi_file.rs @@ -1,8 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::path::Path; -use std::path::absolute; use std::sync::Arc; use itertools::Itertools; @@ -14,6 +12,7 @@ use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_err; use vortex::file::multi::MultiFileDataSource; +use vortex::file::multi::parse_uri_or_path; use vortex::io::compat::Compat; use vortex::io::filesystem::FileSystemRef; use vortex::io::object_store::ObjectStoreFileSystem; @@ -26,38 +25,6 @@ use crate::SESSION; use crate::duckdb::BindInputRef; use crate::duckdb::ExtractedValue; -/// Parse a glob string into a [`Url`]. -/// -/// Accepts full URLs (e.g. `s3://bucket/prefix/*.vortex`, `file:///data/*.vortex`) as well as -/// bare file paths. For bare paths, the path is made absolute (without requiring it to exist) -/// so that relative paths such as `./data/*.vortex` or `../data/*.vortex` are resolved correctly. -fn parse_glob_url(glob_url_str: &str) -> VortexResult { - Url::parse(glob_url_str).or_else(|_| { - let path = absolute(Path::new(glob_url_str)) - .map_err(|e| vortex_err!("Failed making {glob_url_str} absolute: {e}"))?; - // `absolute()` does not normalize `..` components, so `/a/b/../c` stays as-is. - // Normalizing manually avoids `..` being percent-encoded in the resulting URL. - let path = normalize_path(path); - Url::from_file_path(path).map_err(|_| vortex_err!("Neither URL nor path: {glob_url_str}")) - }) -} - -/// Normalize a path by resolving `.` and `..` components without accessing the filesystem. -fn normalize_path(path: std::path::PathBuf) -> std::path::PathBuf { - use std::path::Component; - let mut normalized = std::path::PathBuf::new(); - for component in path.components() { - match component { - Component::CurDir => {} - Component::ParentDir => { - normalized.pop(); - } - c => normalized.push(c), - } - } - normalized -} - fn resolve_filesystem(base_url: &Url) -> VortexResult { let object_store: Arc = match base_url.scheme() { "file" => Arc::new(LocalFileSystem::new()), @@ -104,7 +71,7 @@ pub fn bind_multi_file_scan(input: &BindInputRef) -> VortexResult = Vec::with_capacity(glob_strings.len()); for glob_str in &glob_strings { - glob_urls.push(parse_glob_url(glob_str)?); + glob_urls.push(parse_uri_or_path(glob_str)?); } // Cache filesystems by base URL to avoid resolving the same filesystem multiple times. @@ -134,146 +101,3 @@ pub fn bind_multi_file_scan(input: &BindInputRef) -> VortexResult VortexResult<()> { - let url = parse_glob_url("s3://my-bucket/prefix/*.vortex")?; - assert_eq!(url.scheme(), "s3"); - assert_eq!(url.host_str(), Some("my-bucket")); - assert_eq!(url.path(), "/prefix/*.vortex"); - Ok(()) - } - - #[test] - fn test_parse_glob_url_file_scheme() -> VortexResult<()> { - let url = parse_glob_url("file:///absolute/path/data.vortex")?; - assert_eq!(url.scheme(), "file"); - assert_eq!(url.path(), "/absolute/path/data.vortex"); - Ok(()) - } - - #[test] - fn test_parse_glob_url_absolute_glob_path() -> VortexResult<()> { - let tmpdir = tempfile::tempdir()?; - let glob = format!("{}/*.vortex", tmpdir.path().display()); - let url = parse_glob_url(&glob)?; - assert_eq!(url.scheme(), "file"); - assert!(url.path().ends_with("/*.vortex")); - Ok(()) - } - - #[test] - fn test_parse_glob_url_absolute_existing_path() -> VortexResult<()> { - let tmpfile = tempfile::NamedTempFile::new()?; - let canonical = std::fs::canonicalize(tmpfile.path())?; - let path_str = canonical - .to_str() - .ok_or_else(|| vortex_err!("canonical path is not valid UTF-8"))?; - let url = parse_glob_url(path_str)?; - assert_eq!(url.scheme(), "file"); - assert_eq!(url.path(), path_str); - Ok(()) - } - - #[test] - fn test_parse_glob_url_relative_path() -> VortexResult<()> { - // Create a tempfile in the current working directory so we can refer to it - // by a relative name (just the filename, without any directory component). - let tmpfile = tempfile::NamedTempFile::new_in(".")?; - let filename = tmpfile - .path() - .file_name() - .ok_or_else(|| vortex_err!("temp file missing file name"))? - .to_str() - .ok_or_else(|| vortex_err!("temp file name is not valid UTF-8"))?; - - let url = parse_glob_url(filename)?; - assert_eq!(url.scheme(), "file"); - // The relative name must have been resolved to an absolute path. - assert!(url.path().ends_with(filename)); - assert!(url.path().starts_with('/')); - Ok(()) - } - - #[test] - fn test_parse_glob_url_relative_glob_path() -> VortexResult<()> { - // A relative path with a glob character (e.g. `./data/*.vortex`) must also resolve - // correctly. - let tmpdir = tempfile::tempdir_in(".")?; - let dir_name = tmpdir - .path() - .file_name() - .ok_or_else(|| vortex_err!("temp dir missing file name"))? - .to_str() - .ok_or_else(|| vortex_err!("temp dir name is not valid UTF-8"))?; - let glob = format!("./{dir_name}/*.vortex"); - let url = parse_glob_url(&glob)?; - assert_eq!(url.scheme(), "file"); - assert!(url.path().starts_with('/')); - assert!(url.path().ends_with("/*.vortex")); - Ok(()) - } - - #[test] - fn test_parse_glob_url_nonexistent_path() -> VortexResult<()> { - // absolute() does not require the path to exist, so a non-existent path succeeds. - let url = parse_glob_url("/nonexistent/path/file.vortex")?; - assert_eq!(url.scheme(), "file"); - assert_eq!(url.path(), "/nonexistent/path/file.vortex"); - Ok(()) - } - - #[test] - fn test_parse_glob_url_parent_relative_path() -> VortexResult<()> { - // A path starting with `..` must be resolved to an absolute path without - // percent-encoding the `..` component in the resulting URL. - let tmpfile = tempfile::NamedTempFile::new_in("..")?; - let filename = tmpfile - .path() - .file_name() - .ok_or_else(|| vortex_err!("temp file missing file name"))? - .to_str() - .ok_or_else(|| vortex_err!("temp file name is not valid UTF-8"))?; - let relative = format!("../{filename}"); - - let url = parse_glob_url(&relative)?; - assert_eq!(url.scheme(), "file"); - // The resolved path must be absolute and must not contain encoded dots. - assert!(url.path().starts_with('/')); - assert!( - !url.path().contains("%2E"), - "path must not contain percent-encoded dots" - ); - assert!(url.path().ends_with(filename)); - Ok(()) - } - - // Use absolute paths so the expected result is cwd-independent. - #[rstest] - #[case("/a/./b", "/a/b")] - #[case("/a/b/./c", "/a/b/c")] - #[case("/a/../b", "/b")] - #[case("/a/b/../c", "/a/c")] - #[case("/a/b/../../c", "/c")] - #[case("/a/./b/.././c", "/a/c")] - #[case("/a/b/../..", "/")] - fn test_parse_glob_url_dot_normalization( - #[case] input: &str, - #[case] expected_path: &str, - ) -> VortexResult<()> { - let url = parse_glob_url(input)?; - assert_eq!(url.scheme(), "file"); - assert_eq!( - url.path(), - expected_path, - "input {input:?} should normalize to {expected_path:?}" - ); - Ok(()) - } -} diff --git a/vortex-duckdb/src/projection.rs b/vortex-duckdb/src/projection.rs index 4521115666c..df0481250bd 100644 --- a/vortex-duckdb/src/projection.rs +++ b/vortex-duckdb/src/projection.rs @@ -17,12 +17,14 @@ use vortex::expr::root; use vortex::expr::select; use vortex::layout::layouts::row_idx::row_idx; use vortex::scan::selection::Selection; +use vortex_utils::aliases::hash_set::HashSet; use crate::convert::try_from_table_filter; use crate::convert::try_from_virtual_column_filter; use crate::duckdb::LogicalType; use crate::duckdb::TableFilterClass; use crate::duckdb::TableFilterSetRef; +use crate::table_function::ColumnAggregate; // See MultiFileReader for constants @@ -187,6 +189,28 @@ impl Projection { file_row_number_column_pos, } } + + // Create a projection for aggregate scan + pub fn new_aggregate(aggregates: &[ColumnAggregate], fields: &[DuckdbField]) -> Self { + let mut names = Vec::with_capacity(aggregates.len()); + let mut seen: HashSet = HashSet::with_capacity(aggregates.len()); + for aggregate in aggregates { + let ColumnAggregate::Real { projection_id, .. } = aggregate else { + continue; + }; + if seen.contains(projection_id) { + continue; + } + seen.insert(*projection_id); + let projection_id: usize = projection_id.as_(); + names.push(fields[projection_id].name.as_str()); + } + Projection { + projection: select(names, root()), + file_index_column_pos: None, + file_row_number_column_pos: None, + } + } } pub struct Filter { diff --git a/vortex-duckdb/src/table_function.rs b/vortex-duckdb/src/table_function.rs index 5151e47a464..5509f020ccd 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -4,24 +4,34 @@ use std::cmp::max; use std::fmt::Formatter; use std::fmt::{self}; +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; use custom_labels::CURRENT_LABELSET; +use futures::FutureExt; +use futures::Stream; use futures::StreamExt; +use futures::future::BoxFuture; use itertools::Itertools; use num_traits::AsPrimitive; +use parking_lot::Mutex; use static_assertions::assert_impl_all; use tracing::debug; +use vortex::aggregate_fn::DynAccumulator; use vortex::array::ArrayRef; use vortex::array::Canonical; +use vortex::array::ExecutionCtx; use vortex::array::VortexSessionExecute as _; use vortex::array::arrays::ScalarFn; use vortex::array::arrays::Struct; use vortex::array::arrays::StructArray; use vortex::array::arrays::scalar_fn::ScalarFnArrayExt; +use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::optimizer::ArrayOptimizer; use vortex::error::VortexExpect; use vortex::error::VortexResult; @@ -39,14 +49,19 @@ use vortex::scalar_fn::fns::operators::Operator; use vortex::scalar_fn::fns::pack::Pack; use vortex::scan::DataSource; use vortex::scan::ScanRequest; +use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::parallelism::get_available_parallelism; use crate::RUNTIME; use crate::SESSION; use crate::column_statistics::ColumnStatistics; use crate::column_statistics::ColumnStatisticsAggregate; +use crate::convert::PushedAggregate; use crate::convert::try_from_bound_expression; +use crate::convert::try_from_projection_aggregate; use crate::convert::try_from_projection_expression; +use crate::duckdb::AggregateExpression; +use crate::duckdb::AggregatePushdownInputRef; use crate::duckdb::BindInputRef; use crate::duckdb::BindResultRef; use crate::duckdb::DataChunkRef; @@ -62,6 +77,9 @@ use crate::projection::Filter; use crate::projection::Projection; use crate::projection::extract_schema_from_dtype; +// Aggregate projection index for count(*). See cpp/aggregate_fn_pushdown.cpp +pub const COUNT_STAR_PROJ_IDX: u64 = u64::MAX; + pub struct TableFunctionBind { data_source: Arc, filter_exprs: Vec, @@ -69,6 +87,8 @@ pub struct TableFunctionBind { // There exists at least one non-optional table filter or at least one // complex filter is pushed down. has_non_optional_filter: AtomicBool, + // Non-empty iff this scan is aggregate + aggregates: Vec, } assert_impl_all!(TableFunctionBind: Send, Clone); @@ -82,6 +102,7 @@ impl Clone for TableFunctionBind { has_non_optional_filter: AtomicBool::new( self.has_non_optional_filter.load(Ordering::Relaxed), ), + aggregates: self.aggregates.clone(), } } } @@ -108,7 +129,8 @@ impl<'a> TableInitInput<'a> { } } -type DataSourceIterator = ThreadSafeIterator)>>; +type ScanItem = VortexResult<(ArrayRef, Arc)>; +type DataSourceIterator = ThreadSafeIterator; pub struct TableFunctionGlobal { iterator: DataSourceIterator, @@ -117,6 +139,16 @@ pub struct TableFunctionGlobal { bytes_read: AtomicU64, file_index_column_pos: Option, file_row_number_column_pos: Option, + + // Following 4 fields are used only in aggregate scans. + /// ArrayRef's scanned but not aggregated in "partials". + /// 0 means all arrays have been aggregated but output is not written. + /// u64::MAX means arrays have been aggregated and we've written output row + pending: Arc, + aggregates: Vec, + // Accumulated partials + partials: Mutex>>, + row_count: AtomicU64, } assert_impl_all!(TableFunctionGlobal: Send, Sync); @@ -126,6 +158,8 @@ pub struct TableFunctionLocal { exporter: Option, partition_index: u64, file_index: usize, + // Aggregate scan accumulated partials. Empty for non-aggregate scan + partials: Vec>, } pub struct PartitionData { @@ -134,6 +168,15 @@ pub struct PartitionData { pub file_index: usize, } +#[derive(Clone)] +pub(crate) enum ColumnAggregate { + Real { + projection_id: u64, + aggregate: PushedAggregate, + }, + CountStar, +} + #[derive(Debug)] pub enum Cardinality { /// Unknown number of rows @@ -155,6 +198,7 @@ pub fn bind(input: &BindInputRef, result: &mut BindResultRef) -> VortexResult VortexResult VortexResult VortexResult partition, @@ -254,6 +306,7 @@ pub fn init_global(init_input: &TableInitInput) -> VortexResult VortexResult()).detach(); + let iterator = RUNTIME.block_on_stream_thread_safe(|_handle| scan_driver_stream(stream, rx)); - let iterator = RUNTIME.block_on_stream_thread_safe(|_handle| rx.into_stream()); + let aggregates = bind_data.aggregates.clone(); + let partials = build_partials(&aggregates, &bind_data.column_fields)?; Ok(TableFunctionGlobal { iterator, @@ -280,10 +333,69 @@ pub fn init_global(init_input: &TableInitInput) -> VortexResult TableFunctionLocal { +fn scan_driver_stream(stream: S, rx: kanal::AsyncReceiver) -> ScanDriverStream +where + S: Stream + Send + 'static, +{ + ScanDriverStream { + driver: Some(stream.collect::<()>().boxed()), + rx: rx.into_stream().boxed(), + } +} + +struct ScanDriverStream { + driver: Option>, + rx: futures::stream::BoxStream<'static, ScanItem>, +} + +impl Stream for ScanDriverStream { + type Item = ScanItem; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if let Some(driver) = this.driver.as_mut() + && driver.as_mut().poll(cx).is_ready() + { + this.driver = None; + } + + match this.rx.as_mut().poll_next(cx) { + Poll::Ready(None) if this.driver.is_some() => Poll::Pending, + poll => poll, + } + } +} + +fn build_partials( + aggregates: &[ColumnAggregate], + fields: &[DuckdbField], +) -> VortexResult>> { + aggregates + .iter() + .filter_map(|spec| match spec { + ColumnAggregate::Real { + projection_id, + aggregate, + } => { + let projection_id: usize = projection_id.as_(); + Some(aggregate.build(fields[projection_id].dtype.clone())) + } + ColumnAggregate::CountStar => None, + }) + .collect() +} + +pub fn init_local( + bind_data: &TableFunctionBind, + global: &TableFunctionGlobal, +) -> TableFunctionLocal { unsafe { use custom_labels::sys; @@ -299,11 +411,109 @@ pub fn init_local(global: &TableFunctionGlobal) -> TableFunctionLocal { CURRENT_LABELSET.set(key, value); } + let partials = build_partials(&global.aggregates, &bind_data.column_fields) + // if aggregate initialization produced an error, it would error in + // init_global, see "partials" initialization there + .vortex_expect("local state aggregate initialization failed"); + TableFunctionLocal { iterator: global.iterator.clone(), exporter: None, partition_index: 0, file_index: 0, + partials, + } +} + +fn convert_result(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let array_result = array.optimize_recursive(ctx.session())?; + Ok(if let Some(array) = array_result.as_opt::() { + array.into_owned() + } else if let Some(array) = array_result.as_opt::() + && let Some(pack_options) = array.scalar_fn().as_opt::() + { + StructArray::new( + pack_options.names.clone(), + array.children(), + array.len(), + pack_options.nullability.into(), + ) + } else { + array_result.execute::(ctx)?.into_struct() + }) +} + +fn scan_aggregate( + local_state: &mut TableFunctionLocal, + global_state: &TableFunctionGlobal, + chunk: &mut DataChunkRef, +) -> VortexResult<()> { + let aggregates_len = global_state.aggregates.len(); + // seen[k] = output column for requested column k. + // If min(x), max(x), avg(y) are requested, seen = { 0: 0, 1: 1} + let mut seen: HashMap = HashMap::with_capacity(aggregates_len); + // positions[k] = column id for accumulator k + // If min(x), max(x), avg(y) are requested, positions = [0, 0, 1] + let mut positions: Vec = Vec::with_capacity(aggregates_len); + + for aggregate in &global_state.aggregates { + let ColumnAggregate::Real { projection_id, .. } = aggregate else { + continue; + }; + let len = seen.len(); + let pos = seen.entry_ref(projection_id).or_insert(len); + positions.push(*pos); + } + let has_count_star = local_state.partials.len() < aggregates_len; + + let mut ctx = SESSION.create_execution_ctx(); + loop { + let Some(result) = local_state.iterator.next() else { + // 0 means we're the last thread, u64::MAX means output is written. + // is_err() means CAS didn't succeed + if global_state + .pending + .compare_exchange(0, u64::MAX, Ordering::AcqRel, Ordering::Relaxed) + .is_err() + { + return Ok(()); + } + + let mut accumulators = global_state.partials.lock(); + let row_count = global_state.row_count.load(Ordering::Acquire) as i64; + let mut accum_iter = accumulators.iter_mut(); + for (idx, aggregate) in global_state.aggregates.iter().enumerate() { + let value = match aggregate { + ColumnAggregate::Real { .. } => { + let accum = accum_iter.next().vortex_expect("partial for real agg"); + Value::try_from(accum.finish()?)? + } + ColumnAggregate::CountStar => Value::from(row_count), + }; + chunk.get_vector_mut(idx).reference_value(&value); + } + chunk.set_len(1); + return Ok(()); + }; + let array = convert_result(result?.0, &mut ctx)?; + + for (i, partial) in positions.iter().zip(local_state.partials.iter_mut()) { + partial.accumulate(array.unmasked_field(*i), &mut ctx)?; + } + + { + let mut partials = global_state.partials.lock(); + for (global, local) in partials.iter_mut().zip(&mut local_state.partials) { + global.combine_partials(local.flush()?)?; + } + } + + if has_count_star { + global_state + .row_count + .fetch_add(array.len() as u64, Ordering::Relaxed); + } + global_state.pending.fetch_sub(1, Ordering::Release); } } @@ -312,6 +522,10 @@ pub fn scan( global_state: &TableFunctionGlobal, chunk: &mut DataChunkRef, ) -> VortexResult<()> { + if !local_state.partials.is_empty() { + return scan_aggregate(local_state, global_state, chunk); + } + loop { if local_state.exporter.is_none() { let mut ctx = SESSION.create_execution_ctx(); @@ -319,23 +533,8 @@ pub fn scan( return Ok(()); }; let (array_result, conversion_cache) = result?; - let array_result = array_result.optimize_recursive(ctx.session())?; local_state.file_index = conversion_cache.file_index; - - let array_result: StructArray = if let Some(array) = array_result.as_opt::() { - array.into_owned() - } else if let Some(array) = array_result.as_opt::() - && let Some(pack_options) = array.scalar_fn().as_opt::() - { - StructArray::new( - pack_options.names.clone(), - array.children(), - array.len(), - pack_options.nullability.into(), - ) - } else { - array_result.execute::(&mut ctx)?.into_struct() - }; + let array_result = convert_result(array_result, &mut ctx)?; local_state.exporter = Some(ArrayExporter::try_new( &array_result, @@ -391,7 +590,7 @@ pub fn pushdown_complex_filter( ) -> VortexResult { debug!(%expr, "pushing down expression"); - let Some(expr) = try_from_bound_expression(expr)? else { + let Some(expr) = try_from_bound_expression(expr, &bind_data.column_fields)? else { debug!(%expr, "failed to push down expression"); return Ok(false); }; @@ -447,6 +646,46 @@ pub fn pushdown_projection_expression( } } +/// Turn a scan into an aggregate scan. Input is N aggregations, possibly over +/// same columns. If we return true, optimized pass expands output to N columns, +/// e.g. min(x), max(x) turns into min(x0), max(x1), 2 columns in output. +pub fn pushdown_projection_aggregates( + bind_data: &mut TableFunctionBind, + input: &AggregatePushdownInputRef, +) -> VortexResult { + let len = input.len(); + let mut aggregates = Vec::with_capacity(len); + let mut has_non_count_star = false; + + debug!(%len, "pushing down projection aggregates"); + for i in 0..len { + let AggregateExpression { + expr, + projection_id, + } = input.get(i); + if projection_id == COUNT_STAR_PROJ_IDX { + aggregates.push(ColumnAggregate::CountStar); + continue; + } + let Some(aggregate) = try_from_projection_aggregate(expr)? else { + debug!(%expr, %i, "failed to push down projection aggregate"); + return Ok(false); + }; + debug!(%expr, %projection_id, %i, "pushed down projection aggregate"); + aggregates.push(ColumnAggregate::Real { + projection_id, + aggregate, + }); + has_non_count_star = true; + } + // DuckDB computes just count(*) faster than us + if !has_non_count_star { + return Ok(false); + } + bind_data.aggregates = aggregates; + Ok(true) +} + /// Get column-wise statistics. Available only if we're reading a single file. pub fn statistics(bind_data: &TableFunctionBind, column_index: usize) -> Option { let children = bind_data.data_source.children(); @@ -464,8 +703,16 @@ pub fn statistics(bind_data: &TableFunctionBind, column_index: usize) -> Option< Some(inner) => inner.file_stats().stats_sets(), None => return None, }; - let stats_aggregate = ColumnStatisticsAggregate::new(&stats_sets[column_index]); - let dtype = bind_data.column_fields[column_index].dtype.clone(); + let source_id = if bind_data.aggregates.is_empty() { + column_index + } else { + match bind_data.aggregates[column_index] { + ColumnAggregate::Real { projection_id, .. } => projection_id.as_(), + ColumnAggregate::CountStar => return None, + } + }; + let dtype = bind_data.column_fields[source_id].dtype.clone(); + let stats_aggregate = ColumnStatisticsAggregate::new(&stats_sets[source_id]); Some(ColumnStatistics::from(&stats_aggregate, dtype)) } @@ -480,6 +727,9 @@ pub fn statistics(bind_data: &TableFunctionBind, column_index: usize) -> Option< /// here. const DEFAULT_SELECTIVITY: f64 = 0.2; pub fn cardinality(bind_data: &TableFunctionBind) -> Cardinality { + // If we're doing an aggregate scan, we don't change output cardinality to + // 1 as we want duckdb to do our aggregation in parallel. That may look + // counterintuitive in the plan, though. let has_non_optional_filter = bind_data.has_non_optional_filter.load(Ordering::Relaxed); match bind_data.data_source.row_count() { Precision::Exact(v) => { @@ -519,6 +769,31 @@ pub fn to_string(bind_data: &TableFunctionBind, map: &mut DuckdbStringMapRef) { let mut filters = bind_data.filter_exprs.iter().map(|f| format!("{f}")); map.push("Filters", &filters.join("\n")); } + + if !bind_data.aggregates.is_empty() { + let aggregations = bind_data + .aggregates + .iter() + .map(|agg| match agg { + ColumnAggregate::Real { + projection_id, + aggregate, + } => { + let projection_id: usize = projection_id.as_(); + format!( + "{aggregate}({})", + bind_data.column_fields[projection_id].name + ) + } + ColumnAggregate::CountStar => "count(*)".to_string(), + }) + .join("\n"); + if !aggregations.is_empty() { + map.push("Aggregations", &aggregations); + } + return; + } + let projections = bind_data .column_fields .iter() @@ -545,8 +820,11 @@ fn progress(bytes_read: &AtomicU64, bytes_total: &AtomicU64) -> f64 { mod tests { use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering::Relaxed; + use std::task::Poll; + use crate::RUNTIME; use crate::table_function::progress; + use crate::table_function::scan_driver_stream; #[test] fn test_table_scan_progress() { @@ -561,4 +839,26 @@ mod tests { bytes_total.fetch_add(100, Relaxed); assert!((progress(&bytes_read, &bytes_total) - 50.).abs() < f64::EPSILON); } + + #[test] + fn scan_driver_panic_propagates_through_iterator() { + let (tx, rx) = kanal::bounded_async(1); + let _tx = tx; + let stream = futures::stream::poll_fn(|_| -> Poll> { + panic!("duckdb scan driver panic"); + }); + + let mut iter = + RUNTIME.block_on_stream_thread_safe(|_handle| scan_driver_stream(stream, rx)); + let panic = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| iter.next())) { + Ok(_) => panic!("driver panic must propagate through iterator"), + Err(panic) => panic, + }; + let message = panic + .downcast_ref::<&'static str>() + .copied() + .or_else(|| panic.downcast_ref::().map(String::as_str)) + .unwrap_or(""); + assert!(message.contains("duckdb scan driver panic")); + } } diff --git a/vortex-ffi/Cargo.toml b/vortex-ffi/Cargo.toml index 0fddf153d2b..a64aabac03f 100644 --- a/vortex-ffi/Cargo.toml +++ b/vortex-ffi/Cargo.toml @@ -20,20 +20,17 @@ categories = { workspace = true } all-features = true [dependencies] -arrow-array = { workspace = true, features = ["ffi"] } +arrow-array = { workspace = true } arrow-schema = { workspace = true } async-fs = { workspace = true } bytes = { workspace = true } futures = { workspace = true } -itertools = { workspace = true } mimalloc = { workspace = true, optional = true } -object_store = { workspace = true, features = ["aws", "azure", "gcp"] } paste = { workspace = true } -prost = { workspace = true } -tracing = { workspace = true, features = ["std", "log"] } +tracing = { workspace = true, features = ["log"] } tracing-subscriber = { workspace = true, features = ["env-filter"] } -url = { workspace = true, features = [] } vortex = { workspace = true, features = ["object_store"] } +vortex-arrow = { workspace = true } [dev-dependencies] rand = { workspace = true } diff --git a/vortex-ffi/README.md b/vortex-ffi/README.md index f1273ae5527..ff3b6c94ff2 100644 --- a/vortex-ffi/README.md +++ b/vortex-ffi/README.md @@ -1,4 +1,4 @@ -# Vortex C interface +# Vortex C bindings ## Updating Headers @@ -20,7 +20,7 @@ target_link_libraries(my_target, vortex_ffi_shared) # or target_link_libraries(my_target, vortex_ffi) ``` -## Running C examples: +## Running C examples ```sh cmake -Bbuild -DBUILD_EXAMPLES=1 @@ -99,7 +99,7 @@ cargo +nightly build -Zbuild-std --target= \ 2. Build tests with target triple: ```sh -cmake -Bbuild -DWITH_ASAN=1 -DTARGET_TRIPLE= +cmake -Bbuild -DSANITIZER=asan -DTARGET_TRIPLE= ``` 3. Run the tests (ctest doesn't output failures in detail): diff --git a/vortex-ffi/cbindgen.toml b/vortex-ffi/cbindgen.toml index d1b4ffa5f48..a0ee3aaed7d 100644 --- a/vortex-ffi/cbindgen.toml +++ b/vortex-ffi/cbindgen.toml @@ -11,9 +11,11 @@ header = """ #pragma once #include -// // THIS FILE IS AUTO-GENERATED, DO NOT MAKE EDITS DIRECTLY -// + +// All operations return owned types which need to be freed by calling a +// matching _free() function. This includes all arrays, data sources, scans, +// errors, error messages, and other allocated objects. // https://arrow.apache.org/docs/format/CDataInterface.html#structure-definitions // If you want to use your own Arrow library like nanoarrow, define this macro @@ -63,6 +65,21 @@ typedef struct ArrowArrayStream FFI_ArrowArrayStream; #endif """ +trailer = """ +#include + +/** + * Create a view over a null-terminated C string. + * View is valid as long as "str" is valid + */ +static inline vx_view vx_view_from_cstr(const char* str) { + vx_view s; + s.ptr = str; + s.len = strlen(str); + return s; +} +""" + [export.rename] "f16" = "uint16_t" diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index 0203202e437..1c2caa6fad8 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -3,9 +3,11 @@ #pragma once #include -// // THIS FILE IS AUTO-GENERATED, DO NOT MAKE EDITS DIRECTLY -// + +// All operations return owned types which need to be freed by calling a +// matching _free() function. This includes all arrays, data sources, scans, +// errors, error messages, and other allocated objects. // https://arrow.apache.org/docs/format/CDataInterface.html#structure-definitions // If you want to use your own Arrow library like nanoarrow, define this macro @@ -199,6 +201,52 @@ typedef enum { VX_ESTIMATE_INEXACT = 2, } vx_estimate_type; +/** + * Error category for vx_error. + */ +typedef enum { + /** + * All other errors + */ + VX_ERROR_CODE_OTHER = 0, + /** + * Index out of bounds + */ + VX_ERROR_CODE_OUT_OF_BOUNDS = 1, + /** + * Compute kernel execute error + */ + VX_ERROR_CODE_COMPUTE = 2, + /** + * An invalid argument was provided. + */ + VX_ERROR_CODE_INVALID_ARGUMENT = 3, + /** + * Serialization/deserialization error + */ + VX_ERROR_CODE_SERIALIZATION = 4, + /** + * Unimplemented function + */ + VX_ERROR_CODE_NOT_IMPLEMENTED = 5, + /** + * Type mismatch + */ + VX_ERROR_CODE_MISMATCHED_TYPES = 6, + /** + * Assertion failed + */ + VX_ERROR_CODE_ASSERTION_FAILED = 7, + /** + * IO error + */ + VX_ERROR_CODE_IO = 8, + /** + * Panic inside FFI + */ + VX_ERROR_CODE_PANIC = 9, +} vx_error_code; + /** * Equalities, inequalities, and boolean operations over possibly null values. * For most operations, if either side is null, the result is null. @@ -410,7 +458,7 @@ typedef struct Primitive Primitive; * array is a cheap operation. * * Unless stated explicitly, all operations with vx_array don't take - * ownership of it, and thus it must be freed by the caller. + * ownership of it, and thus the array must be freed by the caller. */ typedef struct vx_array vx_array; @@ -442,11 +490,6 @@ typedef struct vx_array_iterator vx_array_iterator; */ typedef struct vx_array_sink vx_array_sink; -/** - * Strings for use within Vortex. - */ -typedef struct vx_binary vx_binary; - /** * A reference to one or more (possibly remote) paths. * Creating vx_data_source opens the first matched path to read the schema. @@ -475,8 +518,6 @@ typedef struct vx_error vx_error; * data. Each expression consists of an encoding (vtable), heap-allocated * metadata, and child expressions. * - * Unless stated explicitly, all expressions returned are owned and must - * be freed by the caller. * Unless stated explicitly, if an operation on const vx_expression* is * passed NULL, NULL is returned. * Operations on expressions don't take ownership of input values, and so @@ -516,11 +557,6 @@ typedef struct vx_scan vx_scan; */ typedef struct vx_session vx_session; -/** - * Strings for use within Vortex. - */ -typedef struct vx_string vx_string; - typedef struct vx_struct_column_builder vx_struct_column_builder; /** @@ -549,17 +585,34 @@ typedef struct { const vx_array *array; } vx_validity; +/** + * A non owning view over a byte range. + */ +typedef struct { + /** + * NULL "ptr" requires len == 0 + */ + const char *ptr; + /** + * Length in bytes. + */ + size_t len; +} vx_view; + /** * Options for creating a data source. */ typedef struct { /** - * Required: paths to files, tables, or layout trees. - * May be a glob pattern like "*.vortex". - * If you want to include multiple paths, concat them with a comma: - * "file1.vortex,../file2.vortex". + * Required: paths to files, tables, or layout trees. Each entry may be a + * glob pattern like "*.vortex". Must point to an array of size + * "paths_len". paths bytes are copied. + */ + const vx_view *paths; + /** + * Number of entries in `paths`. */ - const char *paths; + size_t paths_len; } vx_data_source_options; /** @@ -629,13 +682,12 @@ extern "C" { #endif // __cplusplus /** - * Clone a borrowed [`vx_array`], returning an owned [`vx_array`]. - * Must be released with [`vx_array_free`]. + * Increase reference count on vx_array */ const vx_array *vx_array_clone(const vx_array *ptr); /** - * Free an owned [`vx_array`] object. + * Decrease reference count on vx_array or free if there are no other references */ void vx_array_free(const vx_array *ptr); @@ -677,10 +729,7 @@ void vx_array_get_validity(const vx_array *array, vx_validity *validity, vx_erro size_t vx_array_len(const vx_array *array); /** - * Get the [`struct@crate::dtype::vx_dtype`] of the array. - * - * The returned pointer is valid as long as the array is valid. - * Do NOT free the returned dtype pointer - it shares the lifetime of the array. + * Get array's dtype */ const vx_dtype *vx_array_dtype(const vx_array *array); @@ -738,9 +787,6 @@ const vx_array *vx_array_new_primitive(vx_ptype ptype, * `nullable` controls the top-level nullability of the resulting array's dtype. For an Arrow * record batch (which has no top-level validity) pass `false`. * - * The imported buffers are referenced zero-copy where possible; the returned array keeps the - * Arrow data alive until it is freed with [`vx_array_free`]. - * * On error, returns NULL and sets `error_out`. * * Example: @@ -799,81 +845,98 @@ double vx_array_get_f64(const vx_array *array, size_t index); double vx_array_get_storage_f64(const vx_array *array, size_t index); /** - * Return the utf-8 string at `index` in the array. The pointer will be null if the value at `index` is null. - * The caller must free the returned pointer. - */ -const vx_string *vx_array_get_utf8(const vx_array *array, uint32_t index); - -/** - * Return the binary at `index` in the array. The pointer will be null if the value at `index` is null. - * The caller must free the returned pointer. + * Return UTF-8 string at "index" in a canonical Utf8 array. + * + * For invalid elements the returned value is unspecified, check validity via + * vx_array_get_validity. + * Returned view is valid as long as "array" is valid. + * Errors if index is out of bounds or array is not a canonical Utf8 array. */ -const vx_binary *vx_array_get_binary(const vx_array *array, uint32_t index); +vx_view vx_array_utf8_at(const vx_array *array, size_t index, vx_error **error_out); /** - * Apply the expression to the array, wrapping it with a ScalarFnArray. - * This operation takes constant time as it doesn't execute the underlying - * array. Executing the underlying array still takes O(n) time. + * Return a binary string at "index" in a canonical Binary array. + * + * For invalid elements the returned value is unspecified, check validity via + * vx_array_get_validity. + * Returned view is valid as long as "array" is valid. + * Errors if index is out of bounds or array is not a canonical Binary array. */ -const vx_array *vx_array_apply(const vx_array *array, const vx_expression *expression, vx_error **error); +vx_view vx_array_binary_at(const vx_array *array, size_t index, vx_error **error_out); /** - * Free an owned [`vx_array_iterator`] object. + * For a canonical Bool array, return bool at "index". + * For invalid elements returned value is unspecified, check validity via + * vx_array_get_validity. + * + * Panics if "array" is not canonical - call vx_array_canonicalize first. + * Panics if "array" is not a Bool array. + * Panics if "index" is out of bounds. */ -void vx_array_iterator_free(vx_array_iterator *ptr); +bool vx_array_get_bool(const vx_array *array, size_t index); /** - * Attempt to advance the `current` pointer of the iterator. - * - * A return value of `true` indicates that another element was pulled from the iterator, and a return - * of `false` indicates that the iterator is finished. + * Decode array into its canonical form. * - * It is an error to call this function again after the iterator is finished. + * On error returns NULL and "sets error_out". */ -const vx_array *vx_array_iterator_next(vx_array_iterator *iter, vx_error **error_out); +const vx_array *vx_array_canonicalize(const vx_session *session, const vx_array *array, vx_error **error_out); /** - * Clone a borrowed [`vx_binary`], returning an owned [`vx_binary`]. - * Must be released with [`vx_binary_free`]. + * Return a pointer to the values buffer of a canonical Primitive array. + * Pointer is valid as long as "array" is valid. + * + * Errors if array is not a canonical Primitive. */ -const vx_binary *vx_binary_clone(const vx_binary *ptr); +const void *vx_array_data_ptr_primitive(const vx_array *array, vx_error **error_out); /** - * Free an owned [`vx_binary`] object. + * Return a pointer to the bitpacked buffer of a canonical Bool array. + * Pointer is valid as long as "array" is valid. + * + * Writes bit offset of the first element into "bit_offset_out". + * "bit_offset_out" must not be NULL. + * + * Errors if array is not a canonical Bool. */ -void vx_binary_free(const vx_binary *ptr); +const void *vx_array_data_ptr_bool(const vx_array *array, size_t *bit_offset_out, vx_error **error_out); /** - * Create a new Vortex UTF-8 string by copying from a pointer and length. + * Apply the expression to the array, wrapping it with a ScalarFnArray. + * This operation takes constant time as it doesn't execute the underlying + * array. Executing the underlying array still takes O(n) time. */ -const vx_binary *vx_binary_new(const char *ptr, size_t len); +const vx_array *vx_array_apply(const vx_array *array, const vx_expression *expression, vx_error **error); /** - * Return the length of the string in bytes. + * Free a vx_array_iterator */ -size_t vx_binary_len(const vx_binary *ptr); +void vx_array_iterator_free(const vx_array_iterator *ptr); /** - * Return the pointer to the string data. + * Attempt to advance the `current` pointer of the iterator. + * + * A return value of `true` indicates that another element was pulled from the iterator, and a return + * of `false` indicates that the iterator is finished. + * + * It is an error to call this function again after the iterator is finished. */ -const char *vx_binary_ptr(const vx_binary *ptr); +const vx_array *vx_array_iterator_next(vx_array_iterator *iter, vx_error **error_out); /** - * Clone a borrowed [`vx_data_source`], returning an owned [`vx_data_source`]. - * Must be released with [`vx_data_source_free`]. + * Increase reference count on vx_data_source */ const vx_data_source *vx_data_source_clone(const vx_data_source *ptr); /** - * Free an owned [`vx_data_source`] object. + * Decrease reference count on vx_data_source or free if there are no other references */ void vx_data_source_free(const vx_data_source *ptr); /** * Create a data source. * The first matched file is opened eagerly. to read the schema. All other I/O - * is deferred until a scan is requested. The returned pointer is owned by the - * caller and must be freed with vx_data_source_free. + * is deferred until a scan is requested. * * On error, returns NULL and sets "err". */ @@ -887,17 +950,13 @@ vx_data_source_new(const vx_session *session, const vx_data_source_options *opti * The bytes are borrowed, not copied: the caller must keep "buffer" alive and * unmodified until the data source is freed. * - * The returned pointer is owned by the caller and must be freed with - * vx_data_source_free. - * * On error, returns NULL and sets "err". */ const vx_data_source * vx_data_source_new_buffer(const vx_session *session, const void *buffer, size_t buffer_len, vx_error **err); /** - * Return the schema of the data source as a non-owned dtype. - * The returned pointer is valid as long as "ds" is alive. Do not free it. + * Return data source's dtype */ const vx_dtype *vx_data_source_dtype(const vx_data_source *ds); @@ -907,13 +966,12 @@ const vx_dtype *vx_data_source_dtype(const vx_data_source *ds); void vx_data_source_get_row_count(const vx_data_source *ds, vx_estimate *row_count); /** - * Clone a borrowed [`vx_dtype`], returning an owned [`vx_dtype`]. - * Must be released with [`vx_dtype_free`]. + * Increase reference count on vx_dtype */ const vx_dtype *vx_dtype_clone(const vx_dtype *ptr); /** - * Free an owned [`vx_dtype`] object. + * Decrease reference count on vx_dtype or free if there are no other references */ void vx_dtype_free(const vx_dtype *ptr); @@ -994,26 +1052,21 @@ uint8_t vx_dtype_decimal_precision(const vx_dtype *dtype); int8_t vx_dtype_decimal_scale(const vx_dtype *dtype); /** - * Return a borrowed reference to the [`vx_struct_fields`] of a struct. - * - * The returned pointer is valid as long as the struct dtype is valid. - * Do NOT free the returned pointer - it shares the lifetime of the struct dtype. + * If "dtype" is DTYPE_STRUCT, return owned vx_struct_fields for this struct, + * return NULL otherwise. Returned vx_struct_fields must be released with + * vx_struct_fields_free. */ const vx_struct_fields *vx_dtype_struct_dtype(const vx_dtype *dtype); /** - * Returns the element type of a list. - * - * The returned pointer is valid as long as the list dtype is valid. - * Do NOT free the returned dtype pointer - it shares the lifetime of the list dtype. + * If "dtype" is DTYPE_LIST, return its owned element dtype, return NULL + * otherwise. Returned dtype must be released with vx_dtype_free. */ const vx_dtype *vx_dtype_list_element(const vx_dtype *dtype); /** - * Returns the element type of a fixed-size list. - * - * The returned pointer is valid as long as the fixed-size list dtype is valid. - * Do NOT free the returned dtype pointer - it shares the lifetime of the fixed-size list dtype. + * If "dtype" is DTYPE_FIXED_SIZE_LIST, return its owned element dtype, return + * NULL otherwise. Returned dtype must be released with vx_dtype_free. */ const vx_dtype *vx_dtype_fixed_size_list_element(const vx_dtype *dtype); @@ -1043,9 +1096,12 @@ bool vx_dtype_is_timestamp(const DType *dtype); uint8_t vx_dtype_time_unit(const DType *dtype); /** - * Returns the time zone, assuming the type is time. Caller is responsible for freeing the returned pointer. + * Return time zone assuming "dtype" is time. + * Returns {NULL, 0} when timestamp has no time zone. + * + * Returned view is valid as long as "dtype" is valid. */ -const vx_string *vx_dtype_time_zone(const DType *dtype); +vx_view vx_dtype_time_zone(const DType *dtype); /** * Convert a dtype to ArrowSchema. @@ -1067,22 +1123,25 @@ int vx_dtype_to_arrow_schema(const vx_dtype *dtype, FFI_ArrowSchema *schema, vx_ const vx_dtype *vx_dtype_from_arrow_schema(FFI_ArrowSchema *schema, vx_error **err); /** - * Free an owned [`vx_error`] object. + * Free a vx_error */ -void vx_error_free(vx_error *ptr); +void vx_error_free(const vx_error *ptr); /** - * Returns the error message from the given Vortex error. - * - * The returned pointer is valid as long as the error is valid. - * Do NOT free the returned string pointer - it shares the lifetime of the error. + * Return error message for this error. + * Returned view is valid while "error" is valid. + */ +vx_view vx_error_message(const vx_error *error); + +/** + * Return category code for "error". */ -const vx_string *vx_error_get_message(const vx_error *error); +vx_error_code vx_error_get_code(const vx_error *error); /** - * Free an owned [`vx_expression`] object. + * Free a vx_expression */ -void vx_expression_free(vx_expression *ptr); +void vx_expression_free(const vx_expression *ptr); /** * Create a root expression. A root expression, applied to an array in @@ -1102,6 +1161,11 @@ void vx_expression_free(vx_expression *ptr); */ vx_expression *vx_expression_root(void); +/** + * Reference-clone a vx_expression + */ +vx_expression *vx_expression_clone(const vx_expression *ptr); + /** * Create a literal expression from a scalar. * @@ -1150,7 +1214,7 @@ vx_expression *vx_expression_literal(const vx_scalar *scalar, vx_error **err); * vx_expression_free(select); * vx_expression_free(root); */ -vx_expression *vx_expression_select(const char *const *names, size_t len, const vx_expression *child); +vx_expression *vx_expression_select(const vx_view *names, size_t len, const vx_expression *child); /** * Create an AND expression for multiple child expressions. @@ -1194,7 +1258,7 @@ vx_expression_binary(vx_binary_operator operator_, const vx_expression *lhs, con * * Returns the logical negation of the input boolean expression. */ -const vx_expression *vx_expression_not(const vx_expression *child); +vx_expression *vx_expression_not(const vx_expression *child); /** * Create an expression that checks for null values. @@ -1212,8 +1276,11 @@ vx_expression *vx_expression_is_null(const vx_expression *child); * * Example: if child is Struct { name=u8, age=u16 } and we do * vx_expression_get_item("name", child), output type will be DTYPE_U8 + * + * "item" is copied. Returns NULL if "child" is NULL or "item" is not valid + * UTF-8. */ -vx_expression *vx_expression_get_item(const char *item, const vx_expression *child); +vx_expression *vx_expression_get_item(vx_view item, const vx_expression *child); /** * Create an expression that checks if a value is contained in a list. @@ -1222,22 +1289,6 @@ vx_expression *vx_expression_get_item(const char *item, const vx_expression *chi */ vx_expression *vx_expression_list_contains(const vx_expression *list, const vx_expression *value); -/** - * Clone a borrowed [`vx_file`], returning an owned [`vx_file`]. - * Must be released with [`vx_file_free`]. - */ -const vx_file *vx_file_clone(const vx_file *ptr); - -/** - * Free an owned [`vx_file`] object. - */ -void vx_file_free(const vx_file *ptr); - -void vx_file_write_array(const vx_session *session, - const char *path, - const vx_array *array, - vx_error **error_out); - /** * Set the stderr logger to output at the specified level. * @@ -1246,24 +1297,19 @@ void vx_file_write_array(const vx_session *session, void vx_set_log_level(vx_log_level level); /** - * Free an owned [`vx_scalar`] object. + * Free a vx_scalar */ -void vx_scalar_free(vx_scalar *ptr); +void vx_scalar_free(const vx_scalar *ptr); /** - * Clone a borrowed scalar handle. - * - * The input scalar handle is not consumed. The returned scalar handle must be - * released with vx_scalar_free. Returns NULL when given a NULL scalar handle. + * Clone a scalar handle. + * If scalar is NULL, returns NULL. */ vx_scalar *vx_scalar_clone(const vx_scalar *scalar); /** - * Return the data type of a scalar. - * - * The returned data type handle borrows storage from the scalar handle, so its - * lifetime is bound to the scalar handle. It MUST NOT be freed separately. - * Returns NULL when given a NULL scalar handle. + * Return scalar's dtype. + * If scalar is NULL, returns NULL. */ const vx_dtype *vx_scalar_dtype(const vx_scalar *scalar); @@ -1340,11 +1386,10 @@ vx_scalar *vx_scalar_new_f16_bits(uint16_t bits, bool is_nullable); /** * Create a UTF-8 scalar. * - * The byte range is copied into the scalar. A NULL data pointer is allowed only - * for an empty byte range. Invalid UTF-8 returns NULL and writes the error - * output. + * The string bytes are copied into the scalar. Invalid UTF-8 returns NULL and + * writes the error output. */ -vx_scalar *vx_scalar_new_utf8(const char *ptr, size_t len, bool is_nullable, vx_error **err); +vx_scalar *vx_scalar_new_utf8(vx_view value, bool is_nullable, vx_error **err); /** * Create a binary scalar. @@ -1358,9 +1403,11 @@ vx_scalar *vx_scalar_new_binary(const uint8_t *ptr, size_t len, bool is_nullable /** * Create a typed null scalar. * - * The data type handle is borrowed, not consumed. The returned scalar uses a - * nullable copy of that logical type, regardless of the input type's top-level - * nullability. A NULL data type handle returns NULL and writes the error output. + * "dtype" is not consumed, you can use it after calling this function. Returned + * scalar uses a nullable copy of that logical type, regardless of the input + * type's top-level nullability. + * + * Returns NULL and sets "err" on error or NULL dtype. */ vx_scalar *vx_scalar_new_null(const vx_dtype *dtype, vx_error **err); @@ -1435,10 +1482,8 @@ vx_scalar *vx_scalar_new_decimal_i256_le(const uint8_t *bytes32, /** * Create a list scalar. * - * The element data type handle is borrowed, not consumed. Child scalar handles - * are cloned into the list value, so the caller keeps ownership of the handle - * array and each scalar in it. A NULL child handle array is allowed only for an - * empty list. Child values are validated against the element logical type. + * "element_dtype" and "elements" are not consumed, you can use them after + * calling this function. If len is 0, you can pass NULL to "elements". */ vx_scalar *vx_scalar_new_list(const vx_dtype *element_dtype, const vx_scalar *const *elements, @@ -1449,27 +1494,21 @@ vx_scalar *vx_scalar_new_list(const vx_dtype *element_dtype, /** * Create a fixed-size list scalar. * - * The element data type handle is borrowed, not consumed. The number of child - * scalars becomes the fixed-size list width and must fit in a 32-bit unsigned - * integer. Child scalar handles are cloned into the list value, so the caller - * keeps ownership of the handle array and each scalar in it. A NULL child - * handle array is allowed only for an empty list. Child values are validated - * against the element logical type. + * "element_dtype" and "elements" are not consumed, you can use them after + * calling this function. If len is 0, you can pass NULL to "elements". + * "len" must fit in uint32_t. */ vx_scalar *vx_scalar_new_fixed_size_list(const vx_dtype *element_dtype, const vx_scalar *const *elements, - size_t len, + uint32_t len, bool is_nullable, vx_error **err); /** * Create a struct scalar. * - * The struct data type handle is borrowed, not consumed. Field scalar handles - * are cloned into the struct value, so the caller keeps ownership of the handle - * array and each scalar in it. Field count and field logical types are validated - * against the struct logical type. A NULL field handle array is allowed only for - * an empty struct value. + * "struct_dtype" and "fields" are not consumed, you can use them after calling + * this function. If len is 0, you can pass NULL to "fields". */ vx_scalar *vx_scalar_new_struct(const vx_dtype *struct_dtype, const vx_scalar *const *fields, @@ -1477,21 +1516,19 @@ vx_scalar *vx_scalar_new_struct(const vx_dtype *struct_dtype, vx_error **err); /** - * Free an owned [`vx_scan`] object. + * Free a vx_scan */ -void vx_scan_free(vx_scan *ptr); +void vx_scan_free(const vx_scan *ptr); /** - * Free an owned [`vx_partition`] object. + * Free a vx_partition */ -void vx_partition_free(vx_partition *ptr); +void vx_partition_free(const vx_partition *ptr); /** * Scan a data source. * - * Return an owned scan that must be freed with vx_scan_free. A scan may be - * consumed only once. - * + * A scan may be consumed only once. * "options" and "estimate" may be NULL. * * If "options" is NULL, all rows and columns are returned. @@ -1506,17 +1543,14 @@ vx_scan *vx_data_source_scan(const vx_data_source *data_source, vx_error **err); /** - * Return borrowed vx_scan's dtype. + * Return scan's dtype. * This function will fail if called after vx_scan_next_partition. - * Called must not free the returned pointer as its lifetime is bound to the - * lifetime of the scan. * On error returns NULL and sets "err". */ const vx_dtype *vx_scan_dtype(const vx_scan *scan, vx_error **err); /** - * Return an owned partition from a scan. - * The returned partition must be freed with vx_partition_free. + * Return an partition from a scan. * * On success returns a partition. * On exhaustion (no more partitions in scan) returns NULL but doesn't set @@ -1555,8 +1589,7 @@ int vx_partition_scan_arrow(const vx_session *session, vx_error **err); /** - * Return an owned owned array from a partition. - * The returned array must be freed with vx_array_free. + * Return an array from a partition. * * On success returns an array. * On exhaustion (no more arrays in partition) returns NULL but doesn't set @@ -1568,9 +1601,9 @@ int vx_partition_scan_arrow(const vx_session *session, const vx_array *vx_partition_next(vx_partition *partition, vx_error **err); /** - * Free an owned [`vx_session`] object. + * Free a vx_session */ -void vx_session_free(vx_session *ptr); +void vx_session_free(const vx_session *ptr); /** * Create a new Vortex session. @@ -1586,18 +1619,29 @@ vx_session *vx_session_new(void); */ vx_session *vx_session_clone(const vx_session *session); +/** + * Increase reference count on vx_file + */ +const vx_file *vx_file_clone(const vx_file *ptr); + +/** + * Decrease reference count on vx_file or free if there are no other references + */ +void vx_file_free(const vx_file *ptr); + /** * Opens a writable array stream, where sink is used to push values into the stream. * To close the stream close the sink with `vx_array_sink_close`. + * "path" is copied. */ -vx_array_sink *vx_array_sink_open_file(const vx_session *session, - const char *path, - const vx_dtype *dtype, - vx_error **error_out); +vx_array_sink * +vx_array_sink_open_file(const vx_session *session, vx_view path, const vx_dtype *dtype, vx_error **error_out); /** * Push an array into a file sink. - * Does not take ownership of array + * Does not take ownership of array. + * + * Errors if array's DType doesn't match sink's DType. */ void vx_array_sink_push(vx_array_sink *sink, const vx_array *array, vx_error **error_out); @@ -1608,40 +1652,15 @@ void vx_array_sink_push(vx_array_sink *sink, const vx_array *array, vx_error **e void vx_array_sink_close(vx_array_sink *sink, vx_error **error_out); /** - * Clone a borrowed [`vx_string`], returning an owned [`vx_string`]. - * Must be released with [`vx_string_free`]. - */ -const vx_string *vx_string_clone(const vx_string *ptr); - -/** - * Free an owned [`vx_string`] object. + * Abort an array sink. File footer is not written, and file is left invalid. + * Don't use sink after this call. */ -void vx_string_free(const vx_string *ptr); +void vx_array_sink_abort(vx_array_sink *sink); /** - * Create a new Vortex UTF-8 string by copying from a pointer and length. + * Free a vx_struct_column_builder */ -const vx_string *vx_string_new(const char *ptr, size_t len); - -/** - * Create a new Vortex UTF-8 string by copying from a null-terminated C-style string. - */ -const vx_string *vx_string_new_from_cstr(const char *ptr); - -/** - * Return the length of the string in bytes. - */ -size_t vx_string_len(const vx_string *ptr); - -/** - * Return the pointer to the string data. - */ -const char *vx_string_ptr(const vx_string *ptr); - -/** - * Free an owned [`vx_struct_column_builder`] object. - */ -void vx_struct_column_builder_free(vx_struct_column_builder *ptr); +void vx_struct_column_builder_free(const vx_struct_column_builder *ptr); /** * Create a new column-wise struct array builder with given validity and a @@ -1660,7 +1679,7 @@ vx_struct_column_builder *vx_struct_column_builder_new(const vx_validity *validi * deallocate it using vx_struct_column_builder_free. */ void vx_struct_column_builder_add_field(vx_struct_column_builder *builder, - const char *name, + vx_view name, const vx_array *field, vx_error **error); @@ -1690,9 +1709,9 @@ void vx_struct_column_builder_add_field(vx_struct_column_builder *builder, const vx_array *vx_struct_column_builder_finalize(vx_struct_column_builder *builder, vx_error **error); /** - * Free an owned [`vx_struct_fields`] object. + * Free a vx_struct_fields */ -void vx_struct_fields_free(vx_struct_fields *ptr); +void vx_struct_fields_free(const vx_struct_fields *ptr); /** * Return the number of fields in the struct dtype. @@ -1700,28 +1719,23 @@ void vx_struct_fields_free(vx_struct_fields *ptr); uint64_t vx_struct_fields_nfields(const vx_struct_fields *dtype); /** - * Return a borrowed reference to the name of the field at the given index. + * Return field name at a given index. + * If index is out of bounds, returns {NULL, 0}. * - * The returned pointer is valid as long as the struct fields is valid. - * Do NOT free the returned string pointer - it shares the lifetime of the struct fields. - * Returns null if the index is out of bounds. + * Returned view is valid as long as "dtype" is valid. */ -const vx_string *vx_struct_fields_field_name(const vx_struct_fields *dtype, size_t idx); +vx_view vx_struct_fields_field_name(const vx_struct_fields *dtype, size_t idx); /** - * Returns an *owned* reference to the dtype of the field at the given index. - * - * The return type is owned since struct dtypes can be lazily parsed from a binary format, in - * which case it's not possible to return a borrowed reference to the field dtype. - * - * Returns null if the index is out of bounds or if the field dtype cannot be parsed. + * Return an owned dtype of the field at a given index. + * Returns NULL if index is out of bounds or if dtype cannot be parsed. */ const vx_dtype *vx_struct_fields_field_dtype(const vx_struct_fields *dtype, size_t idx); /** - * Free an owned [`vx_struct_fields_builder`] object. + * Free a vx_struct_fields_builder */ -void vx_struct_fields_builder_free(vx_struct_fields_builder *ptr); +void vx_struct_fields_builder_free(const vx_struct_fields_builder *ptr); /** * Create a new struct dtype builder. @@ -1731,12 +1745,13 @@ vx_struct_fields_builder *vx_struct_fields_builder_new(void); /** * Add a field to the struct dtype builder. * - * Takes ownership of both the `name` and `dtype` pointers. - * Must either free or finalize the builder. + * "name" is copied. Takes ownership of "dtype". + * Caller must free or finalize the builder. */ void vx_struct_fields_builder_add_field(vx_struct_fields_builder *builder, - const vx_string *name, - const vx_dtype *dtype); + vx_view name, + const vx_dtype *dtype, + vx_error **error_out); /** * Finalize the struct dtype builder, returning a new `vx_struct_fields`. @@ -1748,3 +1763,16 @@ vx_struct_fields *vx_struct_fields_builder_finalize(vx_struct_fields_builder *bu #ifdef __cplusplus } // extern "C" #endif // __cplusplus + +#include + +/** + * Create a view over a null-terminated C string. + * View is valid as long as "str" is valid + */ +static inline vx_view vx_view_from_cstr(const char *str) { + vx_view s; + s.ptr = str; + s.len = strlen(str); + return s; +} diff --git a/vortex-ffi/examples/dtype.c b/vortex-ffi/examples/dtype.c index 007c3993e33..eb60cb72c4d 100644 --- a/vortex-ffi/examples/dtype.c +++ b/vortex-ffi/examples/dtype.c @@ -58,12 +58,14 @@ void print_struct_dtype(const vx_dtype *dtype) { printf("struct(\n"); for (uint64_t i = 0; i < vx_struct_fields_nfields(fields); ++i) { const vx_dtype *field_dtype = vx_struct_fields_field_dtype(fields, i); - const vx_string *field_name = vx_struct_fields_field_name(fields, i); - printf(" %.*s = ", (int)vx_string_len(field_name), vx_string_ptr(field_name)); + const vx_view field_name = vx_struct_fields_field_name(fields, i); + printf(" %.*s = ", (int)field_name.len, field_name.ptr); print_dtype(field_dtype); vx_dtype_free(field_dtype); } printf(")"); + + vx_struct_fields_free(fields); } void print_list_dtype(const vx_dtype *dtype) { @@ -121,8 +123,8 @@ void print_dtype(const vx_dtype *dtype) { } void print_error(const char *what, const vx_error *error) { - const vx_string *str = vx_error_get_message(error); - fprintf(stderr, "%s: %.*s\n", what, (int)vx_string_len(str), vx_string_ptr(str)); + const vx_view str = vx_error_message(error); + fprintf(stderr, "%s: %.*s\n", what, (int)str.len, str.ptr); } int main(int argc, char **argv) { @@ -138,7 +140,8 @@ int main(int argc, char **argv) { return 1; } - vx_data_source_options ds_options = {.paths = argv[1]}; + vx_view path = vx_view_from_cstr(argv[1]); + vx_data_source_options ds_options = {.paths = &path, .paths_len = 1}; const vx_data_source *data_source = vx_data_source_new(session, &ds_options, &error); if (data_source == NULL) { print_error("Failed to create data source", error); @@ -148,7 +151,9 @@ int main(int argc, char **argv) { } printf("dtype: "); - print_dtype(vx_data_source_dtype(data_source)); + const vx_dtype *dtype = vx_data_source_dtype(data_source); + print_dtype(dtype); + vx_dtype_free(dtype); vx_data_source_free(data_source); vx_session_free(session); diff --git a/vortex-ffi/examples/scan.c b/vortex-ffi/examples/scan.c index 43023642474..67df597c69f 100644 --- a/vortex-ffi/examples/scan.c +++ b/vortex-ffi/examples/scan.c @@ -25,8 +25,8 @@ void print_estimate(const char *what, const vx_estimate *estimate) { } void print_error(const char *what, const vx_error *error) { - const vx_string *str = vx_error_get_message(error); - fprintf(stderr, "%s: %.*s\n", what, (int)vx_string_len(str), vx_string_ptr(str)); + const vx_view str = vx_error_message(error); + fprintf(stderr, "%s: %.*s\n", what, (int)str.len, str.ptr); } struct scan_thread_info { @@ -162,7 +162,8 @@ int main(int argc, char *argv[]) { // A datasource is a reference to some files. // We can request multiple scans from a data source. - vx_data_source_options ds_options = {.paths = paths}; + vx_view path = vx_view_from_cstr(paths); + vx_data_source_options ds_options = {.paths = &path, .paths_len = 1}; vx_error *error = NULL; const vx_data_source *data_source = vx_data_source_new(session, &ds_options, &error); if (data_source == NULL) { diff --git a/vortex-ffi/examples/scan_to_arrow.c b/vortex-ffi/examples/scan_to_arrow.c index ffa83d0b123..a8679158de0 100644 --- a/vortex-ffi/examples/scan_to_arrow.c +++ b/vortex-ffi/examples/scan_to_arrow.c @@ -16,14 +16,13 @@ const char *usage = "Scan vortex files to Arrow\n" "Usage: scan_to_arrow \n"; void print_error(const char *what, const vx_error *error) { - const vx_string *str = vx_error_get_message(error); - fprintf(stderr, "%s: %.*s\n", what, (int)vx_string_len(str), vx_string_ptr(str)); + const vx_view str = vx_error_message(error); + fprintf(stderr, "%s: %.*s\n", what, (int)str.len, str.ptr); } void execute_scan(vx_session *session, vx_scan *scan) { vx_error *error = NULL; - // Returned dtype is owned and mustn't be freed const vx_dtype *dtype = vx_scan_dtype(scan, &error); if (dtype == NULL) { print_error("Failed to get scan dtype", error); @@ -37,6 +36,7 @@ void execute_scan(vx_session *session, vx_scan *scan) { vx_error_free(error); return; } + vx_dtype_free(dtype); char schema_buf[1024 * 10]; const int schema_len = ArrowSchemaToString(&schema, schema_buf, sizeof schema_buf, 1); @@ -97,7 +97,8 @@ int main(int argc, char *argv[]) { return -1; } - vx_data_source_options ds_options = {.paths = paths}; + vx_view path = vx_view_from_cstr(paths); + vx_data_source_options ds_options = {.paths = &path, .paths_len = 1}; vx_error *error = NULL; const vx_data_source *data_source = vx_data_source_new(session, &ds_options, &error); if (data_source == NULL) { diff --git a/vortex-ffi/examples/write_sample.c b/vortex-ffi/examples/write_sample.c index e807b21b5a9..456d3bba166 100644 --- a/vortex-ffi/examples/write_sample.c +++ b/vortex-ffi/examples/write_sample.c @@ -3,7 +3,6 @@ #include "vortex.h" #include #include -#include #define SAMPLE_ROWS 200 @@ -15,22 +14,22 @@ const vx_dtype *sample_dtype(void) { vx_struct_fields_builder *builder = vx_struct_fields_builder_new(); const char *age = "age"; - const vx_string *age_name = vx_string_new(age, strlen(age)); + const vx_dtype *age_type = vx_dtype_new_primitive(PTYPE_U8, false); - vx_struct_fields_builder_add_field(builder, age_name, age_type); + vx_struct_fields_builder_add_field(builder, vx_view_from_cstr(age), age_type, NULL); const char *height = "height"; - const vx_string *height_name = vx_string_new(height, strlen(height)); + const vx_dtype *height_type = vx_dtype_new_primitive(PTYPE_U16, true); - vx_struct_fields_builder_add_field(builder, height_name, height_type); + vx_struct_fields_builder_add_field(builder, vx_view_from_cstr(height), height_type, NULL); vx_struct_fields *fields = vx_struct_fields_builder_finalize(builder); return vx_dtype_new_struct(fields, false); } void print_error(const char *what, const vx_error *error) { - const vx_string *str = vx_error_get_message(error); - fprintf(stderr, "%s: %.*s\n", what, (int)vx_string_len(str), vx_string_ptr(str)); + const vx_view str = vx_error_message(error); + fprintf(stderr, "%s: %.*s\n", what, (int)str.len, str.ptr); } const vx_array *sample_array(void) { @@ -53,7 +52,7 @@ const vx_array *sample_array(void) { return NULL; } - vx_struct_column_builder_add_field(builder, "age", age_array, &error); + vx_struct_column_builder_add_field(builder, vx_view_from_cstr("age"), age_array, &error); vx_array_free(age_array); if (error != NULL) { print_error("Error adding age array field to root array", error); @@ -72,7 +71,7 @@ const vx_array *sample_array(void) { return NULL; } - vx_struct_column_builder_add_field(builder, "height", height_array, &error); + vx_struct_column_builder_add_field(builder, vx_view_from_cstr("height"), height_array, &error); vx_array_free(height_array); if (error != NULL) { print_error("Error adding height array field to root array", error); @@ -107,7 +106,7 @@ int main(int argc, char *argv[]) { const vx_dtype *dtype = sample_dtype(); vx_error *error = NULL; - vx_array_sink *sink = vx_array_sink_open_file(session, output, dtype, &error); + vx_array_sink *sink = vx_array_sink_open_file(session, vx_view_from_cstr(output), dtype, &error); vx_dtype_free(dtype); if (error != NULL) { diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index a87178446de..7c20bb72d46 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -5,6 +5,7 @@ //! FFI interface for working with Vortex Arrays. use std::ffi::c_void; use std::ptr; +use std::ptr::NonNull; use std::sync::Arc; use arrow_array::array::make_array; @@ -13,33 +14,42 @@ use arrow_array::ffi::FFI_ArrowSchema; use arrow_array::ffi::from_ffi; use paste::paste; 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::Bool; use vortex::array::arrays::NullArray; +use vortex::array::arrays::Primitive; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; +use vortex::array::arrays::VarBinView; +use vortex::array::arrays::bool::BoolArrayExt; use vortex::array::arrays::struct_::StructArrayExt; -use vortex::array::arrow::FromArrowArray; +use vortex::array::legacy_session; use vortex::array::validity::Validity; use vortex::buffer::Buffer; use vortex::dtype::DType; use vortex::dtype::half::f16; 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_arrow::FromArrowArray; use crate::arc_wrapper; -use crate::binary::vx_binary; use crate::dtype::vx_dtype; use crate::dtype::vx_dtype_variant; +use crate::error::try_or; use crate::error::try_or_default; use crate::error::vx_error; use crate::error::write_error; use crate::expression::vx_expression; use crate::ptype::vx_ptype; -use crate::string::vx_string; +use crate::session::vx_session; +use crate::session::vx_session_ref; +use crate::string::vx_view; arc_wrapper!( /// Arrays are reference-counted handles to owned memory buffers that hold @@ -55,7 +65,7 @@ arc_wrapper!( /// array is a cheap operation. /// /// Unless stated explicitly, all operations with vx_array don't take - /// ownership of it, and thus it must be freed by the caller. + /// ownership of it, and thus the array must be freed by the caller. ArrayRef, vx_array ); @@ -206,19 +216,17 @@ pub unsafe extern "C-unwind" fn vx_array_len(array: *const vx_array) -> usize { vx_array::as_ref(array).len() } -/// Get the [`struct@crate::dtype::vx_dtype`] of the array. -/// -/// The returned pointer is valid as long as the array is valid. -/// Do NOT free the returned dtype pointer - it shares the lifetime of the array. +/// Get array's dtype #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_array_dtype(array: *const vx_array) -> *const vx_dtype { - vx_dtype::new_ref(vx_array::as_ref(array).dtype()) + vx_dtype::new(Arc::new(vx_array::as_ref(array).dtype().clone())) } -// Return an owned field for array at index. +// Return a field for array at index. // Returns NULL and sets error_out if index is out of bounds or array doesn't // have dtype DTYPE_STRUCT. #[unsafe(no_mangle)] +#[allow(clippy::disallowed_methods)] pub unsafe extern "C-unwind" fn vx_array_get_field( array: *const vx_array, index: usize, @@ -227,7 +235,7 @@ pub unsafe extern "C-unwind" fn vx_array_get_field( try_or_default(error_out, || { let array = vx_array::as_ref(array); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = legacy_session().create_execution_ctx(); let field_array = array .clone() .execute::(&mut ctx)? @@ -258,6 +266,7 @@ pub unsafe extern "C-unwind" fn vx_array_slice( /// validity array. Sets error if index is out of bounds or underlying validity /// array is corrupted. #[unsafe(no_mangle)] +#[allow(clippy::disallowed_methods)] pub unsafe extern "C-unwind" fn vx_array_element_is_invalid( array: *const vx_array, index: usize, @@ -265,12 +274,13 @@ pub unsafe extern "C-unwind" fn vx_array_element_is_invalid( ) -> bool { try_or_default(error, || { vortex_ensure!(!array.is_null()); - vx_array::as_ref(array).is_invalid(index, &mut LEGACY_SESSION.create_execution_ctx()) + vx_array::as_ref(array).is_invalid(index, &mut legacy_session().create_execution_ctx()) }) } /// Check how many items in the array are invalid (null). #[unsafe(no_mangle)] +#[allow(clippy::disallowed_methods)] pub unsafe extern "C-unwind" fn vx_array_invalid_count( array: *const vx_array, error_out: *mut *mut vx_error, @@ -278,7 +288,7 @@ pub unsafe extern "C-unwind" fn vx_array_invalid_count( try_or_default(error_out, || { vortex_ensure!(!array.is_null()); let array = vx_array::as_ref(array); - array.invalid_count(&mut LEGACY_SESSION.create_execution_ctx()) + array.invalid_count(&mut legacy_session().create_execution_ctx()) }) } @@ -295,11 +305,18 @@ unsafe fn primitive_from_raw( ptr: *const T, len: usize, validity: &vx_validity, + error: *mut *mut vx_error, ) -> *const vx_array { - let slice = unsafe { std::slice::from_raw_parts(ptr, len) }; - let buffer = Buffer::copy_from(slice); - let array = PrimitiveArray::new(buffer, validity.into()); - vx_array::new(Arc::new(array.into_array())) + try_or_default(error, || { + let slice = if ptr.is_null() { + unsafe { std::slice::from_raw_parts(NonNull::dangling().as_ptr(), len) } + } else { + unsafe { std::slice::from_raw_parts(ptr, len) } + }; + let buffer = Buffer::copy_from(slice); + let array = PrimitiveArray::try_new(buffer, validity.into())?; + Ok(vx_array::new(Arc::new(array.into_array()))) + }) } /// Create a new primitive array from an existing buffer. @@ -331,17 +348,35 @@ pub extern "C-unwind" fn vx_array_new_primitive( let validity = unsafe { &*validity }; match ptype { - vx_ptype::PTYPE_U8 => unsafe { primitive_from_raw(ptr as *const u8, len, validity) }, - vx_ptype::PTYPE_U16 => unsafe { primitive_from_raw(ptr as *const u16, len, validity) }, - vx_ptype::PTYPE_U32 => unsafe { primitive_from_raw(ptr as *const u32, len, validity) }, - vx_ptype::PTYPE_U64 => unsafe { primitive_from_raw(ptr as *const u64, len, validity) }, - vx_ptype::PTYPE_I8 => unsafe { primitive_from_raw(ptr as *const i8, len, validity) }, - vx_ptype::PTYPE_I16 => unsafe { primitive_from_raw(ptr as *const i16, len, validity) }, - vx_ptype::PTYPE_I32 => unsafe { primitive_from_raw(ptr as *const i32, len, validity) }, - vx_ptype::PTYPE_I64 => unsafe { primitive_from_raw(ptr as *const i64, len, validity) }, - vx_ptype::PTYPE_F16 => unsafe { primitive_from_raw(ptr as *const f16, len, validity) }, - vx_ptype::PTYPE_F32 => unsafe { primitive_from_raw(ptr as *const f32, len, validity) }, - vx_ptype::PTYPE_F64 => unsafe { primitive_from_raw(ptr as *const f64, len, validity) }, + vx_ptype::PTYPE_U8 => unsafe { primitive_from_raw(ptr as *const u8, len, validity, error) }, + vx_ptype::PTYPE_U16 => unsafe { + primitive_from_raw(ptr as *const u16, len, validity, error) + }, + vx_ptype::PTYPE_U32 => unsafe { + primitive_from_raw(ptr as *const u32, len, validity, error) + }, + vx_ptype::PTYPE_U64 => unsafe { + primitive_from_raw(ptr as *const u64, len, validity, error) + }, + vx_ptype::PTYPE_I8 => unsafe { primitive_from_raw(ptr as *const i8, len, validity, error) }, + vx_ptype::PTYPE_I16 => unsafe { + primitive_from_raw(ptr as *const i16, len, validity, error) + }, + vx_ptype::PTYPE_I32 => unsafe { + primitive_from_raw(ptr as *const i32, len, validity, error) + }, + vx_ptype::PTYPE_I64 => unsafe { + primitive_from_raw(ptr as *const i64, len, validity, error) + }, + vx_ptype::PTYPE_F16 => unsafe { + primitive_from_raw(ptr as *const f16, len, validity, error) + }, + vx_ptype::PTYPE_F32 => unsafe { + primitive_from_raw(ptr as *const f32, len, validity, error) + }, + vx_ptype::PTYPE_F64 => unsafe { + primitive_from_raw(ptr as *const f64, len, validity, error) + }, } } @@ -355,9 +390,6 @@ pub extern "C-unwind" fn vx_array_new_primitive( /// `nullable` controls the top-level nullability of the resulting array's dtype. For an Arrow /// record batch (which has no top-level validity) pass `false`. /// -/// The imported buffers are referenced zero-copy where possible; the returned array keeps the -/// Arrow data alive until it is freed with [`vx_array_free`]. -/// /// On error, returns NULL and sets `error_out`. /// /// Example: @@ -394,8 +426,9 @@ macro_rules! ffiarray_get_ptype { pub unsafe extern "C-unwind" fn [](array: *const vx_array, index: usize) -> $ptype { let array = vx_array::as_ref(array); // TODO(joe): propagate this error up instead of expecting + #[allow(clippy::disallowed_methods)] let value = array - .execute_scalar(index, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(index, &mut legacy_session().create_execution_ctx()) .vortex_expect("scalar_at failed"); // TODO(joe): propagate this error up instead of expecting value.as_primitive() @@ -407,8 +440,9 @@ macro_rules! ffiarray_get_ptype { pub unsafe extern "C-unwind" fn [](array: *const vx_array, index: usize) -> $ptype { let array = vx_array::as_ref(array); // TODO(joe): propagate this error up instead of expecting + #[allow(clippy::disallowed_methods)] let value = array - .execute_scalar(index, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(index, &mut legacy_session().create_execution_ctx()) .vortex_expect("scalar_at failed"); // TODO(joe): propagate this error up instead of expecting value.as_extension() @@ -433,44 +467,156 @@ ffiarray_get_ptype!(f16); ffiarray_get_ptype!(f32); ffiarray_get_ptype!(f64); -/// Return the utf-8 string at `index` in the array. The pointer will be null if the value at `index` is null. -/// The caller must free the returned pointer. +/// SAFETY: "array" must be null or a valid "vx_array" +unsafe fn varbinview_at( + array: *const vx_array, + index: usize, + want_utf8: bool, + error_out: *mut *mut vx_error, +) -> vx_view { + try_or(error_out, vx_view::null(), || { + let array = unsafe { vx_array_ref(array) }?; + vortex_ensure!(index < array.len(), "index {index} out of bounds"); + let dtype_matches = if want_utf8 { + matches!(array.dtype(), DType::Utf8(_)) + } else { + matches!(array.dtype(), DType::Binary(_)) + }; + vortex_ensure!( + dtype_matches, + "expected a {} array, got {}", + if want_utf8 { "Utf8" } else { "Binary" }, + array.dtype() + ); + let Some(views) = array.as_opt::() else { + vortex_bail!("expected a canonical array, got {}", array.encoding_id()); + }; + Ok(vx_view::from_bytes(views.bytes_at(index).as_slice())) + }) +} + +/// Return UTF-8 string at "index" in a canonical Utf8 array. +/// +/// For invalid elements the returned value is unspecified, check validity via +/// vx_array_get_validity. +/// Returned view is valid as long as "array" is valid. +/// Errors if index is out of bounds or array is not a canonical Utf8 array. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_array_get_utf8( +pub unsafe extern "C-unwind" fn vx_array_utf8_at( array: *const vx_array, - index: u32, -) -> *const vx_string { - let array = vx_array::as_ref(array); - // TODO(joe): propagate this error up instead of expecting - let value = array - .execute_scalar(index as usize, &mut LEGACY_SESSION.create_execution_ctx()) - .vortex_expect("scalar_at failed"); - let utf8_scalar = value.as_utf8(); - if let Some(buffer) = utf8_scalar.value() { - vx_string::new(Arc::from(buffer.as_str())) - } else { - ptr::null() - } + index: usize, + error_out: *mut *mut vx_error, +) -> vx_view { + unsafe { varbinview_at(array, index, true, error_out) } } -/// Return the binary at `index` in the array. The pointer will be null if the value at `index` is null. -/// The caller must free the returned pointer. +/// Return a binary string at "index" in a canonical Binary array. +/// +/// For invalid elements the returned value is unspecified, check validity via +/// vx_array_get_validity. +/// Returned view is valid as long as "array" is valid. +/// Errors if index is out of bounds or array is not a canonical Binary array. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_array_get_binary( +pub unsafe extern "C-unwind" fn vx_array_binary_at( array: *const vx_array, - index: u32, -) -> *const vx_binary { + index: usize, + error_out: *mut *mut vx_error, +) -> vx_view { + unsafe { varbinview_at(array, index, false, error_out) } +} + +/// For a canonical Bool array, return bool at "index". +/// For invalid elements returned value is unspecified, check validity via +/// vx_array_get_validity. +/// +/// Panics if "array" is not canonical - call vx_array_canonicalize first. +/// Panics if "array" is not a Bool array. +/// Panics if "index" is out of bounds. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_array_get_bool(array: *const vx_array, index: usize) -> bool { let array = vx_array::as_ref(array); - // TODO(joe): propagate this error up instead of expecting - let value = array - .execute_scalar(index as usize, &mut LEGACY_SESSION.create_execution_ctx()) - .vortex_expect("scalar_at failed"); - let binary_scalar = value.as_binary(); - if let Some(bytes) = binary_scalar.value() { - vx_binary::new(Arc::from(bytes.as_bytes())) - } else { - ptr::null() + let bool_array = array + .as_opt::() + .vortex_expect("vx_array_get_bool requires a canonical Bool array"); + let bits = bool_array.to_bit_buffer(); + if index >= bits.len() { + vortex_panic!( + "index {index} out of bounds for array of length {}", + bits.len() + ); } + bits.value(index) +} + +/// Decode array into its canonical form. +/// +/// On error returns NULL and "sets error_out". +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_array_canonicalize( + session: *const vx_session, + array: *const vx_array, + error_out: *mut *mut vx_error, +) -> *const vx_array { + try_or_default(error_out, || { + let session = unsafe { vx_session_ref(session) }?; + let array = unsafe { vx_array_ref(array) }?; + let mut ctx = session.create_execution_ctx(); + let canonical = array.clone().execute::(&mut ctx)?; + Ok(vx_array::new(Arc::new(canonical.into_array()))) + }) +} + +/// Return a pointer to the values buffer of a canonical Primitive array. +/// Pointer is valid as long as "array" is valid. +/// +/// Errors if array is not a canonical Primitive. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_array_data_ptr_primitive( + array: *const vx_array, + error_out: *mut *mut vx_error, +) -> *const c_void { + try_or(error_out, ptr::null(), || { + let array = unsafe { vx_array_ref(array) }?; + let primitive = array.as_opt::().ok_or_else(|| { + vortex_err!( + "vx_array_data_ptr_primitive requires a canonical Primitive array, got {}", + array.encoding_id() + ) + })?; + let bytes = primitive + .buffer_handle() + .as_host_opt() + .ok_or_else(|| vortex_err!("array buffer is not in host memory"))?; + Ok(bytes.as_ptr().cast()) + }) +} + +/// Return a pointer to the bitpacked buffer of a canonical Bool array. +/// Pointer is valid as long as "array" is valid. +/// +/// Writes bit offset of the first element into "bit_offset_out". +/// "bit_offset_out" must not be NULL. +/// +/// Errors if array is not a canonical Bool. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_array_data_ptr_bool( + array: *const vx_array, + bit_offset_out: *mut usize, + error_out: *mut *mut vx_error, +) -> *const c_void { + try_or(error_out, ptr::null(), || { + let array = unsafe { vx_array_ref(array) }?; + vortex_ensure!(!bit_offset_out.is_null(), "null bit_offset_out"); + let bool_array = array.as_opt::().ok_or_else(|| { + vortex_err!( + "vx_array_data_ptr_bool requires a canonical Bool array, got {}", + array.encoding_id() + ) + })?; + let bits = bool_array.to_bit_buffer(); + unsafe { bit_offset_out.write(bits.offset()) }; + Ok(bits.inner().as_ptr().cast()) + }) } /// Apply the expression to the array, wrapping it with a ScalarFnArray. @@ -494,6 +640,7 @@ pub unsafe extern "C" fn vx_array_apply( #[cfg(test)] mod tests { use std::ptr; + use std::slice::from_raw_parts; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; @@ -512,12 +659,14 @@ mod tests { use vortex::expr::root; use crate::array::*; - use crate::binary::vx_binary_free; + use crate::dtype::vx_dtype_free; use crate::dtype::vx_dtype_get_variant; use crate::dtype::vx_dtype_variant; use crate::error::vx_error_free; use crate::expression::vx_expression_free; - use crate::string::vx_string_free; + use crate::session::vx_session_free; + use crate::session::vx_session_new; + use crate::tests::assert_error; use crate::tests::assert_no_error; #[test] @@ -540,6 +689,7 @@ mod tests { assert_eq!(vx_array_get_i32(ffi_array, 1), 2); assert_eq!(vx_array_get_i32(ffi_array, 2), 3); + vx_dtype_free(array_dtype); vx_array_free(ffi_array); } } @@ -735,56 +885,43 @@ mod tests { #[test] // TODO(joe): enable once this is fixed https://github.com/Amanieu/parking_lot/issues/477 #[cfg_attr(miri, ignore)] - fn test_get_utf8() { + fn test_utf8_binary_at() { unsafe { - let utf8_array = VarBinViewArray::from_iter_str(["hello", "world", "test"]); + let long = "a string that is longer than twelve bytes"; + let utf8_array = + VarBinViewArray::from_iter_nullable_str([Some("hello"), None, Some(long)]); let ffi_array = vx_array::new(Arc::new(utf8_array.into_array())); - assert!(vx_array_has_dtype(ffi_array, vx_dtype_variant::DTYPE_UTF8)); - let vx_str1 = vx_array_get_utf8(ffi_array, 0); - assert_eq!(vx_string::as_str(vx_str1), "hello"); - vx_string_free(vx_str1); - - let vx_str2 = vx_array_get_utf8(ffi_array, 1); - assert_eq!(vx_string::as_str(vx_str2), "world"); - vx_string_free(vx_str2); - - let vx_str3 = vx_array_get_utf8(ffi_array, 2); - assert_eq!(vx_string::as_str(vx_str3), "test"); - vx_string_free(vx_str3); + let mut error = ptr::null_mut(); + let inlined = vx_array_utf8_at(ffi_array, 0, &raw mut error); + assert!(error.is_null()); + assert_eq!(inlined.as_str().unwrap(), "hello"); - vx_array_free(ffi_array); - } - } + vx_array_utf8_at(ffi_array, 1, &raw mut error); + assert!(error.is_null()); - #[test] - // TODO(joe): enable once this is fixed https://github.com/Amanieu/parking_lot/issues/477 - #[cfg_attr(miri, ignore)] - fn test_get_binary() { - unsafe { - let binary_array = VarBinViewArray::from_iter_bin(vec![ - vec![0x01, 0x02, 0x03], - vec![0xFF, 0xEE], - vec![0xAA, 0xBB, 0xCC, 0xDD], - ]); - let ffi_array = vx_array::new(Arc::new(binary_array.into_array())); - assert!(vx_array_has_dtype( - ffi_array, - vx_dtype_variant::DTYPE_BINARY - )); + let buffered = vx_array_utf8_at(ffi_array, 2, &raw mut error); + assert!(error.is_null()); + assert_eq!(buffered.as_str().unwrap(), long); - let vx_bin1 = vx_array_get_binary(ffi_array, 0); - assert_eq!(vx_binary::as_slice(vx_bin1), &[0x01, 0x02, 0x03]); - vx_binary_free(vx_bin1); + vx_array_utf8_at(ffi_array, 3, &raw mut error); + assert_error(error); - let vx_bin2 = vx_array_get_binary(ffi_array, 1); - assert_eq!(vx_binary::as_slice(vx_bin2), &[0xFF, 0xEE]); - vx_binary_free(vx_bin2); + vx_array_free(ffi_array); - let vx_bin3 = vx_array_get_binary(ffi_array, 2); - assert_eq!(vx_binary::as_slice(vx_bin3), &[0xAA, 0xBB, 0xCC, 0xDD]); - vx_binary_free(vx_bin3); + let numbers = + PrimitiveArray::new(buffer![1i32, 2i32], Validity::NonNullable).into_array(); + let ffi_array = vx_array::new(Arc::new(numbers)); + let value = vx_array_utf8_at(ffi_array, 0, &raw mut error); + assert!(value.ptr.is_null()); + assert_error(error); + vx_array_free(ffi_array); + let binary_array = VarBinViewArray::from_iter_bin(vec![vec![0x01, 0x02, 0x03]]); + let ffi_array = vx_array::new(Arc::new(binary_array.into_array())); + let bin = vx_array_binary_at(ffi_array, 0, &raw mut error); + assert!(error.is_null()); + assert_eq!(bin.as_bytes().unwrap(), &[0x01, 0x02, 0x03]); vx_array_free(ffi_array); } } @@ -858,17 +995,14 @@ mod tests { let vx_arr = vx_array::new(Arc::new(array)); assert!(unsafe { vx_array_has_dtype(vx_arr, vx_dtype_variant::DTYPE_STRUCT) }); - // Get dtype reference - this is valid as long as array lives let dtype_ptr = unsafe { vx_array_dtype(vx_arr) }; let variant = unsafe { vx_dtype_get_variant(dtype_ptr) }; assert_eq!(variant, vx_dtype_variant::DTYPE_STRUCT); - // Proper usage: use dtype while array is still alive - // This demonstrates the correct lifetime pattern unsafe { vx_array_free(vx_arr) }; - - // Note: dtype_ptr is now invalid - this test documents the lifetime pattern - // In real usage, don't access dtype_ptr after freeing the array + unsafe { + vx_dtype_free(dtype_ptr); + } } #[test] @@ -933,4 +1067,152 @@ mod tests { vx_array_free(vx); } } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_get_bool() { + let bools = BoolArray::from_iter([true, false, true]); + unsafe { + let array = vx_array::new(Arc::new(bools.into_array())); + assert!(vx_array_get_bool(array, 0)); + assert!(!vx_array_get_bool(array, 1)); + assert!(vx_array_get_bool(array, 2)); + vx_array_free(array); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_canonicalize_nullable() { + let primitive = PrimitiveArray::new( + buffer![10i32, 20i32, 30i32], + Validity::from_iter([true, false, true]), + ); + + unsafe { + let session = vx_session_new(); + let mut error = ptr::null_mut(); + let mut bit_offset = usize::MAX; + + let array = vx_array::new(Arc::new(primitive.into_array())); + let canonical = vx_array_canonicalize(session, array, &raw mut error); + assert_no_error(error); + let data = vx_array_data_ptr_primitive(canonical, &raw mut error); + assert_no_error(error); + + assert_eq!(from_raw_parts(data.cast::(), 3), [10, 20, 30]); + let mut validity = vx_validity { + r#type: vx_validity_type::VX_VALIDITY_NON_NULLABLE, + array: ptr::null(), + }; + + vx_array_get_validity(canonical, &raw mut validity, &raw mut error); + assert_no_error(error); + assert!(matches!( + validity.r#type, + vx_validity_type::VX_VALIDITY_ARRAY + )); + let validity_bools = vx_array_canonicalize(session, validity.array, &raw mut error); + assert_no_error(error); + let bits = vx_array_data_ptr_bool(validity_bools, &raw mut bit_offset, &raw mut error); + assert_no_error(error); + assert_eq!( + *bits.cast::().add(bit_offset / 8) >> (bit_offset % 8) & 0b111, + 0b101 + ); + + vx_array_free(validity_bools); + vx_array_free(validity.array); + vx_array_free(canonical); + vx_array_free(array); + vx_session_free(session); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_canonicalize() { + let primitive = PrimitiveArray::new(buffer![1u8, 2u8], Validity::NonNullable); + + unsafe { + let session = vx_session_new(); + let mut error = ptr::null_mut(); + let mut validity = vx_validity { + r#type: vx_validity_type::VX_VALIDITY_NON_NULLABLE, + array: ptr::null(), + }; + + let array = vx_array::new(Arc::new(primitive.into_array())); + vx_array_get_validity(array, &raw mut validity, &raw mut error); + assert_no_error(error); + assert!(matches!( + validity.r#type, + vx_validity_type::VX_VALIDITY_NON_NULLABLE + )); + assert!(validity.array.is_null()); + vx_array_free(array); + + vx_session_free(session); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_data_ptr_bool() { + let bools = BoolArray::from_iter([ + true, true, false, true, false, true, true, false, true, true, + ]); + + unsafe { + let session = vx_session_new(); + let mut error = ptr::null_mut(); + let mut bit_offset = usize::MAX; + + let array = vx_array::new(Arc::new(bools.into_array())); + let sliced = vx_array_slice(array, 3, 10, &raw mut error); + assert_no_error(error); + + let canonical = vx_array_canonicalize(session, sliced, &raw mut error); + assert_no_error(error); + + let bits = vx_array_data_ptr_bool(canonical, &raw mut bit_offset, &raw mut error); + assert_no_error(error); + assert!(bit_offset < 8); + for (i, expected) in [true, false, true, true, false, true, true] + .into_iter() + .enumerate() + { + let bit = bit_offset + i; + let actual = (*bits.cast::().add(bit / 8) >> (bit % 8)) & 1 == 1; + assert_eq!(actual, expected, "bit {i}"); + } + vx_array_free(canonical); + vx_array_free(sliced); + vx_array_free(array); + + vx_session_free(session); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_data_ptr_error() { + let strings = VarBinViewArray::from_iter_str(["a", "b"]); + + unsafe { + let mut error = ptr::null_mut(); + + let array = vx_array::new(Arc::new(strings.into_array())); + let data = vx_array_data_ptr_primitive(array, &raw mut error); + assert!(data.is_null()); + assert_error(error); + + let mut bit_offset = usize::MAX; + let bits = vx_array_data_ptr_bool(array, &raw mut bit_offset, &raw mut error); + assert!(bits.is_null()); + assert_error(error); + + vx_array_free(array); + } + } } diff --git a/vortex-ffi/src/binary.rs b/vortex-ffi/src/binary.rs deleted file mode 100644 index a13379787a7..00000000000 --- a/vortex-ffi/src/binary.rs +++ /dev/null @@ -1,104 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::ffi::c_char; -use std::slice; - -use crate::arc_dyn_wrapper; - -arc_dyn_wrapper!( - /// Strings for use within Vortex. - [u8], - vx_binary -); - -impl vx_binary { - #[allow(dead_code)] - pub(crate) fn as_slice(ptr: *const vx_binary) -> &'static [u8] { - unsafe { slice::from_raw_parts(vx_binary_ptr(ptr).cast(), vx_binary_len(ptr)) } - } -} - -/// Create a new Vortex UTF-8 string by copying from a pointer and length. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_binary_new(ptr: *const c_char, len: usize) -> *const vx_binary { - let slice = unsafe { slice::from_raw_parts(ptr.cast(), len) }; - vx_binary::new(slice.into()) -} - -/// Return the length of the string in bytes. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_binary_len(ptr: *const vx_binary) -> usize { - vx_binary::as_ref(ptr).len() -} - -/// Return the pointer to the string data. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_binary_ptr(ptr: *const vx_binary) -> *const c_char { - vx_binary::as_ref(ptr).as_ptr().cast() -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_string_new() { - unsafe { - let test_str = "hello world"; - let ptr = test_str.as_ptr().cast(); - let len = test_str.len(); - - let vx_str = vx_binary_new(ptr, len); - assert_eq!(vx_binary_len(vx_str), 11); - assert_eq!(vx_binary::as_slice(vx_str), "hello world".as_bytes()); - - vx_binary_free(vx_str); - } - } - - #[test] - fn test_string_ptr() { - unsafe { - let test_str = "testing".as_bytes(); - let vx_str = vx_binary::new(test_str.into()); - - let ptr = vx_binary_ptr(vx_str); - let len = vx_binary_len(vx_str); - - let slice = slice::from_raw_parts(ptr.cast::(), len); - assert_eq!(slice, "testing".as_bytes()); - - vx_binary_free(vx_str); - } - } - - #[test] - fn test_empty_string() { - unsafe { - let empty = ""; - let ptr = empty.as_ptr().cast(); - let vx_str = vx_binary_new(ptr, 0); - - assert_eq!(vx_binary_len(vx_str), 0); - assert_eq!(vx_binary::as_slice(vx_str), "".as_bytes()); - - vx_binary_free(vx_str); - } - } - - #[test] - fn test_unicode_string() { - unsafe { - let unicode_str = "Hello 世界 🌍"; - let ptr = unicode_str.as_ptr().cast(); - let len = unicode_str.len(); - - let vx_str = vx_binary_new(ptr, len); - assert_eq!(vx_binary_len(vx_str), unicode_str.len()); - assert_eq!(vx_binary::as_slice(vx_str), unicode_str.as_bytes()); - - vx_binary_free(vx_str); - } - } -} diff --git a/vortex-ffi/src/data_source.rs b/vortex-ffi/src/data_source.rs index 8485be285bd..96395e90969 100644 --- a/vortex-ffi/src/data_source.rs +++ b/vortex-ffi/src/data_source.rs @@ -3,8 +3,6 @@ #![allow(non_camel_case_types)] #![deny(missing_docs)] -use std::ffi::CStr; -use std::ffi::c_char; use std::ffi::c_void; use std::ptr; use std::slice; @@ -31,6 +29,7 @@ use crate::error::vx_error; use crate::scan::vx_estimate; use crate::scan::vx_estimate_type; use crate::session::vx_session; +use crate::string::vx_view; crate::arc_dyn_wrapper!( /// A reference to one or more (possibly remote) paths. @@ -42,13 +41,23 @@ crate::arc_dyn_wrapper!( /// Options for creating a data source. #[repr(C)] -#[cfg_attr(test, derive(Default))] pub struct vx_data_source_options { - /// Required: paths to files, tables, or layout trees. - /// May be a glob pattern like "*.vortex". - /// If you want to include multiple paths, concat them with a comma: - /// "file1.vortex,../file2.vortex". - pub paths: *const c_char, + /// Required: paths to files, tables, or layout trees. Each entry may be a + /// glob pattern like "*.vortex". Must point to an array of size + /// "paths_len". paths bytes are copied. + pub paths: *const vx_view, + /// Number of entries in `paths`. + pub paths_len: usize, +} + +#[cfg(test)] +impl Default for vx_data_source_options { + fn default() -> Self { + vx_data_source_options { + paths: ptr::null(), + paths_len: 0, + } + } } #[cfg(vortex_asan)] @@ -68,13 +77,12 @@ unsafe fn data_source_new( let opts = unsafe { &*opts }; vortex_ensure!(!opts.paths.is_null()); + vortex_ensure!(opts.paths_len > 0, "empty paths"); - let glob = unsafe { CStr::from_ptr(opts.paths) } - .to_string_lossy() - .into_owned(); + let paths = unsafe { slice::from_raw_parts(opts.paths, opts.paths_len) }; let mut data_source = MultiFileDataSource::new(session.clone()); - for glob in glob.split(',') { - data_source = data_source.with_glob(glob, None); + for path in paths { + data_source = data_source.with_glob(unsafe { path.as_str() }?, None); } let data_source = RUNTIME.block_on(async { @@ -95,8 +103,7 @@ unsafe fn data_source_new( /// Create a data source. /// The first matched file is opened eagerly. to read the schema. All other I/O -/// is deferred until a scan is requested. The returned pointer is owned by the -/// caller and must be freed with vx_data_source_free. +/// is deferred until a scan is requested. /// /// On error, returns NULL and sets "err". #[unsafe(no_mangle)] @@ -116,9 +123,6 @@ pub unsafe extern "C-unwind" fn vx_data_source_new( /// The bytes are borrowed, not copied: the caller must keep "buffer" alive and /// unmodified until the data source is freed. /// -/// The returned pointer is owned by the caller and must be freed with -/// vx_data_source_free. -/// /// On error, returns NULL and sets "err". #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_data_source_new_buffer( @@ -147,11 +151,10 @@ pub unsafe extern "C-unwind" fn vx_data_source_new_buffer( }) } -/// Return the schema of the data source as a non-owned dtype. -/// The returned pointer is valid as long as "ds" is alive. Do not free it. +/// Return data source's dtype #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_data_source_dtype(ds: *const vx_data_source) -> *const vx_dtype { - vx_dtype::new_ref(vx_data_source::as_ref(ds).dtype()) + vx_dtype::new(Arc::new(vx_data_source::as_ref(ds).dtype().clone())) } /// Write data source's row count estimate into "row_count". @@ -182,7 +185,6 @@ pub unsafe extern "C-unwind" fn vx_data_source_get_row_count( #[cfg(not(windows))] #[cfg(test)] mod tests { - use std::ffi::CString; use std::ffi::c_void; use std::fs::read; use std::ptr; @@ -194,10 +196,12 @@ mod tests { use crate::data_source::vx_data_source_new_buffer; use crate::data_source::vx_data_source_options; use crate::dtype::vx_dtype; + use crate::dtype::vx_dtype_free; use crate::scan::vx_estimate; use crate::scan::vx_estimate_type; use crate::session::vx_session_free; use crate::session::vx_session_new; + use crate::string::vx_view; use crate::tests::SAMPLE_ROWS; use crate::tests::assert_error; use crate::tests::assert_no_error; @@ -223,12 +227,15 @@ mod tests { assert_error(error); assert!(ds.is_null()); - opts.paths = c"test.vortex".as_ptr(); + let missing = vx_view::from_str("test.vortex"); + opts.paths = &raw const missing; + opts.paths_len = 1; let ds = vx_data_source_new(session, &raw const opts, &raw mut error); assert_error(error); assert!(ds.is_null()); - opts.paths = c"definitely-missing-dir/*.vortex".as_ptr(); + let missing_glob = vx_view::from_str("definitely-missing-dir/*.vortex"); + opts.paths = &raw const missing_glob; let ds = vx_data_source_new(session, &raw const opts, &raw mut error); assert_error(error); assert!(ds.is_null()); @@ -244,9 +251,10 @@ mod tests { let session = vx_session_new(); let (sample, struct_array) = write_sample(session); - let path = CString::new(sample.path().to_str().unwrap()).unwrap(); + let path = vx_view::from_str(sample.path().to_str().unwrap()); let opts = vx_data_source_options { - paths: path.as_ptr(), + paths: &raw const path, + paths_len: 1, }; let mut error = ptr::null_mut(); @@ -254,7 +262,8 @@ mod tests { assert_no_error(error); assert!(!ds.is_null()); - let dtype = vx_dtype::as_ref(vx_data_source_dtype(ds)); + let ffi_dtype = vx_data_source_dtype(ds); + let dtype = vx_dtype::as_ref(ffi_dtype); assert_eq!(dtype, struct_array.dtype()); let mut row_count = vx_estimate::default(); @@ -262,6 +271,42 @@ mod tests { assert_eq!(row_count.r#type, vx_estimate_type::VX_ESTIMATE_EXACT); assert_eq!(row_count.estimate, SAMPLE_ROWS as u64); + vx_dtype_free(ffi_dtype); + vx_data_source_free(ds); + vx_session_free(session); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_many_paths() { + let dir = tempfile::tempdir().unwrap(); + + unsafe { + let session = vx_session_new(); + let (sample, _) = write_sample(session); + + let comma_path = dir.path().join("with,comma.vortex"); + std::fs::copy(sample.path(), &comma_path).unwrap(); + + let paths = [ + vx_view::from_str(sample.path().to_str().unwrap()), + vx_view::from_str(comma_path.to_str().unwrap()), + ]; + let opts = vx_data_source_options { + paths: paths.as_ptr(), + paths_len: paths.len(), + }; + + let mut error = ptr::null_mut(); + let ds = vx_data_source_new(session, &raw const opts, &raw mut error); + assert_no_error(error); + assert!(!ds.is_null()); + + let mut row_count = vx_estimate::default(); + vx_data_source_get_row_count(ds, &raw mut row_count); + assert_eq!(row_count.estimate, 2 * SAMPLE_ROWS as u64); + vx_data_source_free(ds); vx_session_free(session); } @@ -289,7 +334,8 @@ mod tests { assert_no_error(error); assert!(!ds.is_null()); - let dtype = vx_dtype::as_ref(vx_data_source_dtype(ds)); + let ffi_dtype = vx_data_source_dtype(ds); + let dtype = vx_dtype::as_ref(ffi_dtype); assert_eq!(dtype, struct_array.dtype()); let mut row_count = vx_estimate::default(); @@ -297,6 +343,7 @@ mod tests { assert_eq!(row_count.r#type, vx_estimate_type::VX_ESTIMATE_EXACT); assert_eq!(row_count.estimate, SAMPLE_ROWS as u64); + vx_dtype_free(ffi_dtype); vx_data_source_free(ds); vx_session_free(session); } diff --git a/vortex-ffi/src/dtype.rs b/vortex-ffi/src/dtype.rs index c74fa71b6ed..d38a8c62ec8 100644 --- a/vortex-ffi/src/dtype.rs +++ b/vortex-ffi/src/dtype.rs @@ -9,7 +9,6 @@ use arrow_array::ffi::FFI_ArrowSchema; use arrow_schema::Schema; use vortex::dtype::DType; use vortex::dtype::DecimalDType; -use vortex::dtype::arrow::FromArrowType; use vortex::error::VortexExpect; use vortex::error::vortex_ensure; use vortex::error::vortex_panic; @@ -17,13 +16,15 @@ use vortex::extension::datetime::AnyTemporal; use vortex::extension::datetime::Date; use vortex::extension::datetime::Time; use vortex::extension::datetime::Timestamp; +use vortex_arrow::FromArrowType; +use vortex_arrow::ToArrowType; use crate::arc_wrapper; use crate::error::try_or; use crate::error::try_or_default; use crate::error::vx_error; use crate::ptype::vx_ptype; -use crate::string::vx_string; +use crate::string::vx_view; use crate::struct_fields::vx_struct_fields; arc_wrapper!( @@ -208,10 +209,9 @@ pub unsafe extern "C-unwind" fn vx_dtype_decimal_scale(dtype: *const vx_dtype) - .scale() } -/// Return a borrowed reference to the [`vx_struct_fields`] of a struct. -/// -/// The returned pointer is valid as long as the struct dtype is valid. -/// Do NOT free the returned pointer - it shares the lifetime of the struct dtype. +/// If "dtype" is DTYPE_STRUCT, return owned vx_struct_fields for this struct, +/// return NULL otherwise. Returned vx_struct_fields must be released with +/// vx_struct_fields_free. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_dtype_struct_dtype( dtype: *const vx_dtype, @@ -219,34 +219,29 @@ pub unsafe extern "C-unwind" fn vx_dtype_struct_dtype( let Some(struct_dtype) = vx_dtype::as_ref(dtype).as_struct_fields_opt() else { return ptr::null(); }; - vx_struct_fields::new_ref(struct_dtype) + vx_struct_fields::new(struct_dtype.clone()) } -/// Returns the element type of a list. -/// -/// The returned pointer is valid as long as the list dtype is valid. -/// Do NOT free the returned dtype pointer - it shares the lifetime of the list dtype. +/// If "dtype" is DTYPE_LIST, return its owned element dtype, return NULL +/// otherwise. Returned dtype must be released with vx_dtype_free. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_dtype_list_element(dtype: *const vx_dtype) -> *const vx_dtype { let Some(element_dtype) = vx_dtype::as_ref(dtype).as_list_element_opt() else { return ptr::null(); }; - vx_dtype::new_ref(element_dtype) + vx_dtype::new(Arc::clone(element_dtype)) } -/// Returns the element type of a fixed-size list. -/// -/// The returned pointer is valid as long as the fixed-size list dtype is valid. -/// Do NOT free the returned dtype pointer - it shares the lifetime of the fixed-size list dtype. +/// If "dtype" is DTYPE_FIXED_SIZE_LIST, return its owned element dtype, return +/// NULL otherwise. Returned dtype must be released with vx_dtype_free. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_dtype_fixed_size_list_element( dtype: *const vx_dtype, ) -> *const vx_dtype { - // TODO(joe): propagate this error up instead of expecting - let element_dtype = vx_dtype::as_ref(dtype) - .as_fixed_size_list_element_opt() - .vortex_expect("not a fixed-size list dtype"); - vx_dtype::new_ref(element_dtype) + let Some(element_dtype) = vx_dtype::as_ref(dtype).as_fixed_size_list_element_opt() else { + return ptr::null(); + }; + vx_dtype::new(Arc::clone(element_dtype)) } /// Returns the size of a fixed-size list. @@ -312,9 +307,12 @@ pub unsafe extern "C-unwind" fn vx_dtype_time_unit(dtype: *const DType) -> u8 { opts.time_unit().into() } -/// Returns the time zone, assuming the type is time. Caller is responsible for freeing the returned pointer. +/// Return time zone assuming "dtype" is time. +/// Returns {NULL, 0} when timestamp has no time zone. +/// +/// Returned view is valid as long as "dtype" is valid. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_dtype_time_zone(dtype: *const DType) -> *const vx_string { +pub unsafe extern "C-unwind" fn vx_dtype_time_zone(dtype: *const DType) -> vx_view { // TODO(joe): propagate this error up instead of expecting let dtype = unsafe { dtype.as_ref() }.vortex_expect("dtype null"); @@ -328,8 +326,8 @@ pub unsafe extern "C-unwind" fn vx_dtype_time_zone(dtype: *const DType) -> *cons }; match opts.tz.as_ref() { - Some(zone) => vx_string::new(Arc::clone(zone)), - None => ptr::null(), + Some(zone) => vx_view::from_str(zone), + None => vx_view::null(), } } @@ -376,7 +374,6 @@ pub unsafe extern "C-unwind" fn vx_dtype_from_arrow_schema( #[cfg(test)] #[expect(clippy::cast_possible_truncation)] mod tests { - use std::slice; use vortex::array::ArrayRef; use vortex::array::IntoArray; @@ -397,9 +394,7 @@ mod tests { use crate::dtype::vx_dtype_new_utf8; use crate::dtype::vx_dtype_variant; use crate::ptype::vx_ptype; - use crate::string::vx_string; - use crate::string::vx_string_len; - use crate::string::vx_string_ptr; + use crate::string::vx_view; use crate::struct_fields::vx_struct_fields_builder_add_field; use crate::struct_fields::vx_struct_fields_builder_finalize; use crate::struct_fields::vx_struct_fields_builder_new; @@ -431,23 +426,28 @@ mod tests { fn test_struct() { unsafe { let builder = vx_struct_fields_builder_new(); + let mut error = ptr::null_mut(); vx_struct_fields_builder_add_field( builder, - vx_string::new("name".into()), + vx_view::from_str("name"), vx_dtype_new_utf8(false), + &raw mut error, ); + assert!(error.is_null()); vx_struct_fields_builder_add_field( builder, - vx_string::new("age".into()), + vx_view::from_str("age"), vx_dtype_new_primitive(vx_ptype::PTYPE_U8, true), + &raw mut error, ); + assert!(error.is_null()); let person = vx_struct_fields_builder_finalize(builder); assert_eq!(vx_struct_fields_nfields(person), 2); let name = vx_struct_fields_field_name(person, 0); - assert_eq!(vx_string::as_str(name), "name"); + assert_eq!(name.as_str().unwrap(), "name"); let age = vx_struct_fields_field_name(person, 1); - assert_eq!(vx_string::as_str(age), "age"); + assert_eq!(age.as_str().unwrap(), "age"); let dtype0 = vx_struct_fields_field_dtype(person, 0); let dtype1 = vx_struct_fields_field_dtype(person, 1); @@ -457,13 +457,9 @@ mod tests { vx_dtype_variant::DTYPE_PRIMITIVE ); - // Field names are now borrowed references - do not free them - - // Free field dtypes (owned references) vx_dtype_free(dtype0); vx_dtype_free(dtype1); - // Free struct fields vx_struct_fields_free(person); } } @@ -517,6 +513,7 @@ mod tests { vx_dtype_variant::DTYPE_PRIMITIVE ); assert_eq!(vx_dtype_primitive_ptype(element), vx_ptype::PTYPE_I32); + vx_dtype_free(element); vx_dtype_free(list_dtype); } @@ -534,15 +531,14 @@ mod tests { ); assert!(vx_dtype_is_nullable(fsl_dtype)); - // Test element accessor let element = vx_dtype_fixed_size_list_element(fsl_dtype); assert_eq!( vx_dtype_get_variant(element), vx_dtype_variant::DTYPE_PRIMITIVE ); assert_eq!(vx_dtype_primitive_ptype(element), vx_ptype::PTYPE_F64); + vx_dtype_free(element); - // Test size accessor let size = vx_dtype_fixed_size_list_size(fsl_dtype); assert_eq!(size, 3); @@ -565,6 +561,7 @@ mod tests { let element = vx_dtype_fixed_size_list_element(fsl_dtype); assert_eq!(vx_dtype_get_variant(element), vx_dtype_variant::DTYPE_UTF8); assert!(vx_dtype_is_nullable(element)); + vx_dtype_free(element); let size = vx_dtype_fixed_size_list_size(fsl_dtype); assert_eq!(size, 10); @@ -607,6 +604,8 @@ mod tests { ); assert_eq!(vx_dtype_primitive_ptype(innermost), vx_ptype::PTYPE_I32); + vx_dtype_free(innermost); + vx_dtype_free(inner); vx_dtype_free(outer_fsl); } } @@ -703,8 +702,9 @@ mod tests { let n_fields = unsafe { vx_struct_fields_nfields(struct_fields_ptr) }; assert_eq!(n_fields, 2); - // Cleanup in reverse order - this is the safest order unsafe { + vx_struct_fields_free(struct_fields_ptr); + vx_dtype_free(dtype_ptr); vx_array_free(vx_arr); } } @@ -720,18 +720,13 @@ mod tests { let struct_fields_ptr = unsafe { vx_dtype_struct_dtype(dtype_ptr) }; // Test field name access - let field_name_ptr = unsafe { vx_struct_fields_field_name(struct_fields_ptr, 0) }; - assert!(!field_name_ptr.is_null()); - - let name_len = unsafe { vx_string_len(field_name_ptr) }; - let name_ptr = unsafe { vx_string_ptr(field_name_ptr) }; - let name_slice = unsafe { slice::from_raw_parts(name_ptr.cast::(), name_len) }; - let name_str = str::from_utf8(name_slice).unwrap(); - assert_eq!(name_str, "nums"); + let field_name = unsafe { vx_struct_fields_field_name(struct_fields_ptr, 0) }; + assert!(!field_name.ptr.is_null()); + assert_eq!(unsafe { field_name.as_str() }.unwrap(), "nums"); - // Cleanup in careful order unsafe { - // Field name is now a borrowed reference - do not free it + vx_struct_fields_free(struct_fields_ptr); + vx_dtype_free(dtype_ptr); vx_array_free(vx_arr); } } @@ -750,23 +745,16 @@ mod tests { // Test both field names for i in 0..n_fields { - let field_name_ptr = - unsafe { vx_struct_fields_field_name(struct_fields_ptr, i as usize) }; - assert!(!field_name_ptr.is_null()); - - let name_len = unsafe { vx_string_len(field_name_ptr) }; - let name_ptr = unsafe { vx_string_ptr(field_name_ptr) }; - let name_slice = unsafe { slice::from_raw_parts(name_ptr.cast::(), name_len) }; - let name_str = str::from_utf8(name_slice).unwrap(); + let field_name = unsafe { vx_struct_fields_field_name(struct_fields_ptr, i as usize) }; + assert!(!field_name.ptr.is_null()); let expected_name = if i == 0 { "nums" } else { "floats" }; - assert_eq!(name_str, expected_name); - - // Field name is now a borrowed reference - do not free it + assert_eq!(unsafe { field_name.as_str() }.unwrap(), expected_name); } - // Cleanup unsafe { + vx_struct_fields_free(struct_fields_ptr); + vx_dtype_free(dtype_ptr); vx_array_free(vx_arr); } } @@ -800,6 +788,7 @@ mod tests { let f1 = vx_struct_fields_field_dtype(fields, 1); assert_eq!(vx_dtype_get_variant(f1), vx_dtype_variant::DTYPE_UTF8); assert!(vx_dtype_is_nullable(f1)); + vx_struct_fields_free(fields); vx_dtype_free(f1); vx_dtype_free(dtype); } diff --git a/vortex-ffi/src/error.rs b/vortex-ffi/src/error.rs index 3ee02c99e21..179c0c85f79 100644 --- a/vortex-ffi/src/error.rs +++ b/vortex-ffi/src/error.rs @@ -7,35 +7,84 @@ use std::panic::catch_unwind; use std::ptr; use std::sync::Arc; +use vortex::error::VortexError; use vortex::error::VortexResult; use crate::box_wrapper; -use crate::string::vx_string; +use crate::string::vx_view; + +/// Error category for vx_error. +#[repr(C)] +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +#[expect(non_camel_case_types)] +pub enum vx_error_code { + /// All other errors + VX_ERROR_CODE_OTHER = 0, + /// Index out of bounds + VX_ERROR_CODE_OUT_OF_BOUNDS = 1, + /// Compute kernel execute error + VX_ERROR_CODE_COMPUTE = 2, + /// An invalid argument was provided. + VX_ERROR_CODE_INVALID_ARGUMENT = 3, + /// Serialization/deserialization error + VX_ERROR_CODE_SERIALIZATION = 4, + /// Unimplemented function + VX_ERROR_CODE_NOT_IMPLEMENTED = 5, + /// Type mismatch + VX_ERROR_CODE_MISMATCHED_TYPES = 6, + /// Assertion failed + VX_ERROR_CODE_ASSERTION_FAILED = 7, + /// IO error + VX_ERROR_CODE_IO = 8, + /// Panic inside FFI + VX_ERROR_CODE_PANIC = 9, +} + +fn error_code(error: &VortexError) -> vx_error_code { + match error { + VortexError::OutOfBounds(..) => vx_error_code::VX_ERROR_CODE_OUT_OF_BOUNDS, + VortexError::Compute(..) => vx_error_code::VX_ERROR_CODE_COMPUTE, + VortexError::InvalidArgument(..) => vx_error_code::VX_ERROR_CODE_INVALID_ARGUMENT, + VortexError::Serde(..) => vx_error_code::VX_ERROR_CODE_SERIALIZATION, + VortexError::NotImplemented(..) => vx_error_code::VX_ERROR_CODE_NOT_IMPLEMENTED, + VortexError::MismatchedTypes(..) => vx_error_code::VX_ERROR_CODE_MISMATCHED_TYPES, + VortexError::AssertionFailed(..) => vx_error_code::VX_ERROR_CODE_ASSERTION_FAILED, + VortexError::Io(..) => vx_error_code::VX_ERROR_CODE_IO, + VortexError::Context(_, inner) => error_code(inner), + VortexError::Shared(inner) => error_code(inner), + _ => vx_error_code::VX_ERROR_CODE_OTHER, + } +} -pub(crate) struct VortexError { +pub(crate) struct VortexFFIError { message: Arc, + code: vx_error_code, } box_wrapper!( /// The error structure populated by fallible Vortex C functions. - VortexError, + VortexFFIError, vx_error ); -/// Create an owned Vortex FFI error from a message. -pub(crate) fn vx_error_new(message: &str) -> *mut vx_error { - vx_error::new(VortexError { +fn vx_error_new_with_code(message: &str, code: vx_error_code) -> *mut vx_error { + vx_error::new(VortexFFIError { message: message.into(), + code, }) } /// Write an error message to `error` which has not been populated before. /// A null `error` pointer discards the message. pub(crate) fn write_error(error: *mut *mut vx_error, message: &str) { + write_error_with_code(error, message, vx_error_code::VX_ERROR_CODE_OTHER); +} + +fn write_error_with_code(error: *mut *mut vx_error, message: &str, code: vx_error_code) { if error.is_null() { return; } - unsafe { error.write(vx_error_new(message)) }; + unsafe { error.write(vx_error_new_with_code(message, code)) }; } /// Clear `*error_out` to null unless `error_out` itself is null. @@ -68,11 +117,15 @@ pub fn try_or_default( value } Ok(Err(err)) => { - write_error(error_out, &err.to_string()); + write_error_with_code(error_out, &err.to_string(), error_code(&err)); T::default() } Err(payload) => { - write_error(error_out, &panic_message(payload.as_ref())); + write_error_with_code( + error_out, + &panic_message(payload.as_ref()), + vx_error_code::VX_ERROR_CODE_PANIC, + ); T::default() } } @@ -93,23 +146,31 @@ pub fn try_or( value } Ok(Err(err)) => { - write_error(error_out, &err.to_string()); + write_error_with_code(error_out, &err.to_string(), error_code(&err)); error_value } Err(payload) => { - write_error(error_out, &panic_message(payload.as_ref())); + write_error_with_code( + error_out, + &panic_message(payload.as_ref()), + vx_error_code::VX_ERROR_CODE_PANIC, + ); error_value } } } -/// Returns the error message from the given Vortex error. -/// -/// The returned pointer is valid as long as the error is valid. -/// Do NOT free the returned string pointer - it shares the lifetime of the error. +/// Return error message for this error. +/// Returned view is valid while "error" is valid. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_error_message(error: *const vx_error) -> vx_view { + vx_view::from_str(&vx_error::as_ref(error).message) +} + +/// Return category code for "error". #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_error_get_message(error: *const vx_error) -> *const vx_string { - vx_string::new_ref(&vx_error::as_ref(error).message) +pub unsafe extern "C-unwind" fn vx_error_get_code(error: *const vx_error) -> vx_error_code { + vx_error::as_ref(error).code } #[cfg(test)] @@ -155,11 +216,35 @@ mod tests { assert_eq!(try_or(&raw mut error, -1, || panic!("boom")), -1); assert!(!error.is_null()); - let message = unsafe { vx_error_get_message(error) }; + let message = unsafe { vx_error_message(error) }; assert_eq!( - vx_string::as_ref(message).as_ref(), + unsafe { message.as_str() }.unwrap(), "panic in Vortex FFI function: boom" ); unsafe { vx_error_free(error) }; } + + #[test] + fn test_error_codes() { + let mut error: *mut vx_error = ptr::null_mut(); + + assert_eq!( + try_or(&raw mut error, -1, || Err::(vortex_err!( + OutOfBounds: 5, 0, 3 + ))), + -1 + ); + assert_eq!( + unsafe { vx_error_get_code(error) }, + vx_error_code::VX_ERROR_CODE_OUT_OF_BOUNDS + ); + unsafe { vx_error_free(error) }; + + assert_eq!(try_or(&raw mut error, -1, || panic!("panic")), -1); + assert_eq!( + unsafe { vx_error_get_code(error) }, + vx_error_code::VX_ERROR_CODE_PANIC + ); + unsafe { vx_error_free(error) }; + } } diff --git a/vortex-ffi/src/expression.rs b/vortex-ffi/src/expression.rs index 269e1e5d941..e724ea08ce5 100644 --- a/vortex-ffi/src/expression.rs +++ b/vortex-ffi/src/expression.rs @@ -2,8 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #![expect(non_camel_case_types)] -use std::ffi::CStr; -use std::ffi::c_char; use std::ptr; use std::slice; use std::sync::Arc; @@ -28,6 +26,7 @@ use vortex::scalar_fn::fns::operators::Operator; use crate::error::try_or; use crate::error::vx_error; use crate::scalar::vx_scalar; +use crate::string::vx_view; use crate::to_field_names; // Expressions are Arc'ed inside @@ -38,8 +37,6 @@ crate::box_wrapper!( /// data. Each expression consists of an encoding (vtable), heap-allocated /// metadata, and child expressions. /// - /// Unless stated explicitly, all expressions returned are owned and must - /// be freed by the caller. /// Unless stated explicitly, if an operation on const vx_expression* is /// passed NULL, NULL is returned. /// Operations on expressions don't take ownership of input values, and so @@ -66,6 +63,17 @@ pub unsafe extern "C" fn vx_expression_root() -> *mut vx_expression { vx_expression::new(root()) } +/// Reference-clone a vx_expression +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_expression_clone( + ptr: *const vx_expression, +) -> *mut vx_expression { + if ptr.is_null() { + return ptr::null_mut(); + } + vx_expression::new(vx_expression::as_ref(ptr).clone()) +} + /// Create a literal expression from a scalar. /// /// Literal expressions are useful for constants in expression trees, especially scan @@ -121,7 +129,7 @@ pub unsafe extern "C-unwind" fn vx_expression_literal( /// vx_expression_free(root); #[unsafe(no_mangle)] pub unsafe extern "C" fn vx_expression_select( - names: *const *const c_char, + names: *const vx_view, len: usize, child: *const vx_expression, ) -> *mut vx_expression { @@ -263,9 +271,9 @@ pub unsafe extern "C" fn vx_expression_binary( /// /// Returns the logical negation of the input boolean expression. #[unsafe(no_mangle)] -pub unsafe extern "C" fn vx_expression_not(child: *const vx_expression) -> *const vx_expression { +pub unsafe extern "C" fn vx_expression_not(child: *const vx_expression) -> *mut vx_expression { if child.is_null() { - return child; + return ptr::null_mut(); } vx_expression::new(not(vx_expression::as_ref(child).clone())) } @@ -289,22 +297,19 @@ pub unsafe extern "C" fn vx_expression_is_null(child: *const vx_expression) -> * /// /// Example: if child is Struct { name=u8, age=u16 } and we do /// vx_expression_get_item("name", child), output type will be DTYPE_U8 +/// +/// "item" is copied. Returns NULL if "child" is NULL or "item" is not valid +/// UTF-8. #[unsafe(no_mangle)] pub unsafe extern "C" fn vx_expression_get_item( - item: *const c_char, + item: vx_view, child: *const vx_expression, ) -> *mut vx_expression { if child.is_null() { return ptr::null_mut(); } - if item.is_null() { + let Ok(item) = (unsafe { item.as_str() }) else { return ptr::null_mut(); - } - #[expect(clippy::expect_used)] - let item = unsafe { - CStr::from_ptr(item) - .to_str() - .expect("converting pointer to str") }; let item: Arc = Arc::from(item); let item: FieldName = item.into(); @@ -367,6 +372,7 @@ mod tests { use crate::expression::vx_expression_select; use crate::scalar::vx_scalar_free; use crate::scalar::vx_scalar_new_i32; + use crate::string::vx_view; #[test] #[cfg_attr(miri, ignore)] @@ -397,7 +403,7 @@ mod tests { let (array, names_array, ages_array) = struct_array(); unsafe { let root = vx_expression_root(); - let column = vx_expression_get_item(c"age".as_ptr(), root); + let column = vx_expression_get_item(vx_view::from_str("age"), root); assert_ne!(column, ptr::null_mut()); let array = vx_array::new(Arc::new(array.into_array())); @@ -419,7 +425,7 @@ mod tests { vx_expression_free(column); - let column = vx_expression_get_item(c"ololo".as_ptr(), root); + let column = vx_expression_get_item(vx_view::from_str("ololo"), root); assert_ne!(column, ptr::null_mut()); let applied_array = vx_array_apply(array, column, &raw mut error); @@ -496,7 +502,7 @@ mod tests { let array = vx_array::new(Arc::new(array.into_array())); - let columns = [c"name".as_ptr(), c"age".as_ptr()]; + let columns = [vx_view::from_str("name"), vx_view::from_str("age")]; let column = vx_expression_select(columns.as_ptr(), 2, root); assert_ne!(column, ptr::null_mut()); @@ -512,7 +518,7 @@ mod tests { vx_array_free(applied_array); vx_expression_free(column); - let columns = [c"age".as_ptr(), c"ololo".as_ptr()]; + let columns = [vx_view::from_str("age"), vx_view::from_str("ololo")]; let column = vx_expression_select(columns.as_ptr(), 2, root); let applied_array = vx_array_apply(array, column, &raw mut error); assert!(applied_array.is_null()); @@ -540,9 +546,9 @@ mod tests { let array = vx_array::new(Arc::new(array.unwrap().into_array())); let root = vx_expression_root(); - let expression_col1 = vx_expression_get_item(c"col1".as_ptr(), root); - let expression_col2 = vx_expression_get_item(c"col2".as_ptr(), root); - let expression_col3 = vx_expression_get_item(c"col3".as_ptr(), root); + let expression_col1 = vx_expression_get_item(vx_view::from_str("col1"), root); + let expression_col2 = vx_expression_get_item(vx_view::from_str("col2"), root); + let expression_col3 = vx_expression_get_item(vx_view::from_str("col3"), root); let expression_12 = vx_expression_binary( vx_binary_operator::VX_OPERATOR_EQ, expression_col1, diff --git a/vortex-ffi/src/file.rs b/vortex-ffi/src/file.rs deleted file mode 100644 index ff228a6c5c9..00000000000 --- a/vortex-ffi/src/file.rs +++ /dev/null @@ -1,50 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::ffi::CStr; -use std::ffi::c_char; - -use vortex::error::vortex_err; -use vortex::file::VortexFile; -use vortex::file::WriteOptionsSessionExt; -use vortex::io::runtime::BlockingRuntime; - -use crate::RUNTIME; -use crate::arc_wrapper; -use crate::array::vx_array; -use crate::error::try_or_default; -use crate::error::vx_error; -use crate::session::vx_session; - -arc_wrapper!( - /// A handle to a Vortex file encapsulating the footer and logic for instantiating a reader. - VortexFile, - vx_file -); - -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_file_write_array( - session: *const vx_session, - path: *const c_char, - array: *const vx_array, - error_out: *mut *mut vx_error, -) { - let session = vx_session::as_ref(session); - let options = session.write_options(); - let array = vx_array::as_ref(array); - try_or_default(error_out, || { - let path = unsafe { CStr::from_ptr(path) } - .to_str() - .map_err(|e| vortex_err!("invalid utf-8: {e}"))?; - - RUNTIME.block_on(async move { - options - .write( - &mut async_fs::File::create(path).await?, - array.to_array_stream(), - ) - .await?; - Ok(()) - }) - }); -} diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index 43e108e6c28..151359096d4 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -7,12 +7,10 @@ mod array; mod array_iterator; -mod binary; mod data_source; mod dtype; mod error; mod expression; -mod file; mod log; mod macros; mod ptype; @@ -24,13 +22,12 @@ mod string; mod struct_array; mod struct_fields; -use std::ffi::CStr; -use std::ffi::c_char; use std::sync::Arc; use std::sync::LazyLock; pub use array::vx_array; pub use array::vx_array_ref; +pub use dtype::vx_dtype; pub use error::try_or; pub use error::vx_error; pub use error::vx_error_free; @@ -41,10 +38,12 @@ pub use session::vx_session; pub use session::vx_session_free; pub use session::vx_session_new_with; pub use session::vx_session_ref; +pub use sink::vx_array_sink; +pub use sink::vx_array_sink_open_file_with_strategy; +pub use string::vx_view; use vortex::dtype::FieldName; -use vortex::error::VortexExpect; use vortex::error::VortexResult; -use vortex::error::vortex_err; +use vortex::error::vortex_ensure; use vortex::io::runtime::current::CurrentThreadRuntime; #[cfg(all(feature = "mimalloc", not(miri)))] @@ -65,31 +64,25 @@ pub fn ffi_runtime() -> &'static CurrentThreadRuntime { &RUNTIME } -/// SAFETY: name must be a non-NULL pointer -pub(crate) unsafe fn to_field_name(name: *const c_char) -> VortexResult { - let name = unsafe { CStr::from_ptr(name) } - .to_str() - .map_err(|e| vortex_err!("{e}"))?; - let name: Arc = Arc::from(name); +/// SAFETY: name must be a vx_view with non-NULL pointer +pub(crate) unsafe fn to_field_name(name: vx_view) -> VortexResult { + let name: Arc = Arc::from(unsafe { name.as_str() }?); Ok(name.into()) } /// SAFETY: names must be a non-NULL pointer valid for reads up to len. pub(crate) unsafe fn to_field_names( - names: *const *const c_char, + names: *const vx_view, len: usize, ) -> VortexResult> { + vortex_ensure!(!names.is_null() || len == 0, "null names pointer"); (0..len) - .map(|i| unsafe { - let name = *names.offset(i.try_into().vortex_expect("pointer offset overflow")); - to_field_name(name) - }) + .map(|i| unsafe { to_field_name(*names.add(i)) }) .collect() } #[cfg(test)] mod tests { - use std::ffi::CString; use std::ptr; use std::sync::Arc; @@ -107,12 +100,12 @@ mod tests { use crate::dtype::vx_dtype_free; use crate::error::vx_error; use crate::error::vx_error_free; - use crate::error::vx_error_get_message; + use crate::error::vx_error_message; use crate::session::vx_session; use crate::sink::vx_array_sink_close; use crate::sink::vx_array_sink_open_file; use crate::sink::vx_array_sink_push; - use crate::string::vx_string; + use crate::string::vx_view; /// Panic if error is NULL. Free the error if it's not pub(crate) fn assert_error(error: *mut vx_error) { @@ -125,7 +118,7 @@ mod tests { if !error.is_null() { let message; unsafe { - message = vx_string::as_str(vx_error_get_message(error)).to_owned(); + message = vx_error_message(error).as_str().unwrap().to_owned(); vx_error_free(error); } panic!("{message}"); @@ -167,14 +160,13 @@ mod tests { .unwrap(); let file = NamedTempFile::new().unwrap(); - let path = CString::new(file.path().to_str().unwrap()).unwrap(); + let path = vx_view::from_str(file.path().to_str().unwrap()); let dtype = struct_array.dtype(); unsafe { let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype.clone())); let mut error = ptr::null_mut(); - let sink = - vx_array_sink_open_file(session, path.as_ptr(), vx_dtype_ptr, &raw mut error); + let sink = vx_array_sink_open_file(session, path, vx_dtype_ptr, &raw mut error); let array = vx_array::new(Arc::new(struct_array.clone().into_array())); vx_array_sink_push(sink, array, &raw mut error); vx_array_sink_close(sink, &raw mut error); diff --git a/vortex-ffi/src/macros.rs b/vortex-ffi/src/macros.rs index d4e58288b30..09f751c2a9e 100644 --- a/vortex-ffi/src/macros.rs +++ b/vortex-ffi/src/macros.rs @@ -21,7 +21,7 @@ //! Each macro provides a `free` function, and the `Arc` variants also provide a `clone` function. //! //! Converting between the raw pointer and the wrapped type is done using the generated `new`, -//! `new_ref`, `as_ref`, `as_mut`, `into_box`, and `into_arc` methods, which internally check for +//! `as_ref`, `as_mut`, `into_box`, and `into_arc` methods, which internally check for //! null pointers. //! //! ## Internals @@ -59,11 +59,6 @@ macro_rules! arc_dyn_wrapper { Box::into_raw(Box::new($ffi_ident(obj))).cast_const() } - /// Wrap a borrowed object into a raw pointer. - pub(crate) fn new_ref(obj: &std::sync::Arc<$T>) -> *const $ffi_ident { - obj as *const std::sync::Arc<$T> as *const $ffi_ident - } - /// Extract a borrowed reference from a const pointer. pub(crate) fn as_ref<'a>(ptr: *const $ffi_ident) -> &'a std::sync::Arc<$T> { use vortex::error::VortexExpect; @@ -82,8 +77,7 @@ macro_rules! arc_dyn_wrapper { } } - #[doc = r" Clone a borrowed [`" $ffi_ident "`], returning an owned [`" $ffi_ident "`].\n\n"] - #[doc = r" Must be released with [`" $ffi_ident "_free`]."] + #[doc = r" Increase reference count on " $ffi_ident] #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn [<$ffi_ident _clone>](ptr: *const $ffi_ident) -> *const $ffi_ident { if ptr.is_null() { @@ -93,13 +87,12 @@ macro_rules! arc_dyn_wrapper { $ffi_ident::new($ffi_ident::as_ref(ptr).clone()) } - #[doc = r" Free an owned [`" $ffi_ident "`] object."] + #[doc = r" Decrease reference count on " $ffi_ident " or free if there are no other references"] #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn [<$ffi_ident _free>](ptr: *const $ffi_ident) { - if ptr.is_null() { - vortex::error::vortex_panic!("null pointer"); + if !ptr.is_null() { + drop($ffi_ident::into_arc(ptr)) } - drop($ffi_ident::into_arc(ptr)) } } }; @@ -121,11 +114,6 @@ macro_rules! arc_wrapper { std::sync::Arc::into_raw(obj).cast::<$ffi_ident>() } - /// Wrap a borrowed object into a raw pointer. - pub(crate) fn new_ref(obj: &$T) -> *const $ffi_ident { - obj as *const $T as *const $ffi_ident - } - /// Extract a borrowed reference from a const pointer. pub(crate) fn as_ref(ptr: *const $ffi_ident) -> &'static $T { use vortex::error::VortexExpect; @@ -144,8 +132,7 @@ macro_rules! arc_wrapper { } } - #[doc = r" Clone a borrowed [`" $ffi_ident "`], returning an owned [`" $ffi_ident "`].\n\n"] - #[doc = r" Must be released with [`" $ffi_ident "_free`]."] + #[doc = r" Increase reference count on " $ffi_ident] #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn [<$ffi_ident _clone>](ptr: *const $ffi_ident) -> *const $ffi_ident { if ptr.is_null() { @@ -155,13 +142,12 @@ macro_rules! arc_wrapper { ptr } - #[doc = r" Free an owned [`" $ffi_ident "`] object."] + #[doc = r" Decrease reference count on " $ffi_ident " or free if there are no other references"] #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn [<$ffi_ident _free>](ptr: *const $ffi_ident) { - if ptr.is_null() { - vortex::error::vortex_panic!("null pointer"); + if !ptr.is_null() { + unsafe { std::sync::Arc::decrement_strong_count(ptr) }; } - unsafe { std::sync::Arc::decrement_strong_count(ptr) }; } } }; @@ -186,21 +172,6 @@ macro_rules! box_dyn_wrapper { Box::into_raw(Box::new($ffi_ident(obj))) } - /// Wrap a borrowed object into a raw pointer. - pub(crate) fn new_ref(obj: &$T) -> *const $ffi_ident { - obj as *const $T as *const $ffi_ident - } - - /// Extract a borrowed reference from a const pointer. - pub(crate) fn as_ref<'a>(ptr: *const $ffi_ident) -> &'a $T { - use vortex::error::VortexExpect; - // TODO(joe): propagate this error up instead of expecting - unsafe { ptr.as_ref() } - .vortex_expect("null pointer") - .0 - .as_ref() - } - /// Extract a borrowed mutable reference from a mut pointer. pub(crate) fn as_mut<'a>(ptr: *mut $ffi_ident) -> &'a mut $T { use vortex::error::VortexExpect; @@ -220,13 +191,12 @@ macro_rules! box_dyn_wrapper { } } - #[doc = r" Free an owned [`" $ffi_ident "`] object."] + #[doc = r" Free a " $ffi_ident] #[unsafe(no_mangle)] - pub unsafe extern "C-unwind" fn [<$ffi_ident _free>](ptr: *mut $ffi_ident) { - if ptr.is_null() { - vortex::error::vortex_panic!("null pointer"); + pub unsafe extern "C-unwind" fn [<$ffi_ident _free>](ptr: *const $ffi_ident) { + if !ptr.is_null() { + drop($ffi_ident::into_box(ptr.cast_mut())) } - drop($ffi_ident::into_box(ptr)) } } }; @@ -253,11 +223,6 @@ macro_rules! box_wrapper { Box::into_raw(Box::new(obj)).cast::<$ffi_ident>() } - /// Wrap a borrowed object into a raw pointer. - pub(crate) fn new_ref(obj: &$T) -> *const $ffi_ident { - obj as *const $T as *const $ffi_ident - } - /// Extract a borrowed reference from a const pointer. pub(crate) fn as_ref<'a>(ptr: *const $ffi_ident) -> &'a $T { use vortex::error::VortexExpect; @@ -285,17 +250,16 @@ macro_rules! box_wrapper { } } - #[doc = r" Free an owned [`" $ffi_ident "`] object."] + #[doc = r" Free a " $ffi_ident] // These allows only matter once the destructor is re-exported (e.g. `vx_error_free`): // its `# Safety` lives at the C boundary, and its doc links a private wrapper type. #[allow(clippy::missing_safety_doc)] #[allow(rustdoc::private_intra_doc_links)] #[unsafe(no_mangle)] - pub unsafe extern "C-unwind" fn [<$ffi_ident _free>](ptr: *mut $ffi_ident) { - if ptr.is_null() { - vortex::error::vortex_panic!("null pointer"); + pub unsafe extern "C-unwind" fn [<$ffi_ident _free>](ptr: *const $ffi_ident) { + if !ptr.is_null() { + std::mem::drop(unsafe { Box::from_raw(ptr.cast::<$T>().cast_mut()) }) } - std::mem::drop(unsafe { Box::from_raw(ptr.cast::<$T>()) }) } } }; diff --git a/vortex-ffi/src/scalar.rs b/vortex-ffi/src/scalar.rs index 2471c2c4c9f..9f01f799050 100644 --- a/vortex-ffi/src/scalar.rs +++ b/vortex-ffi/src/scalar.rs @@ -3,7 +3,6 @@ //! FFI interface for working with Vortex scalar values. -use std::ffi::c_char; use std::ptr; use std::slice; use std::str; @@ -17,7 +16,6 @@ use vortex::dtype::i256; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_ensure; -use vortex::error::vortex_err; use vortex::scalar::DecimalValue; use vortex::scalar::Scalar; use vortex::scalar::ScalarValue; @@ -25,6 +23,7 @@ use vortex::scalar::ScalarValue; use crate::dtype::vx_dtype; use crate::error::try_or; use crate::error::vx_error; +use crate::string::vx_view; crate::box_wrapper!( /// A typed scalar value. @@ -37,10 +36,8 @@ crate::box_wrapper!( vx_scalar ); -/// Clone a borrowed scalar handle. -/// -/// The input scalar handle is not consumed. The returned scalar handle must be -/// released with vx_scalar_free. Returns NULL when given a NULL scalar handle. +/// Clone a scalar handle. +/// If scalar is NULL, returns NULL. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_scalar_clone(scalar: *const vx_scalar) -> *mut vx_scalar { if scalar.is_null() { @@ -49,17 +46,14 @@ pub unsafe extern "C-unwind" fn vx_scalar_clone(scalar: *const vx_scalar) -> *mu vx_scalar::new(vx_scalar::as_ref(scalar).clone()) } -/// Return the data type of a scalar. -/// -/// The returned data type handle borrows storage from the scalar handle, so its -/// lifetime is bound to the scalar handle. It MUST NOT be freed separately. -/// Returns NULL when given a NULL scalar handle. +/// Return scalar's dtype. +/// If scalar is NULL, returns NULL. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_scalar_dtype(scalar: *const vx_scalar) -> *const vx_dtype { if scalar.is_null() { return ptr::null(); } - vx_dtype::new_ref(vx_scalar::as_ref(scalar).dtype()) + vx_dtype::new(Arc::new(vx_scalar::as_ref(scalar).dtype().clone())) } /// Return whether the scalar is a typed null value. @@ -159,19 +153,16 @@ pub unsafe extern "C-unwind" fn vx_scalar_new_f16_bits( /// Create a UTF-8 scalar. /// -/// The byte range is copied into the scalar. A NULL data pointer is allowed only -/// for an empty byte range. Invalid UTF-8 returns NULL and writes the error -/// output. +/// The string bytes are copied into the scalar. Invalid UTF-8 returns NULL and +/// writes the error output. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_scalar_new_utf8( - ptr: *const c_char, - len: usize, + value: vx_view, is_nullable: bool, err: *mut *mut vx_error, ) -> *mut vx_scalar { try_or(err, ptr::null_mut(), || { - let bytes = bytes_from_raw(ptr.cast(), len, "utf8")?; - let value = str::from_utf8(bytes).map_err(|e| vortex_err!("invalid utf-8: {e}"))?; + let value = unsafe { value.as_str() }?; Ok(vx_scalar::new(Scalar::utf8( value.to_owned(), Nullability::from(is_nullable), @@ -202,9 +193,11 @@ pub unsafe extern "C-unwind" fn vx_scalar_new_binary( /// Create a typed null scalar. /// -/// The data type handle is borrowed, not consumed. The returned scalar uses a -/// nullable copy of that logical type, regardless of the input type's top-level -/// nullability. A NULL data type handle returns NULL and writes the error output. +/// "dtype" is not consumed, you can use it after calling this function. Returned +/// scalar uses a nullable copy of that logical type, regardless of the input +/// type's top-level nullability. +/// +/// Returns NULL and sets "err" on error or NULL dtype. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_scalar_new_null( dtype: *const vx_dtype, @@ -342,10 +335,8 @@ pub unsafe extern "C-unwind" fn vx_scalar_new_decimal_i256_le( /// Create a list scalar. /// -/// The element data type handle is borrowed, not consumed. Child scalar handles -/// are cloned into the list value, so the caller keeps ownership of the handle -/// array and each scalar in it. A NULL child handle array is allowed only for an -/// empty list. Child values are validated against the element logical type. +/// "element_dtype" and "elements" are not consumed, you can use them after +/// calling this function. If len is 0, you can pass NULL to "elements". #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_scalar_new_list( element_dtype: *const vx_dtype, @@ -370,30 +361,25 @@ pub unsafe extern "C-unwind" fn vx_scalar_new_list( /// Create a fixed-size list scalar. /// -/// The element data type handle is borrowed, not consumed. The number of child -/// scalars becomes the fixed-size list width and must fit in a 32-bit unsigned -/// integer. Child scalar handles are cloned into the list value, so the caller -/// keeps ownership of the handle array and each scalar in it. A NULL child -/// handle array is allowed only for an empty list. Child values are validated -/// against the element logical type. +/// "element_dtype" and "elements" are not consumed, you can use them after +/// calling this function. If len is 0, you can pass NULL to "elements". +/// "len" must fit in uint32_t. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_scalar_new_fixed_size_list( element_dtype: *const vx_dtype, elements: *const *const vx_scalar, - len: usize, + len: u32, is_nullable: bool, err: *mut *mut vx_error, ) -> *mut vx_scalar { try_or(err, ptr::null_mut(), || { vortex_ensure!(!element_dtype.is_null(), "element dtype is null"); - let size = u32::try_from(len) - .map_err(|_| vortex_err!("fixed-size list length {len} exceeds u32::MAX"))?; let dtype = DType::FixedSizeList( Arc::new(vx_dtype::as_ref(element_dtype).clone()), - size, + len, Nullability::from(is_nullable), ); - let values = scalar_values_from_raw(elements, len)?; + let values = scalar_values_from_raw(elements, len as usize)?; Ok(vx_scalar::new(Scalar::try_new( dtype, Some(ScalarValue::Tuple(values)), @@ -403,11 +389,8 @@ pub unsafe extern "C-unwind" fn vx_scalar_new_fixed_size_list( /// Create a struct scalar. /// -/// The struct data type handle is borrowed, not consumed. Field scalar handles -/// are cloned into the struct value, so the caller keeps ownership of the handle -/// array and each scalar in it. Field count and field logical types are validated -/// against the struct logical type. A NULL field handle array is allowed only for -/// an empty struct value. +/// "struct_dtype" and "fields" are not consumed, you can use them after calling +/// this function. If len is 0, you can pass NULL to "fields". #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_scalar_new_struct( struct_dtype: *const vx_dtype, @@ -522,6 +505,7 @@ mod tests { use crate::scalar::vx_scalar_new_u32; use crate::scalar::vx_scalar_new_u64; use crate::scalar::vx_scalar_new_utf8; + use crate::string::vx_view; use crate::tests::assert_error; use crate::tests::assert_no_error; @@ -593,18 +577,14 @@ mod tests { let mut error = ptr::null_mut(); let value = "literal"; assert_scalar( - vx_scalar_new_utf8(value.as_ptr().cast(), value.len(), false, &raw mut error), + vx_scalar_new_utf8(vx_view::from_str(value), false, &raw mut error), Scalar::utf8(value, Nullability::NonNullable), ); assert_no_error(error); let invalid_utf8 = [0xffu8]; - let scalar = vx_scalar_new_utf8( - invalid_utf8.as_ptr().cast(), - invalid_utf8.len(), - false, - &raw mut error, - ); + let scalar = + vx_scalar_new_utf8(vx_view::from_bytes(&invalid_utf8), false, &raw mut error); assert!(scalar.is_null()); assert_error(error); @@ -620,10 +600,12 @@ mod tests { vx_dtype_free(dtype); assert_no_error(error); assert!(vx_scalar_is_null(null_scalar)); + let scalar_dtype = vx_scalar_dtype(null_scalar); assert_eq!( - vx_dtype::as_ref(vx_scalar_dtype(null_scalar)), + vx_dtype::as_ref(scalar_dtype), &DType::Primitive(PType::I32, Nullability::Nullable) ); + vx_dtype_free(scalar_dtype); vx_scalar_free(null_scalar); } } @@ -743,11 +725,12 @@ mod tests { ); assert_no_error(error); + let len = u32::try_from(children.len()).unwrap(); assert_scalar( vx_scalar_new_fixed_size_list( element_dtype, children.as_ptr(), - children.len(), + len, false, &raw mut error, ), diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 619f82e3f3e..b21f654634a 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -18,7 +18,6 @@ use futures::StreamExt; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::VortexSessionExecute; -use vortex::array::arrow::ArrowSessionExt; use vortex::array::expr::stats::Precision; use vortex::array::stream::SendableArrayStream; use vortex::buffer::Buffer; @@ -33,6 +32,8 @@ use vortex::scan::Partition; use vortex::scan::PartitionStream; use vortex::scan::ScanRequest; use vortex::scan::selection::Selection; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ToArrowType; use crate::RUNTIME; use crate::array::vx_array; @@ -225,9 +226,7 @@ fn write_estimate>(estimate: Precision, out: &mut vx_estimate) { /// Scan a data source. /// -/// Return an owned scan that must be freed with vx_scan_free. A scan may be -/// consumed only once. -/// +/// A scan may be consumed only once. /// "options" and "estimate" may be NULL. /// /// If "options" is NULL, all rows and columns are returned. @@ -256,10 +255,8 @@ pub unsafe extern "C-unwind" fn vx_data_source_scan( }) } -/// Return borrowed vx_scan's dtype. +/// Return scan's dtype. /// This function will fail if called after vx_scan_next_partition. -/// Called must not free the returned pointer as its lifetime is bound to the -/// lifetime of the scan. /// On error returns NULL and sets "err". #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_scan_dtype( @@ -271,12 +268,11 @@ pub unsafe extern "C-unwind" fn vx_scan_dtype( let VxScan::Pending(scan) = scan else { vortex_bail!("dtype unavailable: scan already started"); }; - Ok(vx_dtype::new_ref(scan.dtype())) + Ok(vx_dtype::new(Arc::new(scan.dtype().clone()))) }) } -/// Return an owned partition from a scan. -/// The returned partition must be freed with vx_partition_free. +/// Return an partition from a scan. /// /// On success returns a partition. /// On exhaustion (no more partitions in scan) returns NULL but doesn't set @@ -398,8 +394,7 @@ pub unsafe extern "C-unwind" fn vx_partition_scan_arrow( }) } -/// Return an owned owned array from a partition. -/// The returned array must be freed with vx_array_free. +/// Return an array from a partition. /// /// On success returns an array. /// On exhaustion (no more arrays in partition) returns NULL but doesn't set @@ -447,7 +442,6 @@ pub unsafe extern "C-unwind" fn vx_partition_next( #[cfg(not(windows))] #[cfg(test)] mod tests { - use std::ffi::CString; use std::ptr; use vortex::VortexSessionDefault; @@ -483,6 +477,7 @@ mod tests { use crate::scan::vx_scan_selection_include; use crate::session::vx_session_free; use crate::session::vx_session_new; + use crate::string::vx_view; use crate::tests::SAMPLE_ROWS; use crate::tests::assert_no_error; use crate::tests::write_sample; @@ -493,9 +488,10 @@ mod tests { unsafe { let session = vx_session_new(); let (sample, struct_array) = write_sample(session); - let path = CString::new(sample.path().to_str().unwrap()).unwrap(); + let path = vx_view::from_str(sample.path().to_str().unwrap()); let ds_options = vx_data_source_options { - paths: path.as_ptr(), + paths: &raw const path, + paths_len: 1, }; let mut error = ptr::null_mut(); @@ -557,8 +553,8 @@ mod tests { let root = vx_expression_root(); let mut opts = vx_scan_options::default(); - for (field, c_field) in [("age", c"age"), ("height", c"height"), ("name", c"name")] { - let field_expr = vx_expression_get_item(c_field.as_ptr(), root); + for field in ["age", "height", "name"] { + let field_expr = vx_expression_get_item(vx_view::from_str(field), root); assert!(!field_expr.is_null()); opts.projection = field_expr; let (array, struct_array) = scan(&raw const opts); @@ -583,8 +579,8 @@ mod tests { let root = vx_expression_root(); let mut opts = vx_scan_options::default(); - let expr_age = vx_expression_get_item(c"age".as_ptr(), root); - let expr_height = vx_expression_get_item(c"height".as_ptr(), root); + let expr_age = vx_expression_get_item(vx_view::from_str("age"), root); + let expr_height = vx_expression_get_item(vx_view::from_str("height"), root); let expr_sum = vx_expression_binary(vx_binary_operator::VX_OPERATOR_ADD, expr_age, expr_height); @@ -614,7 +610,7 @@ mod tests { fn test_filter() { unsafe { let root = vx_expression_root(); - let age_expr = vx_expression_get_item(c"age".as_ptr(), root); + let age_expr = vx_expression_get_item(vx_view::from_str("age"), root); let value = vx_scalar_new_u64(100, false); let mut error = ptr::null_mut(); let lit_100 = vx_expression_literal(value, &raw mut error); @@ -643,7 +639,7 @@ mod tests { fn test_filter_project() { unsafe { let root = vx_expression_root(); - let age_expr = vx_expression_get_item(c"age".as_ptr(), root); + let age_expr = vx_expression_get_item(vx_view::from_str("age"), root); let value = vx_scalar_new_u64(100, false); let mut error = ptr::null_mut(); let lit_100 = vx_expression_literal(value, &raw mut error); @@ -651,7 +647,7 @@ mod tests { vx_scalar_free(value); let filter = vx_expression_binary(vx_binary_operator::VX_OPERATOR_GTE, age_expr, lit_100); - let projection = vx_expression_get_item(c"age".as_ptr(), root); + let projection = vx_expression_get_item(vx_view::from_str("age"), root); let opts = vx_scan_options { projection, @@ -731,9 +727,10 @@ mod tests { unsafe { let session = vx_session_new(); let (sample, _) = write_sample(session); - let path = CString::new(sample.path().to_str().unwrap()).unwrap(); + let path = vx_view::from_str(sample.path().to_str().unwrap()); let ds_options = vx_data_source_options { - paths: path.as_ptr(), + paths: &raw const path, + paths_len: 1, }; let mut error = ptr::null_mut(); diff --git a/vortex-ffi/src/sink.rs b/vortex-ffi/src/sink.rs index 639d22196e3..e7a53a0c8dd 100644 --- a/vortex-ffi/src/sink.rs +++ b/vortex-ffi/src/sink.rs @@ -1,8 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::ffi::CStr; -use std::ffi::c_char; +use std::sync::Arc; use futures::SinkExt; use futures::TryStreamExt; @@ -10,22 +9,34 @@ use futures::channel::mpsc; use futures::channel::mpsc::Sender; use vortex::array::ArrayRef; use vortex::array::stream::ArrayStreamAdapter; +use vortex::dtype::DType; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_ensure; use vortex::error::vortex_err; +use vortex::file::VortexFile; use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; use vortex::file::WriteSummary; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::Task; use vortex::io::session::RuntimeSessionExt; +use vortex::layout::LayoutStrategy; use crate::RUNTIME; +use crate::arc_wrapper; use crate::array::vx_array; use crate::dtype::vx_dtype; use crate::error::try_or_default; use crate::error::vx_error; use crate::session::vx_session; +use crate::string::vx_view; + +arc_wrapper!( + /// A handle to a Vortex file encapsulating the footer and logic for instantiating a reader. + VortexFile, + vx_file +); #[expect(non_camel_case_types)] /// The `sink` interface is used to collect array chunks and place them into a resource @@ -44,43 +55,75 @@ use crate::session::vx_session; pub struct vx_array_sink { sink: Sender>, writer: Task>, + dtype: DType, +} + +/// Open a file sink with an explicit layout strategy. +/// +/// This is a Rust API for FFI crates layered on top of `vortex-ffi`; the base C API continues to +/// use the default write strategy. File creation and write errors are reported when the returned +/// sink is closed. +/// +/// # Safety +/// +/// `session`, `path`, and `dtype` must satisfy the same requirements as +/// `vx_array_sink_open_file`. +pub unsafe fn vx_array_sink_open_file_with_strategy( + session: *const vx_session, + path: vx_view, + dtype: *const vx_dtype, + strategy: Arc, +) -> VortexResult<*mut vx_array_sink> { + let session = vx_session::as_ref(session).clone(); + + if path.ptr.is_null() { + vortex_bail!("null path"); + } + let path = unsafe { path.as_str() }?.to_string(); + + let file_dtype = vx_dtype::as_ref(dtype); + // The channel size 32 was chosen arbitrarily. + let (sink, rx) = mpsc::channel(32); + let dtype = file_dtype.clone(); + let array_stream = ArrayStreamAdapter::new(dtype.clone(), rx.into_stream()); + + let writer_session = session.clone(); + let writer = session.handle().spawn(async move { + let mut file = async_fs::File::create(path).await?; + writer_session + .write_options() + .with_strategy(strategy) + .write(&mut file, array_stream) + .await + }); + + Ok(Box::into_raw(Box::new(vx_array_sink { + sink, + writer, + dtype, + }))) } /// Opens a writable array stream, where sink is used to push values into the stream. /// To close the stream close the sink with `vx_array_sink_close`. +/// "path" is copied. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_array_sink_open_file( session: *const vx_session, - path: *const c_char, + path: vx_view, dtype: *const vx_dtype, error_out: *mut *mut vx_error, ) -> *mut vx_array_sink { try_or_default(error_out, || { - let session = vx_session::as_ref(session); - - if path.is_null() { - vortex_bail!("null path"); - } - let path = unsafe { CStr::from_ptr(path) } - .to_string_lossy() - .to_string(); - - let file_dtype = vx_dtype::as_ref(dtype); - // The channel size 32 was chosen arbitrarily. - let (sink, rx) = mpsc::channel(32); - let array_stream = ArrayStreamAdapter::new(file_dtype.clone(), rx.into_stream()); - - let writer = session.handle().spawn(async move { - let mut file = async_fs::File::create(path).await?; - session.write_options().write(&mut file, array_stream).await - }); - - Ok(Box::into_raw(Box::new(vx_array_sink { sink, writer }))) + let strategy = WriteStrategyBuilder::default().build(); + unsafe { vx_array_sink_open_file_with_strategy(session, path, dtype, strategy) } }) } /// Push an array into a file sink. -/// Does not take ownership of array +/// Does not take ownership of array. +/// +/// Errors if array's DType doesn't match sink's DType. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_array_sink_push( sink: *mut vx_array_sink, @@ -93,6 +136,13 @@ pub unsafe extern "C-unwind" fn vx_array_sink_push( let array = vx_array::as_ref(array); let sink = unsafe { &mut *sink }; + + vortex_ensure!( + *array.dtype() == sink.dtype, + "array dtype {} does not match sink dtype {}", + array.dtype(), + sink.dtype + ); RUNTIME .block_on(sink.sink.send(Ok(array.clone()))) .map_err(|e| vortex_err!("Send error: {e}")) @@ -107,7 +157,11 @@ pub unsafe extern "C-unwind" fn vx_array_sink_close( error_out: *mut *mut vx_error, ) { try_or_default(error_out, || { - let vx_array_sink { sink, writer } = *unsafe { Box::from_raw(sink) }; + let vx_array_sink { + sink, + writer, + dtype: _, + } = *unsafe { Box::from_raw(sink) }; drop(sink); RUNTIME.block_on(async { @@ -119,9 +173,18 @@ pub unsafe extern "C-unwind" fn vx_array_sink_close( }) } +/// Abort an array sink. File footer is not written, and file is left invalid. +/// Don't use sink after this call. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_array_sink_abort(sink: *mut vx_array_sink) { + if sink.is_null() { + return; + } + drop(unsafe { Box::from_raw(sink) }); +} + #[cfg(test)] mod tests { - use std::ffi::CString; use std::sync::Arc; use tempfile::NamedTempFile; @@ -134,6 +197,8 @@ mod tests { use super::*; use crate::array::vx_array; use crate::array::vx_array_free; + use crate::data_source::vx_data_source_new; + use crate::data_source::vx_data_source_options; use crate::dtype::vx_dtype; use crate::dtype::vx_dtype_free; use crate::error::vx_error_free; @@ -147,14 +212,13 @@ mod tests { let session = vx_session_new(); let temp_file = NamedTempFile::new().unwrap(); - let path = CString::new(temp_file.path().to_str().unwrap()).unwrap(); + let path = vx_view::from_str(temp_file.path().to_str().unwrap()); let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype)); let mut error = std::ptr::null_mut(); - let sink = - vx_array_sink_open_file(session, path.as_ptr(), vx_dtype_ptr, &raw mut error); + let sink = vx_array_sink_open_file(session, path, vx_dtype_ptr, &raw mut error); assert!(error.is_null()); assert!(!sink.is_null()); @@ -183,14 +247,13 @@ mod tests { let session = vx_session_new(); let temp_file = NamedTempFile::new().unwrap(); - let path = CString::new(temp_file.path().to_str().unwrap()).unwrap(); + let path = vx_view::from_str(temp_file.path().to_str().unwrap()); let dtype = DType::Primitive(vortex::dtype::PType::U64, false.into()); let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype)); let mut error = std::ptr::null_mut(); - let sink = - vx_array_sink_open_file(session, path.as_ptr(), vx_dtype_ptr, &raw mut error); + let sink = vx_array_sink_open_file(session, path, vx_dtype_ptr, &raw mut error); assert!(error.is_null()); // Push multiple arrays @@ -223,17 +286,12 @@ mod tests { let session = vx_session_new(); // Use a path that will fail during file creation (read-only directory on most systems) - let invalid_path = CString::new("/dev/null/invalid.vortex").unwrap(); + let invalid_path = vx_view::from_str("/dev/null/invalid.vortex"); let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype)); let mut error = std::ptr::null_mut(); - let sink = vx_array_sink_open_file( - session, - invalid_path.as_ptr(), - vx_dtype_ptr, - &raw mut error, - ); + let sink = vx_array_sink_open_file(session, invalid_path, vx_dtype_ptr, &raw mut error); // The sink creation may succeed but close should fail due to invalid path if !sink.is_null() { @@ -261,6 +319,44 @@ mod tests { } } + #[test] + #[cfg_attr(miri, ignore)] + fn test_sink_abort() { + let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); + let array = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); + + let file = NamedTempFile::new().unwrap(); + let path = vx_view::from_str(file.path().to_str().unwrap()); + + unsafe { + let session = vx_session_new(); + let dtype = vx_dtype::new(Arc::new(dtype)); + + let mut error = std::ptr::null_mut(); + let sink = vx_array_sink_open_file(session, path, dtype, &raw mut error); + assert!(error.is_null()); + + let array = vx_array::new(Arc::new(array.into_array())); + vx_array_sink_push(sink, array, &raw mut error); + assert!(error.is_null()); + vx_array_free(array); + + vx_array_sink_abort(sink); + + let opts = vx_data_source_options { + paths: &raw const path, + paths_len: 1, + }; + let ds = vx_data_source_new(session, &raw const opts, &raw mut error); + assert!(ds.is_null()); + assert!(!error.is_null()); + vx_error_free(error); + + vx_dtype_free(dtype); + vx_session_free(session); + } + } + #[test] #[cfg_attr(miri, ignore)] fn test_sink_null_path() { @@ -273,7 +369,7 @@ mod tests { let mut error = std::ptr::null_mut(); // This should return null and set error due to null path let sink = - vx_array_sink_open_file(session, std::ptr::null(), vx_dtype_ptr, &raw mut error); + vx_array_sink_open_file(session, vx_view::null(), vx_dtype_ptr, &raw mut error); assert!(sink.is_null()); assert!(!error.is_null()); diff --git a/vortex-ffi/src/string.rs b/vortex-ffi/src/string.rs index a3dc0937702..8e0dfd71900 100644 --- a/vortex-ffi/src/string.rs +++ b/vortex-ffi/src/string.rs @@ -1,144 +1,91 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::ffi::CStr; use std::ffi::c_char; +use std::ptr; use std::slice; -use vortex::error::VortexExpect; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; use vortex::error::vortex_err; -use crate::arc_dyn_wrapper; - -arc_dyn_wrapper!( - /// Strings for use within Vortex. - str, - vx_string -); - -impl vx_string { - #[allow(dead_code)] - pub(crate) fn as_str(ptr: *const vx_string) -> &'static str { - unsafe { - str::from_utf8_unchecked(slice::from_raw_parts( - vx_string_ptr(ptr).cast(), - vx_string_len(ptr), - )) - } - } -} - -/// Create a new Vortex UTF-8 string by copying from a pointer and length. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_string_new(ptr: *const c_char, len: usize) -> *const vx_string { - let slice = unsafe { slice::from_raw_parts(ptr.cast(), len) }; - // TODO(joe): propagate this error up instead of expecting - let string = String::from_utf8(slice.to_vec()) - .map_err(|e| vortex_err!("invalid utf-8: {e}")) - .vortex_expect("CString creation should succeed"); - vx_string::new(string.into()) -} - -/// Create a new Vortex UTF-8 string by copying from a null-terminated C-style string. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_string_new_from_cstr(ptr: *const c_char) -> *const vx_string { - // TODO(joe): propagate this error up instead of expecting - let string = unsafe { CStr::from_ptr(ptr) } - .to_str() - .map_err(|e| vortex_err!("invalid utf-8: {e}")) - .vortex_expect("CString creation should succeed"); - vx_string::new(string.into()) -} - -/// Return the length of the string in bytes. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_string_len(ptr: *const vx_string) -> usize { - vx_string::as_ref(ptr).len() +/// A non owning view over a byte range. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct vx_view { + /// NULL "ptr" requires len == 0 + pub ptr: *const c_char, + /// Length in bytes. + pub len: usize, } -/// Return the pointer to the string data. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_string_ptr(ptr: *const vx_string) -> *const c_char { - vx_string::as_ref(ptr).as_ptr().cast() -} - -#[cfg(test)] -mod tests { - use std::ffi::CString; - - use super::*; - - #[test] - fn test_string_new() { - unsafe { - let test_str = "hello world"; - let ptr = test_str.as_ptr().cast(); - let len = test_str.len(); - - let vx_str = vx_string_new(ptr, len); - assert_eq!(vx_string_len(vx_str), 11); - assert_eq!(vx_string::as_str(vx_str), "hello world"); - - vx_string_free(vx_str); +impl vx_view { + /// {NULL, 0} for absent values + pub(crate) const fn null() -> vx_view { + vx_view { + ptr: ptr::null(), + len: 0, } } - #[test] - fn test_string_new_from_cstr() { - unsafe { - let c_string = CString::new("test string").unwrap(); - let vx_str = vx_string_new_from_cstr(c_string.as_ptr()); - - assert_eq!(vx_string_len(vx_str), 11); - assert_eq!(vx_string::as_str(vx_str), "test string"); - - vx_string_free(vx_str); + /// Borrow a Rust string + pub(crate) fn from_str(value: &str) -> vx_view { + vx_view { + ptr: value.as_ptr().cast(), + len: value.len(), } } - #[test] - fn test_string_ptr() { - unsafe { - let test_str = "testing"; - let vx_str = vx_string::new(test_str.into()); - - let ptr = vx_string_ptr(vx_str); - let len = vx_string_len(vx_str); - - let slice = slice::from_raw_parts(ptr.cast::(), len); - let recovered = str::from_utf8_unchecked(slice); - assert_eq!(recovered, "testing"); - - vx_string_free(vx_str); + /// Borrow a Rust byte slice + pub(crate) fn from_bytes(value: &[u8]) -> vx_view { + vx_view { + ptr: value.as_ptr().cast(), + len: value.len(), } } - #[test] - fn test_empty_string() { - unsafe { - let empty = ""; - let ptr = empty.as_ptr().cast(); - let vx_str = vx_string_new(ptr, 0); - - assert_eq!(vx_string_len(vx_str), 0); - assert_eq!(vx_string::as_str(vx_str), ""); - - vx_string_free(vx_str); + /// View vx_view as bytes + /// + /// # Safety + /// + /// "ptr" must be valid for "len" reads or NULL with "len == 0". + pub(crate) unsafe fn as_bytes<'a>(&self) -> VortexResult<&'a [u8]> { + if self.ptr.is_null() { + vortex_ensure!(self.len == 0, "null vx_view pointer with non-zero length"); + return Ok(&[]); } + Ok(unsafe { slice::from_raw_parts(self.ptr.cast(), self.len) }) } - #[test] - fn test_unicode_string() { - unsafe { - let unicode_str = "Hello 世界 🌍"; - let ptr = unicode_str.as_ptr().cast(); - let len = unicode_str.len(); + /// View vx_view as UTF-8. + /// + /// # Safety + /// + /// `self.ptr` must be valid for `self.len` reads, or null when `self.len` is zero. + pub unsafe fn as_str<'a>(&self) -> VortexResult<&'a str> { + str::from_utf8(unsafe { self.as_bytes() }?).map_err(|e| vortex_err!("invalid utf-8: {e}")) + } +} - let vx_str = vx_string_new(ptr, len); - assert_eq!(vx_string_len(vx_str), unicode_str.len()); - assert_eq!(vx_string::as_str(vx_str), unicode_str); +#[cfg(test)] +mod tests { + use super::*; - vx_string_free(vx_str); - } + #[test] + fn test_vx_view() -> VortexResult<()> { + let source = "Hello 世界 🌍"; + let view = vx_view::from_str(source); + assert_eq!(unsafe { view.as_str() }?, source); + assert_eq!(unsafe { view.as_bytes() }?, source.as_bytes()); + + assert_eq!(unsafe { vx_view::null().as_str() }?, ""); + let bad = vx_view { + ptr: ptr::null(), + len: 3, + }; + assert!(unsafe { bad.as_str() }.is_err()); + + assert!(unsafe { vx_view::from_bytes(&[0xFFu8, 0xFE]).as_str() }.is_err()); + Ok(()) } } diff --git a/vortex-ffi/src/struct_array.rs b/vortex-ffi/src/struct_array.rs index 765864a6176..21fab2fda7a 100644 --- a/vortex-ffi/src/struct_array.rs +++ b/vortex-ffi/src/struct_array.rs @@ -1,6 +1,5 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::ffi::c_char; use std::ptr; use std::sync::Arc; @@ -16,6 +15,7 @@ use crate::array::vx_array; use crate::array::vx_validity; use crate::error::try_or_default; use crate::error::vx_error; +use crate::string::vx_view; use crate::to_field_name; pub(crate) struct StructBuilder { @@ -59,13 +59,12 @@ pub unsafe extern "C" fn vx_struct_column_builder_new( #[unsafe(no_mangle)] pub unsafe extern "C" fn vx_struct_column_builder_add_field( builder: *mut vx_struct_column_builder, - name: *const c_char, + name: vx_view, field: *const vx_array, error: *mut *mut vx_error, ) { try_or_default(error, || { vortex_ensure!(!builder.is_null()); - vortex_ensure!(!name.is_null()); vortex_ensure!(!field.is_null()); let builder = vx_struct_column_builder::as_mut(builder); @@ -149,6 +148,7 @@ mod tests { use crate::array::vx_validity_type; use crate::error::vx_error_free; use crate::ptype::vx_ptype; + use crate::string::vx_view; use crate::struct_array::vx_struct_column_builder_add_field; use crate::struct_array::vx_struct_column_builder_finalize; use crate::struct_array::vx_struct_column_builder_free; @@ -205,7 +205,7 @@ mod tests { vx_struct_column_builder_add_field( builder, - c"age".as_ptr(), + vx_view::from_str("age"), ffi_age_field, &raw mut error, ); @@ -215,7 +215,7 @@ mod tests { let ffi_null_field = vx_array_new_null(5); vx_struct_column_builder_add_field( builder, - c"null".as_ptr(), + vx_view::from_str("null"), ffi_null_field, &raw mut error, ); @@ -227,7 +227,7 @@ mod tests { let ffi_name_field = vx_array::new(Arc::new(name_field.into_array())); vx_struct_column_builder_add_field( builder, - c"name".as_ptr(), + vx_view::from_str("name"), ffi_name_field, &raw mut error, ); diff --git a/vortex-ffi/src/struct_fields.rs b/vortex-ffi/src/struct_fields.rs index fa8f10aca83..4cb22bdb521 100644 --- a/vortex-ffi/src/struct_fields.rs +++ b/vortex-ffi/src/struct_fields.rs @@ -11,7 +11,9 @@ use vortex::error::VortexExpect; use crate::box_wrapper; use crate::dtype::vx_dtype; -use crate::string::vx_string; +use crate::error::try_or_default; +use crate::error::vx_error; +use crate::string::vx_view; box_wrapper!( /// Represents a Vortex struct data type, without top-level nullability. @@ -29,33 +31,26 @@ pub unsafe extern "C-unwind" fn vx_struct_fields_nfields(dtype: *const vx_struct .nfields() as u64 } -/// Return a borrowed reference to the name of the field at the given index. +/// Return field name at a given index. +/// If index is out of bounds, returns {NULL, 0}. /// -/// The returned pointer is valid as long as the struct fields is valid. -/// Do NOT free the returned string pointer - it shares the lifetime of the struct fields. -/// Returns null if the index is out of bounds. +/// Returned view is valid as long as "dtype" is valid. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_struct_fields_field_name( dtype: *const vx_struct_fields, idx: usize, -) -> *const vx_string { +) -> vx_view { // TODO(joe): propagate this error up instead of expecting let ptr = unsafe { dtype.as_ref() }.vortex_expect("null ptr"); let struct_dtype = &ptr.0; if idx >= struct_dtype.nfields() { - return ptr::null(); + return vx_view::null(); } - let name = struct_dtype.names()[idx].inner(); - vx_string::new_ref(name) + vx_view::from_str(struct_dtype.names()[idx].inner()) } -/// Returns an *owned* reference to the dtype of the field at the given index. -/// -/// The return type is owned since struct dtypes can be lazily parsed from a binary format, in -/// which case it's not possible to return a borrowed reference to the field dtype. -/// -/// Returns null if the index is out of bounds or if the field dtype cannot be parsed. -// TODO(ngates): should StructDType cache owned fields internally? +/// Return an owned dtype of the field at a given index. +/// Returns NULL if index is out of bounds or if dtype cannot be parsed. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_struct_fields_field_dtype( dtype: *const vx_struct_fields, @@ -97,19 +92,23 @@ pub unsafe extern "C-unwind" fn vx_struct_fields_builder_new() -> *mut vx_struct /// Add a field to the struct dtype builder. /// -/// Takes ownership of both the `name` and `dtype` pointers. -/// Must either free or finalize the builder. +/// "name" is copied. Takes ownership of "dtype". +/// Caller must free or finalize the builder. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_struct_fields_builder_add_field( builder: *mut vx_struct_fields_builder, - name: *const vx_string, + name: vx_view, dtype: *const vx_dtype, + error_out: *mut *mut vx_error, ) { - let builder = vx_struct_fields_builder::as_mut(builder); - builder.names.push(vx_string::into_arc(name)); - builder - .fields - .push(vx_dtype::into_arc(dtype).deref().clone()); + try_or_default(error_out, || { + let builder = vx_struct_fields_builder::as_mut(builder); + let field = vx_dtype::into_arc(dtype).deref().clone(); + let name = Arc::from(unsafe { name.as_str() }?); + builder.fields.push(field); + builder.names.push(name); + Ok(()) + }) } /// Finalize the struct dtype builder, returning a new `vx_struct_fields`. diff --git a/vortex-ffi/test/array.cpp b/vortex-ffi/test/array.cpp index 4fde064a339..f8db2661b8d 100644 --- a/vortex-ffi/test/array.cpp +++ b/vortex-ffi/test/array.cpp @@ -9,7 +9,11 @@ TEST_CASE("Null array creation", "[array]") { REQUIRE(array != nullptr); REQUIRE(vx_array_is_nullable(array)); REQUIRE(vx_array_has_dtype(array, DTYPE_NULL)); - REQUIRE(vx_dtype_get_variant(vx_array_dtype(array)) == DTYPE_NULL); + const vx_dtype *dtype = vx_array_dtype(array); + defer { + vx_dtype_free(dtype); + }; + REQUIRE(vx_dtype_get_variant(dtype) == DTYPE_NULL); REQUIRE(vx_array_len(array) == 1999); vx_array_free(array); } @@ -26,7 +30,11 @@ TEST_CASE("Primitive array creation", "[array]") { require_no_error(error); REQUIRE(array != nullptr); REQUIRE(vx_array_has_dtype(array, DTYPE_PRIMITIVE)); - REQUIRE(vx_dtype_get_variant(vx_array_dtype(array)) == DTYPE_PRIMITIVE); + const vx_dtype *dtype = vx_array_dtype(array); + REQUIRE(vx_dtype_get_variant(dtype) == DTYPE_PRIMITIVE); + defer { + vx_dtype_free(dtype); + }; REQUIRE(vx_array_is_primitive(array, PTYPE_U8)); REQUIRE(vx_array_len(array) == buffer.size()); @@ -54,7 +62,7 @@ TEST_CASE("Struct array creation", "[array]") { vx_struct_column_builder *builder = vx_struct_column_builder_new(&validity, 2); CHECK(builder != nullptr); - vx_struct_column_builder_add_field(builder, "age", field_array, &error); + vx_struct_column_builder_add_field(builder, vx_view_from_cstr("age"), field_array, &error); vx_array_free(field_array); SECTION("Struct array builder free") { diff --git a/vortex-ffi/test/common.h b/vortex-ffi/test/common.h index 29087e7a48d..81974b65bb9 100644 --- a/vortex-ffi/test/common.h +++ b/vortex-ffi/test/common.h @@ -6,23 +6,19 @@ #include "vortex.h" inline std::string to_string(vx_error *err) { - const vx_string *msg = vx_error_get_message(err); - return {vx_string_ptr(msg), vx_string_len(msg)}; + const vx_view msg = vx_error_message(err); + return {msg.ptr, msg.len}; } -inline std::string_view to_string_view(const vx_string *msg) { - return {vx_string_ptr(msg), vx_string_len(msg)}; -} - -inline std::string_view to_string_view(vx_error *err) { - return to_string_view(vx_error_get_message(err)); +inline std::string_view to_string_view(vx_view str) { + return {str.ptr, str.len}; } inline void require_no_error(vx_error *error, bool assert = true) { if (!error) { return; } - auto message = to_string(error); + std::string message = to_string(error); vx_error_free(error); if (assert) { FAIL(message); diff --git a/vortex-ffi/test/main.cpp b/vortex-ffi/test/main.cpp index 188e8739d9f..9df02b9f798 100644 --- a/vortex-ffi/test/main.cpp +++ b/vortex-ffi/test/main.cpp @@ -17,25 +17,12 @@ TEST_CASE("Session creation", "[session]") { vx_session_free(session2); } -TEST_CASE("Creating and iterating binaries", "[binary]") { - for (std::string_view str : {"ololo"sv, "Широкая строка"sv, "مرحبا بالعالم"sv}) { - const vx_binary *binary = vx_binary_new(str.data(), str.size()); - - REQUIRE(binary != nullptr); - const size_t len = vx_binary_len(binary); - REQUIRE(len == str.size()); - - const char *ptr = vx_binary_ptr(binary); - REQUIRE(std::string_view {ptr, len} == str); - - const vx_binary *binary2 = vx_binary_clone(binary); - vx_binary_free(binary); - - ptr = vx_binary_ptr(binary2); - REQUIRE(std::string_view {ptr, len} == str); - - vx_binary_free(binary2); - } +TEST_CASE("vx_view from C string", "[str]") { + const std::string_view str = "Широкая строка"sv; + const std::string owned {str}; + const vx_view view = vx_view_from_cstr(owned.c_str()); + REQUIRE(view.len == str.size()); + REQUIRE(std::string_view {view.ptr, view.len} == str); } TEST_CASE("Creating dtypes", "[dtype]") { diff --git a/vortex-ffi/test/scan.cpp b/vortex-ffi/test/scan.cpp index 6857ea82bc1..e1f629dfde0 100644 --- a/vortex-ffi/test/scan.cpp +++ b/vortex-ffi/test/scan.cpp @@ -57,15 +57,14 @@ struct TempPath : fs::path { [[nodiscard]] const vx_dtype *sample_dtype() { vx_struct_fields_builder *builder = vx_struct_fields_builder_new(); - constexpr auto age = "age"sv; - const vx_string *age_name = vx_string_new(age.data(), age.size()); + vx_error *error = nullptr; const vx_dtype *age_type = vx_dtype_new_primitive(PTYPE_U8, false); - vx_struct_fields_builder_add_field(builder, age_name, age_type); + vx_struct_fields_builder_add_field(builder, vx_view_from_cstr("age"), age_type, &error); + require_no_error(error); - constexpr auto height = "height"sv; - const vx_string *height_name = vx_string_new(height.data(), height.size()); const vx_dtype *height_type = vx_dtype_new_primitive(PTYPE_U16, true); - vx_struct_fields_builder_add_field(builder, height_name, height_type); + vx_struct_fields_builder_add_field(builder, vx_view_from_cstr("height"), height_type, &error); + require_no_error(error); vx_struct_fields *fields = vx_struct_fields_builder_finalize(builder); return vx_dtype_new_struct(fields, false); @@ -103,7 +102,7 @@ std::vector sample_height() { }; require_no_error(error); - vx_struct_column_builder_add_field(builder, "age", age_array, &error); + vx_struct_column_builder_add_field(builder, vx_view_from_cstr("age"), age_array, &error); require_no_error(error); std::vector height_buffer = sample_height(); @@ -115,7 +114,7 @@ std::vector sample_height() { }; require_no_error(error); - vx_struct_column_builder_add_field(builder, "height", height_array, &error); + vx_struct_column_builder_add_field(builder, vx_view_from_cstr("height"), height_array, &error); require_no_error(error); const vx_array *array = vx_struct_column_builder_finalize(builder, &error); @@ -168,7 +167,7 @@ UniqueArrayStream sample_array_stream() { }; vx_error *error = nullptr; - vx_array_sink *sink = vx_array_sink_open_file(session, path.c_str(), dtype, &error); + vx_array_sink *sink = vx_array_sink_open_file(session, vx_view_from_cstr(path.c_str()), dtype, &error); REQUIRE(sink != nullptr); require_no_error(error); @@ -206,14 +205,16 @@ TEST_CASE("Creating datasources", "[datasource]") { vx_error_free(error); // First file is opened eagerly - opts.paths = "nonexistent"; + vx_view path_view = vx_view_from_cstr("nonexistent"); + opts.paths = &path_view; + opts.paths_len = 1; ds = vx_data_source_new(session, &opts, &error); REQUIRE(ds == nullptr); REQUIRE(error != nullptr); REQUIRE_THAT(to_string(error), ContainsSubstring("No files matched the glob pattern")); vx_error_free(error); - opts.paths = "/tmp2/*.vortex"; + path_view = vx_view_from_cstr("/tmp2/*.vortex"); ds = vx_data_source_new(session, &opts, &error); REQUIRE(ds == nullptr); REQUIRE(error != nullptr); @@ -225,7 +226,7 @@ TEST_CASE("Creating datasources", "[datasource]") { vx_error_free(error); TempPath file = write_sample(session); - opts.paths = file.c_str(); + path_view = vx_view_from_cstr(file.c_str()); ds = vx_data_source_new(session, &opts, &error); require_no_error(error); REQUIRE(ds != nullptr); @@ -249,7 +250,9 @@ TEST_CASE("Write file and read dtypes", "[datasource]") { vx_error *error = nullptr; vx_data_source_options opts = {}; - opts.paths = path.c_str(); + const vx_view opts_path = vx_view_from_cstr(path.c_str()); + opts.paths = &opts_path; + opts.paths_len = 1; const vx_data_source *ds = vx_data_source_new(session, &opts, &error); require_no_error(error); @@ -265,27 +268,34 @@ TEST_CASE("Write file and read dtypes", "[datasource]") { CHECK(row_count.estimate == SAMPLE_ROWS); const vx_dtype *data_source_dtype = vx_data_source_dtype(ds); + defer { + vx_dtype_free(data_source_dtype); + }; REQUIRE(vx_dtype_get_variant(data_source_dtype) == DTYPE_STRUCT); const vx_struct_fields *fields = vx_dtype_struct_dtype(data_source_dtype); + defer { + vx_struct_fields_free(fields); + }; const size_t len = vx_struct_fields_nfields(fields); REQUIRE(len == 2); const vx_dtype *age_dtype = vx_struct_fields_field_dtype(fields, 0); + const vx_view age_name = vx_struct_fields_field_name(fields, 0); defer { vx_dtype_free(age_dtype); }; - const vx_string *age_name = vx_struct_fields_field_name(fields, 0); + REQUIRE(vx_dtype_get_variant(age_dtype) == DTYPE_PRIMITIVE); REQUIRE(vx_dtype_primitive_ptype(age_dtype) == PTYPE_U8); REQUIRE_FALSE(vx_dtype_is_nullable(age_dtype)); REQUIRE(to_string_view(age_name) == "age"); const vx_dtype *height_dtype = vx_struct_fields_field_dtype(fields, 1); + const vx_view height_name = vx_struct_fields_field_name(fields, 1); defer { vx_dtype_free(height_dtype); }; - const vx_string *height_name = vx_struct_fields_field_name(fields, 1); REQUIRE(vx_dtype_get_variant(height_dtype) == DTYPE_PRIMITIVE); REQUIRE(vx_dtype_primitive_ptype(height_dtype) == PTYPE_U16); REQUIRE(vx_dtype_is_nullable(height_dtype)); @@ -294,7 +304,11 @@ TEST_CASE("Write file and read dtypes", "[datasource]") { void verify_age_field(const vx_array *age_field) { REQUIRE(vx_array_has_dtype(age_field, DTYPE_PRIMITIVE)); - REQUIRE(vx_dtype_primitive_ptype(vx_array_dtype(age_field)) == PTYPE_U8); + const vx_dtype *dtype = vx_array_dtype(age_field); + defer { + vx_dtype_free(dtype); + }; + REQUIRE(vx_dtype_primitive_ptype(dtype) == PTYPE_U8); REQUIRE(vx_array_len(age_field) == SAMPLE_ROWS); for (size_t i = 0; i < SAMPLE_ROWS; ++i) { REQUIRE(vx_array_get_u8(age_field, i) == i); @@ -303,7 +317,11 @@ void verify_age_field(const vx_array *age_field) { void verify_height_field(const vx_array *height_field) { REQUIRE(vx_array_has_dtype(height_field, DTYPE_PRIMITIVE)); - REQUIRE(vx_dtype_primitive_ptype(vx_array_dtype(height_field)) == PTYPE_U16); + const vx_dtype *dtype = vx_array_dtype(height_field); + defer { + vx_dtype_free(dtype); + }; + REQUIRE(vx_dtype_primitive_ptype(dtype) == PTYPE_U16); REQUIRE(vx_array_len(height_field) == SAMPLE_ROWS); for (size_t i = 0; i < SAMPLE_ROWS; ++i) { REQUIRE(vx_array_get_u16(height_field, i) > 0); @@ -314,7 +332,9 @@ void verify_sample_array(const vx_array *array) { REQUIRE(vx_array_len(array) == SAMPLE_ROWS); REQUIRE(vx_array_has_dtype(array, DTYPE_STRUCT)); - const vx_struct_fields *fields = vx_dtype_struct_dtype(vx_array_dtype(array)); + const vx_dtype *dtype = vx_array_dtype(array); + const vx_struct_fields *fields = vx_dtype_struct_dtype(dtype); + vx_dtype_free(dtype); size_t len = vx_struct_fields_nfields(fields); REQUIRE(len == 2); @@ -323,16 +343,18 @@ void verify_sample_array(const vx_array *array) { REQUIRE(vx_dtype_get_variant(age_dtype) == DTYPE_PRIMITIVE); REQUIRE(vx_dtype_primitive_ptype(age_dtype) == PTYPE_U8); vx_dtype_free(age_dtype); - const vx_string *age_name = vx_struct_fields_field_name(fields, 0); + const vx_view age_name = vx_struct_fields_field_name(fields, 0); REQUIRE(to_string_view(age_name) == "age"); const vx_dtype *height_dtype = vx_struct_fields_field_dtype(fields, 1); REQUIRE(vx_dtype_get_variant(height_dtype) == DTYPE_PRIMITIVE); REQUIRE(vx_dtype_primitive_ptype(height_dtype) == PTYPE_U16); vx_dtype_free(height_dtype); - const vx_string *height_name = vx_struct_fields_field_name(fields, 1); + const vx_view height_name = vx_struct_fields_field_name(fields, 1); REQUIRE(to_string_view(height_name) == "height"); + vx_struct_fields_free(fields); + vx_error *error = nullptr; vx_validity validity = {}; vx_array_get_validity(array, &validity, &error); @@ -362,7 +384,9 @@ TEST_CASE("Requesting scans", "[datasource]") { TempPath path = write_sample(session); vx_data_source_options ds_options = {}; - ds_options.paths = path.c_str(); + const vx_view ds_path = vx_view_from_cstr(path.c_str()); + ds_options.paths = &ds_path; + ds_options.paths_len = 1; vx_error *error = nullptr; const vx_data_source *ds = vx_data_source_new(session, &ds_options, &error); @@ -438,7 +462,9 @@ TEST_CASE("Basic scan", "[datasource]") { vx_error *error = nullptr; vx_data_source_options ds_options = {}; - ds_options.paths = path.c_str(); + const vx_view ds_path = vx_view_from_cstr(path.c_str()); + ds_options.paths = &ds_path; + ds_options.paths_len = 1; const vx_data_source *ds = vx_data_source_new(session, &ds_options, &error); require_no_error(error); REQUIRE(ds != nullptr); @@ -481,18 +507,15 @@ TEST_CASE("Multithreaded scan", "[datasource]") { constexpr size_t NUM_FILES = 10; std::vector paths(NUM_FILES); - std::string paths_str; + std::vector path_views(NUM_FILES); for (size_t i = 0; i < NUM_FILES; ++i) { paths[i] = write_sample(session); - if (i == 0) { - paths_str = paths[i].c_str(); - } else { - paths_str += ","s + paths[i].c_str(); - } + path_views[i] = vx_view_from_cstr(paths[i].c_str()); } vx_data_source_options ds_options = {}; - ds_options.paths = paths_str.c_str(); + ds_options.paths = path_views.data(); + ds_options.paths_len = path_views.size(); vx_error *error = nullptr; const vx_data_source *ds = vx_data_source_new(session, &ds_options, &error); @@ -580,7 +603,9 @@ const vx_array *scan_with_options(vx_scan_options &options) { vx_error *error = nullptr; vx_data_source_options ds_options = {}; - ds_options.paths = path.c_str(); + const vx_view ds_path = vx_view_from_cstr(path.c_str()); + ds_options.paths = &ds_path; + ds_options.paths_len = 1; const vx_data_source *ds = vx_data_source_new(session, &ds_options, &error); require_no_error(error); @@ -640,7 +665,7 @@ TEST_CASE("Project single field", "[projection]") { }; vx_scan_options opts = {}; - vx_expression *age_field = vx_expression_get_item("age", root); + vx_expression *age_field = vx_expression_get_item(vx_view_from_cstr("age"), root); REQUIRE(age_field != nullptr); defer { vx_expression_free(age_field); @@ -655,7 +680,7 @@ TEST_CASE("Project single field", "[projection]") { verify_age_field(array); } - vx_expression *height_field = vx_expression_get_item("height", root); + vx_expression *height_field = vx_expression_get_item(vx_view_from_cstr("height"), root); REQUIRE(height_field != nullptr); defer { vx_expression_free(height_field); @@ -677,7 +702,7 @@ TEST_CASE("Filter with literal expression", "[filter]") { vx_expression_free(root); }; - vx_expression *age_field = vx_expression_get_item("age", root); + vx_expression *age_field = vx_expression_get_item(vx_view_from_cstr("age"), root); REQUIRE(age_field != nullptr); defer { vx_expression_free(age_field); @@ -729,7 +754,8 @@ TEST_CASE("Filter with literal expression", "[filter]") { TEST_CASE("Project UTF-8 literal expression", "[projection]") { constexpr auto value = "constant"sv; vx_error *scalar_error = nullptr; - vx_scalar *literal_scalar = vx_scalar_new_utf8(value.data(), value.size(), false, &scalar_error); + vx_scalar *literal_scalar = + vx_scalar_new_utf8(vx_view {value.data(), value.size()}, false, &scalar_error); require_no_error(scalar_error); REQUIRE(literal_scalar != nullptr); defer { @@ -753,12 +779,21 @@ TEST_CASE("Project UTF-8 literal expression", "[projection]") { REQUIRE(vx_array_len(array) == SAMPLE_ROWS); + vx_session *session = vx_session_new(); + defer { + vx_session_free(session); + }; + vx_error *canon_error = nullptr; + const vx_array *canonical = vx_array_canonicalize(session, array, &canon_error); + require_no_error(canon_error); + defer { + vx_array_free(canonical); + }; for (size_t i : {size_t {0}, SAMPLE_ROWS - 1}) { - const vx_string *actual = vx_array_get_utf8(array, static_cast(i)); - REQUIRE(actual != nullptr); - defer { - vx_string_free(actual); - }; + vx_error *at_error = nullptr; + const vx_view actual = vx_array_utf8_at(canonical, i, &at_error); + require_no_error(at_error); + REQUIRE(actual.ptr != nullptr); REQUIRE(to_string_view(actual) == value); } } @@ -838,7 +873,9 @@ TEST_CASE("Scan Arrow schema", "[scan]") { vx_error *error = nullptr; vx_data_source_options ds_options = {}; - ds_options.paths = path.c_str(); + const vx_view ds_path = vx_view_from_cstr(path.c_str()); + ds_options.paths = &ds_path; + ds_options.paths_len = 1; const vx_data_source *ds = vx_data_source_new(session, &ds_options, &error); require_no_error(error); @@ -880,7 +917,9 @@ TEST_CASE("Scan to Arrow", "[scan]") { vx_error *error = nullptr; vx_data_source_options ds_options = {}; - ds_options.paths = path.c_str(); + const vx_view ds_path = vx_view_from_cstr(path.c_str()); + ds_options.paths = &ds_path; + ds_options.paths_len = 1; const vx_data_source *ds = vx_data_source_new(session, &ds_options, &error); require_no_error(error); @@ -920,7 +959,9 @@ TEST_CASE("Broken scan with DType mismatch in filter", "[filter]") { vx_error *error = nullptr; vx_data_source_options ds_opts = {}; - ds_opts.paths = path.c_str(); + const vx_view ds_path = vx_view_from_cstr(path.c_str()); + ds_opts.paths = &ds_path; + ds_opts.paths_len = 1; const vx_data_source *ds = vx_data_source_new(session, &ds_opts, &error); require_no_error(error); defer { @@ -932,7 +973,7 @@ TEST_CASE("Broken scan with DType mismatch in filter", "[filter]") { vx_expression_free(root); }; - vx_expression *age_col = vx_expression_get_item("age", root); + vx_expression *age_col = vx_expression_get_item(vx_view_from_cstr("age"), root); REQUIRE(age_col != nullptr); defer { vx_expression_free(age_col); diff --git a/vortex-ffi/test/struct.cpp b/vortex-ffi/test/struct.cpp index 6ccdb04c037..2a45fccf56f 100644 --- a/vortex-ffi/test/struct.cpp +++ b/vortex-ffi/test/struct.cpp @@ -3,21 +3,22 @@ #include #include +#include "common.h" + using namespace std::string_view_literals; using namespace std::string_literals; TEST_CASE("Struct builder", "[struct]") { vx_struct_fields_builder *builder = vx_struct_fields_builder_new(); + vx_error *error = nullptr; - constexpr auto col1 = "col1"sv; - const vx_string *col1_name = vx_string_new(col1.data(), col1.size()); const vx_dtype *col1_dtype = vx_dtype_new_primitive(PTYPE_U8, false); - vx_struct_fields_builder_add_field(builder, col1_name, col1_dtype); + vx_struct_fields_builder_add_field(builder, vx_view_from_cstr("col1"), col1_dtype, &error); + require_no_error(error); - constexpr auto col2 = "col2"sv; - const vx_string *col2_name = vx_string_new(col2.data(), col2.size()); const vx_dtype *col2_dtype = vx_dtype_new_binary(true); - vx_struct_fields_builder_add_field(builder, col2_name, col2_dtype); + vx_struct_fields_builder_add_field(builder, vx_view_from_cstr("col2"), col2_dtype, &error); + require_no_error(error); SECTION("Struct builder free") { vx_struct_fields_builder_free(builder); @@ -41,12 +42,13 @@ constexpr size_t STRUCT_LEN = 10; TEST_CASE("Creating structs", "[struct]") { vx_struct_fields_builder *builder = vx_struct_fields_builder_new(); REQUIRE(builder != nullptr); + vx_error *error = nullptr; for (size_t i = 0; i < STRUCT_LEN; ++i) { const std::string target_name = "name"s + std::to_string(i); - const vx_string *name = vx_string_new(target_name.data(), target_name.size()); const vx_dtype *dtype = i % 2 ? vx_dtype_new_binary(false) : vx_dtype_new_primitive(PTYPE_F32, true); - vx_struct_fields_builder_add_field(builder, name, dtype); + vx_struct_fields_builder_add_field(builder, vx_view_from_cstr(target_name.c_str()), dtype, &error); + require_no_error(error); } vx_struct_fields *fields = vx_struct_fields_builder_finalize(builder); REQUIRE(fields != nullptr); @@ -54,15 +56,11 @@ TEST_CASE("Creating structs", "[struct]") { const size_t len = vx_struct_fields_nfields(fields); CHECK(len == STRUCT_LEN); for (size_t i = 0; i < len; ++i) { - // borrowed - const vx_string *name = vx_struct_fields_field_name(fields, i); - // owned TODO(myrrc): that's weird API + const vx_view name = vx_struct_fields_field_name(fields, i); const vx_dtype *dtype = vx_struct_fields_field_dtype(fields, i); - std::string_view name_view {vx_string_ptr(name), vx_string_len(name)}; std::string target_name = "name"s + std::to_string(i); - - CHECK(name_view == target_name); + CHECK(to_string_view(name) == target_name); if (i % 2) { CHECK_FALSE(vx_dtype_is_nullable(dtype)); diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index 752ccbc753d..a7af1b293dd 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -30,6 +30,7 @@ parking_lot = { workspace = true } pin-project-lite = { workspace = true } tokio = { workspace = true, features = ["rt"], optional = true } tracing = { workspace = true } +url = { workspace = true } vortex-alp = { workspace = true } vortex-array = { workspace = true } vortex-btrblocks = { workspace = true } diff --git a/vortex-file/src/counting.rs b/vortex-file/src/counting.rs index 78f8a8741a9..003e3eb9465 100644 --- a/vortex-file/src/counting.rs +++ b/vortex-file/src/counting.rs @@ -9,13 +9,14 @@ use std::sync::atomic::Ordering; use vortex_io::IoBuf; use vortex_io::VortexWrite; -/// A wrapper around an `VortexWrite` that counts the number of bytes written. -pub(crate) struct CountingVortexWrite { +/// A wrapper around a [`VortexWrite`] that counts the number of bytes written. +pub struct CountingVortexWrite { inner: W, bytes_written: Arc, } impl CountingVortexWrite { + /// Wrap a writer with a new byte counter. pub fn new(inner: W) -> Self { Self { inner, @@ -23,6 +24,7 @@ impl CountingVortexWrite { } } + /// Returns the shared byte counter updated by this writer. pub fn counter(&self) -> Arc { Arc::clone(&self.bytes_written) } diff --git a/vortex-file/src/footer/field_sizes.rs b/vortex-file/src/footer/field_sizes.rs new file mode 100644 index 00000000000..5d7d2ce7c02 --- /dev/null +++ b/vortex-file/src/footer/field_sizes.rs @@ -0,0 +1,243 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::dtype::Field; +use vortex_array::dtype::FieldPath; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_layout::LayoutChildType; +use vortex_layout::LayoutRef; +use vortex_layout::segments::SegmentId; +use vortex_utils::aliases::hash_map::Entry; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::footer::SegmentSpec; + +/// Compressed sizes of a file's fields, keyed by field path. +/// +/// Sizes are derived from the file's layout tree and segment map: walking the tree from the root, +/// each physical segment is attributed to the deepest [`FieldPath`] whose layout subtree +/// references it. A segment referenced by more than one field is attributed to their longest +/// common ancestor, so a segment is counted exactly once and the size of a field always equals +/// the sum of its child field sizes plus the bytes attributed directly to it (e.g. struct +/// validity, or auxiliary segments such as zone maps and dictionaries written above the field +/// split). +/// +/// Fields nested inside non-struct types (e.g. structs within lists) are not split into separate +/// layout subtrees by the writer, so their bytes are attributed to the enclosing field. +/// +/// File-level metadata, alignment padding, and the footer are not part of any segment and are +/// never attributed. +#[derive(Debug, Clone)] +pub struct CompressedFieldSizes { + sizes: HashMap, +} + +impl CompressedFieldSizes { + pub(crate) fn try_new(root: &LayoutRef, segments: &[SegmentSpec]) -> VortexResult { + // Every segment is attributed to the deepest field path referencing it, collapsing to the + // longest common ancestor if multiple fields reference the same segment. + let mut attribution = HashMap::::default(); + // Seed every field path present in the layout tree so empty fields report a zero size. + let mut sizes = HashMap::::default(); + sizes.insert(FieldPath::root(), 0); + + let mut stack = vec![(LayoutRef::clone(root), FieldPath::root())]; + while let Some((layout, path)) = stack.pop() { + for segment_id in layout.segment_ids() { + match attribution.entry(segment_id) { + Entry::Occupied(mut entry) => { + let ancestor = common_prefix(entry.get(), &path); + entry.insert(ancestor); + } + Entry::Vacant(entry) => { + entry.insert(path.clone()); + } + } + } + + for idx in 0..layout.nchildren() { + let child_path = match layout.child_type(idx) { + LayoutChildType::Field(name) => { + let child_path = path.clone().push(name); + sizes.entry(child_path.clone()).or_insert(0); + child_path + } + _ => path.clone(), + }; + stack.push((layout.child(idx)?, child_path)); + } + } + + for (segment_id, path) in attribution { + let segment = segments.get(*segment_id as usize).ok_or_else(|| { + vortex_err!( + "layout references missing segment {} (segment count: {})", + *segment_id, + segments.len() + ) + })?; + // A segment counts towards its attributed field and every enclosing field. + for len in 0..=path.parts().len() { + *sizes + .entry(FieldPath::from(path.parts()[..len].to_vec())) + .or_insert(0) += u64::from(segment.length); + } + } + + Ok(Self { sizes }) + } + + /// Returns the compressed size in bytes of the field at the given path, including all of its + /// descendants, or `None` if the layout tree contains no such field. + /// + /// [`FieldPath::root`] returns the same value as [`total`][Self::total]. + pub fn get(&self, path: &FieldPath) -> Option { + self.sizes.get(path).copied() + } + + /// Returns the total compressed size in bytes of all segments referenced by the layout tree. + pub fn total(&self) -> u64 { + self.get(&FieldPath::root()).unwrap_or_default() + } + + /// Iterates over all field paths in the layout tree and their compressed sizes. + pub fn iter(&self) -> impl Iterator { + self.sizes.iter().map(|(path, size)| (path, *size)) + } +} + +fn common_prefix(left: &FieldPath, right: &FieldPath) -> FieldPath { + left.parts() + .iter() + .zip(right.parts()) + .take_while(|(l, r)| l == r) + .map(|(l, _)| l.clone()) + .collect::>() + .into() +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_array::IntoArray; + use vortex_array::array_session; + use vortex_array::arrays::StructArray; + use vortex_array::dtype::FieldPath; + use vortex_array::field_path; + use vortex_buffer::ByteBufferMut; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_io::session::RuntimeSession; + use vortex_layout::session::LayoutSession; + use vortex_session::VortexSession; + + use crate::WriteOptionsSessionExt; + use crate::writer::WriteSummary; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session() + .with::() + .with::(); + crate::register_default_encodings(&session); + session + }); + + async fn write_nested() -> VortexResult { + let inner = StructArray::from_fields(&[ + ("c", buffer![10i64, 20, 30, 40].into_array()), + ( + "d", + StructArray::from_fields(&[("e", buffer![1.0f64, 2.0, 3.0, 4.0].into_array())])? + .into_array(), + ), + ])?; + let root = StructArray::from_fields(&[ + ("a", buffer![1u32, 2, 3, 4].into_array()), + ("b", inner.into_array()), + ])?; + + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .write(&mut buf, root.into_array().to_array_stream()) + .await + } + + #[tokio::test] + async fn nested_field_sizes() -> VortexResult<()> { + let summary = write_nested().await?; + let sizes = summary.footer().compressed_field_sizes()?; + + for path in [ + field_path!(a), + field_path!(b), + field_path!(b.c), + field_path!(b.d), + field_path!(b.d.e), + ] { + assert!( + sizes.get(&path).is_some_and(|size| size > 0), + "expected non-zero size for {path}, got {:?}", + sizes.get(&path) + ); + } + assert_eq!(sizes.get(&field_path!(missing)), None); + assert_eq!(sizes.get(&field_path!(b.d.e.missing)), None); + + // Attribution partitions segments: a field's size is the sum of its children plus its own + // directly-attributed segments (none here, as the structs are non-nullable). + let size = |path: FieldPath| sizes.get(&path).unwrap_or_default(); + assert_eq!( + size(field_path!(b)), + size(field_path!(b.c)) + size(field_path!(b.d)) + ); + assert_eq!(size(field_path!(b.d)), size(field_path!(b.d.e))); + assert_eq!(sizes.total(), size(field_path!(a)) + size(field_path!(b))); + + // Every segment in the file is attributed exactly once. + let segment_total: u64 = summary + .footer() + .segment_map() + .iter() + .map(|segment| u64::from(segment.length)) + .sum(); + assert_eq!(sizes.total(), segment_total); + + Ok(()) + } + + #[tokio::test] + async fn column_sizes_in_schema_order() -> VortexResult<()> { + let summary = write_nested().await?; + let sizes = summary.footer().compressed_field_sizes()?; + + assert_eq!( + summary.compressed_column_sizes()?, + vec![ + sizes.get(&field_path!(a)).unwrap_or_default(), + sizes.get(&field_path!(b)).unwrap_or_default(), + ] + ); + Ok(()) + } + + #[tokio::test] + async fn non_struct_file_reports_single_column() -> VortexResult<()> { + let mut buf = ByteBufferMut::empty(); + let summary = SESSION + .write_options() + .write( + &mut buf, + buffer![1i32, 2, 3, 4].into_array().to_array_stream(), + ) + .await?; + + let sizes = summary.footer().compressed_field_sizes()?; + assert!(sizes.total() > 0); + assert_eq!(sizes.get(&FieldPath::root()), Some(sizes.total())); + assert_eq!(summary.compressed_column_sizes()?, vec![sizes.total()]); + Ok(()) + } +} diff --git a/vortex-file/src/footer/mod.rs b/vortex-file/src/footer/mod.rs index f4bf7e19760..fda02182e1b 100644 --- a/vortex-file/src/footer/mod.rs +++ b/vortex-file/src/footer/mod.rs @@ -8,6 +8,7 @@ //! //! The byte-level footer and postscript layout is part of the file-format spec; this module exposes //! the structured Rust representation and serializer/deserializer state machine. +mod field_sizes; mod file_layout; mod file_statistics; mod postscript; @@ -19,6 +20,7 @@ mod serializer; pub use serializer::*; mod deserializer; pub use deserializer::*; +pub use field_sizes::CompressedFieldSizes; pub use file_statistics::FileStatistics; use flatbuffers::root; use itertools::Itertools; @@ -146,6 +148,15 @@ impl Footer { self.statistics.as_ref() } + /// Computes the compressed size in bytes of every field in the file, keyed by field path. + /// + /// Sizes are derived by attributing each segment in the [segment map][Self::segment_map] to a + /// field in the [layout tree][Self::layout]; see [`CompressedFieldSizes`] for the exact + /// attribution semantics. No IO is performed. + pub fn compressed_field_sizes(&self) -> VortexResult { + CompressedFieldSizes::try_new(&self.root_layout, &self.segments) + } + /// Returns the [`DType`] of the file. pub fn dtype(&self) -> &DType { self.root_layout.dtype() diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index e9c7741af93..261ba7191d6 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -113,6 +113,7 @@ mod tests; pub mod v2; mod writer; +pub use counting::CountingVortexWrite; pub use file::*; pub use footer::*; pub use forever_constant::*; diff --git a/vortex-file/src/multi/mod.rs b/vortex-file/src/multi/mod.rs index 215331f0540..587768f17d0 100644 --- a/vortex-file/src/multi/mod.rs +++ b/vortex-file/src/multi/mod.rs @@ -4,6 +4,7 @@ //! Builder for constructing a [`MultiLayoutDataSource`] from multiple Vortex files. mod session; +mod uri; use std::sync::Arc; @@ -14,6 +15,7 @@ use futures::stream; pub use session::MultiFileSession; use session::MultiFileSessionExt; use tracing::debug; +pub use uri::parse_uri_or_path; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; diff --git a/vortex-file/src/multi/uri.rs b/vortex-file/src/multi/uri.rs new file mode 100644 index 00000000000..821eb008e7f --- /dev/null +++ b/vortex-file/src/multi/uri.rs @@ -0,0 +1,203 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Parsing of user-supplied URI-or-path strings into URLs, shared by the language bindings +//! that construct a [`MultiFileDataSource`](super::MultiFileDataSource). + +#[cfg(any( + unix, + windows, + target_os = "redox", + target_os = "wasi", + target_os = "hermit" +))] +use std::path::Component; +#[cfg(any( + unix, + windows, + target_os = "redox", + target_os = "wasi", + target_os = "hermit" +))] +use std::path::Path; +#[cfg(any( + unix, + windows, + target_os = "redox", + target_os = "wasi", + target_os = "hermit" +))] +use std::path::PathBuf; +#[cfg(any( + unix, + windows, + target_os = "redox", + target_os = "wasi", + target_os = "hermit" +))] +use std::path::absolute; + +use url::Url; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +/// Parse a URI-or-path string into a [`Url`]: +/// * full URLs (`s3://...`, `file:///...`) are used as-is, +/// * bare (relative or absolute) file paths are made absolute, have `.`/`..` components +/// normalized, and become `file://` URLs. On targets without a filesystem (e.g. +/// `wasm32-unknown-unknown`), bare paths are rejected instead. +/// +/// Glob characters are preserved, so glob patterns like `/data/*.vortex` parse as expected. +pub fn parse_uri_or_path(uri_or_path: &str) -> VortexResult { + // `Url::parse` accepts Windows absolute paths like `C:\foo` as a URL with a + // single-letter scheme (`c`). No real URL scheme is one character, so treat any + // single-letter scheme as a filesystem path instead. + if let Ok(url) = Url::parse(uri_or_path) + && url.scheme().len() > 1 + { + return Ok(url); + } + file_path_to_url(uri_or_path) +} + +/// Convert a bare filesystem path into a `file://` URL, absolutizing it and normalizing +/// `.`/`..` components. The cfg predicate matches `Url::from_file_path`'s availability in +/// the `url` crate. +#[cfg(any( + unix, + windows, + target_os = "redox", + target_os = "wasi", + target_os = "hermit" +))] +fn file_path_to_url(path: &str) -> VortexResult { + let abs = + absolute(Path::new(path)).map_err(|e| vortex_err!("failed to absolutize {path}: {e}"))?; + Url::from_file_path(normalize_path(abs)) + .map_err(|_| vortex_err!("neither URL nor path: {path}")) +} + +/// `Url::from_file_path` does not exist on targets without a filesystem, such as +/// `wasm32-unknown-unknown`, so bare paths cannot be interpreted there. +#[cfg(not(any( + unix, + windows, + target_os = "redox", + target_os = "wasi", + target_os = "hermit" +)))] +fn file_path_to_url(path: &str) -> VortexResult { + Err(vortex_err!( + "bare file paths are not supported on this platform: {path}" + )) +} + +/// Normalize `.` and `..` without touching the filesystem. +#[cfg(any( + unix, + windows, + target_os = "redox", + target_os = "wasi", + target_os = "hermit" +))] +fn normalize_path(path: PathBuf) -> PathBuf { + let mut out = PathBuf::new(); + for component in path.components() { + match component { + Component::CurDir => {} + Component::ParentDir => { + out.pop(); + } + c => out.push(c), + } + } + out +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[test] + fn test_full_url() -> VortexResult<()> { + let url = parse_uri_or_path("s3://bucket/prefix/*.vortex")?; + assert_eq!(url.scheme(), "s3"); + assert_eq!(url.host_str(), Some("bucket")); + assert_eq!(url.path(), "/prefix/*.vortex"); + Ok(()) + } + + #[test] + fn test_file_scheme() -> VortexResult<()> { + let url = parse_uri_or_path("file:///absolute/path/data.vortex")?; + assert_eq!(url.scheme(), "file"); + assert_eq!(url.path(), "/absolute/path/data.vortex"); + Ok(()) + } + + #[test] + fn test_absolute_path() -> VortexResult<()> { + // Use a drive-prefixed input on Windows so `absolute()` doesn't inject the cwd drive + // and the expected URL path is predictable. The path does not need to exist. + #[cfg(unix)] + let (input, expected_path) = ("/tmp/data/*.vortex", "/tmp/data/*.vortex"); + #[cfg(windows)] + let (input, expected_path) = (r"C:\tmp\data\*.vortex", "/C:/tmp/data/*.vortex"); + let url = parse_uri_or_path(input)?; + assert_eq!(url.scheme(), "file"); + assert_eq!(url.path(), expected_path); + Ok(()) + } + + #[test] + fn test_relative_path_resolves_against_cwd() -> VortexResult<()> { + for input in ["data/table", "./data/table", "../data/table"] { + let url = parse_uri_or_path(input)?; + assert_eq!(url.scheme(), "file"); + assert!(url.path().starts_with('/')); + assert!(url.path().ends_with("/data/table")); + assert!( + !url.path().contains("%2E"), + "path must not contain percent-encoded dots" + ); + } + Ok(()) + } + + #[test] + fn test_single_letter_scheme_is_path() -> VortexResult<()> { + // Regression: `Url::parse("C:\\tmp")` succeeds with scheme="c"; the function must + // treat that as a filesystem path, not a URL. Exercised on all platforms because + // the check lives in `parse_uri_or_path`, not in an OS-specific branch. + let url = parse_uri_or_path(r"C:\tmp\data\*.vortex")?; + assert_eq!(url.scheme(), "file"); + assert_ne!(url.scheme(), "c"); + Ok(()) + } + + // Use absolute paths so the expected result is cwd-independent. + #[cfg(unix)] + #[rstest] + #[case("/a/./b", "/a/b")] + #[case("/a/b/./c", "/a/b/c")] + #[case("/a/../b", "/b")] + #[case("/a/b/../c", "/a/c")] + #[case("/a/b/../../c", "/c")] + #[case("/a/./b/.././c", "/a/c")] + #[case("/a/b/../..", "/")] + fn test_dot_normalization( + #[case] input: &str, + #[case] expected_path: &str, + ) -> VortexResult<()> { + let url = parse_uri_or_path(input)?; + assert_eq!(url.scheme(), "file"); + assert_eq!( + url.path(), + expected_path, + "input {input:?} should normalize to {expected_path:?}" + ); + Ok(()) + } +} diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index e7c9dc3b222..3af33362b05 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; +use std::future::Future; use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::AtomicUsize; @@ -12,6 +14,9 @@ use futures::FutureExt; use futures::StreamExt; use futures::channel::mpsc; use futures::future; +use futures::future::BoxFuture; +use futures::future::Shared; +use parking_lot::Mutex; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; @@ -21,6 +26,7 @@ use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_io::VortexReadAt; use vortex_io::runtime::Handle; +use vortex_io::runtime::JoinOutcome; use vortex_layout::segments::SegmentFuture; use vortex_layout::segments::SegmentId; use vortex_layout::segments::SegmentSource; @@ -67,10 +73,26 @@ pub enum ReadEvent { /// /// I/O requests will be processed in the order they are `registered`, however coalescing may mean /// other registered requests are lumped together into a single I/O operation. +/// A cloneable handle to the background read driver, shared by every in-flight [`ReadFuture`]. +/// +/// [`Shared`] fans a single completion out to all readers with correct waker bookkeeping, so a +/// reader is always woken when the driver finishes — even if another reader polled the driver more +/// recently and was then dropped. Its output is `()`; a driver panic is carried out of band in +/// [`DriverPanic`] so it can be re-raised on the reader side. +type SharedDriver = Shared>; + +/// Slot holding the driver's panic payload, if it panicked while driving reads. The first reader to +/// observe completion takes the payload and re-raises it; later readers report a graceful error. +type DriverPanic = Arc>>>; + pub struct FileSegmentSource { segments: Arc<[SegmentSpec]>, /// A queue for sending read request events to the I/O stream. events: mpsc::UnboundedSender, + /// Background request driver, joined by readers to surface a driver panic. + driver: SharedDriver, + /// Panic payload captured if the driver panicked while driving reads. + driver_panic: DriverPanic, /// The next read request ID. next_id: Arc, } @@ -144,11 +166,30 @@ impl FileSegmentSource { .await }; - handle.spawn(drive_fut).detach(); + // Spawn the driver so the runtime makes I/O progress independently of any reader. Readers + // join it (below) only to surface a panic raised while driving reads. + let mut task = handle.spawn(drive_fut); + let driver_panic: DriverPanic = Arc::new(Mutex::new(None)); + let driver = { + let driver_panic = Arc::clone(&driver_panic); + async move { + // Poll for the terminal outcome without re-raising: a benign abort (runtime + // teardown) resolves to `()` so readers report a graceful error, while a panic is + // stashed for the first reader to re-raise. + if let JoinOutcome::Panicked(panic) = future::poll_fn(|cx| task.poll_join(cx)).await + { + *driver_panic.lock() = Some(panic); + } + } + .boxed() + .shared() + }; Self { segments, events: send, + driver, + driver_panic, next_id: Arc::new(AtomicUsize::new(0)), } } @@ -192,6 +233,8 @@ impl SegmentSource for FileSegmentSource { polled: false, finished: false, events: self.events.clone(), + driver: self.driver.clone(), + driver_panic: Arc::clone(&self.driver_panic), }; // One allocation: we only box the returned SegmentFuture, not the inner ReadFuture. @@ -209,6 +252,8 @@ struct ReadFuture { polled: bool, finished: bool, events: mpsc::UnboundedSender, + driver: SharedDriver, + driver_panic: DriverPanic, } impl Future for ReadFuture { @@ -216,17 +261,28 @@ impl Future for ReadFuture { fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { match self.recv.poll_unpin(cx) { - Poll::Ready(result) => { + // note: we are skipping polled and dropped events for this if the future is ready on + // the first poll, that means this request was completed before it was polled, + // as part of a coalesced request. + Poll::Ready(Ok(result)) => { self.finished = true; - // note: we are skipping polled and dropped events for this if the future - // is ready on the first poll, that means this request was completed - // before it was polled, as part of a coalesced request. - Poll::Ready( - result.unwrap_or_else(|e| { - Err(vortex_err!("ReadRequest dropped by runtime: {e}")) - }), - ) + Poll::Ready(result) } + // The request's sender was dropped, so the driver has finished. Join it so a panic + // raised while driving reads is re-raised here rather than surfacing as a generic + // error. Only report the dropped error once the driver has finished. + Poll::Ready(Err(e)) => match self.driver.poll_unpin(cx) { + Poll::Ready(()) => { + self.finished = true; + // Re-raise the driver panic on the first reader to observe it; later readers + // fall through to the graceful dropped error. + if let Some(panic) = self.driver_panic.lock().take() { + std::panic::resume_unwind(panic); + } + Poll::Ready(Err(vortex_err!("ReadRequest dropped by runtime: {e}"))) + } + Poll::Pending => Poll::Pending, + }, Poll::Pending if !self.polled => { self.polled = true; // Notify the I/O stream that this request has been polled. @@ -321,3 +377,198 @@ impl SegmentSource for BufferSegmentSource { future::ready(Ok(BufferHandle::new_host(slice))).boxed() } } + +#[cfg(test)] +mod tests { + use std::panic::AssertUnwindSafe; + + use futures::future::BoxFuture; + use vortex_io::runtime::tokio::TokioRuntime; + use vortex_layout::segments::SegmentSource; + use vortex_metrics::DefaultMetricsRegistry; + + use super::*; + + #[derive(Clone)] + struct PanickingReadAt; + + impl VortexReadAt for PanickingReadAt { + fn concurrency(&self) -> usize { + 1 + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + async { Ok(4) }.boxed() + } + + fn read_at( + &self, + _offset: u64, + _length: usize, + _alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + async { + panic!("read-at panic"); + } + .boxed() + } + } + + fn panicking_source() -> FileSegmentSource { + let segments: Arc<[SegmentSpec]> = Arc::from([SegmentSpec { + offset: 0, + length: 4, + alignment: Alignment::none(), + }]); + let metrics = DefaultMetricsRegistry::default(); + FileSegmentSource::open( + segments, + PanickingReadAt, + TokioRuntime::current(), + RequestMetrics::new(&metrics, vec![]), + ) + } + + fn panic_message(payload: &(dyn Any + Send)) -> &str { + payload + .downcast_ref::<&str>() + .copied() + .or_else(|| payload.downcast_ref::().map(String::as_str)) + .unwrap_or("") + } + + #[tokio::test] + #[should_panic(expected = "read-at panic")] + async fn file_segment_source_propagates_read_driver_panic() { + let source = panicking_source(); + let _result = source.request(SegmentId::from(0)).await; + } + + // A read-driver panic must propagate on *every* run rather than sometimes surfacing as a + // generic "dropped by runtime" error, which is nondeterministic. + #[tokio::test] + async fn file_segment_source_read_driver_panic_propagates_deterministically() { + for i in 0..100 { + let source = panicking_source(); + let outcome = AssertUnwindSafe(source.request(SegmentId::from(0))) + .catch_unwind() + .await; + assert!( + outcome.is_err(), + "read-driver panic was not propagated on iteration {i}; \ + the request resolved instead of panicking" + ); + } + } + + // A driver panic must surface to every concurrent reader sharing one driver: exactly one + // reader re-raises the original panic, the rest report a graceful error, and none hang. This + // also covers the fan-out invariant — `Shared` wakes every reader on completion, so a reader + // dropped mid-flight can never swallow another reader's wake-up. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn file_segment_source_panic_propagates_to_all_concurrent_readers() { + let source = Arc::new(panicking_source()); + let handle = TokioRuntime::current(); + let reader_count = 8; + + let readers: Vec<_> = (0..reader_count) + .map(|_| { + let source = Arc::clone(&source); + handle.spawn(async move { + AssertUnwindSafe(source.request(SegmentId::from(0))) + .catch_unwind() + .await + }) + }) + .collect(); + + let joined = + tokio::time::timeout(std::time::Duration::from_secs(5), future::join_all(readers)) + .await + .expect("a reader hung instead of observing the driver panic"); + + let mut original_panics = 0; + for reader in joined { + match reader { + // The first reader to observe completion re-raises the original panic. + Err(payload) => { + assert!( + panic_message(&*payload).contains("read-at panic"), + "got: {:?}", + panic_message(&*payload) + ); + original_panics += 1; + } + // Every other reader reports a graceful dropped-by-runtime error. + Ok(result) => assert!(result.is_err(), "expected a dropped-by-runtime error"), + } + } + + assert_eq!( + original_panics, 1, + "exactly one reader should re-raise the original driver panic" + ); + } + + #[derive(Clone)] + struct SlowErrReadAt; + + impl VortexReadAt for SlowErrReadAt { + fn concurrency(&self) -> usize { + 4 + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + async { Ok(1024) }.boxed() + } + + fn read_at( + &self, + offset: u64, + _length: usize, + _alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + async move { + // Stagger completions so some reads finish while others are still in flight. + for _ in 0..(offset as usize % 5 + 1) { + tokio::task::yield_now().await; + } + vortex_bail!("slow read done") + } + .boxed() + } + } + + // Many segment reads driven on separate tasks must all make progress and complete. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn read_driver_concurrent_reads_make_progress() { + let n = 16u32; + let segments: Arc<[SegmentSpec]> = (0..n) + .map(|i| SegmentSpec { + offset: u64::from(i) * 4, + length: 4, + alignment: Alignment::none(), + }) + .collect(); + let metrics = DefaultMetricsRegistry::default(); + let source = Arc::new(FileSegmentSource::open( + segments, + SlowErrReadAt, + TokioRuntime::current(), + RequestMetrics::new(&metrics, vec![]), + )); + + let handle = TokioRuntime::current(); + let tasks: Vec<_> = (0..n) + .map(|i| { + let source = Arc::clone(&source); + handle.spawn(async move { source.request(SegmentId::from(i)).await }) + }) + .collect(); + + let joined = + tokio::time::timeout(std::time::Duration::from_secs(5), future::join_all(tasks)).await; + + assert!(joined.is_ok(), "concurrent reads stalled before completing"); + } +} diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index bb5d1881bbc..ba2ceac7291 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -51,9 +51,11 @@ use vortex_layout::layouts::compressed::CompressingStrategy; use vortex_layout::layouts::compressed::CompressorPlugin; use vortex_layout::layouts::dict::writer::DictStrategy; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_layout::layouts::list::writer::ListLayoutStrategy; use vortex_layout::layouts::repartition::RepartitionStrategy; use vortex_layout::layouts::repartition::RepartitionWriterOptions; use vortex_layout::layouts::table::TableStrategy; +use vortex_layout::layouts::table::use_experimental_list_layout; use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions; use vortex_layout::layouts::zoned::writer::ZonedStrategy; #[cfg(feature = "unstable_encodings")] @@ -169,6 +171,11 @@ pub struct WriteStrategyBuilder { field_writers: HashMap>, allow_encodings: Option>, flat_strategy: Option>, + probe_compressor: Option>, + /// Whether to write list fields using [`ListLayoutStrategy`]. + /// + /// [`ListLayoutStrategy`]: vortex_layout::layouts::list::writer::ListLayoutStrategy + use_list_layout: bool, } impl std::fmt::Debug for WriteStrategyBuilder { @@ -204,6 +211,8 @@ impl Default for WriteStrategyBuilder { field_writers: HashMap::new(), allow_encodings: Some(ALLOWED_ENCODINGS.clone()), flat_strategy: None, + probe_compressor: None, + use_list_layout: use_experimental_list_layout(), } } } @@ -218,6 +227,17 @@ impl WriteStrategyBuilder { self } + /// Enable writing list fields with [`ListLayoutStrategy`]. + /// + /// **Note**: this is an unstable and experimental layout that is expected to change. + /// Using it may lead to unreadable files in the future. + /// + /// [`ListLayoutStrategy`]: vortex_layout::layouts::list::writer::ListLayoutStrategy + pub fn with_list_layout(mut self) -> Self { + self.use_list_layout = true; + self + } + /// Override the write layout for a specific field somewhere in the nested schema tree. /// /// The field path is matched after the root struct is split into columns. This is useful when a @@ -267,6 +287,12 @@ impl WriteStrategyBuilder { self } + /// Override the compressor used to probe whether a column is dict-eligible. + pub fn with_probe_compressor(mut self, compressor: C) -> Self { + self.probe_compressor = Some(Arc::new(compressor)); + self + } + /// Builds the canonical [`LayoutStrategy`] implementation, with the configured overrides /// applied. pub fn build(self) -> Arc { @@ -320,22 +346,30 @@ impl WriteStrategyBuilder { CompressorConfig::BtrBlocks(builder) => Arc::new(builder.build()), CompressorConfig::Opaque(compressor) => compressor, }; - let compress_then_flat = CompressingStrategy::new(flat, stats_compressor); + let compress_then_flat = CompressingStrategy::new(flat, Arc::clone(&stats_compressor)); // 3. apply dict encoding or fallback + let probe_compressor = if let Some(probe_compressor) = self.probe_compressor { + probe_compressor + } else { + Arc::clone(&stats_compressor) + }; let dict = DictStrategy::new( coalescing.clone(), compress_then_flat.clone(), coalescing, Default::default(), + probe_compressor, ); + let row_block_size = NonZeroUsize::new(self.row_block_size).vortex_expect("must be non 0"); + // 2. calculate stats for each row group let stats = ZonedStrategy::new( dict, compress_then_flat.clone(), ZonedLayoutOptions { - block_size: NonZeroUsize::new(self.row_block_size).vortex_expect("must be non 0"), + block_size: row_block_size, ..Default::default() }, ); @@ -354,11 +388,37 @@ impl WriteStrategyBuilder { ); // 0. start with splitting columns - let validity_strategy = CollectStrategy::new(compress_then_flat); + let validity_strategy = CollectStrategy::new(compress_then_flat.clone()); // Take any field overrides from the builder and apply them to the final strategy. - let table_strategy = TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition)) - .with_field_writers(self.field_writers); + let mut table_strategy = + TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition)) + .with_field_writers(self.field_writers); + + if self.use_list_layout { + // We need a closure here to enable recursive application of list layout. + table_strategy = table_strategy.with_list_layout_factory( + move |list_layout: ListLayoutStrategy| -> Arc { + let zoned = ZonedStrategy::new( + list_layout, + compress_then_flat.clone(), + ZonedLayoutOptions { + block_size: row_block_size, + ..Default::default() + }, + ); + Arc::new(RepartitionStrategy::new( + zoned, + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: row_block_size.get(), + block_size_target: None, + canonicalize: false, + }, + )) + }, + ); + } Arc::new(table_strategy) } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index b196f68dc96..e665ab18d50 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -15,6 +15,7 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; +use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; @@ -58,6 +59,9 @@ use vortex_array::stats::PRUNING_STATS; use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; +use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::schemes::string::StringDictScheme; use vortex_buffer::Buffer; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; @@ -237,12 +241,14 @@ async fn test_read_simple_with_spawn() { vec![vec![11, 12], vec![21, 22], vec![31, 32], vec![41, 42]], Arc::new(I32.into()), ) - .unwrap(), + .unwrap() + .into_array(), ListArray::from_iter_slow::( vec![vec![51, 52], vec![61, 62], vec![71, 72], vec![81, 82]], Arc::new(I32.into()), ) - .unwrap(), + .unwrap() + .into_array(), ]) .into_array(); @@ -1666,6 +1672,75 @@ async fn test_writer_with_complex_types() -> VortexResult<()> { Ok(()) } +/// Write `array` with list decomposition forced on (through the full compress/zone pipeline) and +/// read the whole thing back. +async fn write_read_roundtrip(array: ArrayRef) -> VortexResult { + let strategy = crate::strategy::WriteStrategyBuilder::default() + .with_list_layout() + .build(); + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .with_strategy(strategy) + .write(&mut buf, array.to_array_stream()) + .await?; + SESSION + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .read_all() + .await +} + +/// A `list>` column round-trips through the `TableStrategy` dispatcher, exercising list +/// decomposition recursing into itself (the outer list's `elements` are themselves lists). +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn nested_list_of_list_roundtrip() -> VortexResult<()> { + let inner = ListArray::try_new( + buffer![1i32, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 5, 5, 6].into_array(), + Validity::NonNullable, + )? + .into_array(); + let outer = ListArray::try_new( + inner, + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let st = StructArray::from_fields(&[("nested", outer)])?.into_array(); + + let result = write_read_roundtrip(st.clone()).await?; + assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx()); + Ok(()) +} + +/// A `struct<{ items: list>? }>` column round-trips, exercising list decomposition +/// recursing into struct decomposition (list `elements` are structs) plus a nullable list validity +/// child. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn nested_struct_list_struct_roundtrip() -> VortexResult<()> { + let inner_struct = StructArray::from_fields(&[ + ("a", buffer![1i32, 2, 3, 4, 5].into_array()), + ("b", buffer![10i32, 20, 30, 40, 50].into_array()), + ])? + .into_array(); + let items = ListArray::try_new( + inner_struct, + buffer![0u32, 2, 5, 5].into_array(), + Validity::Array(BoolArray::from_iter([true, false, true]).into_array()), + )? + .into_array(); + let st = StructArray::from_fields(&[("items", items)])?.into_array(); + + let result = write_read_roundtrip(st.clone()).await?; + assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx()); + Ok(()) +} + #[tokio::test] async fn test_writer_with_statistics() -> VortexResult<()> { let array = StructArray::from_fields(&[("numbers", buffer![1u32, 2, 3, 4, 5].into_array())])? @@ -1728,15 +1803,12 @@ async fn timestamp_unit_mismatch() -> Result<(), Box> { /// Regression test: filtering a milliseconds timestamp column with a seconds scalar should /// always error, regardless of how the internal children of `DateTimePartsArray` are encoded. /// -/// This test forces `ConstantArray` encoding for the seconds/subseconds children by using a -/// compressor with Dict excluded (which triggers distinct-value computation, letting -/// `ConstantScheme` win for `[0, 0, 0]`). The scanner should still detect the time unit +/// The compressor's built-in constant detection encodes the seconds/subseconds children +/// (`[0, 0, 0]`) as `ConstantArray`s. The scanner should still detect the time unit /// mismatch and error, not silently return wrong results. #[tokio::test] async fn timestamp_unit_mismatch_errors_with_constant_children() -> Result<(), Box> { - // Build a compressor where ConstantScheme wins for [0, 0, 0] by including Dict - // (which enables distinct-value computation). let compressor = vortex_btrblocks::BtrBlocksCompressor::default(); // Write file with MILLISECONDS timestamps using this compressor. @@ -1818,6 +1890,16 @@ fn assert_offsets_ordered(before: &[u64], after: &[u64], context: &str) { } } +/// Whether any node in the layout tree is a dict layout. +fn layout_has_dict(layout: &dyn Layout) -> bool { + layout.encoding_id().as_ref() == "vortex.dict" + || layout + .children() + .unwrap() + .iter() + .any(|child| layout_has_dict(child.as_ref())) +} + /// Mirrors the (private) `IDEAL_SPLIT_SIZE` that `SplitBy::Layout` uses to sub-divide wide /// chunk-boundary spans: layout splits are never wider than this many rows. const MAX_SPLIT_ROWS: u64 = 100_000; @@ -2004,6 +2086,75 @@ async fn test_segment_ordering_dict_codes_before_values() -> VortexResult<()> { Ok(()) } +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn dict_probe_honours_configured_compressor() -> VortexResult<()> { + // Low-cardinality strings so the default cascade picks a dictionary. + let n = 32_768; + let values: Vec<&str> = (0..n).map(|i| ["alpha", "beta", "gamma"][i % 3]).collect(); + let strings = VarBinArray::from(values).into_array(); + + let mut buf = ByteBufferMut::empty(); + let summary = SESSION + .write_options() + .with_strategy(crate::strategy::WriteStrategyBuilder::default().build()) + .write(&mut buf, strings.clone().to_array_stream()) + .await?; + assert!( + layout_has_dict(summary.footer().layout().as_ref()), + "default builder should produce a dict layout for low-cardinality strings" + ); + + let no_string_dict = + BtrBlocksCompressorBuilder::default().exclude_schemes([StringDictScheme.id()]); + let mut buf = ByteBufferMut::empty(); + let summary = SESSION + .write_options() + .with_strategy( + crate::strategy::WriteStrategyBuilder::default() + .with_btrblocks_builder(no_string_dict) + .build(), + ) + .write(&mut buf, strings.to_array_stream()) + .await?; + assert!( + !layout_has_dict(summary.footer().layout().as_ref()), + "excluding StringDict from the configured compressor should disable the dict layout" + ); + + Ok(()) +} + +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn probe_compressor_override_is_independent() -> VortexResult<()> { + // Low-cardinality strings the default cascade would dict-encode. + let n = 32_768; + let values: Vec<&str> = (0..n).map(|i| ["alpha", "beta", "gamma"][i % 3]).collect(); + let strings = VarBinArray::from(values).into_array(); + + let probe_without_dict = BtrBlocksCompressorBuilder::default() + .exclude_schemes([StringDictScheme.id()]) + .build(); + + let mut buf = ByteBufferMut::empty(); + let summary = SESSION + .write_options() + .with_strategy( + crate::strategy::WriteStrategyBuilder::default() + .with_probe_compressor(probe_without_dict) + .build(), + ) + .write(&mut buf, strings.to_array_stream()) + .await?; + assert!( + !layout_has_dict(summary.footer().layout().as_ref()), + "probe override should disable the dict layout independently of the data/stats compressor" + ); + + Ok(()) +} + #[tokio::test] #[cfg_attr(miri, ignore)] async fn test_segment_ordering_zonemaps_after_data() -> VortexResult<()> { diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index eddc8f262cd..2debbfcbf88 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -19,6 +19,7 @@ use itertools::Itertools; use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::dtype::DType; +use vortex_array::dtype::FieldPath; use vortex_array::expr::stats::Stat; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorExt; @@ -515,6 +516,31 @@ impl WriteSummary { pub fn row_count(&self) -> u64 { self.footer.row_count() } + + /// Returns the compressed size in bytes of each top-level column in schema order. + /// + /// A column's size includes every physical segment attributed to its layout subtree, + /// including auxiliary segments such as zone maps and dictionaries; see + /// [`Footer::compressed_field_sizes`] for the exact attribution semantics and for sizes of + /// nested fields. Bytes not attributable to a specific column (e.g. top-level struct + /// validity) are not included in any column's size. + /// + /// For a non-struct file, the returned vector contains a single entry for the root column. + pub fn compressed_column_sizes(&self) -> VortexResult> { + let sizes = self.footer.compressed_field_sizes()?; + let Some(fields) = self.footer.dtype().as_struct_fields_opt() else { + return Ok(vec![sizes.total()]); + }; + Ok(fields + .names() + .iter() + .map(|name| { + sizes + .get(&FieldPath::from_name(name.clone())) + .unwrap_or_default() + }) + .collect()) + } } #[cfg(test)] diff --git a/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs b/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs index 7eb6cc9fb39..fb3b19513de 100644 --- a/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs +++ b/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs @@ -68,6 +68,9 @@ table Variant { } table Union { + names: [string]; + dtypes: [DType]; + type_ids: [byte]; // length must equal dtypes.len(); interpreted as unsigned nullable: bool; } diff --git a/vortex-flatbuffers/src/generated/dtype.rs b/vortex-flatbuffers/src/generated/dtype.rs index e58679fb816..56536ad0b94 100644 --- a/vortex-flatbuffers/src/generated/dtype.rs +++ b/vortex-flatbuffers/src/generated/dtype.rs @@ -1477,7 +1477,10 @@ impl<'a> ::flatbuffers::Follow<'a> for Union<'a> { } impl<'a> Union<'a> { - pub const VT_NULLABLE: ::flatbuffers::VOffsetT = 4; + pub const VT_NAMES: ::flatbuffers::VOffsetT = 4; + pub const VT_DTYPES: ::flatbuffers::VOffsetT = 6; + pub const VT_TYPE_IDS: ::flatbuffers::VOffsetT = 8; + pub const VT_NULLABLE: ::flatbuffers::VOffsetT = 10; #[inline] pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { @@ -1486,14 +1489,38 @@ impl<'a> Union<'a> { #[allow(unused_mut)] pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, - args: &'args UnionArgs + args: &'args UnionArgs<'args> ) -> ::flatbuffers::WIPOffset> { let mut builder = UnionBuilder::new(_fbb); + if let Some(x) = args.type_ids { builder.add_type_ids(x); } + if let Some(x) = args.dtypes { builder.add_dtypes(x); } + if let Some(x) = args.names { builder.add_names(x); } builder.add_nullable(args.nullable); builder.finish() } + #[inline] + pub fn names(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset<&'a str>>> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset<&'a str>>>>(Union::VT_NAMES, None)} + } + #[inline] + pub fn dtypes(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(Union::VT_DTYPES, None)} + } + #[inline] + pub fn type_ids(&self) -> Option<::flatbuffers::Vector<'a, i8>> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, i8>>>(Union::VT_TYPE_IDS, None)} + } #[inline] pub fn nullable(&self) -> bool { // Safety: @@ -1509,18 +1536,27 @@ impl ::flatbuffers::Verifiable for Union<'_> { v: &mut ::flatbuffers::Verifier, pos: usize ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { v.visit_table(pos)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset<&'_ str>>>>("names", Self::VT_NAMES, false)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("dtypes", Self::VT_DTYPES, false)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, i8>>>("type_ids", Self::VT_TYPE_IDS, false)? .visit_field::("nullable", Self::VT_NULLABLE, false)? .finish(); Ok(()) } } -pub struct UnionArgs { +pub struct UnionArgs<'a> { + pub names: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset<&'a str>>>>, + pub dtypes: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, + pub type_ids: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, i8>>>, pub nullable: bool, } -impl<'a> Default for UnionArgs { +impl<'a> Default for UnionArgs<'a> { #[inline] fn default() -> Self { UnionArgs { + names: None, + dtypes: None, + type_ids: None, nullable: false, } } @@ -1531,6 +1567,18 @@ pub struct UnionBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, } impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> UnionBuilder<'a, 'b, A> { + #[inline] + pub fn add_names(&mut self, names: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset<&'b str>>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(Union::VT_NAMES, names); + } + #[inline] + pub fn add_dtypes(&mut self, dtypes: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset>>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(Union::VT_DTYPES, dtypes); + } + #[inline] + pub fn add_type_ids(&mut self, type_ids: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , i8>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(Union::VT_TYPE_IDS, type_ids); + } #[inline] pub fn add_nullable(&mut self, nullable: bool) { self.fbb_.push_slot::(Union::VT_NULLABLE, nullable, false); @@ -1553,6 +1601,9 @@ impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> UnionBuilder<'a, 'b, A> { impl ::core::fmt::Debug for Union<'_> { fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { let mut ds = f.debug_struct("Union"); + ds.field("names", &self.names()); + ds.field("dtypes", &self.dtypes()); + ds.field("type_ids", &self.type_ids()); ds.field("nullable", &self.nullable()); ds.finish() } diff --git a/vortex-flatbuffers/src/generated/message.rs b/vortex-flatbuffers/src/generated/message.rs index b3e4be76ef7..9cc48fe6be6 100644 --- a/vortex-flatbuffers/src/generated/message.rs +++ b/vortex-flatbuffers/src/generated/message.rs @@ -2,8 +2,8 @@ // @generated extern crate alloc; -use crate::array::*; use crate::dtype::*; +use crate::array::*; #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] pub const ENUM_MIN_MESSAGE_VERSION: u8 = 0; diff --git a/vortex-geo/Cargo.toml b/vortex-geo/Cargo.toml index e2f7e4dc10f..9a7980bc631 100644 --- a/vortex-geo/Cargo.toml +++ b/vortex-geo/Cargo.toml @@ -20,14 +20,18 @@ geo = { workspace = true } geo-traits = { workspace = true } geo-types = { workspace = true } geoarrow = { workspace = true } +geoarrow-cast = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } +vortex-arrow = { workspace = true } vortex-error = { workspace = true } vortex-session = { workspace = true } wkb = { workspace = true } [dev-dependencies] rstest = { workspace = true } +vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-layout = { workspace = true } [lints] workspace = true diff --git a/vortex-geo/src/aggregate_fn/aabb.rs b/vortex-geo/src/aggregate_fn/aabb.rs new file mode 100644 index 00000000000..f217494fc1a --- /dev/null +++ b/vortex-geo/src/aggregate_fn/aabb.rs @@ -0,0 +1,464 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The 2D axis-aligned bounding-box (AABB) aggregate for native geometry columns. + +use geo::Rect as GeoRect; +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::scalar::Scalar; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::GeoMetadata; +use crate::extension::Rect; +use crate::extension::box_storage_dtype; +use crate::extension::coordinate::Dimension; +use crate::extension::flatten_coordinates; +use crate::extension::is_native_geometry; + +/// Aggregates a native geometry column's 2D axis-aligned bounding box (AABB) as a native +/// `geoarrow.box`, the spatial analogue of min/max. Also the default zone statistic for such +/// columns (via `zone_stat_default`). +#[derive(Clone, Debug)] +pub struct GeometryAabb; + +/// Running union of geometry AABBs, or `None` until the first row. A transient +/// `geo::Rect` value - the persisted stat is the native box (see `to_scalar`). +pub struct AabbPartial { + rect: Option>, +} + +impl AabbPartial { + /// Grow the accumulated box to also cover `other`. + fn merge(&mut self, other: GeoRect) { + self.rect = Some(self.rect.map_or(other, |cur| { + GeoRect::new( + ( + cur.min().x.min(other.min().x), + cur.min().y.min(other.min().y), + ), + ( + cur.max().x.max(other.max().x), + cur.max().y.max(other.max().y), + ), + ) + })); + } +} + +/// The stat's type: the native `geoarrow.box` (2D), nullable so an empty group is a null box. +fn aabb_dtype() -> DType { + DType::Extension( + ExtDType::::try_new(GeoMetadata::default(), aabb_storage_dtype()) + .vortex_expect("2D box storage is a valid Rect") + .erased(), + ) +} + +/// The `Rect` storage `Struct` backing the zone statistic. +fn aabb_storage_dtype() -> DType { + box_storage_dtype(Dimension::Xy, Nullability::Nullable) +} + +/// The AABB of the raw `x`/`y` slices, or `None` when empty. +fn aabb_of(xs: &[f64], ys: &[f64]) -> Option> { + if xs.is_empty() { + return None; + } + let min_max = |vals: &[f64]| { + vals.iter() + .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), &v| { + (lo.min(v), hi.max(v)) + }) + }; + let (xmin, xmax) = min_max(xs); + let (ymin, ymax) = min_max(ys); + Some(GeoRect::new((xmin, ymin), (xmax, ymax))) +} + +/// Read an AABB stat scalar (a nullable native `geoarrow.box`) into a [`GeoRect`], or `None` when +/// the scalar is null (an empty group). +fn rect_from_storage(scalar: &Scalar) -> VortexResult>> { + if scalar.is_null() { + return Ok(None); + } + let storage = scalar.as_extension().to_storage_scalar(); + let fields = storage.as_struct(); + let read = |name: &str| -> VortexResult { + f64::try_from( + &fields + .field(name) + .ok_or_else(|| vortex_err!("AABB missing {name}"))?, + ) + }; + Ok(Some(GeoRect::new( + (read("xmin")?, read("ymin")?), + (read("xmax")?, read("ymax")?), + ))) +} + +/// Serialize a [`GeoRect`] as a native `geoarrow.box` stat scalar (inverse of [`rect_from_storage`]). +fn rect_to_storage(rect: GeoRect) -> Scalar { + let storage = Scalar::struct_( + aabb_storage_dtype(), + vec![ + Scalar::primitive(rect.min().x, Nullability::NonNullable), + Scalar::primitive(rect.min().y, Nullability::NonNullable), + Scalar::primitive(rect.max().x, Nullability::NonNullable), + Scalar::primitive(rect.max().y, Nullability::NonNullable), + ], + ); + Scalar::extension::(GeoMetadata::default(), storage) +} + +impl AggregateFnVTable for GeometryAabb { + type Options = EmptyOptions; + type Partial = AabbPartial; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.geo.aabb"); + *ID + } + + // Serializable so the zoned writer can persist this as a per-chunk stat. No options to encode. + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) + } + + fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { + is_native_geometry(input_dtype).then(aabb_dtype) + } + + fn zone_stat_default(&self, input_dtype: &DType) -> Option { + is_native_geometry(input_dtype).then(|| self.bind(EmptyOptions)) + } + + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + self.return_dtype(options, input_dtype) + } + + fn empty_partial( + &self, + _options: &Self::Options, + _input_dtype: &DType, + ) -> VortexResult { + Ok(AabbPartial { rect: None }) + } + + fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { + if let Some(rect) = rect_from_storage(&other)? { + partial.merge(rect); + } + Ok(()) + } + + fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + Ok(match partial.rect { + Some(rect) => rect_to_storage(rect), + None => Scalar::null(aabb_dtype()), + }) + } + + fn reset(&self, partial: &mut Self::Partial) { + partial.rect = None; + } + + fn is_saturated(&self, _partial: &Self::Partial) -> bool { + // An AABB can always grow, so it is never saturated. + false + } + + fn accumulate( + &self, + partial: &mut Self::Partial, + batch: &Columnar, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let array = match batch { + Columnar::Canonical(canonical) => canonical.clone().into_array(), + Columnar::Constant(constant) => constant.clone().into_array(), + }; + // Min/max the raw x/y buffers directly - cheap, and avoids `to_geometry`'s panic on empty + // points (which decoding each geometry would hit). + let coords = flatten_coordinates(&array, ctx)?; + let xs = coords + .unmasked_field_by_name("x")? + .clone() + .execute::(ctx)?; + let ys = coords + .unmasked_field_by_name("y")? + .clone() + .execute::(ctx)?; + if let Some(rect) = aabb_of(xs.as_slice::(), ys.as_slice::()) { + partial.merge(rect); + } + Ok(()) + } + + fn finalize(&self, partials: ArrayRef) -> VortexResult { + // The stored partial is already the AABB struct, so finalizing is the identity. + Ok(partials) + } + + fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + self.to_scalar(partial) + } +} + +#[cfg(test)] +mod tests { + use geo::Rect as GeoRect; + use vortex_array::ArrayRef; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::aggregate_fn::session::AggregateFnSessionExt; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_error::VortexResult; + + use super::AabbPartial; + use super::GeometryAabb; + use super::aabb_dtype; + use super::rect_from_storage; + use crate::test_harness::geo_session; + use crate::test_harness::linestring_column; + use crate::test_harness::multilinestring_column; + use crate::test_harness::multipoint_column; + use crate::test_harness::multipolygon_column; + use crate::test_harness::point_column; + use crate::test_harness::polygon_column; + + /// One column of every native geometry type over the same `(x, y)` vertex set. + fn every_native_column(vertices: &[(f64, f64)]) -> VortexResult> { + let (xs, ys): (Vec, Vec) = vertices.iter().copied().unzip(); + let flat = vertices.to_vec(); + Ok(vec![ + point_column(xs, ys)?, + linestring_column(vec![flat.clone()])?, + multipoint_column(vec![flat.clone()])?, + polygon_column(vec![vec![flat.clone()]])?, + multilinestring_column(vec![vec![flat.clone()]])?, + multipolygon_column(vec![vec![vec![flat]]])?, + ]) + } + + /// The aggregate must be serializable so the zoned writer can persist its zone-stat descriptor. + #[test] + fn serializes_for_zone_storage() -> VortexResult<()> { + let session = vortex_array::array_session(); + let metadata = GeometryAabb + .serialize(&EmptyOptions)? + .expect("GeometryAabb must be serializable to be stored as a zone statistic"); + GeometryAabb.deserialize(&metadata, &session)?; + Ok(()) + } + + /// The AABB result's corners as `(xmin, ymin, xmax, ymax)`. + fn aabb(result: &Scalar) -> VortexResult<(f64, f64, f64, f64)> { + let rect = rect_from_storage(result)?.expect("non-null AABB"); + Ok((rect.min().x, rect.min().y, rect.max().x, rect.max().y)) + } + + /// The AABB of a Point column is the min/max of its coordinates, accumulated across batches. + #[test] + fn point_aabb_across_batches() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let mut acc = Accumulator::try_new(GeometryAabb, EmptyOptions, dtype)?; + + acc.accumulate(&point_column(vec![1.0, 3.0], vec![2.0, 4.0])?, &mut ctx)?; + acc.accumulate(&point_column(vec![-1.0], vec![5.0])?, &mut ctx)?; + + assert_eq!(aabb(&acc.finish()?)?, (-1.0, 2.0, 3.0, 5.0)); + Ok(()) + } + + /// The AABB of a Polygon column unions every ring vertex - exercising the `List>` + /// unwrap, not just the bare Point struct. + #[test] + fn polygon_aabb_union_all_vertices() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + // Two rectangles: (0,0)-(2,3) and (5,5)-(7,8). The chunk AABB is their union: (0,0)-(7,8). + let polygons = polygon_column(vec![ + vec![vec![(0.0, 0.0), (2.0, 0.0), (2.0, 3.0), (0.0, 3.0)]], + vec![vec![(5.0, 5.0), (7.0, 5.0), (7.0, 8.0), (5.0, 8.0)]], + ])?; + let dtype = polygons.dtype().clone(); + let mut acc = Accumulator::try_new(GeometryAabb, EmptyOptions, dtype)?; + acc.accumulate(&polygons, &mut ctx)?; + + assert_eq!(aabb(&acc.finish()?)?, (0.0, 0.0, 7.0, 8.0)); + Ok(()) + } + + /// Every native geometry type over the same vertex set yields the same AABB - the zone stat + /// covers the whole type family. + #[test] + fn aabb_covers_every_native_geometry_type() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + for column in every_native_column(&[(1.0, 2.0), (-1.0, 5.0), (3.0, 4.0)])? { + let mut acc = Accumulator::try_new(GeometryAabb, EmptyOptions, column.dtype().clone())?; + acc.accumulate(&column, &mut ctx)?; + assert_eq!( + aabb(&acc.finish()?)?, + (-1.0, 2.0, 3.0, 5.0), + "AABB mismatch for {}", + column.dtype() + ); + } + Ok(()) + } + + /// The AABB of a MultiPolygon column unions every vertex of every polygon's rings - exercising + /// the triple-`List` unwrap. + #[test] + fn multipolygon_aabb_union_all_vertices() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + // Multipolygon 0: squares (0,0)-(1,1) and (4,4)-(5,5); multipolygon 1: square (-3,7)-(-2,9). + let multipolygons = multipolygon_column(vec![ + vec![ + vec![vec![(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 1.0)]], + vec![vec![(4.0, 4.0), (5.0, 4.0), (5.0, 5.0), (4.0, 5.0)]], + ], + vec![vec![vec![ + (-3.0, 7.0), + (-2.0, 7.0), + (-2.0, 9.0), + (-3.0, 9.0), + ]]], + ])?; + let dtype = multipolygons.dtype().clone(); + let mut acc = Accumulator::try_new(GeometryAabb, EmptyOptions, dtype)?; + acc.accumulate(&multipolygons, &mut ctx)?; + + assert_eq!(aabb(&acc.finish()?)?, (-3.0, 0.0, 5.0, 9.0)); + Ok(()) + } + + /// `combine_partials` unions partial boxes - the path the zoned writer takes when a zone's + /// array is chunked. + #[test] + fn combine_partials_unions_boxes() -> VortexResult<()> { + let bbox = |xmin, ymin, xmax, ymax| AabbPartial { + rect: Some(GeoRect::new((xmin, ymin), (xmax, ymax))), + }; + let mut partial = AabbPartial { rect: None }; + GeometryAabb.combine_partials( + &mut partial, + GeometryAabb.to_scalar(&bbox(0.0, 0.0, 1.0, 1.0))?, + )?; + GeometryAabb.combine_partials( + &mut partial, + GeometryAabb.to_scalar(&bbox(5.0, -2.0, 7.0, 3.0))?, + )?; + assert_eq!( + aabb(&GeometryAabb.to_scalar(&partial)?)?, + (0.0, -2.0, 7.0, 3.0) + ); + Ok(()) + } + + /// A null partial (an empty group's AABB) is a no-op in `combine_partials`. + #[test] + fn combine_partials_ignores_null() -> VortexResult<()> { + let mut partial = AabbPartial { + rect: Some(GeoRect::new((0.0, 0.0), (1.0, 1.0))), + }; + GeometryAabb.combine_partials(&mut partial, Scalar::null(aabb_dtype()))?; + assert_eq!( + aabb(&GeometryAabb.to_scalar(&partial)?)?, + (0.0, 0.0, 1.0, 1.0) + ); + Ok(()) + } + + /// All-NaN coordinates: `f64::min`/`max` skip the NaNs and `geo::Rect` normalizes the result to + /// a valid (whole-plane) box, so such a chunk is always kept. Sound - NaN-coordinate rows can + /// never satisfy `distance <= r` anyway. + #[test] + fn all_nan_coordinates_kept() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let column = point_column(vec![f64::NAN, f64::NAN], vec![f64::NAN, f64::NAN])?; + let mut acc = Accumulator::try_new(GeometryAabb, EmptyOptions, column.dtype().clone())?; + acc.accumulate(&column, &mut ctx)?; + + let (xmin, ymin, xmax, ymax) = aabb(&acc.finish()?)?; + assert!(xmin <= xmax && ymin <= ymax); + Ok(()) + } + + /// An empty group yields a null AABB. + #[test] + fn empty_group_is_null() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let mut acc = Accumulator::try_new(GeometryAabb, EmptyOptions, dtype)?; + assert!(acc.finish()?.is_null()); + Ok(()) + } + + /// After `initialize`, the registry yields a default zone statistic for every native geometry + /// type (so the zoned writer stores it) but none for ordinary numeric columns. + #[test] + fn registered_as_geometry_zone_default() -> VortexResult<()> { + let session = geo_session(); + + for column in every_native_column(&[(0.0, 0.0), (1.0, 1.0)])? { + assert!( + !session + .aggregate_fns() + .zone_stat_defaults(column.dtype()) + .is_empty(), + "a geometry zone-stat default should be discovered for {}", + column.dtype() + ); + } + let i32_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + assert!( + session + .aggregate_fns() + .zone_stat_defaults(&i32_dtype) + .is_empty(), + "no geometry zone-stat default should apply to numeric columns" + ); + Ok(()) + } +} diff --git a/vortex-array/src/arrow/executor/validity.rs b/vortex-geo/src/aggregate_fn/mod.rs similarity index 57% rename from vortex-array/src/arrow/executor/validity.rs rename to vortex-geo/src/aggregate_fn/mod.rs index fa848dfc79a..70831f02905 100644 --- a/vortex-array/src/arrow/executor/validity.rs +++ b/vortex-geo/src/aggregate_fn/mod.rs @@ -1,4 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -pub(super) use crate::arrow::null_buffer::to_arrow_null_buffer; +//! Aggregate functions over geometry columns. + +mod aabb; + +pub use aabb::*; diff --git a/vortex-geo/src/extension/linestring.rs b/vortex-geo/src/extension/linestring.rs new file mode 100644 index 00000000000..274d20f28ba --- /dev/null +++ b/vortex-geo/src/extension/linestring.rs @@ -0,0 +1,341 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`LineString`] geometry extension type (`vortex.geo.linestring`): an ordered path of the +//! [`Point`](super::Point) coordinate struct, stored as `List>` and tagged +//! with [`GeoMetadata`] (CRS). + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType; +use geo_traits::to_geo::ToGeoGeometry; +use geo_types::Geometry; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::IntoArrow; +use geoarrow::array::LineStringArray; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::LineStringType; +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtId; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::scalar::ScalarValue; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; +use vortex_arrow::FromArrowType; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; +use vortex_session::registry::Id; + +use super::GeoMetadata; +use super::coordinate::Dimension; +use super::coordinate::coordinate_dimension; +use super::coordinate::coordinate_storage_dtype; +use super::geo_metadata_from_arrow; +use super::geoarrow_metadata; +use super::geoarrow_to_wkb; + +/// A line string: `geoarrow.linestring`, stored as `List>` (an ordered path +/// of vertices). +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct LineString; + +impl ExtVTable for LineString { + type Metadata = GeoMetadata; + // No cheap owned value like Point's `Coordinate`; expose the raw storage scalar. + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> ExtId { + static ID: CachedId = CachedId::new("vortex.geo.linestring"); + *ID + } + + fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { + Ok(metadata.encode_to_vec()) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + Ok(GeoMetadata::decode(metadata)?) + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + linestring_dimension(ext_dtype.storage_dtype()).map(|_| ()) + } + + fn unpack_native<'a>( + _ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult<&'a ScalarValue> { + Ok(storage_value) + } +} + +/// Canonical line-string storage: a list of the coordinate `Struct`. +pub(crate) fn linestring_storage_dtype(dim: Dimension, nullability: Nullability) -> DType { + let coords = coordinate_storage_dtype(dim, Nullability::NonNullable); + DType::List(Arc::new(coords), nullability) +} + +/// Validate `dtype` is `List` and return its [`Dimension`]. +pub(crate) fn linestring_dimension(dtype: &DType) -> VortexResult { + let DType::List(coords, _) = dtype else { + vortex_bail!("linestring storage must be a List of coordinates, was {dtype}"); + }; + coordinate_dimension(coords) +} + +static ARROW_LINESTRING: CachedId = CachedId::new(LineStringType::NAME); + +/// The `geoarrow.linestring` extension type for `dimension`, with separated (struct) coordinates +/// matching `LineString` storage. +fn linestring_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> LineStringType { + LineStringType::new(dimension.into(), geoarrow_metadata(geo_metadata)) +} + +/// Decode `LineString` storage (`List`) to `geo_types` line strings, for the geo scalar +/// functions. CRS does not affect planar geometry ops, so default metadata is used. +pub(crate) fn linestring_geometries( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + linestring_array(storage, ctx)? + .iter() + .map(|geometry| -> VortexResult> { + Ok(geometry + .ok_or_else(|| vortex_err!("geo: null geometry is not supported"))? + .map_err(|e| vortex_err!("geo: geometry access failed: {e}"))? + .to_geometry()) + }) + .collect() +} + +/// Build a geoarrow `LineStringArray` from a `LineString`'s `List` storage. +fn linestring_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let linestring_type = linestring_type( + &GeoMetadata::default(), + linestring_dimension(storage.dtype())?, + ); + let session = ctx.session().clone(); + let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + LineStringArray::try_from((arrow.as_ref(), linestring_type)) + .map_err(|e| vortex_err!("failed to construct LineStringArray: {e}")) +} + +/// A validated `LineString` array (`try_from` checks the extension type). +pub struct LineStringData(ExtensionArray); + +impl TryFrom for LineStringData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a LineString extension array" + ); + Ok(LineStringData(ext)) + } +} + +impl LineStringData { + /// Serialize line strings to WKB (a view array) — the form DuckDB `GEOMETRY` takes. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + geoarrow_to_wkb(&linestring_array(self.0.storage_array(), ctx)?) + } +} + +impl ArrowExportVTable for LineString { + fn arrow_ext_id(&self) -> Id { + *ARROW_LINESTRING + } + + fn vortex_id(&self) -> Id { + self.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let ext_type = dtype.as_extension(); + let geo_metadata = ext_type.metadata::(); + let dimension = linestring_dimension(ext_type.storage_dtype())?; + + let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + field.try_with_extension_type(linestring_type(geo_metadata, dimension))?; + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let is_linestring = array + .dtype() + .as_extension_opt() + .map(|ext| ext.is::()) + .unwrap_or(false); + if !is_linestring { + return Ok(ArrowExport::Unsupported(array)); + } + + let Ok(linestring_meta) = target.try_extension_type::() else { + return Ok(ArrowExport::Unsupported(array)); + }; + if linestring_meta.coord_type() != CoordType::Separated { + return Ok(ArrowExport::Unsupported(array)); + } + + let executed = array.execute::(ctx)?; + let storage = executed.storage_array().clone(); + + let storage_field = Field::new( + String::new(), + target.data_type().clone(), + target.is_nullable(), + ); + let session = ctx.session().clone(); + let arrow_storage = session + .arrow() + .execute_arrow(storage, Some(&storage_field), ctx)?; + + // Round-trip through GeoArrow's line-string array; `into_arrow` is concrete, so wrap in `Arc`. + let linestrings = LineStringArray::try_from((arrow_storage.as_ref(), linestring_meta)) + .map_err(|e| vortex_err!("failed to construct LineStringArray: {e}"))?; + + Ok(ArrowExport::Exported(Arc::new(linestrings.into_arrow()))) + } +} + +impl ArrowImportVTable for LineString { + fn arrow_ext_id(&self) -> Id { + *ARROW_LINESTRING + } + + /// Import a `geoarrow.linestring` field as the [`LineString`] dtype. Keyed off the standard + /// GeoArrow name, so any producer resolves here. Accepts the full `LineStringType` extension, or + /// — for a metadata-less geometry literal — the name alone, inferring the dimension from the + /// coordinate field names. + fn from_arrow_field(&self, field: &Field) -> VortexResult> { + let (dimension, metadata) = + if let Ok(linestring_meta) = field.try_extension_type::() { + vortex_ensure!( + linestring_meta.coord_type() == CoordType::Separated, + "geoarrow.linestring with interleaved coordinates is not supported; \ + re-encode with separated (struct) coordinates" + ); + ( + linestring_meta.dimension().into(), + geo_metadata_from_arrow(linestring_meta.metadata()), + ) + } else { + // Literal: peel the `List` layer to the coordinate struct and read its dimension from + // the field names (the canonical check rejects nullable coordinates). + if field.extension_type_name() != Some(LineStringType::NAME) { + return Ok(None); + } + let DType::List(coords, _) = DType::from_arrow(field) else { + return Ok(None); + }; + let DType::Struct(fields, _) = coords.as_ref() else { + return Ok(None); + }; + let Ok(dimension) = Dimension::from_field_names(fields.names()) else { + return Ok(None); + }; + (dimension, GeoMetadata::default()) + }; + + let storage_dtype = linestring_storage_dtype(dimension, field.is_nullable().into()); + Ok(Some(DType::Extension( + ExtDType::try_with_vtable(LineString, metadata, storage_dtype)?.erased(), + ))) + } + + fn from_arrow_array( + &self, + array: ArrowArrayRef, + field: &Field, + dtype: &DType, + ) -> VortexResult { + let Some(ext_dtype) = dtype.as_extension_opt() else { + return Ok(ArrowImport::Unsupported(array)); + }; + if !ext_dtype.is::() + || field.try_extension_type::().is_err() + || !matches!(array.data_type(), DataType::List(_)) + { + return Ok(ArrowImport::Unsupported(array)); + } + + let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + Ok(ArrowImport::Imported( + ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + )) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_error::VortexResult; + + use super::LineString; + use super::linestring_storage_dtype; + use crate::extension::GeoMetadata; + use crate::extension::coordinate::Dimension; + + fn geo_meta() -> GeoMetadata { + GeoMetadata { + crs: Some("EPSG:4326".to_string()), + } + } + + /// `LineString` accepts the canonical `List` storage of every dimension. + #[rstest] + #[case::xy(Dimension::Xy)] + #[case::xyz(Dimension::Xyz)] + #[case::xym(Dimension::Xym)] + #[case::xyzm(Dimension::Xyzm)] + fn linestring_validates_every_dimension(#[case] dim: Dimension) -> VortexResult<()> { + let storage = linestring_storage_dtype(dim, Nullability::NonNullable); + ExtDType::::try_new(geo_meta(), storage)?; + Ok(()) + } + + /// Non-linestring storage is rejected at dtype construction: a bare coordinate struct (point) is + /// not a list of coordinates. + #[test] + fn linestring_rejects_invalid_storage() -> VortexResult<()> { + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), primitive).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index 684c83bade0..4bd8ae3800f 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -2,30 +2,119 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors pub(crate) mod coordinate; +mod linestring; +mod multilinestring; +mod multipoint; +mod multipolygon; mod point; mod polygon; +mod rect; mod wkb; use std::fmt::Display; use std::sync::Arc; +use ::wkb::reader::GeometryType; +use arrow_array::BinaryArray; use geo_types::Geometry; +use geoarrow::array::GenericWkbArray; +use geoarrow::array::GeoArrowArray; +use geoarrow::datatypes::CoordType; use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Dimension; +use geoarrow::datatypes::GeoArrowType; +use geoarrow::datatypes::LineStringType; use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::MultiLineStringType; +use geoarrow::datatypes::MultiPointType; +use geoarrow::datatypes::MultiPolygonType; +use geoarrow::datatypes::PointType; +use geoarrow::datatypes::PolygonType; +use geoarrow::datatypes::WkbType; +use geoarrow_cast::cast::cast; +pub use linestring::*; +pub use multilinestring::*; +pub use multipoint::*; +pub use multipolygon::*; pub use point::*; pub use polygon::*; +pub use rect::*; use vortex_array::ArrayRef; +use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::StructArray; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::Scalar; +use vortex_arrow::FromArrowArray; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; pub use wkb::*; +/// Whether `dtype` is one of the native geometry extension types the geo kernels operate on. +pub(crate) fn is_native_geometry(dtype: &DType) -> bool { + dtype.as_extension_opt().is_some_and(|ext| { + ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + }) +} + +/// Validate the operands of a geo scalar function: each must be a native geometry type (so the +/// kernel can decode it) and non-nullable (geometry arrays never carry nulls). +pub(crate) fn validate_geometry_operands(dtypes: &[DType]) -> VortexResult<()> { + for dtype in dtypes { + vortex_ensure!( + is_native_geometry(dtype), + "geo: operand {dtype} is not a native geometry type" + ); + vortex_ensure!( + !dtype.is_nullable(), + "geo: nullable operand {dtype} is unsupported" + ); + } + Ok(()) +} + +/// Flatten a native geometry column into a single coordinate `Struct` containing +/// every vertex of every geometry. +pub(crate) fn flatten_coordinates( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if !is_native_geometry(array.dtype()) { + vortex_bail!( + "geo: operand is not a native geometry extension type, was {}", + array.dtype() + ); + } + let mut node = array + .clone() + .execute::(ctx)? + .storage_array() + .clone(); + while matches!(node.dtype(), DType::List(..)) { + node = node + .execute::(ctx)? + .into_listview() + .elements() + .clone(); + } + node.execute::(ctx) +} + /// Decode a native geometry column to `geo_types`. A non-geometry operand is an error. pub(crate) fn geometries( array: &ArrayRef, @@ -44,8 +133,18 @@ pub(crate) fn geometries( .clone(); if ext.is::() { point_geometries(&storage, ctx) + } else if ext.is::() { + linestring_geometries(&storage, ctx) + } else if ext.is::() { + multipoint_geometries(&storage, ctx) } else if ext.is::() { polygon_geometries(&storage, ctx) + } else if ext.is::() { + multilinestring_geometries(&storage, ctx) + } else if ext.is::() { + multipolygon_geometries(&storage, ctx) + } else if ext.is::() { + rect_geometries(&storage, ctx) } else { vortex_bail!("geo: unsupported geometry extension {}", array.dtype()) } @@ -63,6 +162,82 @@ pub(crate) fn single_geometry( .ok_or_else(|| vortex_err!("geo: constant operand decoded to no geometry")) } +/// Decode a WKB geometry literal (DuckDB's wire form for `GEOMETRY` constants) to its native +/// `Point`/`Polygon`/`MultiPolygon` scalar. `None` for unsupported types. Plan-time, one value only. +pub fn native_geometry_scalar_from_wkb(bytes: &[u8]) -> VortexResult> { + let metadata = geoarrow_metadata(&GeoMetadata::default()); + let binary = BinaryArray::from(vec![Some(bytes)]); + let wkb = GenericWkbArray::::try_from(( + &binary as &dyn arrow_array::Array, + WkbType::new(Arc::clone(&metadata)), + )) + .map_err(|e| vortex_err!("failed to read WKB literal: {e}"))?; + + // Cast the WKB value to `target`, import its native storage as a Vortex array. + let to_storage = |target: &GeoArrowType| -> VortexResult { + let native = + cast(&wkb, target).map_err(|e| vortex_err!("failed to cast WKB literal: {e}"))?; + ArrayRef::from_arrow(native.to_array_ref().as_ref(), false) + }; + + let scalar = match Wkb::try_from_bytes(bytes)?.geometry_type() { + GeometryType::Point => { + let target = GeoArrowType::Point( + PointType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), + ); + geo_ext_scalar(Point, to_storage(&target)?)? + } + GeometryType::LineString => { + let target = GeoArrowType::LineString( + LineStringType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), + ); + geo_ext_scalar(LineString, to_storage(&target)?)? + } + GeometryType::Polygon => { + let target = GeoArrowType::Polygon( + PolygonType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), + ); + geo_ext_scalar(Polygon, to_storage(&target)?)? + } + GeometryType::MultiPoint => { + let target = GeoArrowType::MultiPoint( + MultiPointType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), + ); + geo_ext_scalar(MultiPoint, to_storage(&target)?)? + } + GeometryType::MultiLineString => { + let target = GeoArrowType::MultiLineString( + MultiLineStringType::new(Dimension::XY, metadata) + .with_coord_type(CoordType::Separated), + ); + geo_ext_scalar(MultiLineString, to_storage(&target)?)? + } + GeometryType::MultiPolygon => { + let target = GeoArrowType::MultiPolygon( + MultiPolygonType::new(Dimension::XY, metadata) + .with_coord_type(CoordType::Separated), + ); + geo_ext_scalar(MultiPolygon, to_storage(&target)?)? + } + _ => return Ok(None), + }; + Ok(Some(scalar)) +} + +/// Wrap cast-from-WKB `storage` in its `vtable` extension type and pull out the single scalar. +// `scalar_at` is deprecated for `execute_scalar`, but there is no execution context at plan time. +#[allow(deprecated)] +fn geo_ext_scalar>( + vtable: V, + storage: ArrayRef, +) -> VortexResult { + let ext = ExtDType::try_with_vtable(vtable, GeoMetadata::default(), storage.dtype().clone())? + .erased(); + ExtensionArray::try_new(ext, storage)? + .into_array() + .scalar_at(0) +} + /// Extension metadata that is common to all the geospatial extension types. /// /// Currently, this is just the coordinate reference system (CRS). @@ -95,6 +270,15 @@ pub(crate) fn geoarrow_metadata(geo_metadata: &GeoMetadata) -> Arc { )) } +/// Serialize a native geometry array to WKB (a `WkbView` array) via geoarrow's cast. +/// Shared by the `to_wkb` methods on the geometry extension types. +pub(crate) fn geoarrow_to_wkb(geo_array: &dyn GeoArrowArray) -> VortexResult { + let wkb_type = GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&GeoMetadata::default()))); + let wkb = cast(geo_array, &wkb_type) + .map_err(|e| vortex_err!("failed to cast geometry to WKB: {e}"))?; + ArrayRef::from_arrow(wkb.to_array_ref().as_ref(), false) +} + /// Recover [`GeoMetadata`] from GeoArrow metadata. pub(crate) fn geo_metadata_from_arrow(metadata: &Metadata) -> GeoMetadata { let crs = metadata.crs().crs_value().map(|value| { @@ -112,7 +296,16 @@ pub(crate) fn geo_metadata_from_arrow(metadata: &Metadata) -> GeoMetadata { #[cfg(test)] mod tests { use prost::Message; + use vortex_array::dtype::DType; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use super::LineString; + use super::MultiLineString; + use super::MultiPoint; + use super::Point; + use super::Polygon; + use super::native_geometry_scalar_from_wkb; use crate::extension::GeoMetadata; #[test] @@ -127,4 +320,115 @@ mod tests { let decoded = GeoMetadata::decode(bytes.as_slice()).unwrap(); assert_eq!(decoded, meta); } + + /// A little-endian WKB `POINT` literal decodes to the native `Point` extension scalar. + #[test] + fn decodes_wkb_point_to_native() -> VortexResult<()> { + let mut wkb = vec![1u8]; // little-endian byte order + wkb.extend_from_slice(&1u32.to_le_bytes()); // geometry type: point + wkb.extend_from_slice(&1.0f64.to_le_bytes()); // x + wkb.extend_from_slice(&2.0f64.to_le_bytes()); // y + + let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a point scalar"); + let DType::Extension(ext) = scalar.dtype() else { + panic!("expected an extension dtype, got {}", scalar.dtype()); + }; + assert!(ext.is::()); + Ok(()) + } + + /// A little-endian WKB `POLYGON` literal decodes to the native `Polygon` extension scalar. + #[test] + fn decodes_wkb_polygon_to_native() -> VortexResult<()> { + let ring = [(0.0, 0.0), (1.0, 0.0), (0.0, 1.0), (0.0, 0.0)]; + let mut wkb = vec![1u8]; // little-endian byte order + wkb.extend_from_slice(&3u32.to_le_bytes()); // geometry type: polygon + wkb.extend_from_slice(&1u32.to_le_bytes()); // one ring + let ring_len = u32::try_from(ring.len()).map_err(|e| vortex_err!("{e}"))?; + wkb.extend_from_slice(&ring_len.to_le_bytes()); + for (x, y) in ring { + wkb.extend_from_slice(&f64::to_le_bytes(x)); + wkb.extend_from_slice(&f64::to_le_bytes(y)); + } + + let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a polygon scalar"); + let DType::Extension(ext) = scalar.dtype() else { + panic!("expected an extension dtype, got {}", scalar.dtype()); + }; + assert!(ext.is::()); + Ok(()) + } + + /// A little-endian WKB `LINESTRING` literal decodes to the native `LineString` extension scalar. + #[test] + fn decodes_wkb_linestring_to_native() -> VortexResult<()> { + let points = [(0.0, 0.0), (1.0, 1.0)]; + let mut wkb = vec![1u8]; // little-endian byte order + wkb.extend_from_slice(&2u32.to_le_bytes()); // geometry type: linestring + let len = u32::try_from(points.len()).map_err(|e| vortex_err!("{e}"))?; + wkb.extend_from_slice(&len.to_le_bytes()); + for (x, y) in points { + wkb.extend_from_slice(&f64::to_le_bytes(x)); + wkb.extend_from_slice(&f64::to_le_bytes(y)); + } + + let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a linestring scalar"); + let DType::Extension(ext) = scalar.dtype() else { + panic!("expected an extension dtype, got {}", scalar.dtype()); + }; + assert!(ext.is::()); + Ok(()) + } + + /// A little-endian WKB `MULTIPOINT` literal decodes to the native `MultiPoint` extension scalar. + #[test] + fn decodes_wkb_multipoint_to_native() -> VortexResult<()> { + let points = [(0.0, 0.0), (1.0, 1.0)]; + let mut wkb = vec![1u8]; // little-endian byte order + wkb.extend_from_slice(&4u32.to_le_bytes()); // geometry type: multipoint + let len = u32::try_from(points.len()).map_err(|e| vortex_err!("{e}"))?; + wkb.extend_from_slice(&len.to_le_bytes()); + for (x, y) in points { + // each member is a full WKB point + wkb.push(1u8); + wkb.extend_from_slice(&1u32.to_le_bytes()); + wkb.extend_from_slice(&f64::to_le_bytes(x)); + wkb.extend_from_slice(&f64::to_le_bytes(y)); + } + + let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a multipoint scalar"); + let DType::Extension(ext) = scalar.dtype() else { + panic!("expected an extension dtype, got {}", scalar.dtype()); + }; + assert!(ext.is::()); + Ok(()) + } + + /// A little-endian WKB `MULTILINESTRING` literal decodes to the native `MultiLineString` scalar. + #[test] + fn decodes_wkb_multilinestring_to_native() -> VortexResult<()> { + let lines = [[(0.0, 0.0), (1.0, 1.0)], [(2.0, 2.0), (3.0, 3.0)]]; + let mut wkb = vec![1u8]; // little-endian byte order + wkb.extend_from_slice(&5u32.to_le_bytes()); // geometry type: multilinestring + let num_lines = u32::try_from(lines.len()).map_err(|e| vortex_err!("{e}"))?; + wkb.extend_from_slice(&num_lines.to_le_bytes()); + for line in lines { + // each member is a full WKB linestring + wkb.push(1u8); + wkb.extend_from_slice(&2u32.to_le_bytes()); + let len = u32::try_from(line.len()).map_err(|e| vortex_err!("{e}"))?; + wkb.extend_from_slice(&len.to_le_bytes()); + for (x, y) in line { + wkb.extend_from_slice(&f64::to_le_bytes(x)); + wkb.extend_from_slice(&f64::to_le_bytes(y)); + } + } + + let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a multilinestring scalar"); + let DType::Extension(ext) = scalar.dtype() else { + panic!("expected an extension dtype, got {}", scalar.dtype()); + }; + assert!(ext.is::()); + Ok(()) + } } diff --git a/vortex-geo/src/extension/multilinestring.rs b/vortex-geo/src/extension/multilinestring.rs new file mode 100644 index 00000000000..6ba4c4c5324 --- /dev/null +++ b/vortex-geo/src/extension/multilinestring.rs @@ -0,0 +1,359 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`MultiLineString`] extension type (`vortex.geo.multilinestring`), stored as +//! `List>>` (line strings → coordinates) and tagged with +//! [`GeoMetadata`]. The storage layout matches [`Polygon`](super::Polygon); the two are +//! distinguished by their GeoArrow extension name, not their shape. + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType; +use geo_traits::to_geo::ToGeoGeometry; +use geo_types::Geometry; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::IntoArrow; +use geoarrow::array::MultiLineStringArray; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::MultiLineStringType; +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtId; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::scalar::ScalarValue; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; +use vortex_arrow::FromArrowType; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; +use vortex_session::registry::Id; + +use super::GeoMetadata; +use super::coordinate::Dimension; +use super::coordinate::coordinate_dimension; +use super::coordinate::coordinate_storage_dtype; +use super::geo_metadata_from_arrow; +use super::geoarrow_metadata; +use super::geoarrow_to_wkb; + +/// A multilinestring: `geoarrow.multilinestring`, stored as `List>>` +/// (line strings of vertices). +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct MultiLineString; + +impl ExtVTable for MultiLineString { + type Metadata = GeoMetadata; + // No cheap owned value like Point's `Coordinate`; expose the raw storage scalar. + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> ExtId { + static ID: CachedId = CachedId::new("vortex.geo.multilinestring"); + *ID + } + + fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { + Ok(metadata.encode_to_vec()) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + Ok(GeoMetadata::decode(metadata)?) + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + multilinestring_dimension(ext_dtype.storage_dtype()).map(|_| ()) + } + + fn unpack_native<'a>( + _ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult<&'a ScalarValue> { + Ok(storage_value) + } +} + +/// Canonical multilinestring storage: an outer list of line strings, each a list of the coordinate +/// `Struct`. +pub(crate) fn multilinestring_storage_dtype(dim: Dimension, nullability: Nullability) -> DType { + let coords = coordinate_storage_dtype(dim, Nullability::NonNullable); + let line = DType::List(Arc::new(coords), Nullability::NonNullable); + DType::List(Arc::new(line), nullability) +} + +/// Validate `dtype` is `List>` and return its [`Dimension`]. +pub(crate) fn multilinestring_dimension(dtype: &DType) -> VortexResult { + let DType::List(line, _) = dtype else { + vortex_bail!("multilinestring storage must be a List of line strings, was {dtype}"); + }; + let DType::List(coords, _) = line.as_ref() else { + vortex_bail!("multilinestring line storage must be a List of coordinates, was {line}"); + }; + coordinate_dimension(coords) +} + +static ARROW_MULTILINESTRING: CachedId = CachedId::new(MultiLineStringType::NAME); + +/// The `geoarrow.multilinestring` type for `dimension`, with separated (struct) coordinates. +fn multilinestring_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> MultiLineStringType { + MultiLineStringType::new(dimension.into(), geoarrow_metadata(geo_metadata)) +} + +/// Decode storage to `geo_types` for the geo scalar functions (CRS is irrelevant to planar ops). +pub(crate) fn multilinestring_geometries( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + multilinestring_array(storage, ctx)? + .iter() + .map(|geometry| -> VortexResult> { + Ok(geometry + .ok_or_else(|| vortex_err!("geo: null geometry is not supported"))? + .map_err(|e| vortex_err!("geo: geometry access failed: {e}"))? + .to_geometry()) + }) + .collect() +} + +/// Build a geoarrow `MultiLineStringArray` from the `MultiLineString` storage. +fn multilinestring_array( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let multilinestring_type = multilinestring_type( + &GeoMetadata::default(), + multilinestring_dimension(storage.dtype())?, + ); + let session = ctx.session().clone(); + let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + MultiLineStringArray::try_from((arrow.as_ref(), multilinestring_type)) + .map_err(|e| vortex_err!("failed to construct MultiLineStringArray: {e}")) +} + +/// A validated `MultiLineString` array (`try_from` checks the extension type). +pub struct MultiLineStringData(ExtensionArray); + +impl TryFrom for MultiLineStringData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a MultiLineString extension array" + ); + Ok(MultiLineStringData(ext)) + } +} + +impl MultiLineStringData { + /// Serialize multilinestrings to WKB (a view array) — the form DuckDB `GEOMETRY` takes. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + geoarrow_to_wkb(&multilinestring_array(self.0.storage_array(), ctx)?) + } +} + +impl ArrowExportVTable for MultiLineString { + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTILINESTRING + } + + fn vortex_id(&self) -> Id { + self.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let ext_type = dtype.as_extension(); + let geo_metadata = ext_type.metadata::(); + let dimension = multilinestring_dimension(ext_type.storage_dtype())?; + + let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + field.try_with_extension_type(multilinestring_type(geo_metadata, dimension))?; + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let is_multilinestring = array + .dtype() + .as_extension_opt() + .map(|ext| ext.is::()) + .unwrap_or(false); + if !is_multilinestring { + return Ok(ArrowExport::Unsupported(array)); + } + + let Ok(multilinestring_meta) = target.try_extension_type::() else { + return Ok(ArrowExport::Unsupported(array)); + }; + if multilinestring_meta.coord_type() != CoordType::Separated { + return Ok(ArrowExport::Unsupported(array)); + } + + let executed = array.execute::(ctx)?; + let storage = executed.storage_array().clone(); + + let storage_field = Field::new( + String::new(), + target.data_type().clone(), + target.is_nullable(), + ); + let session = ctx.session().clone(); + let arrow_storage = session + .arrow() + .execute_arrow(storage, Some(&storage_field), ctx)?; + + let multilinestrings = + MultiLineStringArray::try_from((arrow_storage.as_ref(), multilinestring_meta)) + .map_err(|e| vortex_err!("failed to construct MultiLineStringArray: {e}"))?; + + Ok(ArrowExport::Exported(Arc::new( + multilinestrings.into_arrow(), + ))) + } +} + +impl ArrowImportVTable for MultiLineString { + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTILINESTRING + } + + /// Import a `geoarrow.multilinestring` field (matched by GeoArrow name). Accepts the full + /// `MultiLineStringType`, or a metadata-less literal (name only), inferring the dimension. + fn from_arrow_field(&self, field: &Field) -> VortexResult> { + let (dimension, metadata) = + if let Ok(multilinestring_meta) = field.try_extension_type::() { + vortex_ensure!( + multilinestring_meta.coord_type() == CoordType::Separated, + "geoarrow.multilinestring with interleaved coordinates is not supported; \ + re-encode with separated (struct) coordinates" + ); + ( + multilinestring_meta.dimension().into(), + geo_metadata_from_arrow(multilinestring_meta.metadata()), + ) + } else { + // Literal: peel the two `List` layers to the coordinate struct and read its dimension + // from the field names (the canonical check rejects nullable coordinates). + if field.extension_type_name() != Some(MultiLineStringType::NAME) { + return Ok(None); + } + let DType::List(line, _) = DType::from_arrow(field) else { + return Ok(None); + }; + let DType::List(coords, _) = line.as_ref() else { + return Ok(None); + }; + let DType::Struct(fields, _) = coords.as_ref() else { + return Ok(None); + }; + let Ok(dimension) = Dimension::from_field_names(fields.names()) else { + return Ok(None); + }; + (dimension, GeoMetadata::default()) + }; + + let storage_dtype = multilinestring_storage_dtype(dimension, field.is_nullable().into()); + Ok(Some(DType::Extension( + ExtDType::try_with_vtable(MultiLineString, metadata, storage_dtype)?.erased(), + ))) + } + + fn from_arrow_array( + &self, + array: ArrowArrayRef, + field: &Field, + dtype: &DType, + ) -> VortexResult { + let Some(ext_dtype) = dtype.as_extension_opt() else { + return Ok(ArrowImport::Unsupported(array)); + }; + if !ext_dtype.is::() + || field.try_extension_type::().is_err() + || !matches!(array.data_type(), DataType::List(_)) + { + return Ok(ArrowImport::Unsupported(array)); + } + + let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + Ok(ArrowImport::Imported( + ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + )) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_error::VortexResult; + + use super::MultiLineString; + use super::multilinestring_storage_dtype; + use crate::extension::GeoMetadata; + use crate::extension::coordinate::Dimension; + use crate::extension::coordinate::coordinate_storage_dtype; + + fn geo_meta() -> GeoMetadata { + GeoMetadata { + crs: Some("EPSG:4326".to_string()), + } + } + + /// `MultiLineString` accepts the canonical `List>` storage of every + /// dimension. + #[rstest] + #[case::xy(Dimension::Xy)] + #[case::xyz(Dimension::Xyz)] + #[case::xym(Dimension::Xym)] + #[case::xyzm(Dimension::Xyzm)] + fn multilinestring_validates_every_dimension(#[case] dim: Dimension) -> VortexResult<()> { + let storage = multilinestring_storage_dtype(dim, Nullability::NonNullable); + ExtDType::::try_new(geo_meta(), storage)?; + Ok(()) + } + + /// Non-multilinestring storage is rejected: a bare coordinate struct (point) fails. + #[test] + fn multilinestring_rejects_invalid_storage() -> VortexResult<()> { + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), primitive).is_err()); + + // A bare list of coordinates is a single line string, not a multilinestring. + let coords = coordinate_storage_dtype(Dimension::Xy, Nullability::NonNullable); + let line = DType::List(Arc::new(coords), Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), line).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/extension/multipoint.rs b/vortex-geo/src/extension/multipoint.rs new file mode 100644 index 00000000000..7d8e429dcc0 --- /dev/null +++ b/vortex-geo/src/extension/multipoint.rs @@ -0,0 +1,336 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`MultiPoint`] geometry extension type (`vortex.geo.multipoint`): an unordered set of the +//! [`Point`](super::Point) coordinate struct, stored as `List>` and tagged +//! with [`GeoMetadata`] (CRS). The storage layout matches [`LineString`](super::LineString); the +//! two are distinguished by their GeoArrow extension name, not their shape. + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType; +use geo_traits::to_geo::ToGeoGeometry; +use geo_types::Geometry; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::IntoArrow; +use geoarrow::array::MultiPointArray; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::MultiPointType; +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtId; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::scalar::ScalarValue; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; +use vortex_arrow::FromArrowType; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; +use vortex_session::registry::Id; + +use super::GeoMetadata; +use super::coordinate::Dimension; +use super::coordinate::coordinate_dimension; +use super::coordinate::coordinate_storage_dtype; +use super::geo_metadata_from_arrow; +use super::geoarrow_metadata; +use super::geoarrow_to_wkb; + +/// A multipoint: `geoarrow.multipoint`, stored as `List>` (a set of points). +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct MultiPoint; + +impl ExtVTable for MultiPoint { + type Metadata = GeoMetadata; + // No cheap owned value like Point's `Coordinate`; expose the raw storage scalar. + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> ExtId { + static ID: CachedId = CachedId::new("vortex.geo.multipoint"); + *ID + } + + fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { + Ok(metadata.encode_to_vec()) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + Ok(GeoMetadata::decode(metadata)?) + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + multipoint_dimension(ext_dtype.storage_dtype()).map(|_| ()) + } + + fn unpack_native<'a>( + _ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult<&'a ScalarValue> { + Ok(storage_value) + } +} + +/// Canonical multipoint storage: a list of the coordinate `Struct`. +pub(crate) fn multipoint_storage_dtype(dim: Dimension, nullability: Nullability) -> DType { + let coords = coordinate_storage_dtype(dim, Nullability::NonNullable); + DType::List(Arc::new(coords), nullability) +} + +/// Validate `dtype` is `List` and return its [`Dimension`]. +pub(crate) fn multipoint_dimension(dtype: &DType) -> VortexResult { + let DType::List(coords, _) = dtype else { + vortex_bail!("multipoint storage must be a List of coordinates, was {dtype}"); + }; + coordinate_dimension(coords) +} + +static ARROW_MULTIPOINT: CachedId = CachedId::new(MultiPointType::NAME); + +/// The `geoarrow.multipoint` extension type for `dimension`, with separated (struct) coordinates. +fn multipoint_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> MultiPointType { + MultiPointType::new(dimension.into(), geoarrow_metadata(geo_metadata)) +} + +/// Decode `MultiPoint` storage (`List`) to `geo_types`, for the geo scalar functions. +pub(crate) fn multipoint_geometries( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + multipoint_array(storage, ctx)? + .iter() + .map(|geometry| -> VortexResult> { + Ok(geometry + .ok_or_else(|| vortex_err!("geo: null geometry is not supported"))? + .map_err(|e| vortex_err!("geo: geometry access failed: {e}"))? + .to_geometry()) + }) + .collect() +} + +/// Build a geoarrow `MultiPointArray` from a `MultiPoint`'s `List` storage. +fn multipoint_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let multipoint_type = multipoint_type( + &GeoMetadata::default(), + multipoint_dimension(storage.dtype())?, + ); + let session = ctx.session().clone(); + let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + MultiPointArray::try_from((arrow.as_ref(), multipoint_type)) + .map_err(|e| vortex_err!("failed to construct MultiPointArray: {e}")) +} + +/// A validated `MultiPoint` array (`try_from` checks the extension type). +pub struct MultiPointData(ExtensionArray); + +impl TryFrom for MultiPointData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a MultiPoint extension array" + ); + Ok(MultiPointData(ext)) + } +} + +impl MultiPointData { + /// Serialize multipoints to WKB (a view array) — the form DuckDB `GEOMETRY` takes. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + geoarrow_to_wkb(&multipoint_array(self.0.storage_array(), ctx)?) + } +} + +impl ArrowExportVTable for MultiPoint { + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTIPOINT + } + + fn vortex_id(&self) -> Id { + self.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let ext_type = dtype.as_extension(); + let geo_metadata = ext_type.metadata::(); + let dimension = multipoint_dimension(ext_type.storage_dtype())?; + + let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + field.try_with_extension_type(multipoint_type(geo_metadata, dimension))?; + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let is_multipoint = array + .dtype() + .as_extension_opt() + .map(|ext| ext.is::()) + .unwrap_or(false); + if !is_multipoint { + return Ok(ArrowExport::Unsupported(array)); + } + + let Ok(multipoint_meta) = target.try_extension_type::() else { + return Ok(ArrowExport::Unsupported(array)); + }; + if multipoint_meta.coord_type() != CoordType::Separated { + return Ok(ArrowExport::Unsupported(array)); + } + + let executed = array.execute::(ctx)?; + let storage = executed.storage_array().clone(); + + let storage_field = Field::new( + String::new(), + target.data_type().clone(), + target.is_nullable(), + ); + let session = ctx.session().clone(); + let arrow_storage = session + .arrow() + .execute_arrow(storage, Some(&storage_field), ctx)?; + + // Round-trip through GeoArrow's multipoint array; `into_arrow` is concrete, so wrap in `Arc`. + let multipoints = MultiPointArray::try_from((arrow_storage.as_ref(), multipoint_meta)) + .map_err(|e| vortex_err!("failed to construct MultiPointArray: {e}"))?; + + Ok(ArrowExport::Exported(Arc::new(multipoints.into_arrow()))) + } +} + +impl ArrowImportVTable for MultiPoint { + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTIPOINT + } + + /// Import a `geoarrow.multipoint` field as the [`MultiPoint`] dtype (matched by GeoArrow name). + /// Accepts the full `MultiPointType`, or a metadata-less literal (name only), inferring the + /// dimension from the coordinate field names. + fn from_arrow_field(&self, field: &Field) -> VortexResult> { + let (dimension, metadata) = + if let Ok(multipoint_meta) = field.try_extension_type::() { + vortex_ensure!( + multipoint_meta.coord_type() == CoordType::Separated, + "geoarrow.multipoint with interleaved coordinates is not supported; \ + re-encode with separated (struct) coordinates" + ); + ( + multipoint_meta.dimension().into(), + geo_metadata_from_arrow(multipoint_meta.metadata()), + ) + } else { + if field.extension_type_name() != Some(MultiPointType::NAME) { + return Ok(None); + } + let DType::List(coords, _) = DType::from_arrow(field) else { + return Ok(None); + }; + let DType::Struct(fields, _) = coords.as_ref() else { + return Ok(None); + }; + let Ok(dimension) = Dimension::from_field_names(fields.names()) else { + return Ok(None); + }; + (dimension, GeoMetadata::default()) + }; + + let storage_dtype = multipoint_storage_dtype(dimension, field.is_nullable().into()); + Ok(Some(DType::Extension( + ExtDType::try_with_vtable(MultiPoint, metadata, storage_dtype)?.erased(), + ))) + } + + fn from_arrow_array( + &self, + array: ArrowArrayRef, + field: &Field, + dtype: &DType, + ) -> VortexResult { + let Some(ext_dtype) = dtype.as_extension_opt() else { + return Ok(ArrowImport::Unsupported(array)); + }; + if !ext_dtype.is::() + || field.try_extension_type::().is_err() + || !matches!(array.data_type(), DataType::List(_)) + { + return Ok(ArrowImport::Unsupported(array)); + } + + let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + Ok(ArrowImport::Imported( + ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + )) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_error::VortexResult; + + use super::MultiPoint; + use super::multipoint_storage_dtype; + use crate::extension::GeoMetadata; + use crate::extension::coordinate::Dimension; + + fn geo_meta() -> GeoMetadata { + GeoMetadata { + crs: Some("EPSG:4326".to_string()), + } + } + + /// `MultiPoint` accepts the canonical `List` storage of every dimension. + #[rstest] + #[case::xy(Dimension::Xy)] + #[case::xyz(Dimension::Xyz)] + #[case::xym(Dimension::Xym)] + #[case::xyzm(Dimension::Xyzm)] + fn multipoint_validates_every_dimension(#[case] dim: Dimension) -> VortexResult<()> { + let storage = multipoint_storage_dtype(dim, Nullability::NonNullable); + ExtDType::::try_new(geo_meta(), storage)?; + Ok(()) + } + + /// Non-multipoint storage is rejected at dtype construction: a bare coordinate struct (point) is + /// not a list of coordinates. + #[test] + fn multipoint_rejects_invalid_storage() -> VortexResult<()> { + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), primitive).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/extension/multipolygon.rs b/vortex-geo/src/extension/multipolygon.rs new file mode 100644 index 00000000000..524e470749c --- /dev/null +++ b/vortex-geo/src/extension/multipolygon.rs @@ -0,0 +1,363 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`MultiPolygon`] extension type (`vortex.geo.multipolygon`), stored as +//! `List>>>` (polygons → rings → coordinates) and tagged with +//! [`GeoMetadata`]. A single `Polygon` is a one-element multipolygon. + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType; +use geo_traits::to_geo::ToGeoGeometry; +use geo_types::Geometry; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::IntoArrow; +use geoarrow::array::MultiPolygonArray; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::MultiPolygonType; +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtId; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::scalar::ScalarValue; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; +use vortex_arrow::FromArrowType; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; +use vortex_session::registry::Id; + +use super::GeoMetadata; +use super::coordinate::Dimension; +use super::coordinate::coordinate_dimension; +use super::coordinate::coordinate_storage_dtype; +use super::geo_metadata_from_arrow; +use super::geoarrow_metadata; +use super::geoarrow_to_wkb; + +/// A multipolygon (`geoarrow.multipolygon`); a single `Polygon` is a one-element multipolygon. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct MultiPolygon; + +impl ExtVTable for MultiPolygon { + type Metadata = GeoMetadata; + // No cheap owned value like Point's `Coordinate`; expose the raw storage scalar. + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> ExtId { + static ID: CachedId = CachedId::new("vortex.geo.multipolygon"); + *ID + } + + fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { + Ok(metadata.encode_to_vec()) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + Ok(GeoMetadata::decode(metadata)?) + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + multipolygon_dimension(ext_dtype.storage_dtype()).map(|_| ()) + } + + fn unpack_native<'a>( + _ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult<&'a ScalarValue> { + Ok(storage_value) + } +} + +/// Storage `List>>`: polygons → rings → coordinates. +pub(crate) fn multipolygon_storage_dtype(dim: Dimension, nullability: Nullability) -> DType { + let coords = coordinate_storage_dtype(dim, Nullability::NonNullable); + let ring = DType::List(Arc::new(coords), Nullability::NonNullable); + let polygon = DType::List(Arc::new(ring), Nullability::NonNullable); + DType::List(Arc::new(polygon), nullability) +} + +/// Validate `dtype` is `List>>` and return its [`Dimension`]. +pub(crate) fn multipolygon_dimension(dtype: &DType) -> VortexResult { + let DType::List(polygon, _) = dtype else { + vortex_bail!("multipolygon storage must be a List of polygons, was {dtype}"); + }; + let DType::List(ring, _) = polygon.as_ref() else { + vortex_bail!("multipolygon polygon storage must be a List of rings, was {polygon}"); + }; + let DType::List(coords, _) = ring.as_ref() else { + vortex_bail!("multipolygon ring storage must be a List of coordinates, was {ring}"); + }; + coordinate_dimension(coords) +} + +static ARROW_MULTIPOLYGON: CachedId = CachedId::new(MultiPolygonType::NAME); + +/// The `geoarrow.multipolygon` type for `dimension`, with separated (struct) coordinates. +fn multipolygon_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> MultiPolygonType { + MultiPolygonType::new(dimension.into(), geoarrow_metadata(geo_metadata)) +} + +/// Decode storage to `geo_types` for the geo scalar functions (CRS is irrelevant to planar ops). +pub(crate) fn multipolygon_geometries( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + multipolygon_array(storage, ctx)? + .iter() + .map(|geometry| -> VortexResult> { + Ok(geometry + .ok_or_else(|| vortex_err!("geo: null geometry is not supported"))? + .map_err(|e| vortex_err!("geo: geometry access failed: {e}"))? + .to_geometry()) + }) + .collect() +} + +/// Build a geoarrow `MultiPolygonArray` from the `MultiPolygon` storage. +fn multipolygon_array( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let multipolygon_type = multipolygon_type( + &GeoMetadata::default(), + multipolygon_dimension(storage.dtype())?, + ); + let session = ctx.session().clone(); + let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + MultiPolygonArray::try_from((arrow.as_ref(), multipolygon_type)) + .map_err(|e| vortex_err!("failed to construct MultiPolygonArray: {e}")) +} + +/// A validated `MultiPolygon` array (`try_from` checks the extension type). +pub struct MultiPolygonData(ExtensionArray); + +impl TryFrom for MultiPolygonData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a MultiPolygon extension array" + ); + Ok(MultiPolygonData(ext)) + } +} + +impl MultiPolygonData { + /// Serialize multipolygons to WKB (a view array) — the form DuckDB `GEOMETRY` takes. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + geoarrow_to_wkb(&multipolygon_array(self.0.storage_array(), ctx)?) + } +} + +impl ArrowExportVTable for MultiPolygon { + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTIPOLYGON + } + + fn vortex_id(&self) -> Id { + self.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let ext_type = dtype.as_extension(); + let geo_metadata = ext_type.metadata::(); + let dimension = multipolygon_dimension(ext_type.storage_dtype())?; + + let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + field.try_with_extension_type(multipolygon_type(geo_metadata, dimension))?; + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let is_multipolygon = array + .dtype() + .as_extension_opt() + .map(|ext| ext.is::()) + .unwrap_or(false); + if !is_multipolygon { + return Ok(ArrowExport::Unsupported(array)); + } + + let Ok(multipolygon_meta) = target.try_extension_type::() else { + return Ok(ArrowExport::Unsupported(array)); + }; + if multipolygon_meta.coord_type() != CoordType::Separated { + return Ok(ArrowExport::Unsupported(array)); + } + + let executed = array.execute::(ctx)?; + let storage = executed.storage_array().clone(); + + let storage_field = Field::new( + String::new(), + target.data_type().clone(), + target.is_nullable(), + ); + let session = ctx.session().clone(); + let arrow_storage = session + .arrow() + .execute_arrow(storage, Some(&storage_field), ctx)?; + + let multipolygons = + MultiPolygonArray::try_from((arrow_storage.as_ref(), multipolygon_meta)) + .map_err(|e| vortex_err!("failed to construct MultiPolygonArray: {e}"))?; + + Ok(ArrowExport::Exported(Arc::new(multipolygons.into_arrow()))) + } +} + +impl ArrowImportVTable for MultiPolygon { + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTIPOLYGON + } + + /// Import a `geoarrow.multipolygon` field (matched by GeoArrow name). Accepts the full + /// `MultiPolygonType`, or a metadata-less literal (name only), inferring the dimension. + fn from_arrow_field(&self, field: &Field) -> VortexResult> { + let (dimension, metadata) = + if let Ok(multipolygon_meta) = field.try_extension_type::() { + vortex_ensure!( + multipolygon_meta.coord_type() == CoordType::Separated, + "geoarrow.multipolygon with interleaved coordinates is not supported; \ + re-encode with separated (struct) coordinates" + ); + ( + multipolygon_meta.dimension().into(), + geo_metadata_from_arrow(multipolygon_meta.metadata()), + ) + } else { + // Literal: peel the three `List` layers to the coordinate struct and read its + // dimension from the field names (the canonical check rejects nullable coordinates). + if field.extension_type_name() != Some(MultiPolygonType::NAME) { + return Ok(None); + } + let DType::List(polygon, _) = DType::from_arrow(field) else { + return Ok(None); + }; + let DType::List(ring, _) = polygon.as_ref() else { + return Ok(None); + }; + let DType::List(coords, _) = ring.as_ref() else { + return Ok(None); + }; + let DType::Struct(fields, _) = coords.as_ref() else { + return Ok(None); + }; + let Ok(dimension) = Dimension::from_field_names(fields.names()) else { + return Ok(None); + }; + (dimension, GeoMetadata::default()) + }; + + let storage_dtype = multipolygon_storage_dtype(dimension, field.is_nullable().into()); + Ok(Some(DType::Extension( + ExtDType::try_with_vtable(MultiPolygon, metadata, storage_dtype)?.erased(), + ))) + } + + fn from_arrow_array( + &self, + array: ArrowArrayRef, + field: &Field, + dtype: &DType, + ) -> VortexResult { + let Some(ext_dtype) = dtype.as_extension_opt() else { + return Ok(ArrowImport::Unsupported(array)); + }; + if !ext_dtype.is::() + || field.try_extension_type::().is_err() + || !matches!(array.data_type(), DataType::List(_)) + { + return Ok(ArrowImport::Unsupported(array)); + } + + let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + Ok(ArrowImport::Imported( + ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + )) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_error::VortexResult; + + use super::MultiPolygon; + use super::multipolygon_storage_dtype; + use crate::extension::GeoMetadata; + use crate::extension::coordinate::Dimension; + use crate::extension::coordinate::coordinate_storage_dtype; + + fn geo_meta() -> GeoMetadata { + GeoMetadata { + crs: Some("EPSG:4326".to_string()), + } + } + + /// `MultiPolygon` accepts the canonical `List>>` storage of every + /// dimension. + #[rstest] + #[case::xy(Dimension::Xy)] + #[case::xyz(Dimension::Xyz)] + #[case::xym(Dimension::Xym)] + #[case::xyzm(Dimension::Xyzm)] + fn multipolygon_validates_every_dimension(#[case] dim: Dimension) -> VortexResult<()> { + let storage = multipolygon_storage_dtype(dim, Nullability::NonNullable); + ExtDType::::try_new(geo_meta(), storage)?; + Ok(()) + } + + /// Non-multipolygon storage is rejected at dtype construction: a bare struct (point) and a + /// double list (polygon) both fail. + #[test] + fn multipolygon_rejects_invalid_storage() -> VortexResult<()> { + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), primitive).is_err()); + + // A double list (polygon) is not a multipolygon. + let coords = coordinate_storage_dtype(Dimension::Xy, Nullability::NonNullable); + let ring = DType::List(Arc::new(coords), Nullability::NonNullable); + let polygon = DType::List(Arc::new(ring), Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), polygon).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/extension/point.rs b/vortex-geo/src/extension/point.rs index 389da97ad2b..cfebecee461 100644 --- a/vortex-geo/src/extension/point.rs +++ b/vortex-geo/src/extension/point.rs @@ -23,20 +23,21 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::arrow::ArrowExport; -use vortex_array::arrow::ArrowExportVTable; -use vortex_array::arrow::ArrowImport; -use vortex_array::arrow::ArrowImportVTable; -use vortex_array::arrow::ArrowSession; -use vortex_array::arrow::ArrowSessionExt; -use vortex_array::arrow::FromArrowArray; use vortex_array::dtype::DType; -use vortex_array::dtype::arrow::FromArrowType; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; +use vortex_arrow::FromArrowType; +use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -51,6 +52,7 @@ use super::coordinate::coordinate_from_struct; use super::coordinate::coordinate_storage_dtype; use super::geo_metadata_from_arrow; use super::geoarrow_metadata; +use super::geoarrow_to_wkb; /// A single location: `geoarrow.point`, stored as `Struct` of non-nullable `f64`. #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] @@ -97,20 +99,46 @@ fn point_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> PointType { PointType::new(dimension.into(), geoarrow_metadata(geo_metadata)) } -/// Decode `Point` storage to `geo_types` points, for the geo scalar functions. -pub(crate) fn point_geometries( - storage: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult>> { +pub struct PointData(ExtensionArray); + +impl TryFrom for PointData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a Point extension array" + ); + Ok(PointData(ext)) + } +} + +impl PointData { + /// Serialize points to WKB (a view array) — the form DuckDB `GEOMETRY` takes. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + geoarrow_to_wkb(&point_array(self.0.storage_array(), ctx)?) + } +} + +/// Build a geoarrow `PointArray` from a `Point`'s `Struct` storage, shared by WKB export +/// and `geo_types` decoding. +fn point_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let point_type = point_type( &GeoMetadata::default(), coordinate_dimension(storage.dtype())?, ); let session = ctx.session().clone(); let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; - let points = PointArray::try_from((arrow.as_ref(), point_type)) - .map_err(|e| vortex_err!("failed to construct PointArray: {e}"))?; - points + PointArray::try_from((arrow.as_ref(), point_type)) + .map_err(|e| vortex_err!("failed to construct PointArray: {e}")) +} + +/// Decode `Point` storage to `geo_types` points, for the geo scalar functions. +pub(crate) fn point_geometries( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + point_array(storage, ctx)? .iter() .map(|geometry| -> VortexResult> { Ok(geometry diff --git a/vortex-geo/src/extension/polygon.rs b/vortex-geo/src/extension/polygon.rs index 83ec6601158..9a74c3ce7b3 100644 --- a/vortex-geo/src/extension/polygon.rs +++ b/vortex-geo/src/extension/polygon.rs @@ -24,20 +24,21 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::arrow::ArrowExport; -use vortex_array::arrow::ArrowExportVTable; -use vortex_array::arrow::ArrowImport; -use vortex_array::arrow::ArrowImportVTable; -use vortex_array::arrow::ArrowSession; -use vortex_array::arrow::ArrowSessionExt; -use vortex_array::arrow::FromArrowArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::dtype::arrow::FromArrowType; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; +use vortex_arrow::FromArrowType; +use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -51,6 +52,7 @@ use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; use super::geo_metadata_from_arrow; use super::geoarrow_metadata; +use super::geoarrow_to_wkb; /// A polygon: `geoarrow.polygon`, stored as `List>>` (rings of vertices). #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] @@ -118,12 +120,7 @@ pub(crate) fn polygon_geometries( storage: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult>> { - let polygon_type = polygon_type(&GeoMetadata::default(), polygon_dimension(storage.dtype())?); - let session = ctx.session().clone(); - let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; - let polygons = PolygonArray::try_from((arrow.as_ref(), polygon_type)) - .map_err(|e| vortex_err!("failed to construct PolygonArray: {e}"))?; - polygons + polygon_array(storage, ctx)? .iter() .map(|geometry| -> VortexResult> { Ok(geometry @@ -134,6 +131,37 @@ pub(crate) fn polygon_geometries( .collect() } +/// Build a geoarrow `PolygonArray` from a `Polygon`'s `List>` storage. +fn polygon_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let polygon_type = polygon_type(&GeoMetadata::default(), polygon_dimension(storage.dtype())?); + let session = ctx.session().clone(); + let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + PolygonArray::try_from((arrow.as_ref(), polygon_type)) + .map_err(|e| vortex_err!("failed to construct PolygonArray: {e}")) +} + +/// A validated `Polygon` array (`try_from` checks the extension type). +pub struct PolygonData(ExtensionArray); + +impl TryFrom for PolygonData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a Polygon extension array" + ); + Ok(PolygonData(ext)) + } +} + +impl PolygonData { + /// Serialize polygons to WKB (a view array) — the form DuckDB `GEOMETRY` takes. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + geoarrow_to_wkb(&polygon_array(self.0.storage_array(), ctx)?) + } +} + impl ArrowExportVTable for Polygon { fn arrow_ext_id(&self) -> Id { *ARROW_POLYGON diff --git a/vortex-geo/src/extension/rect.rs b/vortex-geo/src/extension/rect.rs new file mode 100644 index 00000000000..3576966368c --- /dev/null +++ b/vortex-geo/src/extension/rect.rs @@ -0,0 +1,392 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`Rect`] bounding-box extension type (`vortex.geo.box`): an axis-aligned envelope stored +//! columnarly as `Struct` of non-nullable +//! `f64` — the lower corner's ordinates followed by the upper corner's — tagged with +//! [`GeoMetadata`] (CRS). Its GeoArrow wire type is `geoarrow.box`. +//! +//! Decoding to `geo_types` yields a 2D [`Geometry::Rect`]; any `z`/`m` +//! bounds are dropped, as for the other geometry types. + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType; +use geo_traits::to_geo::ToGeoGeometry; +use geo_types::Geometry; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::IntoArrow; +use geoarrow::array::RectArray; +use geoarrow::datatypes::BoxType; +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtId; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::scalar::ScalarValue; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; +use vortex_arrow::FromArrowType; +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 vortex_session::registry::Id; + +use super::GeoMetadata; +use super::coordinate::Dimension; +use super::geo_metadata_from_arrow; +use super::geoarrow_metadata; + +/// An axis-aligned bounding box (`geoarrow.box`), stored as `Struct`. +// Named `Rect`, not `Box`: matches `geo::Rect` / geoarrow-rs `RectArray`, and `Box` is a std name. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct Rect; + +impl ExtVTable for Rect { + type Metadata = GeoMetadata; + // No cheap owned value; expose the raw storage scalar like `Polygon`. + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> ExtId { + static ID: CachedId = CachedId::new("vortex.geo.box"); + *ID + } + + fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { + Ok(metadata.encode_to_vec()) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + Ok(GeoMetadata::decode(metadata)?) + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + box_dimension(ext_dtype.storage_dtype()).map(|_| ()) + } + + fn unpack_native<'a>( + _ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult<&'a ScalarValue> { + Ok(storage_value) + } +} + +/// The `geoarrow.box` storage field names for `dim`: the lower corner's ordinates followed by the +/// upper corner's. +fn box_field_names(dim: Dimension) -> &'static [&'static str] { + match dim { + Dimension::Xy => &["xmin", "ymin", "xmax", "ymax"], + Dimension::Xyz => &["xmin", "ymin", "zmin", "xmax", "ymax", "zmax"], + Dimension::Xym => &["xmin", "ymin", "mmin", "xmax", "ymax", "mmax"], + Dimension::Xyzm => &[ + "xmin", "ymin", "zmin", "mmin", "xmax", "ymax", "zmax", "mmax", + ], + } +} + +/// The dimension whose box field names match `names`, or `None` if none do. +fn box_dimension_from_names(names: &FieldNames) -> Option { + [ + Dimension::Xy, + Dimension::Xyz, + Dimension::Xym, + Dimension::Xyzm, + ] + .into_iter() + .find(|&dim| { + names + .iter() + .map(AsRef::as_ref) + .eq(box_field_names(dim).iter().copied()) + }) +} + +/// Canonical box storage: a `Struct` of non-nullable `f64` min/max fields, with `nullability` at +/// the struct (per-box) level. +pub(crate) fn box_storage_dtype(dim: Dimension, nullability: Nullability) -> DType { + let names = box_field_names(dim); + let fields = std::iter::repeat_n( + DType::Primitive(PType::F64, Nullability::NonNullable), + names.len(), + ) + .collect::>(); + DType::Struct( + StructFields::new(FieldNames::from(names), fields), + nullability, + ) +} + +/// Validate `dtype` is a box `Struct` of non-nullable `f64` min/max fields, returning its +/// [`Dimension`]. +pub(crate) fn box_dimension(dtype: &DType) -> VortexResult { + let DType::Struct(fields, _) = dtype else { + vortex_bail!("box storage must be a Struct, was {dtype}"); + }; + for (name, field) in fields.names().iter().zip(fields.fields()) { + vortex_ensure!( + matches!( + field, + DType::Primitive(PType::F64, Nullability::NonNullable) + ), + "box field {name} must be non-nullable f64, was {field}" + ); + } + box_dimension_from_names(fields.names()) + .ok_or_else(|| vortex_err!("not a valid geoarrow.box dimension: {:?}", fields.names())) +} + +static ARROW_BOX: CachedId = CachedId::new(BoxType::NAME); + +/// The `geoarrow.box` extension type for `dimension`. +fn box_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> BoxType { + BoxType::new(dimension.into(), geoarrow_metadata(geo_metadata)) +} + +/// Build a geoarrow `RectArray` from a `Rect`'s box `Struct` storage. +fn rect_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let box_type = box_type(&GeoMetadata::default(), box_dimension(storage.dtype())?); + let session = ctx.session().clone(); + let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + RectArray::try_from((arrow.as_ref(), box_type)) + .map_err(|e| vortex_err!("failed to construct RectArray: {e}")) +} + +/// Decode `Rect` storage to `geo_types` (2D [`Geometry::Rect`]), for the geo scalar functions. +pub(crate) fn rect_geometries( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + rect_array(storage, ctx)? + .iter() + .map(|geometry| -> VortexResult> { + Ok(geometry + .ok_or_else(|| vortex_err!("geo: null geometry is not supported"))? + .map_err(|e| vortex_err!("geo: geometry access failed: {e}"))? + .to_geometry()) + }) + .collect() +} + +impl ArrowExportVTable for Rect { + fn arrow_ext_id(&self) -> Id { + *ARROW_BOX + } + + fn vortex_id(&self) -> Id { + self.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let ext_type = dtype.as_extension(); + let geo_metadata = ext_type.metadata::(); + let dimension = box_dimension(ext_type.storage_dtype())?; + + let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + field.try_with_extension_type(box_type(geo_metadata, dimension))?; + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let is_box = array + .dtype() + .as_extension_opt() + .map(|ext| ext.is::()) + .unwrap_or(false); + if !is_box { + return Ok(ArrowExport::Unsupported(array)); + } + + let Ok(box_meta) = target.try_extension_type::() else { + return Ok(ArrowExport::Unsupported(array)); + }; + + let executed = array.execute::(ctx)?; + let storage = executed.storage_array().clone(); + + let storage_field = Field::new( + String::new(), + target.data_type().clone(), + target.is_nullable(), + ); + let session = ctx.session().clone(); + let arrow_storage = session + .arrow() + .execute_arrow(storage, Some(&storage_field), ctx)?; + + // Round-trip through GeoArrow's rect array; `into_arrow` is concrete, so wrap in `Arc`. + let rects = RectArray::try_from((arrow_storage.as_ref(), box_meta)) + .map_err(|e| vortex_err!("failed to construct RectArray: {e}"))?; + + Ok(ArrowExport::Exported(Arc::new(rects.into_arrow()))) + } +} + +impl ArrowImportVTable for Rect { + fn arrow_ext_id(&self) -> Id { + *ARROW_BOX + } + + /// Import a `geoarrow.box` field as the [`Rect`] dtype. Keyed off the standard GeoArrow name, + /// so any producer resolves here. Accepts the full `BoxType` extension, or — for a + /// metadata-less literal — the name alone, inferring the dimension from the field names. + fn from_arrow_field(&self, field: &Field) -> VortexResult> { + let (dimension, metadata) = if let Ok(box_meta) = field.try_extension_type::() { + ( + box_meta.dimension().into(), + geo_metadata_from_arrow(box_meta.metadata()), + ) + } else { + if field.extension_type_name() != Some(BoxType::NAME) { + return Ok(None); + } + // Infer from field names, not the strict `box_dimension` check: a literal's fields may + // be nullable, which that check rejects. + let DType::Struct(fields, _) = DType::from_arrow(field) else { + return Ok(None); + }; + let Some(dimension) = box_dimension_from_names(fields.names()) else { + return Ok(None); + }; + (dimension, GeoMetadata::default()) + }; + + let storage_dtype = box_storage_dtype(dimension, field.is_nullable().into()); + Ok(Some(DType::Extension( + ExtDType::try_with_vtable(Rect, metadata, storage_dtype)?.erased(), + ))) + } + + fn from_arrow_array( + &self, + array: ArrowArrayRef, + field: &Field, + dtype: &DType, + ) -> VortexResult { + let Some(ext_dtype) = dtype.as_extension_opt() else { + return Ok(ArrowImport::Unsupported(array)); + }; + if !ext_dtype.is::() + || field.try_extension_type::().is_err() + || !matches!(array.data_type(), DataType::Struct(_)) + { + return Ok(ArrowImport::Unsupported(array)); + } + + let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + Ok(ArrowImport::Imported( + ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + )) + } +} + +#[cfg(test)] +mod tests { + use geo_types::Geometry; + use rstest::rstest; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::extension::ExtensionArrayExt; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_error::VortexResult; + + use super::Rect; + use super::box_dimension; + use super::box_storage_dtype; + use super::rect_geometries; + use crate::extension::GeoMetadata; + use crate::extension::coordinate::Dimension; + + fn geo_meta() -> GeoMetadata { + GeoMetadata { + crs: Some("EPSG:4326".to_string()), + } + } + + /// `Rect` accepts the canonical box storage of every GeoArrow dimension, and the storage dtype + /// round-trips back to that dimension. + #[rstest] + #[case::xy(Dimension::Xy)] + #[case::xyz(Dimension::Xyz)] + #[case::xym(Dimension::Xym)] + #[case::xyzm(Dimension::Xyzm)] + fn box_validates_every_dimension(#[case] dim: Dimension) -> VortexResult<()> { + let storage = box_storage_dtype(dim, Nullability::NonNullable); + ExtDType::::try_new(geo_meta(), storage.clone())?; + assert_eq!(box_dimension(&storage)?, dim); + Ok(()) + } + + /// Invalid storage is rejected at dtype construction: non-struct storage, a struct with the + /// wrong field names, and a struct with a nullable field. + #[test] + fn box_rejects_invalid_storage() -> VortexResult<()> { + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), primitive).is_err()); + + let point_like = box_storage_dtype(Dimension::Xy, Nullability::NonNullable); + let DType::Struct(fields, _) = &point_like else { + unreachable!("box storage is a struct"); + }; + // Rename a field so it no longer matches the geoarrow.box layout. + let renamed = DType::Struct( + vortex_array::dtype::StructFields::new( + ["x", "ymin", "xmax", "ymax"].into(), + fields.fields().collect(), + ), + Nullability::NonNullable, + ); + assert!(ExtDType::::try_new(geo_meta(), renamed).is_err()); + Ok(()) + } + + /// A `Rect` column decodes to 2D `geo_types` rectangles. + #[test] + fn box_decodes_to_geo_rect() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let rects = crate::test_harness::rect_column(vec![(0.0, 0.0, 2.0, 3.0)])?; + let storage = rects + .execute::(&mut ctx)? + .storage_array() + .clone(); + + let geometries = rect_geometries(&storage, &mut ctx)?; + assert!(matches!(geometries.as_slice(), [Geometry::Rect(_)])); + Ok(()) + } +} diff --git a/vortex-geo/src/extension/wkb.rs b/vortex-geo/src/extension/wkb.rs index a45f1b309a1..fb8c443a7b3 100644 --- a/vortex-geo/src/extension/wkb.rs +++ b/vortex-geo/src/extension/wkb.rs @@ -19,18 +19,18 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::arrow::ArrowExport; -use vortex_array::arrow::ArrowExportVTable; -use vortex_array::arrow::ArrowImport; -use vortex_array::arrow::ArrowImportVTable; -use vortex_array::arrow::ArrowSession; -use vortex_array::arrow::ArrowSessionExt; -use vortex_array::arrow::FromArrowArray; use vortex_array::dtype::DType; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; diff --git a/vortex-geo/src/lib.rs b/vortex-geo/src/lib.rs index 951d93b7b4f..3a0c52f1eb9 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -3,17 +3,31 @@ use std::sync::Arc; -use vortex_array::arrow::ArrowSessionExt; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::scalar_fn::session::ScalarFnSessionExt; +use vortex_array::stats::session::StatsSessionExt; +use vortex_arrow::ArrowSessionExt; use vortex_session::VortexSession; +use crate::aggregate_fn::GeometryAabb; +use crate::extension::LineString; +use crate::extension::MultiLineString; +use crate::extension::MultiPoint; +use crate::extension::MultiPolygon; use crate::extension::Point; use crate::extension::Polygon; +use crate::extension::Rect; use crate::extension::WellKnownBinary; +use crate::prune::GeoDistancePrune; +use crate::prune::GeoIntersectsPrune; +use crate::scalar_fn::contains::GeoContains; use crate::scalar_fn::distance::GeoDistance; +use crate::scalar_fn::intersects::GeoIntersects; +pub mod aggregate_fn; pub mod extension; +pub mod prune; pub mod scalar_fn; #[cfg(test)] mod test_harness; @@ -29,10 +43,35 @@ pub fn initialize(session: &VortexSession) { session.dtypes().register(Point); session.arrow().register_exporter(Arc::new(Point)); session.arrow().register_importer(Arc::new(Point)); + session.dtypes().register(LineString); + session.arrow().register_exporter(Arc::new(LineString)); + session.arrow().register_importer(Arc::new(LineString)); + session.dtypes().register(MultiPoint); + session.arrow().register_exporter(Arc::new(MultiPoint)); + session.arrow().register_importer(Arc::new(MultiPoint)); session.dtypes().register(Polygon); session.arrow().register_exporter(Arc::new(Polygon)); session.arrow().register_importer(Arc::new(Polygon)); + session.dtypes().register(MultiLineString); + session.arrow().register_exporter(Arc::new(MultiLineString)); + session.arrow().register_importer(Arc::new(MultiLineString)); + session.dtypes().register(MultiPolygon); + session.arrow().register_exporter(Arc::new(MultiPolygon)); + session.arrow().register_importer(Arc::new(MultiPolygon)); + session.dtypes().register(Rect); + session.arrow().register_exporter(Arc::new(Rect)); + session.arrow().register_importer(Arc::new(Rect)); // Register the geometry scalar functions. + session.scalar_fns().register(GeoContains); session.scalar_fns().register(GeoDistance); + session.scalar_fns().register(GeoIntersects); + + // The axis-aligned bounding-box (AABB) aggregate; self-declares as a per-chunk zone stat for + // geometry columns. + session.aggregate_fns().register(GeometryAabb); + + // Register the spatial pruning rules that use that AABB. + session.stats().register_rewrite(GeoDistancePrune); + session.stats().register_rewrite(GeoIntersectsPrune); } diff --git a/vortex-geo/src/prune/distance.rs b/vortex-geo/src/prune/distance.rs new file mode 100644 index 00000000000..380571c11f9 --- /dev/null +++ b/vortex-geo/src/prune/distance.rs @@ -0,0 +1,377 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Distance(geom, const) radius` pruning. + +use geo::Rect as GeoRect; +use vortex_array::expr::Expression; +use vortex_array::expr::gt; +use vortex_array::expr::gt_eq; +use vortex_array::expr::lit; +use vortex_array::expr::lt; +use vortex_array::expr::lt_eq; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::literal::Literal; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::stats::rewrite::StatsRewriteCtx; +use vortex_array::stats::rewrite::StatsRewriteRule; +use vortex_error::VortexResult; + +use super::aabb_stat; +use super::geometry_and_constant; +use super::max_dist_sq; +use super::min_dist_sq; +use super::query_aabb; +use crate::scalar_fn::distance::GeoDistance; + +/// Prunes chunks for `ST_Distance(geom, const) r` filters. +/// +/// All four comparisons prune: `<= r` / `< r` skip a chunk whose box is wholly beyond `r` (box +/// min-distance); `>= r` / `> r` skip one wholly within `r` (box max-distance). `==` / `!=` don't +/// prune. +#[derive(Debug)] +pub struct GeoDistancePrune; + +impl StatsRewriteRule for GeoDistancePrune { + fn scalar_fn_id(&self) -> ScalarFnId { + // The predicate root is the comparison, not `GeoDistance`, so key on `Binary`. + Binary.id() + } + + fn falsify( + &self, + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + // Only the ordered comparisons prune today. `== r` could prune in the future (a chunk is + // provably empty when `r` lies outside its box's [min, max] distance interval), it's just + // not implemented. `!= r` cannot: pruning would need every row's distance to equal `r`, + // which an AABB can't prove. + let op = *expr.as_::(); + if !matches!( + op, + Operator::Lte | Operator::Lt | Operator::Gte | Operator::Gt + ) { + return Ok(None); + } + + // The left operand must be `GeoDistance(geom, const)`; the right, the radius literal. + let distance = expr.child(0); + if distance.as_opt::().is_none() { + return Ok(None); + } + let Some(radius) = expr.child(1).as_opt::() else { + return Ok(None); + }; + // Casts any primitive radius (filters arrive uncoerced, so integer literals are + // legitimate); a null or non-primitive (e.g. extension-typed) literal has no value to + // reason about, decline and scan, never error. + let Ok(radius) = f64::try_from(radius) else { + return Ok(None); + }; + // A NaN radius has no sound proof: `distance NaN` is not a total order, so the chunk + // must be scanned, not pruned. + if radius.is_nan() { + return Ok(None); + } + + let Some((geom, constant)) = geometry_and_constant(distance, ctx)? else { + return Ok(None); + }; + let Some(query) = query_aabb(constant, ctx)? else { + return Ok(None); + }; + Ok(distance_prune_proof(geom, query, op, radius)) + } +} + +/// Build the prune proof for `ST_Distance(geom, const) radius`, or `None` when this +/// operator/radius cannot prune. The box-to-box distance bounds every row's true distance for any +/// geometry types: `<=` / `<` prune a chunk whose box is wholly beyond `radius` (min +/// box-distance); `>=` / `>` prune one wholly within it (max box-distance). +/// +/// A distance is always `>= 0`, which decides the degenerate radii up front. +fn distance_prune_proof( + geom: &Expression, + query: GeoRect, + op: Operator, + radius: f64, +) -> Option { + // A distance is always non-negative, so degenerate radii resolve without touching the box. + match op { + // `<= r` / `< r` with a negative radius (or zero, for `<`) match nothing: prune every chunk. + Operator::Lte if radius < 0.0 => return Some(lit(true)), + Operator::Lt if radius <= 0.0 => return Some(lit(true)), + // `>= r` / `> r` with a negative radius (or zero, for `>=`) match all rows: never prune. + Operator::Gte if radius <= 0.0 => return None, + Operator::Gt if radius < 0.0 => return None, + _ => {} + } + // Compared squared to avoid a `sqrt`; all operands are `>= 0`. + let aabb = aabb_stat(geom); + let r2 = lit(radius * radius); + Some(match op { + // Beyond the threshold: even the nearest the box can be exceeds `r`. + Operator::Lte => gt(min_dist_sq(&aabb, query), r2), + Operator::Lt => gt_eq(min_dist_sq(&aabb, query), r2), + // Within the threshold: even the farthest the box can be is below `r`. + Operator::Gte => lt(max_dist_sq(&aabb, query), r2), + Operator::Gt => lt_eq(max_dist_sq(&aabb, query), r2), + _ => return None, + }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::VortexSessionExecute; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::gt_eq; + use vortex_array::expr::lit; + use vortex_array::expr::lt_eq; + use vortex_array::expr::root; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTableExt; + use vortex_array::scalar_fn::fns::binary::Binary; + use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::stats::rewrite::StatsRewriteCtx; + use vortex_array::stats::rewrite::StatsRewriteRule; + use vortex_error::VortexResult; + + use super::GeoDistancePrune; + use crate::prune::test_harness::aabb_zone_map; + use crate::prune::test_harness::empty_zone_map; + use crate::scalar_fn::distance::GeoDistance; + use crate::test_harness::geo_session; + use crate::test_harness::point_column; + + /// Run the rule against `GeoDistance(root, origin) radius`, operands swapped when + /// `geom_first` is false. The radius is any literal scalar, matching the uncoerced filter + /// expressions the rule sees in production. + fn falsify_distance( + operator: Operator, + geom_first: bool, + radius: impl Into, + ) -> VortexResult> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let operands = if geom_first { + [root(), lit(origin)] + } else { + [lit(origin), root()] + }; + let distance = GeoDistance.new_expr(EmptyOptions, operands); + let predicate = Binary.new_expr(operator, [distance, lit(radius.into())]); + + GeoDistancePrune.falsify(&predicate, &StatsRewriteCtx::new(&session, &scope)) + } + + /// All four distance comparisons prune (`<=`/`<` via min-distance, `>=`/`>` via max-distance); + /// `==`/`!=` are left to the scan. + #[rstest] + #[case(Operator::Lte, true)] + #[case(Operator::Lt, true)] + #[case(Operator::Gt, true)] + #[case(Operator::Gte, true)] + #[case(Operator::Eq, false)] + #[case(Operator::NotEq, false)] + fn prunes_distance_comparisons( + #[case] operator: Operator, + #[case] prunes: bool, + ) -> VortexResult<()> { + assert_eq!(falsify_distance(operator, true, 0.5)?.is_some(), prunes); + Ok(()) + } + + /// Distance is symmetric: both operand orders produce a proof. + #[rstest] + #[case(true)] + #[case(false)] + fn falsifies_either_operand_order(#[case] geom_first: bool) -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, geom_first, 0.5)?.is_some()); + Ok(()) + } + + /// A NaN radius must not prune - the scan's total-order compare treats `dist <= NaN` as true. + #[test] + fn nan_radius_never_prunes() -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, true, f64::NAN)?.is_none()); + Ok(()) + } + + /// A negative radius prunes every zone - vacuously sound: `distance <= r < 0` matches no row. + #[test] + fn negative_radius_prunes_vacuously() -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, true, -0.5)?.is_some()); + Ok(()) + } + + /// Filter expressions arrive uncoerced, so `distance <= 10` may carry an integer literal - + /// it casts and prunes like an f64 radius. + #[test] + fn integer_radius_prunes() -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, true, 10i64)?.is_some()); + Ok(()) + } + + /// A null radius has no value to reason about; the rule declines and the chunk is scanned. + #[test] + fn null_radius_never_prunes() -> VortexResult<()> { + let radius = Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)); + assert!(falsify_distance(Operator::Lte, true, radius)?.is_none()); + Ok(()) + } + + /// An extension-typed radius passes `Binary`'s typecheck (extension operands are exempt) but + /// has no numeric value - the rule declines rather than erroring, and the chunk is scanned. + #[test] + fn extension_radius_never_prunes() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let geometry = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + assert!(falsify_distance(Operator::Lte, true, geometry)?.is_none()); + Ok(()) + } + + /// A scope dtype without `GeometryAabb` support gets no proof - the stat reference would + /// fail to bind at prune time. + #[test] + fn unsupported_scope_is_not_pruned() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let scope = DType::Primitive(PType::F64, Nullability::NonNullable); + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let predicate = lt_eq(distance, lit(0.5f64)); + + let ctx = StatsRewriteCtx::new(&session, &scope); + assert!(GeoDistancePrune.falsify(&predicate, &ctx)?.is_none()); + Ok(()) + } + + /// A comparison that does not wrap `GeoDistance` is left untouched. + #[test] + fn ignores_non_distance_comparison() -> VortexResult<()> { + let session = geo_session(); + let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + + let predicate = lt_eq(lit(1.0f64), lit(2.0f64)); + let ctx = StatsRewriteCtx::new(&session, &scope); + assert!(GeoDistancePrune.falsify(&predicate, &ctx)?.is_none()); + Ok(()) + } + + /// End-to-end over a hand-built zone map: the far chunk is skipped, the near one kept. + #[test] + fn prunes_far_chunk_keeps_near() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + // Chunk 0 near the origin (0,0..1,1), chunk 1 far away (100,100..101,101). + let zone_map = aabb_zone_map( + &point_dtype, + &[[0.0, 0.0, 1.0, 1.0], [100.0, 100.0, 101.0, 101.0]], + )?; + + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let predicate = lt_eq(distance, lit(0.5f64)); + let proof = predicate + .falsify(&point_dtype, &session)? + .expect("distance filter should be falsifiable"); + + // `true` means the zone is pruned: chunk 0 (near origin) is kept, chunk 1 (far) is skipped. + let mask = zone_map.prune(&proof, &session)?; + assert_eq!(mask.iter().collect::>(), vec![false, true]); + Ok(()) + } + + /// The true-distance prune skips a chunk that is *diagonally* farther than `r`, even though + /// neither axis alone exceeds `r`, the case a per-axis box-overlap test would wrongly keep. + #[test] + fn prunes_diagonally_distant_chunk() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + // One chunk, AABB (0.8,0.8)..(0.9,0.9): each axis is only 0.8 from the origin (<= r = 1), but + // the near corner is sqrt(0.8^2 + 0.8^2) ~= 1.13 away (> 1), so no point in the box is within 1. + let zone_map = aabb_zone_map(&point_dtype, &[[0.8, 0.8, 0.9, 0.9]])?; + + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let predicate = lt_eq(distance, lit(1.0f64)); + let proof = predicate + .falsify(&point_dtype, &session)? + .expect("distance filter should be falsifiable"); + + assert_eq!( + zone_map + .prune(&proof, &session)? + .iter() + .collect::>(), + vec![true], + ); + Ok(()) + } + + /// A `>= r` filter prunes a chunk lying wholly *within* `r` (every row nearer than `r`, so none + /// satisfy `>= r`) via the box max-distance, while a chunk beyond `r` is kept. + #[test] + fn prunes_within_chunk_for_far_filter() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + // Chunk 0 (AABB 0,0..0.5,0.5, farthest corner ~= 0.707) is entirely within 2 of the origin; + // chunk 1 (100,100..101,101) is entirely beyond it. + let zone_map = aabb_zone_map( + &point_dtype, + &[[0.0, 0.0, 0.5, 0.5], [100.0, 100.0, 101.0, 101.0]], + )?; + + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let proof = gt_eq(distance, lit(2.0f64)) + .falsify(&point_dtype, &session)? + .expect("distance filter should be falsifiable"); + + // Chunk 0 (within 2) is pruned for `>= 2`; chunk 1 (beyond 2) is kept. + let mask = zone_map.prune(&proof, &session)?; + assert_eq!(mask.iter().collect::>(), vec![true, false]); + Ok(()) + } + + /// Backward compat: a zone map written without the `GeometryAabb` stat (an older file) keeps + /// every zone, the missing stat binds to null and `null_as_false` retains the zone. + #[test] + fn missing_aabb_stat_keeps_all_zones() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let zone_map = empty_zone_map(&point_dtype)?; + + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let proof = lt_eq(distance, lit(0.5f64)) + .falsify(&point_dtype, &session)? + .expect("distance filter should be falsifiable"); + + let mask = zone_map.prune(&proof, &session)?; + assert_eq!(mask.iter().collect::>(), vec![false, false]); + Ok(()) + } +} diff --git a/vortex-geo/src/prune/intersects.rs b/vortex-geo/src/prune/intersects.rs new file mode 100644 index 00000000000..9af6c5b4d9e --- /dev/null +++ b/vortex-geo/src/prune/intersects.rs @@ -0,0 +1,166 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Intersects(geom, const)` pruning. + +use vortex_array::expr::Expression; +use vortex_array::expr::gt; +use vortex_array::expr::lit; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::stats::rewrite::StatsRewriteCtx; +use vortex_array::stats::rewrite::StatsRewriteRule; +use vortex_error::VortexResult; + +use super::aabb_stat; +use super::geometry_and_constant; +use super::min_dist_sq; +use super::query_aabb; +use crate::scalar_fn::intersects::GeoIntersects; + +/// Prunes chunks for `ST_Intersects(geom, const)` filters: a chunk whose box is strictly +/// separated from the constant's bounding box cannot contain an intersecting row. +/// +/// Only the positive form prunes. `NOT ST_Intersects` cannot: it would need every row to provably +/// intersect the constant, and box overlap never proves geometry intersection. +#[derive(Debug)] +pub struct GeoIntersectsPrune; + +impl StatsRewriteRule for GeoIntersectsPrune { + fn scalar_fn_id(&self) -> ScalarFnId { + // Unlike the distance rule, the boolean predicate is itself the expression root. + GeoIntersects.id() + } + + fn falsify( + &self, + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + let Some((geom, constant)) = geometry_and_constant(expr, ctx)? else { + return Ok(None); + }; + let Some(query) = query_aabb(constant, ctx)? else { + return Ok(None); + }; + // Disjoint iff the minimum box-to-box distance is positive. Strictly (`gt`, not `gt_eq`): + // boxes that merely touch must scan, since touching geometries do intersect. + Ok(Some(gt(min_dist_sq(&aabb_stat(geom), query), lit(0.0)))) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::VortexSessionExecute; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::lit; + use vortex_array::expr::root; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTableExt; + use vortex_array::stats::rewrite::StatsRewriteCtx; + use vortex_array::stats::rewrite::StatsRewriteRule; + use vortex_error::VortexResult; + + use super::GeoIntersectsPrune; + use crate::prune::test_harness::aabb_zone_map; + use crate::prune::test_harness::empty_zone_map; + use crate::scalar_fn::intersects::GeoIntersects; + use crate::test_harness::geo_session; + use crate::test_harness::point_column; + + /// Run the intersects rule against `GeoIntersects(root, point(1.0, 0.5))`, operands swapped + /// when `geom_first` is false. + fn falsify_intersects(geom_first: bool) -> VortexResult> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let query = point_column(vec![1.0], vec![0.5])?.execute_scalar(0, &mut ctx)?; + let operands = if geom_first { + [root(), lit(query)] + } else { + [lit(query), root()] + }; + let predicate = GeoIntersects.new_expr(EmptyOptions, operands); + GeoIntersectsPrune.falsify(&predicate, &StatsRewriteCtx::new(&session, &scope)) + } + + /// Intersects is symmetric: both operand orders produce a proof. + #[rstest] + #[case(true)] + #[case(false)] + fn falsifies_either_operand_order(#[case] geom_first: bool) -> VortexResult<()> { + assert!(falsify_intersects(geom_first)?.is_some()); + Ok(()) + } + + /// A scope dtype without `GeometryAabb` support gets no proof, the stat reference would + /// fail to bind at prune time. + #[test] + fn unsupported_scope_is_not_pruned() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let scope = DType::Primitive(PType::F64, Nullability::NonNullable); + let query = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let predicate = GeoIntersects.new_expr(EmptyOptions, [root(), lit(query)]); + + let ctx = StatsRewriteCtx::new(&session, &scope); + assert!(GeoIntersectsPrune.falsify(&predicate, &ctx)?.is_none()); + Ok(()) + } + + /// End-to-end: a zone strictly separated from the query is skipped; zones containing or merely + /// touching the query must scan, touching geometries intersect under OGC semantics. + #[test] + fn prunes_disjoint_keeps_touching_and_containing() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + // Query point (1.0, 0.5): zone 0 contains it, zone 1 only touches it at `x == 1`, zone 2 + // is strictly separated. + let zone_map = aabb_zone_map( + &point_dtype, + &[ + [0.5, 0.0, 1.5, 1.0], + [-1.0, 0.0, 1.0, 1.0], + [3.0, 0.0, 4.0, 1.0], + ], + )?; + + let query = point_column(vec![1.0], vec![0.5])?.execute_scalar(0, &mut ctx)?; + let predicate = GeoIntersects.new_expr(EmptyOptions, [root(), lit(query)]); + let proof = predicate + .falsify(&point_dtype, &session)? + .expect("intersects filter should be falsifiable"); + + let mask = zone_map.prune(&proof, &session)?; + assert_eq!(mask.iter().collect::>(), vec![false, false, true]); + Ok(()) + } + + /// Backward compat: a zone map written without the `GeometryAabb` stat keeps every zone. + #[test] + fn missing_aabb_stat_keeps_all_zones() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let zone_map = empty_zone_map(&point_dtype)?; + + let query = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let proof = GeoIntersects + .new_expr(EmptyOptions, [root(), lit(query)]) + .falsify(&point_dtype, &session)? + .expect("intersects filter should be falsifiable"); + + let mask = zone_map.prune(&proof, &session)?; + assert_eq!(mask.iter().collect::>(), vec![false, false]); + Ok(()) + } +} diff --git a/vortex-geo/src/prune/mod.rs b/vortex-geo/src/prune/mod.rs new file mode 100644 index 00000000000..00200f1500c --- /dev/null +++ b/vortex-geo/src/prune/mod.rs @@ -0,0 +1,164 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Chunk pruning for spatial filters, using the per-chunk [`GeometryAabb`] axis-aligned +//! bounding box (AABB). +//! +//! One module per predicate. To prune a new one: recover the geometry column and the constant +//! (symmetric predicates can use `geometry_and_constant`), reduce the constant to its box with +//! `query_aabb`, build the proof over `aabb_stat` from the box-relation builders (`min_dist_sq`, +//! `max_dist_sq`), and register the rule in [`crate::initialize`]. No new statistic or +//! file-format change is needed. + +mod distance; +mod intersects; +#[cfg(test)] +mod test_harness; + +pub use distance::GeoDistancePrune; +use geo::BoundingRect; +use geo::Rect as GeoRect; +pub use intersects::GeoIntersectsPrune; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::expr::Expression; +use vortex_array::expr::case_when; +use vortex_array::expr::checked_add; +use vortex_array::expr::ext_storage; +use vortex_array::expr::get_item; +use vortex_array::expr::gt; +use vortex_array::expr::is_root; +use vortex_array::expr::lit; +use vortex_array::expr::lt; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::literal::Literal; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::stats::rewrite::StatsRewriteCtx; +use vortex_array::stats::stat; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::aggregate_fn::GeometryAabb; +use crate::extension::is_native_geometry; +use crate::extension::single_geometry; + +/// Splits a symmetric two-operand geo predicate into the scope-rooted geometry column and the +/// constant operand's scalar. +/// +/// `None` means the rule must decline: the expression doesn't have the `f(column, constant)` +/// shape (in either operand order), or the column's dtype carries no [`GeometryAabb`] statistic. +/// An asymmetric predicate (e.g. a future contains) must recover which operand is the column +/// itself instead of calling this. +fn geometry_and_constant<'a>( + expr: &'a Expression, + ctx: &StatsRewriteCtx<'_>, +) -> VortexResult> { + // The predicate is symmetric, so the column (scope root) and the constant may be on either + // side. + let (lhs, rhs) = (expr.child(0), expr.child(1)); + let (geom, constant) = if is_root(lhs) { + (lhs, rhs) + } else if is_root(rhs) { + (rhs, lhs) + } else { + return Ok(None); + }; + + // A `GeometryAabb` stat reference only binds for dtypes it supports; anything else (e.g. a + // WKB column) must fall through to the scan. + if !is_native_geometry(&ctx.return_dtype(geom)?) { + return Ok(None); + } + + Ok(constant.as_opt::().map(|scalar| (geom, scalar))) +} + +/// The 2D bounding box of a constant geometry of any type, or `None` for one without an extent +/// (e.g. an empty geometry). +/// +/// Prove claims against this box rather than the constant itself: the constant lies inside it, +/// so whatever holds for the box holds for the geometry. +fn query_aabb(constant: &Scalar, ctx: &StatsRewriteCtx<'_>) -> VortexResult>> { + // Decoding the constant into a concrete geometry runs through the compute stack, which needs + // an execution context. + let mut exec = ctx.session().create_execution_ctx(); + Ok(single_geometry(constant, &mut exec)?.bounding_rect()) +} + +/// The chunk's AABB statistic, as the storage struct with `xmin`/`ymin`/`xmax`/`ymax` fields. +/// +/// A chunk written without the statistic reads as null here; every proof built on top must let +/// that null propagate to its root, where the zone map keeps the chunk. +fn aabb_stat(geom: &Expression) -> Expression { + // `ext_storage` unwraps the native `geoarrow.box` stat value to its backing struct, so + // proofs can `get_item` the coordinate fields. + ext_storage(stat(geom.clone(), GeometryAabb.bind(EmptyOptions))) +} + +/// Lower bound on every row's squared distance to the query AABB: zero when the boxes overlap or +/// touch, positive iff they are strictly separated. +/// +/// Prunes "near" predicates: `min_dist_sq > r^2` proves every row is farther than `r`. +fn min_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { + let field = |name: &str| get_item(name, aabb.clone()); + // Per axis: gap = max(0, q_lo - aabb_hi, aabb_lo - q_hi), positive only when the intervals + // are separated. The nearest two points of the boxes are one axis-gap apart per axis, so the + // squared distance is gap_x^2 + gap_y^2 (squared throughout to avoid a sqrt). + let gap = |q_lo: f64, q_hi: f64, lo: Expression, hi: Expression| { + maximum( + lit(0.0), + maximum( + binop(Operator::Sub, lit(q_lo), hi), + binop(Operator::Sub, lo, lit(q_hi)), + ), + ) + }; + let dx = gap(query.min().x, query.max().x, field("xmin"), field("xmax")); + let dy = gap(query.min().y, query.max().y, field("ymin"), field("ymax")); + checked_add(square(dx), square(dy)) +} + +/// Upper bound on every row's squared distance to the query AABB. +/// +/// Prunes "far" predicates: `max_dist_sq < r^2` proves every row is within `r`. +fn max_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { + let field = |name: &str| get_item(name, aabb.clone()); + // Per axis: span = max(q_hi, aabb_hi) - min(q_lo, aabb_lo), the farthest two points of the + // boxes can be apart. The nullable AABB field is the second `maximum`/`minimum` argument so + // that `case_when`'s else branch carries the nullability - a missing stat propagates null. + let span = |q_lo: f64, q_hi: f64, lo: Expression, hi: Expression| { + binop( + Operator::Sub, + maximum(lit(q_hi), hi), + minimum(lit(q_lo), lo), + ) + }; + let dx = span(query.min().x, query.max().x, field("xmin"), field("xmax")); + let dy = span(query.min().y, query.max().y, field("ymin"), field("ymax")); + checked_add(square(dx), square(dy)) +} + +/// `a b`. +fn binop(op: Operator, a: Expression, b: Expression) -> Expression { + Binary + .try_new_expr(op, [a, b]) + .vortex_expect("binary expression") +} + +/// `e * e`. +fn square(e: Expression) -> Expression { + binop(Operator::Mul, e.clone(), e) +} + +/// `max(a, b)`. +fn maximum(a: Expression, b: Expression) -> Expression { + case_when(gt(a.clone(), b.clone()), a, b) +} + +/// `min(a, b)`. +fn minimum(a: Expression, b: Expression) -> Expression { + case_when(lt(a.clone(), b.clone()), a, b) +} diff --git a/vortex-geo/src/prune/test_harness.rs b/vortex-geo/src/prune/test_harness.rs new file mode 100644 index 00000000000..fb3e12ef9d5 --- /dev/null +++ b/vortex-geo/src/prune/test_harness.rs @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared zone-map fixtures for the pruning-rule tests; each rule module owns its own cases. + +use std::sync::Arc; + +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_layout::layouts::zoned::zone_map::ZoneMap; + +use crate::aggregate_fn::GeometryAabb; +use crate::extension::GeoMetadata; +use crate::extension::Rect; + +/// A single-column zone map holding one native-box AABB stat row (`[xmin, ymin, xmax, ymax]`) +/// per zone, with default (unreferenced) metadata to match the aggregate's return dtype. +pub(super) fn aabb_zone_map(point_dtype: &DType, boxes: &[[f64; 4]]) -> VortexResult { + let aabb_fn = GeometryAabb.bind(EmptyOptions); + let col = |i: usize| PrimitiveArray::from_iter(boxes.iter().map(move |b| b[i])).into_array(); + let storage = StructArray::try_new( + ["xmin", "ymin", "xmax", "ymax"].into(), + vec![col(0), col(1), col(2), col(3)], + boxes.len(), + Validity::AllValid, + )? + .into_array(); + let box_dtype = + ExtDType::::try_new(GeoMetadata::default(), storage.dtype().clone())?.erased(); + let aabbs = ExtensionArray::try_new(box_dtype, storage)?.into_array(); + let zone_array = StructArray::from_fields(&[(aabb_fn.to_string().as_str(), aabbs)])?; + ZoneMap::try_new( + point_dtype.clone(), + zone_array, + Arc::new([aabb_fn]), + 1, + boxes.len() as u64, + ) +} + +/// A two-zone zone map with no stat columns at all, as written before the AABB stat existed. +pub(super) fn empty_zone_map(point_dtype: &DType) -> VortexResult { + ZoneMap::try_new( + point_dtype.clone(), + StructArray::try_new(FieldNames::empty(), vec![], 2, Validity::NonNullable)?, + Arc::new([]), + 1, + 2, + ) +} diff --git a/vortex-geo/src/scalar_fn/contains.rs b/vortex-geo/src/scalar_fn/contains.rs new file mode 100644 index 00000000000..5638c224e20 --- /dev/null +++ b/vortex-geo/src/scalar_fn/contains.rs @@ -0,0 +1,307 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Contains`: OGC containment test between two native geometries. + +use geo::Contains; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::geometries; +use crate::extension::single_geometry; +use crate::extension::validate_geometry_operands; + +/// OGC `ST_Contains` between two native geometry operands, each a column or a constant +/// literal: true where operand `b` lies completely inside operand `a` (boundary contact alone +/// does not count). Containment is not symmetric; the operand order is significant. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct GeoContains; + +impl GeoContains { + /// A lazy `ScalarFnArray` computing per-row whether operand `a` contains operand `b`; + /// either may be constant. The output length is taken from `a`. + pub fn try_new_array(a: ArrayRef, b: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(GeoContains, EmptyOptions).erased(), + vec![a, b], + ) + } +} + +impl ScalarFnVTable for GeoContains { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.geo.contains"); + *ID + } + + fn serialize(&self, _: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + Ok(EmptyOptions) + } + + fn arity(&self, _: &Self::Options) -> Arity { + Arity::Exact(2) + } + + fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("a"), + 1 => ChildName::from("b"), + _ => unreachable!("contains has exactly two children"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + validate_geometry_operands(dtypes)?; + Ok(DType::Bool(Nullability::NonNullable)) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let a = args.get(0)?; + let b = args.get(1)?; + // Containment is not symmetric: `a` is always the container and `b` the contained. + match (a.as_opt::(), b.as_opt::()) { + (Some(qa), Some(qb)) => { + let ga = single_geometry(qa.scalar(), ctx)?; + let gb = single_geometry(qb.scalar(), ctx)?; + Ok(ConstantArray::new( + Scalar::bool(ga.contains(&gb), Nullability::NonNullable), + a.len(), + ) + .into_array()) + } + (Some(qa), None) => constant_contains_column(qa.scalar(), &b, ctx), + (None, Some(qb)) => column_contains_constant(&a, qb.scalar(), ctx), + (None, None) => { + vortex_ensure_eq!( + a.len(), + b.len(), + "geo contains: operand length mismatch {} vs {}", + a.len(), + b.len() + ); + let ag = geometries(&a, ctx)?; + let bg = geometries(&b, ctx)?; + let hits = ag.iter().zip(&bg).map(|(x, y)| x.contains(y)); + Ok(BoolArray::from_iter(hits).into_array()) + } + } + } +} + +/// Whether the constant `container` contains each row of `contained`. +fn constant_contains_column( + container: &Scalar, + contained: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let container = single_geometry(container, ctx)?; + let geoms = geometries(contained, ctx)?; + let hits = geoms.iter().map(|g| container.contains(g)); + Ok(BoolArray::from_iter(hits).into_array()) +} + +/// Whether each row of `container` contains the constant `contained`. +fn column_contains_constant( + container: &ArrayRef, + contained: &Scalar, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let contained = single_geometry(contained, ctx)?; + let geoms = geometries(container, ctx)?; + let hits = geoms.iter().map(|g| g.contains(&contained)); + Ok(BoolArray::from_iter(hits).into_array()) +} + +#[cfg(test)] +mod tests { + use geo_types::Geometry; + use geo_types::LineString; + use geo_types::Point; + use geo_types::Polygon; + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use wkb::writer::WriteOptions; + + use super::GeoContains; + use crate::test_harness::point_column; + + /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. + fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon { + Polygon::new( + LineString::from(vec![(x0, y0), (x1, y0), (x1, y1), (x0, y1), (x0, y0)]), + vec![], + ) + } + + /// A constant column of length `len`, every row the native form of `geometry`. + fn geometry_constant(geometry: &Geometry, len: usize) -> VortexResult { + let mut buf = Vec::new(); + wkb::writer::write_geometry(&mut buf, geometry, &WriteOptions::default()) + .map_err(|e| vortex_err!("writing WKB failed: {e}"))?; + let scalar = crate::extension::native_geometry_scalar_from_wkb(&buf)? + .ok_or_else(|| vortex_err!("unsupported geometry type"))?; + Ok(ConstantArray::new(scalar, len).into_array()) + } + + /// Materialize `array` so it is no longer a `Constant`, forcing the non-constant kernel + /// paths. + fn materialize(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array.execute::(ctx)?.into_array()) + } + + /// Execute `GeoContains(a, b)` and assert the per-row verdicts equal `expected`. + fn assert_contains( + a: ArrayRef, + b: ArrayRef, + expected: impl IntoIterator, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let contains = GeoContains::try_new_array(a, b)?.into_array(); + assert_arrays_eq!(contains, BoolArray::from_iter(expected), &mut ctx); + Ok(()) + } + + // The tests cover each `execute` dispatch arm in match order, then the edge cases. + + /// Constant vs constant: a polygon contains a nested polygon but not a partially + /// overlapping or disjoint one; every output row carries the same verdict. + #[rstest] + #[case::nested(rect_polygon(1.0, 1.0, 3.0, 3.0), true)] + #[case::overlapping(rect_polygon(2.0, 2.0, 6.0, 6.0), false)] + #[case::disjoint(rect_polygon(20.0, 20.0, 24.0, 24.0), false)] + fn constant_vs_constant_polygons( + #[case] other: Polygon, + #[case] expected: bool, + ) -> VortexResult<()> { + let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; + let other = geometry_constant(&Geometry::Polygon(other), 3)?; + assert_contains(container, other, [expected; 3]) + } + + /// Partially overlapping polygons contain each other in neither direction. + #[test] + fn overlapping_polygons_contain_neither_way() -> VortexResult<()> { + let a = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?; + let b = geometry_constant(&Geometry::Polygon(rect_polygon(2.0, 2.0, 6.0, 6.0)), 2)?; + assert_contains(a.clone(), b.clone(), [false; 2])?; + assert_contains(b, a, [false; 2]) + } + + /// Containment is not symmetric: a polygon contains an interior point, but the point does + /// not contain the polygon. + #[test] + fn contains_is_asymmetric() -> VortexResult<()> { + let polygon = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?; + let point = geometry_constant(&Geometry::Point(Point::new(2.0, 2.0)), 2)?; + assert_contains(polygon.clone(), point.clone(), [true; 2])?; + assert_contains(point, polygon, [false; 2]) + } + + /// Constant polygon vs point column: a strictly interior point is contained; points outside + /// or exactly on the boundary are not (OGC contains excludes the boundary). + #[test] + fn constant_polygon_vs_point_column() -> VortexResult<()> { + let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; + let points = point_column(vec![2.0, 10.0, 0.0], vec![2.0, 10.0, 2.0])?; + assert_contains(container, points, [true, false, false]) + } + + /// Polygon column vs constant point: only the polygon around the point contains it. + #[test] + fn polygon_column_vs_constant_point() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let around = materialize( + geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?, + &mut ctx, + )?; + let away = materialize( + geometry_constant(&Geometry::Polygon(rect_polygon(20.0, 20.0, 24.0, 24.0)), 2)?, + &mut ctx, + )?; + let point = geometry_constant(&Geometry::Point(Point::new(2.0, 2.0)), 2)?; + + assert_contains(around, point.clone(), [true; 2])?; + assert_contains(away, point, [false; 2]) + } + + /// Column vs column pairs rows: each polygon row is tested against the point row at the + /// same position. + #[test] + fn polygon_column_vs_point_column() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let polygons = materialize( + geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?, + &mut ctx, + )?; + let points = point_column(vec![2.0, 10.0], vec![2.0, 10.0])?; + assert_contains(polygons, points, [true, false]) + } + + /// Geometry arrays are never nullable, so a nullable operand dtype is rejected. + #[test] + fn nullable_operand_is_rejected() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let result = GeoContains.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype]); + assert!(result.is_err()); + Ok(()) + } + + /// A non-geometry operand dtype is rejected up front, before execution. + #[test] + fn non_geometry_operand_is_rejected() -> VortexResult<()> { + let geo = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let numeric = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = GeoContains.return_dtype(&EmptyOptions, &[geo, numeric]); + assert!(result.is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs index 6531c1dd8f2..2e196706ec2 100644 --- a/vortex-geo/src/scalar_fn/distance.rs +++ b/vortex-geo/src/scalar_fn/distance.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! `ST_Distance`: planar (Euclidean) distance between two native geometries via the `geo` crate. +//! `ST_Distance`: planar (Euclidean) distance between two native geometries. use geo::Distance; use geo::Euclidean; @@ -24,15 +24,16 @@ use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::extension::geometries; use crate::extension::single_geometry; +use crate::extension::validate_geometry_operands; -/// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry operands. -/// Each is a column or a constant literal; `geo` computes the distance between each pair. +/// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry +/// operands, each a column or a constant literal. #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] pub struct GeoDistance; @@ -75,7 +76,8 @@ impl ScalarFnVTable for GeoDistance { } } - fn return_dtype(&self, _: &Self::Options, _: &[DType]) -> VortexResult { + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + validate_geometry_operands(dtypes)?; Ok(DType::Primitive(PType::F64, Nullability::NonNullable)) } @@ -101,14 +103,15 @@ impl ScalarFnVTable for GeoDistance { (Some(query), None) => distances_to_constant(&b, query.scalar(), ctx), (None, Some(query)) => distances_to_constant(&a, query.scalar(), ctx), (None, None) => { - let ag = geometries(&a, ctx)?; - let bg = geometries(&b, ctx)?; - vortex_ensure!( - ag.len() == bg.len(), + vortex_ensure_eq!( + a.len(), + b.len(), "geo distance: operand length mismatch {} vs {}", - ag.len(), - bg.len() + a.len(), + b.len() ); + let ag = geometries(&a, ctx)?; + let bg = geometries(&b, ctx)?; let distances = ag.iter().zip(&bg).map(|(x, y)| Euclidean.distance(x, y)); Ok(PrimitiveArray::from_iter(distances).into_array()) } @@ -116,8 +119,8 @@ impl ScalarFnVTable for GeoDistance { } } -/// Distance from each row of `operand` to a constant `query` geometry, decoded once and broadcast. -/// Distance is symmetric, so this serves a constant on either side. +/// Distance from each row of `operand` to the constant `query` geometry. Distance is symmetric, +/// so this serves a constant on either side. fn distances_to_constant( operand: &ArrayRef, query: &Scalar, @@ -137,6 +140,11 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; use vortex_error::VortexResult; use super::GeoDistance; @@ -218,4 +226,23 @@ mod tests { assert_eq!(distances(distance, &mut ctx)?, vec![5.0, 5.0, 5.0]); Ok(()) } + + /// Geometry arrays are never nullable, so a nullable operand dtype is rejected. + #[test] + fn nullable_operand_is_rejected() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let result = GeoDistance.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype]); + assert!(result.is_err()); + Ok(()) + } + + /// A non-geometry operand dtype is rejected up front, before execution. + #[test] + fn non_geometry_operand_is_rejected() -> VortexResult<()> { + let geo = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let numeric = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = GeoDistance.return_dtype(&EmptyOptions, &[geo, numeric]); + assert!(result.is_err()); + Ok(()) + } } diff --git a/vortex-geo/src/scalar_fn/intersects.rs b/vortex-geo/src/scalar_fn/intersects.rs new file mode 100644 index 00000000000..801745dc277 --- /dev/null +++ b/vortex-geo/src/scalar_fn/intersects.rs @@ -0,0 +1,334 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Intersects`: OGC intersection test between two native geometries. + +use geo::Intersects; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::geometries; +use crate::extension::single_geometry; +use crate::extension::validate_geometry_operands; + +/// OGC `ST_Intersects` (not disjoint; boundary contact counts) between two native geometry +/// operands, each a column or a constant literal. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct GeoIntersects; + +impl GeoIntersects { + /// A lazy `ScalarFnArray` computing per-row whether operands `a` and `b` intersect; either may + /// be constant. The output length is taken from `a`. + pub fn try_new_array(a: ArrayRef, b: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(GeoIntersects, EmptyOptions).erased(), + vec![a, b], + ) + } +} + +impl ScalarFnVTable for GeoIntersects { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.geo.intersects"); + *ID + } + + fn serialize(&self, _: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + Ok(EmptyOptions) + } + + fn arity(&self, _: &Self::Options) -> Arity { + Arity::Exact(2) + } + + fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("a"), + 1 => ChildName::from("b"), + _ => unreachable!("intersects has exactly two children"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + validate_geometry_operands(dtypes)?; + Ok(DType::Bool(Nullability::NonNullable)) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let a = args.get(0)?; + let b = args.get(1)?; + match (a.as_opt::(), b.as_opt::()) { + (Some(qa), Some(qb)) => { + let ga = single_geometry(qa.scalar(), ctx)?; + let gb = single_geometry(qb.scalar(), ctx)?; + Ok(ConstantArray::new( + Scalar::bool(ga.intersects(&gb), Nullability::NonNullable), + a.len(), + ) + .into_array()) + } + (Some(query), None) => intersects_constant(&b, query.scalar(), ctx), + (None, Some(query)) => intersects_constant(&a, query.scalar(), ctx), + (None, None) => { + vortex_ensure_eq!( + a.len(), + b.len(), + "geo intersects: operand length mismatch {} vs {}", + a.len(), + b.len() + ); + let ag = geometries(&a, ctx)?; + let bg = geometries(&b, ctx)?; + let hits = ag.iter().zip(&bg).map(|(x, y)| x.intersects(y)); + Ok(BoolArray::from_iter(hits).into_array()) + } + } + } +} + +/// Whether each row of `operand` intersects the constant `query` geometry. Intersection is +/// symmetric, so this serves a constant on either side. +fn intersects_constant( + operand: &ArrayRef, + query: &Scalar, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let query = single_geometry(query, ctx)?; + let geoms = geometries(operand, ctx)?; + let hits = geoms.iter().map(|g| g.intersects(&query)); + Ok(BoolArray::from_iter(hits).into_array()) +} + +#[cfg(test)] +mod tests { + use geo_types::Coord; + use geo_types::Geometry; + use geo_types::LineString; + use geo_types::MultiPolygon; + use geo_types::Polygon; + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use wkb::writer::WriteOptions; + + use super::GeoIntersects; + use crate::test_harness::point_column; + + /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. + fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon { + Polygon::new( + LineString::from(vec![(x0, y0), (x1, y0), (x1, y1), (x0, y1), (x0, y0)]), + vec![], + ) + } + + /// The query region shared by the point tests: a `10x10` square with a `4..6` square hole. + fn donut() -> Geometry { + Geometry::Polygon(Polygon::new( + LineString::from(vec![(0.0, 0.0), (10.0, 0.0), (10.0, 10.0), (0.0, 10.0)]), + vec![LineString::from(vec![ + (4.0, 4.0), + (6.0, 4.0), + (6.0, 6.0), + (4.0, 6.0), + ])], + )) + } + + /// Probes for [`donut`], one per verdict class: interior, exterior, on the outer boundary, + /// and inside the hole. + fn donut_probes() -> VortexResult { + point_column(vec![2.0, 20.0, 0.0, 5.0], vec![2.0, 20.0, 5.0, 5.0]) + } + + /// [`donut_probes`]'s verdicts: interior and boundary contact intersect; exterior and + /// in-hole do not. + const DONUT_EXPECTED: [bool; 4] = [true, false, true, false]; + + /// A constant column of length `len`, every row the native form of `geometry`. + fn geometry_constant(geometry: &Geometry, len: usize) -> VortexResult { + let mut buf = Vec::new(); + wkb::writer::write_geometry(&mut buf, geometry, &WriteOptions::default()) + .map_err(|e| vortex_err!("writing WKB failed: {e}"))?; + let scalar = crate::extension::native_geometry_scalar_from_wkb(&buf)? + .ok_or_else(|| vortex_err!("unsupported geometry type"))?; + Ok(ConstantArray::new(scalar, len).into_array()) + } + + /// Materialize `array` so it is no longer a `Constant`, forcing the non-constant kernel paths. + fn materialize(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array.execute::(ctx)?.into_array()) + } + + /// Execute `GeoIntersects(a, b)` and assert the per-row verdicts equal `expected`. + fn assert_intersects( + a: ArrayRef, + b: ArrayRef, + expected: impl IntoIterator, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let intersects = GeoIntersects::try_new_array(a, b)?.into_array(); + assert_arrays_eq!(intersects, BoolArray::from_iter(expected), &mut ctx); + Ok(()) + } + + // The tests cover each `execute` dispatch arm in match order, then the edge cases. + + /// Constant vs constant: overlapping and touching (edge or corner) polygons intersect, + /// disjoint ones do not; every output row carries the same verdict. + #[rstest] + #[case::overlapping(rect_polygon(2.0, 2.0, 6.0, 6.0), true)] + #[case::touching_edge(rect_polygon(4.0, 0.0, 8.0, 4.0), true)] + #[case::touching_corner(rect_polygon(4.0, 4.0, 8.0, 8.0), true)] + #[case::disjoint(rect_polygon(20.0, 20.0, 24.0, 24.0), false)] + fn constant_vs_constant_polygons( + #[case] other: Polygon, + #[case] expected: bool, + ) -> VortexResult<()> { + let a = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; + let b = geometry_constant(&Geometry::Polygon(other), 3)?; + assert_intersects(a, b, [expected; 3]) + } + + /// Point column vs constant polygon, with the constant on either side since intersection + /// is symmetric: the [`DONUT_EXPECTED`] verdicts. + #[test] + fn point_column_vs_constant_polygon() -> VortexResult<()> { + assert_intersects( + donut_probes()?, + geometry_constant(&donut(), 4)?, + DONUT_EXPECTED, + )?; + assert_intersects( + geometry_constant(&donut(), 4)?, + donut_probes()?, + DONUT_EXPECTED, + ) + } + + /// Point column vs constant multipolygon: a point in either part intersects, a point in + /// neither does not. + #[test] + fn point_column_vs_constant_multipolygon() -> VortexResult<()> { + let query = Geometry::MultiPolygon(MultiPolygon::new(vec![ + rect_polygon(0.0, 0.0, 2.0, 2.0), + rect_polygon(10.0, 10.0, 12.0, 12.0), + ])); + let points = point_column(vec![1.0, 11.0, 5.0], vec![1.0, 11.0, 5.0])?; + assert_intersects(points, geometry_constant(&query, 3)?, [true, true, false]) + } + + /// Polygon column vs constant: the same pairs and verdicts as + /// [`constant_vs_constant_polygons`]. + #[rstest] + #[case::overlapping(rect_polygon(2.0, 2.0, 6.0, 6.0), true)] + #[case::touching_edge(rect_polygon(4.0, 0.0, 8.0, 4.0), true)] + #[case::disjoint(rect_polygon(20.0, 20.0, 24.0, 24.0), false)] + fn polygon_column_vs_constant( + #[case] other: Polygon, + #[case] expected: bool, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let column = materialize(geometry_constant(&Geometry::Polygon(other), 2)?, &mut ctx)?; + let query = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?; + assert_intersects(column, query, [expected; 2]) + } + + /// Point column vs point column: rows intersect exactly at equal coordinates. + #[test] + fn point_column_vs_point_column() -> VortexResult<()> { + let a = point_column(vec![0.0, 1.0], vec![0.0, 1.0])?; + let b = point_column(vec![0.0, 2.0], vec![0.0, 1.0])?; + assert_intersects(a, b, [true, false]) + } + + /// Point column vs polygon column pairwise: agrees with point column vs constant on the + /// same data. + #[test] + fn columns_agree_with_constant() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let constant = geometry_constant(&donut(), 4)?; + let column = materialize(constant.clone(), &mut ctx)?; + + let against_constant = + GeoIntersects::try_new_array(donut_probes()?, constant)?.into_array(); + let pairwise = GeoIntersects::try_new_array(donut_probes()?, column)?.into_array(); + + assert_arrays_eq!(against_constant, pairwise, &mut ctx); + Ok(()) + } + + /// An empty constant geometry intersects nothing. + #[test] + fn empty_constant_intersects_nothing() -> VortexResult<()> { + let empty = Geometry::LineString(LineString::new(Vec::::new())); + let points = point_column(vec![0.0, 1.0], vec![0.0, 1.0])?; + assert_intersects(points, geometry_constant(&empty, 2)?, [false, false]) + } + + /// Geometry arrays are never nullable, so a nullable operand dtype is rejected. + #[test] + fn nullable_operand_is_rejected() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let result = GeoIntersects.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype]); + assert!(result.is_err()); + Ok(()) + } + + /// A non-geometry operand dtype is rejected up front, before execution. + #[test] + fn non_geometry_operand_is_rejected() -> VortexResult<()> { + let geo = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let numeric = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = GeoIntersects.return_dtype(&EmptyOptions, &[geo, numeric]); + assert!(result.is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/mod.rs b/vortex-geo/src/scalar_fn/mod.rs index 385208f1991..2246f82621b 100644 --- a/vortex-geo/src/scalar_fn/mod.rs +++ b/vortex-geo/src/scalar_fn/mod.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Geometry scalar functions over the [`Point`](crate::extension::Point) type. +//! Geometry scalar functions over the native geometry extension types. +pub mod contains; pub mod distance; +pub mod intersects; diff --git a/vortex-geo/src/test_harness.rs b/vortex-geo/src/test_harness.rs index 2e9e7f43c27..79dc2d3361d 100644 --- a/vortex-geo/src/test_harness.rs +++ b/vortex-geo/src/test_harness.rs @@ -6,29 +6,175 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; use crate::extension::GeoMetadata; +use crate::extension::LineString; +use crate::extension::MultiLineString; +use crate::extension::MultiPoint; +use crate::extension::MultiPolygon; use crate::extension::Point; +use crate::extension::Polygon; +use crate::extension::Rect; +use crate::extension::box_storage_dtype; use crate::extension::coordinate::Coordinate; +use crate::extension::coordinate::Dimension; use crate::extension::coordinate::coordinate_from_struct; +use crate::extension::linestring_storage_dtype; +use crate::extension::multilinestring_storage_dtype; +use crate::extension::multipoint_storage_dtype; +use crate::extension::multipolygon_storage_dtype; +use crate::extension::polygon_storage_dtype; -/// A `Point` column (CRS `EPSG:4326`) over the given x/y coordinates. -pub(crate) fn point_column(xs: Vec, ys: Vec) -> VortexResult { - let storage = StructArray::from_fields(&[ +/// A fresh session with the geospatial types, functions, and pruning rules registered. +pub(crate) fn geo_session() -> VortexSession { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +} + +/// The WGS 84 (`EPSG:4326`) metadata tagged onto test geometry columns. +fn wgs84() -> GeoMetadata { + GeoMetadata { + crs: Some("EPSG:4326".to_string()), + } +} + +/// A coordinate `Struct` over the parallel x/y buffers. +fn xy_struct(xs: Vec, ys: Vec) -> VortexResult { + Ok(StructArray::from_fields(&[ ("x", PrimitiveArray::from_iter(xs).into_array()), ("y", PrimitiveArray::from_iter(ys).into_array()), ])? - .into_array(); - let metadata = GeoMetadata { - crs: Some("EPSG:4326".to_string()), + .into_array()) +} + +/// A list offset as `i32`. +fn offset(n: usize) -> VortexResult { + i32::try_from(n).map_err(|_| vortex_err!("geometry offset overflow")) +} + +/// Wrap `elements` — built from `rows` flattened one level — in a non-nullable `List` with one +/// entry per row. +fn nest(rows: &[Vec], elements: ArrayRef) -> VortexResult { + let mut offsets = vec![0i32]; + let mut len = 0usize; + for row in rows { + len += row.len(); + offsets.push(offset(len)?); + } + Ok(ListArray::try_new( + elements, + PrimitiveArray::from_iter(offsets).into_array(), + Validity::NonNullable, + )? + .into_array()) +} + +/// `List>` storage: one list of `(x, y)` vertices per row. +fn vertex_lists(rows: &[Vec<(f64, f64)>]) -> VortexResult { + let (xs, ys) = rows.iter().flatten().copied().unzip(); + nest(rows, xy_struct(xs, ys)?) +} + +/// `List>>` storage: one list of vertex lists per row. +fn vertex_list_lists(rows: &[Vec>]) -> VortexResult { + let inner: Vec> = rows.iter().flatten().cloned().collect(); + nest(rows, vertex_lists(&inner)?) +} + +/// Wrap `storage` in the geometry extension `vtable` (CRS `EPSG:4326`) with the canonical +/// `storage_dtype` of that type. +fn geo_column + Default>( + storage: ArrayRef, + storage_dtype: DType, +) -> VortexResult { + let dtype = ExtDType::::try_new(wgs84(), storage_dtype)?; + Ok(ExtensionArray::try_new(dtype.erased(), storage)?.into_array()) +} + +/// A `Point` column over the given x/y coordinates, stored as `Struct`. +pub(crate) fn point_column(xs: Vec, ys: Vec) -> VortexResult { + let storage = xy_struct(xs, ys)?; + let storage_dtype = storage.dtype().clone(); + geo_column::(storage, storage_dtype) +} + +/// A `LineString` column: each line a list of `(x, y)` vertices, stored as `List>`. +pub(crate) fn linestring_column(lines: Vec>) -> VortexResult { + geo_column::( + vertex_lists(&lines)?, + linestring_storage_dtype(Dimension::Xy, Nullability::NonNullable), + ) +} + +/// A `MultiPoint` column: each row a list of `(x, y)` points, stored as `List>`. +pub(crate) fn multipoint_column(points: Vec>) -> VortexResult { + geo_column::( + vertex_lists(&points)?, + multipoint_storage_dtype(Dimension::Xy, Nullability::NonNullable), + ) +} + +/// A `Polygon` column: each polygon a list of rings, each ring a list of `(x, y)` vertices, +/// stored as `List>>`. +pub(crate) fn polygon_column(polygons: Vec>>) -> VortexResult { + geo_column::( + vertex_list_lists(&polygons)?, + polygon_storage_dtype(Dimension::Xy, Nullability::NonNullable), + ) +} + +/// A `MultiLineString` column: each row a list of lines, stored as `List>>`. +pub(crate) fn multilinestring_column( + multilines: Vec>>, +) -> VortexResult { + geo_column::( + vertex_list_lists(&multilines)?, + multilinestring_storage_dtype(Dimension::Xy, Nullability::NonNullable), + ) +} + +/// One multipolygon: polygons → rings → `(x, y)` vertices. +pub(crate) type MultiPolygonRings = Vec>>; + +/// A `MultiPolygon` column, stored as `List>>>`. +pub(crate) fn multipolygon_column(multipolygons: Vec) -> VortexResult { + let polygons: Vec>> = multipolygons.iter().flatten().cloned().collect(); + geo_column::( + nest(&multipolygons, vertex_list_lists(&polygons)?)?, + multipolygon_storage_dtype(Dimension::Xy, Nullability::NonNullable), + ) +} + +/// A 2D `Rect` (`geoarrow.box`) column over `(xmin, ymin, xmax, ymax)` boxes, stored as +/// `Struct`. +pub(crate) fn rect_column(boxes: Vec<(f64, f64, f64, f64)>) -> VortexResult { + let field = |select: fn(&(f64, f64, f64, f64)) -> f64| { + PrimitiveArray::from_iter(boxes.iter().map(select)).into_array() }; - let dtype = ExtDType::::try_new(metadata, storage.dtype().clone())?; - Ok(ExtensionArray::new(dtype.erased(), storage).into_array()) + let storage = StructArray::from_fields(&[ + ("xmin", field(|b| b.0)), + ("ymin", field(|b| b.1)), + ("xmax", field(|b| b.2)), + ("ymax", field(|b| b.3)), + ])? + .into_array(); + geo_column::( + storage, + box_storage_dtype(Dimension::Xy, Nullability::NonNullable), + ) } /// Decode a [`Coordinate`] from an extension-typed point scalar (unwrapped to its coordinate diff --git a/vortex-geo/src/tests/linestring.rs b/vortex-geo/src/tests/linestring.rs new file mode 100644 index 00000000000..46e2be94647 --- /dev/null +++ b/vortex-geo/src/tests/linestring.rs @@ -0,0 +1,88 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow interop for the `vortex.geo.linestring` extension type (`geoarrow.linestring`). + +use std::sync::Arc; + +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType as _; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Dimension as GeoArrowDimension; +use geoarrow::datatypes::LineStringType; +use geoarrow::datatypes::Metadata; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_arrow::ArrowSessionExt; +use vortex_error::VortexResult; + +use super::SESSION; +use crate::extension::LineString; + +/// A `geoarrow.linestring` Arrow field with separated (struct) XY coordinates. +fn linestring_field(name: &str, nullable: bool, crs: Option<&str>) -> Field { + let crs = crs + .map(|crs| Crs::from_unknown_crs_type(crs.to_string())) + .unwrap_or_default(); + let metadata = Arc::new(Metadata::new(crs, None)); + LineStringType::new(GeoArrowDimension::XY, metadata).to_field(name, nullable) +} + +/// An imported `geoarrow.linestring` field maps to the LineString extension dtype, recovering the +/// CRS, the `List>` storage, and nullability. +#[test] +fn import_field_recovers_extension() -> VortexResult<()> { + let field = linestring_field("geom", true, Some("EPSG:4326")); + let dtype = SESSION.arrow().from_arrow_field(&field)?; + + let DType::Extension(ext) = &dtype else { + panic!("expected Extension dtype, got {dtype}"); + }; + assert!(ext.is::()); + assert_eq!( + ext.metadata::().crs.as_deref(), + Some("EPSG:4326") + ); + + // Storage peels one List layer (linestring → coordinates) to the coordinate struct. + let DType::List(coords, nullability) = ext.storage_dtype() else { + panic!("expected List storage, got {}", ext.storage_dtype()); + }; + assert_eq!(*nullability, Nullability::Nullable); + let DType::Struct(fields, _) = coords.as_ref() else { + panic!("expected coordinate Struct"); + }; + let names: Vec<&str> = fields.names().iter().map(|n| n.as_ref()).collect(); + assert_eq!(names, vec!["x", "y"]); + Ok(()) +} + +/// A field with interleaved (`FixedSizeList`) coordinates fails to import. +#[test] +fn import_interleaved_field_fails() { + let linestring_type = LineStringType::new(GeoArrowDimension::XY, Default::default()) + .with_coord_type(CoordType::Interleaved); + let field = linestring_type.to_field("geom", false); + assert!(SESSION.arrow().from_arrow_field(&field).is_err()); +} + +/// A field imported to the LineString dtype and exported back carries the `geoarrow.linestring` +/// extension over its `List` storage. +#[test] +fn export_field_carries_extension() -> VortexResult<()> { + let imported = + SESSION + .arrow() + .from_arrow_field(&linestring_field("geom", false, Some("EPSG:4326")))?; + let field = SESSION.arrow().to_arrow_field("geom", &imported)?; + + assert_eq!(field.extension_type_name(), Some(LineStringType::NAME)); + assert!( + matches!(field.data_type(), DataType::List(_)), + "expected List storage, got {}", + field.data_type() + ); + Ok(()) +} diff --git a/vortex-geo/src/tests/mod.rs b/vortex-geo/src/tests/mod.rs index 546de758eba..3ee1714ce47 100644 --- a/vortex-geo/src/tests/mod.rs +++ b/vortex-geo/src/tests/mod.rs @@ -4,7 +4,12 @@ //! Arrow interop tests for the geospatial extension types, exercising the session wiring set up //! by [`crate::initialize`]. +mod linestring; +mod multilinestring; +mod multipoint; +mod multipolygon; mod point; +mod rect; mod wkb; use std::sync::LazyLock; @@ -12,8 +17,4 @@ use std::sync::LazyLock; use vortex_session::VortexSession; /// A session with the geospatial types and functions registered. -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session -}); +static SESSION: LazyLock = LazyLock::new(crate::test_harness::geo_session); diff --git a/vortex-geo/src/tests/multilinestring.rs b/vortex-geo/src/tests/multilinestring.rs new file mode 100644 index 00000000000..694c77cdf3b --- /dev/null +++ b/vortex-geo/src/tests/multilinestring.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow interop for the `vortex.geo.multilinestring` extension type (`geoarrow.multilinestring`). + +use std::sync::Arc; + +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType as _; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Dimension as GeoArrowDimension; +use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::MultiLineStringType; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_arrow::ArrowSessionExt; +use vortex_error::VortexResult; + +use super::SESSION; +use crate::extension::MultiLineString; + +/// A `geoarrow.multilinestring` Arrow field with separated (struct) XY coordinates. +fn multilinestring_field(name: &str, nullable: bool, crs: Option<&str>) -> Field { + let crs = crs + .map(|crs| Crs::from_unknown_crs_type(crs.to_string())) + .unwrap_or_default(); + let metadata = Arc::new(Metadata::new(crs, None)); + MultiLineStringType::new(GeoArrowDimension::XY, metadata).to_field(name, nullable) +} + +/// An imported `geoarrow.multilinestring` field maps to the MultiLineString extension dtype (not +/// Polygon, despite the identical `List>>` storage), recovering CRS and shape. +#[test] +fn import_field_recovers_extension() -> VortexResult<()> { + let field = multilinestring_field("geom", true, Some("EPSG:4326")); + let dtype = SESSION.arrow().from_arrow_field(&field)?; + + let DType::Extension(ext) = &dtype else { + panic!("expected Extension dtype, got {dtype}"); + }; + assert!(ext.is::()); + assert_eq!( + ext.metadata::().crs.as_deref(), + Some("EPSG:4326") + ); + + // Storage peels two List layers (multilinestring → line strings) to the coordinate struct. + let DType::List(lines, nullability) = ext.storage_dtype() else { + panic!("expected List storage, got {}", ext.storage_dtype()); + }; + assert_eq!(*nullability, Nullability::Nullable); + let DType::List(coords, _) = lines.as_ref() else { + panic!("expected List of line strings"); + }; + let DType::Struct(fields, _) = coords.as_ref() else { + panic!("expected coordinate Struct"); + }; + let names: Vec<&str> = fields.names().iter().map(|n| n.as_ref()).collect(); + assert_eq!(names, vec!["x", "y"]); + Ok(()) +} + +/// A field with interleaved (`FixedSizeList`) coordinates fails to import. +#[test] +fn import_interleaved_field_fails() { + let multilinestring_type = MultiLineStringType::new(GeoArrowDimension::XY, Default::default()) + .with_coord_type(CoordType::Interleaved); + let field = multilinestring_type.to_field("geom", false); + assert!(SESSION.arrow().from_arrow_field(&field).is_err()); +} + +/// A field imported to the MultiLineString dtype and exported back carries the +/// `geoarrow.multilinestring` extension over its `List` storage. +#[test] +fn export_field_carries_extension() -> VortexResult<()> { + let imported = SESSION.arrow().from_arrow_field(&multilinestring_field( + "geom", + false, + Some("EPSG:4326"), + ))?; + let field = SESSION.arrow().to_arrow_field("geom", &imported)?; + + assert_eq!(field.extension_type_name(), Some(MultiLineStringType::NAME)); + assert!( + matches!(field.data_type(), DataType::List(_)), + "expected List storage, got {}", + field.data_type() + ); + Ok(()) +} diff --git a/vortex-geo/src/tests/multipoint.rs b/vortex-geo/src/tests/multipoint.rs new file mode 100644 index 00000000000..366dac73fc2 --- /dev/null +++ b/vortex-geo/src/tests/multipoint.rs @@ -0,0 +1,87 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow interop for the `vortex.geo.multipoint` extension type (`geoarrow.multipoint`). + +use std::sync::Arc; + +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType as _; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Dimension as GeoArrowDimension; +use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::MultiPointType; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_arrow::ArrowSessionExt; +use vortex_error::VortexResult; + +use super::SESSION; +use crate::extension::MultiPoint; + +/// A `geoarrow.multipoint` Arrow field with separated (struct) XY coordinates. +fn multipoint_field(name: &str, nullable: bool, crs: Option<&str>) -> Field { + let crs = crs + .map(|crs| Crs::from_unknown_crs_type(crs.to_string())) + .unwrap_or_default(); + let metadata = Arc::new(Metadata::new(crs, None)); + MultiPointType::new(GeoArrowDimension::XY, metadata).to_field(name, nullable) +} + +/// An imported `geoarrow.multipoint` field maps to the MultiPoint extension dtype (not LineString, +/// despite the identical `List>` storage), recovering the CRS and nullability. +#[test] +fn import_field_recovers_extension() -> VortexResult<()> { + let field = multipoint_field("geom", true, Some("EPSG:4326")); + let dtype = SESSION.arrow().from_arrow_field(&field)?; + + let DType::Extension(ext) = &dtype else { + panic!("expected Extension dtype, got {dtype}"); + }; + assert!(ext.is::()); + assert_eq!( + ext.metadata::().crs.as_deref(), + Some("EPSG:4326") + ); + + let DType::List(coords, nullability) = ext.storage_dtype() else { + panic!("expected List storage, got {}", ext.storage_dtype()); + }; + assert_eq!(*nullability, Nullability::Nullable); + let DType::Struct(fields, _) = coords.as_ref() else { + panic!("expected coordinate Struct"); + }; + let names: Vec<&str> = fields.names().iter().map(|n| n.as_ref()).collect(); + assert_eq!(names, vec!["x", "y"]); + Ok(()) +} + +/// A field with interleaved (`FixedSizeList`) coordinates fails to import. +#[test] +fn import_interleaved_field_fails() { + let multipoint_type = MultiPointType::new(GeoArrowDimension::XY, Default::default()) + .with_coord_type(CoordType::Interleaved); + let field = multipoint_type.to_field("geom", false); + assert!(SESSION.arrow().from_arrow_field(&field).is_err()); +} + +/// A field imported to the MultiPoint dtype and exported back carries the `geoarrow.multipoint` +/// extension over its `List` storage. +#[test] +fn export_field_carries_extension() -> VortexResult<()> { + let imported = + SESSION + .arrow() + .from_arrow_field(&multipoint_field("geom", false, Some("EPSG:4326")))?; + let field = SESSION.arrow().to_arrow_field("geom", &imported)?; + + assert_eq!(field.extension_type_name(), Some(MultiPointType::NAME)); + assert!( + matches!(field.data_type(), DataType::List(_)), + "expected List storage, got {}", + field.data_type() + ); + Ok(()) +} diff --git a/vortex-geo/src/tests/multipolygon.rs b/vortex-geo/src/tests/multipolygon.rs new file mode 100644 index 00000000000..407661f6a43 --- /dev/null +++ b/vortex-geo/src/tests/multipolygon.rs @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow interop for the `vortex.geo.multipolygon` extension type (`geoarrow.multipolygon`). + +use std::sync::Arc; + +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType as _; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Dimension as GeoArrowDimension; +use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::MultiPolygonType; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_arrow::ArrowSessionExt; +use vortex_error::VortexResult; + +use super::SESSION; +use crate::extension::MultiPolygon; + +/// A `geoarrow.multipolygon` Arrow field with separated (struct) XY coordinates. +fn multipolygon_field(name: &str, nullable: bool, crs: Option<&str>) -> Field { + let crs = crs + .map(|crs| Crs::from_unknown_crs_type(crs.to_string())) + .unwrap_or_default(); + let metadata = Arc::new(Metadata::new(crs, None)); + MultiPolygonType::new(GeoArrowDimension::XY, metadata).to_field(name, nullable) +} + +/// An imported `geoarrow.multipolygon` field maps to the MultiPolygon extension dtype, recovering the +/// CRS, the `List>>>` storage, and nullability. +#[test] +fn import_field_recovers_extension() -> VortexResult<()> { + let field = multipolygon_field("geom", true, Some("EPSG:4326")); + let dtype = SESSION.arrow().from_arrow_field(&field)?; + + let DType::Extension(ext) = &dtype else { + panic!("expected Extension dtype, got {dtype}"); + }; + assert!(ext.is::()); + assert_eq!( + ext.metadata::().crs.as_deref(), + Some("EPSG:4326") + ); + + // Storage peels three List layers (multipolygon → polygons → rings) to the coordinate struct. + let DType::List(polygons, nullability) = ext.storage_dtype() else { + panic!("expected List storage, got {}", ext.storage_dtype()); + }; + assert_eq!(*nullability, Nullability::Nullable); + let DType::List(rings, _) = polygons.as_ref() else { + panic!("expected List of polygons"); + }; + let DType::List(coords, _) = rings.as_ref() else { + panic!("expected List of rings"); + }; + let DType::Struct(fields, _) = coords.as_ref() else { + panic!("expected coordinate Struct"); + }; + let names: Vec<&str> = fields.names().iter().map(|n| n.as_ref()).collect(); + assert_eq!(names, vec!["x", "y"]); + Ok(()) +} + +/// A field with interleaved (`FixedSizeList`) coordinates fails to import. +#[test] +fn import_interleaved_field_fails() { + let multipolygon_type = MultiPolygonType::new(GeoArrowDimension::XY, Default::default()) + .with_coord_type(CoordType::Interleaved); + let field = multipolygon_type.to_field("geom", false); + assert!(SESSION.arrow().from_arrow_field(&field).is_err()); +} + +/// A field imported to the MultiPolygon dtype and exported back carries the `geoarrow.multipolygon` +/// extension over its `List` storage. +#[test] +fn export_field_carries_extension() -> VortexResult<()> { + let imported = + SESSION + .arrow() + .from_arrow_field(&multipolygon_field("geom", false, Some("EPSG:4326")))?; + let field = SESSION.arrow().to_arrow_field("geom", &imported)?; + + assert_eq!(field.extension_type_name(), Some(MultiPolygonType::NAME)); + assert!( + matches!(field.data_type(), DataType::List(_)), + "expected List storage, got {}", + field.data_type() + ); + Ok(()) +} diff --git a/vortex-geo/src/tests/point.rs b/vortex-geo/src/tests/point.rs index ff74b01ba01..61a2e98dbe7 100644 --- a/vortex-geo/src/tests/point.rs +++ b/vortex-geo/src/tests/point.rs @@ -20,9 +20,9 @@ use geoarrow::datatypes::Dimension as GeoArrowDimension; use geoarrow::datatypes::Metadata; use geoarrow::datatypes::PointType; use vortex_array::VortexSessionExecute; -use vortex_array::arrow::ArrowSessionExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_err; diff --git a/vortex-geo/src/tests/rect.rs b/vortex-geo/src/tests/rect.rs new file mode 100644 index 00000000000..865892493e7 --- /dev/null +++ b/vortex-geo/src/tests/rect.rs @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow interop for the `vortex.geo.box` extension type (`geoarrow.box`). + +use std::sync::Arc; + +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType as _; +use geoarrow::datatypes::BoxType; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Dimension as GeoArrowDimension; +use geoarrow::datatypes::Metadata; +use vortex_array::VortexSessionExecute; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_arrow::ArrowSessionExt; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::SESSION; +use crate::extension::Rect; +use crate::test_harness::rect_column; + +/// A `geoarrow.box` Arrow field of the given dimension. +fn box_field(name: &str, dim: GeoArrowDimension, nullable: bool, crs: Option<&str>) -> Field { + let crs = crs + .map(|crs| Crs::from_unknown_crs_type(crs.to_string())) + .unwrap_or_default(); + let metadata = Arc::new(Metadata::new(crs, None)); + BoxType::new(dim, metadata).to_field(name, nullable) +} + +/// The exported Arrow field carries the `geoarrow.box` extension over the `Struct` layout. +#[test] +fn export_field_carries_extension() -> VortexResult<()> { + let array = rect_column(vec![(0.0, 0.0, 1.0, 1.0)])?; + let field = SESSION.arrow().to_arrow_field("bbox", array.dtype())?; + + assert_eq!(field.extension_type_name(), Some(BoxType::NAME)); + let DataType::Struct(fields) = field.data_type() else { + panic!("expected Struct, got {}", field.data_type()); + }; + let names: Vec<&str> = fields.iter().map(|f| f.name().as_str()).collect(); + assert_eq!(names, vec!["xmin", "ymin", "xmax", "ymax"]); + Ok(()) +} + +/// An imported `geoarrow.box` field maps to the Rect extension dtype, recovering the CRS, +/// box field names, and nullability. +#[test] +fn import_field_recovers_extension() -> VortexResult<()> { + let field = box_field("bbox", GeoArrowDimension::XY, true, Some("EPSG:4326")); + let dtype = SESSION.arrow().from_arrow_field(&field)?; + + let DType::Extension(ext) = &dtype else { + panic!("expected Extension dtype, got {dtype}"); + }; + assert!(ext.is::()); + assert_eq!(ext.metadata::().crs.as_deref(), Some("EPSG:4326")); + + let DType::Struct(fields, nullability) = ext.storage_dtype() else { + panic!("expected Struct storage, got {}", ext.storage_dtype()); + }; + assert_eq!(*nullability, Nullability::Nullable); + let names: Vec<&str> = fields.names().iter().map(|n| n.as_ref()).collect(); + assert_eq!(names, vec!["xmin", "ymin", "xmax", "ymax"]); + Ok(()) +} + +/// A `Rect` column exported to Arrow and imported back is unchanged, including the CRS and the +/// box corners. +#[test] +fn roundtrips_through_arrow() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let original = rect_column(vec![(0.0, 1.0, 2.0, 3.0), (-5.0, -5.0, 5.0, 5.0)])?; + + let target = box_field("bbox", GeoArrowDimension::XY, false, Some("EPSG:4326")); + let exported = SESSION + .arrow() + .execute_arrow(original, Some(&target), &mut ctx)?; + let reimported = SESSION.arrow().from_arrow_array(exported, &target)?; + + let ext = reimported + .dtype() + .as_extension_opt() + .ok_or_else(|| vortex_err!("expected Extension dtype"))?; + assert!(ext.is::()); + assert_eq!(ext.metadata::().crs.as_deref(), Some("EPSG:4326")); + + let mut corner = |row: usize, name: &str| -> VortexResult { + let scalar = reimported.execute_scalar(row, &mut ctx)?; + let storage = scalar + .as_extension_opt() + .ok_or_else(|| vortex_err!("expected extension scalar"))? + .to_storage_scalar(); + f64::try_from( + &storage + .as_struct() + .field(name) + .ok_or_else(|| vortex_err!("missing {name}"))?, + ) + }; + assert_eq!(corner(0, "xmin")?, 0.0); + assert_eq!(corner(0, "ymax")?, 3.0); + assert_eq!(corner(1, "xmin")?, -5.0); + assert_eq!(corner(1, "ymax")?, 5.0); + Ok(()) +} + +/// The existing geo scalar functions run on a `Rect` operand via the shared `geometries()` decode, +/// producing the same results as the equivalent polygon: a box `(0,0)-(10,10)` against interior +/// point `(5,5)` and exterior point `(20,20)`. +#[test] +fn scalar_functions_run_on_rect() -> VortexResult<()> { + use vortex_array::Canonical; + use vortex_array::IntoArray; + use vortex_array::arrays::BoolArray; + use vortex_array::assert_arrays_eq; + + use crate::scalar_fn::contains::GeoContains; + use crate::scalar_fn::distance::GeoDistance; + use crate::scalar_fn::intersects::GeoIntersects; + use crate::test_harness::point_column; + + let mut ctx = SESSION.create_execution_ctx(); + let bbox = rect_column(vec![(0.0, 0.0, 10.0, 10.0), (0.0, 0.0, 10.0, 10.0)])?; + let points = point_column(vec![5.0, 20.0], vec![5.0, 20.0])?; + + // Distance: 0 to the interior point, >0 to the exterior point. + let distance = GeoDistance::try_new_array(bbox.clone(), points.clone())?.into_array(); + let distance = distance.execute::(&mut ctx)?.into_primitive(); + let distances = distance.as_slice::(); + assert_eq!(distances[0], 0.0); + assert!(distances[1] > 0.0); + + // Intersects / Contains: true for the interior point, false for the exterior one. + let intersects = GeoIntersects::try_new_array(bbox.clone(), points.clone())?.into_array(); + assert_arrays_eq!(intersects, BoolArray::from_iter([true, false]), &mut ctx); + + let contains = GeoContains::try_new_array(bbox, points)?.into_array(); + assert_arrays_eq!(contains, BoolArray::from_iter([true, false]), &mut ctx); + Ok(()) +} diff --git a/vortex-geo/src/tests/wkb.rs b/vortex-geo/src/tests/wkb.rs index 4ef3d6e2ba8..d61c48bdc32 100644 --- a/vortex-geo/src/tests/wkb.rs +++ b/vortex-geo/src/tests/wkb.rs @@ -27,11 +27,11 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::varbin::builder::VarBinBuilder; -use vortex_array::arrow::ArrowSessionExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; +use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_err; use wkb::writer::WriteOptions; diff --git a/vortex-io/src/runtime/current.rs b/vortex-io/src/runtime/current.rs index c8ff91b4aec..93b698b43a2 100644 --- a/vortex-io/src/runtime/current.rs +++ b/vortex-io/src/runtime/current.rs @@ -1,11 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::future::Future; use std::sync::Arc; use futures::Stream; use futures::StreamExt; use futures::stream::BoxStream; +use parking_lot::Mutex; use smol::block_on; use crate::runtime::BlockingRuntime; @@ -67,22 +69,21 @@ impl CurrentThreadRuntime { // This allows multiple worker threads to drive the execution while all waiting for results // on the channel. let (result_tx, result_rx) = kanal::bounded_async(1); - self.executor - .spawn(async move { - futures::pin_mut!(stream); - while let Some(item) = stream.next().await { - // If all receivers are dropped, we stop driving the stream. - if let Err(e) = result_tx.send(item).await { - tracing::trace!("all receivers dropped, stopping stream: {}", e); - break; - } + let driver = self.executor.spawn(async move { + futures::pin_mut!(stream); + while let Some(item) = stream.next().await { + // If all receivers are dropped, we stop driving the stream. + if let Err(e) = result_tx.send(item).await { + tracing::trace!("all receivers dropped, stopping stream: {}", e); + break; } - }) - .detach(); + } + }); ThreadSafeIterator { executor: Arc::clone(&self.executor), results: result_rx, + driver: Arc::new(Mutex::new(Some(driver))), } } } @@ -132,6 +133,10 @@ impl Iterator for CurrentThreadIterator<'_, T> { pub struct ThreadSafeIterator { executor: Arc>, results: kanal::AsyncReceiver, + /// Handle to the task driving the stream. Once the stream ends, the first consumer to + /// observe it joins the task so a panic raised while driving the stream is re-raised rather + /// than silently ending the iterator. + driver: Arc>>>, } // Manual clone implementation since `T` does not need to be `Clone`. @@ -140,6 +145,7 @@ impl Clone for ThreadSafeIterator { Self { executor: Arc::clone(&self.executor), results: self.results.clone(), + driver: Arc::clone(&self.driver), } } } @@ -148,17 +154,32 @@ impl Iterator for ThreadSafeIterator { type Item = T; fn next(&mut self) -> Option { - block_on(self.executor.run(self.results.recv())).ok() + match block_on(self.executor.run(self.results.recv())) { + Ok(item) => Some(item), + // The result channel closes when the driver task finishes. Join the task so a panic + // raised while driving the stream is re-raised here instead of being lost. The first + // consumer to observe closure joins it; any later consumer just sees the stream end. + Err(_) => { + let task = self.driver.lock().take(); + if let Some(task) = task { + block_on(self.executor.run(task)); + } + None + } + } } } #[expect(clippy::if_then_some_else_none)] // Clippy is wrong when if/else has await. #[cfg(test)] mod tests { + use std::any::Any; + use std::panic::AssertUnwindSafe; use std::sync::Arc; use std::sync::Barrier; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + use std::task::Poll; use std::thread; use std::time::Duration; @@ -261,6 +282,139 @@ mod tests { assert_eq!(collected, (0..total_items).collect::>()); } + #[test] + fn test_block_on_stream_thread_safe_propagates_driver_panic() { + let runtime = CurrentThreadRuntime::new(); + let mut iter = runtime.block_on_stream_thread_safe(|_h| { + stream::poll_fn(|_| -> Poll> { + panic!("stream driver panic"); + }) + .boxed() + }); + + let panic = std::panic::catch_unwind(AssertUnwindSafe(|| iter.next())) + .expect_err("stream panic must propagate through iterator"); + let message = panic + .downcast_ref::<&'static str>() + .copied() + .or_else(|| panic.downcast_ref::().map(String::as_str)) + .unwrap_or(""); + assert!(message.contains("stream driver panic")); + } + + fn panic_message(panic: &(dyn Any + Send)) -> &str { + panic + .downcast_ref::<&'static str>() + .copied() + .or_else(|| panic.downcast_ref::().map(String::as_str)) + .unwrap_or("") + } + + // A driver panic must propagate on *every* run, regardless of executor scheduling. Running + // the scenario many times guards against a return to timing-dependent propagation. + #[test] + fn test_block_on_stream_thread_safe_panic_propagation_is_deterministic() { + for i in 0..2000 { + let mut iter = CurrentThreadRuntime::new().block_on_stream_thread_safe(|_h| { + stream::poll_fn(|_| -> Poll> { + panic!("deterministic driver panic"); + }) + .boxed() + }); + + let outcome = std::panic::catch_unwind(AssertUnwindSafe(|| iter.next())); + assert!( + outcome.is_err(), + "driver panic was swallowed on iteration {i}: next() returned {:?}", + outcome.ok().flatten(), + ); + } + } + + // A panic after some items were already produced must still surface, not be seen as a clean + // end of stream. + #[test] + fn test_block_on_stream_thread_safe_panic_after_items() { + let mut emitted = 0usize; + let iter = CurrentThreadRuntime::new().block_on_stream_thread_safe(move |_h| { + stream::poll_fn(move |_| -> Poll> { + if emitted < 3 { + emitted += 1; + Poll::Ready(Some(emitted)) + } else { + panic!("driver panic after items"); + } + }) + .boxed() + }); + + // Drain the iterator. The terminal event must be a propagated panic, never a clean + // `None`. This avoids depending on exactly how many buffered items survive channel close. + let outcome = std::panic::catch_unwind(AssertUnwindSafe(move || iter.collect::>())); + match outcome { + Ok(items) => panic!("driver panic was swallowed; stream ended cleanly with {items:?}"), + Err(panic) => assert!(panic_message(&*panic).contains("driver panic after items")), + } + } + + // With multiple consumers, a driver panic must reach *every* consumer that observes the end + // of the stream; it must never be swallowed by all of them. Exactly one consumer joins the + // driver and observes the panic; the rest see the stream end. + #[test] + fn test_block_on_stream_thread_safe_multi_consumer_panic_surfaced() { + let iter = CurrentThreadRuntime::new().block_on_stream_thread_safe(|_h| { + stream::poll_fn(|_| -> Poll> { + panic!("multi consumer driver panic"); + }) + .boxed() + }); + + let num_threads = 4; + let barrier = Arc::new(Barrier::new(num_threads)); + let panics = Arc::new(AtomicUsize::new(0)); + + let handles: Vec<_> = (0..num_threads) + .map(|_| { + let mut iter = iter.clone(); + let barrier = Arc::clone(&barrier); + let panics = Arc::clone(&panics); + thread::spawn(move || { + barrier.wait(); + match std::panic::catch_unwind(AssertUnwindSafe(|| iter.next())) { + // The driver panicked before producing anything, so a clean end is the + // only non-panic outcome a consumer may observe. + Ok(None) => {} + Ok(Some(_)) => panic!("no item was produced before the driver panicked"), + Err(panic) => { + assert!(panic_message(&*panic).contains("multi consumer driver panic")); + panics.fetch_add(1, Ordering::SeqCst); + } + } + }) + }) + .collect(); + + for handle in handles { + handle.join().expect("consumer thread panicked uncaught"); + } + + // The panic surfaces to exactly one consumer and is never swallowed by all of them. + assert_eq!(panics.load(Ordering::SeqCst), 1); + } + + // Clean completion must return `None` with no spurious panic. + #[test] + fn test_block_on_stream_thread_safe_clean_completion_returns_none() { + let mut iter = CurrentThreadRuntime::new() + .block_on_stream_thread_safe(|_h| stream::iter(vec![1usize, 2, 3]).boxed()); + + assert_eq!(iter.next(), Some(1)); + assert_eq!(iter.next(), Some(2)); + assert_eq!(iter.next(), Some(3)); + assert_eq!(iter.next(), None); + assert_eq!(iter.next(), None); + } + #[test] fn test_block_on_stream_concurrent_clone_and_drive() { let num_items = 50; diff --git a/vortex-io/src/runtime/handle.rs b/vortex-io/src/runtime/handle.rs index aa9f0c2100a..bebbaac5f97 100644 --- a/vortex-io/src/runtime/handle.rs +++ b/vortex-io/src/runtime/handle.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; +use std::panic::AssertUnwindSafe; use std::pin::Pin; use std::sync::Arc; use std::sync::Weak; @@ -77,8 +79,11 @@ impl Handle { let span = tracing::trace_span!(target: "vortex_io::spawn", "spawn"); let abort_handle = self.runtime().spawn( async move { + // Catch a panic so it can be re-raised on the joining side (see `Task::poll`) + // rather than being lost, matching `tokio::JoinError` / `smol::Task` semantics. + let output = AssertUnwindSafe(f).catch_unwind().await; // Task::detach allows the receiver to be dropped, so we ignore send errors. - drop(send.send(f.await)); + drop(send.send(output)); } .instrument(span) .boxed(), @@ -116,8 +121,10 @@ impl Handle { let span = tracing::trace_span!(target: "vortex_io::spawn_io", "spawn_io"); let abort_handle = self.runtime().spawn_io( async move { + // See `spawn`: catch a panic so it re-raises on the joining side. + let output = AssertUnwindSafe(f).catch_unwind().await; // Task::detach allows the receiver to be dropped, so we ignore send errors. - drop(send.send(f.await)); + drop(send.send(output)); } .instrument(span) .boxed(), @@ -148,8 +155,10 @@ impl Handle { let _guard = span.enter(); // Optimistically avoid the work if the result won't be used. if !send.is_closed() { + // Catch a panic so it re-raises on the joining side (see `Task::poll`). + let output = std::panic::catch_unwind(AssertUnwindSafe(f)); // Task::detach allows the receiver to be dropped, so we ignore send errors. - drop(send.send(f())); + drop(send.send(output)); } })); Task { @@ -170,8 +179,10 @@ impl Handle { let _guard = span.enter(); // Optimistically avoid the work if the result won't be used. if !send.is_closed() { + // Catch a panic so it re-raises on the joining side (see `Task::poll`). + let output = std::panic::catch_unwind(AssertUnwindSafe(f)); // Task::detach allows the receiver to be dropped, so we ignore send errors. - drop(send.send(f())); + drop(send.send(output)); } })); Task { @@ -181,13 +192,30 @@ impl Handle { } } +/// The value carried from a spawned task back to its [`Task`] handle: either the task's output, +/// or the panic payload if the task panicked, so it can be re-raised on the joining side. +type TaskOutput = Result>; + +/// The terminal outcome of joining a spawned [`Task`] with [`Task::poll_join`], without +/// re-raising a panic. +pub enum JoinOutcome { + /// The task ran to completion and produced this value. + Completed(T), + /// The task panicked. Carries the original panic payload, ready to be re-raised with + /// [`std::panic::resume_unwind`]. + Panicked(Box), + /// The runtime dropped the task's future before it completed, for example because the task + /// was aborted or the runtime shut down. No value or panic payload is available. + Aborted, +} + /// A handle to a spawned Task. /// /// If this handle is dropped, the task is cancelled where possible. In order to allow the task to /// continue running in the background, call [`Task::detach`]. #[must_use = "When a Task is dropped without being awaited, it is cancelled"] pub struct Task { - recv: oneshot::AsyncReceiver, + recv: oneshot::AsyncReceiver>, abort_handle: Option, } @@ -197,24 +225,43 @@ impl Task { pub fn detach(mut self) { drop(self.abort_handle.take()); } + + /// Poll the task to completion, reporting the terminal state as a [`JoinOutcome`] instead of + /// re-raising a panic or panicking on abort. + /// + /// This lets a caller distinguish a task panic (which it may re-raise via + /// [`std::panic::resume_unwind`]) from the runtime dropping the task before it completed — for + /// example a runtime shutting down while work is in flight — which is benign. The [`Future`] + /// implementation, by contrast, re-raises a panic and itself panics on abort. + pub fn poll_join(&mut self, cx: &mut Context<'_>) -> Poll> { + match ready!(self.recv.poll_unpin(cx)) { + Ok(Ok(output)) => Poll::Ready(JoinOutcome::Completed(output)), + Ok(Err(panic)) => Poll::Ready(JoinOutcome::Panicked(panic)), + // The result channel closed without delivering a value: the runtime dropped the task's + // future before it completed (for example the task was aborted or the runtime shut + // down). + Err(_recv_err) => Poll::Ready(JoinOutcome::Aborted), + } + } } impl Future for Task { type Output = T; #[expect(clippy::panic)] - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - match ready!(self.recv.poll_unpin(cx)) { - Ok(result) => Poll::Ready(result), - Err(_recv_err) => { - // If the other end of the channel was dropped, it means the runtime dropped - // the future without ever completing it. If the caller aborted this task by - // dropping it, then they wouldn't be able to poll it anymore. - // So we consider a closed channel to be a Runtime programming error and therefore - // we panic. + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + match ready!(self.get_mut().poll_join(cx)) { + JoinOutcome::Completed(output) => Poll::Ready(output), + // The task panicked: re-raise the original panic on the joining side, preserving its + // message and payload rather than collapsing it into a generic error. + JoinOutcome::Panicked(panic) => std::panic::resume_unwind(panic), + JoinOutcome::Aborted => { + // The runtime dropped the task's future before it completed. If the caller aborted + // this task by dropping it, they wouldn't be able to poll it anymore, so we + // consider a closed channel a runtime programming error and panic. // NOTE(ngates): we don't use vortex_panic to avoid printing a useless backtrace. - panic!("Runtime dropped task without completing it, likely it panicked") + panic!("Runtime dropped task without completing it") } } } @@ -228,3 +275,52 @@ impl Drop for Task { } } } + +#[cfg(test)] +mod tests { + use futures::task::noop_waker; + + use super::*; + + // A task whose result channel closed without a value — the runtime dropped its future before + // it sent a result — must report `Aborted`, so callers can treat a benign teardown as a + // recoverable stop rather than a propagated panic. + #[test] + fn poll_join_reports_aborted_when_channel_closed_without_value() { + let (send, recv) = oneshot::channel::>(); + // Simulate the runtime dropping the task's future before it sent a result. + drop(send); + + let mut task = Task::<()> { + recv: recv.into_future(), + abort_handle: None, + }; + + let waker = noop_waker(); + let mut cx = Context::from_waker(&waker); + assert!(matches!( + task.poll_join(&mut cx), + Poll::Ready(JoinOutcome::Aborted) + )); + } + + // A completed task reports its value through `poll_join` rather than re-raising or panicking. + #[test] + fn poll_join_reports_completed_value() { + let (send, recv) = oneshot::channel::>(); + // Ignore the send result: the payload type is not `Debug`, and the receiver is alive. + drop(send.send(Ok(7))); + + let mut task = Task:: { + recv: recv.into_future(), + abort_handle: None, + }; + + let waker = noop_waker(); + let mut cx = Context::from_waker(&waker); + assert!(matches!( + task.poll_join(&mut cx), + Poll::Ready(JoinOutcome::Completed(7)) + )); + } +} diff --git a/vortex-io/src/runtime/tests.rs b/vortex-io/src/runtime/tests.rs index e35d9db6f78..940566b4fc1 100644 --- a/vortex-io/src/runtime/tests.rs +++ b/vortex-io/src/runtime/tests.rs @@ -4,6 +4,7 @@ #![cfg(feature = "tokio")] #![expect(clippy::cast_possible_truncation)] +use std::panic::AssertUnwindSafe; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -18,6 +19,7 @@ use vortex_buffer::ByteBufferMut; use vortex_error::VortexResult; use crate::VortexReadAt; +use crate::runtime::Task; use crate::runtime::single::block_on; use crate::runtime::tokio::TokioRuntime; use crate::std_file::FileReadAt; @@ -225,6 +227,51 @@ async fn test_handle_spawn_cpu() { assert_eq!(counter.load(Ordering::SeqCst), 1); } +fn panic_message(payload: &(dyn std::any::Any + Send)) -> &str { + payload + .downcast_ref::<&str>() + .copied() + .or_else(|| payload.downcast_ref::().map(String::as_str)) + .unwrap_or("") +} + +// Joining a spawned future that panicked must re-raise the *original* panic, not a generic +// "Runtime dropped task" message. +#[test] +fn test_spawned_future_panic_reraises_original() { + let outcome = block_on(|handle| { + async move { + let task: Task<()> = handle.spawn(async move { panic!("original spawn panic") }); + AssertUnwindSafe(task).catch_unwind().await + } + .boxed_local() + }); + + let payload = outcome.expect_err("panicking task must propagate a panic"); + assert!( + panic_message(&*payload).contains("original spawn panic"), + "got: {:?}", + panic_message(&*payload) + ); +} + +// Same for CPU tasks. +#[tokio::test] +async fn test_spawned_cpu_task_panic_reraises_original() { + let handle = TokioRuntime::current(); + let task: Task<()> = handle.spawn_cpu(|| panic!("original cpu panic")); + + let payload = AssertUnwindSafe(task) + .catch_unwind() + .await + .expect_err("panicking cpu task must propagate a panic"); + assert!( + panic_message(&*payload).contains("original cpu panic"), + "got: {:?}", + panic_message(&*payload) + ); +} + // ============================================================================ // Test custom VortexRead implementation // ============================================================================ diff --git a/vortex-jni/Cargo.toml b/vortex-jni/Cargo.toml index d5f5cea61aa..6bb50cf097d 100644 --- a/vortex-jni/Cargo.toml +++ b/vortex-jni/Cargo.toml @@ -20,6 +20,8 @@ categories = { workspace = true } arrow-array = { workspace = true, features = ["ffi"] } arrow-schema = { workspace = true } async-fs = { workspace = true } +async-lock = { workspace = true } +async-trait = { workspace = true } futures = { workspace = true } jni = { workspace = true } object_store = { workspace = true, features = ["aws", "azure", "gcp"] } @@ -30,13 +32,14 @@ tracing = { workspace = true, features = ["std", "log"] } tracing-subscriber = { workspace = true, features = ["env-filter"] } url = { workspace = true } vortex = { workspace = true, features = ["object_store", "files"] } +vortex-arrow = { workspace = true } vortex-parquet-variant = { workspace = true } [dev-dependencies] jni = { workspace = true, features = ["invocation"] } [lib] -crate-type = ["staticlib", "cdylib"] +crate-type = ["cdylib"] [lints] workspace = true diff --git a/vortex-jni/src/data_source.rs b/vortex-jni/src/data_source.rs index 5f244998f67..47d9d83c577 100644 --- a/vortex-jni/src/data_source.rs +++ b/vortex-jni/src/data_source.rs @@ -4,16 +4,10 @@ //! JNI bindings for [`vortex::scan::DataSource`] (see the equivalent types in //! `vortex-ffi/src/data_source.rs`). //! -//! Glob handling mirrors `vortex-duckdb`'s `VortexMultiFileScan`: -//! * full URLs (`s3://...`, `file:///...`) are used as-is, -//! * bare file paths are made absolute and have `.`/`..` components normalized, -//! * filesystems are cached per base URL so repeated globs against the same bucket share -//! a single client. +//! Globs are parsed with [`parse_uri_or_path`], so full URLs (`s3://...`, `file:///...`) +//! and bare file paths are both accepted. Filesystems are cached per base URL so repeated +//! globs against the same bucket share a single client. -use std::path::Component; -use std::path::Path; -use std::path::PathBuf; -use std::path::absolute; use std::sync::Arc; use jni::EnvUnowned; @@ -22,12 +16,14 @@ use jni::objects::JLongArray; use jni::objects::JObject; use jni::objects::JObjectArray; use jni::objects::JString; +use jni::sys::jint; use jni::sys::jlong; use url::Url; use vortex::error::VortexResult; use vortex::error::vortex_err; use vortex::expr::stats::Precision; use vortex::file::multi::MultiFileDataSource; +use vortex::file::multi::parse_uri_or_path; use vortex::io::filesystem::FileSystemRef; use vortex::io::runtime::BlockingRuntime; use vortex::io::session::RuntimeSessionExt; @@ -38,6 +34,7 @@ use crate::RUNTIME; use crate::dtype::export_dtype_to_arrow; use crate::errors::try_or_throw; use crate::file::extract_properties; +use crate::io::JavaFileSystem; use crate::object_store::object_store_fs; use crate::session::session_ref; @@ -91,7 +88,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_open( let glob_urls: Vec = glob_strings .iter() - .map(|g| parse_glob_url(g.as_str())) + .map(|g| parse_uri_or_path(g.as_str())) .collect::>()?; let mut fs_cache: HashMap = HashMap::new(); @@ -120,36 +117,75 @@ pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_open( }) } -/// Parse a glob string into a [`Url`]. Accepts full URLs and bare (relative or absolute) -/// file paths — see the module docs for details. -fn parse_glob_url(glob: &str) -> VortexResult { - // `Url::parse` accepts Windows absolute paths like `C:\foo` as a URL with a - // single-letter scheme (`c`). No real URL scheme is one character, so treat any - // single-letter scheme as a filesystem path instead. - if let Ok(url) = Url::parse(glob) - && url.scheme().len() > 1 - { - return Ok(url); - } - let path = - absolute(Path::new(glob)).map_err(|e| vortex_err!("failed to absolutize {glob}: {e}"))?; - let path = normalize_path(path); - Url::from_file_path(path).map_err(|_| vortex_err!("neither URL nor path: {glob}")) -} +/// Open a data source over caller-provided `dev.vortex.io.NativeReadable` objects. +/// +/// Unlike [`Java_dev_vortex_jni_NativeDataSource_open`], no storage client is created +/// on the native side: every read is an upcall into the corresponding Java object. +/// `paths` are opaque identifiers (typically the original file locations) used for +/// debugging and deduplication, `lengths` are the known file sizes in bytes. +/// `read_concurrency` caps in-flight `readFully` upcalls across *all* files of the +/// data source; values `<= 0` select the library default. +#[unsafe(no_mangle)] +pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_openFiles( + mut env: EnvUnowned, + _class: JClass, + session_ptr: jlong, + readables: JObjectArray, + paths: JObjectArray, + lengths: JLongArray, + read_concurrency: jint, +) -> jlong { + try_or_throw(&mut env, |env| { + let session = unsafe { session_ref(session_ptr) }; -/// Normalize `.` and `..` without touching the filesystem. -fn normalize_path(path: PathBuf) -> PathBuf { - let mut out = PathBuf::new(); - for component in path.components() { - match component { - Component::CurDir => {} - Component::ParentDir => { - out.pop(); + let count = readables.len(env)?; + if count == 0 { + throw_runtime!("no readables provided"); + } + if paths.len(env)? != count || lengths.len(env)? != count { + throw_runtime!("readables, paths, and lengths must have equal length"); + } + + let mut sizes = vec![0 as jlong; count]; + lengths.get_region(env, 0, &mut sizes)?; + + let vm = env.get_java_vm()?; + let concurrency = usize::try_from(read_concurrency).ok().filter(|c| *c > 0); + let mut fs = JavaFileSystem::new(vm, session.handle(), concurrency); + let mut ordered_paths = Vec::with_capacity(count); + for idx in 0..count { + let path_obj = paths.get_element(env, idx)?; + let path: String = env.cast_local::(path_obj)?.try_to_string(env)?; + if path.contains(['*', '?', '[']) { + throw_runtime!("path '{path}' contains glob characters, which are unsupported"); + } + let size = u64::try_from(sizes[idx]) + .map_err(|_| vortex_err!("negative length for path '{path}'"))?; + + let readable = readables.get_element(env, idx)?; + if readable.is_null() { + throw_runtime!("null readable for path '{path}'"); } - c => out.push(c), + let readable = Arc::new(env.new_global_ref(&readable)?); + + // `MultiFileDataSource::with_glob` strips leading slashes (object-store paths + // are bucket-relative), so key the registry by the same normalized form. + let key = path.trim_start_matches('/').to_string(); + fs.insert(key.clone(), readable, size)?; + ordered_paths.push(key); } - } - out + + let fs: FileSystemRef = Arc::new(fs); + let mut builder = MultiFileDataSource::new(session.clone()); + for path in ordered_paths { + builder = builder.with_glob(path, Some(Arc::clone(&fs))); + } + + let inner = RUNTIME + .block_on(builder.build()) + .map(|ds| Arc::new(ds) as DataSourceRef)?; + Ok(Box::new(NativeDataSource { inner }).into_raw()) + }) } /// URL with the path cleared, used as a cache key for filesystem reuse. @@ -236,47 +272,6 @@ pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_byteSize( mod tests { use super::*; - #[test] - fn test_parse_glob_url_full_url() { - let url = parse_glob_url("s3://bucket/prefix/*.vortex").unwrap(); - assert_eq!(url.scheme(), "s3"); - assert_eq!(url.host_str(), Some("bucket")); - assert_eq!(url.path(), "/prefix/*.vortex"); - } - - #[test] - fn test_parse_glob_url_absolute_path() { - // Use a drive-prefixed input on Windows so `absolute()` doesn't inject the cwd drive - // and the expected URL path is predictable. - #[cfg(unix)] - let (input, expected_path) = ("/tmp/data/*.vortex", "/tmp/data/*.vortex"); - #[cfg(windows)] - let (input, expected_path) = (r"C:\tmp\data\*.vortex", "/C:/tmp/data/*.vortex"); - let url = parse_glob_url(input).unwrap(); - assert_eq!(url.scheme(), "file"); - assert_eq!(url.path(), expected_path); - } - - #[test] - fn test_parse_glob_url_normalizes_dots() { - #[cfg(unix)] - let (input, expected_path) = ("/a/b/../c/./d", "/a/c/d"); - #[cfg(windows)] - let (input, expected_path) = (r"C:\a\b\..\c\.\d", "/C:/a/c/d"); - let url = parse_glob_url(input).unwrap(); - assert_eq!(url.path(), expected_path); - } - - #[test] - fn test_parse_glob_url_single_letter_scheme_is_path() { - // Regression: `Url::parse("C:\\tmp")` succeeds with scheme="c"; the function must - // treat that as a filesystem path, not a URL. Exercised on all platforms because - // the check lives in `parse_glob_url`, not in an OS-specific branch. - let url = parse_glob_url(r"C:\tmp\data\*.vortex").unwrap(); - assert_eq!(url.scheme(), "file"); - assert_ne!(url.scheme(), "c"); - } - #[test] fn test_base_url_strips_path() { let url = Url::parse("s3://bucket/a/b/c").unwrap(); diff --git a/vortex-jni/src/dtype.rs b/vortex-jni/src/dtype.rs index 5a0b54d2c14..1014ae07613 100644 --- a/vortex-jni/src/dtype.rs +++ b/vortex-jni/src/dtype.rs @@ -7,69 +7,23 @@ use std::ptr; use arrow_array::ffi::FFI_ArrowSchema; -use arrow_schema::DataType; -use arrow_schema::FieldRef; -use arrow_schema::Fields; use arrow_schema::Schema; use vortex::dtype::DType; use vortex::error::VortexResult; +use vortex_arrow::ToArrowType; -/// Export a Vortex [`DType`] to the Arrow C Data Interface struct at `schema_addr`. Views -/// (Utf8View/BinaryView) are downgraded to regular Utf8/Binary so Spark and other consumers -/// without view support can read them. +/// Export a Vortex [`DType`] to the Arrow C Data Interface struct at `schema_addr`. String and +/// binary columns are exported as their native view types (Utf8View/BinaryView); consumers are +/// expected to handle them. pub(crate) fn export_dtype_to_arrow(dtype: &DType, schema_addr: i64) -> VortexResult<()> { let arrow_schema = dtype.to_arrow_schema()?; - let viewless = strip_views(DataType::Struct(arrow_schema.fields().clone())); - let fields = match viewless { - DataType::Struct(fields) => fields, - _ => unreachable!("Vortex DType always exports as a struct"), - }; - let schema = Schema::new(fields); - let ffi_schema = FFI_ArrowSchema::try_from(&schema)?; + let ffi_schema = FFI_ArrowSchema::try_from(&arrow_schema)?; unsafe { ptr::write(schema_addr as *mut FFI_ArrowSchema, ffi_schema); } Ok(()) } -/// Replace view-based Arrow types with their non-view counterparts throughout the tree. -pub(crate) fn strip_views(data_type: DataType) -> DataType { - match data_type { - DataType::BinaryView => DataType::Binary, - DataType::Utf8View => DataType::Utf8, - DataType::List(inner) | DataType::ListView(inner) => { - let new_inner = (*inner) - .clone() - .with_data_type(strip_views(inner.data_type().clone())); - DataType::List(FieldRef::new(new_inner)) - } - DataType::LargeList(inner) | DataType::LargeListView(inner) => { - let new_inner = (*inner) - .clone() - .with_data_type(strip_views(inner.data_type().clone())); - DataType::LargeList(FieldRef::new(new_inner)) - } - DataType::Struct(fields) => { - let viewless_fields: Vec = fields - .iter() - .map(|field_ref| { - let field = (**field_ref).clone(); - let data_type = field.data_type().clone(); - FieldRef::new(field.with_data_type(strip_views(data_type))) - }) - .collect(); - DataType::Struct(Fields::from(viewless_fields)) - } - DataType::FixedSizeList(inner, size) => { - let new_inner = (*inner) - .clone() - .with_data_type(strip_views(inner.data_type().clone())); - DataType::FixedSizeList(FieldRef::new(new_inner), size) - } - dt => dt, - } -} - /// Decode an [`FFI_ArrowSchema`] pointed to by `schema_addr` into an Arrow [`Schema`]. pub(crate) fn import_arrow_schema(schema_addr: i64) -> VortexResult { let ffi_schema = unsafe { &*(schema_addr as *const FFI_ArrowSchema) }; diff --git a/vortex-jni/src/errors.rs b/vortex-jni/src/errors.rs index 4cf45a14667..0e903406ae0 100644 --- a/vortex-jni/src/errors.rs +++ b/vortex-jni/src/errors.rs @@ -11,6 +11,7 @@ use jni::sys::JNI_FALSE; use jni::sys::jboolean; use jni::sys::jobject; use vortex::error::VortexError; +use vortex::error::vortex_err; #[derive(Debug, thiserror::Error)] pub enum JNIError { @@ -38,6 +39,15 @@ impl From for JNIError { } } +impl From for VortexError { + fn from(error: JNIError) -> Self { + match error { + JNIError::Vortex(error) => error, + JNIError::Custom(error) => vortex_err!("JNI: {error}"), + } + } +} + /// Types that have a reasonable default value to use /// across the FFI. pub trait JNIDefault { diff --git a/vortex-jni/src/file.rs b/vortex-jni/src/file.rs index 3abb6ef9ffb..343215e533a 100644 --- a/vortex-jni/src/file.rs +++ b/vortex-jni/src/file.rs @@ -13,9 +13,9 @@ use jni::objects::JString; use jni::sys::jlong; use jni::sys::jobject; use object_store::path::Path; -use url::Url; use vortex::error::VortexResult; use vortex::error::vortex_err; +use vortex::file::multi::parse_uri_or_path; use vortex::io::runtime::BlockingRuntime; use vortex::io::session::RuntimeSessionExt; use vortex::utils::aliases::hash_map::HashMap; @@ -58,10 +58,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeFiles_listFiles( try_or_throw(&mut env, |env| { let session = unsafe { session_ref(session_ptr) }; let root_path: String = path.try_to_string(env)?; - - let Ok(url) = Url::parse(&root_path) else { - throw_runtime!("invalid URL: {root_path}"); - }; + let url = parse_uri_or_path(&root_path)?; let properties = extract_properties(env, &options)?; @@ -122,7 +119,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeFiles_delete( return Ok(()); } - let store_url = Url::parse(&delete_uris[0]).map_err(|e| vortex_err!(External: e))?; + let store_url = parse_uri_or_path(&delete_uris[0])?; let properties = extract_properties(env, &options)?; @@ -130,7 +127,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeFiles_delete( RUNTIME.block_on(async { for uri in delete_uris { - let url = Url::parse(&uri).map_err(|e| vortex_err!(External: e))?; + let url = parse_uri_or_path(&uri)?; fs.delete(url.path()).await?; } VortexResult::Ok(()) diff --git a/vortex-jni/src/io/mod.rs b/vortex-jni/src/io/mod.rs new file mode 100644 index 00000000000..fa20a6a4685 --- /dev/null +++ b/vortex-jni/src/io/mod.rs @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Rust implementations of Vortex I/O traits backed by Java objects (upcalls). +//! +//! These bridges let Java callers supply their own I/O — e.g. Iceberg's `FileIO` +//! streams — instead of having the native side open storage itself. Java objects +//! are held as JNI global references and their methods are invoked from whichever +//! runtime thread executes the I/O. +//! +//! # Threading +//! +//! Vortex runtime threads (smol workers and the `blocking` pool used by +//! [`Handle::spawn_blocking`](vortex::io::runtime::Handle::spawn_blocking)) are not +//! JVM threads. Every upcall goes through [`with_jvm`], which attaches the current +//! thread on first use (detached automatically at thread exit) — re-attachment is a +//! cheap thread-local lookup. +//! +//! # JavaVM lifetime +//! +//! Each bridge struct stores the process-wide [`JavaVM`] pointer, captured once at +//! construction from the [`Env`] of the JNI entry point that created it. +//! This is safe with respect to serialization: only *Java* objects are serialized +//! (e.g. Iceberg's `FileIO` shipped to executors), and they reconstruct their native +//! state through fresh JNI calls after deserialization, at which point the entry +//! point's `Env` provides the VM again. Native bridge objects never outlive the +//! process, and JNI guarantees a single VM per process for its entire lifetime. + +mod read_at; +mod write; + +use jni::Env; +use jni::JavaVM; +pub(crate) use read_at::JavaFileSystem; +use vortex::error::VortexError; +use vortex::error::VortexResult; +pub(crate) use write::JavaWrite; + +use crate::errors::JNIError; + +/// Run `f` on the current thread with a JVM attachment. +/// +/// Attaching is a thread-local lookup when the thread is already attached; threads +/// attached here stay attached until they exit. A Java exception thrown inside `f` +/// is caught, cleared, and surfaced as a `VortexError` carrying the exception's +/// class, message, and stack trace. +pub(crate) fn with_jvm( + vm: &JavaVM, + f: impl FnOnce(&mut Env) -> Result, +) -> VortexResult { + vm.attach_current_thread(f).map_err(VortexError::from) +} + +#[cfg(test)] +pub(crate) mod tests { + use std::sync::Arc; + use std::sync::LazyLock; + use std::time::Duration; + + use jni::InitArgsBuilder; + use jni::JValue; + use jni::JavaVM; + use jni::errors::StartJvmError; + use jni::objects::JObject; + use jni::refs::Global; + use vortex::error::VortexResult; + use vortex::error::vortex_bail; + use vortex::io::filesystem::FileSystem; + use vortex::io::runtime::BlockingRuntime; + use vortex::io::runtime::current::CurrentThreadRuntime; + + use super::JavaFileSystem; + use super::with_jvm; + + /// Launch (or reuse) the process-wide JVM backing upcall tests. Returns `None` + /// when no JVM installation can be located, so tests skip on machines without a + /// JDK; any other launch failure panics. + pub(crate) fn test_vm() -> Option<&'static JavaVM> { + static VM: LazyLock> = LazyLock::new(|| { + let args = InitArgsBuilder::new() + .option("-Xcheck:jni") + .build() + .expect("valid JVM init args"); + match JavaVM::new(args) { + Ok(vm) => Some(vm), + Err(StartJvmError::NotFound(_)) => { + eprintln!("skipping JNI upcall test: no JVM found (set JAVA_HOME to run it)"); + None + } + Err(e) => panic!("failed to launch test JVM: {e}"), + } + }); + VM.as_ref() + } + + /// A fresh `java.lang.Object` global ref plus a `java.lang.ref.WeakReference` + /// observing the same object, so tests can detect when the global ref is deleted. + pub(crate) fn object_with_weak_ref( + vm: &JavaVM, + ) -> VortexResult<(Global>, Global>)> { + with_jvm(vm, |env| { + let obj = + env.new_object(jni::jni_str!("java/lang/Object"), jni::jni_sig!("()V"), &[])?; + let weak = env.new_object( + jni::jni_str!("java/lang/ref/WeakReference"), + jni::jni_sig!("(Ljava/lang/Object;)V"), + &[JValue::Object(&obj)], + )?; + Ok((env.new_global_ref(&obj)?, env.new_global_ref(&weak)?)) + }) + } + + /// Assert that the referent observed by `weak` becomes collectible, i.e. that no + /// JNI global reference pins it anymore. + pub(crate) fn assert_weak_ref_clears( + vm: &JavaVM, + weak: &Global>, + ) -> VortexResult<()> { + for _ in 0..100 { + let cleared = with_jvm(vm, |env| { + env.call_static_method( + jni::jni_str!("java/lang/System"), + jni::jni_str!("gc"), + jni::jni_sig!("()V"), + &[], + )?; + let referent = env + .call_method( + weak.as_ref(), + jni::jni_str!("get"), + jni::jni_sig!("()Ljava/lang/Object;"), + &[], + )? + .l()?; + Ok(referent.is_null()) + })?; + if cleared { + return Ok(()); + } + std::thread::sleep(Duration::from_millis(10)); + } + vortex_bail!("JNI global reference leaked: weak reference still has a referent") + } + + #[test] + fn test_file_system_drop_releases_global_refs() -> VortexResult<()> { + let Some(vm) = test_vm() else { + return Ok(()); + }; + let (readable, weak) = object_with_weak_ref(vm)?; + + let runtime = CurrentThreadRuntime::new(); + let mut fs = JavaFileSystem::new(vm.clone(), runtime.handle(), None); + fs.insert("data/file.vortex".to_string(), Arc::new(readable), 4)?; + // `open_read` clones the global ref into a `JavaReadable`; both it and the + // file system must let go of the Java object once dropped. + let reader = runtime.block_on(fs.open_read("data/file.vortex"))?; + drop(fs); + drop(reader); + + assert_weak_ref_clears(vm, &weak) + } +} diff --git a/vortex-jni/src/io/read_at.rs b/vortex-jni/src/io/read_at.rs new file mode 100644 index 00000000000..a08652c7cc1 --- /dev/null +++ b/vortex-jni/src/io/read_at.rs @@ -0,0 +1,293 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Debug; +use std::sync::Arc; + +use async_lock::Semaphore; +use async_trait::async_trait; +use futures::FutureExt; +use futures::StreamExt; +use futures::future::BoxFuture; +use futures::stream; +use futures::stream::BoxStream; +use jni::JValue; +use jni::JavaVM; +use jni::objects::JObject; +use jni::refs::Global; +use vortex::array::buffer::BufferHandle; +use vortex::buffer::Alignment; +use vortex::buffer::ByteBufferMut; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_err; +use vortex::io::CoalesceConfig; +use vortex::io::VortexReadAt; +use vortex::io::filesystem::FileListing; +use vortex::io::filesystem::FileSystem; +use vortex::io::runtime::Handle; +use vortex::utils::aliases::hash_map::EntryRef; +use vortex::utils::aliases::hash_map::HashMap; + +use crate::io::with_jvm; + +/// Default number of concurrent `readFully` upcalls to allow across all files of one +/// [`JavaFileSystem`]. Matches the object-store default since the backing storage is +/// typically remote. +const DEFAULT_CONCURRENCY: usize = 192; + +/// Shared cap on in-flight `readFully` upcalls across every file of one +/// [`JavaFileSystem`], so a wide scan cannot pin `files x concurrency` blocking +/// threads and Java streams. +#[derive(Clone)] +struct UpcallLimiter { + semaphore: Arc, + concurrency: usize, +} + +impl UpcallLimiter { + fn new(concurrency: usize) -> Self { + Self { + semaphore: Arc::new(Semaphore::new(concurrency)), + concurrency, + } + } +} + +/// A [`VortexReadAt`] backed by a Java object implementing `dev.vortex.io.NativeReadable`. +/// +/// Positional reads are forwarded as blocking `readFully(long, ByteBuffer)` upcalls +/// executed on the runtime's blocking pool. The destination is a direct `ByteBuffer` +/// wrapping the Rust-side allocation, so Java writes land in the buffer handed to the +/// scan without a copy through a heap `byte[]`. The file size is captured at +/// construction time (Java callers know it from metadata), so `size()` never crosses +/// the JNI boundary. +pub(crate) struct JavaReadable { + vm: JavaVM, + readable: Arc>>, + len: u64, + handle: Handle, + limiter: UpcallLimiter, +} + +impl JavaReadable { + fn new( + vm: JavaVM, + readable: Arc>>, + len: u64, + handle: Handle, + limiter: UpcallLimiter, + ) -> Self { + Self { + vm, + readable, + len, + handle, + limiter, + } + } +} + +impl VortexReadAt for JavaReadable { + fn coalesce_config(&self) -> Option { + // Upcalls have a fixed JNI overhead and the backing storage is usually + // remote, so favor fewer, larger reads. + Some(CoalesceConfig::object_storage()) + } + + fn concurrency(&self) -> usize { + self.limiter.concurrency + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + let len = self.len; + async move { Ok(len) }.boxed() + } + + fn read_at( + &self, + offset: u64, + length: usize, + alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + let vm = self.vm.clone(); + let readable = Arc::clone(&self.readable); + let len = self.len; + let handle = self.handle.clone(); + let semaphore = Arc::clone(&self.limiter.semaphore); + + async move { + // Take a permit before occupying a blocking thread. The lock-free + // `try_acquire_arc` fast path deliberately barges ahead of queued + // waiters: slight unfairness is fine, the limiter must never become + // the bottleneck itself. + let permit = match semaphore.try_acquire_arc() { + Some(permit) => permit, + None => semaphore.acquire_arc().await, + }; + + handle + .spawn_blocking(move || { + // Keep the permit with the blocking work. Dropping the read future can + // cancel its task handle, but it cannot interrupt a `readFully` upcall + // that has already started. + let _permit = permit; + let end = offset + .checked_add(length as u64) + .ok_or_else(|| vortex_err!("read {offset}+{length} overflows u64"))?; + if end > len { + vortex_bail!("read {offset}..{end} out of bounds for file of length {len}"); + } + // `java.nio.Buffer` capacities are Java ints. + if i32::try_from(length).is_err() { + vortex_bail!("read length {length} exceeds ByteBuffer limit"); + } + let joffset = i64::try_from(offset) + .map_err(|_| vortex_err!("read offset {offset} exceeds i64"))?; + + let mut buffer = ByteBufferMut::with_capacity_aligned(length, alignment); + with_jvm(&vm, |env| { + // SAFETY: the pointer covers `length` bytes of live allocation, and + // the direct buffer wrapping it does not outlive this call — + // `readFully` is synchronous and the interface forbids retaining + // the buffer after it returns. + let dst = unsafe { + env.new_direct_byte_buffer( + buffer.spare_capacity_mut().as_mut_ptr().cast(), + length, + )? + }; + env.call_method( + readable.as_ref(), + jni::jni_str!("readFully"), + jni::jni_sig!("(JLjava/nio/ByteBuffer;)V"), + &[JValue::Long(joffset), JValue::Object(dst.as_ref())], + )?; + // Guard against implementations that return without filling the + // buffer, which would hand uninitialized memory to the scan. + let remaining = env + .call_method( + &dst, + jni::jni_str!("remaining"), + jni::jni_sig!("()I"), + &[], + )? + .i()?; + if remaining != 0 { + return Err(vortex_err!( + "readFully returned with {remaining} of {length} bytes unfilled" + ) + .into()); + } + Ok(()) + }) + .map_err(|e| e.with_context("readFully upcall failed"))?; + + // SAFETY: `readFully` wrote all `length` bytes through the direct + // buffer, verified by the `remaining()` check above. + unsafe { buffer.set_len(length) }; + Ok(BufferHandle::new_host(buffer.freeze())) + }) + .await + } + .boxed() + } +} + +struct JavaFileEntry { + readable: Arc>>, + size: u64, +} + +/// A [`FileSystem`] over a fixed set of Java-provided readables, keyed by path. +/// +/// Built by `NativeDataSource.openFiles`: every file the data source may touch is +/// registered up front together with its size, so `head` (and therefore exact-path +/// glob resolution) is answered without any upcall. `open_read` wraps the registered +/// Java object in a [`JavaReadable`]. +pub(crate) struct JavaFileSystem { + vm: JavaVM, + files: HashMap, + handle: Handle, + limiter: UpcallLimiter, +} + +impl JavaFileSystem { + pub(crate) fn new(vm: JavaVM, handle: Handle, concurrency: Option) -> Self { + Self { + vm, + files: HashMap::new(), + handle, + limiter: UpcallLimiter::new(concurrency.unwrap_or(DEFAULT_CONCURRENCY)), + } + } + + /// Register a Java readable for `path` with a known size. + pub(crate) fn insert( + &mut self, + path: String, + readable: Arc>>, + size: u64, + ) -> VortexResult<()> { + match self.files.entry_ref(&path) { + EntryRef::Occupied(_) => { + vortex_bail!("multiple Java readables normalize to path '{path}'"); + } + EntryRef::Vacant(v) => v.insert(JavaFileEntry { readable, size }), + }; + + Ok(()) + } +} + +impl Debug for JavaFileSystem { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("JavaFileSystem") + .field("files", &self.files.keys()) + .finish() + } +} + +#[async_trait] +impl FileSystem for JavaFileSystem { + fn list(&self, prefix: &str) -> BoxStream<'_, VortexResult> { + let listings: Vec> = self + .files + .iter() + .filter(|(path, _)| path.starts_with(prefix)) + .map(|(path, entry)| { + Ok(FileListing { + path: path.clone(), + size: Some(entry.size), + }) + }) + .collect(); + stream::iter(listings).boxed() + } + + async fn head(&self, path: &str) -> VortexResult> { + Ok(self.files.get(path).map(|entry| FileListing { + path: path.to_string(), + size: Some(entry.size), + })) + } + + async fn open_read(&self, path: &str) -> VortexResult> { + let entry = self + .files + .get(path) + .ok_or_else(|| vortex_err!("no Java readable registered for path '{path}'"))?; + Ok(Arc::new(JavaReadable::new( + self.vm.clone(), + Arc::clone(&entry.readable), + entry.size, + self.handle.clone(), + self.limiter.clone(), + ))) + } + + async fn delete(&self, path: &str) -> VortexResult<()> { + vortex_bail!("delete('{path}') is not supported by a Java-readable file system") + } +} diff --git a/vortex-jni/src/io/write.rs b/vortex-jni/src/io/write.rs new file mode 100644 index 00000000000..ffadf400e17 --- /dev/null +++ b/vortex-jni/src/io/write.rs @@ -0,0 +1,166 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::io; +use std::sync::Arc; + +use jni::JValue; +use jni::JavaVM; +use jni::objects::JObject; +use jni::refs::Global; +use vortex::io::IoBuf; +use vortex::io::VortexWrite; + +use crate::io::with_jvm; + +/// Largest chunk pushed through a single `write` upcall. Java arrays are indexed by +/// `int`, and keeping chunks bounded also keeps per-call array allocations modest. +const MAX_WRITE_CHUNK: usize = 1 << 30; + +/// A [`VortexWrite`] backed by a Java object implementing `dev.vortex.io.NativeWritable`. +/// +/// Bytes are forwarded as blocking `write(byte[], int, int)` upcalls, `flush` maps to +/// `flush()`, and `shutdown` maps to `flush()` as well: the Java caller created the +/// underlying stream and remains responsible for closing it once the writer finishes. +/// Writes run inline on the runtime thread driving the write task, which is attached +/// to the JVM on first use. +pub(crate) struct JavaWrite { + vm: JavaVM, + writable: Arc>>, +} + +impl JavaWrite { + pub(crate) fn new(vm: JavaVM, writable: Arc>>) -> Self { + Self { vm, writable } + } + + fn write_slice(&self, bytes: &[u8]) -> io::Result<()> { + for chunk in bytes.chunks(MAX_WRITE_CHUNK) { + let jlen = i32::try_from(chunk.len()).map_err(io::Error::other)?; + with_jvm(&self.vm, |env| { + let array = env.byte_array_from_slice(chunk)?; + env.call_method( + self.writable.as_ref(), + jni::jni_str!("write"), + jni::jni_sig!("([BII)V"), + &[ + JValue::Object(array.as_ref()), + JValue::Int(0), + JValue::Int(jlen), + ], + )?; + Ok(()) + }) + .map_err(io::Error::other)?; + } + Ok(()) + } + + fn flush_upcall(&self) -> io::Result<()> { + with_jvm(&self.vm, |env| { + env.call_method( + self.writable.as_ref(), + jni::jni_str!("flush"), + jni::jni_sig!("()V"), + &[], + )?; + Ok(()) + }) + .map_err(io::Error::other) + } +} + +impl VortexWrite for JavaWrite { + async fn write_all(&mut self, buffer: B) -> io::Result { + self.write_slice(buffer.as_slice())?; + Ok(buffer) + } + + async fn flush(&mut self) -> io::Result<()> { + self.flush_upcall() + } + + async fn shutdown(&mut self) -> io::Result<()> { + self.flush_upcall() + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use jni::JValue; + use jni::objects::JByteArray; + use vortex::error::VortexResult; + use vortex::error::vortex_err; + + use super::JavaWrite; + use crate::io::tests::assert_weak_ref_clears; + use crate::io::tests::test_vm; + use crate::io::with_jvm; + + /// Write upcalls from a fresh native thread must reach the Java object, and once + /// the writer (and every other native reference) is dropped, its JNI global ref + /// must be deleted so the Java object becomes collectible. + #[test] + fn test_write_upcalls_release_global_ref() -> VortexResult<()> { + let Some(vm) = test_vm() else { + return Ok(()); + }; + + // `ByteArrayOutputStream` has the same `write([BII)V`/`flush()V` shape as + // `dev.vortex.io.NativeWritable`, so it stands in for a caller-provided sink. + let (writable, content, weak) = with_jvm(vm, |env| { + let sink = env.new_object( + jni::jni_str!("java/io/ByteArrayOutputStream"), + jni::jni_sig!("()V"), + &[], + )?; + let weak = env.new_object( + jni::jni_str!("java/lang/ref/WeakReference"), + jni::jni_sig!("(Ljava/lang/Object;)V"), + &[JValue::Object(&sink)], + )?; + Ok(( + env.new_global_ref(&sink)?, + env.new_global_ref(&sink)?, + env.new_global_ref(&weak)?, + )) + })?; + + let thread_vm = vm.clone(); + std::thread::spawn(move || { + let write = JavaWrite::new(thread_vm.clone(), Arc::new(writable)); + write.write_slice(b"vortex")?; + write.flush_upcall()?; + // `with_jvm` attaches permanently: the thread stays attached until it + // exits, at which point the attachment is torn down automatically. + let attached = thread_vm + .is_thread_attached() + .map_err(|e| vortex_err!("is_thread_attached failed: {e}"))?; + assert!(attached, "write upcalls should leave the thread attached"); + VortexResult::Ok(()) + }) + .join() + .map_err(|_| vortex_err!("writer thread panicked"))??; + + let written = with_jvm(vm, |env| { + let array = env + .call_method( + content.as_ref(), + jni::jni_str!("toByteArray"), + jni::jni_sig!("()[B"), + &[], + )? + .l()?; + let array = env.cast_local::(array)?; + let mut bytes = vec![0i8; array.len(env)?]; + array.get_region(env, 0, &mut bytes)?; + Ok(bytes.into_iter().map(|b| b as u8).collect::>()) + })?; + assert_eq!(written, b"vortex"); + + drop(content); + assert_weak_ref_clears(vm, &weak) + } +} diff --git a/vortex-jni/src/lib.rs b/vortex-jni/src/lib.rs index a7d988c3c6a..c8369c28bff 100644 --- a/vortex-jni/src/lib.rs +++ b/vortex-jni/src/lib.rs @@ -23,6 +23,7 @@ mod dtype; mod errors; mod expression; mod file; +mod io; mod logging; mod object_store; mod runtime; diff --git a/vortex-jni/src/scan.rs b/vortex-jni/src/scan.rs index 7f923af1ac9..4b067706eba 100644 --- a/vortex-jni/src/scan.rs +++ b/vortex-jni/src/scan.rs @@ -19,7 +19,6 @@ use arrow_array::RecordBatch; use arrow_array::cast::AsArray; use arrow_array::ffi_stream::FFI_ArrowArrayStream; use arrow_schema::ArrowError; -use arrow_schema::DataType; use arrow_schema::Field; use futures::StreamExt; use jni::EnvUnowned; @@ -29,7 +28,6 @@ use jni::objects::JLongArray; use jni::sys::jboolean; use jni::sys::jlong; use vortex::array::VortexSessionExecute; -use vortex::array::arrow::ArrowSessionExt; use vortex::array::stream::SendableArrayStream; use vortex::buffer::Buffer; use vortex::error::VortexResult; @@ -44,11 +42,13 @@ use vortex::scan::PartitionRef; use vortex::scan::PartitionStream; use vortex::scan::ScanRequest; use vortex::scan::selection::Selection; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ToArrowType; use crate::POOL; use crate::RUNTIME; use crate::data_source::NativeDataSource; -use crate::dtype::strip_views; +use crate::dtype::export_dtype_to_arrow; use crate::errors::try_or_throw; use crate::session::session_ref; @@ -215,7 +215,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeScan_arrowSchema( let NativeScan::Pending(scan) = scan else { throw_runtime!("schema unavailable: scan already started"); }; - crate::dtype::export_dtype_to_arrow(scan.dtype(), schema_addr)?; + export_dtype_to_arrow(scan.dtype(), schema_addr)?; Ok(()) }); } @@ -343,13 +343,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativePartition_scanArrow( let array_stream = partition.execute()?; let dtype = array_stream.dtype().clone(); - let raw_schema = dtype.to_arrow_schema()?; - let viewless = strip_views(DataType::Struct(raw_schema.fields().clone())); - let fields = match viewless { - DataType::Struct(fields) => fields, - _ => unreachable!("Vortex DType always exports as a struct"), - }; - let schema = Arc::new(arrow_schema::Schema::new(fields)); + let schema = Arc::new(dtype.to_arrow_schema()?); let target = Arc::new(Field::new_struct("", schema.fields().clone(), false)); let session = unsafe { session_ref(session_ptr) }; diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 95d3ae1c400..c1ffea2f14f 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -8,6 +8,8 @@ use std::path::PathBuf; use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use arrow_array::RecordBatch; use arrow_array::StructArray; @@ -17,29 +19,39 @@ use arrow_schema::SchemaRef; use async_fs::File; use futures::SinkExt; use futures::channel::mpsc; +use jni::Env; use jni::EnvUnowned; use jni::objects::JClass; use jni::objects::JObject; use jni::objects::JString; +use jni::objects::JValue; use jni::sys::JNI_FALSE; use jni::sys::JNI_TRUE; use jni::sys::jboolean; use jni::sys::jlong; +use jni::sys::jobject; use object_store::ObjectStore; use object_store::path::Path as ObjectStorePath; -use url::Url; use vortex::array::ArrayRef; use vortex::array::VTable; -use vortex::array::arrow::ArrowSessionExt; +use vortex::array::scalar::PValue; +use vortex::array::scalar::Scalar; +use vortex::array::scalar::ScalarValue; +use vortex::array::stats::StatsSet; use vortex::array::stream::ArrayStreamAdapter; use vortex::dtype::DType; use vortex::dtype::Field as DTypeField; use vortex::dtype::FieldPath; +use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::error::vortex_err; +use vortex::expr::stats::Stat; +use vortex::expr::stats::StatsProvider; +use vortex::file::CountingVortexWrite; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; use vortex::file::WriteSummary; +use vortex::file::multi::parse_uri_or_path; use vortex::io::VortexWrite; use vortex::io::compat::Compat; use vortex::io::object_store::ObjectStoreWrite; @@ -48,6 +60,7 @@ use vortex::io::runtime::Task; use vortex::io::session::RuntimeSessionExt; use vortex::session::VortexSession; use vortex::utils::aliases::hash_map::HashMap; +use vortex_arrow::ArrowSessionExt; use vortex_parquet_variant::ParquetVariant; use crate::RUNTIME; @@ -55,6 +68,7 @@ use crate::dtype::import_arrow_schema; use crate::errors::JNIError; use crate::errors::try_or_throw; use crate::file::extract_properties; +use crate::io::JavaWrite; use crate::object_store::make_object_store; use crate::session::session_ref; @@ -71,20 +85,17 @@ fn resolve_store( url_or_path: &str, properties: &HashMap, ) -> VortexResult { - match Url::parse(url_or_path) { - Ok(url) if url.scheme() == "file" => { - let path = url - .to_file_path() - .map_err(|_| vortex_err!("invalid file URL: {url_or_path}"))?; - Ok(ResolvedStore::Path(path)) - } - Ok(url) => { - let path = ObjectStorePath::from_url_path(url.path()) - .map_err(|_| vortex_err!("invalid object_store path: {}", url.path()))?; - let store = make_object_store(&url, properties)?; - Ok(ResolvedStore::ObjectStore(store, path)) - } - Err(_) => Ok(ResolvedStore::Path(PathBuf::from(url_or_path))), + let url = parse_uri_or_path(url_or_path)?; + if url.scheme() == "file" { + let path = url + .to_file_path() + .map_err(|_| vortex_err!("invalid file URL: {url_or_path}"))?; + Ok(ResolvedStore::Path(path)) + } else { + let path = ObjectStorePath::from_url_path(url.path()) + .map_err(|_| vortex_err!("invalid object_store path: {}", url.path()))?; + let store = make_object_store(&url, properties)?; + Ok(ResolvedStore::ObjectStore(store, path)) } } @@ -133,6 +144,7 @@ pub struct NativeWriter { session: VortexSession, arrow_schema: SchemaRef, write_schema: DType, + bytes_written: Arc, sender: mpsc::Sender>, } @@ -141,6 +153,7 @@ impl NativeWriter { session: VortexSession, arrow_schema: SchemaRef, write_schema: DType, + bytes_written: Arc, handle: Task>, sender: mpsc::Sender>, ) -> Self { @@ -149,6 +162,7 @@ impl NativeWriter { session, arrow_schema, write_schema, + bytes_written, sender, } } @@ -186,17 +200,189 @@ impl NativeWriter { .map_err(|e| vortex_err!("failed to send batch: {e}")) } - fn close(mut self) -> VortexResult<()> { + fn bytes_written(&self) -> u64 { + self.bytes_written.load(Ordering::Relaxed) + } + + fn close(mut self) -> VortexResult { self.sender.disconnect(); let handle = self .handle .take() .ok_or_else(|| vortex_err!("writer already closed"))?; - RUNTIME.block_on(async { - handle.await?; - VortexResult::Ok(()) - }) + RUNTIME.block_on(handle) + } +} + +fn checked_jlong(value: u64, name: &str) -> VortexResult { + jlong::try_from(value).map_err(|_| vortex_err!("{name} exceeds Java long range: {value}")) +} + +fn exact_count_jlong( + stats: Option<&StatsSet>, + dtype: Option<&DType>, + stat: Stat, +) -> VortexResult { + stats + .zip(dtype.and_then(|dt| stat.dtype(dt))) + .and_then(|(stats, dt)| stats.get_as::(stat, &dt).as_exact()) + .map(|value| checked_jlong(value, stat.name())) + .transpose() + .map(|value| value.unwrap_or(-1)) +} + +fn big_integer<'local>( + env: &mut Env<'local>, + value: impl ToString, +) -> Result, JNIError> { + let string = env.new_string(value.to_string())?; + Ok(env.new_object( + jni::jni_str!("java/math/BigInteger"), + jni::jni_sig!("(Ljava/lang/String;)V"), + &[JValue::Object(string.as_ref())], + )?) +} + +fn scalar_to_java<'local>( + env: &mut Env<'local>, + scalar: Scalar, +) -> Result, JNIError> { + if scalar.is_null() { + return Ok(JObject::null()); + } + if scalar.dtype().is_extension() { + return scalar_to_java(env, scalar.as_extension().to_storage_scalar()); + } + + let Some(value) = scalar.value() else { + return Ok(JObject::null()); + }; + match value { + ScalarValue::Bool(value) => Ok(env.new_object( + jni::jni_str!("java/lang/Boolean"), + jni::jni_sig!("(Z)V"), + &[JValue::Bool(if *value { JNI_TRUE } else { JNI_FALSE })], + )?), + ScalarValue::Primitive(value) => match value { + PValue::U8(_) | PValue::U16(_) | PValue::I8(_) | PValue::I16(_) | PValue::I32(_) => { + Ok(env.new_object( + jni::jni_str!("java/lang/Integer"), + jni::jni_sig!("(I)V"), + &[JValue::Int(value.cast::()?)], + )?) + } + PValue::U32(_) | PValue::I64(_) => Ok(env.new_object( + jni::jni_str!("java/lang/Long"), + jni::jni_sig!("(J)V"), + &[JValue::Long(value.cast::()?)], + )?), + PValue::U64(value) => big_integer(env, value), + PValue::F16(_) | PValue::F32(_) => Ok(env.new_object( + jni::jni_str!("java/lang/Float"), + jni::jni_sig!("(F)V"), + &[JValue::Float(value.cast::()?)], + )?), + PValue::F64(value) => Ok(env.new_object( + jni::jni_str!("java/lang/Double"), + jni::jni_sig!("(D)V"), + &[JValue::Double(*value)], + )?), + }, + ScalarValue::Decimal(value) => { + let DType::Decimal(decimal_dtype, _) = scalar.dtype() else { + return Err(JNIError::Vortex(vortex_err!( + "decimal statistic has non-decimal dtype {}", + scalar.dtype() + ))); + }; + let unscaled = big_integer(env, value.as_i256())?; + Ok(env.new_object( + jni::jni_str!("java/math/BigDecimal"), + jni::jni_sig!("(Ljava/math/BigInteger;I)V"), + &[ + JValue::Object(&unscaled), + JValue::Int(i32::from(decimal_dtype.scale())), + ], + )?) + } + ScalarValue::Utf8(value) => Ok(env.new_string(value.as_str())?.into()), + ScalarValue::Binary(value) => Ok(env.byte_array_from_slice(value.as_slice())?.into()), + ScalarValue::Tuple(_) | ScalarValue::Variant(_) => Err(JNIError::Vortex(vortex_err!( + "cannot return nested scalar write statistic with dtype {} to Java", + scalar.dtype() + ))), + } +} + +fn write_summary_to_java<'local>( + env: &mut Env<'local>, + summary: &WriteSummary, +) -> Result, JNIError> { + let column_sizes = summary.compressed_column_sizes()?; + let file_stats = summary.footer().statistics(); + let columns = env.new_object_array( + i32::try_from(column_sizes.len()) + .map_err(|_| vortex_err!("column count exceeds Java array range"))?, + jni::jni_str!("dev/vortex/api/VortexColumnStatistics"), + JObject::null(), + )?; + + for (column_index, compressed_size) in column_sizes.into_iter().enumerate() { + let (stats, dtype) = file_stats + .and_then(|all_stats| { + all_stats + .stats_sets() + .get(column_index) + .zip(all_stats.dtypes().get(column_index)) + }) + .map_or((None, None), |(stats, dtype)| (Some(stats), Some(dtype))); + let null_count = exact_count_jlong(stats, dtype, Stat::NullCount)?; + let nan_count = exact_count_jlong(stats, dtype, Stat::NaNCount)?; + let lower_bound = stats + .zip(dtype) + .and_then(|(stats, dtype)| stats.as_typed_ref(dtype).get(Stat::Min).into_inner()); + let upper_bound = stats + .zip(dtype) + .and_then(|(stats, dtype)| stats.as_typed_ref(dtype).get(Stat::Max).into_inner()); + let column = env.with_local_frame_returning_local::<_, JObject, JNIError>(16, |env| { + let lower_bound = match lower_bound { + Some(value) => scalar_to_java(env, value)?, + None => JObject::null(), + }; + let upper_bound = match upper_bound { + Some(value) => scalar_to_java(env, value)?, + None => JObject::null(), + }; + Ok(env.new_object( + jni::jni_str!("dev/vortex/api/VortexColumnStatistics"), + jni::jni_sig!("(IJJJJLjava/lang/Object;Ljava/lang/Object;)V"), + &[ + JValue::Int( + i32::try_from(column_index) + .map_err(|_| vortex_err!("column index exceeds Java int range"))?, + ), + JValue::Long(checked_jlong(compressed_size, "compressed column size")?), + JValue::Long(checked_jlong(summary.row_count(), "row count")?), + JValue::Long(null_count), + JValue::Long(nan_count), + JValue::Object(&lower_bound), + JValue::Object(&upper_bound), + ], + )?) + })?; + columns.set_element(env, column_index, &column)?; + env.delete_local_ref(column); } + + Ok(env.new_object( + jni::jni_str!("dev/vortex/api/VortexWriteSummary"), + jni::jni_sig!("(JJ[Ldev/vortex/api/VortexColumnStatistics;)V"), + &[ + JValue::Long(checked_jlong(summary.size(), "file size")?), + JValue::Long(checked_jlong(summary.row_count(), "row count")?), + JValue::Object(columns.as_ref()), + ], + )?) } #[unsafe(no_mangle)] @@ -227,31 +413,97 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( let stream = ArrayStreamAdapter::new(write_schema.clone(), rx); let write_options = write_options_for_schema(session, &write_schema); - let handle = session.handle().spawn(async move { - match resolved { - ResolvedStore::Path(path) => { + let (bytes_written, handle) = match resolved { + ResolvedStore::Path(path) => { + let file = RUNTIME.block_on(async { if let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) { async_fs::create_dir_all(parent).await?; } - let mut file = File::create(path).await?; - let summary = write_options.write(&mut file, stream).await?; - file.shutdown().await?; + Ok::<_, VortexError>(File::create(path).await?) + })?; + let mut write = CountingVortexWrite::new(file); + let bytes_written = write.counter(); + let handle = session.handle().spawn(async move { + let summary = write_options.write(&mut write, stream).await?; + write.shutdown().await?; Ok(summary) - } - ResolvedStore::ObjectStore(store, path) => { - let mut write = - ObjectStoreWrite::new(Arc::new(Compat::new(store)), &path).await?; + }); + (bytes_written, handle) + } + ResolvedStore::ObjectStore(store, path) => { + let object_write = + RUNTIME.block_on(ObjectStoreWrite::new(Arc::new(Compat::new(store)), &path))?; + let mut write = CountingVortexWrite::new(object_write); + let bytes_written = write.counter(); + let handle = session.handle().spawn(async move { let summary = write_options.write(&mut write, stream).await?; write.shutdown().await?; Ok(summary) - } + }); + (bytes_written, handle) } + }; + + Ok(Box::new(NativeWriter::new( + session.clone(), + arrow_schema, + write_schema, + bytes_written, + handle, + tx, + )) + .into_raw()) + }) +} + +/// Create a writer that streams the file into a caller-provided +/// `dev.vortex.io.NativeWritable` instead of a native storage client. +/// +/// Bytes are pushed through blocking `write`/`flush` upcalls on the runtime thread +/// driving the write task. The Java caller owns the underlying stream and must close +/// it after `NativeWriter.close` returns; the native side only flushes. +#[unsafe(no_mangle)] +pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream( + mut env: EnvUnowned, + _class: JClass, + session_ptr: jlong, + writable: JObject, + arrow_schema_addr: jlong, +) -> jlong { + try_or_throw(&mut env, |env| { + if session_ptr == 0 { + throw_runtime!("null session pointer"); + } + if arrow_schema_addr == 0 { + throw_runtime!("null arrow schema address"); + } + if writable.is_null() { + throw_runtime!("null writable"); + } + let session = unsafe { session_ref(session_ptr) }; + + let arrow_schema = Arc::new(import_arrow_schema(arrow_schema_addr)?); + let write_schema = session.arrow().from_arrow_schema(arrow_schema.as_ref())?; + + let vm = env.get_java_vm()?; + let writable = Arc::new(env.new_global_ref(&writable)?); + let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY); + let stream = ArrayStreamAdapter::new(write_schema.clone(), rx); + let write_options = write_options_for_schema(session, &write_schema); + + let mut write = CountingVortexWrite::new(JavaWrite::new(vm, writable)); + let bytes_written = write.counter(); + let handle = session.handle().spawn(async move { + let summary = write_options.write(&mut write, stream).await?; + write.shutdown().await?; + Ok(summary) }); Ok(Box::new(NativeWriter::new( session.clone(), arrow_schema, write_schema, + bytes_written, handle, tx, )) @@ -288,6 +540,38 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_writeBatch( }) } +#[unsafe(no_mangle)] +pub extern "system" fn Java_dev_vortex_jni_NativeWriter_bytesWritten( + mut env: EnvUnowned, + _class: JClass, + writer_ptr: jlong, +) -> jlong { + if writer_ptr <= 0 { + return -1; + } + + try_or_throw(&mut env, |_env| { + let writer = unsafe { NativeWriter::from_ptr(writer_ptr) }; + Ok(checked_jlong(writer.bytes_written(), "bytes written")?) + }) +} + +#[unsafe(no_mangle)] +pub extern "system" fn Java_dev_vortex_jni_NativeWriter_finish( + mut env: EnvUnowned, + _class: JClass, + writer_ptr: jlong, +) -> jobject { + if writer_ptr <= 0 { + return JObject::null().into_raw(); + } + let writer = unsafe { NativeWriter::from_raw(writer_ptr) }; + try_or_throw(&mut env, |env| { + let summary = writer.close()?; + Ok(write_summary_to_java(env, &summary)?.into_raw()) + }) +} + #[unsafe(no_mangle)] pub extern "system" fn Java_dev_vortex_jni_NativeWriter_close( mut env: EnvUnowned, diff --git a/vortex-json/Cargo.toml b/vortex-json/Cargo.toml index 71eb015b5ee..b839c0ac76c 100644 --- a/vortex-json/Cargo.toml +++ b/vortex-json/Cargo.toml @@ -21,6 +21,7 @@ arrow-array = { workspace = true } arrow-schema = { workspace = true, features = ["canonical_extension_types"] } prost = { workspace = true } vortex-array = { workspace = true, default-features = false } +vortex-arrow = { workspace = true } vortex-error = { workspace = true, default-features = false } vortex-proto = { workspace = true, features = ["expr"] } vortex-session = { workspace = true } diff --git a/vortex-json/src/arrow.rs b/vortex-json/src/arrow.rs index 9db94d86743..987d6e16347 100644 --- a/vortex-json/src/arrow.rs +++ b/vortex-json/src/arrow.rs @@ -13,16 +13,16 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::arrow::ArrowExport; -use vortex_array::arrow::ArrowExportVTable; -use vortex_array::arrow::ArrowImport; -use vortex_array::arrow::ArrowImportVTable; -use vortex_array::arrow::ArrowSession; -use vortex_array::arrow::ArrowSessionExt; -use vortex_array::arrow::FromArrowArray; use vortex_array::dtype::DType; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::registry::CachedId; @@ -156,9 +156,9 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::VarBinArray; - use vortex_array::arrow::ArrowSessionExt; use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; + use vortex_arrow::ArrowSessionExt; use vortex_error::VortexExpect; use vortex_error::VortexResult; diff --git a/vortex-json/src/lib.rs b/vortex-json/src/lib.rs index 21f86ca9a7c..7a6dc8a780e 100644 --- a/vortex-json/src/lib.rs +++ b/vortex-json/src/lib.rs @@ -20,9 +20,9 @@ pub use json_to_variant::JsonToVariant; pub use json_to_variant::JsonToVariantOptions; pub use json_to_variant::ShreddingSpec; pub use json_to_variant::json_to_variant; -use vortex_array::arrow::ArrowSessionExt; use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::scalar_fn::session::ScalarFnSessionExt; +use vortex_arrow::ArrowSessionExt; use vortex_session::VortexSession; /// Register JSON extension support with a session. diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index 61b1253ef43..f772b9ab639 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -40,6 +40,7 @@ termtree = { workspace = true } tokio = { workspace = true, features = ["rt"], optional = true } tracing = { workspace = true } vortex-array = { workspace = true } +vortex-arrow = { workspace = true } vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } @@ -55,6 +56,7 @@ vortex-utils = { workspace = true, features = ["dashmap"] } [dev-dependencies] futures = { workspace = true, features = ["executor"] } +insta = { workspace = true } rstest = { workspace = true } temp-env = { workspace = true } tokio = { workspace = true, features = ["rt", "macros"] } diff --git a/vortex-layout/src/children.rs b/vortex-layout/src/children.rs index dade2cbc4a4..4903265890e 100644 --- a/vortex-layout/src/children.rs +++ b/vortex-layout/src/children.rs @@ -218,15 +218,14 @@ impl LayoutChildren for ViewedLayoutChildren { let encoding_id = self .layout_read_ctx .resolve(fb_child.encoding()) - .ok_or_else(|| vortex_err!("Encoding not found: {}", fb_child.encoding()))?; + .ok_or_else(|| { + vortex_err!("Unknown layout encoding index: {}", fb_child.encoding()) + })?; let Some(encoding) = self.layouts.find(&encoding_id) else { if self.allow_unknown { return viewed_children.foreign_layout_from_fb(fb_child, dtype); } - return Err(vortex_err!( - "Encoding not found in registry: {}", - fb_child.encoding() - )); + vortex_bail!("Unknown layout encoding: {encoding_id}"); }; let build_ctx = LayoutBuildContext { diff --git a/vortex-layout/src/layouts/chunked/mod.rs b/vortex-layout/src/layouts/chunked/mod.rs index a21b5605b31..336de3235b1 100644 --- a/vortex-layout/src/layouts/chunked/mod.rs +++ b/vortex-layout/src/layouts/chunked/mod.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -mod reader; +pub(crate) mod reader; pub mod writer; use std::sync::Arc; diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 489d8823d39..30172eb45d2 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -360,6 +360,7 @@ mod tests { use vortex_array::expr::pack; use vortex_array::expr::root; use vortex_array::validity::Validity; + use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::runtime::Handle; @@ -401,6 +402,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), + Arc::new(BtrBlocksCompressor::default()), ); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -430,6 +432,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), + Arc::new(BtrBlocksCompressor::default()), ); let array = VarBinArray::from_iter( @@ -530,6 +533,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), + Arc::new(BtrBlocksCompressor::default()), ); let array = @@ -582,6 +586,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), + Arc::new(BtrBlocksCompressor::default()), ); let array = VarBinArray::from_iter( diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index 1a376452984..83483c92662 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -30,7 +30,6 @@ use vortex_array::builders::dict::dict_encoder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -44,6 +43,7 @@ use crate::LayoutRef; use crate::LayoutStrategy; use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; +use crate::layouts::compressed::CompressorPlugin; use crate::layouts::dict::DictLayout; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; @@ -108,6 +108,7 @@ pub struct DictStrategy { values: Arc, fallback: Arc, options: DictLayoutOptions, + probe_compressor: Arc, } impl DictStrategy { @@ -116,12 +117,14 @@ impl DictStrategy { values: Values, fallback: Fallback, options: DictLayoutOptions, + probe_compressor: Arc, ) -> Self { Self { codes: Arc::new(codes), values: Arc::new(values), fallback: Arc::new(fallback), options, + probe_compressor, } } } @@ -155,7 +158,9 @@ impl LayoutStrategy for DictStrategy { None => true, // empty stream Some(chunk) => { let mut exec_ctx = session.create_execution_ctx(); - let compressed = BtrBlocksCompressor::default().compress(&chunk, &mut exec_ctx)?; + let compressed = self + .probe_compressor + .compress_chunk(&chunk, &mut exec_ctx)?; !compressed.is::() } }; diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index 983f9e59a58..c4cd032c686 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -11,7 +11,6 @@ use parking_lot::Mutex; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::arrays::ConstantArray; @@ -115,7 +114,7 @@ impl StatsAccumulator { Ok(()) } - fn as_array(&mut self) -> VortexResult> { + fn as_array(&mut self, ctx: &mut ExecutionCtx) -> VortexResult> { let mut names = Vec::new(); let mut fields = Vec::new(); @@ -128,7 +127,7 @@ impl StatsAccumulator { let values = builder.finish(); // We drop any all-null stats columns. - if values.all_invalid()? { + if values.all_invalid(ctx)? { continue; } @@ -146,7 +145,7 @@ impl StatsAccumulator { /// Returns an aggregated stats set for the table. fn as_stats_set(&mut self, stats: &[Stat], ctx: &mut ExecutionCtx) -> VortexResult { let mut stats_set = StatsSet::default(); - let Some(array) = self.as_array()? else { + let Some(array) = self.as_array(ctx)? else { return Ok(stats_set); }; @@ -236,8 +235,8 @@ struct NamedArrays { } impl NamedArrays { - fn all_invalid(&self) -> VortexResult { - self.arrays[0].all_invalid(&mut LEGACY_SESSION.create_execution_ctx()) + fn all_invalid(&self, ctx: &mut ExecutionCtx) -> VortexResult { + self.arrays[0].all_invalid(ctx) } } @@ -524,7 +523,10 @@ mod tests { .vortex_expect("push_chunk should succeed for test data"); acc.push_chunk(&builder2.finish(), &mut ctx) .vortex_expect("push_chunk should succeed for test data"); - let stats_table = acc.as_array().unwrap().expect("Must have stats table"); + let stats_table = acc + .as_array(&mut ctx) + .unwrap() + .expect("Must have stats table"); assert_eq!( stats_table.names().as_ref(), &[ @@ -561,7 +563,10 @@ mod tests { let mut acc = StatsAccumulator::new(array.dtype(), &[Stat::Max, Stat::Min, Stat::Sum], 12); acc.push_chunk(&array, &mut ctx) .vortex_expect("push_chunk should succeed for test array"); - let stats_table = acc.as_array().unwrap().expect("Must have stats table"); + let stats_table = acc + .as_array(&mut ctx) + .unwrap() + .expect("Must have stats table"); assert_eq!( stats_table.names().as_ref(), &[ diff --git a/vortex-layout/src/layouts/list/expr.rs b/vortex-layout/src/layouts/list/expr.rs new file mode 100644 index 00000000000..152b7ea5dce --- /dev/null +++ b/vortex-layout/src/layouts/list/expr.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::expr::Expression; +use vortex_array::expr::is_root; +use vortex_array::expr::not; +use vortex_array::expr::root; +use vortex_array::scalar_fn::fns::is_not_null::IsNotNull; +use vortex_array::scalar_fn::fns::is_null::IsNull; +use vortex_array::scalar_fn::fns::list_length::ListLength; +use vortex_error::VortexResult; + +/// The minimal set of list children an expression needs for evaluation. +/// +/// For example: +/// - `is_null(root())` only needs the validity child. +/// - `list_length(root())` only needs the offsets and validity children. +/// - `root()` needs elements, offsets, and validity children. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub(super) enum ListChildrenNeeded { + /// Only the validity child is needed (`is_null` / `is_not_null`). + Validity, + /// Only the offsets and validity children are needed (`list_length`). + OffsetsAndValidity, + /// All children are needed. + All, +} + +/// The minimal set of list children needed to evaluate `expr`, where `root()` is a field with list dtype. +pub(super) fn get_necessary_list_children(expr: &Expression) -> ListChildrenNeeded { + if is_null_root(expr) { + return ListChildrenNeeded::Validity; + } + + if is_list_length_root(expr) { + return ListChildrenNeeded::OffsetsAndValidity; + } + + if is_root(expr) { + return ListChildrenNeeded::All; + } + + // Otherwise the requirement is the max over the operands. Childless expressions that never + // touch the list, such as literals, fall back to the cheapest usable child. + expr.children() + .iter() + .map(get_necessary_list_children) + .max() + .unwrap_or(ListChildrenNeeded::Validity) +} + +fn is_null_root(expr: &Expression) -> bool { + (expr.is::() || expr.is::()) + && expr.children().len() == 1 + && is_root(expr.child(0)) +} + +fn is_list_length_root(expr: &Expression) -> bool { + expr.is::() && expr.children().len() == 1 && is_root(expr.child(0)) +} + +/// Rewrite a validity-class expression so it can be evaluated against the list's validity bool +/// array (`true` == valid row): `is_not_null(root())` becomes `root()` and `is_null(root())` +/// becomes `not(root())`. All other nodes are rebuilt with rewritten children. +pub(super) fn rewrite_validity_expr(expr: &Expression) -> VortexResult { + if expr.is::() && expr.children().len() == 1 && is_root(expr.child(0)) { + return Ok(root()); + } + if expr.is::() && expr.children().len() == 1 && is_root(expr.child(0)) { + return Ok(not(root())); + } + let children = expr + .children() + .iter() + .map(rewrite_validity_expr) + .collect::>>()?; + expr.clone().with_children(children) +} + +/// Rewrite an offsets-class expression so it can be evaluated against an array of list lengths. +/// `list_length(root())` becomes `root()`. Other references to `root()` are left intact: for +/// offsets-class expressions they can only be validity checks, and the lengths array carries the +/// same validity as the original list. +pub(super) fn rewrite_offsets_expr(expr: &Expression) -> VortexResult { + if is_list_length_root(expr) { + return Ok(root()); + } + + let children = expr + .children() + .iter() + .map(rewrite_offsets_expr) + .collect::>>()?; + expr.clone().with_children(children) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::cast; + use vortex_array::expr::eq; + use vortex_array::expr::gt; + use vortex_array::expr::is_not_null; + use vortex_array::expr::is_null; + use vortex_array::expr::list_length; + use vortex_array::expr::lit; + use vortex_array::expr::not; + use vortex_array::expr::root; + + use super::*; + + /// `get_necessary_list_children` keys off the deepest list child an expression touches; `All` + /// is the always-correct default for anything not specifically recognized. + #[rstest] + // `is_null` / `is_not_null` of the list itself need only validity. + #[case::is_null(is_null(root()), ListChildrenNeeded::Validity)] + #[case::is_not_null(is_not_null(root()), ListChildrenNeeded::Validity)] + // Compound over validity-only operands stays validity. + #[case::not_is_null(not(is_null(root())), ListChildrenNeeded::Validity)] + // A list-independent (constant) expression falls to the cheapest usable child. + #[case::constant(lit(5), ListChildrenNeeded::Validity)] + // `list_length(root())` needs offsets and validity, but not elements. + #[case::list_length(list_length(root()), ListChildrenNeeded::OffsetsAndValidity)] + // Compound over offsets-only operands stays offsets. + #[case::list_length_filter( + gt(list_length(root()), lit(1u64)), + ListChildrenNeeded::OffsetsAndValidity + )] + #[case::cast_list_length( + cast( + list_length(root()), + DType::Primitive(PType::I64, Nullability::Nullable), + ), + ListChildrenNeeded::OffsetsAndValidity + )] + // A bare list reference needs the elements. + #[case::bare_root(root(), ListChildrenNeeded::All)] + // Any other fn over the list needs the elements. + #[case::not_root(not(root()), ListChildrenNeeded::All)] + // `is_null` only short-circuits to validity when its argument is the list itself. + #[case::is_null_of_derived(is_null(not(root())), ListChildrenNeeded::All)] + // Max over operands: validity + elements => elements. + #[case::validity_and_elements(eq(is_null(root()), root()), ListChildrenNeeded::All)] + fn classify_expr_class(#[case] expr: Expression, #[case] expected: ListChildrenNeeded) { + assert_eq!(get_necessary_list_children(&expr), expected); + } +} diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs new file mode 100644 index 00000000000..0d47ad5266d --- /dev/null +++ b/vortex-layout/src/layouts/list/mod.rs @@ -0,0 +1,284 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! An experimental structural layout for list-typed columns. Note that this is expected to change. +//! +//! [`ListLayout`] decomposes a list column into independently configurable child layouts: +//! `elements`, `offsets`, and, for nullable lists, `validity`. Keeping the children independent allows +//! each child to use its own configurable layout and lets nested list elements (e.g. with `List`) be decomposed recursively. +//! +//! This provides benefits such as: +//! * Reading only the children needed to evaluate an expression. For example, `ListLength` does +//! not need to read list elements. +//! * Restricting element reads to the range covered by the selected outer rows, avoiding elements +//! belonging exclusively to unselected leading or trailing lists. +//! * Allowing each child to use its own compression, chunking, and pruning strategy. + +mod expr; +mod reader; +pub mod writer; + +use std::sync::Arc; + +use reader::ListReader; +use vortex_array::DeserializeMetadata; +use vortex_array::ProstMetadata; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::LayoutBuildContext; +use crate::LayoutChildType; +use crate::LayoutEncodingRef; +use crate::LayoutId; +use crate::LayoutReaderContext; +use crate::LayoutReaderRef; +use crate::LayoutRef; +use crate::VTable; +use crate::children::LayoutChildren; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; +use crate::vtable; + +/// Child index of the `elements` layout. +pub const ELEMENTS_CHILD_INDEX: usize = 0; +/// Child index of the `offsets` layout. +pub const OFFSETS_CHILD_INDEX: usize = 1; +/// Child index of the `validity` layout (only present when the list dtype is nullable). +pub const VALIDITY_CHILD_INDEX: usize = 2; + +/// Number of children when the list dtype is non-nullable. +pub const NUM_CHILDREN_NON_NULLABLE: usize = 2; + +vtable!(List); + +impl VTable for List { + type Layout = ListLayout; + type Encoding = ListLayoutEncoding; + type Metadata = ProstMetadata; + + fn id(_encoding: &Self::Encoding) -> LayoutId { + static ID: CachedId = CachedId::new("vortex.list"); + *ID + } + + fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef { + LayoutEncodingRef::new_ref(ListLayoutEncoding.as_ref()) + } + + fn row_count(layout: &Self::Layout) -> u64 { + layout.row_count() + } + + fn dtype(layout: &Self::Layout) -> &DType { + &layout.dtype + } + + fn metadata(layout: &Self::Layout) -> Self::Metadata { + ProstMetadata(ListLayoutMetadata::new(layout.offsets_ptype())) + } + + fn segment_ids(_layout: &Self::Layout) -> Vec { + vec![] + } + + fn nchildren(layout: &Self::Layout) -> usize { + if layout.dtype.is_nullable() { + NUM_CHILDREN_NON_NULLABLE + 1 + } else { + NUM_CHILDREN_NON_NULLABLE + } + } + + fn child(layout: &Self::Layout, idx: usize) -> VortexResult { + match (idx, layout.validity.as_ref()) { + (ELEMENTS_CHILD_INDEX, _) => Ok(Arc::clone(&layout.elements)), + (OFFSETS_CHILD_INDEX, _) => Ok(Arc::clone(&layout.offsets)), + (VALIDITY_CHILD_INDEX, Some(validity)) => Ok(Arc::clone(validity)), + _ => vortex_bail!("Invalid child index {idx} for ListLayout"), + } + } + + fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType { + match (idx, layout.validity.is_some()) { + (ELEMENTS_CHILD_INDEX, _) => LayoutChildType::Auxiliary("elements".into()), + (OFFSETS_CHILD_INDEX, _) => LayoutChildType::Auxiliary("offsets".into()), + (VALIDITY_CHILD_INDEX, true) => LayoutChildType::Auxiliary("validity".into()), + _ => vortex_panic!("Invalid child index {idx} for ListLayout"), + } + } + + fn new_reader( + layout: &Self::Layout, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult { + Ok(Arc::new(ListReader::try_new( + layout.clone(), + name, + segment_source, + session.clone(), + ctx, + )?)) + } + + fn build( + _encoding: &Self::Encoding, + dtype: &DType, + _row_count: u64, + metadata: &::Output, + _segment_ids: Vec, + children: &dyn LayoutChildren, + _ctx: &LayoutBuildContext<'_>, + ) -> VortexResult { + validate_children(dtype, children.nchildren())?; + + let elements_dtype = dtype + .as_list_element_opt() + .ok_or_else(|| vortex_err!("ListLayout requires a List dtype, got {dtype}"))?; + let elements = children.child(ELEMENTS_CHILD_INDEX, elements_dtype.as_ref())?; + + let offsets_dtype = DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable); + let offsets = children.child(OFFSETS_CHILD_INDEX, &offsets_dtype)?; + + let validity = dtype + .is_nullable() + .then(|| children.child(VALIDITY_CHILD_INDEX, &DType::Bool(Nullability::NonNullable))) + .transpose()?; + + Ok(ListLayout { + dtype: dtype.clone(), + elements, + offsets, + validity, + }) + } + + fn with_children(layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { + validate_children(layout.dtype(), children.len())?; + + let mut iter = children.into_iter(); + layout.elements = iter + .next() + .ok_or_else(|| vortex_err!("missing elements child"))?; + layout.offsets = iter + .next() + .ok_or_else(|| vortex_err!("missing offsets child"))?; + layout.validity = layout + .dtype + .is_nullable() + .then(|| { + iter.next() + .ok_or_else(|| vortex_err!("missing validity child")) + }) + .transpose()?; + Ok(()) + } +} + +/// Validates expected number of children based on `dtype` +fn validate_children(dtype: &DType, n_children: usize) -> VortexResult<()> { + let expected = if dtype.is_nullable() { + NUM_CHILDREN_NON_NULLABLE + 1 + } else { + NUM_CHILDREN_NON_NULLABLE + }; + + vortex_ensure_eq!(n_children, expected); + Ok(()) +} + +#[derive(Debug)] +pub struct ListLayoutEncoding; + +/// Stores a list-typed array by shredding `elements`, `offsets`, and optional `validity` children. +#[derive(Clone, Debug)] +pub struct ListLayout { + dtype: DType, + elements: LayoutRef, + offsets: LayoutRef, + validity: Option, +} + +impl ListLayout { + /// Construct a new `ListLayout` from its components. + /// + /// # Invariants + /// + /// - `dtype` must be a [`DType::List`]. + /// - `validity` must be `Some` iff `dtype.is_nullable()`. + /// - `offsets.dtype()` must be a non-nullable integer. + /// - `offsets.row_count()` is the Arrow-canonical `n+1` for `n` lists (or `0` for empty). + /// - When present, `validity.row_count() == offsets.row_count().saturating_sub(1)`. + pub fn new( + dtype: DType, + elements: LayoutRef, + offsets: LayoutRef, + validity: Option, + ) -> Self { + Self { + dtype, + elements, + offsets, + validity, + } + } + + /// Number of lists in this layout. + #[inline] + pub fn row_count(&self) -> u64 { + self.offsets.row_count().saturating_sub(1) + } + + #[inline] + pub fn elements(&self) -> &LayoutRef { + &self.elements + } + + #[inline] + pub fn offsets(&self) -> &LayoutRef { + &self.offsets + } + + #[inline] + pub fn validity(&self) -> Option<&LayoutRef> { + self.validity.as_ref() + } + + /// The integer type used for the `offsets` child layout. + #[inline] + pub fn offsets_ptype(&self) -> PType { + self.offsets.dtype().as_ptype() + } + + /// The dtype of the inner elements column. + pub fn elements_dtype(&self) -> &DType { + self.dtype + .as_list_element_opt() + .vortex_expect("ListLayout dtype must be a List") + } +} + +#[derive(prost::Message)] +pub struct ListLayoutMetadata { + #[prost(enumeration = "PType", tag = "1")] + offsets_ptype: i32, +} + +impl ListLayoutMetadata { + pub fn new(offsets_ptype: PType) -> Self { + let mut metadata = Self::default(); + metadata.set_offsets_ptype(offsets_ptype); + metadata + } +} diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs new file mode 100644 index 00000000000..0d33f1f4ac5 --- /dev/null +++ b/vortex-layout/src/layouts/list/reader.rs @@ -0,0 +1,1180 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::Arc; + +use futures::FutureExt; +use futures::future::BoxFuture; +use futures::try_join; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ListArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldMask; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::root; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::ArrayFuture; +use crate::LayoutReader; +use crate::LayoutReaderContext; +use crate::LayoutReaderRef; +use crate::RowSplits; +use crate::SplitRange; +use crate::layouts::list::ListLayout; +use crate::layouts::list::expr::ListChildrenNeeded; +use crate::layouts::list::expr::get_necessary_list_children; +use crate::layouts::list::expr::rewrite_offsets_expr; +use crate::layouts::list::expr::rewrite_validity_expr; +use crate::segments::SegmentSource; + +type OptionalArrayFuture = BoxFuture<'static, VortexResult>>; + +/// The threshold of mask density below which we push the input mask into projection evaluation, +/// and above which we evaluate the expression over all rows and intersect afterward. +const EXPR_EVAL_THRESHOLD: f64 = 0.2; + +/// Maximum number of outer-row scan ranges contributed by one list layout. +const MAX_LIST_SPLIT_COUNT: u64 = 64; + +/// Reader for [`ListLayout`]. +#[derive(Clone)] +pub struct ListReader { + layout: ListLayout, + name: Arc, + session: VortexSession, + elements: LayoutReaderRef, + offsets: LayoutReaderRef, + validity: Option, +} + +impl ListReader { + pub(super) fn try_new( + layout: ListLayout, + name: Arc, + segment_source: Arc, + session: VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult { + let elements = layout.elements().new_reader( + format!("{name}.elements").into(), + Arc::clone(&segment_source), + &session, + ctx, + )?; + let offsets = layout.offsets().new_reader( + format!("{name}.offsets").into(), + Arc::clone(&segment_source), + &session, + ctx, + )?; + let validity = layout + .validity() + .map(|v| { + v.new_reader( + format!("{name}.validity").into(), + Arc::clone(&segment_source), + &session, + ctx, + ) + }) + .transpose()?; + + Ok(Self { + layout, + name, + session, + elements, + offsets, + validity, + }) + } + + /// Projection for [`ListChildrenNeeded::Validity`] expressions. Reads only the validity child, + /// synthesizing all-valid for a non-nullable list, and never touches the offsets or elements. + fn project_validity( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let validity_reader = self.validity.clone(); + let nullability = self.layout.dtype().nullability(); + let row_range = row_range.clone(); + // Evaluate the rewritten expression against the validity bool array (true == valid row). + let rewritten = rewrite_validity_expr(expr)?; + + Ok(async move { + let mask = mask.await?; + let row_count = usize::try_from(row_range.end - row_range.start)?; + let out_len = if mask.all_true() { + row_count + } else { + mask.true_count() + }; + + let validity_array = match validity_reader.as_ref() { + Some(v) => Some( + v.projection_evaluation(&row_range, &root(), MaskFuture::ready(mask))? + .await?, + ), + None => None, + }; + + let validity = create_validity(validity_array, nullability).to_array(out_len); + + validity.apply(&rewritten) + } + .boxed()) + } + + /// Projection for [`ListChildrenNeeded::All`] expressions. + /// + /// An all-true mask over the full local range reads every child concurrently. Otherwise, the + /// read is bounded to the first and last selected list. + fn project_all( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let is_full_range = row_range.start == 0 && row_range.end == self.layout.row_count(); + let reader = self.clone(); + let row_range = row_range.clone(); + let expr = expr.clone(); + Ok(async move { + let mask = mask.await?; + if is_full_range && mask.all_true() { + reader.project_all_full(&expr)?.await + } else { + reader.project_all_bounded(&row_range, &expr, mask)?.await + } + } + .boxed()) + } + + /// Fetch the complete `elements`, `offsets`, and `validity` children concurrently. + fn project_all_full(&self, expr: &Expression) -> VortexResult { + let row_count = self.layout.row_count(); + let elements_row_count = self.elements.row_count(); + let nullability = self.layout.dtype().nullability(); + let expr = expr.clone(); + + let offsets_fut = self.fetch_raw_offsets(&(0..row_count))?; + let elements_fut = self.fetch_raw_elements(&(0..elements_row_count))?; + let validity_fut = fetch_validity( + self.validity.as_ref(), + &(0..row_count), + MaskFuture::new_true(usize::try_from(row_count)?), + )?; + + Ok(async move { + let (offsets, elements, validity) = try_join!(offsets_fut, elements_fut, validity_fut)?; + // SAFETY: ListLayout is constructed from a valid ListArray and reading its children + // without transformation preserves the list invariants. + let list = unsafe { + ListArray::new_unchecked(elements, offsets, create_validity(validity, nullability)) + } + .into_array(); + list.apply(&expr) + } + .boxed()) + } + + /// Bounded read for a sub-range or selective mask. + /// + /// Crops leading and trailing unselected lists, reads their offsets, and translates the first + /// and last offset into the element-row range to fetch. Any holes in the selection are filtered + /// after reconstructing the list array. + fn project_all_bounded( + &self, + row_range: &Range, + expr: &Expression, + mask: Mask, + ) -> VortexResult { + // Crop to the smallest contiguous row range containing every selected list. + let Some(selected_rows) = selected_row_range(&mask) else { + let empty = Canonical::empty(self.layout.dtype()).into_array(); + let expr = expr.clone(); + return Ok(async move { empty.apply(&expr) }.boxed()); + }; + + let selected_mask = mask.slice(selected_rows.clone()); + let selected_row_range = (row_range.start + u64::try_from(selected_rows.start)?) + ..(row_range.start + u64::try_from(selected_rows.end)?); + + let nullability = self.layout.dtype().nullability(); + let expr = expr.clone(); + let reader = self.clone(); + let offsets_fut = self.fetch_raw_offsets(&selected_row_range)?; + + Ok(async move { + let offsets = offsets_fut.await?; + + let elements_range = elements_range_from_offsets(&offsets, &reader.session)?; + let elements_fut = reader.fetch_raw_elements(&elements_range)?; + let validity_fut = fetch_validity( + reader.validity.as_ref(), + &selected_row_range, + MaskFuture::new_true(selected_mask.len()), + )?; + let (elements, validity) = try_join!(elements_fut, validity_fut)?; + + let offsets = rebase_offsets(offsets, elements_range.start)?; + // SAFETY: the selected offsets remain monotonically increasing, rebasing them against + // the selected element range preserves their lengths, and validity covers the same + // cropped list rows. + let list = unsafe { + ListArray::new_unchecked(elements, offsets, create_validity(validity, nullability)) + } + .into_array(); + let list = if selected_mask.all_true() { + list + } else { + list.filter(selected_mask)? + }; + list.apply(&expr) + } + .boxed()) + } + + /// Projection for [`ListChildrenNeeded::OffsetsAndValidity`] expressions. Only reads offsets and validity children. + fn project_offsets_validity( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let offsets = self.fetch_raw_offsets(row_range)?; + let reader = self.clone(); + let row_range = row_range.clone(); + let rewritten = rewrite_offsets_expr(expr)?; + + Ok(async move { + let mask = mask.await?; + let row_count = usize::try_from(row_range.end - row_range.start)?; + let nullability = reader.layout.dtype().nullability(); + + let validity_mask = if mask.all_true() { + MaskFuture::new_true(row_count) + } else { + MaskFuture::ready(mask.clone()) + }; + let validity_fut = fetch_validity(reader.validity.as_ref(), &row_range, validity_mask)?; + + let offsets = offsets.await?; + let lengths = list_lengths_from_offsets(offsets)?; + let lengths = if mask.all_true() { + lengths + } else { + lengths.filter(mask)? + }; + let validity = validity_fut.await?; + let lengths = apply_lengths_validity(lengths, validity, nullability)?; + + lengths.apply(&rewritten) + } + .boxed()) + } + + /// Fire the offsets read for `row_range` in list row space. The offsets child has an extra entry, so reading + /// `row_range` maps to offsets in `[row_range.start..row_range.end + 1)`. + /// + /// No mask or expression is applied. + fn fetch_raw_offsets(&self, row_range: &Range) -> VortexResult { + let offsets_range = row_range.start..(row_range.end + 1); + let offsets_count = usize::try_from(offsets_range.end - offsets_range.start)?; + self.offsets.projection_evaluation( + &offsets_range, + &root(), + MaskFuture::new_true(offsets_count), + ) + } + + /// Fire the elements read for `row_range` in element space. + /// + /// No mask or expression is applied. + fn fetch_raw_elements(&self, row_range: &Range) -> VortexResult { + let row_count = usize::try_from(row_range.end - row_range.start)?; + self.elements + .projection_evaluation(row_range, &root(), MaskFuture::new_true(row_count)) + } +} + +fn selected_row_range(mask: &Mask) -> Option> { + Some(mask.first()?..mask.last()? + 1) +} + +fn create_validity(validity_array: Option, nullability: Nullability) -> Validity { + match validity_array { + Some(arr) => Validity::Array(arr), + None => match nullability { + Nullability::Nullable => Validity::AllValid, + Nullability::NonNullable => Validity::NonNullable, + }, + } +} + +impl LayoutReader for ListReader { + fn name(&self) -> &Arc { + &self.name + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn dtype(&self) -> &DType { + self.layout.dtype() + } + + fn row_count(&self) -> u64 { + self.layout.row_count() + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + split_range.check_bounds(self.layout.row_count())?; + + // Splits are difficult to calculate because all children live in different row coordinate spaces. + // List elements typically comprise the majority of the data in a list, and validity/offsets can be treated + // as metadata. We therefore want to parallelize the scan based on element work. + // + // Scan splits must be expressed in the list layout's outer-row space, but the elements child + // reports its natural boundaries in element-row space. So we translate the element splits using a + // heuristic to outer-row space. + + let element_row_count = self.elements.row_count(); + if element_row_count != 0 { + let mut element_splits = RowSplits::new_capacity(128); + self.elements.register_splits( + &[FieldMask::All], + &SplitRange::root(0..element_row_count)?, + &mut element_splits, + )?; + + let row_range = split_range.row_range(); + let mut last_split = None; + for element_split in element_splits.into_sorted_deduped() { + let Some(split) = map_element_split_to_outer_grid( + element_split, + element_row_count, + self.layout.row_count(), + MAX_LIST_SPLIT_COUNT, + ) else { + continue; + }; + if split <= row_range.start { + continue; + } + if split >= row_range.end { + break; + } + if last_split == Some(split) { + continue; + } + splits.push( + split_range + .row_offset() + .checked_add(split) + .vortex_expect("List layout split offset overflow"), + ); + last_split = Some(split); + } + } + splits.push(split_range.root_row_range().end); + Ok(()) + } + + // TODO(mk): handle zones for list elements + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let len = mask.len(); + let reader = self.clone(); + let row_range = row_range.clone(); + let expr = expr.clone(); + let session = self.session.clone(); + + Ok(MaskFuture::new(len, async move { + let mask = mask.await?; + + if mask.all_false() { + return Ok(mask); + } + + if mask.density() < EXPR_EVAL_THRESHOLD { + let predicate = reader + .projection_evaluation(&row_range, &expr, MaskFuture::ready(mask.clone()))? + .await?; + let predicate_mask = predicate_array_to_mask(predicate, &session)?; + Ok(mask.intersect_by_rank(&predicate_mask)) + } else { + let predicate = reader + .projection_evaluation(&row_range, &expr, MaskFuture::new_true(len))? + .await?; + let predicate_mask = predicate_array_to_mask(predicate, &session)?; + Ok(mask & &predicate_mask) + } + })) + } + + /// Reads only the list children needed to evaluate `expr`. + /// + /// Validity-only expressions avoid offsets and elements, length expressions read offsets and + /// validity, and general expressions reconstruct a list from all applicable children. + fn projection_evaluation( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + // Read as little as possible based on which list children the expression needs. + match get_necessary_list_children(expr) { + ListChildrenNeeded::Validity => self.project_validity(row_range, expr, mask), + ListChildrenNeeded::OffsetsAndValidity => { + self.project_offsets_validity(row_range, expr, mask) + } + ListChildrenNeeded::All => self.project_all(row_range, expr, mask), + } + } +} + +/// Converts a natural boundary from element-row space into an approximate outer-row scan split. +/// +/// Scan splits must be expressed in the list layout's outer-row space, but the elements child +/// reports its natural boundaries in element-row space. Translating a boundary exactly would +/// require consulting the list offsets, so this function instead preserves its relative position: +/// +/// ```text +/// element_split / element_row_count ≈ outer_split / outer_row_count +/// ``` +/// +/// The relative position is first rounded onto a grid containing at most `max_split_count` scan +/// ranges, then mapped into outer-row space. Multiple element boundaries may therefore map to the +/// same outer split and are deduplicated by the caller. With a grid size of 64, at most 63 interior +/// splits—and therefore 64 scan ranges—can be produced. +/// +/// The result is only a task-sizing hint. It is always between outer rows, but it is not guaranteed +/// to correspond exactly to the original physical element boundary. Endpoint boundaries are +/// omitted because they do not subdivide the scan +fn map_element_split_to_outer_grid( + element_split: u64, + element_row_count: u64, + outer_row_count: u64, + max_split_count: u64, +) -> Option { + if element_split == 0 + || element_split >= element_row_count + || outer_row_count == 0 + || max_split_count < 2 + { + return None; + } + debug_assert!(max_split_count.is_power_of_two()); + + let grid_index = (u128::from(element_split) * u128::from(max_split_count) + + u128::from(element_row_count / 2)) + / u128::from(element_row_count); + if grid_index == 0 || grid_index >= u128::from(max_split_count) { + return None; + } + + let outer_split = grid_index * u128::from(outer_row_count) / u128::from(max_split_count); + let outer_split = u64::try_from(outer_split) + .vortex_expect("Outer split is bounded by the list layout row count"); + (outer_split != 0 && outer_split < outer_row_count).then_some(outer_split) +} + +/// Fetch the validity child for `row_range` under `mask`, yielding `None` for a non-nullable list +/// (which has no validity child). +fn fetch_validity( + validity: Option<&LayoutReaderRef>, + row_range: &Range, + mask: MaskFuture, +) -> VortexResult { + let fut = validity + .map(|v| v.projection_evaluation(row_range, &root(), mask)) + .transpose()?; + Ok(async move { + match fut { + Some(f) => f.await.map(Some), + None => Ok(None), + } + } + .boxed()) +} + +/// Read `offsets[0]` and `offsets[-1]` and return the elements range they bound. +fn elements_range_from_offsets( + offsets: &ArrayRef, + session: &VortexSession, +) -> VortexResult> { + if offsets.is_empty() { + return Ok(0..0); + } + let mut exec_ctx = session.create_execution_ctx(); + let start = offsets + .execute_scalar(0, &mut exec_ctx)? + .as_primitive() + .as_::() + .vortex_expect("offset value fits in u64"); + let end = offsets + .execute_scalar(offsets.len() - 1, &mut exec_ctx)? + .as_primitive() + .as_::() + .vortex_expect("offset value fits in u64"); + Ok(start..end) +} + +/// Subtract `first` from every offset so they index into a sliced `elements[first..]` buffer that +/// starts at zero. +fn rebase_offsets(offsets: ArrayRef, first: u64) -> VortexResult { + if first == 0 { + return Ok(offsets); + } + let constant = ConstantArray::new(first, offsets.len()) + .into_array() + .cast(offsets.dtype().clone())?; + offsets.binary(constant, Operator::Sub) +} + +/// Compute `offsets[i + 1] - offsets[i]` as the unmasked list length values. +fn list_lengths_from_offsets(offsets: ArrayRef) -> VortexResult { + let len = offsets.len().saturating_sub(1); + offsets + .slice(1..offsets.len())? + .binary(offsets.slice(0..len)?, Operator::Sub) +} + +fn apply_lengths_validity( + lengths: ArrayRef, + validity: Option, + nullability: Nullability, +) -> VortexResult { + let len = lengths.len(); + let lengths = lengths.cast(DType::Primitive(PType::U64, nullability))?; + + if matches!(nullability, Nullability::Nullable) { + lengths.mask(create_validity(validity, nullability).to_array(len)) + } else { + Ok(lengths) + } +} + +fn predicate_array_to_mask(array: ArrayRef, session: &VortexSession) -> VortexResult { + let mut ctx = session.create_execution_ctx(); + array.null_as_false().execute(&mut ctx) +} + +#[cfg(test)] +mod tests { + use std::ops::Range; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use rstest::rstest; + use vortex_array::ArrayContext; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::expr::cast; + use vortex_array::expr::gt; + use vortex_array::expr::is_not_null; + use vortex_array::expr::is_null; + use vortex_array::expr::list_length; + use vortex_array::expr::lit; + use vortex_buffer::buffer; + use vortex_io::session::RuntimeSession; + use vortex_io::session::RuntimeSessionExt; + + use super::*; + use crate::LayoutRef; + use crate::LayoutStrategy; + use crate::layouts::chunked::writer::ChunkedLayoutStrategy; + use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::layouts::list::writer::ListLayoutStrategy; + use crate::layouts::repartition::RepartitionStrategy; + use crate::layouts::repartition::RepartitionWriterOptions; + use crate::scan::split_by::SplitBy; + use crate::segments::SegmentFuture; + use crate::segments::SegmentSource; + use crate::segments::TestSegments; + use crate::sequence::SequenceId; + use crate::sequence::SequentialArrayStreamExt; + use crate::session::LayoutSession; + use crate::test::SESSION; + + /// Validity-class projections (`is_null` / `is_not_null` of the list) round-trip through the + /// validity-only read path, for both nullable and non-nullable lists. + #[rstest] + // `create_basic_list_array(true)` has validity `[true, false, true]`. + #[case::nullable(true, vec![true, false, true])] + #[case::non_nullable(false, vec![true, true, true])] + #[tokio::test] + async fn projection_validity_class( + #[case] nullable: bool, + #[case] valid: Vec, + ) -> VortexResult<()> { + let list = create_basic_list_array(nullable); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let not_null = reader + .projection_evaluation(&(0..3), &is_not_null(root()), MaskFuture::new_true(3))? + .await?; + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(not_null, BoolArray::from_iter(valid.clone()), &mut exec_ctx); + + let is_null_res = reader + .projection_evaluation(&(0..3), &is_null(root()), MaskFuture::new_true(3))? + .await?; + assert_arrays_eq!( + is_null_res, + BoolArray::from_iter(valid.iter().map(|v| !v).collect::>()), + &mut exec_ctx + ); + + Ok(()) + } + + #[tokio::test] + async fn projection_list_length_reads_offsets() -> VortexResult<()> { + let list = create_basic_list_array(false); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .projection_evaluation(&(0..3), &list_length(root()), MaskFuture::new_true(3))? + .await?; + + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, buffer![2u64, 2, 1].into_array(), &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn projection_list_length_preserves_validity() -> VortexResult<()> { + let list = create_basic_list_array(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .projection_evaluation(&(0..3), &list_length(root()), MaskFuture::new_true(3))? + .await?; + + let expected = + PrimitiveArray::from_option_iter::([Some(2), None, Some(1)]).into_array(); + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn projection_list_length_applies_sparse_mask() -> VortexResult<()> { + let list = create_basic_list_array(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let mask = Mask::from_iter([false, true, true]); + let result = reader + .projection_evaluation(&(0..3), &list_length(root()), MaskFuture::ready(mask))? + .await?; + + let expected = PrimitiveArray::from_option_iter::([None, Some(1)]).into_array(); + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn projection_cast_list_length() -> VortexResult<()> { + let list = create_basic_list_array(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let expr = cast( + list_length(root()), + DType::Primitive(PType::I64, Nullability::Nullable), + ); + let result = reader + .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? + .await?; + + let expected = + PrimitiveArray::from_option_iter::([Some(2), None, Some(1)]).into_array(); + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn filter_evaluation_list_length() -> VortexResult<()> { + let list = create_basic_list_array(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .filter_evaluation( + &(0..3), + >(list_length(root()), lit(1u64)), + MaskFuture::new_true(3), + )? + .await?; + + assert_eq!(result, Mask::from_iter([true, false, false])); + Ok(()) + } + + #[rstest] + #[case::is_not_null_nullable(true, is_not_null(root()), Mask::from_iter([true, false, true]))] + #[case::is_not_null_non_nullable(false, is_not_null(root()), Mask::new_true(3))] + #[case::is_null_nullable(true, is_null(root()), Mask::from_iter([false, true, false]))] + #[case::is_null_non_nullable(false, is_null(root()), Mask::new_false(3))] + #[tokio::test] + async fn filter_evaluation_validity_class( + #[case] nullable: bool, + #[case] expr: Expression, + #[case] expected: Mask, + ) -> VortexResult<()> { + let list = create_basic_list_array(nullable); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .filter_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? + .await?; + + assert_eq!(result, expected); + Ok(()) + } + + #[tokio::test] + async fn filter_evaluation_intersects_with_input_mask() -> VortexResult<()> { + let list = create_basic_list_array(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let input_mask = Mask::from_iter([true, true, false]); + let result = reader + .filter_evaluation(&(0..3), &is_not_null(root()), MaskFuture::ready(input_mask))? + .await?; + + assert_eq!(result, Mask::from_iter([true, false, false])); + Ok(()) + } + + #[tokio::test] + async fn filter_evaluation_sparse_mask_maps_by_rank() -> VortexResult<()> { + let list = create_six_list_array(); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let input_mask = Mask::from_iter([false, false, false, false, true, false]); + let result = reader + .filter_evaluation(&(0..6), &is_not_null(root()), MaskFuture::ready(input_mask))? + .await?; + + assert_eq!( + result, + Mask::from_iter([false, false, false, false, true, false]) + ); + Ok(()) + } + + fn flat_list_strategy() -> ListLayoutStrategy { + ListLayoutStrategy::default() + } + + fn layout_test_session() -> VortexSession { + vortex_array::array_session() + .with::() + .with::() + } + + async fn write_layout( + strategy: &S, + array: ArrayRef, + ) -> VortexResult<(Arc, LayoutRef, VortexSession)> { + let session = layout_test_session().with_tokio(); + let segments = Arc::new(TestSegments::default()); + let segments_ref: Arc = Arc::::clone(&segments); + let (ptr, eof) = SequenceId::root().split(); + let stream = array.to_array_stream().sequenced(ptr); + let layout = strategy + .write_stream(ArrayContext::empty(), segments, stream, eof, &session) + .await?; + Ok((segments_ref, layout, session)) + } + + fn materialize_u64_array(array: ArrayRef) -> Vec { + let mut ctx = SESSION.create_execution_ctx(); + array + .execute::(&mut ctx) + .unwrap() + .as_slice::() + .to_vec() + } + + fn create_basic_list_array(nullable: bool) -> ArrayRef { + let validity = if nullable { + Validity::Array(BoolArray::from_iter([true, false, true]).into_array()) + } else { + Validity::NonNullable + }; + + ListArray::try_new( + buffer![1i32, 2, 3, 4, 5].into_array(), + buffer![0u32, 2, 4, 5].into_array(), + validity, + ) + .expect("array is valid") + .into_array() + } + + fn create_six_list_array() -> ArrayRef { + let validity = Validity::Array( + BoolArray::from_iter([true, false, true, false, true, true]).into_array(), + ); + + ListArray::try_new( + buffer![1i32, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 1, 2, 3, 4, 5, 6].into_array(), + validity, + ) + .expect("array is valid") + .into_array() + } + + #[tokio::test] + async fn fetch_offsets_includes_extra_endpoint() -> VortexResult<()> { + let list = create_basic_list_array(false); + + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let ctx = LayoutReaderContext::new(); + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + let reader = reader + .as_any() + .downcast_ref::() + .expect("ListReader"); + + let offsets = reader.fetch_raw_offsets(&(1..3))?.await?; + assert_eq!(materialize_u64_array(offsets), vec![2u64, 4, 5]); + + Ok(()) + } + + #[rstest] + #[case::full_range(0..3, false)] + #[case::partial_start(0..2, false)] + #[case::partial_end(1..3, false)] + #[case::middle_single(1..2, false)] + #[case::empty_range(1..1, false)] + #[case::full_range_null(0..3, true)] + #[tokio::test] + async fn projection_evaluation_round_trips( + #[case] row_range: Range, + #[case] nullable: bool, + ) -> VortexResult<()> { + let list = create_basic_list_array(nullable); + let ctx = LayoutReaderContext::new(); + + let len = usize::try_from(row_range.end - row_range.start)?; + let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .projection_evaluation(&row_range, &root(), MaskFuture::new_true(len))? + .await?; + + let expected = + list.slice(usize::try_from(row_range.start)?..usize::try_from(row_range.end)?)?; + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn projection_evaluation_applies_mask() -> VortexResult<()> { + let list = create_basic_list_array(false); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let mask = Mask::from_iter([true, false, true]); + let result = reader + .projection_evaluation(&(0..3), &root(), MaskFuture::ready(mask.clone()))? + .await?; + + let expected = list.filter(mask)?; + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + /// Build a list with 5 rows and lengths [2, 3, 0, 3, 2]. + fn create_wider_list_array(nullable: bool) -> ArrayRef { + let validity = if nullable { + Validity::Array(BoolArray::from_iter([true, true, false, true, true]).into_array()) + } else { + Validity::NonNullable + }; + ListArray::try_new( + buffer![10i32, 11, 20, 21, 22, 30, 31, 32, 40, 41].into_array(), + buffer![0u32, 2, 5, 5, 8, 10].into_array(), + validity, + ) + .expect("array is valid") + .into_array() + } + + /// Sparse row masks over a whole chunk round-trip: the whole-chunk read is filtered by the + /// mask in memory. + #[rstest] + #[case::single_middle(Mask::from_iter([false, false, false, true, false]), false)] + #[case::two_far_apart(Mask::from_iter([true, false, false, true, false]), false)] + #[case::boundaries(Mask::from_iter([true, false, false, false, true]), false)] + #[case::kept_empty_row(Mask::from_iter([false, false, true, false, false]), false)] + #[case::sparse_nullable(Mask::from_iter([true, false, true, false, true]), true)] + #[case::all_false(Mask::new_false(5), false)] + #[tokio::test] + async fn projection_evaluation_sparse_mask_round_trips( + #[case] mask: Mask, + #[case] nullable: bool, + ) -> VortexResult<()> { + let list = create_wider_list_array(nullable); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .projection_evaluation(&(0..5), &root(), MaskFuture::ready(mask.clone()))? + .await?; + + let expected = list.filter(mask)?; + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + /// A partial range against a single flat elements segment takes the bounded path. The flat + /// reader may still fetch its whole segment, but the list reader reconstructs only the requested + /// element range. + #[tokio::test] + async fn projection_evaluation_partial_range_bounded() -> VortexResult<()> { + let list = create_wider_list_array(false); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .projection_evaluation(&(1..4), &root(), MaskFuture::new_true(3))? + .await?; + + let expected = list.slice(1..4)?; + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[test] + fn maps_element_splits_to_outer_grid() { + assert_eq!(map_element_split_to_outer_grid(0, 100, 100, 8), None); + assert_eq!(map_element_split_to_outer_grid(20, 100, 100, 8), Some(25)); + assert_eq!(map_element_split_to_outer_grid(25, 100, 100, 8), Some(25)); + assert_eq!(map_element_split_to_outer_grid(50, 100, 100, 8), Some(50)); + assert_eq!(map_element_split_to_outer_grid(75, 100, 100, 8), Some(75)); + assert_eq!(map_element_split_to_outer_grid(100, 100, 100, 8), None); + + let mut splits = (1..1_000) + .filter_map(|split| { + map_element_split_to_outer_grid(split, 1_000, 100_000, MAX_LIST_SPLIT_COUNT) + }) + .collect::>(); + splits.dedup(); + + let expected = (1..MAX_LIST_SPLIT_COUNT) + .map(|grid_index| grid_index * 100_000 / MAX_LIST_SPLIT_COUNT) + .collect::>(); + assert_eq!(splits, expected); + } + + #[tokio::test] + async fn nested_list_propagates_element_splits() -> VortexResult<()> { + let inner = ListArray::try_new( + PrimitiveArray::from_iter(0..128_i32).into_array(), + PrimitiveArray::from_iter((0..=8_u32).map(|idx| idx * 16)).into_array(), + Validity::NonNullable, + )? + .into_array(); + let outer = ListArray::try_new( + inner, + buffer![0u32, 2, 4, 6, 8].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let inner_strategy = + ListLayoutStrategy::default().with_elements(chunked_elements_strategy()); + let strategy = ListLayoutStrategy::default().with_elements(Arc::new(inner_strategy)); + let (segments, layout, session) = write_layout(&strategy, outer).await?; + let reader = + layout.new_reader("".into(), segments, &session, &LayoutReaderContext::new())?; + + let splits = SplitBy::Layout.splits(reader.as_ref(), &(0..4), &[FieldMask::All])?; + assert_eq!(splits, vec![0, 1, 2, 3, 4]); + Ok(()) + } + + /// A list strategy whose `elements` child is repartitioned into two-element chunks, so the + /// reader takes the bounded (chunk-skipping) path for strict sub-ranges. Offsets stay flat. + fn chunked_elements_list_strategy() -> ListLayoutStrategy { + ListLayoutStrategy::default().with_elements(chunked_elements_strategy()) + } + + fn chunked_elements_strategy() -> Arc { + Arc::new(RepartitionStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: 2, + block_size_target: None, + canonicalize: true, + }, + )) + } + + struct CountingSegmentSource { + inner: Arc, + request_count: Arc, + } + + impl SegmentSource for CountingSegmentSource { + fn request(&self, id: crate::segments::SegmentId) -> SegmentFuture { + self.request_count.fetch_add(1, Ordering::Relaxed); + self.inner.request(id) + } + } + + /// The chunked-elements strategy must actually produce a chunked `elements` layout, otherwise + /// the reader would silently take the whole-chunk path and the bounded read would be untested. + #[tokio::test] + async fn chunked_elements_produces_chunked_layout() -> VortexResult<()> { + let list = create_wider_list_array(false); + let (_segments, layout, _session) = + write_layout(&chunked_elements_list_strategy(), list).await?; + let tree = layout.display_tree().to_string(); + assert!( + tree.contains("elements: vortex.chunked"), + "elements should be chunked:\n{tree}" + ); + Ok(()) + } + + /// A sparse mask must remain selective even when the requested row range covers the complete + /// local list layout. The selected first list touches one element chunk plus the offsets + /// segment; reading all five element chunks would make the count six. + #[tokio::test] + async fn full_range_sparse_mask_crops_element_read() -> VortexResult<()> { + let list = create_wider_list_array(false); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = + write_layout(&chunked_elements_list_strategy(), list.clone()).await?; + let request_count = Arc::new(AtomicUsize::new(0)); + let source = Arc::new(CountingSegmentSource { + inner: segments, + request_count: Arc::clone(&request_count), + }); + let reader = layout.new_reader("".into(), source, &session, &ctx)?; + + let mask = Mask::from_iter([true, false, false, false, false]); + let result = reader + .projection_evaluation(&(0..5), &root(), MaskFuture::ready(mask.clone()))? + .await?; + + let expected = list.filter(mask)?; + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + assert_eq!(request_count.load(Ordering::Relaxed), 2); + Ok(()) + } + + /// With chunked elements, sub-range projections take the bounded read path. Every + /// range/mask/nullability combination must match the same projection over the ground-truth + /// array (`list.slice(range).filter(mask)`). + #[rstest] + #[case::full_all_true(0..5, Mask::new_true(5), false)] + #[case::subrange_all_true(1..4, Mask::new_true(3), false)] + #[case::subrange_sparse(1..4, Mask::from_iter([true, false, true]), false)] + #[case::partial_start(0..2, Mask::new_true(2), false)] + #[case::partial_end(2..5, Mask::new_true(3), false)] + #[case::single_non_empty(0..1, Mask::new_true(1), false)] + #[case::single_empty_row(2..3, Mask::from_iter([true]), false)] + #[case::empty_range(2..2, Mask::new_true(0), false)] + #[case::subrange_all_false(1..4, Mask::new_false(3), false)] + #[case::subrange_single_interior(0..4, Mask::from_iter([false, true, false, false]), false)] + #[case::subrange_sparse_nullable(1..4, Mask::from_iter([true, false, true]), true)] + #[case::partial_end_nullable(2..5, Mask::new_true(3), true)] + #[tokio::test] + async fn chunked_elements_round_trips( + #[case] row_range: Range, + #[case] mask: Mask, + #[case] nullable: bool, + ) -> VortexResult<()> { + let list = create_wider_list_array(nullable); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = + write_layout(&chunked_elements_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .projection_evaluation(&row_range, &root(), MaskFuture::ready(mask.clone()))? + .await?; + + let sliced = + list.slice(usize::try_from(row_range.start)?..usize::try_from(row_range.end)?)?; + let expected = sliced.filter(mask)?; + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs new file mode 100644 index 00000000000..9f33c1eee3e --- /dev/null +++ b/vortex-layout/src/layouts/list/writer.rs @@ -0,0 +1,578 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use async_trait::async_trait; +use futures::StreamExt; +use futures::future::try_join; +use futures::future::try_join_all; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::List; +use vortex_array::arrays::ListView; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::list::ListDataParts; +use vortex_array::arrays::listview::list_from_list_view; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::matcher::Matcher; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_io::kanal_ext::KanalExt; +use vortex_io::session::RuntimeSessionExt; +use vortex_session::VortexSession; + +use crate::IntoLayout; +use crate::LayoutRef; +use crate::LayoutStrategy; +use crate::layouts::flat::writer::FlatLayoutStrategy; +use crate::layouts::list::ListLayout; +use crate::segments::SegmentSinkRef; +use crate::sequence::SendableSequentialStream; +use crate::sequence::SequenceId; +use crate::sequence::SequencePointer; +use crate::sequence::SequentialStream; +use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequentialStreamExt; + +/// Item carried on each child sub-stream: a sequenced, materialized chunk. +type ChildChunk = VortexResult<(SequenceId, ArrayRef)>; + +/// Strategy for writing list-typed arrays, with a fallback for non-list dtypes. +/// +/// This is a *structural* writer that decomposes a list column into independent `elements`, +/// `offsets`, and (when nullable) `validity` sub-columns, each written through its own downstream +/// strategy, producing a single [`ListLayout`]. +/// +/// For list-typed input the strategy transposes the whole column stream into three sub-streams: +/// 1. Each chunk is canonicalized to a [`ListArray`] (rebuilding a [`ListView`] via +/// [`list_from_list_view`] when necessary). +/// 2. `offsets` are rebased to global `u64` positions (cumulative across chunks) so the single +/// `offsets` child indexes into the concatenated `elements` child. +/// 3. `elements`, `offsets`, and `validity` are streamed to their child strategies concurrently. +/// +/// For input whose dtype is not [`DType::List`], the stream is forwarded unchanged to the +/// configured `fallback` strategy. +/// +/// [`ListArray`]: vortex_array::arrays::ListArray +#[derive(Clone)] +pub struct ListLayoutStrategy { + elements: Arc, + offsets: Arc, + validity: Arc, + fallback: Arc, +} + +impl Default for ListLayoutStrategy { + /// Routes every child (elements, offsets, validity) and the non-list fallback through + /// [`FlatLayoutStrategy`]. Override individual children with the `with_*` builder methods. + fn default() -> Self { + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + Self { + elements: Arc::clone(&flat), + offsets: Arc::clone(&flat), + validity: Arc::clone(&flat), + fallback: flat, + } + } +} + +impl ListLayoutStrategy { + /// Strategy for the `elements` child. + pub fn with_elements(mut self, elements: Arc) -> Self { + self.elements = elements; + self + } + + /// Strategy for the `offsets` child. + pub fn with_offsets(mut self, offsets: Arc) -> Self { + self.offsets = offsets; + self + } + + /// Strategy for the `validity` child (written only when the list is nullable). + pub fn with_validity(mut self, validity: Arc) -> Self { + self.validity = validity; + self + } + + /// Strategy for non-list input, which is forwarded through this strategy unchanged. + pub fn with_fallback(mut self, fallback: Arc) -> Self { + self.fallback = fallback; + self + } +} + +#[async_trait] +impl LayoutStrategy for ListLayoutStrategy { + async fn write_stream( + &self, + ctx: ArrayContext, + segment_sink: SegmentSinkRef, + stream: SendableSequentialStream, + mut eof: SequencePointer, + session: &VortexSession, + ) -> VortexResult { + let dtype = stream.dtype().clone(); + if !dtype.is_list() { + return self + .fallback + .write_stream(ctx, segment_sink, stream, eof, session) + .await; + } + + let is_nullable = dtype.is_nullable(); + let element_dtype = dtype + .as_list_element_opt() + .vortex_expect("DType is List") + .as_ref() + .clone(); + // Global (whole-column) offsets are cumulative and may exceed the input offset width, + // so definsively widen. + let offsets_dtype = DType::Primitive(PType::U64, Nullability::NonNullable); + + // One bounded sub-stream per child: elements, offsets, and (when nullable) validity. + let (elements_tx, elements_rx) = kanal::bounded_async::(1); + let (offsets_tx, offsets_rx) = kanal::bounded_async::(1); + let (validity_tx, validity_rx) = if is_nullable { + let (tx, rx) = kanal::bounded_async::(1); + (Some(tx), Some(rx)) + } else { + (None, None) + }; + + // Transpose the list column into its child sub-streams and rebase offsets to global + // positions. Kept joined with the child writers below so producer errors surface rather + // than being hidden as an early channel close. + let fanout_fut = transpose_list_column( + stream, + session.clone(), + elements_tx, + offsets_tx, + validity_tx, + ); + + // Spawn a writer per child sub-stream, concurrently. + let handle = session.handle(); + let mut child_specs: Vec<( + DType, + Arc, + kanal::AsyncReceiver, + )> = vec![ + (element_dtype, Arc::clone(&self.elements), elements_rx), + (offsets_dtype, Arc::clone(&self.offsets), offsets_rx), + ]; + if let Some(validity_rx) = validity_rx { + child_specs.push(( + DType::Bool(Nullability::NonNullable), + Arc::clone(&self.validity), + validity_rx, + )); + } + + let layout_futures: Vec<_> = child_specs + .into_iter() + .map(|(child_dtype, strategy, rx)| { + let child_stream = + SequentialStreamAdapter::new(child_dtype, rx.into_stream().boxed()).sendable(); + let child_eof = eof.split_off(); + let ctx = ctx.clone(); + let segment_sink = Arc::clone(&segment_sink); + let session = session.clone(); + handle.spawn_nested(move |h| async move { + let session = session.with_handle(h); + strategy + .write_stream(ctx, segment_sink, child_stream, child_eof, &session) + .await + }) + }) + .collect(); + + let (_, layouts) = try_join(fanout_fut, try_join_all(layout_futures)).await?; + let mut layouts = layouts.into_iter(); + let elements_layout = layouts.next().vortex_expect("elements layout present"); + let offsets_layout = layouts.next().vortex_expect("offsets layout present"); + let validity_layout = + is_nullable.then(|| layouts.next().vortex_expect("validity layout present")); + + Ok(ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout).into_layout()) + } + + fn buffered_bytes(&self) -> u64 { + let list_bytes = self.elements.buffered_bytes() + + self.offsets.buffered_bytes() + + self.validity.buffered_bytes(); + list_bytes.max(self.fallback.buffered_bytes()) + } +} + +/// Transpose a list column into its `elements`, `offsets`, and (when present) `validity` child +/// sub-streams, rebasing each chunk's local `offsets` to global `u64` positions so the single +/// `offsets` child indexes into the concatenated `elements` child. +/// +/// `validity_tx` is `Some` exactly when the list is nullable. Errors surface to the caller, which +/// joins this against the child writers, rather than being hidden as an early channel close. +async fn transpose_list_column( + mut stream: SendableSequentialStream, + session: VortexSession, + elements_tx: kanal::AsyncSender, + offsets_tx: kanal::AsyncSender, + validity_tx: Option>, +) -> VortexResult<()> { + let mut exec_ctx = session.create_execution_ctx(); + let mut element_base: u64 = 0; + let mut first = true; + let mut saw_chunk = false; + while let Some(chunk) = stream.next().await { + let (sequence_id, array) = chunk?; + saw_chunk = true; + let mut sp = sequence_id.descend(); + let ListDataParts { + elements, + offsets, + validity, + .. + } = canonicalize_to_list_parts(array, &mut exec_ctx)?; + let n_elements = elements.len() as u64; + let row_count = offsets.len().saturating_sub(1); + let offsets = global_offsets(offsets, element_base, first, &mut exec_ctx)?; + element_base += n_elements; + first = false; + + if elements_tx + .send(Ok((sp.advance(), elements))) + .await + .is_err() + || offsets_tx.send(Ok((sp.advance(), offsets))).await.is_err() + { + vortex_bail!("list child writer finished before all chunks were sent"); + } + if let Some(validity_tx) = &validity_tx { + let validity = validity + .execute_mask(row_count, &mut exec_ctx)? + .into_array(); + if validity_tx + .send(Ok((sp.advance(), validity))) + .await + .is_err() + { + vortex_bail!("list validity writer finished before all chunks were sent"); + } + } + } + if !saw_chunk { + vortex_bail!("ListLayoutStrategy needs at least one chunk"); + } + Ok(()) +} + +/// Canonicalize a list-dtype array into [`ListDataParts`]. +fn canonicalize_to_list_parts( + array: ArrayRef, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let canonical = array.execute_until::(exec_ctx)?; + if let Some(list) = canonical.as_opt::() { + Ok(list.into_owned().into_data_parts()) + } else if let Some(view) = canonical.as_opt::() { + Ok(list_from_list_view(view.into_owned(), exec_ctx)?.into_data_parts()) + } else { + unreachable!("AnyList matcher guarantees List or ListView") + } +} + +/// Rebase a chunk's local `offsets` into global `u64` positions for the whole-column `offsets` +/// child. Each chunk's offsets are shifted by `element_base` (the number of elements already +/// emitted) so they index into the concatenated `elements`. The duplicated boundary offset is +/// dropped on every chunk after the first, so the concatenation of all chunks' contributions is a +/// single monotonic `[0, .., total_elements]` array of length `row_count + 1`. +fn global_offsets( + offsets: ArrayRef, + element_base: u64, + first: bool, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let widened = offsets.cast(DType::Primitive(PType::U64, Nullability::NonNullable))?; + let based = if element_base == 0 { + widened + } else { + let base = ConstantArray::new(element_base, widened.len()).into_array(); + widened.binary(base, Operator::Add)? + }; + let based = if first { + based + } else { + based.slice(1..based.len())? + }; + // Materialize so the child sub-stream carries a concrete array rather than a lazy expression. + Ok(based.execute::(exec_ctx)?.into_array()) +} + +/// Matcher for `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()).then_some(()) + } +} + +#[cfg(test)] +mod tests { + use futures::stream; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::StructArray; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + use vortex_io::session::RuntimeSession; + + use super::*; + use crate::layouts::chunked::writer::ChunkedLayoutStrategy; + use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::layouts::table::TableStrategy; + use crate::segments::TestSegments; + use crate::sequence::SequentialArrayStreamExt; + use crate::session::LayoutSession; + + fn layout_test_session() -> VortexSession { + vortex_array::array_session() + .with::() + .with::() + .with_tokio() + } + + fn flat_list_strategy() -> ListLayoutStrategy { + ListLayoutStrategy::default() + } + + async fn write(strategy: &S, array: ArrayRef) -> VortexResult { + let session = layout_test_session(); + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let stream = array.to_array_stream().sequenced(ptr); + strategy + .write_stream(ArrayContext::empty(), segments, stream, eof, &session) + .await + } + + fn i32_list_dtype(nullable: bool) -> DType { + DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + if nullable { + Nullability::Nullable + } else { + Nullability::NonNullable + }, + ) + } + + fn create_basic_list(validity: Validity) -> ArrayRef { + ListArray::try_new( + buffer![1i32, 2, 3, 4, 5].into_array(), + buffer![0u32, 2, 5, 5].into_array(), + validity, + ) + .unwrap() + .into_array() + } + + #[tokio::test] + async fn basic_non_nullable_input() -> VortexResult<()> { + let list = create_basic_list(Validity::NonNullable); + + let layout = write(&flat_list_strategy(), list).await?; + assert_eq!(layout.row_count(), 3); + + insta::assert_snapshot!(layout.display_tree(), @" + vortex.list, dtype: list(i32), children: 2 + ├── elements: vortex.flat, dtype: i32, segment: 0 + └── offsets: vortex.flat, dtype: u64, segment: 1 + "); + Ok(()) + } + + #[tokio::test] + async fn basic_nullable_input() -> VortexResult<()> { + let list = create_basic_list(Validity::Array( + BoolArray::from_iter([true, false, true]).into_array(), + )); + + let layout = write(&flat_list_strategy(), list).await?; + assert_eq!(layout.row_count(), 3); + + insta::assert_snapshot!(layout.display_tree(), @" + vortex.list, dtype: list(i32)?, children: 3 + ├── elements: vortex.flat, dtype: i32, segment: 0 + ├── offsets: vortex.flat, dtype: u64, segment: 1 + └── validity: vortex.flat, dtype: bool, segment: 2 + "); + Ok(()) + } + + /// Non-list input dispatches to the fallback strategy unchanged. + #[tokio::test] + async fn non_list_input_routes_to_fallback() -> VortexResult<()> { + let primitive = buffer![1i32, 2, 3].into_array(); + let layout = write(&flat_list_strategy(), primitive).await?; + insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment: 0"); + Ok(()) + } + + #[tokio::test] + async fn empty_stream_errors() { + let segments = Arc::new(TestSegments::default()); + let (_, eof) = SequenceId::root().split(); + let empty = stream::empty::>().boxed(); + let stream = SequentialStreamAdapter::new(i32_list_dtype(false), empty).sendable(); + let session = layout_test_session(); + + let res = flat_list_strategy() + .write_stream(ArrayContext::empty(), segments, stream, eof, &session) + .await; + assert!(res.is_err()) + } + + #[tokio::test] + async fn list_of_struct_tree() -> VortexResult<()> { + let struct_array = StructArray::from_fields( + [ + ("a", buffer![1i32, 2, 3, 4, 5].into_array()), + ("b", buffer![10i32, 20, 30, 40, 50].into_array()), + ] + .as_slice(), + )? + .into_array(); + let list = ListArray::try_new( + struct_array, + buffer![0u32, 2, 5, 5].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let table_strategy: Arc = + Arc::new(TableStrategy::new(Arc::clone(&flat), Arc::clone(&flat))); + let writer = ListLayoutStrategy::default().with_elements(table_strategy); + + let layout = write(&writer, list).await?; + insta::assert_snapshot!(layout.display_tree(), @" + vortex.list, dtype: list({a=i32, b=i32}), children: 2 + ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2 + │ ├── a: vortex.flat, dtype: i32, segment: 1 + │ └── b: vortex.flat, dtype: i32, segment: 2 + └── offsets: vortex.flat, dtype: u64, segment: 0 + "); + Ok(()) + } + + #[tokio::test] + async fn list_of_list_tree() -> VortexResult<()> { + let inner_list = ListArray::try_new( + buffer![1i32, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 5, 5, 6].into_array(), + Validity::NonNullable, + )? + .into_array(); + let list = ListArray::try_new( + inner_list, + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let writer = + ListLayoutStrategy::default().with_elements(Arc::new(ListLayoutStrategy::default())); + let layout = write(&writer, list).await?; + insta::assert_snapshot!(layout.display_tree(), @" + vortex.list, dtype: list(list(i32)), children: 2 + ├── elements: vortex.list, dtype: list(i32), children: 2 + │ ├── elements: vortex.flat, dtype: i32, segment: 1 + │ └── offsets: vortex.flat, dtype: u64, segment: 2 + └── offsets: vortex.flat, dtype: u64, segment: 0 + "); + Ok(()) + } + + #[tokio::test] + async fn list_of_list_of_list_tree() -> VortexResult<()> { + let innermost = ListArray::try_new( + buffer![1i32, 2, 3, 4].into_array(), + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let middle = ListArray::try_new( + innermost, + buffer![0u32, 2].into_array(), + Validity::NonNullable, + )? + .into_array(); + let outer = + ListArray::try_new(middle, buffer![0u32, 1].into_array(), Validity::NonNullable)? + .into_array(); + + let writer = ListLayoutStrategy::default().with_elements(Arc::new( + ListLayoutStrategy::default().with_elements(Arc::new(ListLayoutStrategy::default())), + )); + let layout = write(&writer, outer).await?; + insta::assert_snapshot!(layout.display_tree(), @" + vortex.list, dtype: list(list(list(i32))), children: 2 + ├── elements: vortex.list, dtype: list(list(i32)), children: 2 + │ ├── elements: vortex.list, dtype: list(i32), children: 2 + │ │ ├── elements: vortex.flat, dtype: i32, segment: 2 + │ │ └── offsets: vortex.flat, dtype: u64, segment: 3 + │ └── offsets: vortex.flat, dtype: u64, segment: 1 + └── offsets: vortex.flat, dtype: u64, segment: 0 + "); + Ok(()) + } + + #[tokio::test] + async fn chunked_list_input_with_chunked_strategy_succeeds() -> VortexResult<()> { + let chunk0 = ListArray::try_new( + buffer![1i32, 2, 3].into_array(), + buffer![0u32, 2, 3].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + let chunk1 = ListArray::try_new( + buffer![4i32, 5, 6, 7].into_array(), + buffer![0u32, 1, 4].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + + let chunked = + ChunkedArray::try_new(vec![chunk0, chunk1], i32_list_dtype(false))?.into_array(); + + let layout = write(&ChunkedLayoutStrategy::new(flat_list_strategy()), chunked).await?; + + insta::assert_snapshot!(layout.display_tree(), @" + vortex.chunked, dtype: list(i32), children: 2 + ├── [0]: vortex.list, dtype: list(i32), children: 2 + │ ├── elements: vortex.flat, dtype: i32, segment: 0 + │ └── offsets: vortex.flat, dtype: u64, segment: 1 + └── [1]: vortex.list, dtype: list(i32), children: 2 + ├── elements: vortex.flat, dtype: i32, segment: 2 + └── offsets: vortex.flat, dtype: u64, segment: 3 + "); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/mod.rs b/vortex-layout/src/layouts/mod.rs index 18df5b8f347..47fa31aa3d9 100644 --- a/vortex-layout/src/layouts/mod.rs +++ b/vortex-layout/src/layouts/mod.rs @@ -16,6 +16,7 @@ pub mod dict; pub mod file_stats; pub mod flat; pub(crate) mod foreign; +pub mod list; pub(crate) mod partitioned; pub mod repartition; pub mod row_idx; diff --git a/vortex-layout/src/layouts/struct_/mod.rs b/vortex-layout/src/layouts/struct_/mod.rs index db056f4c57b..853a6793f4c 100644 --- a/vortex-layout/src/layouts/struct_/mod.rs +++ b/vortex-layout/src/layouts/struct_/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod reader; +pub mod writer; use std::sync::Arc; @@ -21,6 +22,7 @@ use vortex_error::vortex_err; use vortex_session::SessionExt; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +pub use writer::StructStrategy; use crate::LayoutBuildContext; use crate::LayoutChildType; diff --git a/vortex-layout/src/layouts/struct_/writer.rs b/vortex-layout/src/layouts/struct_/writer.rs new file mode 100644 index 00000000000..3076c3ed39c --- /dev/null +++ b/vortex-layout/src/layouts/struct_/writer.rs @@ -0,0 +1,261 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A writer strategy for struct-typed arrays. +//! +//! [`StructStrategy`] transposes a stream of struct chunks into one ordered stream per field +//! (plus a validity stream when the struct is nullable) and writes each through a configurable +//! child strategy, producing a single [`StructLayout`]. It is a *structural* writer: it does not +//! inspect child dtypes or resolve field-path overrides itself. Dispatching a child to the right +//! layout kind is the job of the caller (see [`TableStrategy`]). +//! +//! [`TableStrategy`]: crate::layouts::table::TableStrategy + +use std::sync::Arc; + +use async_trait::async_trait; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::future::try_join; +use futures::future::try_join_all; +use futures::pin_mut; +use itertools::Itertools; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::Nullability; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_io::kanal_ext::KanalExt; +use vortex_io::session::RuntimeSessionExt; +use vortex_session::VortexSession; +use vortex_utils::aliases::DefaultHashBuilder; +use vortex_utils::aliases::hash_map::HashMap; +use vortex_utils::aliases::hash_set::HashSet; + +use crate::IntoLayout; +use crate::LayoutRef; +use crate::LayoutStrategy; +use crate::layouts::struct_::StructLayout; +use crate::segments::SegmentSinkRef; +use crate::sequence::SendableSequentialStream; +use crate::sequence::SequenceId; +use crate::sequence::SequencePointer; +use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequentialStreamExt; + +/// Writes struct-typed arrays into a [`StructLayout`], one child layout per field. +/// +/// Each field is written through a strategy resolved by direct field name: an entry in +/// `field_writers` if present, otherwise `default`. When the struct is nullable, its validity +/// bitmap is written through `validity`. +/// +/// `StructStrategy` is intentionally unaware of nested dtypes and field-path overrides. To write +/// arbitrarily nested struct trees with per-path overrides, drive it from +/// [`TableStrategy`][crate::layouts::table::TableStrategy], which dispatches on dtype and resolves +/// the per-field child strategies before handing them here. +#[derive(Clone)] +pub struct StructStrategy { + /// Per-field child strategies, keyed by direct field name. Fields without an entry use + /// `default`. + field_writers: HashMap>, + /// Strategy for fields that have no entry in `field_writers`. + default: Arc, + /// Strategy for the struct's own validity bitmap, used only when the struct is nullable. + validity: Arc, +} + +impl StructStrategy { + /// Create a new struct writer that writes every field through `default` and the validity + /// bitmap (when present) through `validity`. + pub fn new(validity: Arc, default: Arc) -> Self { + Self { + field_writers: HashMap::default(), + default, + validity, + } + } + + /// Override the strategy for a single field by name. + pub fn with_field_writer( + mut self, + name: impl Into, + writer: Arc, + ) -> Self { + self.field_writers.insert(name.into(), writer); + self + } + + /// Override the strategy for several fields by name at once. + pub fn with_field_writers( + mut self, + writers: impl IntoIterator)>, + ) -> Self { + self.field_writers.extend(writers); + self + } +} + +#[async_trait] +impl LayoutStrategy for StructStrategy { + async fn write_stream( + &self, + ctx: ArrayContext, + segment_sink: SegmentSinkRef, + stream: SendableSequentialStream, + mut eof: SequencePointer, + session: &VortexSession, + ) -> VortexResult { + let dtype = stream.dtype().clone(); + + let Some(struct_dtype) = dtype.as_struct_fields_opt() else { + vortex_bail!("StructStrategy can only write struct-typed streams, got {dtype}"); + }; + + // Check for unique field names at write time. + if HashSet::<_, DefaultHashBuilder>::from_iter(struct_dtype.names().iter()).len() + != struct_dtype.names().len() + { + vortex_bail!("StructLayout must have unique field names"); + } + let is_nullable = dtype.is_nullable(); + + // Optimization: when there are no fields, don't spawn any work and just write a trivial + // StructLayout. + if struct_dtype.nfields() == 0 && !is_nullable { + let row_count = stream + .try_fold( + 0u64, + |acc, (_, arr)| async move { Ok(acc + arr.len() as u64) }, + ) + .await?; + return Ok(StructLayout::new(row_count, dtype, vec![]).into_layout()); + } + + // stream -> stream> + let columns_session = session.clone(); + let columns_vec_stream = stream.map(move |chunk| { + let (sequence_id, chunk) = chunk?; + let mut sequence_pointer = sequence_id.descend(); + let mut ctx = columns_session.create_execution_ctx(); + let struct_chunk = chunk.clone().execute::(&mut ctx)?; + let mut columns: Vec<(SequenceId, ArrayRef)> = Vec::new(); + if is_nullable { + columns.push(( + sequence_pointer.advance(), + chunk + .validity()? + .execute_mask(chunk.len(), &mut ctx)? + .into_array(), + )); + } + + columns.extend( + struct_chunk + .iter_unmasked_fields() + .map(|field| (sequence_pointer.advance(), field.clone())), + ); + + Ok(columns) + }); + + let mut stream_count = struct_dtype.nfields(); + if is_nullable { + stream_count += 1; + } + + let (column_streams_tx, column_streams_rx): (Vec<_>, Vec<_>) = + (0..stream_count).map(|_| kanal::bounded_async(1)).unzip(); + + // Fan out column chunks to their respective transposed streams. Keep this future joined + // with the column writers so producer panics/errors cannot be hidden as channel EOF. + let handle = session.handle(); + let fanout_fut = async move { + pin_mut!(columns_vec_stream); + while let Some(result) = columns_vec_stream.next().await { + match result { + Ok(columns) => { + for (tx, column) in column_streams_tx.iter().zip_eq(columns.into_iter()) { + if tx.send(Ok(column)).await.is_err() { + vortex_bail!( + "struct column writer finished before all chunks were sent" + ); + } + } + } + Err(e) => { + let e: Arc = Arc::new(e); + for tx in column_streams_tx.iter() { + let _ = tx.send(Err(VortexError::from(Arc::clone(&e)))).await; + } + return Err(VortexError::from(e)); + } + } + } + Ok(()) + }; + + // First child column is the validity, subsequent children are the individual struct fields + let column_dtypes: Vec = if is_nullable { + std::iter::once(DType::Bool(Nullability::NonNullable)) + .chain(struct_dtype.fields()) + .collect() + } else { + struct_dtype.fields().collect() + }; + + let column_names: Vec = if is_nullable { + std::iter::once(FieldName::from("__validity")) + .chain(struct_dtype.names().iter().cloned()) + .collect() + } else { + struct_dtype.names().iter().cloned().collect() + }; + + let layout_futures: Vec<_> = column_dtypes + .into_iter() + .zip_eq(column_streams_rx) + .zip_eq(column_names) + .enumerate() + .map(move |(index, ((dtype, recv), name))| { + let column_stream = + SequentialStreamAdapter::new(dtype, recv.into_stream().boxed()).sendable(); + let child_eof = eof.split_off(); + let session = session.clone(); + let ctx = ctx.clone(); + let segment_sink = Arc::clone(&segment_sink); + handle.spawn_nested(move |h| { + // Validity is written through the validity strategy; every other field + // resolves to its named override or the default strategy. + let writer = if index == 0 && is_nullable { + Arc::clone(&self.validity) + } else { + self.field_writers + .get(&name) + .cloned() + .unwrap_or_else(|| Arc::clone(&self.default)) + }; + let session = session.with_handle(h); + + async move { + writer + .write_stream(ctx, segment_sink, column_stream, child_eof, &session) + .await + } + }) + }) + .collect(); + + let (_success, column_layouts) = try_join(fanout_fut, try_join_all(layout_futures)).await?; + // TODO(os): transposed stream could count row counts as well, + // This must hold though, all columns must have the same row count of the struct layout + let row_count = column_layouts.first().map(|l| l.row_count()).unwrap_or(0); + Ok(StructLayout::new(row_count, dtype, column_layouts).into_layout()) + } +} diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 28b3af52c87..0853478c14a 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -3,64 +3,85 @@ //! A configurable writer strategy for tabular data. //! -//! Allows the caller to override specific leaf fields with custom layout strategies. +//! [`TableStrategy`] is a *dispatcher*: it inspects the dtype of the stream it is handed and +//! routes struct columns to [`StructStrategy`], list columns to [`ListLayoutStrategy`], and +//! everything else to the configured leaf strategy. Because it hands *itself* (suitably descended) +//! to those structural writers as the strategy for their children, arbitrarily nested struct/list +//! trees are written with no manual wiring. +//! +//! The dispatcher also owns field-path overrides, letting callers force a specific leaf field — +//! at any depth — onto a custom strategy. +use std::env; use std::sync::Arc; +use std::sync::LazyLock; use async_trait::async_trait; -use futures::StreamExt; -use futures::TryStreamExt; -use futures::future::try_join_all; -use futures::pin_mut; -use itertools::Itertools; use vortex_array::ArrayContext; -use vortex_array::ArrayRef; -use vortex_array::IntoArray; -use vortex_array::VortexSessionExecute; -use vortex_array::arrays::StructArray; -use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::dtype::DType; use vortex_array::dtype::Field; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldPath; -use vortex_array::dtype::Nullability; -use vortex_error::VortexError; use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_io::kanal_ext::KanalExt; -use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; -use vortex_utils::aliases::DefaultHashBuilder; use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; -use crate::IntoLayout; use crate::LayoutRef; use crate::LayoutStrategy; -use crate::layouts::struct_::StructLayout; +use crate::layouts::list::writer::ListLayoutStrategy; +use crate::layouts::struct_::StructStrategy; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; -use crate::sequence::SequenceId; use crate::sequence::SequencePointer; -use crate::sequence::SequentialStreamAdapter; -use crate::sequence::SequentialStreamExt; -/// A configurable strategy for writing tables with nested field columns, allowing -/// overrides for specific leaf columns. +/// Whether [`TableStrategy`] writes list fields using a [`ListLayoutStrategy`] by +/// default. Disabled unless the environment variable `VORTEX_EXPERIMENTAL_LIST_LAYOUT` +/// is set to `1`. +/// +/// [`ListLayoutStrategy`]: crate::layouts::list::writer::ListLayoutStrategy +pub fn use_experimental_list_layout() -> bool { + static USE_EXPERIMENTAL_LIST_LAYOUT: LazyLock = + LazyLock::new(|| env::var("VORTEX_EXPERIMENTAL_LIST_LAYOUT").is_ok_and(|v| v == "1")); + *USE_EXPERIMENTAL_LIST_LAYOUT +} + +type ListLayoutFactory = Arc Arc + Send + Sync>; + +/// A configurable strategy for writing nested tabular data, dispatching each (sub)stream to the +/// structural writer for its dtype. +/// +/// Dispatch rules, applied to the dtype of the stream handed to [`write_stream`]: +/// - **struct** → [`StructStrategy`], with each field written by its override (if any) or by a +/// descended copy of this dispatcher. +/// - **list** → [`ListLayoutStrategy`], with `elements` written by a descended copy of this +/// dispatcher (so nested structs/lists recurse) and `offsets`/`validity` by the leaf/validity +/// strategies. Gated: only when list decomposition is enabled via +/// [`with_list_layout`][Self::with_list_layout] (off by default); otherwise a list falls through +/// to the leaf strategy. +/// - **anything else** → the leaf strategy. +/// +/// [`write_stream`]: LayoutStrategy::write_stream pub struct TableStrategy { - /// A set of leaf field overrides, e.g. to force one column to be compact-compressed. + /// A set of field-path overrides, e.g. to force one column to be compact-compressed. Keys are + /// paths relative to the level this dispatcher sits at. leaf_writers: HashMap>, - /// The writer for any validity arrays that may be present + /// The writer for any validity arrays that may be present, at any level of the tree. validity: Arc, - /// The fallback writer for any fields that do not have an explicit writer set in `leaf_writers` - fallback: Arc, + /// The writer for leaf fields, i.e. anything that is not a struct. + leaf: Arc, + /// Optional factory applied to each dynamically constructed [`ListLayoutStrategy`]. + /// Its presence also enables list decomposition. + /// + /// [`ListLayoutStrategy`]: crate::layouts::list::writer::ListLayoutStrategy + list_layout_factory: Option, } impl TableStrategy { - /// Create a new writer with the specified write strategies for validity, and for all leaf - /// fields, with no overrides. + /// Create a new dispatcher with the given `validity` strategy and `fallback` leaf strategy and + /// no overrides. /// - /// Additional overrides can be configured using the `with_leaf_strategy` method. + /// Additional per-field overrides can be configured with + /// [`with_field_writer`][Self::with_field_writer]. /// /// ## Example /// @@ -78,7 +99,8 @@ impl TableStrategy { Self { leaf_writers: Default::default(), validity, - fallback, + leaf: fallback, + list_layout_factory: None, } } @@ -136,7 +158,7 @@ impl TableStrategy { /// Override the default strategy for leaf columns that don't have overrides. pub fn with_default_strategy(mut self, default: Arc) -> Self { - self.fallback = default; + self.leaf = default; self } @@ -145,18 +167,86 @@ impl TableStrategy { self.validity = validity; self } + + /// Enable writing list fields with [`ListLayoutStrategy`]. + /// + /// **Note**: this is an unstable and experimental layout that is expected to change. + /// Using it may lead to unreadable files in the future. + pub fn with_list_layout(self) -> Self { + self.with_list_layout_factory(|strategy| Arc::new(strategy)) + } + + /// Enable writing list fields with [`ListLayoutStrategy`] and wrap each list writer. This + /// allows repartitioning or zoning to operate in the list's outer-row space before shredding. + /// + /// **Note**: this is an unstable and experimental layout that is expected to change. + /// Using it may lead to unreadable files in the future. + pub fn with_list_layout_factory( + mut self, + factory: impl Fn(ListLayoutStrategy) -> Arc + Send + Sync + 'static, + ) -> Self { + self.list_layout_factory = Some(Arc::new(factory)); + self + } } impl TableStrategy { - /// Descend into a subfield for the writer. + /// Build the [`StructStrategy`] used to write a struct-typed stream at this level. + /// + /// Each field that has an override (or a deeper override beneath it) is resolved up front; + /// every other field falls through to a clean descended dispatcher. + fn struct_strategy(&self) -> StructStrategy { + let mut field_writers: HashMap> = HashMap::default(); + + // The distinct named first-segments of our override paths are the only fields that need + // anything other than the default dispatcher. + let mut named_first: HashSet = HashSet::default(); + for path in self.leaf_writers.keys() { + if let Some(Field::Name(name)) = path.parts().first() { + named_first.insert(name.clone()); + } + } + + for name in named_first { + // `validate_path` forbids overlapping overrides, so a name has *either* an exact + // single-segment override *or* deeper overrides, never both. + let writer = match self.leaf_writers.get(&FieldPath::from_name(name.clone())) { + Some(exact) => Arc::clone(exact), + None => { + Arc::new(self.descend(&Field::Name(name.clone()))) as Arc + } + }; + field_writers.insert(name, writer); + } + + StructStrategy::new(Arc::clone(&self.validity), Arc::new(self.descend_clean())) + .with_field_writers(field_writers) + } + + /// Build the [`ListLayoutStrategy`] used to write a list field stream at this level. + /// + /// The `elements` sub-column is routed back through a clean descended dispatcher so nested + /// structs/lists recurse; `offsets` go straight to the leaf (they are always a primitive + /// column); and `validity` uses the shared validity strategy. + fn list_strategy(&self) -> Option> { + let factory = self.list_layout_factory.as_ref()?; + let list_layout = ListLayoutStrategy::default() + .with_elements(Arc::new(self.descend_clean())) + .with_offsets(Arc::clone(&self.leaf)) + .with_validity(Arc::clone(&self.validity)) + .with_fallback(Arc::clone(&self.leaf)); + Some(factory(list_layout)) + } + + /// Descend into a subfield, retaining only the overrides that apply beneath it (rebased to be + /// relative to the child). fn descend(&self, field: &Field) -> Self { - // Start with the existing set of overrides, then only retain the ones that contain - // the current field let mut new_writers = HashMap::with_capacity(self.leaf_writers.len()); for (field_path, strategy) in &self.leaf_writers { - if field_path.starts_with_field(field) + if field_path.parts().first() == Some(field) && let Some(subpath) = field_path.clone().step_into() + && !subpath.is_root() { new_writers.insert(subpath, Arc::clone(strategy)); } @@ -165,7 +255,19 @@ impl TableStrategy { Self { leaf_writers: new_writers, validity: Arc::clone(&self.validity), - fallback: Arc::clone(&self.fallback), + leaf: Arc::clone(&self.leaf), + list_layout_factory: self.list_layout_factory.clone(), + } + } + + /// A copy of this dispatcher with no overrides, used as the default child strategy for fields + /// that carry no override. + fn descend_clean(&self) -> Self { + Self { + leaf_writers: HashMap::default(), + validity: Arc::clone(&self.validity), + leaf: Arc::clone(&self.leaf), + list_layout_factory: self.list_layout_factory.clone(), } } @@ -188,7 +290,7 @@ impl TableStrategy { } } -/// Specialized strategy for when we exactly know the input schema. +/// Dispatches each stream to the structural writer for its dtype. #[async_trait] impl LayoutStrategy for TableStrategy { async fn write_stream( @@ -196,186 +298,334 @@ impl LayoutStrategy for TableStrategy { ctx: ArrayContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, - mut eof: SequencePointer, + eof: SequencePointer, session: &VortexSession, ) -> VortexResult { let dtype = stream.dtype().clone(); - // Fallback: if the array is not a struct, fallback to writing a single array. - if !dtype.is_struct() { + if dtype.is_struct() { return self - .fallback + .struct_strategy() .write_stream(ctx, segment_sink, stream, eof, session) .await; } - let struct_dtype = dtype.as_struct_fields(); - - // Check for unique field names at write time. - if HashSet::<_, DefaultHashBuilder>::from_iter(struct_dtype.names().iter()).len() - != struct_dtype.names().len() + if dtype.is_list() + && let Some(list_strategy) = self.list_strategy() { - vortex_bail!("StructLayout must have unique field names"); - } - let is_nullable = dtype.is_nullable(); - - // Optimization: when there are no fields, don't spawn any work and just write a trivial - // StructLayout. - if struct_dtype.nfields() == 0 && !is_nullable { - let row_count = stream - .try_fold( - 0u64, - |acc, (_, arr)| async move { Ok(acc + arr.len() as u64) }, - ) - .await?; - return Ok(StructLayout::new(row_count, dtype, vec![]).into_layout()); - } - - // stream -> stream> - let columns_session = session.clone(); - let columns_vec_stream = stream.map(move |chunk| { - let (sequence_id, chunk) = chunk?; - let mut sequence_pointer = sequence_id.descend(); - let mut ctx = columns_session.create_execution_ctx(); - let struct_chunk = chunk.clone().execute::(&mut ctx)?; - let mut columns: Vec<(SequenceId, ArrayRef)> = Vec::new(); - if is_nullable { - columns.push(( - sequence_pointer.advance(), - chunk - .validity()? - .execute_mask(chunk.len(), &mut ctx)? - .into_array(), - )); - } - - columns.extend( - struct_chunk - .iter_unmasked_fields() - .map(|field| (sequence_pointer.advance(), field.clone())), - ); - - Ok(columns) - }); - - let mut stream_count = struct_dtype.nfields(); - if is_nullable { - stream_count += 1; + return list_strategy + .write_stream(ctx, segment_sink, stream, eof, session) + .await; } - let (column_streams_tx, column_streams_rx): (Vec<_>, Vec<_>) = - (0..stream_count).map(|_| kanal::bounded_async(1)).unzip(); - - // Spawn a task to fan out column chunks to their respective transposed streams - let handle = session.handle(); - handle - .spawn(async move { - pin_mut!(columns_vec_stream); - while let Some(result) = columns_vec_stream.next().await { - match result { - Ok(columns) => { - for (tx, column) in column_streams_tx.iter().zip_eq(columns.into_iter()) - { - let _ = tx.send(Ok(column)).await; - } - } - Err(e) => { - let e: Arc = Arc::new(e); - for tx in column_streams_tx.iter() { - let _ = tx.send(Err(VortexError::from(Arc::clone(&e)))).await; - } - break; - } - } - } - }) - .detach(); - - // First child column is the validity, subsequence children are the individual struct fields - let column_dtypes: Vec = if is_nullable { - std::iter::once(DType::Bool(Nullability::NonNullable)) - .chain(struct_dtype.fields()) - .collect() - } else { - struct_dtype.fields().collect() - }; - - let column_names: Vec = if is_nullable { - std::iter::once(FieldName::from("__validity")) - .chain(struct_dtype.names().iter().cloned()) - .collect() - } else { - struct_dtype.names().iter().cloned().collect() - }; - - let layout_futures: Vec<_> = column_dtypes - .into_iter() - .zip_eq(column_streams_rx) - .zip_eq(column_names) - .enumerate() - .map(move |(index, ((dtype, recv), name))| { - let column_stream = - SequentialStreamAdapter::new(dtype.clone(), recv.into_stream().boxed()) - .sendable(); - let child_eof = eof.split_off(); - let field = Field::Name(name.clone()); - let session = session.clone(); - let ctx = ctx.clone(); - let segment_sink = Arc::clone(&segment_sink); - handle.spawn_nested(move |h| { - let validity = Arc::clone(&self.validity); - // descend further and try with new fields - let writer = self - .leaf_writers - .get(&FieldPath::from_name(name)) - .cloned() - .unwrap_or_else(|| { - if dtype.is_struct() { - // Step into the field path for struct columns - Arc::new(self.descend(&field)) - } else { - // Use fallback for leaf columns - Arc::clone(&self.fallback) - } - }); - let session = session.with_handle(h); - - async move { - // If we have a matching writer, we use it. - // Otherwise, we descend into a new modified one. - // Write validity stream - if index == 0 && is_nullable { - validity - .write_stream(ctx, segment_sink, column_stream, child_eof, &session) - .await - } else { - // Use the underlying writer, otherwise use the fallback writer. - writer - .write_stream(ctx, segment_sink, column_stream, child_eof, &session) - .await - } - } - }) - }) - .collect(); - - let column_layouts = try_join_all(layout_futures).await?; - // TODO(os): transposed stream could count row counts as well, - // This must hold though, all columns must have the same row count of the struct layout - let row_count = column_layouts.first().map(|l| l.row_count()).unwrap_or(0); - Ok(StructLayout::new(row_count, dtype, column_layouts).into_layout()) + // Leaf: hand off to the leaf strategy. + self.leaf + .write_stream(ctx, segment_sink, stream, eof, session) + .await } } #[cfg(test)] mod tests { + use std::num::NonZeroUsize; use std::sync::Arc; - + use std::task::Poll; + + use vortex_array::ArrayContext; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::StructFields; use vortex_array::field_path; - + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_io::runtime::single::block_on; + use vortex_io::session::RuntimeSessionExt; + + use crate::LayoutRef; + use crate::LayoutStrategy; + use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::layouts::list::List; + use crate::layouts::repartition::RepartitionStrategy; + use crate::layouts::repartition::RepartitionWriterOptions; use crate::layouts::table::TableStrategy; + use crate::layouts::zoned::Zoned; + use crate::layouts::zoned::writer::ZonedLayoutOptions; + use crate::layouts::zoned::writer::ZonedStrategy; + use crate::segments::TestSegments; + use crate::sequence::SequenceId; + use crate::sequence::SequentialArrayStreamExt; + use crate::sequence::SequentialStreamAdapter; + use crate::sequence::SequentialStreamExt; + use crate::test::SESSION; + + async fn write(strategy: &S, array: ArrayRef) -> VortexResult { + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let stream = array.to_array_stream().sequenced(ptr); + strategy + .write_stream(ArrayContext::empty(), segments, stream, eof, &SESSION) + .await + } + + /// A plain table dispatcher with no overrides. `flat` here is both the validity and leaf + /// strategy. + fn flat_table() -> TableStrategy { + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + TableStrategy::new(Arc::clone(&flat), flat) + } + + /// The dispatcher shreds a top-level struct into one child per field. + #[tokio::test] + async fn dispatches_struct() -> VortexResult<()> { + let struct_array = StructArray::from_fields( + [ + ("a", buffer![1i32, 2, 3].into_array()), + ("b", buffer![10i32, 20, 30].into_array()), + ] + .as_slice(), + )? + .into_array(); + + let layout = write(&flat_table(), struct_array).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.struct, dtype: {a=i32, b=i32}, children: 2 + ├── a: vortex.flat, dtype: i32, segment: 0 + └── b: vortex.flat, dtype: i32, segment: 1 + "); + Ok(()) + } + + /// A `list>` column: the dispatcher recurses into itself so the outer list's + /// `elements` are decomposed as a nested `ListLayout`. + #[tokio::test] + async fn dispatches_nested_list() -> VortexResult<()> { + let inner = ListArray::try_new( + buffer![1i32, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 5, 5, 6].into_array(), + Validity::NonNullable, + )? + .into_array(); + let outer = ListArray::try_new( + inner, + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let layout = write(&flat_table().with_list_layout(), outer).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.list, dtype: list(list(i32)), children: 2 + ├── elements: vortex.list, dtype: list(i32), children: 2 + │ ├── elements: vortex.flat, dtype: i32, segment: 1 + │ └── offsets: vortex.flat, dtype: u64, segment: 2 + └── offsets: vortex.flat, dtype: u64, segment: 0 + "); + Ok(()) + } + + /// A `struct<{ items: list>? }>` column: list decomposition recurses into struct + /// decomposition for the elements, and a nullable list writes a validity child. + #[tokio::test] + async fn dispatches_struct_list_struct() -> VortexResult<()> { + let inner_struct = StructArray::from_fields( + [ + ("a", buffer![1i32, 2, 3, 4, 5].into_array()), + ("b", buffer![10i32, 20, 30, 40, 50].into_array()), + ] + .as_slice(), + )? + .into_array(); + let items = ListArray::try_new( + inner_struct, + buffer![0u32, 2, 5, 5].into_array(), + Validity::Array(BoolArray::from_iter([true, false, true]).into_array()), + )? + .into_array(); + let st = StructArray::from_fields([("items", items)].as_slice())?.into_array(); + + let layout = write(&flat_table().with_list_layout(), st).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.struct, dtype: {items=list({a=i32, b=i32})?}, children: 1 + └── items: vortex.list, dtype: list({a=i32, b=i32})?, children: 3 + ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2 + │ ├── a: vortex.flat, dtype: i32, segment: 2 + │ └── b: vortex.flat, dtype: i32, segment: 3 + ├── offsets: vortex.flat, dtype: u64, segment: 0 + └── validity: vortex.flat, dtype: bool, segment: 1 + "); + Ok(()) + } + + /// A multi-chunk `list` written with a chunked leaf: each sub-column (`elements`, + /// `offsets`) becomes its own `ChunkedLayout`, so elements are chunked independently of rows. + /// This is the "list-of-chunkeds" topology top-level decomposition unlocks. + #[tokio::test] + async fn dispatches_chunked_list() -> VortexResult<()> { + let chunk0 = ListArray::try_new( + buffer![1i32, 2, 3].into_array(), + buffer![0u32, 2, 3].into_array(), + Validity::NonNullable, + )? + .into_array(); + let chunk1 = ListArray::try_new( + buffer![4i32, 5, 6, 7].into_array(), + buffer![0u32, 1, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let dtype = chunk0.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk0, chunk1], dtype)?.into_array(); + + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let dispatcher = TableStrategy::new( + Arc::clone(&flat), + Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())), + ) + .with_list_layout(); + let layout = write(&dispatcher, chunked).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.list, dtype: list(i32), children: 2 + ├── elements: vortex.chunked, dtype: i32, children: 2 + │ ├── [0]: vortex.flat, dtype: i32, segment: 0 + │ └── [1]: vortex.flat, dtype: i32, segment: 1 + └── offsets: vortex.chunked, dtype: u64, children: 2 + ├── [0]: vortex.flat, dtype: u64, segment: 2 + └── [1]: vortex.flat, dtype: u64, segment: 3 + "); + Ok(()) + } + + /// A wrapper can repartition and zone lists in outer-row space before decomposition. + #[tokio::test] + async fn wraps_list_strategy_before_decomposition() -> VortexResult<()> { + let list = ListArray::try_new( + PrimitiveArray::from_iter(0..9_i32).into_array(), + PrimitiveArray::from_iter(0..=9_u32).into_array(), + Validity::NonNullable, + )? + .into_array(); + + let row_block_size = NonZeroUsize::new(4).vortex_expect("4 is non-zero"); + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let stats = Arc::clone(&flat); + let chunked: Arc = + Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())); + let dispatcher = TableStrategy::new(Arc::clone(&flat), chunked).with_list_layout_factory( + move |list_layout| { + let zoned = ZonedStrategy::new( + list_layout, + Arc::clone(&stats), + ZonedLayoutOptions { + block_size: row_block_size, + ..Default::default() + }, + ); + Arc::new(RepartitionStrategy::new( + zoned, + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: row_block_size.get(), + block_size_target: None, + canonicalize: false, + }, + )) as Arc + }, + ); + + let layout = write(&dispatcher, list).await?; + let zoned = layout.as_::(); + assert_eq!(zoned.zone_len(), 4); + assert_eq!(zoned.nzones(), 3); + + let data = layout.child(0)?; + assert!(data.is::()); + assert_eq!(data.row_count(), 9); + Ok(()) + } + + /// A non-struct stream is not shredded; it is handed straight to the leaf strategy. + #[tokio::test] + async fn non_struct_input_uses_leaf() -> VortexResult<()> { + let primitive = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); + let layout = write(&flat_table(), primitive).await?; + insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment: 0"); + Ok(()) + } + + /// A multi-chunk struct is transposed per field; each column is written by a chunked leaf. + #[tokio::test] + async fn chunked_struct() -> VortexResult<()> { + let validity: Arc = Arc::new(FlatLayoutStrategy::default()); + let chunked_flat: Arc = + Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())); + let dispatcher = TableStrategy::new(validity, chunked_flat); + + let c0 = StructArray::from_fields( + [ + ("a", buffer![1i32, 2].into_array()), + ("b", buffer![10i32, 20].into_array()), + ] + .as_slice(), + )? + .into_array(); + let c1 = StructArray::from_fields( + [ + ("a", buffer![3i32].into_array()), + ("b", buffer![30i32].into_array()), + ] + .as_slice(), + )? + .into_array(); + let dtype = c0.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![c0, c1], dtype)?.into_array(); + + let layout = write(&dispatcher, chunked).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.struct, dtype: {a=i32, b=i32}, children: 2 + ├── a: vortex.chunked, dtype: i32, children: 2 + │ ├── [0]: vortex.flat, dtype: i32, segment: 0 + │ └── [1]: vortex.flat, dtype: i32, segment: 1 + └── b: vortex.chunked, dtype: i32, children: 2 + ├── [0]: vortex.flat, dtype: i32, segment: 2 + └── [1]: vortex.flat, dtype: i32, segment: 3 + "); + Ok(()) + } + + /// A field override on a struct field is honored ahead of the default leaf strategy. + #[tokio::test] + async fn field_override_is_used() -> VortexResult<()> { + let struct_array = StructArray::from_fields( + [ + ("a", buffer![1i32, 2, 3].into_array()), + ("b", buffer![10i32, 20, 30].into_array()), + ] + .as_slice(), + )? + .into_array(); + + let strategy = + flat_table().with_field_writer(field_path!(a), Arc::new(FlatLayoutStrategy::default())); + let layout = write(&strategy, struct_array).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.struct, dtype: {a=i32, b=i32}, children: 2 + ├── a: vortex.flat, dtype: i32, segment: 0 + └── b: vortex.flat, dtype: i32, segment: 1 + "); + Ok(()) + } #[test] #[should_panic( @@ -407,4 +657,41 @@ mod tests { ) .with_field_writer(FieldPath::root(), flat); } + + #[test] + #[should_panic(expected = "panic while transposing table stream")] + fn table_fanout_panic_propagates() { + let ctx = ArrayContext::empty(); + let segments = Arc::new(TestSegments::default()); + let (_, eof) = SequenceId::root().split(); + let dtype = DType::Struct( + StructFields::from_iter([( + "a", + DType::Primitive(PType::I32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ); + let stream = + futures::stream::poll_fn(|_| -> Poll>> { + panic!("panic while transposing table stream"); + }); + let strategy = TableStrategy::new( + Arc::new(FlatLayoutStrategy::default()), + Arc::new(FlatLayoutStrategy::default()), + ); + + block_on(|handle| async move { + let session = SESSION.clone().with_handle(handle); + strategy + .write_stream( + ctx, + segments, + SequentialStreamAdapter::new(dtype, stream).sendable(), + eof, + &session, + ) + .await + .unwrap(); + }); + } } diff --git a/vortex-layout/src/layouts/zoned/builder.rs b/vortex-layout/src/layouts/zoned/builder.rs index bbe359f7485..3c86f371fb6 100644 --- a/vortex-layout/src/layouts/zoned/builder.rs +++ b/vortex-layout/src/layouts/zoned/builder.rs @@ -8,8 +8,6 @@ use std::sync::Arc; use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::LEGACY_SESSION; -use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::arrays::StructArray; use vortex_array::builders::ArrayBuilder; @@ -74,6 +72,7 @@ impl AggregateStatsAccumulator { pub(crate) fn as_array( &mut self, + ctx: &mut ExecutionCtx, ) -> VortexResult)>> { let mut names = Vec::new(); let mut fields = Vec::new(); @@ -86,7 +85,7 @@ impl AggregateStatsAccumulator { { let values = builder.finish(); - if values.all_invalid()? { + if values.all_invalid(ctx)? { continue; } @@ -153,8 +152,8 @@ struct NamedArrays { } impl NamedArrays { - fn all_invalid(&self) -> VortexResult { + fn all_invalid(&self, ctx: &mut ExecutionCtx) -> VortexResult { // By convention the first array is the logical validity signal for the stat column. - self.arrays[0].all_invalid(&mut LEGACY_SESSION.create_execution_ctx()) + self.arrays[0].all_invalid(ctx) } } diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 6ed7b541bcb..048905dade6 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -26,6 +26,7 @@ use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::nan_count::NanCount; use vortex_array::aggregate_fn::fns::null_count::NullCount; use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -109,7 +110,7 @@ impl LayoutStrategy for ZonedStrategy { .options .aggregate_fns .clone() - .unwrap_or_else(|| default_zoned_aggregate_fns(stream.dtype())); + .unwrap_or_else(|| default_zoned_aggregate_fns(stream.dtype(), session)); let compute_session = session.clone(); let stats_accumulator = Arc::new(Mutex::new(AggregateStatsAccumulator::new( @@ -164,7 +165,10 @@ impl LayoutStrategy for ZonedStrategy { ) .await?; - let Some((stats_array, aggregate_fns)) = stats_accumulator.lock().as_array()? else { + let mut exec_ctx = session.create_execution_ctx(); + let Some((stats_array, aggregate_fns)) = + stats_accumulator.lock().as_array(&mut exec_ctx)? + else { // If we have no stats (e.g. the DType doesn't support them), then we just return the // child layout. return Ok(data_layout); @@ -192,7 +196,7 @@ impl LayoutStrategy for ZonedStrategy { } } -fn default_zoned_aggregate_fns(dtype: &DType) -> Arc<[AggregateFnRef]> { +fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[AggregateFnRef]> { let (max, min) = match dtype { DType::Utf8(_) | DType::Binary(_) => ( BoundedMax.bind(BoundedMaxOptions { @@ -218,6 +222,9 @@ fn default_zoned_aggregate_fns(dtype: &DType) -> Arc<[AggregateFnRef]> { aggregate_fns.push(NanCount.bind(EmptyOptions)); aggregate_fns.push(NullCount.bind(EmptyOptions)); + // Stats from geo extension types are discovered from the registry at runtime instead. + aggregate_fns.extend(session.aggregate_fns().zone_stat_defaults(dtype)); + aggregate_fns.into() } @@ -237,7 +244,10 @@ mod tests { #[test] fn default_aggregates_bound_variable_length_min_max() { - let aggregate_fns = default_zoned_aggregate_fns(&DType::Utf8(Nullability::NonNullable)); + let aggregate_fns = default_zoned_aggregate_fns( + &DType::Utf8(Nullability::NonNullable), + &vortex_array::array_session(), + ); assert_eq!( aggregate_fns[0].as_::().max_bytes, @@ -251,7 +261,8 @@ mod tests { #[test] fn default_aggregates_keep_fixed_width_min_max_exact() { - let aggregate_fns = default_zoned_aggregate_fns(&PType::I32.into()); + let aggregate_fns = + default_zoned_aggregate_fns(&PType::I32.into(), &vortex_array::array_session()); assert!(aggregate_fns[0].is::()); assert!(aggregate_fns[1].is::()); @@ -263,7 +274,7 @@ mod tests { let dtype = DType::Extension( Timestamp::new(TimeUnit::Microseconds, Nullability::Nullable).erased(), ); - let aggregate_fns = default_zoned_aggregate_fns(&dtype); + let aggregate_fns = default_zoned_aggregate_fns(&dtype, &vortex_array::array_session()); assert!( aggregate_fns diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index 485c2685934..663b29d9501 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -14,7 +14,7 @@ use futures::TryStreamExt; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; -use vortex_array::arrow::ArrowSessionExt; +use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_io::runtime::BlockingRuntime; @@ -128,7 +128,7 @@ mod tests { use arrow_schema::Schema; use vortex_array::ArrayRef; use vortex_array::VortexSessionExecute; - use vortex_array::arrow::FromArrowArray; + use vortex_arrow::FromArrowArray; use vortex_error::VortexResult; use super::*; diff --git a/vortex-layout/src/session.rs b/vortex-layout/src/session.rs index be938bdc192..6bb52b1d1e8 100644 --- a/vortex-layout/src/session.rs +++ b/vortex-layout/src/session.rs @@ -12,6 +12,7 @@ use crate::LayoutEncodingRef; use crate::layouts::chunked::ChunkedLayoutEncoding; use crate::layouts::dict::DictLayoutEncoding; use crate::layouts::flat::FlatLayoutEncoding; +use crate::layouts::list::ListLayoutEncoding; use crate::layouts::struct_::StructLayoutEncoding; use crate::layouts::zoned::LegacyStatsLayoutEncoding; use crate::layouts::zoned::ZonedLayoutEncoding; @@ -57,6 +58,7 @@ impl Default for LayoutSession { LegacyStatsLayoutEncoding.as_ref(), ); layouts.register(DictLayoutEncoding.id(), DictLayoutEncoding.as_ref()); + layouts.register(ListLayoutEncoding.id(), ListLayoutEncoding.as_ref()); Self { registry: layouts } } diff --git a/vortex-mask/Cargo.toml b/vortex-mask/Cargo.toml index 7aa5bb2ef85..cf681101d43 100644 --- a/vortex-mask/Cargo.toml +++ b/vortex-mask/Cargo.toml @@ -37,5 +37,13 @@ harness = false name = "rank" harness = false +[[bench]] +name = "valid_counts" +harness = false + +[[bench]] +name = "mask_iteration" +harness = false + [lints] workspace = true diff --git a/vortex-mask/benches/mask_iteration.rs b/vortex-mask/benches/mask_iteration.rs new file mode 100644 index 00000000000..84ca37991ef --- /dev/null +++ b/vortex-mask/benches/mask_iteration.rs @@ -0,0 +1,153 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Strategy comparison for "process the valid elements of a mask". +//! +//! The kernel is `sum of values[i] for every set bit i` — a stand-in for any +//! validity-gated loop. We compare: +//! +//! * `per_element_value` — `for i in 0..n { if buf.value(i) {..} }` (what callers +//! typically write after materializing validity into a mask) +//! * `bit_iterator` — arrow's `BitIterator` (tracks a word internally) +//! * `word_trailing_zeros`— iterate `u64` words, fast-path all-set/all-unset, and +//! walk set bits with `trailing_zeros` / `w &= w - 1` +//! * `set_slices` — iterate contiguous true runs (`set_slices`) +//! * `set_indices` — iterate set positions (`set_indices`) +//! +//! Run across densities to expose the dense-vs-sparse crossover. + +#![allow(clippy::unwrap_used, clippy::cast_possible_truncation)] + +use divan::Bencher; +use vortex_buffer::BitBuffer; + +fn main() { + divan::main(); +} + +const ARGS: &[(usize, f64)] = &[ + (16_384, 0.01), + (16_384, 0.10), + (16_384, 0.50), + (16_384, 0.90), + (16_384, 0.99), +]; + +fn make(len: usize, density: f64) -> (BitBuffer, Vec) { + let threshold = (density * 1000.0) as usize; + let buf = BitBuffer::from_iter((0..len).map(|i| (i * 7 + 13) % 1000 < threshold)); + let values = (0..len as u64).collect(); + (buf, values) +} + +#[divan::bench(args = ARGS)] +fn per_element_value(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + for i in 0..len { + if buf.value(i) { + acc = acc.wrapping_add(values[i]); + } + } + acc + }); +} + +#[divan::bench(args = ARGS)] +fn bit_iterator(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + for (i, set) in buf.iter().enumerate() { + if set { + acc = acc.wrapping_add(values[i]); + } + } + acc + }); +} + +#[divan::bench(args = ARGS)] +fn word_trailing_zeros(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + let chunks = buf.chunks(); + let mut base = 0usize; + for word in chunks.iter() { + if word == u64::MAX { + for k in 0..64 { + acc = acc.wrapping_add(values[base + k]); + } + } else { + let mut w = word; + while w != 0 { + let b = w.trailing_zeros() as usize; + acc = acc.wrapping_add(values[base + b]); + w &= w - 1; + } + } + base += 64; + } + let mut w = chunks.remainder_bits(); + let rem = buf.len() - base; + while w != 0 { + let b = w.trailing_zeros() as usize; + if b >= rem { + break; + } + acc = acc.wrapping_add(values[base + b]); + w &= w - 1; + } + acc + }); +} + +#[divan::bench(args = ARGS)] +fn for_each_set_index(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + buf.for_each_set_index(|i| acc = acc.wrapping_add(values[i])); + acc + }); +} + +#[divan::bench(args = ARGS)] +fn set_slices(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + for (start, end) in buf.set_slices() { + for i in start..end { + acc = acc.wrapping_add(values[i]); + } + } + acc + }); +} + +#[divan::bench(args = ARGS)] +fn set_indices(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + for i in buf.set_indices() { + acc = acc.wrapping_add(values[i]); + } + acc + }); +} diff --git a/vortex-mask/benches/rank.rs b/vortex-mask/benches/rank.rs index cd22f9ffd1f..fa2b820f75d 100644 --- a/vortex-mask/benches/rank.rs +++ b/vortex-mask/benches/rank.rs @@ -10,6 +10,11 @@ use vortex_buffer::BitBuffer; use vortex_mask::Mask; fn main() { + // Pre-resolve the one-time CpuKernel selections on the rank/select path so no + // benchmark iteration pays the first-call cost. + let warm = create_mask(512, 0.5); + let _ = warm.rank(warm.true_count() / 2); + divan::main(); } diff --git a/vortex-mask/benches/valid_counts.rs b/vortex-mask/benches/valid_counts.rs new file mode 100644 index 00000000000..425a021fa6f --- /dev/null +++ b/vortex-mask/benches/valid_counts.rs @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks for `Mask::valid_counts_for_indices`. +//! +//! This mirrors the hot path in the pco/zstd slice decoders, which call +//! `valid_counts_for_indices(&[slice_start, slice_stop])` to translate a row +//! range into a value range. The cost is dominated by counting set bits in the +//! prefix `[0, slice_stop)`, so a SIMD popcount over the bit buffer should beat +//! a bit-by-bit walk handily for large `slice_stop`. + +#![allow(clippy::unwrap_used, clippy::cast_possible_truncation)] + +use divan::Bencher; +use vortex_buffer::BitBuffer; +use vortex_mask::Mask; + +fn main() { + divan::main(); +} + +const BENCH_SIZES: &[(usize, f64)] = &[(4_096, 0.1), (4_096, 0.9), (16_384, 0.1), (16_384, 0.9)]; + +fn create_mask(len: usize, density: f64) -> Mask { + Mask::from_buffer(BitBuffer::from_iter((0..len).map(|i| { + let threshold = (density * 1000.0) as usize; + (i * 7 + 13) % 1000 < threshold + }))) +} + +/// The pco/zstd slice pattern: two indices bracketing a sub-range. The prefix +/// up to `slice_stop` must be counted, so `slice_stop` near the end is worst case. +#[divan::bench(args = BENCH_SIZES)] +fn slice_bounds(bencher: Bencher, (len, density): (usize, f64)) { + let mask = create_mask(len, density); + let indices = [len / 4, len - len / 8]; + bencher + .with_inputs(|| (&mask, indices)) + .bench_refs(|(mask, indices)| mask.valid_counts_for_indices(indices)); +} + +/// Many monotonically increasing indices spread across the whole mask. +#[divan::bench(args = BENCH_SIZES)] +fn many_indices(bencher: Bencher, (len, density): (usize, f64)) { + let mask = create_mask(len, density); + let stride = (len / 256).max(1); + let indices: Vec = (0..len).step_by(stride).collect(); + bencher + .with_inputs(|| (&mask, indices.as_slice())) + .bench_refs(|(mask, indices)| mask.valid_counts_for_indices(indices)); +} diff --git a/vortex-mask/src/intersect_by_rank.rs b/vortex-mask/src/intersect_by_rank.rs index 9b7da4a35dc..e5a8d6a08e5 100644 --- a/vortex-mask/src/intersect_by_rank.rs +++ b/vortex-mask/src/intersect_by_rank.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use vortex_buffer::BitBuffer; use vortex_buffer::BitChunkIterator; use vortex_buffer::BufferMut; +use vortex_buffer::CpuKernel; use vortex_error::VortexExpect; use crate::Mask; @@ -396,22 +397,32 @@ fn intersect_bit_buffers_dispatch( mask_buffer: &BitBuffer, true_count: usize, ) -> Mask { - #[cfg(target_arch = "x86_64")] - if std::arch::is_x86_feature_detected!("bmi2") { - return intersect_bit_buffers::(self_buffer, mask_buffer, true_count); - } - - intersect_bit_buffers::(self_buffer, mask_buffer, true_count) + type IntersectBuffers = fn(&BitBuffer, &BitBuffer, usize) -> Mask; + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + if std::arch::is_x86_feature_detected!("bmi2") { + return intersect_bit_buffers::; + } + } + intersect_bit_buffers:: + }); + KERNEL.get()(self_buffer, mask_buffer, true_count) } #[inline] fn intersect_rank_indices_dispatch(self_buffer: &BitBuffer, mask_indices: &[usize]) -> Mask { - #[cfg(target_arch = "x86_64")] - if std::arch::is_x86_feature_detected!("bmi2") { - return intersect_bit_buffer_by_rank_indices::(self_buffer, mask_indices); - } - - intersect_bit_buffer_by_rank_indices::(self_buffer, mask_indices) + type IntersectRankIndices = fn(&BitBuffer, &[usize]) -> Mask; + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + if std::arch::is_x86_feature_detected!("bmi2") { + return intersect_bit_buffer_by_rank_indices::; + } + } + intersect_bit_buffer_by_rank_indices:: + }); + KERNEL.get()(self_buffer, mask_indices) } #[inline] diff --git a/vortex-mask/src/lib.rs b/vortex-mask/src/lib.rs index 401077f0b34..1f41c367947 100644 --- a/vortex-mask/src/lib.rs +++ b/vortex-mask/src/lib.rs @@ -628,23 +628,24 @@ impl Mask { /// Given monotonically increasing `indices` in [0, n_rows], returns the /// count of valid elements up to each index. /// - /// This is O(n_rows). + /// This is O(n_rows), but the per-gap counts are computed with a SIMD + /// popcount over the underlying bit buffer rather than walking bit-by-bit. pub fn valid_counts_for_indices(&self, indices: &[usize]) -> Vec { match self { Self::AllTrue(_) => indices.to_vec(), Self::AllFalse(_) => vec![0; indices.len()], Self::Values(values) => { - let mut bool_iter = values.bit_buffer().iter(); + let buffer = values.bit_buffer(); let mut valid_counts = Vec::with_capacity(indices.len()); let mut valid_count = 0; - let mut idx = 0; + let mut prev = 0; for &next_idx in indices { - while idx < next_idx { - idx += 1; - valid_count += bool_iter - .next() - .unwrap_or_else(|| vortex_panic!("Row indices exceed array length")) - as usize; + assert!(next_idx <= buffer.len(), "Row indices exceed array length"); + // `indices` is monotonically increasing, so each gap is counted once; + // the total work across all gaps scans the prefix `[0, last_idx)` once. + if next_idx > prev { + valid_count += buffer.count_range(prev, next_idx); + prev = next_idx; } valid_counts.push(valid_count); } @@ -779,7 +780,10 @@ impl MaskValues { } let mut indices = Vec::with_capacity(self.true_count); - indices.extend(self.buffer.set_indices()); + // Word-at-a-time set-bit walk; faster than collecting `set_indices()`, + // whose per-`next` iterator state inlines less well (see + // `vortex-mask/benches/mask_iteration.rs`). + self.buffer.for_each_set_index(|i| indices.push(i)); debug_assert!(indices.is_sorted()); assert_eq!(indices.len(), self.true_count); indices diff --git a/vortex-proto/Cargo.toml b/vortex-proto/Cargo.toml index e47865ef299..ba0202c2ea0 100644 --- a/vortex-proto/Cargo.toml +++ b/vortex-proto/Cargo.toml @@ -19,6 +19,9 @@ version = { workspace = true } [package.metadata.docs.rs] all-features = true +[package.metadata.cargo-shear] +ignored = ["prost-types"] + [dependencies] prost = { workspace = true } prost-types = { workspace = true } diff --git a/vortex-proto/proto/dtype.proto b/vortex-proto/proto/dtype.proto index 51b99217efb..11feef63a8c 100644 --- a/vortex-proto/proto/dtype.proto +++ b/vortex-proto/proto/dtype.proto @@ -75,6 +75,9 @@ message Variant { } message Union { + repeated string names = 1; + repeated DType dtypes = 2; + repeated int32 type_ids = 3; // length must equal dtypes.len(); each value must fit in uint8 bool nullable = 4; } diff --git a/vortex-proto/src/generated/vortex.dtype.rs b/vortex-proto/src/generated/vortex.dtype.rs index ffbac88cea4..16188687ee5 100644 --- a/vortex-proto/src/generated/vortex.dtype.rs +++ b/vortex-proto/src/generated/vortex.dtype.rs @@ -71,8 +71,15 @@ pub struct Variant { #[prost(bool, tag = "1")] pub nullable: bool, } -#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct Union { + #[prost(string, repeated, tag = "1")] + pub names: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + #[prost(message, repeated, tag = "2")] + pub dtypes: ::prost::alloc::vec::Vec, + /// length must equal dtypes.len(); each value must fit in uint8 + #[prost(int32, repeated, tag = "3")] + pub type_ids: ::prost::alloc::vec::Vec, #[prost(bool, tag = "4")] pub nullable: bool, } diff --git a/vortex-python-cuda/README.md b/vortex-python-cuda/README.md new file mode 100644 index 00000000000..808c10ccbb0 --- /dev/null +++ b/vortex-python-cuda/README.md @@ -0,0 +1,62 @@ +# vortex-data-cuda + +CUDA extension for [Vortex](https://vortex.dev). Exports a `vortex.Array` to +[RAPIDS cuDF](https://docs.rapids.ai/api/cudf/stable/) or any +[Arrow C Device](https://arrow.apache.org/docs/format/CDeviceDataInterface.html) consumer, on the +GPU. Imported as `vortex_cuda`. + +## Install + +```bash +pip install vortex-data vortex-data-cuda # versions must match; CUDA device required +``` + +`to_cudf` also needs RAPIDS `cudf` and `pylibcudf` in the environment. + +## Export to cuDF + +`to_cudf` converts via the Arrow C Device interface: struct arrays become a `cudf.DataFrame`, +everything else a `cudf.Series`. Importing `vortex_cuda` installs it as `vortex.Array.to_cudf`. + +```python +import vortex, vortex_cuda +import pyarrow as pa + +s = vortex.array([1, None, 3]).to_cudf() # -> cudf.Series +df = vortex_cuda.to_cudf( # struct -> cudf.DataFrame + vortex.Array.from_arrow(pa.table({"x": [1, None, 3], "y": [4.0, 5.0, 6.0]})) +) +``` + +Buffers are imported zero-copy; host arrays are moved to the GPU as part of the export. cuDF keeps +shared ownership for the lifetime of the result and any view derived from it, so no extra +bookkeeping is needed. + +Signature: `to_cudf(obj, *, fallback="error")`. Only `fallback="error"` is supported +(`NotImplementedError` otherwise); raises `TypeError` for a non-`vortex.Array`, `RuntimeError` +without a CUDA device, `ImportError` if cuDF/pylibcudf are missing. + +## Export an Arrow C Device array + +`vortex.Array` exposes the standard `__arrow_c_device_array__` protocol (installed when CUDA is +available), so any Arrow-C-Device consumer can ingest it zero-copy: + +```python +import vortex, vortex_cuda, pylibcudf + +array = vortex.array([1, None, 3]) +column = pylibcudf.Column.from_arrow(array) # via the protocol + +schema_capsule, device_array_capsule = vortex_cuda.export_device_array(array) # raw capsules +``` + +`export_device_array` returns `PyCapsule`s named `"arrow_schema"` and `"arrow_device_array"`. The +consumer owns the exported structs and runs the Arrow release callbacks when done (libcudf does +this automatically); Vortex's device buffers stay alive until then. + +## Notes + +- Integer, float, bool, and string arrays (incl. nullable) are supported; nulls are preserved. +- Struct arrays without top-level nulls are supported as cuDF DataFrames. Nullable top-level + structs are rejected because cuDF DataFrames cannot represent a separate row-level struct mask. +- A CUDA device is required; there is no CPU fallback. diff --git a/vortex-python-cuda/python/vortex_cuda/__init__.py b/vortex-python-cuda/python/vortex_cuda/__init__.py index c5d1610724e..985516e6204 100644 --- a/vortex-python-cuda/python/vortex_cuda/__init__.py +++ b/vortex-python-cuda/python/vortex_cuda/__init__.py @@ -1,12 +1,93 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright the Vortex contributors -# pyright: reportMissingModuleSource=false, reportPrivateUsage=false +# pyright: reportAttributeAccessIssue=false, reportMissingModuleSource=false, reportPrivateUsage=false, reportUnknownMemberType=false, reportUnknownVariableType=false + +import importlib from . import _lib +# Private debug hooks used by CUDA bridge tests. _debug_array_metadata_dtype = _lib._debug_array_metadata_dtype _debug_array_metadata_display_values = _lib._debug_array_metadata_display_values +_debug_arrow_device_array_capsule_summary = _lib._debug_arrow_device_array_capsule_summary +_debug_consume_arrow_device_array_capsules = _lib._debug_consume_arrow_device_array_capsules + +# Public native bindings exposed by this extension module. cuda_available = _lib.cuda_available export_device_array = _lib.export_device_array -__all__ = ["cuda_available", "export_device_array"] +_SUPPORTED_FALLBACKS = frozenset({"error"}) + + +def _Array_to_cudf(self: object, *, fallback: str = "error") -> object: + return to_cudf(self, fallback=fallback) + + +def _Array___arrow_c_device_array__( + self: object, + requested_schema: object | None = None, + **kwargs: object, +) -> tuple[object, object]: + return export_device_array(self, requested_schema, **kwargs) + + +def _install_vortex_array_methods() -> None: + import vortex + + setattr(vortex.Array, "to_cudf", _Array_to_cudf) + if cuda_available(): + setattr(vortex.Array, "__arrow_c_device_array__", _Array___arrow_c_device_array__) + + +def _import_cudf_modules() -> tuple[object, object]: + try: + cudf = importlib.import_module("cudf") + pylibcudf = importlib.import_module("pylibcudf") + except ImportError as err: + raise ImportError("vortex_cuda.to_cudf requires RAPIDS cuDF and pylibcudf to be installed") from err + return cudf, pylibcudf + + +def to_cudf(obj: object, *, fallback: str = "error") -> object: + """Convert a Vortex array to a cuDF object through the Arrow Device interface. + + pylibcudf imports the exported Arrow Device array zero-copy and keeps shared ownership of + Vortex's device buffers (via libcudf's ``arrow_column``) for the lifetime of the returned + cuDF object and any view derived from it, so no extra keepalive is required here. + + ``fallback`` is reserved for future policy choices. The initial implementation + supports only ``fallback="error"`` and never falls back to host Arrow conversion. + """ + if fallback not in _SUPPORTED_FALLBACKS: + raise NotImplementedError("vortex_cuda.to_cudf currently supports only fallback='error'") + + import vortex + + if not isinstance(obj, vortex.Array): + raise TypeError(f"vortex_cuda.to_cudf expected a vortex.Array, got {type(obj).__name__}") + + if not cuda_available(): + raise RuntimeError("CUDA is not available; vortex_cuda.to_cudf requires a CUDA device") + + dtype = obj.dtype + if isinstance(dtype, vortex.StructDType) and dtype.is_nullable(): + raise NotImplementedError( + "vortex_cuda.to_cudf cannot preserve top-level nulls for struct arrays as a cuDF DataFrame" + ) + + cudf, pylibcudf = _import_cudf_modules() + + if isinstance(dtype, vortex.StructDType): + table = pylibcudf.Table.from_arrow(obj) + dataframe = cudf.DataFrame.from_pylibcudf(table) + dataframe.columns = dtype.names() + return dataframe + + column = pylibcudf.Column.from_arrow(obj) + return cudf.Series.from_pylibcudf(column) + + +_install_vortex_array_methods() + + +__all__ = ["cuda_available", "export_device_array", "to_cudf"] diff --git a/vortex-python-cuda/python/vortex_cuda/_lib.pyi b/vortex-python-cuda/python/vortex_cuda/_lib.pyi index dad7628363d..51b4e100fe7 100644 --- a/vortex-python-cuda/python/vortex_cuda/_lib.pyi +++ b/vortex-python-cuda/python/vortex_cuda/_lib.pyi @@ -3,6 +3,10 @@ def _debug_array_metadata_dtype(array: object) -> str: ... def _debug_array_metadata_display_values(array: object) -> str: ... +def _debug_arrow_device_array_capsule_summary(schema: object, device_array: object) -> dict[str, object]: ... +def _debug_consume_arrow_device_array_capsules( + schema: object, device_array: object +) -> tuple[bool, bool, bool, bool, bool, bool]: ... def cuda_available() -> bool: ... def export_device_array( array: object, requested_schema: object | None = None, **kwargs: object diff --git a/vortex-python-cuda/src/lib.rs b/vortex-python-cuda/src/lib.rs index 13107a39acb..2ac3166b103 100644 --- a/vortex-python-cuda/src/lib.rs +++ b/vortex-python-cuda/src/lib.rs @@ -462,6 +462,84 @@ fn release_exported(exported: &mut ArrowDeviceArrayWithSchema) { release_device_array(&mut exported.array); } +/// Return non-owning details from Arrow Device capsules for Python-side smoke consumers. +#[pyfunction] +fn _debug_arrow_device_array_capsule_summary<'py>( + py: Python<'py>, + schema: Bound<'py, PyCapsule>, + device_array: Bound<'py, PyCapsule>, +) -> PyResult> { + let schema = unsafe { + schema + .pointer_checked(Some(ARROW_SCHEMA_CAPSULE_NAME))? + .cast::() + .as_ref() + }; + let device_array = unsafe { + device_array + .pointer_checked(Some(ARROW_DEVICE_ARRAY_CAPSULE_NAME))? + .cast::() + .as_ref() + }; + + let summary = PyDict::new(py); + summary.set_item("schema_live", schema.release.is_some())?; + summary.set_item("array_live", device_array.array.release.is_some())?; + summary.set_item("is_cuda", device_array.device_type == ARROW_DEVICE_CUDA)?; + summary.set_item("device_type", device_array.device_type)?; + summary.set_item("device_id", device_array.device_id)?; + summary.set_item("length", device_array.array.length)?; + summary.set_item("null_count", device_array.array.null_count)?; + summary.set_item("n_buffers", device_array.array.n_buffers)?; + summary.set_item("n_children", device_array.array.n_children)?; + Ok(summary) +} + +/// Simulate a Python Arrow Device consumer taking ownership from the returned capsules. +#[pyfunction] +fn _debug_consume_arrow_device_array_capsules( + schema: Bound<'_, PyCapsule>, + device_array: Bound<'_, PyCapsule>, +) -> PyResult<(bool, bool, bool, bool, bool, bool)> { + let mut schema_ptr = schema + .pointer_checked(Some(ARROW_SCHEMA_CAPSULE_NAME))? + .cast::(); + let mut device_array_ptr = device_array + .pointer_checked(Some(ARROW_DEVICE_ARRAY_CAPSULE_NAME))? + .cast::(); + + let schema_ref = unsafe { schema_ptr.as_mut() }; + let device_array_ref = unsafe { device_array_ptr.as_mut() }; + let schema_had_release = schema_ref.release.is_some(); + let array_had_release = device_array_ref.array.release.is_some(); + + release_schema(schema_ref); + release_device_array(device_array_ref); + + let schema_release_cleared = schema_ref.release.is_none(); + let array_release_cleared = device_array_ref.array.release.is_none(); + + set_capsule_name(&schema, USED_ARROW_SCHEMA_CAPSULE_NAME)?; + set_capsule_name(&device_array, USED_ARROW_DEVICE_ARRAY_CAPSULE_NAME)?; + + Ok(( + schema_had_release, + array_had_release, + schema_release_cleared, + array_release_cleared, + schema.is_valid_checked(Some(USED_ARROW_SCHEMA_CAPSULE_NAME)), + device_array.is_valid_checked(Some(USED_ARROW_DEVICE_ARRAY_CAPSULE_NAME)), + )) +} + +fn set_capsule_name(capsule: &Bound<'_, PyCapsule>, name: &CStr) -> PyResult<()> { + let result = unsafe { ffi::PyCapsule_SetName(capsule.as_ptr(), name.as_ptr()) }; + if result != 0 { + return Err(PyErr::fetch(capsule.py())); + } + Ok(()) +} + fn schema_capsule<'py>( py: Python<'py>, schema: FFI_ArrowSchema, @@ -574,6 +652,14 @@ fn _lib(m: &Bound) -> PyResult<()> { m.add_function(wrap_pyfunction!(cuda_available, m)?)?; m.add_function(wrap_pyfunction!(_debug_array_metadata_dtype, m)?)?; m.add_function(wrap_pyfunction!(_debug_array_metadata_display_values, m)?)?; + m.add_function(wrap_pyfunction!( + _debug_arrow_device_array_capsule_summary, + m + )?)?; + m.add_function(wrap_pyfunction!( + _debug_consume_arrow_device_array_capsules, + m + )?)?; m.add_function(wrap_pyfunction!(export_device_array, m)?)?; Ok(()) } diff --git a/vortex-python-cuda/test/test_cuda.py b/vortex-python-cuda/test/test_cuda.py index de269164317..56678ab1f41 100644 --- a/vortex-python-cuda/test/test_cuda.py +++ b/vortex-python-cuda/test/test_cuda.py @@ -1,10 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright the Vortex contributors +# pyright: reportAny=false, reportExplicitAny=false +import gc +import sys import tomllib +import types from pathlib import Path -from typing import cast +from typing import Any, cast +import pytest import vortex_cuda import vortex @@ -27,3 +32,204 @@ def test_extension_exact_pins_base_package(): pyproject = tomllib.loads((Path(__file__).parents[1] / "pyproject.toml").read_text()) assert pyproject["project"]["dependencies"] == [f"vortex-data=={workspace_version()}"] + + +def test_to_cudf_is_exported(): + assert "to_cudf" in vortex_cuda.__all__ + + +def test_import_installs_array_to_cudf(monkeypatch: pytest.MonkeyPatch): + array = vortex.Array.from_range(range(0, 3)) + calls: list[tuple[object, str]] = [] + + def fake_to_cudf(obj: object, *, fallback: str = "error") -> object: + calls.append((obj, fallback)) + return "cudf-object" + + monkeypatch.setattr(vortex_cuda, "to_cudf", fake_to_cudf) + + assert cast(Any, array).to_cudf(fallback="error") == "cudf-object" + assert calls == [(array, "error")] + + +def test_import_installs_array_arrow_c_device_array_on_cuda(monkeypatch: pytest.MonkeyPatch): + array = vortex.Array.from_range(range(0, 3)) + + if not vortex_cuda.cuda_available(): + assert not hasattr(array, "__arrow_c_device_array__") + return + + calls: list[tuple[object, object | None, dict[str, object]]] = [] + + def fake_export_device_array( + exported_array: object, + requested_schema: object | None = None, + **kwargs: object, + ) -> tuple[object, object]: + calls.append((exported_array, requested_schema, kwargs)) + return "schema", "device_array" + + monkeypatch.setattr(vortex_cuda, "export_device_array", fake_export_device_array) + + assert cast(Any, array).__arrow_c_device_array__("requested", future=None) == ("schema", "device_array") + assert calls == [(array, "requested", {"future": None})] + + +def test_to_cudf_rejects_unsupported_fallback(): + with pytest.raises(NotImplementedError, match="fallback='error'"): + _ = vortex_cuda.to_cudf(vortex.Array.from_range(range(0, 3)), fallback="host") + + +def test_to_cudf_rejects_non_vortex_array(): + with pytest.raises(TypeError, match="vortex.Array"): + _ = vortex_cuda.to_cudf(object()) + + +def test_to_cudf_cuda_unavailable_rejects_without_importing_cudf(monkeypatch: pytest.MonkeyPatch): + def fail_import_cudf_modules() -> tuple[object, object]: + raise AssertionError("unexpected import of cuDF modules") + + monkeypatch.setattr(vortex_cuda, "cuda_available", lambda: False) + monkeypatch.setattr("vortex_cuda._import_cudf_modules", fail_import_cudf_modules) + + with pytest.raises(RuntimeError, match="CUDA"): + _ = vortex_cuda.to_cudf(vortex.Array.from_range(range(0, 3))) + + +def test_to_cudf_non_struct_uses_pylibcudf_column_from_arrow(monkeypatch: pytest.MonkeyPatch): + array = vortex.Array.from_range(range(0, 3)) + fake_column = object() + + class FakeSeries: + pass + + fake_series = FakeSeries() + + class FakePylibcudfColumn: + @staticmethod + def from_arrow(obj: object) -> object: + assert obj is array + return fake_column + + class FakePylibcudfTable: + @staticmethod + def from_arrow(_obj: object) -> object: + raise AssertionError("non-struct array should not import through pylibcudf.Table") + + class FakeCudfSeries: + @staticmethod + def from_pylibcudf(column: object) -> object: + assert column is fake_column + return fake_series + + monkeypatch.setattr(vortex_cuda, "cuda_available", lambda: True) + monkeypatch.setitem( + sys.modules, + "pylibcudf", + types.SimpleNamespace(Column=FakePylibcudfColumn, Table=FakePylibcudfTable), + ) + monkeypatch.setitem(sys.modules, "cudf", types.SimpleNamespace(Series=FakeCudfSeries)) + + assert vortex_cuda.to_cudf(array) is fake_series + + +def test_to_cudf_struct_uses_pylibcudf_table_from_arrow(monkeypatch: pytest.MonkeyPatch): + import pyarrow as pa + + array = vortex.Array.from_arrow(pa.table({"x": [1, None, 3], "y": [4.0, 5.0, 6.0]})) + fake_table = object() + + class FakeDataFrame: + columns: object = None + + fake_dataframe = FakeDataFrame() + + class FakePylibcudfColumn: + @staticmethod + def from_arrow(_obj: object) -> object: + raise AssertionError("struct array should not import through pylibcudf.Column") + + class FakePylibcudfTable: + @staticmethod + def from_arrow(obj: object) -> object: + assert obj is array + return fake_table + + class FakeCudfDataFrame: + @staticmethod + def from_pylibcudf(table: object) -> object: + assert table is fake_table + return fake_dataframe + + monkeypatch.setattr(vortex_cuda, "cuda_available", lambda: True) + monkeypatch.setitem( + sys.modules, + "pylibcudf", + types.SimpleNamespace(Column=FakePylibcudfColumn, Table=FakePylibcudfTable), + ) + monkeypatch.setitem(sys.modules, "cudf", types.SimpleNamespace(DataFrame=FakeCudfDataFrame)) + + assert vortex_cuda.to_cudf(array) is fake_dataframe + assert fake_dataframe.columns == ["x", "y"] + + +def test_to_cudf_nullable_struct_rejects_without_importing_cudf(monkeypatch: pytest.MonkeyPatch): + import pyarrow as pa + + array = vortex.Array.from_arrow(pa.array([{"x": 1}, None], type=pa.struct([("x", pa.int64())]))) + + def fail_import_cudf_modules() -> tuple[object, object]: + raise AssertionError("unexpected import of cuDF modules") + + monkeypatch.setattr(vortex_cuda, "cuda_available", lambda: True) + monkeypatch.setattr("vortex_cuda._import_cudf_modules", fail_import_cudf_modules) + + with pytest.raises(NotImplementedError, match="top-level nulls"): + _ = vortex_cuda.to_cudf(array) + + +def test_to_cudf_real_cudf_smoke(): + cudf_module = cast(object, pytest.importorskip("cudf")) + pylibcudf_module = cast(object, pytest.importorskip("pylibcudf")) + assert cudf_module is not None + assert pylibcudf_module is not None + + if not vortex_cuda.cuda_available(): + pytest.skip("CUDA device is not available") + + import pyarrow as pa + + int_series = cast(Any, vortex.array([1, 2, 3])).to_cudf() + assert type(int_series).__name__ == "Series" + assert int_series.to_arrow().to_pylist() == [1, 2, 3] + + nullable_int_series = cast(Any, vortex.array([1, None, 3])).to_cudf() + assert type(nullable_int_series).__name__ == "Series" + assert nullable_int_series.to_arrow().to_pylist() == [1, None, 3] + assert "" in repr(nullable_int_series) + + nullable_bool_series = cast(Any, vortex.array([True, None, False])).to_cudf() + assert type(nullable_bool_series).__name__ == "Series" + assert nullable_bool_series.to_arrow().to_pylist() == [True, None, False] + assert "" in repr(nullable_bool_series) + + string_series = cast(Any, vortex.array(["alpha", "beta", "gamma"])).to_cudf() + assert type(string_series).__name__ == "Series" + assert string_series.to_arrow().to_pylist() == ["alpha", "beta", "gamma"] + + struct_array = vortex.Array.from_arrow(pa.table({"x": [1, None, 3], "y": [4.0, 5.0, 6.0]})) + dataframe = cast(Any, struct_array).to_cudf() + assert type(dataframe).__name__ == "DataFrame" + assert list(dataframe.columns) == ["x", "y"] + assert len(dataframe) == 3 + assert "" in repr(dataframe) + assert {name: str(dtype) for name, dtype in dataframe.dtypes.items()} == { + "x": "int64", + "y": "float64", + } + + x_series = dataframe["x"] + del dataframe + _ = gc.collect() + assert x_series.to_arrow().to_pylist() == [1, None, 3] + assert "" in repr(x_series) diff --git a/vortex-python-cuda/test/test_native_bridge.py b/vortex-python-cuda/test/test_native_bridge.py index c98e70a11d5..960467ea25c 100644 --- a/vortex-python-cuda/test/test_native_bridge.py +++ b/vortex-python-cuda/test/test_native_bridge.py @@ -2,12 +2,38 @@ # SPDX-FileCopyrightText: Copyright the Vortex contributors # pyright: reportPrivateUsage=false +import gc + import pytest import vortex_cuda import vortex +def _require_cuda() -> None: + if not vortex_cuda.cuda_available(): + pytest.skip("CUDA device is not available") + + +def _assert_exported_device_array( + array: object, *, length: int, null_count: int, n_children: int +) -> tuple[object, object]: + schema, device_array = vortex_cuda.export_device_array(array) + summary = vortex_cuda._debug_arrow_device_array_capsule_summary(schema, device_array) + + assert summary["schema_live"] is True + assert summary["array_live"] is True + assert summary["is_cuda"] is True + assert summary["length"] == length + assert summary["null_count"] == null_count + assert summary["n_children"] == n_children + n_buffers = summary["n_buffers"] + assert isinstance(n_buffers, int) + assert n_buffers >= 0 + + return schema, device_array + + def test_debug_array_metadata_dtype_reads_base_vortex_array(): array = vortex.Array.from_range(range(0, 3)) @@ -64,3 +90,66 @@ def test_export_device_array_returns_capsules_or_clean_cuda_error(): schema, device_array = vortex_cuda.export_device_array(array) assert type(schema).__name__ == "PyCapsule" assert type(device_array).__name__ == "PyCapsule" + + +def test_arrow_device_export_primitive_array(): + _require_cuda() + + _ = _assert_exported_device_array(vortex.array([1, 2, 3]), length=3, null_count=0, n_children=0) + + +def test_arrow_device_export_nullable_primitive_array(): + _require_cuda() + + _ = _assert_exported_device_array(vortex.array([1, None, 3]), length=3, null_count=1, n_children=0) + + +def test_arrow_device_export_nullable_bool_array(): + _require_cuda() + + _ = _assert_exported_device_array(vortex.array([True, None, False]), length=3, null_count=1, n_children=0) + + +def test_arrow_device_export_string_array(): + _require_cuda() + + _ = _assert_exported_device_array( + vortex.array(["alpha", "beta", "a longer string that should use the varbin data buffer"]), + length=3, + null_count=0, + n_children=0, + ) + + +def test_arrow_device_export_struct_array(): + import pyarrow as pa + + _require_cuda() + + arrow_table = pa.table({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]}) + struct_array = vortex.Array.from_arrow( + pa.StructArray.from_arrays( # pyright: ignore[reportUnknownMemberType] + [arrow_table.column("a").combine_chunks(), arrow_table.column("b").combine_chunks()], + names=["a", "b"], + ) + ) + + _ = _assert_exported_device_array(struct_array, length=3, null_count=0, n_children=2) + + +def test_arrow_device_capsules_drop_unconsumed(): + _require_cuda() + + schema, device_array = _assert_exported_device_array(vortex.array([1, 2, 3]), length=3, null_count=0, n_children=0) + del schema, device_array + _ = gc.collect() + + +def test_arrow_device_capsules_consumer_release_and_used_names(): + _require_cuda() + + schema, device_array = _assert_exported_device_array(vortex.array([1, 2, 3]), length=3, null_count=0, n_children=0) + consume_result = vortex_cuda._debug_consume_arrow_device_array_capsules(schema, device_array) + assert consume_result == (True, True, True, True, True, True) + del schema, device_array + _ = gc.collect() diff --git a/vortex-python/Cargo.toml b/vortex-python/Cargo.toml index 491bff86e6a..a87548d69a9 100644 --- a/vortex-python/Cargo.toml +++ b/vortex-python/Cargo.toml @@ -52,6 +52,7 @@ pyo3-object_store = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread"], optional = true } url = { workspace = true } vortex = { workspace = true, features = ["object_store"] } +vortex-arrow = { workspace = true } vortex-python-abi = { path = "../vortex-python-abi" } vortex-tui = { workspace = true, optional = true } diff --git a/vortex-python/python/vortex/_lib/dtype.pyi b/vortex-python/python/vortex/_lib/dtype.pyi index bbbfb9fcf6d..9ba6489b0c5 100644 --- a/vortex-python/python/vortex/_lib/dtype.pyi +++ b/vortex-python/python/vortex/_lib/dtype.pyi @@ -6,6 +6,7 @@ from typing import Literal, final import pyarrow as pa class DType: + def is_nullable(self) -> bool: ... def to_arrow_schema(self) -> pa.Schema: ... def to_arrow_type(self) -> pa.DataType: ... @classmethod diff --git a/vortex-python/src/arrays/from_arrow.rs b/vortex-python/src/arrays/from_arrow.rs index e8c9ca04d60..3964742626e 100644 --- a/vortex-python/src/arrays/from_arrow.rs +++ b/vortex-python/src/arrays/from_arrow.rs @@ -13,11 +13,11 @@ use pyo3::prelude::*; use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::arrays::ChunkedArray; -use vortex::array::arrow::FromArrowArray; use vortex::dtype::DType; -use vortex::dtype::arrow::TryFromArrowType; use vortex::error::VortexError; use vortex::error::VortexResult; +use vortex_arrow::FromArrowArray; +use vortex_arrow::TryFromArrowType; use crate::arrays::PyArrayRef; use crate::arrow::FromPyArrow; diff --git a/vortex-python/src/arrays/mod.rs b/vortex-python/src/arrays/mod.rs index e22e8d95955..ad6575fdbe0 100644 --- a/vortex-python/src/arrays/mod.rs +++ b/vortex-python/src/arrays/mod.rs @@ -39,7 +39,6 @@ use vortex::array::arrays::BoolArray; use vortex::array::arrays::Chunked; use vortex::array::arrays::bool::BoolArrayExt; use vortex::array::arrays::chunked::ChunkedArrayExt; -use vortex::array::arrow::ArrowSessionExt; use vortex::array::buffer::BufferHandle; use vortex::array::builtins::ArrayBuiltins; use vortex::array::match_each_integer_ptype; @@ -52,6 +51,7 @@ use vortex::flatbuffers::WriteFlatBufferExt; use vortex::ipc::messages::EncoderMessage; use vortex::ipc::messages::MessageEncoder; use vortex::scalar_fn::fns::operators::Operator; +use vortex_arrow::ArrowSessionExt; use vortex_python_abi::BUFFER_EXPORT_CAPSULE_NAME; use vortex_python_abi::VORTEX_BUFFER_EXPORT_VERSION; use vortex_python_abi::VORTEX_BUFFER_HOST; diff --git a/vortex-python/src/dataset.rs b/vortex-python/src/dataset.rs index acb8285e5a0..504ed05aeee 100644 --- a/vortex-python/src/dataset.rs +++ b/vortex-python/src/dataset.rs @@ -25,6 +25,7 @@ use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; use vortex::io::runtime::BlockingRuntime; use vortex::layout::scan::split_by::SplitBy; +use vortex_arrow::ToArrowType; use crate::RUNTIME; use crate::arrays::PyArrayRef; diff --git a/vortex-python/src/dtype/mod.rs b/vortex-python/src/dtype/mod.rs index 59cee1fc857..ce18ceb6d83 100644 --- a/vortex-python/src/dtype/mod.rs +++ b/vortex-python/src/dtype/mod.rs @@ -33,9 +33,10 @@ use pyo3::pyclass; use pyo3::pymethods; use pyo3::types::PyType; use pyo3::wrap_pyfunction; -use vortex::array::arrow::ArrowSessionExt; use vortex::dtype::DType; -use vortex::dtype::arrow::TryFromArrowType; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ToArrowType; +use vortex_arrow::TryFromArrowType; use crate::arrow::FromPyArrow; use crate::arrow::ToPyArrow; @@ -190,6 +191,11 @@ impl PyDType { self.0.python_repr().to_string() } + /// Return whether this data type allows null values. + fn is_nullable(&self) -> bool { + self.0.is_nullable() + } + /// Construct a Vortex data type from an Arrow data type. #[classmethod] #[pyo3(signature = (arrow_dtype, *, non_nullable = false))] diff --git a/vortex-python/src/file.rs b/vortex-python/src/file.rs index f6bd1ed9cda..9550067b956 100644 --- a/vortex-python/src/file.rs +++ b/vortex-python/src/file.rs @@ -27,6 +27,7 @@ use vortex::io::runtime::BlockingRuntime; use vortex::layout::scan::scan_builder::ScanBuilder; use vortex::layout::scan::split_by::SplitBy; use vortex::layout::segments::MokaSegmentCache; +use vortex_arrow::ToArrowType; use crate::RUNTIME; use crate::arrays::PyArrayRef; diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index a8d61edfc18..182ab7fbe74 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -12,13 +12,11 @@ use pyo3_object_store::PyObjectStore; use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::IntoArray; -use vortex::array::arrow::FromArrowArray; use vortex::array::iter::ArrayIterator; use vortex::array::iter::ArrayIteratorAdapter; use vortex::array::iter::ArrayIteratorExt; use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::dtype::DType; -use vortex::dtype::arrow::TryFromArrowType; use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::file::WriteOptionsSessionExt; @@ -26,6 +24,8 @@ use vortex::file::WriteStrategyBuilder; use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::io::runtime::BlockingRuntime; +use vortex_arrow::FromArrowArray; +use vortex_arrow::TryFromArrowType; use crate::PyVortex; use crate::RUNTIME; diff --git a/vortex-python/src/iter/mod.rs b/vortex-python/src/iter/mod.rs index 89caab3e90f..757f9910ad6 100644 --- a/vortex-python/src/iter/mod.rs +++ b/vortex-python/src/iter/mod.rs @@ -20,13 +20,13 @@ use pyo3::prelude::*; use pyo3::types::PyIterator; use vortex::array::Canonical; use vortex::array::IntoArray; -use vortex::array::LEGACY_SESSION; use vortex::array::VortexSessionExecute; -use vortex::array::arrow::ArrowSessionExt; use vortex::array::iter::ArrayIterator; use vortex::array::iter::ArrayIteratorAdapter; use vortex::array::iter::ArrayIteratorExt; use vortex::dtype::DType; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::ToArrowType; use crate::arrays::PyArrayRef; use crate::arrow::IntoPyArrow; @@ -114,6 +114,7 @@ impl PyArrayIterator { /// Convert the :class:`vortex.ArrayIterator` into a :class:`pyarrow.RecordBatchReader`. /// /// Note that this performs the conversion on the current thread. + #[allow(clippy::disallowed_methods)] fn to_arrow(slf: Bound) -> PyVortexResult> { let schema = Arc::new(slf.get().dtype().to_arrow_schema()?); let target = Field::new_struct("", schema.fields().clone(), false); @@ -132,7 +133,7 @@ impl PyArrayIterator { session().arrow().execute_arrow( chunk?, Some(&target), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut session().create_execution_ctx(), ) }) .map(|chunk| chunk.map_err(|e| ArrowError::ExternalError(Box::new(e)))) diff --git a/vortex-row/Cargo.toml b/vortex-row/Cargo.toml index 9222c7d6a43..bcb50bc1276 100644 --- a/vortex-row/Cargo.toml +++ b/vortex-row/Cargo.toml @@ -18,7 +18,6 @@ version = { workspace = true } workspace = true [dependencies] -bytes = { workspace = true } smallvec = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } diff --git a/vortex-row/benches/row_encode.rs b/vortex-row/benches/row_encode.rs index 0747c5facac..deea043280b 100644 --- a/vortex-row/benches/row_encode.rs +++ b/vortex-row/benches/row_encode.rs @@ -30,6 +30,7 @@ use rand::distr::Alphanumeric; use rand::rngs::StdRng; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; @@ -41,9 +42,10 @@ static GLOBAL: MiMalloc = MiMalloc; const N: usize = 100_000; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); fn main() { + LazyLock::force(&SESSION); divan::main(); } diff --git a/vortex-row/src/codec.rs b/vortex-row/src/codec.rs index fa0959943db..fd14bb254b8 100644 --- a/vortex-row/src/codec.rs +++ b/vortex-row/src/codec.rs @@ -228,7 +228,7 @@ pub(crate) fn row_width_for_dtype(dtype: &DType) -> VortexResult { DType::Variant(_) => { vortex_bail!("row encoding does not support Variant arrays (no well-defined ordering)") } - DType::Union(_) => vortex_bail!("row encoding does not support Union arrays"), + DType::Union(..) => vortex_bail!("row encoding does not support Union arrays"), dtype => vortex_bail!("row encoding does not support dtype: {dtype:?}"), } } diff --git a/vortex-session/src/lib.rs b/vortex-session/src/lib.rs index aee36dfbadd..f69dbc8b3ee 100644 --- a/vortex-session/src/lib.rs +++ b/vortex-session/src/lib.rs @@ -96,7 +96,7 @@ mod tests { let session = VortexSession::empty(); assert!(!session.allows_unknown()); - let session = session.allow_unknown(); + session.allow_unknown(); assert!(session.allows_unknown()); } } diff --git a/vortex-session/src/session.rs b/vortex-session/src/session.rs index 31cfe9ed995..505f3999a2c 100644 --- a/vortex-session/src/session.rs +++ b/vortex-session/src/session.rs @@ -191,11 +191,11 @@ impl VortexSession { /// Inserts a session variable of type `V`, replacing any existing variable of that type. /// - /// This is the internal copy-on-write insert primitive behind [`with_some`](Self::with_some) and + /// This is the copy-on-write insert primitive behind [`with_some`](Self::with_some) and /// [`get_mut`](SessionExt::get_mut); it is not public, so a variable can only enter the type-map /// through those (or through a default inserted by [`get`](SessionExt::get)). The mutation is /// applied in place to the shared backing store, so it is visible through every clone. - fn register(&self, var: V) { + pub fn register(&self, var: V) { let var: Arc = Arc::new(var); self.0.rcu(|current| { let mut next = SessionVars::clone(current); @@ -256,9 +256,8 @@ impl VortexSession { /// Allow deserializing unknown plugin IDs as non-executable foreign placeholders. /// /// Mutates this session in place and returns it for chaining. - pub fn allow_unknown(self) -> Self { + pub fn allow_unknown(&self) { self.get_mut::().allow_unknown = true; - self } } @@ -445,7 +444,7 @@ mod tests { let session = VortexSession::empty(); assert!(!session.allows_unknown()); - let session = session.allow_unknown(); + session.allow_unknown(); assert!(session.allows_unknown()); } diff --git a/vortex-sqllogictest/Cargo.toml b/vortex-sqllogictest/Cargo.toml index ddfd79ef99e..cf7d3e56660 100644 --- a/vortex-sqllogictest/Cargo.toml +++ b/vortex-sqllogictest/Cargo.toml @@ -17,19 +17,22 @@ version = { workspace = true } anyhow = { workspace = true } async-trait = { workspace = true } bigdecimal = { workspace = true } -datafusion = { workspace = true } -datafusion-functions-nested = { workspace = true } datafusion-sqllogictest = { workspace = true } -indicatif = { workspace = true } regex = { workspace = true } rstest = { workspace = true } sqllogictest = "0.29.1" thiserror = { workspace = true } tokio = { workspace = true, features = ["full"] } vortex = { workspace = true, features = ["tokio"] } -vortex-datafusion = { workspace = true } vortex-duckdb = { workspace = true } +[dev-dependencies] +datafusion = { workspace = true } +datafusion-functions-nested.workspace = true +indicatif = { workspace = true } +tracing-subscriber = { workspace = true, features = ["env-filter"] } +vortex-datafusion = { workspace = true } + [lints] workspace = true diff --git a/vortex-sqllogictest/bin/sqllogictests-runner.rs b/vortex-sqllogictest/bin/sqllogictests-runner.rs index 80335ac8951..fd5aa0cec60 100644 --- a/vortex-sqllogictest/bin/sqllogictests-runner.rs +++ b/vortex-sqllogictest/bin/sqllogictests-runner.rs @@ -198,6 +198,11 @@ fn complete_files( } fn main() -> anyhow::Result { + drop( + tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .try_init(), + ); let mut raw_args: Vec = std::env::args().collect(); // We remove the `--complete` flag that isn't standard before we pass the rest. let complete = { diff --git a/vortex-sqllogictest/slt/duckdb/aggregate_pushdown.slt b/vortex-sqllogictest/slt/duckdb/aggregate_pushdown.slt new file mode 100644 index 00000000000..2bf91d7f007 --- /dev/null +++ b/vortex-sqllogictest/slt/duckdb/aggregate_pushdown.slt @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +include ../setup.slt.no + +statement ok +COPY ( +WITH cte AS (SELECT generate_series AS i FROM generate_series(100000)) +SELECT (CASE WHEN i > 0 THEN i ELSE NULL END) AS i FROM cte +) +TO '${WORK_DIR}/agg-pushdown.vortex'; + +query TT +EXPLAIN +SELECT sum(i), min(i), max(i), avg(i), mean(i), any_value(i), first(i), count(*), count(i) +FROM '${WORK_DIR}/agg-pushdown.vortex'; +---- +:.*UNGROUPED_AGGREGATE.* + +query I +SELECT sum(i) FROM '${WORK_DIR}/agg-pushdown.vortex'; +---- +5000050000 + +query II +SELECT min(i), max(i) FROM '${WORK_DIR}/agg-pushdown.vortex'; +---- +1 100000 + +query RR +SELECT avg(i), mean(i) FROM '${WORK_DIR}/agg-pushdown.vortex'; +---- +50000.5 50000.5 + +query II +SELECT count(*), count(i) FROM '${WORK_DIR}/agg-pushdown.vortex'; +---- +100001 100000 + +# inverse order to test count(*) not being first +query II +SELECT count(i), count(*) FROM '${WORK_DIR}/agg-pushdown.vortex'; +---- +100000 100001 + +# Just count() is rejected because not pushing it down is faster: +# we just set chunk length and don't do any aggregation ourselves +query TT +EXPLAIN SELECT count(*) FROM '${WORK_DIR}/agg-pushdown.vortex'; +---- +:.*UNGROUPED_AGGREGATE.* + +query TT +EXPLAIN SELECT count(), count(), count() FROM '${WORK_DIR}/agg-pushdown.vortex'; +---- +:.*UNGROUPED_AGGREGATE.* + +query III +SELECT count(), count(), count() FROM '${WORK_DIR}/agg-pushdown.vortex'; +---- +100001 100001 100001 + +# aggregate over scalar i.e. SELECT mean(strlen(str)) +# + cte +# + view + recursive view +# + cte referencing any_value() diff --git a/vortex-sqllogictest/slt/duckdb/cast_pushdown.slt b/vortex-sqllogictest/slt/duckdb/cast_pushdown.slt new file mode 100644 index 00000000000..6c95279b94a --- /dev/null +++ b/vortex-sqllogictest/slt/duckdb/cast_pushdown.slt @@ -0,0 +1,465 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +include ../setup.slt.no + +statement ok +COPY (SELECT '2020-01-01 02:20:30' ts) TO '$__TEST_DIR__/cast_pushdown-date.vortex'; + +# Date cast is not pushed down +query TT +EXPLAIN SELECT * FROM '$__TEST_DIR__/cast_pushdown-date.vortex' +WHERE ts::DATE = '2020-01-01'; +---- +:.*FILTER.* + +query T +SELECT ts FROM '$__TEST_DIR__/cast_pushdown-date.vortex' +WHERE ts::DATE = '2020-01-01'; +---- +2020-01-01 02:20:30 + +statement ok +COPY (SELECT * FROM ( + VALUES (1, 0),(1, 0),(1, 0),(1, 0),(1, 0),(1, 0),(2, 0),(3, 0),(3, 0),(3, 257) +) AS t(column00, column01)) +TO '$__TEST_DIR__/cast_pushdown.vortex'; + +# Column is used uncasted +query TT +EXPLAIN SELECT * FROM '$__TEST_DIR__/cast_pushdown.vortex'; +---- +:.*PROJECTION.* + +query I +SELECT DISTINCT column00 FROM '$__TEST_DIR__/cast_pushdown.vortex' ORDER BY 1; +---- +1 +2 +3 + +statement ok +CREATE VIEW v AS SELECT column01 AS x, column00 AS y FROM '$__TEST_DIR__/cast_pushdown.vortex' WHERE x > 0; + +query II +SELECT x::BIGINT, y::SMALLINT FROM v; +---- +257 3 + +statement ok +CREATE VIEW vdup AS SELECT column01 AS x, column01 AS y FROM '$__TEST_DIR__/cast_pushdown.vortex'; + +query II +SELECT x::BIGINT, y::BIGINT FROM vdup ORDER BY 1 DESC LIMIT 1; +---- +257 257 + +# Column is used casted +query TT +EXPLAIN SELECT column00::UTINYINT FROM '$__TEST_DIR__/cast_pushdown.vortex'; +---- +:.*PROJECTION.* + +query I +SELECT DISTINCT column00::UTINYINT FROM '$__TEST_DIR__/cast_pushdown.vortex' ORDER BY 1; +---- +1 +2 +3 + +# Column is used casted and uncasted +query TT +EXPLAIN SELECT column00, column00::UTINYINT FROM '$__TEST_DIR__/cast_pushdown.vortex'; +---- +:.*PROJECTION.* + +query II +SELECT column00, column00::UTINYINT FROM '$__TEST_DIR__/cast_pushdown.vortex' ORDER BY 1, 2; +---- +1 1 +1 1 +1 1 +1 1 +1 1 +1 1 +2 2 +3 3 +3 3 +3 3 + +# Column is used uncasted with filter +query TT +EXPLAIN SELECT column00 FROM '$__TEST_DIR__/cast_pushdown.vortex' +WHERE column00 > 0; +---- +:.*PROJECTION.* + +query I +SELECT column00 FROM '$__TEST_DIR__/cast_pushdown.vortex' +WHERE column00 > 0 ORDER BY 1; +---- +1 +1 +1 +1 +1 +1 +2 +3 +3 +3 + +# Column is used uncasted with filter on casted. +# Cast is pushed in WHERE separately +query TT +EXPLAIN SELECT column00 FROM '$__TEST_DIR__/cast_pushdown.vortex' +WHERE column00::UTINYINT > 0; +---- +:.*FILTER.* + +query I +SELECT column00 FROM '$__TEST_DIR__/cast_pushdown.vortex' +WHERE column00::UTINYINT > 0 ORDER BY 1; +---- +1 +1 +1 +1 +1 +1 +2 +3 +3 +3 + +# Same cast in SELECT and ORDER BY: no conflict +query TT +EXPLAIN SELECT column00::UTINYINT +FROM '$__TEST_DIR__/cast_pushdown.vortex' +ORDER BY column00::UTINYINT; +---- +:.*PROJECTION.* + +query I +SELECT column00::UTINYINT +FROM '$__TEST_DIR__/cast_pushdown.vortex' +ORDER BY column00::UTINYINT; +---- +1 +1 +1 +1 +1 +1 +2 +3 +3 +3 + +# Two different casts of the same column: conflict +query TT +EXPLAIN SELECT column00::UTINYINT, column00::USMALLINT +FROM '$__TEST_DIR__/cast_pushdown.vortex'; +---- +:.*PROJECTION.* + +query II +SELECT column00::UTINYINT, column00::USMALLINT +FROM '$__TEST_DIR__/cast_pushdown.vortex' ORDER BY 1, 2; +---- +1 1 +1 1 +1 1 +1 1 +1 1 +1 1 +2 2 +3 3 +3 3 +3 3 + +# TRY_CAST is not pushed down +query TT +EXPLAIN SELECT TRY_CAST(column00 AS UTINYINT) FROM '$__TEST_DIR__/cast_pushdown.vortex'; +---- +:.*PROJECTION.* + +query I +SELECT TRY_CAST(column00 AS UTINYINT) FROM '$__TEST_DIR__/cast_pushdown.vortex' ORDER BY 1; +---- +1 +1 +1 +1 +1 +1 +2 +3 +3 +3 + +query I +SELECT TRY_CAST(column01 AS UTINYINT) FROM '$__TEST_DIR__/cast_pushdown.vortex' ORDER BY 1 NULLS LAST; +---- +0 +0 +0 +0 +0 +0 +0 +0 +0 +NULL + + +statement error +SELECT CAST(column01 AS UTINYINT) FROM '$__TEST_DIR__/cast_pushdown.vortex'; + +# TRY_CAST is not pushed in WHERE -> conflict +query TT +EXPLAIN SELECT column00 FROM '$__TEST_DIR__/cast_pushdown.vortex' +WHERE TRY_CAST(column00 AS UTINYINT) > 0; +---- +:.*FILTER.* + +query I +SELECT column00 FROM '$__TEST_DIR__/cast_pushdown.vortex' +WHERE TRY_CAST(column00 AS UTINYINT) > 0 ORDER BY 1; +---- +1 +1 +1 +1 +1 +1 +2 +3 +3 +3 + +# Same TRY_CAST in SELECT and WHERE: CAST stays in the plan. +query TT +EXPLAIN SELECT TRY_CAST(column00 AS UTINYINT) +FROM '$__TEST_DIR__/cast_pushdown.vortex' +WHERE TRY_CAST(column00 AS UTINYINT) > 0; +---- +:.*CAST.* + +query I +SELECT TRY_CAST(column00 AS UTINYINT) +FROM '$__TEST_DIR__/cast_pushdown.vortex' +WHERE TRY_CAST(column00 AS UTINYINT) > 0 ORDER BY 1; +---- +1 +1 +1 +1 +1 +1 +2 +3 +3 +3 + +query TT +EXPLAIN SELECT column00::UTINYINT, TRY_CAST(column00 AS UTINYINT) +FROM '$__TEST_DIR__/cast_pushdown.vortex'; +---- +:.*CAST.* + +query TT +EXPLAIN SELECT column00::UTINYINT, TRY_CAST(column00 AS UTINYINT) +FROM '$__TEST_DIR__/cast_pushdown.vortex'; +---- +:.*SELECT projections.* + +query II +SELECT column00::UTINYINT, TRY_CAST(column00 AS UTINYINT) +FROM '$__TEST_DIR__/cast_pushdown.vortex' ORDER BY 1, 2; +---- +1 1 +1 1 +1 1 +1 1 +1 1 +1 1 +2 2 +3 3 +3 3 +3 3 + +# CAST and TRY_CAST of different target types: conflict, pushdown blocked +query TT +EXPLAIN SELECT column00::UTINYINT, TRY_CAST(column00 AS USMALLINT) +FROM '$__TEST_DIR__/cast_pushdown.vortex'; +---- +:.*PROJECTION.* + +query TT +EXPLAIN SELECT column00::UTINYINT, TRY_CAST(column00 AS USMALLINT) +FROM '$__TEST_DIR__/cast_pushdown.vortex'; +---- +:.*SELECT projections.* + +query II +SELECT column00::UTINYINT, TRY_CAST(column00 AS USMALLINT) +FROM '$__TEST_DIR__/cast_pushdown.vortex' ORDER BY 1, 2; +---- +1 1 +1 1 +1 1 +1 1 +1 1 +1 1 +2 2 +3 3 +3 3 +3 3 + +# i128 and u128 casts are not allowed as Vortex doesn't support these types + +query TT +EXPLAIN SELECT column00::HUGEINT FROM '$__TEST_DIR__/cast_pushdown.vortex'; +---- +:.*PROJECTION.* + +query TT +EXPLAIN SELECT column00::UHUGEINT FROM '$__TEST_DIR__/cast_pushdown.vortex'; +---- +:.*PROJECTION.* + +query TT +EXPLAIN SELECT column00 +FROM '$__TEST_DIR__/cast_pushdown.vortex' +WHERE column01::HUGEINT >= column00; +---- +:.*FILTER.* + +query TT +EXPLAIN SELECT column01 +FROM '$__TEST_DIR__/cast_pushdown.vortex' +WHERE column00::UHUGEINT >= column01 +---- +:.*FILTER.* + +statement ok +COPY (SELECT * FROM ( + VALUES ('1993-07-15'::DATE), ('1993-08-20'::DATE), ('1993-10-15'::DATE) +) AS t(d)) +TO '$__TEST_DIR__/cast_pushdown-date-filter.vortex'; + +query D +SELECT d FROM '$__TEST_DIR__/cast_pushdown-date-filter.vortex' +WHERE d >= DATE '1993-07-01' + AND d < DATE '1993-07-01' + INTERVAL '3' MONTH +ORDER BY d; +---- +1993-07-15 +1993-08-20 + +query TT +EXPLAIN SELECT d FROM '$__TEST_DIR__/cast_pushdown-date-filter.vortex' +WHERE d < DATE '1993-07-01' + INTERVAL '3' MONTH; +---- +:.*FILTER.* + +statement ok +CREATE OR REPLACE TABLE tbl(val FLOAT) + +statement ok +INSERT INTO tbl VALUES (1), (-1.2), (-1.5) + +statement ok +COPY (SELECT * FROM tbl) TO '$__TEST_DIR__/cast_pushdown-float.vortex'; + +# floating point casts are not pushed down because Vortex's behaviour, truncate +# to zero, differs from Duckdb's rounding +query TT +EXPLAIN SELECT CAST(val AS INTEGER) FROM '$__TEST_DIR__/cast_pushdown-float.vortex'; +---- +:.*PROJECTION.* + +query I +SELECT CAST(val AS INTEGER) FROM '$__TEST_DIR__/cast_pushdown-float.vortex' ORDER BY 1; +---- +-2 +-1 +1 + +statement ok +COPY (SELECT 1 AS x) TO '$__TEST_DIR__/cast_pushdown-vcol.vortex'; + +query I +SELECT file_row_number::BIGINT FROM '$__TEST_DIR__/cast_pushdown-vcol.vortex'; +---- +0 + +statement ok +CREATE VIEW virt_col AS SELECT x, file_row_number FROM '$__TEST_DIR__/cast_pushdown-vcol.vortex'; + +query I +SELECT file_row_number::BIGINT FROM virt_col; +---- +0 + +statement ok +COPY (SELECT * FROM (VALUES (0), (0), (257)) AS t(col)) +TO '$__TEST_DIR__/cast_pushdown-many.vortex'; + +query TT +EXPLAIN SELECT col::UTINYINT FROM '$__TEST_DIR__/cast_pushdown-many.vortex' +WHERE TRY_CAST(col AS UTINYINT) IS NOT NULL; +---- +:.*FILTER.*TRY_CAST.* + +query TT +EXPLAIN SELECT col::UTINYINT FROM '$__TEST_DIR__/cast_pushdown-many.vortex' +WHERE TRY_CAST(col AS UTINYINT) IS NOT NULL; +---- +:.*SELECT projections.* + +query I +SELECT col::UTINYINT FROM '$__TEST_DIR__/cast_pushdown-many.vortex' +WHERE TRY_CAST(col AS UTINYINT) IS NOT NULL; +---- +0 +0 + +# A cast must not be pushed below a filter that stayed in the plan +statement ok +COPY (SELECT * FROM (VALUES (1, 'a'), (257, 'b')) AS t(n, s)) +TO '$__TEST_DIR__/cast_pushdown-residual-filter.vortex'; + +query TT +EXPLAIN SELECT n::UTINYINT FROM '$__TEST_DIR__/cast_pushdown-residual-filter.vortex' +WHERE s || 'x' = 'ax'; +---- +:.*FILTER.* + +query TT +EXPLAIN SELECT n::UTINYINT FROM '$__TEST_DIR__/cast_pushdown-residual-filter.vortex' +WHERE s || 'x' = 'ax'; +---- +:.*SELECT projections.* + +query I +SELECT n::UTINYINT FROM '$__TEST_DIR__/cast_pushdown-residual-filter.vortex' +WHERE s || 'x' = 'ax'; +---- +1 + +# A fully pushed filter does not block cast pushdown +query TT +EXPLAIN SELECT column00::UTINYINT FROM '$__TEST_DIR__/cast_pushdown.vortex' +WHERE column01 = 0; +---- +:.*(FILTER|PROJECTION).* + +query I +SELECT DISTINCT column00::UTINYINT FROM '$__TEST_DIR__/cast_pushdown.vortex' +WHERE column01 = 0 ORDER BY 1; +---- +1 +2 +3 diff --git a/vortex-sqllogictest/slt/duckdb/projection_expression_pushdown.slt b/vortex-sqllogictest/slt/duckdb/projection_expression_pushdown.slt index 9c431097bd7..36e099004cd 100644 --- a/vortex-sqllogictest/slt/duckdb/projection_expression_pushdown.slt +++ b/vortex-sqllogictest/slt/duckdb/projection_expression_pushdown.slt @@ -351,13 +351,13 @@ ORDER BY strlen(str); 3 5 -# prefix isn't pushed down as a complex filter so WHERE is not pushed down +# || isn't pushed down as a complex filter so WHERE is not pushed down # although this is only usage of str in SELECT, we can't push strlen down # as there is usage in WHERE. # This also tests functions with multiple arguments using "str" inside query I SELECT strlen(str) FROM '${WORK_DIR}/pe-pushdown.vortex' -WHERE prefix("str", 'H') > 0 +WHERE prefix(str || 'O', 'H') > 0 ORDER BY strlen(str); ---- 2 @@ -373,13 +373,6 @@ ORDER BY strlen(str); 3 5 -# explain: prefix()/suffix() in WHERE are multi-arg uses of str, no pushdown -query TT -EXPLAIN (FORMAT JSON) -SELECT strlen(str) FROM '${WORK_DIR}/pe-pushdown.vortex' WHERE prefix("str", 'H') > 0; ----- -:SELECT projections - # conflict with concat(), no pushdown query I SELECT strlen(str) FROM '${WORK_DIR}/pe-pushdown.vortex' diff --git a/vortex-sqllogictest/src/duckdb.rs b/vortex-sqllogictest/src/duckdb.rs index 1052781d430..cbf42feae7a 100644 --- a/vortex-sqllogictest/src/duckdb.rs +++ b/vortex-sqllogictest/src/duckdb.rs @@ -63,11 +63,15 @@ impl DuckDB { fn normalize_column_type(logical_type: &LogicalTypeRef) -> DFColumnType { let type_id = logical_type.as_type_id(); - if type_id == LogicalType::int32().as_type_id() + if type_id == LogicalType::int8().as_type_id() + || type_id == LogicalType::int16().as_type_id() + || type_id == LogicalType::int32().as_type_id() || type_id == LogicalType::int64().as_type_id() + || type_id == LogicalType::int128().as_type_id() + || type_id == LogicalType::uint8().as_type_id() + || type_id == LogicalType::uint16().as_type_id() || type_id == LogicalType::uint32().as_type_id() || type_id == LogicalType::uint64().as_type_id() - || type_id == LogicalType::int128().as_type_id() || type_id == LogicalType::uint128().as_type_id() { DFColumnType::Integer diff --git a/vortex-tensor/Cargo.toml b/vortex-tensor/Cargo.toml index a293e604b5b..a6782eea309 100644 --- a/vortex-tensor/Cargo.toml +++ b/vortex-tensor/Cargo.toml @@ -18,11 +18,11 @@ workspace = true [dependencies] vortex-array = { workspace = true } +vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } vortex-compressor = { workspace = true } vortex-error = { workspace = true } vortex-session = { workspace = true } -vortex-utils = { workspace = true } arrow-array = { workspace = true } arrow-schema = { workspace = true } @@ -32,10 +32,5 @@ num-traits = { workspace = true } prost = { workspace = true } [dev-dependencies] -divan = { workspace = true } -mimalloc = { workspace = true } -rand = { workspace = true } -rand_distr = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } -vortex-btrblocks = { path = "../vortex-btrblocks" } diff --git a/vortex-tensor/src/encodings/l2_denorm.rs b/vortex-tensor/src/encodings/l2_denorm.rs index 6cb4fcb0626..d07734bc6ff 100644 --- a/vortex-tensor/src/encodings/l2_denorm.rs +++ b/vortex-tensor/src/encodings/l2_denorm.rs @@ -7,9 +7,9 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_compressor::CascadingCompressor; -use vortex_compressor::ctx::CompressorContext; -use vortex_compressor::estimate::CompressionEstimate; -use vortex_compressor::estimate::EstimateVerdict; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; +use vortex_compressor::scheme::EstimateVerdict; use vortex_compressor::scheme::Scheme; use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; diff --git a/vortex-tensor/src/lib.rs b/vortex-tensor/src/lib.rs index 17b13b1ba1e..6826f3ad191 100644 --- a/vortex-tensor/src/lib.rs +++ b/vortex-tensor/src/lib.rs @@ -13,10 +13,10 @@ use std::sync::Arc; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; -use vortex_array::arrow::ArrowSessionExt; use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_array::session::ArraySessionExt; +use vortex_arrow::ArrowSessionExt; use vortex_session::VortexSession; use crate::scalar_fns::cosine_similarity::CosineSimilarity; diff --git a/vortex-tensor/src/types/vector/arrow.rs b/vortex-tensor/src/types/vector/arrow.rs index a011dc50c49..05c1890e7e1 100644 --- a/vortex-tensor/src/types/vector/arrow.rs +++ b/vortex-tensor/src/types/vector/arrow.rs @@ -20,17 +20,17 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::arrow::ArrowExport; -use vortex_array::arrow::ArrowExportVTable; -use vortex_array::arrow::ArrowImport; -use vortex_array::arrow::ArrowImportVTable; -use vortex_array::arrow::ArrowSession; -use vortex_array::arrow::ArrowSessionExt; -use vortex_array::arrow::FromArrowArray; use vortex_array::dtype::DType; -use vortex_array::dtype::arrow::FromArrowType; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; +use vortex_arrow::FromArrowType; use vortex_error::VortexResult; use vortex_session::registry::CachedId; use vortex_session::registry::Id; @@ -179,10 +179,6 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrow::ArrowExport; - use vortex_array::arrow::ArrowImport; - use vortex_array::arrow::ArrowSession; - use vortex_array::arrow::ArrowSessionExt; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; @@ -190,6 +186,10 @@ mod tests { use vortex_array::dtype::StructFields; use vortex_array::dtype::extension::ExtDType; use vortex_array::validity::Validity; + use vortex_arrow::ArrowExport; + use vortex_arrow::ArrowImport; + use vortex_arrow::ArrowSession; + use vortex_arrow::ArrowSessionExt; use vortex_buffer::buffer; use super::*; diff --git a/vortex-test/compat-gen/Cargo.toml b/vortex-test/compat-gen/Cargo.toml index b22c49112b8..4fbd90a6273 100644 --- a/vortex-test/compat-gen/Cargo.toml +++ b/vortex-test/compat-gen/Cargo.toml @@ -23,9 +23,9 @@ path = "src/main.rs" # Vortex crates vortex = { workspace = true, features = ["files", "tokio", "zstd"] } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } -vortex-fsst = { workspace = true, features = ["_test-harness"] } vortex-session = { workspace = true } # TPC-H generation diff --git a/vortex-test/compat-gen/src/fixtures/arrays/datasets/clickbench.rs b/vortex-test/compat-gen/src/fixtures/arrays/datasets/clickbench.rs index 7cae6944463..12d86c41241 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/datasets/clickbench.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/datasets/clickbench.rs @@ -11,7 +11,7 @@ use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; -use vortex_array::arrow::FromArrowArray; +use vortex_arrow::FromArrowArray; use vortex_error::VortexResult; use vortex_error::vortex_err; diff --git a/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs b/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs index 9d831bef143..59e97a7e587 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs @@ -8,7 +8,7 @@ use tpchgen_arrow::RecordBatchIterator; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; -use vortex_array::arrow::FromArrowArray; +use vortex_arrow::FromArrowArray; use vortex_error::VortexResult; use crate::fixtures::DatasetFixture; diff --git a/vortex-test/e2e-cuda/src/lib.rs b/vortex-test/e2e-cuda/src/lib.rs index 493d7c7d23c..1c6530ce612 100644 --- a/vortex-test/e2e-cuda/src/lib.rs +++ b/vortex-test/e2e-cuda/src/lib.rs @@ -45,9 +45,9 @@ use vortex::array::arrays::StructArray; use vortex::array::arrays::TemporalArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::varbinview::BinaryView; -use vortex::array::arrow::ArrowSessionExt; use vortex::array::stream::ArrayStreamExt; use vortex::array::validity::Validity; +use vortex::arrow::ArrowSessionExt; use vortex::buffer::Buffer; use vortex::buffer::ByteBuffer; use vortex::dtype::DType; diff --git a/vortex-tui/Cargo.toml b/vortex-tui/Cargo.toml index 33cd1958ba2..34b90ecfc05 100644 --- a/vortex-tui/Cargo.toml +++ b/vortex-tui/Cargo.toml @@ -55,6 +55,7 @@ taffy = { workspace = true } vortex = { version = "0.1.0", path = "../vortex", default-features = false, features = [ "files", ] } +vortex-arrow = { workspace = true } # Native-only dependencies (gated behind "native" feature) arrow-array = { workspace = true, optional = true } @@ -74,7 +75,6 @@ vortex-datafusion = { workspace = true, optional = true } # WASM-only dependencies [target.'cfg(target_arch = "wasm32")'.dependencies] console_error_panic_hook = "0.1.7" -js-sys = "0.3.81" ratzilla = "0.3" wasm-bindgen = "0.2.104" wasm-bindgen-futures = { workspace = true } diff --git a/vortex-tui/src/convert.rs b/vortex-tui/src/convert.rs index f6798413422..17fd5dcbc96 100644 --- a/vortex-tui/src/convert.rs +++ b/vortex-tui/src/convert.rs @@ -13,16 +13,16 @@ use parquet::arrow::ParquetRecordBatchStreamBuilder; use tokio::fs::File; use tokio::io::AsyncWriteExt; use vortex::array::ArrayRef; -use vortex::array::arrow::FromArrowArray; use vortex::array::stream::ArrayStreamAdapter; use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::dtype::DType; -use vortex::dtype::arrow::FromArrowType; use vortex::error::VortexExpect; use vortex::error::vortex_err; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; use vortex::session::VortexSession; +use vortex_arrow::FromArrowArray; +use vortex_arrow::FromArrowType; /// Compression strategy to use when converting Parquet files to Vortex format. #[derive(Clone, Copy, Debug, Default, ValueEnum)] diff --git a/vortex-tui/src/main.rs b/vortex-tui/src/main.rs index cc1e740eb83..f7e4a0e1ff1 100644 --- a/vortex-tui/src/main.rs +++ b/vortex-tui/src/main.rs @@ -8,7 +8,8 @@ use vortex_tui::launch; #[tokio::main] async fn main() -> anyhow::Result<()> { - let session = VortexSession::default().with_tokio().allow_unknown(); + let session = VortexSession::default().with_tokio(); + session.allow_unknown(); if let Err(err) = launch(&session).await { // Defer help/version/usage errors back to clap so their formatting // and exit codes match the standalone-binary convention exactly. diff --git a/vortex-utils/Cargo.toml b/vortex-utils/Cargo.toml index ca6c639ec0d..a9af93d998a 100644 --- a/vortex-utils/Cargo.toml +++ b/vortex-utils/Cargo.toml @@ -16,7 +16,6 @@ version = { workspace = true } [dependencies] dashmap = { workspace = true, optional = true } hashbrown = { workspace = true } -vortex-error = { workspace = true } [lints] workspace = true diff --git a/vortex-utils/src/aliases/hash_map.rs b/vortex-utils/src/aliases/hash_map.rs index 07047562b46..5937e2c1fc1 100644 --- a/vortex-utils/src/aliases/hash_map.rs +++ b/vortex-utils/src/aliases/hash_map.rs @@ -9,6 +9,8 @@ pub type RandomState = hashbrown::DefaultHashBuilder; pub type HashMap = hashbrown::HashMap; /// Entry type for HashMap. pub type Entry<'a, K, V, S> = hashbrown::hash_map::Entry<'a, K, V, S>; +/// Entry type for HashMap. +pub type EntryRef<'a, 'b, K, Q, V, S, A> = hashbrown::hash_map::EntryRef<'a, 'b, K, Q, V, S, A>; /// IntoIter type for HashMap. pub type IntoIter = hashbrown::hash_map::IntoIter; /// HashTable type alias. diff --git a/vortex-web/crate/Cargo.toml b/vortex-web/crate/Cargo.toml index 1b65fe832eb..b4de2003c28 100644 --- a/vortex-web/crate/Cargo.toml +++ b/vortex-web/crate/Cargo.toml @@ -19,6 +19,7 @@ serde_json = { workspace = true } vortex = { path = "../../vortex", default-features = false, features = [ "files", ] } +vortex-arrow = { workspace = true } [target.'cfg(target_arch = "wasm32")'.dependencies] futures = { workspace = true } diff --git a/vortex-web/crate/src/lib.rs b/vortex-web/crate/src/lib.rs index 03b36ee4515..84e1b1d6004 100644 --- a/vortex-web/crate/src/lib.rs +++ b/vortex-web/crate/src/lib.rs @@ -17,7 +17,7 @@ use vortex::session::VortexSession; mod wasm; static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::default() - .with_handle(WasmRuntime::handle()) - .allow_unknown() + let session = VortexSession::default().with_handle(WasmRuntime::handle()); + session.allow_unknown(); + session }); diff --git a/vortex-web/crate/src/wasm.rs b/vortex-web/crate/src/wasm.rs index 04abe8ee6d3..0323a98c271 100644 --- a/vortex-web/crate/src/wasm.rs +++ b/vortex-web/crate/src/wasm.rs @@ -20,7 +20,6 @@ use futures::future::BoxFuture; use serde::Serialize; use vortex::array::ArrayRef; use vortex::array::VortexSessionExecute; -use vortex::array::arrow::ArrowSessionExt; use vortex::array::buffer::BufferHandle; use vortex::array::dtype::DType; use vortex::array::serde::SerializedArray; @@ -40,6 +39,7 @@ use vortex::layout::layouts::flat::Flat; use vortex::layout::scan::scan_builder::ScanBuilder; use vortex::session::VortexSession; use vortex::session::registry::ReadContext; +use vortex_arrow::ArrowSessionExt; use wasm_bindgen::prelude::*; use wasm_bindgen_futures::JsFuture; diff --git a/vortex-web/package-lock.json b/vortex-web/package-lock.json index 0617348f82c..ff8c9ce66af 100644 --- a/vortex-web/package-lock.json +++ b/vortex-web/package-lock.json @@ -1450,9 +1450,9 @@ } }, "node_modules/@oxc-resolver/binding-android-arm-eabi": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-android-arm-eabi/-/binding-android-arm-eabi-11.21.3.tgz", - "integrity": "sha512-eNU11A2WNizh04v3uyaJCootrHIaS0B9aHYXvAvVnPNk4xYSjMUjHnhQ6dewPN2MRYDskV85d1N0Aw0WNWhcyg==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-android-arm-eabi/-/binding-android-arm-eabi-11.23.0.tgz", + "integrity": "sha512-8IJyWRLVAyhTfe9/TIEbQqSQnl5rUqYJrUOS6Dkr+Mq9FGHMxDGeiEmwkBqCvDP5KckpPh/GYSgbag66O6JsCw==", "cpu": [ "arm" ], @@ -1464,9 +1464,9 @@ ] }, "node_modules/@oxc-resolver/binding-android-arm64": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-android-arm64/-/binding-android-arm64-11.21.3.tgz", - "integrity": "sha512-8Q+ZjTLvn2dIcWsrmhdrEihm7q+ag/k+mkry7Z+t0QbbHaVxXQfvH9AewyVMh/WrpEKhQ3DDgx9fYbqeCpeOEw==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-android-arm64/-/binding-android-arm64-11.23.0.tgz", + "integrity": "sha512-pprVojnNhHxupwTT2gdeUlkxll6XEvWWBk3oVicOSNVWQC99OBnDhMQDoirqnzrE1bScQSMS2JgPpqdlrhz/Fg==", "cpu": [ "arm64" ], @@ -1478,9 +1478,9 @@ ] }, "node_modules/@oxc-resolver/binding-darwin-arm64": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-darwin-arm64/-/binding-darwin-arm64-11.21.3.tgz", - "integrity": "sha512-wkh0qKZGHXVUDxFw3oA1TXnU2BDYY/r775oJflGeIr8uDPPoN2pk8gijQIzYRT6hoql/lg3+Tx/SaTn9e2/aGg==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-darwin-arm64/-/binding-darwin-arm64-11.23.0.tgz", + "integrity": "sha512-mbIrWIMAJeytyee36OyUP5XH92TP7FaKaQ2m5AjokKy7STgjrhRt7SMXqpqLjhGm6Xn721Xmsg6H3Rtd9YQETw==", "cpu": [ "arm64" ], @@ -1492,9 +1492,9 @@ ] }, "node_modules/@oxc-resolver/binding-darwin-x64": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-darwin-x64/-/binding-darwin-x64-11.21.3.tgz", - "integrity": "sha512-HbNc23FAQYbuyDV2vBWMez4u4mrsm5RAkniGZAWqr6lYZ3N4beeqIb776jzwRl8qL2zRhHVXpUj97X0QgogVzg==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-darwin-x64/-/binding-darwin-x64-11.23.0.tgz", + "integrity": "sha512-UnIphmZ1LazUCr9DXWaKYWtKDefPMbgLsywaoYxRqVCNHhq4MM6d2q1Nz1i9Vzxt5i+cE2nRUYpAUHr/lijNYA==", "cpu": [ "x64" ], @@ -1506,9 +1506,9 @@ ] }, "node_modules/@oxc-resolver/binding-freebsd-x64": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-freebsd-x64/-/binding-freebsd-x64-11.21.3.tgz", - "integrity": "sha512-K6xNsTUPEUdfrn0+kbMq5nOUB5w1C5pavPQngt4TM2FpN91lP0PBe2srSpamb4d69O7h86oAi/qWX/kZNRSjkw==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-freebsd-x64/-/binding-freebsd-x64-11.23.0.tgz", + "integrity": "sha512-aaZ/cSEYFkSxgS2hOrobT6RQcsWNviOX8dW6CEkVx2/UYkAf9MeHbjl3W0usWV53rVV//ndBdn2nb1y7jsu4lw==", "cpu": [ "x64" ], @@ -1520,9 +1520,9 @@ ] }, "node_modules/@oxc-resolver/binding-linux-arm-gnueabihf": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-11.21.3.tgz", - "integrity": "sha512-VcFmOpcpWX1zoEy8M58tR2M9YxM+Z9RuQhqAx5q0CTmrruaP7Gveejg75hzd/5sg5nk9G3aLALEa3hE2FsmmTQ==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-11.23.0.tgz", + "integrity": "sha512-IoJLvO5SjLSVMaq83BNTrPCb1FppvoJc1IhZ5CoUVl3PykUBku7D+LK1j0GSurhJcIc6zfjghsvaZNpq5ev6Mg==", "cpu": [ "arm" ], @@ -1534,9 +1534,9 @@ ] }, "node_modules/@oxc-resolver/binding-linux-arm-musleabihf": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-arm-musleabihf/-/binding-linux-arm-musleabihf-11.21.3.tgz", - "integrity": "sha512-quVoxFLBy43hWaQbbDtQNRwAX5vX76mv7n64icAtQcJ3eNgVeblqmkupF/hAneNthdqSlnd1sTjb3aQSaDPaCQ==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-arm-musleabihf/-/binding-linux-arm-musleabihf-11.23.0.tgz", + "integrity": "sha512-vskFpwg44T/LFsfjSCnVZ5ygcuqzPC1yUzVEiKa8BgHAQz0+QLQQW3EGWLPVi8EXFghzjR4EtgPBtOhCjU4jdw==", "cpu": [ "arm" ], @@ -1548,9 +1548,9 @@ ] }, "node_modules/@oxc-resolver/binding-linux-arm64-gnu": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-11.21.3.tgz", - "integrity": "sha512-X0AqNZgcD07Q4V3RDK18/vYOj/HQT/FnmEFGYS2jTWqY7JO13ryE3TEs3eAIgUJhBnNkpEaiXqz3VK8M7qQhWQ==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-11.23.0.tgz", + "integrity": "sha512-//TcHVhrChyw5RYtgts6WO7KcWq9387c1Z5Zvhqpk/ktAbyaRYgBZrpSY1GDCFq50ASt6B6jhh+JxB1rB45IAg==", "cpu": [ "arm64" ], @@ -1565,9 +1565,9 @@ ] }, "node_modules/@oxc-resolver/binding-linux-arm64-musl": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-arm64-musl/-/binding-linux-arm64-musl-11.21.3.tgz", - "integrity": "sha512-YkaQnaKYdbuaXvRt5Qd0GpbihzVnyfR6z1SpYfIUC6RTu4NF7lDKPjVkYb+jRI2gedVO2rVpN35Y6akG6ud4Lw==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-arm64-musl/-/binding-linux-arm64-musl-11.23.0.tgz", + "integrity": "sha512-ZFqlwiTf7CXLLSGyAR9tYiO33LiaeIEXW+xm42d8mnUGpDgPltyrCGYtQezyMMEXvjhOgCz1X+i7sbDTJEx+bg==", "cpu": [ "arm64" ], @@ -1582,9 +1582,9 @@ ] }, "node_modules/@oxc-resolver/binding-linux-ppc64-gnu": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-11.21.3.tgz", - "integrity": "sha512-gB9HwhrPiFqUzDeEq+y/CgAijz1YdI6BnXz5GaH2Pa9cWdutchlkGFAiAuGb/PjVQpiK6NFKzFuztxrweoit7A==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-11.23.0.tgz", + "integrity": "sha512-oZ5LeN5+H1R19dRjTAxKrxQguH+AsemHcnthEfFxf4OjmBSty2doHLeSmMunKy3zpTHJQ3lh3Af+dNS+W6dYeA==", "cpu": [ "ppc64" ], @@ -1599,9 +1599,9 @@ ] }, "node_modules/@oxc-resolver/binding-linux-riscv64-gnu": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-riscv64-gnu/-/binding-linux-riscv64-gnu-11.21.3.tgz", - "integrity": "sha512-zjDWBlYk8QGv0H8dsPUWqkfjYIIjG2TvspGkzXL0eImbgxtZorA/klKeHyolevoT3Kvbi+1iMr9Lhrh7jf54Og==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-riscv64-gnu/-/binding-linux-riscv64-gnu-11.23.0.tgz", + "integrity": "sha512-O4ciFDyX5ebQd0qkb1bjAIg8IEfiLT03GbSeylwlwlUMK9KwBWaALwrxSbc0Msaz4U6iPj+T9eRXpD5mxBfmvA==", "cpu": [ "riscv64" ], @@ -1616,9 +1616,9 @@ ] }, "node_modules/@oxc-resolver/binding-linux-riscv64-musl": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-riscv64-musl/-/binding-linux-riscv64-musl-11.21.3.tgz", - "integrity": "sha512-4UfsQvacV388y1zpXL7C1x1FNYaV52JtuNRiuzrfQA2z1z6ElVrsidkGsrvQ5EgeSq1Pj7kaKqrgGkvFuxJ/tw==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-riscv64-musl/-/binding-linux-riscv64-musl-11.23.0.tgz", + "integrity": "sha512-P3o8Y9kISYjcxadmbO+94ThRwLhwGuDAbA7dcdd4+YLpfeF+mmobz8fXf4NmSdfSqjyRSkceJDBRZha9NVYkiQ==", "cpu": [ "riscv64" ], @@ -1633,9 +1633,9 @@ ] }, "node_modules/@oxc-resolver/binding-linux-s390x-gnu": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-11.21.3.tgz", - "integrity": "sha512-b5uH+HKH0MP5mNBYaK75SKsJbw52URqrx2LavYdq6wb0l3ExAG5niYRP9DWUNHdKilpaBVM2bXk9HNWrH3ew7Q==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-11.23.0.tgz", + "integrity": "sha512-oj03m1E3RmTFczKhcKJDzHaEDKJnPIsDcQFVxBJsSdXGSuIPdt5TvcM332FfMQgzI6yDJqyl4InrnFfXrmUTKQ==", "cpu": [ "s390x" ], @@ -1650,9 +1650,9 @@ ] }, "node_modules/@oxc-resolver/binding-linux-x64-gnu": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-x64-gnu/-/binding-linux-x64-gnu-11.21.3.tgz", - "integrity": "sha512-PjYlmilBpNRh2ntXNYAK3Am5w/nPfEpnU/96iNx7CI8EzAn12J4JRiec63wHJTH31nLoCNxBg/829pN+3CfG3Q==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-x64-gnu/-/binding-linux-x64-gnu-11.23.0.tgz", + "integrity": "sha512-BqJxbSC8FdP7mSuSpRePTGHm0hXWV+dfz//f7SjsteZncLaBgWTBmi/OZNv7sX6CyG/Pt/eJkPorP+DkMOhMwQ==", "cpu": [ "x64" ], @@ -1667,9 +1667,9 @@ ] }, "node_modules/@oxc-resolver/binding-linux-x64-musl": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-x64-musl/-/binding-linux-x64-musl-11.21.3.tgz", - "integrity": "sha512-QTBAb7JuHlZ7JUEyM8UiQi2f7m/L4swBhP2TNpYIDc9Wp/wRw1G/8sl6i13aIzQAXH7LKIm294LeOHd0lQR8zA==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-linux-x64-musl/-/binding-linux-x64-musl-11.23.0.tgz", + "integrity": "sha512-utmw+VmUrW4K8LI5/6jhg4aGYKJHOIjQ9syYOOA6pF3w7haKu4r4enTe2U0C04/HbUvkq/Zif43xFsKW1Pnq9w==", "cpu": [ "x64" ], @@ -1684,9 +1684,9 @@ ] }, "node_modules/@oxc-resolver/binding-openharmony-arm64": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-openharmony-arm64/-/binding-openharmony-arm64-11.21.3.tgz", - "integrity": "sha512-4j1DFwjwv36ec9kds0jU/ucQ5Ha4ERO/H95BxR5JFf0kqUUAJ1kwII7XhTc1vZrkdJkvLGC9Q2MbpObpum8RBg==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-openharmony-arm64/-/binding-openharmony-arm64-11.23.0.tgz", + "integrity": "sha512-V6lbRrthHa4TbvsLjPtg+EkXT1tRY+s4I8rYLXUfiHlZzGx3sLv1EH9CEOOevjvUYHLsbe/gqCIc73XnQfPb9A==", "cpu": [ "arm64" ], @@ -1698,9 +1698,9 @@ ] }, "node_modules/@oxc-resolver/binding-wasm32-wasi": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-wasm32-wasi/-/binding-wasm32-wasi-11.21.3.tgz", - "integrity": "sha512-i8oluoel5kru/j1WNrjmQSiA3GQ7wvIYVR1IwIoZtKogAhya2iub+ZKIeSIkcJOrnzQ18Tzl/F+kL3fYOxZLvA==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-wasm32-wasi/-/binding-wasm32-wasi-11.23.0.tgz", + "integrity": "sha512-gRoOxQPdnAmIAjxcuQNBxfihvx+wjTaQM/9/eP12xwnGNawOG/+Zz9RHN4WNSxT45b5CrscK4NB8aPh+oZQXAQ==", "cpu": [ "wasm32" ], @@ -1708,18 +1708,18 @@ "license": "MIT", "optional": true, "dependencies": { - "@emnapi/core": "1.11.0", - "@emnapi/runtime": "1.11.0", - "@napi-rs/wasm-runtime": "^1.1.5" + "@emnapi/core": "1.11.1", + "@emnapi/runtime": "1.11.1", + "@napi-rs/wasm-runtime": "^1.1.6" }, "engines": { "node": ">=14.0.0" } }, "node_modules/@oxc-resolver/binding-wasm32-wasi/node_modules/@emnapi/core": { - "version": "1.11.0", - "resolved": "https://registry.npmjs.org/@emnapi/core/-/core-1.11.0.tgz", - "integrity": "sha512-l9Oo58x0HOP5znGzVhYW9U3e5wVuA4LAZU2AGezTmkhO1CgQRFDhDg4nneHsu/t3WniXg9QrG2nIXL/ZS8ln8Q==", + "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, @@ -1729,9 +1729,9 @@ } }, "node_modules/@oxc-resolver/binding-wasm32-wasi/node_modules/@emnapi/runtime": { - "version": "1.11.0", - "resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.11.0.tgz", - "integrity": "sha512-55coeOFKHv1ywEcUXJtWU5f+Jr/W5tZDvZig8DLKSwUN1JpROQ4rk/SNOQiFWmaR/VKF4zuFyW1B8JduOSv6Pg==", + "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, @@ -1751,9 +1751,9 @@ } }, "node_modules/@oxc-resolver/binding-win32-arm64-msvc": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-11.21.3.tgz", - "integrity": "sha512-M/8dw8dD6aOs+NlPJax401CZB9I7Aut84isQLgALGGwke4Afvw+/7yYhZb94yXf6t2sPLhQLmSmtSV+2FhsOWg==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-11.23.0.tgz", + "integrity": "sha512-CgTGMYsJVe1eUiCdJTpGw21svXw79ITsemN1h0hcNkiswasDbN5MoibSLY+gRMWP5syfEz5iffrjZnwEP8xeUA==", "cpu": [ "arm64" ], @@ -1765,9 +1765,9 @@ ] }, "node_modules/@oxc-resolver/binding-win32-x64-msvc": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-win32-x64-msvc/-/binding-win32-x64-msvc-11.21.3.tgz", - "integrity": "sha512-H7BCt/VnS9hnmMp42eGhZ99izSCRvlnWwy/N71K1/J8QoExwY4262Z8QiEkMDtduRJrztayDxETTckmUuAVL9Q==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/@oxc-resolver/binding-win32-x64-msvc/-/binding-win32-x64-msvc-11.23.0.tgz", + "integrity": "sha512-gUGJpr+Rn6zMxm5juApV0K3U845i8t47o8k+rbO0BHbi4PoJIfSPeQmrE2dgohQm2g5k6iviNFyXCGqvmaYUpw==", "cpu": [ "x64" ], @@ -1779,9 +1779,9 @@ ] }, "node_modules/@rolldown/binding-android-arm64": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-android-arm64/-/binding-android-arm64-1.1.4.tgz", + "integrity": "sha512-EZLpf/8y7GXkkra90ML47kzik/GMP3EMcE9bPyHmRfxLC6z9+aW5A8poCsoxjrT5GfEcNAAvWwUHjvP1pUQkfw==", "cpu": [ "arm64" ], @@ -1796,9 +1796,9 @@ } }, "node_modules/@rolldown/binding-darwin-arm64": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-arm64/-/binding-darwin-arm64-1.1.4.tgz", + "integrity": "sha512-aUi+HBvmYb7j8krl1+qJgkG8C17fO79gk3c+jPw4S8glRFc1DTija9S3EyaTSQUm5GJXYKDAsugBEhFHH2vYiQ==", "cpu": [ "arm64" ], @@ -1813,9 +1813,9 @@ } }, "node_modules/@rolldown/binding-darwin-x64": { - "version": "1.1.3", - "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-x64/-/binding-darwin-x64-1.1.3.tgz", - "integrity": "sha512-YtiBp4disu6V560loT6PjMdiRaWmVvDNrUunAalbiFx2ggeJwxdAsgZMcoGP17uyAsTwAj5V1niksxlHnVQ1Sw==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-x64/-/binding-darwin-x64-1.1.4.tgz", + "integrity": "sha512-F7hHC3gwY11+vByKPRWqwGbeXWVgKmL+pTGCinaEhdihzBV2aQ0fvZOch9cXYUOKuKKq429HeYXOqQLc7wFCEg==", "cpu": [ "x64" ], @@ -1830,9 +1830,9 @@ } }, "node_modules/@rolldown/binding-freebsd-x64": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-freebsd-x64/-/binding-freebsd-x64-1.1.4.tgz", + "integrity": "sha512-sI5yw+7s92SK6odiEhD5lKCBlWcpjHS5qyqpVQbZAJ0fIzEUXrmbl3DH2ybR3PZogulNJF+COLtmA8hUfvkCCQ==", "cpu": [ "x64" ], @@ -1847,9 +1847,9 @@ } }, "node_modules/@rolldown/binding-linux-arm-gnueabihf": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-1.1.4.tgz", + "integrity": "sha512-mCi0OKgEieFircrtVYmQAFGszRtMnZ6fpZAXrxanXAu7lqZcsK1E1RAaZNG0uKAnxox3B1f4EyQNnoyMfN1vAA==", "cpu": [ "arm" ], @@ -1864,9 +1864,9 @@ } }, "node_modules/@rolldown/binding-linux-arm64-gnu": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-1.1.4.tgz", + "integrity": "sha512-B9Ial3Kv5sh0SHnB1g/QWcUQCEvCF6QKGAl4zXypYj65mVI+B4AhFBwPtSN7pDrJeIx8Z7zdy4ntx+wQABom7w==", "cpu": [ "arm64" ], @@ -1881,9 +1881,9 @@ } }, "node_modules/@rolldown/binding-linux-arm64-musl": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-musl/-/binding-linux-arm64-musl-1.1.4.tgz", + "integrity": "sha512-lZVym0PuHE1KZ22gmFTC15lAkrg9iTszR617oYRB/iPY1A56ywoJzVKOJBKaot5RiikCObmur6pogpse3gRcng==", "cpu": [ "arm64" ], @@ -1898,9 +1898,9 @@ } }, "node_modules/@rolldown/binding-linux-ppc64-gnu": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-1.1.4.tgz", + "integrity": "sha512-t2DNiLJWNTbnEHyUzTumldML6ET4/g16467LZoDDJ3tSxGvguL5/NyC2lCsNKuyRycg9XeDQF5SSv+TNOhQEXg==", "cpu": [ "ppc64" ], @@ -1915,9 +1915,9 @@ } }, "node_modules/@rolldown/binding-linux-s390x-gnu": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-1.1.4.tgz", + "integrity": "sha512-0WIRnL1Uw4BvTZRLQt+PVgo6ZKTJadlC2btP+/EOXv2f/DWbY0rEgl+y834mIVwP1FkTlWVTrGGJXf12lru7EQ==", "cpu": [ "s390x" ], @@ -1932,9 +1932,9 @@ } }, "node_modules/@rolldown/binding-linux-x64-gnu": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-gnu/-/binding-linux-x64-gnu-1.1.4.tgz", + "integrity": "sha512-JWtGshGfX+oENAKonoNkqEJX+7hC8yfhi9GUyPX1VX4mdh1y5r+ZiJLR5XzAB0aoP6s/PcILsGjKq8O0mm24bw==", "cpu": [ "x64" ], @@ -1949,9 +1949,9 @@ } }, "node_modules/@rolldown/binding-linux-x64-musl": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-musl/-/binding-linux-x64-musl-1.1.4.tgz", + "integrity": "sha512-rT6yQcxUuXs4CnbofqwHRRV0iem349rLMYpTjkgQGLjrY4ado/eDzwPZPTCgTOlF6Nkp8NEv70yLMTn6qkWxsQ==", "cpu": [ "x64" ], @@ -1966,9 +1966,9 @@ } }, "node_modules/@rolldown/binding-openharmony-arm64": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-openharmony-arm64/-/binding-openharmony-arm64-1.1.4.tgz", + "integrity": "sha512-KXMGoboq5cyaCQjDA4GLuRiOwBQ0EyFnJoVViLeZ45/3rFItRODEr+NdsBcVpll40hhNArlm/speWGRvj08LzA==", "cpu": [ "arm64" ], @@ -1983,9 +1983,9 @@ } }, "node_modules/@rolldown/binding-wasm32-wasi": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-wasm32-wasi/-/binding-wasm32-wasi-1.1.4.tgz", + "integrity": "sha512-5K83rb36oJiY7BCyE9zLZtGcPV4g5wvq+xwdO0XPIwDVZI8cyB/AUjkNXGb92/rnmezEkjMOpgY61rtwjQtFwg==", "cpu": [ "wasm32" ], @@ -2036,9 +2036,9 @@ } }, "node_modules/@rolldown/binding-win32-arm64-msvc": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-1.1.4.tgz", + "integrity": "sha512-PnWBtw3TV5KOg69HQQDR0mnQuyCmSGR2pAB4DC1rPF808fgKeTUMj2EOEyKATpgiuxuR5APQmiDO7PDgEjTFSA==", "cpu": [ "arm64" ], @@ -2053,9 +2053,9 @@ } }, "node_modules/@rolldown/binding-win32-x64-msvc": { - "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==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-x64-msvc/-/binding-win32-x64-msvc-1.1.4.tgz", + "integrity": "sha512-M1lpniBePobTfsa7Ks9a199e1akxsXn+GYBUKsEzv3YFzOm1HJAMNwKI3qr0Zq+mxwx9gOZoTdP1yXRYsZUocQ==", "cpu": [ "x64" ], @@ -2299,9 +2299,9 @@ } }, "node_modules/@tailwindcss/node": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/node/-/node-4.3.1.tgz", - "integrity": "sha512-6NDaqRoAMSXD1mr/RXu0HBvNE9a2n5tHPsxu9XHLws8o4Twes5rBM2205SUUiJ9goAtadrN6xTGX0UDEwp/N4A==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/node/-/node-4.3.2.tgz", + "integrity": "sha512-yWP/sqEcBLaD8JuA6zNwxoYKr75qxTioYwlRwekj5Jr/I5GXnoJfjetH/psLUIv74cYTH2lBUEzBkinthoYcBg==", "dev": true, "license": "MIT", "dependencies": { @@ -2311,37 +2311,37 @@ "lightningcss": "1.32.0", "magic-string": "^0.30.21", "source-map-js": "^1.2.1", - "tailwindcss": "4.3.1" + "tailwindcss": "4.3.2" } }, "node_modules/@tailwindcss/oxide": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide/-/oxide-4.3.1.tgz", - "integrity": "sha512-yVPyo8RNkabVr3O2EhHEE0Rewu7YKzc1DhIqfL46LKveFrmu9XbDazNOJY7/GRuvw1h6u3utWnR29H/p5JPlgA==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide/-/oxide-4.3.2.tgz", + "integrity": "sha512-z8ZgnzX8gdNoWLBLqBPoh/sjnxkwvf9ZuWjnO0l0yIzbLa5/9S+eC5QxGZKRobVHIC3/1BoMWjHblqWjcgFgag==", "dev": true, "license": "MIT", "engines": { "node": ">= 20" }, "optionalDependencies": { - "@tailwindcss/oxide-android-arm64": "4.3.1", - "@tailwindcss/oxide-darwin-arm64": "4.3.1", - "@tailwindcss/oxide-darwin-x64": "4.3.1", - "@tailwindcss/oxide-freebsd-x64": "4.3.1", - "@tailwindcss/oxide-linux-arm-gnueabihf": "4.3.1", - "@tailwindcss/oxide-linux-arm64-gnu": "4.3.1", - "@tailwindcss/oxide-linux-arm64-musl": "4.3.1", - "@tailwindcss/oxide-linux-x64-gnu": "4.3.1", - "@tailwindcss/oxide-linux-x64-musl": "4.3.1", - "@tailwindcss/oxide-wasm32-wasi": "4.3.1", - "@tailwindcss/oxide-win32-arm64-msvc": "4.3.1", - "@tailwindcss/oxide-win32-x64-msvc": "4.3.1" + "@tailwindcss/oxide-android-arm64": "4.3.2", + "@tailwindcss/oxide-darwin-arm64": "4.3.2", + "@tailwindcss/oxide-darwin-x64": "4.3.2", + "@tailwindcss/oxide-freebsd-x64": "4.3.2", + "@tailwindcss/oxide-linux-arm-gnueabihf": "4.3.2", + "@tailwindcss/oxide-linux-arm64-gnu": "4.3.2", + "@tailwindcss/oxide-linux-arm64-musl": "4.3.2", + "@tailwindcss/oxide-linux-x64-gnu": "4.3.2", + "@tailwindcss/oxide-linux-x64-musl": "4.3.2", + "@tailwindcss/oxide-wasm32-wasi": "4.3.2", + "@tailwindcss/oxide-win32-arm64-msvc": "4.3.2", + "@tailwindcss/oxide-win32-x64-msvc": "4.3.2" } }, "node_modules/@tailwindcss/oxide-android-arm64": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-android-arm64/-/oxide-android-arm64-4.3.1.tgz", - "integrity": "sha512-SVlyf61g374l5cHyg8x9kf5xmLcOaxvOTsbsqDnSsDJaKOEFZ7GCvi84VAVGpxojYOs1+3K6M0UjXfqPU8vmOQ==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-android-arm64/-/oxide-android-arm64-4.3.2.tgz", + "integrity": "sha512-WHxqIuHpvZ5VtdX6GTl1Ik/Vp2YuN42Et+0CdeaVd/frQ9jAvGmvR8vLT+jk3e8/Q3x8kECB9+R17pgpp2BulA==", "cpu": [ "arm64" ], @@ -2356,9 +2356,9 @@ } }, "node_modules/@tailwindcss/oxide-darwin-arm64": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-darwin-arm64/-/oxide-darwin-arm64-4.3.1.tgz", - "integrity": "sha512-hVnWLwv+e/l7c4WKyVtHVrIPvYdqWHjRB3MDIqARynzFtnQg85kmQEFCbV9Ja0VVx4xXTIiDWY60Y7iz/iNoDA==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-darwin-arm64/-/oxide-darwin-arm64-4.3.2.tgz", + "integrity": "sha512-GZypeUY/IDJW3877KeM+O67vbXr3MBnbtEL4aYhNErv/JWZhye2vGSWWG9tB6iiqR2MqRNkY8IOUy4NdSZV26w==", "cpu": [ "arm64" ], @@ -2373,9 +2373,9 @@ } }, "node_modules/@tailwindcss/oxide-darwin-x64": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-darwin-x64/-/oxide-darwin-x64-4.3.1.tgz", - "integrity": "sha512-Cf7abu0WVgbhU7ANgPUnSAvm7nCvMweusHb8FnaHlLfv/Caq4GYaEZg7ZImzzmjx4lIAfuS8q+eLIS7A7IzxIg==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-darwin-x64/-/oxide-darwin-x64-4.3.2.tgz", + "integrity": "sha512-UIIzmefR6KO1sDU7MzRqAxC8iBpft/VhkGjTjnhoS6k7Z3rQ9wEgA1ODSiyH/tcSYssulNm4Ci3hOeK1jH7ccQ==", "cpu": [ "x64" ], @@ -2390,9 +2390,9 @@ } }, "node_modules/@tailwindcss/oxide-freebsd-x64": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-freebsd-x64/-/oxide-freebsd-x64-4.3.1.tgz", - "integrity": "sha512-ZZqzX2Y+GXtXXfqSfpJhDm60OoZfvLHLCgm+J7NVqgHHJjG/m9ugZI77RwTsVd4fnBJuCFP6Ae6kTJb71UdS8g==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-freebsd-x64/-/oxide-freebsd-x64-4.3.2.tgz", + "integrity": "sha512-GN+uAmcI6DNspnCDwtOAZrTz6oukJnp337qZvxqCGLd3BHBzJpO0ZbTLRvJNdztOeAmTzewewGIMPb0tk2R4WA==", "cpu": [ "x64" ], @@ -2407,9 +2407,9 @@ } }, "node_modules/@tailwindcss/oxide-linux-arm-gnueabihf": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-linux-arm-gnueabihf/-/oxide-linux-arm-gnueabihf-4.3.1.tgz", - "integrity": "sha512-/Ah/xik0LaMYfv9DZ0S/t4pBlBNYOcqtRwusjgovHkvT8ixueWCLyJjsaF5kQIckjb4IT8Q6K6p/iPmZMixYgg==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-linux-arm-gnueabihf/-/oxide-linux-arm-gnueabihf-4.3.2.tgz", + "integrity": "sha512-4ABn7qSbdHRwTiDiuWNegCyb5+2FJ4vKIKc3DmKrvAFw7MU1Lm11dIkTPwUaFdTzc7IsOpDbqBrlh0x6y36U/w==", "cpu": [ "arm" ], @@ -2424,9 +2424,9 @@ } }, "node_modules/@tailwindcss/oxide-linux-arm64-gnu": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-linux-arm64-gnu/-/oxide-linux-arm64-gnu-4.3.1.tgz", - "integrity": "sha512-gqdFoVJlw444GvpnheZLHmvTzSxI/cOUUh2KSNejQjTcYkW062SVD+En0rUgD+QV91bz1XGIGtt1HJd48xUGbQ==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-linux-arm64-gnu/-/oxide-linux-arm64-gnu-4.3.2.tgz", + "integrity": "sha512-wDgEIGwoM8w8pufh9LVt1PahDgNdKXrLC2qfAnV3vAmococ9RWbxeAw4pxPttd/TsJfwjyLf90Dg1y9y8I6Emw==", "cpu": [ "arm64" ], @@ -2441,9 +2441,9 @@ } }, "node_modules/@tailwindcss/oxide-linux-arm64-musl": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-linux-arm64-musl/-/oxide-linux-arm64-musl-4.3.1.tgz", - "integrity": "sha512-Bwv9KwOvE0VKa86xPFif9b9c3Y1NxOV1P0gLti/IYaWEsQYZXDlxfGEtA8mdDZ7SG3wyNXAWYT5SIn3giL57oA==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-linux-arm64-musl/-/oxide-linux-arm64-musl-4.3.2.tgz", + "integrity": "sha512-J5Nuk0uZQIiMTJj3LEx4sAA9tMFUoXQZFv1J6An+QGYe53HKRJuFDi0rpq/tuouCZeAbOBY3kQ6g8qeD4TUjtA==", "cpu": [ "arm64" ], @@ -2458,9 +2458,9 @@ } }, "node_modules/@tailwindcss/oxide-linux-x64-gnu": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-linux-x64-gnu/-/oxide-linux-x64-gnu-4.3.1.tgz", - "integrity": "sha512-Ymi8O8T15HYQdOUWUtTI6ldN0neHP85FC+Qz32xTcZ7iJXtem/x8ITev0o1e9e5rkqj4lONZfTRLvkmin1+tKg==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-linux-x64-gnu/-/oxide-linux-x64-gnu-4.3.2.tgz", + "integrity": "sha512-kqCZpSKOBEJO4mz7OqWoofBZeXTAwaVGPj0ErAj7CojmhKpWVWVOnrt9dE8odoIraZq4oj3ausM37kXi+Tow8w==", "cpu": [ "x64" ], @@ -2475,9 +2475,9 @@ } }, "node_modules/@tailwindcss/oxide-linux-x64-musl": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-linux-x64-musl/-/oxide-linux-x64-musl-4.3.1.tgz", - "integrity": "sha512-M+P/91qJ6uILLw4k2G93GMDRAXj61SMvFQYt39AqvUqYgExXpLL5aepfns7sj4HiAQeolirQF9E0lzRvdf4zPQ==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-linux-x64-musl/-/oxide-linux-x64-musl-4.3.2.tgz", + "integrity": "sha512-cixpqbh2toJDmkuCRI68nXA8ZxNmdK9Y+9v5h3MC3ZQKy/0BO8AWzlkWyRM7JAFSGBlfig4YVTPsK6MVgqz1uw==", "cpu": [ "x64" ], @@ -2492,9 +2492,9 @@ } }, "node_modules/@tailwindcss/oxide-wasm32-wasi": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-wasm32-wasi/-/oxide-wasm32-wasi-4.3.1.tgz", - "integrity": "sha512-zsM8uOeqvVGHsAXsJxsT28ttosFahLJKCLOTUBqRAtKnVgGSRitds9T432QiT8b77Yga7JIBkulIRRlJPtYhRA==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-wasm32-wasi/-/oxide-wasm32-wasi-4.3.2.tgz", + "integrity": "sha512-4ec2Z/LOmRsAgU23CS4xeJfcJlmRg94A/XrbGRCF1gyU/zdDfRLYDVsS+ynSZCmGNxQ1jQriQOKMQeQxBA3Isw==", "bundleDependencies": [ "@napi-rs/wasm-runtime", "@emnapi/core", @@ -2510,9 +2510,9 @@ "license": "MIT", "optional": true, "dependencies": { - "@emnapi/core": "^1.10.0", - "@emnapi/runtime": "^1.10.0", - "@emnapi/wasi-threads": "^1.2.1", + "@emnapi/core": "^1.11.1", + "@emnapi/runtime": "^1.11.1", + "@emnapi/wasi-threads": "^1.2.2", "@napi-rs/wasm-runtime": "^1.1.4", "@tybys/wasm-util": "^0.10.2", "tslib": "^2.8.1" @@ -2522,18 +2522,18 @@ } }, "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/core": { - "version": "1.10.0", + "version": "1.11.1", "dev": true, "inBundle": true, "license": "MIT", "optional": true, "dependencies": { - "@emnapi/wasi-threads": "1.2.1", + "@emnapi/wasi-threads": "1.2.2", "tslib": "^2.4.0" } }, "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/runtime": { - "version": "1.10.0", + "version": "1.11.1", "dev": true, "inBundle": true, "license": "MIT", @@ -2543,7 +2543,7 @@ } }, "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/wasi-threads": { - "version": "1.2.1", + "version": "1.2.2", "dev": true, "inBundle": true, "license": "MIT", @@ -2588,9 +2588,9 @@ "optional": true }, "node_modules/@tailwindcss/oxide-win32-arm64-msvc": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-win32-arm64-msvc/-/oxide-win32-arm64-msvc-4.3.1.tgz", - "integrity": "sha512-aiNvSq9BsVk8V513lDKlrCFAgf8qBMPZTpgEhInL+NwQqs97mYmupVMrPrgBBSL8Pv/0zXu9MrMF9rMun1ZeNg==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-win32-arm64-msvc/-/oxide-win32-arm64-msvc-4.3.2.tgz", + "integrity": "sha512-Zyr/M0+XcYZu3bZrUytc7TXvrk0ftWfl8gN2MwekNDzhqhKRUucMPSeOzM0o0wH5AWOU49BsKRrfKxI2atCPMQ==", "cpu": [ "arm64" ], @@ -2605,9 +2605,9 @@ } }, "node_modules/@tailwindcss/oxide-win32-x64-msvc": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-win32-x64-msvc/-/oxide-win32-x64-msvc-4.3.1.tgz", - "integrity": "sha512-xDEyu1rg290472FEGaKHnzyDyh5QH+AlWvsU5hMoMtPpzmKlRI0jaYKCgSHDYtaQWZOYbMaduSyCwFwY4n1HmA==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-win32-x64-msvc/-/oxide-win32-x64-msvc-4.3.2.tgz", + "integrity": "sha512-QI9BO7KlNZsp2GuO0jwAAj5jCDABOKXRkCk2XuKTSaNEFSdfzqswYVTtCHBNKHLsqyjFyFkqlDiwkNbTYSssMQ==", "cpu": [ "x64" ], @@ -2622,15 +2622,15 @@ } }, "node_modules/@tailwindcss/vite": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/@tailwindcss/vite/-/vite-4.3.1.tgz", - "integrity": "sha512-hItDHuIIlEV61R+faXu66s1K36aTurO/Qw0e45Vskz57gXl9pWOT6eg3zmcEui6CZXddbN7zd41bwmvag4JGwQ==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/@tailwindcss/vite/-/vite-4.3.2.tgz", + "integrity": "sha512-eHpMeX4JXfVNJDEcsouTeCBubJBTcTLigeaw/NTUW6PB5ATKKXdyonnXgTBX2VuRbjz1hjfz6C5XAhr52ImQXA==", "dev": true, "license": "MIT", "dependencies": { - "@tailwindcss/node": "4.3.1", - "@tailwindcss/oxide": "4.3.1", - "tailwindcss": "4.3.1" + "@tailwindcss/node": "4.3.2", + "@tailwindcss/oxide": "4.3.2", + "tailwindcss": "4.3.2" }, "peerDependencies": { "vite": "^5.2.0 || ^6 || ^7 || ^8" @@ -2657,12 +2657,12 @@ } }, "node_modules/@tanstack/react-virtual": { - "version": "3.14.4", - "resolved": "https://registry.npmjs.org/@tanstack/react-virtual/-/react-virtual-3.14.4.tgz", - "integrity": "sha512-dZzAQP2uCDAd+9sAehqmx/DcU+B91Q4Gb0aDSM7t9bJvWDyGF9sapFNW5r1gNLsHs4wTb6ScZENJeYaHxJLiOw==", + "version": "3.14.5", + "resolved": "https://registry.npmjs.org/@tanstack/react-virtual/-/react-virtual-3.14.5.tgz", + "integrity": "sha512-4EKRXh7zBLkbKbFmG3AUVkircuHd+7OdT1pocJSepxtfBd3qnrJgJ5rtPkRYyo9fmyVb2+pI2xPy5oYvMLQy6A==", "license": "MIT", "dependencies": { - "@tanstack/virtual-core": "3.17.2" + "@tanstack/virtual-core": "3.17.3" }, "funding": { "type": "github", @@ -2687,9 +2687,9 @@ } }, "node_modules/@tanstack/virtual-core": { - "version": "3.17.2", - "resolved": "https://registry.npmjs.org/@tanstack/virtual-core/-/virtual-core-3.17.2.tgz", - "integrity": "sha512-w43MvWvmShpb6kIC9MOoLyUkLmRTLPjt61bHWs+X29hACSpX+n8DvgZ3qM7cUfflKlRRcHR9KVJE6TmcqnQvcA==", + "version": "3.17.3", + "resolved": "https://registry.npmjs.org/@tanstack/virtual-core/-/virtual-core-3.17.3.tgz", + "integrity": "sha512-8Np/TFELpI0ySuJoVmjvOrQYXH/8sTX0Biv9szhFhY39xOdAAY+smrMxjxOum/ux3eM8MUJQsEJ0/R0UpvC8dw==", "license": "MIT", "funding": { "type": "github", @@ -2931,17 +2931,17 @@ "license": "MIT" }, "node_modules/@typescript-eslint/eslint-plugin": { - "version": "8.62.0", - "resolved": "https://registry.npmjs.org/@typescript-eslint/eslint-plugin/-/eslint-plugin-8.62.0.tgz", - "integrity": "sha512-o+mpz7EYiMzXoySXiKmzlabIvTVqUuK5yLrAedRPRDA0IpPFMUV1IXt6OqljIxX/kumN6EjUYp41Hqelh6p/Dw==", + "version": "8.62.1", + "resolved": "https://registry.npmjs.org/@typescript-eslint/eslint-plugin/-/eslint-plugin-8.62.1.tgz", + "integrity": "sha512-4EQM77WgVNxj7OkL/5b/D/xZsw00G577+UriYTC7JF5opcF3T2AuoeY7ueLaZgSVjSgCS6yOAJB5bRGLPSJUzA==", "dev": true, "license": "MIT", "dependencies": { "@eslint-community/regexpp": "^4.12.2", - "@typescript-eslint/scope-manager": "8.62.0", - "@typescript-eslint/type-utils": "8.62.0", - "@typescript-eslint/utils": "8.62.0", - "@typescript-eslint/visitor-keys": "8.62.0", + "@typescript-eslint/scope-manager": "8.62.1", + "@typescript-eslint/type-utils": "8.62.1", + "@typescript-eslint/utils": "8.62.1", + "@typescript-eslint/visitor-keys": "8.62.1", "ignore": "^7.0.5", "natural-compare": "^1.4.0", "ts-api-utils": "^2.5.0" @@ -2954,7 +2954,7 @@ "url": "https://opencollective.com/typescript-eslint" }, "peerDependencies": { - "@typescript-eslint/parser": "^8.62.0", + "@typescript-eslint/parser": "^8.62.1", "eslint": "^8.57.0 || ^9.0.0 || ^10.0.0", "typescript": ">=4.8.4 <6.1.0" } @@ -2970,16 +2970,16 @@ } }, "node_modules/@typescript-eslint/parser": { - "version": "8.62.0", - "resolved": "https://registry.npmjs.org/@typescript-eslint/parser/-/parser-8.62.0.tgz", - "integrity": "sha512-dzHeT2gySzZtLDsuqxU9AkYgIsQoHAHtRBpOqM+Ofzx1Bwrd2RcCjQJ+6iQbsHOIR6NS33bF2W1k3blN1zLDrA==", + "version": "8.62.1", + "resolved": "https://registry.npmjs.org/@typescript-eslint/parser/-/parser-8.62.1.tgz", + "integrity": "sha512-sPhE4iHuJDSvoAiec+Ro8JyXw8f0ql13HFR82P99nCm9GwTEKG0KYLvDe6REk8BCXuit6vJAv/Yxg5ABaNS2rA==", "dev": true, "license": "MIT", "dependencies": { - "@typescript-eslint/scope-manager": "8.62.0", - "@typescript-eslint/types": "8.62.0", - "@typescript-eslint/typescript-estree": "8.62.0", - "@typescript-eslint/visitor-keys": "8.62.0", + "@typescript-eslint/scope-manager": "8.62.1", + "@typescript-eslint/types": "8.62.1", + "@typescript-eslint/typescript-estree": "8.62.1", + "@typescript-eslint/visitor-keys": "8.62.1", "debug": "^4.4.3" }, "engines": { @@ -2995,14 +2995,14 @@ } }, "node_modules/@typescript-eslint/project-service": { - "version": "8.62.0", - "resolved": "https://registry.npmjs.org/@typescript-eslint/project-service/-/project-service-8.62.0.tgz", - "integrity": "sha512-wexnCqiTg7BOGtbLDftYpRWlmLq4xfoMd7BKFR6Y75sZS3QmRKLdN3yWLhmIYgqMmP/OXWpj3H8odkb5nGURCQ==", + "version": "8.62.1", + "resolved": "https://registry.npmjs.org/@typescript-eslint/project-service/-/project-service-8.62.1.tgz", + "integrity": "sha512-yQ3RgY5RkSBpsNS1Bx/JQEcA24FOSdfGktoyprAr5u18390UQdtVcfnEv4nIrIshNnavlVyZBKxQwT1fIAE6cg==", "dev": true, "license": "MIT", "dependencies": { - "@typescript-eslint/tsconfig-utils": "^8.62.0", - "@typescript-eslint/types": "^8.62.0", + "@typescript-eslint/tsconfig-utils": "^8.62.1", + "@typescript-eslint/types": "^8.62.1", "debug": "^4.4.3" }, "engines": { @@ -3017,14 +3017,14 @@ } }, "node_modules/@typescript-eslint/scope-manager": { - "version": "8.62.0", - "resolved": "https://registry.npmjs.org/@typescript-eslint/scope-manager/-/scope-manager-8.62.0.tgz", - "integrity": "sha512-1lX38kNxXIRb8mEc3lbq5mdHq1Pf2+U0nFU65KfT18mtPxxl0fvjuEE92mHuXPuCtElJhOrddOpyMlM3Z0umEA==", + "version": "8.62.1", + "resolved": "https://registry.npmjs.org/@typescript-eslint/scope-manager/-/scope-manager-8.62.1.tgz", + "integrity": "sha512-r4d249KbQ1SFdpeStvob8Ih6aPPIzfqllPVOtvhve6ZcpuVcYo5/7zUWckKpHE7StASX4kTKZTLf0WQm/wPkcg==", "dev": true, "license": "MIT", "dependencies": { - "@typescript-eslint/types": "8.62.0", - "@typescript-eslint/visitor-keys": "8.62.0" + "@typescript-eslint/types": "8.62.1", + "@typescript-eslint/visitor-keys": "8.62.1" }, "engines": { "node": "^18.18.0 || ^20.9.0 || >=21.1.0" @@ -3035,9 +3035,9 @@ } }, "node_modules/@typescript-eslint/tsconfig-utils": { - "version": "8.62.0", - "resolved": "https://registry.npmjs.org/@typescript-eslint/tsconfig-utils/-/tsconfig-utils-8.62.0.tgz", - "integrity": "sha512-y2GAdB6ykaXUvuspbYnizQc4oDDz0Tz/Yc7iWrXf9mx8vm/L/0vLHCe0tS2boG96Zy+DivnVDQ9ZUEWoHqqx1g==", + "version": "8.62.1", + "resolved": "https://registry.npmjs.org/@typescript-eslint/tsconfig-utils/-/tsconfig-utils-8.62.1.tgz", + "integrity": "sha512-xadytJqX9vJVQ2fdQjkcIVigwaOJNWkpjdLt6cEQ+xPnrI1fkp+/jZE/I97k9KUjqtpd25i0HeyZf3T6dutv2g==", "dev": true, "license": "MIT", "engines": { @@ -3052,15 +3052,15 @@ } }, "node_modules/@typescript-eslint/type-utils": { - "version": "8.62.0", - "resolved": "https://registry.npmjs.org/@typescript-eslint/type-utils/-/type-utils-8.62.0.tgz", - "integrity": "sha512-+g5O3j0w2ldzC86Pv6fvbO/xhAonbJFIdf/MKQ1d30gndlsVzUOE83ldfSE15Qrl9fhFjK6AovHs5Wpp6vx86w==", + "version": "8.62.1", + "resolved": "https://registry.npmjs.org/@typescript-eslint/type-utils/-/type-utils-8.62.1.tgz", + "integrity": "sha512-aXM5xlqXiTxPibXB93cLAURfT3rlizf7uMXISCXy66Isr/9hISJx3yDsKl0L7lKa51b8JpFuNKby0/O0pEm9jg==", "dev": true, "license": "MIT", "dependencies": { - "@typescript-eslint/types": "8.62.0", - "@typescript-eslint/typescript-estree": "8.62.0", - "@typescript-eslint/utils": "8.62.0", + "@typescript-eslint/types": "8.62.1", + "@typescript-eslint/typescript-estree": "8.62.1", + "@typescript-eslint/utils": "8.62.1", "debug": "^4.4.3", "ts-api-utils": "^2.5.0" }, @@ -3077,9 +3077,9 @@ } }, "node_modules/@typescript-eslint/types": { - "version": "8.62.0", - "resolved": "https://registry.npmjs.org/@typescript-eslint/types/-/types-8.62.0.tgz", - "integrity": "sha512-KvAclkktORPvM54TgLgA4z9HIV1M8zOgw9ZVNXl9f/8dLYfXYX1wkMXP7qmabpijQRV5bHJLOmoyGQbLMaUYeg==", + "version": "8.62.1", + "resolved": "https://registry.npmjs.org/@typescript-eslint/types/-/types-8.62.1.tgz", + "integrity": "sha512-ooCzJFaf+Hg+uG6fA3NRFGuFjlfNlDhBthbv4ZPU/0elCAFUfnyXUvf/WOpHz/jYwSmvU2GkR2LtyUfy1AxZ1Q==", "dev": true, "license": "MIT", "engines": { @@ -3091,16 +3091,16 @@ } }, "node_modules/@typescript-eslint/typescript-estree": { - "version": "8.62.0", - "resolved": "https://registry.npmjs.org/@typescript-eslint/typescript-estree/-/typescript-estree-8.62.0.tgz", - "integrity": "sha512-+hVbNxtW64pIcZWDPGbyaKF7vp2IBTVY5ma1blwwksrjdsbdqqEKvJWMGbBofei4F6Dovx1M0RJgoFeNu2279A==", + "version": "8.62.1", + "resolved": "https://registry.npmjs.org/@typescript-eslint/typescript-estree/-/typescript-estree-8.62.1.tgz", + "integrity": "sha512-xMcW9oP9u7fAMXYs9A65CVmtLQe2r//oXINHfi8HV+oiqhih17sbLdhXr4540YWlgpDKQdY854OL5ZrdCiQsAA==", "dev": true, "license": "MIT", "dependencies": { - "@typescript-eslint/project-service": "8.62.0", - "@typescript-eslint/tsconfig-utils": "8.62.0", - "@typescript-eslint/types": "8.62.0", - "@typescript-eslint/visitor-keys": "8.62.0", + "@typescript-eslint/project-service": "8.62.1", + "@typescript-eslint/tsconfig-utils": "8.62.1", + "@typescript-eslint/types": "8.62.1", + "@typescript-eslint/visitor-keys": "8.62.1", "debug": "^4.4.3", "minimatch": "^10.2.2", "semver": "^7.7.3", @@ -3132,16 +3132,16 @@ } }, "node_modules/@typescript-eslint/utils": { - "version": "8.62.0", - "resolved": "https://registry.npmjs.org/@typescript-eslint/utils/-/utils-8.62.0.tgz", - "integrity": "sha512-82r66fi9zYwZ+mTq3vKgwjbZ1PVk/DJzrXFLpG6RnBbdvH8TEGVHIs9H4d2drhkOzf0syZuD/OZvvlu6GDbP4g==", + "version": "8.62.1", + "resolved": "https://registry.npmjs.org/@typescript-eslint/utils/-/utils-8.62.1.tgz", + "integrity": "sha512-sHtbPfuKNZCG+ih8SyjjucqRntSVmp8XgL5u6o9mAhiSn8ds5o/M/XdM0abweme2Tln3szOstOrZ9OXitvPh0g==", "dev": true, "license": "MIT", "dependencies": { "@eslint-community/eslint-utils": "^4.9.1", - "@typescript-eslint/scope-manager": "8.62.0", - "@typescript-eslint/types": "8.62.0", - "@typescript-eslint/typescript-estree": "8.62.0" + "@typescript-eslint/scope-manager": "8.62.1", + "@typescript-eslint/types": "8.62.1", + "@typescript-eslint/typescript-estree": "8.62.1" }, "engines": { "node": "^18.18.0 || ^20.9.0 || >=21.1.0" @@ -3156,13 +3156,13 @@ } }, "node_modules/@typescript-eslint/visitor-keys": { - "version": "8.62.0", - "resolved": "https://registry.npmjs.org/@typescript-eslint/visitor-keys/-/visitor-keys-8.62.0.tgz", - "integrity": "sha512-CY3uyFSRbcQv3nnSv8S0+lDftMVz6P963PoRlxrV7ew/Md564g9ut60PYzdLM5qW4jFn93GBF+Soi90ISAN+GQ==", + "version": "8.62.1", + "resolved": "https://registry.npmjs.org/@typescript-eslint/visitor-keys/-/visitor-keys-8.62.1.tgz", + "integrity": "sha512-4g3BLxfdTMy8iZG0MaBkadnlRrCJ74cQiFbyEVMrkwIoqdyaXXQM22cotDvrl4x28wgIZ9rEJRoM+mmhSJpJ1g==", "dev": true, "license": "MIT", "dependencies": { - "@typescript-eslint/types": "8.62.0", + "@typescript-eslint/types": "8.62.1", "eslint-visitor-keys": "^5.0.0" }, "engines": { @@ -3403,9 +3403,9 @@ } }, "node_modules/baseline-browser-mapping": { - "version": "2.10.40", - "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.10.40.tgz", - "integrity": "sha512-BSSLZ9/Cjjv7Gtj5B68ZzXcXUg8iOf3fme+FCuh8rC/Go+Kmh8cox7M3A8dolou16s64QjLPOSdngh7GxXvkSw==", + "version": "2.10.42", + "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.10.42.tgz", + "integrity": "sha512-c/jurFrDLyui7o1J86yLkRu4LMsTYcBohveus7/I2Hzdn9KIP2bdJPTue/lR1KH46enoPbD77GKeSYNdyPoD3Q==", "dev": true, "license": "Apache-2.0", "bin": { @@ -3416,9 +3416,9 @@ } }, "node_modules/brace-expansion": { - "version": "5.0.6", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.6.tgz", - "integrity": "sha512-kLpxurY4Z4r9sgMsyG0Z9uzsBlgiU/EFKhj/h91/8yHu0edo7XuixOIH3VcJ8kkxs6/jPzoI6U9Vj3WqbMQ94g==", + "version": "5.0.7", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.7.tgz", + "integrity": "sha512-7oFy703dxfY3/NLxC1fh2SUCQ0H9rmAY+5EpDVfXjUTTs+HEwR2nYaqLv+GWcTsumwxPfiz6CzCNkwXwBUwqCA==", "dev": true, "license": "MIT", "dependencies": { @@ -3479,9 +3479,9 @@ } }, "node_modules/caniuse-lite": { - "version": "1.0.30001799", - "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001799.tgz", - "integrity": "sha512-hG1bReV+OUU+MOqK4t/ZWI0tZOyz3rqS9XuhOUz1cIcbwBKjOyJEJuw9ER5JuNyqxNk8u/JUVbGibBOL1yrjFw==", + "version": "1.0.30001800", + "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001800.tgz", + "integrity": "sha512-MMHtuAz9Ys840zAY5F4k6fV5GaivZ9sPk+nz0mY+GYVzRBnYkN0mpqkSR92oWRQ19yQWo4HvBV/FnC16AJX8MA==", "dev": true, "funding": [ { @@ -3778,9 +3778,9 @@ "peer": true }, "node_modules/electron-to-chromium": { - "version": "1.5.380", - "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.380.tgz", - "integrity": "sha512-W6d5AbuEoRayO447cqrg6lKJIlscgRnnxOZl/08kfV71BQDoEBC7Wwis68z87LjyK6f4kWyTaubuDbhHKrZkbA==", + "version": "1.5.387", + "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.387.tgz", + "integrity": "sha512-TaxwufTFDufvPEoXdhwVrA3UdFWBeWGkYoJ1K8ldF1xe6gKfth6iRNS5lTQ5JPNOHdGQm8PT1QYKUqFLCiUefQ==", "dev": true, "license": "ISC" }, @@ -5066,34 +5066,34 @@ } }, "node_modules/oxc-resolver": { - "version": "11.21.3", - "resolved": "https://registry.npmjs.org/oxc-resolver/-/oxc-resolver-11.21.3.tgz", - "integrity": "sha512-2Mx3fKQz7+xgrBONjsxOgCGtMHOn38/HxMzW1I5efwXB5a4lRN0Vp40gYUJFBWJslcrvwoofTrqoTnLbwTd3pA==", + "version": "11.23.0", + "resolved": "https://registry.npmjs.org/oxc-resolver/-/oxc-resolver-11.23.0.tgz", + "integrity": "sha512-f0+l598CJMOLnYPXsXxttJALH0ljtivdRMKtvHhxRuWa5FYmw5+qODARl8oYjMC/brpzKcrpdORsOBrTqhBZ9A==", "dev": true, "license": "MIT", "funding": { "url": "https://github.com/sponsors/Boshen" }, "optionalDependencies": { - "@oxc-resolver/binding-android-arm-eabi": "11.21.3", - "@oxc-resolver/binding-android-arm64": "11.21.3", - "@oxc-resolver/binding-darwin-arm64": "11.21.3", - "@oxc-resolver/binding-darwin-x64": "11.21.3", - "@oxc-resolver/binding-freebsd-x64": "11.21.3", - "@oxc-resolver/binding-linux-arm-gnueabihf": "11.21.3", - "@oxc-resolver/binding-linux-arm-musleabihf": "11.21.3", - "@oxc-resolver/binding-linux-arm64-gnu": "11.21.3", - "@oxc-resolver/binding-linux-arm64-musl": "11.21.3", - "@oxc-resolver/binding-linux-ppc64-gnu": "11.21.3", - "@oxc-resolver/binding-linux-riscv64-gnu": "11.21.3", - "@oxc-resolver/binding-linux-riscv64-musl": "11.21.3", - "@oxc-resolver/binding-linux-s390x-gnu": "11.21.3", - "@oxc-resolver/binding-linux-x64-gnu": "11.21.3", - "@oxc-resolver/binding-linux-x64-musl": "11.21.3", - "@oxc-resolver/binding-openharmony-arm64": "11.21.3", - "@oxc-resolver/binding-wasm32-wasi": "11.21.3", - "@oxc-resolver/binding-win32-arm64-msvc": "11.21.3", - "@oxc-resolver/binding-win32-x64-msvc": "11.21.3" + "@oxc-resolver/binding-android-arm-eabi": "11.23.0", + "@oxc-resolver/binding-android-arm64": "11.23.0", + "@oxc-resolver/binding-darwin-arm64": "11.23.0", + "@oxc-resolver/binding-darwin-x64": "11.23.0", + "@oxc-resolver/binding-freebsd-x64": "11.23.0", + "@oxc-resolver/binding-linux-arm-gnueabihf": "11.23.0", + "@oxc-resolver/binding-linux-arm-musleabihf": "11.23.0", + "@oxc-resolver/binding-linux-arm64-gnu": "11.23.0", + "@oxc-resolver/binding-linux-arm64-musl": "11.23.0", + "@oxc-resolver/binding-linux-ppc64-gnu": "11.23.0", + "@oxc-resolver/binding-linux-riscv64-gnu": "11.23.0", + "@oxc-resolver/binding-linux-riscv64-musl": "11.23.0", + "@oxc-resolver/binding-linux-s390x-gnu": "11.23.0", + "@oxc-resolver/binding-linux-x64-gnu": "11.23.0", + "@oxc-resolver/binding-linux-x64-musl": "11.23.0", + "@oxc-resolver/binding-openharmony-arm64": "11.23.0", + "@oxc-resolver/binding-wasm32-wasi": "11.23.0", + "@oxc-resolver/binding-win32-arm64-msvc": "11.23.0", + "@oxc-resolver/binding-win32-x64-msvc": "11.23.0" } }, "node_modules/p-limit": { @@ -5200,9 +5200,9 @@ "license": "ISC" }, "node_modules/picomatch": { - "version": "4.0.4", - "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.4.tgz", - "integrity": "sha512-QP88BAKvMam/3NxH6vj2o21R6MjxZUAd6nlwAS/pnGvN9IVLocLHxGYIzFhg6fUQ+5th6P4dv4eW9jX3DSIj7A==", + "version": "4.0.5", + "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.5.tgz", + "integrity": "sha512-RvwwcruNjI1ncT5xRakeyS9Lf8lcItv34KD+aif+VH9kduAyfYBipGh12274xtenIPZ119/R9BdTBa8gAwSh0A==", "dev": true, "license": "MIT", "engines": { @@ -5252,9 +5252,9 @@ } }, "node_modules/prettier": { - "version": "3.9.1", - "resolved": "https://registry.npmjs.org/prettier/-/prettier-3.9.1.tgz", - "integrity": "sha512-ppiDo2CSwexck1eyZUwJHg/N3nf1+6IRCv7W/VJ5vaLnVCmB7+3CdRfMwoCHBBX6xTrREDTksZ4OZl5SSf4zXA==", + "version": "3.9.4", + "resolved": "https://registry.npmjs.org/prettier/-/prettier-3.9.4.tgz", + "integrity": "sha512-yWG/o/4oJfo036EKAfK6ACAoDOfHeRHx4tuxkfBZiauURiaSmYwlpOr5LQqKtIkRD2z1PLteme2WoxEnj4tHTg==", "dev": true, "license": "MIT", "bin": { @@ -5435,13 +5435,13 @@ } }, "node_modules/rolldown": { - "version": "1.1.3", - "resolved": "https://registry.npmjs.org/rolldown/-/rolldown-1.1.3.tgz", - "integrity": "sha512-1F1eEtUBtFvcGm1HQ9TiUIUHPQG7mSAODrhIzjxoUEFuo8OcbrGLiVLkevNgj84TE4lnHvnumwFjhJO5Eu135g==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/rolldown/-/rolldown-1.1.4.tgz", + "integrity": "sha512-IjZYiLxZwpnhwhdBH2ugdTGVSdhCQUmLxLoqyjiL0JxYjyRst+5a0P3xfrTxJ5F638j4Mvvw5FAX5XE6eHpXbA==", "dev": true, "license": "MIT", "dependencies": { - "@oxc-project/types": "=0.137.0", + "@oxc-project/types": "=0.138.0", "@rolldown/pluginutils": "^1.0.0" }, "bin": { @@ -5451,27 +5451,27 @@ "node": "^20.19.0 || >=22.12.0" }, "optionalDependencies": { - "@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" + "@rolldown/binding-android-arm64": "1.1.4", + "@rolldown/binding-darwin-arm64": "1.1.4", + "@rolldown/binding-darwin-x64": "1.1.4", + "@rolldown/binding-freebsd-x64": "1.1.4", + "@rolldown/binding-linux-arm-gnueabihf": "1.1.4", + "@rolldown/binding-linux-arm64-gnu": "1.1.4", + "@rolldown/binding-linux-arm64-musl": "1.1.4", + "@rolldown/binding-linux-ppc64-gnu": "1.1.4", + "@rolldown/binding-linux-s390x-gnu": "1.1.4", + "@rolldown/binding-linux-x64-gnu": "1.1.4", + "@rolldown/binding-linux-x64-musl": "1.1.4", + "@rolldown/binding-openharmony-arm64": "1.1.4", + "@rolldown/binding-wasm32-wasi": "1.1.4", + "@rolldown/binding-win32-arm64-msvc": "1.1.4", + "@rolldown/binding-win32-x64-msvc": "1.1.4" } }, "node_modules/rolldown/node_modules/@oxc-project/types": { - "version": "0.137.0", - "resolved": "https://registry.npmjs.org/@oxc-project/types/-/types-0.137.0.tgz", - "integrity": "sha512-WT+Gb24i8hmvo85AIv2oEYouEXkRlKAlT9WaCa3TfLgNCN+GhrJOGZuIlMouAh38Qe4QOx26eUOVsq70qXrywA==", + "version": "0.138.0", + "resolved": "https://registry.npmjs.org/@oxc-project/types/-/types-0.138.0.tgz", + "integrity": "sha512-1a7ZKmrRTCoN1XMZ4L0PyyqrMnrNlLyPuOkdSX2MZg7IiIGRUyurNhAm73ptDOraoBcIordsIGKNPKUzy3ZmfA==", "dev": true, "license": "MIT", "funding": { @@ -5672,9 +5672,9 @@ } }, "node_modules/tailwindcss": { - "version": "4.3.1", - "resolved": "https://registry.npmjs.org/tailwindcss/-/tailwindcss-4.3.1.tgz", - "integrity": "sha512-hk+TB1m+K8CYNrP6rjQaq/Y+4Zylwpa87mLYBKCunwnnQ9p+fHb7kmSfGqyEJoxF/O6CDyABWVFEafNSYKll+Q==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/tailwindcss/-/tailwindcss-4.3.2.tgz", + "integrity": "sha512-WtctNNSH8A9jlMIqxzuYumOHU5uGZyRv0Q5svQl+oEPy5w84YpBxdb7MdqyiSPQge5jTJ6zFQLq0PFygdccSBA==", "dev": true, "license": "MIT" }, @@ -5808,16 +5808,16 @@ } }, "node_modules/typescript-eslint": { - "version": "8.62.0", - "resolved": "https://registry.npmjs.org/typescript-eslint/-/typescript-eslint-8.62.0.tgz", - "integrity": "sha512-8QxXi+ZACKX0kaqO4gY8kn0RSD9gFfaHDWwjqtEN48aWCBkX4MJaufWN+c3BzlrXLOxfywDL8CaoqUwcRq4j4Q==", + "version": "8.62.1", + "resolved": "https://registry.npmjs.org/typescript-eslint/-/typescript-eslint-8.62.1.tgz", + "integrity": "sha512-vymnnM5g0AKQDSAyfP12nMIBvgwgA42syg74kkuZ4x1VuTzwQKwc5h9rGxeShCjny5o+zWAb6OEoz7XLgrIkIw==", "dev": true, "license": "MIT", "dependencies": { - "@typescript-eslint/eslint-plugin": "8.62.0", - "@typescript-eslint/parser": "8.62.0", - "@typescript-eslint/typescript-estree": "8.62.0", - "@typescript-eslint/utils": "8.62.0" + "@typescript-eslint/eslint-plugin": "8.62.1", + "@typescript-eslint/parser": "8.62.1", + "@typescript-eslint/typescript-estree": "8.62.1", + "@typescript-eslint/utils": "8.62.1" }, "engines": { "node": "^18.18.0 || ^20.9.0 || >=21.1.0" @@ -5914,16 +5914,16 @@ } }, "node_modules/vite": { - "version": "8.1.0", - "resolved": "https://registry.npmjs.org/vite/-/vite-8.1.0.tgz", - "integrity": "sha512-BuJcQK/56NQTWDGn4ABea3q4SSBdNPWwNZKTkkUpcMPnLoquSYH8llRtSUIgoL1KSCpHt5eghLShn50mH36y7Q==", + "version": "8.1.3", + "resolved": "https://registry.npmjs.org/vite/-/vite-8.1.3.tgz", + "integrity": "sha512-Ds+gBRbj0lwRO2Y5hwnUBdxSwlAve9LeRyU4sNnAr0ewW0gWF0n5bgXgUzbgZ49MV9BVUAQUFYVcDUcilUExMA==", "dev": true, "license": "MIT", "dependencies": { "lightningcss": "^1.32.0", "picomatch": "^4.0.4", - "postcss": "^8.5.15", - "rolldown": "~1.1.2", + "postcss": "^8.5.16", + "rolldown": "~1.1.3", "tinyglobby": "^0.2.17" }, "bin": { diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index adaca75632c..2511893d09e 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -16,12 +16,16 @@ version = { workspace = true } [package.metadata.docs.rs] all-features = true +[package.metadata.cargo-shear] +ignored = ["vortex-tensor"] + [lints] workspace = true [dependencies] vortex-alp = { workspace = true } vortex-array = { workspace = true } +vortex-arrow = { workspace = true } vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } vortex-bytebool = { workspace = true } @@ -39,7 +43,7 @@ vortex-mask = { workspace = true } vortex-metrics = { workspace = true } vortex-pco = { workspace = true } vortex-proto = { workspace = true, default-features = true } -vortex-runend = { workspace = true, features = ["arrow"] } +vortex-runend = { workspace = true } vortex-scan = { workspace = true } vortex-sequence = { workspace = true } vortex-session = { workspace = true } @@ -54,26 +58,27 @@ anyhow = { workspace = true } arrow-array = { workspace = true } divan = { workspace = true } fastlanes = { workspace = true } -futures = { workspace = true } mimalloc = { workspace = true } parquet = { workspace = true } -paste = { workspace = true } rand = { workspace = true } -rand_distr = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true, features = ["full"] } tracing = { workspace = true } tracing-subscriber = { workspace = true } vortex = { path = ".", features = ["tokio"] } -vortex-tensor = { workspace = true } [features] default = ["files", "zstd"] files = ["dep:vortex-file"] memmap2 = ["vortex-buffer/memmap2"] -object_store = ["vortex-file/object_store", "vortex-io/object_store"] -tokio = ["vortex-file/tokio"] -zstd = ["dep:vortex-zstd", "vortex-file/zstd"] +object_store = ["vortex-file?/object_store", "vortex-io/object_store"] +tokio = [ + "vortex-error/tokio", + "vortex-file?/tokio", + "vortex-io/tokio", + "vortex-layout/tokio", +] +zstd = ["dep:vortex-zstd", "vortex-file?/zstd"] pretty = ["vortex-array/table-display"] serde = ["vortex-array/serde", "vortex-buffer/serde", "vortex-mask/serde"] # This feature enabled unstable encodings for which we don't guarantee stability. diff --git a/vortex/benches/common_encoding_tree_throughput.rs b/vortex/benches/common_encoding_tree_throughput.rs index 1d1dcf86956..65c2c1bc040 100644 --- a/vortex/benches/common_encoding_tree_throughput.rs +++ b/vortex/benches/common_encoding_tree_throughput.rs @@ -50,6 +50,7 @@ use vortex_session::VortexSession; static GLOBAL: MiMalloc = MiMalloc; fn main() { + LazyLock::force(&SESSION); divan::main(); } diff --git a/vortex/benches/single_encoding_throughput.rs b/vortex/benches/single_encoding_throughput.rs index b75f832a7a1..ca0a8926233 100644 --- a/vortex/benches/single_encoding_throughput.rs +++ b/vortex/benches/single_encoding_throughput.rs @@ -49,6 +49,7 @@ static GLOBAL: MiMalloc = MiMalloc; static SESSION: LazyLock = LazyLock::new(VortexSession::default); fn main() { + LazyLock::force(&SESSION); divan::main(); } diff --git a/vortex/examples/tracing_vortex.rs b/vortex/examples/tracing_vortex.rs index f1227e2e017..c3876f89e20 100644 --- a/vortex/examples/tracing_vortex.rs +++ b/vortex/examples/tracing_vortex.rs @@ -6,7 +6,7 @@ //! This example demonstrates a real-world use case: implementing a `tracing` subscriber //! that writes all log events and spans to Vortex files. //! -//! Run with: cargo run --example tracing_vortex --features tokio +//! Run with: cargo run --example tracing_vortex --features files,tokio #![expect( clippy::disallowed_types, diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 764c1f3d5fc..2703020906c 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -99,7 +99,6 @@ // vortex::compute is deprecated and will be ported over to expressions. pub use vortex_array::aggregate_fn; use vortex_array::aggregate_fn::session::AggregateFnSession; -use vortex_array::arrow::ArrowSession; pub use vortex_array::compute; use vortex_array::dtype::session::DTypeSession; // vortex::expr is in the process of having its dependencies inverted, and will eventually be @@ -127,6 +126,12 @@ pub mod array { // twice. } +/// Arrow interoperability: conversion between Vortex and Apache Arrow arrays, types, and +/// schemas. +pub mod arrow { + pub use vortex_arrow::*; +} + /// Aligned buffers and byte buffers used by arrays, layouts, IPC, and file IO. pub mod buffer { pub use vortex_buffer::*; @@ -297,9 +302,9 @@ impl VortexSessionDefault for VortexSession { .with::() .with::() .with::() - .with::() .with::() .with::(); + vortex_arrow::initialize(&session); // `MultiFileSession` holds a `moka` cache whose clock reads `std::time::Instant::now()` // when constructed. `Instant` is unsupported on `wasm32` and panics with "time not @@ -355,9 +360,9 @@ mod test { use arrow_array::RecordBatchReader; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use vortex::array::arrays::ChunkedArray; - use vortex::array::arrow::FromArrowArray; + use vortex::arrow::FromArrowArray; + use vortex::arrow::FromArrowType; use vortex::dtype::DType; - use vortex::dtype::arrow::FromArrowType; let reader = ParquetRecordBatchReaderBuilder::try_new(File::open( "../docs/_static/example.parquet", diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml index 865d10ac809..ebbc6987d46 100644 --- a/xtask/Cargo.toml +++ b/xtask/Cargo.toml @@ -17,7 +17,6 @@ version = { workspace = true } [dependencies] anyhow = { workspace = true } -cargo_metadata = { workspace = true } clap = { workspace = true, features = ["derive"] } prost-build = { workspace = true } xshell = { workspace = true }