diff --git a/.bumpversion.toml b/.bumpversion.toml index a8e8d03e983..620f381e275 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "9.0.0-beta.8" +current_version = "9.0.0-beta.21" parse = "(?P\\d+)\\.(?P\\d+)\\.(?P\\d+)(-(?P(beta|rc))\\.(?P\\d+))?" serialize = [ "{major}.{minor}.{patch}-{prerelease}.{prerelease_num}", diff --git a/.coderabbit.yaml b/.coderabbit.yaml new file mode 100644 index 00000000000..030a14c965e --- /dev/null +++ b/.coderabbit.yaml @@ -0,0 +1,16 @@ +# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json +language: en-US +early_access: false +reviews: + profile: assertive + poem: false + review_status: true + auto_review: + enabled: true + auto_incremental_review: true + ignore_title_keywords: + - WIP + - Draft + drafts: false + base_branches: + - main diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml new file mode 100644 index 00000000000..0e7363697b0 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -0,0 +1,79 @@ +name: Bug report +description: Report incorrect, unexpected, or crashing behavior. +labels: ["bug"] +body: + - type: markdown + attributes: + value: | + Thanks for taking the time to file a bug report! Please search + [existing issues](https://github.com/lance-format/lance/issues) first to + avoid filing a duplicate. + - type: textarea + id: description + attributes: + label: Description + description: A clear and concise description of what the bug is. + validations: + required: true + - type: textarea + id: reproduction + attributes: + label: Steps to reproduce + description: >- + A minimal, self-contained code snippet or sequence of steps that + reproduces the problem. The more we can copy-paste and run, the faster + we can fix it. + placeholder: | + 1. Write a dataset with `...` + 2. Scan with filter `...` + 3. Observe `...` + render: python + validations: + required: true + - type: textarea + id: expected + attributes: + label: Expected behavior + description: What did you expect to happen? + validations: + required: false + - type: input + id: lance-version + attributes: + label: Lance version + description: >- + Output of `python -c "import lance; print(lance.__version__)"`, or the + `lance`/`pylance` version from your `Cargo.toml` / `pyproject.toml`. + placeholder: "e.g. 0.40.0" + validations: + required: true + - type: dropdown + id: language + attributes: + label: Language binding + multiple: true + options: + - Python + - Rust + - Java + - Other / not sure + validations: + required: true + - type: input + id: environment + attributes: + label: Environment + description: OS, architecture, and storage backend (local, S3, GCS, Azure, ...). + placeholder: "e.g. Ubuntu 22.04, x86_64, S3" + validations: + required: false + - type: textarea + id: logs + attributes: + label: Logs / traceback + description: >- + Any relevant log output or stack trace. This will be automatically + formatted as code, so no need for backticks. + render: shell + validations: + required: false diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 00000000000..32389799971 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,7 @@ +blank_issues_enabled: true +contact_links: + - name: Question / usage help + url: https://discord.gg/lance + about: >- + For questions and general discussion, please ask in the Lance Discord + rather than opening an issue. diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 00000000000..41d6278fc02 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,11 @@ +--- +name: Feature request +about: Suggest a new capability or an improvement to an existing one. +labels: feature +--- + + diff --git a/.github/ISSUE_TEMPLATE/performance_issue.md b/.github/ISSUE_TEMPLATE/performance_issue.md new file mode 100644 index 00000000000..65fcd0be6f4 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/performance_issue.md @@ -0,0 +1,12 @@ +--- +name: Performance issue +about: Report slow operations, high memory use, or a performance regression. +labels: performance +--- + + diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 8a46f198419..a7b7d11c92b 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -6,6 +6,10 @@ updates: schedule: interval: "weekly" day: "wednesday" + groups: + cargo: + patterns: + - "*" - package-ecosystem: "cargo" directory: "/python" @@ -13,6 +17,10 @@ updates: schedule: interval: "weekly" day: "wednesday" + groups: + cargo: + patterns: + - "*" - package-ecosystem: "cargo" directory: "/java/lance-jni" @@ -20,6 +28,10 @@ updates: schedule: interval: "weekly" day: "wednesday" + groups: + cargo: + patterns: + - "*" - package-ecosystem: "uv" directory: "/python" @@ -27,4 +39,7 @@ updates: schedule: interval: "weekly" day: "wednesday" - + groups: + uv: + patterns: + - "*" diff --git a/.github/labeler-issues.yml b/.github/labeler-issues.yml new file mode 100644 index 00000000000..625b321cd6c --- /dev/null +++ b/.github/labeler-issues.yml @@ -0,0 +1,28 @@ +version: 1 +# Never remove labels a template or a human already applied. +appendOnly: true +# Content-based labels for issues, applied by srvaroa/labeler via +# .github/workflows/issue-labeler.yml. This covers issues that bypass the +# .github/ISSUE_TEMPLATE forms entirely — those filed via `gh issue create`, +# the REST API, or an agent, which never see the web-UI template chooser. +# +# The primary signal is a leading marker in the title (e.g. "bug:", "[feature]", +# "perf -"), which many people type by hand; a modest body keyword match is a +# secondary fallback. Rules are OR'd (one label per matching entry), so an +# unmatched issue gets no label and is left for human / LLM triage rather than +# guessed at. appendOnly means matches stack harmlessly with template labels. +labels: +# --- Primary signal: a leading bug/feature/perf marker in the title --- +- label: bug + title: "(?i)^\\s*[\\[(]?\\s*bug\\b" +- label: feature + title: "(?i)^\\s*[\\[(]?\\s*(feature|feat)\\b" +- label: performance + title: "(?i)^\\s*[\\[(]?\\s*(perf|performance)\\b" +# --- Secondary fallback: keywords anywhere in the body --- +- label: bug + body: "(?i)(panic|segfault|traceback|stack ?trace|crash|corrupt|incorrect result|wrong result)" +- label: performance + body: "(?i)(regression|latency|throughput|\\bOOM\\b|out of memory|memory usage|too slow)" +- label: feature + body: "(?i)(feature request|would be (nice|great|useful)|it would be nice|support for|add .+ support)" diff --git a/.github/workflows/issue-labeler.yml b/.github/workflows/issue-labeler.yml new file mode 100644 index 00000000000..36f5a9cdbca --- /dev/null +++ b/.github/workflows/issue-labeler.yml @@ -0,0 +1,29 @@ +name: Issue Labeler + +# Applies bug / feature / performance labels to issues based on their title and +# body. This complements the .github/ISSUE_TEMPLATE forms, which only apply +# labels for issues opened through the web UI; this catches issues opened via +# `gh issue create`, the REST API, or an agent, which bypass templates. +# +# Issues never originate from a fork, so unlike the PR labelers this uses a +# plain `issues` trigger (no pull_request_target) with a scoped GITHUB_TOKEN. +on: + issues: + types: [opened, edited] + +concurrency: + group: ${{ github.workflow }}-${{ github.event.issue.number }} + cancel-in-progress: true + +jobs: + label: + name: Apply issue labels + permissions: + issues: write + runs-on: ubuntu-latest + steps: + - uses: srvaroa/labeler@bf262763a8a8e191f5847873aecc0f29df84f957 # v1.14.0 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + config_path: .github/labeler-issues.yml diff --git a/.github/workflows/java-publish.yml b/.github/workflows/java-publish.yml index 2b22b60dc92..6351c5be402 100644 --- a/.github/workflows/java-publish.yml +++ b/.github/workflows/java-publish.yml @@ -112,7 +112,7 @@ jobs: export CXX=clang++ ldd --version - cargo build --release + cargo build --release --locked " - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4 with: @@ -135,7 +135,7 @@ jobs: run: brew install protobuf - name: Build native lib working-directory: java/lance-jni - run: cargo build --release + run: cargo build --release --locked - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4 with: name: liblance_jni_darwin_aarch64.zip diff --git a/.github/workflows/java.yml b/.github/workflows/java.yml index 83403988244..85b50e10836 100644 --- a/.github/workflows/java.yml +++ b/.github/workflows/java.yml @@ -47,14 +47,14 @@ jobs: run: cargo fmt --check - name: Rust Clippy working-directory: java/lance-jni - run: cargo clippy --all-targets -- -D warnings + run: cargo clippy --all-targets --locked -- -D warnings build-and-test-java: runs-on: ubuntu-24.04-4x timeout-minutes: 60 strategy: matrix: - java-version: [11, 17, 21] + java-version: [11, 17, 21, 25] name: Build and Test with Java ${{ matrix.java-version }} steps: - name: Checkout repository @@ -79,7 +79,13 @@ jobs: distribution: temurin java-version: ${{ matrix.java-version }} cache: "maven" + # Skip the format gate on Java 25: google-java-format 1.22.0 (pinned via + # spotless) reaches into com.sun.tools.javac internals and throws + # NoSuchMethodError on JDK 25's javac. No google-java-format release runs + # on JDK 25 without also dropping JDK 11/17 support, so we exclude 25 from + # the (JDK-independent, already-enforced-on-11/17/21) style check instead. - name: Running code style check with Java ${{ matrix.java-version }} + if: matrix.java-version != '25' working-directory: java run: | mvn spotless:check @@ -91,4 +97,7 @@ jobs: env: LANCE_INTEGRATION_TEST: "1" run: | - mvn install + # spotless is also bound into the build lifecycle, so on Java 25 the + # format gate must be skipped here too (-Dspotless.skip=true), not just + # in the standalone step above. + mvn install ${{ matrix.java-version == '25' && '-Dspotless.skip=true' || '' }} diff --git a/.github/workflows/license-header-check.yml b/.github/workflows/license-header-check.yml index 488ca65e585..2e6e26bbbeb 100644 --- a/.github/workflows/license-header-check.yml +++ b/.github/workflows/license-header-check.yml @@ -22,11 +22,25 @@ jobs: - name: Check out code uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 - name: Install license-header-checker + env: + # Pin to a specific release and download the asset directly from the + # release CDN. This avoids the flaky dynamic "latest" tag lookup in the + # upstream install.sh (it fetches the releases page and parses tag_name, + # which intermittently returns empty and fails the job). + LHC_VERSION: "1.5.0" run: | set -euo pipefail - curl -sfSL https://raw.githubusercontent.com/lluissm/license-header-checker/master/install.sh -o /tmp/install-lhc.sh - bash /tmp/install-lhc.sh - rm -f /tmp/install-lhc.sh + mkdir -p ./bin + base="https://github.com/lluissm/license-header-checker/releases/download/v${LHC_VERSION}" + asset="license-header-checker_${LHC_VERSION}_linux_amd64.tar.gz" + curl -sfSL "${base}/${asset}" -o "/tmp/${asset}" + curl -sfSL "${base}/checksums.txt" -o /tmp/lhc-checksums.txt + # Verify the download against the release checksums before extracting so a + # corrupted or tampered asset cannot be executed. This restores the + # integrity check the upstream install.sh performed via hash_sha256_verify. + (cd /tmp && sha256sum -c --ignore-missing lhc-checksums.txt) + tar -xzf "/tmp/${asset}" -C ./bin license-header-checker + rm -f "/tmp/${asset}" /tmp/lhc-checksums.txt - name: Check license headers (rust) run: ./bin/license-header-checker -a -v ./rust/license_header.txt rust rs && [[ -z `git status -s` ]] - name: Check license headers (python) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index e3c17671ce3..1bcadeb9199 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -253,6 +253,42 @@ jobs: - name: Check benchmarks run: cargo check --profile ci --benches + qemu-pre-haswell: + # Verifies that lance-linalg's runtime SIMD dispatch still works + # correctly when the binary is built with the lower x86-64-v2 baseline + # (the legacy build path documented in CONTRIBUTING.md). Emulates a + # Nehalem CPU under qemu-user; catches any accidental AVX2/FMA + # instructions that leak past the runtime dispatch. + # + # The published-wheel default baseline (`target-cpu=haswell`) is set in + # `.cargo/config.toml`; this job overrides RUSTFLAGS for one job to + # exercise the legacy path without affecting any other build. + name: pre-Haswell SIGILL check (qemu Nehalem) + runs-on: ubuntu-24.04 + timeout-minutes: 60 + env: + CC: clang + CXX: clang++ + RUSTFLAGS: "-C target-cpu=x86-64-v2" + CARGO_TARGET_X86_64_UNKNOWN_LINUX_GNU_RUNNER: "qemu-x86_64 -cpu Nehalem" + steps: + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 + - name: Setup rust toolchain + run: | + rustup toolchain install stable + rustup default stable + - uses: Swatinem/rust-cache@779680da715d629ac1d338a641029a2f4372abb5 # v2 + - name: Install dependencies + run: | + sudo apt update + sudo apt install -y protobuf-compiler libssl-dev qemu-user + - name: Build lance-linalg lib tests in release mode + run: | + cargo test --release -p lance-linalg --lib --no-run + - name: Run lance-linalg lib tests under qemu Nehalem + run: | + cargo test --release -p lance-linalg --lib + msrv: # Check the minimum supported Rust version name: MSRV Check - Rust v${{ matrix.msrv }} diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 09c956152fe..502f568908a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -15,3 +15,16 @@ repos: rev: v1.26.0 hooks: - id: typos + + # Catch a Cargo.toml change that wasn't reflected in the matching Cargo.lock. + # Runs fully offline (`--frozen` = `--locked` + `--offline`) and does not compile, + # so it stays fast. Checks all three lockfiles since a workspace dep change can + # touch the excluded python/ and java/ locks too. + - repo: local + hooks: + - id: cargo-lock-sync + name: Cargo.lock in sync with Cargo.toml (offline) + entry: bash -c 'for m in Cargo.toml python/Cargo.toml java/lance-jni/Cargo.toml; do cargo metadata --frozen --format-version 1 --manifest-path "$m" >/dev/null 2>&1 || { echo "Cargo.lock is out of date for $m. Refresh it (e.g. cargo check --manifest-path $m) and commit the updated lockfile."; exit 1; }; done' + language: system + files: '(^|/)Cargo\.(toml|lock)$' + pass_filenames: false diff --git a/AGENTS.md b/AGENTS.md index 30c0abea1a7..528e730e5d6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -100,6 +100,7 @@ AWS_DEFAULT_REGION=us-east-1 pytest --run-integration python/tests/test_s3_ddb.p - Prefer implementing functionality with the standard library or existing workspace dependencies before adding new external crates. - Keep `Cargo.lock` changes intentional; revert unrelated dependency bumps. Pin broken deps with a comment linking the upstream issue. +- The repo has three lockfiles: the root `Cargo.lock`, `python/Cargo.lock`, and `java/lance-jni/Cargo.lock` (the latter two are excluded from the workspace). A `workspace.dependencies` change must be reflected in all three — refresh the excluded ones with `cargo check --manifest-path python/Cargo.toml` and `cargo check --manifest-path java/lance-jni/Cargo.toml`, then commit the updated lockfiles. The `cargo-lock-sync` pre-commit hook catches a miss offline. - Gate optional/domain-specific deps behind Cargo feature flags. Prefer separate crates for domain functionality (geo, NLP). ## Testing Standards @@ -133,6 +134,11 @@ AWS_DEFAULT_REGION=us-east-1 pytest --run-integration python/tests/test_s3_ddb.p - Indent content under MkDocs admonition directives (`!!! note`, etc.) with 4 spaces. - Proofread comments and docs for typos before committing. +## Filing Issues + +- When opening an issue with `gh issue create` or the API, classify it and pass the matching label: `--label bug`, `--label feature`, or `--label performance`. These paths bypass the `.github/ISSUE_TEMPLATE` forms, so the label is not applied automatically. +- Prefix the title to match, e.g. `bug: ...`, `feature: ...`, or `perf: ...`. A content-based labeler (`.github/workflows/issue-labeler.yml`) uses this as a fallback signal, but an explicit `--label` is the reliable path. + ## Pull Requests - PR titles must follow the Conventional Commits specification because `.github/workflows/pr-title.yml` validates the PR title and body with commitlint. Use prefixes like `feat:`, `fix:`, `docs:`, `perf:`, `ci:`, `test:`, `build:`, `style:`, or `chore:`; add a scope when useful. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8f3ec285f31..3940ed7b683 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -25,6 +25,16 @@ Currently Lance is implemented in Rust and comes with a Python wrapper. So you'l a. Install pre-commit: https://pre-commit.com/#install b. Run `pre-commit install` in the root of the repo +## Building for legacy x86_64 hosts (pre-Haswell) + +The default workspace build targets `haswell` (AVX2 + FMA + F16C), matching the published wheels. To build a binary that runs on pre-Haswell silicon (Sandy Bridge / Ivy Bridge / Westmere on Intel, Bulldozer / Piledriver / Steamroller on AMD — i.e. CPUs without AVX2), set the baseline yourself at build time: + +```sh +RUSTFLAGS="-C target-cpu=x86-64-v2" cargo build --release +``` + +Runtime SIMD dispatch in `lance-linalg::distance` will then pick the appropriate tier (scalar / AVX / AVX+FMA / AVX2+FMA / AVX-512) based on the host. From Python, use `lance.simd_info()` to verify which tier was selected. + ## Sample Workflow 1. Fork the repo diff --git a/Cargo.lock b/Cargo.lock index 11d2cbe555e..586cce097f9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -163,9 +163,9 @@ dependencies = [ [[package]] name = "arc-swap" -version = "1.9.1" +version = "1.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" +checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" dependencies = [ "rustversion", ] @@ -1239,6 +1239,20 @@ name = "bytemuck" version = "1.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] [[package]] name = "byteorder" @@ -1248,9 +1262,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.12.0" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "bytes-utils" @@ -1279,9 +1293,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", @@ -1734,18 +1748,18 @@ 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", ] [[package]] name = "crossbeam-queue" -version = "0.3.12" +version = "0.3.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115" +checksum = "803d13fb3b09d88be9f4dbc29062c66b19bf7170867ceb746d2a8689bf6c7a26" dependencies = [ "crossbeam-utils", ] @@ -2860,11 +2874,17 @@ dependencies = [ "encoding_rs", ] +[[package]] +name = "endian-type" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c34f04666d835ff5d62e058c3995147c06f42fe86ff053337632bca83e42702d" + [[package]] name = "env_filter" -version = "1.0.1" +version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32e90c2accc4b07a8456ea0debdc2e7587bdd890680d71173a15d4ae604f6eef" +checksum = "900d271a03799a1ee8d1ca9b19893b48ca674a9284fefcfb85f05e74ed314217" dependencies = [ "log", "regex", @@ -2872,9 +2892,9 @@ dependencies = [ [[package]] name = "env_logger" -version = "0.11.10" +version = "0.11.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0621c04f2196ac3f488dd583365b9c09be011a4ab8b9f37248ffcc8f6198b56a" +checksum = "de671bd27a75a797dc9ae289ba1e77276e75e2026408aab65185384e2d5cd3f6" dependencies = [ "anstream", "anstyle", @@ -3076,7 +3096,7 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "rand 0.9.4", @@ -3483,6 +3503,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if 1.0.4", "crunchy", "num-traits", @@ -3714,9 +3735,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[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" @@ -4122,9 +4143,9 @@ checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" [[package]] name = "io-uring" -version = "0.7.12" +version = "0.7.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d09b98f7eace8982db770e4408e7470b028ce513ac28fecdc6bf4c30fe92b62" +checksum = "9080b15e63775b9a2ac7dca720f7050a8b955e092ea0f6020a4a80f69998cdc0" dependencies = [ "bitflags 2.13.0", "cfg-if 1.0.4", @@ -4198,18 +4219,18 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "jieba-macros" -version = "0.10.1" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46adade69b634535a8f495cf87710ed893cff53e1dbc9dd750c2ab81c5defb82" +checksum = "38fc0f3831de71556de69643b80a08a5c8cd260a23c6b8dbeb7cd923c779cac5" dependencies = [ "phf_codegen", ] [[package]] name = "jieba-rs" -version = "0.10.1" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11b53580aaa8ec8b713da271da434f8947409242c537a9ab3f7b76bdbb19e8a9" +checksum = "a813bbf185c8c62eb6fcf54a223177b644824d91612045dfd80bb779acd080eb" dependencies = [ "bytecount", "cedarwood", @@ -4380,7 +4401,7 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "all_asserts", "approx", @@ -4402,7 +4423,6 @@ dependencies = [ "aws-credential-types", "aws-sdk-dynamodb", "aws-sdk-s3", - "bitpacking", "byteorder", "bytes", "chrono", @@ -4428,6 +4448,7 @@ dependencies = [ "humantime", "itertools 0.14.0", "lance-arrow", + "lance-bitpacking", "lance-core", "lance-datafusion", "lance-datagen", @@ -4483,7 +4504,7 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-buffer", @@ -4492,6 +4513,7 @@ dependencies = [ "arrow-ord", "arrow-schema", "arrow-select", + "bytemuck", "bytes", "futures", "getrandom 0.2.17", @@ -4531,16 +4553,18 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrayref", + "bitpacking", + "crunchy", "paste", "seq-macro", ] [[package]] name = "lance-core" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-buffer", @@ -4580,7 +4604,7 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-array", @@ -4613,7 +4637,7 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-array", @@ -4632,7 +4656,7 @@ dependencies = [ [[package]] name = "lance-derive" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "proc-macro2", "quote", @@ -4641,7 +4665,7 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-arith", "arrow-array", @@ -4686,7 +4710,7 @@ dependencies = [ [[package]] name = "lance-examples" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "all_asserts", "arrow", @@ -4712,7 +4736,7 @@ dependencies = [ [[package]] name = "lance-file" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-arith", "arrow-array", @@ -4751,7 +4775,7 @@ dependencies = [ [[package]] name = "lance-geo" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "datafusion", "geo-traits", @@ -4765,7 +4789,7 @@ dependencies = [ [[package]] name = "lance-index" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "approx", "arc-swap", @@ -4778,7 +4802,6 @@ dependencies = [ "async-channel", "async-recursion", "async-trait", - "bitpacking", "bitvec", "bytes", "chrono", @@ -4802,6 +4825,7 @@ dependencies = [ "jsonb", "lance-arrow", "lance-arrow-stats", + "lance-bitpacking", "lance-core", "lance-datafusion", "lance-datagen", @@ -4842,7 +4866,7 @@ dependencies = [ [[package]] name = "lance-io" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-arith", @@ -4868,6 +4892,8 @@ dependencies = [ "lance-namespace", "lance-testing", "log", + "metrics", + "metrics-util", "mock_instant", "mockall", "moka", @@ -4890,7 +4916,7 @@ dependencies = [ [[package]] name = "lance-linalg" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "approx", "arrow-array", @@ -4906,11 +4932,12 @@ dependencies = [ "proptest", "rand 0.9.4", "rayon", + "rstest", ] [[package]] name = "lance-namespace" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "async-trait", @@ -4922,7 +4949,7 @@ dependencies = [ [[package]] name = "lance-namespace-datafusion" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-schema", @@ -4938,7 +4965,7 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-array", @@ -4966,7 +4993,7 @@ dependencies = [ "log", "object_store", "opendal", - "quick-xml 0.38.4", + "quick-xml 0.40.1", "rand 0.9.4", "reqwest 0.12.28", "ring", @@ -5002,7 +5029,7 @@ dependencies = [ [[package]] name = "lance-select" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-buffer", @@ -5020,7 +5047,7 @@ dependencies = [ [[package]] name = "lance-table" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-array", @@ -5066,7 +5093,7 @@ dependencies = [ [[package]] name = "lance-test-macros" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "proc-macro2", "quote", @@ -5075,7 +5102,7 @@ dependencies = [ [[package]] name = "lance-testing" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-schema", @@ -5088,7 +5115,7 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "icu_segmenter", "jieba-rs", @@ -5101,7 +5128,7 @@ dependencies = [ [[package]] name = "lance-tools" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "clap", "lance-core", @@ -5456,6 +5483,36 @@ dependencies = [ "libc", ] +[[package]] +name = "metrics" +version = "0.24.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89550ee9f79e88fef3119de263694973a8adb26c21d75322164fb8c493039fe2" +dependencies = [ + "portable-atomic", + "rapidhash", +] + +[[package]] +name = "metrics-util" +version = "0.19.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8496cc523d1f94c1385dd8f0f0c2c480b2b8aeccb5b7e4485ad6365523ae376" +dependencies = [ + "aho-corasick", + "crossbeam-epoch", + "crossbeam-utils", + "hashbrown 0.15.5", + "indexmap 2.14.0", + "metrics", + "ordered-float 4.6.0", + "quanta", + "radix_trie", + "rand 0.9.4", + "rand_xoshiro", + "sketches-ddsketch", +] + [[package]] name = "mime" version = "0.3.17" @@ -5637,6 +5694,15 @@ dependencies = [ "rawpointer", ] +[[package]] +name = "nibble_vec" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a5d83df9f36fe23f0c3648c6bbb8b0298bb5f1939c8f2704431371f4b84d43" +dependencies = [ + "smallvec", +] + [[package]] name = "nix" version = "0.26.4" @@ -6271,6 +6337,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "ordered-float" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7bb71e1b3fa6ca1c61f383464aaf2bb0e2f8e772a1f01d486832464de363b951" +dependencies = [ + "num-traits", +] + [[package]] name = "ordered-float" version = "5.3.0" @@ -6868,6 +6943,21 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "quanta" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3ab5a9d756f0d97bdc89019bd2e4ea098cf9cde50ee7564dde6b81ccc8f06c7" +dependencies = [ + "crossbeam-utils", + "libc", + "once_cell", + "raw-cpuid", + "wasi 0.11.1+wasi-snapshot-preview1", + "web-sys", + "winapi", +] + [[package]] name = "quick-error" version = "1.2.3" @@ -6883,15 +6973,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "quick-xml" -version = "0.38.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b66c2058c55a409d601666cffe35f04333cf1013010882cec174a7467cd4e21c" -dependencies = [ - "memchr", -] - [[package]] name = "quick-xml" version = "0.39.4" @@ -6995,6 +7076,16 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" +[[package]] +name = "radix_trie" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c069c179fcdc6a2fe24d8d18305cf085fdbd4f922c041943e203685d6a1c58fd" +dependencies = [ + "endian-type", + "nibble_vec", +] + [[package]] name = "rancor" version = "0.1.1" @@ -7114,6 +7205,24 @@ version = "1.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "973443cf09a9c8656b574a866ab68dfa19f0867d0340648c7d2f6a71b8a8ea68" +[[package]] +name = "rapidhash" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32b266a82f4aa99bb5c25e28d11cc44ace63d91adbcbcee4d323e2ae3d49ef37" +dependencies = [ + "rustversion", +] + +[[package]] +name = "raw-cpuid" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags 2.13.0", +] + [[package]] name = "rawpointer" version = "0.2.1" @@ -7663,9 +7772,9 @@ checksum = "b50b8869d9fc858ce7266cce0194bd74df58b9d0e3f6df3a9fc8eb470d95c09d" [[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_version" @@ -7719,9 +7828,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", @@ -8244,6 +8353,12 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" +[[package]] +name = "sketches-ddsketch" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c6f73aeb92d671e0cc4dca167e59b2deb6387c375391bc99ee743f326994a2b" + [[package]] name = "slab" version = "0.4.12" @@ -8628,7 +8743,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.3", "once_cell", "rustix", "windows-sys 0.61.2", @@ -10183,9 +10298,9 @@ checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" [[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" diff --git a/Cargo.toml b/Cargo.toml index 5eeccfa7806..a24116f9516 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,7 +32,7 @@ resolver = "3" [workspace.package] -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" edition = "2024" authors = ["Lance Devs "] license = "Apache-2.0" @@ -57,27 +57,27 @@ rust-version = "1.91.0" [workspace.dependencies] arc-swap = "1.7" libc = "0.2.176" -lance = { version = "=9.0.0-beta.8", path = "./rust/lance", default-features = false } -lance-arrow = { version = "=9.0.0-beta.8", path = "./rust/lance-arrow" } -lance-core = { version = "=9.0.0-beta.8", path = "./rust/lance-core" } -lance-datafusion = { version = "=9.0.0-beta.8", path = "./rust/lance-datafusion" } -lance-datagen = { version = "=9.0.0-beta.8", path = "./rust/lance-datagen" } -lance-derive = { version = "=9.0.0-beta.8", path = "./rust/lance-derive" } -lance-encoding = { version = "=9.0.0-beta.8", path = "./rust/lance-encoding" } -lance-file = { version = "=9.0.0-beta.8", path = "./rust/lance-file" } -lance-geo = { version = "=9.0.0-beta.8", path = "./rust/lance-geo" } -lance-index = { version = "=9.0.0-beta.8", path = "./rust/lance-index" } -lance-io = { version = "=9.0.0-beta.8", path = "./rust/lance-io", default-features = false } -lance-linalg = { version = "=9.0.0-beta.8", path = "./rust/lance-linalg" } -lance-namespace = { version = "=9.0.0-beta.8", path = "./rust/lance-namespace" } -lance-namespace-impls = { version = "=9.0.0-beta.8", path = "./rust/lance-namespace-impls" } +lance = { version = "=9.0.0-beta.21", path = "./rust/lance", default-features = false } +lance-arrow = { version = "=9.0.0-beta.21", path = "./rust/lance-arrow" } +lance-core = { version = "=9.0.0-beta.21", path = "./rust/lance-core" } +lance-datafusion = { version = "=9.0.0-beta.21", path = "./rust/lance-datafusion" } +lance-datagen = { version = "=9.0.0-beta.21", path = "./rust/lance-datagen" } +lance-derive = { version = "=9.0.0-beta.21", path = "./rust/lance-derive" } +lance-encoding = { version = "=9.0.0-beta.21", path = "./rust/lance-encoding" } +lance-file = { version = "=9.0.0-beta.21", path = "./rust/lance-file" } +lance-geo = { version = "=9.0.0-beta.21", path = "./rust/lance-geo" } +lance-index = { version = "=9.0.0-beta.21", path = "./rust/lance-index" } +lance-io = { version = "=9.0.0-beta.21", path = "./rust/lance-io", default-features = false } +lance-linalg = { version = "=9.0.0-beta.21", path = "./rust/lance-linalg" } +lance-namespace = { version = "=9.0.0-beta.21", path = "./rust/lance-namespace" } +lance-namespace-impls = { version = "=9.0.0-beta.21", path = "./rust/lance-namespace-impls" } lance-namespace-datafusion = { version = "=7.0.0-beta.9", path = "./rust/lance-namespace-datafusion" } lance-namespace-reqwest-client = "0.8.6" -lance-select = { version = "=9.0.0-beta.8", path = "./rust/lance-select" } -lance-tokenizer = { version = "=9.0.0-beta.8", path = "./rust/lance-tokenizer" } -lance-table = { version = "=9.0.0-beta.8", path = "./rust/lance-table" } -lance-test-macros = { version = "=9.0.0-beta.8", path = "./rust/lance-test-macros" } -lance-testing = { version = "=9.0.0-beta.8", path = "./rust/lance-testing" } +lance-select = { version = "=9.0.0-beta.21", path = "./rust/lance-select" } +lance-tokenizer = { version = "=9.0.0-beta.21", path = "./rust/lance-tokenizer" } +lance-table = { version = "=9.0.0-beta.21", path = "./rust/lance-table" } +lance-test-macros = { version = "=9.0.0-beta.21", path = "./rust/lance-test-macros" } +lance-testing = { version = "=9.0.0-beta.21", path = "./rust/lance-testing" } approx = "0.5.1" # Note that this one does not include pyarrow arrow = { version = "58.0.0", optional = false, features = ["prettyprint"] } @@ -103,10 +103,14 @@ aws-sdk-s3 = { version = "1.38.0", default-features = false } half = { "version" = "2.1", default-features = false, features = [ "num-traits", "std", + "bytemuck", ] } -lance-bitpacking = { version = "=9.0.0-beta.8", path = "./rust/compression/bitpacking" } +lance-bitpacking = { version = "=9.0.0-beta.21", path = "./rust/compression/bitpacking" } bitpacking = "0.9" bitvec = "1" +bytemuck = { version = "1", default-features = false, features = [ + "extern_crate_alloc", +] } bytes = "1.11.1" byteorder = "1.5" clap = { version = "4", features = ["derive"] } @@ -143,7 +147,7 @@ datafusion-substrait = { version = "53.0.0", default-features = false } dirs = "6.0.0" either = "1.0" fst = { version = "0.4.7", features = ["levenshtein"] } -fsst = { version = "=9.0.0-beta.8", path = "./rust/compression/fsst" } +fsst = { version = "=9.0.0-beta.21", path = "./rust/compression/fsst" } futures = "0.3" geoarrow-array = "0.8" geoarrow-schema = "0.8" @@ -160,6 +164,8 @@ jieba-rs = { version = "0.10.0", default-features = false } jsonb = { version = "0.5.3", default-features = false, features = ["databend"] } libm = "0.2.15" log = "0.4" +metrics = { version = "0.24" } +metrics-util = { version = "0.19" } mockall = { version = "0.14.0" } mock_instant = { version = "0.6.0" } moka = { version = "0.12", features = ["future", "sync"] } diff --git a/LICENSE b/LICENSE index 79de57d6670..cd4c38213a9 100644 --- a/LICENSE +++ b/LICENSE @@ -226,3 +226,30 @@ under the MIT license: SOFTWARE. https://github.com/pola-rs/polars/blob/main/LICENSE + +-------------------------------------------------------------------------------- + +This project includes code adapted from the quickwit-oss/bitpacking crate, which +is licensed under the MIT license: + + Copyright (c) 2016 Paul Masurel + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE. + +https://github.com/quickwit-oss/bitpacking/blob/main/LICENSE diff --git a/README.md b/README.md index 886fd70425e..08716c84ead 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,17 @@ For more details, see the full [Lance format specification](https://lance.org/fo > [!TIP] > Lance is in active development and we welcome contributions. Please see our [contributing guide](https://lance.org/community/contributing/) for more information. +## File format stability + +Lance releases frequently because the SDKs, integrations, and performance work are moving quickly. This does not mean the Lance file format changes incompatibly in every release. The Lance file format is identified by the `data_storage_version` stored in each dataset, and stable storage versions are a long-term compatibility contract. + +* Once a dataset is written with a stable `data_storage_version`, future Lance releases will continue to support reading that storage version. +* SDK and API compatibility is separate from file format compatibility. SDK/API changes follow semantic versioning and are documented in the [migration guide](https://lance.org/guide/migration/). +* Older Lance releases may not understand file format versions introduced later. If you run mixed Lance versions, pin `data_storage_version` for deterministic writes. +* The `next` file format alias is unstable and should only be used for experimentation, never for production data. + +For production, write data with a stable `data_storage_version`. See the [format versioning guide](https://lance.org/format/file/versioning/) for the current compatibility matrix. + ## Quick Start **Installation** diff --git a/deny.toml b/deny.toml index 3a65cdd9ed6..75b92e53447 100644 --- a/deny.toml +++ b/deny.toml @@ -83,6 +83,8 @@ ignore = [ { id = "RUSTSEC-2024-0436", reason = "`paste` is used by datafusion" }, { id = "RUSTSEC-2023-0071", reason = "`rsa` is used by opendal via reqsign" }, { id = "RUSTSEC-2025-0119", reason = "`number_prefix` used by hf-hub in examples" }, + { id = "RUSTSEC-2026-0194", reason = "`quick-xml` <0.41 pulled transitively via object_store (datafusion) and reqsign-aws-v4 (opendal); upstream must upgrade first" }, + { id = "RUSTSEC-2026-0195", reason = "`quick-xml` <0.41 pulled transitively via object_store (datafusion) and reqsign-aws-v4 (opendal); upstream must upgrade first" }, ] # If this is true, then cargo deny will use the git executable to fetch advisory database. # If this is false, then it uses a built-in git library. diff --git a/docs/clean-full-website.sh b/docs/clean-full-website.sh index db8013cd744..2ebe0b31b87 100755 --- a/docs/clean-full-website.sh +++ b/docs/clean-full-website.sh @@ -39,7 +39,7 @@ nav: - Apache DataFusion: datafusion.md - PostgreSQL: https://github.com/lancedb/pglance - PyTorch: pytorch.md - - Tensorflow: tensorflow.md + - TensorFlow: tensorflow.md EOF mkdir -p "$docs_src/format/catalog/dir" diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 8144a42e5f5..d9ea6f37dda 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -51,7 +51,12 @@ markdown_extensions: line_spans: __span pygments_lang_class: true - pymdownx.inlinehilite - - pymdownx.snippets + - pymdownx.snippets: + # Allow snippets to pull in files from the repo root (e.g. shared docs + # kept alongside Rust source). Paths are relative to the docs/ dir. + base_path: + - . + - .. - pymdownx.tabbed: alternate_style: true - attr_list diff --git a/docs/src/format/index.md b/docs/src/format/index.md index c6956453b05..055c9129c5f 100644 --- a/docs/src/format/index.md +++ b/docs/src/format/index.md @@ -1,44 +1,44 @@ -# Lance Lakehouse Format Specification +# Lance Lakehouse Format Specifications -Lance is a lakehouse format designed as a stack of interoperating specifications instead of a single file or metadata layout. The storage-facing layers are the file format, table format, index formats, and catalog specifications, with a unified namespace interface sitting above them. +Lance is a lakehouse format defined as a stack of interoperating specifications, rather than as a single file format or metadata layout. The storage-facing layers cover files, tables, indices, and catalogs. A unified namespace interface sits above those layers and gives engines a consistent way to work with Lance tables across catalog implementations. ## Architecture Overview -Modern lakehouses are built from cooperating layers. Lance keeps those layers intentionally decoupled so that the file format, table metadata, indices, and catalogs can evolve independently without forcing lock-in across the stack. +Modern lakehouses are built from complementary layers. Lance keeps those layers intentionally decoupled so that the file format, table metadata, indices, and catalogs can evolve independently without forcing lock-in across the stack. ![Lakehouse Stack](../images/lakehouse_stack.png) At a high level: -- The **file format** stores column data in large random-access-friendly pages and avoids row groups. +- The **file format** stores column data in large pages optimized for random access and avoids row groups. - The **table format** manages fragments, manifests, deletions, schema evolution, and ACID commits. - The **index formats** define redundant search structures such as scalar, vector, full-text, and system indices. - The **catalog specs** define how tables are discovered, registered, and coordinated across engines and services. -- The **namespace client spec** provides a unified interface for engines to interact with any catalog implementation. +- The **namespace client spec** provides a unified interface for engines to interact with any catalog implementations. -The layers are designed so that only table readers, table writers, and index readers or writers need to know the on-disk Lance file layout. +The layers are designed so that only table readers, table writers, and index readers or writers need to understand the on-disk Lance file layout. ## Design Themes ### File Format -The Lance file format is optimized for cloud object storage and highly selective reads. It avoids Parquet-style row groups, uses structural encodings that support efficient random access, and keeps statistics and search structures out of the file format so those concerns can evolve as independent indices. +The Lance file format is optimized for cloud object storage and highly selective reads. It avoids Parquet-style row groups, uses structural encodings for efficient random access, and keeps statistics and search structures out of the file format so those concerns can evolve independently as indices. ### Table Format -The Lance table format stores data in two dimensions: rows are grouped into fragments, and each fragment can contain multiple data files that each contribute a subset of columns. This makes column additions and backfills metadata-heavy instead of rewrite-heavy, which is especially useful for feature engineering and embedding workflows. +The Lance table format organizes data in two dimensions: rows are grouped into fragments, and each fragment can contain multiple data files, each contributing a subset of columns. This makes column additions and backfills primarily metadata operations instead of data rewrites, which is especially useful for feature engineering and embedding workflows. ### Index Formats -Indices are first-class table objects. Lance tables define how indices are discovered, versioned, and coordinated transactionally, while the index formats themselves remain decoupled from both the file encoding and the table manifest structure. +Indices are first-class table objects. Lance tables define how indices are discovered, versioned, and coordinated transactionally. The index formats themselves remain decoupled from both the file encoding and the table manifest structure. ### Catalog Specs -Lance provides storage-native and service-oriented catalog options. The [Directory Catalog](catalog/dir/index.md) supports zero-infrastructure deployments directly on object stores, while the [REST Catalog](catalog/rest/index.md) standardizes enterprise-facing APIs and can act as an external manifest store. +Lance provides both storage-native and service-oriented catalog options. The [Directory Catalog](catalog/dir/index.md) supports zero-infrastructure deployments directly on object stores, while the [REST Catalog](catalog/rest/index.md) standardizes enterprise-facing APIs and can act as an external manifest store. ### Namespace Client Spec -The [Namespace Client Spec](namespace/index.md) provides a unified interface for engines to interact with any catalog implementation, across both Lance native catalog specs and third-party catalog systems, in any programming language. This abstraction allows applications to switch between directory-based, REST-based, or third-party catalogs without changing their code. +The [Namespace Client Spec](namespace/index.md) provides a language-agnostic interface for engines to interact with any catalog implementation, including Lance-native catalogs and third-party catalog systems. This abstraction allows applications to switch between directory-based, REST-based, and third-party catalogs without changing their code. ## Specifications diff --git a/docs/src/format/index/index.md b/docs/src/format/index/index.md index 4a2e33c60a4..af8fdf595b4 100644 --- a/docs/src/format/index/index.md +++ b/docs/src/format/index/index.md @@ -177,7 +177,7 @@ or updated. These should be filtered out during query execution. -There are three situations to consider: +There are four situations to consider: 1. **A fragment has some deleted rows.** A few of the rows in the fragment have been marked as deleted, but some of the rows are still present. The row addresses from the deletion @@ -188,6 +188,17 @@ There are three situations to consider: 3. **A fragment has had the indexed column updated in place.** This cannot be detected just by examining metadata. To prevent reading invalid data, the engine should filter out any row addresses that are not in the index's current `fragment_bitmap`. +4. **A fragment has an updated value in an [overlay file](../table/data_overlay_file.md).** + This can be detected by checking if any of the fragments in the index's `fragment_bitmap` + have overlay files. For each overlay whose `committed_version` is greater than the index + segment's `dataset_version`, the overlay carries updated values not reflected in the index, + so its covered rows must be excluded from index results. Excluded rows are re-evaluated + against their current (overlaid) values on the flat path — dropping them without + re-evaluation would silently lose rows that match under the new value. Exclusion is + field-aware: only overlays covering the indexed field matter. You may exclude just the + affected rows or the whole fragment; the latter is simpler and safer but re-evaluates more + rows than necessary. See [Data Overlay Files](../table/data_overlay_file.md#index-integration) + for the exclusion set, re-evaluation, and correctness invariant. ## Compaction and remapping diff --git a/docs/src/format/index/scalar/bloom_filter.md b/docs/src/format/index/scalar/bloom_filter.md index 5a0e08e4228..6d924f8705d 100644 --- a/docs/src/format/index/scalar/bloom_filter.md +++ b/docs/src/format/index/scalar/bloom_filter.md @@ -4,6 +4,9 @@ Bloom filters are probabilistic data structures that allow for fast membership t They are space-efficient and can test whether an element is a member of a set. It's an inexact filter - they may include false positives but never false negatives. +In addition, since finding NULLs is a common query pattern, the index also maintains a +bitmap of null rows which allows it to return exact results for IS NULL queries. + ## Index Details ```protobuf @@ -32,6 +35,13 @@ The bloom filter index stores zone-based bloom filters in a single file: |---------------------------|--------|-------------------------------------------------------------| | `bloomfilter_item` | String | Expected number of items per zone (default: "8192") | | `bloomfilter_probability` | String | False positive probability (default: "0.00057", ~1 in 1754) | +| `null_bitmap` | UInt32 | Index of null bitmap global buffer | + +### Global Buffers + +| Metadata Key | Description | +|---------------------|------------------------------------------------------------| +| `null_bitmap` | A serialized RowAddrTreeMap specifying which rows are null | ## Bloom Filter Spec @@ -122,10 +132,11 @@ Offset 60-63: Block 1, Word 7 (32-bit LE) ## Accelerated Queries -The bloom filter index provides inexact results for the following query types: +The bloom filter index provides inexact results for the following query types (nullability queries +return exact results): | Query Type | Description | Operation | Result Type | |------------|---------------------------|-------------------------------------------|-------------| | **Equals** | `column = value` | Tests if value exists in bloom filter | AtMost | | **IsIn** | `column IN (v1, v2, ...)` | Tests if any value exists in bloom filter | AtMost | -| **IsNull** | `column IS NULL` | Returns zones where has_null is true | AtMost | \ No newline at end of file +| **IsNull** | `column IS NULL` | Returns zones where has_null is true | Exact | diff --git a/docs/src/format/index/scalar/fts.md b/docs/src/format/index/scalar/fts.md index adc7f94d65e..d5c75158011 100644 --- a/docs/src/format/index/scalar/fts.md +++ b/docs/src/format/index/scalar/fts.md @@ -43,6 +43,8 @@ An FTS index may contain multiple partitions. Each partition has its own set of | `_length` | UInt32 | false | Number of documents containing the token | | `_compressed_position` | List> | true | Optional compressed position lists for phrase queries | +The posting-list file schema metadata includes `posting_block_size`, the number of documents encoded per compressed posting block. Older indexes that do not have this metadata use the legacy block size `128`. + ### Metadata File Schema The metadata file contains JSON-serialized configuration and partition information: @@ -67,6 +69,7 @@ The metadata file contains JSON-serialized configuration and partition informati | `min_gram` | UInt32 | 2 | Minimum n-gram length (only for ngram tokenizer) | | `max_gram` | UInt32 | 15 | Maximum n-gram length (only for ngram tokenizer) | | `prefix_only` | Boolean | false | Generate only prefix n-grams (only for ngram tokenizer) | +| `block_size` | UInt32 | 128 | Documents per compressed posting block. Must be 128 or 256. Missing values from older indexes read as 128. `256` is experimental and may introduce breaking changes. | ## Tokenizers diff --git a/docs/src/format/index/scalar/zonemap.md b/docs/src/format/index/scalar/zonemap.md index 256edc2671d..552d476cc86 100644 --- a/docs/src/format/index/scalar/zonemap.md +++ b/docs/src/format/index/scalar/zonemap.md @@ -8,6 +8,9 @@ zones that cannot contain matching values. Zone maps are "inexact" filters - they can definitively exclude zones but may include false positives that require rechecking. +In addition, since finding NULLs is a common query pattern, the index also maintains a +bitmap of null rows which allows it to return exact results for IS NULL queries. + ## Index Details ```protobuf @@ -34,17 +37,25 @@ The zone map index stores zone statistics in a single file: ### Schema Metadata -| Key | Type | Description | -|-----------------|--------|-------------------------------------------| -| `rows_per_zone` | String | Number of rows per zone (default: "8192") | +| Key | Type | Description | +|---------------------|--------|-------------------------------------------| +| `rows_per_zone` | String | Number of rows per zone (default: "8192") | +| `null_bitmap` | UInt32 | Index of null bitmap global buffer | + +### Global Buffers + +| Metadata Key | Description | +|---------------------|------------------------------------------------------------| +| `null_bitmap` | A serialized RowAddrTreeMap specifying which rows are null | ## Accelerated Queries -The zone map index provides inexact results for the following query types: +The zone map index provides inexact results for the following query types (nullability queries +return exact results): | Query Type | Description | Operation | Result Type | |------------|---------------------------|---------------------------------------------|-------------| | **Equals** | `column = value` | Includes zones where min ≤ value ≤ max | AtMost | | **Range** | `column BETWEEN a AND b` | Includes zones where ranges overlap | AtMost | | **IsIn** | `column IN (v1, v2, ...)` | Includes zones that could contain any value | AtMost | -| **IsNull** | `column IS NULL` | Includes zones where null_count > 0 | AtMost | \ No newline at end of file +| **IsNull** | `column IS NULL` | Includes zones where null_count > 0 | Exact | diff --git a/docs/src/format/table/.pages b/docs/src/format/table/.pages index 16c20058608..5b0cb0e95e6 100644 --- a/docs/src/format/table/.pages +++ b/docs/src/format/table/.pages @@ -6,4 +6,5 @@ nav: - Layout: layout.md - Branch & Tag: branch_tag.md - Row ID & Lineage: row_id_lineage.md + - Data Overlay Files: data_overlay_file.md - MemTable & WAL: mem_wal.md diff --git a/docs/src/format/table/data_overlay_file.md b/docs/src/format/table/data_overlay_file.md new file mode 100644 index 00000000000..9da739a29e3 --- /dev/null +++ b/docs/src/format/table/data_overlay_file.md @@ -0,0 +1,388 @@ +# Data Overlay Files + +!!! warning "Experimental" + + This feature is currently experimental and not yet supported in any library. + + + +!!! note "Overlay files require feature flag 64 (data overlay files)" + + A reader or writer that does not understand overlay files must refuse a + dataset that uses them. Silently ignoring an overlay would return stale base + values, which is a correctness bug rather than a degraded experience. + +Overlay files supply new values for a subset of `(row offset, field)` cells +within a fragment **without rewriting the fragment's base data files**. They make +updates cheap when only a small fraction of rows and/or columns change: instead +of rewriting whole columns or moving rows to a new fragment, a writer appends a +small file carrying just the changed cells. + +This is Lance's third mechanism for changing data in place, alongside +[deletion files](index.md#deletion-files) (which remove rows) and +[data evolution](index.md#data-evolution) (which adds or rewrites whole columns). +An overlay changes individual cells. + +## Concepts + +### Coverage and resolution + +Each overlay declares which cells it provides through a **coverage** bitmap (or, +for sparse overlays, one bitmap per field). The bitmaps index **physical row +offsets**. They include deleted rows and are stable even as deletion vectors change. + +To resolve a cell `(offset, field)` on read, walk the fragment's overlays from +**newest to oldest**. The first overlay that covers `(offset, field)` wins; its +value is used. If no overlay covers the cell, the value falls through to the base +data file (or is `NULL` if no base data file holds that field). + +Precedence among overlays is determined by: + +1. `committed_version` — higher wins (see [Versioning](#versioning-and-ordering)). +2. Position in `DataFragment.overlays` as a tiebreaker — a later entry is newer. + +A covered offset whose value is `NULL` overrides the cell **to** `NULL`. This is +distinct from an offset that is simply absent from the bitmap, which falls +through to the base. Coverage, not value-nullness, decides whether an overlay +applies. + +### Interaction with deletions + +Deletions take precedence over overlays. If a row offset is marked deleted in the +fragment's deletion file, any overlay value for that offset is dead and is +ignored, regardless of commit order. + +### Physical layout + +An overlay's data file stores **one value column per field**, in the order of +`data_file.fields`. It does **not** store a row-offset key column. The position of +a covered offset's value within its column is the **rank** of that offset in the +field's coverage bitmap — the number of set bits below it. Resolving a cell is a +rank lookup plus one value fetch, with no separate offset column to store or +search. + +Because different fields may cover different offset sets, the value columns of a +single sparse overlay may have **different lengths**. The Lance file format +permits columns of differing item counts within one file, so a sparse overlay is +representable as a single file. (See [Writer support](#writer-support) for the +current implementation status.) + +### Dense vs. sparse overlays + +A single overlay is one of two shapes: + +- **Dense (rectangular).** One `shared_offset_bitmap` applies to every field. Every + covered offset has a value for every field. This is the common case for a plain + `UPDATE`, where one `SET` list is applied to one set of rows. +- **Sparse.** A `FieldCoverage` carries one bitmap per field, used when different + fields cover different offset sets — for example a `MERGE` with multiple + `WHEN MATCHED` branches, where different rows update different columns. A dense + overlay would have to widen to the bounding rectangle and fill the untouched + cells with their current values (post-images), which for wide columns such as + embeddings means re-storing data that did not change. A sparse overlay stores + exactly the changed cells. + +## Protobuf + +
+DataOverlayFile protobuf message + +```protobuf +%%% proto.message.DataOverlayFile %%% +``` + +
+ +
+FieldCoverage protobuf message + +```protobuf +%%% proto.message.FieldCoverage %%% +``` + +
+ +## Versioning and ordering + +Overlays reuse the dataset version as their ordering clock rather than +introducing a separate generation counter. + +`committed_version` is the dataset version at which an overlay **became +effective** — the version of the commit that introduced it, **not** the version +it was read from. It is stamped at commit time and re-stamped if the commit is +retried, in the same way as the created-at / last-updated-at version sequences. + +This single value drives every ordering decision: + +- **Overlay vs. overlay** (read precedence): higher `committed_version` wins. +- **Overlay vs. index** (query correctness): an index records the + `dataset_version` it was built from. An index whose `dataset_version >= + committed_version` already incorporates the overlay. An overlay whose + `committed_version > index.dataset_version` is newer than the index and its + cells must be excluded from index results and re-evaluated. +- **Scheduler signal**: the gap between an overlay's `committed_version` and an + index's `dataset_version`, or between an overlay and the base, is a staleness + measure the compaction scheduler can use. + +!!! note "Why effective version, not read version" + + Suppose an overlay reads version 5 and commits at version 6, while an index + is built reading version 5 (before the overlay) and commits at version 7 with + `dataset_version = 5`. If the overlay stored its *read* version (5), the test + `5 > 5` is false, the row would not be excluded, and the index — which never + saw the overlay — would return a stale result. Storing the *effective* + version (6) makes `6 > 5` true, the cell is excluded and re-evaluated, and the + result is correct. + +## Index integration + +Building an index over a fragment that has overlays does **not** require dropping +the fragment from the index's coverage. The fragment stays indexed, and the query +path reconciles overlays at query time using an **exclusion set**. + +The exclusion set for an index on field `F` is the union of the coverage bitmaps, +restricted to field `F`, of every overlay whose `committed_version > +index.dataset_version`. The exclusion is **field-aware**: an overlay that touches +only unrelated columns does not exclude anything from the index on `F`. + +The query then proceeds as: + +1. Run the index search as usual, producing candidate rows. +2. Remove any candidate in the exclusion set. (Its indexed value may be stale.) +3. **Re-evaluate** the excluded rows against their current values — the same flat + path already used for the unindexed tail of fragments. For a scalar predicate + this re-applies the filter; for a vector query it re-scores the row's current + vector. Rows that still match are added back to the result. + +Step 3 is what makes exclusion correct rather than merely safe: removing a row +from index candidates without re-evaluating it would silently drop a row that +should match under its new value. + +Exclusion is always *sufficient* because a write changes a cell only by adding an +overlay, and that overlay's `committed_version` — the version of the commit that +adds it — necessarily exceeds the `dataset_version` of any pre-existing index. So +every cell a write changes is guaranteed to fall in that index's exclusion set. +Compaction may remove an overlay only if no index still relies on it for exclusion +(see [Compaction](#compaction)). + +## Compaction + +Overlays accumulate read cost — every overlay is a bitmap to test, a possible +file to open, and additional work to interleave values. Compaction bounds that cost in two modes: + +- **Overlay → overlay.** Merge several overlays into fewer, computing the + post-image per `(offset, field)` by walking the merged overlays newest-first. + The merged overlay takes the **maximum** `committed_version` of its inputs, so + the exclusion semantics are preserved. The merged overlays must be **contiguous + in `committed_version`** — with overlays at v10, v30, and v50 you cannot merge + just v10 and v50, because stamping the result v50 would incorrectly promote + v10's values above the intervening v30 for any cell v30 also covers. Indexes can + still be re-used, but they may now need to exclude more rows. This is cheap to + write and does not touch the base. +- **Overlay → base.** Fold overlays into a fresh base data file, computing the + post-image for every covered cell, then clear the overlays. The base is + complete, so every post-image is well defined. Overlay offsets are physical, so + they cannot survive a rewrite that reorders rows; folding therefore materializes + values rather than carrying overlays forward. + +!!! warning "Folding an indexed field must update its index" + + An overlay→base fold removes the overlay, which removes the exclusion signal + that kept an index correct. Folding an overlay that covers an indexed field + `F` is therefore equivalent to a column rewrite of `F` and must, in the same + commit, either rebuild the index to a `dataset_version` at least the folded + overlay's `committed_version`, or remove the fragment from the index's + coverage so the rows fall to the flat path. Otherwise the index would serve + stale values with no overlay to exclude them. This is the same rule that + already governs rewriting a column that an index is built on. + +When a fragment with overlays is compacted by a row-rewriting operation +(`RewriteRows`, which produces new fragments with new row addresses), the +overlays are folded into the new base as part of the rewrite, and existing +[fragment-reuse remapping](row_id_lineage.md) handles the row-address changes as +it does today. + +## Row lineage + +An overlay write updates the `last_updated_at_version` of every covered row, so +change-data-feed and time-travel queries observe the update. Because overlays are +addressed by physical offset, they do **not** require stable row IDs to be +enabled; lineage updates apply only when those features are on. + +## Worked example + +The following example illustrates how overlays function across their lifecycle, to make the rules above concrete. + +A table `users` with stable row IDs enabled and these fields: + +| field id | name | type | +|----------|-----------|-------------------------| +| 1 | id | `int32` (primary key) | +| 2 | name | `utf8` | +| 3 | age | `int32` | +| 4 | embedding | `fixed_size_list`| + +Created at version 1 as a single fragment `0` with one base data file +`data/file0.lance` holding all four columns. `physical_rows = 4`: + +| offset | id | name | age | embedding | +|--------|----|-------|-----|------------------| +| 0 | 1 | Alice | 30 | … | +| 1 | 2 | Bob | 25 | … | +| 2 | 3 | Carol | 40 | … | +| 3 | 4 | Dave | 22 | … | + +A BTree scalar index on `age` is built at version 1, covering fragment `0` +(`dataset_version = 1`). + +### Step 1 — write an overlay + +```sql +UPDATE users SET age = age + 1 WHERE id IN (2, 4); -- Bob (offset 1), Dave (offset 3) +``` + +This touches one field (`age`) for two rows, so the writer emits a dense overlay — +one shared bitmap covering both offsets — and commits it as version 2. Fragment +`0` gains: + +```text +DataOverlayFile { + data_file: { path: "data/overlay-.lance", fields: [3], column_indices: [0] } + coverage: shared_offset_bitmap = {1, 3} + committed_version: 2 +} +``` + +The overlay file stores a single `age` column with two values, `[26, 23]`, at +ranks `{1,3}.rank(1) = 0` and `{1,3}.rank(3) = 1`. `last_updated_at_version` is +set to 2 for offsets 1 and 3. + +### Step 2 — read + +`SELECT id, age FROM users` reads base ages `[30, 25, 40, 22]`. For `age` +(field 3), the overlay covers offsets 1 and 3, so `age[1]` is replaced with the +overlay value at rank `{1,3}.rank(1) = 0` → `26`, and `age[3]` with the value at +rank `{1,3}.rank(3) = 1` → `23`. Result ages: `[30, 26, 40, 23]`. + +### Step 3 — index query + +```sql +SELECT * FROM users WHERE age = 26; +``` + +The `age` index was built at `dataset_version = 1`; the overlay's +`committed_version` is 2. Since `2 > 1`, the overlay's coverage for `age`, `{1, 3}`, +is the exclusion set for this query. + +- The index (built at v1) holds Bob's *old* `age = 25`, so a lookup for `26` + returns nothing from the index. +- The whole exclusion set is re-evaluated on the flat path, not just the rows the + index returned. Offset 1's current `age` (26, via the overlay) matches, so Bob + is returned; offset 3's current `age` (23) does not match and is dropped. + +The mirror case `WHERE age = 25` shows exclusion preventing a stale hit: the index +returns offset 1 (stale `25`), but offset 1 is excluded, re-evaluated to `26`, and +correctly dropped. + +### Step 4 — a second, non-rectangular write + +```sql +MERGE INTO users USING staged ON users.id = staged.id +WHEN MATCHED AND staged.kind = 'rename' THEN UPDATE SET name = staged.name -- Carol(2), Dave(3) +WHEN MATCHED AND staged.kind = 'embed' THEN UPDATE SET embedding = staged.embedding -- Bob(1) +``` + +`name` is updated for offsets `{2, 3}` and `embedding` for offset `{1}` — different +fields over different rows. This is a sparse overlay, committed as version 3: + +```text +DataOverlayFile { + data_file: { path: "data/overlay-.lance", fields: [2, 4], column_indices: [0, 1] } + coverage: field_coverage { offset_bitmaps: [ {2,3}, {1} ] } + // name (field 2) ^ ^ embedding (field 4) + committed_version: 3 +} +``` + +The file's `name` column has **two** values (`["Caroline", "David"]`, at +ranks 0 and 1 of `{2,3}`) and its `embedding` column has **one** value (at rank 0 +of `{1}`) — columns of different lengths in one file. + +### Step 5 — read after the second write + +`SELECT name, age, embedding FROM users` resolves each field independently, +newest overlay first: + +- `name`: the v3 overlay covers `{2,3}` → `["Alice", "Bob", "Caroline", "David"]`. +- `age`: the v3 overlay does not cover `age`; the v2 overlay still applies at + offsets 1 and 3 → `[30, 26, 40, 23]`. +- `embedding`: the v3 overlay covers `{1}` → Bob's vector is the new one, others + from base. + +Overlays from different versions coexist and apply per field. + +### Step 6 — compaction (overlay → base) + +The scheduler folds both overlays into fragment `0` at version 4, computing +post-images for `age`, `name`, and `embedding`, and writing a new base data file +`data/file1.lance` with those columns. In the old file, fields 2, 3, and 4 are +marked with a tombstone (`-2`); field 1 (`id`) remains. The fragment's `overlays` list is +cleared. Row addresses are preserved (a column rewrite, not a row rewrite), so +stable row IDs and the deletion vector are untouched. + +Because the fold removed the overlay that was excluding offsets 1 and 3 from the +`age` index, the commit must drop fragment `0` from its coverage so `age` queries +fall to the flat path. + +## Guidance + +!!! note "This section is a stub." + + The following are implementation considerations, not part of the on-disk + specification. + +### When to overlay vs. rewrite a column vs. move rows + + + +*(To be expanded.)* The choice between appending an overlay, rewriting a full +column (data evolution), and moving updated rows to a new fragment depends on the +fraction of rows changed, the fraction of columns changed, column width, the +presence of indexes on the changed columns, and the accumulated overlay read +cost. Roughly: few rows changed favors overlays; most rows in a few columns +favors a column rewrite; most columns changed favors moving rows to a new +fragment. + +### Writer support + + + +*(To be expanded.)* Dense (rectangular) overlays write with the existing +equal-length file writer today. Sparse overlays stored as a **single** file +require the writer to emit columns of independent lengths, which the current v2 +writer does not yet do (it advances all columns from one global row counter). +Until that support lands, a writer can express a sparse update as multiple dense +overlays in one transaction. + +### Scheduling compaction + + + +*(To be expanded.)* The overlay→overlay and overlay→base modes have very +different costs; a cost/benefit scheduler decides when each is worthwhile, using +the version gap as a staleness signal. + +## Related specifications + +- [Table format overview](index.md) +- [Transactions: DataOverlay operation](transaction.md#dataoverlay) — write path + and conflict semantics +- [Row ID & Lineage](row_id_lineage.md) +- [Index Formats: handling overlay rows](../index/index.md#handling-deleted-and-invalidated-rows) +- [Format Versioning](versioning.md) diff --git a/docs/src/format/table/index.md b/docs/src/format/table/index.md index 94ea4b90dc9..f9da132cf3b 100644 --- a/docs/src/format/table/index.md +++ b/docs/src/format/table/index.md @@ -168,6 +168,29 @@ However, this invalidates row addresses and requires rebuilding indices, which c +## Data Overlay Files + +!!! warning "Experimental" + + This feature is currently experimental and not yet supported in any library. + + + +!!! note "Overlay files require feature flag 64 (data overlay files)" + +Overlay files supply new values for a subset of cells within +a fragment without rewriting the base data files. They make updates cheap when only +a small percentage of rows and/or columns change: a writer appends a small file +carrying just the changed cells instead of rewriting whole columns or moving rows +to a new fragment. + +For the full specification — coverage and resolution rules, dense vs. sparse layout, +versioning, index integration, compaction, and a worked example — see the +[Data Overlay Files Specification](data_overlay_file.md). + + + ## Related Specifications ### Storage Layout diff --git a/docs/src/format/table/mem_wal.md b/docs/src/format/table/mem_wal.md index 8a228721123..92a5c5fce4a 100644 --- a/docs/src/format/table/mem_wal.md +++ b/docs/src/format/table/mem_wal.md @@ -7,684 +7,674 @@ scan, point lookup, vector search and full-text search. ![MemWAL Overview](../../images/mem_wal_overview.png) -A Lance table is called a **base table** under the context of the MemWAL spec. -It may have an [unenforced primary key](index.md#unenforced-primary-key) defined in the table schema. -Primary keys are required for primary-key lookups and last-write-wins upsert semantics, -but append-only MemWAL tables may omit them. +A Lance table is called the **base table** in this document. +The base table may have an [unenforced primary key](index.md#unenforced-primary-key) in its schema. +Primary keys are required for primary-key lookups and last-write-wins upsert semantics. +Append-only MemWAL tables may omit a primary key. -On top of the base table, the MemWAL spec defines a set of shards. -Writers write to shards, and data in each shard is merged into the base table asynchronously. -An index is kept in the base table for readers to quickly discover the state of all shards at a point of time. +MemWAL adds a set of shards on top of the base table. +Writers append to shards. +Each shard keeps recent data in an in-memory MemTable, persists writes to a per-shard WAL, flushes MemTables as small Lance datasets, and later merges those flushed generations into the base table. -### MemWAL Shard - -A **MemWAL Shard** is the main unit to horizontally scale out writes. - -Each shard has exactly one active writer at any time. -Writers claim a shard and then write data to that shard. -Data in each shard is expected to be merged into the base table asynchronously. +The base table manifest contains one MemWAL system index entry named `__lance_mem_wal`. +This index stores MemWAL configuration and global progress metadata inline in `IndexMetadata.index_details`. +Each shard's own manifest remains authoritative for shard-local mutable state. -For tables with a primary key, rows of the same primary key must be written to one and only one shard. -If two shards contain rows with the same primary key, the following scenario can cause data corruption: - -1. Shard A receives a write with primary key `pk=1` at time T1 -2. Shard B receives a write with primary key `pk=1` at time T2 (T2 > T1) -3. The row in shard B is merged into the base table first -4. The row in shard A is merged into the base table second -5. The row from Shard A (older) now overwrites the row from Shard B (newer) +### MemWAL Shard -This violates the expected "last write wins" semantics. -By ensuring each primary key is assigned to exactly one shard via the sharding spec, -merge order between shards becomes irrelevant for correctness. -Append-only tables without a primary key do not rely on last-write-wins conflict resolution -and may shard by any deterministic append key or partitioning column. +A **MemWAL shard** is the unit of horizontal write scaling. +Each shard has exactly one active writer epoch at a time. +Writers claim a shard, append WAL entries, update the in-memory MemTable, and publish flushed MemTable generations by updating the shard manifest. -See [MemWAL Shard Architecture](#shard-architecture) for the complete shard architecture. +For primary-key tables, all rows for the same primary key must map to the same shard. +If one primary key can appear in multiple shards, asynchronous merge order between shards can make an older row overwrite a newer row. +Append-only tables without a primary key do not rely on last-write-wins conflict resolution and may use any deterministic shard assignment suitable for the workload. ### MemWAL Index -A **MemWAL Index** is the centralized structure for all MemWAL metadata on top of a base table. -A table has at most one MemWAL index. It stores: +The MemWAL index is a system index entry on the base table. +It has `name = "__lance_mem_wal"`, no indexed fields, and no index files. +`IndexMetadata.files` is `None`. +All MemWAL index data is stored in the `MemWalIndexDetails` protobuf message in `IndexMetadata.index_details`. -- **Configuration**: Sharding specs defining how rows map to shards, and which indexes to maintain -- **Merge progress**: Last generation merged to base table for each shard -- **Index catchup progress**: Which merged generation each base table index has been rebuilt to cover -- **Shard snapshots**: Point-in-time snapshot of shard states for read optimization +The index stores: -The index is the source of truth for **configuration**, **merge progress** and **index catchup progress** -Writers and mergers read the MemWAL index to get these configurations before writing. +- **Configuration**: `sharding_specs`, `maintained_indexes`, and `writer_config_defaults`. +- **Merge progress**: `merged_generations`, the last generation merged into the base table for each shard. +- **Index catchup progress**: `index_catchup`, the merged generation covered by each base-table index. +- **Shard snapshots**: optional point-in-time snapshot fields for read optimization. -Each [shard's manifest](#shard-manifest) is authoritative for its own state. -Readers may use **shard snapshots** as a read-only optimization to see a point-in-time view of shards without opening each shard manifest. -Readers that need the latest shard set must discover shard directories in storage and read each shard's latest manifest. - -See [MemWAL Index Details](#memwal-index-details) for the complete structure. +Shard snapshots are not authoritative. +Readers that need the latest shard set list `_mem_wal/` and read each shard's latest manifest. ## Shard Architecture ![Shard Architecture](../../images/mem_wal_regional.png) -Within a shard, writes are stored in an **in-memory table (MemTable)**. -It is also written to the shard's **Write-Ahead Log (WAL)** for durability guarantee. -The MemTable is periodically **flushed** to storage based on memory pressure and other conditions. -**Flushed MemTables** in storage are then asynchronously **merged** into the base table. +Within a shard, writes first enter an in-memory **MemTable** and are durably appended to the shard **write-ahead log (WAL)**. +The MemTable is periodically **flushed** to storage as a Lance dataset. +Flushed MemTables are asynchronously **merged** into the base table. ### MemTable -A MemTable holds rows inserted into the shard before flushing to storage. -It serves 2 purposes: +A MemTable holds rows inserted into a shard before those rows are flushed to storage. +It serves two purposes: -1. build up data and related indexes to be flushed to storage as a flushed MemTable -2. allow a reader to potentially access data that is not flushed to storage yet +1. It buffers data and per-MemTable indexes before a flushed generation is written. +2. It lets readers access data that has not been flushed yet when strong consistency is required. -#### MemTable Format +The storage format does not prescribe the in-memory MemTable layout. +Conceptually, a MemTable is an append log of Arrow record batches. +Later appends have larger in-memory row positions. +For primary-key tables, in-memory reads use the largest visible row position as the newest row for a key. -The complete in-memory format of a MemTable is implementation-specific and out of the scope of this spec. -The Lance core Rust SDK maintains one default implementation and is available through all its language binding SDKs, -but integrations are free to build their own MemTable format depending on the specific use cases, -as long as it follows the MemWAL storage layout, reader and writer requirements when flushing MemTable. +### MemTable Generation -Conceptually, because Lance uses [Arrow as its in-memory data exchange format](https://arrow.apache.org/docs/format/index.html), -for the ease of explanation in this spec, we will treat MemTable as a list of Arrow record batches, -and each write into the MemTable is a new Arrow record batch. +Each MemTable has a monotonically increasing generation number starting from 1. +When generation `N` is flushed and discarded, the next MemTable uses generation `N + 1`. -#### MemTable Generation +Generation numbers order data freshness within one shard: -Based on conditions like memory limit and durability requirements, -a MemTable needs to be **flushed** to storage and discarded. -When that happens, new writes go to a new MemTable and the cycle repeats. -Each MemTable is assigned a monotonically increasing generation number starting from 1. -When MemTable of generation `N` is discarded, the next MemTable gets assigned generation `N+1`. +- Base table data has generation 0. +- Higher MemWAL generations are newer. +- Within the active in-memory generation, higher row positions are newer. +- Within a flushed generation, flush-time deletion vectors hide older duplicate primary-key rows, so readers see at most the newest row for each primary key. -### WAL +## WAL -WAL serves as the durable storage of all MemTables in a shard. -It consists of data in MemTables ordered by generation. -Every time we write to the WAL, we call it a **WAL Flush**. +The WAL is the durable append log for a shard. +Every durable WAL append creates one **WAL entry**. -#### WAL Durability +### WAL Entry Positions -When a write is flushed to WAL, the specific write becomes durable. -Otherwise, if the MemTable is lost, data is also lost. +WAL entry positions are 1-based. +The first data entry is position 1. +Position 0 is reserved as the sentinel value meaning no WAL entry has been covered. -Multiple writes can be batched together in a single WAL flush to reduce WAL flush frequency and improve throughput. -The more writes a single WAL flush batches, the longer it takes for a write to be durable. +Writers append WAL entries in increasing position order. +If entry `N` is not fully written, entry `N + 1` must not exist. +Recovery replays from `replay_after_wal_entry_position + 1`. -The whole LSM tree's durability is determined by the durability of the WAL. -For example, if WAL is stored in Amazon S3, it has 99.999999999% durability. -If it is stored in local disk, the data will be lost if the local disk is damaged. +### WAL Entry Format -#### WAL Entry +Each WAL entry is an Apache Arrow IPC stream file. +The Arrow schema metadata includes: -Each time a WAL flush happens, it adds a new **WAL Entry** to the WAL. -In other words, a WAL consists of an ordered list of WAL entries starting from position 0. -Writer must flush WAL entries in sequential order from lower to higher position. -If WAL entry `N` is not flushed fully, WAL entry `N+1` must not exist in storage. +- `writer_epoch`: decimal string containing the writer epoch that created the entry. +- `fence_sentinel`: optional marker for a data-less fence sentinel entry. -#### WAL Replay +A normal WAL entry contains one or more record batches. +A fence sentinel entry contains no batches and is skipped during replay. +Sentinels are used so an older writer collides on the next WAL position and discovers that it has been fenced. -**Replaying** a WAL means to read data in the WAL from a lower to a higher position. -This is commonly used to recover the latest MemTable after it is lost, -by reading from the start position of the latest MemTable generation till the highest position in the WAL, -assuming proper fencing to guard against multiple writers to the same shard. +### WAL Storage Layout -See [Writer Fencing](#writer-fencing) for the full fencing mechanism. +WAL entries live under `_mem_wal/{shard_id}/wal/`. +Filenames use bit-reversed 64-bit binary names with the `.arrow` suffix: -#### WAL Entry Format +```text +_mem_wal/{shard_id}/wal/{bit_reversed_position}.arrow +``` -Each WAL entry is a file in storage following the [Apache Arrow IPC stream format](https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format) to store the batch of writes in the MemTable. -The writer epoch is stored in the stream's Arrow schema metadata with key `writer_epoch` for fencing validation during replay. +The bit-reversal spreads sequential positions across object-store keyspace. +For example, position 5 is encoded as: -#### WAL Storage Layout +```text +1010000000000000000000000000000000000000000000000000000000000000.arrow +``` -Each WAL entry is stored within the WAL directory of the shard located at `_mem_wal/{shard_id}/wal`. +## Flushed MemTable -WAL files use bit-reversed 64-bit binary naming to distribute files evenly across the directory keyspace. -This optimizes S3 throughput by spreading sequential writes across S3's internal partitions, minimizing throttling. -The filename is the bit-reversed binary representation of the entry ID with suffix `.arrow`. -For example, entry ID 5 (binary `000...101`) becomes `1010000000000000000000000000000000000000000000000000000000000000.arrow`. +A flushed MemTable is a persisted MemTable generation. +It is stored as a Lance dataset under its shard directory. -### Flushed MemTable +!!! note + This structure is similar to a sorted string table in other LSM implementations, but MemWAL flushed generations are not sorted by key. -A flushed MemTable is created by flushing the MemTable to storage. -In Lance MemWAL spec, a flushed MemTable must be a Lance table following the Lance table format spec. +### Flushed MemTable Storage Layout -!!!note -This is called Sorted String Table (SSTable) or Sorted Run in many LSM-tree literatures and implementations. -However, since our MemTable is not sorted, we just use the term flushed MemTable to avoid confusion. +Generation `i` is flushed to: + +```text +_mem_wal/{shard_id}/{random8}_gen_{i}/ +``` -#### Flushed MemTable Storage Layout +`{random8}` is an 8-character random hex value generated for each flush attempt. +If a flush attempt fails, a retry writes a different directory instead of reusing a partially written one. +The shard manifest records the successful directory name in `flushed_generations.path`. + +The generation directory is a standard Lance dataset written with the base table's data storage version. +Each flushed generation is written as one fragment. +Additional MemWAL sidecars may be present: + +```text +{random8}_gen_{i}/ +├── _versions/ +│ └── {version}.manifest +├── _deletions/ # Present when within-generation dedup deletes rows +├── _indices/ # Present when maintained user indexes are built +│ └── {index_uuid}/ +├── _pk_index/ # Primary-key sidecar BTree, not a manifest index +└── bloom_filter.bin # Primary-key bloom filter +``` -The MemTable of generation `i` is flushed to `_mem_wal/{shard_id}/{random_hex}_gen_{i}/` directory, -where `{random_hex}` is a random 8-character hex value generated at flush time. -The random hex value is necessary to ensure if one MemTable flush attempt fails, -The retry can use another directory. -The content within the generation directory follows the [Lance table storage layout](layout.md). +The exact Lance dataset internals follow the [Lance table storage layout](layout.md). -#### Merging MemTable to Base Table +### Flushed Row Order -Generation numbers determine merge order of flushed MemTable into base table: -lower numbers represent older data and must be merged to the base table first to preserve correct upsert semantics. +Flushed MemTable rows are written in forward insert order. +Physical row offsets increase with write time. +For a duplicate primary key within one flushed generation, the newest row has the largest physical offset. -Within a single flushed MemTable for a primary-key table, -if there are multiple rows of the same primary key, the row that is last inserted wins. -Append-only tables without a primary key retain all inserted rows. +Primary-key flushed generations use a deletion vector to expose last-write-wins semantics. +During flush, the writer scans rows in forward order, keeps the last occurrence of each primary key, and marks all earlier duplicate offsets deleted. +The deletion vector is attached to fragment 0 in the generation manifest. -### Shard Manifest +Append-only flushed generations without a primary key do not perform primary-key deduplication and retain every row. -Each shard has a manifest file. This is the source of truth for the state of a shard. +### Tombstone Rows -#### Shard Manifest Contents +Delete operations are represented as rows with the internal `_tombstone` column. +Tombstone rows follow the same forward row ordering and deletion-vector rules as ordinary rows. +If the newest row for a primary key is a tombstone, the deletion vector keeps that tombstone row and hides older rows for the key. +Read planning then filters `_tombstone = false`, so the key is absent from query results. -The manifest contains: +### Flushed Primary-Key Sidecars -- **Fencing state**: `writer_epoch` as the latest writer fencing token, see [Writer Fencing](#writer-fencing) for more details. -- **Shard assignment**: `shard_spec_id` and `shard_field_values` record how this shard maps to its sharding spec. `shard_field_values` is a map from shard field id to the raw Arrow scalar bytes of the computed value; the matching `ShardingField.result_type` in the `ShardingSpec` determines how to interpret each entry (e.g., 4 little-endian bytes for int32, raw UTF-8 bytes for utf8). -- **WAL pointers**: `replay_after_wal_entry_position` (last entry position flushed to MemTable, 0-based), `wal_entry_position_last_seen` (last entry position seen at manifest update, 0-based) -- **Generation trackers**: `current_generation` (next generation to flush), `flushed_generations` list of generation number and directory path pairs (e.g., generation 1 at `a1b2c3d4_gen_1`) +Primary-key MemTables maintain an implicit BTree for primary-key deduplication, independent of `maintained_indexes`. +When a primary-key MemTable is flushed, the flushed generation writes two primary-key sidecars: -Note: `wal_entry_position_last_seen` is a hint that may be stale since it's not updated on WAL write. -It is updated opportunistically by any reader that can update the shard manifest. -The manifest itself is atomically written, but recovery must try to get newer WAL files to find the actual state beyond this hint. +- `bloom_filter.bin` stores the generation's primary-key bloom filter and lets point lookups skip generations that cannot contain the queried key. +- `_pk_index/` stores a standalone BTree over primary-key values to forward row ids. -The manifest is serialized as a protobuf binary file using the `ShardManifest` message. +The `_pk_index/` sidecar is not a maintained user index, is not registered in the generation manifest, and has no manifest UUID. +Its identity is its immutable generation path. +Readers open it directly from `{generation_path}/_pk_index`. -
-ShardManifest protobuf message +The `_pk_index/` directory is a Lance scalar BTree index store: -```protobuf -%%% mem_wal.message.ShardManifest %%% +```text +_pk_index/ +├── page_data.lance +└── page_lookup.lance ``` -
+Readers load this directory as a BTree index using `BTreeIndexDetails` with default parameters. +The primary-key index type is the Arrow type of the primary-key column for a single-column primary key, or `Binary` for a composite primary key. -#### Shard Manifest Versioning +The `page_lookup.lance` file has the following schema: -Manifests are versioned starting from 1 and immutable. -Each update creates a new manifest file at the next version number. -Updates use put-if-not-exists or file rename to ensure atomicity depending on the storage system. -If two processes compete, one wins and the other retries. +| Column | Type | Nullable | Description | +|--------------|-------------------------|----------|------------------------------------------------| +| `min` | {PrimaryKeyIndexType} | true | Minimum primary-key index value in the page | +| `max` | {PrimaryKeyIndexType} | true | Maximum primary-key index value in the page | +| `null_count` | UInt32 | false | Number of null values in the page | +| `page_idx` | UInt32 | false | Page number pointing into `page_data.lance` | -To commit a manifest version: +The `page_data.lance` file has the following schema: -1. Compute the next version number -2. Write the manifest to `{bit_reversed_version}.binpb` using put-if-not-exists -3. In parallel best-effort write to `version_hint.json` with `{"version": }` (failure is acceptable) +| Column | Type | Nullable | Description | +|----------|-----------------------|----------|-------------------------------------------------------------------| +| `values` | {PrimaryKeyIndexType} | true | Sorted primary-key index values | +| `ids` | UInt64 | false | Forward row ids corresponding to each primary-key index value | -To read the latest manifest version: +For a single-column primary key, the indexed value stores the primary-key scalar directly. +For a composite primary key, the indexed value stores an order-preserving binary tuple encoding of all primary-key columns in primary-key column order. +Each tuple column is encoded as: -1. Read `version_hint.json` to get the latest version hint. If not found, start from version 1 -2. Check existence for subsequent versions from the starting version -3. Continue until a version is not found -4. The latest version is the last found version +- `0x00` for null. +- `0x01` followed by the non-null value encoding otherwise. -!!!note -This works because the write rate to shard manifests is significantly lower than read rates. Shard manifests are only updated when shard metadata changes (MemTable flush), not on every write. This ensures HEAD requests will eventually terminate and find the latest version. +Supported non-null value encodings are: -#### Shard Manifest Storage Layout +- Signed integers and date values: sign-flipped 8-byte big-endian integer bytes. +- Unsigned integers: 8-byte big-endian unsigned integer bytes. +- Boolean: one byte, `0x00` for false and `0x01` for true. +- UTF-8 and binary values: raw bytes, with each `0x00` byte escaped as `0x00 0xff`, followed by a `0x00 0x00` terminator. -All shard manifest versions are stored in `_mem_wal/{shard_id}/manifest` directory. +This encoding is injective and preserves primary-key tuple ordering under lexicographic byte comparison. +Composite primary-key columns must use one of the supported encodings above. -Each shard manifest version file uses bit-reversed 64-bit binary naming, the same scheme as WAL files. -For example, version 5 becomes `1010000000000000000000000000000000000000000000000000000000000000.binpb`. +The sidecar row ids are in the same forward row-position space as the data files, deletion vector, and maintained user indexes. +The sidecar is used for cross-generation membership and block-list checks. +It is not used to choose the newest row inside the same flushed generation; the deletion vector has already hidden older same-generation duplicates. -## MemWAL Index Details +### Maintained User Indexes -The MemWAL Index uses the [standard index storage](../index/index.md#index-storage) at `_indices/{UUID}/`. +When the MemWAL index lists `maintained_indexes`, flush may build matching indexes inside the flushed generation. +These index files live in the generation's `_indices/{index_uuid}/` directory and are recorded in the generation manifest. +The implicit primary-key BTree sidecar is not included in `maintained_indexes` and does not live under `_indices/`. -The index stores its data in two parts: +These indexes use the same row-position space as the forward-written data files. +If the generation has a primary key, the generation deletion vector masks stale duplicate rows for indexed reads as well. -1. **Index details** (`index_details` in `IndexMetadata`): Contains configuration, merge progress, and snapshot metadata -2. **Shard snapshots**: Stored as a Lance file or inline, depending on shard count +### Merging Flushed Generations -### Index Details +Flushed generations are merged into the base table in ascending generation order within each shard. +Lower generation numbers are older and must merge before higher generation numbers. +The base table merge uses merge-insert semantics so newer rows overwrite older rows for the same primary key. -The `index_details` field in `IndexMetadata` contains a `MemWalIndexDetails` protobuf message with the following key fields: +## Shard Manifest -- **Configuration fields** (`sharding_specs`, `maintained_indexes`) are the source of truth for MemWAL configuration. - Writers read these fields to determine how to partition data and which indexes to maintain. -- **Merge progress** (`merged_generations`) tracks the last generation merged to the base table for each shard. - This field is updated atomically with merge-insert data commits, enabling conflict resolution when multiple mergers operate concurrently. - Each entry contains the shard UUID and generation number. -- **Index catchup progress** (`index_catchup`) tracks which merged generation each base table index has been rebuilt to cover. - When data is merged from a flushed MemTable to the base table, the base table's indexes may be rebuilt asynchronously. - During this window, queries should use the flushed MemTable's pre-built indexes instead of scanning unindexed data in the base table. - See [Indexed Read Plan](#indexed-read-plan) for details. -- **Shard snapshot fields** (`snapshot_ts_millis`, `num_shards`, `inline_snapshots`) provide a snapshot of shard states. - The actual shard manifests remain authoritative for shard state. - When `num_shards` is 0, the `inline_snapshots` field may be `None` or an empty Lance file with 0 rows but proper schema. +Each shard has a versioned manifest. +The latest shard manifest is the source of truth for shard-local state. + +### Shard Manifest Contents + +The manifest contains: + +- **Identity**: `shard_id`, `shard_spec_id`, and `shard_field_entries`. +- **Fencing state**: `writer_epoch`. +- **WAL pointers**: `replay_after_wal_entry_position` and `wal_entry_position_last_seen`. +- **Generation state**: `current_generation` and `flushed_generations`. +- **Lifecycle state**: `status`, either `ACTIVE` or `SEALED`. + +`shard_field_entries` stores computed shard field values as raw Arrow scalar bytes keyed by `ShardingField.field_id`. +The matching `ShardingField.result_type` determines how to decode each value. +For example, `int32` values are four little-endian bytes and `utf8` values are raw UTF-8 bytes. + +`replay_after_wal_entry_position` is the most recent 1-based WAL position covered by a flushed generation. +The default value 0 means no WAL entry has been covered and recovery starts at position 1. + +`wal_entry_position_last_seen` is a best-effort hint for the most recent WAL position observed at manifest update time. +It is not authoritative because it is not updated on every WAL write. +Recovery must still probe or list WAL files to find the actual tail. + +`status = SEALED` marks a reversible in-flight drop-table operation. +Sealed shards refuse new writer claims. + +The manifest is serialized as the `ShardManifest` protobuf message.
-MemWalIndexDetails protobuf message +ShardManifest protobuf message ```protobuf -%%% mem_wal.message.MemWalIndexDetails %%% +%%% mem_wal.message.ShardManifest %%% ```
-### Shard Identifier - -Each shard has a unique UUID identifier within the table. -When a new shard is created, implementations may assign either a random UUID or -a deterministic UUID derived from the shard assignment when deterministic -writer fencing is required. - -### Shard Discovery +### Shard Manifest Versioning -The MemWAL index can store shard snapshots for read optimization, but those snapshots may lag the latest shard set. -Implementations that need to discover the current shard set should list `_mem_wal/` shard directories and read each shard's latest [shard manifest](#shard-manifest). +Manifest versions start at 1. +Each update writes a new immutable protobuf file: -Each shard manifest records the shard UUID, sharding spec ID, and computed shard field values needed to map the shard back to a sharding spec assignment. - -### Sharding Spec +```text +_mem_wal/{shard_id}/manifest/{bit_reversed_version}.binpb +``` -A **Sharding Spec** defines how all rows in a table are logically divided into different shards, -enabling automatic shard assignment and query-time shard pruning. +Writers use put-if-not-exists or atomic rename, depending on storage support. +If two processes race to write the same next version, one wins and the other reloads and retries. -Each sharding spec has: +After a successful version write, the writer best-effort updates: -- **Spec ID**: A positive integer that uniquely identifies this spec within the MemWAL index. IDs are never reused. -- **Sharding fields**: An array of field definitions that determine how to compute shard values. +```json +{"version": } +``` -Each shard is bound to a specific sharding spec ID, recorded in its [manifest](#shard-manifest). -Shards without a spec ID (`spec_id = 0`) are manually-created shards not governed by any spec. +in: -A sharding spec's field array consists of **sharding field** definitions. -Each sharding field has the following properties: +```text +_mem_wal/{shard_id}/manifest/version_hint.json +``` -| Property | Description | -| ------------- | ------------------------------------------------------------------------- | -| `field_id` | Unique string identifier for this sharding field | -| `source_ids` | Array of field IDs referencing source columns in the schema | -| `transform` | A well-known shard expression, specify this or `expression` | -| `expression` | A DataFusion SQL expression for custom logic, specify this or `transform` | -| `result_type` | The output type of the shard value | +Readers use `version_hint.json` as a starting point and then probe subsequent versions until a version is missing. +The latest manifest is the last existing version. -#### Shard Expression +## MemWAL Index Details -A **Shard Expression** is a [DataFusion SQL expression](https://datafusion.apache.org/user-guide/sql/index.html) that derives a shard value from source column(s). -Source columns are referenced as `col0`, `col1`, etc., corresponding to the order of field IDs in `source_ids`. +The MemWAL index is stored inline in the base table's `IndexMetadata`. +It is a system index with no file directory. +The `index_details` field contains a `MemWalIndexDetails` protobuf message. -Shard expressions must satisfy the following requirements: +Important fields: -1. **Deterministic**: The same input value must always produce the same output value. -2. **Stateless**: The expression must not depend on external state (e.g., current time, random values, session variables). -3. **Type-promotion resistant**: The expression must produce the same result for equivalent values regardless of their numeric type (e.g., `int32(5)` and `int64(5)` must yield the same shard value). -4. **Column removal resistant**: If a source field ID is not found in the schema, the column should be interpreted as NULL. -5. **NULL-safe**: The expression should properly handle NULL inputs and have defined behavior (e.g., return NULL if input is NULL for single-column expressions). -6. **Consistent with result type**: The expression's return type must be consistent with `result_type` in non-NULL cases. +- `sharding_specs`: sharding configuration used by writers and shard pruning. +- `maintained_indexes`: names of base-table indexes to maintain in MemTables and flushed generations. +- `writer_config_defaults`: string map of default writer configuration values persisted for all writers. +- `merged_generations`: per-shard merge progress, updated atomically with base-table merge commits. +- `index_catchup`: per-index coverage progress after data has merged to the base table. +- `snapshot_ts_millis`, `num_shards`, and `inline_snapshots`: optional shard snapshot fields for read optimization. -#### Shard Transform +If a shard is absent from `index_catchup` for an index, that index is assumed to be fully caught up for the shard. -A **Shard Transform** is a well-known shard expression with a predefined name. -When a transform is specified, the expression is derived automatically. +Shard snapshots, when present, use the following Lance file schema: -| Transform | Parameters | Shard Expression | Result Type | -| -------------- | ------------- | --------------------------------------------------------- | -------------- | -| `identity` | (none) | `col0` | same as source | -| `year` | (none) | `date_part('year', col0)` | `int32` | -| `month` | (none) | `date_part('month', col0)` | `int32` | -| `day` | (none) | `date_part('day', col0)` | `int32` | -| `hour` | (none) | `date_part('hour', col0)` | `int32` | -| `bucket` | `num_buckets` | `abs(murmur3(col0)) % N` | `int32` | -| `multi_bucket` | `num_buckets` | `abs(murmur3_multi(col0, col1, ...)) % N` | `int32` | -| `truncate` | `width` | `left(col0, W)` (string) or `col0 - (col0 % W)` (numeric) | same as source | +| Column | Type | Nullable | Description | +|----------------------------|------------------------------|----------|--------------------------------------------------------| +| `shard_id` | Utf8 | false | Shard UUID string | +| `shard_spec_id` | UInt32 | false | Sharding spec that produced the shard | +| `shard_field_{field_id}` | `ShardingField.result_type` | false | Computed shard field value for the given sharding field | -The `bucket` and `multi_bucket` transforms use Murmur3 hash functions: +The MemWAL index data is stored inline. +Readers discover the latest shard set by listing `_mem_wal/` shard directories and reading shard manifests. -- **`murmur3(col)`**: Computes the 32-bit Murmur3 hash (x86 variant, seed 0) of a single column. Returns a signed 32-bit integer. Returns NULL if input is NULL. -- **`murmur3_multi(col0, col1, ...)`**: Computes the Murmur3 hash across multiple columns. Returns a signed 32-bit integer. NULL fields are ignored during hashing; returns NULL only if all inputs are NULL. +
+MemWalIndexDetails protobuf message -The hash result is wrapped with `abs()` and modulo `N` to produce a non-negative bucket number in the range `[0, N)`. +```protobuf +%%% mem_wal.message.MemWalIndexDetails %%% +``` -### Shard Snapshot Storage +
-Shard snapshots are stored using one of two strategies based on the number of shards: +## Sharding -| Shard Count | Storage Strategy | Location | -| ------------------ | ------------------- | ----------------------------------------- | -| <= 100 (threshold) | Inline | `inline_snapshots` field in index details | -| > 100 | External Lance file | `_indices/{UUID}/index.lance` | +A **ShardingSpec** defines how rows map to shards. +Each spec has a positive `spec_id` and one or more `ShardingField` entries. +Each shard manifest records the `shard_spec_id` and the computed shard field values for that shard. +`spec_id = 0` means the shard was manually created and is not governed by a sharding spec. -The threshold (100 shards) is implementation-defined and may vary. +Each `ShardingField` contains: -**Inline storage**: For small shard counts, snapshots are serialized as a Lance file and stored in the `inline_snapshots` field. -This keeps the index metadata compact while avoiding an additional file read for common cases. +- `field_id`: stable identifier for the computed shard field. +- `source_ids`: field IDs of source columns in the Lance schema. +- `transform`: well-known transform name, when using built-in transform evaluation. +- `expression`: reserved custom expression text, mutually exclusive with `transform`. +- `result_type`: Arrow type name for the computed value. +- `parameters`: transform-specific string parameters. -**External Lance file**: For large shard counts, snapshots are stored as a Lance file at `_indices/{UUID}/index.lance`. -This file uses standard Lance format with the shard snapshot schema, enabling efficient columnar access and compression. +The supported built-in transforms are: -### Shard Snapshot Arrow Schema +- `unsharded`: takes no source columns, always returns `int32` value 0, and creates one shard. +- `bucket`: takes one source column and `num_buckets`, hashes the value, and returns an `int32` bucket id in `[0, num_buckets)`. +- `identity`: takes one source column and returns the raw scalar value as the shard value. -Shard snapshots are stored as a Lance file with one row per shard. -The snapshot schema is optimized for shard discovery. Full mutable shard state -remains in the authoritative shard manifest files. +`bucket` computes a deterministic 32-bit hash with seed 0 and then computes: -| Column | Type | Description | -| ------------------------ | ------------- | ---------------------------------------------------------------------------------------------------------- | -| `shard_id` | `utf8` | Shard UUID string | -| `shard_spec_id` | `uint32` | Sharding spec ID (0 if manual) | -| `shard_field_{field_id}` | varies | One column per sharding field defined in the sharding spec, typed to match the field's `ShardingField.result_type`. | +```text +(hash & i32::MAX) % num_buckets +``` -For example, with a sharding spec containing a field `user_bucket` of type `int32`: +`num_buckets` must be in `[1, 1024]`. +Null bucket values hash to 0 and therefore map to bucket 0. +See [Appendix 3: Bucket Hashing](#appendix-3-bucket-hashing) for the exact hash algorithm and test vectors. -| Column | Type | Description | -| -------------------------- | ------- | ---------------------------- | -| ... | ... | (base columns above) | -| `shard_field_user_bucket` | `int32` | Bucket value for this shard | +The `bucket` transform supports scalar boolean, integer, floating-point, date32, time, timestamp, utf8, and large_utf8 source types. +The `identity` transform supports scalar boolean, integer, utf8, and large_utf8 source types. -This schema records the fields needed to map each shard back to its sharding spec -assignment. Readers that need fencing epochs, WAL positions, or flushed -generation state must read the latest shard manifests directly. +The `year`, `month`, `day`, `hour`, `multi_bucket`, and `truncate` transform names are not supported MemWAL sharding transforms and must not be used in `ShardingSpec.transform`. ## Storage Layout -Here is a recap of the storage layout with all the files and concepts defined so far: +The MemWAL storage layout is: -``` +```text {table_path}/ +├── _versions/ +│ └── ... # Base table manifests, including __lance_mem_wal index metadata ├── _indices/ -│ └── {index_uuid}/ # MemWAL Index (uses standard index storage) -│ └── index.lance # Serialized shard snapshots (Lance file) -│ +│ └── ... # Ordinary base table index files; MemWAL index has no files └── _mem_wal/ - └── {shard_id}/ # Shard directory (UUID v4) + └── {shard_id}/ ├── manifest/ - │ ├── {bit_reversed_version}.binpb # Serialized shard manifest (bit-reversed naming) - │ └── version_hint.json # Version hint file + │ ├── {bit_reversed_version}.binpb + │ └── version_hint.json ├── wal/ - │ ├── {bit_reversed_entry_id}.arrow # WAL data files (bit-reversed naming) + │ ├── {bit_reversed_position}.arrow │ └── ... - └── {random_hash}_gen_{i}/ # Flushed MemTable (generation i, random prefix) + └── {random8}_gen_{generation}/ ├── _versions/ - │ └── {version}.manifest # Table manifest (V2 naming scheme) - ├── _indices/ # Indexes - │ ├── {vector_index}/ - │ └── {scalar_index}/ - └── bloom_filter.bin # Primary key bloom filter + │ └── {version}.manifest + ├── _deletions/ + ├── _indices/ + │ └── {index_uuid}/ + ├── _pk_index/ + └── bloom_filter.bin ``` -## Implementation Expectation - -This specification describes the storage layout for the LSM tree architecture. Implementations are free to use any approach to fulfill the storage layout requirements. Once data is written to the expected storage layout, the reader and writer expectations apply. - -The specification defines: +Some flushed-generation subdirectories are conditional. +For example, `_deletions/` is present only when the generation manifest references a deletion vector, `_indices/` is present only when maintained user indexes are built, and `_pk_index/` plus `bloom_filter.bin` are meaningful for primary-key tables. -- **Storage layout**: The directory structure, file formats, and naming conventions for WAL entries, flushed MemTables, shard manifests, and the MemWAL index -- **Durability guarantees**: How data is persisted through WAL entries and flushed MemTables -- **Consistency model**: How readers and writers coordinate through manifests and epoch-based fencing +## Implementation Expectation -Implementations may choose different approaches for: +This document specifies the storage layout and observable reader and writer invariants. +Implementations may choose different in-memory structures, buffering policies, background scheduling, and query execution plans. -- In-memory data structures and indexing -- Buffering strategies before WAL flush -- Background task scheduling and concurrency -- Query execution strategies +An implementation is compatible when it: -As long as the storage layout is correct and the documented invariants are maintained, implementations can optimize for their specific use cases. +1. Writes WAL entries, shard manifests, flushed generations, and MemWAL index metadata using the documented layout. +2. Preserves WAL position, writer fencing, and manifest versioning invariants. +3. Exposes last-write-wins semantics for primary-key tables. +4. Preserves append-only semantics for tables without primary keys. +5. Maintains generation ordering when merging flushed MemTables into the base table. ## Writer Expectations -A writer operates on a single shard and is responsible for: +A writer operates on one shard and is responsible for: -1. Claiming the shard using epoch-based fencing -2. Writing data to WAL entries and flushed MemTables following the [storage layout](#storage-layout) -3. Maintaining the shard manifest to track WAL and generation progress +1. Claiming the shard with epoch-based fencing. +2. Appending WAL entries in sequential 1-based positions. +3. Maintaining in-memory MemTable state. +4. Flushing MemTable generations to Lance datasets. +5. Updating the shard manifest after a generation is durably flushed. ### Writer Fencing -Writers use epoch-based fencing to ensure single-writer semantics per shard. +Writers use `writer_epoch` to enforce single-writer semantics per shard. To claim a shard: -1. Load the latest shard manifest -2. Increment `writer_epoch` by one -3. Atomically write a new manifest version -4. If the write fails (another writer claimed the epoch), reload and retry with a higher epoch +1. Load the latest shard manifest. +2. Verify the shard is `ACTIVE`. +3. Increment `writer_epoch`. +4. Atomically write a new manifest version. +5. If the manifest write loses a race, reload and retry. -Before any manifest update, a writer must verify its `writer_epoch` remains valid: +Before a manifest update, a writer verifies its local epoch is still current: -- If `local_writer_epoch == stored_writer_epoch`: The writer is still active and may proceed -- If `local_writer_epoch < stored_writer_epoch`: The writer has been fenced and must abort +- If `local_writer_epoch == stored_writer_epoch`, the writer may proceed. +- If `local_writer_epoch < stored_writer_epoch`, the writer has been fenced and must abort. -For a concrete example, see [Appendix 1: Writer Fencing Example](#appendix-1-writer-fencing-example). +WAL append conflicts also detect fencing. +If an older writer collides with a newer writer's WAL entry at the same position, it reloads the manifest and observes the higher epoch. +Fence sentinel entries make this collision path explicit without storing data batches. ## Background Job Expectations -Background jobs handle merging flushed MemTables to the base table and garbage collection. +Background jobs merge flushed generations into the base table and remove obsolete shard data. ### MemTable Merger -Flushed MemTables must be merged to the base table in **ascending generation order** within each shard. This ordering is essential for correct upsert semantics: newer generations must overwrite older ones. - -The merge uses Lance's merge-insert operation with atomic transaction semantics: +Flushed MemTables must merge into the base table in ascending generation order within each shard. +The merge uses Lance merge-insert semantics and updates `merged_generations[shard_id]` atomically with the base-table commit. -- `merged_generations[shard_id]` is updated atomically with the data commit -- On commit conflict, check the conflicting commit's `merged_generations` to determine if the generation was already merged +On commit conflict, a merger reloads the conflicting base-table version: -For a concrete example, see [Appendix 2: Concurrent Merger Example](#appendix-2-concurrent-merger-example). +- If the committed `merged_generations[shard_id]` is already greater than or equal to the generation being merged, the merger skips that generation. +- Otherwise, the merger retries from the latest base-table version. ### Garbage Collector -The garbage collector removes obsolete data from shard directories. Flushed MemTables and their referenced WAL files may be deleted after: - -1. The generation has been merged to the base table (`generation <= merged_generations[shard_id]`) -2. All maintained indexes have caught up (`generation <= min(index_catchup[I].caught_up_generation)`) -3. No retained base table version references the generation for time travel +The garbage collector may remove obsolete flushed generations after: -!!!warning - Deleting WAL files weakens [writer fencing](#writer-fencing) and can lead to silent acknowledgement of lost writes. +1. The generation has been merged to the base table. +2. Every maintained index has caught up to cover the merged generation, or the generation is no longer needed for indexed reads. +3. No retained base-table version needs the generation for time travel or consistency. - Fencing detects a stalled writer when its `put-if-not-exists` for the next WAL entry collides with a newer writer's entry at the same position — only that collision triggers the epoch check. If GC has already removed the WAL file at that position, the stalled writer's PUT lands on empty space and succeeds against its old `writer_epoch`. The entry is acknowledged to the client, but the new manifest's `replay_after_wal_entry_position` has already advanced past it, so the data is never replayed. +!!! warning + Deleting WAL files can weaken writer fencing. - Implementations that GC WAL files must compensate, for example by re-checking fence state after each successful WAL write, encoding the writer epoch into the WAL filename so positions are partitioned by epoch, or otherwise guaranteeing a stalled writer cannot land at a position that has been or will be GC'd. + Fencing detects a stalled writer when its put-if-not-exists for the next WAL entry collides with a newer writer's entry at the same position. + If garbage collection has removed that WAL file, the stalled writer may write into empty space with an old `writer_epoch`. + Implementations that garbage collect WAL files must compensate by re-checking fence state after WAL writes, partitioning WAL positions by epoch, or otherwise preventing stale writers from landing at positions that have been garbage collected. ## Reader Expectations ### LSM Tree Merging Read -For tables with a primary key, readers **MUST** merge results from multiple data sources -(base table, flushed MemTables, in-memory MemTables) by primary key to ensure correctness. - -When the same primary key exists in multiple sources, the reader must keep only the newest version based on: - -1. **Generation number** (`_gen`): Higher generation wins. The base table has generation 0, MemTables have positive integers starting from 1. -2. **Row address** (`_rowaddr`): Within the same generation, higher row address wins (later writes within a batch overwrite earlier ones). +For primary-key tables, readers merge rows from the base table, flushed MemTables, and optionally in-memory MemTables by primary key. +The newest row wins. -The ordering for "newest" is: highest `_gen` first, then highest `_rowaddr`. +Freshness ordering within one shard is: -This deduplication is essential because: +1. Higher generation wins. +2. Within the active in-memory generation, higher row position wins. +3. Within a flushed generation, the generation's deletion vector has already hidden older duplicate primary-key rows. -- A row updated in a MemTable also exists (with older data) in the base table -- A flushed MemTable that has been merged to the base table may not yet be garbage collected, causing the same row to appear in both -- A single write batch may contain multiple updates to the same primary key - -Without proper merging, queries would return duplicate or stale rows. +The base table has generation 0. +MemWAL generations are positive. +This ordering applies only to sources selected for the same read plan. +Readers must not include a flushed generation that is already covered by the base table according to `merged_generations[shard_id]`, because otherwise the positive MemWAL generation would incorrectly outrank base-table rows during deduplication. +Rows from different shards do not need primary-key deduplication if the sharding spec guarantees that each primary key maps to exactly one shard. Append-only tables without a primary key do not perform primary-key deduplication. -Readers should include the relevant base table, flushed MemTables, and in-memory MemTables -according to the requested consistency level; duplicate values are treated as distinct appended rows. - -### Reader Consistency +Rows from all selected sources are distinct appended rows. -Reader consistency depends on two factors: +### Tombstones -1. access to in-memory MemTables -2. the source of shard metadata (either through MemWAL index or shard manifests) +Readers must treat `_tombstone = true` rows as delete markers. +In flushed generations, deletion vectors first resolve same-generation duplicate primary keys. +Then query planning filters tombstone rows from user-visible results. +In active in-memory MemTables, the newest visible row position for a primary key wins; if that row is a tombstone, the key is absent. -Strong consistency requires access to in-memory MemTables for all shards involved in the query and reading shard manifests directly. -Otherwise, the query is eventually consistent due to missing unflushed data or stale MemWAL Index snapshots. +### Reader Consistency -!!!note -Reading a stale MemWAL Index does not impact correctness, only freshness: +Reader consistency depends on: - - **Merged MemTable still in index**: If a flushed MemTable has been merged to the base table but still shows in the MemWAL index, readers query both. This results in some inefficiency for querying the same data twice, but [LSM-tree merging](#lsm-tree-merging-read) ensures correct results since both contain the same data. The inefficiency is also compensated by the fact that the data is covered by index and we rarely end up scanning both data. - - **Garbage collected MemTable still in index**: If a flushed MemTable has been garbage collected, but is still in the MemWAL index, readers would fail to open it and skip it. This is also safe because if it is garbage collected, the data must already exist in the base table. - - **Newly flushed MemTable not in index**: If a newly flushed MemTable is added after the snapshot was built, it is not queried. The result is eventually consistent but correct for the snapshot's point in time. +1. Whether the reader can access active in-memory MemTables. +2. Whether shard metadata comes from latest shard manifests or from an older MemWAL index snapshot. -### Query Planning +Strong consistency requires active in-memory MemTable access for relevant shards and direct reads of latest shard manifests. +Otherwise, reads are eventually consistent because unflushed data or newly-created shards may be absent from the read plan. -#### MemTable Collection +Reading a stale MemWAL index snapshot does not corrupt last-write-wins ordering, but it can reduce freshness: -The query planner collects datasets from multiple sources and assembles them for unified query execution. -Datasets come from: +- If a merged flushed generation is still listed, readers must skip it when `generation <= merged_generations[shard_id]`. + For primary-key tables, including it would let an older flushed row outrank newer base-table contents because MemWAL generations are positive and the base table is modeled as generation 0. + For append-only tables, including it would return the same append twice. +- If a garbage-collected flushed generation is still listed, readers may skip it after failing to open it because its data must already be in the base table or be filtered out by `merged_generations`. +- If a newly flushed generation is not listed, the read is consistent with the older snapshot but may miss fresher data. -1. base table (representing already-merged data) -2. flushed MemTables (persisted but not yet merged) -3. optionally in-memory MemTables (if accessible). +Readers that require latest shard membership should list `_mem_wal/` and read shard manifests instead of relying only on snapshots. -Each dataset is tagged with a generation number: 0 for the base table, and positive integers for MemTable generations. -Within a shard, the generation number determines data freshness, with higher numbers representing newer data. -For primary-key tables, rows from different shards do not need deduplication -since each primary key maps to exactly one shard. -Append-only tables without a primary key do not require cross-shard primary-key deduplication. - -The planner also collects bloom filters from each generation for staleness detection during search queries. +### Query Planning -#### Shard Pruning +A query planner collects sources from: -Before executing queries, if sharding spec is available, -the planner evaluates filter predicates against sharding specs to determine which shards may contain matching data. -This pruning step reduces the number of shards to scan. +1. The base table. +2. Flushed MemTables that are not yet safely replaceable by base-table indexed reads. +3. Active in-memory MemTables, when available and required by the requested consistency level. -For each filter predicate: +Each source is tagged with its shard and generation. +For primary-key reads, the planner applies LSM deduplication across selected sources. +For append-only reads, the planner concatenates selected sources without primary-key deduplication. -1. Extract predicates on columns used in sharding specs -2. Evaluate which shard values can satisfy the predicate -3. Prune shards whose values cannot match +Bloom filters and `_pk_index/` sidecars help prune flushed generations during point lookups and cross-generation deduplication. -For example, with a sharding spec using `bucket(user_id, 10)` and a filter `user_id = 123`: +### Shard Pruning -1. Compute `bucket(123, 10) = 3` -2. Only scan shards with bucket value 3 -3. Skip all other shards +When sharding specs are available, the planner evaluates query predicates against shard fields and skips shards whose computed shard values cannot match. -Shard pruning applies to both scan queries and prefilters in search queries. +For example, with `bucket(user_id, 10)` and predicate `user_id = 123`: -#### Indexed Read Plan +1. Compute the bucket id for `123`. +2. Scan only shards whose manifest has the same computed bucket value. +3. Skip all other bucket shards. -When data is merged from a flushed MemTable to the base table, the base table's indexes are rebuilt asynchronously by the base table index builders. -During this window, the merged data exists in the base table but is not yet covered by the base table's indexes. +### Indexed Read Plan -Without special handling, indexed queries would fall back to expensive full scans for the unindexed part of the base table. -To maintain indexed read performance, the query planner should use `index_catchup` progress to determine the optimal data source for each query. +When data is merged from a flushed MemTable into the base table, base-table indexes may lag behind the data commit. +`index_catchup` records which merged generation each base-table index covers. -The key insight is that flushed MemTables serve as a bridge between the base table's index catchup and the current merged state. -For a query that requires a specific index for acceleration, when `index_gen < merged_gen`, -the generations in the gap `(index_gen, merged_gen]` have data already merged in the base table but are not covered by the base table's index. -Since flushed MemTables contain pre-built indexes (created during [MemTable flush](#flushed-memtable)), queries can use these indexes instead of scanning unindexed data in the base table. -This ensures all reads remain indexed regardless of how far behind the async index builder is. +If an indexed query needs index `I` and `I` has only caught up to generation `G` while `merged_generations[shard_id]` is higher, the planner should read the gap from flushed-generation indexes instead of scanning unindexed base-table rows. +Once index `I` catches up, the planner can use the base-table index for those merged rows. ## Appendices ### Appendix 1: Writer Fencing Example -This example demonstrates how epoch-based fencing prevents data corruption when two writers compete for the same shard. - -#### Initial State +Initial shard manifest: +```text +version: 1 +writer_epoch: 5 +replay_after_wal_entry_position: 10 +wal_entry_position_last_seen: 12 +status: ACTIVE ``` -Shard manifest (version 1): - writer_epoch: 5 - replay_after_wal_entry_position: 10 - wal_entry_position_last_seen: 12 -``` - -#### Scenario - -| Step | Writer A | Writer B | Manifest State | -| ---- | --------------------------------------------- | ----------------------------------------- | ------------------ | -| 1 | Loads manifest, sees epoch=5 | | epoch=5, version=1 | -| 2 | Increments to epoch=6, writes manifest v2 | | epoch=6, version=2 | -| 3 | Starts writing WAL entries 13, 14, 15 | | | -| 4 | | Loads manifest v2, sees epoch=6 | epoch=6, version=2 | -| 5 | | Increments to epoch=7, writes manifest v3 | epoch=7, version=3 | -| 6 | | Starts writing WAL entries 16, 17 | | -| 7 | Tries to flush MemTable, loads manifest | | | -| 8 | Sees epoch=7, but local epoch=6 | | | -| 9 | **Writer A is fenced!** Aborts all operations | | | -| 10 | | Continues writing normally | epoch=7, version=3 | - -#### What Happens to Writer A's WAL Entries? - -Writer A wrote WAL entries 13, 14, 15 with `writer_epoch=6` in their schema metadata. - -When Writer B performs crash recovery or MemTable flush: - -1. Reads WAL entries sequentially starting from `replay_after_wal_entry_position + 1` (entry 11, since positions are 0-based) -2. For each entry, checks existence using HEAD request on the bit-reversed filename -3. Continues until an entry is not found (e.g., entry 18 doesn't exist) -4. Finds entries 13, 14, 15, 16, 17 -5. Reads each file's `writer_epoch` from schema metadata -6. Entries 13, 14, 15 have `writer_epoch=6` which is <= current epoch (7) -> **valid, will be replayed** -7. Entries 16, 17 have `writer_epoch=7` -> **valid, will be replayed** -#### Key Points +Writer A loads version 1, claims epoch 6, and writes manifest version 2. +It appends WAL entries 13, 14, and 15 with `writer_epoch = 6`. -1. **No data loss**: Writer A's entries are not discarded. They were written with a valid epoch at the time and will be included in recovery. +Writer B then loads version 2, claims epoch 7, and writes manifest version 3. +It appends WAL entries 16 and 17 with `writer_epoch = 7`. -2. **Consistency preserved**: Writer A is prevented from making further writes that could conflict with Writer B. +When Writer A later tries to flush or update the shard manifest, it reloads the manifest and sees stored epoch 7 while its local epoch is 6. +Writer A is fenced and must abort. -3. **Orphaned files are safe**: WAL files from fenced writers remain on storage and are replayed by the new writer. They are only garbage collected after being included in a flushed MemTable that has been merged. - -4. **Epoch validation timing**: Writers check their epoch before manifest updates (MemTable flush), not on every WAL write. This keeps the hot path fast while ensuring consistency at commit boundaries. +Recovery starts from `replay_after_wal_entry_position + 1`, which is entry 11. +Entries 13, 14, 15, 16, and 17 are valid replay inputs because they were written by epochs that were valid at write time and are not greater than the current shard epoch. ### Appendix 2: Concurrent Merger Example -This example demonstrates how MemWAL Index and conflict resolution handle concurrent mergers safely. - -#### Initial State +Initial state: -``` -MemWAL Index: +```text +MemWAL index: merged_generations: {shard: 5} -Shard manifest (version 1): +Shard manifest: current_generation: 8 - flushed_generations: [(6, "abc123_gen_6"), (7, "def456_gen_7")] + flushed_generations: + - generation: 6, path: "abc12345_gen_6" + - generation: 7, path: "def67890_gen_7" ``` -#### Scenario 1: Racing on the Same Generation - -Two mergers both try to merge generation 6 concurrently. - -| Step | Merger A | Merger B | MemWAL Index | -| ---- | ------------------------- | ------------------------------ | ---------------- | -| 1 | Reads index: merged_gen=5 | | merged_gen=5 | -| 2 | Reads shard manifest | | | -| 3 | Starts merging gen 6 | | | -| 4 | | Reads index: merged_gen=5 | merged_gen=5 | -| 5 | | Reads shard manifest | | -| 6 | | Starts merging gen 6 | | -| 7 | Commits (merged_gen=6) | | **merged_gen=6** | -| 8 | | Tries to commit | | -| 9 | | **Conflict**: reads new index | | -| 10 | | Sees merged_gen=6 >= 6, aborts | | -| 11 | | Reloads, continues to gen 7 | | +Two mergers both try to merge generation 6. +Merger A commits first and updates `merged_generations[shard]` to 6 in the same base-table commit as the data. +Merger B then hits a commit conflict, reloads the latest MemWAL index, sees `merged_generations[shard] >= 6`, skips generation 6, and continues with generation 7. -Merger B's conflict resolution detected that generation 6 was already merged by checking the MemWAL Index in the conflicting commit. +The MemWAL index is the authoritative merge-progress record because it is committed atomically with the base-table data changes. -#### Scenario 2: Crash After Table Commit +### Appendix 3: Bucket Hashing -Merger A crashes after committing to the table. +The bucket transform hash uses 32-bit wrapping arithmetic with these mixing functions. +Right shifts in `fmix` are logical shifts of the `u32` bit pattern. -| Step | Merger A | Merger B | MemWAL Index | -| ---- | ------------------------- | -------------------------------- | ---------------- | -| 1 | Reads index: merged_gen=5 | | merged_gen=5 | -| 2 | Merges gen 6, commits | | **merged_gen=6** | -| 3 | **CRASH** | | merged_gen=6 | -| 4 | | Reads index: merged_gen=6 | merged_gen=6 | -| 5 | | Reads shard manifest | | -| 6 | | **Skips gen 6** (already merged) | | -| 7 | | Merges gen 7, commits | **merged_gen=7** | - -The MemWAL Index is the single source of truth. Merger B correctly used it to determine that generation 6 was already merged. - -#### Key Points +```text +mix_k1(k) = rotl32(k * 0xcc9e2d51, 15) * 0x1b873593 +mix_h1(h, k) = rotl32(h ^ k, 13) * 5 + 0xe6546b64 +fmix(h, len) = + h = h ^ len + h = (h ^ (h >> 16)) * 0x85ebca6b + h = (h ^ (h >> 13)) * 0xc2b2ae35 + h ^ (h >> 16) +``` -1. **Single source of truth**: `merged_generations` is the authoritative source for merge progress, updated atomically with data. +Signed and unsigned casts use two's-complement wrapping. +Values are normalized and hashed as follows: + +- `bool`: `false` as `0`, `true` as `1`, then `hash_i32`. +- `int8`, `int16`, `int32`, `uint8`, `uint16`, `uint32`, `date32`, `time32`: cast to `i32`, then `hash_i32`. +- `int64`, `uint64`, `timestamp`, `time64`: cast to `i64`, then `hash_i64`. +- `float32`: `-0.0` and `+0.0` normalize to bits `0`; all NaNs normalize to `0x7fc00000`; other values use IEEE 754 bits cast to `i32`, then `hash_i32`. +- `float64`: `-0.0` and `+0.0` normalize to bits `0`; all NaNs normalize to `0x7ff8000000000000`; other values use IEEE 754 bits cast to `i64`, then `hash_i64`. +- `utf8` and `large_utf8`: hash the UTF-8 bytes with `hash_bytes`. + +The helper hashes are: + +```text +hash_i32(v) = fmix(mix_h1(0, mix_k1(v)), 4) + +hash_i64(v) = + low = low 32 bits of v as i32 + high = high 32 bits of v as i32 + fmix(mix_h1(mix_h1(0, mix_k1(low)), mix_k1(high)), 8) + +hash_bytes(bytes) = + h = 0 + for each complete 4-byte little-endian chunk: + h = mix_h1(h, mix_k1(chunk_as_i32)) + for each remaining byte: + h = mix_h1(h, mix_k1(sign_extend_i8(byte))) + fmix(h, byte_length) +``` -2. **Conflict resolution uses MemWAL Index**: When a commit conflicts, the merger checks the conflicting commit's MemWAL Index. +Test vectors for `num_buckets = 8`: -3. **No progress regression**: Because MemWAL Index is updated atomically with data, concurrent mergers cannot regress the merge progress. +- `int32` or `date32`: `1 -> 2`, `2 -> 7`, `null -> 0`, `3 -> 1`. +- `utf8`: `"a" -> 1`, `"b" -> 5`, `null -> 0`. +- `bool`: `true -> 2`. +- `float32`: `1.25 -> 0`. +- `float64`: `1.25 -> 0`. diff --git a/docs/src/format/table/transaction.md b/docs/src/format/table/transaction.md index 78dd5301fb8..ef3384c6254 100644 --- a/docs/src/format/table/transaction.md +++ b/docs/src/format/table/transaction.md @@ -466,6 +466,64 @@ The following operations are retryable conflicts with DataReplacement: A concurrent Delete or Update that only adds a deletion vector to a target fragment (without removing it) is compatible: the positional column file stays aligned and the rebase preserves the deletion vector. +### DataOverlay + +Attaches [overlay files](data_overlay_file.md) to fragments, supplying new values +for a subset of `(row offset, field)` cells without rewriting the fragments' base +data files. The overlays are appended to each fragment's existing `overlays` list, +so overlays written by concurrent commits are preserved. Each overlay's +`committed_version` is stamped to the new dataset version at commit time (and +re-stamped on retry), like the created-at / last-updated-at version sequences. + +
+DataOverlay protobuf message + +```protobuf +%%% proto.message.DataOverlay %%% + +%%% proto.message.DataOverlayGroup %%% +``` + +
+ +#### DataOverlay Compatibility + +A DataOverlay operation only changes cells within existing fragments and preserves +physical row addresses, so — like DataReplacement — it is intentionally permissive. +Because overlays stack and the higher `committed_version` wins each covered cell, +independent backfills never conflict, and a concurrent Delete simply makes the +overlay value for a deleted offset inert. Here are the operations that conflict +with DataOverlay: + +- Overwrite +- Restore +- UpdateMemWalState + +The following operations are retryable conflicts with DataOverlay: + +- Rewrite (only if overlapping fragments) — row-rewriting compaction or an + overlay→base fold changes physical row addresses or consumes the overlays, so + the overlay's offsets are no longer valid; the writer must re-read the new + fragment, recompute, and retry. +- Merge (always). +- A row-moving Update that touches an overlaid fragment — a delete-and-reinsert + update (any update that is not a `REWRITE_COLUMNS` column rewrite) relocates the + updated rows into new fragments, so the overlay's physical offsets no longer + address them; the writer must re-read and retry. + +DataOverlay is compatible with another DataOverlay (any fields), Append, Delete, a +`REWRITE_COLUMNS` column rewrite, and DataReplacement, because all of these +preserve physical row addresses: overlay offsets stay valid, the overlay is newer +and wins its covered cells, and the version gate excludes those cells from any +rebuilt index. + +When a DataReplacement or a `REWRITE_COLUMNS` update writes new base values for a +field, it supersedes any older overlay on that field: the writer tombstones the +overlay's entry for the rewritten field — replacing the field id with the obsolete +sentinel, as with obsolete base columns — so the fresh base values are not silently +shadowed. Overlay entries for other fields are preserved, and an overlay left with +no live fields is dropped. + ### UpdateMemWalState Updates the state of MemWal indices (write-ahead log based indices). diff --git a/docs/src/guide/.pages b/docs/src/guide/.pages index 46ddd475799..7a0fd817b1c 100644 --- a/docs/src/guide/.pages +++ b/docs/src/guide/.pages @@ -6,6 +6,7 @@ nav: - JSON Support: json.md - Tags and Branches: tags_and_branches.md - Object Store Configuration: object_store.md + - Observability: observability.md - Distributed Write: distributed_write.md - Distributed Indexing: distributed_indexing.md - Migration Guide: migration.md diff --git a/docs/src/guide/arrays.md b/docs/src/guide/arrays.md index f824dd641c1..5765ca7675b 100644 --- a/docs/src/guide/arrays.md +++ b/docs/src/guide/arrays.md @@ -119,11 +119,11 @@ calling `lance.arrow.ImageTensorArray.to_encoded`. A `lance.arrow.EncodedImageArray.to_tensor` method is provided to decode encoded images and return them as `lance.arrow.FixedShapeImageTensorArray`, from -which they can be converted to numpy arrays or TensorFlow tensors. +which they can be converted to numpy arrays. For decoding images, it will first attempt to use a decoder provided via the optional function parameter. If decoder is not provided it will attempt to use -[Pillow](https://pillow.readthedocs.io/en/stable/) and [tensorflow](https://www.tensorflow.org/api_docs/python/tf/io/encode_png) in that -order. If neither library or custom decoder is available an exception will be raised. +[Pillow](https://pillow.readthedocs.io/en/stable/). If neither Pillow nor a custom +decoder is available an exception will be raised. ```python from lance.arrow import ImageURIArray @@ -132,13 +132,16 @@ uris = [os.path.join(os.path.dirname(__file__), "images/1.png")] encoded_images = ImageURIArray.from_uris(uris).read_uris() print(encoded_images.to_tensor()) -def tensorflow_decoder(images): - import tensorflow as tf +def pillow_decoder(images): + import io import numpy as np + from PIL import Image - return np.stack(tf.io.decode_png(img.as_py(), channels=3) for img in images.storage) + return np.stack( + np.asarray(Image.open(io.BytesIO(img.as_py()))) for img in images.storage + ) -print(encoded_images.to_tensor(tensorflow_decoder)) +print(encoded_images.to_tensor(pillow_decoder)) # # [[42, 42, 42, 255]] # @@ -164,8 +167,8 @@ created by calling `lance.arrow.ImageArray.from_array` and passing in a It can be encoded into to `lance.arrow.EncodedImageArray` by calling `lance.arrow.FixedShapeImageTensorArray.to_encoded` and passing custom encoder If encoder is not provided it will attempt to use -[tensorflow](https://www.tensorflow.org/api_docs/python/tf/io/encode_png) and [Pillow](https://pillow.readthedocs.io/en/stable/) in that order. Default encoders will -encode to PNG. If neither library is available it will raise an exception. +[Pillow](https://pillow.readthedocs.io/en/stable/). The default encoder will encode +to PNG. If neither Pillow nor a custom encoder is available it will raise an exception. ```python from lance.arrow import ImageURIArray @@ -176,4 +179,4 @@ tensor_images.to_encoded() # # [... # b'\x89PNG\r\n\x1a...' -``` \ No newline at end of file +``` diff --git a/docs/src/guide/blob.md b/docs/src/guide/blob.md index dd13fcaab34..00bd5d086e3 100644 --- a/docs/src/guide/blob.md +++ b/docs/src/guide/blob.md @@ -1,20 +1,26 @@ # Blob Columns -Lance supports large binary objects (images, videos, audio, model artifacts) through blob columns. -Blob access is lazy: reads return `BlobFile` handles so callers can stream bytes on demand. +Lance can store large binary objects (images, videos, audio, model artifacts) in blob columns, where they are treated like any other column payload in the dataset. +Blob columns support both planned full-payload reads and lazy file-like access. -![Blob](../images/blob.png) +!!! tip "Choosing between `read_blobs` and `take_blobs`" + - For data loaders and batch processing that need complete byte payloads, use `read_blobs`. + - Use `take_blobs` when you need a `BlobFile` handle for streaming, seeking, or partial reads. -## What This Page Covers -This page focuses on Python blob workflows and uses Lance file format terminology. +![Blob v2 overview](../images/blob-v2-overview-light.png#only-light) +![Blob v2 overview](../images/blob-v2-overview-dark.png#only-dark) + +If you're unsure about whether you need a blob column in the first place (and why it's useful), read the "[when to use blob column vs. inline binary](#when-to-use-a-blob-column-vs-inline-binary)" section below. + +## Quick Start: Blob v2 + +This page focuses on blob workflows in Python and uses Lance file format terminology. - `data_storage_version` means the Lance **file format version** of a dataset. - A dataset's `data_storage_version` is fixed once the dataset is created. - If you need a different file format version, write a **new dataset**. -## Quick Start (Blob v2) - ```python import lance import pyarrow as pa @@ -35,23 +41,27 @@ table = pa.table( ds = lance.write_dataset(table, "./blobs_v22.lance", data_storage_version="2.2") -blob = ds.take_blobs("blob", indices=[0])[0] -with blob as f: - assert f.read() == b"hello blob v2" +_row_address, payload = ds.read_blobs("blob", indices=[0])[0] +assert payload == b"hello blob v2" ``` -## Version Compatibility (Single Source of Truth) +## Version Compatibility + +Blob support is tied to the dataset's file format version. Earlier file format versions +(`< 2.2`) stored blobs using the `lance-encoding:blob` metadata field, while Blob +v2 introduces a new storage layout that requires file format `>= 2.2`. + +The two +schemes are mutually exclusive: for file format `>= 2.2`, legacy blob metadata +(`lance-encoding:blob`) is rejected on write. The table below is the single +source of truth for which scheme is supported at each `data_storage_version`. | Dataset `data_storage_version` | Legacy blob metadata (`lance-encoding:blob`) | Blob v2 (`lance.blob.v2`) | |---|---|---| | `0.1`, `2.0`, `2.1` | Supported for write/read | Not supported | | `2.2+` | Not supported for write | Supported for write/read (recommended) | -Important: - -- For file format `>= 2.2`, legacy blob metadata (`lance-encoding:blob`) is rejected on write. - -## Blob v2 Write Patterns +## Blob v2: Write Patterns Use `blob_field` and `blob_array` to build blob v2 columns. @@ -158,10 +168,22 @@ ds = lance.write_dataset( ) ``` -## Blob v2 Read Patterns +## Blob v2: Read Patterns + +Choose the read API based on the payload shape you want: + +| API | Returns | Use When | +|---|---|---| +| `read_blobs` | `List[Tuple[int, bytes]]` | You need complete blob payloads in memory, such as training loaders or batch preprocessing. | +| `take_blobs` | `List[BlobFile]` | You need file-like objects for streaming, seeking, or partial reads. | +| `scanner(..., blob_handling="all_binary")` | Arrow binary columns | You want blob columns in a scan result or `pyarrow.Table`. | -Use `take_blobs` to fetch file-like handles. -Exactly one selector must be provided: `ids`, `indices`, or `addresses`. +Do not wrap `take_blobs` in your own thread pool just to call `read()` or +`readall()` on every blob. Use `read_blobs` instead; it plans and executes +batched blob reads through Lance's scheduler. + +Exactly one selector must be provided to `read_blobs` or `take_blobs`: `ids`, +`indices`, or `addresses`. | Selector | Typical Use | Stability | |---|---|---| @@ -169,19 +191,17 @@ Exactly one selector must be provided: `ids`, `indices`, or `addresses`. | `ids` | Logical row-id based reads | Stable logical identity (when row ids are available) | | `addresses` | Low-level physical reads and debugging | Unstable physical location | -### Read by row indices +### Read complete payloads by row indices ```python import lance ds = lance.dataset("./blobs_v22.lance") -blobs = ds.take_blobs("blob", indices=[0, 1]) - -with blobs[0] as f: - data = f.read() +rows = ds.read_blobs("blob", indices=[0, 1]) +payloads = [payload for _row_address, payload in rows] ``` -### Read by row ids +### Read complete payloads by row ids ```python import lance @@ -189,10 +209,10 @@ import lance ds = lance.dataset("./blobs_v22.lance") row_ids = ds.to_table(columns=[], with_row_id=True).column("_rowid").to_pylist() -blobs = ds.take_blobs("blob", ids=row_ids[:2]) +rows = ds.read_blobs("blob", ids=row_ids[:2]) ``` -### Read by row addresses +### Read complete payloads by row addresses ```python import lance @@ -200,7 +220,29 @@ import lance ds = lance.dataset("./blobs_v22.lance") row_addrs = ds.to_table(columns=[], with_row_address=True).column("_rowaddr").to_pylist() -blobs = ds.take_blobs("blob", addresses=row_addrs[:2]) +rows = ds.read_blobs("blob", addresses=row_addrs[:2]) +``` + +### Read blob columns as Arrow binary + +```python +import lance + +ds = lance.dataset("./blobs_v22.lance") +table = ds.scanner(columns=["blob"], blob_handling="all_binary").to_table() +payloads = table.column("blob").to_pylist() +``` + +### Open file-like blob handles lazily + +```python +import lance + +ds = lance.dataset("./blobs_v22.lance") +blobs = ds.take_blobs("blob", indices=[0, 1]) + +with blobs[0] as f: + header = f.read(1024) ``` ### Example: decode video frames lazily @@ -229,10 +271,10 @@ with av.open(blob) as container: pass ``` -## Legacy Compatibility Appendix (`data_storage_version` <= `2.1`) +## Legacy Compatibility (`data_storage_version` <= `2.1`) If you need to keep writing legacy blob columns, use file format `0.1`, `2.0`, or `2.1` -and mark `LargeBinary` fields with `lance-encoding:blob = true`. +and mark `LargeBinary` fields with a metadata kwarg `"lance-encoding:blob": true`. ```python import lance @@ -262,12 +304,12 @@ ds = lance.write_dataset( ) ``` -This write pattern is invalid for `data_storage_version >= 2.2`. -For new datasets, prefer blob v2. +As mentioned above, this write pattern is invalid for `data_storage_version >= 2.2`. +For new datasets, it's recommended to use Lance file format 2.2, which uses blob v2 by default. ## Rewrite to a New Blob v2 Dataset -If your current dataset is legacy blob and you want blob v2, rewrite into a new dataset with `data_storage_version="2.2"`. +If your current dataset consists of legacy blobs (stored in file formats <2.2) and you want to opt in to blob v2, you must rewrite it as a new dataset with `data_storage_version="2.2"`. ```python import lance @@ -297,39 +339,33 @@ lance.write_dataset( ) ``` -Warning: - -- The example above materializes binary payloads in memory (`blob_handling="all_binary"` and `to_pylist()`). -- For large datasets, prefer chunked/batched rewrite pipelines. - -## Troubleshooting - -### "Blob v2 requires file version >= 2.2" +!!! warning + - The example above materializes binary payloads in memory (`blob_handling="all_binary"` and `to_pylist()`). + - For large datasets, prefer chunked/batched rewrite pipelines. -Cause: +## When to Use a Blob Column vs. Inline Binary -- You are writing blob v2 values into a dataset/file format below `2.2`. +Not every binary column needs to be a blob column. Plain Arrow `binary`/`large_binary` stores bytes *inline*, interleaved with your other columns, which is simplest and fastest for really small blobs (e.g., thumbnail images). Using a blob column to store the binary payload makes sense when either of these holds: -Fix: +- **You need partial or streaming reads.** Inline binary is always read in full; there is no way to fetch a byte range without materializing the entire value. Blob columns expose `take_blobs` → `BlobFile` handles that seek and range-read, so you pay only for the bytes you touch. +- **Your values are large (roughly 1 MB or more on average).** Operations that rewrite entire rows, such as compaction or some updates, must copy the large inline payloads forward into the new version — even when those bytes never changed. The bigger the payload, the more bytes you rewrite per logical change (write amplification). A blob column keeps large payloads in separate `.blob` files that are referenced rather than re-copied, so these operations don't rewrite the heavy bytes. -- Write to a dataset created with `data_storage_version="2.2"` (or newer). +!!! tip + As a rule of thumb, if average payload size is below a few tens of KB and you only ever read whole values, plain inline binary is fine. Above ~1 MB, or any time you want file-like access, prefer a blob column. Blob v2 also tunes this automatically: by default it keeps payloads under 16 KiB inline, packs mid-sized payloads into shared `.blob` sidecars, and gives payloads over 2 MiB their own dedicated `.blob` file. -### "Legacy blob columns ... are not supported for file version >= 2.2" - -Cause: - -- You are using legacy blob metadata (`lance-encoding:blob`) while writing `2.2+` data. - -Fix: - -- Replace legacy metadata-based columns with blob v2 columns (`blob_field` / `blob_array`). +## Troubleshooting -### "Exactly one of ids, indices, or addresses must be specified" +This section contains commonly noticed issues or errors, and explains how to address them. -Cause: +### Blob v2 requires file version >= 2.2 +**Cause**: You are writing blob v2 values into a dataset/file format below `2.2`. +**Fix**: Write to a dataset created with `data_storage_version="2.2"` (or newer). -- `take_blobs` received none or multiple selectors. +### Legacy blob columns ... are not supported for file version >= 2.2 +**Cause**: You are using legacy blob metadata (`lance-encoding:blob`) while writing `2.2+` data. +**Fix**: Replace legacy metadata-based columns with blob v2 columns (`blob_field` / `blob_array`). -Fix: -- Provide exactly one of `ids`, `indices`, or `addresses`. +### Exactly one of ids, indices, or addresses must be specified +**Cause**: `read_blobs` or `take_blobs` received none or multiple selectors. +**Fix**: Provide exactly one of `ids`, `indices`, or `addresses`. diff --git a/docs/src/guide/data_types.md b/docs/src/guide/data_types.md index 178b155485b..06f26da1db3 100644 --- a/docs/src/guide/data_types.md +++ b/docs/src/guide/data_types.md @@ -32,7 +32,7 @@ Lance supports the full Apache Arrow type system. When writing data through Pyth ### Blob Type for Large Binary Objects -Lance provides a specialized **Blob** type for efficiently storing and retrieving very large binary objects such as videos, images, audio files, or other multimedia content. Unlike regular binary columns, blobs support lazy loading, which means you can read portions of the data without loading everything into memory. +Lance provides a specialized **Blob** type for efficiently storing and retrieving very large binary objects such as videos, images, audio files, or other multimedia content. Blob columns support planned full-payload reads as well as lazy file-like access for streaming, seeking, and partial reads. For new datasets, use blob v2 (`lance.blob.v2`) via `blob_field` and `blob_array`. @@ -62,9 +62,7 @@ table = pa.table( ) ds = lance.write_dataset(table, "./videos_v22.lance", data_storage_version="2.2") -blob = ds.take_blobs("video", indices=[0])[0] -with blob as f: - payload = f.read() +_row_address, payload = ds.read_blobs("video", indices=[0])[0] ``` For legacy compatibility (`data_storage_version <= 2.1`), you can still write blob columns using `LargeBinary` with `lance-encoding:blob=true`. @@ -101,11 +99,18 @@ ds = lance.write_dataset( ) ``` -To read blob data, use `take_blobs()` which returns file-like objects for lazy reading: +To read complete blob payloads into memory, use `read_blobs()`: + +```python +rows = ds.read_blobs("video", indices=[0]) +_row_address, payload = rows[0] +``` + +Use `take_blobs()` only when you need file-like objects for lazy, partial, or seek-based reading: ```python # Retrieve blob as a file-like object (lazy loading) -blobs = ds.take_blobs("video", ids=[0]) +blobs = ds.take_blobs("video", indices=[0]) # Use with libraries that accept file-like objects import av # pip install av diff --git a/docs/src/guide/migration.md b/docs/src/guide/migration.md index 32d2d89ee7f..5efd7b26b7f 100644 --- a/docs/src/guide/migration.md +++ b/docs/src/guide/migration.md @@ -6,6 +6,27 @@ stable and breaking changes should generally be communicated (via warnings) for give users a chance to migrate. This page documents the breaking changes between releases and gives advice on how to migrate. +## 9.0.0 + +* Newly created FTS / inverted indexes now default to format v2 instead of v1. + The `LANCE_FTS_FORMAT_VERSION` environment variable no longer controls the + format used for newly created indexes. Users who need a specific index layout + should pass the index creation parameter `format_version` explicitly. + +* This affects users who create FTS / inverted indexes and need those indexes to + be readable by older Lance versions, or who depend on the v1 index layout. In + those cases, pass `format_version=1` when creating the index. Otherwise, newly + created indexes will use v2 by default, and older Lance readers may not be able + to read them. + + ```python + dataset.create_scalar_index("text", "INVERTED", format_version=1) + ``` + +* Existing v1 FTS indexes remain queryable. Operations that maintain an existing + v1 index, including append, incremental indexing, optimize, and mem-wal + maintained-index flush, should continue preserving the v1 format. + ## 7.2.0 * The `IndexSegmentBuilder` API has been removed from Rust, Python, and Java. diff --git a/docs/src/guide/object_store.md b/docs/src/guide/object_store.md index f901d2c2411..9ce4800b800 100644 --- a/docs/src/guide/object_store.md +++ b/docs/src/guide/object_store.md @@ -41,6 +41,35 @@ These options apply to all object stores. | `client_max_retries` | Number of times for the object store client to retry the request. Default, `3`. | | `client_retry_timeout` | Timeout for the object store client to retry the request in seconds. Default, `180`. | +## Per-Base Configuration + +A dataset can register additional base paths that store part of its data, and each +base may live in a different bucket, account, or storage provider. A storage option +key of the form `base_.` applies `` only to the base path with that +manifest id. Every base inherits the unscoped options; base-scoped entries add to or +override them for that base only. + +```python +import lance +ds = lance.dataset( + "az://account-a/path", + storage_options={ + # Shared defaults, used by the primary dataset and inherited by bases + "account_name": "account-a", + "account_key": "key-a", + # Overrides for the base path with id 1 + "base_1.account_name": "account-b", + "base_1.account_key": "key-b", + }, +) +``` + +Base ids are assigned when bases are registered (`initial_bases` ids are assigned +sequentially starting at 1, in order) and can be inspected through the manifest base +paths. Keys that do not match `base_.` exactly (e.g. `base_url`) are treated +as regular storage options. Exact per-base parameter maps (`base_store_params`, +keyed by base path URI) take precedence over base-scoped keys for that base. + ## S3 Configuration S3 (and S3-compatible stores) have additional configuration options that configure diff --git a/docs/src/guide/observability.md b/docs/src/guide/observability.md new file mode 100644 index 00000000000..f2053b59447 --- /dev/null +++ b/docs/src/guide/observability.md @@ -0,0 +1,69 @@ +# Observability + +Lance can publish operational metrics to your monitoring stack. The table below +is the authoritative catalogue of the metrics Lance emits, shared verbatim with +the Rust [`lance::metrics`](https://docs.rs/lance/latest/lance/metrics/) module +documentation. + +--8<-- "rust/lance/src/metrics.md" + +## Collecting metrics + +Lance emits through the [`metrics`](https://docs.rs/metrics) crate facade, so it +is not tied to a specific backend — you install a recorder/exporter and route +the metrics wherever you like. Metrics are available from **both** the Rust and +Python APIs. + +### Rust + +Enable the `metrics` feature on the `lance` crate: + +```toml +lance = { version = "...", features = ["metrics"] } +``` + +Then install any `metrics`-compatible recorder once at startup, before opening +datasets. For example, with +[`metrics-exporter-prometheus`](https://docs.rs/metrics-exporter-prometheus): + +```rust +metrics_exporter_prometheus::PrometheusBuilder::new() + .install() + .expect("install Prometheus recorder"); +``` + +Any recorder works — Prometheus, StatsD, an OpenTelemetry bridge, and so on. +When no recorder is installed, emission is a cheap no-op. + +### Python + +Unlike Rust, the Python bindings do not let you plug in an arbitrary recorder: +bridging one across the FFI boundary into the Rust `metrics` facade would be +complicated and inefficient. Instead `pylance` standardizes on OpenTelemetry, +which has good Python support, as its recorder. + +The `pylance` wheels are built with the `metrics` feature enabled. Install the +OpenTelemetry extra and call `instrument_lance_metrics`, which registers Lance's +metrics as observable instruments on your OpenTelemetry `MeterProvider`: + +```bash +pip install "pylance[otel]" +``` + +```python +from lance.otel import instrument_lance_metrics + +# Uses the global MeterProvider; pass meter_provider=... to target a specific one. +instrument_lance_metrics() +``` + +From there the metrics flow through whatever OpenTelemetry pipeline you have +configured (OTLP, Prometheus, console, …). Because OpenTelemetry has no +asynchronous histogram instrument, histograms are exported Prometheus-style as +three observable counters: `_bucket`, `_count`, and `_sum`. +Each `_bucket` sample carries an `le` ("less than or equal") attribute +giving that bucket's inclusive upper bound in the metric's unit; the bucket +count is cumulative, covering every observation at or below `le`. For example, a +`lance_object_store_request_duration_seconds_bucket` sample with `le="0.5"` +counts all requests that completed in 0.5 seconds or less, while `le="+Inf"` is +the total count. diff --git a/docs/src/guide/performance.md b/docs/src/guide/performance.md index 14181eb69af..acccdf14776 100644 --- a/docs/src/guide/performance.md +++ b/docs/src/guide/performance.md @@ -219,6 +219,37 @@ use cases. For example, S3 can typically get up to 5000 req/s and with these settings we should get there in about 10 seconds. +## Fragment Sizing + +A Lance table is a collection of fragments tracked by a manifest. How you size those fragments +trades off two classes of work: + +- **Manifest-level operations** scale with the *number* of fragments. Every dataset mutation + (appends, metadata updates, schema changes, compactions, etc.) rewrites the manifest, so a + larger fragment list makes every write slower. Reads pay a similar cost up front: opening a + dataset, listing fragments, planning a scan, and resolving transaction conflicts at the + dataset level all walk the manifest. +- **Fragment-level operations** scale with the *size* of a fragment. These include scans + against a matching fragment, compaction, updates, deletes, and `merge_insert`. Conflict + detection for these operations is also done at the fragment level. + +Fewer, larger fragments make manifest-level operations cheap but make each fragment-level +operation heavier and increase the chance of conflicts when many writers target the same +fragment. More, smaller fragments do the reverse. + +Practical guidance: + +- The default of 1M rows per fragment works well up to ~1B rows. Past that, bumping toward + ~100M rows per fragment is reasonable, though fragment-count limits are rarely the bottleneck + in practice. +- Tens of thousands of fragments per table is generally fine. +- Keep individual fragments well under object-store object-size limits (S3 caps at 5 TB, and + stores tend to misbehave well before that). 10 GB–100 GB per fragment is a reasonable upper + range; 1 TB is a hard ceiling. +- If you run many concurrent updates, deletes, or `merge_insert` operations, err toward more + fragments — conflict detection is per-fragment, so too few fragments leads to excess + retries. + ## Conflict Handling Lance supports concurrent operations on the same table using optimistic concurrency control. When two diff --git a/docs/src/guide/read_and_write.md b/docs/src/guide/read_and_write.md index ec6cde5173a..c7feb69144b 100644 --- a/docs/src/guide/read_and_write.md +++ b/docs/src/guide/read_and_write.md @@ -456,3 +456,146 @@ affected files are no longer part of any ANN index if they were before. Because of this, it's recommended to rewrite files before re-building indices. + +### Cleanup old versions + +Lance is an immutable format — every write creates a new version. The new version +only writes the data that changed, so an insert writes the new rows and an update +rewrites the affected columns for the affected rows. Even a delete creates a +small deletion file. However, old versions still reference the previous data +files, so those files are kept on disk until explicitly removed. Over time this +means storage grows with each operation — inserts, updates, and deletes alike. + +Keeping old versions has important benefits: readers that opened an older version +can continue reading it without interference from concurrent writers, providing +snapshot isolation. Old versions also enable time travel queries, letting you +read the dataset as it existed at any prior point in time. + +`cleanup_old_versions` deletes old version metadata and any data files that are +no longer referenced by any version, reclaiming the accumulated storage. + +!!! warning + + Once old versions are cleaned up, time travel queries to those versions are + no longer possible. Choose your retention window (`older_than`) accordingly — + any version removed by cleanup cannot be recovered. + +```python +import lance + +dataset = lance.dataset("./my_dataset.lance") +dataset.cleanup_old_versions() +``` + +By default, versions older than 7 days are removed. You can override this with +the `older_than` parameter (a `timedelta`): + +```python +from datetime import timedelta + +dataset.cleanup_old_versions(older_than=timedelta(days=1)) +``` + +!!! note + + Tagged versions are exempt from cleanup. See [Tags and Branches](tags_and_branches.md) + for details. + +By default, Lance only removes files that it can **verify** are no longer needed. +A file is verified when Lance can see that it was referenced by an older version +and is no longer referenced by any newer version. However, some orphaned files +cannot be verified this way — for example, files left behind by aborted or failed +commits that were never recorded in any version. These files are +indistinguishable from files being written by an in-progress operation. + +Cleanup will never delete the current (active) version. This means passing +`older_than=timedelta(0)` is safe and will delete all versions except the current +one. + +The `delete_unverified` flag enables a more aggressive strategy that will also +delete these unverified files: + +```python +dataset.cleanup_old_versions( + older_than=timedelta(hours=2), + delete_unverified=True, +) +``` + +!!! danger + + Only use `delete_unverified=True` when you are confident that no other + concurrent operation has been in-progress for longer than the `older_than` + duration. Lance uses the file's age to decide whether an unverified file is + safe to remove, so any operation that is still running past the `older_than` + window risks having its files deleted out from under it. + + In particular, combining `delete_unverified=True` with `older_than=timedelta(0)` + is **extremely dangerous** — if any other operation is in-progress at all, + its data files may be deleted, leading to dataset corruption. + +### Automatic cleanup + +Instead of calling `cleanup_old_versions` manually, you can configure Lance to +clean up old versions automatically during writes. When auto cleanup is enabled, +Lance will run cleanup every *N* commits (the **interval**), removing versions +older than a specified duration. + +Auto cleanup can be enabled when creating a new dataset: + +```python +import lance +import pyarrow as pa +from lance.dataset import AutoCleanupConfig + +table = pa.table({"id": range(100)}) +ds = lance.write_dataset( + table, + "./my_dataset.lance", + auto_cleanup_options=AutoCleanupConfig( + interval=20, # run cleanup every 20 commits + older_than_seconds=3600, # remove versions older than 1 hour + ), +) +``` + +Or enabled on an existing dataset: + +```python +ds = lance.dataset("./my_dataset.lance") +ds.optimize.enable_auto_cleanup( + AutoCleanupConfig( + interval=20, + older_than_seconds=3600, + ) +) +``` + +And disabled again: + +```python +ds.optimize.disable_auto_cleanup() +``` + +Auto cleanup parameters can also be set directly via dataset config keys: + +```python +ds.update_config({ + "lance.auto_cleanup.interval": "20", + "lance.auto_cleanup.older_than": "3600s", +}) +``` + +!!! warning + + Auto cleanup runs as part of the commit path. If your writer does not have + delete permissions, or you are doing high-frequency writes where the extra + latency matters, pass `skip_auto_cleanup=True` to `write_dataset` to skip it + on a per-write basis. + +### Other cleanup strategies + +It is common to run cleanup as a periodic background task on a dedicated server +(for example, via a cron job or scheduled workflow). This keeps cleanup off the +write path entirely, avoiding any impact to write latency, but requires setting +up and maintaining additional infrastructure. diff --git a/docs/src/images/blob-v2-overview-dark.png b/docs/src/images/blob-v2-overview-dark.png new file mode 100644 index 00000000000..04853012405 Binary files /dev/null and b/docs/src/images/blob-v2-overview-dark.png differ diff --git a/docs/src/images/blob-v2-overview-light.png b/docs/src/images/blob-v2-overview-light.png new file mode 100644 index 00000000000..7dae4f70cc3 Binary files /dev/null and b/docs/src/images/blob-v2-overview-light.png differ diff --git a/docs/src/images/lakehouse_stack.png b/docs/src/images/lakehouse_stack.png index e34c999efd8..b3e4db709d6 100644 Binary files a/docs/src/images/lakehouse_stack.png and b/docs/src/images/lakehouse_stack.png differ diff --git a/docs/src/integrations/.pages b/docs/src/integrations/.pages index 62feffae067..ba764bf40f0 100644 --- a/docs/src/integrations/.pages +++ b/docs/src/integrations/.pages @@ -3,7 +3,7 @@ nav: - Apache DataFusion: datafusion.md - PostgreSQL: https://github.com/lancedb/pglance - PyTorch: pytorch.md - - Tensorflow: tensorflow.md + - TensorFlow: tensorflow.md - Apache Spark: spark - Ray: ray - Trino: trino diff --git a/docs/src/integrations/index.md b/docs/src/integrations/index.md index 0304f8f5277..eb3b25051cc 100644 --- a/docs/src/integrations/index.md +++ b/docs/src/integrations/index.md @@ -27,7 +27,7 @@ GitHub organization. | Integration | Description | Source | |---|---|---| | [PyTorch](pytorch.md) | Use `lance.torch.data.LanceDataset` as a `torch.utils.data.IterableDataset` for training and inference. | Built-in | -| [TensorFlow](tensorflow.md) | Use `lance.tf.data.from_lance` to stream Lance data into `tf.data.Dataset` pipelines. | Built-in | +| [TensorFlow](tensorflow.md) | Use `lance_tensorflow.from_lance` to stream Lance data into `tf.data.Dataset` pipelines. | [lance-format/lance-tensorflow](https://github.com/lance-format/lance-tensorflow) | | [Ray](ray) | Distributed read/write of Lance datasets with Ray Data. | [lance-format/lance-ray](https://github.com/lance-format/lance-ray) | | [Hugging Face](huggingface) | Convert and load Hugging Face datasets to and from Lance in a single call. | [lance-format/lance-huggingface](https://github.com/lance-format/lance-huggingface) | diff --git a/docs/src/integrations/tensorflow.md b/docs/src/integrations/tensorflow.md index 1c5d6b87157..03a5c5c5eca 100644 --- a/docs/src/integrations/tensorflow.md +++ b/docs/src/integrations/tensorflow.md @@ -1,92 +1,71 @@ -# Tensorflow Integration +--- +title: TensorFlow +description: Stream Lance datasets into TensorFlow tf.data pipelines with lance-tensorflow. +--- -Lance can be used as a regular [tf.data.Dataset](https://www.tensorflow.org/api_docs/python/tf/data/Dataset) -in [Tensorflow](https://www.tensorflow.org/). +# TensorFlow Integration -!!! warning +The TensorFlow integration is maintained in the +[lance-format/lance-tensorflow](https://github.com/lance-format/lance-tensorflow) +project. - This feature is experimental and the APIs may change in the future. +The main Lance Python package no longer includes `lance.tf`. Install +`lance-tensorflow` and import `lance_tensorflow` instead. + +```bash +pip install lance-tensorflow +``` ## Reading from Lance -Using `lance.tf.data.from_lance`, you can create an `tf.data.Dataset` easily. +Use `lance_tensorflow.from_lance` to create a `tf.data.Dataset` from a Lance +dataset. ```python -import tensorflow as tf -import lance +from lance_tensorflow import from_lance -# Create tf dataset -ds = lance.tf.data.from_lance("s3://my-bucket/my-dataset") - -# Chain tf dataset with other tf primitives +ds = from_lance( + "s3://my-bucket/my-dataset", + columns=["image", "label"], + filter="split = 'train'", + batch_size=256, +) -for batch in ds.shuffling(32).map(lambda x: tf.io.decode_png(x["image"])): - print(batch) +for batch in ds: + print(batch["label"]) ``` -Backed by the Lance [columnar format](../format/index.md), using `lance.tf.data.from_lance` supports -efficient column selection, filtering, and more. +## Dataset Convenience Methods + +If you want `tf.data.Dataset.from_lance`, register the convenience methods +explicitly after importing `lance_tensorflow`. ```python -ds = lance.tf.data.from_lance( - "s3://my-bucket/my-dataset", - columns=["image", "label"], - filter="split = 'train' AND collected_time > timestamp '2020-01-01'", - batch_size=256) -``` +import tensorflow as tf +import lance_tensorflow -By default, Lance will infer the Tensor spec from the projected columns. You can also specify `tf.TensorSpec` manually. +lance_tensorflow.register_tensorflow_dataset() -```python -batch_size = 256 -ds = lance.tf.data.from_lance( - "s3://my-bucket/my-dataset", - columns=["image", "labels"], - batch_size=batch_size, - output_signature={ - "image": tf.TensorSpec(shape=(), dtype=tf.string), - "labels": tf.RaggedTensorSpec( - dtype=tf.int32, shape=(batch_size, None), ragged_rank=1), - }, +ds = tf.data.Dataset.from_lance("s3://my-bucket/my-dataset") ``` -## Distributed Training and Shuffling +## Migration -Since [a Lance Dataset is a set of Fragments](../format/index.md), we can distribute and shuffle Fragments to different -workers. +Replace old imports: ```python -import tensorflow as tf -from lance.tf.data import from_lance, lance_fragments +import lance.tf.data -world_size = 32 -rank = 10 -seed = 123 # -epoch = 100 +ds = lance.tf.data.from_lance(uri) +``` -dataset_uri = "s3://my-bucket/my-dataset" +with: -# Shuffle fragments distributedly. -fragments = - lance_fragments("s3://my-bucket/my-dataset") - .shuffling(32, seed=seed) - .repeat(epoch) - .enumerate() - .filter(lambda i, _: i % world_size == rank) - .map(lambda _, fid: fid) +```python +from lance_tensorflow import from_lance -ds = from_lance( - uri, - columns=["image", "label"], - fragments=fragments, - batch_size=32 - ) -for batch in ds: - print(batch) +ds = from_lance(uri) ``` -!!! warning - - For multiprocessing you should probably not use fork as lance is - multi-threaded internally and fork and multi-thread do not work well. - Refer to [this discussion](https://discuss.python.org/t/concerns-regarding-deprecation-of-fork-with-alive-threads/33555). \ No newline at end of file +See the [lance-tensorflow README](https://github.com/lance-format/lance-tensorflow) +for the current installation and compatibility details. diff --git a/docs/src/quickstart/full-text-search.md b/docs/src/quickstart/full-text-search.md index f990b2bd589..7f09c4325b7 100644 --- a/docs/src/quickstart/full-text-search.md +++ b/docs/src/quickstart/full-text-search.md @@ -98,6 +98,7 @@ ds.create_scalar_index( remove_stop_words=True, # Remove stop words (language-dependent) custom_stop_words=None, # Optional additional stop words (only used if remove_stop_words=True) ascii_folding=True, # Fold accents to ASCII when possible (e.g., "é" -> "e") + block_size=128, # Posting block size: 128 or 256; 256 is experimental ) ``` diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index 1801b6cfc35..5be46e293d2 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -124,9 +124,9 @@ dependencies = [ [[package]] name = "arc-swap" -version = "1.9.1" +version = "1.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" +checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" dependencies = [ "rustversion", ] @@ -918,15 +918,6 @@ version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" -[[package]] -name = "bitpacking" -version = "0.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96a7139abd3d9cebf8cd6f920a389cf3dc9576172e32f4563f188cae3c3eb019" -dependencies = [ - "crunchy", -] - [[package]] name = "bitvec" version = "1.1.1" @@ -1020,6 +1011,20 @@ name = "bytemuck" version = "1.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] [[package]] name = "byteorder" @@ -1029,9 +1034,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.12.0" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "bytes-utils" @@ -1054,9 +1059,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.64" +version = "1.2.66" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dad887fd958be91b5098c0248def011f4523ab786cd411be668777e55063501f" +checksum = "f5d6cac793997bd970000024b2934968efe83b382de4fdcf4fcb46b6ee4ad996" dependencies = [ "find-msvc-tools", "jobserver", @@ -1394,18 +1399,18 @@ 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", ] [[package]] name = "crossbeam-queue" -version = "0.3.12" +version = "0.3.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115" +checksum = "803d13fb3b09d88be9f4dbc29062c66b19bf7170867ceb746d2a8689bf6c7a26" dependencies = [ "crossbeam-utils", ] @@ -2330,9 +2335,9 @@ dependencies = [ [[package]] name = "env_filter" -version = "1.0.1" +version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32e90c2accc4b07a8456ea0debdc2e7587bdd890680d71173a15d4ae604f6eef" +checksum = "900d271a03799a1ee8d1ca9b19893b48ca674a9284fefcfb85f05e74ed314217" dependencies = [ "log", "regex", @@ -2340,9 +2345,9 @@ dependencies = [ [[package]] name = "env_logger" -version = "0.11.10" +version = "0.11.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0621c04f2196ac3f488dd583365b9c09be011a4ab8b9f37248ffcc8f6198b56a" +checksum = "de671bd27a75a797dc9ae289ba1e77276e75e2026408aab65185384e2d5cd3f6" dependencies = [ "anstream", "anstyle", @@ -2479,7 +2484,7 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "rand 0.9.4", @@ -2872,6 +2877,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if 1.0.4", "crunchy", "num-traits", @@ -3079,9 +3085,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[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" @@ -3434,9 +3440,9 @@ dependencies = [ [[package]] name = "io-uring" -version = "0.7.12" +version = "0.7.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d09b98f7eace8982db770e4408e7470b028ce513ac28fecdc6bf4c30fe92b62" +checksum = "9080b15e63775b9a2ac7dca720f7050a8b955e092ea0f6020a4a80f69998cdc0" dependencies = [ "bitflags", "cfg-if 1.0.4", @@ -3656,7 +3662,7 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arc-swap", "arrow", @@ -3673,7 +3679,6 @@ dependencies = [ "async-trait", "async_cell", "aws-credential-types", - "bitpacking", "byteorder", "bytes", "chrono", @@ -3692,6 +3697,7 @@ dependencies = [ "humantime", "itertools 0.14.0", "lance-arrow", + "lance-bitpacking", "lance-core", "lance-datafusion", "lance-encoding", @@ -3729,7 +3735,7 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-buffer", @@ -3738,6 +3744,7 @@ dependencies = [ "arrow-ord", "arrow-schema", "arrow-select", + "bytemuck", "bytes", "futures", "getrandom 0.2.17", @@ -3771,16 +3778,17 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrayref", + "crunchy", "paste", "seq-macro", ] [[package]] name = "lance-core" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-buffer", @@ -3818,7 +3826,7 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-array", @@ -3850,7 +3858,7 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-array", @@ -3867,7 +3875,7 @@ dependencies = [ [[package]] name = "lance-derive" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "proc-macro2", "quote", @@ -3876,7 +3884,7 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-arith", "arrow-array", @@ -3911,7 +3919,7 @@ dependencies = [ [[package]] name = "lance-file" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-arith", "arrow-array", @@ -3941,7 +3949,7 @@ dependencies = [ [[package]] name = "lance-geo" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "datafusion", "geo-traits", @@ -3955,7 +3963,7 @@ dependencies = [ [[package]] name = "lance-index" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arc-swap", "arrow", @@ -3967,7 +3975,6 @@ dependencies = [ "async-channel", "async-recursion", "async-trait", - "bitpacking", "bitvec", "bytes", "chrono", @@ -3987,6 +3994,7 @@ dependencies = [ "jsonb", "lance-arrow", "lance-arrow-stats", + "lance-bitpacking", "lance-core", "lance-datafusion", "lance-datagen", @@ -4023,7 +4031,7 @@ dependencies = [ [[package]] name = "lance-io" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-arith", @@ -4064,7 +4072,7 @@ dependencies = [ [[package]] name = "lance-jni" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-array", @@ -4100,7 +4108,7 @@ dependencies = [ [[package]] name = "lance-linalg" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-buffer", @@ -4116,7 +4124,7 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "async-trait", @@ -4128,7 +4136,7 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-ipc", @@ -4177,7 +4185,7 @@ dependencies = [ [[package]] name = "lance-select" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-buffer", @@ -4192,7 +4200,7 @@ dependencies = [ [[package]] name = "lance-table" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-array", @@ -4229,7 +4237,7 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "icu_segmenter", "rust-stemmers", @@ -4369,9 +4377,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.32" +version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "953f07c43838f8e6f9758cab68bf5bed85465e7587ebe0b823f1bcd81978ad3a" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" [[package]] name = "loom" @@ -6115,9 +6123,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_version" @@ -6170,9 +6178,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", @@ -7508,9 +7516,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.23.3" +version = "1.23.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "144d6b123cef80b301b8f72a9e2ca4370ddec21950d0a103dd22c437006d2db7" +checksum = "bf80a72845275afea99e7f2b434723d3bc7e38470fcd1c7ed39a599c73319a53" dependencies = [ "getrandom 0.4.3", "js-sys", @@ -8304,9 +8312,9 @@ checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" [[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 = "yoke" diff --git a/java/lance-jni/Cargo.toml b/java/lance-jni/Cargo.toml index 103800e5cc5..c51df5769c7 100644 --- a/java/lance-jni/Cargo.toml +++ b/java/lance-jni/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lance-jni" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" edition = "2024" authors = ["Lance Devs "] rust-version = "1.91" @@ -9,11 +9,18 @@ repository = "https://github.com/lance-format/lance" readme = "../../README.md" description = "JNI bindings for Lance Columnar format" +# `java/lance-jni/` is excluded from the top-level workspace (see ../../Cargo.toml) and +# keeps its own Cargo.lock. Declaring an empty workspace here makes this crate an explicit +# workspace root, so cargo resolves it standalone even when the checkout is nested inside +# another cargo project (e.g. a git worktree placed under the main checkout). +[workspace] + [lib] crate-type = ["cdylib"] [features] default = [] +backtrace = ["lance/backtrace", "lance-core/backtrace"] [dependencies] lance = { path = "../../rust/lance", features = ["substrait"] } diff --git a/java/lance-jni/src/error.rs b/java/lance-jni/src/error.rs index cdb922a3cef..0affdca8f98 100644 --- a/java/lance-jni/src/error.rs +++ b/java/lance-jni/src/error.rs @@ -181,29 +181,36 @@ impl std::fmt::Display for Error { impl From for Error { fn from(err: LanceError) -> Self { + let backtrace_suffix = err + .backtrace() + .map(|bt| format!("\n\nRust backtrace:\n{}", bt)) + .unwrap_or_default(); + let message = format!("{}{}", err, backtrace_suffix); + match &err { LanceError::DatasetNotFound { .. } | LanceError::DatasetAlreadyExists { .. } | LanceError::CommitConflict { .. } - | LanceError::InvalidInput { .. } => Self::input_error(err.to_string()), - LanceError::IO { .. } => Self::io_error(err.to_string()), - LanceError::Timeout { .. } => Self::timeout_error(err.to_string()), - LanceError::NotSupported { .. } => Self::unsupported_error(err.to_string()), - LanceError::NotFound { .. } => Self::io_error(err.to_string()), + | LanceError::InvalidInput { .. } => Self::input_error(message), + LanceError::IO { .. } => Self::io_error(message), + LanceError::Timeout { .. } => Self::timeout_error(message), + LanceError::NotSupported { .. } => Self::unsupported_error(message), + LanceError::NotFound { .. } => Self::io_error(message), LanceError::Namespace { source, .. } => { // Try to downcast to NamespaceError and get the error code if let Some(ns_err) = source.downcast_ref::() { - Self::namespace_error(ns_err.code().as_u32(), ns_err.to_string()) + let ns_message = format!("{}{}", ns_err, backtrace_suffix); + Self::namespace_error(ns_err.code().as_u32(), ns_message) } else { log::warn!( "Failed to downcast NamespaceError source, falling back to runtime error. \ This may indicate a version mismatch. Source type: {:?}", source ); - Self::runtime_error(err.to_string()) + Self::runtime_error(message) } } - _ => Self::runtime_error(err.to_string()), + _ => Self::runtime_error(message), } } } @@ -241,3 +248,104 @@ impl From for Error { Self::input_error(err.to_string()) } } + +#[cfg(test)] +mod tests { + use super::*; + + // Helper: extract the java_class from an Error via Display output + fn java_class(err: &Error) -> &JavaExceptionClass { + &err.java_class + } + + #[test] + fn test_invalid_input_maps_to_illegal_argument() { + let lance_err = LanceError::invalid_input("bad input"); + let jni_err: Error = lance_err.into(); + assert_eq!( + *java_class(&jni_err), + JavaExceptionClass::IllegalArgumentException + ); + assert!(jni_err.message.contains("bad input")); + } + + #[test] + fn test_dataset_not_found_maps_to_illegal_argument() { + let lance_err = LanceError::dataset_not_found("my_dataset", "not found".to_string().into()); + let jni_err: Error = lance_err.into(); + assert_eq!( + *java_class(&jni_err), + JavaExceptionClass::IllegalArgumentException + ); + assert!(jni_err.message.contains("my_dataset")); + } + + #[test] + fn test_dataset_already_exists_maps_to_illegal_argument() { + let lance_err = LanceError::dataset_already_exists("my_dataset"); + let jni_err: Error = lance_err.into(); + assert_eq!( + *java_class(&jni_err), + JavaExceptionClass::IllegalArgumentException + ); + assert!(jni_err.message.contains("my_dataset")); + } + + #[test] + fn test_commit_conflict_maps_to_illegal_argument() { + let lance_err = LanceError::commit_conflict_source(42, "conflict".to_string().into()); + let jni_err: Error = lance_err.into(); + assert_eq!( + *java_class(&jni_err), + JavaExceptionClass::IllegalArgumentException + ); + } + + #[test] + fn test_io_maps_to_ioexception() { + let lance_err = LanceError::io("disk failure"); + let jni_err: Error = lance_err.into(); + assert_eq!(*java_class(&jni_err), JavaExceptionClass::IOException); + assert!(jni_err.message.contains("disk failure")); + } + + #[test] + fn test_not_supported_maps_to_unsupported() { + let lance_err = LanceError::not_supported("nope"); + let jni_err: Error = lance_err.into(); + assert_eq!( + *java_class(&jni_err), + JavaExceptionClass::UnsupportedOperationException + ); + assert!(jni_err.message.contains("nope")); + } + + #[test] + fn test_not_found_maps_to_ioexception() { + let lance_err = LanceError::not_found("missing_uri"); + let jni_err: Error = lance_err.into(); + assert_eq!(*java_class(&jni_err), JavaExceptionClass::IOException); + assert!(jni_err.message.contains("missing_uri")); + } + + #[test] + fn test_fallthrough_maps_to_runtime() { + let lance_err = LanceError::internal("internal oops"); + let jni_err: Error = lance_err.into(); + assert_eq!(*java_class(&jni_err), JavaExceptionClass::RuntimeException); + assert!(jni_err.message.contains("internal oops")); + } + + #[test] + fn test_no_backtrace_suffix_when_backtrace_is_none() { + // Without the backtrace feature enabled in lance-core default tests, + // backtrace() returns None, so no suffix should be appended. + let lance_err = LanceError::io("clean message"); + let jni_err: Error = lance_err.into(); + assert!( + !jni_err.message.contains("Rust backtrace:"), + "Expected no backtrace suffix, got: {}", + jni_err.message + ); + } +} diff --git a/java/lance-jni/src/fragment.rs b/java/lance-jni/src/fragment.rs index d6603925947..59ce5e553ff 100644 --- a/java/lance-jni/src/fragment.rs +++ b/java/lance-jni/src/fragment.rs @@ -828,6 +828,9 @@ impl FromJObjectWithEnv for JObject<'_> { row_id_meta, created_at_version_meta, last_updated_at_version_meta, + // Overlays are not exposed to Java yet, and the reverse conversion + // does not export them, so this round-trip is overlay-free. + overlays: vec![], }) } } diff --git a/java/lance-jni/src/mem_wal.rs b/java/lance-jni/src/mem_wal.rs index 20404b6a88b..c4f56d7b97d 100644 --- a/java/lance-jni/src/mem_wal.rs +++ b/java/lance-jni/src/mem_wal.rs @@ -27,7 +27,7 @@ use jni::sys::{jdouble, jint, jlong}; use lance::dataset::Dataset as LanceDataset; use lance::dataset::mem_wal::scanner::{ FlushedGeneration, LsmDataSourceCollector, LsmPointLookupPlanner, LsmVectorSearchPlanner, - write_pk_sidecar, + parse_filter_expr as parse_lsm_filter_expr, write_pk_sidecar, }; use lance::dataset::mem_wal::write::{MemTableStats, WriteStatsSnapshot}; use lance::dataset::mem_wal::{ @@ -181,6 +181,29 @@ fn inner_put(env: &mut JNIEnv, this: JObject, stream_addr: jlong) -> Result<()> Ok(()) } +#[unsafe(no_mangle)] +pub extern "system" fn Java_org_lance_memwal_ShardWriter_nativeDelete( + mut env: JNIEnv, + this: JObject, + stream_addr: jlong, +) { + ok_or_throw_without_return!(env, inner_delete(&mut env, this, stream_addr)); +} + +fn inner_delete(env: &mut JNIEnv, this: JObject, stream_addr: jlong) -> Result<()> { + let stream_ptr = stream_addr as *mut FFI_ArrowArrayStream; + let reader = unsafe { ArrowArrayStreamReader::from_raw(stream_ptr) }?; + let batches: Vec = reader.collect::>()?; + if batches.is_empty() { + return Ok(()); + } + + let guard = + unsafe { env.get_rust_field::<_, _, BlockingShardWriter>(&this, NATIVE_SHARD_WRITER) }?; + RT.block_on(guard.writer.delete(batches))?; + Ok(()) +} + /// Test-support: write a primary-key dedup sidecar (`_pk_index/`) for a /// flushed-generation dataset already staged at `gen_path`, mirroring what /// production flush emits. Lets Java tests stage a *faithful* flushed @@ -398,7 +421,7 @@ fn inner_scanner_project(env: &mut JNIEnv, this: JObject, columns: JObject) -> R let columns = env.get_strings(&columns)?; with_lsm_scanner(env, &this, |scanner| { let cols: Vec<&str> = columns.iter().map(String::as_str).collect(); - Ok(scanner.project(&cols)) + Ok(scanner.project(&cols)?) }) } @@ -420,7 +443,7 @@ fn inner_scanner_filter(env: &mut JNIEnv, this: JObject, expr: JString) -> Resul pub extern "system" fn Java_org_lance_memwal_LsmScanner_nativeLimit( mut env: JNIEnv, this: JObject, - limit: jlong, + limit: JObject, offset: JObject, ) { ok_or_throw_without_return!(env, inner_scanner_limit(&mut env, this, limit, offset)); @@ -429,13 +452,12 @@ pub extern "system" fn Java_org_lance_memwal_LsmScanner_nativeLimit( fn inner_scanner_limit( env: &mut JNIEnv, this: JObject, - limit: jlong, + limit: JObject, offset: JObject, ) -> Result<()> { - let offset = env.get_u64_opt(&offset)?.map(|v| v as usize); - with_lsm_scanner(env, &this, |scanner| { - Ok(scanner.limit(limit as usize, offset)) - }) + let limit = env.get_u64_opt(&limit)?.map(|v| v as i64); + let offset = env.get_u64_opt(&offset)?.map(|v| v as i64); + with_lsm_scanner(env, &this, |scanner| Ok(scanner.limit(limit, offset)?)) } #[unsafe(no_mangle)] @@ -825,6 +847,7 @@ pub extern "system" fn Java_org_lance_memwal_LsmVectorSearchPlanner_nativeCreate vector_column: JString, pk_columns: JObject, distance_type: JObject, + filter: JObject, ) { ok_or_throw_without_return!( env, @@ -836,6 +859,7 @@ pub extern "system" fn Java_org_lance_memwal_LsmVectorSearchPlanner_nativeCreate vector_column, pk_columns, distance_type, + filter, ) ); } @@ -849,11 +873,13 @@ fn inner_create_vector_planner( vector_column: JString, pk_columns: JObject, distance_type: JObject, + filter: JObject, ) -> Result<()> { let snapshots = read_shard_snapshots(env, &shard_snapshots)?; let vector_column: String = env.get_string(&vector_column)?.into(); let pk_columns = env.get_strings_opt(&pk_columns)?; let distance_type = env.get_string_opt(&distance_type)?; + let filter = env.get_string_opt(&filter)?; let dataset = { let guard = unsafe { env.get_rust_field::<_, _, BlockingDataset>(&dataset, NATIVE_DATASET) }?; @@ -866,9 +892,13 @@ fn inner_create_vector_planner( let base_schema = Arc::new(ArrowSchema::from(dataset.schema())); let dist_type = parse_distance_type(distance_type.as_deref().unwrap_or("l2"))?; let vector_dim = get_vector_dim(&dataset, &vector_column)?; + let filter = filter + .as_deref() + .map(|filter| parse_lsm_filter_expr(base_schema.as_ref(), filter).map_err(Error::from)) + .transpose()?; let collector = LsmDataSourceCollector::new(dataset.clone(), snapshots); - let planner = LsmVectorSearchPlanner::new( + let mut planner = LsmVectorSearchPlanner::new( collector, pk_columns, base_schema.clone(), @@ -876,6 +906,9 @@ fn inner_create_vector_planner( dist_type, ) .with_dataset(dataset); + if let Some(filter) = filter { + planner = planner.with_filter(Some(filter)); + } let blocking = BlockingLsmVectorSearchPlanner { planner, diff --git a/java/pom.xml b/java/pom.xml index bac93e26278..a074013ee91 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -7,7 +7,7 @@ org.lance lance-core Lance Core - 9.0.0-beta.8 + 9.0.0-beta.21 jar Lance Format Java API @@ -39,6 +39,7 @@ 3.7.5 package false + false org.lance.shaded @@ -396,6 +397,9 @@ lance-jni ${rust.release.build} + + ${rust.features} + ${project.build.directory}/classes/nativelib true @@ -409,6 +413,9 @@ lance-jni ${rust.release.build} + + ${rust.features} + -v diff --git a/java/src/main/java/org/lance/ReadOptions.java b/java/src/main/java/org/lance/ReadOptions.java index a2e60e0e75b..5d6447539f0 100644 --- a/java/src/main/java/org/lance/ReadOptions.java +++ b/java/src/main/java/org/lance/ReadOptions.java @@ -189,6 +189,13 @@ public Builder setMetadataCacheSize(int metadataCacheSize) { * Set storage options. Extra options that make sense for a particular storage connection. This * is used to store connection parameters like credentials, endpoint, etc. * + *

For datasets with additional registered base paths, a key of the form {@code + * base_.} applies {@code } only to the base path with that manifest id, + * overriding the unscoped options that every base inherits. For example {@code + * base_1.account_key = abc} makes base 1 use {@code account_key = abc} while all other options + * are shared. Exact per-base bindings set via {@link #setBaseStoreParams(Map)} take precedence + * over base-scoped keys. + * * @param storageOptions the storage options * @return this builder */ diff --git a/java/src/main/java/org/lance/WriteDatasetBuilder.java b/java/src/main/java/org/lance/WriteDatasetBuilder.java index 96a894fc944..9d406406291 100644 --- a/java/src/main/java/org/lance/WriteDatasetBuilder.java +++ b/java/src/main/java/org/lance/WriteDatasetBuilder.java @@ -70,6 +70,7 @@ public class WriteDatasetBuilder { private WriteParams.WriteMode mode = WriteParams.WriteMode.CREATE; private Schema schema; private Map storageOptions = new HashMap<>(); + private Map properties = new HashMap<>(); private Map> baseStoreParams = new HashMap<>(); private boolean ignoreNamespaceStorageOptions = false; private Optional maxRowsPerFile = Optional.empty(); @@ -207,6 +208,27 @@ public WriteDatasetBuilder storageOptions(Map storageOptions) { return this; } + /** + * Sets the table properties to forward to the namespace on table creation. + * + *

These are Lance-namespace properties: catalog-level key-value metadata stored by the + * namespace outside the Lance table (available even if the table manifest does not exist), as + * distinct from the manifest-stored {@code config} (read/write behavior) and {@code metadata} + * (business metadata), and from non-persisted {@code storageOptions}. They are attached to the + * underlying declareTable request via its {@code properties} field. + * + *

Only used when a namespace client is configured via namespaceClient()+tableId() and the + * write creates the table (CREATE mode). Ignored for direct-URI writes and for APPEND/OVERWRITE + * modes. + * + * @param properties Table properties to forward on declareTable + * @return this builder instance + */ + public WriteDatasetBuilder properties(Map properties) { + this.properties = new HashMap<>(properties); + return this; + } + /** * Sets runtime-only object store parameters for registered base paths. * @@ -410,6 +432,9 @@ private Dataset executeWithNamespaceClient() { if (mode == WriteParams.WriteMode.CREATE) { DeclareTableRequest declareRequest = new DeclareTableRequest(); declareRequest.setId(tableId); + if (properties != null && !properties.isEmpty()) { + declareRequest.setProperties(properties); + } DeclareTableResponse declareResponse = namespaceClient.declareTable(declareRequest); tableUri = declareResponse.getLocation(); diff --git a/java/src/main/java/org/lance/WriteParams.java b/java/src/main/java/org/lance/WriteParams.java index d8757006975..90d48feafb7 100644 --- a/java/src/main/java/org/lance/WriteParams.java +++ b/java/src/main/java/org/lance/WriteParams.java @@ -197,6 +197,17 @@ public Builder withDataStorageVersion(String dataStorageVersion) { return this; } + /** + * Set storage options for the write. + * + *

For writes involving additional registered base paths, a key of the form {@code + * base_.} applies {@code } only to the base path with that id, overriding the + * unscoped options that every base inherits. Exact per-base bindings set via {@link + * #withBaseStoreParams(Map)} take precedence over base-scoped keys. + * + * @param storageOptions the storage options + * @return this builder + */ public Builder withStorageOptions(Map storageOptions) { this.storageOptions = storageOptions; return this; diff --git a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java index f509efd922b..82513a89ee7 100755 --- a/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java +++ b/java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java @@ -53,7 +53,9 @@ public static final class Builder { private Integer minNgramLength; private Integer maxNgramLength; private Boolean prefixOnly; + private Integer blockSize = 128; private Boolean skipMerge; + private Integer formatVersion; /** * Configure the base tokenizer. @@ -225,6 +227,27 @@ public Builder prefixOnly(boolean prefixOnly) { return this; } + /** + * Configure the number of documents in each compressed posting block. + * + *

Supported values are {@code 128} and {@code 256}. New indexes default to {@code 128} when + * this is not set. + * + *

{@code blockSize = 256} is experimental and may introduce breaking changes. Use {@code + * 128} when stable compatibility with the legacy posting layout is required. + * + * @param blockSize posting block size + * @return this builder + * @throws IllegalArgumentException if {@code blockSize} is unsupported + */ + public Builder blockSize(int blockSize) { + if (blockSize != 128 && blockSize != 256) { + throw new IllegalArgumentException("blockSize must be one of 128 or 256"); + } + this.blockSize = blockSize; + return this; + } + /** * Configure whether to skip the partition merge stage after indexing. If true, skip the * partition merge stage after indexing. This can be useful for distributed indexing where merge @@ -238,8 +261,31 @@ public Builder skipMerge(boolean skipMerge) { return this; } + /** + * Configure the on-disk FTS format version to write when creating a new index. + * + *

If unset, Lance writes v2 for {@code blockSize = 128} and v3 for {@code blockSize = 256}. + * {@code formatVersion = 3} is experimental and is only valid with {@code blockSize = 256}. + * + * @param formatVersion FTS format version, must be 1, 2, or 3 + * @return this builder + * @throws IllegalArgumentException + */ + public Builder formatVersion(int formatVersion) { + if (formatVersion != 1 && formatVersion != 2 && formatVersion != 3) { + throw new IllegalArgumentException("formatVersion must be 1, 2, or 3"); + } + this.formatVersion = formatVersion; + return this; + } + /** Build a {@link ScalarIndexParams} instance for an inverted index. */ public ScalarIndexParams build() { + if (formatVersion != null) { + Preconditions.checkArgument( + (blockSize == 256 && formatVersion == 3) || (blockSize == 128 && formatVersion != 3), + "formatVersion 3 requires blockSize 256, and blockSize 256 requires formatVersion 3"); + } Map params = new HashMap<>(); if (baseTokenizer != null) { params.put("base_tokenizer", baseTokenizer); @@ -282,9 +328,15 @@ public ScalarIndexParams build() { if (prefixOnly != null) { params.put("prefix_only", prefixOnly); } + if (blockSize != null) { + params.put("block_size", blockSize); + } if (skipMerge != null) { params.put("skip_merge", skipMerge); } + if (formatVersion != null) { + params.put("format_version", formatVersion); + } String json = JsonUtils.toJson(params); return ScalarIndexParams.create(INDEX_TYPE, json); diff --git a/java/src/main/java/org/lance/ipc/ScanOptions.java b/java/src/main/java/org/lance/ipc/ScanOptions.java index a9aad590c2b..7e1b4222262 100644 --- a/java/src/main/java/org/lance/ipc/ScanOptions.java +++ b/java/src/main/java/org/lance/ipc/ScanOptions.java @@ -134,6 +134,8 @@ public ScanOptions( Preconditions.checkArgument( !(filter.isPresent() && substraitFilter.isPresent()), "cannot set both substrait filter and string filter"); + Preconditions.checkArgument( + batchReadahead > 0, "batchReadahead must be greater than 0, got %s", batchReadahead); this.fragmentIds = fragmentIds; this.batchSize = batchSize; this.columns = columns; diff --git a/java/src/main/java/org/lance/memwal/LsmScanner.java b/java/src/main/java/org/lance/memwal/LsmScanner.java index ddbed76da8e..29bfb097375 100644 --- a/java/src/main/java/org/lance/memwal/LsmScanner.java +++ b/java/src/main/java/org/lance/memwal/LsmScanner.java @@ -101,9 +101,11 @@ public LsmScanner filter(String expr) { private native void nativeFilter(String expr); /** Limit the number of rows returned, optionally skipping {@code offset} rows first. */ - public LsmScanner limit(long limit, Optional offset) { + public LsmScanner limit(Optional limit, Optional offset) { + Preconditions.checkNotNull(limit, "limit must not be null"); Preconditions.checkNotNull(offset, "offset must not be null"); - Preconditions.checkArgument(limit >= 0, "limit must not be negative, got %s", limit); + limit.ifPresent( + l -> Preconditions.checkArgument(l >= 0, "limit must not be negative, got %s", l)); offset.ifPresent( o -> Preconditions.checkArgument(o >= 0, "offset must not be negative, got %s", o)); try (LockManager.WriteLock writeLock = lockManager.acquireWriteLock()) { @@ -113,12 +115,17 @@ public LsmScanner limit(long limit, Optional offset) { } } + /** Limit the number of rows returned, optionally skipping {@code offset} rows first. */ + public LsmScanner limit(long limit, Optional offset) { + return limit(Optional.of(limit), offset); + } + /** Limit the number of rows returned. */ public LsmScanner limit(long limit) { return limit(limit, Optional.empty()); } - private native void nativeLimit(long limit, Optional offset); + private native void nativeLimit(Optional limit, Optional offset); /** Include the {@code _rowaddr} internal column in results. */ public LsmScanner withRowAddress() { diff --git a/java/src/main/java/org/lance/memwal/LsmVectorSearchPlanner.java b/java/src/main/java/org/lance/memwal/LsmVectorSearchPlanner.java index 8517d2b2498..75c0d86c475 100644 --- a/java/src/main/java/org/lance/memwal/LsmVectorSearchPlanner.java +++ b/java/src/main/java/org/lance/memwal/LsmVectorSearchPlanner.java @@ -49,7 +49,7 @@ public class LsmVectorSearchPlanner implements AutoCloseable { */ public LsmVectorSearchPlanner( Dataset dataset, List shardSnapshots, String vectorColumn) { - this(dataset, shardSnapshots, vectorColumn, null, null); + this(dataset, shardSnapshots, vectorColumn, null, null, null); } /** @@ -66,6 +66,26 @@ public LsmVectorSearchPlanner( String vectorColumn, List pkColumns, String distanceType) { + this(dataset, shardSnapshots, vectorColumn, pkColumns, distanceType, null); + } + + /** + * @param dataset the base dataset + * @param shardSnapshots shard snapshots specifying the flushed generations to include + * @param vectorColumn name of the {@code FixedSizeList} vector column + * @param pkColumns primary key column names; inferred from schema metadata when {@code null} + * @param distanceType distance metric, one of {@code "l2"}, {@code "cosine"}, {@code "dot"}, + * {@code "hamming"}; defaults to {@code "l2"} when {@code null} + * @param filter SQL predicate applied as a prefilter to every LSM source before vector candidate + * selection; pass {@code null} for no predicate + */ + public LsmVectorSearchPlanner( + Dataset dataset, + List shardSnapshots, + String vectorColumn, + List pkColumns, + String distanceType, + String filter) { Preconditions.checkNotNull(dataset, "dataset must not be null"); Preconditions.checkNotNull(shardSnapshots, "shardSnapshots must not be null"); Preconditions.checkNotNull(vectorColumn, "vectorColumn must not be null"); @@ -75,7 +95,8 @@ public LsmVectorSearchPlanner( shardSnapshots, vectorColumn, Optional.ofNullable(pkColumns), - Optional.ofNullable(distanceType)); + Optional.ofNullable(distanceType), + Optional.ofNullable(filter)); } private native void nativeCreate( @@ -83,7 +104,8 @@ private native void nativeCreate( List shardSnapshots, String vectorColumn, Optional> pkColumns, - Optional distanceType); + Optional distanceType, + Optional filter); /** * Plan a KNN vector search. @@ -96,17 +118,17 @@ private native void nativeCreate( * @param refineBaseTable when true, the base-table arm re-ranks candidates with exact distances * (refine factor 1). Useful when the base table uses an approximate index (IVF-PQ). * Auto-enabled whenever stale filtering is on (see {@code overfetchFactor}). - * @param overfetchFactor single knob controlling both stale-read filtering and over-fetch: + * @param overfetchFactor over-fetch multiple for sources with rows superseded by newer + * generations. Must be at least {@code 1.0}: *

    - *
  • {@code < 1.0} (e.g. {@code 0.0}): stale filtering off — rows superseded by a newer - * generation may surface (the global primary-key dedup still runs). - *
  • {@code == 1.0}: filtering on, no over-fetch — a source with superseded rows fetches - * exactly {@code k} candidates and may return fewer than {@code k} live rows. - *
  • {@code > 1.0}: filtering on, with over-fetch — such a source fetches {@code ceil(k * - * overfetchFactor)} candidates so dropping the stale ones still leaves {@code k} live - * rows. + *
  • {@code == 1.0}: no over-fetch — an affected source fetches exactly {@code k} + * candidates and may return fewer than {@code k} live rows after stale rows are + * dropped. + *
  • {@code > 1.0}: with over-fetch — such a source fetches {@code ceil(k * + * overfetchFactor)} candidates so dropping stale rows is less likely to under-fill the + * result. *
- * There is no separate on/off flag: over-fetch is only meaningful while filtering. + * * @return an executable plan */ public ExecutionPlan planSearch( @@ -119,6 +141,10 @@ public ExecutionPlan planSearch( Preconditions.checkNotNull(query, "query must not be null"); Preconditions.checkArgument(k > 0, "k must be positive, got %s", k); Preconditions.checkArgument(nprobes > 0, "nprobes must be positive, got %s", nprobes); + Preconditions.checkArgument( + Double.isFinite(overfetchFactor) && overfetchFactor >= 1.0, + "overfetchFactor must be finite and >= 1.0, got %s", + overfetchFactor); try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) { Preconditions.checkArgument( nativeVectorPlannerHandle != 0, "LsmVectorSearchPlanner is closed"); diff --git a/java/src/main/java/org/lance/memwal/ShardWriter.java b/java/src/main/java/org/lance/memwal/ShardWriter.java index a5a3000bdba..3414ebda616 100644 --- a/java/src/main/java/org/lance/memwal/ShardWriter.java +++ b/java/src/main/java/org/lance/memwal/ShardWriter.java @@ -36,6 +36,7 @@ * 
{@code
  * try (ShardWriter writer = dataset.memWalWriter(shardId)) {
  *   writer.put(reader);
+ *   writer.delete(keys);
  * }
  * }
* @@ -99,6 +100,34 @@ public void put(ArrowReader reader) { private native void nativePut(long streamAddress); + /** + * Delete rows from the MemWAL by primary key. + * + *

Each batch in {@code reader} must carry this shard's primary key column(s); other columns + * are ignored. Lance builds a tombstone row per key — the primary key plus {@code _tombstone = + * true} and null in every other column — and appends it like an ordinary write. The tombstone is + * the newest value for its key: it wins newest-per-PK resolution (suppressing the older real row) + * and is then dropped from query results. + * + *

Only supported in memtable mode. Because a tombstone nulls every non-PK column, those + * columns must be nullable in the base schema; deleting against a schema with a non-nullable + * non-PK column errors. Deleting on a shard with no primary key columns also errors. + * + * @param reader the keys to delete; consumed fully by this call + */ + public void delete(ArrowReader reader) { + Preconditions.checkNotNull(reader, "reader must not be null"); + try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) { + Preconditions.checkArgument(nativeShardWriterHandle != 0, "ShardWriter is closed"); + try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { + Data.exportArrayStream(allocator, reader, stream); + nativeDelete(stream.memoryAddress()); + } + } + } + + private native void nativeDelete(long streamAddress); + /** Return a snapshot of cumulative write statistics. */ public WriteStats stats() { try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) { diff --git a/java/src/test/java/org/lance/FileReaderWriterTest.java b/java/src/test/java/org/lance/FileReaderWriterTest.java index a849a87c576..85c7430f087 100644 --- a/java/src/test/java/org/lance/FileReaderWriterTest.java +++ b/java/src/test/java/org/lance/FileReaderWriterTest.java @@ -256,7 +256,9 @@ void testInvalidPath() { LanceFileReader.open("/tmp/does_not_exist.lance", allocator); fail("Expected LanceException to be thrown"); } catch (IOException e) { - assertTrue(e.getMessage().contains("Object at location /tmp/does_not_exist.lance not found")); + String message = e.getMessage(); + assertTrue(message.contains("/tmp/does_not_exist.lance")); + assertTrue(message.toLowerCase().contains("not found")); } try { LanceFileReader.open("", allocator); diff --git a/java/src/test/java/org/lance/ScannerTest.java b/java/src/test/java/org/lance/ScannerTest.java index 00434034b64..0b07026bc68 100644 --- a/java/src/test/java/org/lance/ScannerTest.java +++ b/java/src/test/java/org/lance/ScannerTest.java @@ -422,25 +422,32 @@ void testDatasetScannerBatchReadahead(@TempDir Path tempDir) throws Exception { TestUtils.SimpleTestDataset testDataset = new TestUtils.SimpleTestDataset(allocator, datasetPath); testDataset.createEmptyDataset().close(); - int totalRows = 1000; - int batchSize = 100; - int batchReadahead = 5; - try (Dataset dataset = testDataset.write(1, totalRows)) { + + int totalRows = 2000; + int maxRowsPerFile = 100; // ~20 fragments + List fragments = testDataset.createNewFragment(totalRows, maxRowsPerFile); + assertTrue(fragments.size() > 1, "expected multiple fragments, got " + fragments.size()); + + FragmentOperation.Append append = new FragmentOperation.Append(fragments); + try (Dataset dataset = Dataset.commit(allocator, datasetPath, append, Optional.of(1L))) { + int batchReadahead = 2; // far below the default (num compute CPUs) try (LanceScanner scanner = dataset.newScan( - new ScanOptions.Builder() - .batchSize(batchSize) - .batchReadahead(batchReadahead) - .build())) { - // This test is more about ensuring that the batchReadahead parameter is accepted - // and doesn't cause errors. The actual effect of batchReadahead might not be - // directly observable in this test. + new ScanOptions.Builder().batchSize(50).batchReadahead(batchReadahead).build())) { try (ArrowReader reader = scanner.scanBatches()) { int rowCount = 0; + long idSum = 0; while (reader.loadNextBatch()) { - rowCount += reader.getVectorSchemaRoot().getRowCount(); + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + IntVector ids = (IntVector) root.getVector("id"); + for (int i = 0; i < root.getRowCount(); i++) { + idSum += ids.get(i); + } + rowCount += root.getRowCount(); } assertEquals(totalRows, rowCount); + // ids are the contiguous range [0, totalRows) + assertEquals((long) totalRows * (totalRows - 1) / 2, idSum); } } } diff --git a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java index e5024a95c2a..0ccc429ec57 100644 --- a/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java +++ b/java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java @@ -13,19 +13,66 @@ */ package org.lance.index.scalar; +import org.lance.util.JsonUtils; + import org.junit.jupiter.api.Test; +import java.util.Map; + import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; -public class InvertedIndexParamsTest { +class InvertedIndexParamsTest { @Test - public void testIcuSplitTokenizerVariant() { + void testIcuSplitTokenizerVariant() { ScalarIndexParams params = InvertedIndexParams.builder().baseTokenizer("icu/split").build(); assertEquals("inverted", params.getIndexType()); String jsonParams = params.getJsonParams().orElseThrow(AssertionError::new); assertTrue(jsonParams.contains("\"base_tokenizer\":\"icu/split\"")); } + + @Test + void defaultBlockSizeIsSerialized() { + ScalarIndexParams params = InvertedIndexParams.builder().build(); + + Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); + assertEquals(128, ((Number) json.get("block_size")).intValue()); + } + + @Test + void blockSizeIsSerialized() { + ScalarIndexParams params = InvertedIndexParams.builder().blockSize(128).build(); + + assertEquals("inverted", params.getIndexType()); + Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); + assertEquals(128, ((Number) json.get("block_size")).intValue()); + } + + @Test + void invalidBlockSizeIsRejected() { + assertThrows( + IllegalArgumentException.class, () -> InvertedIndexParams.builder().blockSize(129)); + assertThrows( + IllegalArgumentException.class, () -> InvertedIndexParams.builder().blockSize(512)); + } + + @Test + void formatVersionThreeRequiresBlockSize256() { + ScalarIndexParams params = + InvertedIndexParams.builder().blockSize(256).formatVersion(3).build(); + + Map json = JsonUtils.fromJson(params.getJsonParams().orElseThrow()); + assertEquals(256, ((Number) json.get("block_size")).intValue()); + assertEquals(3, ((Number) json.get("format_version")).intValue()); + + assertThrows( + IllegalArgumentException.class, + () -> InvertedIndexParams.builder().formatVersion(3).build()); + assertThrows( + IllegalArgumentException.class, + () -> InvertedIndexParams.builder().blockSize(256).formatVersion(2).build()); + } } diff --git a/java/src/test/java/org/lance/memwal/MemWalTest.java b/java/src/test/java/org/lance/memwal/MemWalTest.java index 5af3bd3f474..57d0b5b81ad 100644 --- a/java/src/test/java/org/lance/memwal/MemWalTest.java +++ b/java/src/test/java/org/lance/memwal/MemWalTest.java @@ -97,6 +97,22 @@ private static VectorSchemaRoot lookupRoot(BufferAllocator allocator, long[] ids return root; } + /** Build a single-batch root carrying only the {@code id} primary key, for deletes. */ + private static VectorSchemaRoot keysRoot(BufferAllocator allocator, long[] ids) { + VectorSchemaRoot root = + VectorSchemaRoot.create( + new Schema( + Collections.singletonList(Field.nullable("id", new ArrowType.Int(64, true)))), + allocator); + BigIntVector idVector = (BigIntVector) root.getVector("id"); + idVector.allocateNew(ids.length); + for (int i = 0; i < ids.length; i++) { + idVector.set(i, ids[i]); + } + root.setRowCount(ids.length); + return root; + } + /** Build a single-batch append-only root without primary-key metadata. */ private static VectorSchemaRoot appendOnlyRoot( BufferAllocator allocator, long[] ids, String prefix) { @@ -382,6 +398,52 @@ void testShardWriterPutAndLsmScanner(@TempDir Path tempDir) throws Exception { } } + @Test + void testShardWriterDeleteMasksBaseRow(@TempDir Path tempDir) throws Exception { + String path = tempDir.resolve("base").toString(); + String shardId = UUID.randomUUID().toString(); + try (BufferAllocator allocator = new RootAllocator(); + Dataset dataset = writeLookupDataset(allocator, path, new long[] {1, 2, 3}, "base")) { + dataset.initializeMemWal(new InitializeMemWalParams()); + + ShardWriterConfig config = + new ShardWriterConfig() + .withDurableWrite(true) + .withSyncIndexedWrite(true) + .withMaxWalBufferSize(1) + .withMaxWalFlushIntervalMs(10); + + try (ShardWriter writer = dataset.memWalWriter(shardId, config)) { + try (VectorSchemaRoot root = lookupRoot(allocator, new long[] {4}, "writer"); + ArrowReader reader = toReader(allocator, root)) { + writer.put(reader); + } + try (VectorSchemaRoot keys = keysRoot(allocator, new long[] {2}); + ArrowReader reader = toReader(allocator, keys)) { + writer.delete(reader); + } + + Map byId = Collections.emptyMap(); + long deadline = System.currentTimeMillis() + 10_000; + while (System.currentTimeMillis() < deadline) { + try (LsmScanner scanner = writer.lsmScanner(); + ArrowReader reader = scanner.scanBatches()) { + byId = readByName(reader); + } + if (!byId.containsKey(2L) && "writer_4".equals(byId.get(4L))) { + break; + } + Thread.sleep(50); + } + + assertEquals("base_1", byId.get(1L)); + assertFalse(byId.containsKey(2L), "deleted base row should be masked by the tombstone"); + assertEquals("base_3", byId.get(3L)); + assertEquals("writer_4", byId.get(4L)); + } + } + } + @Test void testLsmScannerFromSnapshots(@TempDir Path tempDir) throws Exception { String basePath = tempDir.resolve("base").toString(); @@ -406,6 +468,14 @@ void testLsmScannerFromSnapshots(@TempDir Path tempDir) throws Exception { assertEquals("gen1_2", byId.get(2L), "Flushed generation must win over base"); assertEquals("base_3", byId.get(3L)); } + + try (LsmScanner scanner = + LsmScanner.fromSnapshots(dataset, Collections.singletonList(snapshot)) + .limit(Optional.empty(), Optional.of(1L)); + ArrowReader reader = scanner.scanBatches()) { + Map byId = readByName(reader); + assertEquals(2, byId.size(), "Offset-only LSM scan should not require a limit"); + } } } @@ -570,6 +640,40 @@ void testVectorSearch(@TempDir Path tempDir) throws Exception { double recall = (double) found / queryIds.length; assertTrue(recall >= 0.5, "vector search recall too low: " + recall); } + + try (LsmVectorSearchPlanner planner = + new LsmVectorSearchPlanner( + dataset, Collections.emptyList(), "vec", null, null, "id >= 200"); + Float4Vector query = new Float4Vector("q", allocator)) { + int filteredQueryId = 250; + query.allocateNew(VDIM); + for (int d = 0; d < VDIM; d++) { + query.set(d, (float) (filteredQueryId * VDIM + d)); + } + query.setValueCount(VDIM); + + int filteredRows = 0; + boolean foundFilteredNearest = false; + try (ExecutionPlan plan = + planner.planSearch(query, 20, 2, Collections.singletonList("id"), false); + ArrowReader reader = plan.toReader()) { + while (reader.loadNextBatch()) { + VectorSchemaRoot result = reader.getVectorSchemaRoot(); + IntVector ids = (IntVector) result.getVector("id"); + for (int i = 0; i < result.getRowCount(); i++) { + int id = ids.get(i); + assertTrue(id >= 200, "filtered vector search returned id " + id); + if (id == filteredQueryId) { + foundFilteredNearest = true; + } + filteredRows++; + } + } + } + assertTrue(filteredRows > 0, "filtered vector search should return rows"); + assertTrue( + foundFilteredNearest, "filtered vector search recall missed id " + filteredQueryId); + } } finally { dataset.close(); } diff --git a/protos/index_old.proto b/protos/index_old.proto index 601aa2681da..eb984d6fe29 100644 --- a/protos/index_old.proto +++ b/protos/index_old.proto @@ -39,4 +39,9 @@ message InvertedIndexDetails { uint32 min_ngram_length = 9; uint32 max_ngram_length = 10; bool prefix_only = 11; + // Number of documents per compressed posting block. An absent value means + // the index predates this field and must use the legacy block size of 128. + // A present value records the block size used by the index; 256 is only + // valid with format version 3. + optional uint32 block_size = 12; } diff --git a/protos/table.proto b/protos/table.proto index d298809d5d8..7d4fb19d50b 100644 --- a/protos/table.proto +++ b/protos/table.proto @@ -109,10 +109,15 @@ message Manifest { // should not attempt to read the dataset. // // Known flags: - // * 1: deletion files are present - // * 2: row ids are stable and stored as part of the fragment metadata. - // * 4: use v2 format (deprecated) - // * 8: table config is present + // * 1 << 0: deletion files are present + // * 1 << 1: row ids are stable and stored as part of the fragment metadata. + // * 1 << 2: use v2 format (deprecated) + // * 1 << 3: table config is present + // * 1 << 4: dataset uses multiple base paths + // * 1 << 5: transaction file writes are disabled + // * 1 << 6: data overlay files are present (see DataOverlayFile). Readers that do + // not understand overlays must refuse the dataset, since ignoring an overlay + // would silently return stale base values. uint64 reader_feature_flags = 9; // Feature flags for writers. @@ -311,6 +316,15 @@ message DataFragment { repeated DataFile files = 2; + // Optional overlay files for this fragment, which supply new values for a + // subset of cells without rewriting the base data files. This MUST be empty + // if the data overlay files feature flag (64) is not set in the manifest. + // + // Order is significant: a later entry is newer than an earlier one. When two + // overlays cover the same (offset, field) and share a `committed_version`, the + // later entry wins. See DataOverlayFile for the full resolution rules. + repeated DataOverlayFile overlays = 11; + // File that indicates which rows, if any, should be considered deleted. DeletionFile deletion_file = 3; @@ -433,6 +447,66 @@ message DataFile { optional uint32 base_id = 7; } // DataFile +// An overlay file supplies new values for a subset of (row offset, field) cells +// within a fragment, without rewriting the fragment's base data files. It is +// used for efficient updates when only a small fraction of rows and/or columns +// change. +// +// On read, a cell is resolved by consulting the fragment's overlays from newest +// to oldest: the first overlay that covers that (offset, field) wins; if none +// cover it, the value falls through to the base data file. Because deletions +// take precedence over overlays, an overlay value for an offset that is also +// marked deleted is dead and is ignored. +// +// The overlay's data file does NOT store a row-offset key column. Within a value +// column, the position of a covered offset's value is the rank (0-based count of +// set bits below it) of that offset within the field's coverage bitmap. Because +// fields may cover different offset sets, the value columns of a single overlay +// data file may have different lengths (which the Lance file format permits). +message DataOverlayFile { + // The data file storing the overlay's new cell values, one value column per + // field in `data_file.fields`. No row-offset key column is stored. + DataFile data_file = 1; + + // Which (offset, field) cells this overlay provides values for. + oneof coverage { + // A single 32-bit Roaring bitmap of physical row offsets that applies to + // every field in `data_file.fields` (a "dense" / rectangular overlay). + // Every covered offset has a value for every field. This is the common case + // for a plain UPDATE, where one SET list is applied to one set of rows. + bytes shared_offset_bitmap = 2; + // Per-field coverage for a "sparse" overlay, used when different fields cover + // different offset sets (e.g. a MERGE with multiple WHEN MATCHED branches). + FieldCoverage field_coverage = 4; + } + + // The dataset version at which this overlay became effective: the version of + // the commit that introduced it, NOT the version it was read from. It is + // stamped at commit time and re-stamped if the commit is retried, in the same + // way as the created-at / last-updated-at version sequences. + // + // This drives two orderings: + // * Versus index builds: an index whose `dataset_version` >= this value + // already incorporates this overlay. Otherwise the overlay's covered cells + // are excluded from index results for the affected fields and re-evaluated + // against their current values (see the Data Overlay Files specification). + // * Versus other overlays: when two overlays cover the same (offset, field), + // the one with the higher `committed_version` wins. Overlays that share a + // `committed_version` are ordered by their position in + // `DataFragment.overlays`, where a later entry is newer and wins. + uint64 committed_version = 3; +} + +// Per-field coverage for a sparse overlay. +message FieldCoverage { + // One entry per field in the overlay's `data_file.fields`, in the same order. + // Each is a 32-bit Roaring bitmap of the physical row offsets covered for that + // field. An offset present in a field's bitmap but mapped to a NULL value + // means the cell is overridden to NULL (distinct from an offset that is absent, + // which falls through to the base data file). + repeated bytes offset_bitmaps = 1; +} + // Deletion File // // The path of the deletion file is constructed as: @@ -521,6 +595,16 @@ message FragmentReuseIndexDetails { // MemWAL Index Types // ============================================================================ +// Lifecycle status of a WAL shard. Drives drop-table two-phase commit: +// a SEALED shard refuses new writer claims (reversible) until the drop +// commits (the shard dir is deleted) or rolls back (status -> ACTIVE). +enum ShardStatus { + // Normal: the shard accepts writer claims. + ACTIVE = 0; + // A drop is in flight: claims are refused. Reversible to ACTIVE. + SEALED = 1; +} + // Shard manifest containing epoch-based fencing and WAL state. // Each shard has exactly one active writer at any time. message ShardManifest { @@ -566,6 +650,10 @@ message ShardManifest { // List of flushed MemTable generations and their directory paths. repeated FlushedGeneration flushed_generations = 8; + + // Lifecycle status. Default ACTIVE; SEALED marks an in-flight drop + // (drop-table 2PC). A SEALED manifest refuses claims at claim_epoch. + ShardStatus status = 15; } // A shard field value stored as raw Arrow scalar bytes. diff --git a/protos/transaction.proto b/protos/transaction.proto index e72e95025a4..13d11915af4 100644 --- a/protos/transaction.proto +++ b/protos/transaction.proto @@ -315,6 +315,28 @@ message Transaction { repeated DataReplacementGroup replacements = 1; } + // Overlay files to append to a single fragment, in order (the last entry is + // newest). The overlays are appended to the fragment's existing `overlays` + // list; they do not replace it, so overlays written by concurrent commits are + // preserved. + message DataOverlayGroup { + uint64 fragment_id = 1; + // Each DataOverlayFile.committed_version is left 0 by the writer and stamped + // to the new dataset version at commit time (re-stamped on retry), in the + // same way as the created-at / last-updated-at version sequences. The fields + // touched are read from each overlay's `data_file.fields`. + repeated DataOverlayFile overlays = 2; + } + + // Attach overlay files to fragments, supplying new values for a subset of + // (row offset, field) cells without rewriting the fragments' base data files. + // See the DataOverlayFile message in table.proto for resolution, coverage, and + // versioning rules, and the Data Overlay Files and Transactions specifications + // for the (intentionally permissive) conflict semantics. + message DataOverlay { + repeated DataOverlayGroup groups = 1; + } + // Update the merged generations in MemWAL index. // This operation is used during merge-insert to atomically record which // generations have been merged to the base table. @@ -346,6 +368,7 @@ message Transaction { UpdateMemWalState update_mem_wal_state = 112; Clone clone = 113; UpdateBases update_bases = 114; + DataOverlay data_overlay = 115; } // Fields 200/202 (`blob_append` / `blob_overwrite`) previously represented blob dataset ops. diff --git a/python/AGENTS.md b/python/AGENTS.md index 5f5ee0134c4..2025f4ef9df 100644 --- a/python/AGENTS.md +++ b/python/AGENTS.md @@ -20,6 +20,16 @@ Also see [root AGENTS.md](../AGENTS.md) for cross-language standards. * Lint: `uv run make lint` * Format: `uv run make format` +## Beta Releases + +- Python RC and beta preview wheels are published on fury.io, not only PyPI. When a task needs a beta version such as `7.2.0b4`, use a disposable venv and install with `--pre` and the Lance fury indices: + ```shell + uv venv /path/to/venv + uv pip install --python /path/to/venv/bin/python --pre \ + --extra-index-url https://pypi.fury.io/lance-format/ \ + "pylance==7.2.0b4" + ``` + ## API Design - Keep bindings as thin wrappers — centralize validation and logic in Rust core. diff --git a/python/Cargo.lock b/python/Cargo.lock index 88bfda071ba..d380b266cc7 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -187,9 +187,9 @@ dependencies = [ [[package]] name = "arc-swap" -version = "1.9.1" +version = "1.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" +checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" dependencies = [ "rustversion", ] @@ -1048,15 +1048,6 @@ version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" -[[package]] -name = "bitpacking" -version = "0.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96a7139abd3d9cebf8cd6f920a389cf3dc9576172e32f4563f188cae3c3eb019" -dependencies = [ - "crunchy", -] - [[package]] name = "bitvec" version = "1.1.1" @@ -1200,6 +1191,20 @@ name = "bytemuck" version = "1.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] [[package]] name = "byteorder" @@ -1209,9 +1214,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.12.0" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "bytes-utils" @@ -1234,9 +1239,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.64" +version = "1.2.66" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dad887fd958be91b5098c0248def011f4523ab786cd411be668777e55063501f" +checksum = "f5d6cac793997bd970000024b2934968efe83b382de4fdcf4fcb46b6ee4ad996" dependencies = [ "find-msvc-tools", "jobserver", @@ -1592,18 +1597,18 @@ 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", ] [[package]] name = "crossbeam-queue" -version = "0.3.12" +version = "0.3.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115" +checksum = "803d13fb3b09d88be9f4dbc29062c66b19bf7170867ceb746d2a8689bf6c7a26" dependencies = [ "crossbeam-utils", ] @@ -2707,11 +2712,17 @@ dependencies = [ "encoding_rs", ] +[[package]] +name = "endian-type" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c34f04666d835ff5d62e058c3995147c06f42fe86ff053337632bca83e42702d" + [[package]] name = "env_filter" -version = "1.0.1" +version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32e90c2accc4b07a8456ea0debdc2e7587bdd890680d71173a15d4ae604f6eef" +checksum = "900d271a03799a1ee8d1ca9b19893b48ca674a9284fefcfb85f05e74ed314217" dependencies = [ "log", "regex", @@ -2719,9 +2730,9 @@ dependencies = [ [[package]] name = "env_logger" -version = "0.11.10" +version = "0.11.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0621c04f2196ac3f488dd583365b9c09be011a4ab8b9f37248ffcc8f6198b56a" +checksum = "de671bd27a75a797dc9ae289ba1e77276e75e2026408aab65185384e2d5cd3f6" dependencies = [ "anstream", "anstyle", @@ -2859,7 +2870,7 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "rand 0.9.4", @@ -3261,6 +3272,7 @@ version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ + "bytemuck", "cfg-if 1.0.4", "crunchy", "num-traits", @@ -3468,9 +3480,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[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" @@ -3829,9 +3841,9 @@ checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" [[package]] name = "io-uring" -version = "0.7.12" +version = "0.7.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d09b98f7eace8982db770e4408e7470b028ce513ac28fecdc6bf4c30fe92b62" +checksum = "9080b15e63775b9a2ac7dca720f7050a8b955e092ea0f6020a4a80f69998cdc0" dependencies = [ "bitflags 2.13.0", "cfg-if 1.0.4", @@ -3876,18 +3888,18 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "jieba-macros" -version = "0.10.1" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46adade69b634535a8f495cf87710ed893cff53e1dbc9dd750c2ab81c5defb82" +checksum = "38fc0f3831de71556de69643b80a08a5c8cd260a23c6b8dbeb7cd923c779cac5" dependencies = [ "phf_codegen", ] [[package]] name = "jieba-rs" -version = "0.10.1" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11b53580aaa8ec8b713da271da434f8947409242c537a9ab3f7b76bdbb19e8a9" +checksum = "a813bbf185c8c62eb6fcf54a223177b644824d91612045dfd80bb779acd080eb" dependencies = [ "bytecount", "cedarwood", @@ -4058,7 +4070,7 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arc-swap", "arrow", @@ -4076,7 +4088,6 @@ dependencies = [ "async_cell", "aws-credential-types", "aws-sdk-dynamodb", - "bitpacking", "byteorder", "bytes", "chrono", @@ -4095,6 +4106,7 @@ dependencies = [ "humantime", "itertools 0.14.0", "lance-arrow", + "lance-bitpacking", "lance-core", "lance-datafusion", "lance-encoding", @@ -4132,7 +4144,7 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-buffer", @@ -4141,6 +4153,7 @@ dependencies = [ "arrow-ord", "arrow-schema", "arrow-select", + "bytemuck", "bytes", "futures", "getrandom 0.2.17", @@ -4174,16 +4187,17 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrayref", + "crunchy", "paste", "seq-macro", ] [[package]] name = "lance-core" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-buffer", @@ -4221,7 +4235,7 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-array", @@ -4253,7 +4267,7 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-array", @@ -4270,7 +4284,7 @@ dependencies = [ [[package]] name = "lance-derive" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "proc-macro2", "quote", @@ -4279,7 +4293,7 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-arith", "arrow-array", @@ -4314,7 +4328,7 @@ dependencies = [ [[package]] name = "lance-file" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-arith", "arrow-array", @@ -4344,7 +4358,7 @@ dependencies = [ [[package]] name = "lance-geo" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "datafusion", "geo-traits", @@ -4358,7 +4372,7 @@ dependencies = [ [[package]] name = "lance-index" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arc-swap", "arrow", @@ -4370,7 +4384,6 @@ dependencies = [ "async-channel", "async-recursion", "async-trait", - "bitpacking", "bitvec", "bytes", "chrono", @@ -4391,6 +4404,7 @@ dependencies = [ "jsonb", "lance-arrow", "lance-arrow-stats", + "lance-bitpacking", "lance-core", "lance-datafusion", "lance-datagen", @@ -4427,7 +4441,7 @@ dependencies = [ [[package]] name = "lance-io" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-arith", @@ -4451,6 +4465,7 @@ dependencies = [ "lance-core", "lance-namespace", "log", + "metrics", "moka", "object_store", "object_store_opendal", @@ -4468,7 +4483,7 @@ dependencies = [ [[package]] name = "lance-linalg" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-buffer", @@ -4484,7 +4499,7 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "async-trait", @@ -4496,7 +4511,7 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-ipc", @@ -4545,7 +4560,7 @@ dependencies = [ [[package]] name = "lance-select" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow-array", "arrow-buffer", @@ -4560,7 +4575,7 @@ dependencies = [ [[package]] name = "lance-table" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "arrow", "arrow-array", @@ -4599,7 +4614,7 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "icu_segmenter", "jieba-rs", @@ -4805,9 +4820,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.32" +version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "953f07c43838f8e6f9758cab68bf5bed85465e7587ebe0b823f1bcd81978ad3a" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" [[package]] name = "loom" @@ -4928,6 +4943,36 @@ dependencies = [ "libc", ] +[[package]] +name = "metrics" +version = "0.24.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89550ee9f79e88fef3119de263694973a8adb26c21d75322164fb8c493039fe2" +dependencies = [ + "portable-atomic", + "rapidhash", +] + +[[package]] +name = "metrics-util" +version = "0.19.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8496cc523d1f94c1385dd8f0f0c2c480b2b8aeccb5b7e4485ad6365523ae376" +dependencies = [ + "aho-corasick", + "crossbeam-epoch", + "crossbeam-utils", + "hashbrown 0.15.5", + "indexmap 2.14.0", + "metrics", + "ordered-float 4.6.0", + "quanta", + "radix_trie", + "rand 0.9.4", + "rand_xoshiro", + "sketches-ddsketch", +] + [[package]] name = "mime" version = "0.3.17" @@ -5032,6 +5077,15 @@ dependencies = [ "rawpointer", ] +[[package]] +name = "nibble_vec" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a5d83df9f36fe23f0c3648c6bbb8b0298bb5f1939c8f2704431371f4b84d43" +dependencies = [ + "smallvec", +] + [[package]] name = "nom" version = "8.0.0" @@ -5555,6 +5609,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "ordered-float" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7bb71e1b3fa6ca1c61f383464aaf2bb0e2f8e772a1f01d486832464de363b951" +dependencies = [ + "num-traits", +] + [[package]] name = "ordered-float" version = "5.3.0" @@ -6037,7 +6100,7 @@ dependencies = [ [[package]] name = "pylance" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" dependencies = [ "alloc-stdlib", "arrow", @@ -6069,6 +6132,8 @@ dependencies = [ "lance-table", "libc", "log", + "metrics", + "metrics-util", "object_store", "prost", "prost-types", @@ -6154,6 +6219,21 @@ dependencies = [ "serde", ] +[[package]] +name = "quanta" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3ab5a9d756f0d97bdc89019bd2e4ea098cf9cde50ee7564dde6b81ccc8f06c7" +dependencies = [ + "crossbeam-utils", + "libc", + "once_cell", + "raw-cpuid", + "wasi 0.11.1+wasi-snapshot-preview1", + "web-sys", + "winapi", +] + [[package]] name = "quick-xml" version = "0.39.4" @@ -6257,6 +6337,16 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" +[[package]] +name = "radix_trie" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c069c179fcdc6a2fe24d8d18305cf085fdbd4f922c041943e203685d6a1c58fd" +dependencies = [ + "endian-type", + "nibble_vec", +] + [[package]] name = "rancor" version = "0.1.1" @@ -6366,6 +6456,24 @@ version = "1.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "973443cf09a9c8656b574a866ab68dfa19f0867d0340648c7d2f6a71b8a8ea68" +[[package]] +name = "rapidhash" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32b266a82f4aa99bb5c25e28d11cc44ace63d91adbcbcee4d323e2ae3d49ef37" +dependencies = [ + "rustversion", +] + +[[package]] +name = "raw-cpuid" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags 2.13.0", +] + [[package]] name = "rawpointer" version = "0.2.1" @@ -6860,9 +6968,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_version" @@ -6915,9 +7023,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", @@ -7392,6 +7500,12 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" +[[package]] +name = "sketches-ddsketch" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c6f73aeb92d671e0cc4dca167e59b2deb6387c375391bc99ee743f326994a2b" + [[package]] name = "slab" version = "0.4.12" @@ -8359,9 +8473,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.23.3" +version = "1.23.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "144d6b123cef80b301b8f72a9e2ca4370ddec21950d0a103dd22c437006d2db7" +checksum = "bf80a72845275afea99e7f2b434723d3bc7e38470fcd1c7ed39a599c73319a53" dependencies = [ "getrandom 0.4.3", "js-sys", @@ -9089,9 +9203,9 @@ checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" [[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 = "yoke" diff --git a/python/Cargo.toml b/python/Cargo.toml index f4164000fc5..370271b887d 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "pylance" -version = "9.0.0-beta.8" +version = "9.0.0-beta.21" edition = "2024" authors = ["Lance Devs "] license = "Apache-2.0" @@ -8,6 +8,12 @@ rust-version = "1.91" exclude = ["python/lance/conftest.py"] publish = false +# `python/` is excluded from the top-level workspace (see ../Cargo.toml) and keeps +# its own Cargo.lock. Declaring an empty workspace here makes this crate an explicit +# workspace root, so cargo resolves it standalone even when the checkout is nested +# inside another cargo project (e.g. a git worktree placed under the main checkout). +[workspace] + [lib] name = "lance" crate-type = ["cdylib"] @@ -38,6 +44,7 @@ lance = { path = "../rust/lance", features = [ "goosefs", "dynamodb", "substrait", + "metrics", ] } lance-arrow = { path = "../rust/lance-arrow" } lance-core = { path = "../rust/lance-core" } @@ -48,7 +55,7 @@ lance-index = { path = "../rust/lance-index", features = [ "tokenizer-lindera", "tokenizer-jieba", ] } -lance-io = { path = "../rust/lance-io" } +lance-io = { path = "../rust/lance-io", features = ["metrics"] } lance-linalg = { path = "../rust/lance-linalg" } lance-namespace = { path = "../rust/lance-namespace" } lance-namespace-impls = { path = "../rust/lance-namespace-impls", features = ["rest", "rest-adapter", "dir-goosefs"] } @@ -56,6 +63,8 @@ lance-table = { path = "../rust/lance-table" } lance-datafusion = { path = "../rust/lance-datafusion" } libc = "0.2.176" log = "0.4" +metrics = "0.24" +metrics-util = "0.19" prost = "0.14.1" prost-types = "0.14.1" pyo3 = { version = "0.28", features = [ diff --git a/python/pyproject.toml b/python/pyproject.toml index d863fe38517..4297143c5d1 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -57,13 +57,12 @@ tests = [ "polars[pyarrow,pandas]", "psutil", "pytest", - # Only test tensorflow on linux for now. We will deprecate tensorflow soon. - "tensorflow; sys_platform == 'linux'", "tqdm", "datafusion>=53,<54", ] dev = ["ruff==0.11.2", "pyright"] benchmarks = ["pytest-benchmark"] +otel = ["opentelemetry-api", "opentelemetry-sdk"] torch = ["torch>=2.0"] geo = [ "geoarrow-rust-core", @@ -81,9 +80,9 @@ tests = [ "polars[pyarrow,pandas]==1.34.0", "psutil==7.1.0", "pytest==8.4.2", - "tensorflow==2.20.0; sys_platform == 'linux'", "tqdm==4.67.1", "datafusion==53.0.0", + "opentelemetry-sdk==1.30.0", ] dev = [ "maturin==1.13.3", @@ -114,6 +113,8 @@ include = [ "python/lance/schema.py", "python/lance/file.py", "python/lance/util.py", + "python/lance/arrow.py", + "python/tests/test_arrow.py", ] # Dependencies like pyarrow make this difficult to enforce strictly. reportMissingTypeStubs = "warning" @@ -135,9 +136,6 @@ markers = [ filterwarnings = [ 'error::FutureWarning', 'error::DeprecationWarning', - # TensorFlow import can emit NumPy deprecation FutureWarnings in some environments. - # We keep FutureWarnings as errors generally, but ignore this known-noisy import-time warning. - 'ignore:.*`np\\.object` will be defined as the corresponding NumPy scalar\\..*:FutureWarning', # Boto3 'ignore:.*datetime\.datetime\.utcnow\(\) is deprecated.*:DeprecationWarning', # Hugging Face Hub calls this deprecated hf-xet API internally. @@ -151,9 +149,8 @@ filterwarnings = [ 'ignore:.*the load_module\(\) method is deprecated.*:DeprecationWarning', # Pytorch uses deprecated jit.script_method internally (torch/utils/mkldnn.py) 'ignore:.*torch\.jit\.script_method.*is deprecated.*:DeprecationWarning', + # Pytorch uses the same API internally on Python 3.14+ with a different warning message. + 'ignore:.*torch\.jit\.script_method.*is not supported in Python 3\.14\+.*:DeprecationWarning', # huggingface_hub still calls the deprecated hf_xet.download_files() during Xet downloads 'ignore:.*hf_xet\.download_files\(\) is deprecated.*:DeprecationWarning', - # TensorFlow/Keras import can emit NumPy deprecation FutureWarnings in some environments. - # Keep FutureWarnings as errors generally, but ignore this known-noisy import-time warning. - 'ignore:.*np\.object.*:FutureWarning', ] diff --git a/python/python/lance/__init__.py b/python/python/lance/__init__.py index 4338950c649..3479a177824 100644 --- a/python/python/lance/__init__.py +++ b/python/python/lance/__init__.py @@ -9,7 +9,18 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Union from . import io, log -from .blob import Blob, BlobArray, BlobColumn, BlobFile, blob_array, blob_field +from .blob import ( + Blob, + BlobArray, + BlobColumn, + BlobDescriptor, + BlobDescriptorArrayBuilder, + BlobFile, + DedicatedBlobWriter, + PackedBlobWriter, + blob_array, + blob_field, +) from .dataset import ( DataStatistics, FieldStatistics, @@ -36,6 +47,7 @@ ScanStatistics, bytes_read_counter, iops_counter, + simd_info, ) from .mem_wal import ( ExecutionPlan, @@ -72,6 +84,10 @@ "BlobArray", "BlobColumn", "BlobFile", + "DedicatedBlobWriter", + "BlobDescriptorArrayBuilder", + "PackedBlobWriter", + "BlobDescriptor", "blob_array", "blob_field", "CleanupCandidateFile", @@ -100,6 +116,7 @@ "json_to_schema", "schema_to_json", "set_logger", + "simd_info", "write_dataset", "FFILanceTableProvider", "IndexProgress", @@ -166,6 +183,13 @@ def dataset( storage_options : optional, dict Extra options that make sense for a particular storage connection. This is used to store connection parameters like credentials, endpoint, etc. + + For datasets with additional registered base paths, a key of the form + ``base_.`` applies ```` only to the base path with that + manifest id, overriding the unscoped options that every base inherits. + For example ``{"account_key": "shared", "base_1.account_key": "abc"}`` + makes base 1 use ``account_key = abc`` while all other options are + shared. default_scan_options : optional, dict Default scan options that are used when scanning the dataset. This accepts the same arguments described in :py:meth:`lance.LanceDataset.scanner`. The @@ -203,9 +227,10 @@ def dataset( base_store_params : dict of str to dict, optional Runtime-only object store parameters keyed by base path URI. Each key is a base path URI (e.g., "s3://bucket/path") and each value is a dict - of storage options (credentials, endpoint, etc.) for that base. When a base - has no explicit entry here, the top-level ``storage_options`` is - used as a fallback. + of storage options (credentials, endpoint, etc.) for that base. These + take precedence over ``base_.`` entries in ``storage_options``. + When a base has no explicit entry here, the top-level + ``storage_options`` is used as a fallback. Notes ----- diff --git a/python/python/lance/arrow.py b/python/python/lance/arrow.py index 54da15705f8..5e9f671ab80 100644 --- a/python/python/lance/arrow.py +++ b/python/python/lance/arrow.py @@ -4,10 +4,12 @@ """Extensions to PyArrows.""" import json +import typing from pathlib import Path from typing import Callable, Iterable, Optional, Union import pyarrow as pa +import pyarrow.compute as pc from ._arrow.bf16 import ( # noqa: F401 BFloat16, @@ -19,8 +21,10 @@ from .lance import bfloat16_array __all__ = [ + "BFloat16", "BFloat16Array", "BFloat16Type", + "PandasBFloat16Array", "bfloat16_array", "cast", "EncodedImageType", @@ -220,14 +224,15 @@ def from_uris( ['file::///tmp/1.png'] """ + storage: pa.Array if isinstance(uris, (pa.StringArray, pa.LargeStringArray)): - pass + storage = uris elif isinstance(uris, Iterable): - uris = pa.array((str(uri) for uri in uris), type=pa.string()) + storage = pa.array((str(uri) for uri in uris), type=pa.string()) else: raise TypeError("Cannot build a ImageURIArray from {}".format(type(uris))) - return cls.from_storage(ImageURIType(uris.type), uris) + return cls.from_storage(ImageURIType(storage.type), storage) def read_uris(self, storage_type=pa.binary()) -> "EncodedImageArray": """ @@ -268,7 +273,8 @@ def download(url): print("Failed to reach the server: ", e.reason) elif hasattr(e, "code"): print( - "The server could not fulfill the request. Error code: ", e.code + "The server could not fulfill the request. Error code: ", + getattr(e, "code"), ) images = [] @@ -277,7 +283,9 @@ def download(url): if parsed_uri.scheme in ("http", "https"): images.append(download(uri)) else: - filesystem, path = fs.FileSystem.from_uri(uri.as_py()) + filesystem, path = fs.FileSystem.from_uri( # pyright: ignore[reportPrivateImportUsage] + uri.as_py() + ) with filesystem.open_input_stream(path) as f: images.append(f.read()) @@ -297,21 +305,12 @@ def __repr__(self): def pillow_metadata_decoder(images): import io - from PIL import Image + from PIL import Image # pyright: ignore[reportMissingImports] img = Image.open(io.BytesIO(images[0].as_py())) return img - def tensorflow_metadata_decoder(images): - import tensorflow as tf - - img = tf.io.decode_image(images[0].as_py()) - return img - - decoders = ( - ("tensorflow", tensorflow_metadata_decoder), - ("PIL", pillow_metadata_decoder), - ) + decoders = (("PIL", pillow_metadata_decoder),) decoder = None for libname, metadata_decoder in decoders: @@ -343,7 +342,7 @@ def to_tensor( decoder : Callable[pa.binary()], optional A function that takes a binary array and returns a numpy.ndarray or pa.fixed_shape_tensor. If not provided, will attempt to use - tensorflow and then pillow decoder in that order. + pillow. Returns ------- @@ -365,27 +364,19 @@ def to_tensor( if not decoder: - def pillow_decoder(images): + def pillow_decoder(images) -> "np.ndarray": import io - from PIL import Image + from PIL import Image # pyright: ignore[reportMissingImports] return np.stack( - [Image.open(io.BytesIO(img)) for img in images.to_pylist()] + [ + np.asarray(Image.open(io.BytesIO(img))) + for img in images.to_pylist() + ] ) - def tensorflow_decoder(images): - import tensorflow as tf - - decoded_to_tensor = tuple( - tf.io.decode_image(img) for img in images.to_pylist() - ) - return tf.stack(decoded_to_tensor, axis=0).numpy() - - decoders = [ - ("tensorflow", tensorflow_decoder), - ("PIL", pillow_decoder), - ] + decoders = [("PIL", pillow_decoder)] for libname, decoder_function in decoders: try: __import__(libname) @@ -395,15 +386,16 @@ def tensorflow_decoder(images): pass else: raise ValueError( - "No image decoder available. Please either install one of " - "tensorflow, pillow, or pass a decoder argument." + "No image decoder available. Please install pillow or pass a " + "decoder argument." ) image_array = decoder(self.storage) if isinstance(image_array, pa.FixedShapeTensorType): - shape = image_array.shape - arrow_type = image_array.storage_type - tensor_array = image_array + tensor = typing.cast("pa.Array", image_array) + shape = tensor.shape + arrow_type = tensor.storage_type + tensor_array = tensor else: shape = image_array.shape[1:] arrow_type = pa.from_numpy_dtype(image_array.dtype) @@ -476,7 +468,7 @@ def to_encoded(self, encoder=None, storage_type=pa.binary()) -> "EncodedImageArr def pillow_encoder(x): import io - from PIL import Image + from PIL import Image # pyright: ignore[reportMissingImports] encoded_images = [] for y in x: @@ -485,19 +477,8 @@ def pillow_encoder(x): encoded_images.append(buf.getvalue()) return pa.array(encoded_images, type=storage_type) - def tensorflow_encoder(x): - import tensorflow as tf - - encoded_images = ( - tf.io.encode_png(y).numpy() for y in tf.convert_to_tensor(x) - ) - return pa.array(encoded_images, type=storage_type) - if not encoder: - encoders = ( - ("PIL", pillow_encoder), - ("tensorflow", tensorflow_encoder), - ) + encoders = (("PIL", pillow_encoder),) for libname, encoder_function in encoders: try: __import__(libname) @@ -507,8 +488,8 @@ def tensorflow_encoder(x): pass else: raise ValueError( - "No image encoder available. Please either install one of " - "tensorflow, pillow, or pass an encoder argument." + "No image encoder available. Please install pillow or pass an " + "encoder argument." ) return EncodedImageArray.from_storage( @@ -571,7 +552,8 @@ def cast( + f"got: {target_type}" ) np_arr = arr.to_numpy() - float_arr = np_arr.astype(target_type.to_pandas_dtype()) + float_type = typing.cast("pa.DataType", target_type) + float_arr = np_arr.astype(float_type.to_pandas_dtype()) return pa.array(float_arr) elif isinstance(target_type, BFloat16Type) or target_type in ["bfloat16", "bf16"]: if not pa.types.is_floating(arr.type): @@ -586,15 +568,16 @@ def cast( target_type ): # Casting fixed size list to fixed size list - if arr.type.list_size != target_type.list_size: + list_type = typing.cast("pa.DataType", target_type) + if arr.type.list_size != list_type.list_size: raise ValueError( "Only support casting fixed size list to fixed size list " f"with the same size, got: {arr.type} to {target_type}" ) - values = cast(arr.values, target_type.value_type) + values = cast(arr.values, list_type.value_type) return pa.FixedSizeListArray.from_arrays( - values=values, list_size=target_type.list_size + values=values, list_size=list_type.list_size ) # Fallback to normal cast. - return pa.compute.cast(arr, target_type, *args, **kwargs) + return pc.cast(arr, target_type, *args, **kwargs) diff --git a/python/python/lance/blob.py b/python/python/lance/blob.py index a87c9302736..6d1af6797dc 100644 --- a/python/python/lance/blob.py +++ b/python/python/lance/blob.py @@ -8,12 +8,29 @@ import pyarrow as pa -from .lance import LanceBlobFile +from .lance import ( + BlobDescriptor as BlobDescriptor, +) +from .lance import ( + BlobDescriptorArrayBuilder as BlobDescriptorArrayBuilder, +) +from .lance import ( + DedicatedBlobWriter as DedicatedBlobWriter, +) +from .lance import ( + LanceBlobFile, +) +from .lance import ( + PackedBlobWriter as PackedBlobWriter, +) _BLOB_INLINE_SIZE_THRESHOLD_META_KEY = b"lance-encoding:blob-inline-size-threshold" _BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY = ( b"lance-encoding:blob-dedicated-size-threshold" ) +_BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY = ( + b"lance-encoding:blob-pack-file-size-threshold" +) _MAX_RUST_USIZE = ctypes.c_size_t(-1).value @@ -217,6 +234,7 @@ def blob_field( nullable: bool = True, inline_size_threshold: Optional[int] = None, dedicated_size_threshold: Optional[int] = None, + pack_file_size_threshold: Optional[int] = None, ) -> pa.Field: """ Construct an Arrow field for a Lance blob column. @@ -234,14 +252,24 @@ def blob_field( Maximum payload size in bytes to store in packed blob storage before using dedicated blob storage. This threshold is checked before ``inline_size_threshold``. + pack_file_size_threshold : optional, int + Maximum size in bytes of a single packed blob sidecar (``.pack``) file. + Once a sidecar reaches this size a new one is started. """ _validate_threshold("inline_size_threshold", inline_size_threshold, allow_zero=True) _validate_threshold( "dedicated_size_threshold", dedicated_size_threshold, allow_zero=False ) + _validate_threshold( + "pack_file_size_threshold", pack_file_size_threshold, allow_zero=False + ) field = pa.field(name, BlobType(), nullable=nullable) - if inline_size_threshold is None and dedicated_size_threshold is None: + if ( + inline_size_threshold is None + and dedicated_size_threshold is None + and pack_file_size_threshold is None + ): return field metadata = dict(field.metadata or {}) @@ -253,6 +281,10 @@ def blob_field( metadata[_BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY] = str( dedicated_size_threshold ).encode() + if pack_file_size_threshold is not None: + metadata[_BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY] = str( + pack_file_size_threshold + ).encode() return field.with_metadata(metadata) @@ -273,9 +305,11 @@ class BlobColumn: file-like objects. This can be useful for working with medium-to-small binary objects that need - to interface with APIs that expect file-like objects. For very large binary - objects (4-8MB or more per value) you might be better off creating a blob column - and using :py:meth:`lance.Dataset.take_blobs` to access the blob data. + to interface with APIs that expect file-like objects. For very large binary + objects (4-8MB or more per value) you might be better off creating a blob + column. Use :py:meth:`lance.Dataset.read_blobs` when you need complete blob + bytes, or :py:meth:`lance.Dataset.take_blobs` when you need lazy file-like + access. """ def __init__(self, blob_column: Union[pa.Array, pa.ChunkedArray]): diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 833d0f1321d..a1ce77b82b8 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -75,7 +75,11 @@ from .udf import BatchUDF, normalize_transform from .udf import BatchUDFCheckpoint as BatchUDFCheckpoint from .udf import batch_udf as batch_udf -from .util import _target_partition_size_to_num_partitions, td_to_micros +from .util import ( + _normalize_index_segment_ids, + _target_partition_size_to_num_partitions, + td_to_micros, +) if TYPE_CHECKING: from pyarrow._compute import Expression @@ -536,6 +540,54 @@ def use_index(self, use_index: bool) -> "MergeInsertBuilder": """ return super(MergeInsertBuilder, self).use_index(use_index) + def target_bases(self, bases: List[str]) -> "MergeInsertBuilder": + """ + Write new fragments produced by this merge insert to these bases. + + Each entry references a base path registered in the dataset manifest, + by name or by path URI, like the ``target_bases`` parameter of + :func:`~lance.write_dataset`. An entry equal to the dataset's URI + includes the dataset's primary storage in the rotation, e.g. + ``[ds.uri, "base1", "base2"]`` spreads new data files across primary + storage and both bases. New data files are distributed across the + target bases round-robin. Data files that patch existing fragments + and deletion files are always written to the dataset's primary + storage. + + Parameters + ---------- + bases : List[str] + Base names or path URIs to write new data files to. + + Returns + ------- + MergeInsertBuilder + The builder instance for method chaining. + """ + return super(MergeInsertBuilder, self).target_bases(bases) + + def target_all_bases(self, include_primary: bool = True) -> "MergeInsertBuilder": + """ + Write new fragments to every base registered in the dataset manifest. + + The bases are resolved when the merge insert executes, so bases added + later are picked up automatically. When ``include_primary`` is True + (the default), the dataset's primary storage participates in the + round-robin rotation as the first slot. Cannot be combined with + :meth:`target_bases`. + + Parameters + ---------- + include_primary : bool, default True + Whether the dataset's primary storage is part of the rotation. + + Returns + ------- + MergeInsertBuilder + The builder instance for method chaining. + """ + return super(MergeInsertBuilder, self).target_all_bases(include_primary) + def explain_plan( self, schema: Optional[pa.Schema] = None, verbose: bool = False ) -> str: @@ -2122,6 +2174,10 @@ def take_blobs( this API allows you to open binary blob data as a regular Python file-like object. For more details, see :py:class:`lance.BlobFile`. + If you plan to read each selected blob completely with ``read()`` or + ``readall()``, use :py:meth:`read_blobs` instead. It materializes blob + payloads with Lance's planned batched reader. + Exactly one of ids, addresses, or indices must be specified. Parameters @@ -2170,7 +2226,8 @@ def read_blobs( Unlike :py:meth:`take_blobs`, which returns file-like :py:class:`lance.BlobFile` handles for random access, this API plans and executes batched reads and - returns materialized blob payloads. + returns materialized blob payloads. Use this API for training loaders, + batch preprocessing, and other workflows that need complete blob bytes. Exactly one of ids, addresses, or indices must be specified. @@ -3104,12 +3161,13 @@ def _prepare_scalar_index_request( and not pa.types.is_floating(field_type) and not pa.types.is_boolean(field_type) and not pa.types.is_string(field_type) + and not pa.types.is_large_string(field_type) and not pa.types.is_temporal(field_type) and not pa.types.is_fixed_size_binary(field_type) ): raise TypeError( - f"BTREE/BITMAP index column {column} must be int", - ", float, bool, str, fixed-size-binary, or temporal ", + f"BTREE/BITMAP/ZONEMAP index column {column} must be int", + ", float, bool, str, large_str, fixed-size-binary, or temporal", ) elif index_type == "LABEL_LIST": if not pa.types.is_list(field_type): @@ -3203,6 +3261,7 @@ def create_scalar_index( fragment_ids: Optional[List[int]] = None, index_uuid: Optional[str] = None, progress_callback: Optional[Callable[[IndexProgress], None]] = None, + format_version: Optional[Union[int, str]] = None, **kwargs, ): """Create a scalar index on a column. @@ -3308,6 +3367,13 @@ def create_scalar_index( progress_callback : callable, optional A callback that receives :class:`lance.progress.IndexProgress` events while the index is being built. + format_version: int or str, optional + This is for the ``INVERTED`` / ``FTS`` index. Explicit on-disk FTS + format version to write when creating a new index. Accepts ``1``, + ``2``, ``3``, ``"v1"``, ``"v2"``, or ``"v3"``. If unset, Lance + writes v2 for ``block_size=128`` and v3 for ``block_size=256``. + ``format_version=3`` is experimental and is only valid with + ``block_size=256``. with_position: bool, default False This is for the ``INVERTED`` index. If True, the index will store the @@ -3315,6 +3381,12 @@ def create_scalar_index( query. This will significantly increase the index size. It won't impact the performance of non-phrase queries even if it is set to True. + block_size: int, default 128 + This is for the ``INVERTED`` index. Number of documents per compressed + posting block. Must be one of ``128`` or ``256``. + ``block_size=256`` is experimental and may introduce breaking changes. + Use ``128`` when stable compatibility with the legacy posting layout is + required. memory_limit: int, optional This is for the ``INVERTED`` index. Total build-time memory limit in MiB. If set, Lance divides this budget evenly across the workers. If unset, @@ -3415,6 +3487,8 @@ def create_scalar_index( kwargs["index_uuid"] = index_uuid if progress_callback is not None: kwargs["progress_callback"] = progress_callback + if format_version is not None: + kwargs["format_version"] = format_version self._ds.create_index([column], index_type, name, replace, train, None, kwargs) @@ -3741,13 +3815,17 @@ def _create_index_impl( if _check_for_numpy(pq_codebook) and isinstance( pq_codebook, np.ndarray ): + num_bits = kwargs.get("num_bits", 8) + expected_centroids = 2**num_bits if ( len(pq_codebook.shape) != 3 or pq_codebook.shape[0] != num_sub_vectors - or pq_codebook.shape[1] != 256 + or pq_codebook.shape[1] != expected_centroids ): raise ValueError( - f"PQ codebook must be 3D array: (sub_vectors, 256, dim), " + "PQ codebook must be 3D array: " + f"(sub_vectors, {expected_centroids}, dim) " + f"for num_bits={num_bits}, " f"got {pq_codebook.shape}" ) if pq_codebook.dtype not in [np.float16, np.float32, np.float64]: @@ -3869,10 +3947,9 @@ def create_index( pq_codebook : optional, It can be :py:class:`np.ndarray`, :py:class:`pyarrow.FixedSizeListArray`, or :py:class:`pyarrow.FixedShapeTensorArray`. - A ``num_sub_vectors x (2 ^ nbits * dimensions // num_sub_vectors)`` - array of K-mean centroids for PQ codebook. - - Note: ``nbits`` is always 8 for now. + A ``num_sub_vectors x (2 ^ num_bits) x + (dimensions // num_sub_vectors)`` array of K-mean centroids for PQ + codebook. ``num_bits`` defaults to 8. If not provided, a new PQ model will be trained. num_sub_vectors : int, optional The number of sub-vectors for PQ (Product Quantization). @@ -4229,13 +4306,22 @@ def drop_index(self, name: str): """ return self._ds.drop_index(name) - def prewarm_index(self, name: str, *, with_position: bool = False): + def prewarm_index( + self, + name: str, + *, + with_position: bool = False, + index_segments: Optional[Iterable[Union[str, uuid.UUID]]] = None, + ): """ Prewarm an index - This will load the entire index into memory. This can help avoid cold start - issues with index queries. If the index does not fit in the index cache, then - this will result in wasted I/O. + By default, this will load the entire index into memory. This can help + avoid cold start issues with index queries. If the index does not fit in + the index cache, then this will result in wasted I/O. + + Use ``session().index_cache_size_bytes()`` before and after prewarm to + inspect how much the index cache grew. Parameters ---------- @@ -4245,8 +4331,16 @@ def prewarm_index(self, name: str, *, with_position: bool = False): This is only supported for ``INVERTED`` indices. If True, positions are also loaded into the cache during prewarm so phrase queries do not need a separate lazy positions read. - """ - return self._ds.prewarm_index(name, with_position=with_position) + index_segments: iterable of str or uuid.UUID, default None + If specified, prewarm only these physical index segment UUIDs from the + named logical index. Use :meth:`describe_indices` to inspect logical + indices and obtain segment UUIDs from ``IndexDescription.segments``. + """ + return self._ds.prewarm_index( + name, + with_position=with_position, + index_segments=_normalize_index_segment_ids(index_segments), + ) def merge_index_metadata( self, @@ -6102,18 +6196,14 @@ def io_buffer_size(self, io_buffer_size: int) -> ScannerBuilder: used by the scanner. If the buffer is full then the scanner will block until the buffer is processed. - Generally this should scale with the number of concurrent I/O threads. The - default is 2GiB which comfortably provides enough space for somewhere between - 32 and 256 concurrent I/O threads. + Generally this should scale with the number of concurrent I/O threads. If + unset, v2 scans choose a default based on the object store and + ``LANCE_DEFAULT_IO_BUFFER_SIZE`` can override that default. This value is not a hard cap on the amount of RAM the scanner will use. Some space is used for the compute (which can be controlled by the batch size) and Lance does not keep track of memory after it is returned to the user. - Currently, if there is a single batch of data which is larger than the io buffer - size then the scanner will deadlock. This is a known issue and will be fixed in - a future release. - This parameter is only used when reading v2 files """ self._io_buffer_size = io_buffer_size @@ -6121,10 +6211,12 @@ def io_buffer_size(self, io_buffer_size: int) -> ScannerBuilder: def batch_readahead(self, nbatches: Optional[int] = None) -> ScannerBuilder: """ - This parameter is ignored when reading v2 files + Set the maximum number of batches to decode concurrently. + + This parameter must be greater than zero. """ - if nbatches is not None and int(nbatches) < 0: - raise ValueError("batch_readahead must be non-negative") + if nbatches is not None and int(nbatches) <= 0: + raise ValueError("batch_readahead must be greater than 0") self._batch_readahead = nbatches return self @@ -6333,19 +6425,7 @@ def with_fragments( def with_index_segments( self, index_segments: Optional[Iterable[Union[str, uuid.UUID]]] ) -> ScannerBuilder: - if index_segments is not None: - segment_ids = [] - for segment_id in index_segments: - if isinstance(segment_id, (str, uuid.UUID)): - segment_ids.append(str(segment_id)) - else: - raise TypeError( - "index_segments must be an iterable of str or uuid.UUID. " - f"Got {type(segment_id)} instead." - ) - index_segments = segment_ids - - self._index_segments = index_segments + self._index_segments = _normalize_index_segment_ids(index_segments) return self def nearest( @@ -7175,6 +7255,7 @@ def write_dataset( transaction_properties: Optional[Dict[str, str]] = None, initial_bases: Optional[List[DatasetBasePath]] = None, target_bases: Optional[List[str]] = None, + target_all_bases: Optional[bool] = None, base_store_params: Optional[Dict[str, Dict[str, str]]] = None, external_blob_mode: Literal["reference", "ingest"] = "reference", allow_external_blob_outside_bases: bool = False, @@ -7222,6 +7303,13 @@ def write_dataset( storage_options : optional, dict Extra options that make sense for a particular storage connection. This is used to store connection parameters like credentials, endpoint, etc. + + For writes involving additional base paths, a key of the form + ``base_.`` applies ```` only to the base path with that + id, overriding the unscoped options that every base inherits. For + example ``{"account_key": "shared", "base_1.account_key": "abc"}`` + makes base 1 use ``account_key = abc`` while all other options are + shared. data_storage_version: optional, str, default None The version of the data storage format to use. Newer versions are more efficient but require newer versions of lance to read. The default (None) @@ -7273,12 +7361,19 @@ def write_dataset( **CREATE mode**: References must match bases in `initial_bases` **APPEND/OVERWRITE modes**: References must match bases in the existing manifest + target_all_bases: bool, optional + Write new data files round-robin across every base registered in the + manifest. When True (include primary), the dataset's primary storage + participates in the rotation as the first slot; when False, only the + registered bases are used. Cannot be combined with ``target_bases``. base_store_params : dict of str to dict, optional Runtime-only object store parameters keyed by base path URI. Each key is a base path URI (e.g., "s3://bucket/path") and each value is a dict of storage options (credentials, endpoint, etc.) for that base. These - are not persisted to the manifest. When a base has no explicit entry - here, the top-level ``storage_options`` is used as a fallback. + are not persisted to the manifest. These take precedence over + ``base_.`` entries in ``storage_options``. When a base has no + explicit entry here, the top-level ``storage_options`` is used as a + fallback. external_blob_mode: {"reference", "ingest"}, default "reference" How external blob URIs are handled on write. @@ -7421,6 +7516,7 @@ def write_dataset( "transaction_properties": merged_properties, "initial_bases": initial_bases, "target_bases": target_bases, + "target_all_bases": target_all_bases, "base_store_params": base_store_params, "external_blob_mode": external_blob_mode, "allow_external_blob_outside_bases": allow_external_blob_outside_bases, @@ -7487,8 +7583,7 @@ def _coerce_query_vector(query: QueryVectorLike) -> tuple[pa.Array, int]: if isinstance(query.type, pa.FixedSizeListType): query = query.values elif isinstance(query, (list, tuple)) or ( - _check_for_numpy(query), - isinstance(query, np.ndarray), + _check_for_numpy(query) and isinstance(query, np.ndarray) ): query = np.array(query).astype("float64") # workaround for GH-608 query = pa.FloatingPointArray.from_pandas(query, type=pa.float32()) diff --git a/python/python/lance/dependencies.py b/python/python/lance/dependencies.py index df6b71cb22f..a739b7bc75d 100644 --- a/python/python/lance/dependencies.py +++ b/python/python/lance/dependencies.py @@ -27,7 +27,6 @@ _CAGRA_AVAILABLE = True _RAFT_COMMON_AVAILABLE = True _HUGGING_FACE_AVAILABLE = True -_TENSORFLOW_AVAILABLE = True class _LazyModule(ModuleType): @@ -49,7 +48,6 @@ class _LazyModule(ModuleType): "pandas": "pd.", "polars": "pl.", "torch": "torch.", - "tensorflow": "tf.", } def __init__( @@ -163,8 +161,7 @@ def _lazy_import(module_name: str) -> tuple[ModuleType, bool]: import numpy import pandas import polars - import tensorflow - import torch + import torch # type: ignore[reportMissingImports] else: # heavy/optional third party libs numpy, _NUMPY_AVAILABLE = _lazy_import("numpy") @@ -172,7 +169,6 @@ def _lazy_import(module_name: str) -> tuple[ModuleType, bool]: polars, _POLARS_AVAILABLE = _lazy_import("polars") torch, _TORCH_AVAILABLE = _lazy_import("torch") datasets, _HUGGING_FACE_AVAILABLE = _lazy_import("datasets") - tensorflow, _TENSORFLOW_AVAILABLE = _lazy_import("tensorflow") @lru_cache(maxsize=None) @@ -215,26 +211,18 @@ def _check_for_hugging_face(obj: Any, *, check_type: bool = True) -> bool: ) -def _check_for_tensorflow(obj: Any, *, check_type: bool = True) -> bool: - return _TENSORFLOW_AVAILABLE and _might_be( - cast("Hashable", type(obj) if check_type else obj), "tensorflow" - ) - - __all__ = [ # lazy-load third party libs "datasets", "numpy", "pandas", "polars", - "tensorflow", "torch", # lazy utilities "_check_for_hugging_face", "_check_for_numpy", "_check_for_pandas", "_check_for_polars", - "_check_for_tensorflow", "_check_for_torch", "_LazyModule", # exported flags/guards @@ -243,5 +231,4 @@ def _check_for_tensorflow(obj: Any, *, check_type: bool = True) -> bool: "_POLARS_AVAILABLE", "_TORCH_AVAILABLE", "_HUGGING_FACE_AVAILABLE", - "_TENSORFLOW_AVAILABLE", ] diff --git a/python/python/lance/file.py b/python/python/lance/file.py index 7e2b2933aab..5926241977e 100644 --- a/python/python/lance/file.py +++ b/python/python/lance/file.py @@ -26,6 +26,7 @@ ) if TYPE_CHECKING: + from .blob import DedicatedBlobWriter, PackedBlobWriter from .namespace import LanceNamespace @@ -332,6 +333,24 @@ def open_writer( _inner_writer=inner, ) + def open_packed_blob_writer(self, path: str, blob_id: int) -> "PackedBlobWriter": + """ + Opens a packed blob writer for the given data file path. + + The path will be appended to the base path of the session. + """ + return self._session.open_packed_blob_writer(path, blob_id) + + def open_dedicated_blob_writer( + self, path: str, blob_id: int + ) -> "DedicatedBlobWriter": + """ + Opens a dedicated blob writer for the given data file path. + + The path will be appended to the base path of the session. + """ + return self._session.open_dedicated_blob_writer(path, blob_id) + def contains(self, path: str) -> bool: """ Check if a file exists at the given path (relative to this session's base path). diff --git a/python/python/lance/fragment.py b/python/python/lance/fragment.py index 20972d70090..adf220e59f6 100644 --- a/python/python/lance/fragment.py +++ b/python/python/lance/fragment.py @@ -1029,6 +1029,7 @@ def write_fragments( storage_options: Optional[Dict[str, str]] = None, enable_stable_row_ids: bool = False, target_bases: Optional[List[str]] = None, + target_all_bases: Optional[bool] = None, initial_bases: Optional[List["DatasetBasePath"]] = None, base_store_params: Optional[Dict[str, Dict[str, str]]] = None, external_blob_mode: Literal["reference", "ingest"] = "reference", @@ -1054,6 +1055,7 @@ def write_fragments( storage_options: Optional[Dict[str, str]] = None, enable_stable_row_ids: bool = False, target_bases: Optional[List[str]] = None, + target_all_bases: Optional[bool] = None, initial_bases: Optional[List["DatasetBasePath"]] = None, base_store_params: Optional[Dict[str, Dict[str, str]]] = None, external_blob_mode: Literal["reference", "ingest"] = "reference", @@ -1079,6 +1081,7 @@ def write_fragments( storage_options: Optional[Dict[str, str]] = None, enable_stable_row_ids: bool = False, target_bases: Optional[List[str]] = None, + target_all_bases: Optional[bool] = None, initial_bases: Optional[List["DatasetBasePath"]] = None, base_store_params: Optional[Dict[str, Dict[str, str]]] = None, external_blob_mode: Literal["reference", "ingest"] = "reference", @@ -1149,6 +1152,11 @@ def write_fragments( **CREATE mode**: References must match bases in `initial_bases` **APPEND/OVERWRITE modes**: References must match bases in the existing manifest + target_all_bases : bool, optional + Write new data files round-robin across every base registered in the + manifest, resolved at execution time. When True, the dataset's + primary storage participates as the first slot. Cannot be combined + with `target_bases`. initial_bases : list of DatasetBasePath, optional Base paths to register when creating a new dataset (CREATE mode only). @@ -1254,6 +1262,7 @@ def write_fragments( table_id=table_id, enable_stable_row_ids=enable_stable_row_ids, target_bases=target_bases, + target_all_bases=target_all_bases, initial_bases=initial_bases, base_store_params=base_store_params, external_blob_mode=external_blob_mode, diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index 33b375ce1cd..685bc2eaae1 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/python/python/lance/lance/__init__.pyi @@ -87,6 +87,25 @@ from .trace import capture_trace_events as capture_trace_events from .trace import shutdown_tracing as shutdown_tracing from .trace import trace_to_chrome as trace_to_chrome +class MetricPoint: + name: str + kind: str + attributes: Dict[str, str] + value: Optional[float] + buckets: Optional[List[Tuple[str, int]]] + count: Optional[int] + sum: Optional[float] + +class MetricDescription: + name: str + kind: str + unit: Optional[str] + description: str + +def register_lance_metrics_recorder() -> bool: ... +def lance_metrics_catalog() -> List[MetricDescription]: ... +def snapshot_lance_metrics() -> List[MetricPoint]: ... + class CleanupStats: bytes_removed: int old_versions: int @@ -133,6 +152,39 @@ class LanceFileWriter: def add_schema_metadata(self, key: str, value: str) -> None: ... def add_global_buffer(self, data: bytes) -> int: ... +class BlobDescriptor: + def __repr__(self) -> str: ... + +class PackedBlobWriter: + @property + def blob_id(self) -> int: ... + @property + def path(self) -> str: ... + def write_blob(self, data: bytes) -> None: ... + def finish(self) -> List[BlobDescriptor]: ... + +class DedicatedBlobWriter: + @property + def blob_id(self) -> int: ... + @property + def path(self) -> str: ... + def write(self, data: bytes) -> None: ... + def finish(self) -> BlobDescriptor: ... + +class BlobDescriptorArrayBuilder: + def __init__(self, column: str): ... + @property + def field(self) -> pa.Field: ... + def extend_packed( + self, blob_id: int, offsets: List[int], sizes: List[int] + ) -> None: ... + def append_dedicated(self, blob_id: int, size: int) -> None: ... + def append(self, value: BlobDescriptor) -> None: ... + def extend(self, values: List[BlobDescriptor]) -> None: ... + def append_inline(self, data: bytes) -> None: ... + def append_null(self) -> None: ... + def finish(self) -> pa.Array: ... + class LanceFileSession: def __init__( self, @@ -153,6 +205,10 @@ class LanceFileSession: keep_original_array: Optional[bool] = None, max_page_bytes: Optional[int] = None, ) -> LanceFileWriter: ... + def open_packed_blob_writer(self, path: str, blob_id: int) -> PackedBlobWriter: ... + def open_dedicated_blob_writer( + self, path: str, blob_id: int + ) -> DedicatedBlobWriter: ... def contains(self, path: str) -> bool: ... def list(self, path: Optional[str] = None) -> List[str]: ... def list_with_delimiter( @@ -216,6 +272,7 @@ class LanceColumnStatistics: class _Session: def size_bytes(self) -> int: ... + def index_cache_size_bytes(self) -> int: ... class LanceBlobFile: def close(self): ... @@ -433,7 +490,13 @@ class _Dataset: kwargs: Optional[Dict[str, Any]] = None, ): ... def drop_index(self, name: str): ... - def prewarm_index(self, name: str, *, with_position: bool = False): ... + def prewarm_index( + self, + name: str, + *, + with_position: bool = False, + index_segments: Optional[List[str]] = None, + ): ... def merge_index_metadata( self, index_uuid: str, @@ -505,8 +568,13 @@ class _Dataset: index_name: str, partition_id: int, hamming_threshold: int, + index_segments: Optional[List[str]] = None, ) -> pa.RecordBatchReader: ... - def get_ivf_partition_info(self, index_name: str) -> List[dict]: ... + def get_ivf_partition_info( + self, + index_name: str, + index_segments: Optional[List[str]] = None, + ) -> List[dict]: ... def hamming_clustering_for_sample( self, column: str, @@ -528,6 +596,8 @@ class _MergeInsertBuilder: def when_matched_fail(self) -> Self: ... def when_not_matched_insert_all(self) -> Self: ... def when_not_matched_by_source_delete(self, expr: Optional[str] = None) -> Self: ... + def target_bases(self, bases: list[str]) -> Self: ... + def target_all_bases(self, include_primary: bool = True) -> Self: ... def execute(self, new_data: pa.RecordBatchReader) -> ExecuteResult: ... class _Scanner: @@ -615,6 +685,7 @@ def _write_fragments( table_id: Optional[List[str]], enable_stable_row_ids: bool, target_bases: Optional[List[str]] = None, + target_all_bases: Optional[bool] = None, initial_bases: Optional[List[Any]] = None, base_store_params: Optional[Dict[str, Dict[str, str]]] = None, external_blob_mode: Literal["reference", "ingest"] = "reference", @@ -634,6 +705,7 @@ def _write_fragments_transaction( table_id: Optional[List[str]], enable_stable_row_ids: bool, target_bases: Optional[List[str]] = None, + target_all_bases: Optional[bool] = None, initial_bases: Optional[List[Any]] = None, base_store_params: Optional[Dict[str, Dict[str, str]]] = None, external_blob_mode: Literal["reference", "ingest"] = "reference", @@ -664,6 +736,7 @@ class _ShardSnapshot: class _ShardWriter: shard_id: str def put(self, data: Any) -> None: ... + def delete(self, keys: Any) -> None: ... def close(self) -> None: ... def stats(self) -> Dict[str, Any]: ... def memtable_stats(self) -> Dict[str, Any]: ... @@ -757,7 +830,7 @@ class BFloat16: def __gt__(self, other: BFloat16) -> bool: ... def __ge__(self, other: BFloat16) -> bool: ... -def bfloat16_array(values: List[str | None]) -> BFloat16Array: ... +def bfloat16_array(values: Sequence[float | None]) -> BFloat16Array: ... class PyFullTextQuery: @staticmethod diff --git a/python/python/lance/mem_wal.py b/python/python/lance/mem_wal.py index ed0ce316da8..426a5f7ee50 100644 --- a/python/python/lance/mem_wal.py +++ b/python/python/lance/mem_wal.py @@ -15,6 +15,7 @@ from __future__ import annotations +import math from dataclasses import asdict, dataclass, field from typing import TYPE_CHECKING, Dict, Iterable, List, Mapping, Optional, Union @@ -186,6 +187,7 @@ class ShardWriter: with dataset.mem_wal_writer(shard_id) as writer: writer.put(batch) + writer.delete(pa.table({"id": [1]})) Parameters ---------- @@ -223,6 +225,36 @@ def put(self, data, *, schema: Optional[pa.Schema] = None) -> None: reader = _coerce_reader(data, schema) self._raw.put(reader) + def delete(self, keys, *, schema: Optional[pa.Schema] = None) -> None: + """Delete rows by primary key from the MemWAL. + + Parameters + ---------- + keys : ReaderLike + Any Arrow-compatible data containing this shard's primary key + column(s). Non-primary-key columns, if present, are ignored by + the Rust core delete path. + schema : pa.Schema, optional + Schema hint, needed when *keys* is a generator. + + Raises + ------ + IOError + If delete validation fails, WAL flush fails, or the writer has + already been closed. Delete validation is centralized in Rust and + includes checks for primary-key metadata and tombstone-compatible + nullable non-key columns. + + Examples + -------- + :: + + with dataset.mem_wal_writer(shard_id) as writer: + writer.delete(pa.table({"id": [42]})) + """ + reader = _coerce_reader(keys, schema) + self._raw.delete(reader) + def close(self) -> None: """Flush and close the writer. @@ -337,7 +369,9 @@ def filter(self, expr: str) -> "LsmScanner": self._raw = self._raw.filter(expr) return self - def limit(self, n: int, offset: Optional[int] = None) -> "LsmScanner": + def limit( + self, n: Optional[int] = None, offset: Optional[int] = None + ) -> "LsmScanner": """Limit rows returned, optionally with an offset.""" self._raw = self._raw.limit(n, offset) return self @@ -489,6 +523,9 @@ class LsmVectorSearchPlanner: distance_type : str, optional Distance metric — one of ``"l2"`` (default), ``"cosine"``, ``"dot"``, ``"hamming"``. + filter : str, optional + SQL predicate applied as a prefilter to every LSM source before vector + candidate selection. Examples -------- @@ -506,12 +543,15 @@ def __init__( vector_column: str, pk_columns: Optional[List[str]] = None, distance_type: Optional[str] = None, + filter: Optional[str] = None, ) -> None: kwargs = {} if pk_columns is not None: kwargs["pk_columns"] = pk_columns if distance_type is not None: kwargs["distance_type"] = distance_type + if filter is not None: + kwargs["filter"] = filter self._raw = _LsmVectorSearchPlanner( dataset._ds, [s._raw for s in shard_snapshots], @@ -547,22 +587,11 @@ def plan_search( IVF-PQ). Memtable arms use exact HNSW and are unaffected. Auto-enabled whenever stale filtering is on (see ``overfetch_factor``). overfetch_factor : float, optional - Single knob controlling **both** stale-read filtering and over-fetch - (default: ``1.0``): - - - ``< 1.0`` (e.g. ``0.0``): stale filtering **off**. Rows superseded - by a newer generation may surface. (The global primary-key dedup - still runs, so it suppresses stale copies whenever both the stale - and the fresh row reach it.) - - ``== 1.0``: filtering **on**, no over-fetch. Each source with - superseded rows fetches exactly ``k`` candidates and drops the stale - ones, so it may return fewer than ``k`` live rows. - - ``> 1.0``: filtering **on**, with over-fetch. Such a source fetches - ``ceil(k * overfetch_factor)`` candidates so dropping the stale ones - still leaves ``k`` live rows. - - There is no separate on/off flag: over-fetch is only meaningful while - filtering, so the factor encodes both. + Over-fetch multiple for sources with rows superseded by newer + generations. Must be at least ``1.0``. At ``1.0`` each affected + source fetches exactly ``k`` candidates; above ``1.0`` it fetches + ``ceil(k * overfetch_factor)`` candidates so dropping stale rows is + less likely to under-fill the result. Returns ------- @@ -570,6 +599,14 @@ def plan_search( Physical plan for the vector search. Execute it via `to_table`, `to_reader`, or `to_batches`. """ + if k <= 0: + raise ValueError(f"k must be positive, got {k}") + if nprobes <= 0: + raise ValueError(f"nprobes must be positive, got {nprobes}") + if not math.isfinite(overfetch_factor) or overfetch_factor < 1.0: + raise ValueError( + f"overfetch_factor must be finite and >= 1.0, got {overfetch_factor}" + ) return ExecutionPlan( self._raw.plan_search( query, diff --git a/python/python/lance/namespace.py b/python/python/lance/namespace.py index fec3a1cfb1e..6df5b9aa4bc 100644 --- a/python/python/lance/namespace.py +++ b/python/python/lance/namespace.py @@ -342,12 +342,13 @@ class DirectoryNamespace(LanceNamespace): >>> >>> # With AWS credential vending (requires credential-vendor-aws feature) >>> # Use **dict to pass property names with dots - >>> ns = lance.namespace.DirectoryNamespace(**{ + >>> aws_properties = { ... "root": "s3://my-bucket/data", ... "credential_vendor.enabled": "true", ... "credential_vendor.aws_role_arn": "arn:aws:iam::123456789012:role/MyRole", ... "credential_vendor.aws_duration_millis": "3600000", - ... }) + ... } + >>> # ns = lance.namespace.DirectoryNamespace(**aws_properties) With dynamic context provider: diff --git a/python/python/lance/otel.py b/python/python/lance/otel.py new file mode 100644 index 00000000000..0c9ee1790fd --- /dev/null +++ b/python/python/lance/otel.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Bridge Lance's internal metrics into OpenTelemetry. + +Lance core publishes metrics (currently object store request counts, bytes, +latency, errors, and throttles) through the Rust ``metrics`` facade. This module +installs a process-global recorder that aggregates them and registers +OpenTelemetry observable instruments that report the aggregated values into the +user's ``MeterProvider``. + +The bridge is generic: every metric Lance describes is surfaced automatically, +with no per-metric Python code. Histograms have no asynchronous OpenTelemetry +instrument, so each is exported Prometheus-style as cumulative ``le`` buckets +plus ``_count`` and ``_sum`` observable counters. +""" + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING, Optional + +from .lance import ( + lance_metrics_catalog, + register_lance_metrics_recorder, + snapshot_lance_metrics, +) + +if TYPE_CHECKING: + from opentelemetry.metrics import MeterProvider + +_INSTRUMENTED = False + + +def instrument_lance_metrics(meter_provider: Optional["MeterProvider"] = None) -> bool: + """Register Lance metrics as OpenTelemetry observable instruments. + + Installs a process-global metrics recorder and creates one observable + instrument per Lance metric on the given (or global) ``MeterProvider``. The + user's configured ``MetricReader`` then collects them on its own schedule. + + Counters and gauges map directly to observable counters/gauges. Each + histogram is exported as cumulative ``le`` bucket counts (``_bucket``, + with an ``le`` attribute) plus ``_count`` and ``_sum``. + + Parameters + ---------- + meter_provider : opentelemetry.metrics.MeterProvider, optional + The provider to register instruments on. Defaults to the global provider + from ``opentelemetry.metrics.get_meter_provider()``. + + Returns + ------- + bool + ``True`` if the recorder is installed and instruments are registered. + ``False`` if a different ``metrics`` recorder is already installed in + this process (``metrics`` permits only one global recorder), in which + case a warning is emitted and no instruments are created. + + Notes + ----- + Requires the OpenTelemetry SDK (``pip install pylance[otel]``). Calling this + more than once is safe; instruments are created only on the first successful + call. + """ + global _INSTRUMENTED + + try: + from opentelemetry.metrics import Observation, get_meter_provider + except ImportError as exc: + raise ImportError( + "instrument_lance_metrics requires the OpenTelemetry API/SDK. " + "Install it with `pip install pylance[otel]` or " + "`pip install opentelemetry-sdk`." + ) from exc + + if not register_lance_metrics_recorder(): + warnings.warn( + "Could not install the Lance metrics recorder: another `metrics` " + "recorder is already installed in this process. Lance metrics will " + "not be exported via OpenTelemetry.", + stacklevel=2, + ) + return False + + if _INSTRUMENTED: + return True + + provider = meter_provider or get_meter_provider() + meter = provider.get_meter("lance") + + def scalar_callback(metric_name: str): + def callback(_options): + return [ + Observation(point.value, point.attributes) + for point in snapshot_lance_metrics() + if point.name == metric_name and point.value is not None + ] + + return callback + + def bucket_callback(metric_name: str): + def callback(_options): + observations = [] + for point in snapshot_lance_metrics(): + if point.name != metric_name or point.buckets is None: + continue + for le, cumulative in point.buckets: + attributes = dict(point.attributes) + attributes["le"] = le + observations.append(Observation(cumulative, attributes)) + return observations + + return callback + + def field_callback(metric_name: str, field: str): + def callback(_options): + observations = [] + for point in snapshot_lance_metrics(): + if point.name != metric_name: + continue + value = getattr(point, field) + if value is not None: + observations.append(Observation(value, point.attributes)) + return observations + + return callback + + for desc in lance_metrics_catalog(): + unit = desc.unit or "" + if desc.kind == "counter": + meter.create_observable_counter( + desc.name, + callbacks=[scalar_callback(desc.name)], + unit=unit, + description=desc.description, + ) + elif desc.kind == "gauge": + meter.create_observable_gauge( + desc.name, + callbacks=[scalar_callback(desc.name)], + unit=unit, + description=desc.description, + ) + elif desc.kind == "histogram": + # `_bucket` and `_count` observe cumulative counts, so they are + # unitless; only `_sum` carries the histogram's unit (e.g. seconds). + meter.create_observable_counter( + f"{desc.name}_bucket", + callbacks=[bucket_callback(desc.name)], + description=f"{desc.description} (cumulative buckets)", + ) + meter.create_observable_counter( + f"{desc.name}_count", + callbacks=[field_callback(desc.name, "count")], + description=f"{desc.description} (count)", + ) + meter.create_observable_counter( + f"{desc.name}_sum", + callbacks=[field_callback(desc.name, "sum")], + unit=unit, + description=f"{desc.description} (sum)", + ) + + _INSTRUMENTED = True + return True diff --git a/python/python/lance/tf/__init__.py b/python/python/lance/tf/__init__.py deleted file mode 100644 index 1aa41beb99c..00000000000 --- a/python/python/lance/tf/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright The Lance Authors - -import importlib.util - -if importlib.util.find_spec("tensorflow") is None: - raise ImportError( - "Tensorflow is not installed. Please install tensorflow" - + " to use lance.tf module.", - ) diff --git a/python/python/lance/tf/data.py b/python/python/lance/tf/data.py deleted file mode 100644 index 68280f9211f..00000000000 --- a/python/python/lance/tf/data.py +++ /dev/null @@ -1,410 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright The Lance Authors - - -"""Tensorflow Dataset (`tf.data `_) -implementation for Lance. - -.. warning:: - - Experimental feature. API stability is not guaranteed. -""" - -from __future__ import annotations - -from functools import partial -from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Tuple, Union - -import pyarrow as pa - -import lance -from lance import LanceDataset -from lance.arrow import EncodedImageType, FixedShapeImageTensorType, ImageURIType -from lance.dependencies import _check_for_numpy -from lance.dependencies import numpy as np -from lance.dependencies import tensorflow as tf -from lance.fragment import FragmentMetadata, LanceFragment -from lance.log import LOGGER - -if TYPE_CHECKING: - from pathlib import Path - - from lance import LanceNamespace - - -def arrow_data_type_to_tf(dt: pa.DataType) -> tf.DType: - """Convert Pyarrow DataType to Tensorflow.""" - if pa.types.is_boolean(dt): - return tf.bool - elif pa.types.is_int8(dt): - return tf.int8 - elif pa.types.is_int16(dt): - return tf.int16 - elif pa.types.is_int32(dt): - return tf.int32 - elif pa.types.is_int64(dt): - return tf.int64 - elif pa.types.is_uint8(dt): - return tf.uint8 - elif pa.types.is_uint16(dt): - return tf.uint16 - elif pa.types.is_uint32(dt): - return tf.uint32 - elif pa.types.is_uint64(dt): - return tf.uint64 - elif pa.types.is_float16(dt): - return tf.float16 - elif pa.types.is_float32(dt): - return tf.float32 - elif pa.types.is_float64(dt): - return tf.float64 - elif ( - pa.types.is_string(dt) - or pa.types.is_large_string(dt) - or pa.types.is_binary(dt) - or pa.types.is_large_binary(dt) - ): - return tf.string - - raise TypeError(f"Arrow/Tf conversion: Unsupported arrow data type: {dt}") - - -def data_type_to_tensor_spec(dt: pa.DataType) -> tf.TensorSpec: - """Convert PyArrow DataType to Tensorflow TensorSpec.""" - if ( - pa.types.is_boolean(dt) - or pa.types.is_integer(dt) - or pa.types.is_floating(dt) - or pa.types.is_string(dt) - or pa.types.is_binary(dt) - ): - return tf.TensorSpec(shape=(None,), dtype=arrow_data_type_to_tf(dt)) - elif isinstance(dt, pa.FixedShapeTensorType): - return tf.TensorSpec( - shape=(None, *dt.shape), dtype=arrow_data_type_to_tf(dt.value_type) - ) - elif pa.types.is_fixed_size_list(dt): - return tf.TensorSpec( - shape=(None, dt.list_size), dtype=arrow_data_type_to_tf(dt.value_type) - ) - elif pa.types.is_list(dt) or pa.types.is_large_list(dt): - return tf.TensorSpec( - shape=( - None, - None, - ), - dtype=arrow_data_type_to_tf(dt.value_type), - ) - elif pa.types.is_struct(dt): - return {field.name: data_type_to_tensor_spec(field.type) for field in dt} - elif isinstance(dt, (EncodedImageType, ImageURIType)): - return tf.TensorSpec(shape=(None,), dtype=tf.string) - elif isinstance(dt, FixedShapeImageTensorType): - return tf.TensorSpec( - shape=(None, *dt.shape), dtype=arrow_data_type_to_tf(dt.arrow_type) - ) - - raise TypeError("Unsupported data type: ", dt) - - -def schema_to_spec(schema: pa.Schema) -> tf.TypeSpec: - """Convert PyArrow Schema to Tensorflow output signature.""" - signature = {} - for name in schema.names: - field = schema.field(name) - signature[name] = data_type_to_tensor_spec(field.type) - return signature - - -def column_to_tensor(array: pa.Array, tensor_spec: tf.TensorSpec) -> tf.Tensor: - """Convert a PyArrow array into a TensorFlow tensor.""" - if isinstance(tensor_spec, tf.RaggedTensorSpec): - return tf.ragged.constant(array.to_pylist(), dtype=tensor_spec.dtype) - elif isinstance(array.type, pa.FixedShapeTensorType): - return tf.constant(array.to_numpy_ndarray(), dtype=tensor_spec.dtype) - elif isinstance(array.type, FixedShapeImageTensorType): - return tf.constant(array.to_numpy(), dtype=tensor_spec.dtype) - elif isinstance(array.type, pa.StructType): - return { - field.name: column_to_tensor(array.field(i), tensor_spec[field.name]) - for (i, field) in enumerate(array.type) - } - else: - return tf.constant(array.to_pylist(), dtype=tensor_spec.dtype) - - -def from_lance( - dataset: Optional[Union[str, Path, LanceDataset]] = None, - *, - columns: Optional[Union[List[str], Dict[str, str]]] = None, - batch_size: int = 256, - filter: Optional[str] = None, - fragments: Union[Iterable[int], Iterable[LanceFragment], tf.data.Dataset] = None, - output_signature: Optional[Dict[str, tf.TypeSpec]] = None, - namespace_client: Optional["LanceNamespace"] = None, - table_id: Optional[List[str]] = None, - ignore_namespace_table_storage_options: bool = False, -) -> tf.data.Dataset: - """Create a ``tf.data.Dataset`` from a Lance dataset. - - Parameters - ---------- - dataset : Union[str, Path, LanceDataset], optional - Lance dataset or dataset URI/path. Either ``dataset`` or both - ``namespace_client`` and ``table_id`` must be provided. - columns : Optional[List[str]], optional - List of columns to include in the output dataset. - If not set, all columns will be read. - batch_size : int, optional - Batch size, by default 256 - filter : Optional[str], optional - SQL filter expression, by default None. - fragments : Union[List[LanceFragment], tf.data.Dataset], optional - If provided, only the fragments are read. It can be used to feed - for distributed training. - output_signature : Optional[tf.TypeSpec], optional - Override output signature of the returned tensors. If not provided, - the output signature is inferred from the projection Schema. - namespace_client : Optional[LanceNamespace], optional - Namespace client to resolve the table location when ``table_id`` is - provided. - table_id : Optional[List[str]], optional - Table identifier used together with ``namespace_client`` to locate - the table. - ignore_namespace_table_storage_options : bool, default False - When using ``namespace_client``/``table_id``, ignore storage options - returned by the namespace. - - Examples - -------- - - .. code-block:: python - - import tensorflow as tf - import lance.tf.data - - ds = lance.tf.data.from_lance( - "s3://bucket/path", - columns=["image", "id"], - filter="catalog = 'train' AND split = 'train'", - batch_size=100) - - for batch in ds.repeat(10).shuffle(128).map(io_decode): - print(batch["image"].shape) - - ``from_lance`` can take an iterator or ``tf.data.Dataset`` of - Fragments. So that it can be used to feed for distributed training. - - .. code-block:: python - - import tensorflow as tf - import lance.tf.data - - seed = 200 # seed to shuffle the fragments in distributed machines. - fragments = lance.tf.data.lance_fragments("s3://bucket/path") - repeat(10).shuffle(4, seed=seed) - ds = lance.tf.data.from_lance( - "s3://bucket/path", - columns=["image", "id"], - filter="catalog = 'train' AND split = 'train'", - fragments=fragments, - batch_size=100) - for batch in ds.shuffle(128).map(io_decode): - print(batch["image"].shape) - - """ - if isinstance(dataset, LanceDataset): - if namespace_client is not None or table_id is not None: - raise ValueError( - "Cannot specify 'namespace_client' or 'table_id' when passing " - "a LanceDataset instance" - ) - else: - dataset = lance.dataset( - dataset, - namespace_client=namespace_client, - table_id=table_id, - ignore_namespace_table_storage_options=ignore_namespace_table_storage_options, - ) - - if isinstance(fragments, tf.data.Dataset): - fragments = list(fragments.as_numpy_iterator()) - elif _check_for_numpy(fragments) and isinstance(fragments, np.ndarray): - fragments = list(fragments) - - if fragments is not None: - - def gen_fragments(fragments): - for f in fragments: - if isinstance(f, int) or ( - _check_for_numpy(f) and isinstance(f, np.integer) - ): - yield LanceFragment(dataset, int(f)) - elif isinstance(f, FragmentMetadata): - yield LanceFragment(dataset, f.id) - elif isinstance(f, LanceFragment): - yield f - else: - raise TypeError(f"Invalid type passed to `fragments`: {type(f)}") - - # A Generator of Fragments - fragments = gen_fragments(fragments) - - scanner = dataset.scanner( - filter=filter, columns=columns, batch_size=batch_size, fragments=fragments - ) - - if output_signature is None: - schema = scanner.projected_schema - output_signature = schema_to_spec(schema) - LOGGER.debug("Output signature: %s", output_signature) - - def generator(): - for batch in scanner.to_batches(): - yield { - name: column_to_tensor(batch[name], output_signature[name]) - for name in batch.schema.names - } - - return tf.data.Dataset.from_generator(generator, output_signature=output_signature) - - -def lance_fragments(dataset: Union[str, Path, LanceDataset]) -> tf.data.Dataset: - """Create a ``tf.data.Dataset`` of Lance Fragments in the dataset. - - Parameters - ---------- - dataset : Union[str, Path, LanceDataset] - A Lance Dataset or dataset URI/path. - """ - if not isinstance(dataset, LanceDataset): - dataset = lance.dataset(dataset) - return tf.data.Dataset.from_tensor_slices( - [f.fragment_id for f in dataset.get_fragments()] - ) - - -def _ith_batch(i: int, batch_size: int, total_size: int) -> Tuple[int, int]: - """ - Get the start and end index of the ith batch. - - This takes into account the total_size, the total number of rows in the dataset. - """ - start = i * batch_size - end = tf.math.minimum(start + batch_size, total_size) - return (start, end) - - -def from_lance_batches( - dataset: Union[str, Path, LanceDataset], - *, - shuffle: bool = False, - seed: Optional[int] = None, - batch_size: int = 1024, - skip: int = 0, -) -> tf.data.Dataset: - """ - Create a ``tf.data.Dataset`` of batch indices for a Lance dataset. - - Parameters - ---------- - dataset : Union[str, Path, LanceDataset] - A Lance Dataset or dataset URI/path. - shuffle : bool, optional - Shuffle the batches, by default False - seed : Optional[int], optional - Random seed for shuffling, by default None - batch_size : int, optional - Batch size, by default 1024 - skip : int, optional - Number of batches to skip. - - Returns - ------- - tf.data.Dataset - A tensorflow dataset of batch slice ranges. These can be passed to - :func:`lance_take_batches` to create a Tensorflow dataset of batches. - """ - if not isinstance(dataset, LanceDataset): - dataset = lance.dataset(dataset) - num_rows = dataset.count_rows() - num_batches = (num_rows + batch_size - 1) // batch_size - indices = tf.data.Dataset.range(num_batches, dtype=tf.int64) - if shuffle: - indices = indices.shuffle(num_batches, seed=seed) - if skip > 0: - indices = indices.skip(skip) - return indices.map(partial(_ith_batch, batch_size=batch_size, total_size=num_rows)) - - -def lance_take_batches( - dataset: Union[str, Path, LanceDataset], - batch_ranges: Iterable[Tuple[int, int]], - *, - columns: Optional[List[str]] = None, - output_signature: Optional[Dict[str, tf.TypeSpec]] = None, - batch_readahead: int = 10, -) -> tf.data.Dataset: - """ - Create a ``tf.data.Dataset`` of batches from a Lance dataset. - - Parameters - ---------- - dataset : Union[str, Path, LanceDataset] - A Lance Dataset or dataset URI/path. - batch_ranges : Iterable[Tuple[int, int]] - Iterable of batch indices. - columns : Optional[List[str]], optional - List of columns to include in the output dataset. - If not set, all columns will be read. - output_signature : Optional[tf.TypeSpec], optional - Override output signature of the returned tensors. If not provided, - the output signature is inferred from the projection Schema. - batch_readahead : int, default 10 - The number of batches to read ahead in parallel. - - Examples - -------- - You can compose this with ``from_lance_batches`` to create a randomized Tensorflow - dataset. With ``from_lance_batches``, you can deterministically randomized the - batches by setting ``seed``. - - .. code-block:: python - - batch_iter = from_lance_batches(dataset, batch_size=100, shuffle=True, seed=200) - batch_iter = batch_iter.as_numpy_iterator() - lance_ds = lance_take_batches(dataset, batch_iter) - lance_ds = lance_ds.unbatch().shuffle(500, seed=42).batch(100) - """ - if not isinstance(dataset, LanceDataset): - dataset = lance.dataset(dataset) - - if output_signature is None: - schema = dataset.scanner(columns=columns).projected_schema - output_signature = schema_to_spec(schema) - LOGGER.debug("Output signature: %s", output_signature) - - def gen_ranges(): - for start, end in batch_ranges: - yield (start, end) - - def gen_batches(): - batches = dataset._ds.take_scan( - gen_ranges(), - columns=columns, - batch_readahead=batch_readahead, - ) - for batch in batches: - yield { - name: column_to_tensor(batch[name], output_signature[name]) - for name in batch.schema.names - } - - return tf.data.Dataset.from_generator( - gen_batches, output_signature=output_signature - ) - - -# Register `from_lance` to ``tf.data.Dataset``. -tf.data.Dataset.from_lance = from_lance -tf.data.Dataset.from_lance_batches = from_lance_batches diff --git a/python/python/lance/torch/distance.py b/python/python/lance/torch/distance.py index 81201027c87..ed5292cdb2f 100644 --- a/python/python/lance/torch/distance.py +++ b/python/python/lance/torch/distance.py @@ -150,8 +150,9 @@ def pairwise_l2( if x.dtype != y.dtype or (y2 is not None and x.dtype != y2.dtype): raise ValueError("pairwise_l2 data types do not match") origin_dtype = x.dtype - if x.device == torch.device("cpu") and x.dtype == torch.float16: - # Pytorch does not support `x @ y.T` for float16 on CPU + if x.device == torch.device("cpu") and x.dtype in (torch.float16, torch.bfloat16): + # Pytorch does not support `x @ y.T` for float16 on CPU, and bfloat16 + # can hit backend-specific compiler paths. Accumulate in float32. x = x.type(torch.float32) y = y.type(torch.float32) if y2 is not None: diff --git a/python/python/lance/util.py b/python/python/lance/util.py index 2161c4e0d45..180e2441b43 100644 --- a/python/python/lance/util.py +++ b/python/python/lance/util.py @@ -3,8 +3,18 @@ from __future__ import annotations +import uuid from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Iterator, Literal, Optional, Union, cast +from typing import ( + TYPE_CHECKING, + Iterable, + Iterator, + List, + Literal, + Optional, + Union, + cast, +) import pyarrow as pa @@ -51,6 +61,29 @@ def td_to_micros(td: timedelta) -> int: return round(td / timedelta(microseconds=1)) +def _normalize_index_segment_ids( + index_segments: Optional[Iterable[Union[str, uuid.UUID]]], +) -> Optional[List[str]]: + """Normalize a physical index segment selection to a list of UUID strings.""" + if index_segments is None: + return None + if isinstance(index_segments, (str, uuid.UUID)): + raise TypeError( + "index_segments must be an iterable of str or uuid.UUID, " + f"not a single {type(index_segments)}." + ) + segment_ids = [] + for segment_id in index_segments: + if isinstance(segment_id, (str, uuid.UUID)): + segment_ids.append(str(segment_id)) + else: + raise TypeError( + "index_segments must be an iterable of str or uuid.UUID. " + f"Got {type(segment_id)} instead." + ) + return segment_ids + + class KMeans: """KMean model for clustering. diff --git a/python/python/lance/vector.py b/python/python/lance/vector.py index 5ce5e8b61e5..205001ccacd 100644 --- a/python/python/lance/vector.py +++ b/python/python/lance/vector.py @@ -19,9 +19,10 @@ ) from .dependencies import numpy as np from .log import LOGGER -from .util import MetricType, _normalize_metric_type +from .util import MetricType, _normalize_index_segment_ids, _normalize_metric_type if TYPE_CHECKING: + import uuid from pathlib import Path from . import LanceDataset @@ -761,13 +762,17 @@ def hamming_clustering_for_ivf_partition( index_name: str, partition_id: int, hamming_threshold: int, + *, + index_segments: Optional[Iterable[Union[str, uuid.UUID]]] = None, ) -> pa.RecordBatchReader: """ Perform hamming clustering on a partition of an IVF_FLAT index. - Loads a partition from an IVF_FLAT index on a hash column, computes - pairwise hamming distances between all hashes in the partition, - filters by threshold, and clusters the results using union-find. + Loads a partition from every segment of an IVF_FLAT index on a hash + column, computes pairwise hamming distances between all hashes in the + combined partition, filters by threshold, and clusters the results using + union-find. All segments of the logical index must share the same global + IVF centroids; an error is raised if they do not. Parameters ---------- @@ -779,6 +784,11 @@ def hamming_clustering_for_ivf_partition( The partition ID within the IVF_FLAT index hamming_threshold : int Maximum hamming distance to consider as similar + index_segments : iterable of str or uuid.UUID, optional + If specified, only these physical index segment UUIDs of the named + logical index contribute rows. Use + :meth:`LanceDataset.describe_indices` to obtain segment UUIDs from + ``IndexDescription.segments``. Defaults to all segments. Returns ------- @@ -789,30 +799,43 @@ def hamming_clustering_for_ivf_partition( - 'duplicates': list - List of duplicate row IDs in each cluster """ return dataset._ds.hamming_clustering_for_ivf_partition( - index_name, partition_id, hamming_threshold + index_name, + partition_id, + hamming_threshold, + _normalize_index_segment_ids(index_segments), ) def get_ivf_partition_info( dataset: "LanceDataset", index_name: str, + *, + index_segments: Optional[Iterable[Union[str, uuid.UUID]]] = None, ) -> List[dict]: """ Get partition information for an IVF_FLAT index. + Partition sizes are aggregated across all segments of the logical index + unless a subset is selected via ``index_segments``. + Parameters ---------- dataset : LanceDataset The Lance dataset containing the hash column with an IVF_FLAT index. index_name : str Name of the IVF_FLAT index + index_segments : iterable of str or uuid.UUID, optional + If specified, only these physical index segment UUIDs of the named + logical index contribute to the sizes. Defaults to all segments. Returns ------- list[dict] List of partition info dicts with 'partition_id' and 'size' """ - return dataset._ds.get_ivf_partition_info(index_name) + return dataset._ds.get_ivf_partition_info( + index_name, _normalize_index_segment_ids(index_segments) + ) def hamming_clustering_for_sample( diff --git a/python/python/tests/compat/compat_sequence.py b/python/python/tests/compat/compat_sequence.py index bec216d6a6f..48df46500ca 100644 --- a/python/python/tests/compat/compat_sequence.py +++ b/python/python/tests/compat/compat_sequence.py @@ -18,8 +18,9 @@ "ignore the index" mode to diff against, so its oracle reconstructs ground truth from a full scan: tokenize every live row, then require an FTS search for a spread of sampled terms to return exactly the rows that contain them. The FTS scenarios run under both -on-disk format versions (LANCE_FTS_FORMAT_VERSION 1 and 2), which take different merge -paths. +on-disk format versions (1 and 2), which take different merge paths. New Lance versions +pin this with the create-index `format_version` parameter; old Lance versions still use +`LANCE_FTS_FORMAT_VERSION`. The op vocabulary and bounds are deliberately small so the search is runnable; this is exhaustive over the maintenance-lifecycle grammar up to the configured lengths, not over @@ -63,11 +64,12 @@ def describe(kind, from_ref, to_ref, setup_ops, exercise_ops, fts_version=None): class IndexScenario: """A picklable, kind-parameterized scenario run across a version split.""" - def __init__(self, kind, path, setup_ops, exercise_ops): + def __init__(self, kind, path, setup_ops, exercise_ops, fts_version=None): self.kind = kind self.path = str(path) self.setup_ops = list(setup_ops) self.exercise_ops = list(exercise_ops) + self.fts_version = fts_version self.next_idx = 0 # --- in-venv helpers (only lance + pyarrow available) --- @@ -123,6 +125,8 @@ def _op_W(self): def _op_I(self): kwargs = {"with_position": True} if self.kind == "INVERTED" else {} + if self.kind == "INVERTED" and self.fts_version is not None: + kwargs["format_version"] = int(self.fts_version) self._open().create_scalar_index("key", self._index_type(), **kwargs) def _op_D(self): @@ -234,9 +238,10 @@ def search( """Search index-maintenance sequences up to `max_length` ops for one `kind`, across (from_ref -> to_ref). Runs only scenarios in this shard (i % num_shards == shard) so the space can be split across parallel workers. For INVERTED, `fts_version` ("1" or - "2") pins the on-disk FTS format (LANCE_FTS_FORMAT_VERSION) on both sides; both are - Fst token sets and exercise distinct merge paths. Returns failures; stops on the - first when `stop_on_first`.""" + "2") pins the on-disk FTS format on both sides. New Lance versions receive this + through the create-index parameter and old Lance versions receive it through + LANCE_FTS_FORMAT_VERSION. Both are Fst token sets and exercise distinct merge paths. + Returns failures; stops on the first when `stop_on_first`.""" from_venv = venv_factory.get_venv(from_ref) to_venv = venv_factory.get_venv(to_ref) env = {} @@ -256,7 +261,7 @@ def search( if key not in snapshots: snap = base / f"snap_{kind}_{len(snapshots)}" shutil.rmtree(snap, ignore_errors=True) - builder = IndexScenario(kind, snap, setup_tail, []) + builder = IndexScenario(kind, snap, setup_tail, [], fts_version) try: next_idx = from_venv.execute_method(builder, "setup", env) snapshots[key] = (snap, next_idx) @@ -277,7 +282,7 @@ def search( ex_path = base / f"ex_{kind}_{i}" shutil.rmtree(ex_path, ignore_errors=True) shutil.copytree(snap, ex_path) - scenario = IndexScenario(kind, ex_path, setup_tail, exercise) + scenario = IndexScenario(kind, ex_path, setup_tail, exercise, fts_version) scenario.next_idx = next_idx label = describe(kind, from_ref, to_ref, setup_tail, exercise, fts_version) try: diff --git a/python/python/tests/compat/test_scalar_indices.py b/python/python/tests/compat/test_scalar_indices.py index 35022df3b12..4f54e9d31c2 100644 --- a/python/python/tests/compat/test_scalar_indices.py +++ b/python/python/tests/compat/test_scalar_indices.py @@ -9,6 +9,7 @@ and written by other versions. """ +import os import shutil from pathlib import Path @@ -320,7 +321,12 @@ def create(self): max_rows_per_file=100, data_storage_version=safe_data_storage_version(self.compat_version), ) - dataset.create_scalar_index("text", "INVERTED", with_position=True) + kwargs = {"with_position": True} + # Downgrade reads use older wheels, so current-created FTS indexes must + # stay on the legacy posting block layout. + if os.environ.get("LANCE_COMPAT_FTS_LEGACY_BLOCK_SIZE") == "1": + kwargs["block_size"] = 128 + dataset.create_scalar_index("text", "INVERTED", format_version=1, **kwargs) def check_read(self): """Verify FTS index can be queried.""" @@ -351,7 +357,15 @@ def skip_downgrade(self, version: str) -> bool: def current_env(self, method_name: str) -> dict[str, str]: if method_name == "create": - return {"LANCE_FTS_FORMAT_VERSION": "1"} + return { + "LANCE_COMPAT_FTS_LEGACY_BLOCK_SIZE": "1", + "LANCE_FTS_FORMAT_VERSION": "1", + } if method_name == "check_write": return {"LANCE_FTS_FORMAT_VERSION": "2"} return {} + + def compat_env(self, version: str, method_name: str) -> dict[str, str]: + if method_name in {"create", "check_write"}: + return {"LANCE_FTS_FORMAT_VERSION": "1"} + return {} diff --git a/python/python/tests/test_arrow.py b/python/python/tests/test_arrow.py index b6e6024e0e7..a0be09846e1 100644 --- a/python/python/tests/test_arrow.py +++ b/python/python/tests/test_arrow.py @@ -7,10 +7,11 @@ from pathlib import Path import lance +import lance.arrow import numpy as np import pandas as pd import pyarrow as pa -import pytest +import pytest # pyright: ignore[reportMissingImports] from lance.arrow import ( BFloat16, BFloat16Array, @@ -272,8 +273,6 @@ def test_image_uri_arrays(tmp_path: Path, png_uris): def test_image_tensor_arrays(tmp_path: Path, png_uris): - tf = pytest.importorskip("tensorflow") - n = 10 encoded_image_array = ImageURIArray.from_uris(png_uris).read_uris() @@ -296,22 +295,22 @@ def test_image_tensor_arrays(tmp_path: Path, png_uris): assert tensor_image_array.storage.type == pa.list_(pa.uint8(), 4) assert tensor_image_array[2].as_py() == [42, 42, 42, 255] - test_tensor = tf.constant( - np.array([42, 42, 42, 255] * n, dtype=np.uint8).reshape((n, 1, 1, 4)) - ) + test_tensor = np.array([42, 42, 42, 255] * n, dtype=np.uint8).reshape((n, 1, 1, 4)) assert test_tensor.shape == (n, 1, 1, 4) - assert tf.math.reduce_all( - tf.convert_to_tensor(tensor_image_array.to_numpy()) == test_tensor - ) + assert np.array_equal(tensor_image_array.to_numpy(), test_tensor) assert tensor_image_array.to_encoded().to_tensor() == tensor_image_array def png_encoder(images): - import tensorflow as tf + import io - encoded_images = ( - tf.io.encode_png(x).numpy() for x in tf.convert_to_tensor(images) - ) + from PIL import Image # pyright: ignore[reportMissingImports] + + encoded_images = [] + for image in images: + with io.BytesIO() as buf: + Image.fromarray(image).save(buf, format="PNG") + encoded_images.append(buf.getvalue()) return pa.array(encoded_images, type=pa.binary()) assert tensor_image_array.to_encoded(png_encoder).to_tensor() == tensor_image_array @@ -323,20 +322,18 @@ def png_encoder(images): uris = [str(Path(x)) for x in uris] encoded_image_array = ImageArray.from_array(uris).read_uris() - with pytest.raises( - tf.errors.InvalidArgumentError, match="Shapes of all inputs must match" - ): + with pytest.raises(ValueError, match="all input arrays must have the same shape"): encoded_image_array.to_tensor() pattern = r"(object at) 0x[\w\d]+(:?>)" repl = r"\1 0x..\2" - assert re.sub(pattern, repl, encoded_image_array.__repr__()) == ( - "\n" - "[, ..]\n" + repr_ = re.sub(pattern, repl, encoded_image_array.__repr__()) + assert repr_.startswith("\n[ + Lance JSON (JSONB LargeBinary) conversion that write_dataset/write_fragments + perform, so the raw UTF-8 string bytes were written into a column whose schema + declared JSONB. Reads then miss-decoded the bytes and returned garbage. + """ + json_type = pa.json_() + data = pa.table( + { + "uid": pa.array(["a", "b", "c", "d"], type=pa.utf8()), + "payload": pa.array( + ['{"x":1}', '{"x":2}', '{"y":3}', '{"y":4}'], + type=json_type, + ), + } + ) + + frag = LanceFragment.create(tmp_path, data) + operation = LanceOperation.Overwrite(data.schema, [frag]) + dataset = LanceDataset.commit(tmp_path, operation) + + result = dataset.to_table() + assert result.column("uid").to_pylist() == ["a", "b", "c", "d"] + payloads = result.column("payload").to_pylist() + assert [json.loads(p) for p in payloads] == [ + {"x": 1}, + {"x": 2}, + {"y": 3}, + {"y": 4}, + ] + + +def test_fragment_update_columns_with_json_column(tmp_path): + """Test that fragment update_columns works with Arrow JSON extension type. + + Previously this would fail with a type mismatch error because the + HashJoiner didn't convert Arrow JSON (Utf8) to Lance JSON (LargeBinary). + """ + # Create initial dataset with a JSON extension type column + json_type = pa.json_() + data = pa.table( + { + "id": pa.array([1, 2, 3, 4, 5], type=pa.int64()), + "name": pa.array(["a", "b", "c", "d", "e"], type=pa.utf8()), + "meta": pa.array( + ['{"x":1}', '{"x":2}', '{"x":3}', '{"x":4}', '{"x":5}'], + type=json_type, + ), + } + ) + dataset_uri = tmp_path / "test_update_cols_json" + dataset = lance.write_dataset(data, dataset_uri) + + # Prepare update data: update the JSON column for some rows + update_data = pa.table( + { + "_rowid": pa.array([1, 3], type=pa.uint64()), + "meta": pa.array( + ['{"updated":true,"id":2}', '{"updated":true,"id":4}'], + type=json_type, + ), + } + ) + + # This should NOT raise a type mismatch error + fragment = dataset.get_fragment(0) + updated_fragment, fields_modified = fragment.update_columns(update_data) + + assert len(fields_modified) > 0 + + # Commit and verify + op = LanceOperation.Update( + updated_fragments=[updated_fragment], + fields_modified=fields_modified, + ) + updated_dataset = lance.LanceDataset.commit( + str(dataset_uri), op, read_version=dataset.version + ) + + result = updated_dataset.to_table() + ids = result.column("id").to_pylist() + metas = result.column("meta").to_pylist() + + for i, (id_val, meta_val) in enumerate(zip(ids, metas)): + meta = json.loads(meta_val) if isinstance(meta_val, str) else meta_val + if id_val == 2 or id_val == 4: + assert "updated" in meta_val or meta.get("updated") is True, ( + f"id={id_val} should be updated, got {meta_val}" + ) + else: + assert "x" in meta_val or "x" in str(meta), ( + f"id={id_val} should have original value, got {meta_val}" + ) diff --git a/python/python/tests/test_json.py b/python/python/tests/test_json.py index ad2ab6c3ada..7a6a11883dd 100644 --- a/python/python/tests/test_json.py +++ b/python/python/tests/test_json.py @@ -861,3 +861,34 @@ def test_json_merge_insert(tmp_path: Path): result = dataset.to_table(filter="json_get_int(data, 'score') >= 35") assert result.num_rows == 2 + + +def test_json_fragment_session_take(tmp_path: Path): + """FragmentSession.take must return JSON columns as pa.json_(), not raw JSONB.""" + + dataset_path = tmp_path / "json_fragment_session.lance" + + table = pa.table( + { + "id": pa.array([1, 2, 3, 4, 5], type=pa.int32()), + "data": pa.array( + [json.dumps({"name": "row", "score": v}) for v in range(5)], + type=pa.json_(), + ), + } + ) + lance.write_dataset(table, dataset_path) + + dataset = lance.dataset(dataset_path) + fragment = dataset.get_fragments()[0] + session = fragment.open_session(columns=["id", "data"]) + result = session.take([0, 2, 4]) + + # JSON column must be the logical pa.json_() type, not raw LargeBinary JSONB. + assert result.schema.field("data").type == pa.json_() + assert result.column("id").to_pylist() == [1, 3, 5] + assert [json.loads(v) for v in result.column("data").to_pylist()] == [ + {"name": "row", "score": 0}, + {"name": "row", "score": 2}, + {"name": "row", "score": 4}, + ] diff --git a/python/python/tests/test_lance.py b/python/python/tests/test_lance.py index 0162e370665..2ad3af2e029 100644 --- a/python/python/tests/test_lance.py +++ b/python/python/tests/test_lance.py @@ -248,6 +248,28 @@ def test_io_counters(tmp_path): assert lance.bytes_read_counter() > starting_bytes +def test_simd_info(): + info = lance.simd_info() + assert info["tier"] in ( + "none", + "sse", + "avx", + "avx_fma", + "avx2", + "avx512", + "avx512_fp16", + "neon", + "lsx", + "lasx", + ) + assert isinstance(info["target_arch"], str) and info["target_arch"] + if info["target_arch"] == "x86_64": + # The x86_64 ABI mandates SSE2. + assert "sse2" in info["host_features"] + else: + assert info["host_features"] == [] + + @pytest.mark.parametrize( "row_param, column_name", [("with_row_id", "_rowid"), ("with_row_address", "_rowaddr")], diff --git a/python/python/tests/test_mem_wal.py b/python/python/tests/test_mem_wal.py index c21e88b2416..2871baad986 100644 --- a/python/python/tests/test_mem_wal.py +++ b/python/python/tests/test_mem_wal.py @@ -11,6 +11,7 @@ from lance.mem_wal import ( LsmPointLookupPlanner, LsmScanner, + LsmVectorSearchPlanner, ShardingField, ShardingSpec, ShardSnapshot, @@ -163,6 +164,11 @@ def test_lsm_scanner_with_memtables(tmp_path): assert name_by_id[2] == "gen1_2", "Flushed gen must overwrite base for id=2" assert name_by_id[3] == "base_3" + offset_table = ( + LsmScanner.from_snapshots(base_ds, [snap]).limit(None, offset=1).to_table() + ) + assert len(offset_table) == 2, "Offset-only LSM scan should not require a limit" + def test_shard_writer_lsm_scanner_includes_own_flushed_generations(tmp_path): ds_path = str(tmp_path / "base") @@ -190,6 +196,34 @@ def test_shard_writer_lsm_scanner_includes_own_flushed_generations(tmp_path): time.sleep(0.05) +def test_shard_writer_delete_binding_masks_base_row(tmp_path): + ds_path = str(tmp_path / "base") + shard_id = str(uuid.uuid4()) + ds = lance.write_dataset( + _lookup_table([1, 2, 3], "base"), ds_path, schema=_LOOKUP_SCHEMA + ) + ds.initialize_mem_wal() + + delete_keys = pa.table({"id": pa.array([2], type=pa.int64())}) + + with ds.mem_wal_writer( + shard_id, + durable_write=True, + sync_indexed_write=True, + max_wal_buffer_size=1, + max_wal_flush_interval_ms=10, + ) as writer: + writer.put(_lookup_table([4], "writer")) + writer.delete(delete_keys) + table = writer.lsm_scanner().to_table() + + rows = {row["id"]: row["name"] for row in table.to_pylist()} + assert rows[1] == "base_1" + assert 2 not in rows, "deleted base row should be masked by the tombstone" + assert rows[3] == "base_3" + assert rows[4] == "writer_4" + + _VDIM = 4 # matches Rust test fixture dimension @@ -221,6 +255,40 @@ def _vector_search_table(ids): ) +def test_lsm_vector_search_filter_binding(tmp_path): + ds_path = str(tmp_path / "vec") + ds = lance.write_dataset( + _vector_search_table(range(400)), ds_path, schema=_vector_search_schema() + ) + ds.create_index( + "vector", + index_type="IVF_PQ", + name="vector_idx", + num_partitions=2, + num_sub_vectors=2, + ) + ds.initialize_mem_wal(maintained_indexes=["vector_idx"]) + + planner = LsmVectorSearchPlanner(ds, [], "vector", filter="id >= 200") + query_id = 250 + query = pa.array([25.0, 25.1, 25.2, 25.3], type=pa.float32()) + table = planner.plan_search(query, k=20, nprobes=2, columns=["id"]).to_table() + + ids = table.column("id").to_pylist() + assert ids, "filtered vector search should return at least one row" + assert min(ids) >= 200 + assert query_id in ids, f"filtered vector search recall missed id={query_id}: {ids}" + + with pytest.raises(ValueError, match="k must be positive"): + planner.plan_search(query, k=0, nprobes=2, columns=["id"]) + with pytest.raises(ValueError, match="nprobes must be positive"): + planner.plan_search(query, k=20, nprobes=0, columns=["id"]) + with pytest.raises(ValueError, match="overfetch_factor must be finite"): + planner.plan_search( + query, k=20, nprobes=2, columns=["id"], overfetch_factor=math.inf + ) + + VECTOR_DIM = 32 ROWS_PER_BATCH = 50 NUM_WRITE_ROUNDS = 3 diff --git a/python/python/tests/test_multi_base.py b/python/python/tests/test_multi_base.py index 113f3224d14..128a8905842 100644 --- a/python/python/tests/test_multi_base.py +++ b/python/python/tests/test_multi_base.py @@ -82,6 +82,60 @@ def test_multi_base_create_and_read(self): data.sort_values("id").reset_index(drop=True), ) + def test_write_target_all_bases(self): + """target_all_bases=True rotates across primary and all initial bases.""" + data = self.create_test_data(300) + + dataset = lance.write_dataset( + data, + self.primary_uri, + mode="create", + initial_bases=[ + DatasetBasePath(self.path1_uri, name="path1"), + DatasetBasePath(self.path2_uri, name="path2"), + ], + target_all_bases=True, + max_rows_per_file=100, + ) + + base_ids = [f.data_files()[0].base_id for f in dataset.get_fragments()] + assert base_ids == [None, 1, 2] + assert len(dataset.to_table()) == 300 + + def test_base_scoped_storage_options(self): + """base_. storage options flow through write and read.""" + data = self.create_test_data(200) + + # Local stores ignore these options; this verifies base-scoped entries + # are resolved per base without breaking the write or read path. + storage_options = { + "shared_option": "shared", + "base_1.scoped_option": "base1-value", + } + + dataset = lance.write_dataset( + data, + self.primary_uri, + mode="create", + initial_bases=[DatasetBasePath(self.path1_uri, name="path1")], + target_bases=["path1"], + max_rows_per_file=100, + storage_options=storage_options, + ) + assert dataset.count_rows() == 200 + + # Data files land in the scoped base, not the primary path + assert list(Path(self.path1_uri).glob("**/*.lance")) + assert not list((Path(self.primary_uri) / "data").glob("*.lance")) + + # Reopen with the same flat options and read through the base store + dataset = lance.dataset(self.primary_uri, storage_options=storage_options) + result = dataset.to_table().to_pandas() + pd.testing.assert_frame_equal( + result.sort_values("id").reset_index(drop=True), + data.sort_values("id").reset_index(drop=True), + ) + def test_multi_base_append_mode(self): """Test appending data to a multi-base dataset.""" # Create initial dataset @@ -1248,3 +1302,307 @@ def test_write_fragments_create_mode_with_initial_bases(self): dataset_root = Path(dataset_uri) data_files_root = list(dataset_root.glob("*.lance")) assert len(data_files_root) == 0, "Should not have data files in root" + + +class TestDataReplacementWithBases: + """DataReplacement must preserve DataFile.base_id so replacement files can + live in (and resolve against) a storage base other than the dataset root.""" + + def setup_method(self): + self.test_dir = tempfile.mkdtemp() + self.primary_uri = str(Path(self.test_dir) / "primary") + self.base1_uri = str(Path(self.test_dir) / "base1") + Path(self.base1_uri).mkdir(parents=True, exist_ok=True) + + def teardown_method(self): + if hasattr(self, "test_dir"): + shutil.rmtree(self.test_dir, ignore_errors=True) + + def _make_two_base_dataset(self) -> "lance.LanceDataset": + """Fragment 0 at the dataset root, fragment 1 in base1; column a.""" + ds = lance.write_dataset( + pa.table({"a": list(range(8))}), + self.primary_uri, + max_rows_per_file=8, + ) + ds.add_bases([DatasetBasePath(self.base1_uri, name="b1", is_dataset_root=True)]) + return lance.write_dataset( + pa.table({"a": list(range(8, 16))}), + self.primary_uri, + mode="append", + max_rows_per_file=8, + target_bases=["b1"], + ) + + def _write_bare_file(self, data_dir: str, data: pa.Table) -> str: + from lance.file import LanceFileWriter + + name = f"{uuid.uuid4()}.lance" + Path(data_dir).mkdir(parents=True, exist_ok=True) + with LanceFileWriter(f"{data_dir}/{name}") as writer: + writer.write_batch(data) + return name + + def _b_data_file(self, name: str, base_id=None) -> "lance.fragment.DataFile": + from lance.file import stable_version + from lance.fragment import DataFile + + return DataFile( + path=name, + fields=[1], # field id of column "b" (a=0, b=1) + column_indices=[0], + file_major_version=int(stable_version().split(".")[0]), + file_minor_version=int(stable_version().split(".")[1]), + base_id=base_id, + ) + + def test_data_replacement_new_column_into_base(self): + """The all-NULL-column special case: the new column's data file for a + base fragment is written into that base and must keep its base_id.""" + ds = self._make_two_base_dataset() + ds.add_columns(pa.field("b", pa.int32())) + ds = lance.dataset(self.primary_uri) + + root_name = self._write_bare_file( + f"{self.primary_uri}/data", + pa.table({"b": pa.array([x * 10 for x in range(8)], pa.int32())}), + ) + base_name = self._write_bare_file( + f"{self.base1_uri}/data", + pa.table({"b": pa.array([x * 10 for x in range(8, 16)], pa.int32())}), + ) + + op = lance.LanceOperation.DataReplacement( + [ + lance.LanceOperation.DataReplacementGroup( + 0, self._b_data_file(root_name) + ), + lance.LanceOperation.DataReplacementGroup( + 1, self._b_data_file(base_name, base_id=1) + ), + ] + ) + ds = lance.LanceDataset.commit(self.primary_uri, op, read_version=ds.version) + + frags = ds.get_fragments() + root_files = {f.path: f.base_id for f in frags[0].data_files()} + base_files = {f.path: f.base_id for f in frags[1].data_files()} + assert root_files[root_name] is None + assert base_files[base_name] == 1 + + table = ds.to_table() + assert table.column("b").to_pylist() == [x * 10 for x in range(16)] + + def test_data_replacement_replace_existing_file_in_base(self): + """The replace-existing-file branch must take the new file's base_id.""" + ds = self._make_two_base_dataset() + ds.add_columns(pa.field("b", pa.int32())) + ds = lance.dataset(self.primary_uri) + + first = self._write_bare_file( + f"{self.base1_uri}/data", + pa.table({"b": pa.array([0] * 8, pa.int32())}), + ) + op = lance.LanceOperation.DataReplacement( + [lance.LanceOperation.DataReplacementGroup(1, self._b_data_file(first, 1))] + ) + ds = lance.LanceDataset.commit(self.primary_uri, op, read_version=ds.version) + + # Replace the same column file again, still in base1. + second = self._write_bare_file( + f"{self.base1_uri}/data", + pa.table({"b": pa.array([x * 10 for x in range(8, 16)], pa.int32())}), + ) + op = lance.LanceOperation.DataReplacement( + [lance.LanceOperation.DataReplacementGroup(1, self._b_data_file(second, 1))] + ) + ds = lance.LanceDataset.commit(self.primary_uri, op, read_version=ds.version) + + base_files = {f.path: f.base_id for f in ds.get_fragments()[1].data_files()} + assert base_files[second] == 1 + assert first not in base_files + assert ds.to_table().column("b").to_pylist()[8:] == [ + x * 10 for x in range(8, 16) + ] + + +class TestMergeInsertMultiBase: + """Test merge insert on multi-base datasets.""" + + def setup_method(self): + """Set up test directories for each test.""" + self.test_dir = tempfile.mkdtemp() + self.primary_uri = str(Path(self.test_dir) / "primary") + self.base1_uri = str(Path(self.test_dir) / "base1") + self.base2_uri = str(Path(self.test_dir) / "base2") + for uri in [self.primary_uri, self.base1_uri, self.base2_uri]: + Path(uri).mkdir(parents=True, exist_ok=True) + + def teardown_method(self): + """Clean up test directories after each test.""" + if hasattr(self, "test_dir"): + shutil.rmtree(self.test_dir, ignore_errors=True) + + def create_dataset(self): + """Dataset with two registered bases and initial data in base1.""" + initial_data = pd.DataFrame( + { + "id": range(100), + "value": [f"initial_{i}" for i in range(100)], + } + ) + return lance.write_dataset( + initial_data, + self.primary_uri, + mode="create", + initial_bases=[ + DatasetBasePath(self.base1_uri, name="base1"), + DatasetBasePath(self.base2_uri, name="base2"), + ], + target_bases=["base1"], + max_rows_per_file=50, + ) + + def base_name_of(self, dataset, data_file): + base_paths = dataset._ds.base_paths() + if data_file.base_id is None: + return None + return base_paths[data_file.base_id].name + + def test_merge_insert_without_target_bases(self): + """Merge insert on a multi-base dataset writes to primary by default.""" + dataset = self.create_dataset() + + new_data = pd.DataFrame( + { + "id": range(50, 150), + "value": [f"updated_{i}" for i in range(50, 150)], + } + ) + stats = ( + dataset.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(new_data) + ) + assert stats["num_updated_rows"] == 50 + assert stats["num_inserted_rows"] == 50 + + result = dataset.to_table().to_pandas().sort_values("id") + assert len(result) == 150 + assert list(result[result["id"] >= 50]["value"]) == [ + f"updated_{i}" for i in range(50, 150) + ] + + # New fragments (merge output) are in primary storage. + max_initial_fragment_id = 1 # 100 rows / 50 per file -> fragments 0, 1 + for fragment in dataset.get_fragments(): + is_initial = fragment.fragment_id <= max_initial_fragment_id + for data_file in fragment.data_files(): + expected = "base1" if is_initial else None + assert self.base_name_of(dataset, data_file) == expected + + def test_merge_insert_with_target_bases(self): + """Merge insert routes new fragments to the requested base.""" + dataset = self.create_dataset() + + new_data = pd.DataFrame( + { + "id": range(50, 150), + "value": [f"updated_{i}" for i in range(50, 150)], + } + ) + stats = ( + dataset.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .target_bases(["base2"]) + .execute(new_data) + ) + assert stats["num_updated_rows"] == 50 + assert stats["num_inserted_rows"] == 50 + + result = dataset.to_table().to_pandas().sort_values("id") + assert len(result) == 150 + assert list(result[result["id"] >= 50]["value"]) == [ + f"updated_{i}" for i in range(50, 150) + ] + + merge_files = 0 + for fragment in dataset.get_fragments(): + if fragment.fragment_id > 1: + for data_file in fragment.data_files(): + assert self.base_name_of(dataset, data_file) == "base2" + merge_files += 1 + assert merge_files > 0 + assert list(Path(self.base2_uri).glob("**/*.lance")) + + # The routed dataset stays readable from a fresh instance. + reloaded = lance.dataset(self.primary_uri) + assert reloaded.count_rows() == 150 + + def test_merge_insert_with_unknown_target_base(self): + """Merge insert referencing an unregistered base fails.""" + dataset = self.create_dataset() + + new_data = pd.DataFrame({"id": [1], "value": ["updated_1"]}) + with pytest.raises(Exception, match="not found in available bases"): + ( + dataset.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .target_bases(["nonexistent"]) + .execute(new_data) + ) + + def test_merge_insert_target_primary_via_uri(self): + """The dataset URI in target_bases selects primary storage.""" + dataset = self.create_dataset() + + new_data = pd.DataFrame({"id": [200], "value": ["inserted_200"]}) + ( + dataset.merge_insert("id") + .when_not_matched_insert_all() + .target_bases([dataset.uri, "base2"]) + .execute(new_data) + ) + assert dataset.count_rows() == 101 + + # The single new file lands in the first slot: primary storage. + for fragment in dataset.get_fragments(): + if fragment.fragment_id > 1: + for data_file in fragment.data_files(): + assert self.base_name_of(dataset, data_file) is None + + def test_merge_insert_target_all_bases(self): + """target_all_bases spreads new files across all bases, primary first.""" + dataset = self.create_dataset() + + new_data = pd.DataFrame({"id": [300], "value": ["inserted_300"]}) + ( + dataset.merge_insert("id") + .when_not_matched_insert_all() + .target_all_bases() + .execute(new_data) + ) + assert dataset.count_rows() == 101 + # A single new file lands in the first slot: primary storage. + newest = max(f.fragment_id for f in dataset.get_fragments()) + for fragment in dataset.get_fragments(): + if fragment.fragment_id == newest: + for data_file in fragment.data_files(): + assert self.base_name_of(dataset, data_file) is None + + new_data = pd.DataFrame({"id": [301], "value": ["inserted_301"]}) + ( + dataset.merge_insert("id") + .when_not_matched_insert_all() + .target_all_bases(include_primary=False) + .execute(new_data) + ) + assert dataset.count_rows() == 102 + newest = max(f.fragment_id for f in dataset.get_fragments()) + for fragment in dataset.get_fragments(): + if fragment.fragment_id == newest: + for data_file in fragment.data_files(): + assert self.base_name_of(dataset, data_file) == "base1" diff --git a/python/python/tests/test_otel.py b/python/python/tests/test_otel.py new file mode 100644 index 00000000000..a96e0125f56 --- /dev/null +++ b/python/python/tests/test_otel.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +import lance +import pyarrow as pa +import pytest + +# The metrics recorder is process-global and installed once, so the whole +# bridge is exercised in a single test to avoid cross-test global-state coupling. + + +def _metrics_by_name(reader): + data = reader.get_metrics_data() + result = {} + for resource_metrics in data.resource_metrics: + for scope_metrics in resource_metrics.scope_metrics: + for metric in scope_metrics.metrics: + result[metric.name] = metric + return result + + +def test_instrument_lance_metrics_exports_object_store_metrics(tmp_path): + pytest.importorskip("opentelemetry.sdk.metrics") + from lance.otel import instrument_lance_metrics + from opentelemetry.sdk.metrics import MeterProvider + from opentelemetry.sdk.metrics.export import InMemoryMetricReader + + reader = InMemoryMetricReader() + provider = MeterProvider(metric_readers=[reader]) + assert instrument_lance_metrics(provider) + + # The catalog is populated once the recorder is installed. + from lance.lance import lance_metrics_catalog + + catalog = {desc.name: desc for desc in lance_metrics_catalog()} + assert "lance_object_store_requests_total" in catalog + assert catalog["lance_object_store_request_duration_seconds"].kind == "histogram" + # Gauges and the retryable counter are described too, so they surface in the + # export even when a plain write doesn't happen to emit them. + assert catalog["lance_object_store_retryable_responses_total"].kind == "counter" + assert catalog["lance_object_store_in_flight_requests"].kind == "gauge" + + # Generate object store activity on the local filesystem (scheme "file"). + table = pa.table({"id": pa.array(range(256))}) + dataset = lance.write_dataset(table, str(tmp_path / "ds.lance")) + assert dataset.to_table().num_rows == 256 + + metrics = _metrics_by_name(reader) + + requests = metrics["lance_object_store_requests_total"] + points = list(requests.data.data_points) + assert points, "expected at least one request data point" + # The `base` label carries the store scheme ("file") by default. + assert all("base" in p.attributes and "operation" in p.attributes for p in points) + assert sum(p.value for p in points) > 0 + + # Histograms are decomposed into bucket / count / sum observable counters. + bucket = metrics["lance_object_store_request_duration_seconds_bucket"] + bucket_points = list(bucket.data.data_points) + assert bucket_points + assert all("le" in p.attributes for p in bucket_points) + # The implicit +Inf bucket must be present and is the cumulative maximum. + assert any(p.attributes["le"] == "+Inf" for p in bucket_points) + + count = metrics["lance_object_store_request_duration_seconds_count"] + assert sum(p.value for p in count.data.data_points) > 0 + + # The `_sum` instrument must also be wired and report positive latency. + duration_sum = metrics["lance_object_store_request_duration_seconds_sum"] + assert sum(p.value for p in duration_sum.data.data_points) > 0 + + +def test_snapshot_empty_before_install_is_safe(): + # snapshot is callable regardless of installation state and never raises. + from lance.lance import snapshot_lance_metrics + + assert isinstance(snapshot_lance_metrics(), list) + + +def test_instrument_warns_when_recorder_unavailable(monkeypatch): + # A foreign `metrics` recorder already installed -> register returns False; + # instrument_lance_metrics must warn and return False without instrumenting. + pytest.importorskip("opentelemetry.sdk.metrics") + import lance.otel as otel + from opentelemetry.sdk.metrics import MeterProvider + from opentelemetry.sdk.metrics.export import InMemoryMetricReader + + monkeypatch.setattr(otel, "register_lance_metrics_recorder", lambda: False) + + reader = InMemoryMetricReader() + provider = MeterProvider(metric_readers=[reader]) + with pytest.warns(UserWarning, match="recorder"): + assert otel.instrument_lance_metrics(provider) is False diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index 4450f6821cc..985072fe5de 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -949,6 +949,39 @@ def test_create_scalar_index_fts_alias(dataset): assert any(idx.index_type == "Inverted" for idx in dataset.describe_indices()) +def test_create_scalar_index_fts_block_size(dataset): + dataset.create_scalar_index( + "doc", index_type="INVERTED", with_position=False, block_size=256 + ) + indices = dataset.describe_indices() + doc_index = next(index for index in indices if index.name == "doc_idx") + assert doc_index.segments[0].index_version == 3 + + row = dataset.take(indices=[0], columns=["doc"]) + query = row.column(0)[0].as_py().split(" ")[0] + results = dataset.scanner(columns=["doc"], full_text_query=query).to_table() + assert results.num_rows > 0 + + with pytest.raises(ValueError, match="block_size"): + dataset.create_scalar_index( + "doc", index_type="INVERTED", name="doc_invalid_129", block_size=129 + ) + + with pytest.raises(ValueError, match="block_size"): + dataset.create_scalar_index( + "doc", index_type="INVERTED", name="doc_invalid_512", block_size=512 + ) + + with pytest.raises(ValueError, match="block_size=256"): + dataset.create_scalar_index( + "doc", + index_type="INVERTED", + name="doc_invalid_v2_256", + block_size=256, + format_version=2, + ) + + def test_multi_index_create(tmp_path): dataset = lance.write_dataset( pa.table({"ints": range(1024)}), tmp_path, max_rows_per_file=100 @@ -2165,6 +2198,87 @@ def scan_stats_callback(stats: lance.ScanStatistics): assert small_bytes_read < large_bytes_read +@pytest.mark.parametrize("index_type", ["ZONEMAP", "BLOOMFILTER"]) +def test_address_domain_index_with_stable_row_ids(tmp_path: Path, index_type): + """Regression test for issue #7434. + + Address-domain scalar indices (zonemap, bloom filter) report matches as + physical row addresses. On a stable-row-id dataset a row's stable id differs + from its physical address in every fragment except fragment 0, so the index + result must be translated back to the row-id domain before it prefilters the + scan. Without that translation the index silently drops matching rows in + fragments other than fragment 0 (often returning an empty result). + """ + + # A single value "a" is placed in non-adjacent fragments (0, 2, 4) so its + # matching rows span multiple fragments and diverge from fragment 0. + def block(v, n): + return [v] * n + + vals = ( + block("a", 5_000) + + block("b", 5_000) + + block("a", 5_000) + + block("c", 5_000) + + block("a", 5_000) + ) + tbl = pa.table({"category": pa.array(vals, pa.string()), "id": range(len(vals))}) + ds = lance.write_dataset( + tbl, tmp_path, max_rows_per_file=5_000, enable_stable_row_ids=True + ) + assert ds.has_stable_row_ids + assert len(ds.get_fragments()) == 5 + + true_count = ds.scanner( + filter="category = 'a'", use_scalar_index=False + ).count_rows() + assert true_count == 15_000 + + ds.create_scalar_index("category", index_type=index_type, replace=True) + + # The index must be consulted and must return the same rows as a full scan. + assert "ScalarIndexQuery" in ds.scanner(filter="category = 'a'").explain_plan() + indexed = ds.to_table(filter="category = 'a'", columns=["id"]) + assert indexed.num_rows == true_count + assert sorted(indexed["id"].to_pylist()) == sorted( + ds.to_table(filter="category = 'a'", columns=["id"], use_scalar_index=False)[ + "id" + ].to_pylist() + ) + + +def test_zonemap_with_stable_row_ids_after_compaction(tmp_path: Path): + """Zonemap results stay correct after a compaction relocates rows under + stable row ids (physical address != stable id for the surviving rows).""" + ds = lance.write_dataset( + pa.table({"x": range(0, 5_000)}), + tmp_path, + max_rows_per_file=5_000, + enable_stable_row_ids=True, + ) + ds = lance.write_dataset( + pa.table({"x": range(5_000, 10_000)}), + tmp_path, + mode="append", + max_rows_per_file=5_000, + enable_stable_row_ids=True, + ) + # Delete part of the first fragment, then compact so surviving rows keep + # their small stable ids but move to a freshly numbered fragment. + ds.delete("x >= 1000 AND x < 2000") + ds.optimize.compact_files(target_rows_per_fragment=100_000) + ds = lance.dataset(tmp_path) + assert len(ds.get_fragments()) == 1 + + ds.create_scalar_index("x", index_type="ZONEMAP") + + filter_expr = "x >= 6000 AND x <= 6500" + assert "ScalarIndexQuery" in ds.scanner(filter=filter_expr).explain_plan() + expected = ds.to_table(filter=filter_expr, use_scalar_index=False)["x"].to_pylist() + actual = ds.to_table(filter=filter_expr)["x"].to_pylist() + assert sorted(actual) == sorted(expected) == list(range(6000, 6501)) + + def test_zonemap_deletion_handling(tmp_path: Path): """Test zonemap deletion handling""" data = pa.table( @@ -2194,6 +2308,33 @@ def test_zonemap_deletion_handling(tmp_path: Path): assert ids == [0, 2, 4, 6, 8] +@pytest.mark.parametrize("index_type", ["ZONEMAP", "BLOOMFILTER"]) +def test_address_domain_index_not_query_with_stable_row_ids(tmp_path: Path, index_type): + """Regression test: != queries return correct results on stable-row-id datasets. + + Address-domain indices (zonemap, bloom filter) search returns physical row + addresses. Translation to row IDs happens at the Query leaf before the NOT + node is evaluated, so the NOT operates on a correctly translated AllowList. + Without the address-to-row-id translation the AllowList contains wrong IDs, + and the subsequent NOT excludes the wrong rows. + """ + vals = list(range(5_000)) + list(range(5_000, 10_000)) + tbl = pa.table({"x": vals}) + ds = lance.write_dataset( + tbl, tmp_path, max_rows_per_file=5_000, enable_stable_row_ids=True + ) + assert ds.has_stable_row_ids + + ds.create_scalar_index("x", index_type=index_type, replace=True) + + # Without address translation the NOT excludes the wrong rows, producing an + # incorrect result set rather than crashing. + expected = ds.to_table(filter="x != 42", use_scalar_index=False)["x"].to_pylist() + actual = ds.to_table(filter="x != 42")["x"].to_pylist() + assert sorted(actual) == sorted(expected) + assert len(actual) == 9_999 + + def test_zonemap_index_remapping(tmp_path: Path): """Test zonemap index remapping after compaction and optimization""" # Create a dataset with 5 fragments by writing data in chunks @@ -2793,6 +2934,15 @@ def scan_stats_callback(stats: lance.ScanStatistics): cache_entries_after_query = ds._ds.index_cache_entry_count() assert cache_entries_after_query == cache_entries_after_prewarm + segment_uuid = ds.describe_indices()[0].segments[0].uuid + ds = lance.dataset(phrase_path) + ds.prewarm_index("fts_idx", with_position=True, index_segments=[segment_uuid]) + cache_entries_after_prewarm = ds._ds.index_cache_entry_count() + results = ds.to_table(full_text_query=PhraseQuery("word word", "fts")) + assert results.num_rows == test_table_size + cache_entries_after_query = ds._ds.index_cache_entry_count() + assert cache_entries_after_query == cache_entries_after_prewarm + with pytest.raises( TypeError, match="takes 2 positional arguments", @@ -4719,6 +4869,77 @@ def test_nested_field_fts_index(tmp_path): assert results.num_rows == 50 +def test_multiple_nested_field_fts_indices_e2e(tmp_path): + """Test FTS queries against multiple indexed nested string fields.""" + + def make_table(ids, text_values, summary_values): + return pa.table( + { + "id": ids, + "data": pa.StructArray.from_arrays( + [ + pa.array(text_values, type=pa.string()), + pa.array(summary_values, type=pa.string()), + ], + names=["text", "summary"], + ), + } + ) + + def result_ids(query): + return sorted(ds.to_table(full_text_query=query)["id"].to_pylist()) + + ds = lance.write_dataset( + make_table( + [0, 1, 2, 3], + [ + "lance nested alpha", + "plain text", + None, + "phrase target here", + ], + [ + "metadata only", + "database nested beta", + "lance beta", + "other", + ], + ), + tmp_path, + ) + + ds.create_scalar_index("data.text", index_type="INVERTED", with_position=True) + ds.create_scalar_index("data.summary", index_type="INVERTED", with_position=False) + + indexed_fields = { + tuple(index.field_names) + for index in ds.describe_indices() + if index.index_type == "Inverted" + } + assert indexed_fields == {("data.text",), ("data.summary",)} + + assert result_ids(MatchQuery("alpha", "data.text")) == [0] + assert result_ids(MatchQuery("beta", "data.summary")) == [1, 2] + assert result_ids("lance") == [0, 2] + assert result_ids(MultiMatchQuery("nested", ["data.text", "data.summary"])) == [ + 0, + 1, + ] + assert result_ids(PhraseQuery("phrase target", "data.text")) == [3] + + ds = lance.write_dataset( + make_table( + [4, 5], + ["fresh lance append", "plain append"], + ["other", "fresh beta append"], + ), + tmp_path, + mode="append", + ) + + assert result_ids("fresh") == [4, 5] + + def test_nested_field_bitmap_index(tmp_path): """Test BITMAP index creation and querying on nested fields""" # Create dataset with nested categorical field @@ -4851,9 +5072,11 @@ def test_json_inverted_match_query(tmp_path): assert results.num_rows == 1 -@pytest.mark.parametrize("fts_format_version", ["1", "2"]) -def test_describe_indices(tmp_path, monkeypatch, fts_format_version): - monkeypatch.setenv("LANCE_FTS_FORMAT_VERSION", fts_format_version) +@pytest.mark.parametrize( + ("format_version", "expected_format_version"), + [(1, 1), (2, 2), ("v1", 1), ("v2", 2)], +) +def test_describe_indices(tmp_path, format_version, expected_format_version): data = pa.table( { "id": range(100), @@ -4869,7 +5092,7 @@ def test_describe_indices(tmp_path, monkeypatch, fts_format_version): } ) ds = lance.write_dataset(data, tmp_path) - ds.create_scalar_index("text", index_type="INVERTED") + ds.create_scalar_index("text", index_type="INVERTED", format_version=format_version) indices = ds.describe_indices() assert len(indices) == 1 @@ -4883,7 +5106,7 @@ def test_describe_indices(tmp_path, monkeypatch, fts_format_version): assert indices[0].segments[0].uuid is not None assert indices[0].segments[0].fragment_ids == {0} assert indices[0].segments[0].dataset_version_at_last_update == 1 - assert indices[0].segments[0].index_version == int(fts_format_version) + assert indices[0].segments[0].index_version == expected_format_version assert indices[0].segments[0].created_at is not None assert isinstance(indices[0].segments[0].created_at, datetime) assert indices[0].segments[0].size_bytes is not None @@ -4982,6 +5205,28 @@ def test_describe_indices(tmp_path, monkeypatch, fts_format_version): assert index.num_rows_indexed == 50 +def test_create_inverted_index_defaults_to_v2_and_ignores_env(tmp_path, monkeypatch): + monkeypatch.setenv("LANCE_FTS_FORMAT_VERSION", "1") + data = pa.table({"text": ["document about lance database"]}) + ds = lance.write_dataset(data, tmp_path) + + ds.create_scalar_index("text", index_type="INVERTED") + + indices = ds.describe_indices() + assert indices[0].segments[0].index_version == 2 + + +def test_create_inverted_index_rejects_invalid_format_version(tmp_path): + data = pa.table({"text": ["document about lance database"]}) + ds = lance.write_dataset(data, tmp_path) + + with pytest.raises(ValueError, match="unsupported FTS format version"): + ds.create_scalar_index("text", index_type="INVERTED", format_version="v4") + + with pytest.raises(ValueError, match="format_version=3"): + ds.create_scalar_index("text", index_type="INVERTED", format_version="v3") + + def test_vector_filter_fts_search(tmp_path): # Create test data ids = list(range(1, 301)) @@ -5099,3 +5344,28 @@ def test_vector_filter_fts_search(tmp_path): ) with pytest.raises(ValueError): scanner.to_table() + + +@pytest.mark.parametrize("index_type", ["BTREE", "BITMAP", "ZONEMAP"]) +def test_large_string_scalar_index(tmp_path, index_type): + """large_string (LargeUtf8) must be accepted by BTREE, BITMAP, and ZONEMAP.""" + table = pa.table( + { + "id": pa.array([1, 2, 3], pa.int32()), + "category": pa.array(["alpha", "beta", "alpha"], pa.large_string()), + } + ) + ds = lance.write_dataset(table, tmp_path) + assert ds.schema.field("category").type == pa.large_string() + + # Must not raise TypeError + ds.create_scalar_index("category", index_type=index_type) + + indices = ds.describe_indices() + assert any("category" in idx.field_names for idx in indices), ( + f"{index_type} index for large_string column not found in describe_indices()" + ) + + result = ds.scanner(filter="category = 'alpha'").to_table() + assert result.num_rows == 2 + assert set(result.column("id").to_pylist()) == {1, 3} diff --git a/python/python/tests/test_tf.py b/python/python/tests/test_tf.py deleted file mode 100644 index 3652df0e938..00000000000 --- a/python/python/tests/test_tf.py +++ /dev/null @@ -1,351 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright The Lance Authors - -import os -import warnings - -import lance -import numpy as np -import pandas as pd -import pyarrow as pa -import pytest -from lance.arrow import ImageArray -from lance.fragment import LanceFragment - -pytest.skip("Skip tensorflow tests", allow_module_level=True) - -try: - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", category=DeprecationWarning) - import tensorflow as tf # noqa: F401 -except ImportError: - pytest.skip( - "Tensorflow is not installed. Please install tensorflow to " - + "test lance.tf module.", - allow_module_level=True, - ) - -from lance.tf.data import ( # noqa: E402 - from_lance, - from_lance_batches, - lance_fragments, - lance_take_batches, -) - - -@pytest.fixture -def tf_dataset(tmp_path): - df = pd.DataFrame( - { - "a": range(10000), - "s": [f"val-{i}" for i in range(10000)], - "vec": [[i * 0.2] * 128 for i in range(10000)], - } - ) - - schema = pa.schema( - [ - pa.field("a", pa.int64()), - pa.field("s", pa.string()), - pa.field("vec", pa.list_(pa.float32(), 128)), - ] - ) - tbl = pa.Table.from_pandas(df, schema=schema) - uri = tmp_path / "dataset.lance" - lance.write_dataset( - tbl, - uri, - schema=tbl.schema, - max_rows_per_group=100, - max_rows_per_file=1000, - ) - return uri - - -def test_fragment_dataset(tf_dataset): - ds = from_lance(tf_dataset, batch_size=100) - for idx, batch in enumerate(ds): - assert batch["a"].numpy()[0] == idx * 100 - assert batch["s"].numpy()[0] == f"val-{idx * 100}".encode("utf-8") - assert batch["a"].shape == (100,) - assert batch["vec"].shape == ( - 100, - 128, - ) # Fixed size list - - -def test_projection(tf_dataset): - ds = from_lance(tf_dataset, batch_size=100, columns=["a"]) - - for idx, batch in enumerate(ds): - assert list(batch.keys()) == ["a"] - assert batch["a"].numpy()[0] == idx * 100 - assert batch["a"].shape == (100,) - - -def test_filter(tf_dataset): - ds = from_lance(tf_dataset, batch_size=100, filter="a >= 5000") - - for idx, batch in enumerate(ds): - assert batch["a"].numpy()[0] == idx * 100 + 5000 - assert batch["a"].shape == (100,) - - -def test_namespace_table_id(monkeypatch): - calls = {} - - class DummyScanner: - def __init__(self): - self._batch = pa.record_batch([pa.array([1, 2])], names=["a"]) - self.projected_schema = self._batch.schema - - def to_batches(self): - yield self._batch - - class DummyDataset: - def scanner(self, **kwargs): - return DummyScanner() - - def fake_dataset(uri=None, **kwargs): - calls["uri"] = uri - calls["kwargs"] = kwargs - return DummyDataset() - - monkeypatch.setattr(lance, "dataset", fake_dataset) - - ns = object() - ds = from_lance( - None, - namespace_client=ns, - table_id=["tbl"], - ignore_namespace_table_storage_options=True, - ) - - assert calls["kwargs"]["namespace_client"] is ns - assert calls["kwargs"]["table_id"] == ["tbl"] - assert calls["kwargs"]["ignore_namespace_table_storage_options"] is True - - batches = list(ds) - assert [b["a"].numpy().tolist() for b in batches] == [[1, 2]] - - -def test_scan_use_tf_data(tf_dataset): - ds = tf.data.Dataset.from_lance(tf_dataset) - for idx, batch in enumerate(ds): - assert batch["a"].numpy()[0] == idx * 100 - assert batch["s"].numpy()[0] == f"val-{idx * 100}".encode("utf-8") - assert batch["a"].shape == (100,) - assert batch["vec"].shape == ( - 100, - 128, - ) # Fixed size list - - -def test_pass_fragments(tf_dataset): - # Can pass fragments directly to from_lance - dataset = lance.dataset(tf_dataset) - ds = from_lance(tf_dataset, fragments=dataset.get_fragments(), batch_size=100) - ds_default = from_lance(tf_dataset, batch_size=100) - for batch, batch_default in zip(ds, ds_default): - assert batch["a"].numpy()[0] == batch_default["a"].numpy()[0] - assert batch["a"].numpy().shape == (100,) - assert batch["vec"].shape == ( - 100, - 128, - ) - - # Can pass ids directly to from_lance - ds = from_lance(tf_dataset, fragments=[0, 1, 2], batch_size=100) - for idx, batch in enumerate(ds): - assert batch["a"].numpy()[0] == idx * 100 - assert batch["a"].numpy().shape == (100,) - - -def test_shuffle(tf_dataset): - fragments = lance_fragments(tf_dataset).shuffle(4, seed=20).take(3) - - ds = from_lance(tf_dataset, fragments=fragments, batch_size=100) - raw_ds = lance.dataset(tf_dataset) - scanner = raw_ds.scanner( - fragments=[LanceFragment(raw_ds, fid) for fid in [0, 3, 1]], batch_size=100 - ) - - for batch, raw_batch in zip(ds, scanner.to_batches()): - assert batch["a"].numpy()[0] == raw_batch.to_pydict()["a"][0] - assert batch["a"].numpy().shape == (100,) - assert batch["vec"].shape == ( - 100, - 128, - ) # Fixed size list - - -def test_dataset_batches(tf_dataset): - tf_dataset = lance.dataset(tf_dataset) - batch_size = 300 - batches = list( - from_lance_batches(tf_dataset, batch_size=batch_size).as_numpy_iterator() - ) - assert tf_dataset.count_rows() // batch_size + 1 == len(batches) - assert all(end - start == batch_size for start, end in batches[:-2]) - assert batches[-1][1] - batches[-1][0] == tf_dataset.count_rows() % batch_size - - skip = 5 - batches_skipped = list( - from_lance_batches( - tf_dataset, batch_size=batch_size, skip=skip - ).as_numpy_iterator() - ) - assert batches_skipped == batches[skip:] - - batches_shuffled = list( - from_lance_batches( - tf_dataset, batch_size=batch_size, shuffle=True, seed=42 - ).as_numpy_iterator() - ) - # make sure it does a shuffle - assert batches_shuffled != batches - batches_shuffled2 = list( - from_lance_batches( - tf_dataset, batch_size=batch_size, shuffle=True, seed=42 - ).as_numpy_iterator() - ) - # make sure the shuffle can be deterministic - assert batches_shuffled == batches_shuffled2 - - -def test_take_dataset(tf_dataset): - tf_dataset = lance.dataset(tf_dataset) - batch_ds = from_lance_batches( - tf_dataset, batch_size=100, shuffle=True, seed=42 - ).as_numpy_iterator() - lance_ds = lance_take_batches(tf_dataset, batch_ds) - lance_ds = lance_ds.unbatch().shuffle(400, seed=42).batch(100) - - for batch in lance_ds: - assert batch["a"].numpy().shape == (100,) - - batches = [(0, 200), (100, 200)] - lance_ds = lance_take_batches(tf_dataset, batches, columns=["a"]) - for (start, end), batch in zip(batches, lance_ds): - assert batch["a"].numpy().tolist() == np.arange(start, end).tolist() - assert batch.keys() == {"a"} - - -def test_var_length_list(tmp_path): - """Treat var length list as RaggedTensor.""" - df = pd.DataFrame( - { - "a": range(200), - "l": [[i] * (i % 5 + 1) for i in range(200)], - } - ) - - schema = pa.schema( - [ - pa.field("a", pa.int64()), - pa.field("l", pa.list_(pa.int32())), - ] - ) - tbl = pa.Table.from_pandas(df, schema=schema) - - uri = tmp_path / "dataset.lance" - lance.write_dataset( - tbl, - uri, - schema=tbl.schema, - ) - - output_signature = { - "a": tf.TensorSpec(shape=(None,), dtype=tf.int64), - "l": tf.RaggedTensorSpec(dtype=tf.dtypes.int32, shape=(8, None), ragged_rank=1), - } - - ds = tf.data.Dataset.from_lance( - uri, - batch_size=8, - output_signature=output_signature, - ) - for idx, batch in enumerate(ds): - assert batch["a"].numpy()[0] == idx * 8 - assert batch["l"].shape == (8, None) - assert isinstance(batch["l"], tf.RaggedTensor) - - -def test_nested_struct(tmp_path): - table = pa.table( - { - "x": pa.array( - [ - { - "a": 1, - "json": {"b": "hello", "x": b"abc"}, - }, - { - "a": 24, - "json": {"b": "world", "x": b"def"}, - }, - ] - ) - } - ) - uri = tmp_path / "dataset.lance" - dataset = lance.write_dataset(table, uri) - - ds = tf.data.Dataset.from_lance( - dataset, - batch_size=8, - ) - - for batch in ds: - tf.debugging.assert_equal(batch["x"]["a"], tf.constant([1, 24], dtype=tf.int64)) - tf.debugging.assert_equal( - batch["x"]["json"]["b"], tf.constant(["hello", "world"]) - ) - tf.debugging.assert_equal( - batch["x"]["json"]["x"], tf.constant([b"abc", b"def"]) - ) - - -def test_tensor(tmp_path): - arr = np.array([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]], dtype=np.float32) - table = pa.table({"x": pa.FixedShapeTensorArray.from_numpy_ndarray(arr)}) - - uri = tmp_path / "dataset.lance" - dataset = lance.write_dataset(table, uri) - ds = tf.data.Dataset.from_lance(dataset) - - for batch in ds: - assert batch["x"].shape == (2, 2, 3) - assert batch["x"].dtype == tf.float32 - assert batch["x"].numpy().tolist() == arr.tolist() - - -def test_image_types(tmp_path): - path = [os.path.join(os.path.dirname(__file__), "images/1.png")] - uris = ImageArray.from_array(path * 3) - encoded_images = uris.read_uris() - tensors = encoded_images.to_tensor() - table = pa.table( - { - "uris": uris, - "encoded_images": encoded_images, - "tensor_images": tensors, - } - ) - - uri = tmp_path / "dataset.lance" - dataset = lance.write_dataset(table, uri) - ds = tf.data.Dataset.from_lance(dataset) - - for batch in ds: - assert batch["uris"].shape == (3,) - assert batch["uris"].dtype == tf.string - assert batch["uris"].numpy().astype("str").tolist() == uris.tolist() - - assert batch["encoded_images"].shape == (3,) - assert batch["encoded_images"].dtype == tf.string - assert batch["encoded_images"].numpy().tolist() == encoded_images.tolist() - - assert batch["tensor_images"].shape == (3, 1, 1, 4) - assert batch["tensor_images"].dtype == tf.uint8 - assert batch["tensor_images"].numpy().tolist() == tensors.to_numpy().tolist() diff --git a/python/python/tests/test_vector.py b/python/python/tests/test_vector.py index 4ea4e7d425e..3ec889d5127 100644 --- a/python/python/tests/test_vector.py +++ b/python/python/tests/test_vector.py @@ -5,7 +5,12 @@ import numpy as np import pyarrow as pa import pytest -from lance.vector import hamming_clustering_for_sample, vec_to_table +from lance.vector import ( + get_ivf_partition_info, + hamming_clustering_for_ivf_partition, + hamming_clustering_for_sample, + vec_to_table, +) def test_dict(): @@ -182,3 +187,72 @@ def test_hamming_clustering_for_sample(tmp_path): } # Singleton row 5 is not emitted as a cluster. assert clusters == {0: [1, 2], 3: [4]} + + +def test_hamming_clustering_multi_segment(tmp_path): + # 25 distinct hash values, two copies each; the same table is written to + # fragment 0 and appended as fragment 1. + values = [((i // 2) * 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF for i in range(50)] + table = _hash_table([list(value.to_bytes(8, "little")) for value in values]) + dataset = lance.write_dataset(table, tmp_path / "hashes") + dataset.create_index( + "hash", index_type="IVF_FLAT", num_partitions=4, metric="hamming" + ) + dataset = lance.write_dataset(table, tmp_path / "hashes", mode="append") + # Optimizing with merge disabled creates a delta segment for fragment 1. + dataset.optimize.optimize_indices(num_indices_to_merge=0) + + index = dataset.describe_indices()[0] + assert len(index.segments) == 2 + + infos = get_ivf_partition_info(dataset, index.name) + assert sum(info["size"] for info in infos) == 100 + + # All four copies of each value cluster together across both fragments. + frag1_start = 1 << 32 + clusters = [] + for info in infos: + result = hamming_clustering_for_ivf_partition( + dataset, index.name, info["partition_id"], 0 + ).read_all() + clusters.extend( + zip( + result["representative"].to_pylist(), + result["duplicates"].to_pylist(), + ) + ) + assert len(clusters) == 25 + for representative, duplicates in clusters: + assert representative < frag1_start + assert len(duplicates) == 3 + assert any(dup >= frag1_start for dup in duplicates) + + # Selecting the fragment-0 segment reproduces the single-segment scope. + first_segment = next( + segment for segment in index.segments if segment.fragment_ids == {0} + ) + infos = get_ivf_partition_info( + dataset, index.name, index_segments=[first_segment.uuid] + ) + assert sum(info["size"] for info in infos) == 50 + num_selected_clusters = 0 + for info in infos: + result = hamming_clustering_for_ivf_partition( + dataset, + index.name, + info["partition_id"], + 0, + index_segments=[first_segment.uuid], + ).read_all() + for duplicates in result["duplicates"].to_pylist(): + num_selected_clusters += 1 + assert duplicates == [dup for dup in duplicates if dup < frag1_start] + assert len(duplicates) == 1 + assert num_selected_clusters == 25 + + with pytest.raises(ValueError, match="invalid index segment uuid"): + get_ivf_partition_info(dataset, index.name, index_segments=["not-a-uuid"]) + with pytest.raises(TypeError, match="str or uuid.UUID"): + get_ivf_partition_info(dataset, index.name, index_segments=[123]) + with pytest.raises(TypeError, match="not a single"): + get_ivf_partition_info(dataset, index.name, index_segments=first_segment.uuid) diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 4e3addfedb8..fa2c4047cd0 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -576,6 +576,8 @@ def test_index_with_pq_codebook(tmp_path): ivf_centroids=np.random.randn(1, 128).astype(np.float32), pq_codebook=pq_codebook, ) + index = dataset.stats.index_stats("vector_idx") + assert index["indices"][0]["sub_index"]["nbits"] == 8 validate_vector_index(dataset, "vector", refine_factor=10, pass_threshold=0.99) pq_codebook = pa.FixedShapeTensorArray.from_numpy_ndarray(pq_codebook) @@ -592,6 +594,54 @@ def test_index_with_pq_codebook(tmp_path): validate_vector_index(dataset, "vector", refine_factor=10, pass_threshold=0.99) +def test_index_with_4bit_numpy_pq_codebook(tmp_path): + tbl = create_table(nvec=1024, ndim=128) + dataset = lance.write_dataset(tbl, tmp_path) + pq_codebook = np.random.randn(4, 16, 128 // 4).astype(np.float32) + + dataset = dataset.create_index( + "vector", + index_type="IVF_PQ", + num_partitions=1, + num_sub_vectors=4, + num_bits=4, + ivf_centroids=np.random.randn(1, 128).astype(np.float32), + pq_codebook=pq_codebook, + ) + + index = dataset.stats.index_stats("vector_idx") + assert index["indices"][0]["sub_index"]["nbits"] == 4 + + result = dataset.to_table( + nearest={ + "column": "vector", + "q": np.random.randn(128).astype(np.float32), + "k": 10, + } + ) + assert result.num_rows == 10 + + +def test_index_with_pq_codebook_rejects_wrong_num_bits_shape(tmp_path): + tbl = create_table(nvec=8, ndim=128) + dataset = lance.write_dataset(tbl, tmp_path) + pq_codebook = np.random.randn(4, 256, 128 // 4).astype(np.float32) + + with pytest.raises( + ValueError, + match=r"\(sub_vectors, 16, dim\) for num_bits=4, got \(4, 256, 32\)", + ): + dataset.create_index( + "vector", + index_type="IVF_PQ", + num_partitions=1, + num_sub_vectors=4, + num_bits=4, + ivf_centroids=np.random.randn(1, 128).astype(np.float32), + pq_codebook=pq_codebook, + ) + + @pytest.mark.cuda @pytest.mark.parametrize("nullify", [False, True]) def test_create_index_using_cuda(tmp_path, nullify): diff --git a/python/src/blob.rs b/python/src/blob.rs new file mode 100644 index 00000000000..82e8a01ae8a --- /dev/null +++ b/python/src/blob.rs @@ -0,0 +1,267 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use crate::{error::PythonErrorExt, rt}; +use arrow::pyarrow::ToPyArrow; +use bytes::Bytes; +use lance::{ + BlobDescriptor, BlobDescriptorArrayBuilder, BlobRange, DedicatedBlobWriter, PackedBlobWriter, +}; +use pyo3::{ + Bound, PyResult, + exceptions::PyValueError, + pyclass, pymethods, + types::{PyAny, PyAnyMethods, PyDict, PyList, PyListMethods, PyModule}, +}; +use std::sync::Arc; + +#[pyclass(name = "BlobDescriptor", skip_from_py_object)] +#[derive(Clone)] +pub struct PyBlobDescriptor { + inner: BlobDescriptor, +} + +impl From for PyBlobDescriptor { + fn from(inner: BlobDescriptor) -> Self { + Self { inner } + } +} + +#[pymethods] +impl PyBlobDescriptor { + fn __repr__(&self) -> String { + format!("{:?}", self.inner) + } +} + +#[pyclass(name = "BlobDescriptorArrayBuilder", skip_from_py_object)] +pub struct PyBlobDescriptorArrayBuilder { + field: arrow_schema::Field, + inner: Option, +} + +impl PyBlobDescriptorArrayBuilder { + fn inner_mut(&mut self) -> PyResult<&mut BlobDescriptorArrayBuilder> { + self.inner + .as_mut() + .ok_or_else(|| PyValueError::new_err("BlobDescriptorArrayBuilder is already finished")) + } +} + +#[pymethods] +impl PyBlobDescriptorArrayBuilder { + #[new] + pub fn new(column: String) -> Self { + let inner = BlobDescriptorArrayBuilder::new(column); + Self { + field: inner.field().clone(), + inner: Some(inner), + } + } + + #[getter] + pub fn field<'py>(&self, py: pyo3::Python<'py>) -> PyResult> { + let pyarrow = PyModule::import(py, "pyarrow")?; + let child_fields = PyList::empty(py); + for (name, type_fn) in [ + ("kind", "uint8"), + ("data", "large_binary"), + ("uri", "utf8"), + ("blob_id", "uint32"), + ("blob_size", "uint64"), + ("position", "uint64"), + ] { + let data_type = pyarrow.getattr(type_fn)?.call0()?; + let child = pyarrow.call_method1("field", (name, data_type, true))?; + child_fields.append(child)?; + } + let data_type = pyarrow.call_method1("struct", (child_fields,))?; + let metadata = PyDict::new(py); + metadata.set_item("ARROW:extension:name", "lance.blob.v2")?; + let kwargs = PyDict::new(py); + kwargs.set_item("nullable", self.field.is_nullable())?; + kwargs.set_item("metadata", metadata)?; + pyarrow.call_method( + "field", + (self.field.name().as_str(), data_type), + Some(&kwargs), + ) + } + + pub fn extend_packed( + &mut self, + blob_id: u32, + offsets: Vec, + sizes: Vec, + ) -> PyResult<()> { + if offsets.len() != sizes.len() { + return Err(PyValueError::new_err(format!( + "offsets and sizes must have the same length, got {} and {}", + offsets.len(), + sizes.len() + ))); + } + let ranges = offsets + .into_iter() + .zip(sizes) + .map(|(offset, size)| BlobRange { offset, size }) + .collect::>(); + self.inner_mut()? + .extend_packed(blob_id, ranges) + .infer_error() + } + + pub fn append_dedicated(&mut self, blob_id: u32, size: u64) -> PyResult<()> { + self.inner_mut()? + .push_dedicated(blob_id, size) + .infer_error() + } + + pub fn append(&mut self, value: &Bound<'_, PyAny>) -> PyResult<()> { + let value = value.extract::>()?; + self.inner_mut()?.push(value.inner.clone()).infer_error() + } + + pub fn extend(&mut self, values: &Bound<'_, PyAny>) -> PyResult<()> { + let iter = values.try_iter()?; + for value in iter { + let value = value?.extract::>()?; + self.inner_mut()?.push(value.inner.clone()).infer_error()?; + } + Ok(()) + } + + pub fn append_inline(&mut self, data: Vec) -> PyResult<()> { + self.inner_mut()? + .push_inline(Bytes::from(data)) + .infer_error() + } + + pub fn append_null(&mut self) -> PyResult<()> { + self.inner_mut()?.push_null().infer_error() + } + + pub fn finish<'py>(&mut self, py: pyo3::Python<'py>) -> PyResult> { + let inner = self.inner.take().ok_or_else(|| { + PyValueError::new_err("BlobDescriptorArrayBuilder is already finished") + })?; + let column = inner.finish().infer_error()?; + column.array().to_data().to_pyarrow(py) + } +} + +#[pyclass(name = "PackedBlobWriter", skip_from_py_object, unsendable)] +pub struct PyPackedBlobWriter { + inner: Option, +} + +impl PyPackedBlobWriter { + pub(crate) async fn try_new( + object_store: Arc, + data_file_path: object_store::path::Path, + blob_id: u32, + ) -> PyResult { + let inner = + PackedBlobWriter::try_new(object_store.as_ref().clone(), data_file_path, blob_id) + .await + .infer_error()?; + Ok(Self { inner: Some(inner) }) + } + + fn inner(&self) -> PyResult<&PackedBlobWriter> { + self.inner + .as_ref() + .ok_or_else(|| PyValueError::new_err("PackedBlobWriter is already finished")) + } + + fn inner_mut(&mut self) -> PyResult<&mut PackedBlobWriter> { + self.inner + .as_mut() + .ok_or_else(|| PyValueError::new_err("PackedBlobWriter is already finished")) + } +} + +#[pymethods] +impl PyPackedBlobWriter { + #[getter] + pub fn blob_id(&self) -> PyResult { + Ok(self.inner()?.blob_id()) + } + + #[getter] + pub fn path(&self) -> PyResult { + Ok(self.inner()?.path().to_string()) + } + + pub fn write_blob(&mut self, data: Vec) -> PyResult<()> { + rt().block_on(None, self.inner_mut()?.write_blob(data))? + .infer_error() + } + + pub fn finish(&mut self) -> PyResult> { + let inner = self + .inner + .take() + .ok_or_else(|| PyValueError::new_err("PackedBlobWriter is already finished"))?; + let values = rt().block_on(None, inner.finish())?.infer_error()?; + Ok(values.into_iter().map(Into::into).collect()) + } +} + +#[pyclass(name = "DedicatedBlobWriter", skip_from_py_object, unsendable)] +pub struct PyDedicatedBlobWriter { + inner: Option, +} + +impl PyDedicatedBlobWriter { + pub(crate) async fn try_new( + object_store: Arc, + data_file_path: object_store::path::Path, + blob_id: u32, + ) -> PyResult { + let inner = + DedicatedBlobWriter::try_new(object_store.as_ref().clone(), data_file_path, blob_id) + .await + .infer_error()?; + Ok(Self { inner: Some(inner) }) + } + + fn inner(&self) -> PyResult<&DedicatedBlobWriter> { + self.inner + .as_ref() + .ok_or_else(|| PyValueError::new_err("DedicatedBlobWriter is already finished")) + } + + fn inner_mut(&mut self) -> PyResult<&mut DedicatedBlobWriter> { + self.inner + .as_mut() + .ok_or_else(|| PyValueError::new_err("DedicatedBlobWriter is already finished")) + } +} + +#[pymethods] +impl PyDedicatedBlobWriter { + #[getter] + pub fn blob_id(&self) -> PyResult { + Ok(self.inner()?.blob_id()) + } + + #[getter] + pub fn path(&self) -> PyResult { + Ok(self.inner()?.path().to_string()) + } + + pub fn write(&mut self, data: Vec) -> PyResult<()> { + rt().block_on(None, self.inner_mut()?.write(data))? + .infer_error() + } + + pub fn finish(&mut self) -> PyResult { + let inner = self + .inner + .take() + .ok_or_else(|| PyValueError::new_err("DedicatedBlobWriter is already finished"))?; + let value = rt().block_on(None, inner.finish())?.infer_error()?; + Ok(value.into()) + } +} diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 7c525b0aefe..0803246b7bf 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -78,6 +78,7 @@ use lance_index::{ FtsPrewarmOptions, IndexParams, IndexType, PrewarmOptions, optimize::OptimizeOptions, progress::{IndexBuildProgress, NoopIndexBuildProgress}, + scalar::inverted::InvertedListFormatVersion, scalar::{FullTextSearchQuery, InvertedIndexParams, ScalarIndexParams}, vector::{ ApproxMode, DEFAULT_QUERY_PARALLELISM, Query as VectorQuery, @@ -395,6 +396,23 @@ impl MergeInsertBuilder { Ok(slf) } + pub fn target_bases( + mut slf: PyRefMut<'_, Self>, + bases: Vec, + ) -> PyResult> { + slf.builder.target_base_names_or_paths(bases); + Ok(slf) + } + + #[pyo3(signature = (include_primary = true))] + pub fn target_all_bases( + mut slf: PyRefMut<'_, Self>, + include_primary: bool, + ) -> PyResult> { + slf.builder.target_all_bases(include_primary); + Ok(slf) + } + pub fn execute(&mut self, new_data: &Bound) -> PyResult> { let py = new_data.py(); let new_data = convert_reader(new_data)?; @@ -531,6 +549,23 @@ fn extract_index_segments(segments: &Bound<'_, PyAny>) -> PyResult>) -> PyResult>> { + index_segments + .map(|segments| { + segments + .into_iter() + .map(|segment| { + Uuid::parse_str(&segment).map_err(|err| { + PyValueError::new_err(format!( + "invalid index segment uuid '{segment}': {err}" + )) + }) + }) + .collect::>>() + }) + .transpose() +} + impl MergeInsertBuilder { fn build_stats<'a>(stats: &MergeStats, py: Python<'a>) -> PyResult> { let dict = PyDict::new(py); @@ -2429,12 +2464,34 @@ impl Dataset { if let Some(prefix_only) = kwargs.get_item("prefix_only")? { params = params.ngram_prefix_only(prefix_only.extract()?); } + if let Some(block_size) = kwargs.get_item("block_size")? { + params = params + .block_size(block_size.extract()?) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + } if let Some(memory_limit) = kwargs.get_item("memory_limit")? { params = params.memory_limit_mb(memory_limit.extract()?); } if let Some(num_workers) = kwargs.get_item("num_workers")? { params = params.num_workers(num_workers.extract()?); } + if let Some(format_version) = kwargs.get_item("format_version")? + && !format_version.is_none() + { + let value = if let Ok(value) = format_version.cast::() { + value.to_string_lossy().to_string() + } else if let Ok(value) = format_version.extract::() { + value.to_string() + } else { + return Err(PyValueError::new_err( + "format_version must be 1, 2, 3, 'v1', 'v2', or 'v3'", + )); + }; + let format_version = value + .parse::() + .map_err(|err| PyValueError::new_err(err.to_string()))?; + params = params.format_version(format_version); + } } Box::new(params) } @@ -2565,18 +2622,31 @@ impl Dataset { Ok(()) } - #[pyo3(signature = (name, *, with_position = false))] - fn prewarm_index(&self, name: &str, with_position: bool) -> PyResult<()> { + #[pyo3(signature = (name, *, with_position = false, index_segments = None))] + fn prewarm_index( + &self, + name: &str, + with_position: bool, + index_segments: Option>, + ) -> PyResult<()> { + let index_segments = parse_index_segment_ids(index_segments)?; + rt().block_on(None, async { if with_position { - self.ds - .prewarm_index_with_options( - name, - &PrewarmOptions::Fts(FtsPrewarmOptions::new().with_position(true)), - ) - .await + let options = PrewarmOptions::Fts(FtsPrewarmOptions::new().with_position(true)); + if let Some(index_segments) = index_segments.as_deref() { + self.ds + .prewarm_index_segments_with_options(name, index_segments, &options) + .await + } else { + self.ds.prewarm_index_with_options(name, &options).await + } } else { - self.ds.prewarm_index(name).await + if let Some(index_segments) = index_segments.as_deref() { + self.ds.prewarm_index_segments(name, index_segments).await + } else { + self.ds.prewarm_index(name).await + } } })? .infer_error() @@ -3568,9 +3638,10 @@ impl Dataset { /// Perform pairwise hamming distance clustering on a partition of an IVF_FLAT index. /// - /// This function loads a specific partition from an IVF_FLAT index on a hash column, - /// computes pairwise hamming distances between all hashes in the partition, - /// filters by threshold, and clusters the results using union-find. + /// This function loads a specific partition from every segment of an IVF_FLAT + /// index on a hash column, computes pairwise hamming distances between all + /// hashes in the combined partition, filters by threshold, and clusters the + /// results using union-find. /// /// Parameters /// ---------- @@ -3580,6 +3651,9 @@ impl Dataset { /// The partition ID within the IVF_FLAT index /// hamming_threshold : int /// Maximum hamming distance to consider as similar + /// index_segments : list of str, optional + /// If specified, only these physical index segment UUIDs of the named + /// logical index contribute rows. Defaults to all segments. /// /// Returns /// ------- @@ -3587,27 +3661,45 @@ impl Dataset { /// A reader yielding batches with columns: /// - 'representative': uint64 - The representative row ID for each cluster /// - 'duplicates': list - List of duplicate row IDs in each cluster - #[pyo3(signature = (index_name, partition_id, hamming_threshold))] + #[pyo3(signature = (index_name, partition_id, hamming_threshold, index_segments=None))] fn hamming_clustering_for_ivf_partition( &self, py: Python<'_>, index_name: &str, partition_id: usize, hamming_threshold: u32, + index_segments: Option>, ) -> PyResult>> { - use lance::index::vector::hamming::hamming_clustering_for_ivf_partition; + use lance::index::vector::hamming::{ + hamming_clustering_for_ivf_partition, hamming_clustering_for_ivf_partition_segments, + }; + let segment_ids = parse_index_segment_ids(index_segments)?; let ds = self.ds.as_ref(); let reader = rt() - .block_on( - Some(py), - hamming_clustering_for_ivf_partition( - ds, - index_name, - partition_id, - hamming_threshold, - ), - )? + .block_on(Some(py), async { + match segment_ids.as_deref() { + Some(segment_ids) => { + hamming_clustering_for_ivf_partition_segments( + ds, + index_name, + segment_ids, + partition_id, + hamming_threshold, + ) + .await + } + None => { + hamming_clustering_for_ivf_partition( + ds, + index_name, + partition_id, + hamming_threshold, + ) + .await + } + } + })? .map_err(|err| PyValueError::new_err(err.to_string()))?; Ok(PyArrowType(reader)) @@ -3615,26 +3707,43 @@ impl Dataset { /// Get partition information for an IVF_FLAT index. /// + /// Partition sizes are aggregated across all segments of the logical index + /// unless a subset is selected via ``index_segments``. + /// /// Parameters /// ---------- /// index_name : str /// Name of the IVF_FLAT index + /// index_segments : list of str, optional + /// If specified, only these physical index segment UUIDs of the named + /// logical index contribute to the sizes. Defaults to all segments. /// /// Returns /// ------- /// List[dict] /// List of partition info dicts with 'partition_id' and 'size' - #[pyo3(signature = (index_name))] + #[pyo3(signature = (index_name, index_segments=None))] fn get_ivf_partition_info( &self, py: Python<'_>, index_name: &str, + index_segments: Option>, ) -> PyResult>> { - use lance::index::vector::hamming::get_ivf_partition_info; + use lance::index::vector::hamming::{ + get_ivf_partition_info, get_ivf_partition_info_segments, + }; + let segment_ids = parse_index_segment_ids(index_segments)?; let ds = self.ds.as_ref(); let result = rt() - .block_on(Some(py), get_ivf_partition_info(ds, index_name))? + .block_on(Some(py), async { + match segment_ids.as_deref() { + Some(segment_ids) => { + get_ivf_partition_info_segments(ds, index_name, segment_ids).await + } + None => get_ivf_partition_info(ds, index_name).await, + } + })? .map_err(|err| PyValueError::new_err(err.to_string()))?; let partitions: PyResult> = result @@ -4503,6 +4612,11 @@ pub fn get_write_params( p = p.with_target_base_names_or_paths(target_bases_list); } + // Handle target_all_bases parameter (bool: include primary storage) + if let Some(target_all_bases) = get_dict_opt::(options, "target_all_bases")? { + p = p.with_target_all_bases(target_all_bases); + } + // Handle base_store_params: per-base storage options keyed by base path URI if let Some(base_store_params) = get_dict_opt::>>(options, "base_store_params")? diff --git a/python/src/file.rs b/python/src/file.rs index e1906b9f4f2..2a3dd09e17f 100644 --- a/python/src/file.rs +++ b/python/src/file.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::blob::{PyDedicatedBlobWriter, PyPackedBlobWriter}; use crate::namespace::extract_namespace_arc; use crate::{error::PythonErrorExt, rt}; use arrow::pyarrow::PyArrowType; @@ -524,6 +525,30 @@ impl LanceFileSession { )? } + pub fn open_packed_blob_writer( + &self, + path: String, + blob_id: u32, + ) -> PyResult { + let path = self.base_path.child_path(&Path::from(path)); + rt().block_on( + None, + PyPackedBlobWriter::try_new(self.object_store.clone(), path, blob_id), + )? + } + + pub fn open_dedicated_blob_writer( + &self, + path: String, + blob_id: u32, + ) -> PyResult { + let path = self.base_path.child_path(&Path::from(path)); + rt().block_on( + None, + PyDedicatedBlobWriter::try_new(self.object_store.clone(), path, blob_id), + )? + } + pub fn contains(&self, path: String) -> PyResult { let full_path = self.base_path.child_path(&Path::from(path)); rt().block_on(None, async { diff --git a/python/src/fragment.rs b/python/src/fragment.rs index dbe5c426903..336c6a4cf51 100644 --- a/python/src/fragment.rs +++ b/python/src/fragment.rs @@ -825,6 +825,9 @@ impl FromPyObject<'_, '_> for PyLance { row_id_meta, last_updated_at_version_meta, created_at_version_meta, + // Overlays are not exposed to Python yet, and the reverse conversion + // does not export them, so this round-trip is overlay-free. + overlays: vec![], })) } } diff --git a/python/src/lib.rs b/python/src/lib.rs index 3275849a681..c5631d26f10 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -62,6 +62,7 @@ use std::ffi::CString; use std::ptr::NonNull; pub(crate) mod arrow; +pub(crate) mod blob; #[cfg(feature = "datagen")] pub(crate) mod datagen; pub(crate) mod dataset; @@ -73,6 +74,7 @@ pub(crate) mod fragment; pub(crate) mod indices; pub(crate) mod mem_wal; pub(crate) mod namespace; +pub(crate) mod otel; pub(crate) mod reader; pub(crate) mod scanner; pub(crate) mod schema; @@ -96,6 +98,9 @@ pub use indices::register_indices; pub use reader::LanceReader; pub use scanner::Scanner; +use crate::blob::{ + PyBlobDescriptor, PyBlobDescriptorArrayBuilder, PyDedicatedBlobWriter, PyPackedBlobWriter, +}; use crate::executor::BackgroundExecutor; const CLIENT_VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -257,6 +262,10 @@ fn lance(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; @@ -310,10 +319,17 @@ fn lance(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(trace_to_chrome))?; m.add_wrapped(wrap_pyfunction!(capture_trace_events))?; m.add_wrapped(wrap_pyfunction!(shutdown_tracing))?; + // OpenTelemetry metrics bridge + m.add_class::()?; + m.add_class::()?; + m.add_wrapped(wrap_pyfunction!(otel::register_lance_metrics_recorder))?; + m.add_wrapped(wrap_pyfunction!(otel::lance_metrics_catalog))?; + m.add_wrapped(wrap_pyfunction!(otel::snapshot_lance_metrics))?; m.add_wrapped(wrap_pyfunction!(manifest_needs_migration))?; m.add_wrapped(wrap_pyfunction!(language_model_home))?; m.add_wrapped(wrap_pyfunction!(bytes_read_counter))?; m.add_wrapped(wrap_pyfunction!(iops_counter))?; + m.add_wrapped(wrap_pyfunction!(simd_info))?; m.add_wrapped(wrap_pyfunction!(stable_version))?; // Debug functions m.add_wrapped(wrap_pyfunction!(debug::format_schema))?; @@ -332,6 +348,37 @@ fn iops_counter() -> PyResult { Ok(::lance::io::iops_counter()) } +/// Returns a dict describing which SIMD tier the lance runtime dispatches to +/// on this host, plus the raw CPU feature flags it detected. +/// +/// Mirrors `pyarrow.runtime_info()`: a cheap, transparent way to verify that +/// the host is hitting the expected SIMD tier (e.g., `"avx512_fp16"`, +/// `"avx2"`) when debugging vector-search performance. +/// +/// Returns: +/// { +/// "tier": str, # e.g. "avx2", "avx_fma", "neon", "none" +/// "target_arch": str, # e.g. "x86_64", "aarch64", "loongarch64" +/// "host_features": list[str], # raw CPU feature flags (x86_64 only) +/// } +/// +/// Examples: +/// >>> import lance +/// >>> info = lance.simd_info() +/// >>> sorted(info) +/// ['host_features', 'target_arch', 'tier'] +/// >>> isinstance(info["tier"], str) +/// True +#[pyfunction] +pub fn simd_info(py: Python<'_>) -> PyResult> { + let info = lance_core::utils::cpu::simd_info(); + let dict = pyo3::types::PyDict::new(py); + dict.set_item("tier", info.tier.to_string())?; + dict.set_item("target_arch", info.target_arch)?; + dict.set_item("host_features", info.host_features)?; + Ok(dict.into()) +} + #[pyfunction(name = "bytes_read_counter")] fn bytes_read_counter() -> PyResult { Ok(::lance::io::bytes_read_counter()) diff --git a/python/src/mem_wal.rs b/python/src/mem_wal.rs index dc9718c0dce..e415acfc38d 100644 --- a/python/src/mem_wal.rs +++ b/python/src/mem_wal.rs @@ -19,6 +19,7 @@ use futures::TryStreamExt; use lance::dataset::Dataset as LanceDataset; use lance::dataset::mem_wal::scanner::{ FlushedGeneration, LsmDataSourceCollector, LsmPointLookupPlanner, LsmVectorSearchPlanner, + parse_filter_expr as parse_lsm_filter_expr, }; use lance::dataset::mem_wal::write::{MemTableStats, WriteStatsSnapshot}; use lance::dataset::mem_wal::{LsmScanner, ShardSnapshot, ShardWriter, evaluate_sharding_spec}; @@ -237,6 +238,14 @@ struct ClosedShardWriterState { memtable_stats: MemTableStats, } +fn collect_record_batches(data: &Bound<'_, PyAny>) -> PyResult> { + let reader = ArrowArrayStreamReader::from_pyarrow_bound(data) + .map_err(|e| PyValueError::new_err(format!("Cannot read data as Arrow: {}", e)))?; + reader + .collect::>() + .map_err(|e| PyIOError::new_err(format!("Failed to read batches: {}", e))) +} + #[pymethods] impl PyShardWriter { /// Write data batches to the MemWAL. @@ -244,11 +253,7 @@ impl PyShardWriter { /// Accepts any PyArrow-compatible data source (RecordBatch, Table, /// or an Arrow stream reader). pub fn put(&self, py: Python<'_>, data: &Bound<'_, PyAny>) -> PyResult<()> { - let reader = ArrowArrayStreamReader::from_pyarrow_bound(data) - .map_err(|e| PyValueError::new_err(format!("Cannot read data as Arrow: {}", e)))?; - let batches: Vec = reader - .collect::>() - .map_err(|e| PyIOError::new_err(format!("Failed to read batches: {}", e)))?; + let batches = collect_record_batches(data)?; if batches.is_empty() { return Ok(()); @@ -267,6 +272,31 @@ impl PyShardWriter { .map_err(|e: lance::Error| PyIOError::new_err(e.to_string())) } + /// Delete rows from the MemWAL by primary key. + /// + /// Accepts any PyArrow-compatible data source carrying the shard's primary + /// key column(s). Rust core validates that primary keys exist and builds the + /// tombstone rows. + pub fn delete(&self, py: Python<'_>, keys: &Bound<'_, PyAny>) -> PyResult<()> { + let batches = collect_record_batches(keys)?; + + if batches.is_empty() { + return Ok(()); + } + + let inner = self.inner.clone(); + rt().block_on(Some(py), async move { + let guard = inner.lock().await; + match guard.as_ref() { + Some(writer) => writer.delete(batches).await.map(|_| ()), + None => Err(lance_core::Error::invalid_input( + "ShardWriter is already closed", + )), + } + })? + .map_err(|e: lance::Error| PyIOError::new_err(e.to_string())) + } + /// Flush pending data and close the writer. /// /// After close(), calling put() will raise an error. @@ -545,7 +575,11 @@ impl PyLsmScanner { .take() .ok_or_else(|| PyRuntimeError::new_err("Scanner has already been consumed"))?; let cols: Vec<&str> = columns.iter().map(|s| s.as_str()).collect(); - slf.inner = Some(scanner.project(&cols)); + slf.inner = Some( + scanner + .project(&cols) + .map_err(|e| PyValueError::new_err(e.to_string()))?, + ); Ok(slf) } @@ -563,18 +597,22 @@ impl PyLsmScanner { Ok(slf) } - /// Limit the number of rows returned. - #[pyo3(signature = (n, offset=None))] + /// Limit the number of rows returned, optionally with an offset. + #[pyo3(signature = (n=None, offset=None))] pub fn limit( mut slf: PyRefMut<'_, Self>, - n: usize, + n: Option, offset: Option, ) -> PyResult> { let scanner = slf .inner .take() .ok_or_else(|| PyRuntimeError::new_err("Scanner has already been consumed"))?; - slf.inner = Some(scanner.limit(n, offset)); + slf.inner = Some( + scanner + .limit(n.map(|n| n as i64), offset.map(|o| o as i64)) + .map_err(|e| PyValueError::new_err(e.to_string()))?, + ); Ok(slf) } @@ -729,13 +767,14 @@ pub struct PyLsmVectorSearchPlanner { #[pymethods] impl PyLsmVectorSearchPlanner { #[new] - #[pyo3(signature = (dataset, shard_snapshots, vector_column, pk_columns=None, distance_type=None))] + #[pyo3(signature = (dataset, shard_snapshots, vector_column, pk_columns=None, distance_type=None, filter=None))] pub fn new( dataset: &Bound<'_, PyDataset>, shard_snapshots: Vec>, vector_column: String, pk_columns: Option>, distance_type: Option, + filter: Option, ) -> PyResult { let ds = dataset.borrow().ds.clone(); let snapshots: Vec = shard_snapshots @@ -751,9 +790,16 @@ impl PyLsmVectorSearchPlanner { let dist_type = parse_distance_type(distance_type.as_deref().unwrap_or("l2"))?; let vector_dim = get_vector_dim(&ds, &vector_column)?; + let filter = filter + .as_deref() + .map(|filter| { + parse_lsm_filter_expr(base_schema.as_ref(), filter) + .map_err(|e| PyValueError::new_err(e.to_string())) + }) + .transpose()?; let collector = LsmDataSourceCollector::new(ds.clone(), snapshots); - let planner = LsmVectorSearchPlanner::new( + let mut planner = LsmVectorSearchPlanner::new( collector, pk_cols, base_schema.clone(), @@ -761,6 +807,9 @@ impl PyLsmVectorSearchPlanner { dist_type, ) .with_dataset(ds); + if let Some(filter) = filter { + planner = planner.with_filter(Some(filter)); + } Ok(Self { planner, diff --git a/python/src/otel.rs b/python/src/otel.rs new file mode 100644 index 00000000000..71d191d7db9 --- /dev/null +++ b/python/src/otel.rs @@ -0,0 +1,640 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Bridge from the [`metrics`] crate facade to Python OpenTelemetry. +//! +//! Lance core publishes metrics through the global [`metrics`] facade without +//! choosing a backend. This module installs a process-global [`Recorder`] that +//! aggregates those metrics into lock-free cumulative storage, and exposes that +//! state to Python so the bindings can feed it into the user's OpenTelemetry +//! `MeterProvider`. +//! +//! While it targets OpenTelemetry on the Python side, the recorder is agnostic +//! to the metric *source*: it records any metric emitted through the facade, +//! keyed by name and labels. Object store metrics are the first producer, but +//! nothing here is specific to them. New metrics flow through automatically; +//! they only need to be described (see [`describe_all`]) so the Python layer can +//! discover their name, kind, and unit up front. +//! +//! ## Why pull, not push +//! +//! OpenTelemetry collects on its own schedule and invokes observable-instrument +//! callbacks at collection time. Cumulative counters map directly onto OTel's +//! `ObservableCounter` semantics. So the bridge aggregates in Rust and lets the +//! Python collection thread pull a [`snapshot`](snapshot_lance_metrics). The +//! snapshot is lock-free, but it still walks every registered series and +//! allocates owned copies of their names and labels, so it runs with the GIL +//! released to avoid stalling other Python threads during collection. +//! +//! ## Histograms +//! +//! OpenTelemetry has no asynchronous histogram instrument, so histograms cannot +//! be pulled as-is. Instead each histogram is aggregated into fixed buckets +//! (Prometheus style) and exposed as cumulative `le` bucket counts plus a count +//! and sum, which the Python layer surfaces as observable counters. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, LazyLock, Mutex, OnceLock, RwLock}; + +use metrics::{Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, SharedString, Unit}; +use metrics_util::registry::{Registry, Storage}; +use pyo3::prelude::*; + +/// Bucket boundaries used when a histogram has no registered bounds. Covers a +/// broad latency range so unknown histograms still produce useful buckets. +const DEFAULT_BOUNDS: &[f64] = &[ + 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, +]; + +/// The kind of a metric, mirroring the three `metrics` instrument types. +#[derive(Clone, Copy)] +enum MetricKind { + Counter, + Gauge, + Histogram, +} + +impl MetricKind { + fn as_str(self) -> &'static str { + match self { + Self::Counter => "counter", + Self::Gauge => "gauge", + Self::Histogram => "histogram", + } + } +} + +/// Description of a metric, populated by the recorder's `describe_*` methods. +struct MetricDescription { + kind: MetricKind, + unit: Option, + description: String, +} + +/// Catalog of described metrics, keyed by metric name. The Python layer reads +/// this to create one OpenTelemetry instrument per metric up front. +static CATALOG: LazyLock>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +/// Per-metric histogram bucket boundaries, keyed by metric name. Producers +/// register their recommended bounds before any metric is recorded. +static HISTOGRAM_BOUNDS: LazyLock>>> = + LazyLock::new(|| RwLock::new(HashMap::new())); + +/// The installed recorder's registry, available once installation succeeds. +static REGISTRY: OnceLock>> = OnceLock::new(); + +fn bounds_for(name: &str) -> Arc<[f64]> { + HISTOGRAM_BOUNDS + .read() + .unwrap() + .get(name) + .cloned() + .unwrap_or_else(|| Arc::from(DEFAULT_BOUNDS)) +} + +/// A histogram that buckets samples at record time into fixed boundaries, +/// keeping a cumulative count and sum. Bucketing eagerly keeps memory bounded +/// (unlike retaining raw samples) and produces Prometheus-style `le` buckets. +struct BucketedHistogram { + /// Sorted, finite upper bounds. A sample `v` falls in the first bucket whose + /// bound is `>= v`; samples above all bounds fall in the implicit `+Inf` + /// bucket stored as the final entry of `counts`. + bounds: Arc<[f64]>, + /// Per-bucket (non-cumulative) counts; length is `bounds.len() + 1`. + counts: Box<[AtomicU64]>, + count: AtomicU64, + /// Running sum of recorded values, stored as `f64` bits (there is no atomic + /// f64, so the bit pattern is held in a `u64`; see [`Self::add_to_sum`]). + sum_bits: AtomicU64, +} + +// All atomics here use `Ordering::Relaxed`: each metric counter is independent, +// so no happens-before relationship is needed between them, and a snapshot +// reader tolerates slightly stale values. This matches `metrics_util`'s +// `AtomicStorage`. + +impl BucketedHistogram { + fn new(bounds: Arc<[f64]>) -> Self { + let counts = (0..bounds.len() + 1) + .map(|_| AtomicU64::new(0)) + .collect::>() + .into_boxed_slice(); + Self { + bounds, + counts, + count: AtomicU64::new(0), + sum_bits: AtomicU64::new(0), + } + } + + fn add_to_sum(&self, value: f64) { + // No atomic offers an f64 add, so read the current bit pattern, add in + // float space, and CAS it back, retrying if another thread won the race. + let mut current = self.sum_bits.load(Ordering::Relaxed); + loop { + let updated = (f64::from_bits(current) + value).to_bits(); + match self.sum_bits.compare_exchange_weak( + current, + updated, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => break, + Err(actual) => current = actual, + } + } + } + + /// Cumulative `le` buckets, total count, and sum at this instant. + fn snapshot(&self) -> MetricValue { + let mut cumulative = 0u64; + let mut buckets = Vec::with_capacity(self.bounds.len() + 1); + for (i, bound) in self.bounds.iter().enumerate() { + cumulative += self.counts[i].load(Ordering::Relaxed); + buckets.push((format!("{}", bound), cumulative)); + } + cumulative += self.counts[self.bounds.len()].load(Ordering::Relaxed); + buckets.push(("+Inf".to_string(), cumulative)); + MetricValue::Histogram { + buckets, + count: self.count.load(Ordering::Relaxed), + sum: f64::from_bits(self.sum_bits.load(Ordering::Relaxed)), + } + } +} + +impl metrics::HistogramFn for BucketedHistogram { + fn record(&self, value: f64) { + let idx = self.bounds.partition_point(|&bound| bound < value); + self.counts[idx].fetch_add(1, Ordering::Relaxed); + self.count.fetch_add(1, Ordering::Relaxed); + self.add_to_sum(value); + } +} + +/// Storage backing the registry. Counters and gauges are plain atomics (as in +/// `metrics_util`'s `AtomicStorage`); histograms use [`BucketedHistogram`]. +struct LanceStorage; + +impl Storage for LanceStorage { + type Counter = Arc; + type Gauge = Arc; + type Histogram = Arc; + + fn counter(&self, _key: &Key) -> Self::Counter { + Arc::new(AtomicU64::new(0)) + } + + fn gauge(&self, _key: &Key) -> Self::Gauge { + // The `metrics` facade writes the f64 bit pattern into this `u64` (the + // snapshot decodes it with `f64::from_bits`), matching `AtomicStorage`. + // `0` decodes to `0.0`, the correct initial value. + Arc::new(AtomicU64::new(0)) + } + + fn histogram(&self, key: &Key) -> Self::Histogram { + Arc::new(BucketedHistogram::new(bounds_for(key.name()))) + } +} + +struct LanceRecorder { + registry: Arc>, +} + +impl LanceRecorder { + fn describe( + &self, + key: KeyName, + kind: MetricKind, + unit: Option, + description: SharedString, + ) { + CATALOG.lock().unwrap().insert( + key.as_str().to_string(), + MetricDescription { + kind, + unit: unit.map(|u| u.as_canonical_label().to_string()), + description: description.into_owned(), + }, + ); + } +} + +impl Recorder for LanceRecorder { + fn describe_counter(&self, key: KeyName, unit: Option, description: SharedString) { + self.describe(key, MetricKind::Counter, unit, description); + } + + fn describe_gauge(&self, key: KeyName, unit: Option, description: SharedString) { + self.describe(key, MetricKind::Gauge, unit, description); + } + + fn describe_histogram(&self, key: KeyName, unit: Option, description: SharedString) { + self.describe(key, MetricKind::Histogram, unit, description); + } + + fn register_counter(&self, key: &Key, _metadata: &Metadata<'_>) -> Counter { + self.registry + .get_or_create_counter(key, |c| Counter::from_arc(c.clone())) + } + + fn register_gauge(&self, key: &Key, _metadata: &Metadata<'_>) -> Gauge { + self.registry + .get_or_create_gauge(key, |g| Gauge::from_arc(g.clone())) + } + + fn register_histogram(&self, key: &Key, _metadata: &Metadata<'_>) -> Histogram { + self.registry + .get_or_create_histogram(key, |h| Histogram::from_arc(h.clone())) + } +} + +/// Register the recommended histogram bounds for every metric-emitting +/// subsystem. New subsystems add their `histogram_bounds()` here. +fn register_bounds() { + let mut bounds = HISTOGRAM_BOUNDS.write().unwrap(); + for (name, values) in lance_io::object_store::metrics::histogram_bounds() { + bounds.insert((*name).to_string(), Arc::from(*values)); + } +} + +/// Describe every metric-emitting subsystem so the catalog is populated. Must +/// run after the recorder is installed. New subsystems add their +/// `describe_metrics()` here. +fn describe_all() { + lance_io::object_store::metrics::describe_metrics(); +} + +enum MetricValue { + Scalar(f64), + Histogram { + buckets: Vec<(String, u64)>, + count: u64, + sum: f64, + }, +} + +struct MetricPoint { + name: String, + kind: &'static str, + attributes: HashMap, + value: MetricValue, +} + +fn labels(key: &Key) -> HashMap { + key.labels() + .map(|label| (label.key().to_string(), label.value().to_string())) + .collect() +} + +fn collect_points(registry: &Registry) -> Vec { + let mut points = Vec::new(); + for (key, handle) in registry.get_counter_handles() { + points.push(MetricPoint { + name: key.name().to_string(), + kind: "counter", + attributes: labels(&key), + // OpenTelemetry observations are float; counts stay well within the + // f64-exact integer range (2^53), so this cast is lossless in practice. + value: MetricValue::Scalar(handle.load(Ordering::Relaxed) as f64), + }); + } + for (key, handle) in registry.get_gauge_handles() { + points.push(MetricPoint { + name: key.name().to_string(), + kind: "gauge", + attributes: labels(&key), + value: MetricValue::Scalar(f64::from_bits(handle.load(Ordering::Relaxed))), + }); + } + for (key, handle) in registry.get_histogram_handles() { + points.push(MetricPoint { + name: key.name().to_string(), + kind: "histogram", + attributes: labels(&key), + value: handle.snapshot(), + }); + } + points +} + +/// One metric data point exposed to Python. For counters and gauges only +/// `value` is set; for histograms `buckets` (cumulative `le` counts), `count`, +/// and `sum` are set. +#[pyclass(name = "MetricPoint", get_all)] +pub struct PyMetricPoint { + name: String, + kind: String, + attributes: HashMap, + value: Option, + buckets: Option>, + count: Option, + sum: Option, +} + +impl From for PyMetricPoint { + fn from(point: MetricPoint) -> Self { + let (value, buckets, count, sum) = match point.value { + MetricValue::Scalar(v) => (Some(v), None, None, None), + MetricValue::Histogram { + buckets, + count, + sum, + } => (None, Some(buckets), Some(count), Some(sum)), + }; + Self { + name: point.name, + kind: point.kind.to_string(), + attributes: point.attributes, + value, + buckets, + count, + sum, + } + } +} + +/// A described metric, used by the Python layer to create instruments up front. +#[pyclass(name = "MetricDescription", get_all)] +pub struct PyMetricDescription { + name: String, + kind: String, + unit: Option, + description: String, +} + +/// Install the Lance metrics recorder as the process-global `metrics` recorder. +/// +/// Returns `True` if the recorder is installed (now or previously). Returns +/// `False` if a *different* recorder is already installed — `metrics` allows +/// only one global recorder per process, so Lance cannot coexist with another. +#[pyfunction] +pub fn register_lance_metrics_recorder() -> bool { + if REGISTRY.get().is_some() { + return true; + } + let registry = Arc::new(Registry::new(LanceStorage)); + let recorder = LanceRecorder { + registry: registry.clone(), + }; + // Register histogram bounds *before* installing the recorder. Once the + // recorder is global, another thread can emit a metric and create the + // histogram handle concurrently; if the bounds aren't registered yet that + // handle would be built with the fallback bounds and keep them for the + // life of the process. `register_bounds()` doesn't need the recorder. + register_bounds(); + match metrics::set_global_recorder(recorder) { + Ok(()) => { + let _ = REGISTRY.set(registry); + describe_all(); + true + } + Err(_) => false, + } +} + +/// The catalog of described Lance metrics. Empty until the recorder is installed. +#[pyfunction] +pub fn lance_metrics_catalog() -> Vec { + CATALOG + .lock() + .unwrap() + .iter() + .map(|(name, desc)| PyMetricDescription { + name: name.clone(), + kind: desc.kind.as_str().to_string(), + unit: desc.unit.clone(), + description: desc.description.clone(), + }) + .collect() +} + +/// A point-in-time snapshot of every recorded metric. Empty until the recorder +/// is installed. The read is lock-free but walks every series and allocates, so +/// it runs with the GIL released. +#[pyfunction] +pub fn snapshot_lance_metrics(py: Python<'_>) -> Vec { + let Some(registry) = REGISTRY.get() else { + return Vec::new(); + }; + let points = py.detach(|| collect_points(registry)); + points.into_iter().map(PyMetricPoint::from).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use metrics::HistogramFn; + + fn bucket_count(buckets: &[(String, u64)], le: &str) -> u64 { + buckets + .iter() + .find(|(b, _)| b == le) + .map(|(_, c)| *c) + .unwrap_or_else(|| panic!("no bucket with le={le}")) + } + + #[test] + fn bucketed_histogram_records_cumulative_buckets() { + let hist = BucketedHistogram::new(Arc::from([0.1f64, 1.0, 10.0].as_slice())); + hist.record(0.05); // le=0.1 + hist.record(0.5); // le=1 + hist.record(0.5); // le=1 + hist.record(50.0); // +Inf + + let MetricValue::Histogram { + buckets, + count, + sum, + } = hist.snapshot() + else { + panic!("expected histogram"); + }; + + // Buckets are cumulative (Prometheus `le` semantics). + assert_eq!(bucket_count(&buckets, "0.1"), 1); + assert_eq!(bucket_count(&buckets, "1"), 3); + assert_eq!(bucket_count(&buckets, "10"), 3); + assert_eq!(bucket_count(&buckets, "+Inf"), 4); + assert_eq!(count, 4); + assert!((sum - 51.05).abs() < 1e-9); + } + + #[test] + fn bucketed_histogram_boundary_is_inclusive() { + let hist = BucketedHistogram::new(Arc::from([1.0f64].as_slice())); + hist.record(1.0); // exactly the bound -> le=1, not +Inf + let MetricValue::Histogram { buckets, .. } = hist.snapshot() else { + panic!("expected histogram"); + }; + assert_eq!(bucket_count(&buckets, "1"), 1); + assert_eq!(bucket_count(&buckets, "+Inf"), 1); + } + + #[test] + fn bucketed_histogram_boundary_is_inclusive_mid_range() { + // A value equal to a middle bound lands in that bucket, not the next. + let hist = BucketedHistogram::new(Arc::from([0.1f64, 1.0, 10.0].as_slice())); + hist.record(1.0); + let MetricValue::Histogram { buckets, .. } = hist.snapshot() else { + panic!("expected histogram"); + }; + assert_eq!(bucket_count(&buckets, "0.1"), 0); + assert_eq!(bucket_count(&buckets, "1"), 1); + assert_eq!(bucket_count(&buckets, "10"), 1); // cumulative, so still 1 + assert_eq!(bucket_count(&buckets, "+Inf"), 1); + } + + #[test] + fn recorder_aggregates_counters_with_labels() { + let registry = Arc::new(Registry::new(LanceStorage)); + let recorder = LanceRecorder { + registry: registry.clone(), + }; + metrics::with_local_recorder(&recorder, || { + metrics::counter!("test_requests_total", "operation" => "get", "scheme" => "s3") + .increment(2); + metrics::counter!("test_requests_total", "operation" => "get", "scheme" => "s3") + .increment(3); + // A distinct label set must produce a separate point, not merge. + metrics::counter!("test_requests_total", "operation" => "put", "scheme" => "gs") + .increment(7); + }); + + let scalar = |attrs: &[(&str, &str)]| { + let points = collect_points(®istry); + let point = points + .into_iter() + .find(|p| { + p.name == "test_requests_total" + && attrs + .iter() + .all(|(k, v)| p.attributes.get(*k).map(String::as_str) == Some(*v)) + }) + .expect("counter recorded for label set"); + assert_eq!(point.kind, "counter"); + match point.value { + MetricValue::Scalar(v) => v, + _ => panic!("expected scalar"), + } + }; + + // Same labels aggregate; distinct labels stay separate. + assert!((scalar(&[("operation", "get"), ("scheme", "s3")]) - 5.0).abs() < 1e-9); + assert!((scalar(&[("operation", "put"), ("scheme", "gs")]) - 7.0).abs() < 1e-9); + } + + #[test] + fn recorder_records_gauges() { + let registry = Arc::new(Registry::new(LanceStorage)); + let recorder = LanceRecorder { + registry: registry.clone(), + }; + // Gauges store the f64 bit pattern in a u64; the snapshot must decode it. + metrics::with_local_recorder(&recorder, || { + metrics::gauge!("test_gauge", "scheme" => "s3").set(3.5); + }); + + let points = collect_points(®istry); + let point = points + .iter() + .find(|p| p.name == "test_gauge") + .expect("gauge recorded"); + assert_eq!(point.kind, "gauge"); + assert!(matches!(point.value, MetricValue::Scalar(v) if (v - 3.5).abs() < 1e-9)); + } + + #[test] + fn recorder_falls_back_to_default_bounds() { + // A histogram with no registered bounds uses DEFAULT_BOUNDS. + let name = "test_unregistered_histogram"; + assert!(!HISTOGRAM_BOUNDS.read().unwrap().contains_key(name)); + + let registry = Arc::new(Registry::new(LanceStorage)); + let recorder = LanceRecorder { + registry: registry.clone(), + }; + metrics::with_local_recorder(&recorder, || { + metrics::histogram!(name).record(0.02); + }); + + let points = collect_points(®istry); + let point = points.iter().find(|p| p.name == name).expect("recorded"); + let MetricValue::Histogram { buckets, count, .. } = &point.value else { + panic!("expected histogram"); + }; + assert_eq!(*count, 1); + // DEFAULT_BOUNDS yields one bucket per bound plus the implicit `+Inf`. + assert_eq!(buckets.len(), DEFAULT_BOUNDS.len() + 1); + // 0.02 falls in the le=0.025 bucket (the third DEFAULT_BOUNDS entry). + assert_eq!(bucket_count(buckets, "0.025"), 1); + assert_eq!(bucket_count(buckets, "0.01"), 0); + assert_eq!(bucket_count(buckets, "+Inf"), 1); + } + + #[test] + fn recorder_uses_registered_histogram_bounds() { + let name = "test_recorder_bounds_seconds"; + HISTOGRAM_BOUNDS + .write() + .unwrap() + .insert(name.to_string(), Arc::from([0.1f64, 1.0].as_slice())); + + let registry = Arc::new(Registry::new(LanceStorage)); + let recorder = LanceRecorder { + registry: registry.clone(), + }; + metrics::with_local_recorder(&recorder, || { + metrics::histogram!(name).record(0.05); + metrics::histogram!(name).record(5.0); + }); + + let points = collect_points(®istry); + let point = points.iter().find(|p| p.name == name).expect("recorded"); + let MetricValue::Histogram { buckets, count, .. } = &point.value else { + panic!("expected histogram"); + }; + assert_eq!(*count, 2); + assert_eq!(bucket_count(buckets, "0.1"), 1); + assert_eq!(bucket_count(buckets, "+Inf"), 2); + } + + #[test] + fn describe_populates_catalog() { + let name = "test_describe_catalog_total"; + let registry = Arc::new(Registry::new(LanceStorage)); + let recorder = LanceRecorder { registry }; + metrics::with_local_recorder(&recorder, || { + metrics::describe_counter!(name, Unit::Count, "a test counter"); + }); + + let catalog = CATALOG.lock().unwrap(); + let desc = catalog.get(name).expect("described"); + assert!(matches!(desc.kind, MetricKind::Counter)); + assert_eq!(desc.description, "a test counter"); + } + + #[test] + fn describe_all_covers_object_store_metrics() { + use lance_io::object_store::metrics as os; + + let registry = Arc::new(Registry::new(LanceStorage)); + let recorder = LanceRecorder { registry }; + metrics::with_local_recorder(&recorder, describe_all); + + let catalog = CATALOG.lock().unwrap(); + let kind = |name: &str| catalog.get(name).expect("described").kind; + // Every emitted object store metric must be described so the OTel + // bridge can create an instrument for it, including the gauge and the + // retryable counter that a plain request path might never emit. + assert!(matches!(kind(os::METRIC_REQUESTS), MetricKind::Counter)); + assert!(matches!(kind(os::METRIC_BYTES), MetricKind::Counter)); + assert!(matches!(kind(os::METRIC_ERRORS), MetricKind::Counter)); + assert!(matches!(kind(os::METRIC_THROTTLE), MetricKind::Counter)); + assert!(matches!(kind(os::METRIC_RETRYABLE), MetricKind::Counter)); + assert!(matches!(kind(os::METRIC_IN_FLIGHT), MetricKind::Gauge)); + assert!(matches!(kind(os::METRIC_DURATION), MetricKind::Histogram)); + } +} diff --git a/python/src/session.rs b/python/src/session.rs index c91329ec1ee..da3174a9bf9 100644 --- a/python/src/session.rs +++ b/python/src/session.rs @@ -3,7 +3,7 @@ use std::sync::Arc; -use pyo3::{pyclass, pymethods}; +use pyo3::{PyResult, pyclass, pymethods}; use lance::dataset::{DEFAULT_INDEX_CACHE_SIZE, DEFAULT_METADATA_CACHE_SIZE}; use lance::session::Session as LanceSession; @@ -63,6 +63,13 @@ impl Session { self.inner.size_bytes() } + /// Return the current size of the index cache in bytes. + pub fn index_cache_size_bytes(&self) -> PyResult { + rt().block_on(None, async move { + self.inner.index_cache_stats().await.size_bytes as u64 + }) + } + /// Return whether the other session is the same as this one. pub fn is_same_as(&self, other: &Self) -> bool { Arc::ptr_eq(&self.inner, &other.inner) diff --git a/python/uv.lock b/python/uv.lock index a8a7febeea4..ee3d2f872ce 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -9,15 +9,6 @@ resolution-markers = [ "python_full_version < '3.11'", ] -[[package]] -name = "absl-py" -version = "2.3.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/10/2a/c93173ffa1b39c1d0395b7e842bbdc62e556ca9d8d3b5572926f3e4ca752/absl_py-2.3.1.tar.gz", hash = "sha256:a97820526f7fbfd2ec1bce83f3f25e3a14840dac0d8e02a0b71cd75db3f77fc9", size = 116588, upload-time = "2025-07-03T09:31:44.05Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8f/aa/ba0014cc4659328dc818a28827be78e6d97312ab0cb98105a770924dc11e/absl_py-2.3.1-py3-none-any.whl", hash = "sha256:eeecf07f0c2a93ace0772c92e596ace6d3d3996c042b2128459aaae2a76de11d", size = 135811, upload-time = "2025-07-03T09:31:42.253Z" }, -] - [[package]] name = "aiohappyeyeballs" version = "2.6.1" @@ -248,19 +239,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0b/fd/4f8dac58ea17e05978bf35cb9a3e485b1ff3cdd6e2cc29deb08f54080de4/arro3_core-0.6.5-pp310-pypy310_pp73-musllinux_1_2_x86_64.whl", hash = "sha256:9a58acbc61480b533aa84d735db04b1e68fc7f6807ab694d606c03b5e694d83d", size = 2954405, upload-time = "2025-10-13T23:12:35.328Z" }, ] -[[package]] -name = "astunparse" -version = "1.6.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "six" }, - { name = "wheel" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f3/af/4182184d3c338792894f34a62672919db7ca008c89abee9b564dd34d8029/astunparse-1.6.3.tar.gz", hash = "sha256:5ad93a8456f0d084c3456d059fd9a92cce667963232cbf763eac3bc5b7940872", size = 18290, upload-time = "2019-12-22T18:12:13.129Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2b/03/13dde6512ad7b4557eb792fbcf0c653af6076b81e5941d36ec61f7ce6028/astunparse-1.6.3-py2.py3-none-any.whl", hash = "sha256:c2652417f2c8b5bb325c885ae329bdf3f86424075c4fd1a128674bc6fba4b8e8", size = 12732, upload-time = "2019-12-22T18:12:11.297Z" }, -] - [[package]] name = "async-timeout" version = "5.0.1" @@ -421,42 +399,51 @@ wheels = [ [[package]] name = "cuda-toolkit" -version = "13.0.2" +version = "13.0.3.0" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl", hash = "sha256:b198824cf2f54003f50d64ada3a0f184b42ca0846c1c94192fa269ecd97a66eb", size = 2364, upload-time = "2025-12-19T23:24:07.328Z" }, + { url = "https://files.pythonhosted.org/packages/d1/c7/a79086a62c98befcdb8349656c6f114e2db3b8b2422f6e25c97a7f2a9a3c/cuda_toolkit-13.0.3.0-py2.py3-none-any.whl", hash = "sha256:d693caaa261214ddd7dbb60d68e71cbed884e68c2be7509778f3051da0b91c3f", size = 2512, upload-time = "2026-04-14T00:50:08.173Z" }, ] [package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, +] cudart = [ - { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, ] cufft = [ - { name = "nvidia-cufft", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cufft", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, ] cufile = [ - { name = "nvidia-cufile", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cufile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] cupti = [ - { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cuda-cupti", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, ] curand = [ - { name = "nvidia-curand", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-curand", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, ] cusolver = [ - { name = "nvidia-cusolver", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "nvidia-cusolver", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, ] cusparse = [ - { name = "nvidia-cusparse", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, ] nvtx = [ - { name = "nvidia-nvtx", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "nvidia-nvtx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, ] [[package]] @@ -501,6 +488,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f4/c8/09012ac195a0aab58755800d2efdc0e7d5905053509f12cb5d136c911cda/datasets-4.1.1-py3-none-any.whl", hash = "sha256:62e4f6899a36be9ec74a7e759a6951253cc85b3fcfa0a759b0efa8353b149dac", size = 503623, upload-time = "2025-09-18T13:14:25.111Z" }, ] +[[package]] +name = "deprecated" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/49/85/12f0a49a7c4ffb70572b6c2ef13c90c88fd190debda93b23f026b25f9634/deprecated-1.3.1.tar.gz", hash = "sha256:b1b50e0ff0c1fddaa5708a2c6b0a6588bb09b892825ab2b214ac9ea9d92a5223", size = 2932523, upload-time = "2025-10-30T08:19:02.757Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/d0/205d54408c08b13550c733c4b85429e7ead111c7f0014309637425520a9a/deprecated-1.3.1-py2.py3-none-any.whl", hash = "sha256:597bfef186b6f60181535a29fbe44865ce137a5079f295b479886c82729d5f3f", size = 11298, upload-time = "2025-10-30T08:19:00.758Z" }, +] + [[package]] name = "dill" version = "0.4.0" @@ -563,15 +562,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/42/14/42b2651a2f46b022ccd948bca9f2d5af0fd8929c4eec235b8d6d844fbe67/filelock-3.19.1-py3-none-any.whl", hash = "sha256:d38e30481def20772f5baf097c122c3babc4fcdb7e14e57049eb9d88c6dc017d", size = 15988, upload-time = "2025-08-14T16:56:01.633Z" }, ] -[[package]] -name = "flatbuffers" -version = "25.9.23" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/9d/1f/3ee70b0a55137442038f2a33469cc5fddd7e0ad2abf83d7497c18a2b6923/flatbuffers-25.9.23.tar.gz", hash = "sha256:676f9fa62750bb50cf531b42a0a2a118ad8f7f797a511eda12881c016f093b12", size = 22067, upload-time = "2025-09-24T05:25:30.106Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ee/1b/00a78aa2e8fbd63f9af08c9c19e6deb3d5d66b4dda677a0f61654680ee89/flatbuffers-25.9.23-py2.py3-none-any.whl", hash = "sha256:255538574d6cb6d0a79a17ec8bc0d30985913b87513a01cce8bcdb6b4c44d0e2", size = 30869, upload-time = "2025-09-24T05:25:28.912Z" }, -] - [[package]] name = "frozenlist" version = "1.7.0" @@ -680,15 +670,6 @@ http = [ { name = "aiohttp" }, ] -[[package]] -name = "gast" -version = "0.6.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/3c/14/c566f5ca00c115db7725263408ff952b8ae6d6a4e792ef9c84e77d9af7a1/gast-0.6.0.tar.gz", hash = "sha256:88fc5300d32c7ac6ca7b515310862f71e6fdf2c029bbec7c66c0f5dd47b6b1fb", size = 27708, upload-time = "2024-06-27T20:31:49.527Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a3/61/8001b38461d751cd1a0c3a6ae84346796a5758123f3ed97a1b121dfbf4f3/gast-0.6.0-py3-none-any.whl", hash = "sha256:52b182313f7330389f72b069ba00f174cfe2a06411099547288839c6cbafbd54", size = 21173, upload-time = "2024-07-09T13:15:15.615Z" }, -] - [[package]] name = "geoarrow-rust-core" version = "0.6.3" @@ -805,84 +786,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6b/7e/6196a7b6c63c0875474a2c2319f2a2d92bb4acd4a8d260e1e10726ccff2b/geoarrow_rust_io-0.6.1-pp310-pypy310_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:209ddc68c06a2f8577deaf4d744eac21696872f21d367a3ec0b15dc7cf824d5b", size = 10404698, upload-time = "2025-11-21T02:10:23.556Z" }, ] -[[package]] -name = "google-pasta" -version = "0.2.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "six" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/35/4a/0bd53b36ff0323d10d5f24ebd67af2de10a1117f5cf4d7add90df92756f1/google-pasta-0.2.0.tar.gz", hash = "sha256:c9f2c8dfc8f96d0d5808299920721be30c9eec37f2389f28904f454565c8a16e", size = 40430, upload-time = "2020-03-13T18:57:50.34Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a3/de/c648ef6835192e6e2cc03f40b19eeda4382c49b5bafb43d88b931c4c74ac/google_pasta-0.2.0-py3-none-any.whl", hash = "sha256:b32482794a366b5366a32c92a9a9201b107821889935a02b3e51f6b432ea84ed", size = 57471, upload-time = "2020-03-13T18:57:48.872Z" }, -] - -[[package]] -name = "grpcio" -version = "1.75.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/9d/f7/8963848164c7604efb3a3e6ee457fdb3a469653e19002bd24742473254f8/grpcio-1.75.1.tar.gz", hash = "sha256:3e81d89ece99b9ace23a6916880baca613c03a799925afb2857887efa8b1b3d2", size = 12731327, upload-time = "2025-09-26T09:03:36.887Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/51/57/89fd829fb00a6d0bee3fbcb2c8a7aa0252d908949b6ab58bfae99d39d77e/grpcio-1.75.1-cp310-cp310-linux_armv7l.whl", hash = "sha256:1712b5890b22547dd29f3215c5788d8fc759ce6dd0b85a6ba6e2731f2d04c088", size = 5705534, upload-time = "2025-09-26T09:00:52.225Z" }, - { url = "https://files.pythonhosted.org/packages/9a/3d/affe2fb897804c98d56361138e73786af8f4dd876b9d9851cfe6342b53c8/grpcio-1.75.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:683cfc70be0c1383449097cba637317e4737a357cfc185d887fd984206380403", size = 6289953, upload-time = "2025-09-26T09:01:03.699Z" }, - { url = "https://files.pythonhosted.org/packages/87/aa/0f40b7f47a0ff10d7e482bc3af22dac767c7ff27205915f08962d5ca87a2/grpcio-1.75.1-cp310-cp310-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:491444c081a54dcd5e6ada57314321ae526377f498d4aa09d975c3241c5b9e1c", size = 6949785, upload-time = "2025-09-26T09:01:07.504Z" }, - { url = "https://files.pythonhosted.org/packages/a5/45/b04407e44050781821c84f26df71b3f7bc469923f92f9f8bc27f1406dbcc/grpcio-1.75.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ce08d4e112d0d38487c2b631ec8723deac9bc404e9c7b1011426af50a79999e4", size = 6465708, upload-time = "2025-09-26T09:01:11.028Z" }, - { url = "https://files.pythonhosted.org/packages/09/3e/4ae3ec0a4d20dcaafbb6e597defcde06399ccdc5b342f607323f3b47f0a3/grpcio-1.75.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:5a2acda37fc926ccc4547977ac3e56b1df48fe200de968e8c8421f6e3093df6c", size = 7100912, upload-time = "2025-09-26T09:01:14.393Z" }, - { url = "https://files.pythonhosted.org/packages/34/3f/a9085dab5c313bb0cb853f222d095e2477b9b8490a03634cdd8d19daa5c3/grpcio-1.75.1-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:745c5fe6bf05df6a04bf2d11552c7d867a2690759e7ab6b05c318a772739bd75", size = 8042497, upload-time = "2025-09-26T09:01:17.759Z" }, - { url = "https://files.pythonhosted.org/packages/c3/87/ea54eba931ab9ed3f999ba95f5d8d01a20221b664725bab2fe93e3dee848/grpcio-1.75.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:259526a7159d39e2db40d566fe3e8f8e034d0fb2db5bf9c00e09aace655a4c2b", size = 7493284, upload-time = "2025-09-26T09:01:20.896Z" }, - { url = "https://files.pythonhosted.org/packages/0c/3c/35ca9747473a306bfad0cee04504953f7098527cd112a4ab55c55af9e7bd/grpcio-1.75.1-cp311-cp311-linux_armv7l.whl", hash = "sha256:573855ca2e58e35032aff30bfbd1ee103fbcf4472e4b28d4010757700918e326", size = 5709761, upload-time = "2025-09-26T09:01:28.528Z" }, - { url = "https://files.pythonhosted.org/packages/81/40/bc07aee2911f0d426fa53fe636216100c31a8ea65a400894f280274cb023/grpcio-1.75.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b1ea1bbe77ecbc1be00af2769f4ae4a88ce93be57a4f3eebd91087898ed749f9", size = 6296084, upload-time = "2025-09-26T09:01:34.596Z" }, - { url = "https://files.pythonhosted.org/packages/b8/d1/10c067f6c67396cbf46448b80f27583b5e8c4b46cdfbe18a2a02c2c2f290/grpcio-1.75.1-cp311-cp311-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:e5b425aee54cc5e3e3c58f00731e8a33f5567965d478d516d35ef99fd648ab68", size = 6950403, upload-time = "2025-09-26T09:01:36.736Z" }, - { url = "https://files.pythonhosted.org/packages/3f/42/5f628abe360b84dfe8dd8f32be6b0606dc31dc04d3358eef27db791ea4d5/grpcio-1.75.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0049a7bf547dafaeeb1db17079ce79596c298bfe308fc084d023c8907a845b9a", size = 6470166, upload-time = "2025-09-26T09:01:39.474Z" }, - { url = "https://files.pythonhosted.org/packages/c3/93/a24035080251324019882ee2265cfde642d6476c0cf8eb207fc693fcebdc/grpcio-1.75.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:5b8ea230c7f77c0a1a3208a04a1eda164633fb0767b4cefd65a01079b65e5b1f", size = 7107828, upload-time = "2025-09-26T09:01:41.782Z" }, - { url = "https://files.pythonhosted.org/packages/e4/f8/d18b984c1c9ba0318e3628dbbeb6af77a5007f02abc378c845070f2d3edd/grpcio-1.75.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:36990d629c3c9fb41e546414e5af52d0a7af37ce7113d9682c46d7e2919e4cca", size = 8045421, upload-time = "2025-09-26T09:01:45.835Z" }, - { url = "https://files.pythonhosted.org/packages/7e/b6/4bf9aacff45deca5eac5562547ed212556b831064da77971a4e632917da3/grpcio-1.75.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:b10ad908118d38c2453ade7ff790e5bce36580c3742919007a2a78e3a1e521ca", size = 7503290, upload-time = "2025-09-26T09:01:49.28Z" }, - { url = "https://files.pythonhosted.org/packages/3a/81/42be79e73a50aaa20af66731c2defeb0e8c9008d9935a64dd8ea8e8c44eb/grpcio-1.75.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:7b888b33cd14085d86176b1628ad2fcbff94cfbbe7809465097aa0132e58b018", size = 5668314, upload-time = "2025-09-26T09:01:55.424Z" }, - { url = "https://files.pythonhosted.org/packages/14/85/21c71d674f03345ab183c634ecd889d3330177e27baea8d5d247a89b6442/grpcio-1.75.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bb658f703468d7fbb5dcc4037c65391b7dc34f808ac46ed9136c24fc5eeb041d", size = 6246335, upload-time = "2025-09-26T09:02:00.76Z" }, - { url = "https://files.pythonhosted.org/packages/fd/db/3beb661bc56a385ae4fa6b0e70f6b91ac99d47afb726fe76aaff87ebb116/grpcio-1.75.1-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:4b7177a1cdb3c51b02b0c0a256b0a72fdab719600a693e0e9037949efffb200b", size = 6916309, upload-time = "2025-09-26T09:02:02.894Z" }, - { url = "https://files.pythonhosted.org/packages/1e/9c/eda9fe57f2b84343d44c1b66cf3831c973ba29b078b16a27d4587a1fdd47/grpcio-1.75.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7d4fa6ccc3ec2e68a04f7b883d354d7fea22a34c44ce535a2f0c0049cf626ddf", size = 6435419, upload-time = "2025-09-26T09:02:05.055Z" }, - { url = "https://files.pythonhosted.org/packages/c3/b8/090c98983e0a9d602e3f919a6e2d4e470a8b489452905f9a0fa472cac059/grpcio-1.75.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3d86880ecaeb5b2f0a8afa63824de93adb8ebe4e49d0e51442532f4e08add7d6", size = 7064893, upload-time = "2025-09-26T09:02:07.275Z" }, - { url = "https://files.pythonhosted.org/packages/ec/c0/6d53d4dbbd00f8bd81571f5478d8a95528b716e0eddb4217cc7cb45aae5f/grpcio-1.75.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:a8041d2f9e8a742aeae96f4b047ee44e73619f4f9d24565e84d5446c623673b6", size = 8011922, upload-time = "2025-09-26T09:02:09.527Z" }, - { url = "https://files.pythonhosted.org/packages/f2/7c/48455b2d0c5949678d6982c3e31ea4d89df4e16131b03f7d5c590811cbe9/grpcio-1.75.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:3652516048bf4c314ce12be37423c79829f46efffb390ad64149a10c6071e8de", size = 7466181, upload-time = "2025-09-26T09:02:12.279Z" }, - { url = "https://files.pythonhosted.org/packages/46/74/bac4ab9f7722164afdf263ae31ba97b8174c667153510322a5eba4194c32/grpcio-1.75.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:3bed22e750d91d53d9e31e0af35a7b0b51367e974e14a4ff229db5b207647884", size = 5672779, upload-time = "2025-09-26T09:02:19.11Z" }, - { url = "https://files.pythonhosted.org/packages/cf/e4/d1954dce2972e32384db6a30273275e8c8ea5a44b80347f9055589333b3f/grpcio-1.75.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5bf4001d3293e3414d0cf99ff9b1139106e57c3a66dfff0c5f60b2a6286ec133", size = 6248838, upload-time = "2025-09-26T09:02:26.426Z" }, - { url = "https://files.pythonhosted.org/packages/06/43/073363bf63826ba8077c335d797a8d026f129dc0912b69c42feaf8f0cd26/grpcio-1.75.1-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:9f82ff474103e26351dacfe8d50214e7c9322960d8d07ba7fa1d05ff981c8b2d", size = 6922663, upload-time = "2025-09-26T09:02:28.724Z" }, - { url = "https://files.pythonhosted.org/packages/c2/6f/076ac0df6c359117676cacfa8a377e2abcecec6a6599a15a672d331f6680/grpcio-1.75.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0ee119f4f88d9f75414217823d21d75bfe0e6ed40135b0cbbfc6376bc9f7757d", size = 6436149, upload-time = "2025-09-26T09:02:30.971Z" }, - { url = "https://files.pythonhosted.org/packages/6b/27/1d08824f1d573fcb1fa35ede40d6020e68a04391709939e1c6f4193b445f/grpcio-1.75.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:664eecc3abe6d916fa6cf8dd6b778e62fb264a70f3430a3180995bf2da935446", size = 7067989, upload-time = "2025-09-26T09:02:33.233Z" }, - { url = "https://files.pythonhosted.org/packages/c6/98/98594cf97b8713feb06a8cb04eeef60b4757e3e2fb91aa0d9161da769843/grpcio-1.75.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:c32193fa08b2fbebf08fe08e84f8a0aad32d87c3ad42999c65e9449871b1c66e", size = 8010717, upload-time = "2025-09-26T09:02:36.011Z" }, - { url = "https://files.pythonhosted.org/packages/8c/7e/bb80b1bba03c12158f9254762cdf5cced4a9bc2e8ed51ed335915a5a06ef/grpcio-1.75.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5cebe13088b9254f6e615bcf1da9131d46cfa4e88039454aca9cb65f639bd3bc", size = 7463822, upload-time = "2025-09-26T09:02:38.26Z" }, - { url = "https://files.pythonhosted.org/packages/f2/1b/9a0a5cecd24302b9fdbcd55d15ed6267e5f3d5b898ff9ac8cbe17ee76129/grpcio-1.75.1-cp314-cp314-linux_armv7l.whl", hash = "sha256:c05da79068dd96723793bffc8d0e64c45f316248417515f28d22204d9dae51c7", size = 5673319, upload-time = "2025-09-26T09:02:44.742Z" }, - { url = "https://files.pythonhosted.org/packages/09/7a/26da709e42c4565c3d7bf999a9569da96243ce34a8271a968dee810a7cf1/grpcio-1.75.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4484f4b7287bdaa7a5b3980f3c7224c3c622669405d20f69549f5fb956ad0421", size = 6254706, upload-time = "2025-09-26T09:02:50.4Z" }, - { url = "https://files.pythonhosted.org/packages/f1/08/dcb26a319d3725f199c97e671d904d84ee5680de57d74c566a991cfab632/grpcio-1.75.1-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:2720c239c1180eee69f7883c1d4c83fc1a495a2535b5fa322887c70bf02b16e8", size = 6922501, upload-time = "2025-09-26T09:02:52.711Z" }, - { url = "https://files.pythonhosted.org/packages/78/66/044d412c98408a5e23cb348845979a2d17a2e2b6c3c34c1ec91b920f49d0/grpcio-1.75.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:07a554fa31c668cf0e7a188678ceeca3cb8fead29bbe455352e712ec33ca701c", size = 6437492, upload-time = "2025-09-26T09:02:55.542Z" }, - { url = "https://files.pythonhosted.org/packages/4e/9d/5e3e362815152aa1afd8b26ea613effa005962f9da0eec6e0e4527e7a7d1/grpcio-1.75.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:3e71a2105210366bfc398eef7f57a664df99194f3520edb88b9c3a7e46ee0d64", size = 7081061, upload-time = "2025-09-26T09:02:58.261Z" }, - { url = "https://files.pythonhosted.org/packages/1e/1a/46615682a19e100f46e31ddba9ebc297c5a5ab9ddb47b35443ffadb8776c/grpcio-1.75.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:8679aa8a5b67976776d3c6b0521e99d1c34db8a312a12bcfd78a7085cb9b604e", size = 8010849, upload-time = "2025-09-26T09:03:00.548Z" }, - { url = "https://files.pythonhosted.org/packages/67/8e/3204b94ac30b0f675ab1c06540ab5578660dc8b690db71854d3116f20d00/grpcio-1.75.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:aad1c774f4ebf0696a7f148a56d39a3432550612597331792528895258966dc0", size = 7464478, upload-time = "2025-09-26T09:03:03.096Z" }, -] - -[[package]] -name = "h5py" -version = "3.14.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/5d/57/dfb3c5c3f1bf5f5ef2e59a22dec4ff1f3d7408b55bfcefcfb0ea69ef21c6/h5py-3.14.0.tar.gz", hash = "sha256:2372116b2e0d5d3e5e705b7f663f7c8d96fa79a4052d250484ef91d24d6a08f4", size = 424323, upload-time = "2025-06-06T14:06:15.01Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fa/cd/3dd38cdb7cc9266dc4d85f27f0261680cb62f553f1523167ad7454e32b11/h5py-3.14.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:016e89d3be4c44f8d5e115fab60548e518ecd9efe9fa5c5324505a90773e6f03", size = 4324677, upload-time = "2025-06-06T14:04:23.438Z" }, - { url = "https://files.pythonhosted.org/packages/b1/45/e1a754dc7cd465ba35e438e28557119221ac89b20aaebef48282654e3dc7/h5py-3.14.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1223b902ef0b5d90bcc8a4778218d6d6cd0f5561861611eda59fa6c52b922f4d", size = 4557272, upload-time = "2025-06-06T14:04:28.863Z" }, - { url = "https://files.pythonhosted.org/packages/08/0c/5e6aaf221557314bc15ba0e0da92e40b24af97ab162076c8ae009320a42b/h5py-3.14.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8c497600c0496548810047257e36360ff551df8b59156d3a4181072eed47d8ad", size = 4298002, upload-time = "2025-06-06T14:04:47.106Z" }, - { url = "https://files.pythonhosted.org/packages/21/d4/d461649cafd5137088fb7f8e78fdc6621bb0c4ff2c090a389f68e8edc136/h5py-3.14.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:723a40ee6505bd354bfd26385f2dae7bbfa87655f4e61bab175a49d72ebfc06b", size = 4516618, upload-time = "2025-06-06T14:04:52.467Z" }, - { url = "https://files.pythonhosted.org/packages/3f/19/c8bfe8543bfdd7ccfafd46d8cfd96fce53d6c33e9c7921f375530ee1d39a/h5py-3.14.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:554ef0ced3571366d4d383427c00c966c360e178b5fb5ee5bb31a435c424db0c", size = 4708455, upload-time = "2025-06-06T14:05:11.528Z" }, - { url = "https://files.pythonhosted.org/packages/86/f9/f00de11c82c88bfc1ef22633557bfba9e271e0cb3189ad704183fc4a2644/h5py-3.14.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0cbd41f4e3761f150aa5b662df991868ca533872c95467216f2bec5fcad84882", size = 4929422, upload-time = "2025-06-06T14:05:18.399Z" }, - { url = "https://files.pythonhosted.org/packages/0d/ce/3a21d87896bc7e3e9255e0ad5583ae31ae9e6b4b00e0bcb2a67e2b6acdbc/h5py-3.14.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e8cbaf6910fa3983c46172666b0b8da7b7bd90d764399ca983236f2400436eeb", size = 4700675, upload-time = "2025-06-06T14:05:37.38Z" }, - { url = "https://files.pythonhosted.org/packages/e7/ec/86f59025306dcc6deee5fda54d980d077075b8d9889aac80f158bd585f1b/h5py-3.14.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d90e6445ab7c146d7f7981b11895d70bc1dd91278a4f9f9028bc0c95e4a53f13", size = 4921632, upload-time = "2025-06-06T14:05:43.464Z" }, -] - [[package]] name = "hf-xet" version = "1.1.10" @@ -926,6 +829,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d2/23/408243171aa9aaba178d3e2559159c24c1171a641aa83b67bdd3394ead8e/idna-3.15-py3-none-any.whl", hash = "sha256:048adeaf8c2d788c40fee287673ccaa74c24ffd8dcf09ffa555a2fbb59f10ac8", size = 72340, upload-time = "2026-05-12T22:45:55.733Z" }, ] +[[package]] +name = "importlib-metadata" +version = "8.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "zipp" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cd/12/33e59336dca5be0c398a7482335911a33aa0e20776128f038019f1a95f1b/importlib_metadata-8.5.0.tar.gz", hash = "sha256:71522656f0abace1d072b9e5481a48f07c138e00f079c38c8f883823f9c26bd7", size = 55304, upload-time = "2024-09-11T14:56:08.937Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/d9/a1e041c5e7caa9a05c925f4bdbdfb7f006d1f74996af53467bc394c97be7/importlib_metadata-8.5.0-py3-none-any.whl", hash = "sha256:45e54197d28b7a7f1559e60b95e7c567032b602131fbd588f1497f47880aa68b", size = 26514, upload-time = "2024-09-11T14:56:07.019Z" }, +] + [[package]] name = "iniconfig" version = "2.1.0" @@ -956,26 +871,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/31/b4/b9b800c45527aadd64d5b442f9b932b00648617eb5d63d2c7a6587b7cafc/jmespath-1.0.1-py3-none-any.whl", hash = "sha256:02e2e4cc71b5bcab88332eebf907519190dd9e6e82107fa7f83b1003a6252980", size = 20256, upload-time = "2022-06-17T18:00:10.251Z" }, ] -[[package]] -name = "keras" -version = "3.11.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "absl-py" }, - { name = "h5py" }, - { name = "ml-dtypes" }, - { name = "namex" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "optree" }, - { name = "packaging" }, - { name = "rich" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/6a/89/646425fe9a46f9053430e1271f817c36041c6f33469950a3caafc3d2591e/keras-3.11.3.tar.gz", hash = "sha256:efda616835c31b7d916d72303ef9adec1257320bc9fd4b2b0138840fc65fb5b7", size = 1065906, upload-time = "2025-08-21T22:08:57.643Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/94/5b/4c778cc921ce4b864b238f63f8e3ff6e954ab19b80c9fa680593ad8093d4/keras-3.11.3-py3-none-any.whl", hash = "sha256:f484f050e05ee400455b05ec8c36ed35edc34de94256b6073f56cfe68f65491f", size = 1408438, upload-time = "2025-08-21T22:08:55.858Z" }, -] - [[package]] name = "lance-namespace" version = "0.8.6" @@ -1003,39 +898,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c5/90/1e27de15cd1b16785a1c7312beb0a59e75c8344a815f600f58173a565bd1/lance_namespace_urllib3_client-0.8.6-py3-none-any.whl", hash = "sha256:9d78249c3fb15aa3d15d668f78f04a275af3d08d800a7027492f37996ac4968b", size = 369950, upload-time = "2026-06-12T17:36:40.438Z" }, ] -[[package]] -name = "libclang" -version = "18.1.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/6e/5c/ca35e19a4f142adffa27e3d652196b7362fa612243e2b916845d801454fc/libclang-18.1.1.tar.gz", hash = "sha256:a1214966d08d73d971287fc3ead8dfaf82eb07fb197680d8b3859dbbbbf78250", size = 39612, upload-time = "2024-03-17T16:04:37.434Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1d/fc/716c1e62e512ef1c160e7984a73a5fc7df45166f2ff3f254e71c58076f7c/libclang-18.1.1-py2.py3-none-manylinux2010_x86_64.whl", hash = "sha256:c533091d8a3bbf7460a00cb6c1a71da93bffe148f172c7d03b1c31fbf8aa2a0b", size = 24515943, upload-time = "2024-03-17T16:03:45.942Z" }, - { url = "https://files.pythonhosted.org/packages/3c/3d/f0ac1150280d8d20d059608cf2d5ff61b7c3b7f7bcf9c0f425ab92df769a/libclang-18.1.1-py2.py3-none-manylinux2014_aarch64.whl", hash = "sha256:54dda940a4a0491a9d1532bf071ea3ef26e6dbaf03b5000ed94dd7174e8f9592", size = 23784972, upload-time = "2024-03-17T16:12:47.677Z" }, - { url = "https://files.pythonhosted.org/packages/fe/2f/d920822c2b1ce9326a4c78c0c2b4aa3fde610c7ee9f631b600acb5376c26/libclang-18.1.1-py2.py3-none-manylinux2014_armv7l.whl", hash = "sha256:cf4a99b05376513717ab5d82a0db832c56ccea4fd61a69dbb7bccf2dfb207dbe", size = 20259606, upload-time = "2024-03-17T16:17:42.437Z" }, - { url = "https://files.pythonhosted.org/packages/2d/c2/de1db8c6d413597076a4259cea409b83459b2db997c003578affdd32bf66/libclang-18.1.1-py2.py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:69f8eb8f65c279e765ffd28aaa7e9e364c776c17618af8bff22a8df58677ff4f", size = 24921494, upload-time = "2024-03-17T16:14:20.132Z" }, -] - -[[package]] -name = "markdown" -version = "3.9" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/8d/37/02347f6d6d8279247a5837082ebc26fc0d5aaeaf75aa013fcbb433c777ab/markdown-3.9.tar.gz", hash = "sha256:d2900fe1782bd33bdbbd56859defef70c2e78fc46668f8eb9df3128138f2cb6a", size = 364585, upload-time = "2025-09-04T20:25:22.885Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/70/ae/44c4a6a4cbb496d93c6257954260fe3a6e91b7bed2240e5dad2a717f5111/markdown-3.9-py3-none-any.whl", hash = "sha256:9f4d91ed810864ea88a6f32c07ba8bee1346c0cc1f6b1f9f6c822f2a9667d280", size = 107441, upload-time = "2025-09-04T20:25:21.784Z" }, -] - -[[package]] -name = "markdown-it-py" -version = "4.0.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "mdurl" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/5b/f5/4ec618ed16cc4f8fb3b701563655a69816155e79e24a17b651541804721d/markdown_it_py-4.0.0.tar.gz", hash = "sha256:cb0a2b4aa34f932c007117b194e945bd74e0ec24133ceb5bac59009cda1cb9f3", size = 73070, upload-time = "2025-08-11T12:57:52.854Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" }, -] - [[package]] name = "markupsafe" version = "3.0.3" @@ -1145,15 +1007,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6c/93/e32e79333f0902ba292b996f504f5f06be59587f7d02ab8d5ed1e3066445/maturin-1.13.3-py3-none-win_arm64.whl", hash = "sha256:2389fe92d017cea9d94e521fa0175314a4c52f79a1057b901fbc9f8686ef7d0b", size = 9706562, upload-time = "2026-05-11T07:43:31.743Z" }, ] -[[package]] -name = "mdurl" -version = "0.1.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, -] - [[package]] name = "ml-dtypes" version = "0.5.3" @@ -1325,15 +1178,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/d9/f7f9379981e39b8c2511c9e0326d212accacb82f12fbfdc1aa2ce2a7b2b6/multiprocess-0.70.16-py39-none-any.whl", hash = "sha256:a0bafd3ae1b732eac64be2e72038231c1ba97724b60b09400d68f229fcc2fbf3", size = 133351, upload-time = "2024-01-28T18:52:31.981Z" }, ] -[[package]] -name = "namex" -version = "0.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0c/c0/ee95b28f029c73f8d49d8f52edaed02a1d4a9acb8b69355737fdb1faa191/namex-0.1.0.tar.gz", hash = "sha256:117f03ccd302cc48e3f5c58a296838f6b89c83455ab8683a1e85f2a430aa4306", size = 6649, upload-time = "2025-05-26T23:17:38.918Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b2/bc/465daf1de06409cdd4532082806770ee0d8d7df434da79c76564d0f69741/namex-0.1.0-py3-none-any.whl", hash = "sha256:e2012a474502f1e2251267062aae3114611f07df4224b6e06334c57b0f2ce87c", size = 5905, upload-time = "2025-05-26T23:17:37.695Z" }, -] - [[package]] name = "networkx" version = "3.4.2" @@ -1674,6 +1518,46 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, ] +[[package]] +name = "opentelemetry-api" +version = "1.30.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "deprecated" }, + { name = "importlib-metadata" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2b/6d/bbbf879826b7f3c89a45252010b5796fb1f1a0d45d9dc4709db0ef9a06c8/opentelemetry_api-1.30.0.tar.gz", hash = "sha256:375893400c1435bf623f7dfb3bcd44825fe6b56c34d0667c542ea8257b1a1240", size = 63703, upload-time = "2025-02-04T18:17:13.789Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/36/0a/eea862fae6413d8181b23acf8e13489c90a45f17986ee9cf4eab8a0b9ad9/opentelemetry_api-1.30.0-py3-none-any.whl", hash = "sha256:d5f5284890d73fdf47f843dda3210edf37a38d66f44f2b5aedc1e89ed455dc09", size = 64955, upload-time = "2025-02-04T18:16:46.167Z" }, +] + +[[package]] +name = "opentelemetry-sdk" +version = "1.30.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/93/ee/d710062e8a862433d1be0b85920d0c653abe318878fef2d14dfe2c62ff7b/opentelemetry_sdk-1.30.0.tar.gz", hash = "sha256:c9287a9e4a7614b9946e933a67168450b9ab35f08797eb9bc77d998fa480fa18", size = 158633, upload-time = "2025-02-04T18:17:28.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/97/28/64d781d6adc6bda2260067ce2902bd030cf45aec657e02e28c5b4480b976/opentelemetry_sdk-1.30.0-py3-none-any.whl", hash = "sha256:14fe7afc090caad881addb6926cec967129bd9260c4d33ae6a217359f6b61091", size = 118717, upload-time = "2025-02-04T18:17:09.353Z" }, +] + +[[package]] +name = "opentelemetry-semantic-conventions" +version = "0.51b0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "deprecated" }, + { name = "opentelemetry-api" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1e/c0/0f9ef4605fea7f2b83d55dd0b0d7aebe8feead247cd6facd232b30907b4f/opentelemetry_semantic_conventions-0.51b0.tar.gz", hash = "sha256:3fabf47f35d1fd9aebcdca7e6802d86bd5ebc3bc3408b7e3248dde6e87a18c47", size = 107191, upload-time = "2025-02-04T18:17:29.903Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/75/d7bdbb6fd8630b4cafb883482b75c4fc276b6426619539d266e32ac53266/opentelemetry_semantic_conventions-0.51b0-py3-none-any.whl", hash = "sha256:fdc777359418e8d06c86012c3dc92c88a6453ba662e941593adb062e48c2eeae", size = 177416, upload-time = "2025-02-04T18:17:11.305Z" }, +] + [[package]] name = "opt-einsum" version = "3.4.0" @@ -2041,17 +1925,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cc/35/cc0aaecf278bb4575b8555f2b137de5ab821595ddae9da9d3cd1da4072c7/propcache-0.3.2-py3-none-any.whl", hash = "sha256:98f1ec44fb675f5052cccc8e609c46ed23a35a1cfd18545ad4e29002d858a43f", size = 12663, upload-time = "2025-06-09T22:56:04.484Z" }, ] -[[package]] -name = "protobuf" -version = "6.32.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/fa/a4/cc17347aa2897568beece2e674674359f911d6fe21b0b8d6268cd42727ac/protobuf-6.32.1.tar.gz", hash = "sha256:ee2469e4a021474ab9baafea6cd070e5bf27c7d29433504ddea1a4ee5850f68d", size = 440635, upload-time = "2025-09-11T21:38:42.935Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3f/be/8dd0a927c559b37d7a6c8ab79034fd167dcc1f851595f2e641ad62be8643/protobuf-6.32.1-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:2f5b80a49e1eb7b86d85fcd23fe92df154b9730a725c3b38c4e43b9d77018bf4", size = 322874, upload-time = "2025-09-11T21:38:35.509Z" }, - { url = "https://files.pythonhosted.org/packages/5c/f6/88d77011b605ef979aace37b7703e4eefad066f7e84d935e5a696515c2dd/protobuf-6.32.1-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:b1864818300c297265c83a4982fd3169f97122c299f56a56e2445c3698d34710", size = 322013, upload-time = "2025-09-11T21:38:37.017Z" }, - { url = "https://files.pythonhosted.org/packages/97/b7/15cc7d93443d6c6a84626ae3258a91f4c6ac8c0edd5df35ea7658f71b79c/protobuf-6.32.1-py3-none-any.whl", hash = "sha256:2601b779fc7d32a866c6b4404f9d42a3f67c5b9f3f15b4db3cccabe06b95c346", size = 169289, upload-time = "2025-09-11T21:38:41.234Z" }, -] - [[package]] name = "psutil" version = "7.1.0" @@ -2298,6 +2171,10 @@ geo = [ { name = "geoarrow-rust-core" }, { name = "geoarrow-rust-io" }, ] +otel = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-sdk" }, +] tests = [ { name = "boto3" }, { name = "datafusion" }, @@ -2309,7 +2186,6 @@ tests = [ { name = "polars", extra = ["pandas", "pyarrow"] }, { name = "psutil" }, { name = "pytest" }, - { name = "tensorflow", marker = "sys_platform == 'linux'" }, { name = "tqdm" }, ] torch = [ @@ -2331,12 +2207,12 @@ tests = [ { name = "datasets" }, { name = "duckdb" }, { name = "ml-dtypes" }, + { name = "opentelemetry-sdk" }, { name = "pandas" }, { name = "pillow" }, { name = "polars", extra = ["pandas", "pyarrow"] }, { name = "psutil" }, { name = "pytest" }, - { name = "tensorflow", marker = "sys_platform == 'linux'" }, { name = "tqdm" }, ] @@ -2351,6 +2227,8 @@ requires-dist = [ { name = "lance-namespace", specifier = ">=0.8.5,<0.9" }, { name = "ml-dtypes", marker = "extra == 'tests'" }, { name = "numpy", specifier = ">=1.22" }, + { name = "opentelemetry-api", marker = "extra == 'otel'" }, + { name = "opentelemetry-sdk", marker = "extra == 'otel'" }, { name = "pandas", marker = "extra == 'tests'" }, { name = "pillow", marker = "extra == 'tests'" }, { name = "polars", extras = ["pyarrow", "pandas"], marker = "extra == 'tests'" }, @@ -2360,11 +2238,10 @@ requires-dist = [ { name = "pytest", marker = "extra == 'tests'" }, { name = "pytest-benchmark", marker = "extra == 'benchmarks'" }, { name = "ruff", marker = "extra == 'dev'", specifier = "==0.11.2" }, - { name = "tensorflow", marker = "sys_platform == 'linux' and extra == 'tests'" }, { name = "torch", marker = "extra == 'torch'", specifier = ">=2.0" }, { name = "tqdm", marker = "extra == 'tests'" }, ] -provides-extras = ["benchmarks", "dev", "geo", "tests", "torch"] +provides-extras = ["benchmarks", "dev", "geo", "otel", "tests", "torch"] [package.metadata.requires-dev] benchmarks = [{ name = "pytest-benchmark", specifier = "==5.1.0" }] @@ -2379,12 +2256,12 @@ tests = [ { name = "datasets", specifier = "==4.1.1" }, { name = "duckdb", specifier = "==1.4.0" }, { name = "ml-dtypes", specifier = "==0.5.3" }, + { name = "opentelemetry-sdk", specifier = "==1.30.0" }, { name = "pandas", specifier = "==2.3.3" }, { name = "pillow", specifier = "==11.3.0" }, { name = "polars", extras = ["pyarrow", "pandas"], specifier = "==1.34.0" }, { name = "psutil", specifier = "==7.1.0" }, { name = "pytest", specifier = "==8.4.2" }, - { name = "tensorflow", marker = "sys_platform == 'linux'", specifier = "==2.20.0" }, { name = "tqdm", specifier = "==4.67.1" }, ] @@ -2649,19 +2526,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/56/5d/c814546c2333ceea4ba42262d8c4d55763003e767fa169adc693bd524478/requests-2.33.0-py3-none-any.whl", hash = "sha256:3324635456fa185245e24865e810cecec7b4caf933d7eb133dcde67d48cee69b", size = 65017, upload-time = "2026-03-25T15:10:40.382Z" }, ] -[[package]] -name = "rich" -version = "14.1.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "markdown-it-py" }, - { name = "pygments" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/fe/75/af448d8e52bf1d8fa6a9d089ca6c07ff4453d86c65c145d0a300bb073b9b/rich-14.1.0.tar.gz", hash = "sha256:e497a48b844b0320d45007cdebfeaeed8db2a4f4bcf49f15e455cfc4af11eaa8", size = 224441, upload-time = "2025-07-25T07:32:58.125Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e3/30/3c4d035596d3cf444529e0b2953ad0466f6049528a879d27534700580395/rich-14.1.0-py3-none-any.whl", hash = "sha256:536f5f1785986d6dbdea3c75205c473f970777b4a0d6c6dd1b696aa05a3fa04f", size = 243368, upload-time = "2025-07-25T07:32:56.73Z" }, -] - [[package]] name = "ruff" version = "0.11.2" @@ -2729,84 +2593,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, ] -[[package]] -name = "tensorboard" -version = "2.20.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "absl-py" }, - { name = "grpcio" }, - { name = "markdown" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "packaging" }, - { name = "pillow" }, - { name = "protobuf" }, - { name = "setuptools" }, - { name = "tensorboard-data-server" }, - { name = "werkzeug" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/9c/d9/a5db55f88f258ac669a92858b70a714bbbd5acd993820b41ec4a96a4d77f/tensorboard-2.20.0-py3-none-any.whl", hash = "sha256:9dc9f978cb84c0723acf9a345d96c184f0293d18f166bb8d59ee098e6cfaaba6", size = 5525680, upload-time = "2025-07-17T19:20:49.638Z" }, -] - -[[package]] -name = "tensorboard-data-server" -version = "0.7.2" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7a/13/e503968fefabd4c6b2650af21e110aa8466fe21432cd7c43a84577a89438/tensorboard_data_server-0.7.2-py3-none-any.whl", hash = "sha256:7e0610d205889588983836ec05dc098e80f97b7e7bbff7e994ebb78f578d0ddb", size = 2356, upload-time = "2023-10-23T21:23:32.16Z" }, - { url = "https://files.pythonhosted.org/packages/73/c6/825dab04195756cf8ff2e12698f22513b3db2f64925bdd41671bfb33aaa5/tensorboard_data_server-0.7.2-py3-none-manylinux_2_31_x86_64.whl", hash = "sha256:ef687163c24185ae9754ed5650eb5bc4d84ff257aabdc33f0cc6f74d8ba54530", size = 6590363, upload-time = "2023-10-23T21:23:35.583Z" }, -] - -[[package]] -name = "tensorflow" -version = "2.20.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "absl-py" }, - { name = "astunparse" }, - { name = "flatbuffers" }, - { name = "gast" }, - { name = "google-pasta" }, - { name = "grpcio" }, - { name = "h5py" }, - { name = "keras" }, - { name = "libclang" }, - { name = "ml-dtypes" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "opt-einsum" }, - { name = "packaging" }, - { name = "protobuf" }, - { name = "requests" }, - { name = "setuptools" }, - { name = "six" }, - { name = "tensorboard" }, - { name = "termcolor" }, - { name = "typing-extensions" }, - { name = "wrapt" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/ff/07/ea91ac67a9fd36d3372099f5a3e69860ded544f877f5f2117802388f4212/tensorflow-2.20.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:02a0293d94f5c8b7125b66abf622cc4854a33ae9d618a0d41309f95e091bbaea", size = 259307122, upload-time = "2025-08-13T16:50:47.909Z" }, - { url = "https://files.pythonhosted.org/packages/e5/9e/0d57922cf46b9e91de636cd5b5e0d7a424ebe98f3245380a713f1f6c2a0b/tensorflow-2.20.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7abd7f3a010e0d354dc804182372779a722d474c4d8a3db8f4a3f5baef2a591e", size = 620425510, upload-time = "2025-08-13T16:51:02.608Z" }, - { url = "https://files.pythonhosted.org/packages/f1/b7/a3d455db88ab5b35ce53ab885ec0dd9f28d905a86a2250423048bc8cafa0/tensorflow-2.20.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3e9568c8efcb05c0266be223e3269c62ebf7ad3498f156438311735f6fa5ced5", size = 259465882, upload-time = "2025-08-13T16:51:39.546Z" }, - { url = "https://files.pythonhosted.org/packages/ff/0c/7df285ee8a88139fab0b237003634d90690759fae9c18f55ddb7c04656ec/tensorflow-2.20.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:481499fd0f824583de8945be61d5e827898cdaa4f5ea1bc2cc28ca2ccff8229e", size = 620570129, upload-time = "2025-08-13T16:51:55.104Z" }, - { url = "https://files.pythonhosted.org/packages/ec/b4/f028a5de27d0fda10ba6145bc76e40c37ff6d2d1e95b601adb5ae17d635e/tensorflow-2.20.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2bfbfb3dd0e22bffc45fe1e922390d27753e99261fab8a882e802cf98a0e078f", size = 259533109, upload-time = "2025-08-13T16:52:31.513Z" }, - { url = "https://files.pythonhosted.org/packages/9c/d1/6aa15085d672056d5f08b5f28b1c7ce01c4e12149a23b0c98e3c79d04441/tensorflow-2.20.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:25265b0bc527e0d54b1e9cc60c44a24f44a809fe27666b905f0466471f9c52ec", size = 620682547, upload-time = "2025-08-13T16:52:46.396Z" }, - { url = "https://files.pythonhosted.org/packages/ea/4c/c1aa90c5cc92e9f7f9c78421e121ef25bae7d378f8d1d4cbad46c6308836/tensorflow-2.20.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:47c88e05a07f1ead4977b4894b3ecd4d8075c40191065afc4fd9355c9db3d926", size = 259663776, upload-time = "2025-08-13T16:53:24.507Z" }, - { url = "https://files.pythonhosted.org/packages/43/fb/8be8547c128613d82a2b006004026d86ed0bd672e913029a98153af4ffab/tensorflow-2.20.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5fa3729b0126f75a99882b89fb7d536515721eda8014a63e259e780ba0a37372", size = 620815537, upload-time = "2025-08-13T16:53:42.577Z" }, -] - -[[package]] -name = "termcolor" -version = "3.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ca/6c/3d75c196ac07ac8749600b60b03f4f6094d54e132c4d94ebac6ee0e0add0/termcolor-3.1.0.tar.gz", hash = "sha256:6a6dd7fbee581909eeec6a756cff1d7f7c376063b14e4a298dc4980309e55970", size = 14324, upload-time = "2025-04-30T11:37:53.791Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/4f/bd/de8d508070629b6d84a30d01d57e4a65c69aa7f5abe7560b8fad3b50ea59/termcolor-3.1.0-py3-none-any.whl", hash = "sha256:591dd26b5c2ce03b9e43f391264626557873ce1d379019786f99b0c2bee140aa", size = 7684, upload-time = "2025-04-30T11:37:52.382Z" }, -] - [[package]] name = "tomli" version = "2.2.1" @@ -2848,51 +2634,50 @@ wheels = [ [[package]] name = "torch" -version = "2.12.1" +version = "2.13.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, - { name = "cuda-toolkit", extra = ["cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, + { name = "cuda-bindings", marker = "python_full_version < '3.15' and sys_platform == 'linux'" }, + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, { name = "filelock" }, { name = "fsspec" }, { name = "jinja2" }, { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "nvidia-cublas", marker = "sys_platform == 'linux'" }, { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, { name = "setuptools" }, { name = "sympy" }, - { name = "triton", marker = "sys_platform == 'linux'" }, + { name = "triton", marker = "python_full_version < '3.15' and sys_platform == 'linux'" }, { name = "typing-extensions" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/db/ed/ff0c4f8cef63977a646dc80e40c05cae873f4097b12dc87e1cd7e1cecf42/torch-2.12.1-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:ec56e82be6a8b0c036771a77f7d32ad3c299770571af9815b3dafe61434389d5", size = 87967927, upload-time = "2026-06-17T21:08:43.16Z" }, - { url = "https://files.pythonhosted.org/packages/85/1b/c8ecf60c9dba535f9ea341c359c600c0bd877a7ca14b3296f13316321847/torch-2.12.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:42cd7339bf266f14944710e8274be63e7e012bb937834a8d85a8327a9860eba6", size = 426366829, upload-time = "2026-06-17T21:07:18.574Z" }, - { url = "https://files.pythonhosted.org/packages/ab/d6/73d4a3f27e00526e98086f3a64ab609af1345cca62367749fbc3c8e4b83c/torch-2.12.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:a7817f0f89a796d9de239d06f69faf5d7e19a6a5db6710a5ead777c912f9f50a", size = 532144834, upload-time = "2026-06-17T21:08:00.633Z" }, - { url = "https://files.pythonhosted.org/packages/e3/51/4010c8fa6f9d1f42c054a321970ca95ec58e4e4494f5b53a34c3f3c9e310/torch-2.12.1-cp310-cp310-win_amd64.whl", hash = "sha256:2af3d9cc866e0a15ae7635ff0a9c61d6624a353ad657f5bcd8d86c26cdc64693", size = 122949863, upload-time = "2026-06-17T21:08:39.016Z" }, - { url = "https://files.pythonhosted.org/packages/59/38/7028d3be540f1dcdf41660a2b01d0c51d2cb73915fe370d84e4d277a6d47/torch-2.12.1-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:ef81f503912effea2ce3d9b12a2e3a6ed488943e91271c90c7a829f60baf6aa2", size = 87975425, upload-time = "2026-06-17T21:08:34.094Z" }, - { url = "https://files.pythonhosted.org/packages/5a/e3/750b3e3548635ceac03ba255daa26dbc7ed66ca3484dc4b4d955ab7f4501/torch-2.12.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:107df6888624bdea41508f9aeb6149d9333c737a5530ceecb56c904e811369ae", size = 426379894, upload-time = "2026-06-17T21:06:55.077Z" }, - { url = "https://files.pythonhosted.org/packages/dc/ca/ed24783da629ff3e640ba3f70a7639e9045d3d88b93ee6bc47b8a28a1f2c/torch-2.12.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:6e29e7e74d05bda7d955c75e99459f878ebd970ef851b4057edbd3b34a5eb4a3", size = 532169264, upload-time = "2026-06-17T21:08:17.65Z" }, - { url = "https://files.pythonhosted.org/packages/46/61/c63f0158446f3a98ea672b004d761b848911eba567ea4a624c7db5aadc04/torch-2.12.1-cp311-cp311-win_amd64.whl", hash = "sha256:a513506cfda3c1c78dabeb6574c1597538c0254b3d39af174dde35d8177f4ce3", size = 122953086, upload-time = "2026-06-17T21:08:27.69Z" }, - { url = "https://files.pythonhosted.org/packages/f0/54/efb7ebca77970012b0cc21687a55d70eb2ba514b2c2b8e18d9fb1222f3be/torch-2.12.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:d2dd0f2c5f7ccbddaf34cade0deaf476808368f902b9cdb7f36a2ab42301bc0e", size = 87991951, upload-time = "2026-06-17T21:07:49.309Z" }, - { url = "https://files.pythonhosted.org/packages/1e/00/4210d76ca7424981f04033ebe7e48816ab83287a62538747a58825db770c/torch-2.12.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:2de4e19b88a481482c6c75291f2d6a52eda3ce51f311b29aa9b68499c830c07c", size = 426382721, upload-time = "2026-06-17T21:06:41.842Z" }, - { url = "https://files.pythonhosted.org/packages/76/1f/bc9f5a5aa569307076365f25afcebacb22e9c754b1bcfbaaa146627c7fda/torch-2.12.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:649e4ced014ba646f76f8cb9c9726735a6323eb321b7919f942790a923f90921", size = 532261322, upload-time = "2026-06-17T21:06:06.673Z" }, - { url = "https://files.pythonhosted.org/packages/9e/49/c549461daa008159d006a76a991fbc2f26fa8bac27a4030c858463dcb20f/torch-2.12.1-cp312-cp312-win_amd64.whl", hash = "sha256:e86550597877fb272ddc52db2f85b82cb601ea7bd932576a0340152cae2200b3", size = 122988095, upload-time = "2026-06-17T21:07:44.9Z" }, - { url = "https://files.pythonhosted.org/packages/ff/4a/0300261818e1560d72cc160ac826005507e8b7ca0a35788b591436d05b4a/torch-2.12.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:c75e93173c700bccd6bfcc4a9d19ce242ab6dacd1f1781483027a16239b9e650", size = 87992358, upload-time = "2026-06-17T21:07:40.299Z" }, - { url = "https://files.pythonhosted.org/packages/30/a7/874a5ca05e8f159211dca7921060f7057acc1adb26431e119fd150623efc/torch-2.12.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:fcb61ccd20784b62bdd78ec84238a5cfb383b4994902e03bac95505ab360884c", size = 426386134, upload-time = "2026-06-17T21:07:31.481Z" }, - { url = "https://files.pythonhosted.org/packages/e1/75/20bb8fe9c1ad6538cce8cd0391b51927ae5af0b17ed1eab44b8824465dc1/torch-2.12.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:f4afc8083dff08719edbea346644476e3cec0cf40ebe256be0ee5d5b7c7e8c0d", size = 532268019, upload-time = "2026-06-17T21:05:37.925Z" }, - { url = "https://files.pythonhosted.org/packages/d1/fa/824ddb662af55b2eabc0dbb7b57c7c0b1bcd93693754a2b8509ec4d16490/torch-2.12.1-cp313-cp313-win_amd64.whl", hash = "sha256:f92609e3b3ce72f25e2eb780d043ced2480c1a86c47c852604fc7a9108648386", size = 122987777, upload-time = "2026-06-17T21:07:09.49Z" }, - { url = "https://files.pythonhosted.org/packages/63/b7/1b49fe7086ea36839cc80abc43174c43d0ab6f676c0891c871c162f44fe3/torch-2.12.1-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:e9b6f7d2dd66ea87a3ae620069d31335d594c06effb1a383bdd21cfe61e44ece", size = 88010025, upload-time = "2026-06-17T21:07:03.934Z" }, - { url = "https://files.pythonhosted.org/packages/d7/06/5b44063a6545036dcc680d2d303b137d9176cfb2cc1e1863e3ef94abeb52/torch-2.12.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:7973ccd3d2cd35c74449213f7bded199bec6c6247e705cbeda7407af79703d91", size = 426392891, upload-time = "2026-06-17T21:05:52.261Z" }, - { url = "https://files.pythonhosted.org/packages/f8/dd/c9ce9a4b0eb3c5bb92d9ea56766e2c22559f0b45171149188494edcce80f/torch-2.12.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:c64ac4aac16be5e296dcd912305605804b203333c690bf98c55bc09494ee92ad", size = 532272494, upload-time = "2026-06-17T21:06:22.72Z" }, - { url = "https://files.pythonhosted.org/packages/21/7c/f3a601fc1b1f663ff269bfe553654e638651939aa6563e8daa7167c33098/torch-2.12.1-cp314-cp314-win_amd64.whl", hash = "sha256:f6dc4caf7eb4adb38a2d9f536b51db56310fdd1254e69a2d96767e1367c892b3", size = 122987254, upload-time = "2026-06-17T21:06:33.199Z" }, - { url = "https://files.pythonhosted.org/packages/e6/8c/b8087556cf81ddd808dbeb34afb8396d7ae7a1694ab489f08b1a0004e7d0/torch-2.12.1-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:2afbb2bdaa8a95040e733f05492ddf133c3967c9b7ce0abd218d704b6cab437d", size = 88303173, upload-time = "2026-06-17T21:05:06.603Z" }, - { url = "https://files.pythonhosted.org/packages/4a/07/fe09d1699fbed2afa10ebc692ff2b99d113f2605b6748cea633989e2789a/torch-2.12.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:97eba061fcb042fed191400b15568990073d67eaacaa6ee9b7ca01dd8b790fe9", size = 426404009, upload-time = "2026-06-17T21:04:57.557Z" }, - { url = "https://files.pythonhosted.org/packages/2e/f7/0ce4f6c1962c60ded7270e0a9eb560fb615c92b89d332cf9e3dff36d5ecc/torch-2.12.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:3867b861391701012adb2df93360efb88494dca245a185e3bb7624495cfe3f33", size = 532184292, upload-time = "2026-06-17T21:05:17.526Z" }, - { url = "https://files.pythonhosted.org/packages/70/db/e384c12aba30320ca92aaaf557456cbcb26f04b4df307728bb8f019f5000/torch-2.12.1-cp314-cp314t-win_amd64.whl", hash = "sha256:dd15595f8fc764cffde8c6361a3beb6ef69a028c851b1b3e70e077f615980d4e", size = 123231142, upload-time = "2026-06-17T21:05:27.061Z" }, + { url = "https://files.pythonhosted.org/packages/7f/e7/19894fdb51c7dbaf94f5a79bb0871da0992e8e4241e579cb006da46d2e58/torch-2.13.0-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:94f0de129916f77b8dc2c7a8eff644cfeddfe59e39c9f55e9f6e17543410281d", size = 111178962, upload-time = "2026-07-08T16:05:49.855Z" }, + { url = "https://files.pythonhosted.org/packages/d1/5c/b1d5de470c54e339b30a92d96683a71bcebd78f5f2a7fc714cd6dc6bbd68/torch-2.13.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:0ab4b69f3ee03a62a002cfbf77b1ca5e88aceb4ea64cb4388bb28f638ddbb045", size = 427198333, upload-time = "2026-07-08T16:05:36.847Z" }, + { url = "https://files.pythonhosted.org/packages/50/c0/68a84105e1fcb8970144b388ff3d3e5dc15a3be28c1e247841f7d7247e41/torch-2.13.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:c78b7b4d04461855a764cf01bae9a462bb88bc93defcfa11235cbc8fdf3e12c4", size = 526555154, upload-time = "2026-07-08T16:05:06.507Z" }, + { url = "https://files.pythonhosted.org/packages/2e/c9/0bb9d097b03cbaf96bb75b15e867347b8e41bfcdfe0539452d17d9e63993/torch-2.13.0-cp310-cp310-win_amd64.whl", hash = "sha256:2bd30b6b730d987fa386ce3898933762c5cb8cc82eb0535211d787cc3ce2dfeb", size = 122015602, upload-time = "2026-07-08T16:05:45.25Z" }, + { url = "https://files.pythonhosted.org/packages/5b/fe/cba54dc58523434919b66f13a667e36e436deddd77ca519e96553617d4ec/torch-2.13.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:e76f9bcecc52b8ff711239a2f7547d5353df95878ab232f0773c1d95928b92f8", size = 111187938, upload-time = "2026-07-08T16:05:17.065Z" }, + { url = "https://files.pythonhosted.org/packages/c2/59/1e3160e18e12aa3038390efab3ce02b36a9d4d6a527ecdd8520dca2e68d8/torch-2.13.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:092790c696a760c729fd5722835f50b9d81fd7c8f141571f3f3cf4081a8f664c", size = 427199369, upload-time = "2026-07-08T16:04:51.054Z" }, + { url = "https://files.pythonhosted.org/packages/01/79/1f2d34ad7034ee1c7ffc1cf8bf0f8213af2a81df6ecdb3997ecec107c09d/torch-2.13.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:60fcdcb2f3876e21146cb4524ef06397d727ca9ad5f020818547e25075fe3cb7", size = 526574961, upload-time = "2026-07-08T16:04:07.075Z" }, + { url = "https://files.pythonhosted.org/packages/6c/fd/0f2ce40f58aefbdb3392f9acce3c8171940943ae2d661f70558bfa73befb/torch-2.13.0-cp311-cp311-win_amd64.whl", hash = "sha256:a0d8b11f16a48d60e2015d8213aa0390744cbebb98e58b62b3514dddc656e330", size = 122015870, upload-time = "2026-07-08T16:05:27.59Z" }, + { url = "https://files.pythonhosted.org/packages/c4/3a/ed0f4d4d1dcde03bced7aac9a28e800abcdc0cbd06b6775044c9fbd877b7/torch-2.13.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:2fe228aba290d14b9f31b049be550dbd469c3fd3013d7a19705b30454da97027", size = 111213045, upload-time = "2026-07-08T16:05:22.997Z" }, + { url = "https://files.pythonhosted.org/packages/df/a9/f6a2a4d763ff1df02e9a64c477029db614295bc9367f4131223791ccc243/torch-2.13.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:572df8be8ffb4599c88cbd6a0726f1f854f4da65d2e3c09f0e2c2283333cd6d4", size = 427210998, upload-time = "2026-07-08T16:04:37.708Z" }, + { url = "https://files.pythonhosted.org/packages/f3/82/fea946351658e6534db52d2cc12bc53087cbf87f9440c5f180f367c1950b/torch-2.13.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:796633c4cdf0fe2cdced72d8f88f22e73dbcfce83132763162f6d4bff13b820b", size = 526605292, upload-time = "2026-07-08T16:04:22.81Z" }, + { url = "https://files.pythonhosted.org/packages/21/d6/e8f3c6f7e01f626f77259de9860d2a78bc84c40539e28e79b7e98b0bb659/torch-2.13.0-cp312-cp312-win_amd64.whl", hash = "sha256:024c6cc0c1b085f2f91f20a3dc27b0471d021c31ce84b81be3afdc39f791fd9d", size = 122057313, upload-time = "2026-07-08T16:03:53.43Z" }, + { url = "https://files.pythonhosted.org/packages/0d/fa/c1c10b7aff4a9a3e8956d4f0a5f468fa6db7abc3208805719076772b4833/torch-2.13.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:33449899ce5496c1b84b4853179d94fd102028ae1407314d9fb956bb79e70d09", size = 111213743, upload-time = "2026-07-08T16:03:28.579Z" }, + { url = "https://files.pythonhosted.org/packages/11/18/9ecb37b56293a0be8d80f810bf672a72fe7e02f8b475d5ef1b9bf8a0d748/torch-2.13.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:1e09d6a722504957c694faceca843acde562786df1144ebcc5a74075ec7f6005", size = 427213008, upload-time = "2026-07-08T16:03:44.106Z" }, + { url = "https://files.pythonhosted.org/packages/d4/5a/7c50ba1b7b713d71d34669c6d13dab0a11531a3eceb0307a5162dbfec0f7/torch-2.13.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:a3a9a21312872af8a26950b2c15680335a386a1f56ed03e780653d78b9607e9e", size = 526602329, upload-time = "2026-07-08T16:03:12.649Z" }, + { url = "https://files.pythonhosted.org/packages/91/3d/e7adcc6aaf36961cd18f56cf8ad0f3058c3a5c84ccf391762176c94581b8/torch-2.13.0-cp313-cp313-win_amd64.whl", hash = "sha256:49b58f1e2c52440abb6f17c28f0335fe6c6d01ad1a7f55b0183b81e4b34d64e6", size = 122057920, upload-time = "2026-07-08T16:03:01.808Z" }, + { url = "https://files.pythonhosted.org/packages/36/76/6dcc7f0c07052102dd36f83cbc5800842a909c8c3fbf1a7f8a5844954de9/torch-2.13.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:d849b390e07d8d333ce8ecaf91b273c656c598379a19c9acf1318a883f6b391c", size = 111227066, upload-time = "2026-07-08T16:03:33.6Z" }, + { url = "https://files.pythonhosted.org/packages/e9/09/2c10e8cd0e00fa5d23c052df6ce467eaa7182399f5e0f824f1e4ff42ccae/torch-2.13.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:a3893dc2da0a972a8ca5d698c85a9f967559ac5f8ee1797b77408aa8734d073c", size = 427226309, upload-time = "2026-07-08T16:02:53.127Z" }, + { url = "https://files.pythonhosted.org/packages/76/c6/22c2102bbef14ca6a6cb4c20e42f088e49c5f812be4e160ae57502e325f9/torch-2.13.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:49f1ea385c754e54919408a9bb3b5a72b0b755bbe2c916c1d6f70afbec4908a2", size = 526614507, upload-time = "2026-07-08T16:02:16.441Z" }, + { url = "https://files.pythonhosted.org/packages/2b/0c/7d1deb6bce5bc3e6042caf39100ac768eba3b9a098e1dddd16f75bd6489b/torch-2.13.0-cp314-cp314-win_amd64.whl", hash = "sha256:4f8573e3ce9ebcd53fe922f01077a6085ccdfbe5f12fd215883a9d87d7a744fd", size = 122051871, upload-time = "2026-07-08T16:03:23.521Z" }, + { url = "https://files.pythonhosted.org/packages/f4/ce/aa8b7f9949d32e0f2f624f342bc3b48112c1b8a130288465938bc83bcbf9/torch-2.13.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:c28def70706c2f9ecc752574766e8ae4da9b810ab6676b611166761a78a9f1e1", size = 111537025, upload-time = "2026-07-08T16:02:44.28Z" }, + { url = "https://files.pythonhosted.org/packages/69/d1/491e3a0389430946145888b0203f2b6a759ce2a61481b96a85c2da4f2ced/torch-2.13.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:31061ff56ed8fbf26c749806905aeb749ebeb819810fd5d52508aa5afd90dddc", size = 427219769, upload-time = "2026-07-08T16:02:31.18Z" }, + { url = "https://files.pythonhosted.org/packages/9a/1d/38006e045bf0a1fc28ef01e757c554e59e59a8770c284bc4f47b14e60441/torch-2.13.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:cc26eead4cf51d0b544e31e364dcf000846549c273bd148936fe9d24d29acb92", size = 526571320, upload-time = "2026-07-08T16:01:59.348Z" }, + { url = "https://files.pythonhosted.org/packages/56/94/655c91992a882bd5071aa0b5d22a07dbb130d801e872be97c0b627a7c693/torch-2.13.0-cp314-cp314t-win_amd64.whl", hash = "sha256:a7de8a313090dc5c7d7ba4bfe5c3be222528f9a4dba1acc83bddb1157360c4b8", size = 122306773, upload-time = "2026-07-08T16:02:39.832Z" }, ] [[package]] @@ -2992,30 +2777,66 @@ version = "1.17.3" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/95/8f/aeb76c5b46e273670962298c23e7ddde79916cb74db802131d49a85e4b7d/wrapt-1.17.3.tar.gz", hash = "sha256:f66eb08feaa410fe4eebd17f2a2c8e2e46d3476e9f8c783daa8e09e0faa666d0", size = 55547, upload-time = "2025-08-12T05:53:21.714Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/23/bb82321b86411eb51e5a5db3fb8f8032fd30bd7c2d74bfe936136b2fa1d6/wrapt-1.17.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:88bbae4d40d5a46142e70d58bf664a89b6b4befaea7b2ecc14e03cedb8e06c04", size = 53482, upload-time = "2025-08-12T05:51:44.467Z" }, + { url = "https://files.pythonhosted.org/packages/45/69/f3c47642b79485a30a59c63f6d739ed779fb4cc8323205d047d741d55220/wrapt-1.17.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e6b13af258d6a9ad602d57d889f83b9d5543acd471eee12eb51f5b01f8eb1bc2", size = 38676, upload-time = "2025-08-12T05:51:32.636Z" }, + { url = "https://files.pythonhosted.org/packages/d1/71/e7e7f5670c1eafd9e990438e69d8fb46fa91a50785332e06b560c869454f/wrapt-1.17.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:fd341868a4b6714a5962c1af0bd44f7c404ef78720c7de4892901e540417111c", size = 38957, upload-time = "2025-08-12T05:51:54.655Z" }, { url = "https://files.pythonhosted.org/packages/de/17/9f8f86755c191d6779d7ddead1a53c7a8aa18bccb7cea8e7e72dfa6a8a09/wrapt-1.17.3-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:f9b2601381be482f70e5d1051a5965c25fb3625455a2bf520b5a077b22afb775", size = 81975, upload-time = "2025-08-12T05:52:30.109Z" }, { url = "https://files.pythonhosted.org/packages/f2/15/dd576273491f9f43dd09fce517f6c2ce6eb4fe21681726068db0d0467096/wrapt-1.17.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:343e44b2a8e60e06a7e0d29c1671a0d9951f59174f3709962b5143f60a2a98bd", size = 83149, upload-time = "2025-08-12T05:52:09.316Z" }, { url = "https://files.pythonhosted.org/packages/0c/c4/5eb4ce0d4814521fee7aa806264bf7a114e748ad05110441cd5b8a5c744b/wrapt-1.17.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:33486899acd2d7d3066156b03465b949da3fd41a5da6e394ec49d271baefcf05", size = 82209, upload-time = "2025-08-12T05:52:10.331Z" }, { url = "https://files.pythonhosted.org/packages/31/4b/819e9e0eb5c8dc86f60dfc42aa4e2c0d6c3db8732bce93cc752e604bb5f5/wrapt-1.17.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:e6f40a8aa5a92f150bdb3e1c44b7e98fb7113955b2e5394122fa5532fec4b418", size = 81551, upload-time = "2025-08-12T05:52:31.137Z" }, + { url = "https://files.pythonhosted.org/packages/f8/83/ed6baf89ba3a56694700139698cf703aac9f0f9eb03dab92f57551bd5385/wrapt-1.17.3-cp310-cp310-win32.whl", hash = "sha256:a36692b8491d30a8c75f1dfee65bef119d6f39ea84ee04d9f9311f83c5ad9390", size = 36464, upload-time = "2025-08-12T05:53:01.204Z" }, + { url = "https://files.pythonhosted.org/packages/2f/90/ee61d36862340ad7e9d15a02529df6b948676b9a5829fd5e16640156627d/wrapt-1.17.3-cp310-cp310-win_amd64.whl", hash = "sha256:afd964fd43b10c12213574db492cb8f73b2f0826c8df07a68288f8f19af2ebe6", size = 38748, upload-time = "2025-08-12T05:53:00.209Z" }, + { url = "https://files.pythonhosted.org/packages/bd/c3/cefe0bd330d389c9983ced15d326f45373f4073c9f4a8c2f99b50bfea329/wrapt-1.17.3-cp310-cp310-win_arm64.whl", hash = "sha256:af338aa93554be859173c39c85243970dc6a289fa907402289eeae7543e1ae18", size = 36810, upload-time = "2025-08-12T05:52:51.906Z" }, + { url = "https://files.pythonhosted.org/packages/52/db/00e2a219213856074a213503fdac0511203dceefff26e1daa15250cc01a0/wrapt-1.17.3-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:273a736c4645e63ac582c60a56b0acb529ef07f78e08dc6bfadf6a46b19c0da7", size = 53482, upload-time = "2025-08-12T05:51:45.79Z" }, + { url = "https://files.pythonhosted.org/packages/5e/30/ca3c4a5eba478408572096fe9ce36e6e915994dd26a4e9e98b4f729c06d9/wrapt-1.17.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5531d911795e3f935a9c23eb1c8c03c211661a5060aab167065896bbf62a5f85", size = 38674, upload-time = "2025-08-12T05:51:34.629Z" }, + { url = "https://files.pythonhosted.org/packages/31/25/3e8cc2c46b5329c5957cec959cb76a10718e1a513309c31399a4dad07eb3/wrapt-1.17.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:0610b46293c59a3adbae3dee552b648b984176f8562ee0dba099a56cfbe4df1f", size = 38959, upload-time = "2025-08-12T05:51:56.074Z" }, { url = "https://files.pythonhosted.org/packages/5d/8f/a32a99fc03e4b37e31b57cb9cefc65050ea08147a8ce12f288616b05ef54/wrapt-1.17.3-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:b32888aad8b6e68f83a8fdccbf3165f5469702a7544472bdf41f582970ed3311", size = 82376, upload-time = "2025-08-12T05:52:32.134Z" }, { url = "https://files.pythonhosted.org/packages/31/57/4930cb8d9d70d59c27ee1332a318c20291749b4fba31f113c2f8ac49a72e/wrapt-1.17.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8cccf4f81371f257440c88faed6b74f1053eef90807b77e31ca057b2db74edb1", size = 83604, upload-time = "2025-08-12T05:52:11.663Z" }, { url = "https://files.pythonhosted.org/packages/a8/f3/1afd48de81d63dd66e01b263a6fbb86e1b5053b419b9b33d13e1f6d0f7d0/wrapt-1.17.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d8a210b158a34164de8bb68b0e7780041a903d7b00c87e906fb69928bf7890d5", size = 82782, upload-time = "2025-08-12T05:52:12.626Z" }, { url = "https://files.pythonhosted.org/packages/1e/d7/4ad5327612173b144998232f98a85bb24b60c352afb73bc48e3e0d2bdc4e/wrapt-1.17.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:79573c24a46ce11aab457b472efd8d125e5a51da2d1d24387666cd85f54c05b2", size = 82076, upload-time = "2025-08-12T05:52:33.168Z" }, + { url = "https://files.pythonhosted.org/packages/bb/59/e0adfc831674a65694f18ea6dc821f9fcb9ec82c2ce7e3d73a88ba2e8718/wrapt-1.17.3-cp311-cp311-win32.whl", hash = "sha256:c31eebe420a9a5d2887b13000b043ff6ca27c452a9a22fa71f35f118e8d4bf89", size = 36457, upload-time = "2025-08-12T05:53:03.936Z" }, + { url = "https://files.pythonhosted.org/packages/83/88/16b7231ba49861b6f75fc309b11012ede4d6b0a9c90969d9e0db8d991aeb/wrapt-1.17.3-cp311-cp311-win_amd64.whl", hash = "sha256:0b1831115c97f0663cb77aa27d381237e73ad4f721391a9bfb2fe8bc25fa6e77", size = 38745, upload-time = "2025-08-12T05:53:02.885Z" }, + { url = "https://files.pythonhosted.org/packages/9a/1e/c4d4f3398ec073012c51d1c8d87f715f56765444e1a4b11e5180577b7e6e/wrapt-1.17.3-cp311-cp311-win_arm64.whl", hash = "sha256:5a7b3c1ee8265eb4c8f1b7d29943f195c00673f5ab60c192eba2d4a7eae5f46a", size = 36806, upload-time = "2025-08-12T05:52:53.368Z" }, + { url = "https://files.pythonhosted.org/packages/9f/41/cad1aba93e752f1f9268c77270da3c469883d56e2798e7df6240dcb2287b/wrapt-1.17.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:ab232e7fdb44cdfbf55fc3afa31bcdb0d8980b9b95c38b6405df2acb672af0e0", size = 53998, upload-time = "2025-08-12T05:51:47.138Z" }, + { url = "https://files.pythonhosted.org/packages/60/f8/096a7cc13097a1869fe44efe68dace40d2a16ecb853141394047f0780b96/wrapt-1.17.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:9baa544e6acc91130e926e8c802a17f3b16fbea0fd441b5a60f5cf2cc5c3deba", size = 39020, upload-time = "2025-08-12T05:51:35.906Z" }, + { url = "https://files.pythonhosted.org/packages/33/df/bdf864b8997aab4febb96a9ae5c124f700a5abd9b5e13d2a3214ec4be705/wrapt-1.17.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6b538e31eca1a7ea4605e44f81a48aa24c4632a277431a6ed3f328835901f4fd", size = 39098, upload-time = "2025-08-12T05:51:57.474Z" }, { url = "https://files.pythonhosted.org/packages/9f/81/5d931d78d0eb732b95dc3ddaeeb71c8bb572fb01356e9133916cd729ecdd/wrapt-1.17.3-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:042ec3bb8f319c147b1301f2393bc19dba6e176b7da446853406d041c36c7828", size = 88036, upload-time = "2025-08-12T05:52:34.784Z" }, { url = "https://files.pythonhosted.org/packages/ca/38/2e1785df03b3d72d34fc6252d91d9d12dc27a5c89caef3335a1bbb8908ca/wrapt-1.17.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3af60380ba0b7b5aeb329bc4e402acd25bd877e98b3727b0135cb5c2efdaefe9", size = 88156, upload-time = "2025-08-12T05:52:13.599Z" }, { url = "https://files.pythonhosted.org/packages/b3/8b/48cdb60fe0603e34e05cffda0b2a4adab81fd43718e11111a4b0100fd7c1/wrapt-1.17.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:0b02e424deef65c9f7326d8c19220a2c9040c51dc165cddb732f16198c168396", size = 87102, upload-time = "2025-08-12T05:52:14.56Z" }, { url = "https://files.pythonhosted.org/packages/3c/51/d81abca783b58f40a154f1b2c56db1d2d9e0d04fa2d4224e357529f57a57/wrapt-1.17.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:74afa28374a3c3a11b3b5e5fca0ae03bef8450d6aa3ab3a1e2c30e3a75d023dc", size = 87732, upload-time = "2025-08-12T05:52:36.165Z" }, + { url = "https://files.pythonhosted.org/packages/9e/b1/43b286ca1392a006d5336412d41663eeef1ad57485f3e52c767376ba7e5a/wrapt-1.17.3-cp312-cp312-win32.whl", hash = "sha256:4da9f45279fff3543c371d5ababc57a0384f70be244de7759c85a7f989cb4ebe", size = 36705, upload-time = "2025-08-12T05:53:07.123Z" }, + { url = "https://files.pythonhosted.org/packages/28/de/49493f962bd3c586ab4b88066e967aa2e0703d6ef2c43aa28cb83bf7b507/wrapt-1.17.3-cp312-cp312-win_amd64.whl", hash = "sha256:e71d5c6ebac14875668a1e90baf2ea0ef5b7ac7918355850c0908ae82bcb297c", size = 38877, upload-time = "2025-08-12T05:53:05.436Z" }, + { url = "https://files.pythonhosted.org/packages/f1/48/0f7102fe9cb1e8a5a77f80d4f0956d62d97034bbe88d33e94699f99d181d/wrapt-1.17.3-cp312-cp312-win_arm64.whl", hash = "sha256:604d076c55e2fdd4c1c03d06dc1a31b95130010517b5019db15365ec4a405fc6", size = 36885, upload-time = "2025-08-12T05:52:54.367Z" }, + { url = "https://files.pythonhosted.org/packages/fc/f6/759ece88472157acb55fc195e5b116e06730f1b651b5b314c66291729193/wrapt-1.17.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:a47681378a0439215912ef542c45a783484d4dd82bac412b71e59cf9c0e1cea0", size = 54003, upload-time = "2025-08-12T05:51:48.627Z" }, + { url = "https://files.pythonhosted.org/packages/4f/a9/49940b9dc6d47027dc850c116d79b4155f15c08547d04db0f07121499347/wrapt-1.17.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:54a30837587c6ee3cd1a4d1c2ec5d24e77984d44e2f34547e2323ddb4e22eb77", size = 39025, upload-time = "2025-08-12T05:51:37.156Z" }, + { url = "https://files.pythonhosted.org/packages/45/35/6a08de0f2c96dcdd7fe464d7420ddb9a7655a6561150e5fc4da9356aeaab/wrapt-1.17.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:16ecf15d6af39246fe33e507105d67e4b81d8f8d2c6598ff7e3ca1b8a37213f7", size = 39108, upload-time = "2025-08-12T05:51:58.425Z" }, { url = "https://files.pythonhosted.org/packages/0c/37/6faf15cfa41bf1f3dba80cd3f5ccc6622dfccb660ab26ed79f0178c7497f/wrapt-1.17.3-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:6fd1ad24dc235e4ab88cda009e19bf347aabb975e44fd5c2fb22a3f6e4141277", size = 88072, upload-time = "2025-08-12T05:52:37.53Z" }, { url = "https://files.pythonhosted.org/packages/78/f2/efe19ada4a38e4e15b6dff39c3e3f3f73f5decf901f66e6f72fe79623a06/wrapt-1.17.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ed61b7c2d49cee3c027372df5809a59d60cf1b6c2f81ee980a091f3afed6a2d", size = 88214, upload-time = "2025-08-12T05:52:15.886Z" }, { url = "https://files.pythonhosted.org/packages/40/90/ca86701e9de1622b16e09689fc24b76f69b06bb0150990f6f4e8b0eeb576/wrapt-1.17.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:423ed5420ad5f5529db9ce89eac09c8a2f97da18eb1c870237e84c5a5c2d60aa", size = 87105, upload-time = "2025-08-12T05:52:17.914Z" }, { url = "https://files.pythonhosted.org/packages/fd/e0/d10bd257c9a3e15cbf5523025252cc14d77468e8ed644aafb2d6f54cb95d/wrapt-1.17.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e01375f275f010fcbf7f643b4279896d04e571889b8a5b3f848423d91bf07050", size = 87766, upload-time = "2025-08-12T05:52:39.243Z" }, + { url = "https://files.pythonhosted.org/packages/e8/cf/7d848740203c7b4b27eb55dbfede11aca974a51c3d894f6cc4b865f42f58/wrapt-1.17.3-cp313-cp313-win32.whl", hash = "sha256:53e5e39ff71b3fc484df8a522c933ea2b7cdd0d5d15ae82e5b23fde87d44cbd8", size = 36711, upload-time = "2025-08-12T05:53:10.074Z" }, + { url = "https://files.pythonhosted.org/packages/57/54/35a84d0a4d23ea675994104e667ceff49227ce473ba6a59ba2c84f250b74/wrapt-1.17.3-cp313-cp313-win_amd64.whl", hash = "sha256:1f0b2f40cf341ee8cc1a97d51ff50dddb9fcc73241b9143ec74b30fc4f44f6cb", size = 38885, upload-time = "2025-08-12T05:53:08.695Z" }, + { url = "https://files.pythonhosted.org/packages/01/77/66e54407c59d7b02a3c4e0af3783168fff8e5d61def52cda8728439d86bc/wrapt-1.17.3-cp313-cp313-win_arm64.whl", hash = "sha256:7425ac3c54430f5fc5e7b6f41d41e704db073309acfc09305816bc6a0b26bb16", size = 36896, upload-time = "2025-08-12T05:52:55.34Z" }, + { url = "https://files.pythonhosted.org/packages/02/a2/cd864b2a14f20d14f4c496fab97802001560f9f41554eef6df201cd7f76c/wrapt-1.17.3-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:cf30f6e3c077c8e6a9a7809c94551203c8843e74ba0c960f4a98cd80d4665d39", size = 54132, upload-time = "2025-08-12T05:51:49.864Z" }, + { url = "https://files.pythonhosted.org/packages/d5/46/d011725b0c89e853dc44cceb738a307cde5d240d023d6d40a82d1b4e1182/wrapt-1.17.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:e228514a06843cae89621384cfe3a80418f3c04aadf8a3b14e46a7be704e4235", size = 39091, upload-time = "2025-08-12T05:51:38.935Z" }, + { url = "https://files.pythonhosted.org/packages/2e/9e/3ad852d77c35aae7ddebdbc3b6d35ec8013af7d7dddad0ad911f3d891dae/wrapt-1.17.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:5ea5eb3c0c071862997d6f3e02af1d055f381b1d25b286b9d6644b79db77657c", size = 39172, upload-time = "2025-08-12T05:51:59.365Z" }, { url = "https://files.pythonhosted.org/packages/c3/f7/c983d2762bcce2326c317c26a6a1e7016f7eb039c27cdf5c4e30f4160f31/wrapt-1.17.3-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:281262213373b6d5e4bb4353bc36d1ba4084e6d6b5d242863721ef2bf2c2930b", size = 87163, upload-time = "2025-08-12T05:52:40.965Z" }, { url = "https://files.pythonhosted.org/packages/e4/0f/f673f75d489c7f22d17fe0193e84b41540d962f75fce579cf6873167c29b/wrapt-1.17.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:dc4a8d2b25efb6681ecacad42fca8859f88092d8732b170de6a5dddd80a1c8fa", size = 87963, upload-time = "2025-08-12T05:52:20.326Z" }, { url = "https://files.pythonhosted.org/packages/df/61/515ad6caca68995da2fac7a6af97faab8f78ebe3bf4f761e1b77efbc47b5/wrapt-1.17.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:373342dd05b1d07d752cecbec0c41817231f29f3a89aa8b8843f7b95992ed0c7", size = 86945, upload-time = "2025-08-12T05:52:21.581Z" }, { url = "https://files.pythonhosted.org/packages/d3/bd/4e70162ce398462a467bc09e768bee112f1412e563620adc353de9055d33/wrapt-1.17.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d40770d7c0fd5cbed9d84b2c3f2e156431a12c9a37dc6284060fb4bec0b7ffd4", size = 86857, upload-time = "2025-08-12T05:52:43.043Z" }, + { url = "https://files.pythonhosted.org/packages/2b/b8/da8560695e9284810b8d3df8a19396a6e40e7518059584a1a394a2b35e0a/wrapt-1.17.3-cp314-cp314-win32.whl", hash = "sha256:fbd3c8319de8e1dc79d346929cd71d523622da527cca14e0c1d257e31c2b8b10", size = 37178, upload-time = "2025-08-12T05:53:12.605Z" }, + { url = "https://files.pythonhosted.org/packages/db/c8/b71eeb192c440d67a5a0449aaee2310a1a1e8eca41676046f99ed2487e9f/wrapt-1.17.3-cp314-cp314-win_amd64.whl", hash = "sha256:e1a4120ae5705f673727d3253de3ed0e016f7cd78dc463db1b31e2463e1f3cf6", size = 39310, upload-time = "2025-08-12T05:53:11.106Z" }, + { url = "https://files.pythonhosted.org/packages/45/20/2cda20fd4865fa40f86f6c46ed37a2a8356a7a2fde0773269311f2af56c7/wrapt-1.17.3-cp314-cp314-win_arm64.whl", hash = "sha256:507553480670cab08a800b9463bdb881b2edeed77dc677b0a5915e6106e91a58", size = 37266, upload-time = "2025-08-12T05:52:56.531Z" }, + { url = "https://files.pythonhosted.org/packages/77/ed/dd5cf21aec36c80443c6f900449260b80e2a65cf963668eaef3b9accce36/wrapt-1.17.3-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:ed7c635ae45cfbc1a7371f708727bf74690daedc49b4dba310590ca0bd28aa8a", size = 56544, upload-time = "2025-08-12T05:51:51.109Z" }, + { url = "https://files.pythonhosted.org/packages/8d/96/450c651cc753877ad100c7949ab4d2e2ecc4d97157e00fa8f45df682456a/wrapt-1.17.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:249f88ed15503f6492a71f01442abddd73856a0032ae860de6d75ca62eed8067", size = 40283, upload-time = "2025-08-12T05:51:39.912Z" }, + { url = "https://files.pythonhosted.org/packages/d1/86/2fcad95994d9b572db57632acb6f900695a648c3e063f2cd344b3f5c5a37/wrapt-1.17.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5a03a38adec8066d5a37bea22f2ba6bbf39fcdefbe2d91419ab864c3fb515454", size = 40366, upload-time = "2025-08-12T05:52:00.693Z" }, { url = "https://files.pythonhosted.org/packages/64/0e/f4472f2fdde2d4617975144311f8800ef73677a159be7fe61fa50997d6c0/wrapt-1.17.3-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:5d4478d72eb61c36e5b446e375bbc49ed002430d17cdec3cecb36993398e1a9e", size = 108571, upload-time = "2025-08-12T05:52:44.521Z" }, { url = "https://files.pythonhosted.org/packages/cc/01/9b85a99996b0a97c8a17484684f206cbb6ba73c1ce6890ac668bcf3838fb/wrapt-1.17.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:223db574bb38637e8230eb14b185565023ab624474df94d2af18f1cdb625216f", size = 113094, upload-time = "2025-08-12T05:52:22.618Z" }, { url = "https://files.pythonhosted.org/packages/25/02/78926c1efddcc7b3aa0bc3d6b33a822f7d898059f7cd9ace8c8318e559ef/wrapt-1.17.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e405adefb53a435f01efa7ccdec012c016b5a1d3f35459990afc39b6be4d5056", size = 110659, upload-time = "2025-08-12T05:52:24.057Z" }, { url = "https://files.pythonhosted.org/packages/dc/ee/c414501ad518ac3e6fe184753632fe5e5ecacdcf0effc23f31c1e4f7bfcf/wrapt-1.17.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:88547535b787a6c9ce4086917b6e1d291aa8ed914fdd3a838b3539dc95c12804", size = 106946, upload-time = "2025-08-12T05:52:45.976Z" }, + { url = "https://files.pythonhosted.org/packages/be/44/a1bd64b723d13bb151d6cc91b986146a1952385e0392a78567e12149c7b4/wrapt-1.17.3-cp314-cp314t-win32.whl", hash = "sha256:41b1d2bc74c2cac6f9074df52b2efbef2b30bdfe5f40cb78f8ca22963bc62977", size = 38717, upload-time = "2025-08-12T05:53:15.214Z" }, + { url = "https://files.pythonhosted.org/packages/79/d9/7cfd5a312760ac4dd8bf0184a6ee9e43c33e47f3dadc303032ce012b8fa3/wrapt-1.17.3-cp314-cp314t-win_amd64.whl", hash = "sha256:73d496de46cd2cdbdbcce4ae4bcdb4afb6a11234a1df9c085249d55166b95116", size = 41334, upload-time = "2025-08-12T05:53:14.178Z" }, + { url = "https://files.pythonhosted.org/packages/46/78/10ad9781128ed2f99dbc474f43283b13fea8ba58723e98844367531c18e9/wrapt-1.17.3-cp314-cp314t-win_arm64.whl", hash = "sha256:f38e60678850c42461d4202739f9bf1e3a737c7ad283638251e79cc49effb6b6", size = 38471, upload-time = "2025-08-12T05:52:57.784Z" }, { url = "https://files.pythonhosted.org/packages/1f/f6/a933bd70f98e9cf3e08167fc5cd7aaaca49147e48411c0bd5ae701bb2194/wrapt-1.17.3-py3-none-any.whl", hash = "sha256:7171ae35d2c33d326ac19dd8facb1e82e5fd04ef8c6c0e394d7af55a55051c22", size = 23591, upload-time = "2025-08-12T05:53:20.674Z" }, ] @@ -3235,3 +3056,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/94/c3/b2e9f38bc3e11191981d57ea08cab2166e74ea770024a646617c9cddd9f6/yarl-1.20.1-cp313-cp313t-win_amd64.whl", hash = "sha256:541d050a355bbbc27e55d906bc91cb6fe42f96c01413dd0f4ed5a5240513874f", size = 93003, upload-time = "2025-06-10T00:45:27.752Z" }, { url = "https://files.pythonhosted.org/packages/b4/2d/2345fce04cfd4bee161bf1e7d9cdc702e3e16109021035dbb24db654a622/yarl-1.20.1-py3-none-any.whl", hash = "sha256:83b8eb083fe4683c6115795d9fc1cfaf2cbbefb19b3a1cb68f6527460f483a77", size = 46542, upload-time = "2025-06-10T00:46:07.521Z" }, ] + +[[package]] +name = "zipp" +version = "4.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b9/d8/eab98a517c14134c0b2eb4e2387bc5f457334293ec5d2dd3857ec2966802/zipp-4.1.0.tar.gz", hash = "sha256:4cb57381f544315db7688e976e922a2b18cdb513d21cc194eb42232ba2a3e602", size = 26214, upload-time = "2026-05-18T20:08:57.967Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/13/547360d81e6d88d58492968ffda9f9542854f11310ee556fef14260cc886/zipp-4.1.0-py3-none-any.whl", hash = "sha256:25ad4e16390cd314347dd8f1de67a2ac538ae658ed4ab9db16029c07c188e97f", size = 10238, upload-time = "2026-05-18T20:08:57.045Z" }, +] diff --git a/rust-toolchain.toml b/rust-toolchain.toml index 5699bd4d536..c4058187023 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,5 +1,5 @@ # We keep this pinned to keep clippy and rustfmt in sync between local and CI. # Feel free to upgrade to bring in new lints. [toolchain] -channel = "1.94.0" +channel = "1.97.0" components = ["rustfmt", "clippy", "rust-analyzer"] diff --git a/rust/AGENTS.md b/rust/AGENTS.md index 6b2729c6692..70a803c6c76 100644 --- a/rust/AGENTS.md +++ b/rust/AGENTS.md @@ -17,6 +17,10 @@ Also see [root AGENTS.md](../AGENTS.md) for cross-language standards. - Delete obsolete internal (`pub(crate)` / private) methods in the same PR that introduces their replacements. For public API methods, follow the deprecation path in root AGENTS.md instead. - Choose log levels by audience: `debug!` for routine/high-frequency ops, `info!` for infrequent operator-visible state changes, `warn!` for unexpected conditions. +## Concurrency + +- The closure passed to `spawn_cpu()` must only consume CPU and return — it must **never** wait on anything: **no channels** (blocking send/recv), **no I/O**, **no locks**, and no `block_on`/`.blocking_*`. The CPU pool can collapse to a single worker in resource-constrained environments (`<= 3` CPUs), so a parked closure can deadlock the whole pool with a silent 0% hang. Keep the waiting in surrounding async code and hand only the pure-CPU work to `spawn_cpu()`. Only dispatch substantial work (rule of thumb: ~100µs+ of CPU); below that the pool overhead outweighs the benefit and the work is better left inline. See the doc comment on `spawn_cpu` for the rationale. + ## API Design - Use `with_`-prefixed builder methods for optional config (e.g., `MyStruct::new(required).with_option(v)`) — don't create separate constructor variants. diff --git a/rust/compression/bitpacking/Cargo.toml b/rust/compression/bitpacking/Cargo.toml index f086ede8baa..56db321ac6d 100644 --- a/rust/compression/bitpacking/Cargo.toml +++ b/rust/compression/bitpacking/Cargo.toml @@ -13,8 +13,12 @@ rust-version.workspace = true [dependencies] arrayref = "0.3" +crunchy = "0.2" paste = "1" seq-macro = "0.3" +[dev-dependencies] +bitpacking = { workspace = true, features = ["bitpacker4x"] } + [lints] workspace = true diff --git a/rust/compression/bitpacking/src/bitpacker_internal/bitpacker4x.rs b/rust/compression/bitpacking/src/bitpacker_internal/bitpacker4x.rs new file mode 100644 index 00000000000..ca3c68908ef --- /dev/null +++ b/rust/compression/bitpacking/src/bitpacker_internal/bitpacker4x.rs @@ -0,0 +1,820 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use super::BitPacker; + +use crate::bitpacker_internal::{Available, UnsafeBitPacker}; + +const BLOCK_LEN: usize = 32 * 4; + +#[cfg(target_arch = "x86_64")] +mod sse3 { + + use super::BLOCK_LEN; + use crate::bitpacker_internal::Available; + + use std::arch::x86_64::__m128i as DataType; + use std::arch::x86_64::_mm_and_si128 as op_and; + use std::arch::x86_64::_mm_lddqu_si128 as load_unaligned; + use std::arch::x86_64::_mm_or_si128 as op_or; + use std::arch::x86_64::_mm_set1_epi32 as set1; + use std::arch::x86_64::_mm_slli_epi32 as left_shift_32; + use std::arch::x86_64::_mm_srli_epi32 as right_shift_32; + use std::arch::x86_64::_mm_storeu_si128 as store_unaligned; + use std::arch::x86_64::{ + _mm_add_epi32, _mm_cvtsi128_si32, _mm_shuffle_epi32, _mm_slli_si128, _mm_srli_si128, + _mm_sub_epi32, + }; + + #[allow(non_snake_case)] + #[inline] + unsafe fn or_collapse_to_u32(accumulator: DataType) -> u32 { + let a__b__c__d_ = accumulator; + let ______a__b_ = _mm_srli_si128(a__b__c__d_, 8); + let a__b__ca_db = op_or(a__b__c__d_, ______a__b_); + let ___a__b__ca = _mm_srli_si128(a__b__ca_db, 4); + let _______cadb = op_or(a__b__ca_db, ___a__b__ca); + _mm_cvtsi128_si32(_______cadb) as u32 + } + + #[target_feature(enable = "sse3")] + unsafe fn compute_delta(curr: DataType, prev: DataType) -> DataType { + _mm_sub_epi32( + curr, + op_or(_mm_slli_si128(curr, 4), _mm_srli_si128(prev, 12)), + ) + } + + #[target_feature(enable = "sse3")] + #[allow(non_snake_case)] + #[inline] + unsafe fn integrate_delta(prev: DataType, delta: DataType) -> DataType { + let offset = _mm_shuffle_epi32(prev, 0xff); + let a__b__c__d_ = delta; + let ______a__b_ = _mm_slli_si128(delta, 8); + let a__b__ca_db = _mm_add_epi32(______a__b_, a__b__c__d_); + let ___a__b__ca = _mm_slli_si128(a__b__ca_db, 4); + let a_ab_abc_abcd: DataType = _mm_add_epi32(___a__b__ca, a__b__ca_db); + _mm_add_epi32(offset, a_ab_abc_abcd) + } + + #[target_feature(enable = "sse3")] + #[inline] + unsafe fn add(left: DataType, right: DataType) -> DataType { + _mm_add_epi32(left, right) + } + + unsafe fn sub(left: DataType, right: DataType) -> DataType { + _mm_sub_epi32(left, right) + } + + declare_bitpacker!(target_feature(enable = "sse3")); + + impl Available for UnsafeBitPackerImpl { + fn available() -> bool { + is_x86_feature_detected!("sse3") + } + } +} + +#[cfg(all(target_arch = "aarch64", target_endian = "little"))] +mod neon { + + use super::BLOCK_LEN; + use crate::bitpacker_internal::Available; + use std::arch::aarch64::{ + uint32x4_t, vaddq_u32, vandq_u32, vdupq_n_u32, vextq_u32, vgetq_lane_u32, vld1q_u32, + vorrq_u32, vshlq_n_u32, vshrq_n_u32, vst1q_u32, vsubq_u32, + }; + + pub(crate) type DataType = uint32x4_t; + + #[inline] + /// Creates a vector with all elements set to `el`. + unsafe fn set1(el: i32) -> DataType { + vdupq_n_u32(el as u32) + } + + #[inline] + unsafe fn right_shift_32(el: DataType) -> DataType { + const { + assert!(N >= 0); + assert!(N <= 32); + } + + // We unroll here because vshrq_n_u32 only accepts constants from 1 to 32. + match N { + 0 => el, + 1 => vshrq_n_u32::<1>(el), + 2 => vshrq_n_u32::<2>(el), + 3 => vshrq_n_u32::<3>(el), + 4 => vshrq_n_u32::<4>(el), + 5 => vshrq_n_u32::<5>(el), + 6 => vshrq_n_u32::<6>(el), + 7 => vshrq_n_u32::<7>(el), + 8 => vshrq_n_u32::<8>(el), + 9 => vshrq_n_u32::<9>(el), + 10 => vshrq_n_u32::<10>(el), + 11 => vshrq_n_u32::<11>(el), + 12 => vshrq_n_u32::<12>(el), + 13 => vshrq_n_u32::<13>(el), + 14 => vshrq_n_u32::<14>(el), + 15 => vshrq_n_u32::<15>(el), + 16 => vshrq_n_u32::<16>(el), + 17 => vshrq_n_u32::<17>(el), + 18 => vshrq_n_u32::<18>(el), + 19 => vshrq_n_u32::<19>(el), + 20 => vshrq_n_u32::<20>(el), + 21 => vshrq_n_u32::<21>(el), + 22 => vshrq_n_u32::<22>(el), + 23 => vshrq_n_u32::<23>(el), + 24 => vshrq_n_u32::<24>(el), + 25 => vshrq_n_u32::<25>(el), + 26 => vshrq_n_u32::<26>(el), + 27 => vshrq_n_u32::<27>(el), + 28 => vshrq_n_u32::<28>(el), + 29 => vshrq_n_u32::<29>(el), + 30 => vshrq_n_u32::<30>(el), + 31 => vshrq_n_u32::<31>(el), + 32 => vdupq_n_u32(0), + _ => core::hint::unreachable_unchecked(), + } + } + + #[inline] + unsafe fn left_shift_32(el: DataType) -> DataType { + const { + assert!(N >= 0); + assert!(N <= 32); + } + + // We unroll here because vshlq_n_u32 only accepts constants from 0 to 31. + match N { + 0 => el, + 1 => vshlq_n_u32::<1>(el), + 2 => vshlq_n_u32::<2>(el), + 3 => vshlq_n_u32::<3>(el), + 4 => vshlq_n_u32::<4>(el), + 5 => vshlq_n_u32::<5>(el), + 6 => vshlq_n_u32::<6>(el), + 7 => vshlq_n_u32::<7>(el), + 8 => vshlq_n_u32::<8>(el), + 9 => vshlq_n_u32::<9>(el), + 10 => vshlq_n_u32::<10>(el), + 11 => vshlq_n_u32::<11>(el), + 12 => vshlq_n_u32::<12>(el), + 13 => vshlq_n_u32::<13>(el), + 14 => vshlq_n_u32::<14>(el), + 15 => vshlq_n_u32::<15>(el), + 16 => vshlq_n_u32::<16>(el), + 17 => vshlq_n_u32::<17>(el), + 18 => vshlq_n_u32::<18>(el), + 19 => vshlq_n_u32::<19>(el), + 20 => vshlq_n_u32::<20>(el), + 21 => vshlq_n_u32::<21>(el), + 22 => vshlq_n_u32::<22>(el), + 23 => vshlq_n_u32::<23>(el), + 24 => vshlq_n_u32::<24>(el), + 25 => vshlq_n_u32::<25>(el), + 26 => vshlq_n_u32::<26>(el), + 27 => vshlq_n_u32::<27>(el), + 28 => vshlq_n_u32::<28>(el), + 29 => vshlq_n_u32::<29>(el), + 30 => vshlq_n_u32::<30>(el), + 31 => vshlq_n_u32::<31>(el), + 32 => vdupq_n_u32(0), + _ => core::hint::unreachable_unchecked(), + } + } + + use vorrq_u32 as op_or; + + #[inline] + unsafe fn op_and(left: DataType, right: DataType) -> DataType { + vandq_u32(left, right) + } + + #[inline] + unsafe fn load_unaligned(addr: *const DataType) -> DataType { + vld1q_u32(addr.cast::()) + } + + #[inline] + unsafe fn store_unaligned(addr: *mut DataType, data: DataType) { + vst1q_u32(addr.cast::(), data); + } + + #[inline] + /// Collapses the vector by performing a bitwise OR across all lanes + unsafe fn or_collapse_to_u32(acc: DataType) -> u32 { + vgetq_lane_u32(acc, 0) + | vgetq_lane_u32(acc, 1) + | vgetq_lane_u32(acc, 2) + | vgetq_lane_u32(acc, 3) + } + + unsafe fn compute_delta(curr: DataType, prev: DataType) -> DataType { + // Build a vector with [prev[3], curr[0], curr[1], curr[2]] + let prev_shifted = vextq_u32(prev, curr, 3); + vsubq_u32(curr, prev_shifted) + } + + #[allow(non_snake_case)] + #[inline] + unsafe fn integrate_delta(prev: DataType, delta: DataType) -> DataType { + let base = vdupq_n_u32(vgetq_lane_u32(prev, 3)); + let zero = vdupq_n_u32(0); + let a__b__c__d_ = delta; + let ______a__b_ = vextq_u32(zero, a__b__c__d_, 2); + let a__b__ca_db = vaddq_u32(______a__b_, a__b__c__d_); + let ___a__b__ca = vextq_u32(zero, a__b__ca_db, 3); + let a_ab_abc_abcd = vaddq_u32(___a__b__ca, a__b__ca_db); + vaddq_u32(base, a_ab_abc_abcd) + } + + #[inline] + unsafe fn add(left: DataType, right: DataType) -> DataType { + vaddq_u32(left, right) + } + + #[inline] + unsafe fn sub(left: DataType, right: DataType) -> DataType { + vsubq_u32(left, right) + } + + declare_bitpacker!(target_feature(enable = "neon")); + + impl Available for UnsafeBitPackerImpl { + fn available() -> bool { + std::arch::is_aarch64_feature_detected!("neon") + } + } +} + +mod scalar { + + use super::BLOCK_LEN; + use crate::bitpacker_internal::Available; + use std::ptr; + + pub(crate) type DataType = [u32; 4]; + + pub(crate) fn set1(el: i32) -> DataType { + [el as u32; 4] + } + + pub(crate) fn right_shift_32(el: DataType) -> DataType { + [el[0] >> N, el[1] >> N, el[2] >> N, el[3] >> N] + } + + pub(crate) fn left_shift_32(el: DataType) -> DataType { + [el[0] << N, el[1] << N, el[2] << N, el[3] << N] + } + + pub(crate) fn op_or(left: DataType, right: DataType) -> DataType { + [ + left[0] | right[0], + left[1] | right[1], + left[2] | right[2], + left[3] | right[3], + ] + } + + pub(crate) fn op_and(left: DataType, right: DataType) -> DataType { + [ + left[0] & right[0], + left[1] & right[1], + left[2] & right[2], + left[3] & right[3], + ] + } + + pub(crate) unsafe fn load_unaligned(addr: *const DataType) -> DataType { + ptr::read_unaligned(addr) + } + + pub(crate) unsafe fn store_unaligned(addr: *mut DataType, data: DataType) { + ptr::write_unaligned(addr, data); + } + + pub(crate) fn or_collapse_to_u32(accumulator: DataType) -> u32 { + (accumulator[0] | accumulator[1]) | (accumulator[2] | accumulator[3]) + } + + fn compute_delta(curr: DataType, prev: DataType) -> DataType { + [ + curr[0].wrapping_sub(prev[3]), + curr[1].wrapping_sub(curr[0]), + curr[2].wrapping_sub(curr[1]), + curr[3].wrapping_sub(curr[2]), + ] + } + + fn integrate_delta(offset: DataType, delta: DataType) -> DataType { + let el0 = offset[3].wrapping_add(delta[0]); + let el1 = el0.wrapping_add(delta[1]); + let el2 = el1.wrapping_add(delta[2]); + let el3 = el2.wrapping_add(delta[3]); + [el0, el1, el2, el3] + } + + pub(crate) fn add(left: DataType, right: DataType) -> DataType { + [ + left[0].wrapping_add(right[0]), + left[1].wrapping_add(right[1]), + left[2].wrapping_add(right[2]), + left[3].wrapping_add(right[3]), + ] + } + + pub(crate) fn sub(left: DataType, right: DataType) -> DataType { + [ + left[0].wrapping_sub(right[0]), + left[1].wrapping_sub(right[1]), + left[2].wrapping_sub(right[2]), + left[3].wrapping_sub(right[3]), + ] + } + + // The `allow(unused)` is here to put an attribute that has no effect. + // + // For other bitpacker, we enable specific CPU instruction set, but for the + // scalar bitpacker none is required. + declare_bitpacker!(allow(unused)); + + impl Available for UnsafeBitPackerImpl { + fn available() -> bool { + true + } + } +} + +#[derive(Clone, Copy)] +enum InstructionSet { + #[cfg(target_arch = "x86_64")] + SSE3, + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + NEON, + Scalar, +} + +/// `BitPacker4x` packs integers in groups of 4. This gives an opportunity +/// to leverage `SSE3` instructions to encode and decode the stream. +/// +/// One block must contain `128 integers`. +#[derive(Clone, Copy)] +pub struct BitPacker4x(InstructionSet); + +impl BitPacker4x { + #[cfg(target_arch = "x86_64")] + pub(crate) fn new_sse() -> Option { + sse3::UnsafeBitPackerImpl::available().then_some(BitPacker4x(InstructionSet::SSE3)) + } + + #[cfg(not(target_arch = "x86_64"))] + pub(crate) fn new_sse() -> Option { + None + } + + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + pub(crate) fn new_neon() -> Option { + neon::UnsafeBitPackerImpl::available().then_some(BitPacker4x(InstructionSet::NEON)) + } + + #[cfg(not(all(target_arch = "aarch64", target_endian = "little")))] + pub(crate) fn new_neon() -> Option { + None + } + + pub(crate) fn new_scalar() -> Self { + BitPacker4x(InstructionSet::Scalar) + } +} + +impl BitPacker for BitPacker4x { + const BLOCK_LEN: usize = BLOCK_LEN; + + fn new() -> Self { + Self::new_sse() + .or_else(Self::new_neon) + .unwrap_or_else(Self::new_scalar) + } + + fn compress(&self, decompressed: &[u32], compressed: &mut [u8], num_bits: u8) -> usize { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::SSE3 => { + sse3::UnsafeBitPackerImpl::compress(decompressed, compressed, num_bits) + } + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => { + neon::UnsafeBitPackerImpl::compress(decompressed, compressed, num_bits) + } + InstructionSet::Scalar => { + scalar::UnsafeBitPackerImpl::compress(decompressed, compressed, num_bits) + } + } + } + } + + fn compress_sorted( + &self, + initial: u32, + decompressed: &[u32], + compressed: &mut [u8], + num_bits: u8, + ) -> usize { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::SSE3 => sse3::UnsafeBitPackerImpl::compress_sorted( + initial, + decompressed, + compressed, + num_bits, + ), + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => neon::UnsafeBitPackerImpl::compress_sorted( + initial, + decompressed, + compressed, + num_bits, + ), + InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::compress_sorted( + initial, + decompressed, + compressed, + num_bits, + ), + } + } + } + + fn compress_strictly_sorted( + &self, + initial: Option, + decompressed: &[u32], + compressed: &mut [u8], + num_bits: u8, + ) -> usize { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::SSE3 => sse3::UnsafeBitPackerImpl::compress_strictly_sorted( + initial, + decompressed, + compressed, + num_bits, + ), + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => neon::UnsafeBitPackerImpl::compress_strictly_sorted( + initial, + decompressed, + compressed, + num_bits, + ), + InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::compress_strictly_sorted( + initial, + decompressed, + compressed, + num_bits, + ), + } + } + } + + fn decompress(&self, compressed: &[u8], decompressed: &mut [u32], num_bits: u8) -> usize { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::SSE3 => { + sse3::UnsafeBitPackerImpl::decompress(compressed, decompressed, num_bits) + } + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => { + neon::UnsafeBitPackerImpl::decompress(compressed, decompressed, num_bits) + } + InstructionSet::Scalar => { + scalar::UnsafeBitPackerImpl::decompress(compressed, decompressed, num_bits) + } + } + } + } + + fn decompress_strictly_sorted( + &self, + initial: Option, + compressed: &[u8], + decompressed: &mut [u32], + num_bits: u8, + ) -> usize { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::SSE3 => sse3::UnsafeBitPackerImpl::decompress_strictly_sorted( + initial, + compressed, + decompressed, + num_bits, + ), + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => neon::UnsafeBitPackerImpl::decompress_strictly_sorted( + initial, + compressed, + decompressed, + num_bits, + ), + InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::decompress_strictly_sorted( + initial, + compressed, + decompressed, + num_bits, + ), + } + } + } + + fn decompress_sorted( + &self, + initial: u32, + compressed: &[u8], + decompressed: &mut [u32], + num_bits: u8, + ) -> usize { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::SSE3 => sse3::UnsafeBitPackerImpl::decompress_sorted( + initial, + compressed, + decompressed, + num_bits, + ), + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => neon::UnsafeBitPackerImpl::decompress_sorted( + initial, + compressed, + decompressed, + num_bits, + ), + InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::decompress_sorted( + initial, + compressed, + decompressed, + num_bits, + ), + } + } + } + + fn num_bits(&self, decompressed: &[u32]) -> u8 { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::SSE3 => sse3::UnsafeBitPackerImpl::num_bits(decompressed), + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => neon::UnsafeBitPackerImpl::num_bits(decompressed), + InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::num_bits(decompressed), + } + } + } + + fn num_bits_sorted(&self, initial: u32, decompressed: &[u32]) -> u8 { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::SSE3 => { + sse3::UnsafeBitPackerImpl::num_bits_sorted(initial, decompressed) + } + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => { + neon::UnsafeBitPackerImpl::num_bits_sorted(initial, decompressed) + } + InstructionSet::Scalar => { + scalar::UnsafeBitPackerImpl::num_bits_sorted(initial, decompressed) + } + } + } + } + + fn num_bits_strictly_sorted(&self, initial: Option, decompressed: &[u32]) -> u8 { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::SSE3 => { + sse3::UnsafeBitPackerImpl::num_bits_strictly_sorted(initial, decompressed) + } + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => { + neon::UnsafeBitPackerImpl::num_bits_strictly_sorted(initial, decompressed) + } + InstructionSet::Scalar => { + scalar::UnsafeBitPackerImpl::num_bits_strictly_sorted(initial, decompressed) + } + } + } + } +} + +#[cfg(any( + target_arch = "x86_64", + all(target_arch = "aarch64", target_endian = "little") +))] +#[cfg(any())] +mod tests { + use super::BLOCK_LEN; + use super::scalar; + use crate::bitpacker_internal::Available; + use crate::tests::test_util_compatible; + use crate::{BitPacker, BitPacker4x}; + + #[cfg(target_arch = "x86_64")] + #[test] + fn test_compatible_sse3() { + use super::sse3; + if sse3::UnsafeBitPackerImpl::available() { + test_util_compatible::( + BLOCK_LEN, + ); + } + } + + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[test] + fn test_compatible_neon() { + use super::neon; + if neon::UnsafeBitPackerImpl::available() { + test_util_compatible::( + BLOCK_LEN, + ); + } + } + + #[test] + fn test_delta_bit_width_32() { + let values = vec![i32::max_value() as u32 + 1; BitPacker4x::BLOCK_LEN]; + let bit_packer = BitPacker4x::new(); + let bit_width = bit_packer.num_bits_sorted(0, &values); + assert_eq!(bit_width, 32); + + let mut block = vec![0u8; BitPacker4x::compressed_block_size(bit_width)]; + bit_packer.compress_sorted(0, &values, &mut block, bit_width); + + let mut decoded_values = vec![0x10101010; BitPacker4x::BLOCK_LEN]; + bit_packer.decompress_sorted(0, &block, &mut decoded_values, bit_width); + + assert_eq!(values, decoded_values); + } + + #[test] + fn test_bit_width_32() { + let mut values = vec![i32::max_value() as u32 + 1; BitPacker4x::BLOCK_LEN]; + values[0] = 0; + let bit_packer = BitPacker4x::new(); + let bit_width = bit_packer.num_bits(&values); + assert_eq!(bit_width, 32); + + let mut block = vec![0u8; BitPacker4x::compressed_block_size(bit_width)]; + bit_packer.compress(&values, &mut block, bit_width); + + let mut decoded_values = vec![0x10101010; BitPacker4x::BLOCK_LEN]; + bit_packer.decompress(&block, &mut decoded_values, bit_width); + + assert_eq!(values, decoded_values); + } +} + +#[cfg(test)] +mod tests { + use super::{BLOCK_LEN, BitPacker4x}; + use crate::bitpacker_internal::BitPacker; + use bitpacking::{BitPacker as ExternalBitPacker, BitPacker4x as ExternalBitPacker4x}; + + fn mask_for_width(width: u8) -> u32 { + match width { + 0 => 0, + 32 => u32::MAX, + _ => (1u32 << width) - 1, + } + } + + fn raw_values(width: u8) -> Vec { + let mask = mask_for_width(width); + (0..BLOCK_LEN) + .map(|idx| ((idx * 17 + 3) as u32) & mask) + .collect() + } + + fn sorted_values(width: u8) -> (u32, Vec) { + if width == 0 { + return (11, vec![11; BLOCK_LEN]); + } + if width == 32 { + return (0, vec![u32::MAX; BLOCK_LEN]); + } + + let mask = mask_for_width(width).min(127); + let initial = 11u32; + let mut current = initial; + let values = (0..BLOCK_LEN) + .map(|idx| { + current += (idx as u32 * 7 + 1) & mask; + current + }) + .collect(); + (initial, values) + } + + fn strictly_sorted_values(width: u8) -> (Option, Vec) { + let mask = mask_for_width(width).min(127); + let mut current = 0u32; + let values = (0..BLOCK_LEN) + .map(|idx| { + if idx == 0 { + current = 0; + } else { + current += 1 + ((idx as u32 * 5) & mask); + } + current + }) + .collect(); + (None, values) + } + + #[test] + fn scalar_backend_matches_external_bitpacker4x() { + let scalar = BitPacker4x::new_scalar(); + let external = ExternalBitPacker4x::new(); + + for width in 0..=32 { + let values = raw_values(width); + assert_eq!(scalar.num_bits(&values), external.num_bits(&values)); + + let mut actual = vec![0u8; BitPacker4x::compressed_block_size(width)]; + let actual_len = scalar.compress(&values, &mut actual, width); + + let mut expected = vec![0u8; ExternalBitPacker4x::compressed_block_size(width)]; + let expected_len = external.compress(&values, &mut expected, width); + + assert_eq!(actual_len, expected_len); + assert_eq!(actual, expected, "raw width {width}"); + + let mut decoded = vec![0u32; BLOCK_LEN]; + assert_eq!(scalar.decompress(&actual, &mut decoded, width), actual_len); + assert_eq!(decoded, values); + + let (initial, values) = sorted_values(width); + assert_eq!( + scalar.num_bits_sorted(initial, &values), + external.num_bits_sorted(initial, &values) + ); + + let mut actual = vec![0u8; BitPacker4x::compressed_block_size(width)]; + let actual_len = scalar.compress_sorted(initial, &values, &mut actual, width); + + let mut expected = vec![0u8; ExternalBitPacker4x::compressed_block_size(width)]; + let expected_len = external.compress_sorted(initial, &values, &mut expected, width); + + assert_eq!(actual_len, expected_len); + assert_eq!(actual, expected, "sorted width {width}"); + + let mut decoded = vec![0u32; BLOCK_LEN]; + assert_eq!( + scalar.decompress_sorted(initial, &actual, &mut decoded, width), + actual_len + ); + assert_eq!(decoded, values); + } + } + + #[test] + fn scalar_backend_matches_external_strictly_sorted_bitpacker4x() { + let scalar = BitPacker4x::new_scalar(); + let external = ExternalBitPacker4x::new(); + + for width in 0..=16 { + let (initial, values) = strictly_sorted_values(width); + let num_bits = external.num_bits_strictly_sorted(initial, &values); + assert_eq!(scalar.num_bits_strictly_sorted(initial, &values), num_bits); + + let mut actual = vec![0u8; BitPacker4x::compressed_block_size(num_bits)]; + let actual_len = + scalar.compress_strictly_sorted(initial, &values, &mut actual, num_bits); + + let mut expected = vec![0u8; ExternalBitPacker4x::compressed_block_size(num_bits)]; + let expected_len = + external.compress_strictly_sorted(initial, &values, &mut expected, num_bits); + + assert_eq!(actual_len, expected_len); + assert_eq!(actual, expected, "strict width {width}"); + + let mut decoded = vec![0u32; BLOCK_LEN]; + assert_eq!( + scalar.decompress_strictly_sorted(initial, &actual, &mut decoded, num_bits), + actual_len + ); + assert_eq!(decoded, values); + } + } +} diff --git a/rust/compression/bitpacking/src/bitpacker_internal/bitpacker8x.rs b/rust/compression/bitpacking/src/bitpacker_internal/bitpacker8x.rs new file mode 100644 index 00000000000..b17edacabb2 --- /dev/null +++ b/rust/compression/bitpacking/src/bitpacker_internal/bitpacker8x.rs @@ -0,0 +1,872 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use super::BitPacker; + +use crate::bitpacker_internal::{Available, UnsafeBitPacker}; + +const BLOCK_LEN: usize = 32 * 8; + +#[cfg(target_arch = "x86_64")] +mod avx2 { + use super::BLOCK_LEN; + use crate::bitpacker_internal::Available; + + use std::arch::x86_64::__m256i as DataType; + use std::arch::x86_64::_mm256_and_si256 as op_and; + use std::arch::x86_64::_mm256_lddqu_si256 as load_unaligned; + use std::arch::x86_64::_mm256_or_si256 as op_or; + use std::arch::x86_64::_mm256_set1_epi32 as set1; + use std::arch::x86_64::_mm256_slli_epi32 as left_shift_32; + use std::arch::x86_64::_mm256_srli_epi32 as right_shift_32; + use std::arch::x86_64::_mm256_storeu_si256 as store_unaligned; + + use std::arch::x86_64::{ + _mm256_add_epi32, _mm256_extract_epi32, _mm256_permute2f128_si256, _mm256_shuffle_epi32, + _mm256_slli_si256, _mm256_srli_si256, _mm256_sub_epi32, + }; + + #[allow(non_snake_case)] + unsafe fn or_collapse_to_u32(accumulator: DataType) -> u32 { + let a__b__c__d__e__f__g__h_ = accumulator; + let ______a__b________e__f = _mm256_srli_si256(a__b__c__d__e__f__g__h_, 8); + let a__b__ca_db_e__f__ge_hf = op_or(a__b__c__d__e__f__g__h_, ______a__b________e__f); + let ___a__b__ca____e__f__ge = _mm256_srli_si256(a__b__ca_db_e__f__ge_hf, 4); + let _________cadb______gehf = op_or(a__b__ca_db_e__f__ge_hf, ___a__b__ca____e__f__ge); + let cadb = _mm256_extract_epi32(_________cadb______gehf, 0); + let gehf = _mm256_extract_epi32(_________cadb______gehf, 4); + (cadb | gehf) as u32 + } + + unsafe fn compute_delta(curr: DataType, prev: DataType) -> DataType { + let left_shift = _mm256_slli_si256(curr, 4); + let curr_shift = _mm256_srli_si256(curr, 12); + let curr_right_only = _mm256_permute2f128_si256(curr_shift, curr_shift, 8); + let prev_shift = _mm256_srli_si256(prev, 12); + let sub_left = _mm256_permute2f128_si256(prev_shift, prev_shift, 3 | (8 << 4)); + let diff = op_or(left_shift, op_or(curr_right_only, sub_left)); + _mm256_sub_epi32(curr, diff) + } + + #[allow(non_snake_case)] + unsafe fn integrate_delta(prev: DataType, delta: DataType) -> DataType { + let offset_repeat = _mm256_shuffle_epi32(prev, 0xff); + let offset = _mm256_permute2f128_si256(offset_repeat, offset_repeat, 3 | (8 << 4)); + let a__b__c__d__e__f__g__h__ = delta; + let ______a__b________e__f__ = _mm256_slli_si256(delta, 8); + let a__b__ca_db_e__f__ge_fh_ = + _mm256_add_epi32(a__b__c__d__e__f__g__h__, ______a__b________e__f__); + let ___a__b__ca____e__f__ge_ = _mm256_slli_si256(a__b__ca_db_e__f__ge_fh_, 4); + let halved_prefix_sum = + _mm256_add_epi32(___a__b__ca____e__f__ge_, a__b__ca_db_e__f__ge_fh_); + let offsetted_halved_prefix_sum = _mm256_add_epi32(halved_prefix_sum, offset); + let select_last_low = _mm256_shuffle_epi32(offsetted_halved_prefix_sum, 0xff); + let high_offset = _mm256_permute2f128_si256(select_last_low, select_last_low, 8); + _mm256_add_epi32(high_offset, offsetted_halved_prefix_sum) + } + + unsafe fn add(left: DataType, right: DataType) -> DataType { + _mm256_add_epi32(left, right) + } + + unsafe fn sub(left: DataType, right: DataType) -> DataType { + _mm256_sub_epi32(left, right) + } + + declare_bitpacker!(target_feature(enable = "avx2")); + + impl Available for UnsafeBitPackerImpl { + fn available() -> bool { + is_x86_feature_detected!("avx2") + } + } +} + +#[cfg(all(target_arch = "aarch64", target_endian = "little"))] +mod neon { + use super::BLOCK_LEN; + use crate::bitpacker_internal::Available; + use std::arch::aarch64::{ + uint32x4_t, vaddq_u32, vandq_u32, vdupq_n_u32, vextq_u32, vgetq_lane_u32, vld1q_u32, + vorrq_u32, vshlq_n_u32, vshrq_n_u32, vst1q_u32, vsubq_u32, + }; + + pub(crate) type DataType = [uint32x4_t; 2]; + + #[inline] + unsafe fn set1(el: i32) -> DataType { + let lanes = vdupq_n_u32(el as u32); + [lanes, lanes] + } + + #[inline] + unsafe fn right_shift_32(el: DataType) -> DataType { + const { + assert!(N >= 0); + assert!(N <= 32); + } + + match N { + 0 => el, + 1 => [vshrq_n_u32::<1>(el[0]), vshrq_n_u32::<1>(el[1])], + 2 => [vshrq_n_u32::<2>(el[0]), vshrq_n_u32::<2>(el[1])], + 3 => [vshrq_n_u32::<3>(el[0]), vshrq_n_u32::<3>(el[1])], + 4 => [vshrq_n_u32::<4>(el[0]), vshrq_n_u32::<4>(el[1])], + 5 => [vshrq_n_u32::<5>(el[0]), vshrq_n_u32::<5>(el[1])], + 6 => [vshrq_n_u32::<6>(el[0]), vshrq_n_u32::<6>(el[1])], + 7 => [vshrq_n_u32::<7>(el[0]), vshrq_n_u32::<7>(el[1])], + 8 => [vshrq_n_u32::<8>(el[0]), vshrq_n_u32::<8>(el[1])], + 9 => [vshrq_n_u32::<9>(el[0]), vshrq_n_u32::<9>(el[1])], + 10 => [vshrq_n_u32::<10>(el[0]), vshrq_n_u32::<10>(el[1])], + 11 => [vshrq_n_u32::<11>(el[0]), vshrq_n_u32::<11>(el[1])], + 12 => [vshrq_n_u32::<12>(el[0]), vshrq_n_u32::<12>(el[1])], + 13 => [vshrq_n_u32::<13>(el[0]), vshrq_n_u32::<13>(el[1])], + 14 => [vshrq_n_u32::<14>(el[0]), vshrq_n_u32::<14>(el[1])], + 15 => [vshrq_n_u32::<15>(el[0]), vshrq_n_u32::<15>(el[1])], + 16 => [vshrq_n_u32::<16>(el[0]), vshrq_n_u32::<16>(el[1])], + 17 => [vshrq_n_u32::<17>(el[0]), vshrq_n_u32::<17>(el[1])], + 18 => [vshrq_n_u32::<18>(el[0]), vshrq_n_u32::<18>(el[1])], + 19 => [vshrq_n_u32::<19>(el[0]), vshrq_n_u32::<19>(el[1])], + 20 => [vshrq_n_u32::<20>(el[0]), vshrq_n_u32::<20>(el[1])], + 21 => [vshrq_n_u32::<21>(el[0]), vshrq_n_u32::<21>(el[1])], + 22 => [vshrq_n_u32::<22>(el[0]), vshrq_n_u32::<22>(el[1])], + 23 => [vshrq_n_u32::<23>(el[0]), vshrq_n_u32::<23>(el[1])], + 24 => [vshrq_n_u32::<24>(el[0]), vshrq_n_u32::<24>(el[1])], + 25 => [vshrq_n_u32::<25>(el[0]), vshrq_n_u32::<25>(el[1])], + 26 => [vshrq_n_u32::<26>(el[0]), vshrq_n_u32::<26>(el[1])], + 27 => [vshrq_n_u32::<27>(el[0]), vshrq_n_u32::<27>(el[1])], + 28 => [vshrq_n_u32::<28>(el[0]), vshrq_n_u32::<28>(el[1])], + 29 => [vshrq_n_u32::<29>(el[0]), vshrq_n_u32::<29>(el[1])], + 30 => [vshrq_n_u32::<30>(el[0]), vshrq_n_u32::<30>(el[1])], + 31 => [vshrq_n_u32::<31>(el[0]), vshrq_n_u32::<31>(el[1])], + 32 => set1(0), + _ => core::hint::unreachable_unchecked(), + } + } + + #[inline] + unsafe fn left_shift_32(el: DataType) -> DataType { + const { + assert!(N >= 0); + assert!(N <= 32); + } + + match N { + 0 => el, + 1 => [vshlq_n_u32::<1>(el[0]), vshlq_n_u32::<1>(el[1])], + 2 => [vshlq_n_u32::<2>(el[0]), vshlq_n_u32::<2>(el[1])], + 3 => [vshlq_n_u32::<3>(el[0]), vshlq_n_u32::<3>(el[1])], + 4 => [vshlq_n_u32::<4>(el[0]), vshlq_n_u32::<4>(el[1])], + 5 => [vshlq_n_u32::<5>(el[0]), vshlq_n_u32::<5>(el[1])], + 6 => [vshlq_n_u32::<6>(el[0]), vshlq_n_u32::<6>(el[1])], + 7 => [vshlq_n_u32::<7>(el[0]), vshlq_n_u32::<7>(el[1])], + 8 => [vshlq_n_u32::<8>(el[0]), vshlq_n_u32::<8>(el[1])], + 9 => [vshlq_n_u32::<9>(el[0]), vshlq_n_u32::<9>(el[1])], + 10 => [vshlq_n_u32::<10>(el[0]), vshlq_n_u32::<10>(el[1])], + 11 => [vshlq_n_u32::<11>(el[0]), vshlq_n_u32::<11>(el[1])], + 12 => [vshlq_n_u32::<12>(el[0]), vshlq_n_u32::<12>(el[1])], + 13 => [vshlq_n_u32::<13>(el[0]), vshlq_n_u32::<13>(el[1])], + 14 => [vshlq_n_u32::<14>(el[0]), vshlq_n_u32::<14>(el[1])], + 15 => [vshlq_n_u32::<15>(el[0]), vshlq_n_u32::<15>(el[1])], + 16 => [vshlq_n_u32::<16>(el[0]), vshlq_n_u32::<16>(el[1])], + 17 => [vshlq_n_u32::<17>(el[0]), vshlq_n_u32::<17>(el[1])], + 18 => [vshlq_n_u32::<18>(el[0]), vshlq_n_u32::<18>(el[1])], + 19 => [vshlq_n_u32::<19>(el[0]), vshlq_n_u32::<19>(el[1])], + 20 => [vshlq_n_u32::<20>(el[0]), vshlq_n_u32::<20>(el[1])], + 21 => [vshlq_n_u32::<21>(el[0]), vshlq_n_u32::<21>(el[1])], + 22 => [vshlq_n_u32::<22>(el[0]), vshlq_n_u32::<22>(el[1])], + 23 => [vshlq_n_u32::<23>(el[0]), vshlq_n_u32::<23>(el[1])], + 24 => [vshlq_n_u32::<24>(el[0]), vshlq_n_u32::<24>(el[1])], + 25 => [vshlq_n_u32::<25>(el[0]), vshlq_n_u32::<25>(el[1])], + 26 => [vshlq_n_u32::<26>(el[0]), vshlq_n_u32::<26>(el[1])], + 27 => [vshlq_n_u32::<27>(el[0]), vshlq_n_u32::<27>(el[1])], + 28 => [vshlq_n_u32::<28>(el[0]), vshlq_n_u32::<28>(el[1])], + 29 => [vshlq_n_u32::<29>(el[0]), vshlq_n_u32::<29>(el[1])], + 30 => [vshlq_n_u32::<30>(el[0]), vshlq_n_u32::<30>(el[1])], + 31 => [vshlq_n_u32::<31>(el[0]), vshlq_n_u32::<31>(el[1])], + 32 => set1(0), + _ => core::hint::unreachable_unchecked(), + } + } + + #[inline] + unsafe fn op_or(left: DataType, right: DataType) -> DataType { + [vorrq_u32(left[0], right[0]), vorrq_u32(left[1], right[1])] + } + + #[inline] + unsafe fn op_and(left: DataType, right: DataType) -> DataType { + [vandq_u32(left[0], right[0]), vandq_u32(left[1], right[1])] + } + + #[inline] + unsafe fn load_unaligned(addr: *const DataType) -> DataType { + let ptr = addr.cast::(); + [vld1q_u32(ptr), vld1q_u32(ptr.add(4))] + } + + #[inline] + unsafe fn store_unaligned(addr: *mut DataType, data: DataType) { + let ptr = addr.cast::(); + vst1q_u32(ptr, data[0]); + vst1q_u32(ptr.add(4), data[1]); + } + + #[inline] + unsafe fn or_collapse_to_u32(accumulator: DataType) -> u32 { + vgetq_lane_u32(accumulator[0], 0) + | vgetq_lane_u32(accumulator[0], 1) + | vgetq_lane_u32(accumulator[0], 2) + | vgetq_lane_u32(accumulator[0], 3) + | vgetq_lane_u32(accumulator[1], 0) + | vgetq_lane_u32(accumulator[1], 1) + | vgetq_lane_u32(accumulator[1], 2) + | vgetq_lane_u32(accumulator[1], 3) + } + + #[inline] + unsafe fn compute_delta(curr: DataType, prev: DataType) -> DataType { + [ + vsubq_u32(curr[0], vextq_u32(prev[1], curr[0], 3)), + vsubq_u32(curr[1], vextq_u32(curr[0], curr[1], 3)), + ] + } + + #[allow(non_snake_case)] + #[inline] + unsafe fn integrate_half(base: u32, delta: uint32x4_t) -> uint32x4_t { + let base = vdupq_n_u32(base); + let zero = vdupq_n_u32(0); + let a__b__c__d_ = delta; + let ______a__b_ = vextq_u32(zero, a__b__c__d_, 2); + let a__b__ca_db = vaddq_u32(______a__b_, a__b__c__d_); + let ___a__b__ca = vextq_u32(zero, a__b__ca_db, 3); + let a_ab_abc_abcd = vaddq_u32(___a__b__ca, a__b__ca_db); + vaddq_u32(base, a_ab_abc_abcd) + } + + #[inline] + unsafe fn integrate_delta(prev: DataType, delta: DataType) -> DataType { + let low = integrate_half(vgetq_lane_u32(prev[1], 3), delta[0]); + let high = integrate_half(vgetq_lane_u32(low, 3), delta[1]); + [low, high] + } + + #[inline] + unsafe fn add(left: DataType, right: DataType) -> DataType { + [vaddq_u32(left[0], right[0]), vaddq_u32(left[1], right[1])] + } + + #[inline] + unsafe fn sub(left: DataType, right: DataType) -> DataType { + [vsubq_u32(left[0], right[0]), vsubq_u32(left[1], right[1])] + } + + declare_bitpacker!(target_feature(enable = "neon")); + + impl Available for UnsafeBitPackerImpl { + fn available() -> bool { + std::arch::is_aarch64_feature_detected!("neon") + } + } +} + +mod scalar { + use super::BLOCK_LEN; + use crate::bitpacker_internal::Available; + use std::ptr; + + pub(crate) type DataType = [u32; 8]; + + pub(crate) fn set1(el: i32) -> DataType { + [el as u32; 8] + } + + pub(crate) fn right_shift_32(el: DataType) -> DataType { + [ + el[0] >> N, + el[1] >> N, + el[2] >> N, + el[3] >> N, + el[4] >> N, + el[5] >> N, + el[6] >> N, + el[7] >> N, + ] + } + + pub(crate) fn left_shift_32(el: DataType) -> DataType { + [ + el[0] << N, + el[1] << N, + el[2] << N, + el[3] << N, + el[4] << N, + el[5] << N, + el[6] << N, + el[7] << N, + ] + } + + pub(crate) fn op_or(left: DataType, right: DataType) -> DataType { + [ + left[0] | right[0], + left[1] | right[1], + left[2] | right[2], + left[3] | right[3], + left[4] | right[4], + left[5] | right[5], + left[6] | right[6], + left[7] | right[7], + ] + } + + pub(crate) fn op_and(left: DataType, right: DataType) -> DataType { + [ + left[0] & right[0], + left[1] & right[1], + left[2] & right[2], + left[3] & right[3], + left[4] & right[4], + left[5] & right[5], + left[6] & right[6], + left[7] & right[7], + ] + } + + pub(crate) unsafe fn load_unaligned(addr: *const DataType) -> DataType { + ptr::read_unaligned(addr) + } + + pub(crate) unsafe fn store_unaligned(addr: *mut DataType, data: DataType) { + ptr::write_unaligned(addr, data); + } + + pub(crate) fn or_collapse_to_u32(accumulator: DataType) -> u32 { + ((accumulator[0] | accumulator[1]) | (accumulator[2] | accumulator[3])) + | ((accumulator[4] | accumulator[5]) | (accumulator[6] | accumulator[7])) + } + + fn compute_delta(curr: DataType, prev: DataType) -> DataType { + [ + curr[0].wrapping_sub(prev[7]), + curr[1].wrapping_sub(curr[0]), + curr[2].wrapping_sub(curr[1]), + curr[3].wrapping_sub(curr[2]), + curr[4].wrapping_sub(curr[3]), + curr[5].wrapping_sub(curr[4]), + curr[6].wrapping_sub(curr[5]), + curr[7].wrapping_sub(curr[6]), + ] + } + + fn integrate_delta(offset: DataType, delta: DataType) -> DataType { + let el0 = offset[7].wrapping_add(delta[0]); + let el1 = el0.wrapping_add(delta[1]); + let el2 = el1.wrapping_add(delta[2]); + let el3 = el2.wrapping_add(delta[3]); + let el4 = el3.wrapping_add(delta[4]); + let el5 = el4.wrapping_add(delta[5]); + let el6 = el5.wrapping_add(delta[6]); + let el7 = el6.wrapping_add(delta[7]); + [el0, el1, el2, el3, el4, el5, el6, el7] + } + + pub(crate) fn add(left: DataType, right: DataType) -> DataType { + [ + left[0].wrapping_add(right[0]), + left[1].wrapping_add(right[1]), + left[2].wrapping_add(right[2]), + left[3].wrapping_add(right[3]), + left[4].wrapping_add(right[4]), + left[5].wrapping_add(right[5]), + left[6].wrapping_add(right[6]), + left[7].wrapping_add(right[7]), + ] + } + + pub(crate) fn sub(left: DataType, right: DataType) -> DataType { + [ + left[0].wrapping_sub(right[0]), + left[1].wrapping_sub(right[1]), + left[2].wrapping_sub(right[2]), + left[3].wrapping_sub(right[3]), + left[4].wrapping_sub(right[4]), + left[5].wrapping_sub(right[5]), + left[6].wrapping_sub(right[6]), + left[7].wrapping_sub(right[7]), + ] + } + + // The `allow(unused)` is here to put an attribute that has no effect. + // + // For other bitpackers, we enable a specific CPU instruction set, but for + // the scalar bitpacker none is required. + declare_bitpacker!(allow(unused)); + + impl Available for UnsafeBitPackerImpl { + fn available() -> bool { + true + } + } +} + +#[derive(Clone, Copy)] +enum InstructionSet { + #[cfg(target_arch = "x86_64")] + AVX2, + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + NEON, + Scalar, +} + +/// 8-wide bitpacker implementation. +/// +/// One block contains 256 integers. +#[derive(Clone, Copy)] +pub struct BitPacker8x(InstructionSet); + +impl BitPacker8x { + #[cfg(target_arch = "x86_64")] + pub(crate) fn new_avx2() -> Option { + avx2::UnsafeBitPackerImpl::available().then_some(BitPacker8x(InstructionSet::AVX2)) + } + + #[cfg(not(target_arch = "x86_64"))] + pub(crate) fn new_avx2() -> Option { + None + } + + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + pub(crate) fn new_neon() -> Option { + neon::UnsafeBitPackerImpl::available().then_some(BitPacker8x(InstructionSet::NEON)) + } + + #[cfg(not(all(target_arch = "aarch64", target_endian = "little")))] + pub(crate) fn new_neon() -> Option { + None + } + + pub(crate) fn new_scalar() -> Self { + BitPacker8x(InstructionSet::Scalar) + } +} + +impl BitPacker for BitPacker8x { + const BLOCK_LEN: usize = BLOCK_LEN; + + fn new() -> Self { + Self::new_avx2() + .or_else(Self::new_neon) + .unwrap_or_else(Self::new_scalar) + } + + fn compress(&self, decompressed: &[u32], compressed: &mut [u8], num_bits: u8) -> usize { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::AVX2 => { + avx2::UnsafeBitPackerImpl::compress(decompressed, compressed, num_bits) + } + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => { + neon::UnsafeBitPackerImpl::compress(decompressed, compressed, num_bits) + } + InstructionSet::Scalar => { + scalar::UnsafeBitPackerImpl::compress(decompressed, compressed, num_bits) + } + } + } + } + + fn compress_sorted( + &self, + initial: u32, + decompressed: &[u32], + compressed: &mut [u8], + num_bits: u8, + ) -> usize { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::AVX2 => avx2::UnsafeBitPackerImpl::compress_sorted( + initial, + decompressed, + compressed, + num_bits, + ), + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => neon::UnsafeBitPackerImpl::compress_sorted( + initial, + decompressed, + compressed, + num_bits, + ), + InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::compress_sorted( + initial, + decompressed, + compressed, + num_bits, + ), + } + } + } + + fn compress_strictly_sorted( + &self, + initial: Option, + decompressed: &[u32], + compressed: &mut [u8], + num_bits: u8, + ) -> usize { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::AVX2 => avx2::UnsafeBitPackerImpl::compress_strictly_sorted( + initial, + decompressed, + compressed, + num_bits, + ), + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => neon::UnsafeBitPackerImpl::compress_strictly_sorted( + initial, + decompressed, + compressed, + num_bits, + ), + InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::compress_strictly_sorted( + initial, + decompressed, + compressed, + num_bits, + ), + } + } + } + + fn decompress(&self, compressed: &[u8], decompressed: &mut [u32], num_bits: u8) -> usize { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::AVX2 => { + avx2::UnsafeBitPackerImpl::decompress(compressed, decompressed, num_bits) + } + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => { + neon::UnsafeBitPackerImpl::decompress(compressed, decompressed, num_bits) + } + InstructionSet::Scalar => { + scalar::UnsafeBitPackerImpl::decompress(compressed, decompressed, num_bits) + } + } + } + } + + fn decompress_sorted( + &self, + initial: u32, + compressed: &[u8], + decompressed: &mut [u32], + num_bits: u8, + ) -> usize { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::AVX2 => avx2::UnsafeBitPackerImpl::decompress_sorted( + initial, + compressed, + decompressed, + num_bits, + ), + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => neon::UnsafeBitPackerImpl::decompress_sorted( + initial, + compressed, + decompressed, + num_bits, + ), + InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::decompress_sorted( + initial, + compressed, + decompressed, + num_bits, + ), + } + } + } + + fn decompress_strictly_sorted( + &self, + initial: Option, + compressed: &[u8], + decompressed: &mut [u32], + num_bits: u8, + ) -> usize { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::AVX2 => avx2::UnsafeBitPackerImpl::decompress_strictly_sorted( + initial, + compressed, + decompressed, + num_bits, + ), + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => neon::UnsafeBitPackerImpl::decompress_strictly_sorted( + initial, + compressed, + decompressed, + num_bits, + ), + InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::decompress_strictly_sorted( + initial, + compressed, + decompressed, + num_bits, + ), + } + } + } + + fn num_bits(&self, decompressed: &[u32]) -> u8 { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::AVX2 => avx2::UnsafeBitPackerImpl::num_bits(decompressed), + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => neon::UnsafeBitPackerImpl::num_bits(decompressed), + InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::num_bits(decompressed), + } + } + } + + fn num_bits_sorted(&self, initial: u32, decompressed: &[u32]) -> u8 { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::AVX2 => { + avx2::UnsafeBitPackerImpl::num_bits_sorted(initial, decompressed) + } + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => { + neon::UnsafeBitPackerImpl::num_bits_sorted(initial, decompressed) + } + InstructionSet::Scalar => { + scalar::UnsafeBitPackerImpl::num_bits_sorted(initial, decompressed) + } + } + } + } + + fn num_bits_strictly_sorted(&self, initial: Option, decompressed: &[u32]) -> u8 { + unsafe { + match self.0 { + #[cfg(target_arch = "x86_64")] + InstructionSet::AVX2 => { + avx2::UnsafeBitPackerImpl::num_bits_strictly_sorted(initial, decompressed) + } + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + InstructionSet::NEON => { + neon::UnsafeBitPackerImpl::num_bits_strictly_sorted(initial, decompressed) + } + InstructionSet::Scalar => { + scalar::UnsafeBitPackerImpl::num_bits_strictly_sorted(initial, decompressed) + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::BitPacker8x; + use crate::bitpacker_internal::BitPacker; + use bitpacking::{BitPacker as ExternalBitPacker, BitPacker8x as ExternalBitPacker8x}; + + fn mask_for_width(width: u8) -> u32 { + match width { + 0 => 0, + 32 => u32::MAX, + _ => (1u32 << width) - 1, + } + } + + fn raw_values(width: u8, seed: u64) -> Vec { + let mask = mask_for_width(width); + let mut state = seed; + (0..BitPacker8x::BLOCK_LEN) + .map(|idx| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + match seed % 4 { + 0 => 0, + 1 => mask, + 2 => idx as u32 & mask, + _ => state as u32 & mask, + } + }) + .collect() + } + + fn sorted_values(width: u8, seed: u64) -> (u32, Vec) { + if width == 0 { + return (17, vec![17; BitPacker8x::BLOCK_LEN]); + } + if width == 32 { + return (0, vec![u32::MAX; BitPacker8x::BLOCK_LEN]); + } + + let mask = mask_for_width(width).min(127); + let mut state = seed; + let mut current = 17u32; + let values = (0..BitPacker8x::BLOCK_LEN) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + current += state as u32 & mask; + current + }) + .collect(); + (17, values) + } + + fn strictly_sorted_values(width: u8, seed: u64) -> (Option, Vec) { + let mask = mask_for_width(width).min(127); + let mut state = seed; + let mut current = 0u32; + let values = (0..BitPacker8x::BLOCK_LEN) + .map(|idx| { + if idx == 0 { + current = 0; + } else { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + current += 1 + (state as u32 & mask); + } + current + }) + .collect(); + (None, values) + } + + fn assert_raw_compatible(ours: BitPacker8x, external: ExternalBitPacker8x) { + for width in 0..=32 { + for seed in [0, 1, 2, 123456789] { + let values = raw_values(width, seed); + assert_eq!(ours.num_bits(&values), external.num_bits(&values)); + + let mut actual = vec![0u8; BitPacker8x::compressed_block_size(width)]; + let actual_len = ours.compress(&values, &mut actual, width); + + let mut expected = vec![0u8; ExternalBitPacker8x::compressed_block_size(width)]; + let expected_len = external.compress(&values, &mut expected, width); + + assert_eq!(actual_len, expected_len); + assert_eq!(actual, expected, "raw width {width} seed {seed}"); + + let mut decoded = vec![0u32; BitPacker8x::BLOCK_LEN]; + assert_eq!(ours.decompress(&actual, &mut decoded, width), actual_len); + assert_eq!(decoded, values); + } + } + } + + fn assert_sorted_compatible(ours: BitPacker8x, external: ExternalBitPacker8x) { + for width in 0..=32 { + for seed in [0, 1, 2, 123456789] { + let (initial, values) = sorted_values(width, seed); + assert_eq!( + ours.num_bits_sorted(initial, &values), + external.num_bits_sorted(initial, &values) + ); + + let mut actual = vec![0u8; BitPacker8x::compressed_block_size(width)]; + let actual_len = ours.compress_sorted(initial, &values, &mut actual, width); + + let mut expected = vec![0u8; ExternalBitPacker8x::compressed_block_size(width)]; + let expected_len = external.compress_sorted(initial, &values, &mut expected, width); + + assert_eq!(actual_len, expected_len); + assert_eq!(actual, expected, "sorted width {width} seed {seed}"); + + let mut decoded = vec![0u32; BitPacker8x::BLOCK_LEN]; + assert_eq!( + ours.decompress_sorted(initial, &actual, &mut decoded, width), + actual_len + ); + assert_eq!(decoded, values); + } + } + } + + fn assert_strictly_sorted_compatible(ours: BitPacker8x, external: ExternalBitPacker8x) { + for width in 0..=16 { + for seed in [0, 1, 2, 123456789] { + let (initial, values) = strictly_sorted_values(width, seed); + let num_bits = external.num_bits_strictly_sorted(initial, &values); + assert_eq!(ours.num_bits_strictly_sorted(initial, &values), num_bits); + + let mut actual = vec![0u8; BitPacker8x::compressed_block_size(num_bits)]; + let actual_len = + ours.compress_strictly_sorted(initial, &values, &mut actual, num_bits); + + let mut expected = vec![0u8; ExternalBitPacker8x::compressed_block_size(num_bits)]; + let expected_len = + external.compress_strictly_sorted(initial, &values, &mut expected, num_bits); + + assert_eq!(actual_len, expected_len); + assert_eq!(actual, expected, "strict width {width} seed {seed}"); + + let mut decoded = vec![0u32; BitPacker8x::BLOCK_LEN]; + assert_eq!( + ours.decompress_strictly_sorted(initial, &actual, &mut decoded, num_bits), + actual_len + ); + assert_eq!(decoded, values); + } + } + } + + #[test] + fn bitpacker8x_raw_compatible_with_external_bitpacking() { + assert_raw_compatible(BitPacker8x::new(), ExternalBitPacker8x::new()); + } + + #[test] + fn bitpacker8x_sorted_compatible_with_external_bitpacking() { + assert_sorted_compatible(BitPacker8x::new(), ExternalBitPacker8x::new()); + } + + #[test] + fn scalar_backend_matches_external_bitpacker8x() { + let scalar = BitPacker8x::new_scalar(); + let external = ExternalBitPacker8x::new(); + + assert_raw_compatible(scalar, external); + assert_sorted_compatible(scalar, external); + } + + #[test] + fn scalar_backend_matches_external_strictly_sorted_bitpacker8x() { + let scalar = BitPacker8x::new_scalar(); + let external = ExternalBitPacker8x::new(); + + assert_strictly_sorted_compatible(scalar, external); + } + + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + #[test] + fn neon_backend_matches_external_bitpacker8x() { + if let Some(neon) = BitPacker8x::new_neon() { + let external = ExternalBitPacker8x::new(); + + assert_raw_compatible(neon, external); + assert_sorted_compatible(neon, external); + assert_strictly_sorted_compatible(neon, external); + } + } +} diff --git a/rust/compression/bitpacking/src/bitpacker_internal/macros.rs b/rust/compression/bitpacking/src/bitpacker_internal/macros.rs new file mode 100644 index 00000000000..e79686fe9af --- /dev/null +++ b/rust/compression/bitpacking/src/bitpacker_internal/macros.rs @@ -0,0 +1,669 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +macro_rules! pack_unpack_with_bits { + + ($name:ident, $n:expr, $cpufeature:meta) => { + + + mod $name { + + use crunchy::unroll; + use super::BLOCK_LEN; + use super::{Sink, Transformer}; + use super::{DataType, + set1, + right_shift_32, + left_shift_32, + op_or, + op_and, + load_unaligned, + store_unaligned}; + + const NUM_BITS: usize = $n; + const NUM_BYTES_PER_BLOCK: usize = NUM_BITS * BLOCK_LEN / 8; + + #[$cpufeature] + pub(crate) unsafe fn pack(input_arr: &[u32], output_arr: &mut [u8], mut delta_computer: TDeltaComputer) -> usize { + assert_eq!(input_arr.len(), BLOCK_LEN, "Input block too small {}, (expected {})", input_arr.len(), BLOCK_LEN); + assert!(output_arr.len() >= NUM_BYTES_PER_BLOCK, "Output array too small (numbits {}). {} <= {}", NUM_BITS, output_arr.len(), NUM_BYTES_PER_BLOCK); + + let input_ptr = input_arr.as_ptr().cast::(); + let mut output_ptr = output_arr.as_mut_ptr().cast::(); + let mut out_register: DataType = delta_computer.transform(load_unaligned(input_ptr)); + + unroll! { + for iter in 0..30 { + const i: usize = 1 + iter; + + const bits_filled: usize = i * NUM_BITS; + const inner_cursor: usize = bits_filled % 32; + const remaining: usize = 32 - inner_cursor; + + let offset_ptr = input_ptr.add(i); + let in_register: DataType = delta_computer.transform(load_unaligned(offset_ptr)); + + out_register = + if inner_cursor > 0 { + op_or(out_register, left_shift_32::<{inner_cursor as i32}>(in_register)) + } else { + in_register + }; + + if remaining <= NUM_BITS { + store_unaligned(output_ptr, out_register); + output_ptr = output_ptr.add(1); + if 0 < remaining && remaining < NUM_BITS { + out_register = right_shift_32::<{remaining as i32}>(in_register); + } + } + } + } + let in_register: DataType = delta_computer.transform(load_unaligned(input_ptr.add(31))); + out_register = if 32 - NUM_BITS > 0 { + op_or(out_register, left_shift_32::<{32 - NUM_BITS as i32}>(in_register)) + } else { + op_or(out_register, in_register) + }; + store_unaligned(output_ptr, out_register); + + NUM_BYTES_PER_BLOCK + } + + #[$cpufeature] + pub(crate) unsafe fn unpack(compressed: &[u8], mut output: Output) -> usize { + + assert!(compressed.len() >= NUM_BYTES_PER_BLOCK, "Compressed array seems too small. ({} < {}) ", compressed.len(), NUM_BYTES_PER_BLOCK); + + let mut input_ptr = compressed.as_ptr().cast::(); + + let mask_scalar: u32 = ((1u64 << NUM_BITS) - 1u64) as u32; + let mask = set1(mask_scalar as i32); + + let mut in_register: DataType = load_unaligned(input_ptr); + + let out_register = op_and(in_register, mask); + output.process(out_register); + + unroll! { + for iter in 0..31 { + const i: usize = iter + 1; + + const inner_cursor: usize = (i * NUM_BITS) % 32; + const inner_capacity: usize = 32 - inner_cursor; + + let shifted_in_register = if inner_cursor != 0 { + right_shift_32::<{inner_cursor as i32}>(in_register) + } else { + in_register + }; + let mut out_register: DataType = op_and(shifted_in_register, mask); + + // We consumed our current quadruplets entirely. + // We therefore read another one. + if inner_capacity <= NUM_BITS && i != 31 { + input_ptr = input_ptr.add(1); + in_register = load_unaligned(input_ptr); + + // This quadruplets is actually cutting one of + // our `DataType`. We need to read the next one. + if inner_capacity < NUM_BITS { + let shifted = if inner_capacity != 0 { + left_shift_32::<{inner_capacity as i32}>(in_register) + } else { + in_register + }; + let masked = op_and(shifted, mask); + out_register = op_or(out_register, masked); + } + } + + output.process(out_register); + } + } + + + NUM_BYTES_PER_BLOCK + } + } + } +} + +macro_rules! pack_unpack_with_bits_32 { + ($cpufeature:meta) => { + mod pack_unpack_with_bits_32 { + use super::BLOCK_LEN; + use super::{DataType, load_unaligned, store_unaligned}; + use super::{Sink, Transformer}; + use crunchy::unroll; + + const NUM_BITS: usize = 32; + const NUM_BYTES_PER_BLOCK: usize = NUM_BITS * BLOCK_LEN / 8; + + #[$cpufeature] + pub(crate) unsafe fn pack( + input_arr: &[u32], + output_arr: &mut [u8], + mut delta_computer: TDeltaComputer, + ) -> usize { + assert_eq!( + input_arr.len(), + BLOCK_LEN, + "Input block too small {}, (expected {})", + input_arr.len(), + BLOCK_LEN + ); + assert!( + output_arr.len() >= NUM_BYTES_PER_BLOCK, + "Output array too small (numbits {}). {} <= {}", + NUM_BITS, + output_arr.len(), + NUM_BYTES_PER_BLOCK + ); + + let input_ptr: *const DataType = input_arr.as_ptr().cast::(); + let output_ptr = output_arr.as_mut_ptr().cast::(); + unroll! { + for i in 0..32 { + let input_offset_ptr = input_ptr.add(i); + let output_offset_ptr = output_ptr.add(i); + let input_register = load_unaligned(input_offset_ptr); + let output_register = delta_computer.transform(input_register); + store_unaligned(output_offset_ptr, output_register); + } + } + NUM_BYTES_PER_BLOCK + } + + #[$cpufeature] + pub(crate) unsafe fn unpack( + compressed: &[u8], + mut output: Output, + ) -> usize { + assert!( + compressed.len() >= NUM_BYTES_PER_BLOCK, + "Compressed array seems too small. ({} < {}) ", + compressed.len(), + NUM_BYTES_PER_BLOCK + ); + let input_ptr = compressed.as_ptr().cast::(); + for i in 0..32 { + let input_offset_ptr = input_ptr.add(i); + let in_register: DataType = load_unaligned(input_offset_ptr); + output.process(in_register); + } + NUM_BYTES_PER_BLOCK + } + } + }; +} + +macro_rules! declare_bitpacker { + ($cpufeature:meta) => { + use super::super::UnsafeBitPacker; + use super::super::most_significant_bit; + use crunchy::unroll; + + pack_unpack_with_bits!(pack_unpack_with_bits_1, 1, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_2, 2, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_3, 3, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_4, 4, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_5, 5, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_6, 6, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_7, 7, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_8, 8, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_9, 9, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_10, 10, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_11, 11, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_12, 12, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_13, 13, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_14, 14, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_15, 15, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_16, 16, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_17, 17, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_18, 18, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_19, 19, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_20, 20, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_21, 21, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_22, 22, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_23, 23, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_24, 24, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_25, 25, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_26, 26, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_27, 27, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_28, 28, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_29, 29, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_30, 30, $cpufeature); + pack_unpack_with_bits!(pack_unpack_with_bits_31, 31, $cpufeature); + pack_unpack_with_bits_32!($cpufeature); + + unsafe fn compress_generic( + decompressed: &[u32], + compressed: &mut [u8], + num_bits: u8, + delta_computer: DeltaComputer, + ) -> usize { + match num_bits { + 0 => 0, + 1 => pack_unpack_with_bits_1::pack(decompressed, compressed, delta_computer), + 2 => pack_unpack_with_bits_2::pack(decompressed, compressed, delta_computer), + 3 => pack_unpack_with_bits_3::pack(decompressed, compressed, delta_computer), + 4 => pack_unpack_with_bits_4::pack(decompressed, compressed, delta_computer), + 5 => pack_unpack_with_bits_5::pack(decompressed, compressed, delta_computer), + 6 => pack_unpack_with_bits_6::pack(decompressed, compressed, delta_computer), + 7 => pack_unpack_with_bits_7::pack(decompressed, compressed, delta_computer), + 8 => pack_unpack_with_bits_8::pack(decompressed, compressed, delta_computer), + 9 => pack_unpack_with_bits_9::pack(decompressed, compressed, delta_computer), + 10 => pack_unpack_with_bits_10::pack(decompressed, compressed, delta_computer), + 11 => pack_unpack_with_bits_11::pack(decompressed, compressed, delta_computer), + 12 => pack_unpack_with_bits_12::pack(decompressed, compressed, delta_computer), + 13 => pack_unpack_with_bits_13::pack(decompressed, compressed, delta_computer), + 14 => pack_unpack_with_bits_14::pack(decompressed, compressed, delta_computer), + 15 => pack_unpack_with_bits_15::pack(decompressed, compressed, delta_computer), + 16 => pack_unpack_with_bits_16::pack(decompressed, compressed, delta_computer), + 17 => pack_unpack_with_bits_17::pack(decompressed, compressed, delta_computer), + 18 => pack_unpack_with_bits_18::pack(decompressed, compressed, delta_computer), + 19 => pack_unpack_with_bits_19::pack(decompressed, compressed, delta_computer), + 20 => pack_unpack_with_bits_20::pack(decompressed, compressed, delta_computer), + 21 => pack_unpack_with_bits_21::pack(decompressed, compressed, delta_computer), + 22 => pack_unpack_with_bits_22::pack(decompressed, compressed, delta_computer), + 23 => pack_unpack_with_bits_23::pack(decompressed, compressed, delta_computer), + 24 => pack_unpack_with_bits_24::pack(decompressed, compressed, delta_computer), + 25 => pack_unpack_with_bits_25::pack(decompressed, compressed, delta_computer), + 26 => pack_unpack_with_bits_26::pack(decompressed, compressed, delta_computer), + 27 => pack_unpack_with_bits_27::pack(decompressed, compressed, delta_computer), + 28 => pack_unpack_with_bits_28::pack(decompressed, compressed, delta_computer), + 29 => pack_unpack_with_bits_29::pack(decompressed, compressed, delta_computer), + 30 => pack_unpack_with_bits_30::pack(decompressed, compressed, delta_computer), + 31 => pack_unpack_with_bits_31::pack(decompressed, compressed, delta_computer), + 32 => pack_unpack_with_bits_32::pack(decompressed, compressed, delta_computer), + _ => { + panic!("Num bits must be <= 32. Was {}.", num_bits); + } + } + } + + pub trait Transformer { + unsafe fn transform(&mut self, data: DataType) -> DataType; + } + + struct NoDelta; + + impl Transformer for NoDelta { + #[inline] + unsafe fn transform(&mut self, current: DataType) -> DataType { + current + } + } + + struct DeltaComputer { + pub previous: DataType, + } + + impl Transformer for DeltaComputer { + #[inline] + unsafe fn transform(&mut self, current: DataType) -> DataType { + let result = compute_delta(current, self.previous); + self.previous = current; + result + } + } + + struct StrictDeltaComputer { + pub previous: DataType, + } + + impl Transformer for StrictDeltaComputer { + #[inline] + unsafe fn transform(&mut self, current: DataType) -> DataType { + let result = compute_delta(current, self.previous); + self.previous = current; + sub(result, set1(1)) + } + } + + pub trait Sink { + unsafe fn process(&mut self, data_type: DataType); + } + + struct Store { + output_ptr: *mut DataType, + } + + impl Store { + fn new(output_ptr: *mut DataType) -> Store { + Store { output_ptr } + } + } + + struct DeltaIntegrate { + current: DataType, + output_ptr: *mut DataType, + } + + impl DeltaIntegrate { + unsafe fn new(initial: u32, output_ptr: *mut DataType) -> DeltaIntegrate { + DeltaIntegrate { + current: set1(initial as i32), + output_ptr, + } + } + } + + impl Sink for DeltaIntegrate { + #[inline] + unsafe fn process(&mut self, delta: DataType) { + self.current = integrate_delta(self.current, delta); + store_unaligned(self.output_ptr, self.current); + self.output_ptr = self.output_ptr.add(1); + } + } + + struct StrictDeltaIntegrate { + current: DataType, + output_ptr: *mut DataType, + } + + impl StrictDeltaIntegrate { + unsafe fn new(initial: u32, output_ptr: *mut DataType) -> StrictDeltaIntegrate { + StrictDeltaIntegrate { + current: set1(initial as i32), + output_ptr, + } + } + } + + impl Sink for StrictDeltaIntegrate { + #[inline] + unsafe fn process(&mut self, delta: DataType) { + self.current = integrate_delta(self.current, add(delta, set1(1))); + store_unaligned(self.output_ptr, self.current); + self.output_ptr = self.output_ptr.add(1); + } + } + + impl Sink for Store { + #[inline] + unsafe fn process(&mut self, out_register: DataType) { + store_unaligned(self.output_ptr, out_register); + self.output_ptr = self.output_ptr.add(1); + } + } + + #[inline] + unsafe fn decompress_to( + compressed: &[u8], + mut sink: Output, + num_bits: u8, + ) -> usize { + match num_bits { + 0 => { + let zero = set1(0i32); + for _ in 0..32 { + sink.process(zero); + } + 0 + } + 1 => pack_unpack_with_bits_1::unpack(compressed, sink), + 2 => pack_unpack_with_bits_2::unpack(compressed, sink), + 3 => pack_unpack_with_bits_3::unpack(compressed, sink), + 4 => pack_unpack_with_bits_4::unpack(compressed, sink), + 5 => pack_unpack_with_bits_5::unpack(compressed, sink), + 6 => pack_unpack_with_bits_6::unpack(compressed, sink), + 7 => pack_unpack_with_bits_7::unpack(compressed, sink), + 8 => pack_unpack_with_bits_8::unpack(compressed, sink), + 9 => pack_unpack_with_bits_9::unpack(compressed, sink), + 10 => pack_unpack_with_bits_10::unpack(compressed, sink), + 11 => pack_unpack_with_bits_11::unpack(compressed, sink), + 12 => pack_unpack_with_bits_12::unpack(compressed, sink), + 13 => pack_unpack_with_bits_13::unpack(compressed, sink), + 14 => pack_unpack_with_bits_14::unpack(compressed, sink), + 15 => pack_unpack_with_bits_15::unpack(compressed, sink), + 16 => pack_unpack_with_bits_16::unpack(compressed, sink), + 17 => pack_unpack_with_bits_17::unpack(compressed, sink), + 18 => pack_unpack_with_bits_18::unpack(compressed, sink), + 19 => pack_unpack_with_bits_19::unpack(compressed, sink), + 20 => pack_unpack_with_bits_20::unpack(compressed, sink), + 21 => pack_unpack_with_bits_21::unpack(compressed, sink), + 22 => pack_unpack_with_bits_22::unpack(compressed, sink), + 23 => pack_unpack_with_bits_23::unpack(compressed, sink), + 24 => pack_unpack_with_bits_24::unpack(compressed, sink), + 25 => pack_unpack_with_bits_25::unpack(compressed, sink), + 26 => pack_unpack_with_bits_26::unpack(compressed, sink), + 27 => pack_unpack_with_bits_27::unpack(compressed, sink), + 28 => pack_unpack_with_bits_28::unpack(compressed, sink), + 29 => pack_unpack_with_bits_29::unpack(compressed, sink), + 30 => pack_unpack_with_bits_30::unpack(compressed, sink), + 31 => pack_unpack_with_bits_31::unpack(compressed, sink), + 32 => pack_unpack_with_bits_32::unpack(compressed, sink), + _ => { + panic!("Num bits must be <= 32. Was {}.", num_bits); + } + } + } + + pub struct UnsafeBitPackerImpl; + + impl UnsafeBitPacker for UnsafeBitPackerImpl { + const BLOCK_LEN: usize = BLOCK_LEN; + + #[$cpufeature] + unsafe fn compress(decompressed: &[u32], compressed: &mut [u8], num_bits: u8) -> usize { + compress_generic(decompressed, compressed, num_bits, NoDelta) + } + + #[$cpufeature] + unsafe fn compress_sorted( + initial: u32, + decompressed: &[u32], + compressed: &mut [u8], + num_bits: u8, + ) -> usize { + let delta_computer = DeltaComputer { + previous: set1(initial as i32), + }; + compress_generic(decompressed, compressed, num_bits, delta_computer) + } + + #[$cpufeature] + unsafe fn compress_strictly_sorted( + initial: Option, + decompressed: &[u32], + compressed: &mut [u8], + num_bits: u8, + ) -> usize { + // to allow encoding [0, 1, 2, ..], we need to permit an initial value "lower" than + // zero. To get a clean api, that value is None, but in practice, as we work on + // wrapping integers, u32::MAX/-1 does the job just fine. + let initial = initial.unwrap_or(u32::MAX); + let delta_computer = StrictDeltaComputer { + previous: set1(initial as i32), + }; + compress_generic(decompressed, compressed, num_bits, delta_computer) + } + + #[$cpufeature] + unsafe fn decompress( + compressed: &[u8], + decompressed: &mut [u32], + num_bits: u8, + ) -> usize { + assert!( + decompressed.len() >= BLOCK_LEN, + "The output array is not large enough : ({} >= {})", + decompressed.len(), + BLOCK_LEN + ); + let output_ptr = decompressed.as_mut_ptr().cast::(); + let output = Store::new(output_ptr); + decompress_to(compressed, output, num_bits) + } + + #[$cpufeature] + unsafe fn decompress_sorted( + initial: u32, + compressed: &[u8], + decompressed: &mut [u32], + num_bits: u8, + ) -> usize { + assert!( + decompressed.len() >= BLOCK_LEN, + "The output array is not large enough : ({} >= {})", + decompressed.len(), + BLOCK_LEN + ); + let output_ptr = decompressed.as_mut_ptr().cast::(); + let output = DeltaIntegrate::new(initial, output_ptr); + decompress_to(compressed, output, num_bits) + } + + #[$cpufeature] + unsafe fn decompress_strictly_sorted( + initial: Option, + compressed: &[u8], + decompressed: &mut [u32], + num_bits: u8, + ) -> usize { + assert!( + decompressed.len() >= BLOCK_LEN, + "The output array is not large enough : ({} >= {})", + decompressed.len(), + BLOCK_LEN + ); + let initial = initial.unwrap_or(u32::MAX); + let output_ptr = decompressed.as_mut_ptr().cast::(); + let output = StrictDeltaIntegrate::new(initial, output_ptr); + decompress_to(compressed, output, num_bits) + } + + #[$cpufeature] + unsafe fn num_bits(decompressed: &[u32]) -> u8 { + assert_eq!( + decompressed.len(), + BLOCK_LEN, + "`decompressed`'s len is not `BLOCK_LEN={}`", + BLOCK_LEN + ); + let data: *const DataType = decompressed.as_ptr().cast::(); + let mut accumulator = load_unaligned(data); + unroll! { + for iter in 0..31 { + let i = iter + 1; + let newvec = load_unaligned(data.add(i)); + accumulator = op_or(accumulator, newvec); + } + } + most_significant_bit(or_collapse_to_u32(accumulator)) + } + + #[$cpufeature] + unsafe fn num_bits_sorted(initial: u32, decompressed: &[u32]) -> u8 { + assert_eq!( + decompressed.len(), + BLOCK_LEN, + "`decompressed`'s len is not `BLOCK_LEN={}`", + BLOCK_LEN + ); + let initial_vec = set1(initial as i32); + let data: *const DataType = decompressed.as_ptr().cast::(); + + let first = load_unaligned(data); + let mut accumulator = compute_delta(load_unaligned(data), initial_vec); + let mut previous = first; + + unroll! { + for iter in 0..30 { + let i = iter + 1; + let current = load_unaligned(data.add(i)); + let delta = compute_delta(current, previous); + accumulator = op_or(accumulator, delta); + previous = current; + } + } + let current = load_unaligned(data.add(31)); + let delta = compute_delta(current, previous); + accumulator = op_or(accumulator, delta); + most_significant_bit(or_collapse_to_u32(accumulator)) + } + + #[$cpufeature] + unsafe fn num_bits_strictly_sorted(initial: Option, decompressed: &[u32]) -> u8 { + assert_eq!( + decompressed.len(), + BLOCK_LEN, + "`decompressed`'s len is not `BLOCK_LEN={}`", + BLOCK_LEN + ); + let initial = initial.unwrap_or(u32::MAX); + let initial_vec = set1(initial as i32); + let one = set1(1); + let data: *const DataType = decompressed.as_ptr().cast::(); + + let first = load_unaligned(data); + let mut accumulator = sub(compute_delta(load_unaligned(data), initial_vec), one); + let mut previous = first; + + unroll! { + for iter in 0..30 { + let i = iter + 1; + let current = load_unaligned(data.add(i)); + let delta = sub(compute_delta(current, previous), one); + accumulator = op_or(accumulator, delta); + previous = current; + } + } + let current = load_unaligned(data.add(31)); + let delta = sub(compute_delta(current, previous), one); + accumulator = op_or(accumulator, delta); + most_significant_bit(or_collapse_to_u32(accumulator)) + } + } + + #[cfg(any())] + mod tests { + use super::super::UnsafeBitPacker; + use super::UnsafeBitPackerImpl; + use crate::Available; + use crate::tests::{DeltaKind, test_suite_compress_decompress}; + + #[test] + fn test_num_bits() { + if UnsafeBitPackerImpl::available() { + for num_bits in 0..32 { + for pos in 0..32 { + let mut vals = [0u32; UnsafeBitPackerImpl::BLOCK_LEN]; + if num_bits > 0 { + vals[pos] = 1 << (num_bits - 1); + } + assert_eq!( + unsafe { UnsafeBitPackerImpl::num_bits(&vals[..]) }, + num_bits + ); + } + } + } + } + + #[test] + fn test_bitpacker() { + if UnsafeBitPackerImpl::available() { + test_suite_compress_decompress::(DeltaKind::NoDelta); + } + } + + #[test] + fn test_bitpacker_delta() { + if UnsafeBitPackerImpl::available() { + test_suite_compress_decompress::(DeltaKind::Delta); + } + } + + #[test] + fn test_bitpacker_strict_delta() { + if UnsafeBitPackerImpl::available() { + test_suite_compress_decompress::(DeltaKind::StrictDelta); + } + } + } + }; +} diff --git a/rust/compression/bitpacking/src/bitpacker_internal/mod.rs b/rust/compression/bitpacking/src/bitpacker_internal/mod.rs new file mode 100644 index 00000000000..80803e50ec8 --- /dev/null +++ b/rust/compression/bitpacking/src/bitpacker_internal/mod.rs @@ -0,0 +1,147 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +// Lance-owned u32 SIMD bitpacking kernels. +// +// This is adapted from the MIT-licensed `bitpacking` crate so Lance can keep the +// hot FTS posting-list bitpacking implementation inside lance-bitpacking while +// preserving byte compatibility with the existing 4x format. + +#![allow(dead_code)] +#![allow(unsafe_op_in_unsafe_fn)] +#![allow(clippy::redundant_pub_crate)] +#![allow(clippy::upper_case_acronyms)] +#![allow(clippy::use_self)] + +#[macro_use] +mod macros; + +mod bitpacker4x; +mod bitpacker8x; + +pub use bitpacker4x::BitPacker4x; +pub use bitpacker8x::BitPacker8x; + +pub(crate) trait Available { + fn available() -> bool; +} + +pub(crate) trait UnsafeBitPacker { + const BLOCK_LEN: usize; + + unsafe fn compress(decompressed: &[u32], compressed: &mut [u8], num_bits: u8) -> usize; + + unsafe fn compress_sorted( + initial: u32, + decompressed: &[u32], + compressed: &mut [u8], + num_bits: u8, + ) -> usize; + + unsafe fn compress_strictly_sorted( + initial: Option, + decompressed: &[u32], + compressed: &mut [u8], + num_bits: u8, + ) -> usize; + + unsafe fn decompress(compressed: &[u8], decompressed: &mut [u32], num_bits: u8) -> usize; + + unsafe fn decompress_sorted( + initial: u32, + compressed: &[u8], + decompressed: &mut [u32], + num_bits: u8, + ) -> usize; + + unsafe fn decompress_strictly_sorted( + initial: Option, + compressed: &[u8], + decompressed: &mut [u32], + num_bits: u8, + ) -> usize; + + unsafe fn num_bits(decompressed: &[u32]) -> u8; + + unsafe fn num_bits_sorted(initial: u32, decompressed: &[u32]) -> u8; + + unsafe fn num_bits_strictly_sorted(initial: Option, decompressed: &[u32]) -> u8; +} + +/// Block bitpacker for fixed-size `u32` blocks. +/// +/// Implementations own runtime SIMD dispatch and use caller-provided buffers. +/// Packed bytes are stable for a given implementation and bit width. +pub trait BitPacker: Sized + Clone + Copy { + /// Number of `u32` values in one physical block. + const BLOCK_LEN: usize; + + /// Select the best supported implementation for the current CPU. + /// + /// Lance uses SIMD backends when available and falls back to a scalar + /// backend otherwise, matching the existing allocation-free call shape used + /// by the upstream `bitpacking` crate. + fn new() -> Self; + + /// Compress one full block of raw values into `compressed`. + fn compress(&self, decompressed: &[u32], compressed: &mut [u8], num_bits: u8) -> usize; + + /// Delta-compress one full non-decreasing block into `compressed`. + fn compress_sorted( + &self, + initial: u32, + decompressed: &[u32], + compressed: &mut [u8], + num_bits: u8, + ) -> usize; + + /// Delta-compress one full strictly increasing block into `compressed`. + fn compress_strictly_sorted( + &self, + initial: Option, + decompressed: &[u32], + compressed: &mut [u8], + num_bits: u8, + ) -> usize; + + /// Decompress one raw block into `decompressed`. + fn decompress(&self, compressed: &[u8], decompressed: &mut [u32], num_bits: u8) -> usize; + + /// Decompress one delta-compressed non-decreasing block. + fn decompress_sorted( + &self, + initial: u32, + compressed: &[u8], + decompressed: &mut [u32], + num_bits: u8, + ) -> usize; + + /// Decompress one delta-compressed strictly increasing block. + fn decompress_strictly_sorted( + &self, + initial: Option, + compressed: &[u8], + decompressed: &mut [u32], + num_bits: u8, + ) -> usize; + + /// Return the minimum bit width needed to represent a full raw block. + fn num_bits(&self, decompressed: &[u32]) -> u8; + + /// Return the minimum bit width needed to represent deltas in a full block. + fn num_bits_sorted(&self, initial: u32, decompressed: &[u32]) -> u8; + + /// Return the minimum bit width needed to represent strict deltas in a full block. + fn num_bits_strictly_sorted(&self, initial: Option, decompressed: &[u32]) -> u8; + + /// Return the byte size of one compressed block at `num_bits`. + #[must_use] + fn compressed_block_size(num_bits: u8) -> usize { + Self::BLOCK_LEN * num_bits as usize / 8 + } +} + +#[inline] +fn most_significant_bit(value: u32) -> u8 { + (u32::BITS - value.leading_zeros()) as u8 +} diff --git a/rust/compression/bitpacking/src/lib.rs b/rust/compression/bitpacking/src/lib.rs index 0101c4a1df0..c6aa6d75ad0 100644 --- a/rust/compression/bitpacking/src/lib.rs +++ b/rust/compression/bitpacking/src/lib.rs @@ -16,6 +16,10 @@ use arrayref::{array_mut_ref, array_ref}; use core::mem::size_of; +mod bitpacker_internal; + +pub use bitpacker_internal::{BitPacker, BitPacker4x, BitPacker8x}; + pub const FL_ORDER: [usize; 8] = [0, 4, 2, 6, 1, 5, 3, 7]; pub trait FastLanes: Sized + Copy { @@ -1709,6 +1713,7 @@ pack_64!(pack_64_64, 64); #[cfg(test)] mod test { use super::*; + use bitpacking::{BitPacker as ExternalBitPacker, BitPacker4x as ExternalBitPacker4x}; use core::array; // a fast random number generator pub struct XorShift { @@ -1730,6 +1735,104 @@ mod test { } } + fn mask_for_width(width: u8) -> u32 { + match width { + 0 => 0, + 32 => u32::MAX, + _ => (1u32 << width) - 1, + } + } + + fn raw_bitpacker4x_case(width: u8, seed: u64) -> Vec { + let mask = mask_for_width(width); + let mut rng = XorShift::new(seed); + (0..BitPacker4x::BLOCK_LEN) + .map(|idx| match seed % 4 { + 0 => 0, + 1 => mask, + 2 => idx as u32 & mask, + _ => (rng.next() as u32) & mask, + }) + .collect() + } + + fn sorted_bitpacker4x_case(width: u8, seed: u64) -> (u32, Vec) { + if width == 0 { + return (17, vec![17; BitPacker4x::BLOCK_LEN]); + } + if width == 32 { + return (0, vec![u32::MAX; BitPacker4x::BLOCK_LEN]); + } + + let mask = mask_for_width(width).min(127); + let mut rng = XorShift::new(seed); + let mut current = 17u32; + let values = (0..BitPacker4x::BLOCK_LEN) + .map(|_| { + current += (rng.next() as u32) & mask; + current + }) + .collect(); + (17, values) + } + + #[test] + fn test_bitpacker4x_raw_compatible_with_external_bitpacking() { + let ours = BitPacker4x::new(); + let external = ExternalBitPacker4x::new(); + + for width in 0..=32 { + for seed in [0, 1, 2, 123456789] { + let values = raw_bitpacker4x_case(width, seed); + assert_eq!(ours.num_bits(&values), external.num_bits(&values)); + + let mut actual = vec![0u8; BitPacker4x::compressed_block_size(width)]; + let actual_len = ours.compress(&values, &mut actual, width); + + let mut expected = vec![0u8; ExternalBitPacker4x::compressed_block_size(width)]; + let expected_len = external.compress(&values, &mut expected, width); + + assert_eq!(actual_len, expected_len); + assert_eq!(actual, expected, "width {width} seed {seed}"); + + let mut decoded = vec![0u32; BitPacker4x::BLOCK_LEN]; + let consumed = ours.decompress(&actual, &mut decoded, width); + assert_eq!(consumed, actual_len); + assert_eq!(decoded, values); + } + } + } + + #[test] + fn test_bitpacker4x_sorted_compatible_with_external_bitpacking() { + let ours = BitPacker4x::new(); + let external = ExternalBitPacker4x::new(); + + for width in 0..=32 { + for seed in [0, 1, 2, 123456789] { + let (initial, values) = sorted_bitpacker4x_case(width, seed); + assert_eq!( + ours.num_bits_sorted(initial, &values), + external.num_bits_sorted(initial, &values) + ); + + let mut actual = vec![0u8; BitPacker4x::compressed_block_size(width)]; + let actual_len = ours.compress_sorted(initial, &values, &mut actual, width); + + let mut expected = vec![0u8; ExternalBitPacker4x::compressed_block_size(width)]; + let expected_len = external.compress_sorted(initial, &values, &mut expected, width); + + assert_eq!(actual_len, expected_len); + assert_eq!(actual, expected, "width {width} seed {seed}"); + + let mut decoded = vec![0u32; BitPacker4x::BLOCK_LEN]; + let consumed = ours.decompress_sorted(initial, &actual, &mut decoded, width); + assert_eq!(consumed, actual_len); + assert_eq!(decoded, values); + } + } + } + // a macro version of this function generalize u8, u16, u32, u64 takes very long time for a test build, so I // write it for each type separately fn pack_unpack_u8(bit_width: usize) { @@ -2025,6 +2128,10 @@ mod test { for value in &mut values { *value = rng.next(); } + } else { + for value in &mut values { + *value = rng.next() % (1 << bit_width); + } } let mut packed = vec![0; 1024 * bit_width / u64::T]; unsafe { diff --git a/rust/compression/fsst/Cargo.toml b/rust/compression/fsst/Cargo.toml index da5d8f01d04..7056896e3b9 100644 --- a/rust/compression/fsst/Cargo.toml +++ b/rust/compression/fsst/Cargo.toml @@ -16,7 +16,6 @@ arrow-array.workspace = true rand.workspace = true [dev-dependencies] -arrow-array.workspace = true test-log.workspace = true tokio.workspace = true diff --git a/rust/compression/fsst/examples/benchmark.rs b/rust/compression/fsst/examples/benchmark.rs index c442243e112..f71abeefedd 100644 --- a/rust/compression/fsst/examples/benchmark.rs +++ b/rust/compression/fsst/examples/benchmark.rs @@ -56,7 +56,10 @@ fn benchmark(file_path: &str) { let mut decompression_out_bufs = vec![]; let mut decompression_out_offsets_bufs = vec![]; for _ in 0..TEST_NUM { - let this_decom_out_buf = vec![0u8; BUFFER_SIZE * 3]; + // `decompress` requires the output buffer to be at least 8x the compressed input (a 1-byte + // code can expand to an 8-byte symbol). The compressed buffer is at most `BUFFER_SIZE`, so + // `BUFFER_SIZE * 8` is a safe upper bound. + let this_decom_out_buf = vec![0u8; BUFFER_SIZE * 8]; let this_decom_out_offsets_buf = vec![0i32; BUFFER_SIZE * 3]; decompression_out_bufs.push(this_decom_out_buf); decompression_out_offsets_bufs.push(this_decom_out_offsets_buf); diff --git a/rust/compression/fsst/src/fsst.rs b/rust/compression/fsst/src/fsst.rs index 0a2bf1d03d7..d00a6ed806b 100644 --- a/rust/compression/fsst/src/fsst.rs +++ b/rust/compression/fsst/src/fsst.rs @@ -57,6 +57,12 @@ use std::ptr; #[inline] fn fsst_unaligned_load_unchecked(v: *const u8) -> u64 { + // SAFETY: the caller must guarantee that `v` points to at least 8 readable bytes. All callers + // uphold this: `compress_bulk` loads from a 520-byte stack buffer at an offset < 511 (leaving + // >= 8 bytes), `build_symbol_table` guards the load with `word.len() > 7 && curr < word.len() - 7`, + // `find_longest_symbol_from_char_slice` copies into a stack `[u8; 8]` before loading, and + // `FsstDecoder::init` reads symbols from a `symbol_table` buffer already validated to be + // exactly `FSST_SYMBOL_TABLE_SIZE` bytes. unsafe { ptr::read_unaligned(v as *const u64) } } @@ -812,12 +818,31 @@ fn decompress_bulk( ) -> io::Result<()> { let symbols = decoder.symbols; let lens = decoder.lens; + // SAFETY invariant shared by every `unsafe` block in this closure: + // - `out` is sized to at least 8x `compressed_strs` (checked in `FsstDecoder::init`, which the + // sole public entry point always runs before reaching this function). Each code advances + // `out_curr` by `lens[code]`, which is 1..=8 for a well-formed symbol table, and each + // consumed input byte yields at most 8 output bytes, so `out_curr + 8 <= out.len()` at every + // 8-byte write, including the final one. This is why we can `write_unaligned` a full 8-byte + // word per code and advance by only the length. + // NOTE: `lens` is loaded verbatim from the (untrusted) symbol table and is NOT re-validated + // to be <= 8 on decode, and offsets (below) are likewise trusted. A corrupted table or + // offset buffer can violate these bounds; callers must supply structures produced by + // `compress` (or otherwise trusted). Hardening the decoder against corrupt input is a + // separate concern, not addressed here. + // - The only unchecked read is `read_unaligned::`, gated by `in_curr + 4 <= in_end`; the + // scalar paths use bounds-checked indexing. `in_end` is a caller-provided offset into + // `compressed_strs`; the read is sound only if `in_end <= compressed_strs.len()`, which is a + // trusted precondition (holds for encoder-produced offsets; not validated here). let mut decompress = |mut in_curr: usize, in_end: usize, out_curr: &mut usize| { // Do SIMD operation here by 4 bytes while in_curr + 4 <= in_end { let next_block; let mut code; let mut len; + // SAFETY: the loop guard proves `in_curr + 4 <= in_end`. Per the closure-level + // invariant, `in_end <= compressed_strs.len()` is a trusted precondition (not checked + // here), so the 4-byte read is in bounds for well-formed input. unsafe { next_block = ptr::read_unaligned(compressed_strs.as_ptr().add(in_curr) as *const u32); @@ -983,6 +1008,10 @@ fn decompress_bulk( if in_curr < in_end { // last code cannot be an escape code let code = compressed_strs[in_curr] as usize; + // SAFETY: see the closure-level invariant. This is the final write and has no + // subsequent write to cover its slack, so it is the tightest case: `out_curr` is at + // most 8*(consumed_input_bytes - 1), and the 8-byte store lands within `out.len()` + // precisely because the caller sized `out` to 8x the input. unsafe { let src = symbols[code]; ptr::write_unaligned(out.as_mut_ptr().add(*out_curr) as *mut u64, src); @@ -1188,8 +1217,19 @@ impl FsstDecoder { } self.decoder_switch_on = (st_info & (1 << 24)) != 0; - // when decoder_switch_on is true, we make sure the out_buf is at least 3 times the size of the in_buf, - if self.decoder_switch_on && in_buf.len() * 3 > out_buf.len() { + // A single 1-byte code can decode to a symbol of up to MAX_SYMBOL_LENGTH (8) bytes, so the + // decoded output can be up to 8x the input. `decompress_bulk` also relies on this bound: it + // writes a full 8-byte word per code (advancing only by the symbol length), so the output + // buffer must be large enough that even the final write stays in bounds. Require out_buf to + // be at least 8x in_buf. `checked_mul` guards against `in_buf.len() * 8` wrapping on 32-bit + // targets (an input >= 512 MiB would otherwise bypass the check); treat overflow as too + // small. + if self.decoder_switch_on + && in_buf + .len() + .checked_mul(8) + .is_none_or(|needed| needed > out_buf.len()) + { return Err(io::Error::new( io::ErrorKind::InvalidInput, "output buffer too small for FSST decoder", @@ -1288,8 +1328,10 @@ pub fn compress( // the following 32 bits after FSST_MAGIC contains information about FSST encoding, such as decoder_switch_on, suffix_lim, terminator, n_symbols // when the decoder_switch_on is off in the in_buf header, `decompress` first make sure the out_buf is at least the same size as the in_buf, then simply copy the // input data to the output -// when the decoder_switch_on is on, `decompress` first make sure the out_buf is at least 3 times the size of the in_buf, then start decoding the -// data using the symbol table +// when the decoder_switch_on is on, `decompress` first make sure the out_buf is at least 8 times the size of the in_buf, then start decoding the +// data using the symbol table. The 8x bound is required for correctness: a 1-byte code can expand to +// an 8-byte symbol, and the decode loop writes a full 8-byte word per code, so a smaller buffer can +// be written out of bounds. // the out_offsets_buf should be at least the same size as the in_offsets_buf, otherwise an error is returned // the symbol_table is the same symbol table created by `compression` pub fn decompress( @@ -1644,4 +1686,59 @@ But exactly how the acquaintance and friendship came about, we cannot say."; ); } } + + // Build a genuinely FSST-compressed (decoder_switch_on) buffer to exercise the decode-side + // output-buffer size contract. Returns (symbol_table, compressed_bytes, compressed_offsets). + fn compress_paragraph() -> ([u8; FSST_SYMBOL_TABLE_SIZE], Vec, Vec) { + let test_input = TEST_PARAGRAPH.repeat((1024 * 1024) / TEST_PARAGRAPH.len()); + let lines_vec = test_input.lines().collect::>(); + let string_array = StringArray::from(lines_vec); + let mut compress_output_buf: Vec = vec![0; string_array.value_data().len()]; + let mut compress_offset_buf: Vec = vec![0; string_array.value_offsets().len()]; + let mut symbol_table = [0; FSST_SYMBOL_TABLE_SIZE]; + compress( + symbol_table.as_mut(), + string_array.value_data(), + string_array.value_offsets(), + &mut compress_output_buf, + &mut compress_offset_buf, + ) + .unwrap(); + (symbol_table, compress_output_buf, compress_offset_buf) + } + + // The decoder writes a full 8-byte word per code, so the output buffer must be at least 8x the + // compressed input. A buffer sized below 8x must be rejected rather than written out of bounds. + #[test_log::test(tokio::test)] + async fn test_decompress_rejects_undersized_output_buffer() { + let (symbol_table, compressed, compressed_offsets) = compress_paragraph(); + // Sanity check: this input actually engaged FSST compression (decoder_switch_on). + let st_info = u64::from_ne_bytes(symbol_table[..8].try_into().unwrap()); + assert!(st_info & (1 << 24) != 0, "expected decoder_switch_on input"); + + // One byte short of the 8x requirement must be rejected. + let mut too_small = vec![0u8; compressed.len() * 8 - 1]; + let mut out_offsets = vec![0i32; compressed_offsets.len()]; + let err = decompress( + &symbol_table, + &compressed, + &compressed_offsets, + &mut too_small, + &mut out_offsets, + ) + .unwrap_err(); + assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput); + + // Exactly 8x is the tight bound and must succeed. + let mut exact = vec![0u8; compressed.len() * 8]; + let mut out_offsets = vec![0i32; compressed_offsets.len()]; + decompress( + &symbol_table, + &compressed, + &compressed_offsets, + &mut exact, + &mut out_offsets, + ) + .unwrap(); + } } diff --git a/rust/lance-arrow/Cargo.toml b/rust/lance-arrow/Cargo.toml index afbc9c26ed3..21513e638d8 100644 --- a/rust/lance-arrow/Cargo.toml +++ b/rust/lance-arrow/Cargo.toml @@ -21,6 +21,7 @@ arrow-ipc = { workspace = true } arrow-ord = { workspace = true } arrow-schema = { workspace = true } arrow-select = { workspace = true } +bytemuck = { workspace = true } bytes = { workspace = true } futures = { workspace = true } half = { workspace = true } diff --git a/rust/lance-arrow/src/bfloat16.rs b/rust/lance-arrow/src/bfloat16.rs index 3049be1117e..74c7f259a40 100644 --- a/rust/lance-arrow/src/bfloat16.rs +++ b/rust/lance-arrow/src/bfloat16.rs @@ -7,7 +7,7 @@ use std::fmt::Formatter; use std::slice; use arrow_array::{Array, FixedSizeBinaryArray, builder::BooleanBufferBuilder}; -use arrow_buffer::MutableBuffer; +use arrow_buffer::{Buffer, MutableBuffer}; use arrow_data::ArrayData; use arrow_schema::{ArrowError, DataType, Field as ArrowField}; use half::bf16; @@ -144,6 +144,11 @@ impl FromIterator> for BFloat16Array { .len(len) .add_buffer(buffer.into()) .null_bit_buffer(null_buffer); + // SAFETY: the value buffer contains exactly `2 * len` bytes (two bytes + // pushed per iteration of the loop above, including the zero-fill for + // null slots), which matches the `FixedSizeBinary(2)` storage layout. + // The null bit buffer, when present, has `len` bits appended above, so + // its length covers the array's logical range. let array_data = unsafe { array_data.build_unchecked() }; Self { inner: FixedSizeBinaryArray::from(array_data), @@ -159,17 +164,20 @@ impl FromIterator for BFloat16Array { impl From> for BFloat16Array { fn from(data: Vec) -> Self { - let mut buffer = MutableBuffer::with_capacity(data.len() * 2); - - let bytes = data.iter().flat_map(|val| { - let bytes = val.to_bits().to_le_bytes(); - bytes.to_vec() - }); - - buffer.extend(bytes); + let len = data.len(); + // Zero-copy: `bf16` is `#[repr(transparent)]` over `u16` and derives + // `bytemuck::Pod`, so `cast_vec` reinterprets the allocation in place — + // no per-element copy or heap alloc. The crate-root `compile_error!` + // pins `target_endian = "little"`, so the resulting bytes match the + // `FixedSizeBinary(2)` on-disk order Lance writes elsewhere. + let raw: Vec = bytemuck::cast_vec(data); let array_data = ArrayData::builder(DataType::FixedSizeBinary(2)) - .len(data.len()) - .add_buffer(buffer.into()); + .len(len) + .add_buffer(Buffer::from_vec(raw)); + // SAFETY: the value buffer contains exactly `2 * len` bytes — one + // `u16` per element after the layout-compatible cast — matching the + // `FixedSizeBinary(2)` storage layout. No null buffer is attached, so + // every element is logically valid. let array_data = unsafe { array_data.build_unchecked() }; Self { inner: FixedSizeBinaryArray::from(array_data), @@ -268,12 +276,64 @@ mod from_arrow { impl FloatArray for FixedSizeBinaryArray { type FloatType = BFloat16Type; + /// Returns the underlying `bf16` values as a borrowed slice. + /// + /// # Preconditions + /// + /// - `value_length()` must be 2 (the `FixedSizeBinary(2)` storage shape + /// used by [`BFloat16Array`]). Asserted at entry. + /// - The value buffer must be at least 2-byte aligned. Lance's in-tree + /// constructors always satisfy this: value buffers are built either via + /// `MutableBuffer` (aligned to arrow-buffer's `ALIGNMENT` constant, ≥32 + /// bytes) or via `Buffer::from_vec::` (aligned to `align_of::()` + /// == 2); both meet `bf16`'s 2-byte requirement. Externally-built + /// `FixedSizeBinaryArray`s arriving via FFI, IPC, or + /// `Buffer::from_custom_allocation` are not required by arrow-rs to be + /// aligned beyond a single byte; passing one to this method violates the + /// precondition. A `debug_assert` below catches such inputs in debug and + /// test builds. + /// + /// # Endianness + /// + /// `lance-arrow` is gated on `target_endian = "little"` at the crate root, + /// so this method always returns values in the same byte order Lance writes + /// (see [`BFloat16Array::value`] and the [`FromIterator`] impls). fn as_slice(&self) -> &[bf16] { assert_eq!( self.value_length(), 2, "BFloat16 arrays must use FixedSizeBinary(2) storage" ); + debug_assert_eq!( + (self.value_data().as_ptr() as usize) % std::mem::align_of::(), + 0, + "BFloat16 value buffer must be at least 2-byte aligned" + ); + // SAFETY: + // - The assert above pins `value_size == 2`, so `value_data().len() / 2` + // equals the array's logical element count. + // `FixedSizeBinaryArray::From` constructs its value buffer + // as `buffers[0].slice_with_length(offset * 2, len * 2)` (arrow-array + // `fixed_size_binary_array.rs`), so `value_data()` already returns + // the offset-adjusted slice. Do not replace `value_data()` with an + // accessor that returns the un-sliced backing buffer. + // - `bf16` is `#[repr(transparent)]` over `u16` (size 2, alignment 2); + // every `u16` bit pattern is a valid `bf16`, so any byte content + // yields a defined value — never UB. + // - Alignment is the caller's responsibility per the precondition + // documented above. The `debug_assert_eq!` immediately preceding this + // block catches violations in debug and test builds only — release + // builds rely on callers honoring the precondition. arrow-rs + // declares `FixedSizeBinary(n)`'s + // `BufferSpec::FixedWidth { alignment: align_of::() == 1 }` + // (arrow-data `data.rs`), so arrow-rs alone does not guarantee + // 2-byte alignment. Lance's in-tree construction paths build value + // buffers via `MutableBuffer` (arrow-buffer `ALIGNMENT` constant, + // ≥32 bytes) or `Buffer::from_vec::` (2-byte aligned), both of + // which satisfy `bf16`'s 2-byte requirement. + // - The returned slice borrows from `self`; the underlying ref-counted, + // immutable Arrow buffer cannot be mutated or freed for the slice's + // lifetime. unsafe { slice::from_raw_parts( self.value_data().as_ptr() as *const bf16, @@ -301,6 +361,16 @@ mod tests { assert_eq!(array, array2); assert_eq!(array.len(), 3); + // Pin the raw little-endian bytes emitted by `From>` (rewritten to + // reinterpret the Vec via `bytemuck::cast_vec`), so a layout/byte-order + // regression is caught directly rather than only through Debug formatting. + // bf16 is the high 16 bits of the f32: 1.0->0x3F80, 2.0->0x4000, 3.0->0x4040. + let inner = array2.clone().into_inner(); + let raw_bytes: Vec = (0..inner.len()) + .flat_map(|i| inner.value(i).to_vec()) + .collect(); + assert_eq!(raw_bytes, vec![0x80, 0x3F, 0x00, 0x40, 0x40, 0x40]); + let expected_fmt = "BFloat16Array\n[\n 1.0,\n 2.0,\n 3.0,\n]"; assert_eq!(expected_fmt, format!("{:?}", array)); diff --git a/rust/lance-arrow/src/deepcopy.rs b/rust/lance-arrow/src/deepcopy.rs index b747b24d466..afeb853350b 100644 --- a/rust/lance-arrow/src/deepcopy.rs +++ b/rust/lance-arrow/src/deepcopy.rs @@ -14,6 +14,15 @@ pub fn deep_copy_buffer(buffer: &Buffer) -> Buffer { pub fn deep_copy_nulls(nulls: Option<&NullBuffer>) -> Option { let nulls = nulls?; let bit_buffer = deep_copy_buffer(nulls.inner().inner()); + // SAFETY: `null_count` is taken from the source `NullBuffer`, which already + // upheld `NullBuffer::new_unchecked`'s invariant — the unset-bit count over + // the logical bit slice `[bit_offset, bit_offset + bit_len)`. `NullBuffer::slice` + // adjusts only `BooleanBuffer::bit_offset` / `bit_len` and never byte-advances + // the inner `Buffer`, so `deep_copy_buffer` (which copies the source `Buffer`'s + // `as_slice()` view from byte 0) reproduces the exact bit pattern at the same + // bit offsets; the unset-bit count is therefore preserved. `BooleanBuffer::new` + // panics (does not UB) if `bit_offset + bit_len > 8 * buffer.len()`, and the + // copy has the same length, so that check still passes. Some(unsafe { NullBuffer::new_unchecked( BooleanBuffer::new(bit_buffer, nulls.offset(), nulls.len()), @@ -37,6 +46,19 @@ pub fn deep_copy_array_data(data: &ArrayData) -> ArrayData { .iter() .map(deep_copy_array_data) .collect::>(); + // SAFETY: `build_unchecked` inherits `ArrayData::new_unchecked`'s contract — + // `(data_type, len, offset, nulls, buffers, child_data)` must form a valid + // Arrow array. This call reproduces `data` structurally: `data_type`, `len`, + // and `offset` are forwarded unchanged; each buffer is replaced by a byte- + // identical copy of its offset-applied `as_slice()` view (the output buffer + // is `MutableBuffer`-allocated, at least as aligned as the source); `nulls` + // is deep-copied with the same bit offset/length and unset-bit count (see + // `deep_copy_nulls`); `child_data` is recursively cloned with the same + // guarantee. Every value-level invariant the source upheld — UTF-8 validity, + // monotonic offsets, in-bounds dictionary indices, run-end monotonicity, + // struct child-length matching — therefore transfers to the copy. If the + // source `ArrayData` was itself constructed via `new_unchecked` with an + // invalid payload, this function faithfully reproduces that invalidity. unsafe { ArrayDataBuilder::new(data_type) .len(len) diff --git a/rust/lance-arrow/src/floats.rs b/rust/lance-arrow/src/floats.rs index 054f4418e0b..ad5f3768980 100644 --- a/rust/lance-arrow/src/floats.rs +++ b/rust/lance-arrow/src/floats.rs @@ -184,6 +184,22 @@ pub trait FloatArray: Array + Clone + 'static { type FloatType: ArrowFloatType; /// Returns a reference to the underlying data as a slice. + /// + /// # Panics + /// + /// Implementations may panic if the array's storage shape does not match + /// the expected element layout. In particular, the `bf16` impl panics if + /// `value_length() != 2` (the `FixedSizeBinary(2)` shape required by + /// `BFloat16Array`). + /// + /// # Preconditions + /// + /// Implementations may impose additional invariants on the underlying + /// buffer. The `bf16` impl requires the value buffer to be at least + /// 2-byte aligned — satisfied automatically by every in-tree Lance + /// constructor, but external callers passing externally-built arrays + /// (FFI, IPC, `Buffer::from_custom_allocation`) must ensure alignment. + /// See the impl's docstring for details. fn as_slice(&self) -> &[T::Native]; /// Construct an array from a vector of values. diff --git a/rust/lance-arrow/src/lib.rs b/rust/lance-arrow/src/lib.rs index 34a67600543..2ac962aa1aa 100644 --- a/rust/lance-arrow/src/lib.rs +++ b/rust/lance-arrow/src/lib.rs @@ -5,6 +5,16 @@ //! //! To improve Arrow-RS ergonomic +#![warn(clippy::undocumented_unsafe_blocks)] + +// lance-arrow reinterprets value bytes as native numeric types in +// `FloatArray::as_slice` for `bf16` (rust/lance-arrow/src/bfloat16.rs), which +// requires the host byte order to match the on-disk byte order Lance writes. +// Lance writes little-endian; building on a big-endian target would silently +// produce wrong numeric values. +#[cfg(not(target_endian = "little"))] +compile_error!("lance-arrow only supports little-endian targets"); + use std::sync::Arc; use std::{collections::HashMap, ptr::NonNull}; @@ -54,6 +64,9 @@ pub const BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY: &str = "lance-encoding:blob-dedicated-size-threshold"; /// Metadata key for overriding the inline blob size threshold (in bytes) pub const BLOB_INLINE_SIZE_THRESHOLD_META_KEY: &str = "lance-encoding:blob-inline-size-threshold"; +/// Metadata key for overriding the maximum size (in bytes) of a packed blob sidecar file +pub const BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY: &str = + "lance-encoding:blob-pack-file-size-threshold"; type Result = std::result::Result; @@ -291,7 +304,7 @@ impl FixedSizeListArrayExt for FixedSizeListArray { field.is_nullable(), )), *size, - Arc::new(Float32Array::from_iter_values( + Arc::new(Float32Array::from_iter( self.values() .as_any() .downcast_ref::() @@ -299,7 +312,7 @@ impl FixedSizeListArrayExt for FixedSizeListArray { "Fail to cast primitive array to Int8Type".to_string(), ))? .into_iter() - .filter_map(|x| x.map(|y| y as f32)), + .map(|x| x.map(|y| y as f32)), )), self.nulls().cloned(), )), @@ -310,7 +323,7 @@ impl FixedSizeListArrayExt for FixedSizeListArray { field.is_nullable(), )), *size, - Arc::new(Float32Array::from_iter_values( + Arc::new(Float32Array::from_iter( self.values() .as_any() .downcast_ref::() @@ -318,7 +331,7 @@ impl FixedSizeListArrayExt for FixedSizeListArray { "Fail to cast primitive array to Int16Type".to_string(), ))? .into_iter() - .filter_map(|x| x.map(|y| y as f32)), + .map(|x| x.map(|y| y as f32)), )), self.nulls().cloned(), )), @@ -329,7 +342,7 @@ impl FixedSizeListArrayExt for FixedSizeListArray { field.is_nullable(), )), *size, - Arc::new(Float32Array::from_iter_values( + Arc::new(Float32Array::from_iter( self.values() .as_any() .downcast_ref::() @@ -337,7 +350,7 @@ impl FixedSizeListArrayExt for FixedSizeListArray { "Fail to cast primitive array to Int32Type".to_string(), ))? .into_iter() - .filter_map(|x| x.map(|y| y as f32)), + .map(|x| x.map(|y| y as f32)), )), self.nulls().cloned(), )), @@ -348,7 +361,7 @@ impl FixedSizeListArrayExt for FixedSizeListArray { field.is_nullable(), )), *size, - Arc::new(Float64Array::from_iter_values( + Arc::new(Float64Array::from_iter( self.values() .as_any() .downcast_ref::() @@ -356,7 +369,7 @@ impl FixedSizeListArrayExt for FixedSizeListArray { "Fail to cast primitive array to Int64Type".to_string(), ))? .into_iter() - .filter_map(|x| x.map(|y| y as f64)), + .map(|x| x.map(|y| y as f64)), )), self.nulls().cloned(), )), @@ -367,7 +380,7 @@ impl FixedSizeListArrayExt for FixedSizeListArray { field.is_nullable(), )), *size, - Arc::new(Float64Array::from_iter_values( + Arc::new(Float64Array::from_iter( self.values() .as_any() .downcast_ref::() @@ -375,7 +388,7 @@ impl FixedSizeListArrayExt for FixedSizeListArray { "Fail to cast primitive array to UInt8Type".to_string(), ))? .into_iter() - .filter_map(|x| x.map(|y| y as f64)), + .map(|x| x.map(|y| y as f64)), )), self.nulls().cloned(), )), @@ -386,7 +399,7 @@ impl FixedSizeListArrayExt for FixedSizeListArray { field.is_nullable(), )), *size, - Arc::new(Float64Array::from_iter_values( + Arc::new(Float64Array::from_iter( self.values() .as_any() .downcast_ref::() @@ -394,7 +407,7 @@ impl FixedSizeListArrayExt for FixedSizeListArray { "Fail to cast primitive array to UInt32Type".to_string(), ))? .into_iter() - .filter_map(|x| x.map(|y| y as f64)), + .map(|x| x.map(|y| y as f64)), )), self.nulls().cloned(), )), @@ -1009,13 +1022,7 @@ fn merge_list_struct(left: &dyn Array, right: &dyn Array) -> Arc { fn normalize_validity( validity: Option<&arrow_buffer::NullBuffer>, ) -> Option<&arrow_buffer::NullBuffer> { - validity.and_then(|v| { - if v.null_count() == v.len() { - None - } else { - Some(v) - } - }) + validity.filter(|v| v.null_count() != v.len()) } /// Helper function to merge validity buffers from two struct arrays @@ -1372,9 +1379,12 @@ fn merge_with_schema( ); let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls()); + // `trimmed_values` starts at the first used value, so offsets + // must be shifted to match or `ListArray::new` panics when the + // input list was sliced (e.g. from a filtered batch). let merged_list = ListArray::new( child_field.clone(), - left_list.offsets().clone(), + left_list.trimmed_offsets(), merged_values, merged_validity, ); @@ -1399,7 +1409,7 @@ fn merge_with_schema( merge_struct_validity(left_list.nulls(), right_list.nulls()); let merged_list = LargeListArray::new( child_field.clone(), - left_list.offsets().clone(), + left_list.trimmed_offsets(), merged_values, merged_validity, ); @@ -1563,6 +1573,63 @@ mod tests { use arrow_array::{ListArray, StringArray, new_empty_array, new_null_array}; use arrow_buffer::OffsetBuffer; + #[test] + fn test_convert_to_floating_point_preserves_inner_nulls() { + // A FixedSizeList with a null inner element must convert to a + // FixedSizeList with the null kept in place. Dropping it would + // shorten the values array and shift every later element (and, when the + // remaining count is not a multiple of the list size, panic). + let values = Int8Array::from(vec![Some(1), None, Some(3), Some(4)]); + let fsl = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Int8, true)), + 2, + Arc::new(values), + None, + ); + + let converted = fsl.convert_to_floating_point().unwrap(); + + assert_eq!(converted.len(), 2); + let conv_values = converted + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(conv_values.len(), 4); + assert_eq!(conv_values.value(0), 1.0); + assert!(conv_values.is_null(1)); + assert_eq!(conv_values.value(2), 3.0); + assert_eq!(conv_values.value(3), 4.0); + } + + #[test] + fn test_convert_to_floating_point_preserves_inner_nulls_f64_arm() { + // The Float64-producing arms (Int64/UInt8/UInt32) share the same fix as the + // Float32 arms; cover one representative (UInt8 -> Float64) so both branch + // families are exercised. + let values = UInt8Array::from(vec![Some(10u8), None, Some(30), Some(40)]); + let fsl = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::UInt8, true)), + 2, + Arc::new(values), + None, + ); + + let converted = fsl.convert_to_floating_point().unwrap(); + + assert_eq!(converted.len(), 2); + let conv_values = converted + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(conv_values.len(), 4); + assert_eq!(conv_values.value(0), 10.0); + assert!(conv_values.is_null(1)); + assert_eq!(conv_values.value(2), 30.0); + assert_eq!(conv_values.value(3), 40.0); + } + #[test] fn test_merge_recursive() { let a_array = Int32Array::from(vec![Some(1), Some(2), Some(3)]); @@ -2310,6 +2377,118 @@ mod tests { assert_eq!(merged_array.len(), 2); } + #[test] + fn test_merge_with_schema_sliced_list_struct() { + test_merge_with_schema_sliced_list_struct_generic::(); + } + + #[test] + fn test_merge_with_schema_sliced_large_list_struct() { + test_merge_with_schema_sliced_list_struct_generic::(); + } + + // Regression for #6580: merge_with_schema panicked when the left list was a + // sliced view whose offsets did not start at zero (common after a filtered + // scan). Cloning those offsets alongside `trimmed_values` produced offsets + // larger than the trimmed child, panicking in `(Large)ListArray::new`. + fn test_merge_with_schema_sliced_list_struct_generic() { + let make_list_dtype = |item_field: Arc| { + if O::IS_LARGE { + DataType::LargeList(item_field) + } else { + DataType::List(item_field) + } + }; + + // Build a List with two rows of 5 items each, then slice away + // the first row so the remaining list's offsets start at 5, not 0. + let struct_fields_a = Fields::from(vec![Field::new("a", DataType::Int32, true)]); + let left_values = Arc::new(StructArray::new( + struct_fields_a.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef], + None, + )); + let full_list = GenericListArray::::new( + Arc::new(Field::new("item", DataType::Struct(struct_fields_a), true)), + OffsetBuffer::::from_lengths([5, 5]), + left_values, + None, + ); + let sliced_left = full_list.slice(1, 1); + assert_eq!(sliced_left.offsets()[0].as_usize(), 5); + assert_eq!(sliced_left.offsets()[1].as_usize(), 10); + + let struct_fields_b = Fields::from(vec![Field::new("b", DataType::Int32, true)]); + let right_values = Arc::new(StructArray::new( + struct_fields_b.clone(), + vec![Arc::new(Int32Array::from_iter_values(100..105)) as ArrayRef], + None, + )); + let right_list = GenericListArray::::new( + Arc::new(Field::new("item", DataType::Struct(struct_fields_b), true)), + OffsetBuffer::::from_lengths([5]), + right_values, + None, + ); + + let target_item_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])), + true, + )); + let target_fields = Fields::from(vec![Field::new( + "items", + make_list_dtype(target_item_field), + true, + )]); + + let left_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "items", + sliced_left.data_type().clone(), + true, + )])), + vec![Arc::new(sliced_left) as ArrayRef], + ) + .unwrap(); + let right_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "items", + right_list.data_type().clone(), + true, + )])), + vec![Arc::new(right_list) as ArrayRef], + ) + .unwrap(); + + let merged = left_batch + .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec())) + .unwrap(); + + let merged_list = merged + .column_by_name("items") + .unwrap() + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(merged_list.len(), 1); + assert_eq!(merged_list.value_length(0).as_usize(), 5); + let merged_struct = merged_list.values().as_struct(); + assert_eq!(merged_struct.num_columns(), 2); + let a = merged_struct + .column_by_name("a") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + // After shifting offsets to zero, values 5..10 should be first. + let a_vals: Vec = a.iter().map(|v| v.unwrap()).collect(); + assert_eq!(a_vals, vec![5, 6, 7, 8, 9]); + } + #[test] fn test_project_by_schema_list_struct_reorder() { // Test that project_by_schema correctly reorders fields inside List diff --git a/rust/lance-arrow/src/list.rs b/rust/lance-arrow/src/list.rs index 0c24fc579da..06b0fc592cf 100644 --- a/rust/lance-arrow/src/list.rs +++ b/rust/lance-arrow/src/list.rs @@ -23,6 +23,16 @@ pub trait ListArrayExt { /// behaves similarly to `values()` except it slices the array so that it starts at /// the first list offset and ends at the last list offset. fn trimmed_values(&self) -> Arc; + /// The offset type of the underlying list array. + type Offset: OffsetSizeTrait; + /// Returns offsets shifted so the first offset is zero, matching + /// [`Self::trimmed_values`]. + /// + /// Sliced list arrays (e.g. a filtered batch) keep offsets that reference the + /// original values buffer, so combining them with trimmed values produces + /// offsets that exceed the values length. Use this together with + /// `trimmed_values` when constructing a new list array. + fn trimmed_offsets(&self) -> OffsetBuffer; } impl ListArrayExt for GenericListArray { @@ -90,6 +100,20 @@ impl ListArrayExt for GenericListArray .unwrap_or(0); self.values().slice(first_value, last_value - first_value) } + + type Offset = OffsetSize; + + fn trimmed_offsets(&self) -> OffsetBuffer { + let offsets = self.offsets(); + let Some(&first) = offsets.first() else { + return offsets.clone(); + }; + if first == OffsetSize::zero() { + return offsets.clone(); + } + let shifted: Vec = offsets.iter().map(|&o| o - first).collect(); + OffsetBuffer::new(ScalarBuffer::from(shifted)) + } } #[cfg(test)] diff --git a/rust/lance-core/Cargo.toml b/rust/lance-core/Cargo.toml index 7f956c70430..2f6183be8a1 100644 --- a/rust/lance-core/Cargo.toml +++ b/rust/lance-core/Cargo.toml @@ -55,6 +55,10 @@ proptest.workspace = true rstest.workspace = true [features] +# Capture Rust backtraces in error types. When disabled (the default), +# the backtrace field is zero-sized with no overhead. At runtime, capture +# is still gated by RUST_BACKTRACE=1. +backtrace = [] datafusion = ["dep:datafusion-common", "dep:datafusion-sql"] [lints] diff --git a/rust/lance-core/src/cache/mod.rs b/rust/lance-core/src/cache/mod.rs index 07038c6e9d5..4f93f261bea 100644 --- a/rust/lance-core/src/cache/mod.rs +++ b/rust/lance-core/src/cache/mod.rs @@ -688,6 +688,39 @@ mod tests { assert!(cache.size_bytes().await <= capacity); } + #[tokio::test] + async fn test_cache_weighs_key_footprint() { + // Weighted size charges the key's unique bytes, not just the value. + let cache = LanceCache::with_capacity(usize::MAX); + let key = "k".repeat(10_000); + let value = Arc::new(vec![1_i32]); + let expected = + std::mem::size_of::() + key.len() + cache_entry_size(&*value); + cache + .insert_with_key(&TestKey::>::new(&key), value) + .await; + assert_eq!(cache.size_bytes().await, expected); + } + + #[tokio::test] + async fn test_cache_shared_prefix_not_charged_per_entry() { + // The shared prefix contributes nothing per entry (it isn't freed on a + // single eviction); only struct + unique key + value are charged. + let cache = LanceCache::with_capacity(usize::MAX).with_key_prefix(&"p".repeat(10_000)); + for i in 0..100 { + cache + .insert_with_key( + &TestKey::>::new(&i.to_string()), + Arc::new(vec![1_i32]), + ) + .await; + } + let value_cost = cache_entry_size(&vec![1_i32]); + let key_bytes: usize = (0..100).map(|i| i.to_string().len()).sum(); + let expected = 100 * (std::mem::size_of::() + value_cost) + key_bytes; + assert_eq!(cache.size_bytes().await, expected); + } + #[tokio::test] async fn test_cache_trait_objects() { #[derive(Debug, DeepSizeOf)] diff --git a/rust/lance-core/src/cache/moka.rs b/rust/lance-core/src/cache/moka.rs index a3956c1720c..fd86b064f6b 100644 --- a/rust/lance-core/src/cache/moka.rs +++ b/rust/lance-core/src/cache/moka.rs @@ -20,6 +20,12 @@ struct MokaCacheEntry { size_bytes: usize, } +/// Per-entry key cost for eviction: the struct plus the unique `key` bytes. +/// Excludes the shared `prefix` `Arc`, which isn't freed per eviction. +fn key_footprint(key: &InternalCacheKey) -> usize { + std::mem::size_of::() + key.key().len() +} + /// Default [`CacheBackend`] backed by a [moka](https://crates.io/crates/moka) cache. /// /// Provides weighted-capacity eviction and concurrent-load deduplication @@ -40,7 +46,12 @@ impl MokaCacheBackend { pub fn with_capacity(capacity: usize) -> Self { let cache = moka::future::Cache::builder() .max_capacity(capacity as u64) - .weigher(|_, v: &MokaCacheEntry| v.size_bytes.try_into().unwrap_or(u32::MAX)) + .weigher(|key: &InternalCacheKey, entry: &MokaCacheEntry| { + key_footprint(key) + .saturating_add(entry.size_bytes) + .try_into() + .unwrap_or(u32::MAX) + }) .support_invalidation_closures() .build(); Self { cache } @@ -148,6 +159,9 @@ impl CacheBackend for MokaCacheBackend { // Iterate rather than using `weighted_size()` because moka's // weighted_size can be stale without `run_pending_tasks()`, which // is async and can't be called from this synchronous context. - self.cache.iter().map(|(_, v)| v.size_bytes).sum() + self.cache + .iter() + .map(|(key, entry)| key_footprint(key.as_ref()) + entry.size_bytes) + .sum() } } diff --git a/rust/lance-core/src/datatypes/field.rs b/rust/lance-core/src/datatypes/field.rs index 9f06d421949..d5eb89dccb0 100644 --- a/rust/lance-core/src/datatypes/field.rs +++ b/rust/lance-core/src/datatypes/field.rs @@ -58,6 +58,8 @@ pub const LANCE_UNENFORCED_CLUSTERING_KEY_POSITION: &str = /// The value should be non-negative i32 value. Any negative value will be seen as -1. pub const LANCE_FIELD_ID_KEY: &str = "lance:field_id"; +const PACKED_KEYS: [&str; 2] = ["packed", "lance-encoding:packed"]; + fn has_blob_v2_extension(field: &ArrowField) -> bool { field .metadata() @@ -269,24 +271,15 @@ impl Field { } pub fn apply_projection(&self, projection: &Projection) -> Option { - // Map fields encode their physical layout as a single child entries - // struct (`Struct`) whose presence is required for the - // parent to be readable — we never want to filter into that subtree. - // But the parent field itself is still subject to selection: if the - // caller didn't ask for this Map column, drop it like any other - // non-selected leaf. Without this early return the unconditional - // children clone would keep `children.is_empty() == false` forever - // and every Map column in the schema would survive every projection, - // pulling tens-of-bytes-per-row of unrelated data through downstream - // operators (notably `SortExec` in scalar-index training, where it - // was responsible for >100 GiB external-sort spills on real-world - // tables). - if self.logical_type.is_map() && !projection.contains_field_id(self.id) { + // Maps and blob descriptors are atomic physical layouts. Map children + // must remain together, while projected blob descriptor children may + // have synthetic IDs that cannot be selected independently. + let is_atomic_layout = self.logical_type.is_map() || self.is_blob(); + if is_atomic_layout && !projection.contains_field_id(self.id) { return None; } - let children = if self.logical_type.is_map() { - // Map field is selected: keep all children intact. + let children = if is_atomic_layout { self.children.clone() } else { self.children @@ -549,6 +542,20 @@ impl Field { .get(ARROW_EXT_NAME_KEY) .map(|name| name == BLOB_V2_EXT_NAME) .unwrap_or(false) + || self.is_blob_v2_descriptor() + } + + fn is_blob_v2_descriptor(&self) -> bool { + self.metadata.contains_key(BLOB_META_KEY) + && self.logical_type == BLOB_V2_DESC_LANCE_FIELD.logical_type + && self.children.len() == BLOB_V2_DESC_LANCE_FIELD.children.len() + && self + .children + .iter() + .zip(BLOB_V2_DESC_LANCE_FIELD.children.iter()) + .all(|(child, expected)| { + child.name == expected.name && child.data_type() == expected.data_type() + }) } // Blob columns intentionally have two schema representations: @@ -575,6 +582,31 @@ impl Field { } } + /// Convert a blob field to the materialized binary payload view. + /// + /// The field keeps its name and id but uses `LargeBinary` with no children. + /// Blob v2 fields retain their extension marker internally so scan planning + /// can recognize the binary view before exposing a plain Arrow binary field. + pub fn binary_blob_mut(&mut self) { + if !self.is_blob() { + return; + } + let is_blob_v2 = self.is_blob_v2(); + + self.logical_type = LogicalType::try_from(&DataType::LargeBinary) + .expect("LargeBinary is always a valid logical type"); + self.children.clear(); + self.encoding = Some(Encoding::VarBinary); + if is_blob_v2 { + self.metadata.remove(BLOB_META_KEY); + for key in PACKED_KEYS { + self.metadata.remove(key); + } + self.metadata + .insert(ARROW_EXT_NAME_KEY.to_string(), BLOB_V2_EXT_NAME.to_string()); + } + } + /// Convert blob v2 fields in this field tree to their descriptor view. pub fn unload_blobs_recursive(&mut self) { if self.is_blob_v2() { @@ -812,6 +844,13 @@ impl Field { } if self.is_blob() != other.is_blob() { + if ignore_types { + return Ok(if self.id >= 0 { + self.clone() + } else { + other.clone() + }); + } return Err(Error::arrow(format!( "Attempt to intersect blob and non-blob field: {}", self.name @@ -847,7 +886,7 @@ impl Field { .iter() .filter_map(|c| { if let Some(other_child) = other.child(&c.name) { - let intersection = c.intersection(other_child).ok()?; + let intersection = c.do_intersection(other_child, ignore_types).ok()?; Some(intersection) } else { None @@ -1038,7 +1077,6 @@ impl Field { // Check if field has metadata `packed` set to true, this check is case insensitive. pub fn is_packed_struct(&self) -> bool { - const PACKED_KEYS: [&str; 2] = ["packed", "lance-encoding:packed"]; PACKED_KEYS.iter().any(|key| { self.metadata .get(*key) @@ -1847,6 +1885,14 @@ mod tests { #[test] fn blob_unloaded_mut_selects_layout_from_metadata() { let metadata = HashMap::from([(BLOB_META_KEY.to_string(), "true".to_string())]); + let mut binary_field: Field = ArrowField::new("blob", DataType::LargeBinary, true) + .with_metadata(metadata.clone()) + .try_into() + .unwrap(); + binary_field.binary_blob_mut(); + assert!(binary_field.metadata.contains_key(BLOB_META_KEY)); + assert!(!binary_field.is_blob_v2()); + let mut field: Field = ArrowField::new("blob", DataType::LargeBinary, true) .with_metadata(metadata) .try_into() @@ -1854,6 +1900,12 @@ mod tests { field.unloaded_mut(); assert_eq!(field.children.len(), 2); assert_eq!(field.logical_type, BLOB_DESC_LANCE_FIELD.logical_type); + assert!(field.is_blob()); + assert!(!field.is_blob_v2()); + field.unloaded_mut(); + assert_eq!(field.children.len(), 2); + assert_eq!(field.logical_type, BLOB_DESC_LANCE_FIELD.logical_type); + assert!(!field.is_blob_v2()); let metadata = HashMap::from([(ARROW_EXT_NAME_KEY.to_string(), BLOB_V2_EXT_NAME.to_string())]); @@ -1874,6 +1926,12 @@ mod tests { field.unloaded_mut(); assert_eq!(field.children.len(), 5); assert_eq!(field.logical_type, BLOB_V2_DESC_LANCE_FIELD.logical_type); + assert!(!field.metadata.contains_key(ARROW_EXT_NAME_KEY)); + assert!(field.is_blob_v2()); + field.unloaded_mut(); + assert_eq!(field.children.len(), 5); + assert_eq!(field.logical_type, BLOB_V2_DESC_LANCE_FIELD.logical_type); + assert!(!field.metadata.contains_key(ARROW_EXT_NAME_KEY)); } #[test] @@ -1944,4 +2002,43 @@ mod tests { .unwrap(); assert_eq!(unloaded_projected, unloaded); } + + #[test] + fn blob_descriptor_projection_preserves_synthetic_children() { + let metadata = + HashMap::from([(ARROW_EXT_NAME_KEY.to_string(), BLOB_V2_EXT_NAME.to_string())]); + let mut blob: Field = ArrowField::new( + "blob", + DataType::Struct( + vec![ + ArrowField::new("data", DataType::LargeBinary, true), + ArrowField::new("uri", DataType::Utf8, true), + ] + .into(), + ), + true, + ) + .with_metadata(metadata) + .try_into() + .unwrap(); + let mut next_id = 0; + blob.set_id(-1, &mut next_id); + + let schema = Arc::new(crate::datatypes::Schema { + fields: vec![blob], + metadata: HashMap::new(), + }); + let descriptor_schema = Projection::full(schema) + .with_blob_handling(crate::datatypes::BlobHandling::BlobsDescriptions) + .to_bare_schema(); + assert!( + descriptor_schema.fields[0] + .children + .iter() + .all(|child| child.id == -1) + ); + + let projected = Projection::full(Arc::new(descriptor_schema)).to_bare_schema(); + assert_eq!(projected.fields[0].children.len(), 5); + } } diff --git a/rust/lance-core/src/datatypes/schema.rs b/rust/lance-core/src/datatypes/schema.rs index d13eb476359..6f9fc61b334 100644 --- a/rust/lance-core/src/datatypes/schema.rs +++ b/rust/lance-core/src/datatypes/schema.rs @@ -760,10 +760,6 @@ impl Schema { Ok(schema) } - pub fn all_fields_nullable(&self) -> bool { - SchemaFieldIterPreOrder::new(self).all(|f| f.nullable) - } - /// Returns the properly formatted path from root to the field. /// Field names containing dots are quoted (e.g., struct.`field.with.dot`) pub fn field_path(&self, field_id: i32) -> Result { @@ -1075,6 +1071,17 @@ pub enum BlobHandling { } impl BlobHandling { + fn should_load_binary(&self, field: &Field) -> bool { + if !field.is_blob() { + return false; + } + match self { + Self::AllBinary => true, + Self::SomeBlobsBinary(set) | Self::SomeBinary(set) => set.contains(&(field.id as u32)), + Self::BlobsDescriptions | Self::AllDescriptions => false, + } + } + fn should_unload(&self, field: &Field) -> bool { // Blob v2 columns are Structs, so we need to treat any blob-marked field as unloadable // even if the physical data type is not binary-like. @@ -1090,10 +1097,34 @@ impl BlobHandling { } } + /// Whether `field` will be projected as a lightweight blob *description* + /// (offset + size) rather than its full binary value under this handling. + /// + /// A description is tiny and cheap to read eagerly; the full binary value is + /// not. Materialization heuristics use this to decide early vs late loading. + pub fn returns_description(&self, field: &Field) -> bool { + self.should_unload(field) + } + + /// Apply this blob handling policy to a projected field tree. + /// + /// Blob descriptor modes convert blob leaves to descriptor views. Binary + /// modes convert selected blob leaves to `LargeBinary`. Non-blob nested + /// fields are preserved while their children are handled recursively. pub fn unload_if_needed(&self, mut field: Field) -> Field { + if self.should_load_binary(&field) { + field.binary_blob_mut(); + return field; + } if self.should_unload(&field) { field.unloaded_mut(); + return field; } + field.children = field + .children + .into_iter() + .map(|child| self.unload_if_needed(child)) + .collect(); field } } @@ -2474,86 +2505,6 @@ mod tests { assert!(res.is_none(), "Expected None, got {:?}", res); } - #[test] - pub fn test_all_fields_nullable() { - let test_cases = vec![ - ( - vec![], // empty schema - true, - ), - ( - vec![ - Field::new_arrow("a", DataType::Int32, true).unwrap(), - Field::new_arrow("b", DataType::Utf8, true).unwrap(), - ], // basic case - true, - ), - ( - vec![ - Field::new_arrow("a", DataType::Int32, false).unwrap(), - Field::new_arrow("b", DataType::Utf8, true).unwrap(), - ], - false, - ), - ( - // check nested schema, parent is nullable - vec![ - Field::new_arrow( - "struct", - DataType::Struct(ArrowFields::from(vec![ArrowField::new( - "a", - DataType::Int32, - false, - )])), - true, - ) - .unwrap(), - ], - false, - ), - ( - // check nested schema, child is nullable - vec![ - Field::new_arrow( - "struct", - DataType::Struct(ArrowFields::from(vec![ArrowField::new( - "a", - DataType::Int32, - true, - )])), - false, - ) - .unwrap(), - ], - false, - ), - ( - // check nested schema, all is nullable - vec![ - Field::new_arrow( - "struct", - DataType::Struct(ArrowFields::from(vec![ArrowField::new( - "a", - DataType::Int32, - true, - )])), - true, - ) - .unwrap(), - ], - true, - ), - ]; - - for (fields, expected) in test_cases { - let schema = Schema { - fields, - metadata: Default::default(), - }; - assert_eq!(schema.all_fields_nullable(), expected); - } - } - #[test] fn test_schema_unenforced_primary_key() { let cases = vec![ diff --git a/rust/lance-core/src/error.rs b/rust/lance-core/src/error.rs index 516dc695a1d..a85f08cb741 100644 --- a/rust/lance-core/src/error.rs +++ b/rust/lance-core/src/error.rs @@ -8,6 +8,52 @@ use snafu::{IntoError as _, Location, Snafu}; type BoxedError = Box; +#[cfg(feature = "backtrace")] +mod backtrace_support { + use std::backtrace::Backtrace; + + use snafu::{AsBacktrace, GenerateImplicitData}; + + #[derive(Debug)] + pub struct MaybeBacktrace(pub Option); + + impl GenerateImplicitData for MaybeBacktrace { + fn generate() -> Self { + Self(>::generate()) + } + } + + impl AsBacktrace for MaybeBacktrace { + fn as_backtrace(&self) -> Option<&Backtrace> { + self.0.as_ref() + } + } +} + +#[cfg(not(feature = "backtrace"))] +mod backtrace_support { + use std::backtrace::Backtrace; + + use snafu::{AsBacktrace, GenerateImplicitData}; + + #[derive(Debug)] + pub struct MaybeBacktrace; + + impl GenerateImplicitData for MaybeBacktrace { + fn generate() -> Self { + Self + } + } + + impl AsBacktrace for MaybeBacktrace { + fn as_backtrace(&self) -> Option<&Backtrace> { + None + } + } +} + +use backtrace_support::MaybeBacktrace; + /// Error for when a requested field is not found in a schema. /// /// This error computes suggestions lazily (only when displayed) to avoid @@ -51,6 +97,28 @@ pub fn box_error(e: impl std::error::Error + Send + Sync + 'static) -> BoxedErro Box::new(e) } +/// Why a writer is fenced. Both reasons are terminal, but callers must tell them +/// apart (a peer takeover vs. our own failure) rather than parse the message. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum FenceReason { + /// A successor writer claimed a higher epoch; this writer lost ownership. + PeerClaimedEpoch, + /// Our own WAL persistence failed, so in-memory state may have diverged from + /// the durable WAL. The writer must be reopened to replay. + PersistenceFailure, +} + +impl std::fmt::Display for FenceReason { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Stable strings — surfaced in error messages. + let s = match self { + Self::PeerClaimedEpoch => "peer claimed epoch", + Self::PersistenceFailure => "persistence failure", + }; + f.write_str(s) + } +} + #[derive(Debug, Snafu)] #[snafu(visibility(pub))] pub enum Error { @@ -59,18 +127,24 @@ pub enum Error { source: BoxedError, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Dataset already exists: {uri}, {location}"))] DatasetAlreadyExists { uri: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Append with different schema: {difference}, location: {location}"))] SchemaMismatch { difference: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Dataset at path {path} was not found: {source}, {location}"))] DatasetNotFound { @@ -78,6 +152,8 @@ pub enum Error { source: BoxedError, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Encountered corrupt file {path}: {source}, {location}"))] CorruptFile { @@ -85,13 +161,16 @@ pub enum Error { source: BoxedError, #[snafu(implicit)] location: Location, - // TODO: add backtrace? + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Not supported: {source}, {location}"))] NotSupported { source: BoxedError, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Commit conflict for version {version}: {source}, {location}"))] CommitConflict { @@ -99,12 +178,16 @@ pub enum Error { source: BoxedError, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Incompatible transaction: {source}, {location}"))] IncompatibleTransaction { source: BoxedError, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Retryable commit conflict for version {version}: {source}, {location}"))] RetryableCommitConflict { @@ -112,12 +195,16 @@ pub enum Error { source: BoxedError, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Too many concurrent writers. {message}, {location}"))] TooMuchWriteContention { message: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Operation timed out: {message}, {location}"))] Timeout { @@ -132,54 +219,72 @@ pub enum Error { message: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("A prerequisite task failed: {message}, {location}"))] PrerequisiteFailed { message: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Unprocessable: {message}, {location}"))] Unprocessable { message: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("LanceError(Arrow): {message}, {location}"))] Arrow { message: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("LanceError(Schema): {message}, {location}"))] Schema { message: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Not found: {uri}, {location}"))] NotFound { uri: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("LanceError(IO): {source}, {location}"))] IO { source: BoxedError, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("LanceError(Index): {message}, {location}"))] Index { message: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Lance index not found: {identity}, {location}"))] IndexNotFound { identity: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Cannot infer storage location from: {message}"))] InvalidTableLocation { message: String }, @@ -190,18 +295,24 @@ pub enum Error { error: BoxedError, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Cloned error: {message}, {location}"))] Cloned { message: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Query Execution error: {message}, {location}"))] Execution { message: String, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Ref is invalid: {message}"))] InvalidRef { message: String }, @@ -220,12 +331,16 @@ pub enum Error { minor_version: u16, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, #[snafu(display("Namespace error: {source}, {location}"))] Namespace { source: BoxedError, #[snafu(implicit)] location: Location, + #[snafu(implicit)] + backtrace: MaybeBacktrace, }, /// External error passed through from user code. /// @@ -248,9 +363,78 @@ pub enum Error { #[snafu(implicit)] location: Location, }, + /// A writer has been fenced and must stop (see [`FenceReason`]). The message + /// keeps the `Writer fenced` prefix for legacy string consumers; new code + /// should match on [`Error::fence_reason`]. + #[snafu(display("Writer fenced ({reason}): {message}, {location}"))] + Fenced { + reason: FenceReason, + message: String, + #[snafu(implicit)] + location: Location, + }, } impl Error { + /// Returns the captured Rust backtrace, if available. + /// + /// Requires the `backtrace` feature to be enabled at compile time + /// and `RUST_BACKTRACE=1` at runtime. + #[cfg(feature = "backtrace")] + pub fn backtrace(&self) -> Option<&std::backtrace::Backtrace> { + match self { + Self::InvalidInput { backtrace, .. } + | Self::DatasetAlreadyExists { backtrace, .. } + | Self::SchemaMismatch { backtrace, .. } + | Self::DatasetNotFound { backtrace, .. } + | Self::CorruptFile { backtrace, .. } + | Self::NotSupported { backtrace, .. } + | Self::CommitConflict { backtrace, .. } + | Self::IncompatibleTransaction { backtrace, .. } + | Self::RetryableCommitConflict { backtrace, .. } + | Self::TooMuchWriteContention { backtrace, .. } + | Self::Internal { backtrace, .. } + | Self::PrerequisiteFailed { backtrace, .. } + | Self::Unprocessable { backtrace, .. } + | Self::Arrow { backtrace, .. } + | Self::Schema { backtrace, .. } + | Self::NotFound { backtrace, .. } + | Self::IO { backtrace, .. } + | Self::Index { backtrace, .. } + | Self::IndexNotFound { backtrace, .. } + | Self::Wrapped { backtrace, .. } + | Self::Cloned { backtrace, .. } + | Self::Execution { backtrace, .. } + | Self::VersionConflict { backtrace, .. } + | Self::Namespace { backtrace, .. } => { + use snafu::AsBacktrace; + backtrace.as_backtrace() + } + // Variants without a backtrace field — listed explicitly so that + // adding a new variant with a backtrace field triggers a compiler error. + Self::InvalidTableLocation { .. } + | Self::Stop + | Self::InvalidRef { .. } + | Self::RefConflict { .. } + | Self::RefNotFound { .. } + | Self::Cleanup { .. } + | Self::VersionNotFound { .. } + | Self::External { .. } + | Self::FieldNotFound { .. } + | Self::Timeout { .. } + | Self::DiskCapExceeded { .. } + | Self::Fenced { .. } => None, + } + } + + /// Returns the captured Rust backtrace, if available. + /// + /// Always returns `None` when the `backtrace` feature is not enabled. + #[cfg(not(feature = "backtrace"))] + pub fn backtrace(&self) -> Option<&std::backtrace::Backtrace> { + None + } + #[track_caller] pub fn corrupt_file(path: object_store::path::Path, message: impl Into) -> Self { CorruptFileSnafu { path }.into_error(message.into().into()) @@ -271,6 +455,36 @@ impl Error { IOSnafu.into_error(message.into().into()) } + /// A successor writer claimed a higher epoch; this writer lost ownership. + #[track_caller] + pub fn fenced_by_peer(message: impl Into) -> Self { + FencedSnafu { + reason: FenceReason::PeerClaimedEpoch, + message: message.into(), + } + .build() + } + + /// Our WAL persistence failed; in-memory state may have diverged from the + /// durable WAL, so the writer must be reopened to replay. + #[track_caller] + pub fn writer_poisoned(message: impl Into) -> Self { + FencedSnafu { + reason: FenceReason::PersistenceFailure, + message: message.into(), + } + .build() + } + + /// The [`FenceReason`] if this is [`Error::Fenced`], else `None`. Prefer this + /// over matching the error message to decide how to react to a fence. + pub fn fence_reason(&self) -> Option { + match self { + Self::Fenced { reason, .. } => Some(*reason), + _ => None, + } + } + #[track_caller] pub fn io_source(source: BoxedError) -> Self { IOSnafu.into_error(source) @@ -549,7 +763,11 @@ impl From for Error { impl From for Error { #[track_caller] fn from(e: object_store::Error) -> Self { - Self::io_source(box_error(e)) + match e { + // source intentionally dropped; Error::NotFound carries only the path + object_store::Error::NotFound { path, .. } => Self::not_found(path), + other => Self::io_source(box_error(other)), + } } } @@ -734,6 +952,41 @@ mod test { } } + #[test] + fn test_caller_location_capture_not_found() { + let current_fn = get_caller_location(); + let f: Box Result<()>> = Box::new(|| { + Err(object_store::Error::NotFound { + path: "some/path".to_string(), + source: "not found".into(), + })?; + Ok(()) + }); + match f().unwrap_err() { + Error::NotFound { location, .. } => { + // +2 is the beginning of object_store::Error::NotFound... + assert_eq!(location.line(), current_fn.line() + 2, "{}", location) + } + #[allow(unreachable_patterns)] + other => panic!("expected NotFound, got {:?}", other), + } + } + + #[test] + fn test_object_store_not_found_converts_to_not_found() { + let os_err = object_store::Error::NotFound { + path: "test/path".to_string(), + source: "no such file".into(), + }; + let lance_err: Error = os_err.into(); + match lance_err { + Error::NotFound { uri, .. } => { + assert_eq!(uri, "test/path"); + } + other => panic!("Expected NotFound, got {:?}", other), + } + } + #[derive(Debug)] struct MyCustomError { code: i32, @@ -986,4 +1239,68 @@ mod test { _ => panic!("Expected InvalidInput variant, got {:?}", recovered), } } + + #[test] + fn test_backtrace_accessor() { + // Verify that backtrace() returns the expected result based on feature state + let err = Error::io("test backtrace"); + let bt = err.backtrace(); + #[cfg(feature = "backtrace")] + { + // With the backtrace feature enabled, whether a backtrace is captured + // depends on the RUST_BACKTRACE env var at runtime. We just verify + // the accessor doesn't panic and returns a valid Option. + let _ = bt; + } + #[cfg(not(feature = "backtrace"))] + { + // Without the backtrace feature, this must always be None. + assert!(bt.is_none()); + } + } + + #[test] + fn test_backtrace_captured_when_feature_enabled() { + // Test that backtrace is actually captured when the feature is on and + // RUST_BACKTRACE=1 is set in the environment before the process starts. + // + // NOTE: std::backtrace::Backtrace caches the RUST_BACKTRACE env check, + // so set_var at runtime does not reliably enable capture. This test + // verifies the accessor works correctly in both cases: + // - If RUST_BACKTRACE=1 was set before the test binary started, we get Some. + // - If not, we get None (even with the feature on), which is expected. + #[cfg(feature = "backtrace")] + { + let err = Error::io("backtrace capture test"); + if std::env::var("RUST_BACKTRACE").is_ok() { + assert!( + err.backtrace().is_some(), + "Expected a backtrace when RUST_BACKTRACE=1 and backtrace feature is enabled" + ); + } + // When RUST_BACKTRACE is not set, backtrace() may return None even + // with the feature enabled — this is correct runtime gating behavior. + } + #[cfg(not(feature = "backtrace"))] + { + let err = Error::io("backtrace capture test"); + assert!(err.backtrace().is_none()); + } + } + + #[test] + fn test_backtrace_returns_none_for_variants_without_location() { + let err = Error::InvalidTableLocation { + message: "test".to_string(), + }; + assert!(err.backtrace().is_none()); + + let err = Error::InvalidRef { + message: "test".to_string(), + }; + assert!(err.backtrace().is_none()); + + let err = Error::Stop; + assert!(err.backtrace().is_none()); + } } diff --git a/rust/lance-core/src/lib.rs b/rust/lance-core/src/lib.rs index 8379fa74c4d..32fb34ad5fe 100644 --- a/rust/lance-core/src/lib.rs +++ b/rust/lance-core/src/lib.rs @@ -17,7 +17,7 @@ pub mod levenshtein; pub mod traits; pub mod utils; -pub use error::{ArrowResult, Error, Result, box_error}; +pub use error::{ArrowResult, Error, FenceReason, Result, box_error}; /// Wildcard to indicate all non-system columns pub const WILDCARD: &str = "*"; diff --git a/rust/lance-core/src/utils.rs b/rust/lance-core/src/utils.rs index c202329838c..51d72c3dae3 100644 --- a/rust/lance-core/src/utils.rs +++ b/rust/lance-core/src/utils.rs @@ -15,6 +15,7 @@ pub mod hash; pub mod io_stats; pub mod parse; pub mod path; +pub mod row_addr_remap; pub mod tempfile; pub mod testing; pub mod tokio; diff --git a/rust/lance-core/src/utils/cpu.rs b/rust/lance-core/src/utils/cpu.rs index 4e7ab01871d..c4d5a976cbb 100644 --- a/rust/lance-core/src/utils/cpu.rs +++ b/rust/lance-core/src/utils/cpu.rs @@ -1,14 +1,29 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use std::fmt; use std::sync::LazyLock; -/// A level of SIMD support for some feature +/// A level of SIMD support for some feature. +/// +/// `#[non_exhaustive]` so future tiers (e.g. AVX-512 BF16, AMX) can be added +/// without breaking external `match` consumers. #[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] pub enum SimdSupport { None, Neon, Sse, + /// AVX (256-bit float ops) but no FMA and no AVX2. + /// Intel Sandy Bridge / Ivy Bridge. + Avx, + /// AVX + FMA but no AVX2. + /// AMD Piledriver / Steamroller / FX-7500. + AvxFma, + /// AVX2 + FMA. Intel Haswell / AMD Excavator and later. + /// + /// Selecting this tier asserts FMA is present: the kernels it dispatches to + /// are `#[target_feature(enable = "avx,fma")]`. Avx2, Avx512, Avx512FP16, @@ -16,6 +31,132 @@ pub enum SimdSupport { Lasx, } +impl fmt::Display for SimdSupport { + /// Formats the tier name in lowercase, matching pyarrow's + /// `runtime_info().simd_level` convention. + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let name = match self { + Self::None => "none", + Self::Neon => "neon", + Self::Sse => "sse", + Self::Avx => "avx", + Self::AvxFma => "avx_fma", + Self::Avx2 => "avx2", + Self::Avx512 => "avx512", + Self::Avx512FP16 => "avx512_fp16", + Self::Lsx => "lsx", + Self::Lasx => "lasx", + }; + f.write_str(name) + } +} + +/// Snapshot of the SIMD tier lance dispatches to on the current host, plus the +/// raw CPU features detected for diagnostic purposes. +/// +/// Mirrors the role of `pyarrow.runtime_info()`: a single, cheap call users can +/// make to verify which SIMD tier the runtime selected and what underlying +/// features the host advertises. Obtain one with [`simd_info()`]. +#[derive(Debug, Clone)] +pub struct SimdInfo { + /// The SIMD tier lance dispatches to at runtime on this host. + pub tier: SimdSupport, + /// The architecture name (e.g. "x86_64", "aarch64", "loongarch64"). + pub target_arch: &'static str, + /// Raw CPU feature flags detected on this host (x86_64 only; empty on + /// other architectures). Each entry is a feature name like "avx2", + /// "fma", "avx512f", "popcnt", etc. + pub host_features: Vec<&'static str>, +} + +/// Returns a snapshot of the SIMD tier lance is using on this host along with +/// the raw CPU feature flags that drove the decision. +/// +/// Useful for performance debugging and giving users a way to verify which +/// dispatch tier they are hitting without rebuilding lance. See [`SimdInfo`] +/// for the meaning of each field and [`SimdSupport`] for the tier values. +/// +/// # Examples +/// +/// ``` +/// use lance_core::utils::cpu::simd_info; +/// +/// let info = simd_info(); +/// println!("dispatching to {} on {}", info.tier, info.target_arch); +/// ``` +pub fn simd_info() -> SimdInfo { + SimdInfo { + tier: *SIMD_SUPPORT, + target_arch: std::env::consts::ARCH, + host_features: detect_host_features(), + } +} + +#[cfg(target_arch = "x86_64")] +fn detect_host_features() -> Vec<&'static str> { + // Each call must be inline: `is_x86_feature_detected!` does its own custom + // input parsing and rejects feature names received via a `macro_rules!` + // `:literal` metavariable on some toolchains. + let mut features = Vec::with_capacity(17); + if is_x86_feature_detected!("sse2") { + features.push("sse2"); + } + if is_x86_feature_detected!("sse3") { + features.push("sse3"); + } + if is_x86_feature_detected!("ssse3") { + features.push("ssse3"); + } + if is_x86_feature_detected!("sse4.1") { + features.push("sse4.1"); + } + if is_x86_feature_detected!("sse4.2") { + features.push("sse4.2"); + } + if is_x86_feature_detected!("popcnt") { + features.push("popcnt"); + } + if is_x86_feature_detected!("avx") { + features.push("avx"); + } + if is_x86_feature_detected!("avx2") { + features.push("avx2"); + } + if is_x86_feature_detected!("fma") { + features.push("fma"); + } + if is_x86_feature_detected!("f16c") { + features.push("f16c"); + } + if is_x86_feature_detected!("bmi1") { + features.push("bmi1"); + } + if is_x86_feature_detected!("bmi2") { + features.push("bmi2"); + } + if is_x86_feature_detected!("avx512f") { + features.push("avx512f"); + } + if is_x86_feature_detected!("avx512bw") { + features.push("avx512bw"); + } + if is_x86_feature_detected!("avx512cd") { + features.push("avx512cd"); + } + if is_x86_feature_detected!("avx512dq") { + features.push("avx512dq"); + } + if is_x86_feature_detected!("avx512vl") { + features.push("avx512vl"); + } + features +} + +#[cfg(not(target_arch = "x86_64"))] +fn detect_host_features() -> Vec<&'static str> { + Vec::new() +} + /// Support for SIMD operations pub static SIMD_SUPPORT: LazyLock = LazyLock::new(|| { #[cfg(all(target_arch = "aarch64", any(target_os = "ios", target_os = "tvos")))] @@ -42,8 +183,19 @@ pub static SIMD_SUPPORT: LazyLock = LazyLock::new(|| { } else { SimdSupport::Avx512 } - } else if is_x86_feature_detected!("avx2") { + } else if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") { + // FMA is checked explicitly: every kernel selected for this tier is + // `#[target_feature(enable = "avx,fma")]`, and AVX2 does not imply + // FMA in the ISA. Every shipping AVX2 part has FMA, so this only + // guards against a host that would otherwise take an FMA kernel + // without FMA. SimdSupport::Avx2 + } else if is_x86_feature_detected!("avx") && is_x86_feature_detected!("fma") { + // AMD Piledriver / Steamroller / FX-7500: 256-bit float ops + FMA but no AVX2. + SimdSupport::AvxFma + } else if is_x86_feature_detected!("avx") { + // Intel Sandy Bridge / Ivy Bridge: 256-bit float ops without FMA. + SimdSupport::Avx } else { SimdSupport::None } @@ -138,3 +290,68 @@ mod aarch64 { false } } + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[test] + fn simd_info_exposes_tier() { + let info = simd_info(); + assert_eq!(info.target_arch, std::env::consts::ARCH); + // Tier should match the detected SIMD support. + assert_eq!(info.tier, *SIMD_SUPPORT); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn simd_info_features_include_baseline() { + let info = simd_info(); + // The x86_64 ABI mandates SSE2, so it must always be present on this + // architecture. + assert!(info.host_features.contains(&"sse2")); + } + + #[cfg(not(target_arch = "x86_64"))] + #[test] + fn simd_info_features_empty_off_x86_64() { + let info = simd_info(); + assert!(info.host_features.is_empty()); + } + + /// The `Avx2` and `AvxFma` tiers both dispatch to kernels declared + /// `#[target_feature(enable = "avx,fma")]`, so neither may be selected on a + /// host without FMA. AVX2 does not imply FMA in the ISA, so the detection + /// checks it explicitly. (`Avx512*` is excluded: its kernels declare + /// `avx512f`, which is what `has_avx512` verifies.) + #[cfg(target_arch = "x86_64")] + #[test] + fn avx_fma_tiers_are_only_selected_when_fma_is_detected() { + if matches!(*SIMD_SUPPORT, SimdSupport::Avx2 | SimdSupport::AvxFma) { + assert!( + is_x86_feature_detected!("fma"), + "tier {} dispatches to avx,fma kernels but the host has no FMA", + *SIMD_SUPPORT + ); + } + } + + #[rstest] + #[case::none(SimdSupport::None, "none")] + #[case::neon(SimdSupport::Neon, "neon")] + #[case::sse(SimdSupport::Sse, "sse")] + #[case::avx(SimdSupport::Avx, "avx")] + #[case::avx_fma(SimdSupport::AvxFma, "avx_fma")] + #[case::avx2(SimdSupport::Avx2, "avx2")] + #[case::avx512(SimdSupport::Avx512, "avx512")] + #[case::avx512_fp16(SimdSupport::Avx512FP16, "avx512_fp16")] + #[case::lsx(SimdSupport::Lsx, "lsx")] + #[case::lasx(SimdSupport::Lasx, "lasx")] + fn simd_support_display_matches_lowercase_convention( + #[case] tier: SimdSupport, + #[case] expected: &str, + ) { + assert_eq!(tier.to_string(), expected); + } +} diff --git a/rust/lance-core/src/utils/row_addr_remap.rs b/rust/lance-core/src/utils/row_addr_remap.rs new file mode 100644 index 00000000000..6f5a6f2aae5 --- /dev/null +++ b/rust/lance-core/src/utils/row_addr_remap.rs @@ -0,0 +1,413 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Compact row-address remapping for compaction. +//! +//! Compaction rewrites rows into new fragments, so indices that store physical +//! row addresses need an old-address to new-address mapping without building an +//! O(total rows) `HashMap>`. +//! +//! Layout: +//! +//! * Old rows: `old_fragment_id -> (old_offsets, old_rows_before)` +//! * `old_offsets`: rewritten old row offsets in this old fragment. +//! * `old_rows_before`: rewritten row count before this old fragment. +//! * New rows: ordered new-fragment ranges +//! `(fragment_id, new_rows_before, physical_rows)` +//! * `new_rows_before`: rewritten row count before this new fragment. +//! +//! Lookup: +//! +//! * An address whose fragment was not rewritten returns `None`. +//! * For an address whose fragment was rewritten: +//! * Read `(old_offsets, old_rows_before)` from the old-row layout. +//! * If `offset` is not in `old_offsets`, return `Some(None)` because the +//! row was deleted. +//! * Otherwise, `old_offsets.rank(offset) - 1` is this row's 0-based +//! position among rewritten old rows in this old fragment. Add +//! `old_rows_before` to get `k`, the row's 0-based position among all +//! rewritten old rows. +//! * In the new-row layout, find the range +//! `(fragment_id, new_rows_before, physical_rows)` where +//! `new_rows_before <= k < new_rows_before + physical_rows`. +//! * The new address is `(fragment_id, k - new_rows_before)`. +//! +//! Ordering: +//! +//! Compact remap does not store each old-to-new row mapping. It computes `k` +//! from the old-row layout, then maps it to the k-th row written to the new +//! fragments. This requires the reader-to-writer pipeline to preserve row order. +//! +//! * `old_frag_ids` must match the order old fragments are read. Within each +//! old fragment, rewritten rows are interpreted by ascending old row offset. +//! * `new_frags` must match the order new rows are written. +//! * Current compaction satisfies this because it scans selected fragments in +//! order and writes the resulting stream without reordering rows. + +use crate::utils::address::RowAddress; +use crate::{Error, Result}; +use roaring::{RoaringBitmap, RoaringTreemap}; +use std::collections::HashMap; + +/// A queryable row-address remapping with the exact semantics of +/// `HashMap>::get(&addr).copied()`: +/// +/// * `None` — the address is not affected by this remap (keep it unchanged) +/// * `Some(None)` — the row was deleted +/// * `Some(Some(addr))` — the row moved to `addr` +#[derive(Clone)] +pub enum RowAddrRemap { + /// Compact, `O(#fragments)` remap built from per-group rewritten-row + /// bitmaps and new-fragment layouts. + Compact(CompactRowAddrRemap), + /// Full materialized old-to-new address map. Uses `O(#rows)` memory. + Direct(HashMap>), +} + +impl RowAddrRemap { + pub fn compact(groups: impl IntoIterator) -> Result { + Ok(Self::Compact(CompactRowAddrRemap::new(groups)?)) + } + + /// Build a remap from a fully materialized old-to-new address map. + pub fn direct(map: HashMap>) -> Self { + Self::Direct(map) + } + + /// An empty remap that leaves every address unchanged. + pub fn empty() -> Self { + Self::Direct(HashMap::new()) + } + + /// Look up `addr`. See [`RowAddrRemap`] for the tri-state return semantics. + #[inline] + pub fn get(&self, addr: u64) -> Option> { + match self { + Self::Compact(c) => c.get(addr), + Self::Direct(m) => m.get(&addr).copied(), + } + } + + pub fn is_empty(&self) -> bool { + match self { + Self::Compact(c) => c.is_empty(), + Self::Direct(m) => m.is_empty(), + } + } + + pub fn affected_fragments(&self) -> RoaringBitmap { + match self { + Self::Compact(c) => RoaringBitmap::from_iter(c.frag_to_group.keys().copied()), + Self::Direct(m) => RoaringBitmap::from_iter(m.keys().map(|addr| (addr >> 32) as u32)), + } + } + + pub fn fully_deleted_fragments(&self) -> Option { + match self { + Self::Compact(c) => c.fully_deleted_fragments(), + Self::Direct(m) => { + if m.values().all(|v| v.is_none()) { + Some(RoaringBitmap::from_iter( + m.keys().map(|addr| (addr >> 32) as u32), + )) + } else { + None + } + } + } + } +} + +/// Input describing one rewrite group: the old row addresses that were +/// rewritten plus the fragment layout before/after the rewrite. +pub struct GroupInput { + /// Old row addresses that were read and re-written into the new fragments. + pub rewritten_old_row_addrs: RoaringTreemap, + /// Old fragment ids covered by this group. + pub old_frag_ids: Vec, + /// New fragments produced by this group, as `(fragment_id, physical_rows)`, + pub new_frags: Vec<(u32, u32)>, +} + +#[derive(Clone)] +struct GroupRemap { + /// Old fragment id -> (rewritten old row offsets in that fragment, + /// rewritten row count before this fragment in the group). + frags: HashMap, + /// New fragment ranges as `(fragment_id, rewritten_rows_before, physical_rows)`, + /// used to map a rewritten row's group-local index to its new address via binary search. + new_frag_row_ranges: Vec<(u32, u64, u32)>, +} + +impl GroupRemap { + fn new(input: GroupInput) -> Result { + // `compute_new_addr` maps a rewritten row's group-local index to a new + // address by accumulating `physical_rows` in `new_frags` order, so that + // order must be the order rows were written. New fragment ids are + // reserved monotonically in write order (see `reserve_fragment_ids` in + // compaction), so ascending id is a proxy for write order; reject any + // input that violates it before it can silently misplace addresses. + let mut new_frag_row_ranges = Vec::with_capacity(input.new_frags.len()); + let mut rewritten_rows_before = 0u64; + let mut prev_frag_id: Option = None; + for (frag_id, physical_rows) in input.new_frags { + if physical_rows == 0 { + continue; + } + if let Some(prev) = prev_frag_id + && frag_id <= prev + { + return Err(Error::invalid_input(format!( + "compaction new fragments must be in ascending id (write) order, but fragment {frag_id} follows {prev}", + ))); + } + prev_frag_id = Some(frag_id); + new_frag_row_ranges.push((frag_id, rewritten_rows_before, physical_rows)); + rewritten_rows_before += physical_rows as u64; + } + let total_new_rows = rewritten_rows_before; + + let mut per_frag: HashMap = input + .rewritten_old_row_addrs + .bitmaps() + .map(|(frag_id, bitmap)| (frag_id, bitmap.clone())) + .collect(); + let mut frags = HashMap::new(); + let mut rewritten_rows_before = 0u64; + for &frag_id in &input.old_frag_ids { + // A fragment with no rewritten rows (fully deleted) contributes + // nothing to the rewritten row sequence. + if let Some(bitmap) = per_frag.remove(&frag_id) { + let num_rewritten_rows = bitmap.len(); + frags.insert(frag_id, (bitmap, rewritten_rows_before)); + rewritten_rows_before += num_rewritten_rows; + } + } + // Rewritten old row addresses must reference only fragments listed in `old_frag_ids`. + if !per_frag.is_empty() { + return Err(Error::invalid_input(format!( + "compaction rewritten old row addresses reference fragments {:?} not in the rewrite group's old fragments {:?}", + per_frag.keys().collect::>(), + input.old_frag_ids, + ))); + } + + // Rewritten old rows are mapped positionally onto the new rows, so the + // two counts must match exactly + let total_rewritten_old_rows = input.rewritten_old_row_addrs.len(); + if total_new_rows != total_rewritten_old_rows { + return Err(Error::invalid_input(format!( + "compaction rewrote {total_rewritten_old_rows} old rows from fragments {:?} but the new fragments hold {total_new_rows} rows", + input.old_frag_ids, + ))); + } + + Ok(Self { + frags, + new_frag_row_ranges, + }) + } + + fn compute_new_addr(&self, rewritten_row_index: u64) -> u64 { + let idx = + match self + .new_frag_row_ranges + .binary_search_by(|(_, rewritten_rows_before, _)| { + rewritten_rows_before.cmp(&rewritten_row_index) + }) { + Ok(i) => i, + Err(i) => i - 1, + }; + let (frag_id, rewritten_rows_before, _rows) = self.new_frag_row_ranges[idx]; + let offset = (rewritten_row_index - rewritten_rows_before) as u32; + u64::from(RowAddress::new_from_parts(frag_id, offset)) + } + + /// Compute the new address for an old row in this group. + /// Returns `None` if the old row was not rewritten. + #[inline] + fn get(&self, frag: u32, offset: u32) -> Option { + match self.frags.get(&frag) { + Some((bitmap, rewritten_rows_before)) if bitmap.contains(offset) => { + let rewritten_row_index = rewritten_rows_before + bitmap.rank(offset) - 1; + Some(self.compute_new_addr(rewritten_row_index)) + } + _ => None, + } + } +} + +/// Compact remap backed by per-group rewritten row bitmaps + new-fragment layouts. +#[derive(Clone)] +pub struct CompactRowAddrRemap { + groups: Vec, + /// Old fragment id -> index into `groups`. Size is O(#fragments), not rows. + frag_to_group: HashMap, +} + +impl CompactRowAddrRemap { + fn new(groups: impl IntoIterator) -> Result { + let mut frag_to_group = HashMap::new(); + let mut group_remaps = Vec::new(); + for input in groups { + let gi = group_remaps.len(); + for &frag_id in &input.old_frag_ids { + frag_to_group.insert(frag_id, gi); + } + group_remaps.push(GroupRemap::new(input)?); + } + Ok(Self { + groups: group_remaps, + frag_to_group, + }) + } + + #[inline] + pub fn get(&self, addr: u64) -> Option> { + let frag = (addr >> 32) as u32; + // Not in any rewrite group -> unaffected by this remap. + let gi = *self.frag_to_group.get(&frag)?; + Some(self.groups[gi].get(frag, addr as u32)) + } + + pub fn is_empty(&self) -> bool { + self.groups.is_empty() + } + + fn fully_deleted_fragments(&self) -> Option { + // A group with any rewritten row moved at least one row. + if self.groups.iter().any(|g| !g.frags.is_empty()) { + return None; + } + Some(RoaringBitmap::from_iter(self.frag_to_group.keys().copied())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn addr(frag: u32, offset: u32) -> u64 { + u64::from(RowAddress::new_from_parts(frag, offset)) + } + + #[test] + fn test_compact_lookup() { + // Group A: out-of-order old frags [4, 3], split new frags (11 empty), + // some deletions. frag 4 (5 rows) keeps 0,2,4; frag 3 keeps 0,1, so the + // rewritten rows (4,0)(4,2)(4,4)(3,0)(3,1) go to new frags 10(2), 12(3). + // Group B is a fully-deleted fragment. + let group_a = GroupInput { + rewritten_old_row_addrs: RoaringTreemap::from_iter([ + addr(4, 0), + addr(4, 2), + addr(4, 4), + addr(3, 0), + addr(3, 1), + ]), + old_frag_ids: vec![4, 3], + new_frags: vec![(10, 2), (11, 0), (12, 3)], + }; + let group_b = GroupInput { + rewritten_old_row_addrs: RoaringTreemap::new(), + old_frag_ids: vec![7], + new_frags: vec![], + }; + let remap = RowAddrRemap::compact([group_a, group_b]).unwrap(); + + // Moves, in rewrite order; frag 4 comes first despite the larger id. + assert_eq!(remap.get(addr(4, 0)), Some(Some(addr(10, 0)))); + assert_eq!(remap.get(addr(4, 2)), Some(Some(addr(10, 1)))); + // Rank 2 skips the zero-row new fragment 11 and lands in fragment 12. + assert_eq!(remap.get(addr(4, 4)), Some(Some(addr(12, 0)))); + assert_eq!(remap.get(addr(3, 0)), Some(Some(addr(12, 1)))); + assert_eq!(remap.get(addr(3, 1)), Some(Some(addr(12, 2)))); + // Deleted offsets inside a rewritten fragment. + assert_eq!(remap.get(addr(4, 1)), Some(None)); + assert_eq!(remap.get(addr(4, 3)), Some(None)); + // Covered but fully-deleted fragment -> Some(None), not None. + assert_eq!(remap.get(addr(7, 0)), Some(None)); + // Fragment in no group -> unaffected. + assert_eq!(remap.get(addr(9, 0)), None); + assert!(!remap.is_empty()); + } + + #[test] + fn test_fragment_sets() { + // No rewritten rows at all: every covered fragment is fully deleted. + let dead = RowAddrRemap::compact([GroupInput { + rewritten_old_row_addrs: RoaringTreemap::new(), + old_frag_ids: vec![3, 7], + new_frags: vec![], + }]) + .unwrap(); + assert_eq!( + dead.fully_deleted_fragments(), + Some(RoaringBitmap::from_iter([3u32, 7u32])) + ); + assert_eq!( + dead.affected_fragments(), + RoaringBitmap::from_iter([3u32, 7u32]) + ); + + // At least one rewritten row -> not fully deleted, but both covered + // fragments (including the fully-deleted frag 1) are still affected. + let alive = RowAddrRemap::compact([GroupInput { + rewritten_old_row_addrs: RoaringTreemap::from_iter([addr(0, 0)]), + old_frag_ids: vec![0, 1], + new_frags: vec![(10, 1)], + }]) + .unwrap(); + assert!(alive.fully_deleted_fragments().is_none()); + assert_eq!( + alive.affected_fragments(), + RoaringBitmap::from_iter([0u32, 1u32]) + ); + } + + #[test] + fn test_compact_rejects_rewritten_addrs_outside_old_frags() { + // Rewritten addresses reference frag 5, not in old_frag_ids. The count + // still matches (2 == 2), so only the per-fragment split catches it. + let input = GroupInput { + rewritten_old_row_addrs: RoaringTreemap::from_iter([addr(0, 0), addr(5, 0)]), + old_frag_ids: vec![0], + new_frags: vec![(10, 2)], + }; + assert!(RowAddrRemap::compact([input]).is_err()); + } + + #[test] + fn test_compact_rejects_new_frags_out_of_write_order() { + // New fragments out of ascending id (write) order would make + // `compute_new_addr` accumulate rows in the wrong order, silently + // misplacing addresses. A zero-row fragment between them is ignored. + let input = GroupInput { + rewritten_old_row_addrs: RoaringTreemap::from_iter([addr(0, 0), addr(0, 1)]), + old_frag_ids: vec![0], + new_frags: vec![(12, 1), (11, 1)], + }; + assert!(RowAddrRemap::compact([input]).is_err()); + } + + #[test] + fn test_direct_and_empty() { + // Direct covers arbitrary maps the compact form can't express. + let mut map = HashMap::new(); + map.insert(addr(2, 0), Some(addr(9, 9))); + map.insert(addr(5, 1), None); + let remap = RowAddrRemap::direct(map); + assert_eq!(remap.get(addr(2, 0)), Some(Some(addr(9, 9)))); + assert_eq!(remap.get(addr(5, 1)), Some(None)); + assert_eq!(remap.get(addr(2, 1)), None); + // affected_fragments over an explicit map: the fragment of every key. + assert_eq!( + remap.affected_fragments(), + RoaringBitmap::from_iter([2u32, 5u32]) + ); + + let empty = RowAddrRemap::empty(); + assert!(empty.is_empty()); + assert_eq!(empty.get(addr(0, 0)), None); + } +} diff --git a/rust/lance-core/src/utils/tokio.rs b/rust/lance-core/src/utils/tokio.rs index 46c9475665b..fd33fc6b216 100644 --- a/rust/lance-core/src/utils/tokio.rs +++ b/rust/lance-core/src/utils/tokio.rs @@ -66,11 +66,15 @@ static RUNTIME_INSTALLED: atomic::AtomicBool = atomic::AtomicBool::new(false); static ATFORK_INSTALLED: atomic::AtomicBool = atomic::AtomicBool::new(false); -fn global_cpu_runtime() -> &'static mut Runtime { +fn global_cpu_runtime() -> &'static Runtime { loop { let ptr = CPU_RUNTIME.load(Ordering::SeqCst); if !ptr.is_null() { - return unsafe { &mut *ptr }; + // SAFETY: `ptr` was produced by `Box::into_raw` below and is only ever + // reset to null by `atfork_tokio_child` in the forked child (single- + // threaded, async-signal context). The `Box` is never reclaimed, so the + // `Runtime` lives for the rest of the process. + return unsafe { &*ptr }; } if !RUNTIME_INSTALLED.fetch_or(true, Ordering::SeqCst) { break; @@ -82,7 +86,9 @@ fn global_cpu_runtime() -> &'static mut Runtime { } let new_ptr = Box::into_raw(Box::new(create_runtime())); CPU_RUNTIME.store(new_ptr, Ordering::SeqCst); - unsafe { &mut *new_ptr } + // SAFETY: `new_ptr` was just obtained from `Box::into_raw`, so it is non-null, + // aligned, and points to a live `Runtime` that is never reclaimed. + unsafe { &*new_ptr } } /// After a fork() operation, force re-creation of the BackgroundExecutor. Note: this function @@ -108,6 +114,43 @@ fn install_atfork() {} /// /// This can also be used to convert a big chunk of synchronous work into a future /// so that it can be run in parallel with something like StreamExt::buffered() +/// +/// # Only hand over substantial CPU work +/// +/// Dispatching to the pool has real overhead (a `spawn_blocking` hop plus a oneshot +/// channel round trip). As a rule of thumb the closure should be expected to do at +/// least ~100µs of CPU work; below that the thread-pool overhead is likely to +/// outweigh any parallelism benefit, and the work is better left inline. +/// +/// # The task must never wait on anything +/// +/// The CPU pool is sized to [`get_num_compute_intensive_cpus`], which is +/// `max(1, num_cpus - LANCE_IO_CORE_RESERVATION)`. On a big host that is plenty of +/// workers (e.g. 62 on a 64-core box), but in resource-constrained environments it can +/// collapse to a **single blocking thread** — on machines with `<= 3` visible CPUs +/// (1-vCPU VMs, CI runners, CPU-limited Kubernetes pods) the pool has exactly one +/// worker. A closure passed to `spawn_cpu` occupies one of these threads for its entire +/// lifetime, including any time it spends *parked*. So the closure must only consume +/// CPU and return; it must +/// **never** block, wait, or park. Concretely, the closure must not, directly or +/// transitively: +/// +/// * **No channels** — no blocking send/recv (`send_blocking`, blocking `recv`, etc.). +/// A full/empty channel parks the thread, and whatever would drain/fill the channel +/// may need the same pool to run. +/// * **No I/O** — no file, network, or object-store reads/writes, and no disk spills. +/// I/O parks the thread while making no progress on CPU work. +/// * **No locks** — no acquiring a contended lock (or any lock that is held across an +/// `.await` elsewhere). Waiting for the lock parks the thread. +/// * **No `block_on` / `.blocking_*`** — never drive or wait on another async task +/// from inside the closure. +/// +/// If any of these hold, the parked thread can starve the exact work that would +/// unblock it, deadlocking the whole pool with no timeout and no error — a silent +/// hang at 0% CPU. (See .) When work +/// needs to wait on a channel/lock/I/O, keep the waiting in an async task and only +/// hand the pure-CPU portion to `spawn_cpu`, e.g. build each batch with `spawn_cpu` +/// and dispatch it with `tx.send(batch).await` in the surrounding async code. pub fn spawn_cpu< E: std::error::Error + Send + 'static, F: FnOnce() -> std::result::Result + Send + 'static, diff --git a/rust/lance-datafusion/src/exec.rs b/rust/lance-datafusion/src/exec.rs index 8f346f45612..a6be584420c 100644 --- a/rust/lance-datafusion/src/exec.rs +++ b/rust/lance-datafusion/src/exec.rs @@ -6,6 +6,7 @@ use std::{ collections::HashMap, fmt::{self, Formatter}, + num::NonZero, sync::{Arc, Mutex, OnceLock}, time::Duration, }; @@ -14,7 +15,6 @@ use chrono::{DateTime, Utc}; use arrow_array::RecordBatch; use arrow_schema::Schema as ArrowSchema; -use datafusion::physical_plan::metrics::MetricType; use datafusion::{ catalog::streaming::StreamingTable, dataframe::DataFrame, @@ -36,6 +36,7 @@ use datafusion::{ streaming::PartitionStream, }, }; +use datafusion::{execution::memory_pool::TrackConsumersPool, physical_plan::metrics::MetricType}; use datafusion_common::{DataFusionError, Statistics}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; @@ -310,7 +311,7 @@ impl std::fmt::Debug for LanceExecutionOptions { } } -const DEFAULT_LANCE_MEM_POOL_SIZE_PER_PARTITION: u64 = 100 * 1024 * 1024; +const DEFAULT_LANCE_MEM_POOL_SIZE_PER_PARTITION: u64 = 150 * 1024 * 1024; const DEFAULT_LANCE_MAX_TEMP_DIRECTORY_SIZE: u64 = 100 * 1024 * 1024 * 1024; // 100GB impl LanceExecutionOptions { @@ -366,12 +367,21 @@ pub fn new_session_context(options: &LanceExecutionOptions) -> SessionContext { session_config = session_config.with_target_partitions(target_partition); } if options.use_spilling() { + // The default 10MB sort spill reservation seems to be too small for many common cases. + // + // There currently is no reasonable guidance provided by DataFusion for setting this value. + // We bump this to 40MB but try a smaller value if the mem pool is small. + let sort_spill_reservation_bytes = + (options.mem_pool_size() / 3).min(40 * 1024 * 1024) as usize; + session_config = + session_config.with_sort_spill_reservation_bytes(sort_spill_reservation_bytes); let disk_manager_builder = DiskManagerBuilder::default() .with_max_temp_directory_size(options.max_temp_directory_size()); runtime_env_builder = runtime_env_builder .with_disk_manager_builder(disk_manager_builder) - .with_memory_pool(Arc::new(FairSpillPool::new( - options.mem_pool_size() as usize + .with_memory_pool(Arc::new(TrackConsumersPool::new( + FairSpillPool::new(options.mem_pool_size() as usize), + NonZero::try_from(16).unwrap(), ))); } let runtime_env = runtime_env_builder.build_arc().unwrap(); diff --git a/rust/lance-datafusion/src/logical_expr.rs b/rust/lance-datafusion/src/logical_expr.rs index 0eed438dae7..db9abd7e204 100644 --- a/rust/lance-datafusion/src/logical_expr.rs +++ b/rust/lance-datafusion/src/logical_expr.rs @@ -51,14 +51,10 @@ pub fn resolve_column_type(expr: &Expr, schema: &Schema) -> Option { field_path.push(c.name.as_str()); break; } - Expr::ScalarFunction(udf) => { - if udf.name() == GetFieldFunc::default().name() { - let name = get_as_string_scalar_opt(&udf.args[1])?; - field_path.push(name); - current_expr = &udf.args[0]; - } else { - return None; - } + Expr::ScalarFunction(udf) if udf.name() == GetFieldFunc::default().name() => { + let name = get_as_string_scalar_opt(&udf.args[1])?; + field_path.push(name); + current_expr = &udf.args[0]; } _ => return None, } @@ -285,8 +281,10 @@ pub fn field_path_to_expr(field_path: &str) -> Result { ))); } - // Build the column expression, handling nested fields - let mut expr = col(&parts[0]); + // Build the column expression, handling nested fields. + let mut expr = Expr::Column(datafusion::common::Column::new_unqualified( + parts[0].clone(), + )); for part in &parts[1..] { expr = expr.field_newstyle(part); } @@ -301,8 +299,26 @@ mod tests { use super::*; use arrow_schema::{Field, Schema as ArrowSchema}; + use datafusion::common::Column; use datafusion_functions::core::expr_ext::FieldAccessor; + #[test] + fn test_field_path_to_expr_preserves_case_sensitive_root_column() { + let expr = field_path_to_expr("VECTOR").unwrap(); + + assert_eq!(expr, Expr::Column(Column::new_unqualified("VECTOR"))); + } + + #[test] + fn test_field_path_to_expr_preserves_case_sensitive_escaped_nested_path() { + let expr = field_path_to_expr("Parent.`Child.With.Dot`").unwrap(); + + assert_eq!( + expr, + Expr::Column(Column::new_unqualified("Parent")).field_newstyle("Child.With.Dot") + ); + } + #[test] fn test_resolve_large_utf8() { let arrow_schema = ArrowSchema::new(vec![Field::new("a", DataType::LargeUtf8, false)]); diff --git a/rust/lance-datafusion/src/planner.rs b/rust/lance-datafusion/src/planner.rs index 3fd1841ada7..31e3c2d9e4c 100644 --- a/rust/lance-datafusion/src/planner.rs +++ b/rust/lance-datafusion/src/planner.rs @@ -329,10 +329,82 @@ impl Planner { }) } + fn is_logical_binary_op(op: &BinaryOperator) -> bool { + matches!(op, BinaryOperator::And | BinaryOperator::Or) + } + + fn is_same_logical_binary_op(left: &BinaryOperator, right: &BinaryOperator) -> bool { + matches!( + (left, right), + (BinaryOperator::And, BinaryOperator::And) | (BinaryOperator::Or, BinaryOperator::Or) + ) + } + + fn flatten_logical_binary_exprs<'a>( + left: &'a SQLExpr, + op: &BinaryOperator, + right: &'a SQLExpr, + ) -> Vec<&'a SQLExpr> { + let mut leaves = Vec::new(); + let mut stack = vec![right, left]; + + while let Some(expr) = stack.pop() { + match expr { + SQLExpr::BinaryOp { + left, + op: child_op, + right, + } if Self::is_same_logical_binary_op(op, child_op) => { + stack.push(right.as_ref()); + stack.push(left.as_ref()); + } + _ => leaves.push(expr), + } + } + + leaves + } + + fn balanced_binary_expr(mut exprs: VecDeque, op: Operator) -> Result { + if exprs.is_empty() { + return Err(Error::invalid_input("Binary expression has no operands")); + } + + while exprs.len() > 1 { + let mut next = VecDeque::with_capacity(exprs.len().div_ceil(2)); + while let Some(left) = exprs.pop_front() { + if let Some(right) = exprs.pop_front() { + next.push_back(Expr::BinaryExpr(BinaryExpr::new( + Box::new(left), + op, + Box::new(right), + ))); + } else { + next.push_back(left); + } + } + exprs = next; + } + + exprs + .pop_front() + .ok_or_else(|| Error::invalid_input("Binary expression has no operands")) + } + fn binary_expr(&self, left: &SQLExpr, op: &BinaryOperator, right: &SQLExpr) -> Result { + let df_op = self.binary_op(op)?; + if Self::is_logical_binary_op(op) { + let leaves = Self::flatten_logical_binary_exprs(left, op, right); + let mut exprs = VecDeque::with_capacity(leaves.len()); + for leaf in leaves { + exprs.push_back(self.parse_sql_expr(leaf)?); + } + return Self::balanced_binary_expr(exprs, df_op); + } + Ok(Expr::BinaryExpr(BinaryExpr::new( Box::new(self.parse_sql_expr(left)?), - self.binary_op(op)?, + df_op, Box::new(self.parse_sql_expr(right)?), ))) } @@ -431,7 +503,7 @@ impl Planner { } _ => Err(Error::invalid_input(format!( "Unsupported function args: {:?}", - &func.args + func.args ))), } } @@ -990,13 +1062,9 @@ impl TreeNodeVisitor<'_> for ColumnCapturingVisitor { self.columns.insert(path); self.current_path.clear(); } - Expr::ScalarFunction(udf) => { - if udf.name() == GetFieldFunc::default().name() { - if let Some(name) = get_as_string_scalar_opt(&udf.args[1]) { - self.current_path.push_front(name.to_string()) - } else { - self.current_path.clear(); - } + Expr::ScalarFunction(udf) if udf.name() == GetFieldFunc::default().name() => { + if let Some(name) = get_as_string_scalar_opt(&udf.args[1]) { + self.current_path.push_front(name.to_string()) } else { self.current_path.clear(); } @@ -1102,6 +1170,22 @@ mod tests { ); } + #[test] + fn test_parse_deep_logical_filter() { + let planner = Planner::new(Arc::new(Schema::empty())); + + for op in ["AND", "OR"] { + let filter = std::iter::repeat_n("true", 1000) + .collect::>() + .join(&format!(" {op} ")); + + let expr = planner.parse_filter(&filter).unwrap(); + let optimized = planner.optimize_expr(expr).unwrap(); + + assert_eq!(optimized, lit(true)); + } + } + #[derive(Debug, Eq, PartialEq, Hash)] struct StrictFloat64Udf { signature: Signature, diff --git a/rust/lance-datafusion/src/sql.rs b/rust/lance-datafusion/src/sql.rs index b9049aa0214..67ce2ea24a2 100644 --- a/rust/lance-datafusion/src/sql.rs +++ b/rust/lance-datafusion/src/sql.rs @@ -45,18 +45,16 @@ pub(crate) fn parse_sql_filter(filter: &str) -> Result { let sql = format!("SELECT 1 FROM t WHERE {filter}"); let statement = parse_statement(&sql)?; - let selection = if let Statement::Query(query) = &statement { - if let SetExpr::Select(s) = query.body.as_ref() { - s.selection.as_ref() - } else { - None - } + if let Statement::Query(query) = statement + && let SetExpr::Select(select) = *query.body + && let Some(expr) = select.selection + { + Ok(expr) } else { - None - }; - let expr = - selection.ok_or_else(|| Error::invalid_input(format!("Filter is not valid: {filter}")))?; - Ok(expr.clone()) + Err(Error::invalid_input(format!( + "Filter is not valid: {filter}" + ))) + } } /// Parse a SQL expression to Expression. This is more lenient than parse_sql_filter @@ -65,22 +63,16 @@ pub(crate) fn parse_sql_expr(expr: &str) -> Result { let sql = format!("SELECT {expr} FROM t"); let statement = parse_statement(&sql)?; - let selection = if let Statement::Query(query) = &statement { - if let SetExpr::Select(s) = query.body.as_ref() { - if let SelectItem::UnnamedExpr(expr) = &s.projection[0] { - Some(expr) - } else { - None - } - } else { - None - } + if let Statement::Query(query) = statement + && let SetExpr::Select(select) = *query.body + && let Some(SelectItem::UnnamedExpr(expr)) = select.projection.into_iter().next() + { + Ok(expr) } else { - None - }; - let expr = selection - .ok_or_else(|| Error::invalid_input(format!("Expression is not valid: {expr}")))?; - Ok(expr.clone()) + Err(Error::invalid_input(format!( + "Expression is not valid: {expr}" + ))) + } } fn parse_statement(statement: &str) -> Result { diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index e5d617e44ef..8d2d28f21b0 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -26,7 +26,10 @@ use crate::{ }, data::{DataBlock, FixedWidthDataBlock, VariableWidthBlock}, encodings::{ - logical::primitive::{fullzip::PerValueCompressor, miniblock::MiniBlockCompressor}, + logical::primitive::{ + fullzip::PerValueCompressor, + miniblock::{MAX_MINIBLOCK_VALUES, MiniBlockCompressor}, + }, physical::{ binary::{ BinaryBlockDecompressor, BinaryMiniBlockDecompressor, BinaryMiniBlockEncoder, @@ -51,7 +54,10 @@ use crate::{ PackedStructVariablePerValueEncoder, VariablePackedStructFieldDecoder, VariablePackedStructFieldKind, }, - rle::{RleDecompressor, RleEncoder}, + rle::{ + RleChildDecompressor, RleDecompressor, RleEncoder, RunLengthWidth, + rle_encoded_size, select_run_length_width, + }, value::{ValueDecompressor, ValueEncoder}, }, }, @@ -78,6 +84,7 @@ const DEFAULT_RLE_COMPRESSION_THRESHOLD: f64 = 0.5; // Minimum block size (32kb) to trigger general block compression const MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION: u64 = 32 * 1024; +const RLE_BLOCK_HEADER_BYTES: u128 = std::mem::size_of::() as u128; /// Trait for compression algorithms that compress an entire block of data into one opaque /// and self-described chunk. @@ -164,7 +171,9 @@ fn try_bss_for_mini_block( fn try_rle_for_mini_block( data: &FixedWidthDataBlock, + version: LanceFileVersion, params: &CompressionFieldParams, + use_rle_v2: bool, ) -> Option> { let bits = data.bits_per_value; if !matches!(bits, 8 | 16 | 32 | 64) { @@ -188,43 +197,88 @@ fn try_rle_for_mini_block( return None; } - // Estimate the encoded size. - // - // RLE stores (value, run_length) pairs. Run lengths are u8 and long runs are split into - // multiple entries of up to 255 values. We don't know the run length distribution here, - // so we conservatively account for splitting with an upper bound. let num_values = data.num_values; - let estimated_pairs = (run_count.saturating_add(num_values / 255)).min(num_values); - let raw_bytes = (num_values as u128) * (type_size as u128); - let rle_bytes = (estimated_pairs as u128) * ((type_size + 1) as u128); + let (run_length_width, rle_bytes) = if use_rle_v2 { + estimate_rle_width_and_size_from_data(data, Some(*MAX_MINIBLOCK_VALUES)).ok()? + } else { + ( + RunLengthWidth::U8, + estimate_rle_size_for_width_from_data( + data, + Some(*MAX_MINIBLOCK_VALUES), + RunLengthWidth::U8, + ) + .ok()?, + ) + }; + + let use_child_encodings = version.resolve() >= LanceFileVersion::V2_3; + let child_compression = if use_child_encodings { + rle_child_compression_config(params) + } else { + None + }; + let use_child_bitpacking = use_child_encodings; + let rle_encoder = || { + if use_child_encodings { + RleEncoder::with_child_encoding( + run_length_width, + child_compression, + child_compression, + use_child_bitpacking, + ) + } else { + RleEncoder::with_run_length_width(run_length_width) + } + }; + + #[cfg(feature = "bitpacking")] + let bitpack_bytes = estimate_inline_bitpacking_bytes(data).map(u128::from); + #[cfg(not(feature = "bitpacking"))] + let bitpack_bytes = None::; + + let mut selected_rle_bytes = rle_bytes; + let should_estimate_child_size = use_child_encodings + && (child_compression.is_some() || cfg!(feature = "bitpacking")) + && (rle_bytes >= raw_bytes || bitpack_bytes.is_some_and(|bytes| bytes < rle_bytes)); + if should_estimate_child_size { + selected_rle_bytes = rle_encoder().selected_payload_size(data).ok()?; + } - if rle_bytes < raw_bytes { - #[cfg(feature = "bitpacking")] + if selected_rle_bytes < raw_bytes { + if let Some(bitpack_bytes) = bitpack_bytes + && bitpack_bytes < selected_rle_bytes { - if let Some(bitpack_bytes) = estimate_inline_bitpacking_bytes(data) - && (bitpack_bytes as u128) < rle_bytes - { - return None; - } + return None; } - return Some(Box::new(RleEncoder::new())); + return Some(Box::new(rle_encoder())); } None } +fn rle_child_compression_config(params: &CompressionFieldParams) -> Option { + let raw = params.compression.as_deref()?; + if matches!(raw, "none" | "fsst") { + return None; + } + let scheme = CompressionScheme::from_str(raw).ok()?; + Some(CompressionConfig::new(scheme, params.compression_level)) +} + fn try_rle_for_block( data: &FixedWidthDataBlock, version: LanceFileVersion, params: &CompressionFieldParams, -) -> Option<(Box, CompressiveEncoding)> { + use_rle_v2: bool, +) -> Result, CompressiveEncoding)>> { if version < LanceFileVersion::V2_2 { - return None; + return Ok(None); } let bits = data.bits_per_value; if !matches!(bits, 8 | 16 | 32 | 64) { - return None; + return Ok(None); } let run_count = data.expect_single_stat::(Stat::RunCount); @@ -232,15 +286,71 @@ fn try_rle_for_block( .rle_threshold .unwrap_or(DEFAULT_RLE_COMPRESSION_THRESHOLD); - if (run_count as f64) < (data.num_values as f64) * threshold { - let compressor = Box::new(RleEncoder::new()); - let encoding = ProtobufUtils21::rle( - ProtobufUtils21::flat(bits, None), - ProtobufUtils21::flat(/*bits_per_value=*/ 8, None), - ); - return Some((compressor, encoding)); + let passes_threshold = match params.rle_threshold { + Some(_) => (run_count as f64) < (data.num_values as f64) * threshold, + None => true, + }; + + if !passes_threshold { + return Ok(None); } - None + + let raw_bytes = (data.num_values as u128) * ((bits / 8) as u128); + let (run_length_width, rle_payload_bytes) = if use_rle_v2 { + estimate_rle_width_and_size_from_data(data, None)? + } else { + ( + RunLengthWidth::U8, + estimate_rle_size_for_width_from_data(data, None, RunLengthWidth::U8)?, + ) + }; + let rle_bytes = rle_payload_bytes.saturating_add(RLE_BLOCK_HEADER_BYTES); + + if rle_bytes >= raw_bytes { + return Ok(None); + } + + #[cfg(feature = "bitpacking")] + { + if let Some(bitpack_bytes) = estimate_block_bitpacking_bytes(data) + && bitpack_bytes < rle_bytes + { + return Ok(None); + } + } + + let compressor = Box::new(RleEncoder::with_run_length_width(run_length_width)); + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::flat(bits, None), + ProtobufUtils21::flat(run_length_width.bits_per_value(), None), + ); + Ok(Some((compressor, encoding))) +} + +fn estimate_rle_width_and_size_from_data( + data: &FixedWidthDataBlock, + max_segment_values: Option, +) -> Result<(RunLengthWidth, u128)> { + select_run_length_width( + &data.data, + data.num_values, + data.bits_per_value, + max_segment_values, + ) +} + +fn estimate_rle_size_for_width_from_data( + data: &FixedWidthDataBlock, + max_segment_values: Option, + run_length_width: RunLengthWidth, +) -> Result { + rle_encoded_size( + &data.data, + data.num_values, + data.bits_per_value, + max_segment_values, + run_length_width, + ) } fn try_bitpack_for_mini_block(_data: &FixedWidthDataBlock) -> Option> { @@ -325,6 +435,67 @@ fn try_bitpack_for_block( } } +#[cfg(feature = "bitpacking")] +fn estimate_block_bitpacking_bytes(data: &FixedWidthDataBlock) -> Option { + let bits = data.bits_per_value; + if !matches!(bits, 8 | 16 | 32 | 64) || data.num_values == 0 { + return None; + } + + let bit_widths = data.expect_stat(Stat::BitWidth); + let widths = bit_widths.as_primitive::(); + let max_bit_width = *widths.values().iter().max()?; + let word_bytes = (bits / 8) as u128; + + let bitpacked_words = if data.num_values <= 1024 { + 1 + (1024u128 * (max_bit_width as u128)) / (bits as u128) + } else { + estimate_out_of_line_bitpacking_words(data.num_values, max_bit_width, bits)? + }; + let bitpacked_bytes = bitpacked_words.saturating_mul(word_bytes); + if bitpacked_bytes >= data.data_size() as u128 { + return None; + } + + Some(bitpacked_bytes) +} + +#[cfg(feature = "bitpacking")] +fn estimate_out_of_line_bitpacking_words( + num_values: u64, + compressed_bits_per_value: u64, + bits_per_value: u64, +) -> Option { + let num_values = usize::try_from(num_values).ok()?; + let compressed_bits_per_value = usize::try_from(compressed_bits_per_value).ok()?; + let bits_per_value = usize::try_from(bits_per_value).ok()?; + if compressed_bits_per_value >= bits_per_value { + return None; + } + + let elems_per_chunk = 1024usize; + let num_chunks = num_values.div_ceil(elems_per_chunk); + let words_per_chunk = (elems_per_chunk * compressed_bits_per_value).div_ceil(bits_per_value); + let last_chunk_is_runt = !num_values.is_multiple_of(elems_per_chunk); + + if !last_chunk_is_runt { + return Some((num_chunks * words_per_chunk) as u128); + } + + let num_whole_chunks = num_chunks - 1; + let remaining_items = num_values - num_whole_chunks * elems_per_chunk; + let tail_bit_savings = bits_per_value - compressed_bits_per_value; + let padding_cost = compressed_bits_per_value * (elems_per_chunk - remaining_items); + let tail_pack_savings = tail_bit_savings * remaining_items; + let tail_words = if padding_cost < tail_pack_savings { + words_per_chunk + } else { + remaining_items + }; + + Some((num_whole_chunks * words_per_chunk + tail_words) as u128) +} + fn maybe_wrap_general_for_mini_block( inner: Box, params: &CompressionFieldParams, @@ -393,6 +564,10 @@ impl DefaultCompressionStrategy { self } + fn use_rle_v2(&self) -> bool { + self.version.resolve() >= LanceFileVersion::V2_3 + } + /// Parse compression parameters from field metadata fn parse_field_metadata(field: &Field, version: &LanceFileVersion) -> CompressionFieldParams { let mut params = CompressionFieldParams::default(); @@ -456,7 +631,7 @@ impl DefaultCompressionStrategy { } let base = try_bss_for_mini_block(data, params) - .or_else(|| try_rle_for_mini_block(data, params)) + .or_else(|| try_rle_for_mini_block(data, self.version, params, self.use_rle_v2())) .or_else(|| try_bitpack_for_mini_block(data)) .unwrap_or_else(|| Box::new(ValueEncoder::default())); @@ -620,6 +795,13 @@ impl CompressionStrategy for DefaultCompressionStrategy { if (max_len > 32 * 1024 || per_value_requested) && data_size >= FSST_LEAST_INPUT_SIZE as u64 { + if compression == Some("zstd") { + let config = CompressionConfig::new( + CompressionScheme::Zstd, + field_params.compression_level, + ); + return Ok(Box::new(CompressedBufferEncoder::try_new(config)?)); + } return Ok(Box::new(CompressedBufferEncoder::default())); } @@ -664,7 +846,7 @@ impl CompressionStrategy for DefaultCompressionStrategy { match data { DataBlock::FixedWidth(fixed_width) => { if let Some((compressor, encoding)) = - try_rle_for_block(fixed_width, self.version, &field_params) + try_rle_for_block(fixed_width, self.version, &field_params, self.use_rle_v2())? { return Ok((compressor, encoding)); } @@ -814,10 +996,10 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { // compression. Ok(Box::new(ValueDecompressor::from_fsl(fsl))) } - Compression::Rle(rle) => { - let bits_per_value = validate_rle_compression(rle)?; - Ok(Box::new(RleDecompressor::new(bits_per_value))) - } + Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor( + rle, + decompression_strategy, + )?)), Compression::ByteStreamSplit(bss) => { let Compression::Flat(values) = bss.values.as_ref().unwrap().compression.as_ref().unwrap() @@ -1004,16 +1186,15 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { Ok(Box::new(general_decompressor)) } - Compression::Rle(rle) => { - let bits_per_value = validate_rle_compression(rle)?; - Ok(Box::new(RleDecompressor::new(bits_per_value))) - } + Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor(rle, self)?)), _ => todo!(), } } } -/// Validates RLE compression format and extracts bits_per_value -fn validate_rle_compression(rle: &crate::format::pb21::Rle) -> Result { +fn create_rle_decompressor( + rle: &crate::format::pb21::Rle, + decompression_strategy: &dyn DecompressionStrategy, +) -> Result { let values = rle .values .as_ref() @@ -1023,34 +1204,162 @@ fn validate_rle_compression(rle: &crate::format::pb21::Rle) -> Result { .as_ref() .ok_or_else(|| Error::invalid_input("RLE compression missing run lengths encoding"))?; - let values = values - .compression - .as_ref() - .ok_or_else(|| Error::invalid_input("RLE compression missing values compression"))?; - let Compression::Flat(values) = values else { + let values = create_rle_child_decompressor(values, "values", decompression_strategy)?; + let run_lengths = + create_rle_child_decompressor(run_lengths, "run lengths", decompression_strategy)?; + + if !matches!(values.bits_per_value(), 8 | 16 | 32 | 64) { + return Err(Error::invalid_input(format!( + "RLE compression only supports 8, 16, 32, or 64-bit values, got {}", + values.bits_per_value() + ))); + } + + let run_length_width = + RunLengthWidth::from_bits(run_lengths.bits_per_value()).ok_or_else(|| { + Error::invalid_input(format!( + "RLE compression only supports 8, 16, or 32-bit run lengths, got {}", + run_lengths.bits_per_value() + )) + })?; + + if values.requires_num_values() && run_lengths.requires_num_values() { return Err(Error::invalid_input( - "RLE compression only supports flat values", + "RLE values and run lengths child encodings cannot both require the run count", )); - }; + } + + if values.is_identity() && run_lengths.is_identity() { + return Ok(RleDecompressor::with_run_length_width( + values.bits_per_value(), + run_length_width, + )); + } + + Ok(RleDecompressor::with_child_decompressors( + values.bits_per_value(), + run_length_width, + values, + run_lengths, + )) +} - let run_lengths = run_lengths +fn create_rle_child_decompressor( + encoding: &CompressiveEncoding, + role: &str, + decompression_strategy: &dyn DecompressionStrategy, +) -> Result { + let compression = encoding .compression .as_ref() - .ok_or_else(|| Error::invalid_input("RLE compression missing run lengths compression"))?; - let Compression::Flat(run_lengths) = run_lengths else { - return Err(Error::invalid_input( - "RLE compression only supports flat run lengths", - )); - }; + .ok_or_else(|| Error::invalid_input(format!("RLE {role} missing child compression")))?; + let (bits_per_value, requires_num_values, needs_decompressor) = + validate_rle_child_compression(compression, role)?; - if run_lengths.bits_per_value != 8 { - return Err(Error::invalid_input(format!( - "RLE compression only supports 8-bit run lengths, got {}", - run_lengths.bits_per_value - ))); + if needs_decompressor { + Ok(RleChildDecompressor::block( + bits_per_value, + decompression_strategy.create_block_decompressor(encoding)?, + requires_num_values, + )) + } else { + Ok(RleChildDecompressor::flat(bits_per_value)) + } +} + +fn validate_rle_child_compression( + compression: &Compression, + role: &str, +) -> Result<(u64, bool, bool)> { + match compression { + Compression::Flat(flat) => Ok((flat.bits_per_value, false, false)), + Compression::General(general) => { + general.compression.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} general child missing compression config" + )) + })?; + let values = general.values.as_ref().ok_or_else(|| { + Error::invalid_input(format!("RLE {role} general child missing inner encoding")) + })?; + let inner = values.compression.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} general child missing inner compression" + )) + })?; + let (bits_per_value, requires_num_values) = + validate_rle_block_child_inner(inner, role)?; + Ok((bits_per_value, requires_num_values, true)) + } + Compression::OutOfLineBitpacking(out_of_line) => { + let values = out_of_line.values.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} bitpacking child missing values encoding" + )) + })?; + let Compression::Flat(_) = values.compression.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} bitpacking child missing values compression" + )) + })? + else { + return Err(Error::invalid_input(format!( + "RLE {role} bitpacking child only supports flat values" + ))); + }; + Ok((out_of_line.uncompressed_bits_per_value, true, true)) + } + other => Err(Error::invalid_input(format!( + "RLE {role} only supports flat, general, or out-of-line bitpacking child encodings, got {}", + compression_name(other) + ))), } +} + +fn validate_rle_block_child_inner(compression: &Compression, role: &str) -> Result<(u64, bool)> { + match compression { + Compression::Flat(flat) => Ok((flat.bits_per_value, false)), + Compression::OutOfLineBitpacking(out_of_line) => { + let values = out_of_line.values.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} bitpacking child missing values encoding" + )) + })?; + let Compression::Flat(_) = values.compression.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} bitpacking child missing values compression" + )) + })? + else { + return Err(Error::invalid_input(format!( + "RLE {role} bitpacking child only supports flat values" + ))); + }; + Ok((out_of_line.uncompressed_bits_per_value, true)) + } + other => Err(Error::invalid_input(format!( + "RLE {role} general child only supports flat or out-of-line bitpacking inner encodings, got {}", + compression_name(other) + ))), + } +} - Ok(values.bits_per_value) +fn compression_name(compression: &Compression) -> &'static str { + match compression { + Compression::Flat(_) => "flat", + Compression::Variable(_) => "variable", + Compression::Fsst(_) => "fsst", + Compression::OutOfLineBitpacking(_) => "out-of-line bitpacking", + Compression::InlineBitpacking(_) => "inline bitpacking", + Compression::General(_) => "general", + Compression::Constant(_) => "constant", + Compression::Dictionary(_) => "dictionary", + Compression::ByteStreamSplit(_) => "byte stream split", + Compression::PackedStruct(_) => "packed struct", + Compression::FixedSizeList(_) => "fixed-size list", + Compression::VariablePackedStruct(_) => "variable packed struct", + Compression::Rle(_) => "rle", + } } #[cfg(test)] @@ -1139,6 +1448,37 @@ mod tests { DataBlock::FixedWidth(block) } + fn rle_run_length_bits(encoding: &CompressiveEncoding) -> u64 { + let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else { + panic!("expected RLE encoding"); + }; + let Compression::Flat(run_lengths) = rle + .run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap() + else { + panic!("expected flat run lengths"); + }; + run_lengths.bits_per_value + } + + fn expect_rle_encoding(encoding: &CompressiveEncoding) -> &crate::format::pb21::Rle { + match encoding.compression.as_ref().unwrap() { + Compression::Rle(rle) => rle, + Compression::General(general) => { + let inner = general.values.as_ref().unwrap(); + let Compression::Rle(rle) = inner.compression.as_ref().unwrap() else { + panic!("expected wrapped RLE encoding"); + }; + rle + } + other => panic!("expected RLE encoding, got {}", compression_name(other)), + } + } + fn create_variable_width_block( bits_per_offset: u8, num_values: u64, @@ -1290,6 +1630,60 @@ mod tests { ); } + #[test] + fn test_rle_block_accounts_for_header_before_selecting() { + let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let field = create_test_field("small_constant", DataType::Int32); + let values = vec![42i32; 2]; + let mut block = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 2, + block_info: BlockInfo::default(), + }; + block.compute_stat(); + let data = DataBlock::FixedWidth(block); + + let (compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); + + assert!(format!("{compressor:?}").contains("ValueEncoder")); + assert!(matches!( + encoding.compression.as_ref(), + Some(Compression::Flat(_)) + )); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_block_prefers_bitpacking_when_smaller() { + let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let field = create_test_field("levels", DataType::UInt16); + + let mut values = Vec::with_capacity(2048); + for run_idx in 0..1024 { + values.extend(std::iter::repeat_n((run_idx % 2) as u16, 2)); + } + let mut block = FixedWidthDataBlock { + bits_per_value: 16, + data: LanceBuffer::reinterpret_vec(values), + num_values: 2048, + block_info: BlockInfo::default(), + }; + block.compute_stat(); + let data = DataBlock::FixedWidth(block); + + let (compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); + let debug_str = format!("{compressor:?}"); + assert!( + debug_str.contains("OutOfLineBitpacking"), + "expected OutOfLineBitpacking, got: {debug_str}" + ); + assert!(matches!( + encoding.compression.as_ref(), + Some(Compression::OutOfLineBitpacking(_)) + )); + } + #[test] #[cfg(feature = "bitpacking")] fn test_low_cardinality_prefers_bitpacking_over_rle() { @@ -1505,6 +1899,30 @@ mod tests { ); } + #[test] + #[cfg(feature = "zstd")] + fn test_compression_level_honored_for_large_per_value() { + let mut params = CompressionParams::new(); + params.columns.insert( + "html".to_string(), + CompressionFieldParams { + compression: Some("zstd".to_string()), + compression_level: Some(19), + ..Default::default() + }, + ); + let strategy = DefaultCompressionStrategy::with_params(params); + let field = create_test_field("html", DataType::Utf8); + let large = create_variable_width_block(32, 64, 40 * 1024); + + let per_value = strategy.create_per_value(&field, &large).unwrap(); + let debug = format!("{per_value:?}"); + assert!( + debug.contains("ZstdBufferCompressor") && debug.contains("compression_level: 19"), + "expected zstd level 19 to reach the per-value compressor, got: {debug}" + ); + } + #[test] fn test_parameter_merge_priority() { let mut params = CompressionParams::new(); @@ -1659,6 +2077,277 @@ mod tests { assert!(debug_str.contains("RleEncoder")); } + #[test] + fn test_rle_v2_miniblock_selects_u16_run_lengths() { + let mut metadata = HashMap::new(); + metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string()); + metadata.insert(BSS_META_KEY.to_string(), "off".to_string()); + let mut field = create_test_field("test_column", DataType::Int32); + field.metadata = metadata; + + let values = vec![7i32; 1000]; + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 1000, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (_compressed, encoding) = compressor.compress(data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 16); + } + + #[test] + fn test_rle_v2_miniblock_keeps_u8_run_lengths_before_v2_3() { + for version in [LanceFileVersion::V2_1, LanceFileVersion::V2_2] { + let mut metadata = HashMap::new(); + metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string()); + metadata.insert(BSS_META_KEY.to_string(), "off".to_string()); + let mut field = create_test_field("test_column", DataType::Int32); + field.metadata = metadata; + + let values = vec![7i32; 1000]; + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 1000, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = DefaultCompressionStrategy::new().with_version(version); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (_compressed, encoding) = compressor.compress(data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 8, "version={version}"); + } + } + + #[test] + fn test_rle_v2_uses_selected_width_cost_before_bitpacking() { + let mut metadata = HashMap::new(); + metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string()); + metadata.insert(BSS_META_KEY.to_string(), "off".to_string()); + let mut field = create_test_field("test_column", DataType::Int32); + field.metadata = metadata; + + let values = vec![0i32; 4096]; + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 4096, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let debug_str = format!("{compressor:?}"); + assert!(debug_str.contains("RleEncoder")); + + let (_compressed, encoding) = compressor.compress(data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 16); + } + + #[test] + fn test_rle_v2_sorted_dictionary_indices_select_u16_run_lengths() { + let field = create_test_field("dict_indices", DataType::Int32); + + let mut values = Vec::with_capacity(1_200); + for value in 0..4 { + values.extend(std::iter::repeat_n(value, 300)); + } + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 1_200, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (_compressed, encoding) = compressor.compress(data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 16); + } + + #[test] + fn test_rle_v2_short_runs_keep_u8_run_lengths() { + let field = create_test_field("dict_indices", DataType::Int32); + + let mut values = Vec::with_capacity(1_280); + for value in 0..10 { + values.extend(std::iter::repeat_n(value, 128)); + } + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 1_280, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (_compressed, encoding) = compressor.compress(data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 8); + } + + #[test] + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_rle_miniblock_released_versions_keep_flat_children_when_compression_requested() { + for version in [LanceFileVersion::V2_1, LanceFileVersion::V2_2] { + let mut params = CompressionParams::new(); + params.columns.insert( + "dict_indices".to_string(), + CompressionFieldParams { + compression: Some( + if cfg!(feature = "lz4") { "lz4" } else { "zstd" }.to_string(), + ), + rle_threshold: Some(1.0), + bss: Some(BssMode::Off), + ..Default::default() + }, + ); + let strategy = DefaultCompressionStrategy::with_params(params).with_version(version); + let field = create_test_field("dict_indices", DataType::UInt32); + + let mut values = Vec::with_capacity(8192 * 4); + for value in 0..8192u32 { + values.extend(std::iter::repeat_n(value, 4)); + } + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 8192 * 4, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (_compressed, encoding) = compressor.compress(data).unwrap(); + let rle = expect_rle_encoding(&encoding); + + assert!( + matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + Compression::Flat(_) + ), + "version={version}" + ); + assert!( + matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + Compression::Flat(_) + ), + "version={version}" + ); + } + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_miniblock_strategy_bitpacks_child_values_when_smaller() { + let field = create_test_field("dict_indices", DataType::Int32); + + let mut values = Vec::with_capacity(8192 * 4); + for value in 0..8192 { + values.extend(std::iter::repeat_n(value, 4)); + } + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 8192 * 4, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let debug_str = format!("{compressor:?}"); + assert!(debug_str.contains("RleEncoder")); + + let (_compressed, encoding) = compressor.compress(data).unwrap(); + let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else { + panic!("expected RLE encoding"); + }; + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + Compression::OutOfLineBitpacking(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + Compression::Flat(_) + )); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_miniblock_keeps_child_bitpacked_rle_when_smaller_than_inline_bitpacking() { + let field = create_test_field("int_score", DataType::UInt64); + + let mut values = Vec::with_capacity(8192 * 8); + for run_idx in 0..8192 { + let value = match run_idx % 3 { + 0 => 3u64, + 1 => 4u64, + _ => 5u64, + }; + values.extend(std::iter::repeat_n(value, 8)); + } + let mut data = FixedWidthDataBlock { + bits_per_value: 64, + data: LanceBuffer::reinterpret_vec(values), + num_values: 8192 * 8, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let debug_str = format!("{compressor:?}"); + assert!( + debug_str.contains("RleEncoder"), + "expected RLE to beat inline bitpacking after child selection, got: {debug_str}" + ); + + let (_compressed, encoding) = compressor.compress(data).unwrap(); + let rle = expect_rle_encoding(&encoding); + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + Compression::OutOfLineBitpacking(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + Compression::Flat(_) + )); + } + #[test] fn test_field_metadata_override_params() { // Set up params with one configuration @@ -1875,6 +2564,60 @@ mod tests { ); } + #[test] + fn test_rle_v2_block_selects_u32_run_lengths() { + let field = create_test_field("dict_indices", DataType::Int32); + let expected_values = vec![42i32; 70_000]; + let mut block = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(expected_values.clone()), + num_values: expected_values.len() as u64, + block_info: BlockInfo::default(), + }; + block.compute_stat(); + let data = DataBlock::FixedWidth(block); + + let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) + .with_version(LanceFileVersion::V2_3); + let (compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 32); + + let compressed = compressor.compress(data).unwrap(); + let decompressor = DefaultDecompressionStrategy::default() + .create_block_decompressor(&encoding) + .unwrap(); + let decoded = decompressor + .decompress(compressed, expected_values.len() as u64) + .unwrap(); + + match decoded { + DataBlock::FixedWidth(block) => { + let values = block.data.borrow_to_typed_slice::(); + assert_eq!(values.as_ref(), expected_values); + } + _ => panic!("expected fixed-width block"), + } + } + + #[test] + fn test_rle_v2_block_keeps_u8_run_lengths_for_v2_2() { + let field = create_test_field("dict_indices", DataType::Int32); + let values = vec![42i32; 70_000]; + let mut block = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 70_000, + block_info: BlockInfo::default(), + }; + block.compute_stat(); + let data = DataBlock::FixedWidth(block); + + let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) + .with_version(LanceFileVersion::V2_2); + let (_compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); + assert_eq!(rle_run_length_bits(&encoding), 8); + } + #[test] fn test_rle_block_used_for_version_v2_2() { let field = create_test_field("test_repdef", DataType::UInt16); diff --git a/rust/lance-encoding/src/decoder.rs b/rust/lance-encoding/src/decoder.rs index bb46e725098..aea3575dcb1 100644 --- a/rust/lance-encoding/src/decoder.rs +++ b/rust/lance-encoding/src/decoder.rs @@ -1290,7 +1290,7 @@ impl DecodeBatchScheduler { /// * `ranges` - The ranges of rows to load /// * `sink` - A channel to send the decode tasks /// * `scheduler` An I/O scheduler to issue I/O requests - #[instrument(skip_all)] + #[instrument(level = "debug", skip_all)] pub fn schedule_ranges( &mut self, ranges: &[Range], @@ -1326,7 +1326,7 @@ impl DecodeBatchScheduler { /// * `range` - The range of rows to load /// * `sink` - A channel to send the decode tasks /// * `scheduler` An I/O scheduler to issue I/O requests - #[instrument(skip_all)] + #[instrument(level = "debug", skip_all)] pub fn schedule_range( &mut self, range: Range, @@ -1508,34 +1508,41 @@ impl BatchDecodeStream { pub fn into_stream(self) -> BoxStream<'static, ReadBatchTask> { let stream = futures::stream::unfold(self, |mut slf| async move { - let next_task = slf.next_batch_task().await; - let next_task = next_task.transpose().map(|next_task| { - let num_rows = next_task.as_ref().map(|t| t.num_rows).unwrap_or(0); - let emitted_batch_size_warning = slf.emitted_batch_size_warning.clone(); - let task = async move { - let next_task = next_task?; - // Real decode work happens inside into_batch, which can block the current - // thread for a long time. By spawning it as a new task, we allow Tokio's - // worker threads to keep making progress. - let (batch, _data_size) = - tokio::spawn( - async move { next_task.into_batch(emitted_batch_size_warning) }, - ) + let next_task = match slf.next_batch_task().await { + Ok(Some(next_task)) => next_task, + Ok(None) => return None, + Err(err) => { + slf.rows_remaining = 0; + return Some(( + ReadBatchTask { + task: async move { Err(err) }.boxed(), + num_rows: 0, + }, + slf, + )); + } + }; + let num_rows = next_task.num_rows; + let emitted_batch_size_warning = slf.emitted_batch_size_warning.clone(); + let task = async move { + // Real decode work happens inside into_batch, which can block the current + // thread for a long time. By spawning it as a new task, we allow Tokio's + // worker threads to keep making progress. + let (batch, _data_size) = + tokio::spawn(async move { next_task.into_batch(emitted_batch_size_warning) }) .await .map_err(|err| Error::wrapped(err.into()))??; - Ok(batch) - }; - (task, num_rows) - }); - next_task.map(|(task, num_rows)| { - // This should be true since batch size is u32 - debug_assert!(num_rows <= u32::MAX as u64); - let next_task = ReadBatchTask { + Ok(batch) + }; + // This should be true since batch size is u32 + debug_assert!(num_rows <= u32::MAX as u64); + Some(( + ReadBatchTask { task: task.boxed(), num_rows: num_rows as u32, - }; - (next_task, slf) - }) + }, + slf, + )) }); stream.boxed() } @@ -1640,7 +1647,7 @@ impl BatchDecodeIterator { /// /// Note that `scheduled_need` is cumulative. E.g. this method /// should be called with 5, 10, 15 and not 5, 5, 5 - #[instrument(skip_all)] + #[instrument(level = "debug", skip_all)] fn wait_for_io(&mut self, scheduled_need: u64, to_take: u64) -> Result { while self.rows_scheduled < scheduled_need && !self.messages.is_empty() { let message = self.messages.pop_front().unwrap()?; @@ -1913,54 +1920,60 @@ impl StructuralBatchDecodeStream { pub fn into_stream(self) -> BoxStream<'static, ReadBatchTask> { let stream = futures::stream::unfold(self, |mut slf| async move { - let next_task = slf.next_batch_task().await; - let next_task = next_task.transpose().map(|next_task| { - let num_rows = next_task.as_ref().map(|t| t.num_rows).unwrap_or(0); - let emitted_batch_size_warning = slf.emitted_batch_size_warning.clone(); - let bytes_per_row_feedback = slf.bytes_per_row_feedback.clone(); - // Capture the per-stream policy once so every emitted batch task follows the - // same throughput-vs-overhead choice made by the scheduler. - let spawn_batch_decode_tasks = slf.spawn_batch_decode_tasks; - let task = async move { - let next_task = next_task?; - let (batch, data_size) = if spawn_batch_decode_tasks { - tokio::spawn( - async move { next_task.into_batch(emitted_batch_size_warning) }, - ) + let next_task = match slf.next_batch_task().await { + Ok(Some(next_task)) => next_task, + Ok(None) => return None, + Err(err) => { + slf.rows_remaining = 0; + return Some(( + ReadBatchTask { + task: async move { Err(err) }.boxed(), + num_rows: 0, + }, + slf, + )); + } + }; + let num_rows = next_task.num_rows; + let emitted_batch_size_warning = slf.emitted_batch_size_warning.clone(); + let bytes_per_row_feedback = slf.bytes_per_row_feedback.clone(); + // Capture the per-stream policy once so every emitted batch task follows the + // same throughput-vs-overhead choice made by the scheduler. + let spawn_batch_decode_tasks = slf.spawn_batch_decode_tasks; + let task = async move { + let (batch, data_size) = if spawn_batch_decode_tasks { + tokio::spawn(async move { next_task.into_batch(emitted_batch_size_warning) }) .await .map_err(|err| Error::wrapped(err.into()))?? + } else { + next_task.into_batch(emitted_batch_size_warning)? + }; + let num_rows = batch.num_rows() as u64; + if let Some(bpr) = data_size.checked_div(num_rows) { + let prev = bytes_per_row_feedback.load(Ordering::Relaxed); + let next = if prev == 0 || bpr >= prev { + // First batch or actual size is larger than estimate: + // adopt immediately to avoid OOM. + bpr } else { - next_task.into_batch(emitted_batch_size_warning)? + // Actual size is smaller: degrade gradually toward + // the true value to avoid over-correcting on a + // single anomalous batch. + (prev + bpr) / 2 }; - let num_rows = batch.num_rows() as u64; - if num_rows > 0 { - let bpr = data_size / num_rows; - let prev = bytes_per_row_feedback.load(Ordering::Relaxed); - let next = if prev == 0 || bpr >= prev { - // First batch or actual size is larger than estimate: - // adopt immediately to avoid OOM. - bpr - } else { - // Actual size is smaller: degrade gradually toward - // the true value to avoid over-correcting on a - // single anomalous batch. - (prev + bpr) / 2 - }; - bytes_per_row_feedback.store(next.max(1), Ordering::Relaxed); - } - Ok(batch) - }; - (task, num_rows) - }); - next_task.map(|(task, num_rows)| { - // This should be true since batch size is u32 - debug_assert!(num_rows <= u32::MAX as u64); - let next_task = ReadBatchTask { + bytes_per_row_feedback.store(next.max(1), Ordering::Relaxed); + } + Ok(batch) + }; + // This should be true since batch size is u32 + debug_assert!(num_rows <= u32::MAX as u64); + Some(( + ReadBatchTask { task: task.boxed(), num_rows: num_rows as u32, - }; - (next_task, slf) - }) + }, + slf, + )) }); stream.boxed() } @@ -2915,6 +2928,60 @@ pub async fn decode_batch( // test coalesce indices to ranges mod tests { use super::*; + use crate::previous::decoder::{DecoderReady, LogicalPageDecoder}; + use std::collections::VecDeque; + + #[derive(Debug)] + struct FailingPageDecoder { + page_data_type: DataType, + total_rows: u64, + load_error_message: &'static str, + } + + impl FailingPageDecoder { + fn new( + page_data_type: DataType, + total_rows: u64, + load_error_message: &'static str, + ) -> Self { + Self { + page_data_type, + total_rows, + load_error_message, + } + } + } + + impl LogicalPageDecoder for FailingPageDecoder { + fn wait_for_loaded(&'_ mut self, _rows_needed: u64) -> BoxFuture<'_, Result<()>> { + let load_error_message = self.load_error_message; + async move { Err(Error::io(load_error_message)) }.boxed() + } + + fn rows_loaded(&self) -> u64 { + 0 + } + + fn num_rows(&self) -> u64 { + self.total_rows + } + + fn rows_drained(&self) -> u64 { + 0 + } + + fn drain(&mut self, requested_rows: u64) -> Result { + Err(Error::internal(format!( + "failing page decoder should not be drained after load error \ + (requested_rows={})", + requested_rows + ))) + } + + fn data_type(&self) -> &DataType { + &self.page_data_type + } + } #[test] fn test_read_zero_dimension_fsl_errors_instead_of_panicking() { @@ -2962,6 +3029,100 @@ mod tests { ); } + #[tokio::test] + async fn test_legacy_stream_stops_on_load_error() { + use arrow_schema::Field as ArrowField; + + let rows_per_batch = 1; + let total_rows = 2; + let scheduled_rows = 1; + let page_rows = 1; + let batch_readahead = 2; + let load_error_message = "simulated page load failure"; + let fields = Fields::from(vec![ArrowField::new("vector", DataType::Float32, true)]); + let root_decoder = SimpleStructDecoder::new(fields, total_rows); + let (tx, rx) = unbounded_channel(); + + tx.send(Ok(DecoderMessage { + scheduled_so_far: scheduled_rows, + decoders: vec![MessageType::DecoderReady(DecoderReady { + decoder: Box::new(FailingPageDecoder::new( + DataType::Float32, + page_rows, + load_error_message, + )), + path: VecDeque::from([0]), + })], + })) + .unwrap(); + drop(tx); + + let stream = + BatchDecodeStream::new(rx, rows_per_batch, total_rows, root_decoder).into_stream(); + let mut batches = stream.map(|task| task.task).buffered(batch_readahead); + + let err = batches + .next() + .await + .expect("stream should emit the legacy page-load error") + .unwrap_err(); + assert!( + err.to_string().contains(load_error_message), + "unexpected error: {}", + err + ); + assert!( + batches.next().await.is_none(), + "stream should stop after the legacy page-load error" + ); + } + + #[tokio::test] + async fn test_structural_stream_stops_on_load_error() { + let rows_per_batch = 1; + let total_rows = 2; + let scheduled_rows = 1; + let batch_readahead = 2; + let load_error_message = "simulated page load failure"; + let fields = Fields::from(vec![ArrowField::new("vector", DataType::Float32, true)]); + let root_decoder = StructuralStructDecoder::new(fields, false, /*is_root=*/ true).unwrap(); + let (tx, rx) = unbounded_channel(); + let failed_page = async move { Err(Error::io(load_error_message)) }.boxed(); + + tx.send(Ok(DecoderMessage { + scheduled_so_far: scheduled_rows, + decoders: vec![MessageType::UnloadedPage(UnloadedPageShard(failed_page))], + })) + .unwrap(); + drop(tx); + + let stream = StructuralBatchDecodeStream::new( + rx, + rows_per_batch, + total_rows, + root_decoder, + /*spawn_batch_decode_tasks=*/ true, + None, + ) + .into_stream(); + let mut batches = stream.map(|task| task.task).buffered(batch_readahead); + + let err = batches + .next() + .await + .expect("stream should emit the page-load error") + .unwrap_err(); + assert!( + err.to_string().contains(load_error_message), + "unexpected error: {}", + err + ); + assert!( + batches.next().await.is_none(), + "stream should stop after the page-load error" + ); + } + #[test] fn test_coalesce_indices_to_ranges_with_single_index() { let indices = vec![1]; diff --git a/rust/lance-encoding/src/encodings/logical/blob.rs b/rust/lance-encoding/src/encodings/logical/blob.rs index cad2112bafe..d1798b1ef69 100644 --- a/rust/lance-encoding/src/encodings/logical/blob.rs +++ b/rust/lance-encoding/src/encodings/logical/blob.rs @@ -267,16 +267,11 @@ impl FieldEncoder for BlobV2StructuralEncoder { &mut self, array: ArrayRef, external_buffers: &mut OutOfLineBuffers, - mut repdef: RepDefBuilder, + repdef: RepDefBuilder, row_number: u64, num_rows: u64, ) -> Result> { let struct_arr = array.as_struct(); - if let Some(validity) = struct_arr.nulls() { - repdef.add_validity_bitmap(validity.clone()); - } else { - repdef.add_no_null(struct_arr.len()); - } let kind_col = struct_arr .column_by_name("kind") @@ -403,7 +398,7 @@ impl FieldEncoder for BlobV2StructuralEncoder { let descriptor_array = Arc::new(StructArray::try_new( BLOB_V2_DESC_FIELDS.clone(), children, - None, + struct_arr.nulls().cloned(), )?) as ArrayRef; self.descriptor_encoder.maybe_encode( @@ -535,6 +530,32 @@ mod tests { .await; } + #[tokio::test] + async fn test_blob_round_trip_empty_values() { + // Empty values share size == 0 with nulls in the descriptor layout + // and schedule no read; each must decode to zero-length bytes without + // consuming the read result of a following non-empty blob. Empties + // are placed before payloads so a misassignment corrupts the output + // instead of only exhausting the read iterator. + let blob_metadata = + HashMap::from([(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())]); + + let val1: &[u8] = &vec![1u8; 1024]; + let val2: &[u8] = &vec![2u8; 10240]; + let empty: &[u8] = &[]; + let array = Arc::new(LargeBinaryArray::from(vec![ + Some(empty), + Some(val1), + None, + Some(empty), + Some(val2), + None, + Some(empty), + ])); + + check_round_trip_encoding_of_data(vec![array], &TestCases::default(), blob_metadata).await; + } + #[tokio::test] async fn test_blob_v2_external_round_trip() { let blob_metadata = HashMap::from([( diff --git a/rust/lance-encoding/src/encodings/logical/list.rs b/rust/lance-encoding/src/encodings/logical/list.rs index 153238719ae..250eb476671 100644 --- a/rust/lance-encoding/src/encodings/logical/list.rs +++ b/rust/lance-encoding/src/encodings/logical/list.rs @@ -428,6 +428,96 @@ mod tests { check_basic_random(field).await; } + /// Regression test: a `List>` column written as MULTIPLE + /// batches (chunks) whose flattened leaf values cross a value-page boundary + /// fails to decode with "Max offset N exceeds length of values M" (Arrow + /// error raised by `ListArray::try_new` in `StructuralListDecodeTask::decode`). + /// + /// The trigger (verified against the production file and pylance 7.0.0b12 / + /// 7.0.0 / 9.0.0-beta.10) requires ALL of: + /// 1. >= 2 list layers (`List>`), + /// 2. a leaf large enough to be chunked into multiple value pages, + /// 3. the column written as more than one batch. + /// A single batch of the identical data round-trips fine — which is why the + /// earlier single-chunk version of this test (and the small `test_nested_list` + /// cases) did not catch it. Found in production on the gaming TransNet + /// `dino_embedding_per_frame` column (rectangular 3 x 768 float per row). + /// + /// Each element of the `vec![..]` passed to `check_round_trip_encoding_of_data` + /// is encoded as a separate batch (its own `RepDefBuilder`), so we split the + /// rows into two chunks to exercise the multi-batch repdef accumulation path. + #[rstest] + #[test_log::test(tokio::test)] + async fn test_multipage_nested_float_list( + #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] + structural_encoding: &str, + ) { + use arrow_array::Float32Array; + + // Production shape: 3 inner lists per row, 768 floats each. + let inner_per_row: usize = 3; + let inner_len: usize = 768; + // Two chunks (batches) -> two pages; a read batch that spans the page + // boundary is where the multi-page outer-offset bug triggered. A single + // [2731] chunk (one page) decodes fine, which is why this needs >= 2. + let chunk_rows: &[usize] = &[1366, 1365]; + + let make_chunk = |start_row: usize, num_rows: usize| -> Arc { + let total_inner = num_rows * inner_per_row; + let total_values = total_inner * inner_len; + let values = Float32Array::from( + (0..total_values) + .map(|i| (start_row + i) as f32) + .collect::>(), + ); + let inner_offsets = ScalarBuffer::::from( + (0..=total_inner) + .map(|i| (i * inner_len) as i32) + .collect::>(), + ); + let inner_list = ListArray::new( + Arc::new(Field::new("item", DataType::Float32, true)), + OffsetBuffer::new(inner_offsets), + Arc::new(values), + None, + ); + let outer_offsets = ScalarBuffer::::from( + (0..=num_rows) + .map(|i| (i * inner_per_row) as i32) + .collect::>(), + ); + Arc::new(ListArray::new( + Arc::new(Field::new( + "item", + DataType::List(Arc::new(Field::new("item", DataType::Float32, true))), + true, + )), + OffsetBuffer::new(outer_offsets), + Arc::new(inner_list), + None, + )) + }; + + let mut start = 0; + let chunks: Vec> = chunk_rows + .iter() + .map(|&n| { + let c = make_chunk(start, n); + start += n; + c + }) + .collect(); + + let mut field_metadata = HashMap::new(); + field_metadata.insert( + STRUCTURAL_ENCODING_META_KEY.to_string(), + structural_encoding.into(), + ); + + let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + check_round_trip_encoding_of_data(chunks, &test_cases, field_metadata).await; + } + #[test_log::test(tokio::test)] async fn test_list_struct_list() { let struct_type = DataType::Struct(Fields::from(vec![Field::new( diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 9b506359e55..6e5cb4cd66b 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -58,7 +58,7 @@ use crate::{ use crate::{ repdef::{ CompositeRepDefUnraveler, ControlWordIterator, ControlWordParser, DefinitionInterpretation, - RepDefSlicer, SerializedRepDefs, StructuralPagePlan, build_control_word_iterator, + MiniBlockRepDefBudget, RepDefSlicer, SerializedRepDefs, build_control_word_iterator, }, utils::accumulation::AccumulationQueue, }; @@ -3660,7 +3660,18 @@ impl StructuralFieldDecoder for StructuralPrimitiveFieldDecoder { let mut remaining = num_rows; let mut tasks = Vec::new(); while remaining > 0 { - let cur_page = self.page_decoders.front_mut().unwrap(); + let queued_pages = self.page_decoders.len(); + let Some(cur_page) = self.page_decoders.front_mut() else { + return Err(Error::internal(format!( + "Primitive decoder missing page decoder while draining field '{}' (data_type={:?}, requested_rows={}, remaining_rows={}, rows_drained_in_current={}, queued_pages={})", + self.field.name(), + self.field.data_type(), + num_rows, + remaining, + self.rows_drained_in_current, + queued_pages + ))); + }; let num_in_page = cur_page.num_rows() - self.rows_drained_in_current; let to_take = num_in_page.min(remaining); @@ -3776,7 +3787,7 @@ struct DictEncodingBudget { max_encoded_size: usize, } -// A primitive page after optional structural splitting. +// A primitive page after applying the dense mini-block rep/def budget. struct PrimitivePageData { // Arrow leaf arrays that contain this page's visible values. arrays: Vec, @@ -3786,8 +3797,8 @@ struct PrimitivePageData { row_number: u64, // Number of top-level rows in this page. num_rows: u64, - // Present when one top-level row is too large for one miniblock rep/def chunk. - unsplittable_miniblock_levels: Option, + // Present when one top-level row is too large for one mini-block rep/def page. + single_row_miniblock_repdef_levels: Option, } // Immutable encoder state shared by per-page encode tasks. @@ -5281,33 +5292,33 @@ impl PrimitiveStructuralEncoder { Ok(sliced) } - fn split_structural_pages_for_miniblock_budget( + fn split_pages_for_miniblock_repdef_budget( arrays: Vec, repdef: SerializedRepDefs, - plan: StructuralPagePlan, + budget: MiniBlockRepDefBudget, row_number: u64, num_rows: u64, ) -> Result> { - if plan == StructuralPagePlan::Fits { + if budget == MiniBlockRepDefBudget::WithinBudget { return Ok(vec![PrimitivePageData { arrays, repdef, row_number, num_rows, - unsplittable_miniblock_levels: None, + single_row_miniblock_repdef_levels: None, }]); } - if let StructuralPagePlan::UnsplittableOverBudget(num_levels) = plan { + if let MiniBlockRepDefBudget::SingleRowOverBudget(num_levels) = budget { return Ok(vec![PrimitivePageData { arrays, repdef, row_number, num_rows, - unsplittable_miniblock_levels: Some(num_levels), + single_row_miniblock_repdef_levels: Some(num_levels), }]); } - let StructuralPagePlan::Split(splits) = plan else { + let MiniBlockRepDefBudget::RequiresPageSplit(splits) = budget else { unreachable!(); }; @@ -5320,7 +5331,7 @@ impl PrimitiveStructuralEncoder { repdef, row_number: row_number + split.row_start, num_rows: split.num_rows, - unsplittable_miniblock_levels: None, + single_row_miniblock_repdef_levels: None, }); } Ok(pages) @@ -5342,7 +5353,7 @@ impl PrimitiveStructuralEncoder { repdef, row_number, num_rows, - unsplittable_miniblock_levels, + single_row_miniblock_repdef_levels, } = page; let num_values = arrays.iter().map(|arr| arr.len() as u64).sum(); @@ -5425,7 +5436,7 @@ impl PrimitiveStructuralEncoder { ); } - if let Some(num_levels) = unsplittable_miniblock_levels { + if let Some(num_levels) = single_row_miniblock_repdef_levels { let requested_encoding = encoding_metadata .get(STRUCTURAL_ENCODING_META_KEY) .map(|requested| requested.to_lowercase()); @@ -5652,16 +5663,17 @@ impl PrimitiveStructuralEncoder { let num_values = arrays.iter().map(|arr| arr.len() as u64).sum(); let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity()); let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty()); - let (repdef, structural_plan) = RepDefBuilder::serialize_with_structural_plan( - repdefs, - miniblock::max_repdef_levels_per_chunk, - num_rows, - num_values, - )?; - let pages = Self::split_structural_pages_for_miniblock_budget( + let (repdef, miniblock_repdef_budget) = + RepDefBuilder::serialize_with_miniblock_repdef_budget( + repdefs, + miniblock::max_repdef_levels_per_chunk, + num_rows, + num_values, + )?; + let pages = Self::split_pages_for_miniblock_repdef_budget( arrays, repdef, - structural_plan, + miniblock_repdef_budget, row_number, num_rows, )?; @@ -5794,9 +5806,9 @@ mod tests { STRUCTURAL_ENCODING_MINIBLOCK, }; use crate::data::BlockInfo; - use crate::decoder::PageEncoding; + use crate::decoder::{PageEncoding, StructuralFieldDecoder}; use crate::encodings::logical::primitive::{ - ChunkDrainInstructions, PrimitiveStructuralEncoder, + ChunkDrainInstructions, PrimitiveStructuralEncoder, StructuralPrimitiveFieldDecoder, }; use crate::format::ProtobufUtils21; use crate::format::pb21; @@ -5805,7 +5817,7 @@ mod tests { use crate::testing::{TestCases, check_round_trip_encoding_of_data}; use crate::version::LanceFileVersion; use arrow_array::{ArrayRef, Int8Array, StringArray}; - use arrow_schema::DataType; + use arrow_schema::{DataType, Field as ArrowField}; use std::collections::HashMap; use std::{collections::VecDeque, sync::Arc}; @@ -5829,6 +5841,33 @@ mod tests { assert!((!PrimitiveStructuralEncoder::is_narrow(&block))); } + #[test] + fn test_primitive_decoder_empty_page_queue_returns_error() { + let field = Arc::new(ArrowField::new("vector", DataType::Float32, true)); + let mut decoder = StructuralPrimitiveFieldDecoder::new(&field, false); + + let err = decoder.drain(1).unwrap_err(); + assert!( + matches!(&err, lance_core::Error::Internal { .. }), + "expected internal error, got: {err:?}" + ); + let message = err.to_string(); + for expected in [ + "Primitive decoder missing page decoder", + "field 'vector'", + "data_type=Float32", + "requested_rows=1", + "remaining_rows=1", + "rows_drained_in_current=0", + "queued_pages=0", + ] { + assert!( + message.contains(expected), + "expected error to contain {expected:?}, got: {message}" + ); + } + } + #[test] fn test_fullzip_fixed_rejects_non_byte_aligned_values() { let fixed = FixedWidthDataBlock { diff --git a/rust/lance-encoding/src/encodings/logical/primitive/blob.rs b/rust/lance-encoding/src/encodings/logical/primitive/blob.rs index eed3e584b7e..52a7039b2b6 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/blob.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/blob.rs @@ -258,7 +258,9 @@ impl BlobPageScheduler { let bytes = read_fut.await?; let mut bytes_iter = bytes.into_iter(); for blob in loaded_blobs.iter_mut() { - if blob.def == 0 { + // Empty values have def == 0 too but scheduled no read; their + // bytes were set at scheduling time. + if blob.def == 0 && blob.bytes.is_none() { blob.set_bytes(bytes_iter.next().expect_ok()?); } } @@ -364,7 +366,17 @@ impl StructuralPageScheduler for BlobPageScheduler { if size == 0 { let rep = (position & 0xFFFF) as u16; let def = ((position >> 16) & 0xFFFF) as u16; - loaded_blobs.push(LoadedBlob::new(rep, def)); + let mut blob = LoadedBlob::new(rep, def); + if def == 0 { + // A size-0 descriptor with definition level 0 is a + // valid, empty value (nulls carry their non-zero + // packed rep/def levels in `position`). No read is + // scheduled for it, so it gets its zero-length bytes + // here rather than consuming another blob's read + // result in the load task. + blob.set_bytes(Bytes::new()); + } + loaded_blobs.push(blob); } else { loaded_blobs.push(LoadedBlob::new(0, 0)); ranges_to_read.push(position..(position + size)); diff --git a/rust/lance-encoding/src/encodings/physical/bitpacking.rs b/rust/lance-encoding/src/encodings/physical/bitpacking.rs index 8ebdcc13c56..be0b747e7dc 100644 --- a/rust/lance-encoding/src/encodings/physical/bitpacking.rs +++ b/rust/lance-encoding/src/encodings/physical/bitpacking.rs @@ -241,6 +241,15 @@ impl MiniBlockDecompressor for InlineBitpacking { fn decompress(&self, data: Vec, num_values: u64) -> Result { assert_eq!(data.len(), 1); let data = data.into_iter().next().unwrap(); + if num_values == 0 { + // Empty mini-blocks have no inline bit-width header to decode. + return Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::empty(), + bits_per_value: self.uncompressed_bit_width, + num_values: 0, + block_info: BlockInfo::new(), + })); + } match self.uncompressed_bit_width { 8 => Self::unchunk::(data, num_values), 16 => Self::unchunk::(data, num_values), @@ -528,15 +537,36 @@ mod test { use arrow_array::{Array, Int8Array, Int64Array}; use arrow_schema::DataType; + use rstest::rstest; - use super::{ELEMS_PER_CHUNK, bitpack_out_of_line, unpack_out_of_line}; + use super::{ELEMS_PER_CHUNK, InlineBitpacking, bitpack_out_of_line, unpack_out_of_line}; use crate::{ buffer::LanceBuffer, - data::{BlockInfo, FixedWidthDataBlock}, + compression::MiniBlockDecompressor, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, testing::{TestCases, check_round_trip_encoding_of_data}, version::LanceFileVersion, }; + #[rstest] + #[case::u8(8)] + #[case::u16(16)] + #[case::u32(32)] + #[case::u64(64)] + fn test_inline_bitpacking_decompress_empty_miniblock(#[case] bit_width: u64) { + let decompressor = InlineBitpacking::new(bit_width); + let decompressed = + MiniBlockDecompressor::decompress(&decompressor, vec![LanceBuffer::empty()], 0) + .unwrap(); + + let DataBlock::FixedWidth(block) = decompressed else { + panic!("Expected FixedWidth block"); + }; + assert_eq!(block.bits_per_value, bit_width); + assert_eq!(block.num_values, 0); + assert_eq!(block.data.len(), 0); + } + #[test_log::test(tokio::test)] async fn test_miniblock_bitpack() { let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); diff --git a/rust/lance-encoding/src/encodings/physical/general.rs b/rust/lance-encoding/src/encodings/physical/general.rs index 4d58f72e71a..53c61928870 100644 --- a/rust/lance-encoding/src/encodings/physical/general.rs +++ b/rust/lance-encoding/src/encodings/physical/general.rs @@ -161,7 +161,7 @@ mod tests { // Small data with RLE - should not compress due to size threshold TestCase { name: "small_rle_data", - inner_encoder: Box::new(RleEncoder), + inner_encoder: Box::new(RleEncoder::new()), compression: CompressionConfig { scheme: CompressionScheme::Lz4, level: None, @@ -173,7 +173,7 @@ mod tests { // Large repeated data with RLE + LZ4 TestCase { name: "large_rle_lz4", - inner_encoder: Box::new(RleEncoder), + inner_encoder: Box::new(RleEncoder::new()), compression: CompressionConfig { scheme: CompressionScheme::Lz4, level: None, @@ -185,7 +185,7 @@ mod tests { // Large repeated data with RLE + Zstd TestCase { name: "large_rle_zstd", - inner_encoder: Box::new(RleEncoder), + inner_encoder: Box::new(RleEncoder::new()), compression: CompressionConfig { scheme: CompressionScheme::Zstd, level: Some(3), @@ -403,7 +403,7 @@ mod tests { // Test that small buffers don't get compressed let small_test = TestCase { name: "small_buffer_no_compression", - inner_encoder: Box::new(RleEncoder), + inner_encoder: Box::new(RleEncoder::new()), compression: CompressionConfig { scheme: CompressionScheme::Lz4, level: None, @@ -496,7 +496,7 @@ mod tests { // RLE produces 2 buffers (values and lengths), test that both are handled correctly let data = create_repeated_i32_block(vec![1; 100]); let compressor = GeneralMiniBlockCompressor::new( - Box::new(RleEncoder), + Box::new(RleEncoder::new()), CompressionConfig { scheme: CompressionScheme::Lz4, level: None, @@ -519,7 +519,7 @@ mod tests { // Test case 1: 32-bit RLE data let test_32 = TestCase { name: "rle_32bit_with_general_wrapper", - inner_encoder: Box::new(RleEncoder), + inner_encoder: Box::new(RleEncoder::new()), compression: CompressionConfig { scheme: CompressionScheme::Lz4, level: None, @@ -532,7 +532,7 @@ mod tests { // For 32-bit RLE, the compression strategy should automatically wrap it // Let's directly test the compressor let compressor = GeneralMiniBlockCompressor::new( - Box::new(RleEncoder), + Box::new(RleEncoder::new()), CompressionConfig { scheme: CompressionScheme::Lz4, level: None, @@ -589,7 +589,7 @@ mod tests { let block_64 = DataBlock::from_array(array_64); let compressor_64 = GeneralMiniBlockCompressor::new( - Box::new(RleEncoder), + Box::new(RleEncoder::new()), CompressionConfig { scheme: CompressionScheme::Lz4, level: None, diff --git a/rust/lance-encoding/src/encodings/physical/packed.rs b/rust/lance-encoding/src/encodings/physical/packed.rs index 3ade6a70818..ad2221dffed 100644 --- a/rust/lance-encoding/src/encodings/physical/packed.rs +++ b/rust/lance-encoding/src/encodings/physical/packed.rs @@ -323,7 +323,7 @@ impl PerValueCompressor for PackedStructVariablePerValueEncoder { let mut field_data = Vec::with_capacity(self.fields.len()); let mut field_metadata = Vec::with_capacity(self.fields.len()); - for (field, child_block) in self.fields.iter().zip(struct_block.children.into_iter()) { + for (field, child_block) in self.fields.iter().zip(struct_block.children) { let compressor = crate::compression::CompressionStrategy::create_per_value( &self.strategy, field, @@ -688,7 +688,7 @@ impl VariablePerValueDecompressor for PackedStructVariablePerValueDecompressor { } let mut children = Vec::with_capacity(self.fields.len()); - for (field, accumulator) in self.fields.iter().zip(accumulators.into_iter()) { + for (field, accumulator) in self.fields.iter().zip(accumulators) { match (field, accumulator) { ( VariablePackedStructFieldDecoder { diff --git a/rust/lance-encoding/src/encodings/physical/rle.rs b/rust/lance-encoding/src/encodings/physical/rle.rs index 88e27bf954f..aaff104f0f8 100644 --- a/rust/lance-encoding/src/encodings/physical/rle.rs +++ b/rust/lance-encoding/src/encodings/physical/rle.rs @@ -10,7 +10,7 @@ //! RLE uses a dual-buffer format to store compressed data: //! //! - **Values Buffer**: Stores unique values in their original data type -//! - **Lengths Buffer**: Stores the repeat count for each value as u8 +//! - **Lengths Buffer**: Stores the repeat count for each value as u8, u16, or u32 //! //! ### Example //! @@ -18,13 +18,13 @@ //! //! Encoded as: //! - Values buffer: `[1, 2, 3]` (3 × 4 bytes for i32) -//! - Lengths buffer: `[3, 2, 4]` (3 × 1 byte for u8) +//! - Lengths buffer: `[3, 2, 4]` (3 × 1 byte for u8 in compatibility mode) //! //! ### Long Run Handling //! -//! When a run exceeds 255 values, it is split into multiple runs of 255 -//! followed by a final run with the remainder. For example, a run of 1000 -//! identical values becomes 4 runs: [255, 255, 255, 235]. +//! In compatibility mode, when a run exceeds 255 values, it is split into multiple +//! runs of 255 followed by a final run with the remainder. RLE v2 can use u16 or +//! u32 run lengths to reduce this splitting. //! //! ## Supported Types //! @@ -65,18 +65,319 @@ use crate::encodings::logical::primitive::miniblock::{ MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressor, }; +use crate::encodings::physical::block::{CompressionConfig, GeneralBufferCompressor}; use crate::format::ProtobufUtils21; use crate::format::pb21::CompressiveEncoding; use lance_core::{Error, Result}; +/// Width used to encode RLE run lengths. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RunLengthWidth { + /// Compatibility mode. Runs longer than 255 values are split. + U8, + /// RLE v2 mode for runs up to 65,535 values per entry. + U16, + /// RLE v2 mode for runs up to 4,294,967,295 values per entry. + U32, +} + +impl RunLengthWidth { + pub(crate) fn from_bits(bits_per_value: u64) -> Option { + match bits_per_value { + 8 => Some(Self::U8), + 16 => Some(Self::U16), + 32 => Some(Self::U32), + _ => None, + } + } + + pub(crate) fn bits_per_value(self) -> u64 { + match self { + Self::U8 => 8, + Self::U16 => 16, + Self::U32 => 32, + } + } + + fn bytes_per_value(self) -> usize { + match self { + Self::U8 => 1, + Self::U16 => 2, + Self::U32 => 4, + } + } + + fn max_run_length(self) -> u64 { + match self { + Self::U8 => u8::MAX as u64, + Self::U16 => u16::MAX as u64, + Self::U32 => u32::MAX as u64, + } + } + + fn write_length(self, length: u64, dst: &mut Vec) { + match self { + Self::U8 => dst.push(length as u8), + Self::U16 => dst.extend_from_slice(&(length as u16).to_le_bytes()), + Self::U32 => dst.extend_from_slice(&(length as u32).to_le_bytes()), + } + } + + fn read_length(self, bytes: &[u8]) -> u64 { + match self { + Self::U8 => bytes[0] as u64, + Self::U16 => u16::from_le_bytes([bytes[0], bytes[1]]) as u64, + Self::U32 => u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as u64, + } + } +} + +const RUN_LENGTH_WIDTHS: [RunLengthWidth; 3] = + [RunLengthWidth::U8, RunLengthWidth::U16, RunLengthWidth::U32]; + +/// Select the lowest-cost run length width from precomputed entry counts. +pub(crate) fn select_run_length_width_from_entries( + entries: &[u64], + bits_per_value: u64, +) -> Result<(RunLengthWidth, u128)> { + if entries.len() != RUN_LENGTH_WIDTHS.len() { + return Err(Error::invalid_input_source( + format!( + "RLE run length entry statistics must have {} values, got {}", + RUN_LENGTH_WIDTHS.len(), + entries.len() + ) + .into(), + )); + } + + if !matches!(bits_per_value, 8 | 16 | 32 | 64) { + return Err(Error::invalid_input_source( + format!("RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}") + .into(), + )); + } + + let mut best_width = RUN_LENGTH_WIDTHS[0]; + let mut best_cost = rle_encoded_size_from_entries(entries[0], bits_per_value, best_width); + for (&width, &entry_count) in RUN_LENGTH_WIDTHS.iter().zip(entries.iter()).skip(1) { + let cost = rle_encoded_size_from_entries(entry_count, bits_per_value, width); + if cost < best_cost { + best_width = width; + best_cost = cost; + } + } + + Ok((best_width, best_cost)) +} + +pub(crate) fn rle_encoded_size_from_entries( + entry_count: u64, + bits_per_value: u64, + run_length_width: RunLengthWidth, +) -> u128 { + let bytes_per_value = (bits_per_value / 8) as u128; + let bytes_per_length = run_length_width.bytes_per_value() as u128; + (entry_count as u128) * (bytes_per_value + bytes_per_length) +} + +pub(crate) fn run_length_width_index(run_length_width: RunLengthWidth) -> usize { + match run_length_width { + RunLengthWidth::U8 => 0, + RunLengthWidth::U16 => 1, + RunLengthWidth::U32 => 2, + } +} + +pub(crate) fn select_run_length_width( + data: &LanceBuffer, + num_values: u64, + bits_per_value: u64, + max_segment_values: Option, +) -> Result<(RunLengthWidth, u128)> { + let entries = collect_run_length_entries(data, num_values, bits_per_value, max_segment_values)?; + select_run_length_width_from_entries(&entries, bits_per_value) +} + +pub(crate) fn rle_encoded_size( + data: &LanceBuffer, + num_values: u64, + bits_per_value: u64, + max_segment_values: Option, + run_length_width: RunLengthWidth, +) -> Result { + let entries = collect_run_length_entries(data, num_values, bits_per_value, max_segment_values)?; + let width_idx = run_length_width_index(run_length_width); + Ok(rle_encoded_size_from_entries( + entries[width_idx], + bits_per_value, + run_length_width, + )) +} + +fn collect_run_length_entries( + data: &LanceBuffer, + num_values: u64, + bits_per_value: u64, + max_segment_values: Option, +) -> Result<[u64; 3]> { + let num_values = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("RLE num_values does not fit in usize: {num_values}").into(), + ) + })?; + + macro_rules! collect_entries { + ($ty:ty) => {{ + let type_size = std::mem::size_of::<$ty>(); + let expected_bytes = num_values.checked_mul(type_size).ok_or_else(|| { + Error::invalid_input_source( + format!( + "RLE input byte length overflow: {num_values} values of {type_size} bytes" + ) + .into(), + ) + })?; + if data.len() != expected_bytes { + return Err(Error::invalid_input_source( + format!( + "RLE input data size mismatch: {} bytes for {} values of {} bytes", + data.len(), + num_values, + type_size + ) + .into(), + )); + } + let values = data.borrow_to_typed_slice::<$ty>(); + let values = values.get(..num_values).ok_or_else(|| { + Error::invalid_input_source( + format!( + "RLE data has {} values but {} were expected", + values.len(), + num_values + ) + .into(), + ) + })?; + Ok(collect_run_length_entries_from_slice( + values, + max_segment_values, + )) + }}; + } + + match bits_per_value { + 8 => collect_entries!(u8), + 16 => collect_entries!(u16), + 32 => collect_entries!(u32), + 64 => collect_entries!(u64), + _ => Err(Error::invalid_input_source( + format!("RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}") + .into(), + )), + } +} + +fn collect_run_length_entries_from_slice( + values: &[T], + max_segment_values: Option, +) -> [u64; 3] { + if values.is_empty() { + return [0; 3]; + } + + let mut entries = [0u64; 3]; + let mut prev = values[0]; + let mut current_length = 1u64; + + for &value in &values[1..] { + if value != prev { + accumulate_run_length_entries(current_length, max_segment_values, &mut entries); + prev = value; + current_length = 1; + } else { + current_length += 1; + } + } + accumulate_run_length_entries(current_length, max_segment_values, &mut entries); + + entries +} + +pub(crate) fn accumulate_run_length_entries( + run_length: u64, + max_segment_values: Option, + entries: &mut [u64; 3], +) { + let max_segment_values = max_segment_values.unwrap_or(run_length).max(1); + let mut remaining = run_length; + while remaining > 0 { + let segment = remaining.min(max_segment_values); + for (idx, width) in RUN_LENGTH_WIDTHS.iter().enumerate() { + let entry_count = segment.div_ceil(width.max_run_length()); + entries[idx] = entries[idx].saturating_add(entry_count); + } + remaining -= segment; + } +} + /// RLE encoder for miniblock format -#[derive(Debug, Default)] -pub struct RleEncoder; +#[derive(Debug)] +pub struct RleEncoder { + run_length_width: RunLengthWidth, + values_compression: Option, + run_lengths_compression: Option, + use_child_bitpacking: bool, +} + +#[derive(Clone)] +struct RleChildCandidate { + encoding: CompressiveEncoding, + data: LanceBuffer, + chunk_sizes: Vec, + size: usize, + requires_num_values: bool, +} + +impl Default for RleEncoder { + fn default() -> Self { + Self::new() + } +} impl RleEncoder { pub fn new() -> Self { - Self + Self { + run_length_width: RunLengthWidth::U8, + values_compression: None, + run_lengths_compression: None, + use_child_bitpacking: false, + } + } + + pub(crate) fn with_run_length_width(run_length_width: RunLengthWidth) -> Self { + Self { + run_length_width, + values_compression: None, + run_lengths_compression: None, + use_child_bitpacking: false, + } + } + + pub(crate) fn with_child_encoding( + run_length_width: RunLengthWidth, + values_compression: Option, + run_lengths_compression: Option, + use_child_bitpacking: bool, + ) -> Self { + Self { + run_length_width, + values_compression, + run_lengths_compression, + use_child_bitpacking, + } } fn encode_data( @@ -89,17 +390,23 @@ impl RleEncoder { return Ok((Vec::new(), Vec::new())); } + let num_values = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("RLE num_values does not fit in usize: {num_values}").into(), + ) + })?; let bytes_per_value = (bits_per_value / 8) as usize; + let bytes_per_length = self.run_length_width.bytes_per_value(); // Pre-allocate global buffers with estimated capacity // Assume average compression ratio of ~10:1 (10 values per run) - let estimated_runs = num_values as usize / 10; + let estimated_runs = num_values / 10; let mut all_values = Vec::with_capacity(estimated_runs * bytes_per_value); - let mut all_lengths = Vec::with_capacity(estimated_runs); + let mut all_lengths = Vec::with_capacity(estimated_runs * bytes_per_length); let mut chunks = Vec::new(); let mut offset = 0usize; - let mut values_remaining = num_values as usize; + let mut values_remaining = num_values; while values_remaining > 0 { let values_start = all_values.len(); @@ -134,7 +441,14 @@ impl RleEncoder { &mut all_values, &mut all_lengths, ), - _ => unreachable!("RLE encoding bits_per_value must be 8, 16, 32 or 64"), + _ => { + return Err(Error::invalid_input_source( + format!( + "RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}" + ) + .into(), + )); + } }; if values_processed == 0 { @@ -175,22 +489,75 @@ impl RleEncoder { )) } - /// Encodes a chunk of data using RLE compression with dynamic boundary detection. - /// - /// This function processes values sequentially, detecting runs (sequences of identical values) - /// and encoding them as (value, length) pairs. It dynamically determines whether this chunk - /// should be the last chunk based on how many values were processed. - /// - /// # Key Features: - /// - Tracks byte usage to ensure we don't exceed MAX_MINIBLOCK_BYTES - /// - Maintains power-of-2 checkpoints for non-last chunks - /// - Splits long runs (>255) into multiple entries - /// - Dynamically determines if this is the last chunk - /// - /// # Returns: - /// - num_runs: Number of runs encoded - /// - values_processed: Number of input values processed - /// - is_last_chunk: Whether this chunk processed all remaining values + fn encode_block_data( + &self, + data: &LanceBuffer, + num_values: u64, + bits_per_value: u64, + ) -> Result> { + match bits_per_value { + 8 => self.encode_block_data_generic::(data, num_values), + 16 => self.encode_block_data_generic::(data, num_values), + 32 => self.encode_block_data_generic::(data, num_values), + 64 => self.encode_block_data_generic::(data, num_values), + _ => Err(Error::invalid_input_source( + format!( + "RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}" + ) + .into(), + )), + } + } + + fn encode_block_data_generic( + &self, + data: &LanceBuffer, + num_values: u64, + ) -> Result> + where + T: bytemuck::Pod + PartialEq + Copy + ArrowNativeType, + { + let num_values = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("RLE num_values does not fit in usize: {num_values}").into(), + ) + })?; + let type_size = std::mem::size_of::(); + let expected_bytes = num_values.checked_mul(type_size).ok_or_else(|| { + Error::invalid_input_source( + format!("RLE input byte length overflow: {num_values} values of {type_size} bytes") + .into(), + ) + })?; + if data.len() != expected_bytes { + return Err(Error::invalid_input_source( + format!( + "RLE input data size mismatch: {} bytes for {} values of {} bytes", + data.len(), + num_values, + type_size + ) + .into(), + )); + } + if num_values == 0 { + return Ok(vec![LanceBuffer::empty(), LanceBuffer::empty()]); + } + + let values_ref = data.borrow_to_typed_slice::(); + let values = values_ref.as_ref(); + let estimated_runs = num_values / 10; + let mut all_values = Vec::with_capacity(estimated_runs * type_size); + let mut all_lengths = + Vec::with_capacity(estimated_runs * self.run_length_width.bytes_per_value()); + self.encode_values(values, &mut all_values, &mut all_lengths); + Ok(vec![ + LanceBuffer::from(all_values), + LanceBuffer::from(all_lengths), + ]) + } + + /// Encodes the largest valid mini-block prefix from `offset`. fn encode_chunk_rolling( &self, data: &LanceBuffer, @@ -203,7 +570,6 @@ impl RleEncoder { T: bytemuck::Pod + PartialEq + Copy + std::fmt::Debug + ArrowNativeType, { let type_size = std::mem::size_of::(); - let chunk_start = offset * type_size; let max_by_count = *MAX_MINIBLOCK_VALUES as usize; let max_values = values_remaining.min(max_by_count); @@ -217,112 +583,101 @@ impl RleEncoder { let chunk_buffer = data.slice_with_length(chunk_start, chunk_len); let typed_data_ref = chunk_buffer.borrow_to_typed_slice::(); let typed_data: &[T] = typed_data_ref.as_ref(); + let max_values = max_values.min(typed_data.len()); if typed_data.is_empty() { return (0, 0, false); } - // Record starting positions for this chunk let values_start = all_values.len(); + let all_remaining_values_fit = values_remaining <= max_by_count; + let encoded_size = self.encoded_size(&typed_data[..max_values]); + let (values_to_encode, is_last_chunk) = if all_remaining_values_fit + && encoded_size <= MAX_MINIBLOCK_BYTES as usize + { + (max_values, true) + } else if let Some(values_to_encode) = self.largest_power_of_two_prefix::(typed_data) { + (values_to_encode, false) + } else { + return (0, 0, false); + }; - let mut current_value = typed_data[0]; - let mut current_length = 1u64; - let mut bytes_used = 0usize; - let mut total_values_encoded = 0usize; // Track total encoded values + self.encode_values(&typed_data[..values_to_encode], all_values, all_lengths); - // Power-of-2 checkpoints for ensuring non-last chunks have valid sizes. - // - // We start from a slightly larger minimum checkpoint for smaller types since - // they encode more compactly and are less likely to hit MAX_MINIBLOCK_BYTES. - let min_checkpoint_log2 = match type_size { - 1 => 8, // 256 - 2 => 7, // 128 - _ => 6, // 64 - }; - let max_checkpoint_log2 = (values_remaining.min(*MAX_MINIBLOCK_VALUES as usize)) - .next_power_of_two() - .ilog2(); - let mut checkpoint_log2 = min_checkpoint_log2; + let num_runs = (all_values.len() - values_start) / type_size; + (num_runs, values_to_encode, is_last_chunk) + } + + fn largest_power_of_two_prefix(&self, values: &[T]) -> Option + where + T: bytemuck::Pod + PartialEq + Copy, + { + let max_prefix = values.len().min(*MAX_MINIBLOCK_VALUES as usize); + let mut prefix = 1usize << max_prefix.ilog2(); + while prefix > 1 { + if self.encoded_size(&values[..prefix]) <= MAX_MINIBLOCK_BYTES as usize { + return Some(prefix); + } + prefix >>= 1; + } + None + } + + fn encoded_size(&self, values: &[T]) -> usize + where + T: bytemuck::Pod + PartialEq + Copy, + { + if values.is_empty() { + return 0; + } - // Save state at checkpoints so we can roll back if needed - let mut last_checkpoint_state = None; + let mut current_value = values[0]; + let mut current_length = 1u64; + let mut encoded_size = 0usize; - for &value in typed_data[1..].iter() { + for &value in values.iter().skip(1) { if value == current_value { current_length += 1; } else { - // Calculate space needed (may need multiple u8s if run > 255) - let run_chunks = current_length.div_ceil(255) as usize; - let bytes_needed = run_chunks * (type_size + 1); - - // Stop if adding this run would exceed byte limit - if bytes_used + bytes_needed > MAX_MINIBLOCK_BYTES as usize { - if let Some((val_pos, len_pos, _, checkpoint_values)) = last_checkpoint_state { - // Roll back to last power-of-2 checkpoint - all_values.truncate(val_pos); - all_lengths.truncate(len_pos); - let num_runs = (val_pos - values_start) / type_size; - return (num_runs, checkpoint_values, false); - } - break; - } - - bytes_used += self.add_run(¤t_value, current_length, all_values, all_lengths); - total_values_encoded += current_length as usize; + encoded_size += self.run_size::(current_length); current_value = value; current_length = 1; } - - // Check if we reached a power-of-2 checkpoint. - while checkpoint_log2 <= max_checkpoint_log2 { - let checkpoint_values = 1usize << checkpoint_log2; - if checkpoint_values > values_remaining || total_values_encoded < checkpoint_values - { - break; - } - last_checkpoint_state = Some(( - all_values.len(), - all_lengths.len(), - bytes_used, - checkpoint_values, - )); - checkpoint_log2 += 1; - } } + encoded_size += self.run_size::(current_length); + encoded_size + } - // After the loop, we always have a pending run that needs to be added - // unless we've exceeded the byte limit - if current_length > 0 { - let run_chunks = current_length.div_ceil(255) as usize; - let bytes_needed = run_chunks * (type_size + 1); + fn run_size(&self, length: u64) -> usize + where + T: bytemuck::Pod, + { + let type_size = std::mem::size_of::(); + let run_chunks = length.div_ceil(self.run_length_width.max_run_length()) as usize; + run_chunks * (type_size + self.run_length_width.bytes_per_value()) + } - if bytes_used + bytes_needed <= MAX_MINIBLOCK_BYTES as usize { - let _ = self.add_run(¤t_value, current_length, all_values, all_lengths); - total_values_encoded += current_length as usize; - } + fn encode_values(&self, values: &[T], all_values: &mut Vec, all_lengths: &mut Vec) + where + T: bytemuck::Pod + PartialEq + Copy, + { + if values.is_empty() { + return; } - // Determine if we've processed all remaining values - let is_last_chunk = total_values_encoded == values_remaining; + let mut current_value = values[0]; + let mut current_length = 1u64; - // Non-last chunks must have power-of-2 values for miniblock format - if !is_last_chunk { - if total_values_encoded.is_power_of_two() { - // Already at power-of-2 boundary - } else if let Some((val_pos, len_pos, _, checkpoint_values)) = last_checkpoint_state { - // Roll back to last valid checkpoint - all_values.truncate(val_pos); - all_lengths.truncate(len_pos); - let num_runs = (val_pos - values_start) / type_size; - return (num_runs, checkpoint_values, false); + for &value in values.iter().skip(1) { + if value == current_value { + current_length += 1; } else { - // No valid checkpoint, can't create a valid chunk - return (0, 0, false); + self.add_run(¤t_value, current_length, all_values, all_lengths); + current_value = value; + current_length = 1; } } - - let num_runs = (all_values.len() - values_start) / type_size; - (num_runs, total_values_encoded, is_last_chunk) + self.add_run(¤t_value, current_length, all_values, all_lengths); } fn add_run( @@ -337,24 +692,338 @@ impl RleEncoder { { let value_bytes = bytemuck::bytes_of(value); let type_size = std::mem::size_of::(); - let num_full_chunks = (length / 255) as usize; - let remainder = (length % 255) as u8; + let max_run_length = self.run_length_width.max_run_length(); + let num_full_chunks = (length / max_run_length) as usize; + let remainder = length % max_run_length; let total_chunks = num_full_chunks + if remainder > 0 { 1 } else { 0 }; all_values.reserve(total_chunks * type_size); - all_lengths.reserve(total_chunks); + all_lengths.reserve(total_chunks * self.run_length_width.bytes_per_value()); for _ in 0..num_full_chunks { all_values.extend_from_slice(value_bytes); - all_lengths.push(255); + self.run_length_width + .write_length(max_run_length, all_lengths); } if remainder > 0 { all_values.extend_from_slice(value_bytes); - all_lengths.push(remainder); + self.run_length_width.write_length(remainder, all_lengths); + } + + total_chunks * (type_size + self.run_length_width.bytes_per_value()) + } + + fn flat_child_candidate( + buffers: &[LanceBuffer], + chunks: &[MiniBlockChunk], + buffer_index: usize, + bits_per_value: u64, + ) -> RleChildCandidate { + RleChildCandidate { + encoding: ProtobufUtils21::flat(bits_per_value, None), + data: buffers[buffer_index].clone(), + chunk_sizes: chunks + .iter() + .map(|chunk| chunk.buffer_sizes[buffer_index]) + .collect(), + size: buffers[buffer_index].len(), + requires_num_values: false, + } + } + + fn general_child_candidate( + buffers: &[LanceBuffer], + chunks: &[MiniBlockChunk], + buffer_index: usize, + bits_per_value: u64, + compression: CompressionConfig, + ) -> Result> { + if buffers.is_empty() || buffers[buffer_index].is_empty() { + return Ok(None); + }; + + let compressor = GeneralBufferCompressor::get_compressor(compression)?; + let original = &buffers[buffer_index]; + let mut compressed = Vec::new(); + let mut offset = 0usize; + let mut total_original_size = 0usize; + let mut compressed_sizes = Vec::with_capacity(chunks.len()); + + for chunk in chunks.iter() { + let chunk_size = chunk.buffer_sizes[buffer_index] as usize; + let end = offset.checked_add(chunk_size).ok_or_else(|| { + Error::invalid_input_source("RLE child buffer offset overflow".into()) + })?; + if end > original.len() { + return Err(Error::invalid_input_source( + format!( + "RLE child buffer {} chunk size exceeds buffer length: end {}, len {}", + buffer_index, + end, + original.len() + ) + .into(), + )); + } + + let start = compressed.len(); + compressor.compress(&original.as_ref()[offset..end], &mut compressed)?; + let compressed_size = compressed.len() - start; + let compressed_size = u32::try_from(compressed_size).map_err(|_| { + Error::invalid_input_source( + format!( + "RLE child buffer {} compressed chunk is too large: {} bytes", + buffer_index, compressed_size + ) + .into(), + ) + })?; + compressed_sizes.push(compressed_size); + total_original_size += chunk_size; + offset = end; + } + + if compressed.len() >= total_original_size { + return Ok(None); } - total_chunks * (type_size + 1) + let encoding = + ProtobufUtils21::wrapped(compression, ProtobufUtils21::flat(bits_per_value, None))?; + Ok(Some( + RleChildCandidate { + encoding, + data: LanceBuffer::from(compressed), + chunk_sizes: compressed_sizes, + size: 0, + requires_num_values: false, + } + .with_size_from_data(), + )) + } + + #[cfg(feature = "bitpacking")] + fn bitpacked_child_candidate( + buffers: &[LanceBuffer], + chunks: &[MiniBlockChunk], + buffer_index: usize, + bits_per_value: u64, + ) -> Result> { + let original = &buffers[buffer_index]; + if original.is_empty() { + return Ok(None); + } + let packed_bits = Self::required_bits(original, bits_per_value)?; + if packed_bits >= bits_per_value { + return Ok(None); + } + + let compressor = crate::encodings::physical::bitpacking::OutOfLineBitpacking::new( + packed_bits, + bits_per_value, + ); + let mut packed = Vec::new(); + let mut offset = 0usize; + let mut packed_sizes = Vec::with_capacity(chunks.len()); + let bytes_per_value = usize::try_from(bits_per_value / 8).map_err(|_| { + Error::invalid_input_source( + format!("RLE child bit width is too large: {bits_per_value}").into(), + ) + })?; + + for chunk in chunks { + let chunk_size = chunk.buffer_sizes[buffer_index] as usize; + let end = offset.checked_add(chunk_size).ok_or_else(|| { + Error::invalid_input_source("RLE child buffer offset overflow".into()) + })?; + if end > original.len() { + return Err(Error::invalid_input_source( + format!( + "RLE child buffer {} chunk size exceeds buffer length: end {}, len {}", + buffer_index, + end, + original.len() + ) + .into(), + )); + } + if bytes_per_value == 0 || !chunk_size.is_multiple_of(bytes_per_value) { + return Err(Error::invalid_input_source( + format!( + "RLE child buffer {} chunk has invalid size {} for {} bits per value", + buffer_index, chunk_size, bits_per_value + ) + .into(), + )); + } + + let child_values = (chunk_size / bytes_per_value) as u64; + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value, + data: original.slice_with_length(offset, chunk_size), + num_values: child_values, + block_info: BlockInfo::default(), + }); + let chunk_packed = BlockCompressor::compress(&compressor, block)?; + let packed_size = u32::try_from(chunk_packed.len()).map_err(|_| { + Error::invalid_input_source( + format!( + "RLE child buffer {} bitpacked chunk is too large: {} bytes", + buffer_index, + chunk_packed.len() + ) + .into(), + ) + })?; + packed_sizes.push(packed_size); + packed.extend_from_slice(chunk_packed.as_ref()); + offset = end; + } + + if packed.len() >= original.len() { + return Ok(None); + } + + Ok(Some( + RleChildCandidate { + encoding: ProtobufUtils21::out_of_line_bitpacking( + bits_per_value, + ProtobufUtils21::flat(packed_bits, None), + ), + data: LanceBuffer::from(packed), + chunk_sizes: packed_sizes, + size: 0, + requires_num_values: true, + } + .with_size_from_data(), + )) + } + + #[cfg(feature = "bitpacking")] + fn required_bits(buffer: &LanceBuffer, bits_per_value: u64) -> Result { + let max_value = match bits_per_value { + 8 => buffer.as_ref().iter().map(|value| *value as u64).max(), + 16 => buffer + .as_ref() + .chunks_exact(2) + .map(|value| u16::from_le_bytes(value.try_into().unwrap()) as u64) + .max(), + 32 => buffer + .as_ref() + .chunks_exact(4) + .map(|value| u32::from_le_bytes(value.try_into().unwrap()) as u64) + .max(), + 64 => buffer + .as_ref() + .chunks_exact(8) + .map(|value| u64::from_le_bytes(value.try_into().unwrap())) + .max(), + _ => { + return Err(Error::invalid_input_source( + format!( + "RLE child bitpacking only supports 8, 16, 32, or 64-bit values, got {bits_per_value}" + ) + .into(), + )); + } + } + .unwrap_or(0); + Ok((u64::BITS - max_value.leading_zeros()).max(1) as u64) + } + + fn child_candidates( + buffers: &[LanceBuffer], + chunks: &[MiniBlockChunk], + buffer_index: usize, + bits_per_value: u64, + compression: Option, + use_child_bitpacking: bool, + ) -> Result> { + #[cfg(not(feature = "bitpacking"))] + let _ = use_child_bitpacking; + let mut candidates = vec![Self::flat_child_candidate( + buffers, + chunks, + buffer_index, + bits_per_value, + )]; + if let Some(compression) = compression + && let Some(candidate) = Self::general_child_candidate( + buffers, + chunks, + buffer_index, + bits_per_value, + compression, + )? + { + candidates.push(candidate); + } + #[cfg(feature = "bitpacking")] + { + if use_child_bitpacking + && let Some(candidate) = + Self::bitpacked_child_candidate(buffers, chunks, buffer_index, bits_per_value)? + { + candidates.push(candidate); + } + } + Ok(candidates) + } + + fn select_child_candidates( + values: Vec, + run_lengths: Vec, + ) -> (RleChildCandidate, RleChildCandidate) { + let mut best: Option<(usize, usize, usize)> = None; + for (value_idx, value) in values.iter().enumerate() { + for (length_idx, length) in run_lengths.iter().enumerate() { + if value.requires_num_values && length.requires_num_values { + continue; + } + let size = value.size + length.size; + if best.is_none_or(|(_, _, best_size)| size < best_size) { + best = Some((value_idx, length_idx, size)); + } + } + } + let (value_idx, length_idx, _) = + best.expect("flat RLE child candidates should always be selectable"); + (values[value_idx].clone(), run_lengths[length_idx].clone()) + } + + pub(crate) fn selected_payload_size(&self, data: &FixedWidthDataBlock) -> Result { + let (all_buffers, chunks) = + self.encode_data(&data.data, data.num_values, data.bits_per_value)?; + if all_buffers.is_empty() { + return Ok(0); + } + + let values_candidates = Self::child_candidates( + &all_buffers, + &chunks, + 0, + data.bits_per_value, + self.values_compression, + self.use_child_bitpacking, + )?; + let run_lengths_candidates = Self::child_candidates( + &all_buffers, + &chunks, + 1, + self.run_length_width.bits_per_value(), + self.run_lengths_compression, + self.use_child_bitpacking, + )?; + let (values, run_lengths) = + Self::select_child_candidates(values_candidates, run_lengths_candidates); + Ok((values.size as u128).saturating_add(run_lengths.size as u128)) + } +} + +impl RleChildCandidate { + fn with_size_from_data(mut self) -> Self { + self.size = self.data.len(); + self } } @@ -367,17 +1036,53 @@ impl MiniBlockCompressor for RleEncoder { let (all_buffers, chunks) = self.encode_data(&fixed_width.data, num_values, bits_per_value)?; + if all_buffers.is_empty() { + let compressed = MiniBlockCompressed { + data: all_buffers, + chunks, + num_values, + }; + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::flat(bits_per_value, None), + ProtobufUtils21::flat(self.run_length_width.bits_per_value(), None), + ); + return Ok((compressed, encoding)); + } + + let values_candidates = Self::child_candidates( + &all_buffers, + &chunks, + 0, + bits_per_value, + self.values_compression, + self.use_child_bitpacking, + )?; + let run_lengths_candidates = Self::child_candidates( + &all_buffers, + &chunks, + 1, + self.run_length_width.bits_per_value(), + self.run_lengths_compression, + self.use_child_bitpacking, + )?; + let (values, run_lengths) = + Self::select_child_candidates(values_candidates, run_lengths_candidates); + let chunks = chunks + .into_iter() + .enumerate() + .map(|(idx, chunk)| MiniBlockChunk { + buffer_sizes: vec![values.chunk_sizes[idx], run_lengths.chunk_sizes[idx]], + log_num_values: chunk.log_num_values, + }) + .collect(); let compressed = MiniBlockCompressed { - data: all_buffers, + data: vec![values.data, run_lengths.data], chunks, num_values, }; - let encoding = ProtobufUtils21::rle( - ProtobufUtils21::flat(bits_per_value, None), - ProtobufUtils21::flat(/*bits_per_value=*/ 8, None), - ); + let encoding = ProtobufUtils21::rle(values.encoding, run_lengths.encoding); Ok((compressed, encoding)) } @@ -396,8 +1101,8 @@ impl BlockCompressor for RleEncoder { let num_values = fixed_width.num_values; let bits_per_value = fixed_width.bits_per_value; - let (all_buffers, _) = - self.encode_data(&fixed_width.data, num_values, bits_per_value)?; + let all_buffers = + self.encode_block_data(&fixed_width.data, num_values, bits_per_value)?; let values_size = all_buffers[0].len() as u64; @@ -418,14 +1123,170 @@ impl BlockCompressor for RleEncoder { #[derive(Debug)] pub struct RleDecompressor { bits_per_value: u64, + run_length_width: RunLengthWidth, + values: RleChildDecompressor, + run_lengths: RleChildDecompressor, +} + +#[derive(Debug)] +pub(crate) struct RleChildDecompressor { + bits_per_value: u64, + inner: RleChildDecompressorInner, +} + +#[derive(Debug)] +enum RleChildDecompressorInner { + Flat, + Block { + decompressor: Box, + requires_num_values: bool, + }, +} + +impl RleChildDecompressor { + pub(crate) fn flat(bits_per_value: u64) -> Self { + Self { + bits_per_value, + inner: RleChildDecompressorInner::Flat, + } + } + + pub(crate) fn block( + bits_per_value: u64, + decompressor: Box, + requires_num_values: bool, + ) -> Self { + Self { + bits_per_value, + inner: RleChildDecompressorInner::Block { + decompressor, + requires_num_values, + }, + } + } + + pub(crate) fn bits_per_value(&self) -> u64 { + self.bits_per_value + } + + pub(crate) fn requires_num_values(&self) -> bool { + match &self.inner { + RleChildDecompressorInner::Flat => false, + RleChildDecompressorInner::Block { + requires_num_values, + .. + } => *requires_num_values, + } + } + + pub(crate) fn is_identity(&self) -> bool { + matches!(self.inner, RleChildDecompressorInner::Flat) + } + + fn decode( + &self, + data: LanceBuffer, + num_values: Option, + label: &str, + ) -> Result { + match &self.inner { + RleChildDecompressorInner::Flat => Ok(data), + RleChildDecompressorInner::Block { + decompressor, + requires_num_values, + } => { + let num_values = if *requires_num_values { + num_values.ok_or_else(|| { + Error::invalid_input_source( + format!("RLE {label} child compression requires the run count").into(), + ) + })? + } else { + num_values.unwrap_or(0) + }; + let decoded = decompressor.decompress(data, num_values)?; + self.extract_fixed_width(decoded, num_values, label) + } + } + } + + fn extract_fixed_width( + &self, + data: DataBlock, + expected_num_values: u64, + label: &str, + ) -> Result { + match data { + DataBlock::FixedWidth(block) => { + if block.bits_per_value != self.bits_per_value { + return Err(Error::invalid_input_source( + format!( + "RLE {label} child decoded {}-bit values, expected {}", + block.bits_per_value, self.bits_per_value + ) + .into(), + )); + } + if expected_num_values != 0 && block.num_values != expected_num_values { + return Err(Error::invalid_input_source( + format!( + "RLE {label} child decoded {} values, expected {}", + block.num_values, expected_num_values + ) + .into(), + )); + } + Ok(block.data) + } + _ => Err(Error::invalid_input_source( + format!("RLE {label} child decoded to a non fixed-width block").into(), + )), + } + } } impl RleDecompressor { pub fn new(bits_per_value: u64) -> Self { - Self { bits_per_value } + Self { + bits_per_value, + run_length_width: RunLengthWidth::U8, + values: RleChildDecompressor::flat(bits_per_value), + run_lengths: RleChildDecompressor::flat(RunLengthWidth::U8.bits_per_value()), + } + } + + pub(crate) fn with_run_length_width( + bits_per_value: u64, + run_length_width: RunLengthWidth, + ) -> Self { + Self { + bits_per_value, + run_length_width, + values: RleChildDecompressor::flat(bits_per_value), + run_lengths: RleChildDecompressor::flat(run_length_width.bits_per_value()), + } + } + + pub(crate) fn with_child_decompressors( + bits_per_value: u64, + run_length_width: RunLengthWidth, + values: RleChildDecompressor, + run_lengths: RleChildDecompressor, + ) -> Self { + Self { + bits_per_value, + run_length_width, + values, + run_lengths, + } } - fn decode_data(&self, data: Vec, num_values: u64) -> Result { + fn decode_data( + &self, + data: Vec, + num_values: u64, + clamp_overflow: bool, + ) -> Result { if num_values == 0 { return Ok(DataBlock::FixedWidth(FixedWidthDataBlock { bits_per_value: self.bits_per_value, @@ -445,15 +1306,46 @@ impl RleDecompressor { )); } - let values_buffer = &data[0]; - let lengths_buffer = &data[1]; + let mut data_iter = data.into_iter(); + let values_buffer = data_iter.next().unwrap(); + let lengths_buffer = data_iter.next().unwrap(); + let (values_buffer, lengths_buffer) = + self.decode_child_buffers(values_buffer, lengths_buffer)?; let decoded_data = match self.bits_per_value { - 8 => self.decode_generic::(values_buffer, lengths_buffer, num_values)?, - 16 => self.decode_generic::(values_buffer, lengths_buffer, num_values)?, - 32 => self.decode_generic::(values_buffer, lengths_buffer, num_values)?, - 64 => self.decode_generic::(values_buffer, lengths_buffer, num_values)?, - _ => unreachable!("RLE decoding bits_per_value must be 8, 16, 32, 64, or 128"), + 8 => self.decode_generic::( + &values_buffer, + &lengths_buffer, + num_values, + clamp_overflow, + )?, + 16 => self.decode_generic::( + &values_buffer, + &lengths_buffer, + num_values, + clamp_overflow, + )?, + 32 => self.decode_generic::( + &values_buffer, + &lengths_buffer, + num_values, + clamp_overflow, + )?, + 64 => self.decode_generic::( + &values_buffer, + &lengths_buffer, + num_values, + clamp_overflow, + )?, + _ => { + return Err(Error::invalid_input_source( + format!( + "RLE decoding bits_per_value must be 8, 16, 32, or 64, got {}", + self.bits_per_value + ) + .into(), + )); + } }; Ok(DataBlock::FixedWidth(FixedWidthDataBlock { @@ -464,16 +1356,80 @@ impl RleDecompressor { })) } + fn decode_child_buffers( + &self, + values_buffer: LanceBuffer, + lengths_buffer: LanceBuffer, + ) -> Result<(LanceBuffer, LanceBuffer)> { + let values_requires_num_runs = self.values.requires_num_values(); + let lengths_requires_num_runs = self.run_lengths.requires_num_values(); + if values_requires_num_runs && lengths_requires_num_runs { + return Err(Error::invalid_input_source( + "RLE values and run lengths child compression both require the run count".into(), + )); + } + + if values_requires_num_runs { + let lengths_buffer = self + .run_lengths + .decode(lengths_buffer, None, "run lengths")?; + let num_runs = Self::num_child_values( + &lengths_buffer, + self.run_lengths.bits_per_value(), + "run lengths", + )?; + let values_buffer = self + .values + .decode(values_buffer, Some(num_runs), "values")?; + Ok((values_buffer, lengths_buffer)) + } else if lengths_requires_num_runs { + let values_buffer = self.values.decode(values_buffer, None, "values")?; + let num_runs = + Self::num_child_values(&values_buffer, self.values.bits_per_value(), "values")?; + let lengths_buffer = + self.run_lengths + .decode(lengths_buffer, Some(num_runs), "run lengths")?; + Ok((values_buffer, lengths_buffer)) + } else { + let values_buffer = self.values.decode(values_buffer, None, "values")?; + let lengths_buffer = self + .run_lengths + .decode(lengths_buffer, None, "run lengths")?; + Ok((values_buffer, lengths_buffer)) + } + } + + fn num_child_values(buffer: &LanceBuffer, bits_per_value: u64, label: &str) -> Result { + let bytes_per_value = usize::try_from(bits_per_value / 8).map_err(|_| { + Error::invalid_input_source( + format!("RLE {label} child bit width is too large: {bits_per_value}").into(), + ) + })?; + if bytes_per_value == 0 || !buffer.len().is_multiple_of(bytes_per_value) { + return Err(Error::invalid_input_source( + format!( + "RLE {label} child decoded to {} bytes, not divisible by {}", + buffer.len(), + bytes_per_value + ) + .into(), + )); + } + Ok((buffer.len() / bytes_per_value) as u64) + } + fn decode_generic( &self, values_buffer: &LanceBuffer, lengths_buffer: &LanceBuffer, num_values: u64, + clamp_overflow: bool, ) -> Result where T: bytemuck::Pod + Copy + std::fmt::Debug + ArrowNativeType, { let type_size = std::mem::size_of::(); + let length_size = self.run_length_width.bytes_per_value(); if values_buffer.is_empty() || lengths_buffer.is_empty() { if num_values == 0 { @@ -485,19 +1441,22 @@ impl RleDecompressor { } } - if !values_buffer.len().is_multiple_of(type_size) || lengths_buffer.is_empty() { + if !values_buffer.len().is_multiple_of(type_size) + || !lengths_buffer.len().is_multiple_of(length_size) + { return Err(Error::invalid_input_source(format!( - "Invalid buffer sizes for RLE {} decoding: values {} bytes (not divisible by {}), lengths {} bytes", + "Invalid buffer sizes for RLE {} decoding: values {} bytes (not divisible by {}), lengths {} bytes (not divisible by {})", std::any::type_name::(), values_buffer.len(), type_size, - lengths_buffer.len() + lengths_buffer.len(), + length_size ) .into())); } let num_runs = values_buffer.len() / type_size; - let num_length_entries = lengths_buffer.len(); + let num_length_entries = lengths_buffer.len() / length_size; if num_runs != num_length_entries { return Err(Error::invalid_input_source( format!( @@ -510,26 +1469,55 @@ impl RleDecompressor { let values_ref = values_buffer.borrow_to_typed_slice::(); let values: &[T] = values_ref.as_ref(); - let lengths: &[u8] = lengths_buffer.as_ref(); - - let expected_value_count = num_values as usize; - let mut decoded: Vec = Vec::with_capacity(expected_value_count); - - for (value, &length) in values.iter().zip(lengths.iter()) { - if decoded.len() == expected_value_count { - break; - } + let lengths = lengths_buffer.as_ref(); + let expected_value_count = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("RLE num_values does not fit in usize: {num_values}").into(), + ) + })?; + // Legacy miniblock encoders rolled back to a power-of-2 checkpoint after a run + // had already crossed it, so a chunk's run lengths can sum past its declared + // value count (the excess values are re-encoded at the start of the next chunk). + // The pre-run-length-width decoder truncated the excess, so miniblock decoding + // clamps rather than rejects to keep those files readable. Block payloads never + // legitimately overflow, so they decode strictly. + let mut decoded: Vec = Vec::new(); + decoded + .try_reserve_exact(expected_value_count) + .map_err(|_| { + Error::invalid_input_source( + format!("RLE decoding cannot allocate {expected_value_count} values").into(), + ) + })?; + for (value, length_bytes) in values.iter().zip(lengths.chunks_exact(length_size)) { + let length = self.run_length_width.read_length(length_bytes); if length == 0 { return Err(Error::invalid_input_source( "RLE decoding encountered a zero run length".into(), )); } - + let length = usize::try_from(length).map_err(|_| { + Error::invalid_input_source( + format!("RLE run length does not fit in usize: {length}").into(), + ) + })?; let remaining = expected_value_count - decoded.len(); - let write_len = (length as usize).min(remaining); - - decoded.resize(decoded.len() + write_len, *value); + if length > remaining { + if !clamp_overflow { + return Err(Error::invalid_input_source( + format!( + "RLE decoding overflowed expected value count: produced at least {}, expected {}", + decoded.len() + length, + expected_value_count + ) + .into(), + )); + } + decoded.resize(expected_value_count, *value); + break; + } + decoded.resize(decoded.len() + length, *value); } if decoded.len() != expected_value_count { @@ -554,7 +1542,7 @@ impl RleDecompressor { impl MiniBlockDecompressor for RleDecompressor { fn decompress(&self, data: Vec, num_values: u64) -> Result { - self.decode_data(data, num_values) + self.decode_data(data, num_values, true) } } @@ -591,16 +1579,21 @@ impl BlockDecompressor for RleDecompressor { let values_buffer = data.slice_with_length(values_start, values_size); let lengths_buffer = data.slice_with_length(lengths_start, data.len() - lengths_start); - self.decode_data(vec![values_buffer, lengths_buffer], num_values) + self.decode_data(vec![values_buffer, lengths_buffer], num_values, false) } } #[cfg(test)] mod tests { use super::*; + use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy}; use crate::data::DataBlock; use crate::encodings::logical::primitive::miniblock::MAX_MINIBLOCK_VALUES; - use crate::{buffer::LanceBuffer, compression::BlockDecompressor}; + use crate::encodings::physical::block::{CompressionConfig, CompressionScheme}; + use crate::{ + buffer::LanceBuffer, + compression::{BlockCompressor, BlockDecompressor}, + }; use arrow_array::Int32Array; // ========== Core Functionality Tests ========== @@ -641,6 +1634,371 @@ mod tests { assert_eq!(lengths_buffer.len(), 6); } + #[test] + fn test_rle_v2_u16_miniblock_encoding() { + let encoder = RleEncoder::with_run_length_width(RunLengthWidth::U16); + + let data = vec![42i32; 1000]; + let array = Int32Array::from(data); + let (compressed, encoding) = + MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + + assert_eq!(compressed.data[0].len(), 4); + assert_eq!(compressed.data[1].len(), 2); + assert_eq!(compressed.data[1].as_ref(), &1000u16.to_le_bytes()); + + let rle = match encoding.compression.as_ref().unwrap() { + crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle, + other => panic!("expected RLE encoding, got {other:?}"), + }; + let run_lengths = rle.run_lengths.as_ref().unwrap(); + let flat = match run_lengths.compression.as_ref().unwrap() { + crate::format::pb21::compressive_encoding::Compression::Flat(flat) => flat, + other => panic!("expected flat run lengths, got {other:?}"), + }; + assert_eq!(flat.bits_per_value, 16); + + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let decompressed = MiniBlockDecompressor::decompress( + &decompressor, + compressed.data, + compressed.num_values, + ) + .unwrap(); + match decompressed { + DataBlock::FixedWidth(block) => { + let values = block.data.borrow_to_typed_slice::(); + assert_eq!(values.as_ref(), vec![42i32; 1000]); + } + _ => panic!("Expected FixedWidth block"), + } + } + + #[test] + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_rle_miniblock_compressed_values_child() { + let compression = test_general_compression(); + let encoder = + RleEncoder::with_child_encoding(RunLengthWidth::U8, Some(compression), None, false); + let array = Int32Array::from(repeating_runs(1024, 4)); + let (compressed, encoding) = + MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + + let rle = expect_rle(&encoding); + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + crate::format::pb21::compressive_encoding::Compression::General(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + crate::format::pb21::compressive_encoding::Compression::Flat(_) + )); + + let decompressor = DefaultDecompressionStrategy::default() + .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default()) + .unwrap(); + let decoded = + MiniBlockDecompressor::decompress(decompressor.as_ref(), compressed.data, 1024 * 4) + .unwrap(); + assert_decoded_i32_eq(decoded, &repeating_runs(1024, 4)); + } + + #[test] + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_rle_miniblock_compressed_run_lengths_child() { + let compression = test_general_compression(); + let encoder = + RleEncoder::with_child_encoding(RunLengthWidth::U8, None, Some(compression), false); + let expected = repeating_runs(1024, 4); + let (compressed, encoding) = MiniBlockCompressor::compress( + &encoder, + DataBlock::from_array(Int32Array::from(expected.clone())), + ) + .unwrap(); + + let rle = expect_rle(&encoding); + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + crate::format::pb21::compressive_encoding::Compression::Flat(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + crate::format::pb21::compressive_encoding::Compression::General(_) + )); + + let decompressor = DefaultDecompressionStrategy::default() + .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default()) + .unwrap(); + let decoded = + MiniBlockDecompressor::decompress(decompressor.as_ref(), compressed.data, 1024 * 4) + .unwrap(); + assert_decoded_i32_eq(decoded, &expected); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_miniblock_bitpacked_run_lengths_child() { + use crate::encodings::physical::bitpacking::OutOfLineBitpacking; + + let expected = repeating_runs(1024, 4); + let (compressed, _) = MiniBlockCompressor::compress( + &RleEncoder::new(), + DataBlock::from_array(Int32Array::from(expected.clone())), + ) + .unwrap(); + let run_lengths = compressed.data[1].clone(); + let num_runs = run_lengths.len() as u64; + let run_lengths_block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 8, + data: run_lengths, + num_values: num_runs, + block_info: BlockInfo::default(), + }); + let bitpacked_run_lengths = + BlockCompressor::compress(&OutOfLineBitpacking::new(3, 8), run_lengths_block).unwrap(); + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::flat(32, None), + ProtobufUtils21::out_of_line_bitpacking(8, ProtobufUtils21::flat(3, None)), + ); + + let decompressor = DefaultDecompressionStrategy::default() + .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default()) + .unwrap(); + let decoded = MiniBlockDecompressor::decompress( + decompressor.as_ref(), + vec![compressed.data[0].clone(), bitpacked_run_lengths], + expected.len() as u64, + ) + .unwrap(); + assert_decoded_i32_eq(decoded, &expected); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_rejects_two_count_dependent_child_encodings() { + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::out_of_line_bitpacking(32, ProtobufUtils21::flat(3, None)), + ProtobufUtils21::out_of_line_bitpacking(8, ProtobufUtils21::flat(3, None)), + ); + + let err = DefaultDecompressionStrategy::default() + .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default()) + .unwrap_err(); + assert!( + err.to_string() + .contains("cannot both require the run count") + ); + } + + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_general_compression() -> CompressionConfig { + if cfg!(feature = "zstd") { + CompressionConfig::new(CompressionScheme::Zstd, Some(3)) + } else { + CompressionConfig::new(CompressionScheme::Lz4, None) + } + } + + fn repeating_runs(num_runs: usize, run_length: usize) -> Vec { + let mut values = Vec::with_capacity(num_runs * run_length); + for run in 0..num_runs { + values.extend(std::iter::repeat_n((run % 8) as i32, run_length)); + } + values + } + + fn expect_rle(encoding: &CompressiveEncoding) -> &crate::format::pb21::Rle { + match encoding.compression.as_ref().unwrap() { + crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle, + other => panic!("expected RLE encoding, got {other:?}"), + } + } + + fn assert_decoded_i32_eq(decoded: DataBlock, expected: &[i32]) { + match decoded { + DataBlock::FixedWidth(block) => { + let values = block.data.borrow_to_typed_slice::(); + assert_eq!(values.as_ref(), expected); + } + _ => panic!("Expected FixedWidth block"), + } + } + + #[test] + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_rle_miniblock_compressed_children_multiple_chunks() { + let compression = test_general_compression(); + let encoder = RleEncoder::with_child_encoding( + RunLengthWidth::U8, + Some(compression), + Some(compression), + false, + ); + let expected = repeating_runs(8192, 4); + let (compressed, encoding) = MiniBlockCompressor::compress( + &encoder, + DataBlock::from_array(Int32Array::from(expected.clone())), + ) + .unwrap(); + + assert!(compressed.chunks.len() > 1); + let rle = expect_rle(&encoding); + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + crate::format::pb21::compressive_encoding::Compression::General(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + crate::format::pb21::compressive_encoding::Compression::General(_) + )); + + let decoded = decompress_i32_chunks(&compressed, &encoding); + assert_eq!(decoded, expected); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_miniblock_bitpacks_values_child_when_smaller() { + let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true); + let expected = monotonic_runs(2048, 4); + let (compressed, encoding) = MiniBlockCompressor::compress( + &encoder, + DataBlock::from_array(Int32Array::from(expected.clone())), + ) + .unwrap(); + + let rle = expect_rle(&encoding); + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + crate::format::pb21::compressive_encoding::Compression::OutOfLineBitpacking(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + crate::format::pb21::compressive_encoding::Compression::Flat(_) + )); + + let decoded = decompress_i32_chunks(&compressed, &encoding); + assert_eq!(decoded, expected); + } + + #[test] + #[cfg(feature = "bitpacking")] + fn test_rle_miniblock_bitpacks_run_lengths_when_values_do_not_shrink() { + let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true); + let expected = high_entropy_runs(2048, 4); + let (compressed, encoding) = MiniBlockCompressor::compress( + &encoder, + DataBlock::from_array(Int32Array::from(expected.clone())), + ) + .unwrap(); + + let rle = expect_rle(&encoding); + assert!(matches!( + rle.values.as_ref().unwrap().compression.as_ref().unwrap(), + crate::format::pb21::compressive_encoding::Compression::Flat(_) + )); + assert!(matches!( + rle.run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap(), + crate::format::pb21::compressive_encoding::Compression::OutOfLineBitpacking(_) + )); + + let decoded = decompress_i32_chunks(&compressed, &encoding); + assert_eq!(decoded, expected); + } + + fn decompress_i32_chunks( + compressed: &MiniBlockCompressed, + encoding: &CompressiveEncoding, + ) -> Vec { + let strategy = DefaultDecompressionStrategy::default(); + let decompressor = strategy + .create_miniblock_decompressor(encoding, &strategy) + .unwrap(); + let mut offsets = vec![0usize; compressed.data.len()]; + let mut values_processed = 0u64; + let mut decoded_values = Vec::new(); + + for chunk in &compressed.chunks { + let chunk_values = chunk.num_values(values_processed, compressed.num_values); + let mut chunk_buffers = Vec::with_capacity(chunk.buffer_sizes.len()); + for (idx, size) in chunk.buffer_sizes.iter().enumerate() { + let size = *size as usize; + chunk_buffers.push(compressed.data[idx].slice_with_length(offsets[idx], size)); + offsets[idx] += size; + } + + let decoded = decompressor + .decompress(chunk_buffers, chunk_values) + .unwrap(); + match decoded { + DataBlock::FixedWidth(block) => { + let values = block.data.borrow_to_typed_slice::(); + decoded_values.extend_from_slice(values.as_ref()); + } + _ => panic!("Expected FixedWidth block"), + } + values_processed += chunk_values; + } + + assert_eq!(values_processed, compressed.num_values); + decoded_values + } + + #[cfg(feature = "bitpacking")] + fn monotonic_runs(num_runs: usize, run_length: usize) -> Vec { + let mut values = Vec::with_capacity(num_runs * run_length); + for run in 0..num_runs { + values.extend(std::iter::repeat_n(run as i32, run_length)); + } + values + } + + #[cfg(feature = "bitpacking")] + fn high_entropy_runs(num_runs: usize, run_length: usize) -> Vec { + let mut values = Vec::with_capacity(num_runs * run_length); + let mut state = 7u64; + for _ in 0..num_runs { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + values.extend(std::iter::repeat_n((state >> 32) as i32, run_length)); + } + values + } + + #[test] + fn test_select_run_length_width_prefers_u16_for_long_runs() { + let mut entries = [0u64; 3]; + accumulate_run_length_entries(300, Some(*MAX_MINIBLOCK_VALUES), &mut entries); + let (width, _) = select_run_length_width_from_entries(&entries, 32).unwrap(); + assert_eq!(width, RunLengthWidth::U16); + } + // ========== Round-trip Tests for Different Types ========== #[test] @@ -760,6 +2118,118 @@ mod tests { ); } + #[test] + fn test_u16_length_buffer_must_be_aligned() { + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let values = LanceBuffer::from(vec![1, 0, 0, 0]); + let lengths = LanceBuffer::from(vec![5]); + let result = MiniBlockDecompressor::decompress(&decompressor, vec![values, lengths], 5); + assert!(matches!(&result, Err(Error::InvalidInput { .. }))); + assert!( + result + .unwrap_err() + .to_string() + .contains("not divisible by 2") + ); + } + + #[test] + fn test_rle_rejects_underflow_and_zero_lengths_and_clamps_overflow() { + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let value = LanceBuffer::from(1i32.to_le_bytes().to_vec()); + + let underflow = MiniBlockDecompressor::decompress( + &decompressor, + vec![ + value.clone(), + LanceBuffer::from(4u16.to_le_bytes().to_vec()), + ], + 5, + ) + .unwrap_err(); + assert!(underflow.to_string().contains("produced 4 values")); + + let overflow = MiniBlockDecompressor::decompress( + &decompressor, + vec![ + value.clone(), + LanceBuffer::from(6u16.to_le_bytes().to_vec()), + ], + 5, + ) + .unwrap(); + match overflow { + DataBlock::FixedWidth(block) => { + assert_eq!(block.num_values, 5); + let decoded = block.data.borrow_to_typed_slice::(); + assert_eq!(decoded.as_ref(), &[1i32; 5]); + } + _ => panic!("Expected FixedWidth block"), + } + + let zero = MiniBlockDecompressor::decompress( + &decompressor, + vec![value, LanceBuffer::from(0u16.to_le_bytes().to_vec())], + 5, + ) + .unwrap_err(); + assert!(zero.to_string().contains("zero run length")); + } + + #[test] + fn test_block_rle_rejects_overflow() { + // Block payloads have no chunk boundaries, so run lengths summing past + // num_values can only be corruption and must stay a hard error. + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let values = 1i32.to_le_bytes(); + let lengths = 6u16.to_le_bytes(); + let mut payload = Vec::new(); + payload.extend_from_slice(&(values.len() as u64).to_le_bytes()); + payload.extend_from_slice(&values); + payload.extend_from_slice(&lengths); + + let error = BlockDecompressor::decompress(&decompressor, LanceBuffer::from(payload), 5) + .unwrap_err(); + assert!(matches!(&error, Error::InvalidInput { .. })); + assert!( + error + .to_string() + .contains("overflowed expected value count") + ); + } + + #[test] + fn test_rle_truncates_legacy_chunk_boundary_overflow() { + // Legacy encoders emitted chunks declaring 2048 values whose final run crossed + // the checkpoint boundary (e.g. run lengths summing to 2080); the excess values + // are duplicated at the start of the next chunk and must be ignored here. + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let mut values = Vec::new(); + values.extend_from_slice(&7i32.to_le_bytes()); + values.extend_from_slice(&8i32.to_le_bytes()); + let mut lengths = Vec::new(); + lengths.extend_from_slice(&2000u16.to_le_bytes()); + lengths.extend_from_slice(&80u16.to_le_bytes()); + + let decoded = MiniBlockDecompressor::decompress( + &decompressor, + vec![LanceBuffer::from(values), LanceBuffer::from(lengths)], + 2048, + ) + .unwrap(); + match decoded { + DataBlock::FixedWidth(block) => { + assert_eq!(block.num_values, 2048); + let decoded = block.data.borrow_to_typed_slice::(); + let decoded = decoded.as_ref(); + assert_eq!(decoded.len(), 2048); + assert!(decoded[..2000].iter().all(|&v| v == 7)); + assert!(decoded[2000..].iter().all(|&v| v == 8)); + } + _ => panic!("Expected FixedWidth block"), + } + } + #[test] fn test_empty_data_handling() { let encoder = RleEncoder::new(); @@ -1098,6 +2568,30 @@ mod tests { ); // 20% variety let arr = Arc::new(Int32Array::from(values)) as Arc; check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await; + + #[cfg(any(feature = "lz4", feature = "zstd"))] + { + let mut metadata = HashMap::new(); + metadata.insert( + "lance-encoding:rle-threshold".to_string(), + "0.8".to_string(), + ); + metadata.insert("lance-encoding:bss".to_string(), "off".to_string()); + metadata.insert( + "lance-encoding:compression".to_string(), + if cfg!(feature = "zstd") { + "zstd".to_string() + } else { + "lz4".to_string() + }, + ); + let mut values = Vec::with_capacity(2048 * 4); + for run in 0..2048 { + values.extend(std::iter::repeat_n(i32::MIN + (run % 8), 4)); + } + let arr = Arc::new(Int32Array::from(values)) as Arc; + check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await; + } } /// Generator that produces repetitive patterns suitable for RLE diff --git a/rust/lance-encoding/src/encodings/physical/value.rs b/rust/lance-encoding/src/encodings/physical/value.rs index c49bbd3efbd..606f49b699a 100644 --- a/rust/lance-encoding/src/encodings/physical/value.rs +++ b/rust/lance-encoding/src/encodings/physical/value.rs @@ -27,7 +27,7 @@ pub struct ValueEncoder {} impl ValueEncoder { /// Use the largest chunk we can smaller than 4KiB - fn find_log_vals_per_chunk(bytes_per_word: u64, values_per_word: u64) -> (u64, u64) { + fn find_log_vals_per_chunk(bytes_per_word: u64, values_per_word: u64) -> Result<(u64, u64)> { let mut size_bytes = 2 * bytes_per_word; let (mut log_num_vals, mut num_vals) = match values_per_word { 1 => (1, 2), @@ -35,8 +35,14 @@ impl ValueEncoder { _ => unreachable!(), }; - // If the type is so wide that we can't even fit 2 values we shouldn't be here - assert!(size_bytes < MAX_MINIBLOCK_BYTES); + if size_bytes >= MAX_MINIBLOCK_BYTES { + let num_values = 2 * values_per_word; + return Err(Error::invalid_input(format!( + "Value is too wide for miniblock encoding: {} values require {} bytes but a \ + miniblock chunk is limited to {} bytes.", + num_values, size_bytes, MAX_MINIBLOCK_BYTES + ))); + } while 2 * size_bytes < MAX_MINIBLOCK_BYTES && 2 * num_vals <= *MAX_MINIBLOCK_VALUES { log_num_vals += 1; @@ -44,10 +50,10 @@ impl ValueEncoder { num_vals *= 2; } - (log_num_vals, num_vals) + Ok((log_num_vals, num_vals)) } - fn chunk_data(data: FixedWidthDataBlock) -> MiniBlockCompressed { + fn chunk_data(data: FixedWidthDataBlock) -> Result { // Usually there are X bytes per value. However, when working with boolean // or FSL we might have some number of bits per value that isn't // divisible by 8. In this case, to avoid chunking in the middle of a byte @@ -60,7 +66,7 @@ impl ValueEncoder { // Aim for 4KiB chunks let (log_vals_per_chunk, vals_per_chunk) = - Self::find_log_vals_per_chunk(bytes_per_word, values_per_word); + Self::find_log_vals_per_chunk(bytes_per_word, values_per_word)?; let num_chunks = bit_util::ceil(data.num_values as usize, vals_per_chunk as usize); debug_assert_eq!(vals_per_chunk % values_per_word, 0); let bytes_per_chunk = bytes_per_word * (vals_per_chunk / values_per_word); @@ -99,11 +105,11 @@ impl ValueEncoder { debug_assert_eq!(chunks.len(), num_chunks); - MiniBlockCompressed { + Ok(MiniBlockCompressed { chunks, data: vec![data_buffer], num_values: data.num_values, - } + }) } } @@ -177,7 +183,7 @@ impl ValueEncoder { data: FixedWidthDataBlock, layers: Vec, num_rows: u64, - ) -> (MiniBlockCompressed, CompressiveEncoding) { + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { // Count size to calculate rows per chunk let mut ceil_bytes_validity = 0; let mut cum_dim = 1; @@ -198,7 +204,7 @@ impl ValueEncoder { }; let est_bytes_per_word = (ceil_bytes_validity * vals_per_word) + cum_bytes_per_word; let (log_rows_per_chunk, rows_per_chunk) = - Self::find_log_vals_per_chunk(est_bytes_per_word, vals_per_word); + Self::find_log_vals_per_chunk(est_bytes_per_word, vals_per_word)?; let num_chunks = num_rows.div_ceil(rows_per_chunk) as usize; @@ -258,17 +264,17 @@ impl ValueEncoder { .chain(std::iter::once(data.data)) .collect::>(); - ( + Ok(( MiniBlockCompressed { chunks, data: buffers, num_values: num_rows, }, encoding, - ) + )) } - fn miniblock_fsl(data: DataBlock) -> (MiniBlockCompressed, CompressiveEncoding) { + fn miniblock_fsl(data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { let num_rows = data.num_values(); let fsl = data.as_fixed_size_list().unwrap(); let mut layers = Vec::new(); @@ -469,9 +475,9 @@ impl MiniBlockCompressor for ValueEncoder { match chunk { DataBlock::FixedWidth(fixed_width) => { let encoding = ProtobufUtils21::flat(fixed_width.bits_per_value, None); - Ok((Self::chunk_data(fixed_width), encoding)) + Ok((Self::chunk_data(fixed_width)?, encoding)) } - DataBlock::FixedSizeList(_) => Ok(Self::miniblock_fsl(chunk)), + DataBlock::FixedSizeList(_) => Self::miniblock_fsl(chunk), _ => Err(Error::invalid_input_source( format!( "Cannot compress a data block of type {} with ValueEncoder", @@ -989,6 +995,59 @@ mod tests { assert_eq!(decompressed.as_ref(), &sample_list); } + fn wide_fixed_size_binary() -> ArrayRef { + let wide_value = vec![0xABu8; 5000]; + Arc::new( + arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size( + std::iter::repeat_n(Some(wide_value.as_slice()), 4), + 5000, + ) + .unwrap(), + ) + } + + fn wide_fixed_size_list_bool() -> ArrayRef { + // A wide FSL is sub-byte, so it chunks eight values per word and the + // smallest unit is 16 values rather than 2. + let dimension = 4095; + let values = arrow_array::BooleanArray::from(vec![false; dimension * 2]); + let field = Arc::new(Field::new("item", DataType::Boolean, true)); + Arc::new(FixedSizeListArray::new( + field, + dimension as i32, + Arc::new(values), + None, + )) + } + + #[rstest::rstest] + #[case::fixed_size_binary(wide_fixed_size_binary(), 2)] + #[case::fixed_size_list_bool(wide_fixed_size_list_bool(), 16)] + fn test_wide_value_miniblock_returns_error( + #[case] array: ArrayRef, + #[case] expected_min_values: u64, + ) { + let starting_data = DataBlock::from_array(array); + + let encoder = ValueEncoder::default(); + let result = MiniBlockCompressor::compress(&encoder, starting_data); + + let err = result.expect_err("wide values should not be encodable as miniblock"); + assert!( + matches!(err, lance_core::Error::InvalidInput { .. }), + "expected InvalidInput, got {err:?}" + ); + let msg = err.to_string(); + assert!( + msg.contains("too wide for miniblock encoding"), + "unexpected error message: {msg}" + ); + assert!( + msg.contains(&format!("{expected_min_values} values require")), + "unexpected error message: {msg}" + ); + } + #[test] fn test_fsl_value_compression_per_value() { let sample_list = create_simple_fsl(); diff --git a/rust/lance-encoding/src/previous/encodings/logical/blob.rs b/rust/lance-encoding/src/previous/encodings/logical/blob.rs index 13fa3b346cb..20f6afc5b36 100644 --- a/rust/lance-encoding/src/previous/encodings/logical/blob.rs +++ b/rust/lance-encoding/src/previous/encodings/logical/blob.rs @@ -318,7 +318,7 @@ impl BlobFieldEncoder { .nulls() .cloned() .unwrap_or(NullBuffer::new_valid(binarray.len())); - for (w, is_valid) in binarray.value_offsets().windows(2).zip(nulls.into_iter()) { + for (w, is_valid) in binarray.value_offsets().windows(2).zip(&nulls) { if is_valid { let start = w[0] as u64; let end = w[1] as u64; diff --git a/rust/lance-encoding/src/repdef.rs b/rust/lance-encoding/src/repdef.rs index bdff69e0e4d..711ab732928 100644 --- a/rust/lance-encoding/src/repdef.rs +++ b/rust/lance-encoding/src/repdef.rs @@ -122,9 +122,9 @@ use crate::buffer::LanceBuffer; pub type LevelBuffer = Vec; -/// A contiguous top-level-row range that can be encoded as one structural page. +/// A top-level-row range whose dense rep/def stream fits one mini-block page. #[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct StructuralPageSplit { +pub(crate) struct MiniBlockRepDefSplit { /// Top-level row offset, relative to the original unsplit page. pub(crate) row_start: u64, /// Number of top-level rows in this split. @@ -137,15 +137,15 @@ pub(crate) struct StructuralPageSplit { pub(crate) num_values: u64, } -/// Planner result for structural page budget handling. +/// Dense mini-block rep/def budget result for one accumulated page. #[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) enum StructuralPagePlan { - /// The original page can be encoded as-is. - Fits, - /// The original page should be split on top-level row boundaries. - Split(Vec), - /// One top-level row is larger than the requested structural page budget. - UnsplittableOverBudget(u64), +pub(crate) enum MiniBlockRepDefBudget { + /// The dense rep/def stream fits one mini-block structural page. + WithinBudget, + /// The dense rep/def stream fits after splitting on top-level row boundaries. + RequiresPageSplit(Vec), + /// A single top-level row has this many rep/def levels and exceeds the budget. + SingleRowOverBudget(u64), } // As we build def levels we add this to special values to indicate that they @@ -806,21 +806,20 @@ impl SerializerContext { max_levels_per_page: Option, num_rows: u64, num_values: u64, - ) -> Result { + ) -> Result { // Extremely sparse lists can have many rep/def levels for very few // visible leaf values. If this ratio becomes too skewed then a - // miniblock structural chunk can exceed its packed rep/def metadata - // budget even though the value buffers are small. We detect that case - // while normalizing special def levels and split the structural page on - // top-level row boundaries so each emitted page stays within the - // miniblock structural budget. + // mini-block rep/def chunk can exceed its packed metadata budget even + // though the value buffers are small. We detect that case while + // normalizing special def levels and split on top-level row boundaries + // so each emitted dense mini-block page stays within the budget. if self.def_levels.is_empty() { - return Ok(StructuralPagePlan::Fits); + return Ok(MiniBlockRepDefBudget::WithinBudget); } if self.rep_levels.is_empty() { self.normalize_specials(); - return Ok(StructuralPagePlan::Fits); + return Ok(MiniBlockRepDefBudget::WithinBudget); } if self.rep_levels.len() != self.def_levels.len() { @@ -833,12 +832,12 @@ impl SerializerContext { let Some(max_levels_per_page) = max_levels_per_page else { self.normalize_specials(); - return Ok(StructuralPagePlan::Fits); + return Ok(MiniBlockRepDefBudget::WithinBudget); }; if num_values == 0 { self.normalize_specials(); - return Ok(StructuralPagePlan::Fits); + return Ok(MiniBlockRepDefBudget::WithinBudget); } let max_schema_rep = def_meaning.iter().filter(|level| level.is_list()).count() as u16; @@ -847,7 +846,7 @@ impl SerializerContext { if !should_plan { self.normalize_specials(); - return Ok(StructuralPagePlan::Fits); + return Ok(MiniBlockRepDefBudget::WithinBudget); } let max_visible_level = max_visible_level.unwrap(); @@ -855,7 +854,7 @@ impl SerializerContext { let mut counted_rows = 0u64; let mut counted_values = 0u64; let mut saw_structural_overhead = false; - let mut unsplittable_over_budget = None; + let mut single_row_over_budget_levels = None; let mut current_row_level_start = None; let mut current_row_num_values = 0u64; @@ -876,14 +875,14 @@ impl SerializerContext { saw_structural_overhead |= row_has_structural_overhead; if row_has_structural_overhead && row_num_levels > max_levels_per_page { - unsplittable_over_budget = Some(row_num_levels); + single_row_over_budget_levels = Some(row_num_levels); } if current_page_num_rows > 0 && (current_page_has_structural_overhead || row_has_structural_overhead) && current_page_num_levels + row_num_levels > max_levels_per_page { - splits.push(StructuralPageSplit { + splits.push(MiniBlockRepDefSplit { row_start: current_page_row_start, num_rows: current_page_num_rows, level_range: current_page_level_start..current_page_level_end, @@ -966,14 +965,14 @@ impl SerializerContext { ))); } if !saw_structural_overhead { - return Ok(StructuralPagePlan::Fits); + return Ok(MiniBlockRepDefBudget::WithinBudget); } - if let Some(row_num_levels) = unsplittable_over_budget { - return Ok(StructuralPagePlan::UnsplittableOverBudget(row_num_levels)); + if let Some(row_num_levels) = single_row_over_budget_levels { + return Ok(MiniBlockRepDefBudget::SingleRowOverBudget(row_num_levels)); } if current_page_num_rows > 0 { - splits.push(StructuralPageSplit { + splits.push(MiniBlockRepDefSplit { row_start: current_page_row_start, num_rows: current_page_num_rows, level_range: current_page_level_start..current_page_level_end, @@ -983,9 +982,9 @@ impl SerializerContext { } if splits.len() > 1 { - Ok(StructuralPagePlan::Split(splits)) + Ok(MiniBlockRepDefBudget::RequiresPageSplit(splits)) } else { - Ok(StructuralPagePlan::Fits) + Ok(MiniBlockRepDefBudget::WithinBudget) } } @@ -1023,12 +1022,12 @@ impl SerializerContext { ) } - fn build_with_structural_plan( + fn build_with_miniblock_repdef_budget( mut self, max_levels_per_page: Option, num_rows: u64, num_values: u64, - ) -> Result<(SerializedRepDefs, StructuralPagePlan)> { + ) -> Result<(SerializedRepDefs, MiniBlockRepDefBudget)> { if self.current_len == 0 { return Ok(( SerializedRepDefs::new_with_fixed_size_list_levels( @@ -1037,7 +1036,7 @@ impl SerializerContext { self.def_meaning, self.has_fsl, ), - StructuralPagePlan::Fits, + MiniBlockRepDefBudget::WithinBudget, )); } @@ -1046,7 +1045,7 @@ impl SerializerContext { .into_iter() .rev() .collect::>(); - let plan = self.normalize_specials_and_plan_splits( + let budget = self.normalize_specials_and_plan_splits( &def_meaning, max_levels_per_page, num_rows, @@ -1071,7 +1070,7 @@ impl SerializerContext { def_meaning, self.has_fsl, ), - plan, + budget, )) } } @@ -1412,15 +1411,15 @@ impl RepDefBuilder { Self::serialize_builders(builders).0.build() } - /// Converts gathered structural buffers into rep/def levels and an encode-time plan. - pub(crate) fn serialize_with_structural_plan( + /// Converts gathered structural buffers into rep/def levels and a mini-block budget result. + pub(crate) fn serialize_with_miniblock_repdef_budget( builders: Vec, max_levels_for_bits: impl FnOnce(u64) -> u64, num_rows: u64, num_values: u64, - ) -> Result<(SerializedRepDefs, StructuralPagePlan)> { + ) -> Result<(SerializedRepDefs, MiniBlockRepDefBudget)> { let (context, bits_per_level) = Self::serialize_builders(builders); - context.build_with_structural_plan( + context.build_with_miniblock_repdef_budget( bits_per_level.map(max_levels_for_bits), num_rows, num_values, @@ -1743,7 +1742,11 @@ impl RepDefUnraveler { } let num_new_lists = offsets.len() - old_offsets_len; offsets.push(to_offset(curlen)?); - rep_levels.truncate(offsets.len() - 1); + // Truncate to the number of lists THIS unraveler produced (write_idx), + // not `offsets.len() - 1` — the latter includes offsets contributed by + // earlier unravelers in a multi-page read, which would leave too many + // rep levels for the next (outer) layer and over-count its lists. + rep_levels.truncate(write_idx); if let Some(validity) = validity { // Even though we don't have validity it is possible another unraveler did and so we need // to push all valids @@ -2959,6 +2962,40 @@ mod tests { assert_eq!(val, None); } + #[test] + fn test_repdef_nested_list_multibatch_matches_single() { + // Single builder: List>, 3 rows. + // outer [0,2,3,5] -> rows have 2,1,2 inner lists + // inner [0,1,3,5,7,9] -> 5 inner lists, lengths 1,2,2,2,2 (9 leaf) + let mut single = RepDefBuilder::default(); + single.add_offsets(offsets_64(&[0, 2, 3, 5]), None); + single.add_offsets(offsets_64(&[0, 1, 3, 5, 7, 9]), None); + single.add_no_null(9); + let single_rep = RepDefBuilder::serialize(vec![single]) + .repetition_levels + .unwrap(); + + // Same logical data split into two batches: + // batch0 = rows 0,1 : outer [0,2,3], inner [0,1,3,5] (3 inner, 5 leaf) + // batch1 = row 2 : outer [0,2], inner [0,2,4] (2 inner, 4 leaf) + let mut b0 = RepDefBuilder::default(); + b0.add_offsets(offsets_64(&[0, 2, 3]), None); + b0.add_offsets(offsets_64(&[0, 1, 3, 5]), None); + b0.add_no_null(5); + let mut b1 = RepDefBuilder::default(); + b1.add_offsets(offsets_64(&[0, 2]), None); + b1.add_offsets(offsets_64(&[0, 2, 4]), None); + b1.add_no_null(4); + let multi_rep = RepDefBuilder::serialize(vec![b0, b1]) + .repetition_levels + .unwrap(); + + assert_eq!( + *single_rep, *multi_rep, + "multi-batch nested-list rep levels must equal single-batch" + ); + } + #[test] fn test_only_empty_lists() { let mut builder = RepDefBuilder::default(); diff --git a/rust/lance-file/src/previous/writer/mod.rs b/rust/lance-file/src/previous/writer/mod.rs index 4b04f722925..ab13e782367 100644 --- a/rust/lance-file/src/previous/writer/mod.rs +++ b/rust/lance-file/src/previous/writer/mod.rs @@ -29,6 +29,7 @@ use tokio::io::AsyncWriteExt; use crate::format::{MAGIC, MAJOR_VERSION, MINOR_VERSION}; use crate::previous::format::metadata::{Metadata, StatisticsMetadata}; use crate::previous::page_table::{PageInfo, PageTable}; +use crate::writer::FileWriteSummary; /// The file format currently includes a "manifest" where it stores the schema for /// self-describing files. Historically this has been a table format manifest that @@ -243,23 +244,29 @@ impl FileWriter { pub async fn finish_with_metadata( &mut self, metadata: &HashMap, - ) -> Result { + ) -> Result { self.schema .metadata .extend(metadata.iter().map(|(k, y)| (k.clone(), y.clone()))); self.finish().await } - pub async fn finish(&mut self) -> Result { + pub async fn finish(&mut self) -> Result { self.write_footer().await?; - Writer::shutdown(self.object_writer.as_mut()).await?; + // `shutdown` flushes the footer and reports the authoritative on-disk + // byte count, so the size is sourced here instead of via a later + // `tell()` that would rely on the cursor surviving shutdown. + let write_result = Writer::shutdown(self.object_writer.as_mut()).await?; let num_rows = self .metadata .batch_offsets .last() .cloned() .unwrap_or_default(); - Ok(num_rows as usize) + Ok(FileWriteSummary { + num_rows: num_rows as u64, + size_bytes: write_result.size as u64, + }) } /// Total records written in this file. diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index 0fc0baf0ee0..cea89235d73 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -519,16 +519,20 @@ const FOOTER_LEN: usize = 40; // How a field maps onto physical columns, shared by the projection-building and // projection-validation walks so they stay in lockstep. In the 2.0 layout every -// field (including structs and lists) has its own column; in 2.1 only leaves do, -// and blob/packed-struct fields are opaque (a single column with no descent). +// ordinary field (including structs and lists) has its own column; in 2.1 only +// leaves do. Blob/packed-struct fields are opaque in all versions: they are a +// single column with no descent, including unloaded blob descriptor schemas. // Returns `(contributes, recurse)`: whether the field has its own column and // whether to walk into its children. The DFS order is the field's own column (if // any) followed by its children, so a field's root (first) column is always the // first entry of its sub-slice. fn field_column_shape(field: &Field, is_structural: bool) -> (bool, bool) { - let contributes = - !is_structural || field.children.is_empty() || field.is_blob() || field.is_packed_struct(); - let recurse = !is_structural || (!field.is_blob() && !field.is_packed_struct()); + if field.is_blob() || field.is_packed_struct() { + return (true, false); + } + + let contributes = !is_structural || field.children.is_empty(); + let recurse = !field.children.is_empty(); (contributes, recurse) } @@ -1935,7 +1939,7 @@ impl FileMetadataProvider { .collect::>(); let metadata_bytes = io.submit_request(ranges, 0).await?; for ((result_index, column_index, _), bytes) in - missing_columns.into_iter().zip(metadata_bytes.into_iter()) + missing_columns.into_iter().zip(metadata_bytes) { let column_metadata = pbfile::ColumnMetadata::decode(bytes)?; let column_info = FileReader::meta_to_col_info( @@ -2574,7 +2578,11 @@ impl EncodedBatchReaderExt for EncodedBatch { #[cfg(test)] mod tests { - use std::{collections::BTreeMap, pin::Pin, sync::Arc}; + use std::{ + collections::{BTreeMap, HashMap}, + pin::Pin, + sync::Arc, + }; use arrow_array::{ RecordBatch, UInt32Array, @@ -2583,7 +2591,7 @@ mod tests { use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema}; use bytes::Bytes; use futures::{StreamExt, prelude::stream::TryStreamExt}; - use lance_arrow::RecordBatchExt; + use lance_arrow::{BLOB_META_KEY, RecordBatchExt}; use lance_core::{ArrowResult, datatypes::Schema}; use lance_datagen::{BatchCount, ByteCount, RowCount, array, gen_batch}; use lance_encoding::{ @@ -3884,6 +3892,38 @@ mod tests { ); } + #[test] + fn test_validate_v2_0_unloaded_blob_projection_is_opaque() { + let metadata = HashMap::from([(BLOB_META_KEY.to_string(), "true".to_string())]); + let arrow = ArrowSchema::new(vec![ + Field::new("blob", DataType::LargeBinary, true).with_metadata(metadata), + ]); + let mut schema = Schema::try_from(&arrow).unwrap(); + schema.fields[0].unloaded_mut(); + let projection = ReaderProjection { + schema: Arc::new(schema), + column_indices: vec![0], + }; + let column_len = |column: usize| { + assert_eq!(column, 0); + Ok(3) + }; + let mut cursor = 0usize; + + let rows = validate_field_length( + &projection.schema.fields[0], + false, + true, + &projection.column_indices, + &mut cursor, + &column_len, + ) + .unwrap(); + + assert_eq!(rows, 3); + assert_eq!(cursor, 1); + } + #[test] fn test_validate_length_list_and_empty_struct() { let validate = |dt: DataType, diff --git a/rust/lance-index/Cargo.toml b/rust/lance-index/Cargo.toml index 85de43c0f9b..b9d1a5fde29 100644 --- a/rust/lance-index/Cargo.toml +++ b/rust/lance-index/Cargo.toml @@ -44,6 +44,7 @@ lance-encoding.workspace = true lance-file.workspace = true lance-geo = { workspace = true, optional = true } lance-io.workspace = true +lance-bitpacking.workspace = true libsais-rs = "0.2" lance-linalg.workspace = true lance-select.workspace = true @@ -70,7 +71,6 @@ bytes.workspace = true chrono.workspace = true uuid.workspace = true async-channel = "2.3.1" -bitpacking = { version = "0.9.2", features = ["bitpacker4x"] } rand_distr.workspace = true lance-datagen.workspace = true rangemap.workspace = true diff --git a/rust/lance-index/benches/geo.rs b/rust/lance-index/benches/geo.rs index a9c403cc6bf..078aaefc6fb 100644 --- a/rust/lance-index/benches/geo.rs +++ b/rust/lance-index/benches/geo.rs @@ -15,7 +15,7 @@ use geoarrow_schema::Dimension; use lance_core::cache::LanceCache; use lance_core::{Error, ROW_ID}; use lance_index::scalar::lance_format::LanceIndexStore; -use lance_index::scalar::registry::ScalarIndexPlugin; +use lance_index::scalar::registry::BasicTrainer; use lance_index::scalar::rtree::{BoundingBox, RTreeIndex, RTreeIndexPlugin, RTreeTrainingRequest}; use lance_index::scalar::{GeoQuery, RelationQuery, ScalarIndex}; use lance_io::object_store::ObjectStore; diff --git a/rust/lance-index/protos-cache/cache.proto b/rust/lance-index/protos-cache/cache.proto index b24a27055d7..a861e2ef652 100644 --- a/rust/lance-index/protos-cache/cache.proto +++ b/rust/lance-index/protos-cache/cache.proto @@ -28,6 +28,11 @@ message CompressedPostingHeader { PositionStorage position_storage = 4; // Only meaningful when position_storage == POSITION_STORAGE_SHARED. PositionStreamCodec position_stream_codec = 5; + // Number of documents in each compressed posting block. Older cache entries + // omit this field and decode as the legacy 128-doc block size. + uint32 block_size = 6; + // Whether an impact IPC section follows the posting/position sections. + bool has_impacts = 7; } // Header for a serialized `PlainPostingList` cache entry. Followed by an Arrow diff --git a/rust/lance-index/src/frag_reuse.rs b/rust/lance-index/src/frag_reuse.rs index d42b41ca9f0..4c70db44094 100644 --- a/rust/lance-index/src/frag_reuse.rs +++ b/rust/lance-index/src/frag_reuse.rs @@ -10,15 +10,40 @@ use std::any::Any; use std::sync::Arc; +use arrow_array::RecordBatch; use async_trait::async_trait; use lance_core::{Error, Result}; -use roaring::RoaringBitmap; +use lance_select::RowAddrTreeMap; +use roaring::{RoaringBitmap, RoaringTreemap}; use serde::Serialize; pub use lance_table::system_index::frag_reuse::*; +use crate::scalar::RowIdRemapper; use crate::{Index, IndexType}; +impl RowIdRemapper for FragReuseIndex { + fn remap_row_id(&self, row_id: u64) -> Option { + self.remap_row_id(row_id) + } + + fn remap_row_addrs_tree_map(&self, row_addrs: &RowAddrTreeMap) -> RowAddrTreeMap { + self.remap_row_addrs_tree_map(row_addrs) + } + + fn remap_row_ids_roaring_tree_map(&self, row_ids: &RoaringTreemap) -> RoaringTreemap { + self.remap_row_ids_roaring_tree_map(row_ids) + } + + fn remap_row_ids_record_batch( + &self, + batch: RecordBatch, + row_id_idx: usize, + ) -> Result { + self.remap_row_ids_record_batch(batch, row_id_idx) + } +} + #[derive(Serialize)] struct FragReuseStatistics { num_versions: usize, diff --git a/rust/lance-index/src/scalar.rs b/rust/lance-index/src/scalar.rs index 448123a1024..21ae7aac71c 100644 --- a/rust/lance-index/src/scalar.rs +++ b/rust/lance-index/src/scalar.rs @@ -13,6 +13,7 @@ use datafusion::functions::string::contains::ContainsFunc; use datafusion::functions_nested::array_has; use datafusion::physical_plan::SendableRecordBatchStream; use datafusion_common::{Column, scalar::ScalarValue}; +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::collections::{HashMap, HashSet}; use std::fmt::Debug; use std::pin::Pin; @@ -27,7 +28,7 @@ use lance_core::deepsize::DeepSizeOf; use lance_core::{Error, Result}; use lance_io::stream::{RecordBatchStream, RecordBatchStreamAdapter}; use lance_select::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; -use roaring::RoaringBitmap; +use roaring::{RoaringBitmap, RoaringTreemap}; use serde::Serialize; use crate::metrics::MetricsCollector; @@ -51,7 +52,6 @@ pub mod rtree; pub mod zoned; pub mod zonemap; -use crate::frag_reuse::FragReuseIndex; pub use inverted::tokenizer::InvertedIndexParams; use lance_datafusion::udf::CONTAINS_TOKENS_UDF; @@ -1103,13 +1103,26 @@ pub trait ScalarIndex: Send + Sync + std::fmt::Debug + Index + DeepSizeOf { metrics: &dyn MetricsCollector, ) -> Result; + /// Returns true if this index reports matches as physical row addresses + /// (`fragment_id << 32 | offset`) rather than row ids + /// + /// Address-domain indices (e.g. zone map, bloom filter) are built over the + /// `_rowaddr` column. On a dataset with stable row ids the address and + /// row-id domains diverge, so these results must be translated back to row + /// ids (via the per-fragment row-id sequences, known only at the dataset + /// layer) before they are combined with row-id results or handed to the + /// scan. The default (row-id domain) needs no translation. + fn results_are_row_addresses(&self) -> bool { + false + } + /// Returns true if the remap operation is supported fn can_remap(&self) -> bool; /// Remap the row ids, creating a new remapped version of this index in `dest_store` async fn remap( &self, - mapping: &HashMap>, + mapping: &RowAddrRemap, dest_store: &dyn IndexStore, ) -> Result; @@ -1132,6 +1145,34 @@ pub trait ScalarIndex: Send + Sync + std::fmt::Debug + Index + DeepSizeOf { /// This returns a ScalarIndexParams that can be used to recreate an index /// with the same configuration on another dataset. fn derive_index_params(&self) -> Result; + + /// Global `[min, max]` of the indexed column from index metadata, without a + /// scan, or `None` if this index type cannot supply a sound bound. When + /// `Some`, the range is a superset of live values (conservative under + /// deletes): safe to prune with, not guaranteed tight. + fn value_range(&self) -> Option<(ScalarValue, ScalarValue)> { + None + } +} + +/// Abstraction over any type that can remap row IDs during index loading. +/// +/// This decouples scalar index plugins from the table-level [`crate::frag_reuse::FragReuseIndex`] +/// type. [`crate::frag_reuse::FragReuseIndex`] implements this trait, but callers may also +/// supply custom implementations for testing or other remapping strategies. +pub trait RowIdRemapper: Send + Sync + std::fmt::Debug { + /// Remap a single row id. Returns `None` if the row was deleted. + fn remap_row_id(&self, row_id: u64) -> Option; + /// Remap all addresses in a [`RowAddrTreeMap`], dropping deleted rows. + fn remap_row_addrs_tree_map(&self, row_addrs: &RowAddrTreeMap) -> RowAddrTreeMap; + /// Remap all row ids in a [`RoaringTreemap`], dropping deleted rows. + fn remap_row_ids_roaring_tree_map(&self, row_ids: &RoaringTreemap) -> RoaringTreemap; + /// Remap the row-id column at `row_id_idx` inside `batch`, dropping deleted rows. + fn remap_row_ids_record_batch( + &self, + batch: RecordBatch, + row_id_idx: usize, + ) -> Result; } #[cfg(test)] diff --git a/rust/lance-index/src/scalar/bitmap.rs b/rust/lance-index/src/scalar/bitmap.rs index c343f6b1063..91595b277f0 100644 --- a/rust/lance-index/src/scalar/bitmap.rs +++ b/rust/lance-index/src/scalar/bitmap.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::{ any::Any, cmp::Reverse, @@ -42,14 +43,13 @@ use super::{ use crate::pbold; use crate::{Index, IndexType, metrics::MetricsCollector}; use crate::{ - frag_reuse::FragReuseIndex, progress::IndexBuildProgress, scalar::{ - CreatedIndex, UpdateCriteria, + CreatedIndex, RowIdRemapper, UpdateCriteria, expression::SargableQueryParser, registry::{ - ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, - VALUE_COLUMN_NAME, + BasicTrainer, ScalarIndexLoad, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, + TrainingRequest, VALUE_COLUMN_NAME, single_flight_open, }, }, }; @@ -125,7 +125,7 @@ pub struct BitmapIndex { index_cache: WeakLanceCache, - frag_reuse_index: Option>, + frag_reuse_index: Option>, lazy_reader: LazyIndexReader, } @@ -196,11 +196,23 @@ impl BitmapIndexState { }) } + fn from_scalar_index(index: &dyn ScalarIndex) -> Result { + let bitmap = index + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::internal( + "BitmapIndexState::from_scalar_index called with a non-bitmap index", + ) + })?; + Self::from_index(bitmap) + } + pub(crate) fn to_bitmap_index( &self, store: Arc, index_cache: &LanceCache, - frag_reuse_index: Option>, + frag_reuse_index: Option>, ) -> Result> { Ok(Arc::new(BitmapIndex::new( self.index_map.clone(), @@ -335,7 +347,7 @@ impl BitmapIndex { value_type: DataType, store: Arc, index_cache: WeakLanceCache, - frag_reuse_index: Option>, + frag_reuse_index: Option>, ) -> Self { let lazy_reader = LazyIndexReader::new(store.clone()); Self { @@ -351,7 +363,7 @@ impl BitmapIndex { pub(crate) async fn load( store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, index_cache: &LanceCache, ) -> Result> { let page_lookup_file = store.open_index_file(BITMAP_LOOKUP_NAME).await?; @@ -784,7 +796,7 @@ impl ScalarIndex for BitmapIndex { /// Remap the row ids, creating a new remapped version of this index in `dest_store` async fn remap( &self, - mapping: &HashMap>, + mapping: &RowAddrRemap, dest_store: &dyn IndexStore, ) -> Result { let state = self.load_bitmap_index_state().await?; @@ -879,8 +891,7 @@ impl BitmapBatchWriter { return Ok(()); } let keys_array = - ScalarValue::iter_to_array(self.keys.drain(..).collect::>().into_iter()) - .unwrap(); + ScalarValue::iter_to_array(self.keys.drain(..).collect::>()).unwrap(); let total_size: usize = self.serialized.iter().map(|b| b.len()).sum(); let mut binary_builder = BinaryBuilder::with_capacity(self.serialized.len(), total_size); for b in self.serialized.drain(..) { @@ -1111,10 +1122,7 @@ async fn drain_same_key_bitmaps( let merged_key = OrderableScalarValue(key); advance_cursor_and_push(cursors, heap, item.shard_idx).await?; - loop { - let Some(Reverse(next_item)) = heap.peek() else { - break; - }; + while let Some(Reverse(next_item)) = heap.peek() { if next_item.key != merged_key { break; } @@ -1268,7 +1276,7 @@ impl BitmapIndexPlugin { let bitmap_size = bytes.len(); if cur_bytes + bitmap_size > MAX_BITMAP_ARRAY_LENGTH { - let keys_array = ScalarValue::iter_to_array(cur_keys.clone().into_iter()).unwrap(); + let keys_array = ScalarValue::iter_to_array(cur_keys.clone()).unwrap(); let mut binary_builder = BinaryBuilder::new(); for b in &cur_bitmaps { binary_builder.append_value(b); @@ -1538,7 +1546,7 @@ impl BitmapIndexPlugin { /// Remaps every bitmap in a materialized bitmap-index state using row-id mappings. pub(crate) fn remap_bitmap_state( state: HashMap, - mapping: &HashMap>, + mapping: &RowAddrRemap, ) -> HashMap { state .into_iter() @@ -1546,10 +1554,7 @@ impl BitmapIndexPlugin { let remapped_bitmap = RowAddrTreeMap::from_iter(bitmap.row_addrs().unwrap().filter_map(|addr| { let addr_as_u64 = u64::from(addr); - mapping - .get(&addr_as_u64) - .copied() - .unwrap_or(Some(addr_as_u64)) + mapping.get(addr_as_u64).unwrap_or(Some(addr_as_u64)) })); (key, remapped_bitmap) }) @@ -1708,11 +1713,7 @@ pub async fn merge_bitmap_indices( } #[async_trait] -impl ScalarIndexPlugin for BitmapIndexPlugin { - fn name(&self) -> &str { - "Bitmap" - } - +impl BasicTrainer for BitmapIndexPlugin { fn new_training_request( &self, params: &str, @@ -1731,27 +1732,6 @@ impl ScalarIndexPlugin for BitmapIndexPlugin { Ok(Box::new(BitmapTrainingRequest::new(params))) } - fn provides_exact_answer(&self) -> bool { - true - } - - fn version(&self) -> u32 { - BITMAP_INDEX_VERSION - } - - fn new_query_parser( - &self, - index_name: String, - _index_details: &prost_types::Any, - ) -> Option> { - // Bitmap indexes cannot answer `LikePrefix` queries (see `search`), so the parser - // is configured to skip them and let such predicates fall back to ordinary filtering. - Some(Box::new( - SargableQueryParser::new(index_name, self.name().to_string(), false) - .without_like_prefix(), - )) - } - async fn train_index( &self, data: SendableRecordBatchStream, @@ -1793,13 +1773,45 @@ impl ScalarIndexPlugin for BitmapIndexPlugin { files: vec![file], }) } +} + +#[async_trait] +impl ScalarIndexPlugin for BitmapIndexPlugin { + fn basic_trainer(&self) -> Option<&dyn BasicTrainer> { + Some(self) + } + + fn name(&self) -> &str { + "Bitmap" + } + + fn provides_exact_answer(&self) -> bool { + true + } + + fn version(&self) -> u32 { + BITMAP_INDEX_VERSION + } + + fn new_query_parser( + &self, + index_name: String, + _index_details: &prost_types::Any, + ) -> Option> { + // Bitmap indexes cannot answer `LikePrefix` queries (see `search`), so the parser + // is configured to skip them and let such predicates fall back to ordinary filtering. + Some(Box::new( + SargableQueryParser::new(index_name, self.name().to_string(), false) + .without_like_prefix(), + )) + } /// Load an index from storage async fn load_index( &self, index_store: Arc, _index_details: &prost_types::Any, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result> { Ok(BitmapIndex::load(index_store, frag_reuse_index, cache).await? as Arc) @@ -1808,7 +1820,7 @@ impl ScalarIndexPlugin for BitmapIndexPlugin { async fn get_from_cache( &self, index_store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result>> { let Some(state) = cache.get_with_key(&BitmapIndexStateKey).await else { @@ -1819,19 +1831,33 @@ impl ScalarIndexPlugin for BitmapIndexPlugin { } async fn put_in_cache(&self, cache: &LanceCache, index: Arc) -> Result<()> { - let bitmap = index - .as_any() - .downcast_ref::() - .ok_or_else(|| { - Error::internal("BitmapIndexPlugin::put_in_cache called with a non-bitmap index") - })?; - let state = BitmapIndexState::from_index(bitmap)?; + let state = BitmapIndexState::from_scalar_index(index.as_ref())?; cache .insert_with_key(&BitmapIndexStateKey, Arc::new(state)) .await; Ok(()) } + async fn get_or_insert_in_cache( + &self, + index_store: Arc, + frag_reuse_index: Option>, + cache: &LanceCache, + load: ScalarIndexLoad<'_>, + ) -> Result> { + single_flight_open( + cache, + BitmapIndexStateKey, + load, + BitmapIndexState::from_scalar_index, + move |state| { + Ok(state.to_bitmap_index(index_store, cache, frag_reuse_index)? + as Arc) + }, + ) + .await + } + async fn load_statistics( &self, index_store: Arc, @@ -1877,7 +1903,7 @@ mod tests { use lance_core::utils::{address::RowAddress, tempfile::TempObjDir}; use lance_io::object_store::ObjectStore; use lance_select::RowSetOps; - use std::collections::HashMap; + use rstest::rstest; fn assert_state_roundtrips(state: &BitmapIndexState) { let mut buf = Vec::new(); @@ -2428,8 +2454,44 @@ mod tests { assert_eq!(red_rows_2, vec![0, 3, 6, 10, 11]); } + // frags 1 and 2 (3 rows each) are compacted into frag 3: the 6 rows are + // rewritten in order to frag 3 offsets 0..6. + fn bitmap_remap_compact() -> RowAddrRemap { + use lance_core::utils::row_addr_remap::GroupInput; + use roaring::RoaringTreemap; + RowAddrRemap::compact([GroupInput { + rewritten_old_row_addrs: RoaringTreemap::from_iter( + (0..3) + .map(|o| u64::from(RowAddress::new_from_parts(1, o))) + .chain((0..3).map(|o| u64::from(RowAddress::new_from_parts(2, o)))), + ), + old_frag_ids: vec![1, 2], + new_frags: vec![(3, 6)], + }]) + .unwrap() + } + + fn bitmap_remap_explicit() -> RowAddrRemap { + // The same mapping, listed out explicitly. + RowAddrRemap::direct( + (0..6u32) + .map(|i| { + let (f, o) = if i < 3 { (1, i) } else { (2, i - 3) }; + ( + u64::from(RowAddress::new_from_parts(f, o)), + Some(u64::from(RowAddress::new_from_parts(3, i))), + ) + }) + .collect(), + ) + } + + // remap must behave identically whether the mapping is compact or explicit. + #[rstest] + #[case(bitmap_remap_compact())] + #[case(bitmap_remap_explicit())] #[tokio::test] - async fn test_remap_bitmap_with_null() { + async fn test_remap_bitmap_with_null(#[case] remap: RowAddrRemap) { use arrow_array::UInt32Array; // Create a temporary store. @@ -2494,38 +2556,8 @@ mod tests { assert_eq!(index.index_map.len(), 2); // 2 non-null values (1 and 2) assert!(!index.null_map.is_empty()); // Should have null values - // Create a remap that simulates compaction of frags 1 and 2 into frag 3 - let mut row_addr_map = HashMap::>::new(); - row_addr_map.insert( - RowAddress::new_from_parts(1, 0).into(), - Some(RowAddress::new_from_parts(3, 0).into()), - ); - row_addr_map.insert( - RowAddress::new_from_parts(1, 1).into(), - Some(RowAddress::new_from_parts(3, 1).into()), - ); - row_addr_map.insert( - RowAddress::new_from_parts(1, 2).into(), - Some(RowAddress::new_from_parts(3, 2).into()), - ); - row_addr_map.insert( - RowAddress::new_from_parts(2, 0).into(), - Some(RowAddress::new_from_parts(3, 3).into()), - ); - row_addr_map.insert( - RowAddress::new_from_parts(2, 1).into(), - Some(RowAddress::new_from_parts(3, 4).into()), - ); - row_addr_map.insert( - RowAddress::new_from_parts(2, 2).into(), - Some(RowAddress::new_from_parts(3, 5).into()), - ); - // Perform remap - index - .remap(&row_addr_map, test_store.as_ref()) - .await - .unwrap(); + index.remap(&remap, test_store.as_ref()).await.unwrap(); // Reload and check let reloaded_idx = BitmapIndex::load(test_store, None, &LanceCache::no_cache()) diff --git a/rust/lance-index/src/scalar/bloomfilter.rs b/rust/lance-index/src/scalar/bloomfilter.rs index 596ea4cc989..4a05cc8fc69 100644 --- a/rust/lance-index/src/scalar/bloomfilter.rs +++ b/rust/lance-index/src/scalar/bloomfilter.rs @@ -9,7 +9,7 @@ use crate::scalar::expression::{BloomFilterQueryParser, ScalarQueryParser}; use crate::scalar::registry::{ - ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, + BasicTrainer, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, }; use crate::scalar::{ BloomFilterQuery, BuiltinIndexType, CreatedIndex, IndexFile, ScalarIndexParams, UpdateCriteria, @@ -20,15 +20,19 @@ use arrow_schema::{DataType, Field}; use lance_arrow_stats::StatisticsAccumulator; use lance_core::utils::bloomfilter::as_bytes; use lance_core::utils::bloomfilter::sbbf::{Sbbf, SbbfBuilder}; +use lance_core::utils::row_addr_remap::RowAddrRemap; +use lance_select::RowAddrTreeMap; use serde::{Deserialize, Serialize}; +use std::collections::HashMap; use std::sync::LazyLock; use datafusion::execution::SendableRecordBatchStream; -use std::{collections::HashMap, sync::Arc}; +use std::sync::Arc; -use crate::scalar::FragReuseIndex; -use crate::scalar::{AnyQuery, IndexStore, MetricsCollector, ScalarIndex, SearchResult}; +use crate::scalar::{ + AnyQuery, IndexStore, MetricsCollector, RowIdRemapper, ScalarIndex, SearchResult, +}; use crate::{Index, IndexType}; use arrow_array::{ArrayRef, RecordBatch}; use async_trait::async_trait; @@ -42,6 +46,7 @@ use super::zoned::{ZoneBound, ZoneProcessor, ZoneTrainer, rebuild_zones, search_ const BLOOMFILTER_FILENAME: &str = "bloomfilter.lance"; const BLOOMFILTER_ITEM_META_KEY: &str = "bloomfilter_item"; +const NULL_BITMAP_META_KEY: &str = "null_bitmap"; const BLOOMFILTER_PROBABILITY_META_KEY: &str = "bloomfilter_probability"; const BLOOMFILTER_INDEX_VERSION: u32 = 0; @@ -78,18 +83,20 @@ pub struct BloomFilterIndex { number_of_items: u64, // Probability of false positives, fraction between 0 and 1 probability: f64, + // Exact set of null row addresses; None for older indices without this bitmap. + null_rows: Option, } impl DeepSizeOf for BloomFilterIndex { fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - self.zones.deep_size_of_children(context) + self.zones.deep_size_of_children(context) + self.null_rows.deep_size_of_children(context) } } impl BloomFilterIndex { async fn load( store: Arc, - _fri: Option>, + _fri: Option>, _index_cache: &LanceCache, ) -> Result> { let index_file = store.open_index_file(BLOOMFILTER_FILENAME).await?; @@ -110,10 +117,21 @@ impl BloomFilterIndex { .and_then(|bs| bs.parse().ok()) .unwrap_or(*DEFAULT_PROBABILITY); + let null_rows = if let Some(idx_str) = file_schema.metadata.get(NULL_BITMAP_META_KEY) { + let idx = idx_str.parse::().map_err(|e| { + Error::invalid_input(format!("invalid null bitmap buffer index: {e}")) + })?; + let bytes = index_file.read_global_buffer(idx).await?; + Some(RowAddrTreeMap::deserialize_from(bytes.as_ref())?) + } else { + None + }; + Ok(Arc::new(Self::try_from_serialized( bloom_data, number_of_items, probability, + null_rows, )?)) } @@ -121,13 +139,14 @@ impl BloomFilterIndex { data: RecordBatch, number_of_items: u64, probability: f64, + null_rows: Option, ) -> Result { if data.num_rows() == 0 { - // Return empty index for empty data return Ok(Self { zones: Vec::new(), number_of_items, probability, + null_rows, }); } @@ -208,6 +227,7 @@ impl BloomFilterIndex { zones: blocks, number_of_items, probability, + null_rows, }) } @@ -413,18 +433,28 @@ impl ScalarIndex for BloomFilterIndex { metrics: &dyn MetricsCollector, ) -> Result { let query = query.as_any().downcast_ref::().unwrap(); + if let BloomFilterQuery::IsNull() = query + && let Some(null_rows) = &self.null_rows + { + return Ok(SearchResult::exact(null_rows.clone())); + } + search_zones(&self.zones, metrics, |block| { self.evaluate_block_against_query(block, query) }) } + fn results_are_row_addresses(&self) -> bool { + true + } + fn can_remap(&self) -> bool { false } async fn remap( &self, - _mapping: &HashMap>, + _mapping: &RowAddrRemap, _dest_store: &dyn IndexStore, ) -> Result { Err(Error::invalid_input_source( @@ -446,18 +476,29 @@ impl ScalarIndex for BloomFilterIndex { let processor = BloomFilterProcessor::new(params.clone())?; let trainer = ZoneTrainer::new(processor, params.number_of_items)?; - let updated_blocks = rebuild_zones(&self.zones, trainer, new_data).await?; + let (updated_blocks, new_null_rows) = rebuild_zones(&self.zones, trainer, new_data).await?; + + // Merge existing and new null rows. If the existing index had no null bitmap + // (legacy format — null positions unknown), preserve that None: updating cannot + // recover the missing information, and claiming the result has zero nulls would + // be a false negative. Only a full retrain produces a fresh, complete bitmap. + let merged_null_rows = self.null_rows.as_ref().map(|existing| { + let mut merged = existing.clone(); + merged |= &new_null_rows; + merged + }); // Write the combined zones back to storage let mut builder = BloomFilterIndexBuilder::try_new(params)?; builder.blocks = updated_blocks; - let file = builder.write_index(dest_store).await?; + builder.null_rows = merged_null_rows; + let files = builder.write_index(dest_store).await?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&pb::BloomFilterIndexDetails::default()) .unwrap(), index_version: BLOOMFILTER_INDEX_VERSION, - files: vec![file], + files, }) } @@ -540,6 +581,10 @@ impl BloomFilterIndexBuilderParams { pub struct BloomFilterIndexBuilder { params: BloomFilterIndexBuilderParams, blocks: Vec, + // None means "legacy index — null positions unknown"; Some means a complete bitmap. + // write_index omits the null-bitmap global buffer when this is None, preserving the + // legacy format so that downstream searches remain conservative. + null_rows: Option, } impl BloomFilterIndexBuilder { @@ -547,6 +592,7 @@ impl BloomFilterIndexBuilder { Ok(Self { params, blocks: Vec::new(), + null_rows: None, }) } @@ -556,7 +602,9 @@ impl BloomFilterIndexBuilder { pub async fn train(&mut self, batches_source: SendableRecordBatchStream) -> Result<()> { let processor = BloomFilterProcessor::new(self.params.clone())?; let trainer = ZoneTrainer::new(processor, self.params.number_of_items)?; - self.blocks = trainer.train(batches_source).await?; + let (blocks, null_rows) = trainer.train(batches_source).await?; + self.blocks = blocks; + self.null_rows = Some(null_rows); Ok(()) } @@ -613,7 +661,7 @@ impl BloomFilterIndexBuilder { Ok(RecordBatch::try_new(schema, columns)?) } - pub async fn write_index(self, index_store: &dyn IndexStore) -> Result { + pub async fn write_index(self, index_store: &dyn IndexStore) -> Result> { let record_batch = self.bloomfilter_stats_as_batch()?; let mut file_schema = record_batch.schema().as_ref().clone(); @@ -621,7 +669,6 @@ impl BloomFilterIndexBuilder { BLOOMFILTER_ITEM_META_KEY.to_string(), self.params.number_of_items.to_string(), ); - file_schema.metadata.insert( BLOOMFILTER_PROBABILITY_META_KEY.to_string(), self.params.probability.to_string(), @@ -631,7 +678,24 @@ impl BloomFilterIndexBuilder { .new_index_file(BLOOMFILTER_FILENAME, Arc::new(file_schema)) .await?; index_file.write_record_batch(record_batch).await?; - index_file.finish().await + + let bloomfilter_file = if let Some(null_rows) = self.null_rows { + let mut null_bitmap_bytes = Vec::with_capacity(null_rows.serialized_size()); + null_rows.serialize_into(&mut null_bitmap_bytes)?; + let null_bitmap_idx = index_file + .add_global_buffer(bytes::Bytes::from(null_bitmap_bytes)) + .await?; + index_file + .finish_with_metadata(HashMap::from([( + NULL_BITMAP_META_KEY.to_string(), + null_bitmap_idx.to_string(), + )])) + .await? + } else { + index_file.finish_with_metadata(HashMap::new()).await? + }; + + Ok(vec![bloomfilter_file]) } } @@ -978,7 +1042,7 @@ impl BloomFilterIndexPlugin { batches_source: SendableRecordBatchStream, index_store: &dyn IndexStore, options: Option, - ) -> Result { + ) -> Result> { let mut builder = BloomFilterIndexBuilder::try_new(options.unwrap_or_default())?; builder.train(batches_source).await?; @@ -988,11 +1052,7 @@ impl BloomFilterIndexPlugin { } #[async_trait] -impl ScalarIndexPlugin for BloomFilterIndexPlugin { - fn name(&self) -> &str { - "BloomFilter" - } - +impl BasicTrainer for BloomFilterIndexPlugin { fn new_training_request( &self, params: &str, @@ -1067,14 +1127,21 @@ impl ScalarIndexPlugin for BloomFilterIndexPlugin { "must provide training request created by new_training_request".into(), ) })?; - let file = Self::train_bloomfilter_index(data, index_store, Some(request.params)).await?; + let files = Self::train_bloomfilter_index(data, index_store, Some(request.params)).await?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&pb::BloomFilterIndexDetails::default()) .unwrap(), index_version: BLOOMFILTER_INDEX_VERSION, - files: vec![file], + files, }) } +} + +#[async_trait] +impl ScalarIndexPlugin for BloomFilterIndexPlugin { + fn basic_trainer(&self) -> Option<&dyn BasicTrainer> { + Some(self) + } fn provides_exact_answer(&self) -> bool { false @@ -1084,6 +1151,10 @@ impl ScalarIndexPlugin for BloomFilterIndexPlugin { BLOOMFILTER_INDEX_VERSION } + fn name(&self) -> &str { + "BloomFilter" + } + fn new_query_parser( &self, index_name: String, @@ -1100,7 +1171,7 @@ impl ScalarIndexPlugin for BloomFilterIndexPlugin { &self, index_store: Arc, _index_details: &prost_types::Any, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result> { Ok( @@ -1160,7 +1231,7 @@ mod tests { use lance_select::RowAddrTreeMap; use crate::scalar::{ - BloomFilterQuery, ScalarIndex, SearchResult, + BloomFilterQuery, IndexStore, ScalarIndex, SearchResult, bloomfilter::{BloomFilterIndex, BloomFilterIndexBuilderParams}, lance_format::LanceIndexStore, }; @@ -2028,10 +2099,10 @@ mod tests { expected.insert_range(500..750); // Should match the zone containing 500 assert_eq!(result, SearchResult::at_most(expected)); - // Test IsNull query + // Test IsNull query (no nulls in data, should return exact empty set) let query = BloomFilterQuery::IsNull(); let result = index.search(&query, &NoOpMetricsCollector).await.unwrap(); - assert_eq!(result, SearchResult::at_most(RowAddrTreeMap::new())); // No nulls in the data + assert_eq!(result, SearchResult::exact(RowAddrTreeMap::new())); // Test IsIn query let query = BloomFilterQuery::IsIn(vec![ @@ -2124,4 +2195,152 @@ mod tests { _ => panic!("Expected AtMost search result from bloomfilter"), } } + + // Writes a bloomfilter file in the legacy format (no null bitmap global buffer), + // simulating an index created before the null bitmap feature was added. + async fn write_legacy_bloomfilter(store: &dyn IndexStore, has_null: bool) { + use crate::scalar::bloomfilter::{ + BLOOMFILTER_FILENAME, BLOOMFILTER_ITEM_META_KEY, BLOOMFILTER_PROBABILITY_META_KEY, + }; + use arrow_array::BooleanArray; + let schema = Arc::new(Schema::new(vec![ + Field::new("fragment_id", DataType::UInt64, false), + Field::new("zone_start", DataType::UInt64, false), + Field::new("zone_length", DataType::UInt64, false), + Field::new("has_null", DataType::Boolean, false), + Field::new("bloom_filter_data", DataType::Binary, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt64Array::from(vec![0u64])) as _, + Arc::new(UInt64Array::from(vec![0u64])) as _, + Arc::new(UInt64Array::from(vec![3u64])) as _, + Arc::new(BooleanArray::from(vec![has_null])) as _, + Arc::new(arrow_array::BinaryArray::from_vec(vec![b"".as_ref()])) as _, + ], + ) + .unwrap(); + let mut file_schema = schema.as_ref().clone(); + file_schema + .metadata + .insert(BLOOMFILTER_ITEM_META_KEY.to_string(), "1000".to_string()); + file_schema.metadata.insert( + BLOOMFILTER_PROBABILITY_META_KEY.to_string(), + "0.01".to_string(), + ); + let mut writer = store + .new_index_file(BLOOMFILTER_FILENAME, Arc::new(file_schema)) + .await + .unwrap(); + writer.write_record_batch(batch).await.unwrap(); + writer.finish().await.unwrap(); + } + + // Updating a legacy (null_rows = None) index must not silently treat None as + // "no nulls". The bug: `self.null_rows.clone().unwrap_or_default()` collapses + // None into an empty RowAddrTreeMap; after the merge the updated index has + // `null_rows = Some(empty)`, so an IsNull search returns `exact(empty)` — a + // false negative even though the legacy zone recorded has_null = true. + #[tokio::test] + async fn test_update_legacy_none_null_rows_not_treated_as_no_nulls() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + // Write a legacy-format index (no null bitmap) with has_null=true in its zone. + write_legacy_bloomfilter(store.as_ref(), true).await; + + let index = BloomFilterIndex::load(store.clone(), None, &LanceCache::no_cache()) + .await + .unwrap(); + assert!( + index.null_rows.is_none(), + "precondition: legacy null_rows is None" + ); + + // Update with new data from fragment 1 (no nulls). The destination is the + // same store so we can reload from it afterwards. + let new_schema = Arc::new(Schema::new(vec![ + Field::new(VALUE_COLUMN_NAME, DataType::Int32, true), + Field::new(ROW_ADDR, DataType::UInt64, false), + ])); + let new_batch = RecordBatch::try_new( + new_schema.clone(), + vec![ + Arc::new(arrow_array::Int32Array::from(vec![ + Some(10i32), + Some(20), + Some(30), + ])) as _, + Arc::new(UInt64Array::from_iter_values( + (0u64..3).map(|i| (1u64 << 32) | i), + )) as _, + ], + ) + .unwrap(); + let new_stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + new_schema, + stream::once(std::future::ready(Ok(new_batch))), + )); + + index + .update(new_stream, store.as_ref(), None) + .await + .unwrap(); + + let updated_index = BloomFilterIndex::load(store.clone(), None, &LanceCache::no_cache()) + .await + .unwrap(); + + // The legacy zone had has_null=true, so there ARE nulls at unknown positions. + // An IsNull search on the updated index must NOT claim "no nulls" (exact empty). + // It must be conservative and return AtMost, falling back to the has_null scan. + let result = updated_index + .search(&BloomFilterQuery::IsNull(), &NoOpMetricsCollector) + .await + .unwrap(); + + // With the bug: null_rows = Some(empty) → returns exact(empty) ← FALSE NEGATIVE + // With the fix: null_rows = None → falls through to has_null scan → AtMost + assert!( + !result.is_exact(), + "IsNull on an updated legacy index must not return exact(empty); \ + the legacy zone had has_null=true so nulls exist at unknown positions" + ); + } + + #[tokio::test] + async fn test_legacy_bloomfilter_no_null_bitmap() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + write_legacy_bloomfilter(store.as_ref(), true).await; + + let index = BloomFilterIndex::load(store, None, &LanceCache::no_cache()) + .await + .expect("failed to load legacy bloomfilter"); + + assert!( + index.null_rows.is_none(), + "legacy index should have no null bitmap" + ); + + // IS NULL should fall back to the has_null zone scan and return AtMost, not Exact. + let result = index + .search(&BloomFilterQuery::IsNull(), &NoOpMetricsCollector) + .await + .unwrap(); + assert!( + !result.is_exact(), + "IS NULL on a legacy index should not be exact" + ); + } } diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index 7668973b44e..bd63f05edd5 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::{ any::Any, cmp::Ordering, @@ -17,17 +18,19 @@ use super::{ }; use crate::cache_pb::{BTreeIndexHeader, RangeToFile}; use crate::{Index, IndexType}; +use crate::{metrics::NoOpMetricsCollector, scalar::registry::TrainingCriteria}; +use crate::{pbold, scalar::btree::flat::FlatIndex}; use crate::{ - frag_reuse::FragReuseIndex, progress::{IndexBuildProgress, noop_progress}, scalar::{ - CreatedIndex, UpdateCriteria, + CreatedIndex, RowIdRemapper, UpdateCriteria, expression::{SargableQueryParser, ScalarQueryParser}, - registry::{ScalarIndexPlugin, TrainingOrdering, TrainingRequest, VALUE_COLUMN_NAME}, + registry::{ + BasicTrainer, ScalarIndexLoad, ScalarIndexPlugin, TrainingOrdering, TrainingRequest, + VALUE_COLUMN_NAME, single_flight_open, + }, }, }; -use crate::{metrics::NoOpMetricsCollector, scalar::registry::TrainingCriteria}; -use crate::{pbold, scalar::btree::flat::FlatIndex}; use arrow_arith::numeric::add; use arrow_array::{ Array, ArrayAccessor, ArrowNativeTypeOp, PrimitiveArray, RecordBatch, UInt32Array, @@ -1389,11 +1392,22 @@ impl DeepSizeOf for BTreeIndexState { } impl BTreeIndexState { + fn from_index(index: &dyn ScalarIndex) -> Result { + let btree = index.as_any().downcast_ref::().ok_or_else(|| { + Error::internal("BTreeIndexState::from_index called with a non-BTree index") + })?; + Ok(Self { + lookup_batch: btree.page_lookup.batch.clone(), + batch_size: btree.batch_size, + ranges_to_files: btree.ranges_to_files.clone(), + }) + } + fn reconstruct( &self, store: Arc, index_cache: &LanceCache, - frag_reuse_index: Option>, + frag_reuse_index: Option>, ) -> Result> { let index = BTreeIndex::try_from_serialized( self.lookup_batch.clone(), @@ -1519,7 +1533,7 @@ pub struct BTreeIndex { /// - The local page_idx is calculated: `142 - 100 = 42`. /// - The system now knows to read page `42` from the file `part_2_page_file.lance`. ranges_to_files: Option>>, - frag_reuse_index: Option>, + frag_reuse_index: Option>, } impl DeepSizeOf for BTreeIndex { @@ -1540,7 +1554,7 @@ impl BTreeIndex { index_cache: WeakLanceCache, batch_size: u64, ranges_to_files: Option>>, - frag_reuse_index: Option>, + frag_reuse_index: Option>, ) -> Self { Self { page_lookup, @@ -1696,7 +1710,7 @@ impl BTreeIndex { index_cache: &LanceCache, batch_size: u64, ranges_to_files: Option>>, - frag_reuse_index: Option>, + frag_reuse_index: Option>, ) -> Result { let data_type = data.column(0).data_type().clone(); let page_lookup = Arc::new(BTreeLookup::try_new(data)?); @@ -1714,7 +1728,7 @@ impl BTreeIndex { async fn load( store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, index_cache: &LanceCache, ) -> Result> { let (page_lookup_file, standalone_partition_page_file) = @@ -1950,7 +1964,7 @@ fn filter_keeps_nothing(filter: &Option) -> bool { fn remap_row_ids( stream: SendableRecordBatchStream, - frag_reuse_index: Arc, + frag_reuse_index: Arc, ) -> SendableRecordBatchStream { let schema = stream.schema(); let remapped = stream.map(move |batch_result| { @@ -2214,7 +2228,7 @@ impl ScalarIndex for BTreeIndex { async fn remap( &self, - mapping: &HashMap>, + mapping: &RowAddrRemap, dest_store: &dyn IndexStore, ) -> Result { // (part_id, path) @@ -2381,7 +2395,7 @@ pub trait BTreeSubIndex: Debug + Send + Sync + DeepSizeOf { async fn remap_subindex( &self, serialized: RecordBatch, - mapping: &HashMap>, + mapping: &RowAddrRemap, ) -> Result; } @@ -3196,11 +3210,7 @@ impl TrainingRequest for BTreeTrainingRequest { pub struct BTreeIndexPlugin; #[async_trait] -impl ScalarIndexPlugin for BTreeIndexPlugin { - fn name(&self) -> &str { - "BTree" - } - +impl BasicTrainer for BTreeIndexPlugin { fn new_training_request( &self, params: &str, @@ -3216,26 +3226,6 @@ impl ScalarIndexPlugin for BTreeIndexPlugin { Ok(Box::new(BTreeTrainingRequest::new(params))) } - fn provides_exact_answer(&self) -> bool { - true - } - - fn version(&self) -> u32 { - BTREE_INDEX_VERSION - } - - fn new_query_parser( - &self, - index_name: String, - _index_details: &prost_types::Any, - ) -> Option> { - Some(Box::new(SargableQueryParser::new( - index_name, - self.name().to_string(), - false, - ))) - } - async fn train_index( &self, data: SendableRecordBatchStream, @@ -3280,12 +3270,43 @@ impl ScalarIndexPlugin for BTreeIndexPlugin { files, }) } +} + +#[async_trait] +impl ScalarIndexPlugin for BTreeIndexPlugin { + fn basic_trainer(&self) -> Option<&dyn BasicTrainer> { + Some(self) + } + + fn name(&self) -> &str { + "BTree" + } + + fn provides_exact_answer(&self) -> bool { + true + } + + fn version(&self) -> u32 { + BTREE_INDEX_VERSION + } + + fn new_query_parser( + &self, + index_name: String, + _index_details: &prost_types::Any, + ) -> Option> { + Some(Box::new(SargableQueryParser::new( + index_name, + self.name().to_string(), + false, + ))) + } async fn load_index( &self, index_store: Arc, _index_details: &prost_types::Any, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result> { Ok(BTreeIndex::load(index_store, frag_reuse_index, cache).await? as Arc) @@ -3294,7 +3315,7 @@ impl ScalarIndexPlugin for BTreeIndexPlugin { async fn get_from_cache( &self, index_store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result>> { let Some(state) = cache.get_with_key(&BTreeIndexStateKey).await else { @@ -3308,23 +3329,34 @@ impl ScalarIndexPlugin for BTreeIndexPlugin { } async fn put_in_cache(&self, cache: &LanceCache, index: Arc) -> Result<()> { - let btree = index.as_any().downcast_ref::().ok_or_else(|| { - Error::internal("BTreeIndexPlugin::put_in_cache called with a non-BTree index") - })?; - let state = BTreeIndexState { - lookup_batch: btree.page_lookup.batch.clone(), - batch_size: btree.batch_size, - ranges_to_files: btree.ranges_to_files.clone(), - }; + let state = BTreeIndexState::from_index(index.as_ref())?; cache .insert_with_key(&BTreeIndexStateKey, Arc::new(state)) .await; Ok(()) } + + async fn get_or_insert_in_cache( + &self, + index_store: Arc, + frag_reuse_index: Option>, + cache: &LanceCache, + load: ScalarIndexLoad<'_>, + ) -> Result> { + single_flight_open( + cache, + BTreeIndexStateKey, + load, + BTreeIndexState::from_index, + move |state| state.reconstruct(index_store, cache, frag_reuse_index), + ) + .await + } } #[cfg(test)] mod tests { + use lance_core::utils::row_addr_remap::RowAddrRemap; use std::sync::atomic::Ordering; use std::{collections::HashMap, sync::Arc}; @@ -3473,7 +3505,7 @@ mod tests { // Remap with a no-op mapping. The remapped index should be identical to the original index - .remap(&HashMap::default(), remap_store.as_ref()) + .remap(&RowAddrRemap::empty(), remap_store.as_ref()) .await .unwrap(); @@ -4984,7 +5016,7 @@ mod tests { // Remap with a no-op mapping. The remapped index should be identical to the original ranged_index - .remap(&HashMap::default(), remap_store.as_ref()) + .remap(&RowAddrRemap::empty(), remap_store.as_ref()) .await .unwrap(); @@ -5295,12 +5327,8 @@ mod tests { mapping.insert(old_id, None); } - let mut new_id_counter = 100_000; - // Remap all other rows - for old_id in (0..1000).chain(10000..15000) { - let new_id = new_id_counter; - new_id_counter += 1; + for (new_id, old_id) in (100_000..).zip((0..1000).chain(10000..15000)) { mapping.insert(old_id, Some(new_id)); } @@ -5312,7 +5340,10 @@ mod tests { )); // Remap the index with our deletion mapping - index.remap(&mapping, remap_store.as_ref()).await.unwrap(); + index + .remap(&RowAddrRemap::direct(mapping), remap_store.as_ref()) + .await + .unwrap(); let remapped_index = BTreeIndex::load(remap_store.clone(), None, &LanceCache::no_cache()) .await @@ -6286,11 +6317,12 @@ mod tests { // Remap row 0 -> row 5000 (outside the original [0, 1000) range so no collision). // Querying for value == 0 should now return row 5000, confirming reconstruct threaded // the FragReuseIndex through to the rebuilt BTreeIndex. - let frag_reuse_index = Arc::new(FragReuseIndex::new( - Uuid::new_v4(), - vec![HashMap::from([(0u64, Some(5000u64))])], - FragReuseIndexDetails { versions: vec![] }, - )); + let frag_reuse_index: Arc = + Arc::new(FragReuseIndex::new( + Uuid::new_v4(), + vec![HashMap::from([(0u64, Some(5000u64))])], + FragReuseIndexDetails { versions: vec![] }, + )); let reconstructed = state .reconstruct( test_store.clone(), diff --git a/rust/lance-index/src/scalar/btree/flat.rs b/rust/lance-index/src/scalar/btree/flat.rs index 744f6a3cb3c..1af5a6a69e9 100644 --- a/rust/lance-index/src/scalar/btree/flat.rs +++ b/rust/lance-index/src/scalar/btree/flat.rs @@ -1,7 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::collections::{BTreeSet, HashMap}; +use lance_core::utils::row_addr_remap::RowAddrRemap; +use std::collections::BTreeSet; use std::{ops::Bound, sync::Arc}; use arrow_array::Array; @@ -133,10 +134,7 @@ impl FlatIndex { NullableRowAddrSet::new(self.all_addrs_map.clone(), Default::default()) } - pub fn remap_batch( - batch: RecordBatch, - mapping: &HashMap>, - ) -> Result { + pub fn remap_batch(batch: RecordBatch, mapping: &RowAddrRemap) -> Result { let row_ids = batch.column(IDS_COL_IDX).as_primitive::(); let val_idx_and_new_id = row_ids .values() @@ -144,8 +142,7 @@ impl FlatIndex { .enumerate() .filter_map(|(idx, old_id)| { mapping - .get(old_id) - .copied() + .get(*old_id) .unwrap_or(Some(*old_id)) .map(|new_id| (idx, new_id)) }) @@ -221,13 +218,13 @@ impl FlatIndex { )); } } - (Bound::Included(lower) | Bound::Excluded(lower), Bound::Unbounded) => { - if lower.is_null() { - return Ok(NullableRowAddrSet::new( - Default::default(), - self.all_addrs_map.clone(), - )); - } + (Bound::Included(lower) | Bound::Excluded(lower), Bound::Unbounded) + if lower.is_null() => + { + return Ok(NullableRowAddrSet::new( + Default::default(), + self.all_addrs_map.clone(), + )); } _ => {} }, @@ -346,7 +343,11 @@ mod tests { use super::*; use arrow_array::{record_batch, types::Int32Type}; use datafusion_common::ScalarValue; + use lance_core::utils::row_addr_remap::GroupInput; use lance_datagen::{RowCount, array, gen_batch}; + use roaring::RoaringTreemap; + use rstest::rstest; + use std::collections::HashMap; fn example_index() -> FlatIndex { let batch = gen_batch() @@ -490,9 +491,10 @@ mod tests { // 3 -> delete // Keep remaining as is let mapping = HashMap::>::from_iter(vec![(0, Some(2000)), (3, None)]); - let remapped = - FlatIndex::try_new(FlatIndex::remap_batch((*index.data).clone(), &mapping).unwrap()) - .unwrap(); + let remapped = FlatIndex::try_new( + FlatIndex::remap_batch((*index.data).clone(), &RowAddrRemap::direct(mapping)).unwrap(), + ) + .unwrap(); let expected = FlatIndex::try_new( gen_batch() @@ -505,18 +507,22 @@ mod tests { assert_eq!(remapped.data, expected.data); } - // It's possible, during compaction, that an entire page of values is deleted. We just serialize - // it as an empty record batch. - #[tokio::test] - async fn test_remap_to_nothing() { + // An entire page (frag 0) is deleted during compaction. remap_batch must + // drop every row regardless of which RowAddrRemap mode expresses it. + // example_index holds row ids 5, 0, 3, 100, all in frag 0. + #[rstest] + #[case::compact(RowAddrRemap::compact([GroupInput { + rewritten_old_row_addrs: RoaringTreemap::new(), + old_frag_ids: vec![0], + new_frags: vec![], + }]) + .unwrap())] + #[case::explicit(RowAddrRemap::direct( + [5u64, 0, 3, 100].into_iter().map(|id| (id, None)).collect(), + ))] + fn test_remap_to_nothing(#[case] remap: RowAddrRemap) { let index = example_index(); - let mapping = HashMap::>::from_iter(vec![ - (5, None), - (0, None), - (3, None), - (100, None), - ]); - let remapped = FlatIndex::remap_batch((*index.data).clone(), &mapping).unwrap(); + let remapped = FlatIndex::remap_batch((*index.data).clone(), &remap).unwrap(); assert_eq!(remapped.num_rows(), 0); } diff --git a/rust/lance-index/src/scalar/expression.rs b/rust/lance-index/src/scalar/expression.rs index 8ef0a7217d5..cd76b68ac66 100644 --- a/rust/lance-index/src/scalar/expression.rs +++ b/rust/lance-index/src/scalar/expression.rs @@ -1402,6 +1402,20 @@ pub trait ScalarIndexLoader: Send + Sync { index_name: &str, metrics: &dyn MetricsCollector, ) -> Result>; + + /// Translate an address-domain index result into the row-id domain + /// + /// Address-domain indices (see [`ScalarIndex::results_are_row_addresses`]) + /// report matches as physical row addresses. The default returns `result` + /// unchanged, which is correct when addresses and row ids coincide (no + /// stable row ids). A dataset with stable row ids overrides this to remap + /// addresses to stable row ids via its per-fragment row-id sequences. + async fn row_addr_result_to_row_ids( + &self, + result: NullableIndexExprResult, + ) -> Result { + Ok(result) + } } /// This represents a search into a scalar index @@ -1462,6 +1476,175 @@ impl PartialEq for ScalarIndexExpr { } } +/// Returns the tighter (more restrictive) lower bound. +/// Priority: Included/Excluded > Unbounded; Excluded > Included for same value. +fn tighter_lower_bound(a: &Bound, b: &Bound) -> Bound { + match (a, b) { + (Bound::Unbounded, other) | (other, Bound::Unbounded) => other.clone(), + (Bound::Included(va), Bound::Included(vb)) => { + if va >= vb { + Bound::Included(va.clone()) + } else { + Bound::Included(vb.clone()) + } + } + (Bound::Excluded(va), Bound::Excluded(vb)) => { + if va >= vb { + Bound::Excluded(va.clone()) + } else { + Bound::Excluded(vb.clone()) + } + } + (Bound::Excluded(va), Bound::Included(vb)) => { + if va >= vb { + Bound::Excluded(va.clone()) + } else { + Bound::Included(vb.clone()) + } + } + (Bound::Included(va), Bound::Excluded(vb)) => { + if vb >= va { + Bound::Excluded(vb.clone()) + } else { + Bound::Included(va.clone()) + } + } + } +} + +/// Returns the tighter (more restrictive) upper bound. +/// Priority: Included/Excluded > Unbounded; Excluded > Included for same value. +fn tighter_upper_bound(a: &Bound, b: &Bound) -> Bound { + match (a, b) { + (Bound::Unbounded, other) | (other, Bound::Unbounded) => other.clone(), + (Bound::Included(va), Bound::Included(vb)) => { + if va <= vb { + Bound::Included(va.clone()) + } else { + Bound::Included(vb.clone()) + } + } + (Bound::Excluded(va), Bound::Excluded(vb)) => { + if va <= vb { + Bound::Excluded(va.clone()) + } else { + Bound::Excluded(vb.clone()) + } + } + (Bound::Excluded(va), Bound::Included(vb)) => { + if va <= vb { + Bound::Excluded(va.clone()) + } else { + Bound::Included(vb.clone()) + } + } + (Bound::Included(va), Bound::Excluded(vb)) => { + if vb <= va { + Bound::Excluded(vb.clone()) + } else { + Bound::Included(va.clone()) + } + } + } +} + +impl ScalarIndexExpr { + /// Optimize the expression tree by merging range queries on the same index. + /// + /// This collects all leaf Range queries from the AND tree, groups them by + /// index name, merges overlapping ranges into a single closed-range query, + /// and rebuilds the tree. This handles the case where `log_time >= X` and + /// `log_time <= Y` end up in different branches of a nested AND tree. + pub fn optimize(self) -> Self { + match self { + Self::And(_, _) => self.optimize_and_tree(), + Self::Or(lhs, rhs) => Self::Or(Box::new(lhs.optimize()), Box::new(rhs.optimize())), + Self::Not(inner) => Self::Not(Box::new(inner.optimize())), + other => other, + } + } + + /// Flatten an AND tree, merge ranges on same index, rebuild. + fn optimize_and_tree(self) -> Self { + let mut leaves = Vec::new(); + self.collect_and_leaves(&mut leaves); + + // Try to merge Range queries on the same index + let mut merged_indices: Vec = vec![false; leaves.len()]; + let mut result_leaves: Vec = Vec::new(); + + for i in 0..leaves.len() { + if merged_indices[i] { + continue; + } + let mut current = leaves[i].clone(); + + // Try to merge with subsequent leaves on the same index + for j in (i + 1)..leaves.len() { + if merged_indices[j] { + continue; + } + if let Some(merged) = try_merge_range_pair(¤t, &leaves[j]) { + current = merged; + merged_indices[j] = true; + } + } + + result_leaves.push(current); + } + + // Rebuild the AND tree from remaining leaves + let mut iter = result_leaves.into_iter(); + let first = iter.next().expect("AND tree must have at least one leaf"); + iter.fold(first, |acc, leaf| Self::And(Box::new(acc), Box::new(leaf))) + } + + /// Recursively collect all leaf nodes from an AND tree. + fn collect_and_leaves(self, leaves: &mut Vec) { + match self { + Self::And(lhs, rhs) => { + lhs.collect_and_leaves(leaves); + rhs.collect_and_leaves(leaves); + } + other => leaves.push(other), + } + } +} + +/// Try to merge two ScalarIndexExpr nodes if they are both Range queries on the same index. +fn try_merge_range_pair(lhs: &ScalarIndexExpr, rhs: &ScalarIndexExpr) -> Option { + let (ScalarIndexExpr::Query(l), ScalarIndexExpr::Query(r)) = (lhs, rhs) else { + return None; + }; + if l.index_name != r.index_name || l.column != r.column { + return None; + } + + let l_query = l.query.as_any().downcast_ref::()?; + let r_query = r.query.as_any().downcast_ref::()?; + + let (SargableQuery::Range(l_low, l_high), SargableQuery::Range(r_low, r_high)) = + (l_query, r_query) + else { + return None; + }; + + let merged_low = tighter_lower_bound(l_low, r_low); + let merged_high = tighter_upper_bound(l_high, r_high); + + Some(ScalarIndexExpr::Query(ScalarIndexSearch { + column: l.column.clone(), + index_name: l.index_name.clone(), + index_type: l.index_type.clone(), + query: Arc::new(SargableQuery::Range(merged_low, merged_high)), + needs_recheck: l.needs_recheck || r.needs_recheck, + // Both queries target the same index (checked above), so they share + // the same fragment coverage; carry it over to keep the merged query + // usable by coverage-dependent optimizer rules. + fragment_bitmap: l.fragment_bitmap.clone(), + })) +} + impl std::fmt::Display for ScalarIndexExpr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { @@ -1524,7 +1707,15 @@ impl ScalarIndexExpr { .load_index(&search.column, &search.index_name, metrics) .await?; let search_result = index.search(search.query.as_ref(), metrics).await?; - Ok(search_result.into()) + let result: NullableIndexExprResult = search_result.into(); + if index.results_are_row_addresses() { + // Translate address-domain results to the row-id domain + // before combining or scanning; otherwise stable-row-id + // datasets silently drop matches (issue #7434). + index_loader.row_addr_result_to_row_ids(result).await + } else { + Ok(result) + } } } } @@ -3630,4 +3821,830 @@ mod tests { // Query against an unindexed path must not bind to either index. check_no_index(&index_info, "json_extract(json, '$.c') = 'foo'"); } + + #[test] + fn test_optimize_nested_and_tree() { + use super::{ScalarIndexExpr, ScalarIndexSearch}; + use crate::scalar::SargableQuery; + use datafusion_common::ScalarValue; + use std::ops::Bound; + use std::sync::Arc; + + // Simulate: AND(AND(fqdn@idx_fqdn, log_time >= X @idx_time), AND(log_time <= Y @idx_time, channel@idx_ch)) + let fqdn_query = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "fqdn".to_string(), + index_name: "fqdn_idx".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Equals(ScalarValue::Utf8(Some( + "eng.bhd.0068".to_string(), + )))), + needs_recheck: false, + fragment_bitmap: None, + }); + let time_gte = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "log_time".to_string(), + index_name: "log_time_idx".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(100))), + Bound::Unbounded, + )), + needs_recheck: false, + fragment_bitmap: None, + }); + let time_lte = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "log_time".to_string(), + index_name: "log_time_idx".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range( + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(200))), + )), + needs_recheck: false, + fragment_bitmap: None, + }); + let channel_query = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "channel".to_string(), + index_name: "channel_idx".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Equals(ScalarValue::Utf8(Some( + "/ados/node/monitor".to_string(), + )))), + needs_recheck: false, + fragment_bitmap: None, + }); + + // Build nested AND: AND(AND(fqdn, time_gte), AND(time_lte, channel)) + let nested = ScalarIndexExpr::And( + Box::new(ScalarIndexExpr::And( + Box::new(fqdn_query), + Box::new(time_gte), + )), + Box::new(ScalarIndexExpr::And( + Box::new(time_lte), + Box::new(channel_query), + )), + ); + + // Optimize should merge time_gte + time_lte into a single Range + let optimized = nested.optimize(); + + // The optimized tree should contain a merged Range(Included(100), Included(200)) + // and the other queries as separate leaves + // Let's verify by collecting leaves + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + + // Should have 3 leaves: fqdn, merged_time, channel + assert_eq!( + leaves.len(), + 3, + "Expected 3 leaves after merge, got: {:?}", + leaves + ); + + // Find the merged time query + let time_query = leaves + .iter() + .find(|l| { + if let ScalarIndexExpr::Query(s) = l { + s.index_name == "log_time_idx" + } else { + false + } + }) + .expect("Should have a log_time query"); + + if let ScalarIndexExpr::Query(s) = time_query { + let range = s.query.as_any().downcast_ref::().unwrap(); + assert_eq!( + *range, + SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(100))), + Bound::Included(ScalarValue::Int64(Some(200))), + ), + "Range should be merged into closed interval" + ); + } + } + + #[test] + fn test_optimize_no_merge_different_indexes() { + use super::{ScalarIndexExpr, ScalarIndexSearch}; + use crate::scalar::SargableQuery; + use datafusion_common::ScalarValue; + use std::ops::Bound; + use std::sync::Arc; + + // Two range queries on DIFFERENT indexes should NOT be merged + let range_a = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "a".to_string(), + index_name: "idx_a".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(10))), + Bound::Unbounded, + )), + needs_recheck: false, + fragment_bitmap: None, + }); + let range_b = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "b".to_string(), + index_name: "idx_b".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range( + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(100))), + )), + needs_recheck: false, + fragment_bitmap: None, + }); + + let expr = ScalarIndexExpr::And(Box::new(range_a), Box::new(range_b)); + let optimized = expr.optimize(); + + // Should remain as two separate leaves (not merged) + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + assert_eq!(leaves.len(), 2); + } + + #[test] + fn test_optimize_no_merge_non_range() { + use super::{ScalarIndexExpr, ScalarIndexSearch}; + use crate::scalar::SargableQuery; + use datafusion_common::ScalarValue; + use std::ops::Bound; + use std::sync::Arc; + + // Equals + Range on same index should NOT merge + let eq_query = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "x".to_string(), + index_name: "idx_x".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Equals(ScalarValue::Int64(Some(42)))), + needs_recheck: false, + fragment_bitmap: None, + }); + let range_query = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "x".to_string(), + index_name: "idx_x".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range( + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(100))), + )), + needs_recheck: false, + fragment_bitmap: None, + }); + + let expr = ScalarIndexExpr::And(Box::new(eq_query), Box::new(range_query)); + let optimized = expr.optimize(); + + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + assert_eq!(leaves.len(), 2, "Equals + Range should not merge"); + } + + #[test] + fn test_optimize_exclusive_bounds() { + use super::{ScalarIndexExpr, ScalarIndexSearch}; + use crate::scalar::SargableQuery; + use datafusion_common::ScalarValue; + use std::ops::Bound; + use std::sync::Arc; + + // x > 10 AND x < 20 → Range(Excluded(10), Excluded(20)) + let gt = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "x".to_string(), + index_name: "idx_x".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range( + Bound::Excluded(ScalarValue::Int64(Some(10))), + Bound::Unbounded, + )), + needs_recheck: false, + fragment_bitmap: None, + }); + let lt = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "x".to_string(), + index_name: "idx_x".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range( + Bound::Unbounded, + Bound::Excluded(ScalarValue::Int64(Some(20))), + )), + needs_recheck: false, + fragment_bitmap: None, + }); + + let expr = ScalarIndexExpr::And(Box::new(gt), Box::new(lt)); + let optimized = expr.optimize(); + + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + assert_eq!(leaves.len(), 1); + + if let ScalarIndexExpr::Query(s) = &leaves[0] { + let range = s.query.as_any().downcast_ref::().unwrap(); + assert_eq!( + *range, + SargableQuery::Range( + Bound::Excluded(ScalarValue::Int64(Some(10))), + Bound::Excluded(ScalarValue::Int64(Some(20))), + ) + ); + } else { + panic!("Expected a Query leaf"); + } + } + + #[test] + fn test_optimize_preserves_or_and_not() { + use super::{ScalarIndexExpr, ScalarIndexSearch}; + use crate::scalar::SargableQuery; + use datafusion_common::ScalarValue; + use std::ops::Bound; + use std::sync::Arc; + + // AND(OR(a, b), Range_gte, Range_lte) + // OR node should be preserved, ranges should merge + let or_node = ScalarIndexExpr::Or( + Box::new(ScalarIndexExpr::Query(ScalarIndexSearch { + column: "c".to_string(), + index_name: "idx_c".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Equals(ScalarValue::Utf8(Some( + "a".to_string(), + )))), + needs_recheck: false, + fragment_bitmap: None, + })), + Box::new(ScalarIndexExpr::Query(ScalarIndexSearch { + column: "c".to_string(), + index_name: "idx_c".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Equals(ScalarValue::Utf8(Some( + "b".to_string(), + )))), + needs_recheck: false, + fragment_bitmap: None, + })), + ); + let range_gte = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "t".to_string(), + index_name: "idx_t".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(5))), + Bound::Unbounded, + )), + needs_recheck: false, + fragment_bitmap: None, + }); + let range_lte = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "t".to_string(), + index_name: "idx_t".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range( + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(50))), + )), + needs_recheck: false, + fragment_bitmap: None, + }); + + let expr = ScalarIndexExpr::And( + Box::new(ScalarIndexExpr::And(Box::new(or_node), Box::new(range_gte))), + Box::new(range_lte), + ); + let optimized = expr.optimize(); + + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + // Should have 2 leaves: OR node (preserved) + merged range + assert_eq!(leaves.len(), 2); + + // One leaf should be the merged range + let merged = leaves + .iter() + .find(|l| matches!(l, ScalarIndexExpr::Query(s) if s.index_name == "idx_t")) + .expect("Should have merged range"); + + if let ScalarIndexExpr::Query(s) = merged { + let range = s.query.as_any().downcast_ref::().unwrap(); + assert_eq!( + *range, + SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(5))), + Bound::Included(ScalarValue::Int64(Some(50))), + ) + ); + } + + // Other leaf should be the OR node + assert!( + leaves + .iter() + .any(|l| matches!(l, ScalarIndexExpr::Or(_, _))) + ); + } + + #[test] + fn test_optimize_needs_recheck_preserved() { + use super::{ScalarIndexExpr, ScalarIndexSearch}; + use crate::scalar::SargableQuery; + use datafusion_common::ScalarValue; + use std::ops::Bound; + use std::sync::Arc; + + // If either range has needs_recheck=true, merged result should too + let range_a = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "x".to_string(), + index_name: "idx_x".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(1))), + Bound::Unbounded, + )), + needs_recheck: true, + fragment_bitmap: None, + }); + let range_b = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "x".to_string(), + index_name: "idx_x".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range( + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(99))), + )), + needs_recheck: false, + fragment_bitmap: None, + }); + + let expr = ScalarIndexExpr::And(Box::new(range_a), Box::new(range_b)); + let optimized = expr.optimize(); + + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + assert_eq!(leaves.len(), 1); + + if let ScalarIndexExpr::Query(s) = &leaves[0] { + assert!( + s.needs_recheck, + "Merged query should preserve needs_recheck=true" + ); + } else { + panic!("Expected a Query leaf"); + } + } + + #[test] + fn test_optimize_single_query_passthrough() { + use super::{ScalarIndexExpr, ScalarIndexSearch}; + use crate::scalar::SargableQuery; + use datafusion_common::ScalarValue; + use std::ops::Bound; + use std::sync::Arc; + + // A single query (not in AND) should pass through unchanged + let single = ScalarIndexExpr::Query(ScalarIndexSearch { + column: "x".to_string(), + index_name: "idx_x".to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(1))), + Bound::Unbounded, + )), + needs_recheck: false, + fragment_bitmap: None, + }); + + let optimized = single.clone().optimize(); + assert_eq!(optimized, single); + } + + #[test] + fn test_optimize_complex_nested_cases() { + use super::{ScalarIndexExpr, ScalarIndexSearch}; + use crate::scalar::SargableQuery; + use datafusion_common::ScalarValue; + use std::ops::Bound; + use std::sync::Arc; + + let make_range = + |col: &str, idx: &str, low: Bound, high: Bound| { + ScalarIndexExpr::Query(ScalarIndexSearch { + column: col.to_string(), + index_name: idx.to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Range(low, high)), + needs_recheck: false, + fragment_bitmap: None, + }) + }; + let make_eq = |col: &str, idx: &str, val: ScalarValue| { + ScalarIndexExpr::Query(ScalarIndexSearch { + column: col.to_string(), + index_name: idx.to_string(), + index_type: "".to_string(), + query: Arc::new(SargableQuery::Equals(val)), + needs_recheck: false, + fragment_bitmap: None, + }) + }; + + // Case 1: Multiple ranges on same index should all merge + // x >= 10 AND x <= 200 AND x >= 50 AND x <= 100 → Range(50, 100) + { + let expr = ScalarIndexExpr::And( + Box::new(ScalarIndexExpr::And( + Box::new(make_range( + "x", + "idx_x", + Bound::Included(ScalarValue::Int64(Some(10))), + Bound::Unbounded, + )), + Box::new(make_range( + "x", + "idx_x", + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(200))), + )), + )), + Box::new(ScalarIndexExpr::And( + Box::new(make_range( + "x", + "idx_x", + Bound::Included(ScalarValue::Int64(Some(50))), + Bound::Unbounded, + )), + Box::new(make_range( + "x", + "idx_x", + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(100))), + )), + )), + ); + + let optimized = expr.optimize(); + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + assert_eq!(leaves.len(), 1, "All 4 ranges should merge into 1"); + + if let ScalarIndexExpr::Query(s) = &leaves[0] { + let range = s.query.as_any().downcast_ref::().unwrap(); + assert_eq!( + *range, + SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(50))), + Bound::Included(ScalarValue::Int64(Some(100))), + ) + ); + } else { + panic!("Expected merged Query"); + } + } + + // Case 2: NOT(range) should NOT be merged with sibling range + // AND(NOT(x >= 10), x <= 20) → stays as 2 leaves + { + let not_node = ScalarIndexExpr::Not(Box::new(make_range( + "x", + "idx_x", + Bound::Included(ScalarValue::Int64(Some(10))), + Bound::Unbounded, + ))); + let range_lte = make_range( + "x", + "idx_x", + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(20))), + ); + + let expr = ScalarIndexExpr::And(Box::new(not_node), Box::new(range_lte)); + let optimized = expr.optimize(); + + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + assert_eq!( + leaves.len(), + 2, + "NOT(range) should not merge with sibling range" + ); + assert!(leaves.iter().any(|l| matches!(l, ScalarIndexExpr::Not(_)))); + } + + // Case 3: Ranges inside OR branches get optimized independently + // OR(AND(x >= 10, x <= 20), AND(x >= 30, x <= 40)) + { + let branch1 = ScalarIndexExpr::And( + Box::new(make_range( + "x", + "idx_x", + Bound::Included(ScalarValue::Int64(Some(10))), + Bound::Unbounded, + )), + Box::new(make_range( + "x", + "idx_x", + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(20))), + )), + ); + let branch2 = ScalarIndexExpr::And( + Box::new(make_range( + "x", + "idx_x", + Bound::Included(ScalarValue::Int64(Some(30))), + Bound::Unbounded, + )), + Box::new(make_range( + "x", + "idx_x", + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(40))), + )), + ); + + let expr = ScalarIndexExpr::Or(Box::new(branch1), Box::new(branch2)); + let optimized = expr.optimize(); + + // Top level should still be OR + if let ScalarIndexExpr::Or(lhs, rhs) = &optimized { + // Each branch should be a single merged range + let mut left_leaves = Vec::new(); + lhs.clone().collect_and_leaves(&mut left_leaves); + assert_eq!( + left_leaves.len(), + 1, + "Left OR branch should have merged range" + ); + if let ScalarIndexExpr::Query(s) = &left_leaves[0] { + let range = s.query.as_any().downcast_ref::().unwrap(); + assert_eq!( + *range, + SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(10))), + Bound::Included(ScalarValue::Int64(Some(20))), + ) + ); + } + + let mut right_leaves = Vec::new(); + rhs.clone().collect_and_leaves(&mut right_leaves); + assert_eq!( + right_leaves.len(), + 1, + "Right OR branch should have merged range" + ); + if let ScalarIndexExpr::Query(s) = &right_leaves[0] { + let range = s.query.as_any().downcast_ref::().unwrap(); + assert_eq!( + *range, + SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(30))), + Bound::Included(ScalarValue::Int64(Some(40))), + ) + ); + } + } else { + panic!("Expected OR at top level"); + } + } + + // Case 4: Multiple indexes, each with its own range pair + // AND(a >= 1, a <= 10, b >= 100, b <= 200) + // → merged into: AND(a in [1,10], b in [100,200]) + { + let expr = ScalarIndexExpr::And( + Box::new(ScalarIndexExpr::And( + Box::new(make_range( + "a", + "idx_a", + Bound::Included(ScalarValue::Int64(Some(1))), + Bound::Unbounded, + )), + Box::new(make_range( + "a", + "idx_a", + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(10))), + )), + )), + Box::new(ScalarIndexExpr::And( + Box::new(make_range( + "b", + "idx_b", + Bound::Included(ScalarValue::Int64(Some(100))), + Bound::Unbounded, + )), + Box::new(make_range( + "b", + "idx_b", + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(200))), + )), + )), + ); + + let optimized = expr.optimize(); + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + assert_eq!( + leaves.len(), + 2, + "Two indexes should each merge independently" + ); + + let a_leaf = leaves + .iter() + .find(|l| matches!(l, ScalarIndexExpr::Query(s) if s.index_name == "idx_a")) + .expect("Should have idx_a"); + if let ScalarIndexExpr::Query(s) = a_leaf { + let range = s.query.as_any().downcast_ref::().unwrap(); + assert_eq!( + *range, + SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(1))), + Bound::Included(ScalarValue::Int64(Some(10))), + ) + ); + } + + let b_leaf = leaves + .iter() + .find(|l| matches!(l, ScalarIndexExpr::Query(s) if s.index_name == "idx_b")) + .expect("Should have idx_b"); + if let ScalarIndexExpr::Query(s) = b_leaf { + let range = s.query.as_any().downcast_ref::().unwrap(); + assert_eq!( + *range, + SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(100))), + Bound::Included(ScalarValue::Int64(Some(200))), + ) + ); + } + } + + // Case 5: Mix of Equals and Range on different indexes + // AND(fqdn = 'x', time >= 100, time <= 200, channel = 'y') + // → AND(fqdn = 'x', time in [100,200], channel = 'y') + { + let expr = ScalarIndexExpr::And( + Box::new(ScalarIndexExpr::And( + Box::new(make_eq( + "fqdn", + "fqdn_idx", + ScalarValue::Utf8(Some("x".to_string())), + )), + Box::new(make_range( + "time", + "time_idx", + Bound::Included(ScalarValue::Int64(Some(100))), + Bound::Unbounded, + )), + )), + Box::new(ScalarIndexExpr::And( + Box::new(make_range( + "time", + "time_idx", + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(200))), + )), + Box::new(make_eq( + "channel", + "ch_idx", + ScalarValue::Utf8(Some("y".to_string())), + )), + )), + ); + + let optimized = expr.optimize(); + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + assert_eq!(leaves.len(), 3, "fqdn + merged_time + channel = 3 leaves"); + + let time_leaf = leaves + .iter() + .find(|l| matches!(l, ScalarIndexExpr::Query(s) if s.index_name == "time_idx")) + .expect("Should have time query"); + if let ScalarIndexExpr::Query(s) = time_leaf { + let range = s.query.as_any().downcast_ref::().unwrap(); + assert_eq!( + *range, + SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(100))), + Bound::Included(ScalarValue::Int64(Some(200))), + ) + ); + } + } + + // Case 6: Overlapping closed ranges → takes intersection + // Range(10, 50) AND Range(30, 80) → Range(30, 50) + { + let expr = ScalarIndexExpr::And( + Box::new(make_range( + "x", + "idx_x", + Bound::Included(ScalarValue::Int64(Some(10))), + Bound::Included(ScalarValue::Int64(Some(50))), + )), + Box::new(make_range( + "x", + "idx_x", + Bound::Included(ScalarValue::Int64(Some(30))), + Bound::Included(ScalarValue::Int64(Some(80))), + )), + ); + + let optimized = expr.optimize(); + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + assert_eq!(leaves.len(), 1); + if let ScalarIndexExpr::Query(s) = &leaves[0] { + let range = s.query.as_any().downcast_ref::().unwrap(); + assert_eq!( + *range, + SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(30))), + Bound::Included(ScalarValue::Int64(Some(50))), + ) + ); + } + } + + // Case 7: Mixed Included/Excluded on same value + // x >= 5 AND x > 5 → Excluded(5) (stricter) + { + let expr = ScalarIndexExpr::And( + Box::new(make_range( + "x", + "idx_x", + Bound::Included(ScalarValue::Int64(Some(5))), + Bound::Unbounded, + )), + Box::new(make_range( + "x", + "idx_x", + Bound::Excluded(ScalarValue::Int64(Some(5))), + Bound::Unbounded, + )), + ); + + let optimized = expr.optimize(); + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + assert_eq!(leaves.len(), 1); + if let ScalarIndexExpr::Query(s) = &leaves[0] { + let range = s.query.as_any().downcast_ref::().unwrap(); + assert_eq!( + *range, + SargableQuery::Range( + Bound::Excluded(ScalarValue::Int64(Some(5))), + Bound::Unbounded, + ) + ); + } + } + + // Case 8: Empty range (lower > upper) — still valid, just produces empty results + // x >= 200 AND x <= 100 → Range(Included(200), Included(100)) + // BTree search will find 0 pages matching this, which is correct. + { + let expr = ScalarIndexExpr::And( + Box::new(make_range( + "x", + "idx_x", + Bound::Included(ScalarValue::Int64(Some(200))), + Bound::Unbounded, + )), + Box::new(make_range( + "x", + "idx_x", + Bound::Unbounded, + Bound::Included(ScalarValue::Int64(Some(100))), + )), + ); + + let optimized = expr.optimize(); + let mut leaves = Vec::new(); + optimized.collect_and_leaves(&mut leaves); + assert_eq!(leaves.len(), 1); + if let ScalarIndexExpr::Query(s) = &leaves[0] { + let range = s.query.as_any().downcast_ref::().unwrap(); + // Merged: lower=Included(200), upper=Included(100) — empty range, but valid + assert_eq!( + *range, + SargableQuery::Range( + Bound::Included(ScalarValue::Int64(Some(200))), + Bound::Included(ScalarValue::Int64(Some(100))), + ) + ); + } + } + } } diff --git a/rust/lance-index/src/scalar/fmindex.rs b/rust/lance-index/src/scalar/fmindex.rs index 2aa59b5c192..c79e4a33274 100644 --- a/rust/lance-index/src/scalar/fmindex.rs +++ b/rust/lance-index/src/scalar/fmindex.rs @@ -22,30 +22,31 @@ use std::cmp::Reverse; use std::collections::{BinaryHeap, HashMap}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, OnceLock}; use arrow_array::RecordBatch; use arrow_schema::{DataType, Field}; use async_trait::async_trait; use datafusion::execution::SendableRecordBatchStream; -use futures::StreamExt; +use futures::{StreamExt, TryStreamExt}; use lance_core::cache::LanceCache; use lance_core::deepsize::DeepSizeOf; +use lance_core::utils::row_addr_remap::RowAddrRemap; use lance_core::utils::tokio::{get_num_compute_intensive_cpus, spawn_cpu}; use lance_core::{Error, ROW_ADDR, Result}; use roaring::RoaringBitmap; -use crate::frag_reuse::FragReuseIndex; use crate::metrics::MetricsCollector; use crate::pb; use crate::scalar::expression::{ScalarQueryParser, TextQueryParser}; use crate::scalar::registry::{ - DefaultTrainingRequest, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, - VALUE_COLUMN_NAME, + BasicTrainer, DefaultTrainingRequest, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, + TrainingRequest, VALUE_COLUMN_NAME, }; use crate::scalar::{ AnyQuery, BuiltinIndexType, CreatedIndex, IndexFile, IndexStore, OldIndexDataFilter, - ScalarIndex, ScalarIndexParams, SearchResult, TextQuery, UpdateCriteria, + RowIdRemapper, ScalarIndex, ScalarIndexParams, SearchResult, TextQuery, UpdateCriteria, }; use crate::{Index, IndexType}; @@ -53,6 +54,9 @@ const FMINDEX_INDEX_VERSION: u32 = 10; const BLOCK_WORDS: usize = 4096; const PARTITION_SIZE: usize = 10_000; const DEFAULT_PARTITION_SIZE_BYTES: usize = 16 * 1024 * 1024; +const DEFAULT_DEMAND_PAGE_TARGET_BYTES: usize = 512 * 1024; +const DEFAULT_PREWARM_CHUNK_TARGET_BYTES: usize = 8 * 1024 * 1024; +const FMINDEX_PARTITION_FINGERPRINT_KEY: &str = "partition_fingerprint"; const SENTINEL_BYTE: u8 = 0xFF; /// SA sampling rate. Store every D-th SA entry. Locate walks at most D LF steps. @@ -83,11 +87,49 @@ static LANCE_FMINDEX_WRITE_QUEUE_SIZE: std::sync::LazyLock = .parse() .expect("failed to parse LANCE_FMINDEX_WRITE_QUEUE_SIZE") }); +static LANCE_FMINDEX_RESUME_EXISTING_PARTITIONS: std::sync::LazyLock = + std::sync::LazyLock::new(|| { + std::env::var("LANCE_FMINDEX_RESUME_EXISTING_PARTITIONS") + .map(|value| { + matches!( + value.as_str(), + "1" | "true" | "TRUE" | "True" | "yes" | "YES" + ) + }) + .unwrap_or(false) + }); +static LANCE_FMINDEX_PREWARM_CHUNK_BYTES: std::sync::LazyLock = + std::sync::LazyLock::new(|| { + std::env::var("LANCE_FMINDEX_PREWARM_CHUNK_BYTES") + .unwrap_or_else(|_| DEFAULT_PREWARM_CHUNK_TARGET_BYTES.to_string()) + .parse() + .expect("failed to parse LANCE_FMINDEX_PREWARM_CHUNK_BYTES") + }); +static LANCE_FMINDEX_DEMAND_PAGE_BYTES: std::sync::LazyLock = + std::sync::LazyLock::new(|| { + std::env::var("LANCE_FMINDEX_DEMAND_PAGE_BYTES") + .unwrap_or_else(|_| DEFAULT_DEMAND_PAGE_TARGET_BYTES.to_string()) + .parse() + .expect("failed to parse LANCE_FMINDEX_DEMAND_PAGE_BYTES") + }); +static LANCE_FMINDEX_PREWARM_CHUNK_CONCURRENCY: std::sync::LazyLock = + std::sync::LazyLock::new(|| { + std::env::var("LANCE_FMINDEX_PREWARM_CHUNK_CONCURRENCY") + .unwrap_or_else(|_| "1".to_string()) + .parse() + .expect("failed to parse LANCE_FMINDEX_PREWARM_CHUNK_CONCURRENCY") + }); fn fmindex_partition_path(partition_id: u64) -> String { format!("part_{partition_id}_fm.lance") } +fn fmindex_partition_id_from_path(path: &str) -> Option { + path.strip_prefix("part_") + .and_then(|r| r.strip_suffix("_fm.lance")) + .and_then(|s| s.parse::().ok()) +} + fn fmindex_num_workers() -> usize { (*LANCE_FMINDEX_NUM_WORKERS).max(1) } @@ -104,6 +146,26 @@ fn fmindex_write_queue_size() -> usize { (*LANCE_FMINDEX_WRITE_QUEUE_SIZE).max(1) } +fn fmindex_resume_existing_partitions() -> bool { + *LANCE_FMINDEX_RESUME_EXISTING_PARTITIONS +} + +fn fmindex_prewarm_chunk_bytes() -> usize { + (*LANCE_FMINDEX_PREWARM_CHUNK_BYTES).max(BLOCK_WORDS * 8) +} + +fn fmindex_demand_page_bytes() -> usize { + (*LANCE_FMINDEX_DEMAND_PAGE_BYTES).max(BLOCK_WORDS * 8) +} + +fn fmindex_demand_page_rows() -> usize { + (fmindex_demand_page_bytes() / (BLOCK_WORDS * 8)).max(1) +} + +fn fmindex_prewarm_chunk_concurrency() -> usize { + (*LANCE_FMINDEX_PREWARM_CHUNK_CONCURRENCY).max(1) +} + // ── Bitvector with O(1) rank ───────────────────────────────────────────────── const SUPERBLOCK_BITS: usize = 512; @@ -474,18 +536,82 @@ impl LazyRankBitVec { } } - /// Pre-load all blocks into memory. Call this before sync rank/access operations - /// to avoid the need for `block_in_place` during queries. - async fn load_all_blocks(&self) -> Result<()> { - for (idx, lock) in self.blocks.iter().enumerate() { - if lock.get().is_none() { - let words = self.load_block(idx).await?; - let _ = lock.set(words); + async fn load_block_if_needed(&self, idx: usize) -> Result<()> { + if idx >= self.blocks.len() { + return Err(Error::index(format!( + "FM-Index block {idx} is out of range for {} blocks", + self.blocks.len() + ))); + } + if self.blocks[idx].get().is_none() { + self.load_page_containing(idx).await?; + } + Ok(()) + } + + async fn load_page_containing(&self, idx: usize) -> Result<()> { + let page_rows = fmindex_demand_page_rows(); + let start = (idx / page_rows) * page_rows; + let end = (start + page_rows).min(self.blocks.len()); + if self.blocks[start..end] + .iter() + .all(|block| block.get().is_some()) + { + return Ok(()); + } + + let batch = self + .reader + .read_range( + self.block_row_offset + start..self.block_row_offset + end, + Some(&["words"]), + ) + .await?; + if batch.num_rows() != end - start { + return Err(Error::index(format!( + "expected {} FM-Index block rows, got {}", + end - start, + batch.num_rows() + ))); + } + + let col = Self::words_column(&batch)?; + for row in 0..batch.num_rows() { + let block_idx = start + row; + if self.blocks[block_idx].get().is_none() { + let words = Self::decode_words(col.value(row)); + let _ = self.blocks[block_idx].set(words); } } Ok(()) } + async fn load_block_for_rank(&self, pos: usize) -> Result<()> { + if pos > self.len { + return Err(Error::invalid_input(format!( + "FM-Index rank position {pos} exceeds bitvector length {}", + self.len + ))); + } + if pos == 0 { + return Ok(()); + } + if pos == self.len && pos.is_multiple_of(BLOCK_BITS) { + if let Some(last_idx) = self.blocks.len().checked_sub(1) { + return self.load_block_if_needed(last_idx).await; + } + return Ok(()); + } + if pos.is_multiple_of(BLOCK_BITS) { + return Ok(()); + } + self.load_block_if_needed(pos / BLOCK_BITS).await + } + + async fn load_block_for_access(&self, pos: usize) -> Result<()> { + self.load_block_if_needed(pos / BLOCK_BITS).await + } + #[inline] fn ensure_block(&self, idx: usize) -> &[u64] { self.blocks[idx].get_or_init(|| { @@ -502,27 +628,60 @@ impl LazyRankBitVec { .reader .read_range(row..row + 1, Some(&["words"])) .await?; - let col = batch + let col = Self::words_column(&batch)?; + Ok(Self::decode_words(col.value(0))) + } + + fn words_column(batch: &RecordBatch) -> Result<&arrow_array::LargeBinaryArray> { + batch .column(0) .as_any() .downcast_ref::() - .ok_or_else(|| Error::invalid_input("expected LargeBinary words column"))?; - Ok(col - .value(0) - .chunks_exact(8) + .ok_or_else(|| Error::invalid_input("expected LargeBinary words column")) + } + + fn decode_words(raw: &[u8]) -> Vec { + raw.chunks_exact(8) .map(|c| u64::from_le_bytes(c.try_into().unwrap())) - .collect()) + .collect() + } + + fn terminal_rank1(&self) -> usize { + if self.len == 0 { + return 0; + } + let Some(last_idx) = self.blocks.len().checked_sub(1) else { + return 0; + }; + let mut count = self.prefix_ranks.get(last_idx).copied().unwrap_or(0) as usize; + let block = self.ensure_block(last_idx); + let local = self.len - last_idx * BLOCK_BITS; + let full_words = local / 64; + let trailing_bits = local % 64; + for w in &block[..full_words] { + count += w.count_ones() as usize; + } + if trailing_bits > 0 { + count += (block[full_words] & ((1u64 << trailing_bits) - 1)).count_ones() as usize; + } + count } #[inline] fn rank1(&self, pos: usize) -> usize { + debug_assert!(pos <= self.len); if pos == 0 { return 0; } let bi = pos / BLOCK_BITS; let local = pos % BLOCK_BITS; if local == 0 { - return self.prefix_ranks[bi] as usize; + if let Some(prefix_rank) = self.prefix_ranks.get(bi) { + return *prefix_rank as usize; + } + if pos == self.len { + return self.terminal_rank1(); + } } let mut count = self.prefix_ranks[bi] as usize; let block = self.ensure_block(bi); @@ -544,6 +703,7 @@ impl LazyRankBitVec { #[inline] fn get(&self, pos: usize) -> bool { + debug_assert!(pos < self.len); let bi = pos / BLOCK_BITS; let local = pos % BLOCK_BITS; let block = self.ensure_block(bi); @@ -579,9 +739,105 @@ impl std::fmt::Debug for LazyHuffmanWaveletTree { impl LazyHuffmanWaveletTree { /// Pre-load all wavelet tree blocks into memory. async fn load_all(&self) -> Result<()> { - for node in &self.nodes { - node.load_all_blocks().await?; + if self.nodes.is_empty() { + return Ok(()); + } + + #[derive(Clone, Copy)] + struct RowRange { + start: usize, + end: usize, + node_idx: usize, } + + let mut ranges = self + .nodes + .iter() + .enumerate() + .filter(|(_, node)| !node.blocks.is_empty()) + .map(|(node_idx, node)| RowRange { + start: node.block_row_offset, + end: node.block_row_offset + node.blocks.len(), + node_idx, + }) + .collect::>(); + ranges.sort_by_key(|range| range.start); + let Some(total_rows) = ranges.iter().map(|range| range.end).max() else { + return Ok(()); + }; + let ranges = Arc::new(ranges); + + let chunk_loaded = |start: usize, end: usize| { + ranges.iter().all(|range| { + let overlap_start = start.max(range.start); + let overlap_end = end.min(range.end); + if overlap_start >= overlap_end { + return true; + } + let node = &self.nodes[range.node_idx]; + node.blocks[overlap_start - range.start..overlap_end - range.start] + .iter() + .all(|block| block.get().is_some()) + }) + }; + + let chunk_rows = (fmindex_prewarm_chunk_bytes() / (BLOCK_WORDS * 8)).max(1); + let reader = Arc::clone(&self.nodes[0].reader); + let chunks = (0..total_rows) + .step_by(chunk_rows) + .map(|start| { + let end = (start + chunk_rows).min(total_rows); + (start, end) + }) + .filter(|(start, end)| !chunk_loaded(*start, *end)) + .collect::>(); + + futures::stream::iter(chunks) + .map(|(start, end)| { + let reader = Arc::clone(&reader); + async move { + let batch = reader.read_range(start..end, Some(&["words"])).await?; + Result::<_>::Ok((start, end, batch)) + } + }) + .buffer_unordered(fmindex_prewarm_chunk_concurrency()) + .try_for_each(|(start, end, batch)| { + let ranges = Arc::clone(&ranges); + let nodes = &self.nodes; + async move { + if batch.num_rows() != end - start { + return Err(Error::index(format!( + "expected {} FM-Index block rows, got {}", + end - start, + batch.num_rows() + ))); + } + + let col = LazyRankBitVec::words_column(&batch)?; + let mut range_idx = ranges.partition_point(|range| range.end <= start); + for row in 0..batch.num_rows() { + let absolute_row = start + row; + while range_idx < ranges.len() && absolute_row >= ranges[range_idx].end { + range_idx += 1; + } + if range_idx == ranges.len() || absolute_row < ranges[range_idx].start { + continue; + } + + let range = ranges[range_idx]; + let block_idx = absolute_row - range.start; + let node = &nodes[range.node_idx]; + if node.blocks[block_idx].get().is_none() { + let words = LazyRankBitVec::decode_words(col.value(row)); + let _ = node.blocks[block_idx].set(words); + } + } + + Ok(()) + } + }) + .await?; + Ok(()) } @@ -610,6 +866,31 @@ impl LazyHuffmanWaveletTree { } } + async fn access_async(&self, mut pos: usize) -> Result { + if self.nodes.is_empty() { + return Ok(0); + } + let mut node_idx = 0; + loop { + self.nodes[node_idx].load_block_for_access(pos).await?; + let bit = self.nodes[node_idx].get(pos); + let (ref left, ref right) = self.children[node_idx]; + if bit { + pos = self.nodes[node_idx].rank1(pos); + match right { + WaveletChild::Leaf(b) => return Ok(*b), + WaveletChild::Node(next) => node_idx = *next, + } + } else { + pos = self.nodes[node_idx].rank0(pos); + match left { + WaveletChild::Leaf(b) => return Ok(*b), + WaveletChild::Node(next) => node_idx = *next, + } + } + } + } + #[inline] fn rank(&self, c: u8, pos: usize) -> usize { let code = &self.codes[c as usize]; @@ -629,6 +910,26 @@ impl LazyHuffmanWaveletTree { hi - lo } + async fn rank_async(&self, c: u8, pos: usize) -> Result { + let code = &self.codes[c as usize]; + if code.length == 0 { + return Ok(0); + } + let (mut lo, mut hi) = (0, pos); + for (level, &nid) in code.node_path.iter().enumerate() { + self.nodes[nid].load_block_for_rank(lo).await?; + self.nodes[nid].load_block_for_rank(hi).await?; + if (code.bits >> (code.length - 1 - level as u8)) & 1 == 0 { + lo = self.nodes[nid].rank0(lo); + hi = self.nodes[nid].rank0(hi); + } else { + lo = self.nodes[nid].rank1(lo); + hi = self.nodes[nid].rank1(hi); + } + } + Ok(hi - lo) + } + #[inline] fn rank_pair(&self, c: u8, lo: usize, hi: usize) -> (usize, usize) { let code = &self.codes[c as usize]; @@ -650,6 +951,29 @@ impl LazyHuffmanWaveletTree { (l - s, h - s) } + async fn rank_pair_async(&self, c: u8, lo: usize, hi: usize) -> Result<(usize, usize)> { + let code = &self.codes[c as usize]; + if code.length == 0 { + return Ok((0, 0)); + } + let (mut s, mut l, mut h) = (0, lo, hi); + for (level, &nid) in code.node_path.iter().enumerate() { + self.nodes[nid].load_block_for_rank(s).await?; + self.nodes[nid].load_block_for_rank(l).await?; + self.nodes[nid].load_block_for_rank(h).await?; + if (code.bits >> (code.length - 1 - level as u8)) & 1 == 0 { + s = self.nodes[nid].rank0(s); + l = self.nodes[nid].rank0(l); + h = self.nodes[nid].rank0(h); + } else { + s = self.nodes[nid].rank1(s); + l = self.nodes[nid].rank1(l); + h = self.nodes[nid].rank1(h); + } + } + Ok((l - s, h - s)) + } + fn deep_size(&self) -> usize { self.nodes.iter().map(|n| n.deep_size()).sum::() + self @@ -998,12 +1322,18 @@ struct LazyFMIndex { sa_samples: Vec, doc_start_positions: Vec, c_table: Vec, + fully_prewarmed: AtomicBool, } impl LazyFMIndex { /// Pre-load all wavelet tree blocks before sync search operations. async fn prewarm(&self) -> Result<()> { - self.wavelet.load_all().await + if self.fully_prewarmed.load(Ordering::Acquire) { + return Ok(()); + } + self.wavelet.load_all().await.inspect(|_| { + self.fully_prewarmed.store(true, Ordering::Release); + }) } fn backward_search(&self, pattern: &[u8]) -> (usize, usize) { @@ -1023,6 +1353,23 @@ impl LazyFMIndex { (lo, hi) } + async fn backward_search_async(&self, pattern: &[u8]) -> Result<(usize, usize)> { + if pattern.is_empty() || self.wavelet.len == 0 { + return Ok((0, 0)); + } + let (mut lo, mut hi) = (0, self.wavelet.len); + for &b in pattern.iter().rev() { + let c = self.c_table[b as usize]; + let (occ_lo, occ_hi) = self.wavelet.rank_pair_async(b, lo, hi).await?; + lo = c + occ_lo; + hi = c + occ_hi; + if lo >= hi { + return Ok((0, 0)); + } + } + Ok((lo, hi)) + } + #[inline] fn locate(&self, mut pos: usize) -> usize { let mut steps = 0; @@ -1042,6 +1389,24 @@ impl LazyFMIndex { } } + async fn locate_async(&self, mut pos: usize) -> Result { + let mut steps = 0; + let n = self.wavelet.len; + loop { + if pos.is_multiple_of(SA_SAMPLE_RATE) && (pos / SA_SAMPLE_RATE) < self.sa_samples.len() + { + return Ok((self.sa_samples[pos / SA_SAMPLE_RATE] as usize + steps) % n); + } + let c = self.wavelet.access_async(pos).await?; + pos = self.c_table[c as usize] + self.wavelet.rank_async(c, pos).await?; + steps += 1; + if steps >= n { + log::warn!("FM-Index SA locate exceeded {n} steps, possible index corruption"); + return Ok(0); + } + } + } + #[inline] fn doc_for_position(&self, text_pos: usize) -> usize { let tp = text_pos as u64; @@ -1085,6 +1450,24 @@ impl LazyFMIndex { result } + async fn search_row_addrs_async(&self, pattern: &[u8]) -> Result> { + let (lo, hi) = self.backward_search_async(pattern).await?; + if lo >= hi { + return Ok(Vec::new()); + } + let mut seen = std::collections::HashSet::new(); + let mut result = Vec::new(); + for i in lo..hi { + let text_pos = self.locate_async(i).await?; + let doc_idx = self.doc_for_position(text_pos); + let row_addr = self.row_ids[doc_idx]; + if seen.insert(row_addr) { + result.push(row_addr); + } + } + Ok(result) + } + #[allow(clippy::too_many_arguments)] async fn from_reader( reader: Arc, @@ -1196,6 +1579,7 @@ impl LazyFMIndex { sa_samples, doc_start_positions, c_table, + fully_prewarmed: AtomicBool::new(false), }) } @@ -1220,6 +1604,7 @@ struct FMIndexPartition { #[derive(Debug)] pub struct FMIndexScalarIndex { partitions: Vec>, + io_parallelism: usize, } impl DeepSizeOf for FMIndexScalarIndex { @@ -1301,18 +1686,13 @@ impl FMIndexScalarIndex { async fn load( store: Arc, - _fri: Option>, + _fri: Option>, _cache: &LanceCache, ) -> Result> { let files = store.list_files_with_sizes().await?; let mut pfiles: Vec<(u64, String)> = Vec::new(); for f in &files { - if let Some(id) = f - .path - .strip_prefix("part_") - .and_then(|r| r.strip_suffix("_fm.lance")) - .and_then(|s| s.parse::().ok()) - { + if let Some(id) = fmindex_partition_id_from_path(&f.path) { pfiles.push((id, f.path.clone())); } } @@ -1320,13 +1700,73 @@ impl FMIndexScalarIndex { return Err(Error::invalid_input("no FM-Index partition files found")); } pfiles.sort_by_key(|(id, _)| *id); - let mut parts = Vec::with_capacity(pfiles.len()); - for (id, name) in &pfiles { - parts.push(Arc::new( - Self::load_partition(store.as_ref(), name, *id).await?, - )); - } - Ok(Arc::new(Self { partitions: parts })) + let io_parallelism = store.io_parallelism().max(1); + let mut parts = futures::stream::iter(pfiles) + .map(|(id, name)| { + let store = Arc::clone(&store); + async move { + let partition = Self::load_partition(store.as_ref(), &name, id).await?; + Result::<_>::Ok((id, Arc::new(partition))) + } + }) + .buffer_unordered(io_parallelism) + .try_collect::>() + .await?; + parts.sort_by_key(|(id, _)| *id); + let parts = parts + .into_iter() + .map(|(_, partition)| partition) + .collect::>(); + Ok(Arc::new(Self { + partitions: parts, + io_parallelism, + })) + } + + fn partition_parallelism(&self) -> usize { + self.io_parallelism.max(1).min(self.partitions.len().max(1)) + } + + async fn prewarm_partitions(&self) -> Result<()> { + futures::stream::iter(self.partitions.iter().cloned()) + .map(|partition| async move { partition.fm.prewarm().await }) + .buffer_unordered(self.partition_parallelism()) + .try_collect::>() + .await?; + Ok(()) + } + + async fn search_string_contains(&self, pattern: &[u8]) -> Result { + use lance_select::RowAddrTreeMap; + + let pattern: Arc<[u8]> = Arc::from(pattern); + let tree = futures::stream::iter(self.partitions.iter().cloned()) + .map(|partition| { + let pattern = Arc::clone(&pattern); + async move { + if partition.fm.fully_prewarmed.load(Ordering::Acquire) { + spawn_cpu(move || { + Result::>::Ok(partition.fm.search_row_addrs(pattern.as_ref())) + }) + .await + } else { + partition.fm.search_row_addrs_async(pattern.as_ref()).await + } + } + }) + .buffer_unordered(self.partition_parallelism()) + .try_fold(RowAddrTreeMap::new(), |mut tree, row_addrs| async move { + for row_addr in row_addrs { + tree.insert(row_addr); + } + Result::Ok(tree) + }) + .await?; + + Ok(SearchResult::Exact(lance_select::NullableRowAddrSet::new( + tree, + Default::default(), + ))) } } @@ -1339,7 +1779,7 @@ impl Index for FMIndexScalarIndex { self } async fn prewarm(&self) -> Result<()> { - Ok(()) + self.prewarm_partitions().await } fn statistics(&self) -> Result { Ok(serde_json::json!({ @@ -1376,19 +1816,7 @@ impl ScalarIndex for FMIndexScalarIndex { .ok_or_else(|| Error::invalid_input("Fm only supports TextQuery"))?; match tq { TextQuery::StringContains(pattern) => { - let pb = pattern.as_bytes(); - use lance_select::RowAddrTreeMap; - let mut tree = RowAddrTreeMap::new(); - for p in &self.partitions { - p.fm.prewarm().await?; - for row_addr in p.fm.search_row_addrs(pb) { - tree.insert(row_addr); - } - } - Ok(SearchResult::Exact(lance_select::NullableRowAddrSet::new( - tree, - Default::default(), - ))) + self.search_string_contains(pattern.as_bytes()).await } // Regex queries are routed only to the ngram index (the FM-index's // query parser advertises `supports_regex = false`), so this is @@ -1401,11 +1829,7 @@ impl ScalarIndex for FMIndexScalarIndex { fn can_remap(&self) -> bool { false } - async fn remap( - &self, - _: &HashMap>, - _: &dyn IndexStore, - ) -> Result { + async fn remap(&self, _: &RowAddrRemap, _: &dyn IndexStore) -> Result { Err(Error::not_supported("Fm does not support remap")) } async fn update( @@ -1445,6 +1869,7 @@ struct FMIndexPartitionConfig { max_rows: usize, max_bytes: usize, queue_size: usize, + resume_existing: bool, } impl FMIndexPartitionConfig { @@ -1454,6 +1879,7 @@ impl FMIndexPartitionConfig { max_rows: fmindex_partition_rows(), max_bytes: fmindex_partition_bytes(), queue_size: fmindex_write_queue_size(), + resume_existing: fmindex_resume_existing_partitions(), } } @@ -1463,6 +1889,7 @@ impl FMIndexPartitionConfig { max_rows: self.max_rows.max(1), max_bytes: self.max_bytes.max(1), queue_size: self.queue_size.max(1), + resume_existing: self.resume_existing, } } } @@ -1493,6 +1920,23 @@ async fn write_partitioned_fmindex_stream_with_config( async_channel::Receiver, ) = async_channel::bounded(config.queue_size); let store = store.clone_arc(); + let mut completed_files = if config.resume_existing { + store + .list_files_with_sizes() + .await? + .into_iter() + .filter_map(|file| fmindex_partition_id_from_path(&file.path).map(|id| (id, file))) + .collect::>() + } else { + HashMap::new() + }; + if !completed_files.is_empty() { + log::info!( + "resuming FMIndex build with {} existing partition files", + completed_files.len() + ); + } + let mut files = Vec::new(); let mut worker_tasks = Vec::with_capacity(config.num_workers); for _ in 0..config.num_workers { let receiver = receiver.clone(); @@ -1541,12 +1985,17 @@ async fn write_partitioned_fmindex_stream_with_config( config.max_rows, config.max_bytes, ) { - send_fmindex_partition( + finish_fmindex_partition( &sender, + store.as_ref(), &mut partition, &mut partition_bytes, partition_id, config.max_rows, + &mut FMIndexPartitionResumeState { + completed_files: &mut completed_files, + output_files: &mut files, + }, ) .await?; partition_id += 1; @@ -1556,12 +2005,17 @@ async fn write_partitioned_fmindex_stream_with_config( } if !partition.is_empty() { - send_fmindex_partition( + finish_fmindex_partition( &sender, + store.as_ref(), &mut partition, &mut partition_bytes, partition_id, config.max_rows, + &mut FMIndexPartitionResumeState { + completed_files: &mut completed_files, + output_files: &mut files, + }, ) .await?; } @@ -1571,7 +2025,6 @@ async fn write_partitioned_fmindex_stream_with_config( .await; drop(sender); - let mut files = Vec::new(); let mut worker_error = None; for worker_task in worker_tasks { match worker_task.await { @@ -1594,6 +2047,14 @@ async fn write_partitioned_fmindex_stream_with_config( return Err(err); } producer_result?; + + for (partition_id, file) in completed_files.drain() { + log::info!( + "deleting stale FMIndex partition {partition_id} from previous build: {}", + file.path + ); + store.delete_index_file(&file.path).await?; + } if files.is_empty() { return Ok(vec![write_empty_fmindex_partition(store.as_ref()).await?]); } @@ -1615,15 +2076,34 @@ fn fmindex_partition_limit_reached( rows >= max_rows || bytes >= max_bytes } -async fn send_fmindex_partition( +struct FMIndexPartitionResumeState<'a> { + completed_files: &'a mut HashMap, + output_files: &'a mut Vec<(u64, IndexFile)>, +} + +async fn finish_fmindex_partition( sender: &async_channel::Sender, + store: &dyn IndexStore, partition: &mut Vec<(u64, Vec)>, partition_bytes: &mut usize, partition_id: u64, max_rows: usize, + resume_state: &mut FMIndexPartitionResumeState<'_>, ) -> Result<()> { let texts = std::mem::replace(partition, Vec::with_capacity(max_rows.min(PARTITION_SIZE))); *partition_bytes = 0; + if let Some(file) = resume_state.completed_files.remove(&partition_id) { + let fingerprint = fmindex_partition_fingerprint(&texts); + if fmindex_partition_matches(store, &file, &fingerprint).await? { + resume_state.output_files.push((partition_id, file)); + return Ok(()); + } + log::info!( + "rebuilding stale FMIndex partition {partition_id}: {}", + file.path + ); + store.delete_index_file(&file.path).await?; + } sender .send(FMIndexPartitionJob { partition_id, @@ -1637,6 +2117,41 @@ async fn send_fmindex_partition( }) } +fn fmindex_partition_fingerprint(texts: &[(u64, Vec)]) -> String { + const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325; + const FNV_PRIME: u64 = 0x0000_0100_0000_01b3; + + fn update(hash: &mut u64, bytes: &[u8]) { + for byte in bytes { + *hash ^= *byte as u64; + *hash = hash.wrapping_mul(FNV_PRIME); + } + } + + let mut hash = FNV_OFFSET; + update(&mut hash, b"lance-fmindex-partition-v1"); + update(&mut hash, &(texts.len() as u64).to_le_bytes()); + for (row_id, bytes) in texts { + update(&mut hash, &row_id.to_le_bytes()); + update(&mut hash, &(bytes.len() as u64).to_le_bytes()); + update(&mut hash, bytes); + } + format!("{hash:016x}") +} + +async fn fmindex_partition_matches( + store: &dyn IndexStore, + file: &IndexFile, + expected_fingerprint: &str, +) -> Result { + let reader = store.open_index_file(&file.path).await?; + Ok(reader + .schema() + .metadata + .get(FMINDEX_PARTITION_FINGERPRINT_KEY) + .is_some_and(|fingerprint| fingerprint == expected_fingerprint)) +} + fn sanitize_text_bytes(bytes: &[u8]) -> Vec { bytes .iter() @@ -1764,7 +2279,12 @@ fn hex_decode(s: &str) -> Result> { /// - Wavelet block rows (BWT nodes) /// - SA sample blocks (packed u64 in LargeBinary) /// - Metadata: c_table, huffman_codes, tree_topology, row_ids, doc_start_positions -async fn write_fmindex(fm: &FMIndex, store: &dyn IndexStore, filename: &str) -> Result { +async fn write_fmindex( + fm: &FMIndex, + store: &dyn IndexStore, + filename: &str, + partition_fingerprint: Option<&str>, +) -> Result { let schema = Arc::new(FMIndex::block_schema()); let mut writer = store.new_index_file(filename, schema.clone()).await?; @@ -1814,6 +2334,12 @@ async fn write_fmindex(fm: &FMIndex, store: &dyn IndexStore, filename: &str) -> metadata.insert("total_wavelet_rows".into(), nw.to_string()); metadata.insert("sa_sample_rate".into(), SA_SAMPLE_RATE.to_string()); metadata.insert("alphabet_size".into(), fm.alphabet_size.to_string()); + if let Some(fingerprint) = partition_fingerprint { + metadata.insert( + FMINDEX_PARTITION_FINGERPRINT_KEY.into(), + fingerprint.to_string(), + ); + } metadata.insert("c_table".into(), hex_encode(&fm.serialize_c_table())); metadata.insert( "huffman_codes".into(), @@ -1858,9 +2384,16 @@ async fn write_fmindex_partition( store: &dyn IndexStore, partition_id: u64, ) -> Result { + let fingerprint = fmindex_partition_fingerprint(texts); let refs: Vec<(u64, &[u8])> = texts.iter().map(|(id, t)| (*id, t.as_slice())).collect(); let fm = FMIndex::build(&refs)?; - write_fmindex(&fm, store, &fmindex_partition_path(partition_id)).await + write_fmindex( + &fm, + store, + &fmindex_partition_path(partition_id), + Some(&fingerprint), + ) + .await } async fn write_fmindex_partition_owned( @@ -1868,18 +2401,26 @@ async fn write_fmindex_partition_owned( store: Arc, partition_id: u64, ) -> Result<(u64, IndexFile)> { + let fingerprint = fmindex_partition_fingerprint(&texts); let fm = spawn_cpu(move || { let refs: Vec<(u64, &[u8])> = texts.iter().map(|(id, t)| (*id, t.as_slice())).collect(); FMIndex::build(&refs) }) .await?; - let file = write_fmindex(&fm, store.as_ref(), &fmindex_partition_path(partition_id)).await?; + let file = write_fmindex( + &fm, + store.as_ref(), + &fmindex_partition_path(partition_id), + Some(&fingerprint), + ) + .await?; Ok((partition_id, file)) } async fn write_empty_fmindex_partition(store: &dyn IndexStore) -> Result { let fm = FMIndex::build(&[])?; - write_fmindex(&fm, store, &fmindex_partition_path(0)).await + let fingerprint = fmindex_partition_fingerprint(&[]); + write_fmindex(&fm, store, &fmindex_partition_path(0), Some(&fingerprint)).await } // ── Plugin ─────────────────────────────────────────────────────────────────── @@ -1888,10 +2429,7 @@ async fn write_empty_fmindex_partition(store: &dyn IndexStore) -> Result &str { - "Fm" - } +impl BasicTrainer for FMIndexPlugin { fn new_training_request( &self, _params: &str, @@ -1925,11 +2463,22 @@ impl ScalarIndexPlugin for FMIndexPlugin { files, }) } - fn provides_exact_answer(&self) -> bool { - true +} + +#[async_trait] +impl ScalarIndexPlugin for FMIndexPlugin { + fn basic_trainer(&self) -> Option<&dyn BasicTrainer> { + Some(self) } - fn version(&self) -> u32 { - FMINDEX_INDEX_VERSION + + fn name(&self) -> &str { + "Fm" + } + fn provides_exact_answer(&self) -> bool { + true + } + fn version(&self) -> u32 { + FMINDEX_INDEX_VERSION } fn new_query_parser( &self, @@ -1949,7 +2498,7 @@ impl ScalarIndexPlugin for FMIndexPlugin { &self, store: Arc, details: &prost_types::Any, - fri: Option>, + fri: Option>, cache: &LanceCache, ) -> Result> { let _ = details.to_msg::().unwrap_or_default(); @@ -1976,6 +2525,7 @@ mod tests { use std::sync::Arc; use crate::scalar::lance_format::LanceIndexStore; + use crate::scalar::registry::BasicTrainer; #[derive(Debug, Clone)] struct FailNewFileStore { @@ -2002,6 +2552,12 @@ mod tests { self.inner.io_parallelism() } + fn with_io_priority(&self, io_priority: u64) -> Arc { + Arc::new(Self { + inner: self.inner.with_io_priority(io_priority), + }) + } + async fn new_index_file( &self, _name: &str, @@ -2037,6 +2593,29 @@ mod tests { } } + fn loaded_wavelet_blocks(index: &FMIndexScalarIndex) -> usize { + index + .partitions + .iter() + .flat_map(|partition| partition.fm.wavelet.nodes.iter()) + .map(|node| { + node.blocks + .iter() + .filter(|block| block.get().is_some()) + .count() + }) + .sum() + } + + fn total_wavelet_blocks(index: &FMIndexScalarIndex) -> usize { + index + .partitions + .iter() + .flat_map(|partition| partition.fm.wavelet.nodes.iter()) + .map(|node| node.blocks.len()) + .sum() + } + #[test] fn test_fmindex_build_and_search() { let texts: Vec<(u64, &[u8])> = vec![ @@ -2328,7 +2907,7 @@ mod tests { )); // Write - write_fmindex(&fm, store.as_ref(), &fmindex_partition_path(0)) + write_fmindex(&fm, store.as_ref(), &fmindex_partition_path(0), None) .await .unwrap(); @@ -2423,6 +3002,144 @@ mod tests { } } + #[tokio::test(flavor = "multi_thread")] + async fn test_public_prewarm_loads_lazy_blocks() { + let docs: Vec> = (0..30) + .map(|i| format!("document {i} hello world test data").into_bytes()) + .collect(); + let texts: Vec<(u64, Vec)> = docs + .into_iter() + .enumerate() + .map(|(i, d)| (i as u64, d)) + .collect(); + + let tempdir = tempfile::tempdir().unwrap(); + let index_dir = Path::from_filesystem_path(tempdir.path()).unwrap(); + let store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + index_dir, + Arc::new(LanceCache::no_cache()), + )); + + write_partitioned_fmindex(&texts, store.as_ref()) + .await + .unwrap(); + + let index = FMIndexScalarIndex::load(store, None, &LanceCache::no_cache()) + .await + .unwrap(); + let total_blocks = total_wavelet_blocks(index.as_ref()); + assert!(total_blocks > 0); + assert_eq!(loaded_wavelet_blocks(index.as_ref()), 0); + + index.prewarm().await.unwrap(); + assert_eq!(loaded_wavelet_blocks(index.as_ref()), total_blocks); + + let r = index + .search( + &TextQuery::StringContains("hello world".to_string()), + &crate::metrics::NoOpMetricsCollector, + ) + .await + .unwrap(); + match r { + SearchResult::Exact(set) => { + assert_eq!(set.len(), Some(30)); + } + _ => panic!("expected exact result"), + } + assert_eq!(loaded_wavelet_blocks(index.as_ref()), total_blocks); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_lazy_rank_terminal_block_boundary() { + let tempdir = tempfile::tempdir().unwrap(); + let index_dir = Path::from_filesystem_path(tempdir.path()).unwrap(); + let store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + index_dir, + Arc::new(LanceCache::no_cache()), + )); + let schema = Arc::new(arrow_schema::Schema::new(vec![Field::new( + "words", + DataType::LargeBinary, + false, + )])); + let words = vec![u64::MAX; BLOCK_WORDS]; + let bytes = FMIndex::u64_to_bytes(&words); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(LargeBinaryArray::from(vec![bytes.as_slice()]))], + ) + .unwrap(); + let mut writer = store.new_index_file("rank.lance", schema).await.unwrap(); + writer.write_record_batch(batch).await.unwrap(); + writer.finish().await.unwrap(); + + let reader = store.open_index_file("rank.lance").await.unwrap(); + let bitvec = LazyRankBitVec::new(vec![0], 1, reader, 0, BLOCK_BITS); + bitvec.load_block_for_rank(BLOCK_BITS).await.unwrap(); + + assert_eq!(bitvec.rank1(BLOCK_BITS), BLOCK_BITS); + assert_eq!(bitvec.rank0(BLOCK_BITS), 0); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_contains_search_does_not_full_prewarm_partitions() { + let docs: Vec> = (0..30) + .map(|i| format!("document {i} hello world test data").into_bytes()) + .collect(); + let texts: Vec<(u64, Vec)> = docs + .into_iter() + .enumerate() + .map(|(i, d)| (i as u64, d)) + .collect(); + + let tempdir = tempfile::tempdir().unwrap(); + let index_dir = Path::from_filesystem_path(tempdir.path()).unwrap(); + let store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + index_dir, + Arc::new(LanceCache::no_cache()), + )); + + write_partitioned_fmindex(&texts, store.as_ref()) + .await + .unwrap(); + + let index = FMIndexScalarIndex::load(store, None, &LanceCache::no_cache()) + .await + .unwrap(); + assert_eq!(loaded_wavelet_blocks(index.as_ref()), 0); + assert!( + index + .partitions + .iter() + .all(|partition| !partition.fm.fully_prewarmed.load(Ordering::Acquire)) + ); + + let r = index + .search( + &TextQuery::StringContains("document 15".to_string()), + &crate::metrics::NoOpMetricsCollector, + ) + .await + .unwrap(); + match r { + SearchResult::Exact(set) => { + assert_eq!(set.len(), Some(1)); + } + _ => panic!("expected exact result"), + } + + assert!( + index + .partitions + .iter() + .all(|partition| !partition.fm.fully_prewarmed.load(Ordering::Acquire)) + ); + } + #[tokio::test(flavor = "multi_thread")] async fn test_plugin_train_and_load() { let docs = vec!["hello world", "hello rust", "goodbye world"]; @@ -2525,6 +3242,7 @@ mod tests { max_rows: 100, max_bytes: 5, queue_size: 1, + resume_existing: false, }, ) .await @@ -2553,6 +3271,209 @@ mod tests { } } + #[tokio::test(flavor = "multi_thread")] + async fn test_stream_train_resumes_existing_partitions() { + let schema = Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new( + crate::scalar::registry::VALUE_COLUMN_NAME, + DataType::Utf8, + false, + ), + arrow_schema::Field::new(ROW_ADDR, DataType::UInt64, false), + ])); + let tempdir = tempfile::tempdir().unwrap(); + let index_dir = Path::from_filesystem_path(tempdir.path()).unwrap(); + let store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + index_dir, + Arc::new(LanceCache::no_cache()), + )); + + let first_batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["first partition"])), + Arc::new(UInt64Array::from(vec![0])), + ], + ) + .unwrap(); + let first_stream = + RecordBatchStreamAdapter::new(schema.clone(), stream::iter(vec![Ok(first_batch)])); + let first_files = write_partitioned_fmindex_stream_with_config( + Box::pin(first_stream), + store.as_ref(), + FMIndexPartitionConfig { + num_workers: 1, + max_rows: 1, + max_bytes: 1024, + queue_size: 1, + resume_existing: false, + }, + ) + .await + .unwrap(); + assert_eq!(first_files.len(), 1); + assert_eq!(first_files[0].path, fmindex_partition_path(0)); + + let resumed_batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec![ + "first partition", + "second partition", + ])), + Arc::new(UInt64Array::from(vec![0, 1])), + ], + ) + .unwrap(); + let resumed_stream = + RecordBatchStreamAdapter::new(schema.clone(), stream::iter(vec![Ok(resumed_batch)])); + let resumed_files = write_partitioned_fmindex_stream_with_config( + Box::pin(resumed_stream), + store.as_ref(), + FMIndexPartitionConfig { + num_workers: 1, + max_rows: 1, + max_bytes: 1024, + queue_size: 1, + resume_existing: true, + }, + ) + .await + .unwrap(); + + let resumed_paths = resumed_files + .iter() + .map(|file| file.path.clone()) + .collect::>(); + assert_eq!( + resumed_paths, + vec![fmindex_partition_path(0), fmindex_partition_path(1)] + ); + + let index = FMIndexScalarIndex::load(store, None, &LanceCache::no_cache()) + .await + .unwrap(); + let r = index + .search( + &TextQuery::StringContains("partition".to_string()), + &crate::metrics::NoOpMetricsCollector, + ) + .await + .unwrap(); + match r { + SearchResult::Exact(set) => { + assert_eq!(set.len(), Some(2)); + } + _ => panic!("expected exact result"), + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_stream_train_rebuilds_stale_resume_partitions() { + let schema = Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new( + crate::scalar::registry::VALUE_COLUMN_NAME, + DataType::Utf8, + false, + ), + arrow_schema::Field::new(ROW_ADDR, DataType::UInt64, false), + ])); + let tempdir = tempfile::tempdir().unwrap(); + let index_dir = Path::from_filesystem_path(tempdir.path()).unwrap(); + let store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + index_dir, + Arc::new(LanceCache::no_cache()), + )); + + let stale_batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["stale zero", "stale one"])), + Arc::new(UInt64Array::from(vec![0, 1])), + ], + ) + .unwrap(); + let stale_stream = + RecordBatchStreamAdapter::new(schema.clone(), stream::iter(vec![Ok(stale_batch)])); + write_partitioned_fmindex_stream_with_config( + Box::pin(stale_stream), + store.as_ref(), + FMIndexPartitionConfig { + num_workers: 1, + max_rows: 1, + max_bytes: 1024, + queue_size: 1, + resume_existing: false, + }, + ) + .await + .unwrap(); + + let fresh_batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["fresh zero"])), + Arc::new(UInt64Array::from(vec![0])), + ], + ) + .unwrap(); + let fresh_stream = + RecordBatchStreamAdapter::new(schema.clone(), stream::iter(vec![Ok(fresh_batch)])); + let fresh_files = write_partitioned_fmindex_stream_with_config( + Box::pin(fresh_stream), + store.as_ref(), + FMIndexPartitionConfig { + num_workers: 1, + max_rows: 10, + max_bytes: 1024, + queue_size: 1, + resume_existing: true, + }, + ) + .await + .unwrap(); + + assert_eq!(fresh_files.len(), 1); + assert_eq!(fresh_files[0].path, fmindex_partition_path(0)); + let remaining_paths = store + .list_files_with_sizes() + .await + .unwrap() + .iter() + .map(|file| file.path.clone()) + .collect::>(); + assert_eq!(remaining_paths, vec![fmindex_partition_path(0)]); + + let index = FMIndexScalarIndex::load(store, None, &LanceCache::no_cache()) + .await + .unwrap(); + let r = index + .search( + &TextQuery::StringContains("fresh zero".to_string()), + &crate::metrics::NoOpMetricsCollector, + ) + .await + .unwrap(); + match r { + SearchResult::Exact(set) => assert_eq!(set.len(), Some(1)), + _ => panic!("expected exact result"), + } + + let r = index + .search( + &TextQuery::StringContains("stale one".to_string()), + &crate::metrics::NoOpMetricsCollector, + ) + .await + .unwrap(); + match r { + SearchResult::Exact(set) => assert_eq!(set.len(), Some(0)), + _ => panic!("expected exact result"), + } + } + #[tokio::test(flavor = "multi_thread")] async fn test_stream_train_propagates_worker_write_error() { let schema = Arc::new(arrow_schema::Schema::new(vec![ @@ -2590,6 +3511,7 @@ mod tests { max_rows: 1, max_bytes: 1024, queue_size: 1, + resume_existing: false, }, ) .await @@ -2598,6 +3520,34 @@ mod tests { assert!(format!("{err}").contains("injected FMIndex write failure")); } + #[tokio::test] + async fn test_fail_new_file_store_with_io_priority_preserves_failure() { + // `with_io_priority` re-wraps in `FailNewFileStore`, so the reprioritized + // store must keep injecting the `new_index_file` failure. + let tempdir = tempfile::tempdir().unwrap(); + let index_dir = Path::from_filesystem_path(tempdir.path()).unwrap(); + let inner = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + index_dir, + Arc::new(LanceCache::no_cache()), + )) as Arc; + let store = FailNewFileStore { inner }; + + let reprioritized = store.with_io_priority(7); + + let schema = Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new( + "x", + DataType::UInt64, + false, + )])); + let err = reprioritized + .new_index_file("test", schema) + .await + .err() + .expect("new_index_file should fail"); + assert!(format!("{err}").contains("injected FMIndex write failure")); + } + #[tokio::test(flavor = "multi_thread")] async fn test_plugin_train_streams_multiple_partitions() { fn training_batch( diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index d0bb0e40d3a..926bfad6a69 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -4,6 +4,7 @@ pub mod builder; mod cache_codec; mod encoding; +mod impact; mod index; mod iter; pub mod json; @@ -91,7 +92,7 @@ pub async fn build_global_bm25_scorer( let (mut total_tokens, mut num_docs, first_token_docs) = first_index.bm25_stats_for_terms(&terms).await?; let mut token_docs = HashMap::with_capacity(terms.len()); - for (term, count) in terms.iter().cloned().zip(first_token_docs.into_iter()) { + for (term, count) in terms.iter().cloned().zip(first_token_docs) { token_docs.insert(term, count); } @@ -100,7 +101,7 @@ pub async fn build_global_bm25_scorer( index.bm25_stats_for_terms(&terms).await?; total_tokens += segment_total_tokens; num_docs += segment_num_docs; - for (term, count) in terms.iter().zip(segment_token_docs.into_iter()) { + for (term, count) in terms.iter().zip(segment_token_docs) { *token_docs .get_mut(term) .expect("global scorer terms should already be initialized") += count; @@ -114,12 +115,11 @@ use lance_core::Error; use crate::pbold; use crate::progress::IndexBuildProgress; -use crate::{ - frag_reuse::FragReuseIndex, - scalar::{ - CreatedIndex, ScalarIndex, - expression::{FtsQueryParser, ScalarQueryParser}, - registry::{ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest}, +use crate::scalar::{ + CreatedIndex, RowIdRemapper, ScalarIndex, + expression::{FtsQueryParser, ScalarQueryParser}, + registry::{ + BasicTrainer, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, }, }; @@ -147,6 +147,8 @@ impl InvertedIndexPlugin { } }); + params.validate_format_version()?; + let format_version = params.resolved_format_version(); let details = pbold::InvertedIndexDetails::try_from(¶ms)?; let mut inverted_index = InvertedIndexBuilder::new_with_fragment_mask(params, fragment_mask) @@ -154,7 +156,7 @@ impl InvertedIndexPlugin { let files = inverted_index.update(data, index_store, None).await?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&details).unwrap(), - index_version: current_fts_format_version().index_version(), + index_version: format_version.index_version(), files, }) } @@ -193,11 +195,7 @@ impl TrainingRequest for InvertedIndexTrainingRequest { } #[async_trait] -impl ScalarIndexPlugin for InvertedIndexPlugin { - fn name(&self) -> &str { - "Inverted" - } - +impl BasicTrainer for InvertedIndexPlugin { fn new_training_request( &self, params: &str, @@ -215,37 +213,10 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { .into())) } - let params = serde_json::from_str::(params)?; + let params = InvertedIndexParams::from_training_json(params)?; Ok(Box::new(InvertedIndexTrainingRequest::new(params))) } - fn provides_exact_answer(&self) -> bool { - false - } - - fn version(&self) -> u32 { - max_supported_fts_format_version().index_version() - } - - fn new_query_parser( - &self, - index_name: String, - _index_details: &prost_types::Any, - ) -> Option> { - let Ok(index_details) = _index_details.to_msg::() else { - return None; - }; - - if Self::can_accelerate_queries(&index_details) { - Some(Box::new(FtsQueryParser::new( - index_name, - self.name().to_string(), - ))) - } else { - None - } - } - /// Train a new index /// /// The provided data must fulfill all the criteria returned by `training_criteria`. @@ -280,6 +251,44 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { ) .await } +} + +#[async_trait] +impl ScalarIndexPlugin for InvertedIndexPlugin { + fn basic_trainer(&self) -> Option<&dyn BasicTrainer> { + Some(self) + } + + fn name(&self) -> &str { + "Inverted" + } + + fn provides_exact_answer(&self) -> bool { + false + } + + fn version(&self) -> u32 { + max_supported_fts_format_version().index_version() + } + + fn new_query_parser( + &self, + index_name: String, + _index_details: &prost_types::Any, + ) -> Option> { + let Ok(index_details) = _index_details.to_msg::() else { + return None; + }; + + if Self::can_accelerate_queries(&index_details) { + Some(Box::new(FtsQueryParser::new( + index_name, + self.name().to_string(), + ))) + } else { + None + } + } /// Load an index from storage /// @@ -289,7 +298,7 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { &self, index_store: Arc, _index_details: &prost_types::Any, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result> { Ok( @@ -308,6 +317,7 @@ impl ScalarIndexPlugin for InvertedIndexPlugin { #[cfg(test)] mod tests { use super::*; + use crate::scalar::{BuiltinIndexType, ScalarIndexParams}; #[test] fn test_plugin_version_tracks_max_supported_format() { @@ -317,4 +327,36 @@ mod tests { max_supported_fts_format_version().index_version() ); } + + #[test] + fn test_new_training_request_defaults_missing_block_size_to_128() { + let plugin = InvertedIndexPlugin; + let field = Field::new("text", DataType::Utf8, true); + + let cases = [ + ( + ScalarIndexParams::for_builtin(BuiltinIndexType::Inverted), + false, + ), + (ScalarIndexParams::new("inverted".to_string()), false), + ( + ScalarIndexParams::new("inverted".to_string()) + .with_params(&serde_json::json!({ "with_position": true })), + true, + ), + ]; + + for (params, expected_with_position) in cases { + let request = plugin + .new_training_request(params.params.as_deref().unwrap_or("{}"), &field) + .unwrap(); + let request = request + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(request.parameters.posting_block_size(), DEFAULT_BLOCK_SIZE); + assert_eq!(request.parameters.has_positions(), expected_with_position); + } + } } diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index a53b6ddd7dc..41804e94fd1 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -1,10 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use super::encoding::encode_group_starts; use super::{InvertedIndexParams, index::*}; use crate::scalar::inverted::document_tokenizer::DocType; use crate::scalar::inverted::json::JsonTextStream; +use crate::scalar::inverted::tokenizer::LEGACY_BLOCK_SIZE; use crate::scalar::inverted::tokenizer::document_tokenizer::LanceTokenizer; #[cfg(test)] use crate::scalar::lance_format::LanceIndexStore; @@ -13,37 +13,35 @@ use crate::vector::graph::OrderedFloat; use crate::{progress::IndexBuildProgress, progress::noop_progress}; use arrow::array::AsArray; use arrow::datatypes; -use arrow_array::{Array, BinaryArray, RecordBatch, UInt64Array}; +use arrow_array::{Array, BinaryArray, RecordBatch}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; -use bitpacking::{BitPacker, BitPacker4x}; -use bytes::Bytes; -use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream}; +use datafusion::execution::SendableRecordBatchStream; use fst::Streamer; -use futures::{Stream, StreamExt, TryStreamExt}; +use futures::{StreamExt, TryStreamExt}; use lance_arrow::json::JSON_EXT_NAME; use lance_arrow::{ARROW_EXT_NAME_KEY, iter_str_array}; +use lance_bitpacking::{BitPacker, BitPacker4x}; use lance_core::cache::LanceCache; use lance_core::deepsize::DeepSizeOf; use lance_core::error::LanceOptionExt; +use lance_core::utils::row_addr_remap::RowAddrRemap; use lance_core::utils::tokio::{IO_CORE_RESERVATION, get_num_compute_intensive_cpus, spawn_cpu}; -use lance_core::{Error, ROW_ID, ROW_ID_FIELD, Result}; +use lance_core::{Error, ROW_ID, Result}; use lance_io::object_store::ObjectStore; use lance_select::RowSetOps; use object_store::path::Path; use roaring::RoaringBitmap; use smallvec::SmallVec; use std::collections::HashMap; -use std::pin::Pin; use std::str::FromStr; use std::sync::Arc; use std::sync::LazyLock; -use std::task::{Context, Poll}; use std::{fmt::Debug, sync::atomic::AtomicU64}; use tracing::instrument; -// the number of elements in each block -// each block contains 128 row ids and 128 frequencies -// WARNING: changing this value will break the compatibility with existing indexes +// The legacy bitpacking block size. Position streams still use this block size; +// FTS posting blocks choose their physical bitpacker from the configured +// InvertedIndexParams::block_size. pub const BLOCK_SIZE: usize = BitPacker4x::BLOCK_LEN; // The default number of workers to use for FTS builds. @@ -74,92 +72,8 @@ static LANCE_FTS_POSTING_BATCH_ROWS: LazyLock = LazyLock::new(|| { .parse() .expect("failed to parse LANCE_FTS_POSTING_BATCH_ROWS") }); -// Target serialized byte size of a posting-list cache group. Consecutive -// posting lists are grouped into a single cache entry until their combined -// serialized size reaches this target, amortizing per-entry overhead across -// small (Zipfian-rare) terms. See issue #7040. -static LANCE_FTS_POSTING_GROUP_TARGET_BYTES: LazyLock = LazyLock::new(|| { - std::env::var("LANCE_FTS_POSTING_GROUP_TARGET_BYTES") - .unwrap_or_else(|_| "4096".to_string()) - .parse() - .expect("failed to parse LANCE_FTS_POSTING_GROUP_TARGET_BYTES") -}); -// Maximum number of posting lists in a single cache group, regardless of byte -// size. Caps the work and memory of a single group read for corpora with many -// tiny terms. -static LANCE_FTS_POSTING_GROUP_MAX_TOKENS: LazyLock = LazyLock::new(|| { - std::env::var("LANCE_FTS_POSTING_GROUP_MAX_TOKENS") - .unwrap_or_else(|_| "256".to_string()) - .parse() - .expect("failed to parse LANCE_FTS_POSTING_GROUP_MAX_TOKENS") -}); const MAX_RETAINED_TOKEN_IDS: usize = 8 * 1024; -/// Write-time configuration controlling how consecutive posting lists are -/// grouped into a single read-path cache entry (issue #7040). Defaults come -/// from the `LANCE_FTS_POSTING_GROUP_*` environment variables. -#[derive(Debug, Clone, Copy)] -pub(crate) struct PostingGroupConfig { - pub(crate) target_bytes: usize, - pub(crate) max_tokens: usize, -} - -impl Default for PostingGroupConfig { - fn default() -> Self { - Self { - target_bytes: (*LANCE_FTS_POSTING_GROUP_TARGET_BYTES).max(1), - max_tokens: (*LANCE_FTS_POSTING_GROUP_MAX_TOKENS).max(1), - } - } -} - -/// Accumulates posting-list group boundaries at write time. Tokens are pushed -/// in row order; a group is cut once its serialized bytes reach -/// `target_bytes` or it holds `max_tokens` posting lists. A posting list -/// larger than the target that *starts* a group occupies that group alone (the -/// clamp case); one encountered mid-group is absorbed and closes that group, so -/// a single term is never split across groups. -#[derive(Debug)] -pub(crate) struct PostingGroupAccumulator { - config: PostingGroupConfig, - starts: Vec, - next_token: u32, - current_bytes: usize, - current_tokens: usize, -} - -impl PostingGroupAccumulator { - pub(crate) fn new(config: PostingGroupConfig) -> Self { - Self { - config, - starts: Vec::new(), - next_token: 0, - current_bytes: 0, - current_tokens: 0, - } - } - - /// Record the next posting list in row order, given its serialized byte size. - pub(crate) fn push(&mut self, posting_bytes: usize) { - if self.current_tokens == 0 { - self.starts.push(self.next_token); - } - self.current_bytes += posting_bytes; - self.current_tokens += 1; - self.next_token += 1; - if self.current_bytes >= self.config.target_bytes - || self.current_tokens >= self.config.max_tokens - { - self.current_bytes = 0; - self.current_tokens = 0; - } - } - - pub(crate) fn into_starts(self) -> Vec { - self.starts - } -} - fn default_num_workers() -> usize { let total_cpus = get_num_compute_intensive_cpus() + *IO_CORE_RESERVATION; std::cmp::max(1, total_cpus / 2) @@ -181,25 +95,44 @@ fn resolve_worker_memory_limit_bytes(params: &InvertedIndexParams, num_workers: .unwrap_or(default_worker_memory_limit_bytes) } -fn merge_all_tail_partitions(tails: Vec) -> Result> { - if tails.is_empty() { - return Ok(None); +/// Merge the workers' leftover tail builders into as few partitions as the +/// memory budget allows. Folding unconditionally would collapse every build +/// whose workers never hit the flush threshold into a single partition, which +/// destroys intra-query parallelism; splitting by the same per-partition +/// budget as the flush path makes the final partition count converge to +/// roughly total_builder_memory / memory_limit_bytes regardless of worker +/// layout. +fn merge_all_tail_partitions( + tails: Vec, + memory_limit_bytes: u64, +) -> Result> { + let mut merged_builders: Vec = Vec::new(); + let mut merged: Option = None; + for tail in tails { + let builder = tail.builder; + if builder.is_empty() { + continue; + } + match &mut merged { + Some(current) => { + let would_exceed_memory = + current.memory_size().saturating_add(builder.memory_size()) + >= memory_limit_bytes; + let would_exceed_doc_ids = + current.docs.len().saturating_add(builder.docs.len()) > u32::MAX as usize; + if would_exceed_memory || would_exceed_doc_ids { + merged_builders.push(std::mem::replace(current, builder)); + } else { + current.merge_from(builder)?; + } + } + None => merged = Some(builder), + } } - merge_tail_partition_group(tails).map(Some) -} - -fn merge_tail_partition_group(group: Vec) -> Result { - let mut group = group.into_iter(); - let mut merged = group - .next() - .ok_or_else(|| { - Error::invalid_input("cannot merge an empty tail partition group".to_owned()) - })? - .builder; - for tail in group { - merged.merge_from(tail.builder)?; - } - Ok(merged) + if let Some(builder) = merged { + merged_builders.push(builder); + } + Ok(merged_builders) } #[derive(Debug)] @@ -246,6 +179,7 @@ impl InvertedIndexBuilder { fragment_mask: Option, deleted_fragments: RoaringBitmap, ) -> Self { + let format_version = params.resolved_format_version(); Self { params, partitions, @@ -253,16 +187,19 @@ impl InvertedIndexBuilder { src_store: store, token_set_format, fragment_mask, - format_version: current_fts_format_version(), - posting_tail_codec: current_fts_format_version().posting_tail_codec(), + format_version, + posting_tail_codec: format_version.posting_tail_codec(), progress: noop_progress(), deleted_fragments, } } pub fn with_posting_tail_codec(mut self, posting_tail_codec: PostingTailCodec) -> Self { - self.format_version = - InvertedListFormatVersion::from_posting_tail_codec(posting_tail_codec); + self.format_version = InvertedListFormatVersion::from_posting_tail_codec_and_block_size( + posting_tail_codec, + self.params.block_size, + ) + .expect("invalid posting tail codec for posting block size"); self.posting_tail_codec = posting_tail_codec; self } @@ -289,6 +226,7 @@ impl InvertedIndexBuilder { dest_store: &dyn IndexStore, old_data_filter: Option, ) -> Result> { + validate_format_version_block_size(self.format_version, self.params.block_size)?; let schema = new_data.schema(); let doc_col = schema.field(0).name(); @@ -323,6 +261,7 @@ impl InvertedIndexBuilder { old_segments: &[Arc], old_data_filter: Option, ) -> Result> { + validate_format_version_block_size(self.format_version, self.params.block_size)?; let schema = new_data.schema(); let doc_col = schema.field(0).name(); @@ -448,6 +387,7 @@ impl InvertedIndexBuilder { fragment_mask: self.fragment_mask, token_set_format: self.token_set_format, worker_memory_limit_bytes, + block_size: self.params.block_size, }; let next_id = self.next_partition_id(); let id_alloc = Arc::new(AtomicU64::new(next_id)); @@ -528,16 +468,28 @@ impl InvertedIndexBuilder { tail_partitions.push(tail_partition); } } - let merged_tail_partitions = - spawn_cpu(move || merge_all_tail_partitions(tail_partitions)).await?; - if let Some(builder) = merged_tail_partitions { - self.new_partitions.push(builder.id()); - let mut builder = builder; - files.extend( - builder - .write_to(dest_store.as_ref(), self.partition_write_target()) - .await?, - ); + let merged_tail_partitions = spawn_cpu(move || { + merge_all_tail_partitions(tail_partitions, worker_memory_limit_bytes) + }) + .await?; + // Tail partitions hold most of the data when workers rarely hit the + // flush threshold; writing them one at a time serializes the + // posting-list compression of nearly the whole index behind a + // single producer thread. Compress and write them concurrently. + let write_target = self.partition_write_target(); + let mut tail_writes = + futures::stream::iter(merged_tail_partitions.into_iter().map(|mut builder| { + let dest_store = dest_store.clone(); + async move { + let partition_id = builder.id(); + let files = builder.write_to(dest_store.as_ref(), write_target).await?; + Result::Ok((partition_id, files)) + } + })) + .buffer_unordered(get_num_compute_intensive_cpus().clamp(1, 16)); + while let Some((partition_id, partition_files)) = tail_writes.try_next().await? { + self.new_partitions.push(partition_id); + files.extend(partition_files); } log::info!("wait workers indexing elapsed: {:?}", start.elapsed()); Result::Ok(files) @@ -548,7 +500,7 @@ impl InvertedIndexBuilder { pub async fn remap( &mut self, - mapping: &HashMap>, + mapping: &RowAddrRemap, src_store: Arc, dest_store: &dyn IndexStore, ) -> Result> { @@ -586,6 +538,7 @@ impl InvertedIndexBuilder { dest_store: &dyn IndexStore, partitions: &[u64], ) -> Result { + validate_format_version_block_size(self.format_version, self.params.block_size)?; let mut serialized_deleted_fragments = Vec::with_capacity(self.deleted_fragments.serialized_size()); self.deleted_fragments @@ -602,6 +555,14 @@ impl InvertedIndexBuilder { POSTING_TAIL_CODEC_KEY.to_owned(), self.posting_tail_codec.as_str().to_owned(), ), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + self.format_version.index_version().to_string(), + ), + ( + POSTING_BLOCK_SIZE_KEY.to_owned(), + self.params.block_size.to_string(), + ), ]); if self.params.with_position && self.format_version.uses_shared_position_stream() { @@ -646,6 +607,7 @@ impl InvertedIndexBuilder { dest_store: &dyn IndexStore, partition: u64, // Modify parameter type ) -> Result { + validate_format_version_block_size(self.format_version, self.params.block_size)?; let partitions = vec![partition]; let mut metadata = HashMap::from_iter(vec![ ("partitions".to_owned(), serde_json::to_string(&partitions)?), @@ -658,6 +620,14 @@ impl InvertedIndexBuilder { POSTING_TAIL_CODEC_KEY.to_owned(), self.posting_tail_codec.as_str().to_owned(), ), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + self.format_version.index_version().to_string(), + ), + ( + POSTING_BLOCK_SIZE_KEY.to_owned(), + self.params.block_size.to_string(), + ), ]); if self.params.with_position && self.format_version.uses_shared_position_stream() { metadata.insert( @@ -794,10 +764,10 @@ pub struct InnerBuilder { token_set_format: TokenSetFormat, format_version: InvertedListFormatVersion, posting_tail_codec: PostingTailCodec, + block_size: usize, pub(crate) tokens: TokenSet, pub(crate) posting_lists: Vec, pub(crate) docs: DocSet, - pub(crate) group_config: PostingGroupConfig, } impl InnerBuilder { @@ -816,16 +786,51 @@ impl InnerBuilder { token_set_format: TokenSetFormat, format_version: InvertedListFormatVersion, ) -> Self { + Self::new_with_format_version_and_block_size( + id, + with_position, + token_set_format, + format_version, + LEGACY_BLOCK_SIZE, + ) + } + + pub fn new_with_block_size( + id: u64, + with_position: bool, + token_set_format: TokenSetFormat, + block_size: usize, + ) -> Self { + let format_version = default_fts_format_version_for_block_size(block_size) + .expect("invalid posting list block size"); + Self::new_with_format_version_and_block_size( + id, + with_position, + token_set_format, + format_version, + block_size, + ) + } + + pub fn new_with_format_version_and_block_size( + id: u64, + with_position: bool, + token_set_format: TokenSetFormat, + format_version: InvertedListFormatVersion, + block_size: usize, + ) -> Self { + validate_format_version_block_size(format_version, block_size) + .expect("invalid FTS format version for posting block size"); Self { id, with_position, token_set_format, format_version, posting_tail_codec: format_version.posting_tail_codec(), + block_size, tokens: TokenSet::default(), posting_lists: Vec::new(), docs: DocSet::default(), - group_config: PostingGroupConfig::default(), } } @@ -835,13 +840,34 @@ impl InnerBuilder { token_set_format: TokenSetFormat, posting_tail_codec: PostingTailCodec, ) -> Self { - let format_version = if posting_tail_codec == PostingTailCodec::Fixed32 { - InvertedListFormatVersion::V1 - } else { - InvertedListFormatVersion::V2 - }; - let mut builder = - Self::new_with_format_version(id, with_position, token_set_format, format_version); + Self::new_with_posting_tail_codec_and_block_size( + id, + with_position, + token_set_format, + posting_tail_codec, + LEGACY_BLOCK_SIZE, + ) + } + + pub fn new_with_posting_tail_codec_and_block_size( + id: u64, + with_position: bool, + token_set_format: TokenSetFormat, + posting_tail_codec: PostingTailCodec, + block_size: usize, + ) -> Self { + let format_version = InvertedListFormatVersion::from_posting_tail_codec_and_block_size( + posting_tail_codec, + block_size, + ) + .expect("invalid posting tail codec for posting block size"); + let mut builder = Self::new_with_format_version_and_block_size( + id, + with_position, + token_set_format, + format_version, + block_size, + ); builder.posting_tail_codec = posting_tail_codec; builder } @@ -873,7 +899,7 @@ impl InnerBuilder { self.posting_lists = posting_lists; } - pub async fn remap(&mut self, mapping: &HashMap>) -> Result<()> { + pub async fn remap(&mut self, mapping: &RowAddrRemap) -> Result<()> { // for the docs, we need to remove the rows that are removed from the doc set, // and update the row ids of the rows that are updated let removed = self.docs.remap(mapping); @@ -912,7 +938,7 @@ impl InnerBuilder { mapping.insert(*row_id, None); } } - self.remap(&mapping).await + self.remap(&RowAddrRemap::direct(mapping)).await } pub fn merge_from(&mut self, other: Self) -> Result<()> { @@ -922,10 +948,10 @@ impl InnerBuilder { token_set_format, format_version, posting_tail_codec, + block_size, tokens, posting_lists, docs, - group_config: _, } = other; if self.with_position != with_position { @@ -952,6 +978,12 @@ impl InnerBuilder { self.posting_tail_codec, posting_tail_codec ))); } + if self.block_size != block_size { + return Err(Error::index(format!( + "cannot merge partitions with mismatched FTS block sizes: {} vs {}", + self.block_size, block_size + ))); + } let mut token_id_map = vec![u32::MAX; posting_lists.len()]; match tokens.tokens { @@ -977,7 +1009,11 @@ impl InnerBuilder { self.docs.append(*row_id, *num_tokens); } self.posting_lists.resize_with(self.tokens.len(), || { - PostingListBuilder::new_with_posting_tail_codec(with_position, self.posting_tail_codec) + PostingListBuilder::new_with_posting_tail_codec_and_block_size( + with_position, + self.posting_tail_codec, + self.block_size, + ) }); for (token_id, posting_list) in posting_lists.into_iter().enumerate() { @@ -1044,7 +1080,11 @@ impl InnerBuilder { let mut writer = store .new_index_file( path, - inverted_list_schema_for_version(self.with_position, self.format_version), + inverted_list_schema_for_version_with_block_size( + self.with_position, + self.format_version, + self.block_size, + ), ) .await?; let posting_lists = std::mem::take(&mut self.posting_lists); @@ -1057,48 +1097,73 @@ impl InnerBuilder { ); let with_position = self.with_position; let format_version = self.format_version; - let group_config = self.group_config; - let schema = inverted_list_schema_for_version(self.with_position, self.format_version); + let schema = inverted_list_schema_for_version_with_block_size( + self.with_position, + self.format_version, + self.block_size, + ); let docs_for_batches = docs.clone(); let schema_for_batches = schema.clone(); let batch_rows = *LANCE_FTS_POSTING_BATCH_ROWS; let (tx, rx) = async_channel::bounded(*LANCE_FTS_WRITE_QUEUE_SIZE); - let producer = spawn_cpu(move || { + // The producer builds posting-list batches on the CPU pool and hands them + // to the writer through a bounded channel. Each batch is built inside its + // own `spawn_cpu` call and dispatched with an async `send().await` instead + // of a blocking send: when the channel is full the producer *yields* its + // task rather than parking the pool thread it is running on. Parking would + // deadlock on hosts whose CPU pool has a single thread: once the consumer + // accumulates enough data to flush an encoded column, the flush also needs + // that pool. The parked producer and the starved consumer would then wait + // on each other forever. + let producer = tokio::spawn(async move { let mut batch_builder = PostingListBatchBuilder::new( - schema_for_batches.clone(), + schema_for_batches, with_position, format_version, batch_rows, - group_config, ); - for posting_list in posting_lists { - posting_list.append_to_batch_with_docs( - &docs_for_batches, - &mut batch_builder, - format_version, - )?; - if batch_builder.len() < batch_rows { - continue; - } - - let batch = batch_builder.finish()?; - if let Err(err) = tx.send_blocking(batch) { - return Err(Error::execution(format!( - "failed to send posting list batch to writer: {err}" - ))); - } - } + let mut posting_lists = posting_lists.into_iter(); + loop { + let docs_for_batches = docs_for_batches.clone(); + // Build the next batch on the CPU pool. The builder and the + // remaining posting lists are moved in and handed back so state + // persists across batches. + let (next_builder, next_posting_lists, batch) = spawn_cpu(move || { + let mut batch_builder = batch_builder; + let mut posting_lists = posting_lists; + let mut batch = None; + for posting_list in posting_lists.by_ref() { + posting_list.append_to_batch_with_docs( + &docs_for_batches, + &mut batch_builder, + format_version, + )?; + if batch_builder.len() >= batch_rows { + batch = Some(batch_builder.finish()?); + break; + } + } + if batch.is_none() && !batch_builder.is_empty() { + batch = Some(batch_builder.finish()?); + } + Result::Ok((batch_builder, posting_lists, batch)) + }) + .await?; + batch_builder = next_builder; + posting_lists = next_posting_lists; - if !batch_builder.is_empty() { - let batch = batch_builder.finish()?; - if let Err(err) = tx.send_blocking(batch) { - return Err(Error::execution(format!( - "failed to send posting list batch to writer: {err}" - ))); + let Some(batch) = batch else { + // No more batches: the posting lists are exhausted. + break; + }; + if tx.send(batch).await.is_err() { + // The receiver is gone, which only happens when the writer + // failed; stop producing and let that error surface there. + break; } } - Result::Ok(batch_builder.into_group_starts()) + Result::Ok(()) }); while let Ok(batch) = rx.recv().await { @@ -1110,22 +1175,8 @@ impl InnerBuilder { } } drop(rx); - let group_starts = producer.await?; - - // Persist the posting-list cache-group boundaries as a global buffer, - // recording its 1-indexed id in schema metadata so the reader can group - // small posting lists into a single cache entry (issue #7040). Empty - // partitions skip this entirely and fall back to the per-token path. - let mut extra_metadata = HashMap::new(); - if !group_starts.is_empty() { - let encoded = encode_group_starts(&group_starts); - let buffer_id = writer.add_global_buffer(Bytes::from(encoded)).await?; - extra_metadata.insert( - POSTING_GROUP_OFFSETS_BUF_KEY.to_owned(), - buffer_id.to_string(), - ); - } - writer.finish_with_metadata(extra_metadata).await + producer.await??; + writer.finish().await } #[instrument(level = "debug", skip_all)] @@ -1207,6 +1258,11 @@ struct WorkerOutput { tail_partition: Option, } +enum DocumentSource<'a> { + Text(&'a str), + StringList(&'a dyn Array), +} + #[derive(Debug, Clone, Copy)] struct IndexWorkerConfig { with_position: bool, @@ -1214,6 +1270,7 @@ struct IndexWorkerConfig { fragment_mask: Option, token_set_format: TokenSetFormat, worker_memory_limit_bytes: u64, + block_size: usize, } impl IndexWorker { @@ -1257,17 +1314,22 @@ impl IndexWorker { id_alloc: Arc, config: IndexWorkerConfig, ) -> Result { - let schema = inverted_list_schema_for_version(config.with_position, config.format_version); + let schema = inverted_list_schema_for_version_with_block_size( + config.with_position, + config.format_version, + config.block_size, + ); Ok(Self { tokenizer, dest_store, - builder: InnerBuilder::new_with_format_version( + builder: InnerBuilder::new_with_format_version_and_block_size( id_alloc.fetch_add(1, std::sync::atomic::Ordering::Relaxed) | config.fragment_mask.unwrap_or(0), config.with_position, config.token_set_format, config.format_version, + config.block_size, ), partitions: Vec::new(), files: Vec::new(), @@ -1292,147 +1354,259 @@ impl IndexWorker { async fn process_batch(&mut self, batch: RecordBatch) -> Result<()> { let doc_col = batch.column(0); - let doc_iter = iter_str_array(doc_col); let row_id_col = batch[ROW_ID].as_primitive::(); - let docs = doc_iter - .zip(row_id_col.values().iter()) - .filter_map(|(doc, row_id)| doc.map(|doc| (doc, *row_id))); + match doc_col.data_type() { + DataType::Utf8 | DataType::LargeUtf8 => { + let docs = iter_str_array(doc_col.as_ref()) + .zip(row_id_col.values().iter()) + .filter_map(|(doc, row_id)| doc.map(|doc| (doc, *row_id))); + + for (doc, row_id) in docs { + self.process_document(row_id, DocumentSource::Text(doc), false) + .await?; + } + } + DataType::List(_) => { + self.process_string_list_batch::(doc_col, row_id_col) + .await?; + } + DataType::LargeList(_) => { + self.process_string_list_batch::(doc_col, row_id_col) + .await?; + } + data_type => { + return Err(Error::index(format!( + "expect data type String, LargeString, List(String), or LargeList(String) but got {}", + data_type + ))); + } + } + + Ok(()) + } + + async fn process_string_list_batch( + &mut self, + doc_col: &Arc, + row_id_col: &arrow_array::PrimitiveArray, + ) -> Result<()> { + let docs = doc_col.as_list::(); + match docs.value_type() { + datatypes::DataType::Utf8 | datatypes::DataType::LargeUtf8 => {} + data_type => { + return Err(Error::index(format!( + "expect list item data type String or LargeString but got {}", + data_type + ))); + } + } + + for (doc, row_id) in docs.iter().zip(row_id_col.values().iter()) { + let Some(doc) = doc else { + continue; + }; + + self.process_document(*row_id, DocumentSource::StringList(doc.as_ref()), true) + .await?; + } + + Ok(()) + } + + fn checked_token_position(row_id: u64, token_position: usize) -> Result { + u32::try_from(token_position).map_err(|_| { + Error::invalid_input(format!( + "token position overflow for row_id={row_id}: token_position={token_position}" + )) + }) + } + fn materialize_string_list(elements: &dyn Array) -> String { + let mut doc = String::new(); + for element in iter_str_array(elements).flatten() { + if !doc.is_empty() { + doc.push(' '); + } + doc.push_str(element); + } + doc + } + + async fn process_document( + &mut self, + row_id: u64, + document: DocumentSource<'_>, + skip_empty_document: bool, + ) -> Result<()> { let with_position = self.has_position(); - for (doc, row_id) in docs { - let builder_was_empty = self.builder.docs.is_empty(); - let old_temporary_memory_size = self.temporary_memory_size(); - let old_token_memory_size = self.builder.tokens.memory_size() as u64; - let doc_id = self.builder.docs.len() as u32; - let mut token_num: u32 = 0; - let mut posting_memory_delta = 0i64; - if with_position { + let builder_was_empty = self.builder.docs.is_empty(); + let old_temporary_memory_size = self.temporary_memory_size(); + let old_token_memory_size = self.builder.tokens.memory_size() as u64; + let doc_id = self.builder.docs.len() as u32; + let mut token_num: u32 = 0; + let mut doc_length_bytes = 0usize; + let mut posting_memory_delta = 0i64; + if with_position { + { if self.token_ids.capacity() < self.last_token_count { self.token_ids .reserve(self.last_token_count - self.token_ids.capacity()); } self.token_ids.clear(); + let tokenizer = &mut self.tokenizer; let builder = &mut self.builder; let token_ids = &mut self.token_ids; let memory_size = &mut self.memory_size; let posting_tail_codec = builder.posting_tail_codec; - let mut token_stream = self.tokenizer.token_stream_for_doc(doc); - while token_stream.advance() { - let token = token_stream.token(); - let token_id = builder.tokens.get_or_add(&token.text); - if token_id as usize == builder.posting_lists.len() { - let old_posting_lists_overhead_size = (builder.posting_lists.capacity() - * std::mem::size_of::()) - as u64; - builder.posting_lists.push( - PostingListBuilder::new_with_posting_tail_codec( - true, - posting_tail_codec, - ), - ); - let new_posting_lists_overhead_size = (builder.posting_lists.capacity() - * std::mem::size_of::()) - as u64; - Self::adjust_tracked_value( - memory_size, - old_posting_lists_overhead_size, - new_posting_lists_overhead_size, - ); + let block_size = builder.block_size; + let mut process_text = |text: &str| -> Result<()> { + doc_length_bytes += text.len(); + let mut token_stream = tokenizer.token_stream_for_doc(text); + while token_stream.advance() { + let token = token_stream.token(); + let position = Self::checked_token_position(row_id, token.position)?; + let token_id = builder.tokens.get_or_add(&token.text); + if token_id as usize == builder.posting_lists.len() { + let old_posting_lists_overhead_size = (builder.posting_lists.capacity() + * std::mem::size_of::()) + as u64; + builder.posting_lists.push( + PostingListBuilder::new_with_posting_tail_codec_and_block_size( + true, + posting_tail_codec, + block_size, + ), + ); + let new_posting_lists_overhead_size = (builder.posting_lists.capacity() + * std::mem::size_of::()) + as u64; + Self::adjust_tracked_value( + memory_size, + old_posting_lists_overhead_size, + new_posting_lists_overhead_size, + ); + } + let posting_list = &mut builder.posting_lists[token_id as usize]; + let old_posting_memory_size = posting_list.size(); + if posting_list.add_occurrence(doc_id, position)? { + token_ids.push(token_id); + } + let new_posting_memory_size = posting_list.size(); + posting_memory_delta += + new_posting_memory_size as i64 - old_posting_memory_size as i64; + token_num += 1; } - let posting_list = &mut builder.posting_lists[token_id as usize]; - let old_posting_memory_size = posting_list.size(); - if posting_list.add_occurrence(doc_id, token.position as u32)? { - token_ids.push(token_id); + Ok(()) + }; + + match document { + DocumentSource::Text(doc) => { + process_text(doc)?; + } + DocumentSource::StringList(elements) => { + let doc = Self::materialize_string_list(elements); + process_text(&doc)?; } - let new_posting_memory_size = posting_list.size(); - posting_memory_delta += - new_posting_memory_size as i64 - old_posting_memory_size as i64; - token_num += 1; } - } else { + } + } else { + { if self.token_ids.capacity() < self.last_token_count { self.token_ids .reserve(self.last_token_count - self.token_ids.capacity()); } self.token_ids.clear(); - let mut token_stream = self.tokenizer.token_stream_for_doc(doc); - while token_stream.advance() { - let token_id = self.builder.tokens.get_or_add(&token_stream.token().text); - self.token_ids.push(token_id); - token_num += 1; - } - } - self.adjust_tracked_memory_size( - old_token_memory_size, - self.builder.tokens.memory_size() as u64, - ); + let tokenizer = &mut self.tokenizer; + let builder = &mut self.builder; + let token_ids = &mut self.token_ids; + let mut process_text = |text: &str| { + doc_length_bytes += text.len(); + let mut token_stream = tokenizer.token_stream_for_doc(text); + while token_stream.advance() { + let token_id = builder.tokens.get_or_add(&token_stream.token().text); + token_ids.push(token_id); + token_num += 1; + } + }; - if !with_position { - let old_posting_lists_overhead_size = self.posting_lists_overhead_size(); - self.builder - .posting_lists - .resize_with(self.builder.tokens.len(), || { - PostingListBuilder::new_with_posting_tail_codec( - false, - self.builder.posting_tail_codec, - ) - }); - let new_posting_lists_overhead_size = self.posting_lists_overhead_size(); - Self::adjust_tracked_value( - &mut self.memory_size, - old_posting_lists_overhead_size, - new_posting_lists_overhead_size, - ); + match document { + DocumentSource::Text(doc) => process_text(doc), + DocumentSource::StringList(elements) => { + let doc = Self::materialize_string_list(elements); + process_text(&doc); + } + } } + } + self.adjust_tracked_memory_size( + old_token_memory_size, + self.builder.tokens.memory_size() as u64, + ); - let old_doc_memory_size = self.builder.docs.memory_size() as u64; - let appended_doc_id = self.builder.docs.append(row_id, token_num); - debug_assert_eq!(appended_doc_id, doc_id); + if skip_empty_document && token_num == 0 { + self.last_token_count = 0; + self.trim_temporary_buffers(); self.adjust_tracked_memory_size( - old_doc_memory_size, - self.builder.docs.memory_size() as u64, + old_temporary_memory_size, + self.temporary_memory_size(), ); - self.total_doc_length += doc.len(); + return Ok(()); + } - if with_position { - for &token_id in &self.token_ids { - let (old_posting_memory_size, new_posting_memory_size) = { - let posting_list = &mut self.builder.posting_lists[token_id as usize]; - let old_posting_memory_size = posting_list.size(); - posting_list.finish_open_doc(doc_id)?; - let new_posting_memory_size = posting_list.size(); - (old_posting_memory_size, new_posting_memory_size) - }; - posting_memory_delta += - new_posting_memory_size as i64 - old_posting_memory_size as i64; - } - Self::apply_delta(&mut self.memory_size, posting_memory_delta); - } else if token_num > 0 { - self.token_ids.sort_unstable(); - let mut iter = self.token_ids.iter(); - let mut current = *iter.next().unwrap(); - let mut count = 1u32; - for &token_id in iter { - if token_id == current { - count += 1; - continue; - } + if !with_position { + let old_posting_lists_overhead_size = self.posting_lists_overhead_size(); + self.builder + .posting_lists + .resize_with(self.builder.tokens.len(), || { + PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + self.builder.posting_tail_codec, + self.builder.block_size, + ) + }); + let new_posting_lists_overhead_size = self.posting_lists_overhead_size(); + Self::adjust_tracked_value( + &mut self.memory_size, + old_posting_lists_overhead_size, + new_posting_lists_overhead_size, + ); + } - let (old_posting_memory_size, new_posting_memory_size) = { - let posting_list = &mut self.builder.posting_lists[current as usize]; - let old_posting_memory_size = posting_list.size(); - posting_list.add(doc_id, PositionRecorder::Count(count)); - let new_posting_memory_size = posting_list.size(); - (old_posting_memory_size, new_posting_memory_size) - }; - posting_memory_delta += - new_posting_memory_size as i64 - old_posting_memory_size as i64; + let old_doc_memory_size = self.builder.docs.memory_size() as u64; + let appended_doc_id = self.builder.docs.append(row_id, token_num); + debug_assert_eq!(appended_doc_id, doc_id); + self.adjust_tracked_memory_size( + old_doc_memory_size, + self.builder.docs.memory_size() as u64, + ); + self.total_doc_length += doc_length_bytes; - current = token_id; - count = 1; + if with_position { + for &token_id in &self.token_ids { + let (old_posting_memory_size, new_posting_memory_size) = { + let posting_list = &mut self.builder.posting_lists[token_id as usize]; + let old_posting_memory_size = posting_list.size(); + posting_list.finish_open_doc(doc_id)?; + let new_posting_memory_size = posting_list.size(); + (old_posting_memory_size, new_posting_memory_size) + }; + posting_memory_delta += + new_posting_memory_size as i64 - old_posting_memory_size as i64; + } + Self::apply_delta(&mut self.memory_size, posting_memory_delta); + } else if token_num > 0 { + self.token_ids.sort_unstable(); + let mut iter = self.token_ids.iter(); + let mut current = *iter.next().unwrap(); + let mut count = 1u32; + for &token_id in iter { + if token_id == current { + count += 1; + continue; } + let (old_posting_memory_size, new_posting_memory_size) = { let posting_list = &mut self.builder.posting_lists[current as usize]; let old_posting_memory_size = posting_list.size(); @@ -1442,27 +1616,35 @@ impl IndexWorker { }; posting_memory_delta += new_posting_memory_size as i64 - old_posting_memory_size as i64; - Self::apply_delta(&mut self.memory_size, posting_memory_delta); - } - self.last_token_count = self.token_ids.len(); - self.trim_temporary_buffers(); - self.adjust_tracked_memory_size( - old_temporary_memory_size, - self.temporary_memory_size(), - ); - if self.builder.docs.len() == 1 && self.memory_size > self.worker_memory_limit_bytes { - return Err(Error::invalid_input(format!( - "single document row_id={} exceeds worker memory limit: {} > {} bytes", - row_id, self.memory_size, self.worker_memory_limit_bytes - ))); + current = token_id; + count = 1; } + let (old_posting_memory_size, new_posting_memory_size) = { + let posting_list = &mut self.builder.posting_lists[current as usize]; + let old_posting_memory_size = posting_list.size(); + posting_list.add(doc_id, PositionRecorder::Count(count)); + let new_posting_memory_size = posting_list.size(); + (old_posting_memory_size, new_posting_memory_size) + }; + posting_memory_delta += new_posting_memory_size as i64 - old_posting_memory_size as i64; + Self::apply_delta(&mut self.memory_size, posting_memory_delta); + } + self.last_token_count = self.token_ids.len(); + self.trim_temporary_buffers(); + self.adjust_tracked_memory_size(old_temporary_memory_size, self.temporary_memory_size()); - if self.builder.docs.len() as u32 == u32::MAX - || (!builder_was_empty && self.memory_size >= self.worker_memory_limit_bytes) - { - self.flush().await?; - } + if self.builder.docs.len() == 1 && self.memory_size > self.worker_memory_limit_bytes { + return Err(Error::invalid_input(format!( + "single document row_id={} exceeds worker memory limit: {} > {} bytes", + row_id, self.memory_size, self.worker_memory_limit_bytes + ))); + } + + if self.builder.docs.len() as u32 == u32::MAX + || (!builder_was_empty && self.memory_size >= self.worker_memory_limit_bytes) + { + self.flush().await?; } Ok(()) @@ -1481,15 +1663,17 @@ impl IndexWorker { self.memory_size = self.temporary_memory_size(); let with_position = self.has_position(); let format_version = self.builder.format_version; + let block_size = self.builder.block_size; let builder = std::mem::replace( &mut self.builder, - InnerBuilder::new_with_format_version( + InnerBuilder::new_with_format_version_and_block_size( self.id_alloc .fetch_add(1, std::sync::atomic::Ordering::Relaxed) | self.fragment_mask.unwrap_or(0), with_position, self.token_set_format, format_version, + block_size, ), ); let written_partition_id = builder.id(); @@ -1609,17 +1793,56 @@ pub fn inverted_list_schema_for_version( with_position: bool, format_version: InvertedListFormatVersion, ) -> SchemaRef { + inverted_list_schema_for_version_with_block_size( + with_position, + format_version, + LEGACY_BLOCK_SIZE, + ) +} + +pub fn inverted_list_schema_for_version_with_block_size( + with_position: bool, + format_version: InvertedListFormatVersion, + block_size: usize, +) -> SchemaRef { + inverted_list_schema_for_version_with_block_size_and_impacts( + with_position, + format_version, + block_size, + true, + ) +} + +pub(crate) fn inverted_list_schema_for_version_with_block_size_and_impacts( + with_position: bool, + format_version: InvertedListFormatVersion, + block_size: usize, + with_impacts: bool, +) -> SchemaRef { + validate_format_version_block_size(format_version, block_size) + .expect("invalid FTS format version for posting block size"); match format_version { - InvertedListFormatVersion::V1 => inverted_list_schema_v1(with_position), - InvertedListFormatVersion::V2 => inverted_list_schema_with_tail_codec_and_position_codec( - with_position, - PostingTailCodec::VarintDelta, - Some(PositionStreamCodec::PackedDelta), - ), + InvertedListFormatVersion::V1 => { + inverted_list_schema_v1(with_position, block_size, with_impacts) + } + InvertedListFormatVersion::V2 | InvertedListFormatVersion::V3 => { + inverted_list_schema_with_tail_codec_and_position_codec( + with_position, + format_version, + PostingTailCodec::VarintDelta, + Some(PositionStreamCodec::PackedDelta), + block_size, + with_impacts, + ) + } } } -fn inverted_list_schema_v1(with_position: bool) -> SchemaRef { +fn inverted_list_schema_v1( + with_position: bool, + block_size: usize, + with_impacts: bool, +) -> SchemaRef { let mut fields = vec![ arrow_schema::Field::new( POSTING_COL, @@ -1633,6 +1856,17 @@ fn inverted_list_schema_v1(with_position: bool) -> SchemaRef { arrow_schema::Field::new(MAX_SCORE_COL, datatypes::DataType::Float32, false), arrow_schema::Field::new(LENGTH_COL, datatypes::DataType::UInt32, false), ]; + if with_impacts { + fields.push(arrow_schema::Field::new( + IMPACT_COL, + datatypes::DataType::List(Arc::new(Field::new( + "item", + datatypes::DataType::LargeBinary, + true, + ))), + false, + )); + } if with_position { fields.push(arrow_schema::Field::new( POSITION_COL, @@ -1648,24 +1882,44 @@ fn inverted_list_schema_v1(with_position: bool) -> SchemaRef { false, )); } - Arc::new(arrow_schema::Schema::new(fields)) + Arc::new(arrow_schema::Schema::new_with_metadata( + fields, + HashMap::from([ + (POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string()), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + InvertedListFormatVersion::V1.index_version().to_string(), + ), + ]), + )) } pub fn inverted_list_schema_with_tail_codec( with_position: bool, posting_tail_codec: PostingTailCodec, ) -> SchemaRef { + let format_version = InvertedListFormatVersion::from_posting_tail_codec_and_block_size( + posting_tail_codec, + LEGACY_BLOCK_SIZE, + ) + .expect("invalid posting tail codec for posting block size"); inverted_list_schema_with_tail_codec_and_position_codec( with_position, + format_version, posting_tail_codec, Some(PositionStreamCodec::PackedDelta), + LEGACY_BLOCK_SIZE, + false, ) } fn inverted_list_schema_with_tail_codec_and_position_codec( with_position: bool, + format_version: InvertedListFormatVersion, posting_tail_codec: PostingTailCodec, position_codec: Option, + block_size: usize, + with_impacts: bool, ) -> SchemaRef { let mut fields = vec![ // we compress the posting lists (including row ids and frequencies), @@ -1682,6 +1936,17 @@ fn inverted_list_schema_with_tail_codec_and_position_codec( arrow_schema::Field::new(MAX_SCORE_COL, datatypes::DataType::Float32, false), arrow_schema::Field::new(LENGTH_COL, datatypes::DataType::UInt32, false), ]; + if with_impacts { + fields.push(arrow_schema::Field::new( + IMPACT_COL, + datatypes::DataType::List(Arc::new(Field::new( + "item", + datatypes::DataType::LargeBinary, + true, + ))), + false, + )); + } if with_position { fields.push(arrow_schema::Field::new( COMPRESSED_POSITION_COL, @@ -1702,6 +1967,11 @@ fn inverted_list_schema_with_tail_codec_and_position_codec( POSTING_TAIL_CODEC_KEY.to_owned(), posting_tail_codec.as_str().to_owned(), )]); + metadata.insert( + FTS_FORMAT_VERSION_KEY.to_owned(), + format_version.index_version().to_string(), + ); + metadata.insert(POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string()); if let Some(position_codec) = position_codec.filter(|_| with_position) { metadata.insert( POSITIONS_LAYOUT_KEY.to_owned(), @@ -1715,129 +1985,6 @@ fn inverted_list_schema_with_tail_codec_and_position_codec( Arc::new(arrow_schema::Schema::new_with_metadata(fields, metadata)) } -/// Flatten the string list stream into a string stream -pub struct FlattenStream { - /// Inner record batch stream with 2 columns: - /// 1. doc_col: List(Utf8) or List(LargeUtf8) - /// 2. row_id_col: UInt64 - inner: SendableRecordBatchStream, - field_type: DataType, - data_type: DataType, -} - -impl FlattenStream { - pub fn new(input: SendableRecordBatchStream) -> Self { - let schema = input.schema(); - let field = schema.field(0); - let data_type = match field.data_type() { - DataType::List(f) if matches!(f.data_type(), DataType::Utf8) => DataType::Utf8, - DataType::List(f) if matches!(f.data_type(), DataType::LargeUtf8) => { - DataType::LargeUtf8 - } - DataType::LargeList(f) if matches!(f.data_type(), DataType::Utf8) => DataType::Utf8, - DataType::LargeList(f) if matches!(f.data_type(), DataType::LargeUtf8) => { - DataType::LargeUtf8 - } - _ => panic!( - "expect data type List(Utf8) or List(LargeUtf8) but got {:?}", - field.data_type() - ), - }; - Self { - inner: input, - field_type: field.data_type().clone(), - data_type, - } - } -} - -impl Stream for FlattenStream { - type Item = datafusion_common::Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - match Pin::new(&mut self.inner).poll_next(cx) { - Poll::Ready(Some(Ok(batch))) => { - let doc_col = batch.column(0); - let batch = match self.field_type { - DataType::List(_) => flatten_string_list::(&batch, doc_col).map_err(|e| { - datafusion_common::error::DataFusionError::Execution(format!( - "flatten string list error: {}", - e - )) - }), - DataType::LargeList(_) => { - flatten_string_list::(&batch, doc_col).map_err(|e| { - datafusion_common::error::DataFusionError::Execution(format!( - "flatten string list error: {}", - e - )) - }) - } - _ => unreachable!( - "expect data type List or LargeList but got {:?}", - self.field_type - ), - }; - Poll::Ready(Some(batch)) - } - Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))), - Poll::Ready(None) => Poll::Ready(None), - Poll::Pending => Poll::Pending, - } - } -} - -impl RecordBatchStream for FlattenStream { - fn schema(&self) -> SchemaRef { - let schema = Schema::new(vec![ - Field::new( - self.inner.schema().field(0).name(), - self.data_type.clone(), - true, - ), - ROW_ID_FIELD.clone(), - ]); - - Arc::new(schema) - } -} - -fn flatten_string_list( - batch: &RecordBatch, - doc_col: &Arc, -) -> Result { - let docs = doc_col.as_list::(); - let row_ids = batch[ROW_ID].as_primitive::(); - - let row_ids = row_ids - .values() - .iter() - .zip(docs.iter()) - .flat_map(|(row_id, doc)| std::iter::repeat_n(*row_id, doc.map(|d| d.len()).unwrap_or(0))); - - let row_ids = Arc::new(UInt64Array::from_iter_values(row_ids)); - let docs = match docs.value_type() { - datatypes::DataType::Utf8 | datatypes::DataType::LargeUtf8 => docs.values().clone(), - _ => { - return Err(Error::index(format!( - "expect data type String or LargeString but got {}", - docs.value_type() - ))); - } - }; - - let schema = Schema::new(vec![ - Field::new( - batch.schema().field(0).name(), - docs.data_type().clone(), - true, - ), - ROW_ID_FIELD.clone(), - ]); - let batch = RecordBatch::try_new(Arc::new(schema), vec![docs, row_ids])?; - Ok(batch) -} - pub(crate) fn token_file_path(partition_id: u64) -> String { format!("part_{}_{}", partition_id, TOKENS_FILE) } @@ -2111,7 +2258,7 @@ pub fn document_input( DataType::List(field) | DataType::LargeList(field) if matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8) => { - Ok(Box::pin(FlattenStream::new(input))) + Ok(input) } DataType::LargeBinary => match field.metadata().get(ARROW_EXT_NAME_KEY) { Some(name) if name.as_str() == JSON_EXT_NAME => { @@ -2556,8 +2703,7 @@ mod tests { } async fn add_global_buffer(&mut self, _data: Bytes) -> Result { - // The posting-list writer stores the group offsets as a global - // buffer; mirror the real writer's 1-indexed return value. + // Mirror the real writer's 1-indexed return value. Ok(1) } @@ -2642,63 +2788,6 @@ mod tests { } } - fn collect_group_starts(config: PostingGroupConfig, sizes: &[usize]) -> Vec { - let mut acc = PostingGroupAccumulator::new(config); - for &size in sizes { - acc.push(size); - } - acc.into_starts() - } - - #[test] - fn test_group_accumulator_cuts_on_target_bytes() { - let config = PostingGroupConfig { - target_bytes: 100, - max_tokens: 1000, - }; - // 40+40 -> cut at 80? no, 80 < 100; third 40 reaches 120 >= 100 -> cut. - // So group 0 = tokens [0,3), then a new group starts at token 3. - let starts = collect_group_starts(config, &[40, 40, 40, 10, 10]); - assert_eq!(starts, vec![0, 3]); - } - - #[test] - fn test_group_accumulator_cuts_on_max_tokens() { - let config = PostingGroupConfig { - target_bytes: 1_000_000, - max_tokens: 2, - }; - // Byte target never reached; cap of 2 forces a cut every 2 tokens. - let starts = collect_group_starts(config, &[1, 1, 1, 1, 1]); - assert_eq!(starts, vec![0, 2, 4]); - } - - #[test] - fn test_group_accumulator_clamps_oversized_term() { - let config = PostingGroupConfig { - target_bytes: 100, - max_tokens: 64, - }; - // A term larger than the target that *starts* a group occupies that - // group alone ([1, 2) here), so a single huge posting list is never - // forced to share a cache entry. Token 0 (==100) closes its own group - // first; the trailing small terms regroup after the big one. - let starts = collect_group_starts(config, &[100, 5000, 10, 10]); - assert_eq!(starts, vec![0, 1, 2]); - - // A huge term encountered mid-group is absorbed and closes that group; - // we never split one term across groups. - let starts = collect_group_starts(config, &[10, 10, 5000, 10, 10]); - assert_eq!(starts, vec![0, 3]); - } - - #[test] - fn test_group_accumulator_empty_and_single() { - let config = PostingGroupConfig::default(); - assert_eq!(collect_group_starts(config, &[]), Vec::::new()); - assert_eq!(collect_group_starts(config, &[10]), vec![0]); - } - #[tokio::test] async fn test_write_posting_lists_batches_multiple_rows() -> Result<()> { let mut builder = InnerBuilder::new(0, false, TokenSetFormat::default()); @@ -3184,6 +3273,7 @@ mod tests { fragment_mask: None, token_set_format, worker_memory_limit_bytes: u64::MAX, + block_size: params.block_size, }, ) .await?; @@ -3207,6 +3297,7 @@ mod tests { fragment_mask: None, token_set_format, worker_memory_limit_bytes: u64::MAX, + block_size: params.block_size, }, ) .await?; @@ -3230,7 +3321,8 @@ mod tests { token_set_format, None, RoaringBitmap::new(), - ); + ) + .with_format_version(InvertedListFormatVersion::V1); builder.write(dest_store.as_ref()).await?; let metadata_reader = dest_store.open_index_file(METADATA_FILE).await?; @@ -3294,6 +3386,14 @@ mod tests { .await?; let index = InvertedIndex::load(src_store, None, &LanceCache::no_cache()).await?; + let derived_params = index.derive_index_params()?; + let derived_params: InvertedIndexParams = + serde_json::from_str(derived_params.params.as_deref().unwrap())?; + assert_eq!( + derived_params.format_version, + Some(InvertedListFormatVersion::V1) + ); + let schema = Arc::new(Schema::new(vec![ Field::new("doc", DataType::Utf8, true), Field::new(ROW_ID, DataType::UInt64, false), @@ -3680,6 +3780,7 @@ mod tests { fragment_mask: None, token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, + block_size: InvertedIndexParams::default().block_size, }, ) .await?; @@ -3711,6 +3812,7 @@ mod tests { fragment_mask: None, token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, + block_size: InvertedIndexParams::default().block_size, }, ) .await?; @@ -3749,6 +3851,7 @@ mod tests { fragment_mask: None, token_set_format: TokenSetFormat::default(), worker_memory_limit_bytes: u64::MAX, + block_size: InvertedIndexParams::default().block_size, }, ) .await?; @@ -3797,27 +3900,208 @@ mod tests { assert_eq!(resolve_worker_memory_limit_bytes(¶ms, 16), 256 << 20); } + fn tail_with_docs(id: u64, num_docs: u64) -> TailPartition { + let mut builder = InnerBuilder::new(id, false, TokenSetFormat::default()); + let token = builder.tokens.add(format!("token{}", id)); + builder + .posting_lists + .resize_with(builder.tokens.len(), || PostingListBuilder::new(false)); + for row in 0..num_docs { + let doc = builder.docs.append(row, 1); + builder.posting_lists[token as usize].add(doc, PositionRecorder::Count(1)); + } + TailPartition { builder } + } + #[test] - fn test_merge_all_tail_partitions_combines_everything() -> Result<()> { - let merged = merge_all_tail_partitions(vec![ - TailPartition { - builder: InnerBuilder::new(0, false, TokenSetFormat::default()), - }, - TailPartition { - builder: InnerBuilder::new(1, false, TokenSetFormat::default()), - }, - TailPartition { - builder: InnerBuilder::new(2, false, TokenSetFormat::default()), - }, - ])?; + fn test_merge_all_tail_partitions_combines_under_budget() -> Result<()> { + let merged = merge_all_tail_partitions( + vec![ + tail_with_docs(0, 4), + tail_with_docs(1, 4), + tail_with_docs(2, 4), + ], + u64::MAX, + )?; + assert_eq!(merged.len(), 1); + assert_eq!(merged[0].id(), 0); + assert_eq!(merged[0].docs.len(), 12); + Ok(()) + } - assert_eq!(merged.expect("merged builder should exist").id(), 0); + #[test] + fn test_merge_all_tail_partitions_splits_on_memory_budget() -> Result<()> { + let tails = vec![ + tail_with_docs(0, 64), + tail_with_docs(1, 64), + tail_with_docs(2, 64), + tail_with_docs(3, 64), + ]; + let single = tails[0].builder.memory_size(); + // A budget below two builders' footprint must keep them separate. + let merged = merge_all_tail_partitions(tails, single + 1)?; + assert_eq!(merged.len(), 4); + assert!(merged.iter().all(|builder| builder.docs.len() == 64)); Ok(()) } #[test] fn test_merge_all_tail_partitions_returns_none_for_empty_input() -> Result<()> { - assert!(merge_all_tail_partitions(Vec::new())?.is_none()); + assert!(merge_all_tail_partitions(Vec::new(), u64::MAX)?.is_empty()); + Ok(()) + } + + /// Build an inverted index over `batches` with an explicit worker/memory + /// layout and return it loaded, so tests can compare query behavior + /// across partition shapes of the same corpus. + async fn build_fuzzy_corpus_index( + batches: Vec, + num_workers: usize, + memory_limit_mb: u64, + ) -> Result<(TempDir, Arc)> { + let index_dir = TempDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + index_dir.obj_path(), + Arc::new(LanceCache::no_cache()), + )); + let schema = batches[0].schema(); + let stream = + RecordBatchStreamAdapter::new(schema, stream::iter(batches.into_iter().map(Ok))); + let params = + InvertedIndexParams::new("whitespace".to_string(), lance_tokenizer::Language::English) + .with_position(false) + .remove_stop_words(false) + .stem(false) + .max_token_length(None) + .num_workers(num_workers) + .memory_limit_mb(memory_limit_mb); + let mut builder = InvertedIndexBuilder::new(params); + builder + .update(Box::pin(stream), store.as_ref(), None) + .await?; + let index = InvertedIndex::load(store, None, &LanceCache::no_cache()).await?; + // The caller keeps the TempDir alive for as long as it queries the + // loaded index. + Ok((index_dir, index)) + } + + /// Regression test for tail-partition splitting vs fuzzy queries + /// (https://github.com/lance-format/lance/pull/7601#pullrequestreview): + /// splitting the leftover tails into budget-sized partitions must not + /// change fuzzy results while `max_expansions` is not the binding + /// constraint. Every doc carries one member of two dense fuzzy families + /// plus unique filler tokens, so a small worker memory budget forces + /// flushes and a tail split, while a fuzzy query matches every doc. + #[tokio::test] + async fn test_tail_partition_split_preserves_fuzzy_results() -> Result<()> { + use std::collections::HashMap; + + use crate::prefilter::NoFilter; + use crate::scalar::inverted::document_tokenizer::DocType; + use crate::scalar::inverted::query::{FtsSearchParams, Operator, Tokens}; + + const NUM_DOCS: usize = 4000; + const DOCS_PER_BATCH: usize = 20; + let alpha_variants = ["alpha", "alphb", "alphc", "alphd", "alphe"]; + let beta_variants = ["beta", "betb", "betc", "betd", "bete"]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("doc", DataType::Utf8, true), + Field::new(ROW_ID, DataType::UInt64, false), + ])); + let batches = (0..NUM_DOCS / DOCS_PER_BATCH) + .map(|batch_idx| { + let mut docs = Vec::with_capacity(DOCS_PER_BATCH); + let mut row_ids = Vec::with_capacity(DOCS_PER_BATCH); + for row in 0..DOCS_PER_BATCH { + let i = batch_idx * DOCS_PER_BATCH + row; + // 8 unique filler tokens per doc grow the vocabulary so + // the corpus comfortably exceeds the small worker budget. + docs.push(format!( + "{} {} f{i}a f{i}b f{i}c f{i}d f{i}e f{i}f f{i}g f{i}h", + alpha_variants[i % alpha_variants.len()], + beta_variants[i % beta_variants.len()], + )); + row_ids.push(i as u64); + } + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(docs)), + Arc::new(UInt64Array::from(row_ids)), + ], + ) + .unwrap() + }) + .collect::>(); + + // Reference shape: everything in one partition. + let (_ref_dir, ref_index) = build_fuzzy_corpus_index(batches.clone(), 1, 1024).await?; + assert_eq!( + ref_index.partitions.len(), + 1, + "reference build should stay in a single partition" + ); + + // Tail-heavy shape: a 1MB budget across 2 workers forces flushes and + // splits the leftover tails by the same budget. + let (_split_dir, split_index) = build_fuzzy_corpus_index(batches, 2, 1).await?; + assert!( + split_index.partitions.len() > 1, + "small worker budget should produce multiple partitions, got {}", + split_index.partitions.len() + ); + + async fn fuzzy_search( + index: &InvertedIndex, + tokens: &[&str], + operator: Operator, + ) -> HashMap { + let tokens = Arc::new(Tokens::new( + tokens.iter().map(|t| t.to_string()).collect(), + DocType::Text, + )); + // max_expansions=50 is far above the 10 family variants, so the + // per-partition vs global cap distinction cannot bind here. + let params = Arc::new( + FtsSearchParams::new() + .with_limit(Some(NUM_DOCS)) + .with_fuzziness(Some(1)) + .with_max_expansions(50), + ); + let (row_ids, scores) = index + .bm25_search( + tokens, + params, + operator, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + row_ids.into_iter().zip(scores).collect() + } + + for (tokens, operator) in [ + (vec!["alphx"], Operator::Or), + (vec!["alphx", "betx"], Operator::Or), + (vec!["alphx", "betx"], Operator::And), + ] { + let reference = fuzzy_search(&ref_index, &tokens, operator).await; + let split = fuzzy_search(&split_index, &tokens, operator).await; + assert_eq!( + reference.len(), + NUM_DOCS, + "every doc carries a family variant, {tokens:?} {operator:?} should match all" + ); + assert_eq!( + reference, split, + "fuzzy {operator:?} results must not depend on the tail partition shape for {tokens:?}" + ); + } + Ok(()) } @@ -3839,10 +4123,15 @@ mod tests { let second_doc = second.docs.append(20, 2); second.posting_lists[world as usize].add(second_doc, PositionRecorder::Count(2)); - let merged = merge_tail_partition_group(vec![ - TailPartition { builder: first }, - TailPartition { builder: second }, - ])?; + let merged = merge_all_tail_partitions( + vec![ + TailPartition { builder: first }, + TailPartition { builder: second }, + ], + u64::MAX, + )?; + assert_eq!(merged.len(), 1); + let merged = &merged[0]; assert_eq!(merged.id(), 0); assert_eq!(merged.docs.len(), 2); @@ -4073,7 +4362,10 @@ mod tests { // Remap the index via the ScalarIndex trait method use crate::scalar::ScalarIndex; let mapping = HashMap::from([(0u64, Some(50 << 32))]); - index.remap(&mapping, dest_store.as_ref()).await.unwrap(); + index + .remap(&RowAddrRemap::direct(mapping), dest_store.as_ref()) + .await + .unwrap(); // Reload from dest and verify deleted fragments are preserved let remapped_index = InvertedIndex::load(dest_store.clone(), None, &LanceCache::no_cache()) diff --git a/rust/lance-index/src/scalar/inverted/cache_codec.rs b/rust/lance-index/src/scalar/inverted/cache_codec.rs index a676455d5c9..e59b22dc501 100644 --- a/rust/lance-index/src/scalar/inverted/cache_codec.rs +++ b/rust/lance-index/src/scalar/inverted/cache_codec.rs @@ -14,9 +14,13 @@ //! - the compressed posting list: an IPC section for `blocks`, then the //! position sections (legacy IPC, or shared block-offsets IPC + a raw blob of //! the [`SharedPositionStream`] byte buffer, which has its own portable -//! encoding); +//! encoding), then an optional impact IPC section; //! - the plain posting list: an IPC section of `(row_ids, frequencies)`, then //! an optional legacy position IPC section; +//! - a packed posting-list group: one IPC section containing the original +//! `List` posting rows and optional impact rows; prewarmed groups +//! omit score/length metadata and inject it from the posting reader into +//! query-local views; //! - the standalone [`Positions`] codec: the position sections alone. //! //! All sections read back zero-copy via [`lance_arrow::ipc`]. This is the FTS @@ -39,10 +43,13 @@ use crate::cache_pb::{ PostingTailCodec as PbPostingTailCodec, }; +use super::impact::ImpactSkipData; use super::index::{ CompressedPositionStorage, CompressedPostingList, PlainPostingList, PositionStreamCodec, - Positions, PostingList, PostingListGroup, PostingTailCodec, SharedPositionStream, + Positions, PostingList, PostingListGroup, PostingListGroupStorage, PostingTailCodec, + SharedPositionStream, }; +use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; // --------------------------------------------------------------------------- // Tags @@ -50,6 +57,8 @@ use super::index::{ const POSTING_VARIANT_PLAIN: u8 = 0; const POSTING_VARIANT_COMPRESSED: u8 = 1; +const GROUP_VARIANT_MATERIALIZED: u8 = 0; +const GROUP_VARIANT_PACKED: u8 = 1; // --------------------------------------------------------------------------- // Codec enum mappings @@ -72,6 +81,23 @@ fn proto_to_posting_tail_codec(c: PbPostingTailCodec) -> PostingTailCodec { } } +fn posting_tail_codec_to_tag(c: PostingTailCodec) -> u8 { + match c { + PostingTailCodec::Fixed32 => 0, + PostingTailCodec::VarintDelta => 1, + } +} + +fn posting_tail_codec_from_tag(tag: u8) -> Result { + match tag { + 0 => Ok(PostingTailCodec::Fixed32), + 1 => Ok(PostingTailCodec::VarintDelta), + other => Err(Error::io(format!( + "unknown packed posting tail codec: {other}" + ))), + } +} + fn position_stream_codec_to_proto(c: PositionStreamCodec) -> PbPositionStreamCodec { match c { PositionStreamCodec::VarintDocDelta => PbPositionStreamCodec::VarintDocDelta, @@ -95,6 +121,7 @@ const BLOCK_OFFSETS_COLUMN: &str = "block_offsets"; const ROW_IDS_COLUMN: &str = "row_ids"; const FREQUENCIES_COLUMN: &str = "frequencies"; const BLOCKS_COLUMN: &str = "blocks"; +const IMPACTS_COLUMN: &str = "impacts"; fn legacy_positions_batch(list: &ListArray) -> Result { let schema = Arc::new(Schema::new(vec![Field::new( @@ -108,7 +135,8 @@ fn legacy_positions_batch(list: &ListArray) -> Result { fn read_legacy_positions(r: &mut CacheEntryReader<'_>) -> Result { let batch = r.read_ipc()?; Ok(batch - .column(0) + .column_by_name(POSITION_LIST_COLUMN) + .ok_or_else(|| Error::io("legacy position column is missing".to_string()))? .as_any() .downcast_ref::() .ok_or_else(|| Error::io("legacy position column is not a ListArray".to_string()))? @@ -156,7 +184,8 @@ fn read_position_sections( PbPositionStorage::Shared => { let batch = r.read_ipc()?; let block_offsets = batch - .column(0) + .column_by_name(BLOCK_OFFSETS_COLUMN) + .ok_or_else(|| Error::io("block_offsets column is missing".to_string()))? .as_primitive_opt::() .ok_or_else(|| Error::io("block_offsets column is not UInt32".to_string()))? .values() @@ -178,7 +207,11 @@ fn read_position_sections( impl CacheCodecImpl for PostingList { const TYPE_ID: &'static str = "lance.fts.PostingList"; - const CURRENT_VERSION: u32 = 1; + // Version 3 adds the optional impact IPC section. Main already used v2 for + // configurable posting block sizes, so impact data needs a distinct + // version to keep older readers from accepting a body with an extra + // section they cannot consume. + const CURRENT_VERSION: u32 = 3; fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { match self { @@ -194,15 +227,24 @@ impl CacheCodecImpl for PostingList { } fn deserialize(r: &mut CacheEntryReader<'_>) -> Result { - let variant = r.read_u8()?; - match variant { - POSTING_VARIANT_PLAIN => Ok(Self::Plain(deserialize_plain(r)?)), - POSTING_VARIANT_COMPRESSED => Ok(Self::Compressed(deserialize_compressed(r)?)), - other => Err(Error::io(format!("unknown PostingList variant: {other}"))), + match r.version() { + 1 | 2 | Self::CURRENT_VERSION => deserialize_posting_list_body(r), + other => Err(Error::io(format!( + "unsupported PostingList cache version: {other}" + ))), } } } +fn deserialize_posting_list_body(r: &mut CacheEntryReader<'_>) -> Result { + let variant = r.read_u8()?; + match variant { + POSTING_VARIANT_PLAIN => Ok(PostingList::Plain(deserialize_plain(r)?)), + POSTING_VARIANT_COMPRESSED => Ok(PostingList::Compressed(deserialize_compressed(r)?)), + other => Err(Error::io(format!("unknown PostingList variant: {other}"))), + } +} + fn serialize_plain(w: &mut CacheEntryWriter<'_>, plain: &PlainPostingList) -> Result<()> { // Plain postings carry only per-doc legacy positions (or none). let position_storage = if plain.positions.is_some() { @@ -236,13 +278,15 @@ fn deserialize_plain(r: &mut CacheEntryReader<'_>) -> Result { let batch = r.read_ipc()?; let row_ids = batch - .column(0) + .column_by_name(ROW_IDS_COLUMN) + .ok_or_else(|| Error::io("row_ids column is missing".to_string()))? .as_primitive_opt::() .ok_or_else(|| Error::io("row_ids column is not UInt64".to_string()))? .values() .clone(); let frequencies = batch - .column(1) + .column_by_name(FREQUENCIES_COLUMN) + .ok_or_else(|| Error::io("frequencies column is missing".to_string()))? .as_primitive_opt::() .ok_or_else(|| Error::io("frequencies column is not Float32".to_string()))? .values() @@ -291,6 +335,8 @@ fn serialize_compressed( posting_tail_codec: posting_tail_codec_to_proto(posting.posting_tail_codec) as i32, position_storage: position_storage as i32, position_stream_codec: position_stream_codec as i32, + block_size: posting.block_size as u32, + has_impacts: posting.impacts.is_some(), }; w.write_header(&header)?; @@ -305,6 +351,15 @@ fn serialize_compressed( if let Some(storage) = &posting.positions { write_position_sections(w, storage)?; } + if let Some(impacts) = &posting.impacts { + let schema = Arc::new(Schema::new(vec![Field::new( + IMPACTS_COLUMN, + DataType::LargeBinary, + false, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(impacts.entries().clone())])?; + w.write_ipc(&batch)?; + } Ok(()) } @@ -314,7 +369,8 @@ fn deserialize_compressed(r: &mut CacheEntryReader<'_>) -> Result() .ok_or_else(|| Error::io("blocks column is not a LargeBinaryArray".to_string()))? @@ -322,13 +378,33 @@ fn deserialize_compressed(r: &mut CacheEntryReader<'_>) -> Result= 3 && header.has_impacts { + let batch = r.read_ipc()?; + let entries = batch + .column_by_name(IMPACTS_COLUMN) + .ok_or_else(|| Error::io("impacts column is missing".to_string()))? + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::io("impacts column is not a LargeBinaryArray".to_string()))? + .clone(); + Some(ImpactSkipData::new(entries, blocks.len())?) + } else { + None + }; Ok(CompressedPostingList::new( blocks, header.max_score, header.length, posting_tail_codec, + block_size, positions, + impacts, )) } @@ -336,34 +412,79 @@ fn deserialize_compressed(r: &mut CacheEntryReader<'_>) -> Result) -> Result<()> { - let count = u32::try_from(self.posting_lists.len()) + let count = u32::try_from(self.len()) .map_err(|_| Error::io("posting list group too large to serialize".to_string()))?; - w.write_header(&PostingListGroupHeader { count })?; - for posting in &self.posting_lists { - posting.serialize(w)?; + match &self.storage { + PostingListGroupStorage::Materialized(posting_lists) => { + w.write_u8(GROUP_VARIANT_MATERIALIZED)?; + w.write_header(&PostingListGroupHeader { count })?; + for posting in posting_lists { + posting.serialize(w)?; + } + } + PostingListGroupStorage::Packed(group) => { + w.write_u8(GROUP_VARIANT_PACKED)?; + w.write_header(&PostingListGroupHeader { count })?; + w.write_u8(posting_tail_codec_to_tag(group.posting_tail_codec))?; + w.write_ipc(&group.batch)?; + } } Ok(()) } fn deserialize(r: &mut CacheEntryReader<'_>) -> Result { - let header: PostingListGroupHeader = r.read_header()?; - let mut posting_lists = Vec::with_capacity(header.count as usize); - for _ in 0..header.count { - posting_lists.push(PostingList::deserialize(r)?); + match r.version() { + 1 => return deserialize_materialized_group(r), + 2 | 3 | Self::CURRENT_VERSION => {} + other => { + return Err(Error::io(format!( + "unsupported PostingListGroup cache version: {other}" + ))); + } } - Ok(Self::new(posting_lists)) + + let variant = r.read_u8()?; + match variant { + GROUP_VARIANT_MATERIALIZED => deserialize_materialized_group(r), + GROUP_VARIANT_PACKED => { + let header: PostingListGroupHeader = r.read_header()?; + let posting_tail_codec = posting_tail_codec_from_tag(r.read_u8()?)?; + let batch = r.read_ipc()?; + if batch.num_rows() != header.count as usize { + return Err(Error::io(format!( + "packed posting group row count {} does not match header count {}", + batch.num_rows(), + header.count + ))); + } + Self::new_packed(batch, posting_tail_codec) + } + other => Err(Error::io(format!( + "unknown PostingListGroup variant: {other}" + ))), + } + } +} + +fn deserialize_materialized_group(r: &mut CacheEntryReader<'_>) -> Result { + let header: PostingListGroupHeader = r.read_header()?; + let mut posting_lists = Vec::with_capacity(header.count as usize); + for _ in 0..header.count { + posting_lists.push(deserialize_posting_list_body(r)?); } + Ok(PostingListGroup::new(posting_lists)) } // --------------------------------------------------------------------------- @@ -409,17 +530,26 @@ impl CacheCodecImpl for Positions { #[cfg(test)] mod tests { + use std::collections::HashMap; + use std::sync::Arc; + use arrow::buffer::ScalarBuffer; - use arrow_array::LargeBinaryArray; - use arrow_array::builder::{Int32Builder, ListBuilder}; + use arrow_array::builder::{Int32Builder, LargeBinaryBuilder, ListBuilder}; + use arrow_array::{Array, LargeBinaryArray, RecordBatch}; + use arrow_schema::{Field, Schema}; use bytes::Bytes; use lance_core::Result; use lance_core::cache::{CacheCodecImpl, CacheEntryReader, CacheEntryWriter}; + use crate::cache_pb::{CompressedPostingHeader, PostingTailCodec as PbPostingTailCodec}; + + use super::super::impact::{ImpactSkipData, ImpactSkipDataBuilder}; use super::super::index::{ - CompressedPositionStorage, CompressedPostingList, PlainPostingList, PositionStreamCodec, - Positions, PostingList, PostingListGroup, PostingTailCodec, SharedPositionStream, + CompressedPositionStorage, CompressedPostingList, IMPACT_COL, POSTING_BLOCK_SIZE_KEY, + POSTING_COL, PlainPostingList, PositionStreamCodec, Positions, PostingList, + PostingListGroup, PostingTailCodec, SharedPositionStream, }; + use super::super::tokenizer::LEGACY_BLOCK_SIZE; fn legacy_positions(rows: &[&[i32]]) -> arrow_array::ListArray { let mut builder = ListBuilder::new(Int32Builder::new()); @@ -432,6 +562,73 @@ mod tests { builder.finish() } + fn packed_batch(postings: &[Vec>], block_size: Option) -> RecordBatch { + let mut builder = ListBuilder::new(LargeBinaryBuilder::new()); + for posting in postings { + for block in posting { + builder.values().append_value(block); + } + builder.append(true); + } + let postings = builder.finish(); + let fields = vec![Field::new(POSTING_COL, postings.data_type().clone(), false)]; + let schema = Arc::new(match block_size { + Some(block_size) => Schema::new_with_metadata( + fields, + HashMap::from([(POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string())]), + ), + None => Schema::new(fields), + }); + RecordBatch::try_new(schema, vec![Arc::new(postings)]).unwrap() + } + + fn packed_group( + postings: &[Vec>], + posting_tail_codec: PostingTailCodec, + block_size: Option, + ) -> PostingListGroup { + PostingListGroup::new_packed(packed_batch(postings, block_size), posting_tail_codec) + .unwrap() + } + + fn packed_group_with_impacts( + postings: &[Vec>], + impacts: &[ImpactSkipData], + posting_tail_codec: PostingTailCodec, + block_size: usize, + ) -> PostingListGroup { + assert_eq!(postings.len(), impacts.len()); + let posting_batch = packed_batch(postings, Some(block_size)); + let mut impacts_builder = ListBuilder::new(LargeBinaryBuilder::new()); + for impacts in impacts { + for entry_idx in 0..impacts.entries().len() { + impacts_builder + .values() + .append_value(impacts.entries().value(entry_idx)); + } + impacts_builder.append(true); + } + let impacts = impacts_builder.finish(); + let fields = vec![ + Field::new( + POSTING_COL, + posting_batch.column(0).data_type().clone(), + false, + ), + Field::new(IMPACT_COL, impacts.data_type().clone(), false), + ]; + let schema = Arc::new(Schema::new_with_metadata( + fields, + posting_batch.schema_ref().metadata().clone(), + )); + let batch = RecordBatch::try_new( + schema, + vec![posting_batch.column(0).clone(), Arc::new(impacts)], + ) + .unwrap(); + PostingListGroup::new_packed(batch, posting_tail_codec).unwrap() + } + fn assert_plain_eq(a: &PlainPostingList, b: &PlainPostingList) { assert_eq!(a.row_ids.as_ref(), b.row_ids.as_ref()); assert_eq!(a.frequencies.as_ref(), b.frequencies.as_ref()); @@ -461,6 +658,20 @@ mod tests { } } + fn impact_skip_data(level0_len: usize, block_size: usize) -> ImpactSkipData { + let mut builder = ImpactSkipDataBuilder::with_capacity(level0_len, block_size); + for block_idx in 0..level0_len { + let doc_base = block_idx as u32 * 10; + builder + .append_block(&[ + (doc_base + 1, block_idx as u32 + 1, 10), + (doc_base + 9, block_idx as u32 + 2, 8), + ]) + .unwrap(); + } + builder.finish().unwrap() + } + /// Serialize a codec body (no envelope) into a standalone buffer. fn body_bytes(entry: &T) -> Bytes { let mut buf = Vec::new(); @@ -475,6 +686,34 @@ mod tests { T::deserialize(&mut r) } + fn from_body_version(data: &Bytes, version: u32) -> Result { + let mut r = CacheEntryReader::new(data, 0, version); + T::deserialize(&mut r) + } + + fn compressed_body_with_ipc_sections( + blocks: &RecordBatch, + impacts: Option<&RecordBatch>, + ) -> Bytes { + let mut buf = Vec::new(); + let mut w = CacheEntryWriter::new(&mut buf); + w.write_u8(super::POSTING_VARIANT_COMPRESSED).unwrap(); + w.write_header(&CompressedPostingHeader { + max_score: 1.0, + length: 1, + posting_tail_codec: PbPostingTailCodec::VarintDelta as i32, + block_size: 256, + has_impacts: impacts.is_some(), + ..Default::default() + }) + .unwrap(); + w.write_ipc(blocks).unwrap(); + if let Some(impacts) = impacts { + w.write_ipc(impacts).unwrap(); + } + Bytes::from(buf) + } + fn roundtrip_posting_list(entry: &PostingList) -> PostingList { from_body::(&body_bytes(entry)).unwrap() } @@ -532,14 +771,22 @@ mod tests { Some(&[1u8, 2, 3, 4, 5][..]), Some(&[6, 7, 8, 9, 10][..]), ]); - let posting = - CompressedPostingList::new(blocks, 3.5, 42, PostingTailCodec::VarintDelta, None); + let posting = CompressedPostingList::new( + blocks, + 3.5, + 42, + PostingTailCodec::VarintDelta, + 256, + None, + None, + ); let entry = PostingList::Compressed(posting.clone()); match roundtrip_posting_list(&entry) { PostingList::Compressed(restored) => { assert_eq!(restored.max_score, posting.max_score); assert_eq!(restored.length, posting.length); assert_eq!(restored.posting_tail_codec, posting.posting_tail_codec); + assert_eq!(restored.block_size, posting.block_size); assert_eq!(restored.blocks, posting.blocks); assert!(restored.positions.is_none()); } @@ -547,6 +794,73 @@ mod tests { } } + #[test] + fn compressed_posting_list_impacts_roundtrip() { + let blocks = LargeBinaryArray::from_opt_vec(vec![ + Some(&[1u8, 2, 3, 4, 5][..]), + Some(&[6, 7, 8, 9, 10][..]), + ]); + let impacts = impact_skip_data(blocks.len(), 256); + let posting = CompressedPostingList::new( + blocks, + 3.5, + 42, + PostingTailCodec::VarintDelta, + 256, + None, + Some(impacts.clone()), + ); + let entry = PostingList::Compressed(posting); + match roundtrip_posting_list(&entry) { + PostingList::Compressed(restored) => { + let restored = restored.impacts.expect("impacts should roundtrip"); + assert_eq!(restored.level0_len(), impacts.level0_len()); + assert_eq!(restored.level1_len(), impacts.level1_len()); + assert_eq!(restored.entries(), impacts.entries()); + } + PostingList::Plain(_) => panic!("expected Compressed variant"), + } + } + + #[test] + fn compressed_posting_list_missing_ipc_columns_returns_error() { + let empty = RecordBatch::new_empty(Arc::new(Schema::empty())); + assert!( + from_body::(&compressed_body_with_ipc_sections(&empty, None)).is_err() + ); + + let blocks = LargeBinaryArray::from_opt_vec(vec![Some(&[1_u8, 2, 3][..])]); + let schema = Arc::new(Schema::new(vec![Field::new( + super::BLOCKS_COLUMN, + blocks.data_type().clone(), + false, + )])); + let blocks = RecordBatch::try_new(schema, vec![Arc::new(blocks)]).unwrap(); + assert!( + from_body::(&compressed_body_with_ipc_sections(&blocks, Some(&empty))) + .is_err() + ); + } + + #[test] + fn compressed_posting_list_v1_cache_without_impacts_decodes() { + let posting = CompressedPostingList::new( + LargeBinaryArray::from_opt_vec(vec![Some(&[1u8, 2, 3][..])]), + 1.25, + 5, + PostingTailCodec::Fixed32, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, + None, + None, + ); + let data = body_bytes(&PostingList::Compressed(posting)); + let restored = from_body_version::(&data, 1).unwrap(); + let PostingList::Compressed(restored) = restored else { + panic!("expected Compressed variant"); + }; + assert!(restored.impacts.is_none()); + } + #[test] fn compressed_posting_list_legacy_positions_roundtrip() { let blocks = LargeBinaryArray::from_opt_vec(vec![Some(&[1u8, 2, 3][..])]); @@ -555,9 +869,11 @@ mod tests { 1.25, 5, PostingTailCodec::Fixed32, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, Some(CompressedPositionStorage::LegacyPerDoc(legacy_positions( &[&[0, 4, 8]], ))), + None, ); let entry = PostingList::Compressed(posting.clone()); match roundtrip_posting_list(&entry) { @@ -589,7 +905,9 @@ mod tests { 7.0, 3, PostingTailCodec::VarintDelta, + 256, Some(CompressedPositionStorage::SharedStream(stream)), + None, ); let entry = PostingList::Compressed(posting.clone()); match roundtrip_posting_list(&entry) { @@ -617,9 +935,11 @@ mod tests { 7.0, 3, PostingTailCodec::VarintDelta, + 256, Some(CompressedPositionStorage::SharedStream( expected_stream.clone(), )), + None, ); let serialized = body_bytes(&PostingList::Compressed(posting)); @@ -651,6 +971,8 @@ mod tests { 2.5, 7, PostingTailCodec::VarintDelta, + 256, + None, None, )); @@ -661,9 +983,11 @@ mod tests { ] { let group = PostingListGroup::new(members.clone()); let restored = from_body::(&body_bytes(&group)).unwrap(); - assert_eq!(restored.posting_lists.len(), members.len()); - for (a, b) in members.iter().zip(restored.posting_lists.iter()) { - match (a, b) { + assert!(!restored.is_packed()); + assert_eq!(restored.len(), members.len()); + for (index, a) in members.iter().enumerate() { + let b = restored.posting_list(index, None, None).unwrap().unwrap(); + match (a, &b) { (PostingList::Plain(x), PostingList::Plain(y)) => assert_plain_eq(x, y), (PostingList::Compressed(x), PostingList::Compressed(y)) => { assert_eq!(x.blocks, y.blocks); @@ -676,6 +1000,167 @@ mod tests { } } + #[test] + fn packed_posting_list_group_roundtrip_and_v1_fallback() { + let group = packed_group( + &[vec![vec![1, 2, 3], vec![4, 5]], vec![vec![7; 16 * 1024]]], + PostingTailCodec::VarintDelta, + Some(256), + ); + let restored = from_body::(&body_bytes(&group)).unwrap(); + assert!(restored.is_packed()); + assert_eq!(restored.len(), 2); + for slot in 0..2 { + let max_score = [1.5, 3.25][slot]; + let length = [3, 4096][slot]; + let expected = group + .posting_list(slot, Some(max_score), Some(length)) + .unwrap() + .unwrap(); + let actual = restored + .posting_list(slot, Some(max_score), Some(length)) + .unwrap() + .unwrap(); + let (PostingList::Compressed(expected), PostingList::Compressed(actual)) = + (expected, actual) + else { + panic!("expected compressed packed posting views"); + }; + assert_eq!(actual.blocks, expected.blocks); + assert_eq!(actual.max_score, expected.max_score); + assert_eq!(actual.length, expected.length); + assert_eq!(actual.posting_tail_codec, expected.posting_tail_codec); + assert_eq!(actual.block_size, 256); + } + + let legacy_packed = + packed_group(&[vec![vec![9, 8, 7]]], PostingTailCodec::VarintDelta, None); + let restored = from_body::(&body_bytes(&legacy_packed)).unwrap(); + let PostingList::Compressed(posting) = restored + .posting_list(0, Some(2.0), Some(3)) + .unwrap() + .unwrap() + else { + panic!("expected compressed legacy packed posting"); + }; + assert_eq!(posting.block_size, LEGACY_BLOCK_SIZE); + + let legacy_member = PostingList::Compressed(CompressedPostingList::new( + LargeBinaryArray::from_opt_vec(vec![Some(&[9u8, 8, 7][..])]), + 2.0, + 3, + PostingTailCodec::VarintDelta, + LEGACY_BLOCK_SIZE, + None, + None, + )); + let mut legacy_body = Vec::new(); + let mut writer = CacheEntryWriter::new(&mut legacy_body); + writer + .write_header(&crate::cache_pb::PostingListGroupHeader { count: 1 }) + .unwrap(); + legacy_member.serialize(&mut writer).unwrap(); + let legacy_body = Bytes::from(legacy_body); + let mut reader = CacheEntryReader::new(&legacy_body, 0, 1); + let restored = PostingListGroup::deserialize(&mut reader).unwrap(); + assert!(!restored.is_packed()); + assert_eq!(restored.len(), 1); + } + + #[test] + fn posting_list_group_impacted_compressed_members_roundtrip() { + let first = CompressedPostingList::new( + LargeBinaryArray::from_opt_vec(vec![Some(&[1u8, 2, 3][..]), Some(&[4u8, 5, 6][..])]), + 3.0, + 256, + PostingTailCodec::VarintDelta, + LEGACY_BLOCK_SIZE, + None, + Some(impact_skip_data(2, LEGACY_BLOCK_SIZE)), + ); + let second = CompressedPostingList::new( + LargeBinaryArray::from_opt_vec(vec![Some(&[7u8, 8, 9][..])]), + 5.0, + 128, + PostingTailCodec::Fixed32, + 256, + Some(CompressedPositionStorage::SharedStream( + SharedPositionStream::new( + PositionStreamCodec::PackedDelta, + vec![0u32, 12], + Bytes::from(vec![0xABu8; 32]), + ), + )), + Some(impact_skip_data(1, 256)), + ); + let members = vec![ + PostingList::Compressed(first.clone()), + PostingList::Compressed(second.clone()), + ]; + let group = PostingListGroup::new(members); + let restored = from_body::(&body_bytes(&group)).unwrap(); + assert!(!restored.is_packed()); + assert_eq!(restored.len(), 2); + + let expected = [&first, &second]; + for (slot, expected) in expected.iter().enumerate() { + let restored = restored.posting_list(slot, None, None).unwrap().unwrap(); + let PostingList::Compressed(restored) = restored else { + panic!("expected compressed member"); + }; + assert_eq!(restored.blocks, expected.blocks); + assert_eq!(restored.length, expected.length); + assert_eq!(restored.max_score, expected.max_score); + assert_eq!(restored.posting_tail_codec, expected.posting_tail_codec); + assert_eq!(restored.block_size, expected.block_size); + assert_eq!( + restored.impacts.as_ref().unwrap().entries(), + expected.impacts.as_ref().unwrap().entries() + ); + match (&expected.positions, &restored.positions) { + (Some(expected), Some(restored)) => { + assert_position_storage_eq(expected, restored); + } + (None, None) => {} + _ => panic!("position storage mismatch"), + } + } + } + + #[test] + fn packed_posting_list_group_impacts_roundtrip() { + let postings = vec![vec![vec![1, 2, 3], vec![4, 5, 6]], vec![vec![7, 8, 9]]]; + let expected_impacts = vec![impact_skip_data(2, 256), impact_skip_data(1, 256)]; + let group = packed_group_with_impacts( + &postings, + &expected_impacts, + PostingTailCodec::VarintDelta, + 256, + ); + + let restored = from_body::(&body_bytes(&group)).unwrap(); + assert!(restored.is_packed()); + assert_eq!(restored.len(), expected_impacts.len()); + for (slot, expected) in expected_impacts.iter().enumerate() { + let posting = restored + .posting_list(slot, Some(3.0), Some(256)) + .unwrap() + .unwrap(); + let PostingList::Compressed(posting) = posting else { + panic!("expected compressed packed posting"); + }; + let actual = posting.impacts.as_ref().expect("impacts should roundtrip"); + assert_eq!(actual.entries(), expected.entries()); + assert_eq!(actual.level0_len(), expected.level0_len()); + assert_eq!( + actual.level1_doc_up_to(0), + expected.level1_doc_up_to(0), + "impact entries should remain decodable with the packed block size", + ); + assert!(actual.level1_doc_up_to(0).is_some()); + } + } + #[test] fn positions_legacy_roundtrip() { let positions = Positions(CompressedPositionStorage::LegacyPerDoc(legacy_positions( @@ -717,26 +1202,130 @@ mod tests { use std::sync::Arc; use arrow_array::Array; - use lance_core::cache::CacheCodec; + use arrow_schema::DataType; + use lance_core::cache::{ + CacheCodec, CacheCodecImpl, CacheDecode, CacheEntryReader, CacheEntryWriter, + CacheMissReason, + }; + use lance_core::{Error, Result}; use prost::Message; + use super::super::{ + BLOCKS_COLUMN, GROUP_VARIANT_PACKED, POSTING_VARIANT_COMPRESSED, + posting_tail_codec_to_tag, + }; use super::*; - use crate::cache_pb::{CompressedPostingHeader, PostingTailCodec as PbPostingTailCodec}; + use crate::cache_pb::{ + CompressedPostingHeader, PostingListGroupHeader, PostingTailCodec as PbPostingTailCodec, + }; type ArcAny = Arc; + struct PostingListV2Codec(PostingList); + + impl CacheCodecImpl for PostingListV2Codec { + const TYPE_ID: &'static str = ::TYPE_ID; + const CURRENT_VERSION: u32 = 2; + + fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { + self.0.serialize(w) + } + + fn deserialize(r: &mut CacheEntryReader<'_>) -> Result { + PostingList::deserialize(r).map(Self) + } + } + + struct PostingListGroupV3Codec(PostingListGroup); + + impl CacheCodecImpl for PostingListGroupV3Codec { + const TYPE_ID: &'static str = ::TYPE_ID; + const CURRENT_VERSION: u32 = 3; + + fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { + self.0.serialize(w) + } + + fn deserialize(r: &mut CacheEntryReader<'_>) -> Result { + PostingListGroup::deserialize(r).map(Self) + } + } + + struct LegacyCompressedPostingV1 { + blocks: LargeBinaryArray, + } + + impl CacheCodecImpl for LegacyCompressedPostingV1 { + const TYPE_ID: &'static str = ::TYPE_ID; + const CURRENT_VERSION: u32 = 1; + + fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { + w.write_u8(POSTING_VARIANT_COMPRESSED)?; + w.write_header(&CompressedPostingHeader { + max_score: 2.0, + length: 3, + posting_tail_codec: PbPostingTailCodec::VarintDelta as i32, + ..Default::default() + })?; + let schema = Arc::new(Schema::new(vec![Field::new( + BLOCKS_COLUMN, + DataType::LargeBinary, + false, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(self.blocks.clone())])?; + w.write_ipc(&batch) + } + + fn deserialize(_r: &mut CacheEntryReader<'_>) -> Result { + Err(Error::io( + "LegacyCompressedPostingV1 is a writer-only test codec".to_string(), + )) + } + } + + struct LegacyPackedGroupV2 { + batch: RecordBatch, + posting_tail_codec: PostingTailCodec, + } + + impl CacheCodecImpl for LegacyPackedGroupV2 { + const TYPE_ID: &'static str = ::TYPE_ID; + const CURRENT_VERSION: u32 = 2; + + fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { + w.write_u8(GROUP_VARIANT_PACKED)?; + let count = u32::try_from(self.batch.num_rows()) + .map_err(|_| Error::io("legacy packed group is too large".to_string()))?; + w.write_header(&PostingListGroupHeader { count })?; + w.write_u8(posting_tail_codec_to_tag(self.posting_tail_codec))?; + w.write_ipc(&self.batch) + } + + fn deserialize(_r: &mut CacheEntryReader<'_>) -> Result { + Err(Error::io( + "LegacyPackedGroupV2 is a writer-only test codec".to_string(), + )) + } + } + fn codec() -> CacheCodec { CacheCodec::from_impl::() } - /// Serialize an entry through the full codec (envelope + body). - fn serialize_entry(entry: PostingList) -> Vec { + fn serialize_typed_entry(entry: T) -> Vec { let any: ArcAny = Arc::new(entry); let mut buf = Vec::new(); - codec().serialize(&any, &mut buf).unwrap(); + CacheCodec::from_impl::() + .serialize(&any, &mut buf) + .unwrap(); buf } + /// Serialize an entry through the full codec (envelope + body). + fn serialize_entry(entry: PostingList) -> Vec { + serialize_typed_entry(entry) + } + /// A `Bytes` whose base address is 64-byte aligned, modelling a backend /// that reads cache entries into an aligned buffer. fn aligned_bytes(payload: &[u8]) -> Bytes { @@ -760,7 +1349,9 @@ mod tests { 7.0, 3, PostingTailCodec::VarintDelta, + 256, Some(CompressedPositionStorage::SharedStream(stream)), + None, )) } @@ -795,25 +1386,17 @@ mod tests { assert!(points_in(stream.bytes().as_ptr() as usize)); } - /// Every member of a `PostingListGroup` must also decode zero-copy. The - /// group writes its members inline so each member's IPC sections stay - /// 64-byte aligned within the entry; embedding members in per-member - /// sub-buffers would land them at arbitrary offsets and force a - /// realigning memcpy on load. + /// A packed group's single IPC batch and all posting views decoded from + /// it must borrow the cache entry's aligned input buffer. #[test] - fn group_member_sections_are_zero_copy_through_envelope() { - let make_member = |fill: u8| { - let blocks = - LargeBinaryArray::from_opt_vec(vec![Some(&[fill; 48][..]), Some(&[fill; 48])]); - PostingList::Compressed(CompressedPostingList::new( - blocks, - 7.0, - 3, - PostingTailCodec::VarintDelta, - None, - )) - }; - let group = PostingListGroup::new(vec![make_member(9), make_member(1)]); + fn packed_group_sections_are_zero_copy_through_envelope() { + let postings = vec![ + vec![vec![9; 48], vec![9; 48]], + vec![vec![1; 48], vec![1; 48]], + ]; + let impacts = vec![impact_skip_data(2, 256), impact_skip_data(2, 256)]; + let group = + packed_group_with_impacts(&postings, &impacts, PostingTailCodec::VarintDelta, 256); let group_codec = CacheCodec::from_impl::(); let any: ArcAny = Arc::new(group); @@ -828,8 +1411,13 @@ mod tests { let end = base + serialized.len(); let points_in = |ptr: usize| ptr >= base && ptr < end; - assert_eq!(restored.posting_lists.len(), 2); - for member in &restored.posting_lists { + assert!(restored.is_packed()); + assert_eq!(restored.len(), 2); + for slot in 0..restored.len() { + let member = restored + .posting_list(slot, Some(7.0), Some(3)) + .unwrap() + .unwrap(); let PostingList::Compressed(member) = member else { panic!("expected Compressed member"); }; @@ -837,7 +1425,17 @@ mod tests { assert!( points_in(buf.as_ptr() as usize), "group member blocks buffer was realigned out of the input — \ - misaligned IPC section", + misaligned IPC section", + ); + } + let impacts = member + .impacts + .as_ref() + .expect("packed impacts should decode"); + for buf in impacts.entries().to_data().buffers() { + assert!( + points_in(buf.as_ptr() as usize), + "group member impact buffer was realigned out of the input", ); } } @@ -915,6 +1513,110 @@ mod tests { assert!(codec().deserialize(&Bytes::from(buf)).hit().is_none()); } + #[test] + fn old_codecs_reject_new_impact_envelopes_as_version_too_new() { + let posting = Bytes::from(serialize_entry(compressed_with_shared_positions())); + match CacheCodec::from_impl::().deserialize(&posting) { + CacheDecode::Miss(reason) => { + assert_eq!(reason, CacheMissReason::VersionTooNew) + } + CacheDecode::Hit(_) => panic!("v2 PostingList codec accepted a v3 envelope"), + } + + let group = packed_group( + &[vec![vec![1, 2, 3]], vec![vec![4, 5, 6]]], + PostingTailCodec::VarintDelta, + Some(256), + ); + let group = Bytes::from(serialize_typed_entry(group)); + match CacheCodec::from_impl::().deserialize(&group) { + CacheDecode::Miss(reason) => { + assert_eq!(reason, CacheMissReason::VersionTooNew) + } + CacheDecode::Hit(_) => { + panic!("v3 PostingListGroup codec accepted a v4 envelope") + } + } + } + + #[test] + fn current_codecs_read_previous_main_versions() { + let previous_posting = PostingListV2Codec(compressed_with_shared_positions()); + let previous_posting = Bytes::from(serialize_typed_entry(previous_posting)); + let restored = codec().deserialize(&previous_posting).hit().unwrap(); + let restored = restored.downcast::().unwrap(); + let PostingList::Compressed(restored) = restored.as_ref() else { + panic!("expected compressed posting"); + }; + assert_eq!(restored.block_size, 256); + assert!(restored.impacts.is_none()); + + let previous_group = PostingListGroupV3Codec(packed_group( + &[vec![vec![1, 2, 3]], vec![vec![4, 5, 6]]], + PostingTailCodec::VarintDelta, + Some(256), + )); + let previous_group = Bytes::from(serialize_typed_entry(previous_group)); + let restored = CacheCodec::from_impl::() + .deserialize(&previous_group) + .hit() + .unwrap() + .downcast::() + .unwrap(); + assert!(restored.is_packed()); + assert_eq!(restored.len(), 2); + let PostingList::Compressed(restored) = restored + .posting_list(0, Some(2.0), Some(3)) + .unwrap() + .unwrap() + else { + panic!("expected compressed packed posting"); + }; + assert_eq!(restored.block_size, 256); + assert!(restored.impacts.is_none()); + } + + #[test] + fn current_codecs_read_legacy_payloads_without_block_size() { + let legacy_posting = LegacyCompressedPostingV1 { + blocks: LargeBinaryArray::from_opt_vec(vec![Some(&[9u8, 8, 7][..])]), + }; + let legacy_posting = Bytes::from(serialize_typed_entry(legacy_posting)); + let restored = codec().deserialize(&legacy_posting).hit().unwrap(); + let restored = restored.downcast::().unwrap(); + let PostingList::Compressed(restored) = restored.as_ref() else { + panic!("expected a compressed legacy posting"); + }; + assert_eq!(restored.block_size, LEGACY_BLOCK_SIZE); + + let legacy_batch = packed_batch(&[vec![vec![1, 2, 3]], vec![vec![4, 5, 6]]], None); + assert!( + !legacy_batch + .schema_ref() + .metadata() + .contains_key(POSTING_BLOCK_SIZE_KEY) + ); + let legacy_group = LegacyPackedGroupV2 { + batch: legacy_batch, + posting_tail_codec: PostingTailCodec::VarintDelta, + }; + let legacy_group = Bytes::from(serialize_typed_entry(legacy_group)); + let restored = CacheCodec::from_impl::() + .deserialize(&legacy_group) + .hit() + .unwrap() + .downcast::() + .unwrap(); + let PostingList::Compressed(restored) = restored + .posting_list(0, Some(2.0), Some(3)) + .unwrap() + .unwrap() + else { + panic!("expected a compressed legacy packed posting"); + }; + assert_eq!(restored.block_size, LEGACY_BLOCK_SIZE); + } + /// A pre-stabilization blob (no magic) self-heals to a miss. #[test] fn pre_stabilization_blob_is_miss() { diff --git a/rust/lance-index/src/scalar/inverted/encoding.rs b/rust/lance-index/src/scalar/inverted/encoding.rs index fabc4ac9261..8fb7f8f4032 100644 --- a/rust/lance-index/src/scalar/inverted/encoding.rs +++ b/rust/lance-index/src/scalar/inverted/encoding.rs @@ -5,10 +5,15 @@ use std::io::Write; use super::builder::BLOCK_SIZE; use super::index::{PositionStreamCodec, PostingTailCodec}; +#[cfg(test)] +use super::tokenizer::LEGACY_BLOCK_SIZE; +use super::tokenizer::validate_block_size; use arrow::array::LargeBinaryBuilder; -use bitpacking::{BitPacker, BitPacker4x}; +use lance_bitpacking::{BitPacker, BitPacker4x, BitPacker8x}; use lance_core::{Error, Result}; +pub const MAX_POSTING_BLOCK_SIZE: usize = BitPacker8x::BLOCK_LEN; + // we compress the posting list to multiple blocks of fixed number of elements (BLOCK_SIZE), // returns a LargeBinaryArray, where each binary is a compressed block (128 row ids + 128 frequencies) // each block is: @@ -41,18 +46,40 @@ pub fn compress_posting_list<'a>( #[cfg(test)] pub fn compress_posting_list_with_tail_codec<'a>( + length: usize, + doc_ids: impl Iterator, + frequencies: impl Iterator, + block_max_scores: impl Iterator, + tail_codec: PostingTailCodec, +) -> Result { + compress_posting_list_with_tail_codec_and_block_size( + length, + doc_ids, + frequencies, + block_max_scores, + tail_codec, + LEGACY_BLOCK_SIZE, + ) +} + +#[cfg(test)] +pub fn compress_posting_list_with_tail_codec_and_block_size<'a>( length: usize, doc_ids: impl Iterator, frequencies: impl Iterator, mut block_max_scores: impl Iterator, tail_codec: PostingTailCodec, + block_size: usize, ) -> Result { - if length < BLOCK_SIZE { + let block_size = validate_block_size(block_size)?; + if length < block_size { // directly do remainder compression to avoid overhead of creating buffer let mut builder = LargeBinaryBuilder::with_capacity(1, length * 4 * 2 + 1); - // write the max score of the block - let max_score = block_max_scores.next().unwrap(); - let _ = builder.write(max_score.to_le_bytes().as_ref())?; + // write the max score of the block (128-doc blocks only) + if posting_block_score_prefix_len(block_size) > 0 { + let max_score = block_max_scores.next().unwrap(); + let _ = builder.write(max_score.to_le_bytes().as_ref())?; + } compress_posting_remainder( doc_ids.copied().collect::>().as_slice(), frequencies.copied().collect::>().as_slice(), @@ -63,27 +90,25 @@ pub fn compress_posting_list_with_tail_codec<'a>( return Ok(builder.finish()); } - let mut builder = LargeBinaryBuilder::with_capacity(length.div_ceil(BLOCK_SIZE), length * 3); - let mut buffer = [0u8; BLOCK_SIZE * 4 + 5]; - let mut doc_id_buffer = Vec::with_capacity(BLOCK_SIZE); - let mut freq_buffer = Vec::with_capacity(BLOCK_SIZE); + let mut builder = LargeBinaryBuilder::with_capacity(length.div_ceil(block_size), length * 3); + let mut doc_id_buffer = Vec::with_capacity(block_size); + let mut freq_buffer = Vec::with_capacity(block_size); for (doc_id, freq) in std::iter::zip(doc_ids, frequencies) { doc_id_buffer.push(*doc_id); freq_buffer.push(*freq); - if doc_id_buffer.len() < BLOCK_SIZE { + if doc_id_buffer.len() < block_size { continue; } - assert_eq!(doc_id_buffer.len(), BLOCK_SIZE); + assert_eq!(doc_id_buffer.len(), block_size); - // write the max score of the block - let max_score = block_max_scores.next().unwrap(); - let _ = builder.write(max_score.to_le_bytes().as_ref())?; - // delta encoding + bitpacking for doc ids - compress_sorted_block(&doc_id_buffer, &mut buffer, &mut builder)?; - // bitpacking for frequencies - compress_block(&freq_buffer, &mut buffer, &mut builder)?; + // write the max score of the block (128-doc blocks only) + if posting_block_score_prefix_len(block_size) > 0 { + let max_score = block_max_scores.next().unwrap(); + let _ = builder.write(max_score.to_le_bytes().as_ref())?; + } + encode_posting_block_payload(&doc_id_buffer, &freq_buffer, &mut builder)?; builder.append_value(""); doc_id_buffer.clear(); freq_buffer.clear(); @@ -91,44 +116,187 @@ pub fn compress_posting_list_with_tail_codec<'a>( // we don't compress the last block if it is not full if !doc_id_buffer.is_empty() { - // write the max score of the block - let max_score = block_max_scores.next().unwrap(); - let _ = builder.write(max_score.to_le_bytes().as_ref())?; + // write the max score of the block (128-doc blocks only) + if posting_block_score_prefix_len(block_size) > 0 { + let max_score = block_max_scores.next().unwrap(); + let _ = builder.write(max_score.to_le_bytes().as_ref())?; + } compress_posting_remainder(&doc_id_buffer, &freq_buffer, tail_codec, &mut builder)?; builder.append_value(""); } Ok(builder.finish()) } +/// Byte length of the block-max-score prefix on posting blocks. 128-doc +/// blocks store a per-block max score, patched in at build time; 256-doc +/// (V3) blocks always carry impact skip data, which supersedes it, so they +/// store none. +#[inline] +pub fn posting_block_score_prefix_len(block_size: usize) -> usize { + if block_size == MAX_POSTING_BLOCK_SIZE { + 0 + } else { + 4 + } +} + pub fn encode_full_posting_block_into( doc_ids: &[u32], frequencies: &[u32], block: &mut Vec, ) -> Result<()> { - debug_assert_eq!(doc_ids.len(), BLOCK_SIZE); - debug_assert_eq!(frequencies.len(), BLOCK_SIZE); - block.extend_from_slice(&0f32.to_le_bytes()); - let mut buffer = [0u8; BLOCK_SIZE * 4 + 5]; - compress_sorted_block(doc_ids, &mut buffer, block)?; - compress_block(frequencies, &mut buffer, block)?; + validate_block_size(doc_ids.len())?; + debug_assert_eq!(doc_ids.len(), frequencies.len()); + if posting_block_score_prefix_len(doc_ids.len()) > 0 { + block.extend_from_slice(&0f32.to_le_bytes()); + } + encode_posting_block_payload(doc_ids, frequencies, block)?; + Ok(()) +} + +fn encode_posting_block_payload( + doc_ids: &[u32], + frequencies: &[u32], + block: &mut impl Write, +) -> Result<()> { + debug_assert_eq!(doc_ids.len(), frequencies.len()); + validate_block_size(doc_ids.len())?; + let mut buffer = [0u8; MAX_POSTING_BLOCK_SIZE * 4 + 5]; + match doc_ids.len() { + BitPacker4x::BLOCK_LEN => { + compress_sorted_block_with::(doc_ids, &mut buffer, block)?; + compress_block_with::(frequencies, &mut buffer, block)?; + } + // 256-doc blocks (format V3) store frequencies with patched FOR: + // outliers no longer widen the whole block, which matters because one + // large tf per block otherwise doubles the frequency payload. + BitPacker8x::BLOCK_LEN => { + compress_sorted_block_with::(doc_ids, &mut buffer, block)?; + compress_pfor_block_with::(frequencies, &mut buffer, block)?; + } + _ => unreachable!("validated posting block size should be supported"), + } + Ok(()) +} + +/// Patched FOR (Lucene PForUtil style): pick the body bit width that +/// minimizes total bytes, pack all values masked to that width, and append +/// up to [`PFOR_MAX_EXCEPTIONS`] exceptions as (index u8, high-bits varint). +const PFOR_MAX_EXCEPTIONS: usize = 31; + +#[inline] +fn u32_bits(value: u32) -> usize { + (32 - value.leading_zeros()) as usize +} + +#[inline] +fn varint_u32_len(value: u32) -> usize { + u32_bits(value).max(1).div_ceil(7) +} + +fn compress_pfor_block_with( + data: &[u32], + buffer: &mut [u8], + builder: &mut impl Write, +) -> Result<()> { + debug_assert_eq!(data.len(), P::BLOCK_LEN); + let max_bits = data.iter().map(|&v| u32_bits(v)).max().unwrap_or(0); + let mut best_width = max_bits; + let mut best_cost = P::BLOCK_LEN * max_bits / 8; + for width in (0..max_bits).rev() { + let mut exceptions = 0usize; + let mut exception_bytes = 0usize; + for &value in data { + if u32_bits(value) > width { + exceptions += 1; + exception_bytes += 1 + varint_u32_len(value >> width); + } + } + if exceptions > PFOR_MAX_EXCEPTIONS { + break; + } + let cost = P::BLOCK_LEN * width / 8 + exception_bytes; + if cost < best_cost { + best_cost = cost; + best_width = width; + } + } + + let mask = if best_width >= 32 { + u32::MAX + } else { + (1u32 << best_width) - 1 + }; + let mut body = [0u32; MAX_POSTING_BLOCK_SIZE]; + let mut exception_buf = Vec::new(); + let mut exception_count = 0u8; + for (index, &value) in data.iter().enumerate() { + body[index] = value & mask; + if u32_bits(value) > best_width { + exception_buf.push(index as u8); + encode_varint_u32(&mut exception_buf, value >> best_width); + exception_count += 1; + } + } + let compressor = P::new(); + let num_bytes = compressor.compress(&body[..P::BLOCK_LEN], buffer, best_width as u8); + let _ = builder.write(&[best_width as u8, exception_count])?; + let _ = builder.write(&buffer[..num_bytes])?; + let _ = builder.write(&exception_buf)?; Ok(()) } +fn decompress_pfor_block_with( + block: &[u8], + buffer: &mut [u32], + res: &mut Vec, +) -> usize { + debug_assert!(buffer.len() >= P::BLOCK_LEN); + let buffer = &mut buffer[..P::BLOCK_LEN]; + let width = block[0]; + let exception_count = block[1] as usize; + let compressor = P::new(); + let num_bytes = compressor.decompress(&block[2..], buffer, width); + let mut offset = 2 + num_bytes; + for _ in 0..exception_count { + let index = block[offset] as usize; + offset += 1; + let high = decode_varint_u32(block, &mut offset) + .expect("pfor exception high bits should be a valid varint"); + buffer[index] |= high << width; + } + res.extend_from_slice(buffer); + offset +} + pub fn encode_remainder_posting_block_into( doc_ids: &[u32], frequencies: &[u32], codec: PostingTailCodec, + block_size: usize, block: &mut Vec, ) -> Result<()> { debug_assert_eq!(doc_ids.len(), frequencies.len()); - block.extend_from_slice(&0f32.to_le_bytes()); + if posting_block_score_prefix_len(block_size) > 0 { + block.extend_from_slice(&0f32.to_le_bytes()); + } compress_posting_remainder(doc_ids, frequencies, codec, block)?; Ok(()) } #[inline] fn compress_sorted_block(data: &[u32], buffer: &mut [u8], builder: &mut impl Write) -> Result<()> { - let compressor = BitPacker4x::new(); + compress_sorted_block_with::(data, buffer, builder) +} + +#[inline] +fn compress_sorted_block_with( + data: &[u32], + buffer: &mut [u8], + builder: &mut impl Write, +) -> Result<()> { + debug_assert_eq!(data.len(), P::BLOCK_LEN); + let compressor = P::new(); let num_bits = compressor.num_bits_sorted(data[0], data); let num_bytes = compressor.compress_sorted(data[0], data, buffer, num_bits); let _ = builder.write(data[0].to_le_bytes().as_ref())?; @@ -139,7 +307,17 @@ fn compress_sorted_block(data: &[u32], buffer: &mut [u8], builder: &mut impl Wri #[inline] fn compress_block(data: &[u32], buffer: &mut [u8], builder: &mut impl Write) -> Result<()> { - let compressor = BitPacker4x::new(); + compress_block_with::(data, buffer, builder) +} + +#[inline] +fn compress_block_with( + data: &[u32], + buffer: &mut [u8], + builder: &mut impl Write, +) -> Result<()> { + debug_assert_eq!(data.len(), P::BLOCK_LEN); + let compressor = P::new(); let num_bits = compressor.num_bits(data); let num_bytes = compressor.compress(data, buffer, num_bits); let _ = builder.write(&[num_bits])?; @@ -236,7 +414,7 @@ pub fn compress_positions(positions: &[u32]) -> Result, mut value: u32) { +pub fn encode_varint_u32(dst: &mut Vec, mut value: u32) { while value >= 0x80 { dst.push((value as u8) | 0x80); value >>= 7; @@ -244,39 +422,6 @@ fn encode_varint_u32(dst: &mut Vec, mut value: u32) { dst.push(value as u8); } -/// Encode a monotonically increasing sequence of `group_starts` (the first -/// row of each posting-list cache group) as varint-encoded deltas. The first -/// value is stored as-is and each subsequent value as its delta from the -/// previous one; since deltas are group sizes (1..=cap) they fit in ~1 byte, -/// keeping the buffer tiny even for indexes with millions of tokens. See -/// issue #7040. -pub(super) fn encode_group_starts(group_starts: &[u32]) -> Vec { - let mut dst = Vec::with_capacity(group_starts.len()); - let mut previous = 0u32; - for &start in group_starts { - debug_assert!(start >= previous, "group_starts must be monotonic"); - encode_varint_u32(&mut dst, start - previous); - previous = start; - } - dst -} - -/// Decode the buffer produced by [`encode_group_starts`] back into the -/// absolute `group_starts` values. -pub(super) fn decode_group_starts(src: &[u8]) -> Result> { - let mut group_starts = Vec::new(); - let mut offset = 0; - let mut previous = 0u32; - while offset < src.len() { - let delta = decode_varint_u32(src, &mut offset)?; - previous = previous - .checked_add(delta) - .ok_or_else(|| Error::index("group_starts delta decode overflowed u32".to_owned()))?; - group_starts.push(previous); - } - Ok(group_starts) -} - #[derive(Debug, Clone, PartialEq, Eq)] pub(super) struct PositionBlockBuilder { codec: PositionStreamCodec, @@ -379,7 +524,7 @@ impl PositionBlockBuilder { } #[inline] -fn decode_varint_u32(src: &[u8], offset: &mut usize) -> Result { +pub fn decode_varint_u32(src: &[u8], offset: &mut usize) -> Result { let mut value = 0u32; let mut shift = 0u32; while *offset < src.len() { @@ -650,23 +795,46 @@ pub fn decompress_posting_list_with_tail_codec( posting_list: &arrow::array::LargeBinaryArray, tail_codec: PostingTailCodec, ) -> Result<(Vec, Vec)> { + decompress_posting_list_with_tail_codec_and_block_size( + num_docs, + posting_list, + tail_codec, + LEGACY_BLOCK_SIZE, + ) +} + +#[cfg(test)] +pub fn decompress_posting_list_with_tail_codec_and_block_size( + num_docs: u32, + posting_list: &arrow::array::LargeBinaryArray, + tail_codec: PostingTailCodec, + block_size: usize, +) -> Result<(Vec, Vec)> { + let block_size = validate_block_size(block_size)?; let mut doc_ids: Vec = Vec::with_capacity(num_docs as usize); let mut frequencies: Vec = Vec::with_capacity(num_docs as usize); - let mut buffer = [0u32; BLOCK_SIZE]; - let bitpacking_blocks = num_docs as usize / BLOCK_SIZE; + let mut buffer = [0u32; MAX_POSTING_BLOCK_SIZE]; + let bitpacking_blocks = num_docs as usize / block_size; for compressed in posting_list.iter().take(bitpacking_blocks) { let compressed = compressed.unwrap(); - decompress_posting_block(compressed, &mut buffer, &mut doc_ids, &mut frequencies); + decompress_posting_block( + compressed, + &mut buffer, + &mut doc_ids, + &mut frequencies, + block_size, + ); } - let remainder = num_docs as usize % BLOCK_SIZE; + let remainder = num_docs as usize % block_size; if remainder > 0 { let compressed = posting_list.value(bitpacking_blocks); decompress_posting_remainder( compressed, remainder, tail_codec, + block_size, &mut doc_ids, &mut frequencies, ); @@ -701,24 +869,46 @@ pub fn read_num_positions(compressed: &arrow::array::LargeBinaryArray) -> u32 { pub fn decompress_posting_block( block: &[u8], - buffer: &mut [u32; BLOCK_SIZE], + buffer: &mut [u32], doc_ids: &mut Vec, frequencies: &mut Vec, + block_size: usize, ) { - // skip the first 4 bytes for the max block score - let block = &block[4..]; - let num_bytes = decompress_sorted_block(block, buffer, doc_ids); - decompress_block(&block[num_bytes..], buffer, frequencies); + debug_assert!(validate_block_size(block_size).is_ok()); + debug_assert!(buffer.len() >= block_size); + // skip the block max score prefix (128-doc blocks only) + let mut block = &block[posting_block_score_prefix_len(block_size)..]; + match block_size { + BitPacker4x::BLOCK_LEN => { + let num_bytes = decompress_sorted_block_with::(block, buffer, doc_ids); + block = &block[num_bytes..]; + let num_bytes = decompress_block_with::(block, buffer, frequencies); + block = &block[num_bytes..]; + } + BitPacker8x::BLOCK_LEN => { + let num_bytes = decompress_sorted_block_with::(block, buffer, doc_ids); + block = &block[num_bytes..]; + let num_bytes = decompress_pfor_block_with::(block, buffer, frequencies); + block = &block[num_bytes..]; + } + _ => unreachable!("validated posting block size should be supported"), + } + debug_assert!( + block.is_empty(), + "posting block has {} trailing bytes after decoding", + block.len() + ); } pub fn decompress_posting_remainder( block: &[u8], n: usize, codec: PostingTailCodec, + block_size: usize, doc_ids: &mut Vec, frequencies: &mut Vec, ) { - let block = &block[4..]; + let block = &block[posting_block_score_prefix_len(block_size)..]; match codec { PostingTailCodec::Fixed32 => { decompress_raw_remainder(block, n, doc_ids); @@ -755,9 +945,14 @@ pub fn decompress_posting_remainder( } } -pub fn decode_full_posting_block(block: &[u8], doc_ids: &mut Vec, frequencies: &mut Vec) { - let mut buffer = [0u32; BLOCK_SIZE]; - decompress_posting_block(block, &mut buffer, doc_ids, frequencies); +pub fn decode_full_posting_block( + block: &[u8], + doc_ids: &mut Vec, + frequencies: &mut Vec, + block_size: usize, +) { + let mut buffer = [0u32; MAX_POSTING_BLOCK_SIZE]; + decompress_posting_block(block, &mut buffer, doc_ids, frequencies, block_size); } pub fn decompress_sorted_block( @@ -765,7 +960,17 @@ pub fn decompress_sorted_block( buffer: &mut [u32; BLOCK_SIZE], res: &mut Vec, ) -> usize { - let compressor = BitPacker4x::new(); + decompress_sorted_block_with::(block, buffer, res) +} + +fn decompress_sorted_block_with( + block: &[u8], + buffer: &mut [u32], + res: &mut Vec, +) -> usize { + debug_assert!(buffer.len() >= P::BLOCK_LEN); + let buffer = &mut buffer[..P::BLOCK_LEN]; + let compressor = P::new(); let initial = u32::from_le_bytes(block[0..4].try_into().unwrap()); let num_bits = block[4]; let num_bytes = compressor.decompress_sorted(initial, &block[5..], buffer, num_bits); @@ -773,11 +978,18 @@ pub fn decompress_sorted_block( 5 + num_bytes } -fn decompress_block(block: &[u8], buffer: &mut [u32; BLOCK_SIZE], res: &mut Vec) { - let compressor = BitPacker4x::new(); +fn decompress_block_with( + block: &[u8], + buffer: &mut [u32], + res: &mut Vec, +) -> usize { + debug_assert!(buffer.len() >= P::BLOCK_LEN); + let buffer = &mut buffer[..P::BLOCK_LEN]; + let compressor = P::new(); let num_bits = block[0]; - compressor.decompress(&block[1..], buffer, num_bits); + let num_bytes = compressor.decompress(&block[1..], buffer, num_bits); res.extend_from_slice(&buffer[..]); + 1 + num_bytes } pub fn decompress_raw_remainder(compressed: &[u8], n: usize, dest: &mut Vec) { @@ -787,11 +999,18 @@ pub fn decompress_raw_remainder(compressed: &[u8], n: usize, dest: &mut Vec } } -pub fn read_posting_tail_first_doc(block: &[u8], codec: PostingTailCodec) -> u32 { +pub fn read_posting_tail_first_doc( + block: &[u8], + codec: PostingTailCodec, + block_size: usize, +) -> u32 { + let prefix = posting_block_score_prefix_len(block_size); match codec { - PostingTailCodec::Fixed32 => u32::from_le_bytes(block[4..8].try_into().unwrap()), + PostingTailCodec::Fixed32 => { + u32::from_le_bytes(block[prefix..prefix + 4].try_into().unwrap()) + } PostingTailCodec::VarintDelta => { - let mut offset = 4usize; + let mut offset = prefix; decode_varint_u32(block, &mut offset) .expect("posting tail block should contain a valid first doc id") } @@ -805,31 +1024,6 @@ mod tests { use itertools::Itertools; use rand::Rng; - #[test] - fn test_group_starts_codec_roundtrip() { - for case in [ - vec![], - vec![0u32], - vec![0, 1, 2, 3], - // realistic: monotonic with varied group sizes and a large jump - vec![0, 64, 128, 129, 4096, 1_000_000], - ] { - let encoded = encode_group_starts(&case); - let decoded = decode_group_starts(&encoded).unwrap(); - assert_eq!(decoded, case, "roundtrip mismatch for {case:?}"); - } - } - - #[test] - fn test_decode_group_starts_rejects_overflow() { - // A crafted buffer whose deltas sum past u32::MAX must error rather - // than wrap. Encodes u32::MAX followed by a +1 delta. - let mut buf = Vec::new(); - encode_varint_u32(&mut buf, u32::MAX); - encode_varint_u32(&mut buf, 1); - assert!(decode_group_starts(&buf).is_err()); - } - #[test] fn test_compress_posting_list() -> Result<()> { let num_rows: usize = BLOCK_SIZE * 1024 - 7; @@ -867,6 +1061,84 @@ mod tests { Ok(()) } + #[test] + fn test_compress_posting_list_supported_block_sizes() -> Result<()> { + for block_size in [128, 256] { + let num_rows: usize = block_size * 2 + 7; + let doc_ids = (0..num_rows as u32).collect::>(); + let frequencies = (0..num_rows as u32) + .map(|value| value % 7 + 1) + .collect::>(); + let block_max_scores = + (0..num_rows.div_ceil(block_size)).map(|value| value as f32 + 1.0); + + let posting_list = compress_posting_list_with_tail_codec_and_block_size( + doc_ids.len(), + doc_ids.iter(), + frequencies.iter(), + block_max_scores, + PostingTailCodec::VarintDelta, + block_size, + )?; + assert_eq!(posting_list.len(), num_rows.div_ceil(block_size)); + + let (decoded_doc_ids, decoded_frequencies) = + decompress_posting_list_with_tail_codec_and_block_size( + num_rows as u32, + &posting_list, + PostingTailCodec::VarintDelta, + block_size, + )?; + assert_eq!(decoded_doc_ids, doc_ids); + assert_eq!(decoded_frequencies, frequencies); + } + Ok(()) + } + + #[test] + fn test_256_posting_block_uses_single_physical_bitpack_chunk() -> Result<()> { + let block_size = BitPacker8x::BLOCK_LEN; + let doc_ids = (0..block_size as u32).collect::>(); + let frequencies = (0..block_size as u32) + .map(|value| value % 13 + 1) + .collect::>(); + + let posting_list = compress_posting_list_with_tail_codec_and_block_size( + doc_ids.len(), + doc_ids.iter(), + frequencies.iter(), + std::iter::once(1.0), + PostingTailCodec::VarintDelta, + block_size, + )?; + assert_eq!(posting_list.len(), 1); + + let block = posting_list.value(0); + // 256-doc blocks carry no block-max-score prefix (impacts supply the + // per-block bound): [first_doc u32][doc num_bits u8][doc payload]... + let doc_num_bits = block[4]; + let doc_bytes = BitPacker8x::compressed_block_size(doc_num_bits); + let freq_header_offset = 5 + doc_bytes; + // 256-doc blocks use patched FOR for frequencies: + // [width u8][exception_count u8][body][exceptions...] + let freq_num_bits = block[freq_header_offset]; + let exception_count = block[freq_header_offset + 1] as usize; + assert_eq!(exception_count, 0, "uniform freqs need no exceptions"); + let freq_bytes = BitPacker8x::compressed_block_size(freq_num_bits); + assert_eq!(block.len(), freq_header_offset + 2 + freq_bytes); + + let (decoded_doc_ids, decoded_frequencies) = + decompress_posting_list_with_tail_codec_and_block_size( + doc_ids.len() as u32, + &posting_list, + PostingTailCodec::VarintDelta, + block_size, + )?; + assert_eq!(decoded_doc_ids, doc_ids); + assert_eq!(decoded_frequencies, frequencies); + Ok(()) + } + #[test] fn test_compress_posting_list_fixed32_tail_still_roundtrips() -> Result<()> { let doc_ids = vec![3_u32, 10_u32, 24_u32]; diff --git a/rust/lance-index/src/scalar/inverted/impact.rs b/rust/lance-index/src/scalar/inverted/impact.rs new file mode 100644 index 00000000000..72b26eb07b2 --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/impact.rs @@ -0,0 +1,1016 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::mem::size_of; +use std::sync::{Arc, Mutex, MutexGuard}; + +use arrow_array::builder::LargeBinaryBuilder; +use arrow_array::{Array, LargeBinaryArray}; +use lance_core::{Error, Result}; + +use super::scorer::Scorer; + +pub const IMPACT_LEVEL1_BLOCKS: usize = 32; +const SMALL_FRONTIER_FREQ_LIMIT: usize = 256; + +/// On-disk encoding of one impact entry, shared by every posting block size. +/// +/// Entries contain `[doc_up_to varint][pair_count varint][pairs...]`. Each +/// pair stores a varint whose high bits are `freq_delta - 1` and whose low bit +/// reports whether a one-byte norm delta follows. The norm itself is the +/// quantized `u8` document-length code; the common `norm_delta == 1` case needs +/// no norm byte. +#[derive(Debug, Clone)] +pub struct ImpactSkipData { + entries: LargeBinaryArray, + level0_len: usize, + // Last doc id covered by each entry (level0 entries then level1 entries), + // decoded once at construction. Level1 markers are fully validated because + // WAND may use them to skip a group; u32::MAX marks malformed entries. + entry_doc_up_tos: Arc<[u32]>, + // The most recently baked bounds with a stable scorer key. Each query holds + // its own Arc in ImpactScoreCache, so replacing this slot for another scorer + // cannot change bounds already in use. Scorers without a key never enter the + // shared slot. Malformed entries bake to INFINITY so pruning stays safe. + last_keyed_bounds: Arc>, +} + +impl PartialEq for ImpactSkipData { + fn eq(&self, other: &Self) -> bool { + self.entries == other.entries && self.level0_len == other.level0_len + } +} + +#[derive(Debug, Clone, Copy)] +pub struct ImpactScore { + pub score: f32, + pub entries_scanned: usize, +} + +#[derive(Debug)] +struct ImpactBounds { + per_entry: Box<[f32]>, + global: f32, +} + +type LastKeyedImpactBounds = Option<(u64, Arc)>; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ImpactBoundsCacheKey { + Keyed(u64), + QueryLocal, +} + +#[derive(Debug, Default, Clone)] +pub struct ImpactScoreCache { + key: Option, + bounds: Option>, +} + +impl ImpactScoreCache { + fn bounds<'a, S: Scorer + ?Sized>( + &'a mut self, + impacts: &ImpactSkipData, + scorer: &S, + ) -> &'a ImpactBounds { + let scorer_key = scorer.doc_weight_cache_key(); + let cache_key = scorer_key + .map(ImpactBoundsCacheKey::Keyed) + .unwrap_or(ImpactBoundsCacheKey::QueryLocal); + if self.key != Some(cache_key) { + self.key = Some(cache_key); + self.bounds = None; + } + + self.bounds + .get_or_insert_with(|| impacts.bounds_for_scorer(scorer, scorer_key)) + } + + fn entry_score( + &mut self, + impacts: &ImpactSkipData, + entry_idx: usize, + query_weight: f32, + scorer: &S, + ) -> f32 { + if query_weight <= 0.0 { + return 0.0; + } + query_weight * self.bounds(impacts, scorer).per_entry[entry_idx] + } +} + +impl ImpactSkipData { + pub fn new(entries: LargeBinaryArray, level0_len: usize) -> Result { + let expected_len = level0_len + level1_len(level0_len); + if entries.len() != expected_len { + return Err(Error::index(format!( + "impact entry count mismatch: got {}, expected {} for {} level0 blocks", + entries.len(), + expected_len, + level0_len + ))); + } + let entry_doc_up_tos = (0..entries.len()) + .map(|entry_idx| { + if entries.is_null(entry_idx) { + return u32::MAX; + } + let bytes = entries.value(entry_idx); + let doc_up_to = if entry_idx < level0_len { + decode_level0_entry_doc_up_to(bytes) + } else { + decode_entry_doc_up_to(bytes) + }; + doc_up_to.unwrap_or(u32::MAX) + }) + .collect::>(); + Ok(Self { + entries, + level0_len, + entry_doc_up_tos, + last_keyed_bounds: Arc::new(Mutex::new(None)), + }) + } + + fn keyed_bounds_guard(&self) -> MutexGuard<'_, LastKeyedImpactBounds> { + match self.last_keyed_bounds.lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + } + } + + fn bounds_for_scorer( + &self, + scorer: &S, + scorer_key: Option, + ) -> Arc { + let Some(scorer_key) = scorer_key else { + return Arc::new(self.compute_bounds(scorer)); + }; + + { + let cached = self.keyed_bounds_guard(); + if let Some((cached_key, bounds)) = cached.as_ref() + && *cached_key == scorer_key + { + return bounds.clone(); + } + } + + // Compute outside the mutex. Concurrent misses may duplicate this work, + // but only the short publication/check below holds the shared lock. + let computed = Arc::new(self.compute_bounds(scorer)); + let mut cached = self.keyed_bounds_guard(); + if let Some((cached_key, bounds)) = cached.as_ref() + && *cached_key == scorer_key + { + return bounds.clone(); + } + *cached = Some((scorer_key, computed.clone())); + computed + } + + fn compute_bounds(&self, scorer: &S) -> ImpactBounds { + let per_entry = (0..self.entries.len()) + .map(|entry_idx| { + if self.entries.is_null(entry_idx) { + return f32::INFINITY; + } + let bytes = self.entries.value(entry_idx); + let mut max_doc_weight = 0.0_f32; + match for_each_entry_pair(bytes, |freq, doc_len| { + max_doc_weight = max_doc_weight.max(scorer.doc_weight(freq, doc_len)); + }) { + Ok(()) => max_doc_weight, + Err(_) => f32::INFINITY, + } + }) + .collect::>(); + // The level1 entries cover every block, so their max is the list-wide + // max doc weight; zero-entry lists fall back to the empty level0 slab. + let global = if per_entry.len() > self.level0_len { + per_entry[self.level0_len..] + .iter() + .copied() + .fold(0.0_f32, f32::max) + } else { + per_entry.iter().copied().fold(0.0_f32, f32::max) + }; + ImpactBounds { per_entry, global } + } + + /// List-wide max doc weight, from the scorer-specific cached bounds. The + /// tightest valid global score bound is `query_weight * this`, matching what + /// the non-impact format stores as `max_score` at build time. + pub fn global_max_doc_weight_cached( + &self, + scorer: &S, + cache: &mut ImpactScoreCache, + ) -> f32 { + cache.bounds(self, scorer).global + } + + pub fn entries(&self) -> &LargeBinaryArray { + &self.entries + } + + /// Conservative heap charge for query-independent derived state and one + /// shared keyed-bound slab, whether or not that slab has been initialized + /// yet. The Arrow impact entries are owned by the enclosing batch and are + /// deliberately excluded so packed-group cache accounting counts them once. + pub(crate) fn derived_cache_bytes(&self) -> usize { + Self::derived_cache_bytes_for_entries(self.entries.len()) + } + + pub(crate) fn derived_cache_bytes_for_entries(entry_count: usize) -> usize { + entry_count * size_of::() + + size_of::>() + + size_of::() + + entry_count * size_of::() + } + + #[cfg(test)] + pub(crate) fn shares_derived_state_with(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.entry_doc_up_tos, &other.entry_doc_up_tos) + && Arc::ptr_eq(&self.last_keyed_bounds, &other.last_keyed_bounds) + } + + #[cfg(test)] + pub fn level0_len(&self) -> usize { + self.level0_len + } + + #[cfg(test)] + pub fn level1_len(&self) -> usize { + level1_len(self.level0_len) + } + + pub(crate) fn level1_doc_up_to(&self, group_idx: usize) -> Option { + if group_idx >= level1_len(self.level0_len) { + return None; + } + match self.entry_doc_up_tos[self.level0_len + group_idx] { + u32::MAX => None, + doc_up_to => Some(doc_up_to), + } + } + + pub fn level0_score_cached( + &self, + block_idx: usize, + query_weight: f32, + scorer: &S, + cache: &mut ImpactScoreCache, + ) -> f32 { + if block_idx >= self.level0_len { + return 0.0; + } + cache.entry_score(self, block_idx, query_weight, scorer) + } + + pub fn max_score_up_to_cached( + &self, + start_block_idx: usize, + up_to: u64, + query_weight: f32, + scorer: &S, + cache: &mut ImpactScoreCache, + ) -> ImpactScore + where + S: Scorer + ?Sized, + { + let mut block_idx = start_block_idx; + let mut max_score = 0.0_f32; + let mut entries_scanned = 0usize; + + while block_idx < self.level0_len { + let group_idx = block_idx / IMPACT_LEVEL1_BLOCKS; + let group_start = group_idx * IMPACT_LEVEL1_BLOCKS; + let group_end = ((group_idx + 1) * IMPACT_LEVEL1_BLOCKS).min(self.level0_len); + if block_idx == group_start { + let level1_entry_idx = self.level0_len + group_idx; + match self.entry_doc_up_tos[level1_entry_idx] { + u32::MAX => { + return ImpactScore { + score: f32::INFINITY, + entries_scanned: entries_scanned + 1, + }; + } + doc_up_to if u64::from(doc_up_to) <= up_to => { + max_score = max_score.max(cache.entry_score( + self, + level1_entry_idx, + query_weight, + scorer, + )); + entries_scanned += 1; + block_idx = group_end; + continue; + } + _ => {} + } + } + + max_score = max_score.max(cache.entry_score(self, block_idx, query_weight, scorer)); + entries_scanned += 1; + match self.entry_doc_up_tos[block_idx] { + u32::MAX => { + return ImpactScore { + score: f32::INFINITY, + entries_scanned, + }; + } + doc_up_to if u64::from(doc_up_to) >= up_to => break, + _ => {} + } + block_idx += 1; + } + + ImpactScore { + score: max_score, + entries_scanned, + } + } +} + +pub struct ImpactSkipDataBuilder { + entries: LargeBinaryBuilder, + level0_len: usize, + level1_entries: Vec>, + level1_docs: Vec<(u32, u32, u32)>, +} + +impl ImpactSkipDataBuilder { + pub fn with_capacity(level0_blocks: usize, block_size: usize) -> Self { + Self { + entries: LargeBinaryBuilder::with_capacity( + level0_blocks + level1_len(level0_blocks), + 0, + ), + level0_len: 0, + level1_entries: Vec::with_capacity(level1_len(level0_blocks)), + level1_docs: Vec::with_capacity(IMPACT_LEVEL1_BLOCKS * block_size), + } + } + + pub fn append_block(&mut self, docs: &[(u32, u32, u32)]) -> Result<()> { + let bytes = encode_impact_entry(docs)?; + self.entries.append_value(bytes.as_slice()); + self.level0_len += 1; + self.level1_docs.extend_from_slice(docs); + if self.level0_len.is_multiple_of(IMPACT_LEVEL1_BLOCKS) { + self.flush_level1()?; + } + Ok(()) + } + + pub fn finish(mut self) -> Result { + if !self.level1_docs.is_empty() { + self.flush_level1()?; + } + for entry in self.level1_entries { + self.entries.append_value(entry.as_slice()); + } + ImpactSkipData::new(self.entries.finish(), self.level0_len) + } + + fn flush_level1(&mut self) -> Result<()> { + let bytes = encode_impact_entry(self.level1_docs.as_slice())?; + self.level1_entries.push(bytes); + self.level1_docs.clear(); + Ok(()) + } +} + +#[cfg(test)] +pub fn build_impact_skip_data(blocks: &[Vec<(u32, u32, u32)>]) -> Result { + let block_size = blocks.iter().map(Vec::len).max().unwrap_or(0).max(1); + let mut builder = ImpactSkipDataBuilder::with_capacity(blocks.len(), block_size); + for block in blocks { + builder.append_block(block)?; + } + builder.finish() +} + +fn encode_impact_entry(docs: &[(u32, u32, u32)]) -> Result> { + if docs.is_empty() { + return Err(Error::index( + "cannot encode an empty impact entry".to_owned(), + )); + } + let doc_up_to = docs + .last() + .map(|(doc_id, _, _)| *doc_id) + .expect("non-empty impact entry was validated above"); + let frontier = quantized_impact_frontier(docs); + let pair_count = u32::try_from(frontier.len()).map_err(|_| { + Error::index("impact frontier too large to encode as u32 pair count".to_string()) + })?; + let mut bytes = Vec::with_capacity(5 + frontier.len() * 2); + super::encoding::encode_varint_u32(&mut bytes, doc_up_to); + super::encoding::encode_varint_u32(&mut bytes, pair_count); + let mut previous_freq = 0u32; + let mut previous_norm = 0u8; + for (pair_idx, (freq, norm)) in frontier.into_iter().enumerate() { + let freq_delta_minus_one = freq + .checked_sub(previous_freq) + .and_then(|delta| delta.checked_sub(1)) + .ok_or_else(|| { + Error::index(format!( + "impact frequencies must be positive and strictly increasing: previous={previous_freq}, current={freq}" + )) + })?; + let norm_delta = norm.checked_sub(previous_norm).ok_or_else(|| { + Error::index(format!( + "impact norms must be non-decreasing: previous={previous_norm}, current={norm}" + )) + })?; + if pair_idx > 0 && norm_delta == 0 { + return Err(Error::index(format!( + "impact norms must be strictly increasing after quantization: norm={norm}" + ))); + } + + let has_explicit_norm_delta = norm_delta != 1; + let packed_freq_delta = + (u64::from(freq_delta_minus_one) << 1) | u64::from(has_explicit_norm_delta); + encode_varint_u64(&mut bytes, packed_freq_delta); + if has_explicit_norm_delta { + bytes.push(norm_delta); + } + previous_freq = freq; + previous_norm = norm; + } + Ok(bytes) +} + +fn decode_entry_doc_up_to(bytes: &[u8]) -> Result { + let mut offset = 0usize; + let doc_up_to = super::encoding::decode_varint_u32(bytes, &mut offset)?; + // Level-1 doc ids drive whole-group skips, so only publish a doc id after + // validating the complete entry. A truncated entry may still have a valid + // first varint. + for_each_entry_pair(bytes, |_, _| {})?; + Ok(doc_up_to) +} + +fn decode_level0_entry_doc_up_to(bytes: &[u8]) -> Result { + // A malformed level0 frontier bakes an INFINITY score before this marker + // can terminate a range scan, so avoid parsing every frontier twice on the + // query-load path. Level1 entries are fully validated. + let mut offset = 0usize; + super::encoding::decode_varint_u32(bytes, &mut offset) +} + +/// Walk an entry's (freq, doc_len) frontier pairs, validating the layout. +fn for_each_entry_pair(bytes: &[u8], mut visit: impl FnMut(u32, u32)) -> Result<()> { + let mut offset = 0usize; + let _doc_up_to = super::encoding::decode_varint_u32(bytes, &mut offset)?; + let pair_count = super::encoding::decode_varint_u32(bytes, &mut offset)?; + if pair_count == 0 { + return Err(Error::index( + "impact entry must contain at least one frontier pair".to_owned(), + )); + } + + let mut previous_freq = 0u32; + let mut previous_norm = 0u16; + for pair_idx in 0..pair_count { + let packed_freq_delta = decode_varint_u64(bytes, &mut offset)?; + let freq_delta_minus_one = u32::try_from(packed_freq_delta >> 1) + .map_err(|_| Error::index("impact freq delta exceeds u32".to_owned()))?; + let freq_delta = freq_delta_minus_one + .checked_add(1) + .ok_or_else(|| Error::index("impact freq delta overflow".to_owned()))?; + let freq = previous_freq + .checked_add(freq_delta) + .ok_or_else(|| Error::index("impact frequency overflow".to_owned()))?; + + let has_explicit_norm_delta = packed_freq_delta & 1 != 0; + let norm_delta = if has_explicit_norm_delta { + let norm_delta = bytes.get(offset).copied().ok_or_else(|| { + Error::index("unexpected EOF while decoding impact norm delta".to_owned()) + })?; + offset += 1; + norm_delta + } else { + 1 + }; + if pair_idx > 0 && norm_delta == 0 { + return Err(Error::index( + "impact norms must be strictly increasing".to_owned(), + )); + } + + let norm = previous_norm + .checked_add(u16::from(norm_delta)) + .filter(|norm| *norm <= u16::from(u8::MAX)) + .ok_or_else(|| Error::index("impact norm delta overflow".to_owned()))?; + let norm = norm as u8; + visit(freq, super::index::dequantize_doc_length(norm)); + previous_freq = freq; + previous_norm = u16::from(norm); + } + if offset != bytes.len() { + return Err(Error::index(format!( + "impact entry has {} trailing bytes", + bytes.len() - offset + ))); + } + Ok(()) +} + +#[inline] +fn encode_varint_u64(dst: &mut Vec, mut value: u64) { + while value >= 0x80 { + dst.push((value as u8) | 0x80); + value >>= 7; + } + dst.push(value as u8); +} + +#[inline] +fn decode_varint_u64(src: &[u8], offset: &mut usize) -> Result { + let mut value = 0u64; + let mut shift = 0u32; + while *offset < src.len() { + let byte = src[*offset]; + *offset += 1; + if shift == 63 && byte & 0xFE != 0 { + return Err(Error::index( + "invalid u64 varint in impact entry".to_owned(), + )); + } + value |= u64::from(byte & 0x7F) << shift; + if byte & 0x80 == 0 { + return Ok(value); + } + shift += 7; + if shift > 63 { + return Err(Error::index( + "invalid u64 varint in impact entry".to_owned(), + )); + } + } + Err(Error::index( + "unexpected EOF while decoding impact entry".to_owned(), + )) +} + +fn quantized_impact_frontier(docs: &[(u32, u32, u32)]) -> Vec<(u32, u8)> { + let raw_frontier = impact_frontier(docs); + let mut frontier: Vec<(u32, u8)> = Vec::with_capacity(raw_frontier.len()); + for (freq, doc_len) in raw_frontier { + let norm = super::index::quantize_doc_length(doc_len); + match frontier.last_mut() { + Some((last_freq, last_norm)) if *last_norm == norm => { + // At the same quantized norm, the larger frequency dominates. + *last_freq = freq; + } + Some((_, last_norm)) => { + debug_assert!( + *last_norm < norm, + "raw impact frontier document lengths must be increasing" + ); + frontier.push((freq, norm)); + } + None => frontier.push((freq, norm)), + } + } + frontier +} + +fn impact_frontier(docs: &[(u32, u32, u32)]) -> Vec<(u32, u32)> { + let max_freq = docs.iter().map(|(_, freq, _)| *freq).max().unwrap_or(0) as usize; + if max_freq <= SMALL_FRONTIER_FREQ_LIMIT { + return impact_frontier_small_freq(docs, max_freq); + } + + impact_frontier_sparse_freq(docs) +} + +fn impact_frontier_small_freq(docs: &[(u32, u32, u32)], max_freq: usize) -> Vec<(u32, u32)> { + let mut min_doc_len_by_freq = [u32::MAX; SMALL_FRONTIER_FREQ_LIMIT + 1]; + for (_, freq, doc_len) in docs { + min_doc_len_by_freq[*freq as usize] = min_doc_len_by_freq[*freq as usize].min(*doc_len); + } + + let min_doc_lens = min_doc_len_by_freq[..=max_freq] + .iter() + .enumerate() + .filter_map(|(freq, doc_len)| (*doc_len != u32::MAX).then_some((freq as u32, *doc_len))) + .collect::>(); + frontier_from_min_doc_lens(min_doc_lens) +} + +fn impact_frontier_sparse_freq(docs: &[(u32, u32, u32)]) -> Vec<(u32, u32)> { + let mut pairs = docs + .iter() + .map(|(_, freq, doc_len)| (*freq, *doc_len)) + .collect::>(); + pairs.sort_unstable_by_key(|(freq, _)| *freq); + + let mut min_doc_lens: Vec<(u32, u32)> = Vec::with_capacity(pairs.len()); + for (freq, doc_len) in pairs { + match min_doc_lens.last_mut() { + Some((last_freq, last_doc_len)) if *last_freq == freq => { + *last_doc_len = (*last_doc_len).min(doc_len); + } + _ => min_doc_lens.push((freq, doc_len)), + } + } + + frontier_from_min_doc_lens(min_doc_lens) +} + +fn frontier_from_min_doc_lens(min_doc_lens: Vec<(u32, u32)>) -> Vec<(u32, u32)> { + let mut best_doc_len = u32::MAX; + let mut frontier = Vec::with_capacity(min_doc_lens.len()); + for (freq, doc_len) in min_doc_lens.into_iter().rev() { + if doc_len < best_doc_len { + frontier.push((freq, doc_len)); + best_doc_len = doc_len; + } + } + frontier.reverse(); + frontier +} + +fn level1_len(level0_len: usize) -> usize { + level0_len.div_ceil(IMPACT_LEVEL1_BLOCKS) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; + + use super::*; + use crate::scalar::inverted::scorer::{MemBM25Scorer, Scorer}; + + struct KeyedCountingScorer { + key: u64, + calls: Arc, + } + + impl Scorer for KeyedCountingScorer { + fn query_weight(&self, _token: &str) -> f32 { + 1.0 + } + + fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32 { + self.calls.fetch_add(1, Ordering::Relaxed); + freq as f32 / doc_tokens as f32 + } + + fn doc_weight_cache_key(&self) -> Option { + Some(self.key) + } + } + + #[test] + fn impact_entry_frontier_drops_dominated_pairs() { + let docs = vec![(0, 1, 10), (1, 1, 8), (2, 2, 9), (3, 3, 20)]; + assert_eq!(impact_frontier(&docs), vec![(1, 8), (2, 9), (3, 20)]); + } + + #[test] + fn impact_entry_frontier_handles_sparse_large_frequencies() { + let docs = vec![ + (0, 1, 100), + (1, 1, 80), + (2, 512, 90), + (3, 1_000, 120), + (4, 1_000, 110), + ]; + assert_eq!( + impact_frontier(&docs), + vec![(1, 80), (512, 90), (1_000, 110)] + ); + } + + #[test] + fn quantized_impact_frontier_drops_equal_norms() { + let docs = vec![(0, 1, 16), (1, 2, 17), (2, 3, 24)]; + assert_eq!( + quantized_impact_frontier(&docs), + vec![ + (2, super::super::index::quantize_doc_length(17)), + (3, super::super::index::quantize_doc_length(24)), + ] + ); + } + + #[test] + fn impact_max_score_can_use_level1_entry() { + let blocks = (0..40) + .map(|block| vec![(block as u32, 1 + block as u32 % 3, 10)]) + .collect::>(); + let impacts = build_impact_skip_data(&blocks).unwrap(); + assert_eq!(impacts.level0_len(), 40); + assert_eq!(impacts.level1_len(), 2); + let scorer = MemBM25Scorer::new(400, 40, HashMap::from([(String::from("token"), 40usize)])); + let mut cache = ImpactScoreCache::default(); + let score = impacts.max_score_up_to_cached(0, 31, 1.0, &scorer, &mut cache); + assert!(score.entries_scanned < IMPACT_LEVEL1_BLOCKS); + assert!(score.score > 0.0); + } + + #[test] + fn impact_level1_doc_up_to_reports_full_and_partial_groups() { + let blocks = (0..40) + .map(|block| vec![(block as u32, 1, 10)]) + .collect::>(); + let impacts = build_impact_skip_data(&blocks).unwrap(); + + assert_eq!( + impacts.level1_doc_up_to(0), + Some((IMPACT_LEVEL1_BLOCKS - 1) as u32) + ); + assert_eq!(impacts.level1_doc_up_to(1), Some(39)); + assert_eq!(impacts.level1_doc_up_to(2), None); + } + + #[test] + fn impact_level1_doc_up_to_returns_none_for_malformed_entry() { + let level0 = encode_impact_entry(&[(0, 1, 10)]).unwrap(); + let malformed_level1 = vec![1, 2, 3]; + let entries = LargeBinaryArray::from_opt_vec(vec![ + Some(level0.as_slice()), + Some(malformed_level1.as_slice()), + ]); + let impacts = ImpactSkipData::new(entries, 1).unwrap(); + + assert_eq!(impacts.level1_doc_up_to(0), None); + } + + #[test] + fn impact_level1_doc_up_to_validates_complete_entry() { + let level0 = encode_impact_entry(&[(0, 1, 10)]).unwrap(); + // A complete first varint is not enough: the pair count and frontier + // are required before this doc id can safely drive a group skip. + let truncated_level1 = [31_u8]; + let entries = LargeBinaryArray::from_opt_vec(vec![ + Some(level0.as_slice()), + Some(truncated_level1.as_slice()), + ]); + let impacts = ImpactSkipData::new(entries, 1).unwrap(); + + assert_eq!(impacts.level1_doc_up_to(0), None); + } + + #[test] + fn empty_impact_frontiers_are_malformed_bounds() { + let scorer = MemBM25Scorer::new(10, 1, HashMap::new()); + let level0 = encode_impact_entry(&[(0, 1, 10)]).unwrap(); + let empty_level1 = [0, 0]; + let entries = LargeBinaryArray::from_opt_vec(vec![ + Some(level0.as_slice()), + Some(empty_level1.as_slice()), + ]); + let impacts = ImpactSkipData::new(entries, 1).unwrap(); + let mut cache = ImpactScoreCache::default(); + + assert_eq!(impacts.level1_doc_up_to(0), None); + assert!( + impacts + .global_max_doc_weight_cached(&scorer, &mut cache) + .is_infinite() + ); + assert!( + ImpactSkipDataBuilder::with_capacity(1, 128) + .append_block(&[]) + .is_err() + ); + } + + #[test] + fn null_impact_entry_is_an_infinite_bound_even_with_hidden_bytes() { + let level0 = encode_impact_entry(&[(0, 1, 10)]).unwrap(); + let level1 = encode_impact_entry(&[(0, 2, 8)]).unwrap(); + let mut values = level0.clone(); + values.extend_from_slice(&level1); + let entries = LargeBinaryArray::new( + OffsetBuffer::new(ScalarBuffer::from(vec![ + 0_i64, + level0.len() as i64, + values.len() as i64, + ])), + Buffer::from_vec(values), + Some(NullBuffer::from(vec![true, false])), + ); + let impacts = ImpactSkipData::new(entries, 1).unwrap(); + let scorer = MemBM25Scorer::new(10, 1, HashMap::new()); + let mut cache = ImpactScoreCache::default(); + + assert_eq!(impacts.level1_doc_up_to(0), None); + assert!( + impacts + .global_max_doc_weight_cached(&scorer, &mut cache) + .is_infinite() + ); + } + + #[test] + fn impact_bounds_follow_changed_bm25_average_doc_length() { + let impacts = build_impact_skip_data(&[vec![(0, 1, 100)]]).unwrap(); + let low_avgdl = MemBM25Scorer::new(1, 1, HashMap::new()); + let high_avgdl = MemBM25Scorer::new(100, 1, HashMap::new()); + let mut low_cache = ImpactScoreCache::default(); + let mut high_cache = ImpactScoreCache::default(); + + let low_bound = impacts.global_max_doc_weight_cached(&low_avgdl, &mut low_cache); + let high_bound = impacts.global_max_doc_weight_cached(&high_avgdl, &mut high_cache); + let quantized_doc_length = super::super::index::dequantize_doc_length( + super::super::index::quantize_doc_length(100), + ); + + assert!((low_bound - low_avgdl.doc_weight(1, quantized_doc_length)).abs() < 1e-6); + assert!((high_bound - high_avgdl.doc_weight(1, quantized_doc_length)).abs() < 1e-6); + assert!(low_bound >= low_avgdl.doc_weight(1, 100)); + assert!(high_bound >= high_avgdl.doc_weight(1, 100)); + assert!( + high_bound > low_bound, + "larger avgdl must recompute a larger bound: low={low_bound}, high={high_bound}" + ); + } + + #[test] + fn impact_bounds_reuse_same_scorer_key_across_queries() { + let impacts = build_impact_skip_data(&[vec![(0, 2, 10)]]).unwrap(); + let cloned = impacts.clone(); + assert!(impacts.shares_derived_state_with(&cloned)); + let calls = Arc::new(AtomicUsize::new(0)); + let scorer = KeyedCountingScorer { + key: 7, + calls: calls.clone(), + }; + let mut first_query = ImpactScoreCache::default(); + let mut second_query = ImpactScoreCache::default(); + + let first = impacts.global_max_doc_weight_cached(&scorer, &mut first_query); + let baked_calls = calls.load(Ordering::Relaxed); + assert!(baked_calls > 0); + let second = cloned.global_max_doc_weight_cached(&scorer, &mut second_query); + + assert_eq!(second, first); + assert_eq!( + calls.load(Ordering::Relaxed), + baked_calls, + "the same keyed bounds should be shared across query caches" + ); + } + + #[test] + fn malformed_unscanned_entry_does_not_poison_range_score() { + let level0_0 = encode_impact_entry(&[(0, 1, 10)]).unwrap(); + let malformed_level0_1 = vec![1, 2, 3]; + let level1 = encode_impact_entry(&[(0, 1, 10), (1, 1, 10)]).unwrap(); + let entries = LargeBinaryArray::from_opt_vec(vec![ + Some(level0_0.as_slice()), + Some(malformed_level0_1.as_slice()), + Some(level1.as_slice()), + ]); + let impacts = ImpactSkipData::new(entries, 2).unwrap(); + let scorer = MemBM25Scorer::new(10, 10, HashMap::from([(String::from("token"), 2usize)])); + let mut cache = ImpactScoreCache::default(); + + let score = impacts.max_score_up_to_cached(0, 0, 1.0, &scorer, &mut cache); + assert!(score.score.is_finite()); + assert_eq!(score.entries_scanned, 1); + + assert_eq!( + impacts.level0_score_cached(1, 1.0, &scorer, &mut cache), + f32::INFINITY + ); + } + + #[test] + fn impact_entries_store_quantized_norm_deltas() { + let docs = vec![(7, 1, 1), (9, 2, 2), (12, 3, 5)]; + let encoded = encode_impact_entry(&docs).unwrap(); + + // doc_up_to=12, pair_count=3, two implicit +1 norm deltas, then an + // explicit +3 norm delta folded behind the frequency varint's low bit. + assert_eq!(encoded, vec![12, 3, 0, 0, 1, 3]); + + let mut pairs = Vec::new(); + for_each_entry_pair(&encoded, |freq, doc_len| pairs.push((freq, doc_len))).unwrap(); + assert_eq!(pairs, vec![(1, 1), (2, 2), (3, 5)]); + } + + #[test] + fn malformed_norm_deltas_are_rejected() { + // doc_up_to=0, pair_count=1, and the pair flag promises a norm byte + // that is not present. + let truncated = [0, 1, 1]; + let error = for_each_entry_pair(&truncated, |_, _| {}).unwrap_err(); + assert!(matches!(&error, Error::Index { .. })); + assert!(error.to_string().contains("impact norm delta")); + + // The first pair reaches norm 255, so an implicit +1 on the second + // pair must fail instead of wrapping back to zero. + let overflowing = [0, 2, 1, 255, 0]; + let error = for_each_entry_pair(&overflowing, |_, _| {}).unwrap_err(); + assert!(matches!(&error, Error::Index { .. })); + assert!(error.to_string().contains("impact norm delta overflow")); + } + + #[test] + fn impact_entries_roundtrip_quantized_frontier() { + let docs = vec![(3, 1, 100), (9, 2, 40), (200, 7, 80), (4095, 130, 900)]; + let encoded = encode_impact_entry(&docs).unwrap(); + assert_eq!(decode_entry_doc_up_to(&encoded).unwrap(), 4095); + let mut decoded_pairs = Vec::new(); + for_each_entry_pair(&encoded, |freq, doc_len| { + decoded_pairs.push((freq, doc_len)) + }) + .unwrap(); + let expected_pairs = quantized_impact_frontier(&docs) + .into_iter() + .map(|(freq, norm)| (freq, super::super::index::dequantize_doc_length(norm))) + .collect::>(); + assert_eq!(decoded_pairs, expected_pairs); + assert!(!decoded_pairs.is_empty()); + + // A 256-doc-block skip data goes through the shared codec end to end. + let blocks: Vec> = (0..3) + .map(|b| (0..256).map(|i| (b * 256 + i, 1 + i % 5, 10)).collect()) + .collect(); + let impacts = build_impact_skip_data(&blocks).unwrap(); + assert_eq!(impacts.level1_doc_up_to(0), Some(767)); + let scorer = MemBM25Scorer::new(400, 768, HashMap::from([(String::from("t"), 768usize)])); + let mut cache = ImpactScoreCache::default(); + assert!( + impacts + .level0_score_cached(0, 1.0, &scorer, &mut cache) + .is_finite() + ); + let level1 = impacts.max_score_up_to_cached(0, 767, 1.0, &scorer, &mut cache); + assert!(level1.score.is_finite() && level1.score > 0.0); + } + + #[test] + fn v2_and_v3_impacts_use_identical_encoding() { + let docs = vec![(3, 1, 100), (9, 2, 40), (200, 7, 80)]; + let mut v2_builder = ImpactSkipDataBuilder::with_capacity(1, 128); + v2_builder.append_block(&docs).unwrap(); + let v2 = v2_builder.finish().unwrap(); + let mut v3_builder = ImpactSkipDataBuilder::with_capacity(1, 256); + v3_builder.append_block(&docs).unwrap(); + let v3 = v3_builder.finish().unwrap(); + + assert_eq!(v2.entries(), v3.entries()); + } + + #[test] + fn impact_upper_bound_covers_real_scores() { + let blocks = vec![ + vec![(0, 1, 100), (3, 2, 40), (7, 4, 80)], + vec![(9, 3, 15), (10, 1, 5), (12, 5, 30)], + vec![(16, 2, 10), (18, 6, 70), (21, 3, 12)], + vec![(24, 1, 4), (28, 7, 100), (30, 2, 8)], + ]; + let impacts = build_impact_skip_data(&blocks).unwrap(); + let scorer = MemBM25Scorer::new(474, 31, HashMap::from([(String::from("token"), 4usize)])); + let query_weight = scorer.query_weight("token"); + let mut cache = ImpactScoreCache::default(); + + for start_block_idx in 0..blocks.len() { + let up_to = blocks + .iter() + .skip(start_block_idx) + .take(2) + .flatten() + .map(|(doc_id, _, _)| *doc_id) + .max() + .unwrap(); + let upper_bound = impacts.max_score_up_to_cached( + start_block_idx, + u64::from(up_to), + query_weight, + &scorer, + &mut cache, + ); + let exact_max = blocks + .iter() + .skip(start_block_idx) + .flatten() + .take_while(|(doc_id, _, _)| *doc_id <= up_to) + .map(|(_, freq, doc_len)| query_weight * scorer.doc_weight(*freq, *doc_len)) + .fold(0.0_f32, f32::max); + assert!( + upper_bound.score + 1e-6 >= exact_max, + "upper bound {} should cover exact max {} from block {} up to doc {}", + upper_bound.score, + exact_max, + start_block_idx, + up_to + ); + } + } +} diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index fc0b02ec209..21aadf7d475 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -1,15 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::fmt::{Debug, Display}; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; +use std::sync::{Arc, OnceLock}; use std::{ cmp::{Reverse, min}, collections::BinaryHeap, }; use std::{ - collections::{BTreeMap, HashMap, HashSet}, + collections::{BTreeMap, BTreeSet, HashMap, HashSet}, ops::Range, time::Instant, }; @@ -50,16 +51,19 @@ use lance_select::{RowAddrMask, RowAddrTreeMap}; use roaring::RoaringBitmap; use std::sync::LazyLock; use tokio::{sync::OnceCell, task::spawn_blocking}; -use tracing::{info, instrument}; +use tracing::{info, instrument, warn}; -use super::encoding::{PositionBlockBuilder, decode_group_starts}; +use super::encoding::{MAX_POSTING_BLOCK_SIZE, PositionBlockBuilder}; +use super::impact::{IMPACT_LEVEL1_BLOCKS, ImpactSkipData, ImpactSkipDataBuilder}; use super::iter::PostingListIterator; use super::lazy_docset::LazyDocSet; +use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; use super::{InvertedIndexBuilder, InvertedIndexParams, wand::*}; use super::{ builder::{ - BLOCK_SIZE, PostingGroupAccumulator, PostingGroupConfig, ScoredDoc, doc_file_path, - inverted_list_schema_for_version, posting_file_path, token_file_path, + BLOCK_SIZE, ScoredDoc, doc_file_path, + inverted_list_schema_for_version_with_block_size_and_impacts, posting_file_path, + token_file_path, }, iter::PlainPostingListIterator, query::*, @@ -69,14 +73,14 @@ use super::{ builder::{InnerBuilder, PositionRecorder}, iter::CompressedPostingListIterator, }; -use crate::frag_reuse::FragReuseIndex; use crate::pbold; use crate::progress::IndexBuildProgress; use crate::scalar::inverted::scorer::MemBM25Scorer; use crate::scalar::inverted::tokenizer::document_tokenizer::LanceTokenizer; use crate::scalar::{ AnyQuery, BuiltinIndexType, CreatedIndex, IndexReader, IndexStore, MetricsCollector, - OldIndexDataFilter, ScalarIndex, ScalarIndexParams, SearchResult, TokenQuery, UpdateCriteria, + OldIndexDataFilter, RowIdRemapper, ScalarIndex, ScalarIndexParams, SearchResult, TokenQuery, + UpdateCriteria, }; use crate::{FtsPrewarmOptions, Index}; use crate::{prefilter::PreFilter, scalar::inverted::iter::take_fst_keys}; @@ -85,8 +89,10 @@ use std::str::FromStr; // Version 0: Arrow TokenSetFormat (legacy) // Version 1: Fst TokenSetFormat with per-doc compressed positions // Version 2: Fst TokenSetFormat with shared posting-list position streams. +// Version 3: Version 2 layout with 256-document physical posting blocks. pub const INVERTED_INDEX_VERSION_V1: u32 = 1; pub const INVERTED_INDEX_VERSION_V2: u32 = 2; +pub const INVERTED_INDEX_VERSION_V3: u32 = 3; pub const TOKENS_FILE: &str = "tokens.lance"; pub const INVERT_LIST_FILE: &str = "invert.lance"; pub const DOCS_FILE: &str = "docs.lance"; @@ -102,6 +108,7 @@ pub const POSITION_COL: &str = "_position"; pub const COMPRESSED_POSITION_COL: &str = "_compressed_position"; pub const POSITION_BLOCK_OFFSET_COL: &str = "_position_block_offset"; pub const POSTING_COL: &str = "_posting"; +pub const IMPACT_COL: &str = "_impacts"; pub const MAX_SCORE_COL: &str = "_max_score"; pub const LENGTH_COL: &str = "_length"; pub const BLOCK_MAX_SCORE_COL: &str = "_block_max_score"; @@ -109,13 +116,10 @@ pub const NUM_TOKEN_COL: &str = "_num_tokens"; pub const SCORE_COL: &str = "_score"; pub const TOKEN_SET_FORMAT_KEY: &str = "token_set_format"; pub const POSTING_TAIL_CODEC_KEY: &str = "posting_tail_codec"; +pub const FTS_FORMAT_VERSION_KEY: &str = "format_version"; pub const POSITIONS_LAYOUT_KEY: &str = "positions_layout"; pub const POSITIONS_CODEC_KEY: &str = "positions_codec"; -/// Schema-metadata key holding the 1-indexed global-buffer id of the -/// varint-delta-encoded posting-list cache-group boundaries (issue #7040). -/// Absent on indexes written before grouping was introduced, which fall back -/// to the per-token cache path. -pub const POSTING_GROUP_OFFSETS_BUF_KEY: &str = "posting_group_offsets_buf"; +pub const POSTING_BLOCK_SIZE_KEY: &str = "posting_block_size"; pub const POSTING_TAIL_CODEC_FIXED32_V1: &str = "fixed32_v1"; pub const POSTING_TAIL_CODEC_VARINT_DELTA_V1: &str = "varint_delta_v1"; pub const POSITIONS_LAYOUT_SHARED_STREAM_V2: &str = "shared_stream_v2"; @@ -133,26 +137,33 @@ pub static FTS_SCHEMA: LazyLock = static ROW_ID_SCHEMA: LazyLock = LazyLock::new(|| Arc::new(Schema::new(vec![ROW_ID_FIELD.clone()]))); -fn resolve_fts_format_version( +pub fn resolve_fts_format_version( value: Option<&str>, ) -> std::result::Result { - value.unwrap_or("1").parse() + match value { + Some(value) => value.parse(), + None => Ok(default_fts_format_version()), + } +} + +pub fn default_fts_format_version() -> InvertedListFormatVersion { + InvertedListFormatVersion::V2 } pub fn current_fts_format_version() -> InvertedListFormatVersion { - resolve_fts_format_version(std::env::var("LANCE_FTS_FORMAT_VERSION").ok().as_deref()) - .expect("failed to parse LANCE_FTS_FORMAT_VERSION") + default_fts_format_version() } pub fn max_supported_fts_format_version() -> InvertedListFormatVersion { - InvertedListFormatVersion::V2 + InvertedListFormatVersion::V3 } #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] pub enum InvertedListFormatVersion { - #[default] V1, + #[default] V2, + V3, } impl InvertedListFormatVersion { @@ -163,29 +174,50 @@ impl InvertedListFormatVersion { } } + pub fn from_posting_tail_codec_and_block_size( + codec: PostingTailCodec, + block_size: usize, + ) -> Result { + validate_block_size(block_size)?; + let format_version = match (codec, block_size) { + (PostingTailCodec::Fixed32, LEGACY_BLOCK_SIZE) => Self::V1, + (PostingTailCodec::VarintDelta, LEGACY_BLOCK_SIZE) => Self::V2, + (PostingTailCodec::VarintDelta, 256) => Self::V3, + (PostingTailCodec::Fixed32, 256) => { + return Err(Error::invalid_input( + "FTS format_version=3 requires the varint-delta posting tail codec".to_string(), + )); + } + _ => unreachable!("validate_block_size limits supported block sizes"), + }; + validate_format_version_block_size(format_version, block_size)?; + Ok(format_version) + } + pub fn index_version(self) -> u32 { match self { Self::V1 => INVERTED_INDEX_VERSION_V1, Self::V2 => INVERTED_INDEX_VERSION_V2, + Self::V3 => INVERTED_INDEX_VERSION_V3, } } pub fn posting_tail_codec(self) -> PostingTailCodec { match self { Self::V1 => PostingTailCodec::Fixed32, - Self::V2 => PostingTailCodec::VarintDelta, + Self::V2 | Self::V3 => PostingTailCodec::VarintDelta, } } pub fn position_codec(self) -> Option { match self { Self::V1 => None, - Self::V2 => Some(PositionStreamCodec::PackedDelta), + Self::V2 | Self::V3 => Some(PositionStreamCodec::PackedDelta), } } pub fn uses_shared_position_stream(self) -> bool { - matches!(self, Self::V2) + matches!(self, Self::V2 | Self::V3) } } @@ -196,14 +228,47 @@ impl FromStr for InvertedListFormatVersion { match s.trim() { "1" | "v1" | "V1" => Ok(Self::V1), "2" | "v2" | "V2" => Ok(Self::V2), + "3" | "v3" | "V3" => Ok(Self::V3), other => Err(Error::index(format!( - "unsupported FTS format version {}, expected 1 or 2", + "unsupported FTS format version {}, expected 1, 2, or 3", other ))), } } } +pub fn default_fts_format_version_for_block_size( + block_size: usize, +) -> Result { + validate_block_size(block_size)?; + match block_size { + LEGACY_BLOCK_SIZE => Ok(InvertedListFormatVersion::V2), + 256 => Ok(InvertedListFormatVersion::V3), + _ => unreachable!("validate_block_size limits supported block sizes"), + } +} + +pub fn validate_format_version_block_size( + format_version: InvertedListFormatVersion, + block_size: usize, +) -> Result<()> { + validate_block_size(block_size)?; + match (format_version, block_size) { + (InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2, LEGACY_BLOCK_SIZE) + | (InvertedListFormatVersion::V3, 256) => Ok(()), + (InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2, 256) => { + Err(Error::invalid_input(format!( + "FTS format_version={} is incompatible with block_size=256; use format_version=3", + format_version.index_version() + ))) + } + (InvertedListFormatVersion::V3, other) => Err(Error::invalid_input(format!( + "FTS format_version=3 requires block_size=256, got {other}" + ))), + _ => unreachable!("validate_block_size limits supported block sizes"), + } +} + #[derive(Debug)] struct PartitionCandidates { tokens_by_position: Vec, @@ -225,6 +290,7 @@ impl PartitionCandidates { struct LoadedPostings { postings: Vec, grouped_expansions: Vec, + impact_safe: bool, } impl LoadedPostings { @@ -232,6 +298,7 @@ impl LoadedPostings { Self { postings: Vec::new(), grouped_expansions: Vec::new(), + impact_safe: false, } } } @@ -365,6 +432,21 @@ pub(super) fn parse_posting_tail_codec( .unwrap_or(PostingTailCodec::Fixed32)) } +pub(super) fn parse_posting_block_size(metadata: &HashMap) -> Result { + metadata + .get(POSTING_BLOCK_SIZE_KEY) + .map(|value| { + let block_size = value.parse::().map_err(|err| { + Error::index(format!( + "invalid {POSTING_BLOCK_SIZE_KEY} metadata value {value:?}: {err}" + )) + })?; + validate_block_size(block_size) + }) + .transpose() + .map(|block_size| block_size.unwrap_or(LEGACY_BLOCK_SIZE)) +} + impl PositionStreamCodec { pub fn as_str(self) -> &'static str { match self { @@ -402,6 +484,25 @@ fn parse_shared_position_codec(metadata: &HashMap) -> Result, ) -> Result { + if let Some(value) = metadata.get(FTS_FORMAT_VERSION_KEY) { + let format_version = InvertedListFormatVersion::from_str(value)?; + validate_format_version_block_size(format_version, parse_posting_block_size(metadata)?)?; + return Ok(format_version); + } + let block_size = parse_posting_block_size(metadata)?; + if block_size == 256 { + if metadata + .get(POSTING_TAIL_CODEC_KEY) + .map(|_| parse_posting_tail_codec(metadata)) + .transpose()? + .is_some_and(|posting_tail_codec| posting_tail_codec != PostingTailCodec::VarintDelta) + { + return Err(Error::index( + "FTS block_size=256 requires the varint-delta posting tail codec".to_string(), + )); + } + return Ok(InvertedListFormatVersion::V3); + } if metadata.contains_key(POSITIONS_CODEC_KEY) || metadata.contains_key(POSITIONS_LAYOUT_KEY) { return Ok(InvertedListFormatVersion::V2); } @@ -418,6 +519,7 @@ pub struct InvertedIndex { store: Arc, tokenizer: Box, token_set_format: TokenSetFormat, + format_version: InvertedListFormatVersion, pub(crate) partitions: Vec>, corpus_stats: Arc>, // Fragments which are contained in the index, but no longer in the dataset. @@ -430,6 +532,7 @@ impl Debug for InvertedIndex { f.debug_struct("InvertedIndex") .field("params", &self.params) .field("token_set_format", &self.token_set_format) + .field("format_version", &self.format_version) .field("partitions", &self.partitions) .field("deleted_fragments", &self.deleted_fragments) .finish() @@ -473,20 +576,16 @@ async fn resolve_deferred_candidates( impl InvertedIndex { fn format_version(&self) -> InvertedListFormatVersion { - self.partitions - .first() - .map(|partition| { - InvertedListFormatVersion::from_posting_tail_codec( - partition.inverted_list.posting_tail_codec(), - ) - }) - .unwrap_or_else(current_fts_format_version) + self.format_version } fn index_version(&self) -> u32 { - match self.token_set_format { - TokenSetFormat::Arrow => 0, - TokenSetFormat::Fst => self.format_version().index_version(), + match (self.token_set_format, self.format_version()) { + ( + TokenSetFormat::Arrow, + InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2, + ) => 0, + (_, format_version) => format_version.index_version(), } } @@ -628,22 +727,25 @@ impl InvertedIndex { query_tokens: &Tokens, params: &FtsSearchParams, ) -> Result { - let (total_tokens, num_docs) = self.aggregate_corpus_stats().await?; - let mut terms: Vec = Vec::new(); - let mut seen = HashSet::new(); if matches!(params.fuzziness, Some(n) if n != 0) { let expanded = self.expand_fuzzy_tokens(query_tokens, params)?; - for idx in 0..expanded.len() { - let token = expanded.get_token(idx); - if seen.insert(token.to_string()) { - terms.push(token.to_string()); - } - } + self.bm25_scorer_for_final_tokens(&expanded).await } else { - for token in query_tokens { - if seen.insert(token.to_string()) { - terms.push(token.to_string()); - } + self.bm25_scorer_for_final_tokens(query_tokens).await + } + } + + /// Scorer for a token list that needs no further fuzzy expansion: dedup + /// the terms and pull their document frequencies. `bm25_search` calls + /// this with the tokens it already expanded, so the expansion runs once + /// per query rather than once for the scorer and once per partition. + async fn bm25_scorer_for_final_tokens(&self, tokens: &Tokens) -> Result { + let (total_tokens, num_docs) = self.aggregate_corpus_stats().await?; + let mut terms: Vec = Vec::new(); + let mut seen = HashSet::new(); + for token in tokens { + if seen.insert(token.to_string()) { + terms.push(token.to_string()); } } let mut token_docs = HashMap::with_capacity(terms.len()); @@ -715,17 +817,46 @@ impl InvertedIndex { } /// Expand fuzzy query tokens against all partitions in this segment. + /// + /// `params.max_expansions` caps the whole query's expansion, not any + /// single partition's: for each query token the per-partition candidates + /// (each streamed in FST key order) merge into one lexicographically + /// ordered set, and the remaining budget takes a prefix of it. The + /// selected terms are a pure function of the segment's vocabulary, so + /// splitting the same corpus into more partitions cannot change which + /// terms a fuzzy query matches. pub fn expand_fuzzy_tokens(&self, tokens: &Tokens, params: &FtsSearchParams) -> Result { let mut expanded_tokens = Vec::new(); let mut expanded_positions = Vec::new(); let mut seen = HashSet::new(); - for partition in &self.partitions { - let expanded = partition.expand_fuzzy(tokens, params)?; - for idx in 0..expanded.len() { - let token = expanded.get_token(idx); - let position = expanded.position(idx); - if seen.insert((token.to_string(), position)) { - expanded_tokens.push(token.to_string()); + for token_idx in 0..tokens.len() { + let remaining = params.max_expansions.saturating_sub(expanded_tokens.len()); + if remaining == 0 { + break; + } + let token = tokens.get_token(token_idx); + let position = tokens.position(token_idx); + // Each partition contributes at most its `remaining` + // lexicographically smallest candidates, so the global + // lex-smallest `remaining` selection below is unaffected by the + // per-partition truncation. + let mut candidates = BTreeSet::new(); + let base_prefix_len = tokens.token_type().prefix_len(token) as u32; + for partition in &self.partitions { + partition.collect_fuzzy_candidates( + token, + base_prefix_len, + params, + remaining, + &mut candidates, + )?; + } + for candidate in candidates { + if expanded_tokens.len() >= params.max_expansions { + break; + } + if seen.insert((candidate.clone(), position)) { + expanded_tokens.push(candidate); expanded_positions.push(position); } } @@ -751,19 +882,40 @@ impl InvertedIndex { metrics: Arc, base_scorer: Option<&MemBM25Scorer>, ) -> Result<(Vec, Vec)> { + // Fuzzy expansion runs once here, with the global `max_expansions` + // budget, instead of once per partition: partitions receive the + // final token list, so the matched terms cannot depend on how the + // corpus happens to be partitioned. + let tokens = if matches!(params.fuzziness, Some(n) if n != 0) { + let expanded = Arc::new(self.expand_fuzzy_tokens(tokens.as_ref(), params.as_ref())?); + if operator == Operator::And || params.phrase_slop.is_some() { + // AND/phrase semantics require every original token position + // to keep at least one expansion; a position that expands to + // nothing anywhere in the segment can never be matched. + let surviving = (0..expanded.len()) + .map(|idx| expanded.position(idx)) + .collect::>(); + if (0..tokens.len()).any(|idx| !surviving.contains(&tokens.position(idx))) { + return Ok((Vec::new(), Vec::new())); + } + } + expanded + } else { + tokens + }; + // The wand only consults `scorer.doc_weight`, which is metadata-free. // The outer aggregation below consults `scorer.query_weight`, which // hits per-token `posting_len`; building a `MemBM25Scorer` with // precomputed per-term IDFs avoids the v2 bulk metadata pull. let local_scorer; - let scorer: &dyn Scorer = if let Some(base_scorer) = base_scorer { + let scorer: &MemBM25Scorer = if let Some(base_scorer) = base_scorer { base_scorer } else { - local_scorer = self - .bm25_base_scorer(tokens.as_ref(), params.as_ref()) - .await?; + local_scorer = self.bm25_scorer_for_final_tokens(tokens.as_ref()).await?; &local_scorer }; + let impact_scorer = Arc::new(scorer.clone()); let limit = params.limit.unwrap_or(usize::MAX); if limit == 0 { @@ -802,8 +954,9 @@ impl InvertedIndex { // Shared top-k floor across this query's partitions. Seeded to -inf so // the first real score wins; each partition publishes its local k-th // and prunes against the running global k-th (a lower bound on the true - // global k-th — see `Wand::shared_threshold`). - let shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + // global k-th - see `Wand::shared_threshold`). + let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + let legacy_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); let parts = self .partitions .iter() @@ -813,19 +966,23 @@ impl InvertedIndex { let params = params.clone(); let mask = mask.clone(); let metrics = metrics.clone(); - let shared_threshold = shared_threshold.clone(); + let impact_scorer = impact_scorer.clone(); + let impact_shared_threshold = impact_shared_threshold.clone(); + let legacy_shared_threshold = legacy_shared_threshold.clone(); async move { let loaded_postings = part .load_posting_lists( tokens.as_ref(), params.as_ref(), operator, + impact_scorer.as_ref(), metrics.as_ref(), ) .await?; let LoadedPostings { postings, grouped_expansions, + impact_safe, } = loaded_postings; if postings.is_empty() { // No hits in this partition; its DocSet stays @@ -849,6 +1006,7 @@ impl InvertedIndex { let metrics = metrics.clone(); let part_for_wand = part.clone(); let has_grouped_expansions = !grouped_expansions.is_empty(); + let use_impact_path = impact_safe && !has_grouped_expansions; let wand_params = if has_grouped_expansions { let mut rescoring_params = params.as_ref().clone(); rescoring_params.limit = @@ -859,9 +1017,12 @@ impl InvertedIndex { }; let partition_threshold = if has_grouped_expansions { Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())) + } else if use_impact_path { + impact_shared_threshold } else { - shared_threshold + legacy_shared_threshold }; + let wand_scorer = use_impact_path.then(|| impact_scorer.clone()); let candidates = spawn_cpu(move || { let candidates = part_for_wand.bm25_search( docs_for_wand.as_ref(), @@ -869,6 +1030,7 @@ impl InvertedIndex { operator, mask, postings, + wand_scorer, metrics.as_ref(), partition_threshold, )?; @@ -977,7 +1139,7 @@ impl InvertedIndex { async fn load_legacy_index( store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, index_cache: &LanceCache, ) -> Result> { log::warn!("loading legacy FTS index"); @@ -1026,6 +1188,7 @@ impl InvertedIndex { store: store.clone(), tokenizer, token_set_format: TokenSetFormat::Arrow, + format_version: InvertedListFormatVersion::V1, partitions: vec![Arc::new(InvertedPartition { id: 0, store, @@ -1045,7 +1208,7 @@ impl InvertedIndex { pub async fn load( store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, index_cache: &LanceCache, ) -> Result> where @@ -1076,6 +1239,7 @@ impl InvertedIndex { .map(|name| TokenSetFormat::from_str(name)) .transpose()? .unwrap_or(TokenSetFormat::Arrow); + let format_version = parse_format_version_from_metadata(&reader.schema().metadata)?; // Load deleted_fragments if present (optional for backward compatibility) let deleted_fragments = if reader.num_rows() > 0 { @@ -1121,6 +1285,7 @@ impl InvertedIndex { store, tokenizer, token_set_format, + format_version, partitions, corpus_stats: Arc::new(OnceCell::new()), deleted_fragments, @@ -1187,6 +1352,89 @@ const PREWARM_MAX_CHUNK_TOKENS: usize = 256 * 1024; /// Floor on token rows per chunk, so a partition always makes progress. const PREWARM_MIN_CHUNK_TOKENS: usize = 1; +/// Maximum number of posting lists in a runtime synthetic cache group. This is +/// deliberately token-count based so grouping works for old v2 indexes without +/// scanning posting lengths or requiring index rebuilds. +static LANCE_FTS_POSTING_GROUP_MAX_TOKENS: LazyLock = LazyLock::new(|| { + std::env::var("LANCE_FTS_POSTING_GROUP_MAX_TOKENS") + .unwrap_or_else(|_| "128".to_string()) + .parse() + .expect("failed to parse LANCE_FTS_POSTING_GROUP_MAX_TOKENS") +}); + +fn runtime_posting_group_tokens() -> usize { + (*LANCE_FTS_POSTING_GROUP_MAX_TOKENS).max(1) +} + +/// Runtime posting-list cache grouping. Non-empty v2 indexes synthesize fixed +/// groups at read time so prewarm and queries share group cache entries without +/// persisted grouping metadata or index rebuilds. +#[derive(Debug, Clone, DeepSizeOf)] +enum PostingGrouping { + /// Leaves legacy or empty partitions ungrouped. + None, + /// Uses a fixed runtime cache group size measured in token rows, not posting bytes. + SyntheticFixed { group_size: u32 }, +} + +impl PostingGrouping { + fn for_reader(is_legacy_layout: bool, token_count: usize) -> Self { + if is_legacy_layout || token_count == 0 { + return Self::None; + } + + let group_size = u32::try_from(runtime_posting_group_tokens()) + .unwrap_or(u32::MAX) + .max(1); + Self::SyntheticFixed { group_size } + } + + fn is_grouped(&self) -> bool { + !matches!(self, Self::None) + } + + fn range_for_token(&self, token_id: u32, token_count: usize) -> Option<(u32, u32)> { + match self { + Self::None => None, + Self::SyntheticFixed { group_size } => { + let token_count = u32::try_from(token_count).unwrap_or(u32::MAX); + let start = (token_id / *group_size) * *group_size; + let end = start.saturating_add(*group_size).min(token_count); + Some((start, end)) + } + } + } + + fn aligned_chunk_end(&self, token_count: usize, tok_start: usize, desired_end: usize) -> usize { + match self { + Self::None => desired_end, + Self::SyntheticFixed { group_size } => synthetic_group_aligned_chunk_end( + usize::try_from(*group_size).unwrap_or(usize::MAX).max(1), + token_count, + tok_start, + desired_end, + ), + } + } + + fn ranges_for_chunk( + &self, + tok_start: usize, + tok_end: usize, + token_count: usize, + ) -> Vec<(u32, u32)> { + match self { + Self::None => Vec::new(), + Self::SyntheticFixed { group_size } => synthetic_group_ranges_for_chunk( + usize::try_from(*group_size).unwrap_or(usize::MAX).max(1), + tok_start, + tok_end, + token_count, + ), + } + } +} + /// Token rows per chunk: byte target / average bytes-per-token, clamped to `[MIN, MAX]`. fn prewarm_chunk_tokens(token_count: usize, file_size_bytes: u64) -> usize { if token_count == 0 { @@ -1197,11 +1445,8 @@ fn prewarm_chunk_tokens(token_count: usize, file_size_bytes: u64) -> usize { by_bytes.clamp(PREWARM_MIN_CHUNK_TOKENS, PREWARM_MAX_CHUNK_TOKENS) } -/// Snap a chunk's exclusive token end back to a posting-group boundary so no group -/// straddles chunks. Returns the largest group boundary in `(tok_start, desired_end]`, -/// or the next boundary past an oversized group so it runs as one solo chunk. -fn group_aligned_chunk_end( - starts: &[u32], +fn synthetic_group_aligned_chunk_end( + group_size: usize, token_count: usize, tok_start: usize, desired_end: usize, @@ -1210,21 +1455,38 @@ fn group_aligned_chunk_end( return token_count; } - let first_after_start = starts.partition_point(|&start| start as usize <= tok_start); - let first_after_desired = starts.partition_point(|&start| start as usize <= desired_end); - if first_after_desired > first_after_start { - return starts[first_after_desired - 1] as usize; + let boundary = desired_end - (desired_end % group_size); + if boundary > tok_start { + boundary + } else { + tok_start.saturating_add(group_size).min(token_count) } +} - // Oversized group: extend to its end so it runs as one chunk. - starts - .get(first_after_start) - .map(|&start| start as usize) - .unwrap_or(token_count) +fn synthetic_group_ranges_for_chunk( + group_size: usize, + tok_start: usize, + tok_end: usize, + token_count: usize, +) -> Vec<(u32, u32)> { + let mut ranges = Vec::new(); + let mut start = tok_start - (tok_start % group_size); + if start < tok_start { + start = start.saturating_add(group_size).min(token_count); + } + while start < tok_end { + let end = start.saturating_add(group_size).min(token_count); + ranges.push(( + u32::try_from(start).unwrap_or(u32::MAX), + u32::try_from(end).unwrap_or(u32::MAX), + )); + start = end; + } + ranges } fn prewarm_chunk_ranges( - group_starts: Option<&[u32]>, + grouping: &PostingGrouping, token_count: usize, chunk_tokens: usize, ) -> Vec<(usize, usize)> { @@ -1233,8 +1495,8 @@ fn prewarm_chunk_ranges( while tok_start < token_count { let mut tok_end = (tok_start + chunk_tokens).min(token_count); // `tok_start` is always a group boundary; snap `tok_end` back to one too. - if let Some(starts) = group_starts { - tok_end = group_aligned_chunk_end(starts, token_count, tok_start, tok_end); + if grouping.is_grouped() { + tok_end = grouping.aligned_chunk_end(token_count, tok_start, tok_end); } ranges.push((tok_start, tok_end)); tok_start = tok_end; @@ -1242,35 +1504,69 @@ fn prewarm_chunk_ranges( ranges } -fn group_start_indices_for_chunk(starts: &[u32], tok_start: usize, tok_end: usize) -> Range { - let first = starts.partition_point(|&start| (start as usize) < tok_start); - let end = starts.partition_point(|&start| (start as usize) < tok_end); - first..end -} - -fn group_range_for_start_index(starts: &[u32], token_count: usize, group_idx: usize) -> (u32, u32) { - let start = starts[group_idx]; - let end = starts - .get(group_idx + 1) - .copied() - .unwrap_or(token_count as u32); - (start, end) -} - impl InvertedIndex { pub async fn prewarm_with_options(&self, options: &FtsPrewarmOptions) -> Result<()> { let with_position = options.with_position; let chunk_concurrency = self.store.io_parallelism().max(1); + let prewarm_started = Instant::now(); + info!( + partition_count = self.partitions.len(), + with_position, chunk_concurrency, "fts index prewarm started" + ); for part in &self.partitions { - part.inverted_list + let partition_started = Instant::now(); + info!( + partition_id = part.id(), + token_count = part.tokens.len(), + with_position, + chunk_concurrency, + "fts partition prewarm started" + ); + if let Err(err) = part + .inverted_list .prewarm_posting_lists(with_position, chunk_concurrency) - .await?; + .await + { + warn!( + partition_id = part.id(), + error = %err, + elapsed_ms = partition_started.elapsed().as_millis() as u64, + "fts partition posting list prewarm failed" + ); + return Err(err); + } + info!( + partition_id = part.id(), + elapsed_ms = partition_started.elapsed().as_millis() as u64, + "fts partition posting lists prewarmed" + ); // Materialize the deferred DocSet too: prewarm's contract is // that subsequent queries do no IO, so the per-doc row_ids / // num_tokens must be resident, not lazily faulted in at query // time. `ensure_loaded` opens, reads, and drops the reader. - part.docs.ensure_loaded().await?; + let docs_started = Instant::now(); + if let Err(err) = part.docs.ensure_loaded().await { + warn!( + partition_id = part.id(), + error = %err, + elapsed_ms = docs_started.elapsed().as_millis() as u64, + total_elapsed_ms = partition_started.elapsed().as_millis() as u64, + "fts partition docset prewarm failed" + ); + return Err(err); + } + info!( + partition_id = part.id(), + docset_elapsed_ms = docs_started.elapsed().as_millis() as u64, + elapsed_ms = partition_started.elapsed().as_millis() as u64, + "fts partition prewarm finished" + ); } + info!( + partition_count = self.partitions.len(), + elapsed_ms = prewarm_started.elapsed().as_millis() as u64, + "fts index prewarm finished" + ); Ok(()) } /// Search docs match the input text. @@ -1329,7 +1625,7 @@ impl ScalarIndex for InvertedIndex { async fn remap( &self, - mapping: &HashMap>, + mapping: &RowAddrRemap, dest_store: &dyn IndexStore, ) -> Result { let files = self @@ -1381,8 +1677,9 @@ impl ScalarIndex for InvertedIndex { // Empty tokenizer metadata only appears in legacy simple-tokenizer indexes. params.base_tokenizer = "simple".to_string(); } + params = params.format_version(self.format_version()); - let params_json = serde_json::to_string(¶ms)?; + let params_json = params.to_training_json()?.to_string(); Ok(ScalarIndexParams { index_type: BuiltinIndexType::Inverted.as_str().to_string(), @@ -1435,7 +1732,7 @@ impl InvertedPartition { pub async fn load( store: Arc, id: u64, - frag_reuse_index: Option>, + frag_reuse_index: Option>, index_cache: &LanceCache, token_set_format: TokenSetFormat, ) -> Result { @@ -1457,6 +1754,8 @@ impl InvertedPartition { num_docs, false, frag_reuse_index, + // V3 (256-doc block) partitions score with quantized doc lengths. + inverted_list.block_size() == MAX_POSTING_BLOCK_SIZE, )); Ok(Self { @@ -1478,47 +1777,29 @@ impl InvertedPartition { let mut new_positions = Vec::with_capacity(new_tokens.capacity()); let mut seen = HashSet::new(); for token_idx in 0..tokens.len() { - if new_tokens.len() >= params.max_expansions { + let remaining = params.max_expansions.saturating_sub(new_tokens.len()); + if remaining == 0 { break; } let token = tokens.get_token(token_idx); let position = tokens.position(token_idx); - let fuzziness = match params.fuzziness { - Some(fuzziness) => fuzziness, - None => MatchQuery::auto_fuzziness(token), - }; - let lev = fst::automaton::Levenshtein::new(token, fuzziness) - .map_err(|e| Error::index(format!("failed to construct the fuzzy query: {}", e)))?; - - let base_len = tokens.token_type().prefix_len(token) as u32; - if let TokenMap::Fst(ref map) = self.tokens.tokens { - let mut expanded = Vec::new(); - let remaining = params.max_expansions - new_tokens.len(); - match base_len + params.prefix_length { - 0 => take_fst_keys(map.search(lev), &mut expanded, remaining), - prefix_length => { - let prefix = &token[..min(prefix_length as usize, token.len())]; - let prefix = fst::automaton::Str::new(prefix).starts_with(); - take_fst_keys( - map.search(lev.intersection(prefix)), - &mut expanded, - remaining, - ) - } + let base_prefix_len = tokens.token_type().prefix_len(token) as u32; + let mut candidates = BTreeSet::new(); + self.collect_fuzzy_candidates( + token, + base_prefix_len, + params, + remaining, + &mut candidates, + )?; + for candidate in candidates { + if new_tokens.len() >= params.max_expansions { + break; } - for token in expanded { - if seen.insert((token.clone(), position)) { - new_tokens.push(token); - new_positions.push(position); - if new_tokens.len() >= params.max_expansions { - break; - } - } + if seen.insert((candidate.clone(), position)) { + new_tokens.push(candidate); + new_positions.push(position); } - } else { - return Err(Error::index( - "tokens is not fst, which is not expected".to_owned(), - )); } } Ok(Tokens::with_positions( @@ -1528,6 +1809,47 @@ impl InvertedPartition { )) } + /// Collect up to `limit` fuzzy candidates for one query token from this + /// partition's token FST, in key (lexicographic) order. Callers merge + /// candidates across partitions and apply the query-wide + /// `max_expansions` budget; truncating each partition at `limit` is + /// lossless for that selection because any term among the merged + /// lexicographically-smallest `limit` is also among its own partition's + /// smallest `limit`. + fn collect_fuzzy_candidates( + &self, + token: &str, + base_prefix_len: u32, + params: &FtsSearchParams, + limit: usize, + candidates: &mut BTreeSet, + ) -> Result<()> { + let fuzziness = match params.fuzziness { + Some(fuzziness) => fuzziness, + None => MatchQuery::auto_fuzziness(token), + }; + let lev = fst::automaton::Levenshtein::new(token, fuzziness) + .map_err(|e| Error::index(format!("failed to construct the fuzzy query: {}", e)))?; + + if let TokenMap::Fst(ref map) = self.tokens.tokens { + let mut expanded = Vec::new(); + match base_prefix_len + params.prefix_length { + 0 => take_fst_keys(map.search(lev), &mut expanded, limit), + prefix_length => { + let prefix = &token[..min(prefix_length as usize, token.len())]; + let prefix = fst::automaton::Str::new(prefix).starts_with(); + take_fst_keys(map.search(lev.intersection(prefix)), &mut expanded, limit) + } + } + candidates.extend(expanded); + Ok(()) + } else { + Err(Error::index( + "tokens is not fst, which is not expected".to_owned(), + )) + } + } + fn union_plain_posting_lists(postings: Vec) -> Result { let mut freqs_by_row_id = BTreeMap::new(); for posting in postings { @@ -1556,6 +1878,13 @@ impl InvertedPartition { postings: Vec, docs: &DocSet, ) -> Result { + let block_size = postings + .iter() + .find_map(|posting| match posting { + PostingList::Compressed(posting) => Some(posting.block_size), + PostingList::Plain(_) => None, + }) + .unwrap_or(LEGACY_BLOCK_SIZE); let mut freqs_by_doc_id = BTreeMap::new(); for posting in postings { for (doc_id, freq, _) in posting.iter() { @@ -1580,7 +1909,7 @@ impl InvertedPartition { ))); } - let mut builder = PostingListBuilder::new(false); + let mut builder = PostingListBuilder::new_with_block_size(false, block_size); let mut doc_ids = Vec::with_capacity(freqs_by_doc_id.len()); let mut frequencies = Vec::with_capacity(freqs_by_doc_id.len()); for (doc_id, freq) in freqs_by_doc_id { @@ -1588,7 +1917,11 @@ impl InvertedPartition { doc_ids.push(doc_id); frequencies.push(freq); } - let block_max_scores = docs.calculate_block_max_scores(doc_ids.iter(), frequencies.iter()); + let block_max_scores = docs.calculate_block_max_scores_with_block_size( + doc_ids.iter(), + frequencies.iter(), + block_size, + ); let batch = builder.to_batch(block_max_scores)?; let max_score = batch[MAX_SCORE_COL].as_primitive::().value(0); let length = batch[LENGTH_COL].as_primitive::().value(0); @@ -1626,6 +1959,7 @@ impl InvertedPartition { tokens: &Tokens, params: &FtsSearchParams, operator: Operator, + impact_scorer: &MemBM25Scorer, metrics: &dyn MetricsCollector, ) -> Result { let is_fuzzy = matches!(params.fuzziness, Some(n) if n != 0); @@ -1636,10 +1970,11 @@ impl InvertedPartition { .map(|index| tokens.position(index)) .collect::>() }); - let tokens = match is_fuzzy { - true => self.expand_fuzzy(tokens, params)?, - false => tokens.clone(), - }; + // Fuzzy expansion already ran once at the index level (see + // `InvertedIndex::bm25_search`) under the global `max_expansions` + // budget; the incoming tokens are final and `is_fuzzy` only drives + // the grouped dedup/scoring semantics below. + let tokens = tokens.clone(); let token_positions = (0..tokens.len()) .map(|index| tokens.position(index)) .collect::>(); @@ -1703,11 +2038,18 @@ impl InvertedPartition { } if !is_fuzzy_and_query { + let impact_safe = loaded_postings + .iter() + .all(|(_, _, _, posting)| posting.has_impacts()); return Ok(LoadedPostings { postings: loaded_postings .into_iter() .map(|(token_id, token, position, posting)| { - let query_weight = idf(posting.len(), num_docs); + let query_weight = if impact_safe { + impact_scorer.query_weight(&token) + } else { + idf(posting.len(), num_docs) + }; PostingIterator::with_query_weight( token, token_id, @@ -1719,6 +2061,7 @@ impl InvertedPartition { }) .collect(), grouped_expansions: Vec::new(), + impact_safe, }); } @@ -1786,6 +2129,7 @@ impl InvertedPartition { Ok(LoadedPostings { postings: grouped_postings, grouped_expansions, + impact_safe: false, }) } @@ -1801,6 +2145,7 @@ impl InvertedPartition { operator: Operator, mask: Arc, postings: Vec, + impact_scorer: Option>, metrics: &dyn MetricsCollector, shared_threshold: Arc, ) -> Result> { @@ -1811,19 +2156,26 @@ impl InvertedPartition { // Caller selects the DocSet shape via `LazyDocSet::docs_for_wand` // and passes it in here; wand uses `docs.has_row_ids()` to // handle the num_tokens-only case. - let scorer = IndexBM25Scorer::new(std::iter::once(self)); - let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) - .with_shared_threshold(shared_threshold); - let hits = wand.search(params, mask, metrics)?; + let hits = if let Some(scorer) = impact_scorer { + let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) + .with_shared_threshold(shared_threshold); + wand.search(params, mask, metrics)? + } else { + let scorer = IndexBM25Scorer::new(std::iter::once(self)); + let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) + .with_shared_threshold(shared_threshold); + wand.search(params, mask, metrics)? + }; Ok(hits) } pub async fn into_builder(self) -> Result { - let mut builder = InnerBuilder::new_with_posting_tail_codec( + let mut builder = InnerBuilder::new_with_posting_tail_codec_and_block_size( self.id, self.inverted_list.has_positions(), self.token_set_format, self.inverted_list.posting_tail_codec(), + self.inverted_list.block_size(), ); builder.tokens = self.tokens.into_mutable(); // into_builder rewrites every doc, so materialize the full @@ -2232,14 +2584,15 @@ pub struct PostingListReader { metadata: PostingMetadata, has_position: bool, + has_impacts: bool, posting_tail_codec: PostingTailCodec, + block_size: usize, positions_layout: PositionsLayout, - /// First row of each posting-list cache group, decoded at open from the - /// global buffer named by [`POSTING_GROUP_OFFSETS_BUF_KEY`] (issue #7040). - /// `None` for indexes written before grouping; those use the per-token - /// cache path. Always present for grouped v2 indexes with `>0` rows. - group_starts: Option>, + /// Runtime posting-list cache grouping. Non-empty v2 indexes use synthetic + /// fixed groups so prewarm can improve cache density without rebuilding the + /// index or relying on persisted grouping metadata. + grouping: PostingGrouping, index_cache: WeakLanceCache, } @@ -2314,7 +2667,7 @@ impl DeepSizeOf for PostingListReader { }) .unwrap_or(0), }; - metadata_size + self.group_starts.deep_size_of_children(context) + metadata_size + self.grouping.deep_size_of_children(context) } } @@ -2331,7 +2684,9 @@ impl PostingListReader { PositionsLayout::None }; let posting_tail_codec = parse_posting_tail_codec(&reader.schema().metadata)?; + let block_size = parse_posting_block_size(&reader.schema().metadata)?; let has_position = positions_layout != PositionsLayout::None; + let has_impacts = reader.schema().field(IMPACT_COL).is_some(); let metadata = if reader.schema().field(POSTING_COL).is_none() { let (offsets, max_scores) = Self::load_metadata(reader.schema())?; PostingMetadata::LegacyV1 { @@ -2344,36 +2699,22 @@ impl PostingListReader { } }; - let group_starts = Self::load_group_starts(reader.as_ref()).await?; + let is_legacy_layout = matches!(&metadata, PostingMetadata::LegacyV1 { .. }); + let grouping = PostingGrouping::for_reader(is_legacy_layout, reader.num_rows()); Ok(Self { reader, metadata, has_position, + has_impacts, posting_tail_codec, + block_size, positions_layout, - group_starts, + grouping, index_cache: WeakLanceCache::from(index_cache), }) } - /// Decode the posting-list cache-group boundaries from the global buffer - /// recorded in schema metadata, if present (issue #7040). Returns `None` - /// for indexes written before grouping was introduced. - async fn load_group_starts(reader: &dyn IndexReader) -> Result>> { - let Some(buf_id) = reader.schema().metadata.get(POSTING_GROUP_OFFSETS_BUF_KEY) else { - return Ok(None); - }; - let buf_id: u32 = buf_id.parse().map_err(|e| { - Error::index(format!( - "invalid {POSTING_GROUP_OFFSETS_BUF_KEY} metadata value {buf_id:?}: {e}" - )) - })?; - let bytes = reader.read_global_buffer(buf_id).await?; - let group_starts = decode_group_starts(&bytes)?; - Ok(Some(group_starts)) - } - // for legacy format // returns the offsets and max scores fn load_metadata( @@ -2413,6 +2754,10 @@ impl PostingListReader { self.posting_tail_codec } + pub(crate) fn block_size(&self) -> usize { + self.block_size + } + fn is_legacy_layout(&self) -> bool { matches!(self.metadata, PostingMetadata::LegacyV1 { .. }) } @@ -2541,7 +2886,7 @@ impl PostingListReader { self.posting_batch_legacy(token_id, with_position).await } else { let token_id = token_id as usize; - let columns = if with_position { + let mut columns = if with_position { match self.positions_layout { PositionsLayout::SharedStream(_) => { vec![ @@ -2556,6 +2901,9 @@ impl PostingListReader { } else { vec![POSTING_COL] }; + if self.has_impacts { + columns.push(IMPACT_COL); + } let batch = self .reader .read_range(token_id..token_id + 1, Some(&columns)) @@ -2600,39 +2948,49 @@ impl PostingListReader { Some((start, end)) => { let group = self .index_cache - .get_or_insert_with_key(PostingListGroupKey { start, end }, || async move { - metrics.record_part_load(); - info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="inverted", part_id=start); - self.load_posting_list_group(start, end).await - }) + .get_or_insert_with_key( + posting_list_group_cache_key(start, end, self.has_impacts), + || async move { + metrics.record_part_load(); + info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="inverted", part_id=start); + self.load_posting_list_group(start, end).await + }, + ) .await?; + let (max_score, length) = if group.needs_external_metadata() { + self.posting_metadata_for_token(token_id).await? + } else { + (None, None) + }; let slot = (token_id - start) as usize; group - .get(slot) + .posting_list(slot, max_score, length)? .ok_or_else(|| { Error::index(format!( "token {token_id} maps to slot {slot} outside posting group [{start}, {end})" )) })? - .clone() } - // Fallback for indexes written before grouping: one cache entry - // per token. + // Fallback for layouts that cannot use row-based groups: one cache + // entry per token. None => self .index_cache - .get_or_insert_with_key(PostingListKey { token_id }, || async move { - metrics.record_part_load(); - info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="inverted", part_id=token_id); - // Fetch the posting batch and this token's (max_score, - // length) in parallel; for cold v2 partitions this is one - // single-row metadata read plus one posting-row read, - // instead of pulling the full per-token metadata table. - let (batch, (max_score, length)) = futures::try_join!( - self.posting_batch(token_id, false), - self.posting_metadata_for_token(token_id), - )?; - self.posting_list_from_batch(&batch, max_score, length) - }) + .get_or_insert_with_key( + posting_list_cache_key(token_id, self.has_impacts), + || async move { + metrics.record_part_load(); + info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="inverted", part_id=token_id); + // Fetch the posting batch and this token's (max_score, + // length) in parallel; for cold v2 partitions this is one + // single-row metadata read plus one posting-row read, + // instead of pulling the full per-token metadata table. + let (batch, (max_score, length)) = futures::try_join!( + self.posting_batch(token_id, false), + self.posting_metadata_for_token(token_id), + )?; + self.posting_list_from_batch(&batch, max_score, length) + }, + ) .await? .as_ref() .clone(), @@ -2648,55 +3006,30 @@ impl PostingListReader { } /// Map a token id to its cache group's row range `[start, end)`, or `None` - /// when grouping is not available (pre-grouping indexes) so the caller - /// falls back to the per-token path. In v2 the token id is the row offset, - /// so the group range is also the physical row range. + /// when grouping is not available so the caller falls back to the per-token + /// path. In v2 the token id is the row offset, so the group range is also + /// the physical row range. fn group_range_for_token(&self, token_id: u32) -> Option<(u32, u32)> { - let starts = self.group_starts.as_ref()?; - // partition_point returns the count of group starts <= token_id, so the - // owning group begins at index k - 1 and the next start (if any) is its - // exclusive end. - let k = starts.partition_point(|&s| s <= token_id); - // k == 0 means token_id precedes the first group start, which cannot - // happen for a valid token in a grouped index (the first group starts - // at row 0); guard anyway and fall back to the per-token path. - if k == 0 { - return None; - } - let start = starts[k - 1]; - // The last group runs to the final posting list. `self.len()` is the - // authoritative posting-list count (offsets length for v1, row count for - // v2), and prewarm derives the same `end` from it — so warm- and - // cold-cache group keys are identical by construction, not by the - // incidental v2 `num_rows == token_count` equality. - let end = starts.get(k).copied().unwrap_or(self.len() as u32); - Some((start, end)) + self.grouping.range_for_token(token_id, self.len()) } - /// Read rows `[start, end)` of the posting file and decode them into a - /// [`PostingListGroup`] cache value (issue #7040). Positions are excluded; - /// phrase queries load them on demand via [`Self::read_positions`]. + /// Read rows `[start, end)` into one compact Arrow-backed cache value. + /// Positions are excluded; phrase queries load them on demand via + /// [`Self::read_positions`]. async fn load_posting_list_group(&self, start: u32, end: u32) -> Result { + let mut columns = vec![POSTING_COL, MAX_SCORE_COL, LENGTH_COL]; + if self.has_impacts { + columns.push(IMPACT_COL); + } let batch = self .reader - .read_range( - start as usize..end as usize, - Some(&[POSTING_COL, MAX_SCORE_COL, LENGTH_COL]), - ) + .read_range(start as usize..end as usize, Some(&columns)) .await?; - let max_scores = batch[MAX_SCORE_COL].as_primitive::(); - let lengths = batch[LENGTH_COL].as_primitive::(); - let mut posting_lists = Vec::with_capacity(batch.num_rows()); - for i in 0..batch.num_rows() { - let row = batch.slice(i, 1); - let posting = self.posting_list_from_batch( - &row, - Some(max_scores.value(i)), - Some(lengths.value(i)), - )?; - posting_lists.push(posting); - } - Ok(PostingListGroup::new(posting_lists)) + PostingListGroup::new_packed_with_block_size( + batch.shrink_to_fit()?, + self.posting_tail_codec, + self.block_size, + ) } fn posting_list_from_batch_parts( @@ -2704,6 +3037,7 @@ impl PostingListReader { max_score: Option, length: Option, posting_tail_codec: PostingTailCodec, + block_size: usize, positions_layout: PositionsLayout, ) -> Result { let posting_list = PostingList::from_batch_with_tail_codec_and_positions_layout( @@ -2711,6 +3045,7 @@ impl PostingListReader { max_score, length, posting_tail_codec, + block_size, positions_layout, )?; Ok(posting_list) @@ -2727,6 +3062,7 @@ impl PostingListReader { max_score, length, self.posting_tail_codec, + self.block_size, self.positions_layout, ) } @@ -2763,6 +3099,7 @@ impl PostingListReader { ctx.max_scores.map(|scores| scores[global]), ctx.lengths.map(|lengths| lengths[global]), ctx.posting_tail_codec, + ctx.block_size, ctx.positions_layout, )?; posting_lists.push((global as u32, posting_list)); @@ -2820,43 +3157,63 @@ impl PostingListReader { )); } - // Make sure max_scores/lengths are populated before we clone them into - // the blocking task; otherwise the v2 branch would unwrap empty - // OnceCells. + // Make max_scores/lengths available for query-local packed views. The + // materialized fallback also clones them into its blocking build task. self.ensure_metadata_loaded().await?; - let state = self.chunk_build_state(); - // With grouping the cache stores one entry per group, so a group's posting - // lists must all be resident at once: align chunk boundaries to whole - // groups. Without grouping, chunks are plain token ranges. - let group_starts = self.group_starts.clone(); + // With grouping the cache stores one entry per group, so a group's + // posting lists must all be resident at once: align chunk boundaries to + // whole groups. Without grouping, chunks are plain token ranges. + let grouping = self.grouping.clone(); + let use_packed_groups = grouping.is_grouped() && !with_position; + // Packed groups reuse the reader's bulk metadata at query time, so they + // do not need the temporary full-partition metadata clones used by the + // materialized fallback. + let state = (!use_packed_groups).then(|| self.chunk_build_state()); let token_count = self.len(); let posting_data_size_bytes = self.posting_data_size_bytes(); let chunk_tokens = chunk_tokens_override .unwrap_or_else(|| prewarm_chunk_tokens(token_count, posting_data_size_bytes)) .max(1); - let chunk_ranges = prewarm_chunk_ranges(group_starts.as_deref(), token_count, chunk_tokens); + let chunk_ranges = prewarm_chunk_ranges(&grouping, token_count, chunk_tokens); let chunk_count = chunk_ranges.len(); let chunk_concurrency = chunk_concurrency.max(1); let read_build_start = Instant::now(); stream::iter(chunk_ranges) .map(|(tok_start, tok_end)| { - let state = &state; - let group_starts = group_starts.as_deref(); + let state = state.as_ref(); + let grouping = &grouping; async move { - let posting_lists = self - .build_chunk_postings(tok_start, tok_end, with_position, state) - .await?; - self.publish_chunk_postings( - posting_lists, - group_starts, - tok_start, - tok_end, - token_count, - with_position, - ) - .await; + if use_packed_groups { + let groups = self + .build_packed_chunk_groups(tok_start, tok_end, token_count, grouping) + .await?; + for (start, end, group) in groups { + self.index_cache + .insert_with_key( + &posting_list_group_cache_key(start, end, self.has_impacts), + Arc::new(group), + ) + .await; + } + } else { + let state = state.expect( + "materialized prewarm must initialize posting-list build state", + ); + let posting_lists = self + .build_chunk_postings(tok_start, tok_end, with_position, state) + .await?; + self.publish_chunk_postings( + posting_lists, + grouping, + tok_start, + tok_end, + token_count, + with_position, + ) + .await; + } Result::Ok(()) } }) @@ -2899,6 +3256,7 @@ impl PostingListReader { max_scores: max_scores.map(Arc::new), lengths: lengths.map(Arc::new), posting_tail_codec: self.posting_tail_codec, + block_size: self.block_size, positions_layout: self.positions_layout, } } @@ -2932,12 +3290,14 @@ impl PostingListReader { let max_scores = state.max_scores.clone(); let lengths = state.lengths.clone(); let posting_tail_codec = state.posting_tail_codec; + let block_size = state.block_size; let positions_layout = state.positions_layout; let posting_lists = spawn_blocking(move || { let ctx = PrewarmBuildCtx { max_scores: max_scores.as_deref().map(|v| v.as_slice()), lengths: lengths.as_deref().map(|v| v.as_slice()), posting_tail_codec, + block_size, positions_layout, }; let chunk = PrewarmChunk { @@ -2966,6 +3326,52 @@ impl PostingListReader { Ok(posting_lists) } + /// Build compact v2 groups directly from one posting-row chunk. Each group + /// slice is deep-copied once, so it owns only its Arrow buffers without + /// materializing a `Vec` or retaining the full chunk. + async fn build_packed_chunk_groups( + &self, + tok_start: usize, + tok_end: usize, + token_count: usize, + grouping: &PostingGrouping, + ) -> Result> { + debug_assert!(grouping.is_grouped()); + debug_assert!(!self.is_legacy_layout()); + + let chunk_batch = self.read_chunk_batch(tok_start, tok_end, false).await?; + let ranges = grouping.ranges_for_chunk(tok_start, tok_end, token_count); + let posting_tail_codec = self.posting_tail_codec; + let block_size = self.block_size; + + spawn_blocking(move || { + let mut groups = Vec::with_capacity(ranges.len()); + for (start, end) in ranges { + let start_usize = start as usize; + let end_usize = end as usize; + let local_start = start_usize - tok_start; + let group_len = end_usize - start_usize; + let group_batch = chunk_batch.slice(local_start, group_len).shrink_to_fit()?; + groups.push(( + start, + end, + PostingListGroup::new_packed_with_block_size( + group_batch, + posting_tail_codec, + block_size, + )?, + )); + } + Result::Ok(groups) + }) + .await + .map_err(|err| { + Error::internal(format!( + "Failed to build packed prewarm posting groups in blocking task: {err}" + )) + })? + } + /// Strip positions into their own per-token cache entries (the posting cache /// holds positions-free lists), then populate the same cache keys the read /// path uses: grouped entries when grouping is active, per-token entries @@ -2973,14 +3379,26 @@ impl PostingListReader { async fn publish_chunk_postings( &self, posting_lists: Vec<(u32, PostingList)>, - group_starts: Option<&[u32]>, + grouping: &PostingGrouping, tok_start: usize, tok_end: usize, token_count: usize, with_position: bool, ) { - match group_starts { - Some(starts) => { + match grouping { + PostingGrouping::None => { + for (token_id, mut posting_list) in posting_lists { + self.cache_positions(&mut posting_list, token_id, with_position) + .await; + self.index_cache + .insert_with_key( + &posting_list_cache_key(token_id, self.has_impacts), + Arc::new(posting_list), + ) + .await; + } + } + PostingGrouping::SyntheticFixed { .. } => { let mut chunk_postings = Vec::with_capacity(posting_lists.len()); for (token_id, mut posting_list) in posting_lists { self.cache_positions(&mut posting_list, token_id, with_position) @@ -2991,23 +3409,16 @@ impl PostingListReader { // in it; `chunk_postings[i]` is token `tok_start + i`. The last // group's `end` derives from `token_count`, matching the read path // so both produce identical `PostingListGroupKey`s. - for group_idx in group_start_indices_for_chunk(starts, tok_start, tok_end) { - let (start, end) = group_range_for_start_index(starts, token_count, group_idx); + for (start, end) in grouping.ranges_for_chunk(tok_start, tok_end, token_count) { let start_usize = start as usize; let lo = start_usize - tok_start; let hi = end as usize - tok_start; let group = PostingListGroup::new(chunk_postings[lo..hi].to_vec()); self.index_cache - .insert_with_key(&PostingListGroupKey { start, end }, Arc::new(group)) - .await; - } - } - None => { - for (token_id, mut posting_list) in posting_lists { - self.cache_positions(&mut posting_list, token_id, with_position) - .await; - self.index_cache - .insert_with_key(&PostingListKey { token_id }, Arc::new(posting_list)) + .insert_with_key( + &posting_list_group_cache_key(start, end, self.has_impacts), + Arc::new(group), + ) .await; } } @@ -3177,6 +3588,9 @@ impl PostingListReader { } } } + if self.has_impacts { + base_columns.push(IMPACT_COL); + } base_columns } } @@ -3188,6 +3602,7 @@ struct ChunkBuildState { max_scores: Option>>, lengths: Option>>, posting_tail_codec: PostingTailCodec, + block_size: usize, positions_layout: PositionsLayout, } @@ -3198,6 +3613,7 @@ struct PrewarmBuildCtx<'a> { max_scores: Option<&'a [f32]>, lengths: Option<&'a [u32]>, posting_tail_codec: PostingTailCodec, + block_size: usize, positions_layout: PositionsLayout, } @@ -3291,8 +3707,8 @@ impl CacheKey for PostingListKey { /// Cache key for a group of consecutive posting lists stored as a single /// entry, covering rows `[start, end)` (issue #7040). The range, not a token -/// id, is the key so that a write-time config change that reshapes groups -/// simply misses old entries instead of serving a differently-shaped group. +/// id, is the key so a runtime group-size change simply misses old entries +/// instead of serving a differently-shaped group. #[derive(Debug, Clone)] pub struct PostingListGroupKey { pub start: u32, @@ -3315,6 +3731,52 @@ impl CacheKey for PostingListGroupKey { } } +/// Internal cache-key decorator that isolates impact-bearing posting values +/// without changing the source-compatible public posting key structs. +#[derive(Debug, Clone)] +struct ImpactAwareCacheKey { + inner: K, + has_impacts: bool, +} + +impl CacheKey for ImpactAwareCacheKey { + type ValueType = K::ValueType; + + fn key(&self) -> std::borrow::Cow<'_, str> { + if self.has_impacts { + format!("{}-impacts", self.inner.key()).into() + } else { + self.inner.key() + } + } + + fn type_name() -> &'static str { + K::type_name() + } + + fn codec() -> Option { + K::codec() + } +} + +fn posting_list_cache_key(token_id: u32, has_impacts: bool) -> ImpactAwareCacheKey { + ImpactAwareCacheKey { + inner: PostingListKey { token_id }, + has_impacts, + } +} + +fn posting_list_group_cache_key( + start: u32, + end: u32, + has_impacts: bool, +) -> ImpactAwareCacheKey { + ImpactAwareCacheKey { + inner: PostingListGroupKey { start, end }, + has_impacts, + } +} + #[derive(Debug, Clone, DeepSizeOf)] struct PostingMetadataValue { max_score: f32, @@ -3429,26 +3891,351 @@ impl SharedPositionStream { } /// A group of consecutive posting lists held in a single cache entry, in row -/// order (issue #7040). `posting_lists[i]` corresponds to row `start + i`, -/// where `start` is the group's first row from [`PostingListGroupKey`]. -#[derive(Debug, Clone, DeepSizeOf)] +/// order (issue #7040). Prewarmed modern groups without positions retain only +/// the compact Arrow posting rows read from `invert.lance`; max-score/length +/// metadata stays in the reader and is injected when a query creates a +/// posting-list view. Cold-loaded groups may keep inline metadata to preserve +/// one-read query loading. Legacy and position-bearing prewarm paths use the +/// materialized fallback. +#[derive(Debug, Clone)] pub struct PostingListGroup { - pub(super) posting_lists: Vec, + pub(super) storage: PostingListGroupStorage, } -impl PostingListGroup { - pub(super) fn new(posting_lists: Vec) -> Self { - Self { posting_lists } - } +#[derive(Debug, Clone)] +pub(super) enum PostingListGroupStorage { + Packed(PackedPostingListGroup), + Materialized(Vec), +} + +#[derive(Debug, Clone)] +pub(super) struct PackedPostingListGroup { + pub(super) batch: RecordBatch, + pub(super) posting_tail_codec: PostingTailCodec, + pub(super) block_size: usize, + first_docs_states: Arc<[OnceLock>]>, + first_docs_state_capacity_bytes: usize, + impact_states: Option>]>>, + impact_state_capacity_bytes: usize, +} + +impl DeepSizeOf for PostingListGroup { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + match &self.storage { + PostingListGroupStorage::Packed(group) => group + .batch + .columns() + .iter() + .map(|column| sliced_cache_bytes(column.as_ref())) + .sum::() + .saturating_add(group.first_docs_state_capacity_bytes) + .saturating_add(group.impact_state_capacity_bytes), + PostingListGroupStorage::Materialized(posting_lists) => { + posting_lists.deep_size_of_children(context) + } + } + } +} + +impl PostingListGroup { + pub(super) fn new(posting_lists: Vec) -> Self { + Self { + storage: PostingListGroupStorage::Materialized(posting_lists), + } + } + + pub(super) fn new_packed( + batch: RecordBatch, + posting_tail_codec: PostingTailCodec, + ) -> Result { + let block_size = parse_posting_block_size(batch.schema_ref().metadata())?; + Self::new_packed_with_block_size(batch, posting_tail_codec, block_size) + } + + fn new_packed_with_block_size( + batch: RecordBatch, + posting_tail_codec: PostingTailCodec, + block_size: usize, + ) -> Result { + validate_block_size(block_size)?; + if let Some(encoded_block_size) = batch.schema_ref().metadata().get(POSTING_BLOCK_SIZE_KEY) + { + let encoded_block_size = encoded_block_size.parse::().map_err(|err| { + Error::index(format!( + "invalid {POSTING_BLOCK_SIZE_KEY} metadata value {encoded_block_size:?}: {err}" + )) + })?; + if encoded_block_size != block_size { + return Err(Error::index(format!( + "packed posting group {POSTING_BLOCK_SIZE_KEY}={encoded_block_size} does not match block_size={block_size}" + ))); + } + } + + // Projected reads may drop schema metadata. Restore the reader's + // validated block size before the batch enters the packed cache so IPC + // roundtrips remain self-describing. Older packed cache entries omit + // the key and enter through new_packed with the legacy 128-doc default. + let mut schema = batch.schema().as_ref().clone(); + schema + .metadata + .insert(POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string()); + let batch = batch.with_schema(Arc::new(schema))?; + let postings = batch + .column_by_name(POSTING_COL) + .and_then(|column| column.as_list_opt::()) + .ok_or_else(|| { + Error::index(format!( + "packed posting group column {POSTING_COL} must be List" + )) + })?; + if postings.values().data_type() != &DataType::LargeBinary { + return Err(Error::index(format!( + "packed posting group column {POSTING_COL} must contain LargeBinary values, got {}", + postings.values().data_type() + ))); + } + if postings.null_count() != 0 { + return Err(Error::index( + "packed posting group column must not contain nulls".to_string(), + )); + } + let total_posting_blocks = (0..batch.num_rows()) + .map(|slot| postings.value_length(slot) as usize) + .sum::(); + let first_docs_states: Arc<[OnceLock>]> = (0..batch.num_rows()) + .map(|_| OnceLock::new()) + .collect::>() + .into(); + // Reserve the compact per-slot state slab and the block-head arrays it + // can lazily retain, so warming these derived values cannot grow the + // cache beyond its admission charge. + let first_docs_state_capacity_bytes = first_docs_states + .len() + .saturating_mul(std::mem::size_of::>>()) + .saturating_add(total_posting_blocks.saturating_mul(std::mem::size_of::())); + let (impact_states, impact_state_capacity_bytes) = if let Some(impacts) = + batch.column_by_name(IMPACT_COL) + { + let impacts = impacts.as_list_opt::().ok_or_else(|| { + Error::index(format!( + "packed posting group column {IMPACT_COL} must be List" + )) + })?; + if impacts.values().data_type() != &DataType::LargeBinary { + return Err(Error::index(format!( + "packed posting group column {IMPACT_COL} must contain LargeBinary values, got {}", + impacts.values().data_type() + ))); + } + if impacts.null_count() != 0 { + return Err(Error::index(format!( + "packed posting group column {IMPACT_COL} must not contain nulls" + ))); + } + let mut derived_cache_bytes = 0usize; + for slot in 0..batch.num_rows() { + let posting_blocks = postings.value_length(slot) as usize; + let impact_entries = impacts.value_length(slot) as usize; + let expected_impact_entries = + posting_blocks.saturating_add(posting_blocks.div_ceil(IMPACT_LEVEL1_BLOCKS)); + if impact_entries != expected_impact_entries { + return Err(Error::index(format!( + "packed posting group impact slot {slot} has {impact_entries} entries, expected {expected_impact_entries} for {posting_blocks} posting blocks" + ))); + } + derived_cache_bytes = derived_cache_bytes.saturating_add( + ImpactSkipData::derived_cache_bytes_for_entries(impact_entries), + ); + } + + let states: Arc<[OnceLock>]> = (0..batch.num_rows()) + .map(|_| OnceLock::new()) + .collect::>() + .into(); + // Account up front for every allocation that the lazy states can + // eventually retain. The impact entry bytes themselves remain in + // `batch` and are already charged exactly once above. + let per_slot_bytes = std::mem::size_of::>>() + .saturating_add(std::mem::size_of::()); + let capacity_bytes = states + .len() + .saturating_mul(per_slot_bytes) + .saturating_add(derived_cache_bytes); + (Some(states), capacity_bytes) + } else { + (None, 0) + }; + match ( + batch.column_by_name(MAX_SCORE_COL), + batch.column_by_name(LENGTH_COL), + ) { + (None, None) => {} + (Some(max_scores), Some(lengths)) => { + let max_scores = max_scores + .as_primitive_opt::() + .ok_or_else(|| { + Error::index(format!( + "packed posting group column {MAX_SCORE_COL} must be Float32" + )) + })?; + let lengths = lengths.as_primitive_opt::().ok_or_else(|| { + Error::index(format!( + "packed posting group column {LENGTH_COL} must be UInt32" + )) + })?; + if max_scores.null_count() != 0 || lengths.null_count() != 0 { + return Err(Error::index( + "packed posting group metadata columns must not contain nulls".to_string(), + )); + } + } + _ => { + return Err(Error::index(format!( + "packed posting group must contain both {MAX_SCORE_COL} and {LENGTH_COL}, or neither" + ))); + } + } + + Ok(Self { + storage: PostingListGroupStorage::Packed(PackedPostingListGroup { + batch, + posting_tail_codec, + block_size, + first_docs_states, + first_docs_state_capacity_bytes, + impact_states, + impact_state_capacity_bytes, + }), + }) + } + + pub(super) fn len(&self) -> usize { + match &self.storage { + PostingListGroupStorage::Packed(group) => group.batch.num_rows(), + PostingListGroupStorage::Materialized(posting_lists) => posting_lists.len(), + } + } + + #[cfg(test)] + pub(super) fn is_packed(&self) -> bool { + matches!(&self.storage, PostingListGroupStorage::Packed(_)) + } + + fn needs_external_metadata(&self) -> bool { + match &self.storage { + PostingListGroupStorage::Packed(group) => { + group.batch.column_by_name(MAX_SCORE_COL).is_none() + } + PostingListGroupStorage::Materialized(_) => false, + } + } - /// Borrow the posting list at offset `slot` within the group (i.e. - /// `token_id - start`). - pub(super) fn get(&self, slot: usize) -> Option<&PostingList> { - self.posting_lists.get(slot) + /// Build an owned posting-list view for `slot`. Packed groups clone only + /// Arrow array metadata; the compressed posting bytes remain shared with + /// the group's `List` child buffers. + pub(super) fn posting_list( + &self, + slot: usize, + max_score: Option, + length: Option, + ) -> Result> { + match &self.storage { + PostingListGroupStorage::Materialized(posting_lists) => { + Ok(posting_lists.get(slot).cloned()) + } + PostingListGroupStorage::Packed(group) => { + if slot >= group.batch.num_rows() { + return Ok(None); + } + let postings = group + .batch + .column_by_name(POSTING_COL) + .and_then(|column| column.as_list_opt::()) + .ok_or_else(|| { + Error::index(format!( + "packed posting group column {POSTING_COL} must be List" + )) + })?; + let blocks = postings.value(slot); + let blocks = blocks.as_binary_opt::().ok_or_else(|| { + Error::index(format!( + "packed posting group slot {slot} is not LargeBinary" + )) + })?; + let max_score = match group.batch.column_by_name(MAX_SCORE_COL) { + Some(column) => column + .as_primitive_opt::() + .expect("packed group metadata was validated at construction") + .value(slot), + None => max_score.ok_or_else(|| { + Error::index("packed posting group requires max-score metadata".to_string()) + })?, + }; + let length = match group.batch.column_by_name(LENGTH_COL) { + Some(column) => column + .as_primitive_opt::() + .expect("packed group metadata was validated at construction") + .value(slot), + None => length.ok_or_else(|| { + Error::index("packed posting group requires length metadata".to_string()) + })?, + }; + let impacts = match ( + group.impact_states.as_ref(), + group.batch.column_by_name(IMPACT_COL), + ) { + (Some(states), Some(column)) => { + let state = states.get(slot).ok_or_else(|| { + Error::index(format!( + "packed posting group impact state missing slot {slot}" + )) + })?; + let impact_lists = column.as_list_opt::().ok_or_else(|| { + Error::index(format!( + "packed posting group column {IMPACT_COL} must be List" + )) + })?; + let entries = impact_lists.value(slot); + let entries = entries.as_binary_opt::().ok_or_else(|| { + Error::index(format!( + "packed posting group impact slot {slot} is not LargeBinary" + )) + })?; + let impacts = + state.get_or_init(|| { + Box::new(ImpactSkipData::new(entries.clone(), blocks.len()).expect( + "packed impact entry count was validated at construction", + )) + }); + Some(impacts.as_ref().clone()) + } + (None, None) => None, + _ => { + return Err(Error::internal( + "packed posting group impact column/state mismatch".to_string(), + )); + } + }; + Ok(Some(PostingList::Compressed( + CompressedPostingList::new( + blocks.clone(), + max_score, + length, + group.posting_tail_codec, + group.block_size, + None, + impacts, + ) + .with_packed_first_docs(group.first_docs_states.clone(), slot), + ))) + } + } } } #[derive(Debug, Clone, DeepSizeOf)] +#[allow(clippy::large_enum_variant)] pub enum PostingList { Plain(PlainPostingList), Compressed(CompressedPostingList), @@ -3461,7 +4248,8 @@ impl PostingList { length: Option, ) -> Result { let posting_tail_codec = parse_posting_tail_codec(batch.schema_ref().metadata())?; - Self::from_batch_with_tail_codec(batch, max_score, length, posting_tail_codec) + let block_size = parse_posting_block_size(batch.schema_ref().metadata())?; + Self::from_batch_with_tail_codec(batch, max_score, length, posting_tail_codec, block_size) } pub fn from_batch_with_tail_codec( @@ -3469,6 +4257,7 @@ impl PostingList { max_score: Option, length: Option, posting_tail_codec: PostingTailCodec, + block_size: usize, ) -> Result { let positions_layout = if batch.column_by_name(COMPRESSED_POSITION_COL).is_some() { PositionsLayout::SharedStream(parse_shared_position_codec( @@ -3484,6 +4273,7 @@ impl PostingList { max_score, length, posting_tail_codec, + block_size, positions_layout, ) } @@ -3493,6 +4283,7 @@ impl PostingList { max_score: Option, length: Option, posting_tail_codec: PostingTailCodec, + block_size: usize, positions_layout: PositionsLayout, ) -> Result { match batch.column_by_name(POSTING_COL) { @@ -3507,8 +4298,9 @@ impl PostingList { max_score.unwrap(), length.unwrap(), posting_tail_codec, + block_size, shared_position_codec, - ); + )?; Ok(Self::Compressed(posting)) } None => { @@ -3529,6 +4321,13 @@ impl PostingList { } } + pub fn has_impacts(&self) -> bool { + match self { + Self::Plain(_) => false, + Self::Compressed(posting) => posting.impacts.is_some(), + } + } + pub fn set_positions(&mut self, positions: CompressedPositionStorage) { match self { Self::Plain(posting) => match positions { @@ -3578,9 +4377,14 @@ impl PostingList { Self::Plain(_) => PostingTailCodec::Fixed32, Self::Compressed(posting) => posting.posting_tail_codec, }; - let mut builder = PostingListBuilder::new_with_posting_tail_codec( + let block_size = match &self { + Self::Plain(_) => LEGACY_BLOCK_SIZE, + Self::Compressed(posting) => posting.block_size, + }; + let mut builder = PostingListBuilder::new_with_posting_tail_codec_and_block_size( self.has_position(), posting_tail_codec, + block_size, ); match self { // legacy format @@ -3733,15 +4537,77 @@ impl PlainPostingList { } } -#[derive(Debug, PartialEq, Clone)] +#[derive(Debug, Clone)] +enum FirstDocsState { + Standalone(Arc>>), + Packed { + states: Arc<[OnceLock>]>, + slot: usize, + }, +} + +impl FirstDocsState { + fn standalone() -> Self { + Self::Standalone(Arc::new(OnceLock::new())) + } + + fn state(&self) -> &OnceLock> { + match self { + Self::Standalone(state) => state, + Self::Packed { states, slot } => &states[*slot], + } + } + + fn get_or_init(&self, initialize: impl FnOnce() -> Box<[u32]>) -> &[u32] { + self.state().get_or_init(initialize) + } + + fn capacity_bytes( + &self, + block_count: usize, + context: &mut lance_core::deepsize::Context, + ) -> usize { + if context.mark_seen(self.state() as *const _ as usize) { + std::mem::size_of::>>() + .saturating_add(block_count.saturating_mul(std::mem::size_of::())) + } else { + 0 + } + } + + #[cfg(test)] + fn shares_state_with(&self, other: &Self) -> bool { + std::ptr::eq(self.state(), other.state()) + } +} + +#[derive(Debug, Clone)] pub struct CompressedPostingList { pub max_score: f32, pub length: u32, // each binary is a block of compressed data - // that contains `BLOCK_SIZE` doc ids and then `BLOCK_SIZE` frequencies + // that contains `block_size` doc ids and then `block_size` frequencies, + // packed by the physical bitpacker matching that block size. pub blocks: LargeBinaryArray, pub posting_tail_codec: PostingTailCodec, + pub block_size: usize, pub positions: Option, + pub(crate) impacts: Option, + // First doc id per block, baked lazily and shared across per-query clones + // of the cached list. See `block_first_docs`. + first_docs: FirstDocsState, +} + +impl PartialEq for CompressedPostingList { + fn eq(&self, other: &Self) -> bool { + self.max_score == other.max_score + && self.length == other.length + && self.blocks == other.blocks + && self.posting_tail_codec == other.posting_tail_codec + && self.block_size == other.block_size + && self.positions == other.positions + && self.impacts == other.impacts + } } impl DeepSizeOf for CompressedPostingList { @@ -3752,33 +4618,68 @@ impl DeepSizeOf for CompressedPostingList { .as_ref() .map(|positions| positions.deep_size_of_children(context)) .unwrap_or(0) + + self + .impacts + .as_ref() + .map(|impacts| { + sliced_cache_bytes(impacts.entries()) + .saturating_add(impacts.derived_cache_bytes()) + }) + .unwrap_or(0) + + self.first_docs.capacity_bytes(self.blocks.len(), context) } } impl CompressedPostingList { - pub fn new( + pub(crate) fn new( blocks: LargeBinaryArray, max_score: f32, length: u32, posting_tail_codec: PostingTailCodec, + block_size: usize, positions: Option, + impacts: Option, ) -> Self { + debug_assert!(block_size.is_power_of_two()); Self { max_score, length, blocks, posting_tail_codec, + block_size, positions, + impacts, + first_docs: FirstDocsState::standalone(), } } + fn with_packed_first_docs(mut self, states: Arc<[OnceLock>]>, slot: usize) -> Self { + debug_assert!(slot < states.len()); + self.first_docs = FirstDocsState::Packed { states, slot }; + self + } + + /// Block sizes are validated powers of two, so per-doc hot loops derive + /// block indices with shift/mask instead of runtime division, which is + /// measurably slower in the iterator advance path. + #[inline] + pub(crate) fn block_shift(&self) -> u32 { + self.block_size.trailing_zeros() + } + + #[inline] + pub(crate) fn block_mask(&self) -> usize { + self.block_size - 1 + } + pub fn from_batch( batch: &RecordBatch, max_score: f32, length: u32, posting_tail_codec: PostingTailCodec, + block_size: usize, shared_position_codec: Option, - ) -> Self { + ) -> Result { debug_assert_eq!(batch.num_rows(), 1); let blocks = batch[POSTING_COL] .as_list::() @@ -3807,14 +4708,24 @@ impl CompressedPostingList { ) }) }; + let impacts = batch + .column_by_name(IMPACT_COL) + .map(|col| { + let entries = col.as_list::().value(0).as_binary::().clone(); + ImpactSkipData::new(entries, blocks.len()) + }) + .transpose()?; - Self { + Ok(Self { max_score, length, blocks, posting_tail_codec, + block_size, positions, - } + impacts, + first_docs: FirstDocsState::standalone(), + }) } pub fn iter(&self) -> CompressedPostingListIterator { @@ -3823,23 +4734,58 @@ impl CompressedPostingList { self.blocks.clone(), self.posting_tail_codec, self.positions.clone(), + self.block_size, ) } pub fn block_max_score(&self, block_idx: usize) -> f32 { + // 256-doc (V3) blocks store no per-block max score: their impact + // skip data supplies the tight per-block bound, so callers on that + // path never reach here. Fall back to the list-level max, which is + // still a valid (looser) bound for any block. + if super::encoding::posting_block_score_prefix_len(self.block_size) == 0 { + return self.max_score; + } let block = self.blocks.value(block_idx); block[0..4].try_into().map(f32::from_le_bytes).unwrap() } + #[inline] pub fn block_least_doc_id(&self, block_idx: usize) -> u32 { - let block = self.blocks.value(block_idx); - let remainder = self.length as usize % BLOCK_SIZE; - let is_remainder_block = remainder > 0 && block_idx + 1 == self.blocks.len(); - if is_remainder_block { - super::encoding::read_posting_tail_first_doc(block, self.posting_tail_codec) - } else { - block[4..8].try_into().map(u32::from_le_bytes).unwrap() - } + self.block_first_docs()[block_idx] + } + + /// First doc id of every block, decoded once per cached list and shared by + /// the per-query clones. Block boundary lookups (window bounds, block + /// binary searches) are hot enough that re-reading the block headers — + /// and re-decoding the tail block — shows up in profiles. + pub(crate) fn block_first_docs(&self) -> &[u32] { + self.first_docs.get_or_init(|| { + (0..self.blocks.len()) + .map(|block_idx| { + let block = self.blocks.value(block_idx); + let remainder = self.length as usize % self.block_size; + if block_idx + 1 == self.blocks.len() && remainder > 0 { + return super::encoding::read_posting_tail_first_doc( + block, + self.posting_tail_codec, + self.block_size, + ); + } + let prefix = super::encoding::posting_block_score_prefix_len(self.block_size); + block[prefix..prefix + 4] + .try_into() + .map(u32::from_le_bytes) + .unwrap() + }) + .collect::>() + .into_boxed_slice() + }) + } + + #[cfg(test)] + fn shares_first_docs_with(&self, other: &Self) -> bool { + self.first_docs.shares_state_with(&other.first_docs) } } @@ -3890,12 +4836,14 @@ impl EncodedBlocks { doc_ids: &[u32], frequencies: &[u32], codec: PostingTailCodec, + block_size: usize, ) -> Result<()> { self.offsets.push(self.bytes.len() as u32); super::encoding::encode_remainder_posting_block_into( doc_ids, frequencies, codec, + block_size, &mut self.bytes, ) } @@ -3964,6 +4912,7 @@ pub struct PostingListBuilder { open_doc_id: Option, open_doc_frequency: u32, open_doc_last_position: Option, + block_size: usize, memory_size_bytes: u32, len: u32, } @@ -3971,14 +4920,11 @@ pub struct PostingListBuilder { pub(super) struct PostingListBatchBuilder { schema: SchemaRef, postings: ListBuilder, + impacts: Option>, max_scores: Float32Builder, lengths: UInt32Builder, positions: BatchPositionsBuilder, len: usize, - /// Tracks posting-list cache-group boundaries in row order across all - /// batches this builder produces (issue #7040). Outlives `finish`, which - /// only resets the per-batch column builders. - group_accumulator: PostingGroupAccumulator, } enum BatchPositionsBuilder { @@ -3993,6 +4939,7 @@ enum BatchPositionsBuilder { struct PostingListParts<'a> { with_positions: bool, posting_tail_codec: PostingTailCodec, + block_size: usize, length: usize, encoded_blocks: EncodedBlocks, encoded_position_blocks: EncodedPositionBlocks, @@ -4006,7 +4953,6 @@ impl PostingListBatchBuilder { with_positions: bool, format_version: InvertedListFormatVersion, capacity: usize, - group_config: PostingGroupConfig, ) -> Self { let positions = if !with_positions { BatchPositionsBuilder::None @@ -4021,14 +4967,18 @@ impl PostingListBatchBuilder { capacity, )) }; + let impacts = schema + .field_with_name(IMPACT_COL) + .ok() + .map(|_| ListBuilder::with_capacity(LargeBinaryBuilder::new(), capacity)); Self { schema, postings: ListBuilder::with_capacity(LargeBinaryBuilder::new(), capacity), + impacts, max_scores: Float32Builder::with_capacity(capacity), lengths: UInt32Builder::with_capacity(capacity), positions, len: 0, - group_accumulator: PostingGroupAccumulator::new(group_config), } } @@ -4043,11 +4993,11 @@ impl PostingListBatchBuilder { fn append( &mut self, compressed: LargeBinaryArray, + impacts: Option<&ImpactSkipData>, max_score: f32, length: u32, positions: Option<&CompressedPositionStorage>, ) -> Result<()> { - let posting_bytes = compressed.value_data().len(); { let values = self.postings.values(); for index in 0..compressed.len() { @@ -4055,7 +5005,19 @@ impl PostingListBatchBuilder { } } self.postings.append(true); - self.group_accumulator.push(posting_bytes); + if let Some(impacts_builder) = &mut self.impacts { + let impacts = impacts.ok_or_else(|| { + Error::index(format!( + "impacts builder missing impact data for posting length {}", + length + )) + })?; + let values = impacts_builder.values(); + for index in 0..impacts.entries().len() { + values.append_value(impacts.entries().value(index)); + } + impacts_builder.append(true); + } self.max_scores.append_value(max_score); self.lengths.append_value(length); @@ -4120,6 +5082,9 @@ impl PostingListBatchBuilder { Arc::new(self.max_scores.finish()) as ArrayRef, Arc::new(self.lengths.finish()) as ArrayRef, ]; + if let Some(impacts) = &mut self.impacts { + columns.push(Arc::new(impacts.finish()) as ArrayRef); + } match &mut self.positions { BatchPositionsBuilder::None => {} BatchPositionsBuilder::Legacy(position_lists) => { @@ -4136,13 +5101,6 @@ impl PostingListBatchBuilder { self.len = 0; RecordBatch::try_new(self.schema.clone(), columns).map_err(Error::from) } - - /// Consume the builder and return the posting-list cache-group boundaries - /// accumulated across all batches (issue #7040). Each entry is the first - /// row of a group; the sequence is monotonically increasing. - pub fn into_group_starts(self) -> Vec { - self.group_accumulator.into_starts() - } } impl PostingListBuilder { @@ -4155,9 +5113,10 @@ impl PostingListBuilder { } pub fn new(with_position: bool) -> Self { - Self::new_with_posting_tail_codec( + Self::new_with_posting_tail_codec_and_block_size( with_position, current_fts_format_version().posting_tail_codec(), + LEGACY_BLOCK_SIZE, ) } @@ -4165,6 +5124,27 @@ impl PostingListBuilder { with_position: bool, posting_tail_codec: PostingTailCodec, ) -> Self { + Self::new_with_posting_tail_codec_and_block_size( + with_position, + posting_tail_codec, + LEGACY_BLOCK_SIZE, + ) + } + + pub fn new_with_block_size(with_position: bool, block_size: usize) -> Self { + Self::new_with_posting_tail_codec_and_block_size( + with_position, + current_fts_format_version().posting_tail_codec(), + block_size, + ) + } + + pub fn new_with_posting_tail_codec_and_block_size( + with_position: bool, + posting_tail_codec: PostingTailCodec, + block_size: usize, + ) -> Self { + validate_block_size(block_size).expect("invalid posting list block size"); Self { with_positions: with_position, posting_tail_codec, @@ -4175,6 +5155,7 @@ impl PostingListBuilder { open_doc_id: None, open_doc_frequency: 0, open_doc_last_position: None, + block_size, len: 0, memory_size_bytes: 0, } @@ -4196,8 +5177,8 @@ impl PostingListBuilder { &self, mut visit: impl FnMut(u32, u32, Option>) -> std::result::Result<(), E>, ) -> std::result::Result<(), E> { - let mut doc_ids = Vec::with_capacity(BLOCK_SIZE); - let mut frequencies = Vec::with_capacity(BLOCK_SIZE); + let mut doc_ids = Vec::with_capacity(self.block_size); + let mut frequencies = Vec::with_capacity(self.block_size); let mut decoded_positions = Vec::new(); let mut position_block_index = 0usize; @@ -4205,7 +5186,12 @@ impl PostingListBuilder { for block in encoded_blocks.iter() { doc_ids.clear(); frequencies.clear(); - super::encoding::decode_full_posting_block(block, &mut doc_ids, &mut frequencies); + super::encoding::decode_full_posting_block( + block, + &mut doc_ids, + &mut frequencies, + self.block_size, + ); decoded_positions.clear(); if self.with_positions { let position_blocks = self @@ -4285,7 +5271,7 @@ impl PostingListBuilder { } self.len += 1; - if self.tail_entries.len() == BLOCK_SIZE { + if self.tail_entries.len() == self.block_size { self.flush_tail_block() .expect("posting list block compression should succeed"); } @@ -4344,7 +5330,7 @@ impl PostingListBuilder { self.open_doc_id = None; self.open_doc_frequency = 0; self.open_doc_last_position = None; - if self.tail_entries.len() == BLOCK_SIZE { + if self.tail_entries.len() == self.block_size { self.flush_tail_block()?; } Ok(()) @@ -4395,13 +5381,17 @@ impl PostingListBuilder { self.open_doc_id.is_none(), "cannot flush a posting block while a document is still open" ); - debug_assert_eq!(self.tail_entries.len(), BLOCK_SIZE); - let mut doc_ids = [0u32; BLOCK_SIZE]; - let mut frequencies = [0u32; BLOCK_SIZE]; - for (index, entry) in self.tail_entries.iter().enumerate() { - doc_ids[index] = entry.doc_id; - frequencies[index] = entry.frequency; - } + debug_assert_eq!(self.tail_entries.len(), self.block_size); + let doc_ids = self + .tail_entries + .iter() + .map(|entry| entry.doc_id) + .collect::>(); + let frequencies = self + .tail_entries + .iter() + .map(|entry| entry.frequency) + .collect::>(); let encoded_blocks_size_before = self .encoded_blocks .as_ref() @@ -4505,6 +5495,7 @@ impl PostingListBuilder { fn build_batch( self, compressed: LargeBinaryArray, + impacts: Option, max_score: f32, schema: SchemaRef, positions: Option, @@ -4523,6 +5514,22 @@ impl PostingListBuilder { length as u32, ))) as ArrayRef, ]; + if schema.field_with_name(IMPACT_COL).is_ok() { + let impacts = impacts.ok_or_else(|| { + Error::index(format!( + "impact column requested without impact data for posting length {}", + length + )) + })?; + let impact_offsets = + OffsetBuffer::new(ScalarBuffer::from(vec![0, impacts.entries().len() as i32])); + columns.push(Arc::new(ListArray::try_new( + Arc::new(Field::new("item", datatypes::DataType::LargeBinary, true)), + impact_offsets, + Arc::new(impacts.entries().clone()), + None, + )?) as ArrayRef); + } columns.extend(Self::build_position_columns(positions)?); let batch = RecordBatch::try_new(schema, columns)?; @@ -4572,6 +5579,7 @@ impl PostingListBuilder { open_doc_id, open_doc_frequency, open_doc_last_position, + block_size, len, .. } = self; @@ -4581,6 +5589,7 @@ impl PostingListBuilder { let parts = PostingListParts { with_positions, posting_tail_codec, + block_size, length: len as usize, encoded_blocks: encoded_blocks .map(|encoded_blocks| *encoded_blocks) @@ -4591,13 +5600,19 @@ impl PostingListBuilder { tail_entries: tail_entries.as_slice(), tail_position_block: with_positions.then(|| tail_positions.finish()), }; - let (compressed, shared_positions, max_score) = + let (compressed, shared_positions, max_score, impacts) = Self::build_compressed_with_scores_from_parts(parts, docs)?; let positions = match legacy_positions { Some(positions) => Some(CompressedPositionStorage::LegacyPerDoc(positions)), None => shared_positions.map(CompressedPositionStorage::SharedStream), }; - batch_builder.append(compressed, max_score, len, positions.as_ref()) + batch_builder.append( + compressed, + Some(&impacts), + max_score, + len, + positions.as_ref(), + ) } fn extend_tail_components( @@ -4614,11 +5629,17 @@ impl PostingListBuilder { fn build_compressed_with_scores_from_parts( parts: PostingListParts<'_>, docs: &DocSet, - ) -> Result<(LargeBinaryArray, Option, f32)> { + ) -> Result<( + LargeBinaryArray, + Option, + f32, + ImpactSkipData, + )> { let PostingListParts { with_positions, posting_tail_codec, length, + block_size, mut encoded_blocks, mut encoded_position_blocks, tail_entries, @@ -4627,41 +5648,58 @@ impl PostingListBuilder { let avgdl = docs.average_length(); let idf_scale = idf(length, docs.len()) * (K1 + 1.0); let mut max_score = f32::MIN; - let mut doc_ids = Vec::with_capacity(BLOCK_SIZE); - let mut frequencies = Vec::with_capacity(BLOCK_SIZE); + let mut doc_ids = Vec::with_capacity(block_size); + let mut frequencies = Vec::with_capacity(block_size); + let mut impact_block = Vec::with_capacity(block_size); + let mut impact_builder = + ImpactSkipDataBuilder::with_capacity(length.div_ceil(block_size), block_size); for index in 0..encoded_blocks.len() { let block = encoded_blocks.block(index); doc_ids.clear(); frequencies.clear(); - super::encoding::decode_full_posting_block(block, &mut doc_ids, &mut frequencies); - let block_score = compute_block_score( + super::encoding::decode_full_posting_block( + block, + &mut doc_ids, + &mut frequencies, + block_size, + ); + let block_score = compute_block_score_and_impact_block( docs, avgdl, idf_scale, doc_ids.iter().copied(), frequencies.iter().copied(), + &mut impact_block, ); + impact_builder.append_block(impact_block.as_slice())?; max_score = max_score.max(block_score); - encoded_blocks.set_block_score(index, block_score); + if super::encoding::posting_block_score_prefix_len(block_size) > 0 { + encoded_blocks.set_block_score(index, block_score); + } } if !tail_entries.is_empty() { Self::extend_tail_components(tail_entries, &mut doc_ids, &mut frequencies); - let block_score = compute_block_score( + let block_score = compute_block_score_and_impact_block( docs, avgdl, idf_scale, doc_ids.iter().copied(), frequencies.iter().copied(), + &mut impact_block, ); + impact_builder.append_block(impact_block.as_slice())?; max_score = max_score.max(block_score); encoded_blocks.append_remainder_block_with_codec( doc_ids.as_slice(), frequencies.as_slice(), posting_tail_codec, + block_size, )?; - encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); + if super::encoding::posting_block_score_prefix_len(block_size) > 0 { + encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); + } if with_positions { encoded_position_blocks.push_encoded_block( tail_position_block @@ -4671,22 +5709,27 @@ impl PostingListBuilder { } } + let impacts = impact_builder.finish()?; Ok(( encoded_blocks.into_array(), with_positions.then(|| encoded_position_blocks.into_stream()), max_score, + impacts, )) } + #[allow(clippy::too_many_arguments)] fn build_compressed_with_block_scores_from_parts( with_positions: bool, posting_tail_codec: PostingTailCodec, + block_size: usize, mut encoded_blocks: EncodedBlocks, mut encoded_position_blocks: EncodedPositionBlocks, tail_entries: &[RawDocInfo], tail_position_block: Option>, mut block_max_scores: impl Iterator, ) -> Result<(LargeBinaryArray, Option, f32)> { + let has_score_prefix = super::encoding::posting_block_score_prefix_len(block_size) > 0; let mut max_score = f32::MIN; let mut doc_ids = Vec::with_capacity(BLOCK_SIZE); let mut frequencies = Vec::with_capacity(BLOCK_SIZE); @@ -4696,7 +5739,9 @@ impl PostingListBuilder { .next() .ok_or_else(|| Error::index("missing block max score".to_owned()))?; max_score = max_score.max(block_score); - encoded_blocks.set_block_score(index, block_score); + if has_score_prefix { + encoded_blocks.set_block_score(index, block_score); + } } if !tail_entries.is_empty() { @@ -4709,8 +5754,11 @@ impl PostingListBuilder { doc_ids.as_slice(), frequencies.as_slice(), posting_tail_codec, + block_size, )?; - encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); + if has_score_prefix { + encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); + } if with_positions { encoded_position_blocks.push_encoded_block( tail_position_block @@ -4728,12 +5776,16 @@ impl PostingListBuilder { } pub fn to_batch(self, block_max_scores: Vec) -> Result { - let format_version = if self.posting_tail_codec == PostingTailCodec::Fixed32 { - InvertedListFormatVersion::V1 - } else { - InvertedListFormatVersion::V2 - }; - let schema = inverted_list_schema_for_version(self.has_positions(), format_version); + let format_version = InvertedListFormatVersion::from_posting_tail_codec_and_block_size( + self.posting_tail_codec, + self.block_size, + )?; + let schema = inverted_list_schema_for_version_with_block_size_and_impacts( + self.has_positions(), + format_version, + self.block_size, + false, + ); let legacy_positions = if self.with_positions && !format_version.uses_shared_position_stream() { Some(self.build_legacy_positions()?) @@ -4750,6 +5802,7 @@ impl PostingListBuilder { open_doc_id, open_doc_frequency, open_doc_last_position, + block_size, len, .. } = self; @@ -4760,6 +5813,7 @@ impl PostingListBuilder { Self::build_compressed_with_block_scores_from_parts( with_positions, posting_tail_codec, + block_size, encoded_blocks .map(|encoded_blocks| *encoded_blocks) .unwrap_or_default(), @@ -4780,6 +5834,7 @@ impl PostingListBuilder { open_doc_id: None, open_doc_frequency: 0, open_doc_last_position: None, + block_size, memory_size_bytes: 0, len, }; @@ -4787,17 +5842,11 @@ impl PostingListBuilder { Some(positions) => Some(CompressedPositionStorage::LegacyPerDoc(positions)), None => shared_positions.map(CompressedPositionStorage::SharedStream), }; - builder.build_batch(compressed, max_score, schema, positions) + builder.build_batch(compressed, None, max_score, schema, positions) } pub fn to_batch_with_docs(self, docs: &DocSet, schema: SchemaRef) -> Result { - let format_version = if schema.column_with_name(POSITION_COL).is_some() - && schema.column_with_name(COMPRESSED_POSITION_COL).is_none() - { - InvertedListFormatVersion::V1 - } else { - InvertedListFormatVersion::V2 - }; + let format_version = parse_format_version_from_metadata(schema.metadata())?; let legacy_positions = if self.with_positions && !format_version.uses_shared_position_stream() { Some(self.build_legacy_positions()?) @@ -4814,6 +5863,7 @@ impl PostingListBuilder { open_doc_id, open_doc_frequency, open_doc_last_position, + block_size, len, .. } = self; @@ -4823,6 +5873,7 @@ impl PostingListBuilder { let parts = PostingListParts { with_positions, posting_tail_codec, + block_size, length: len as usize, encoded_blocks: encoded_blocks .map(|encoded_blocks| *encoded_blocks) @@ -4833,7 +5884,7 @@ impl PostingListBuilder { tail_entries: tail_entries.as_slice(), tail_position_block: with_positions.then(|| tail_positions.finish()), }; - let (compressed, shared_positions, max_score) = + let (compressed, shared_positions, max_score, impacts) = Self::build_compressed_with_scores_from_parts(parts, docs)?; let builder = Self { with_positions, @@ -4845,6 +5896,7 @@ impl PostingListBuilder { open_doc_id: None, open_doc_frequency: 0, open_doc_last_position: None, + block_size, memory_size_bytes: 0, len, }; @@ -4852,13 +5904,16 @@ impl PostingListBuilder { Some(positions) => Some(CompressedPositionStorage::LegacyPerDoc(positions)), None => shared_positions.map(CompressedPositionStorage::SharedStream), }; - builder.build_batch(compressed, max_score, schema, positions) + builder.build_batch(compressed, Some(impacts), max_score, schema, positions) } pub fn remap(&mut self, removed: &[u32]) { let mut cursor = 0; - let mut new_builder = - Self::new_with_posting_tail_codec(self.has_positions(), self.posting_tail_codec); + let mut new_builder = Self::new_with_posting_tail_codec_and_block_size( + self.has_positions(), + self.posting_tail_codec, + self.block_size, + ); for (doc_id, freq, positions) in self.iter() { while cursor < removed.len() && removed[cursor] < doc_id { cursor += 1; @@ -4877,19 +5932,23 @@ impl PostingListBuilder { } } -fn compute_block_score( +fn compute_block_score_and_impact_block( docs: &DocSet, avgdl: f32, idf_scale: f32, doc_ids: impl Iterator, frequencies: impl Iterator, + impact_block: &mut Vec<(u32, u32, u32)>, ) -> f32 { + impact_block.clear(); let mut block_max_score = f32::MIN; for (doc_id, freq) in doc_ids.zip(frequencies) { - let doc_norm = K1 * (1.0 - B + B * docs.num_tokens(doc_id) as f32 / avgdl); - let freq = freq as f32; - let score = freq / (freq + doc_norm); + let doc_len = docs.num_tokens(doc_id); + let doc_norm = K1 * (1.0 - B + B * doc_len as f32 / avgdl); + let freq_f32 = freq as f32; + let score = freq_f32 / (freq_f32 + doc_norm); block_max_score = block_max_score.max(score); + impact_block.push((doc_id, freq, doc_len)); } block_max_score * idf_scale } @@ -5000,9 +6059,54 @@ impl Ord for RawDocInfo { } } +/// Lucene SmallFloat-style doc-length quantization for V3 scoring and impact +/// norms: a 4-mantissa-bit float-like byte code. Values 0-7 are exact; larger +/// values keep their top four significand bits (relative error <= 6.25%) and +/// decode to their bucket floor. The floor only ever shortens a doc, so impact +/// bounds remain conservative for exact-scoring V2 as well as quantized V3. +pub(super) fn quantize_doc_length(value: u32) -> u8 { + let num_bits = 32 - value.leading_zeros(); + if num_bits < 4 { + value as u8 + } else { + let shift = num_bits - 4; + (((value >> shift) as u8) & 0x07) | (((shift + 1) as u8) << 3) + } +} + +#[inline] +pub(super) fn dequantize_doc_length(code: u8) -> u32 { + DEQUANTIZED_DOC_LENGTHS[code as usize] +} + +pub(super) static DEQUANTIZED_DOC_LENGTHS: [u32; 256] = build_dequantized_doc_lengths(); + +const fn build_dequantized_doc_lengths() -> [u32; 256] { + let mut table = [0u32; 256]; + let mut code = 0usize; + while code < 256 { + let bits = (code & 0x07) as u64; + let shift = (code >> 3) as i64 - 1; + let decoded = if shift < 0 { + bits + } else { + (bits | 0x08) << shift + }; + // Codes past the largest u32 encoding are never produced; saturate so + // the table stays total. + table[code] = if decoded > u32::MAX as u64 { + u32::MAX + } else { + decoded as u32 + }; + code += 1; + } + table +} + // DocSet is a mapping from row ids to the number of tokens in the document // It's used to sort the documents by the bm25 score -#[derive(Debug, Clone, Default, DeepSizeOf)] +#[derive(Debug, Clone, Default)] pub struct DocSet { row_ids: Vec, num_tokens: Vec, @@ -5010,13 +6114,33 @@ pub struct DocSet { inv: Vec<(u64, u32)>, total_tokens: u64, + + // V3 (256-doc block) partitions score with quantized doc lengths: the + // flag is set at partition load and the byte-norm slab bakes lazily on + // first scoring use (shared by clones of the loaded set). 128-block + // partitions never set the flag and keep exact scoring. + scoring_quantized: bool, + norms: Arc>>, } -impl DocSet { - #[inline] - pub fn len(&self) -> usize { - // Use num_tokens instead of row_ids so the deferred-row_ids - // scoring path (which constructs a DocSet via +impl DeepSizeOf for DocSet { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.row_ids.deep_size_of_children(context) + + self.num_tokens.deep_size_of_children(context) + + self.inv.deep_size_of_children(context) + + self + .norms + .get() + .map(|slab| std::mem::size_of_val(slab.as_ref())) + .unwrap_or(0) + } +} + +impl DocSet { + #[inline] + pub fn len(&self) -> usize { + // Use num_tokens instead of row_ids so the deferred-row_ids + // scoring path (which constructs a DocSet via // [`Self::from_num_tokens_only`]) still reports the right doc // count. self.num_tokens.len() @@ -5046,12 +6170,11 @@ impl DocSet { /// Resolve a `row_id` to every `doc_id` it owns. /// - /// A scalar column maps each row to a single document, but a - /// `list` column indexes every element as its own document, so a - /// single `row_id` can own several `doc_id`s sharing that key in `inv`. + /// Modern indexes map each row to a single document. Older list indexes + /// may have indexed each list element as its own document, so a single + /// `row_id` can still own several `doc_id`s sharing that key in `inv`. /// The prefilter path (`flat_search`) walks an allow-list of row_ids and - /// must evaluate *all* of a row's documents; resolving to one `doc_id` - /// silently drops matches at non-last list positions (lancedb#3352). + /// must evaluate all legacy documents for that row. pub fn doc_ids(&self, row_id: u64) -> impl Iterator + '_ { if self.inv.is_empty() { // in legacy format, the row id is doc id (one document per row) @@ -5079,9 +6202,19 @@ impl DocSet { doc_ids: impl Iterator, freqs: impl Iterator, ) -> Vec { + self.calculate_block_max_scores_with_block_size(doc_ids, freqs, LEGACY_BLOCK_SIZE) + } + + pub fn calculate_block_max_scores_with_block_size<'a>( + &self, + doc_ids: impl Iterator, + freqs: impl Iterator, + block_size: usize, + ) -> Vec { + validate_block_size(block_size).expect("invalid posting list block size"); let avgdl = self.average_length(); let length = doc_ids.size_hint().0; - let num_blocks = length.div_ceil(BLOCK_SIZE); + let num_blocks = length.div_ceil(block_size); let mut block_max_scores = Vec::with_capacity(num_blocks); let idf_scale = idf(length, self.len()) * (K1 + 1.0); let mut max_score = f32::MIN; @@ -5092,13 +6225,13 @@ impl DocSet { if score > max_score { max_score = score; } - if (i + 1) % BLOCK_SIZE == 0 { + if (i + 1) % block_size == 0 { max_score *= idf_scale; block_max_scores.push(max_score); max_score = f32::MIN; } } - if !length.is_multiple_of(BLOCK_SIZE) { + if !length.is_multiple_of(block_size) { max_score *= idf_scale; block_max_scores.push(max_score); } @@ -5127,7 +6260,7 @@ impl DocSet { pub async fn load( reader: Arc, is_legacy: bool, - frag_reuse_index: Option>, + frag_reuse_index: Option>, ) -> Result { let batch = reader.read_range(0..reader.num_rows(), None).await?; let row_id_col = batch[ROW_ID].as_primitive::(); @@ -5148,6 +6281,8 @@ impl DocSet { num_tokens, inv: Vec::new(), total_tokens, + scoring_quantized: false, + norms: Arc::new(std::sync::OnceLock::new()), } } @@ -5159,7 +6294,7 @@ impl DocSet { row_id_col: &UInt64Array, num_tokens_col: &arrow_array::UInt32Array, is_legacy: bool, - frag_reuse_index: Option>, + frag_reuse_index: Option>, ) -> Result { // for legacy format, the row id is doc id; sorting keeps binary search viable if is_legacy { @@ -5183,6 +6318,8 @@ impl DocSet { num_tokens, inv: Vec::new(), total_tokens, + scoring_quantized: false, + norms: Arc::new(std::sync::OnceLock::new()), }); } @@ -5224,6 +6361,8 @@ impl DocSet { num_tokens, inv, total_tokens, + scoring_quantized: false, + norms: Arc::new(std::sync::OnceLock::new()), }); } @@ -5243,21 +6382,24 @@ impl DocSet { num_tokens, inv, total_tokens, + scoring_quantized: false, + norms: Arc::new(std::sync::OnceLock::new()), }) } // remap the row ids to the new row ids // returns the removed doc ids - pub fn remap(&mut self, mapping: &HashMap>) -> Vec { + pub fn remap(&mut self, mapping: &RowAddrRemap) -> Vec { let mut removed = Vec::new(); let len = self.len(); let row_ids = std::mem::replace(&mut self.row_ids, Vec::with_capacity(len)); let num_tokens = std::mem::replace(&mut self.num_tokens, Vec::with_capacity(len)); + self.invalidate_norms(); self.total_tokens = 0; for (doc_id, (row_id, num_token)) in std::iter::zip(row_ids, num_tokens).enumerate() { - match mapping.get(&row_id) { + match mapping.get(row_id) { Some(Some(new_row_id)) => { - self.row_ids.push(*new_row_id); + self.row_ids.push(new_row_id); self.num_tokens.push(num_token); self.total_tokens += num_token as u64; } @@ -5279,6 +6421,39 @@ impl DocSet { self.num_tokens[doc_id as usize] } + /// Enable quantized doc-length scoring (V3 / 256-doc block partitions). + pub fn set_quantized_scoring(&mut self, quantized: bool) { + self.scoring_quantized = quantized; + } + + /// The quantized doc-length slab when this set scores quantized (V3 + /// partitions), baked on first use; `None` for exact-scoring sets. + pub fn scoring_norms(&self) -> Option<&[u8]> { + if !self.scoring_quantized { + return None; + } + Some( + self.norms + .get_or_init(|| { + self.num_tokens + .iter() + .map(|&n| quantize_doc_length(n)) + .collect() + }) + .as_ref(), + ) + } + + /// Doc length as scoring sees it: the quantized bucket floor for V3 + /// partitions, the exact value otherwise. + #[inline] + pub fn scoring_num_tokens(&self, doc_id: u32) -> u32 { + match self.scoring_norms() { + Some(norms) => dequantize_doc_length(norms[doc_id as usize]), + None => self.num_tokens[doc_id as usize], + } + } + // this can be used only if it's a legacy format, // which store the sorted row ids so that we can use binary search #[inline] @@ -5295,9 +6470,18 @@ impl DocSet { self.row_ids.push(row_id); self.num_tokens.push(num_tokens); self.total_tokens += num_tokens as u64; + self.invalidate_norms(); self.row_ids.len() as u32 - 1 } + // Drop the baked norm slab after a mutation; it re-bakes on the next + // scoring use. + fn invalidate_norms(&mut self) { + if self.norms.get().is_some() { + self.norms = Arc::new(std::sync::OnceLock::new()); + } + } + pub(crate) fn memory_size(&self) -> usize { self.row_ids.capacity() * std::mem::size_of::() + self.num_tokens.capacity() * std::mem::size_of::() @@ -5324,6 +6508,12 @@ pub fn flat_full_text_search( match batches[0][doc_col].data_type() { DataType::Utf8 => do_flat_full_text_search::(batches, doc_col, query, tokenizer), DataType::LargeUtf8 => do_flat_full_text_search::(batches, doc_col, query, tokenizer), + DataType::List(_) => { + do_flat_full_text_search_list::(batches, doc_col, query, tokenizer) + } + DataType::LargeList(_) => { + do_flat_full_text_search_list::(batches, doc_col, query, tokenizer) + } data_type => Err(Error::invalid_input(format!( "unsupported data type {} for inverted index", data_type @@ -5359,6 +6549,46 @@ fn do_flat_full_text_search( Ok(results) } +fn do_flat_full_text_search_list( + batches: &[&RecordBatch], + doc_col: &str, + query: &str, + tokenizer: Option>, +) -> Result> { + let mut results = Vec::new(); + let mut tokenizer = + tokenizer.unwrap_or_else(|| InvertedIndexParams::default().build().unwrap()); + let query_tokens = collect_query_tokens(query, &mut tokenizer); + + for batch in batches { + let row_id_array = batch[ROW_ID].as_primitive::(); + let doc_array = batch[doc_col].as_list::(); + match doc_array.value_type() { + DataType::Utf8 | DataType::LargeUtf8 => {} + data_type => { + return Err(Error::invalid_input(format!( + "unsupported list item data type {} for inverted index", + data_type + ))); + } + } + for i in 0..row_id_array.len() { + if doc_array.is_null(i) { + continue; + } + let elements = doc_array.value(i); + if iter_str_array(elements.as_ref()) + .flatten() + .any(|element| has_query_token(element, &mut tokenizer, &query_tokens)) + { + results.push(row_id_array.value(i)); + } + } + } + + Ok(results) +} + const FLAT_ROW_ID_COL_IDX: usize = 0; const FLAT_ALL_TOKENS_COL_IDX: usize = 1; const FLAT_QUERY_TOKEN_COUNTS_COL_IDX: usize = 2; @@ -5416,6 +6646,8 @@ async fn tokenize_and_count( // thread is invisible to the caller's poll timer otherwise). let start = std::time::Instant::now(); let batch = batch?; + let row_id_array = batch[ROW_ID].as_primitive::(); + let mut row_ids = UInt64Builder::with_capacity(batch.num_rows()); let mut all_token_counts = UInt64Builder::with_capacity(batch.num_rows()); let mut query_token_counts = FixedSizeListBuilder::with_capacity( UInt64Builder::with_capacity(batch.num_rows() * query_tokens.len()), @@ -5423,20 +6655,7 @@ async fn tokenize_and_count( batch.num_rows(), ); let mut temp_query_token_counts = Vec::with_capacity(query_tokens.len()); - let doc_iter = iter_str_array(batch.column(doc_col_idx)); - for doc in doc_iter { - let Some(doc) = doc else { - all_token_counts.append_value(0); - query_token_counts - .values() - .append_value_n(0, query_tokens.len()); - query_token_counts.append(true); - continue; - }; - - temp_query_token_counts.clear(); - temp_query_token_counts.extend(std::iter::repeat_n(0, query_tokens.len())); - + let mut count_text = |doc: &str, temp_query_token_counts: &mut Vec| -> u64 { let mut stream = tokenizer.token_stream_for_doc(doc); let mut all_tokens = 0; while let Some(token) = stream.next() { @@ -5445,20 +6664,67 @@ async fn tokenize_and_count( temp_query_token_counts[token_index] += 1; } } - all_token_counts.append_value(all_tokens); - for count in temp_query_token_counts.iter().copied() { - query_token_counts.values().append_value(count); + all_tokens + }; + let mut append_counts = + |row_id: u64, all_tokens: u64, temp_query_token_counts: &[u64]| { + row_ids.append_value(row_id); + all_token_counts.append_value(all_tokens); + for count in temp_query_token_counts.iter().copied() { + query_token_counts.values().append_value(count); + } + query_token_counts.append(true); + }; + match batch.column(doc_col_idx).data_type() { + DataType::Utf8 | DataType::LargeUtf8 => { + let doc_iter = iter_str_array(batch.column(doc_col_idx)); + for (doc, row_id) in doc_iter.zip(row_id_array.values().iter()) { + temp_query_token_counts.clear(); + temp_query_token_counts + .extend(std::iter::repeat_n(0, query_tokens.len())); + + let Some(doc) = doc else { + append_counts(*row_id, 0, &temp_query_token_counts); + continue; + }; + + let all_tokens = count_text(doc, &mut temp_query_token_counts); + append_counts(*row_id, all_tokens, &temp_query_token_counts); + } + } + DataType::List(_) => { + tokenize_and_count_list::( + batch.column(doc_col_idx), + row_id_array, + &mut count_text, + &mut append_counts, + &mut temp_query_token_counts, + query_tokens.len(), + )?; + } + DataType::LargeList(_) => { + tokenize_and_count_list::( + batch.column(doc_col_idx), + row_id_array, + &mut count_text, + &mut append_counts, + &mut temp_query_token_counts, + query_tokens.len(), + )?; + } + data_type => { + return DataFusionResult::Err(datafusion_common::DataFusionError::Execution( + format!("unsupported data type {} for flat full text search", data_type), + )); } - query_token_counts.append(true); } - let row_ids = batch[ROW_ID].clone(); + let row_ids = row_ids.finish(); let all_token_counts = all_token_counts.finish(); let query_token_counts = query_token_counts.finish(); let result_batch = RecordBatch::try_new( - output_schema, vec![ - row_ids, + Arc::new(row_ids) as ArrayRef, Arc::new(all_token_counts) as ArrayRef, Arc::new(query_token_counts) as ArrayRef, ], @@ -5484,6 +6750,47 @@ async fn tokenize_and_count( )?) } +fn tokenize_and_count_list( + doc_col: &ArrayRef, + row_id_array: &arrow_array::PrimitiveArray, + count_text: &mut impl FnMut(&str, &mut Vec) -> u64, + append_counts: &mut impl FnMut(u64, u64, &[u64]), + temp_query_token_counts: &mut Vec, + query_tokens_len: usize, +) -> DataFusionResult<()> { + let doc_array = doc_col.as_list::(); + match doc_array.value_type() { + DataType::Utf8 | DataType::LargeUtf8 => {} + data_type => { + return Err(datafusion_common::DataFusionError::Execution(format!( + "unsupported list item data type {} for flat full text search", + data_type + ))); + } + } + + for i in 0..row_id_array.len() { + if doc_array.is_null(i) { + continue; + } + + temp_query_token_counts.clear(); + temp_query_token_counts.extend(std::iter::repeat_n(0, query_tokens_len)); + + let elements = doc_array.value(i); + let mut all_tokens = 0; + for element in iter_str_array(elements.as_ref()).flatten() { + all_tokens += count_text(element, temp_query_token_counts); + } + + if all_tokens > 0 { + append_counts(row_id_array.value(i), all_tokens, temp_query_token_counts); + } + } + + Ok(()) +} + /// Initialize the BM25 scorer /// /// In order to calculate BM25 scores we need to know token counts for the entire corpus. We extract these from the @@ -5524,7 +6831,9 @@ fn initialize_scorer( for _ in 0..counted_input.num_rows() { for token_count in all_token_counts.iter_mut() { - *token_count += input_token_counters.next().unwrap_or_default(); + if input_token_counters.next().unwrap_or_default() > 0 { + *token_count += 1; + } } } @@ -5728,7 +7037,10 @@ mod tests { use crate::prefilter::NoFilter; use crate::scalar::ScalarIndex; use crate::scalar::inverted::builder::{ - InnerBuilder, InvertedIndexBuilder, PositionRecorder, inverted_list_schema, + InnerBuilder, InvertedIndexBuilder, PositionRecorder, doc_file_path, inverted_list_schema, + inverted_list_schema_for_version_with_block_size, + inverted_list_schema_for_version_with_block_size_and_impacts, posting_file_path, + token_file_path, }; use crate::scalar::inverted::encoding::{ compress_positions, compress_posting_list_with_tail_codec, @@ -5736,13 +7048,16 @@ mod tests { }; use crate::scalar::inverted::query::{FtsSearchParams, Operator}; use crate::scalar::lance_format::LanceIndexStore; - use arrow::array::{AsArray, Int32Builder, LargeBinaryBuilder, ListBuilder, UInt32Builder}; + use arrow::array::{ + AsArray, GenericListBuilder, GenericStringBuilder, Int32Builder, LargeBinaryBuilder, + ListBuilder, UInt32Builder, + }; use arrow::datatypes::{Float32Type, UInt32Type}; use arrow_array::{ArrayRef, Float32Array, RecordBatch, StringArray, UInt32Array, UInt64Array}; use arrow_schema::{DataType, Field, Schema}; use std::collections::HashMap; use std::sync::Arc; - use std::sync::atomic::Ordering; + use std::sync::atomic::{AtomicU32, Ordering}; use crate::scalar::inverted::tokenizer::document_tokenizer::TextTokenizer; use lance_tokenizer::{Language, SimpleTokenizer, StopWordFilter, TextAnalyzer}; @@ -5756,15 +7071,21 @@ mod tests { token: &str, row_id: u64, ) -> Result> { - let mut partition = InnerBuilder::new_with_format_version( + let block_size = params.posting_block_size(); + let format_version = params.resolved_format_version(); + let mut partition = InnerBuilder::new_with_format_version_and_block_size( 0, false, token_set_format, - InvertedListFormatVersion::V1, + format_version, + block_size, ); partition.tokens.add(token.to_owned()); - let mut posting_list = - PostingListBuilder::new_with_posting_tail_codec(false, PostingTailCodec::Fixed32); + let mut posting_list = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + format_version.posting_tail_codec(), + block_size, + ); posting_list.add(0, PositionRecorder::Count(1)); partition.posting_lists.push(posting_list); partition.docs.append(row_id, 1); @@ -5780,6 +7101,15 @@ mod tests { TOKEN_SET_FORMAT_KEY.to_owned(), token_set_format.to_string(), ), + ( + POSTING_TAIL_CODEC_KEY.to_owned(), + format_version.posting_tail_codec().as_str().to_owned(), + ), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + format_version.index_version().to_string(), + ), + (POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string()), ]); let mut writer = store .new_index_file(METADATA_FILE, Arc::new(arrow_schema::Schema::empty())) @@ -5800,6 +7130,146 @@ mod tests { )) } + #[test] + fn test_posting_block_size_schema_metadata() { + assert_eq!(parse_posting_block_size(&HashMap::new()).unwrap(), 128); + + let metadata = HashMap::from([(POSTING_BLOCK_SIZE_KEY.to_owned(), "512".to_owned())]); + let err = parse_posting_block_size(&metadata).unwrap_err(); + assert!(err.to_string().contains("block_size")); + + let metadata = HashMap::from([(POSTING_BLOCK_SIZE_KEY.to_owned(), "129".to_owned())]); + let err = parse_posting_block_size(&metadata).unwrap_err(); + assert!(err.to_string().contains("block_size")); + } + + #[test] + fn test_posting_builder_writes_impacts_for_supported_block_sizes() { + for block_size in [128, 256] { + let format_version = default_fts_format_version_for_block_size(block_size).unwrap(); + let num_docs = block_size * 33 + 1; + let mut docs = DocSet::default(); + let mut posting = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + format_version.posting_tail_codec(), + block_size, + ); + for doc_id in 0..num_docs { + docs.append(doc_id as u64, (doc_id % 5 + 1) as u32); + posting.add( + doc_id as u32, + PositionRecorder::Count((doc_id % 3 + 1) as u32), + ); + } + let schema = + inverted_list_schema_for_version_with_block_size(false, format_version, block_size); + let batch = posting.to_batch_with_docs(&docs, schema).unwrap(); + assert!(batch.column_by_name(IMPACT_COL).is_some()); + let max_score = batch[MAX_SCORE_COL].as_primitive::().value(0); + let length = batch[LENGTH_COL].as_primitive::().value(0); + let posting = PostingList::from_batch(&batch, Some(max_score), Some(length)).unwrap(); + let PostingList::Compressed(posting) = posting else { + panic!("expected compressed posting list"); + }; + let impacts = posting.impacts.expect("posting should include impacts"); + assert_eq!(impacts.level0_len(), posting.blocks.len()); + assert_eq!(impacts.level1_len(), posting.blocks.len().div_ceil(32)); + assert_eq!( + impacts.entries().len(), + impacts.level0_len() + impacts.level1_len() + ); + } + } + + #[test] + fn test_posting_builder_without_impact_column_roundtrips_without_impacts() { + let mut posting = PostingListBuilder::new(false); + for doc_id in 0..BLOCK_SIZE + 3 { + posting.add(doc_id as u32, PositionRecorder::Count(1)); + } + let batch = posting.to_batch(vec![1.0, 1.0]).unwrap(); + assert!(batch.column_by_name(IMPACT_COL).is_none()); + let posting = + PostingList::from_batch(&batch, Some(1.0), Some((BLOCK_SIZE + 3) as u32)).unwrap(); + assert!(!posting.has_impacts()); + } + + #[tokio::test] + async fn test_build_search_uses_configured_posting_block_size() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + let params = InvertedIndexParams::default().block_size(256).unwrap(); + let format_version = params.resolved_format_version(); + let block_size = params.posting_block_size(); + let num_docs = block_size + 7; + + let mut builder = InnerBuilder::new_with_format_version_and_block_size( + 0, + false, + TokenSetFormat::default(), + format_version, + block_size, + ); + builder.tokens.add("needle".to_owned()); + let mut posting_list = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + format_version.posting_tail_codec(), + block_size, + ); + for doc_id in 0..num_docs { + posting_list.add(doc_id as u32, PositionRecorder::Count(1)); + builder.docs.append(1_000 + doc_id as u64, 1); + } + builder.posting_lists.push(posting_list); + builder.write(store.as_ref()).await.unwrap(); + write_test_metadata(&store, vec![0], params).await; + + let cache = Arc::new(LanceCache::with_capacity(4096)); + let index = InvertedIndex::load(store.clone(), None, cache.as_ref()) + .await + .unwrap(); + assert_eq!(index.partitions[0].inverted_list.block_size(), block_size); + + let posting = index.partitions[0] + .inverted_list + .posting_list(0, false, &NoOpMetricsCollector) + .await + .unwrap(); + let PostingList::Compressed(posting) = posting else { + panic!("expected compressed posting list"); + }; + assert_eq!(posting.block_size, block_size); + assert_eq!(posting.blocks.len(), num_docs.div_ceil(block_size)); + let impacts = posting + .impacts + .as_ref() + .expect("newly written posting list should include impacts"); + assert_eq!(impacts.level0_len(), posting.blocks.len()); + assert_eq!(impacts.level1_len(), posting.blocks.len().div_ceil(32)); + assert_eq!( + impacts.entries().len(), + impacts.level0_len() + impacts.level1_len() + ); + + let tokens = Arc::new(Tokens::new(vec!["needle".to_owned()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(10))); + let prefilter = Arc::new(NoFilter); + let metrics = Arc::new(NoOpMetricsCollector); + let (row_ids, scores) = index + .bm25_search(tokens, params, Operator::Or, prefilter, metrics, None) + .await + .unwrap(); + + assert_eq!(row_ids.len(), 10); + assert_eq!(scores.len(), 10); + assert!(row_ids.iter().all(|row_id| *row_id >= 1_000)); + } + #[tokio::test] async fn test_posting_builder_remap() { let posting_tail_codec = PostingTailCodec::Fixed32; @@ -6030,15 +7500,28 @@ mod tests { } #[test] - fn test_resolve_fts_format_version_defaults_to_v1() { + fn test_resolve_fts_format_version_defaults_to_v2() { assert_eq!( resolve_fts_format_version(None).unwrap(), - InvertedListFormatVersion::V1 + InvertedListFormatVersion::V2 ); assert_eq!( resolve_fts_format_version(Some("2")).unwrap(), InvertedListFormatVersion::V2 ); + assert_eq!( + resolve_fts_format_version(Some("3")).unwrap(), + InvertedListFormatVersion::V3 + ); + } + + #[test] + fn test_block_size_256_metadata_resolves_to_v3() { + let metadata = HashMap::from([(POSTING_BLOCK_SIZE_KEY.to_owned(), "256".to_owned())]); + assert_eq!( + parse_format_version_from_metadata(&metadata).unwrap(), + InvertedListFormatVersion::V3 + ); } #[test] @@ -6313,7 +7796,7 @@ mod tests { let mut builder = index.into_builder().await.unwrap(); let mapping = HashMap::from([(0, None), (2, Some(3))]); - builder.remap(&mapping).await.unwrap(); + builder.remap(&RowAddrRemap::direct(mapping)).await.unwrap(); // after remap, the doc 0 is removed, and the doc 2 is updated to 3 assert_eq!(builder.tokens.len(), 1); @@ -6328,7 +7811,7 @@ mod tests { // remap to delete all docs let mapping = HashMap::from([(1, None), (3, None)]); - builder.remap(&mapping).await.unwrap(); + builder.remap(&RowAddrRemap::direct(mapping)).await.unwrap(); assert_eq!(builder.tokens.len(), 0); assert_eq!(builder.posting_lists.len(), 0); @@ -6463,7 +7946,7 @@ mod tests { } #[tokio::test] - async fn test_modern_prewarm_shrinks_cached_posting_buffers() { + async fn test_modern_prewarm_packs_group_with_shared_posting_buffer() { let tmpdir = TempObjDir::default(); let store = Arc::new(LanceIndexStore::new( ObjectStore::local().into(), @@ -6515,6 +7998,10 @@ mod tests { !inverted_list.is_legacy_layout(), "test should use modern posting layout" ); + assert!( + inverted_list.has_impacts, + "modern posting fixture should include impact skip data" + ); inverted_list.prewarm_posting_lists(false, 2).await.unwrap(); @@ -6523,109 +8010,76 @@ mod tests { let (start, end) = inverted_list.group_range_for_token(0).unwrap(); let group = inverted_list .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&posting_list_group_cache_key( + start, + end, + inverted_list.has_impacts, + )) .await .unwrap(); - let PostingList::Compressed(alpha) = group.get(0).unwrap() else { + assert!( + group.is_packed(), + "no-position prewarm should pack v2 groups" + ); + assert!( + group.needs_external_metadata(), + "prewarmed packed groups must not duplicate reader score/length metadata" + ); + let (alpha_score, alpha_len) = inverted_list.bulk_metadata_for_token(0); + let PostingList::Compressed(alpha) = group + .posting_list(0, alpha_score, alpha_len) + .unwrap() + .unwrap() + else { panic!("expected compressed posting list for token 0"); }; - let PostingList::Compressed(beta) = group.get(1).unwrap() else { - panic!("expected compressed posting list for token 1"); + let PostingList::Compressed(alpha_again) = group + .posting_list(0, alpha_score, alpha_len) + .unwrap() + .unwrap() + else { + panic!("expected compressed posting list for repeated token 0 access"); }; - - assert_ne!( - alpha.blocks.values().as_ptr(), - beta.blocks.values().as_ptr(), - "prewarm should not leave cached posting lists sharing the same values buffer" - ); - } - - #[test] - fn test_group_aligned_chunk_end_boundary_cases() { - let starts = [0, 3, 7, 10]; - let token_count = 13; - - assert_eq!( - group_aligned_chunk_end(&starts, token_count, 0, 5), - 3, - "chunk should snap back to the largest group boundary that fits" - ); - assert_eq!( - group_aligned_chunk_end(&starts, token_count, 3, 6), - 7, - "oversized groups should run as one chunk" - ); - assert_eq!( - group_aligned_chunk_end(&starts, token_count, 7, 10), - 10, - "an exact next group boundary should be selected" - ); - assert_eq!( - group_aligned_chunk_end(&starts, token_count, 10, 12), - 13, - "the last group should extend to token_count" - ); - assert_eq!( - group_aligned_chunk_end(&starts, token_count, 7, 13), - 13, - "token_count should act as the final boundary" - ); - } - - #[test] - fn test_group_start_indices_for_chunk_boundary_cases() { - let starts = [0, 3, 7, 10]; - let token_count = 13; - let ranges_for_chunk = |tok_start, tok_end| { - group_start_indices_for_chunk(&starts, tok_start, tok_end) - .map(|group_idx| group_range_for_start_index(&starts, token_count, group_idx)) - .collect::>() + let (beta_score, beta_len) = inverted_list.bulk_metadata_for_token(1); + let PostingList::Compressed(beta) = group + .posting_list(1, beta_score, beta_len) + .unwrap() + .unwrap() + else { + panic!("expected compressed posting list for token 1"); }; - assert_eq!( - ranges_for_chunk(0, 7), - vec![(0, 3), (3, 7)], - "publish should include only groups that start in the chunk" + assert!( + alpha.impacts.is_some() && beta.impacts.is_some(), + "packed prewarm must preserve impact skip data" ); - assert_eq!( - ranges_for_chunk(7, 13), - vec![(7, 10), (10, 13)], - "publish should include the final group ending at token_count" + assert!( + alpha + .impacts + .as_ref() + .unwrap() + .shares_derived_state_with(alpha_again.impacts.as_ref().unwrap()), + "repeated packed slot access must share decoded impact state" ); - assert_eq!( - ranges_for_chunk(3, 10), - vec![(3, 7), (7, 10)], - "publish selection should work for an interior chunk" + assert!( + alpha.shares_first_docs_with(&alpha_again), + "repeated packed slot access must share decoded block heads" ); - } - - #[test] - fn test_prewarm_chunk_ranges_preserve_group_boundaries() { - let starts = [0, 3, 7, 10]; assert_eq!( - prewarm_chunk_ranges(Some(&starts), 13, 5), - vec![(0, 3), (3, 7), (7, 10), (10, 13)], - "grouped chunk ranges must never split a posting cache group" + alpha.block_first_docs().as_ptr(), + alpha_again.block_first_docs().as_ptr(), + "packed block heads should be decoded only once per slot" ); assert_eq!( - prewarm_chunk_ranges(None, 13, 5), - vec![(0, 5), (5, 10), (10, 13)], - "ungrouped chunk ranges should use plain token ranges" + alpha.blocks.values().as_ptr(), + beta.blocks.values().as_ptr(), + "packed posting views should share the group's values buffer" ); } - /// Prewarming a large partition in multiple chunks must end up holding exactly the - /// same per-token posting lists (doc ids and frequencies) as the whole-file path. - /// Parametrized over layout: the legacy-v1 chunk path rebases global offsets to - /// chunk-local rows, which the v2 one-row-per-token path never exercises. - #[rstest::rstest] - #[case::v1(InvertedListFormatVersion::V1)] - #[case::v2(InvertedListFormatVersion::V2)] #[tokio::test] - async fn test_prewarm_streams_in_chunks_preserves_content( - #[case] format_version: InvertedListFormatVersion, - ) { + async fn test_packed_prewarm_groups_do_not_retain_the_full_chunk() { let tmpdir = TempObjDir::default(); let store = Arc::new(LanceIndexStore::new( ObjectStore::local().into(), @@ -6633,30 +8087,168 @@ mod tests { Arc::new(LanceCache::no_cache()), )); - // One partition with many tokens (so it spans many chunks) and several - // docs per token (so each token is more than one posting row). - const NUM_TOKENS: u32 = 20; - const DOCS_PER_TOKEN: u32 = 3; - let posting_tail_codec = format_version.posting_tail_codec(); - let mut builder = InnerBuilder::new_with_format_version( - 0, - false, - TokenSetFormat::default(), - format_version, - ); - // Small groups so the partition spans several; chunks snap to whole groups, - // so several groups are needed to stream in more than one chunk. - builder.group_config = PostingGroupConfig { - target_bytes: 4096, - max_tokens: 4, + let mut builder = InnerBuilder::new(0, false, TokenSetFormat::default()); + for token_id in 0..4u32 { + builder.tokens.add(format!("t{token_id}")); + let mut posting = PostingListBuilder::new(false); + posting.add(token_id, PositionRecorder::Count(1)); + builder.posting_lists.push(posting); + builder.docs.append(1000 + token_id as u64, 1); + } + builder.write(store.as_ref()).await.unwrap(); + + let reader = store.open_index_file(&posting_file_path(0)).await.unwrap(); + let cache = LanceCache::with_capacity(1 << 20); + let mut posting_reader = PostingListReader::try_new(reader, &cache).await.unwrap(); + posting_reader.grouping = PostingGrouping::SyntheticFixed { group_size: 2 }; + + assert_eq!( + posting_reader + .prewarm_posting_lists_chunked(false, Some(4), 1) + .await + .unwrap(), + 1, + "the test must read both groups in one prewarm chunk" + ); + + let first_group = posting_reader + .index_cache + .get_with_key(&posting_list_group_cache_key( + 0, + 2, + posting_reader.has_impacts, + )) + .await + .unwrap(); + let second_group = posting_reader + .index_cache + .get_with_key(&posting_list_group_cache_key( + 2, + 4, + posting_reader.has_impacts, + )) + .await + .unwrap(); + let (first_score, first_len) = posting_reader.bulk_metadata_for_token(0); + let PostingList::Compressed(first) = first_group + .posting_list(0, first_score, first_len) + .unwrap() + .unwrap() + else { + panic!("expected compressed posting list in first group"); }; + let (neighbor_score, neighbor_len) = posting_reader.bulk_metadata_for_token(1); + let PostingList::Compressed(first_neighbor) = first_group + .posting_list(1, neighbor_score, neighbor_len) + .unwrap() + .unwrap() + else { + panic!("expected compressed posting list in first group"); + }; + let (second_score, second_len) = posting_reader.bulk_metadata_for_token(2); + let PostingList::Compressed(second) = second_group + .posting_list(0, second_score, second_len) + .unwrap() + .unwrap() + else { + panic!("expected compressed posting list in second group"); + }; + + assert_eq!( + first.blocks.values().as_ptr(), + first_neighbor.blocks.values().as_ptr(), + "postings in one group should share the group's values buffer" + ); + assert_ne!( + first.blocks.values().as_ptr(), + second.blocks.values().as_ptr(), + "each group must own a compact buffer instead of retaining the full chunk" + ); + } + + #[test] + fn test_prewarm_chunk_ranges_preserve_group_boundaries() { + let grouping = PostingGrouping::SyntheticFixed { group_size: 4 }; + assert_eq!( + prewarm_chunk_ranges(&grouping, 13, 5), + vec![(0, 4), (4, 8), (8, 13)], + "grouped chunks may contain multiple groups but must never split one" + ); + assert_eq!( + prewarm_chunk_ranges(&PostingGrouping::None, 13, 5), + vec![(0, 5), (5, 10), (10, 13)], + "ungrouped chunk ranges should use plain token ranges" + ); + } + + #[test] + fn test_synthetic_grouping_preserves_fixed_boundaries() { + let grouping = PostingGrouping::SyntheticFixed { group_size: 4 }; + assert_eq!( + grouping.range_for_token(5, 10), + Some((4, 8)), + "synthetic token groups should be fixed-size ranges" + ); + assert_eq!( + grouping.range_for_token(9, 10), + Some((8, 10)), + "the final synthetic group should end at token_count" + ); + assert_eq!( + prewarm_chunk_ranges(&grouping, 10, 6), + vec![(0, 4), (4, 10)], + "prewarm chunks may contain multiple synthetic groups but must not split one" + ); + assert_eq!( + grouping.ranges_for_chunk(4, 10, 10), + vec![(4, 8), (8, 10)], + "publish selection should enumerate synthetic groups in a chunk" + ); + } + + /// Prewarming a large partition in multiple chunks must end up holding exactly the + /// same per-token posting lists (doc ids and frequencies) as the whole-file path. + /// Parametrized over layout: the legacy-v1 chunk path rebases global offsets to + /// chunk-local rows, while the modern one-row-per-token path covers both + /// legacy-sized v2 and 256-doc v3 posting blocks. + #[rstest::rstest] + #[case::v1(InvertedListFormatVersion::V1, LEGACY_BLOCK_SIZE)] + #[case::v2(InvertedListFormatVersion::V2, LEGACY_BLOCK_SIZE)] + #[case::v3(InvertedListFormatVersion::V3, 256)] + #[tokio::test] + async fn test_prewarm_streams_in_chunks_preserves_content( + #[case] format_version: InvertedListFormatVersion, + #[case] block_size: usize, + ) { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + // One partition with enough tokens to span multiple runtime synthetic + // groups and several docs per token. + let num_tokens = runtime_posting_group_tokens() as u32 + 4; + const DOCS_PER_TOKEN: u32 = 3; + let posting_tail_codec = format_version.posting_tail_codec(); + let mut builder = InnerBuilder::new_with_format_version_and_block_size( + 0, + false, + TokenSetFormat::default(), + format_version, + block_size, + ); // expected[token] = [(doc_id, frequency)] in stored (doc-id) order. let mut expected: Vec> = Vec::new(); let mut doc_id = 0u64; - for t in 0..NUM_TOKENS { + for t in 0..num_tokens { builder.tokens.add(format!("tok_{t:03}")); - let mut posting = - PostingListBuilder::new_with_posting_tail_codec(false, posting_tail_codec); + let mut posting = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + posting_tail_codec, + block_size, + ); let mut docs = Vec::new(); for _ in 0..DOCS_PER_TOKEN { posting.add(doc_id as u32, PositionRecorder::Count(1)); @@ -6669,15 +8261,15 @@ mod tests { } builder.write(store.as_ref()).await.unwrap(); + let params = InvertedIndexParams::default() + .block_size(block_size) + .unwrap(); let metadata = std::collections::HashMap::from_iter(vec![ ( "partitions".to_owned(), serde_json::to_string(&vec![0u64]).unwrap(), ), - ( - "params".to_owned(), - serde_json::to_string(&InvertedIndexParams::default()).unwrap(), - ), + ("params".to_owned(), serde_json::to_string(¶ms).unwrap()), ( TOKEN_SET_FORMAT_KEY.to_owned(), TokenSetFormat::default().to_string(), @@ -6686,6 +8278,11 @@ mod tests { POSTING_TAIL_CODEC_KEY.to_owned(), posting_tail_codec.as_str().to_owned(), ), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + format_version.index_version().to_string(), + ), + (POSTING_BLOCK_SIZE_KEY.to_owned(), block_size.to_string()), ]); let mut writer = store .new_index_file(METADATA_FILE, Arc::new(arrow_schema::Schema::empty())) @@ -6698,10 +8295,12 @@ mod tests { .await .unwrap(); let inverted_list = &index.partitions[0].inverted_list; - assert_eq!(inverted_list.len(), NUM_TOKENS as usize); + assert_eq!(inverted_list.len(), num_tokens as usize); + assert_eq!(inverted_list.block_size(), block_size); - // Force a small chunk so the partition deterministically splits; with - // CHUNK_TOKENS < NUM_TOKENS each chunk is bounded below the whole partition. + // Force a small target chunk. Since CHUNK_TOKENS is below the runtime + // group size, synthetic group alignment should still split only at + // group boundaries. const CHUNK_TOKENS: usize = 6; let chunk_count = inverted_list .prewarm_posting_lists_chunked(false, Some(CHUNK_TOKENS), 2) @@ -6716,9 +8315,34 @@ mod tests { "single partition must be streamed in more than one chunk, got {chunk_count}" ); + if format_version == InvertedListFormatVersion::V3 { + let (start, end) = inverted_list.group_range_for_token(0).unwrap(); + let group = inverted_list + .index_cache + .get_with_key(&posting_list_group_cache_key( + start, + end, + inverted_list.has_impacts, + )) + .await + .expect("v3 prewarm should populate the packed group cache"); + assert!(group.is_packed()); + let (max_score, length) = inverted_list.bulk_metadata_for_token(0); + let PostingList::Compressed(posting) = + group.posting_list(0, max_score, length).unwrap().unwrap() + else { + panic!("expected compressed v3 posting list"); + }; + assert_eq!(posting.block_size, 256); + assert!( + posting.impacts.is_some(), + "v3 packed prewarm must preserve impact skip data" + ); + } + // (2) Correctness: every token's posting list round-trips with exactly // the doc ids and frequencies of the whole-file path. - for token_id in 0..NUM_TOKENS { + for token_id in 0..num_tokens { let actual = inverted_list .posting_list(token_id, false, &NoOpMetricsCollector) .await @@ -6747,7 +8371,7 @@ mod tests { let format_version = InvertedListFormatVersion::V2; let posting_tail_codec = format_version.posting_tail_codec(); - const NUM_TOKENS: u32 = 16; + let num_tokens = runtime_posting_group_tokens() as u32 + 4; const DOCS_PER_TOKEN: u32 = 3; let mut builder = InnerBuilder::new_with_format_version( 0, @@ -6755,14 +8379,10 @@ mod tests { TokenSetFormat::default(), format_version, ); - builder.group_config = PostingGroupConfig { - target_bytes: 4096, - max_tokens: 4, - }; // expected[token] = [(doc_id, frequency, positions)]. let mut expected: Vec)>> = Vec::new(); let mut doc_id = 0u64; - for t in 0..NUM_TOKENS { + for t in 0..num_tokens { builder.tokens.add(format!("tok_{t:03}")); let mut posting = PostingListBuilder::new_with_posting_tail_codec(true, posting_tail_codec); @@ -6830,17 +8450,29 @@ mod tests { "partition must be streamed in more than one chunk, got {chunk_count}" ); - for token_id in 0..NUM_TOKENS { + for token_id in 0..num_tokens { // The prewarmed posting cache entry is positions-free. let (start, end) = inverted_list.group_range_for_token(token_id).unwrap(); let group = inverted_list .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&posting_list_group_cache_key( + start, + end, + inverted_list.has_impacts, + )) .await .unwrap(); let slot = (token_id - start) as usize; assert!( - !group.get(slot).unwrap().has_position(), + !group.is_packed(), + "with-position prewarm should retain the materialized fallback" + ); + assert!( + !group + .posting_list(slot, None, None) + .unwrap() + .unwrap() + .has_position(), "token {token_id} posting cache entry must be positions-free after prewarm" ); @@ -7176,8 +8808,13 @@ mod tests { // K adjacent cold tokens shares a single group cache entry: one // read_range bounded by the group size, independent of the partition's // total token count. - let num_tokens = 500; - let queried_tokens: [u32; 4] = [0, 1, 2, 3]; + let runtime_group_size = runtime_posting_group_tokens().max(1); + let queried_token_count = runtime_group_size.min(4); + let queried_tokens = (0..queried_token_count as u32).collect::>(); + let num_tokens = runtime_group_size + .saturating_mul(2) + .max(queried_token_count + 1) + .min(1024); let (index, counter, _tmpdir) = load_counted_v2_index(num_tokens, LanceCache::no_cache()).await; let inverted_list = index.partitions[0].inverted_list.clone(); @@ -7186,15 +8823,18 @@ mod tests { "this test only proves the lazy path for v2 indexes", ); assert!( - inverted_list.group_starts.is_some(), - "freshly written v2 index should carry posting group offsets", + matches!( + &inverted_list.grouping, + PostingGrouping::SyntheticFixed { .. } + ), + "freshly written v2 index should use runtime synthetic groups", ); // This fixture uses a no-op cache, so each call re-reads; that isolates // the per-query read shape. Each posting_list call reads exactly its // own group — bounded by the group size, never the full token table. let metrics = Arc::new(NoOpMetricsCollector); - for token_id in queried_tokens { + for &token_id in &queried_tokens { inverted_list .posting_list(token_id, false, metrics.as_ref()) .await @@ -7204,9 +8844,9 @@ mod tests { let (start, end) = inverted_list.group_range_for_token(0).unwrap(); let group_len = (end - start) as usize; assert!( - (queried_tokens.len()..num_tokens).contains(&group_len), - "group [{start}, {end}) should cover the queried neighborhood but be \ - far smaller than the {num_tokens}-token table", + (queried_tokens.len()..=num_tokens).contains(&group_len), + "group [{start}, {end}) should cover the queried neighborhood and \ + stay bounded by the {num_tokens}-token table", ); assert_eq!( counter.read_range_calls(), @@ -7222,8 +8862,8 @@ mod tests { } /// Build a single-partition v2 index where every token's posting list spans - /// `docs_per_token` docs. Small `docs_per_token` yields tiny posting lists - /// that the writer packs densely into shared cache groups. + /// `docs_per_token` docs. Runtime grouping packs consecutive token rows + /// into shared cache groups. async fn load_v2_index_with_grouped_postings( num_tokens: usize, docs_per_token: usize, @@ -7280,25 +8920,18 @@ mod tests { (index, cache) } - /// The read path decodes a posting-list group by slicing one buffer read for - /// the whole `[start, end)` row range, so every posting list in a cached - /// group shares a single `blocks` buffer. `DeepSizeOf` must count each - /// posting's slice of that buffer, not the whole buffer once per posting — - /// otherwise a group of N postings reports ~N times its real footprint. - #[rstest::rstest] - #[case::single_doc_terms(512, 1)] - #[case::small_terms(512, 4)] - #[case::medium_terms(256, 32)] + /// Packed groups charge their Arrow buffers and contiguous metadata once, + /// avoiding the per-member enum/array object graph of a materialized group. #[tokio::test] - async fn test_read_path_group_size_counts_slices_not_shared_buffer( - #[case] num_tokens: usize, - #[case] docs_per_token: usize, - ) { - let (index, _cache) = load_v2_index_with_grouped_postings(num_tokens, docs_per_token).await; + async fn test_packed_group_deep_size_is_smaller_than_materialized_graph() { + let (index, _cache) = load_v2_index_with_grouped_postings(512, 1).await; let inverted_list = index.partitions[0].inverted_list.clone(); assert!(!inverted_list.is_legacy_layout(), "expected v2 layout"); assert!( - inverted_list.group_starts.is_some(), + matches!( + &inverted_list.grouping, + PostingGrouping::SyntheticFixed { .. } + ), "expected grouped posting lists" ); @@ -7310,22 +8943,31 @@ mod tests { let (start, end) = inverted_list.group_range_for_token(0).unwrap(); let group = inverted_list .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&posting_list_group_cache_key( + start, + end, + inverted_list.has_impacts, + )) .await .unwrap(); + assert!(group.is_packed(), "cold v2 group should use packed storage"); + inverted_list.ensure_metadata_loaded().await.unwrap(); - // Sum what counting the full backing buffer once per posting list would - // charge, and confirm the postings really do share a single buffer. let mut distinct_buffers = std::collections::HashSet::new(); - let mut charged_if_counted_per_posting = 0usize; - for posting in &group.posting_lists { + let mut materialized = Vec::with_capacity(group.len()); + for slot in 0..group.len() { + let (max_score, length) = inverted_list.bulk_metadata_for_token(start + slot as u32); + let posting = group + .posting_list(slot, max_score, length) + .unwrap() + .unwrap(); let PostingList::Compressed(compressed) = posting else { panic!("expected compressed posting lists"); }; - charged_if_counted_per_posting += compressed.blocks.get_buffer_memory_size(); distinct_buffers.insert(compressed.blocks.values().as_ptr()); + materialized.push(PostingList::Compressed(compressed)); } - let posting_count = group.posting_lists.len(); + let posting_count = materialized.len(); assert!( posting_count > 1, @@ -7336,13 +8978,12 @@ mod tests { 1, "read-path postings in a group should share one backing buffer" ); - // With slice-aware accounting the shared buffer is counted ~once, so the - // whole group costs far less than counting it once per posting list. - let reported = group.deep_size_of(); + let packed_size = group.deep_size_of(); + let materialized_size = PostingListGroup::new(materialized).deep_size_of(); assert!( - reported < charged_if_counted_per_posting / 2, - "group deep_size_of {reported}B should not scale with the {posting_count}x-counted \ - shared buffer ({charged_if_counted_per_posting}B)" + packed_size * 4 < materialized_size * 3, + "packed group deep_size_of {packed_size}B should be at least 25% smaller than the \ + {materialized_size}B materialized graph for {posting_count} postings" ); } @@ -7397,6 +9038,8 @@ mod tests { 1.0, SLICE_LEN as u32, PostingTailCodec::Fixed32, + LEGACY_BLOCK_SIZE, + None, None, ); @@ -7544,11 +9187,23 @@ mod tests { let (start, end) = inverted_list.group_range_for_token(0).unwrap(); let group = inverted_list .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&posting_list_group_cache_key( + start, + end, + inverted_list.has_impacts, + )) .await .unwrap(); assert!( - !group.get(0).unwrap().has_position(), + !group.is_packed(), + "with-position prewarm should retain the materialized fallback" + ); + assert!( + !group + .posting_list(0, None, None) + .unwrap() + .unwrap() + .has_position(), "posting cache should remain positions-free after prewarm" ); @@ -7757,6 +9412,7 @@ mod tests { partition_ids: Vec, params: InvertedIndexParams, ) { + let format_version = params.resolved_format_version(); let metadata = HashMap::from([ ( "partitions".to_owned(), @@ -7767,12 +9423,337 @@ mod tests { TOKEN_SET_FORMAT_KEY.to_owned(), TokenSetFormat::default().to_string(), ), + ( + POSTING_TAIL_CODEC_KEY.to_owned(), + format_version.posting_tail_codec().as_str().to_owned(), + ), + ( + FTS_FORMAT_VERSION_KEY.to_owned(), + format_version.index_version().to_string(), + ), + ( + POSTING_BLOCK_SIZE_KEY.to_owned(), + params.posting_block_size().to_string(), + ), ]); let mut writer = store .new_index_file(METADATA_FILE, Arc::new(arrow_schema::Schema::empty())) .await .unwrap(); - writer.finish_with_metadata(metadata).await.unwrap(); + writer.finish_with_metadata(metadata).await.unwrap(); + } + + async fn write_test_partition_with_optional_impacts( + store: &Arc, + partition_id: u64, + mut builder: InnerBuilder, + token_set_format: TokenSetFormat, + with_impacts: bool, + ) { + let format_version = InvertedListFormatVersion::V1; + let block_size = LEGACY_BLOCK_SIZE; + let docs = std::mem::take(&mut builder.docs); + let schema = inverted_list_schema_for_version_with_block_size_and_impacts( + false, + format_version, + block_size, + with_impacts, + ); + + let mut posting_writer = store + .new_index_file(&posting_file_path(partition_id), schema.clone()) + .await + .unwrap(); + for posting_list in std::mem::take(&mut builder.posting_lists) { + let batch = posting_list + .to_batch_with_docs(&docs, schema.clone()) + .unwrap(); + posting_writer.write_record_batch(batch).await.unwrap(); + } + posting_writer.finish().await.unwrap(); + + let token_batch = std::mem::take(&mut builder.tokens) + .to_batch(token_set_format) + .unwrap(); + let mut token_writer = store + .new_index_file(&token_file_path(partition_id), token_batch.schema()) + .await + .unwrap(); + token_writer.write_record_batch(token_batch).await.unwrap(); + token_writer.finish().await.unwrap(); + + let doc_batch = docs.to_batch().unwrap(); + let mut doc_writer = store + .new_index_file(&doc_file_path(partition_id), doc_batch.schema()) + .await + .unwrap(); + doc_writer.write_record_batch(doc_batch).await.unwrap(); + doc_writer.finish().await.unwrap(); + } + + async fn load_global_scoring_test_index( + second_partition_has_impacts: bool, + ) -> (TempObjDir, Arc) { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + let partition_specs = [ + (0, 100, 5_000, 101..111, 5_000, true), + (1, 200, 1_000, 201..301, 1, second_partition_has_impacts), + ]; + for ( + partition_id, + matching_row_id, + matching_doc_length, + other_row_ids, + other_doc_length, + with_impacts, + ) in partition_specs + { + let mut builder = InnerBuilder::new_with_format_version( + partition_id, + false, + TokenSetFormat::default(), + InvertedListFormatVersion::V1, + ); + builder.tokens.add("alpha".to_owned()); + builder + .posting_lists + .push(PostingListBuilder::new_with_posting_tail_codec( + false, + InvertedListFormatVersion::V1.posting_tail_codec(), + )); + builder.posting_lists[0].add(0, PositionRecorder::Count(1)); + builder.docs.append(matching_row_id, matching_doc_length); + for row_id in other_row_ids { + builder.docs.append(row_id, other_doc_length); + } + write_test_partition_with_optional_impacts( + &store, + partition_id, + builder, + TokenSetFormat::default(), + with_impacts, + ) + .await; + } + + write_test_metadata(&store, vec![0, 1], InvertedIndexParams::default()).await; + let cache = LanceCache::with_capacity(4096); + let index = InvertedIndex::load(store, None, &cache).await.unwrap(); + (tmpdir, index) + } + + async fn search_test_impact_partition( + partition: &InvertedPartition, + tokens: &Tokens, + params: &FtsSearchParams, + scorer: Arc, + shared_threshold: Arc, + ) -> Vec { + let LoadedPostings { + postings, + grouped_expansions, + impact_safe, + } = partition + .load_posting_lists( + tokens, + params, + Operator::Or, + scorer.as_ref(), + &NoOpMetricsCollector, + ) + .await + .unwrap(); + assert!(impact_safe); + assert!(grouped_expansions.is_empty()); + + let mask = NoFilter.mask(); + let docs_for_wand = partition.docs.docs_for_wand(mask.as_ref()).await.unwrap(); + let mut candidates = partition + .bm25_search( + docs_for_wand.as_ref(), + params, + Operator::Or, + mask, + postings, + Some(scorer), + &NoOpMetricsCollector, + shared_threshold, + ) + .unwrap(); + resolve_deferred_candidates(&partition.docs, &mut candidates) + .await + .unwrap(); + candidates + } + + #[tokio::test] + async fn test_impact_partitions_share_global_threshold_without_pruning_winner() { + // Partition 0 wins under its local corpus statistics but loses under + // the global statistics. If its local score escapes into the shared + // floor, partition 1 will incorrectly prune the real global winner. + let (_tmpdir, index) = load_global_scoring_test_index(true).await; + let first_partition = index + .partitions + .iter() + .find(|partition| partition.id() == 0) + .unwrap(); + let second_partition = index + .partitions + .iter() + .find(|partition| partition.id() == 1) + .unwrap(); + + let tokens = Arc::new(Tokens::new(vec!["alpha".to_owned()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(1))); + let scorer = Arc::new( + index + .bm25_base_scorer(tokens.as_ref(), params.as_ref()) + .await + .unwrap(), + ); + first_partition + .inverted_list + .ensure_metadata_loaded() + .await + .unwrap(); + second_partition + .inverted_list + .ensure_metadata_loaded() + .await + .unwrap(); + let first_local_scorer = IndexBM25Scorer::new(std::iter::once(first_partition.as_ref())); + let second_local_scorer = IndexBM25Scorer::new(std::iter::once(second_partition.as_ref())); + let first_local_score = + first_local_scorer.query_weight("alpha") * first_local_scorer.doc_weight(1, 5_000); + let second_local_score = + second_local_scorer.query_weight("alpha") * second_local_scorer.doc_weight(1, 1_000); + assert!(first_local_score > second_local_score); + let shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + + // Search sequentially so partition 0 deterministically publishes its + // score before partition 1 evaluates its impact upper bound. + let first_candidates = search_test_impact_partition( + first_partition, + tokens.as_ref(), + params.as_ref(), + scorer.clone(), + shared_threshold.clone(), + ) + .await; + assert_eq!(first_candidates.len(), 1); + assert!(matches!( + first_candidates[0].addr, + CandidateAddr::RowId(100) + )); + let first_score = + scorer.query_weight("alpha") * scorer.doc_weight(1, first_candidates[0].doc_length); + let published_threshold = f32::from_bits(shared_threshold.load(Ordering::Relaxed)); + assert!( + (published_threshold - first_score).abs() < 1e-6, + "published threshold: {published_threshold}, expected global score: {first_score}" + ); + + let second_candidates = search_test_impact_partition( + second_partition, + tokens.as_ref(), + params.as_ref(), + scorer.clone(), + shared_threshold.clone(), + ) + .await; + assert_eq!(second_candidates.len(), 1); + assert!(matches!( + second_candidates[0].addr, + CandidateAddr::RowId(200) + )); + let second_score = + scorer.query_weight("alpha") * scorer.doc_weight(1, second_candidates[0].doc_length); + assert!( + second_score > first_score, + "second score: {second_score}, first score: {first_score}" + ); + assert!( + (f32::from_bits(shared_threshold.load(Ordering::Relaxed)) - second_score).abs() < 1e-6 + ); + + let (row_ids, scores) = index + .bm25_search( + tokens, + params, + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + assert_eq!(row_ids, vec![200]); + assert_eq!(scores.len(), 1); + assert!((scores[0] - second_score).abs() < 1e-6); + } + + #[tokio::test] + async fn test_mixed_impact_and_legacy_partitions_use_global_final_scores() { + let (_tmpdir, index) = load_global_scoring_test_index(false).await; + + let impact_partition = index + .partitions + .iter() + .find(|partition| partition.id() == 0) + .unwrap(); + let legacy_partition = index + .partitions + .iter() + .find(|partition| partition.id() == 1) + .unwrap(); + + let impact_posting = impact_partition + .inverted_list + .posting_list(0, false, &NoOpMetricsCollector) + .await + .unwrap(); + assert!(impact_posting.has_impacts()); + + let legacy_posting = legacy_partition + .inverted_list + .posting_list(0, false, &NoOpMetricsCollector) + .await + .unwrap(); + assert!(!legacy_posting.has_impacts()); + + let tokens = Arc::new(Tokens::new(vec!["alpha".to_string()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(1))); + let (row_ids, scores) = index + .bm25_search( + tokens.clone(), + params.clone(), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + + assert_eq!(row_ids, vec![200]); + assert_eq!(row_ids.len(), scores.len()); + + let scorer = index + .bm25_base_scorer(tokens.as_ref(), params.as_ref()) + .await + .unwrap(); + let expected_score = scorer.query_weight("alpha") * scorer.doc_weight(1, 1_000); + assert!( + (scores[0] - expected_score).abs() < 1e-6, + "score: {}, expected: {}", + scores[0], + expected_score + ); } #[tokio::test] @@ -7832,6 +9813,85 @@ mod tests { ); } + // Enough distinct tokens that `write_posting_lists` emits several posting-list + // batches (the default batch size is 256 rows), exercising the restructured + // producer and async send path. + const MANY_BATCH_TOKENS: u64 = 1000; + const MANY_BATCH_ROW_ID_BASE: u64 = 1000; + + // Writes a single partition whose posting lists span many output batches. Each + // token `tok{i:05}` maps to row id `MANY_BATCH_ROW_ID_BASE + i`. + async fn write_partition_spanning_many_batches(store: &dyn IndexStore) { + let mut builder = InnerBuilder::new(0, false, TokenSetFormat::default()); + for i in 0..MANY_BATCH_TOKENS { + // Zero-padded so tokens are inserted in sorted order, as the set expects. + builder.tokens.add(format!("tok{i:05}")); + let doc_id = builder.docs.append(MANY_BATCH_ROW_ID_BASE + i, 1); + let mut posting_list = PostingListBuilder::new(false); + posting_list.add(doc_id, PositionRecorder::Count(1)); + builder.posting_lists.push(posting_list); + } + builder + .write(store) + .await + .expect("writing posting lists should succeed"); + } + + // Correctness guard for the restructured posting-list writer. The producer now + // builds each batch in its own `spawn_cpu` call, handing the builder and the + // remaining posting lists back out so state (the cross-batch cache-group + // accumulator) is preserved, and dispatches with an async `send().await`. This + // verifies that path over many batches by checking representative tokens after + // they cross the producer/consumer boundary. + // + // Note: this does not reproduce the single-thread-pool deadlock the async send + // fixes -- that requires a 1-thread CPU pool (a process-global singleton) plus + // ~8MB of buffered posting data to trigger a consumer-side encoder flush, which + // is impractical as a lightweight unit test. + #[tokio::test] + async fn test_write_many_posting_list_batches_preserves_all_batches() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + write_partition_spanning_many_batches(store.as_ref()).await; + + write_test_metadata(&store, vec![0], InvertedIndexParams::default()).await; + let cache = Arc::new(LanceCache::with_capacity(4096)); + let index = InvertedIndex::load(store.clone(), None, cache.as_ref()) + .await + .unwrap(); + + // Probe tokens from the first, a middle, and the last batch to confirm they + // remain queryable after crossing the producer/consumer boundary. + for token_idx in [0u64, MANY_BATCH_TOKENS / 2, MANY_BATCH_TOKENS - 1] { + let tokens = Arc::new(Tokens::new( + vec![format!("tok{token_idx:05}")], + DocType::Text, + )); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(10))); + let (row_ids, _) = index + .bm25_search( + tokens, + params, + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + assert_eq!( + row_ids, + vec![MANY_BATCH_ROW_ID_BASE + token_idx], + "token tok{token_idx:05} should map to its single document" + ); + } + } + #[tokio::test] async fn test_and_query_skips_partition_missing_required_term() { let tmpdir = TempObjDir::default(); @@ -8031,6 +10091,132 @@ mod tests { ); } + /// Write one partition holding `variants` in order, with one + /// single-token doc per variant taken from `row_ids`. + async fn write_variant_partition( + store: &Arc, + partition_id: u64, + variants: &[&str], + row_ids: &[u64], + ) { + let mut builder = InnerBuilder::new(partition_id, false, TokenSetFormat::default()); + for token in variants { + builder.tokens.add((*token).to_owned()); + builder.posting_lists.push(PostingListBuilder::new(false)); + } + for (local_idx, row_id) in row_ids.iter().enumerate() { + builder.posting_lists[local_idx].add(local_idx as u32, PositionRecorder::Count(1)); + builder.docs.append(*row_id, 1); + } + builder.write(store.as_ref()).await.unwrap(); + } + + #[tokio::test] + async fn test_fuzzy_expansion_cap_is_global_across_partitions() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + write_variant_partition(&store, 0, &["alpha", "alphb"], &[100, 101]).await; + write_variant_partition(&store, 1, &["alphc", "alphd"], &[102, 103]).await; + write_test_metadata(&store, vec![0, 1], InvertedIndexParams::default()).await; + let cache = Arc::new(LanceCache::with_capacity(4096)); + let index = InvertedIndex::load(store.clone(), None, cache.as_ref()) + .await + .unwrap(); + + let params = FtsSearchParams::new() + .with_fuzziness(Some(1)) + .with_max_expansions(3); + let tokens = Tokens::new(vec!["alphx".to_owned()], DocType::Text); + + let expanded = index.expand_fuzzy_tokens(&tokens, ¶ms).unwrap(); + let expanded_terms = (0..expanded.len()) + .map(|idx| expanded.get_token(idx).to_owned()) + .collect::>(); + assert_eq!( + expanded_terms, + vec!["alpha".to_owned(), "alphb".to_owned(), "alphc".to_owned()], + "max_expansions must cap the whole query across partitions, \ + in lexicographic order" + ); + } + + #[tokio::test] + async fn test_fuzzy_results_independent_of_partition_shape() { + // The same four single-variant docs, laid out as one partition and + // as two. With a binding max_expansions the two shapes must still + // match the same documents with the same scores. + let single_dir = TempObjDir::default(); + let single_store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + single_dir.clone(), + Arc::new(LanceCache::no_cache()), + )); + write_variant_partition( + &single_store, + 0, + &["alpha", "alphb", "alphc", "alphd"], + &[100, 101, 102, 103], + ) + .await; + write_test_metadata(&single_store, vec![0], InvertedIndexParams::default()).await; + + let split_dir = TempObjDir::default(); + let split_store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + split_dir.clone(), + Arc::new(LanceCache::no_cache()), + )); + write_variant_partition(&split_store, 0, &["alpha", "alphb"], &[100, 101]).await; + write_variant_partition(&split_store, 1, &["alphc", "alphd"], &[102, 103]).await; + write_test_metadata(&split_store, vec![0, 1], InvertedIndexParams::default()).await; + + let params = Arc::new( + FtsSearchParams::new() + .with_limit(Some(10)) + .with_fuzziness(Some(1)) + .with_max_expansions(3), + ); + + let mut results = Vec::new(); + for store in [single_store, split_store] { + let cache = LanceCache::with_capacity(4096); + let index = InvertedIndex::load(store, None, &cache).await.unwrap(); + let tokens = Arc::new(Tokens::new(vec!["alphx".to_owned()], DocType::Text)); + let (row_ids, scores) = index + .bm25_search( + tokens, + params.clone(), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + let mut scored = row_ids.into_iter().zip(scores).collect::>(); + scored.sort_unstable_by_key(|(row_id, _)| *row_id); + results.push(scored); + } + + assert_eq!( + results[0] + .iter() + .map(|(row_id, _)| *row_id) + .collect::>(), + vec![100, 101, 102], + "a binding cap keeps the three lexicographically smallest variants" + ); + assert_eq!( + results[0], results[1], + "fuzzy results must not depend on the partition shape" + ); + } + #[tokio::test] async fn test_fuzzy_and_scores_grouped_expansions_by_matched_token() { let tmpdir = TempObjDir::default(); @@ -8458,6 +10644,70 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_block_size_256_writes_v3_metadata_and_index_version() -> Result<()> { + let src_dir = TempObjDir::default(); + let dest_dir = TempObjDir::default(); + let src_store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + src_dir.clone(), + Arc::new(LanceCache::no_cache()), + )); + let dest_store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + dest_dir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + let params = InvertedIndexParams::default().block_size(256)?; + let format_version = params.resolved_format_version(); + assert_eq!(format_version, InvertedListFormatVersion::V3); + + let mut partition = InnerBuilder::new_with_format_version_and_block_size( + 0, + false, + TokenSetFormat::default(), + format_version, + params.posting_block_size(), + ); + partition.tokens.add("hello".to_owned()); + let mut posting_list = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + format_version.posting_tail_codec(), + params.posting_block_size(), + ); + posting_list.add(0, PositionRecorder::Count(1)); + partition.posting_lists.push(posting_list); + partition.docs.append(100, 1); + partition.write(src_store.as_ref()).await?; + + write_test_metadata(&src_store, vec![0], params).await; + + let index = InvertedIndex::load(src_store, None, &LanceCache::no_cache()).await?; + assert_eq!(index.format_version(), InvertedListFormatVersion::V3); + assert_eq!( + index.index_version(), + InvertedListFormatVersion::V3.index_version() + ); + + let created = index + .update(empty_doc_stream(), dest_store.as_ref(), None) + .await?; + assert_eq!( + created.index_version, + InvertedListFormatVersion::V3.index_version() + ); + + let updated = InvertedIndex::load(dest_store, None, &LanceCache::no_cache()).await?; + assert_eq!(updated.format_version(), InvertedListFormatVersion::V3); + assert_eq!( + updated.index_version(), + InvertedListFormatVersion::V3.index_version() + ); + + Ok(()) + } + #[tokio::test] async fn test_merge_segments_preserves_arrow_token_set_format() -> Result<()> { let src_dir = TempObjDir::default(); @@ -8730,58 +10980,68 @@ mod tests { ); } - /// An [`IndexReader`] wrapper that hides the posting-group-offsets schema - /// metadata key, so a [`PostingListReader`] opened on it takes the - /// pre-grouping per-token fallback path (issue #7040). - struct GroupKeyStrippingReader { - inner: Arc, - schema: lance_core::datatypes::Schema, - } + #[tokio::test] + async fn flat_bm25_search_treats_string_lists_as_row_documents() { + let mut docs_builder = + GenericListBuilder::::new(GenericStringBuilder::::new()); + docs_builder.values().append_value("alpha"); + docs_builder.values().append_value("alpha beta"); + docs_builder.append(true); + docs_builder.values().append_value("beta"); + docs_builder.append(true); + docs_builder.append(true); + docs_builder.values().append_null(); + docs_builder.append(true); + docs_builder.append(false); + + let docs = Arc::new(docs_builder.finish()) as ArrayRef; + let schema = Arc::new(Schema::new(vec![ + ROW_ID_FIELD.clone(), + Field::new("text", docs.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt64Array::from(vec![0u64, 1, 2, 3, 4])) as ArrayRef, + docs, + ], + ) + .unwrap(); - impl GroupKeyStrippingReader { - fn new(inner: Arc) -> Self { - let mut schema = inner.schema().clone(); - schema.metadata.remove(POSTING_GROUP_OFFSETS_BUF_KEY); - Self { inner, schema } - } - } + let input: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + schema.clone(), + stream::iter(vec![Ok(batch)]), + )); + let tokenizer: Box = Box::new(TextTokenizer::new( + TextAnalyzer::builder(SimpleTokenizer::default()).build(), + )); - #[async_trait] - impl IndexReader for GroupKeyStrippingReader { - async fn read_record_batch(&self, n: u64, batch_size: u64) -> Result { - self.inner.read_record_batch(n, batch_size).await - } - async fn read_global_buffer(&self, index: u32) -> Result { - self.inner.read_global_buffer(index).await - } - async fn read_range( - &self, - range: std::ops::Range, - projection: Option<&[&str]>, - ) -> Result { - self.inner.read_range(range, projection).await - } - async fn num_batches(&self, batch_size: u64) -> u32 { - self.inner.num_batches(batch_size).await - } - fn num_rows(&self) -> usize { - self.inner.num_rows() - } - fn schema(&self) -> &lance_core::datatypes::Schema { - &self.schema - } + let result_stream = flat_bm25_search_stream_with_metrics( + input, + "text".to_string(), + "alpha".to_string(), + tokenizer, + None, + 100, + None, + ) + .await + .unwrap(); + let batches: Vec<_> = result_stream.try_collect().await.unwrap(); + let scored = arrow::compute::concat_batches(&FTS_SCHEMA, &batches).unwrap(); + let row_ids = scored[ROW_ID].as_primitive::(); + + assert_eq!(row_ids.values(), &[0]); } fn posting_entries(posting: &PostingList) -> Vec<(u64, u32)> { posting.iter().map(|(doc, freq, _)| (doc, freq)).collect() } - /// The grouped read path and the legacy per-token fallback must return - /// identical posting lists for every token, including at group - /// boundaries. Builds a single v2 partition that spans several groups, - /// then reads it both with and without the group offsets present. + /// Runtime synthetic grouping must return correct posting lists for every + /// token, including across synthetic group boundaries. #[tokio::test] - async fn test_posting_list_fallback_matches_grouped() { + async fn test_posting_list_synthetic_grouping_reads_group_boundaries() { let tmpdir = TempObjDir::default(); let store = Arc::new(LanceIndexStore::new( ObjectStore::local().into(), @@ -8789,15 +11049,8 @@ mod tests { Arc::new(LanceCache::no_cache()), )); - // A small token cap forces several groups regardless of the default, - // so the comparison exercises the partition_point math at group - // boundaries. - let num_tokens = 150u32; + let num_tokens = runtime_posting_group_tokens() as u32 + 4; let mut builder = InnerBuilder::new(0, false, TokenSetFormat::default()); - builder.group_config = PostingGroupConfig { - target_bytes: 4096, - max_tokens: 32, - }; for t in 0..num_tokens { builder.tokens.add(format!("t{t}")); let mut pl = PostingListBuilder::new(false); @@ -8809,36 +11062,27 @@ mod tests { let reader = store.open_index_file(&posting_file_path(0)).await.unwrap(); let cache = LanceCache::no_cache(); - let grouped = PostingListReader::try_new(reader.clone(), &cache) - .await - .unwrap(); - assert!( - grouped.group_starts.as_ref().is_some_and(|s| s.len() > 1), - "fixture should span multiple groups", - ); - - let stripped: Arc = Arc::new(GroupKeyStrippingReader::new(reader)); - let fallback = PostingListReader::try_new(stripped, &cache).await.unwrap(); + let posting_reader = PostingListReader::try_new(reader, &cache).await.unwrap(); assert!( - fallback.group_starts.is_none(), - "stripped reader must take the per-token fallback path", + matches!( + &posting_reader.grouping, + PostingGrouping::SyntheticFixed { .. } + ), + "v2 reader must synthesize runtime posting groups", ); let metrics = NoOpMetricsCollector; for token in 0..num_tokens { - let g = grouped.posting_list(token, false, &metrics).await.unwrap(); - let f = fallback.posting_list(token, false, &metrics).await.unwrap(); - assert_eq!( - posting_entries(&g), - posting_entries(&f), - "grouped vs fallback mismatch for token {token}", - ); - assert_eq!(g.len(), f.len(), "length mismatch for token {token}"); + let posting = posting_reader + .posting_list(token, false, &metrics) + .await + .unwrap(); assert_eq!( - g.max_score(), - f.max_score(), - "max_score mismatch for token {token}", + posting_entries(&posting), + vec![(token as u64, 1)], + "synthetic grouping mismatch for token {token}", ); + assert_eq!(posting.len(), 1, "length mismatch for token {token}"); } } @@ -8856,14 +11100,8 @@ mod tests { Arc::new(LanceCache::no_cache()), )); - // Small token cap so the partition spans several groups regardless of - // the default, exercising every group boundary including the last. - let num_tokens = 150u32; + let num_tokens = runtime_posting_group_tokens() as u32 + 4; let mut builder = InnerBuilder::new(0, false, TokenSetFormat::default()); - builder.group_config = PostingGroupConfig { - target_bytes: 4096, - max_tokens: 32, - }; for t in 0..num_tokens { builder.tokens.add(format!("t{t}")); let mut pl = PostingListBuilder::new(false); @@ -8879,11 +11117,11 @@ mod tests { let cache = LanceCache::with_capacity(1 << 20); let posting_reader = PostingListReader::try_new(reader, &cache).await.unwrap(); assert!( - posting_reader - .group_starts - .as_ref() - .is_some_and(|s| s.len() > 1), - "fixture should span multiple groups", + matches!( + &posting_reader.grouping, + PostingGrouping::SyntheticFixed { .. } + ), + "v2 reader should use runtime synthetic groups", ); posting_reader @@ -8896,7 +11134,11 @@ mod tests { assert!( posting_reader .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&posting_list_group_cache_key( + start, + end, + posting_reader.has_impacts, + )) .await .is_some(), "prewarm did not populate group [{start}, {end}) that the read \ @@ -8913,10 +11155,10 @@ mod tests { ); } - /// An empty partition writes no group-offsets buffer, so its reader takes - /// the per-token fallback path (issue #7040). + /// An empty partition has no synthetic groups because there are no token + /// rows to cache. #[tokio::test] - async fn test_empty_partition_has_no_group_offsets() { + async fn test_empty_partition_has_no_synthetic_groups() { let tmpdir = TempObjDir::default(); let store = Arc::new(LanceIndexStore::new( ObjectStore::local().into(), @@ -8928,28 +11170,20 @@ mod tests { builder.write(store.as_ref()).await.unwrap(); let reader = store.open_index_file(&posting_file_path(0)).await.unwrap(); - assert!( - !reader - .schema() - .metadata - .contains_key(POSTING_GROUP_OFFSETS_BUF_KEY), - "empty partition must not write the group-offsets metadata key", - ); - let posting_reader = PostingListReader::try_new(reader, &LanceCache::no_cache()) .await .unwrap(); assert!( - posting_reader.group_starts.is_none(), - "reader for an empty partition must use the per-token fallback path", + matches!(&posting_reader.grouping, PostingGrouping::None), + "reader for an empty partition must not create cache groups", ); assert!(posting_reader.is_empty()); } - /// A posting list that alone exceeds the group target lands in its own - /// `[t, t+1)` group (the clamp case) and reads back intact (issue #7040). + /// A large posting list can share a runtime synthetic group with neighbors; + /// grouping is token-count based and should still read every member intact. #[tokio::test] - async fn test_oversized_term_is_own_group_on_read() { + async fn test_large_posting_reads_inside_synthetic_group() { let tmpdir = TempObjDir::default(); let store = Arc::new(LanceIndexStore::new( ObjectStore::local().into(), @@ -8957,14 +11191,8 @@ mod tests { Arc::new(LanceCache::no_cache()), )); - // A tiny byte target so a modest posting trips the clamp without - // needing a huge fixture; the surrounding tiny terms regroup after it. let mut builder = InnerBuilder::new(0, false, TokenSetFormat::default()); - builder.group_config = PostingGroupConfig { - target_bytes: 50, - max_tokens: 1000, - }; - let big_docs = 30u32; + let big_docs = (BLOCK_SIZE * 3 + 5) as u32; builder.tokens.add("big".to_owned()); let mut big = PostingListBuilder::new(false); for d in 0..big_docs { @@ -8986,11 +11214,12 @@ mod tests { let posting_reader = PostingListReader::try_new(reader, &LanceCache::no_cache()) .await .unwrap(); + let expected_end = runtime_posting_group_tokens().min(5) as u32; assert_eq!( posting_reader.group_range_for_token(0), - Some((0, 1)), - "an oversized term must occupy its own single-row group", + Some((0, expected_end)), + "runtime synthetic grouping should group by token count, not posting bytes", ); let big = posting_reader .posting_list(0, false, &NoOpMetricsCollector) @@ -9005,11 +11234,11 @@ mod tests { assert_eq!(tiny.len(), 1); } - /// When the group offsets are absent, prewarm populates per-token - /// `PostingListKey` entries (the fallback path), matching what the read - /// path then looks up (issue #7040). + /// Non-empty v2 indexes should prewarm synthetic `PostingListGroupKey` + /// entries, matching what the read path then looks up without persisted + /// grouping metadata. #[tokio::test] - async fn test_prewarm_fallback_populates_per_token_entries() { + async fn test_prewarm_synthetic_grouping_populates_group_entries() { let tmpdir = TempObjDir::default(); let store = Arc::new(LanceIndexStore::new( ObjectStore::local().into(), @@ -9029,10 +11258,12 @@ mod tests { builder.write(store.as_ref()).await.unwrap(); let reader = store.open_index_file(&posting_file_path(0)).await.unwrap(); - let stripped: Arc = Arc::new(GroupKeyStrippingReader::new(reader)); let cache = LanceCache::with_capacity(1 << 20); - let posting_reader = PostingListReader::try_new(stripped, &cache).await.unwrap(); - assert!(posting_reader.group_starts.is_none()); + let posting_reader = PostingListReader::try_new(reader, &cache).await.unwrap(); + assert!(matches!( + &posting_reader.grouping, + PostingGrouping::SyntheticFixed { .. } + )); posting_reader .prewarm_posting_lists(false, 2) @@ -9040,13 +11271,34 @@ mod tests { .unwrap(); for token_id in 0..num_tokens { + let (start, end) = posting_reader.group_range_for_token(token_id).unwrap(); + let group = posting_reader + .index_cache + .get_with_key(&posting_list_group_cache_key( + start, + end, + posting_reader.has_impacts, + )) + .await + .unwrap_or_else(|| { + panic!( + "synthetic prewarm should populate group [{start}, {end}) for token {token_id}" + ) + }); + assert!( + group.is_packed(), + "no-position synthetic prewarm should insert a packed group" + ); assert!( posting_reader .index_cache - .get_with_key(&PostingListKey { token_id }) + .get_with_key(&posting_list_cache_key( + token_id, + posting_reader.has_impacts, + )) .await - .is_some(), - "fallback prewarm should populate per-token entry {token_id}", + .is_none(), + "synthetic prewarm should not populate per-token entry {token_id}", ); } } @@ -9063,15 +11315,11 @@ mod tests { Arc::new(LanceCache::no_cache()), )); - // 130 rare tokens (one doc each) plus one common token in every doc; a - // small token cap spreads them across several groups so scoring must - // index into the right group slot. - let num_rare = 130u32; + // Rare tokens (one doc each) plus one common token in every doc. The + // token count exceeds the runtime group size so scoring must index + // into the right synthetic group slot. + let num_rare = runtime_posting_group_tokens() as u32 + 2; let mut builder = InnerBuilder::new(0, false, TokenSetFormat::default()); - builder.group_config = PostingGroupConfig { - target_bytes: 4096, - max_tokens: 32, - }; for t in 0..num_rare { builder.tokens.add(format!("t{t}")); builder.posting_lists.push(PostingListBuilder::new(false)); @@ -9118,7 +11366,7 @@ mod tests { index .bm25_search( Arc::new(Tokens::new(vec![term], DocType::Text)), - Arc::new(FtsSearchParams::new().with_limit(Some(200))), + Arc::new(FtsSearchParams::new().with_limit(Some(num_rare as usize))), Operator::Or, Arc::new(NoFilter), Arc::new(NoOpMetricsCollector), @@ -9129,8 +11377,13 @@ mod tests { } }; - let (rows_70, _) = query("t70").await; - assert_eq!(rows_70, vec![1070], "rare token must map to its single doc"); + let rare_query_id = num_rare / 2; + let (rare_rows, _) = query(&format!("t{rare_query_id}")).await; + assert_eq!( + rare_rows, + vec![1000 + rare_query_id as u64], + "rare token must map to its single doc", + ); // Cold vs warm cache must agree for the common (large) token. let (cold_rows, cold_scores) = query("common").await; diff --git a/rust/lance-index/src/scalar/inverted/iter.rs b/rust/lance-index/src/scalar/inverted/iter.rs index dc07b15769c..3755fa41d80 100644 --- a/rust/lance-index/src/scalar/inverted/iter.rs +++ b/rust/lance-index/src/scalar/inverted/iter.rs @@ -6,10 +6,9 @@ use arrow_array::{Array, LargeBinaryArray}; use super::{ CompressedPositionStorage, PostingList, PostingTailCodec, - builder::BLOCK_SIZE, encoding::{ - decode_position_stream_block, decompress_positions, decompress_posting_block, - decompress_posting_remainder, + MAX_POSTING_BLOCK_SIZE, decode_position_stream_block, decompress_positions, + decompress_posting_block, decompress_posting_remainder, }, }; @@ -28,6 +27,7 @@ impl<'a> PostingListIterator<'a> { posting.blocks.clone(), posting.posting_tail_codec, posting.positions.clone(), + posting.block_size, ))) } } @@ -82,13 +82,14 @@ pub struct CompressedPostingListIterator { next_block_idx: usize, posting_tail_codec: PostingTailCodec, positions: Option, + block_size: usize, idx: usize, doc_ids: Vec, frequencies: Vec, doc_idx_in_block: usize, decoded_positions: Vec, position_offsets: Vec, - buffer: [u32; BLOCK_SIZE], + buffer: [u32; MAX_POSTING_BLOCK_SIZE], } impl CompressedPostingListIterator { @@ -97,10 +98,11 @@ impl CompressedPostingListIterator { blocks: LargeBinaryArray, posting_tail_codec: PostingTailCodec, positions: Option, + block_size: usize, ) -> Self { debug_assert!(length > 0, "length: {}", length); debug_assert_eq!( - length.div_ceil(BLOCK_SIZE), + length.div_ceil(block_size), blocks.len(), "length: {}, num_blocks: {}", length, @@ -108,18 +110,19 @@ impl CompressedPostingListIterator { ); Self { - remainder: length % BLOCK_SIZE, + remainder: length % block_size, blocks, next_block_idx: 0, posting_tail_codec, positions, + block_size, idx: 0, doc_ids: Vec::new(), frequencies: Vec::new(), doc_idx_in_block: 0, decoded_positions: Vec::new(), position_offsets: Vec::new(), - buffer: [0; BLOCK_SIZE], + buffer: [0; MAX_POSTING_BLOCK_SIZE], } } } @@ -163,6 +166,7 @@ impl Iterator for CompressedPostingListIterator { compressed, self.remainder, self.posting_tail_codec, + self.block_size, &mut self.doc_ids, &mut self.frequencies, ); @@ -172,6 +176,7 @@ impl Iterator for CompressedPostingListIterator { &mut self.buffer, &mut self.doc_ids, &mut self.frequencies, + self.block_size, ); } self.doc_idx_in_block = 0; diff --git a/rust/lance-index/src/scalar/inverted/lazy_docset.rs b/rust/lance-index/src/scalar/inverted/lazy_docset.rs index 7a0ee41efd8..a8724513340 100644 --- a/rust/lance-index/src/scalar/inverted/lazy_docset.rs +++ b/rust/lance-index/src/scalar/inverted/lazy_docset.rs @@ -24,7 +24,7 @@ use lance_core::ROW_ID; use lance_core::Result; use tokio::sync::OnceCell; -use crate::frag_reuse::FragReuseIndex; +use crate::scalar::RowIdRemapper; use crate::scalar::inverted::index::{DocSet, NUM_TOKEN_COL}; use crate::scalar::{IndexReader, IndexStore}; use lance_select::mask::RowAddrMask; @@ -63,7 +63,10 @@ pub struct DeferredDocSet { store: Arc, docs_path: String, is_legacy: bool, - frag_reuse_index: Option>, + frag_reuse_index: Option>, + /// V3 (256-doc block) partitions score with quantized doc lengths; the + /// flag is applied to every `DocSet` this deferred set materializes. + quantized_scoring: bool, /// Doc count cached at construction so `len()` stays sync + IO-free. num_rows: usize, /// `sum(num_tokens)` cached on first compute. @@ -74,6 +77,11 @@ pub struct DeferredDocSet { row_ids_col: OnceCell>, /// Full DocSet, materialized on first `ensure_loaded`. full: OnceCell>, + /// num_tokens-only DocSet, materialized on first + /// `ensure_num_tokens_loaded`. Cached because wand scoring calls this + /// once per query per partition; rebuilding it copied the whole + /// num_tokens column (tens of MB per partition) on every query. + tokens_only: OnceCell>, } impl std::fmt::Debug for LazyDocSet { @@ -103,6 +111,10 @@ impl lance_core::deepsize::DeepSizeOf for LazyDocSet { .get() .map(|d| d.deep_size_of_children(ctx)) .unwrap_or(0) + + d.tokens_only + .get() + .map(|d| d.deep_size_of_children(ctx)) + .unwrap_or(0) + d.num_tokens_col .get() .map(|arr| arr.len() * std::mem::size_of::()) @@ -117,23 +129,27 @@ impl lance_core::deepsize::DeepSizeOf for LazyDocSet { } impl LazyDocSet { + #[allow(clippy::too_many_arguments)] pub fn new( store: Arc, docs_path: String, num_rows: usize, is_legacy: bool, - frag_reuse_index: Option>, + frag_reuse_index: Option>, + quantized_scoring: bool, ) -> Self { Self::Deferred(Box::new(DeferredDocSet { store, docs_path, is_legacy, frag_reuse_index, + quantized_scoring, num_rows, total_tokens: OnceCell::new(), num_tokens_col: OnceCell::new(), row_ids_col: OnceCell::new(), full: OnceCell::new(), + tokens_only: OnceCell::new(), })) } @@ -290,7 +306,7 @@ impl DeferredDocSet { .get_or_try_init(|| async { // If the stats path already pulled NUM_TOKEN_COL, // read only ROW_ID and rebuild from the two columns. - let docs = if self.num_tokens_col.get().is_some() { + let mut docs = if self.num_tokens_col.get().is_some() { let num_tokens = self.num_tokens_column().await?; let row_ids = self.row_ids_column().await?; DocSet::from_columns( @@ -307,6 +323,7 @@ impl DeferredDocSet { ) .await? }; + docs.set_quantized_scoring(self.quantized_scoring); Result::Ok(Arc::new(docs)) }) .await? @@ -319,8 +336,16 @@ impl DeferredDocSet { if let Some(full) = self.full.get() { return Ok(full.clone()); } - let num_tokens = self.num_tokens_column().await?; - let docs = Arc::new(DocSet::from_num_tokens_only(num_tokens.as_ref())); + let docs = self + .tokens_only + .get_or_try_init(|| async { + let num_tokens = self.num_tokens_column().await?; + let mut docs = DocSet::from_num_tokens_only(num_tokens.as_ref()); + docs.set_quantized_scoring(self.quantized_scoring); + Result::Ok(Arc::new(docs)) + }) + .await? + .clone(); let _ = self.total_tokens.set(docs.total_tokens_num()); Ok(docs) } diff --git a/rust/lance-index/src/scalar/inverted/scorer.rs b/rust/lance-index/src/scalar/inverted/scorer.rs index eb7d78ff397..35a3be60e0a 100644 --- a/rust/lance-index/src/scalar/inverted/scorer.rs +++ b/rust/lance-index/src/scalar/inverted/scorer.rs @@ -3,6 +3,7 @@ use super::InvertedPartition; use std::collections::HashMap; +use std::sync::Arc; // the Scorer trait is used to calculate the score of a token in a document // in general, the score is calculated as: @@ -10,6 +11,40 @@ use std::collections::HashMap; pub trait Scorer: Send + Sync { fn query_weight(&self, token: &str) -> f32; fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32; + + /// Finite upper bound for every non-negative value returned by + /// [`Self::doc_weight`]. Returning `None` disables score-independent + /// pruning where the posting format has no stored impact bounds. + fn doc_weight_upper_bound(&self) -> Option { + None + } + + /// Stable identity for the corpus-level inputs used by [`Self::doc_weight`]. + /// + /// Implementations should return `Some` only when equal keys guarantee the + /// same document weight for every `(freq, doc_tokens)` pair. Scorers without + /// such an identity keep impact bounds in the query-local cache only. + fn doc_weight_cache_key(&self) -> Option { + None + } +} + +impl Scorer for Arc { + fn query_weight(&self, token: &str) -> f32 { + self.as_ref().query_weight(token) + } + + fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32 { + self.as_ref().doc_weight(freq, doc_tokens) + } + + fn doc_weight_upper_bound(&self) -> Option { + self.as_ref().doc_weight_upper_bound() + } + + fn doc_weight_cache_key(&self) -> Option { + self.as_ref().doc_weight_cache_key() + } } // BM25 parameters @@ -82,6 +117,14 @@ impl Scorer for MemBM25Scorer { let doc_norm = K1 * (1.0 - B + B * doc_tokens / self.avg_doc_length()); (K1 + 1.0) * freq / (freq + doc_norm) } + + fn doc_weight_upper_bound(&self) -> Option { + Some(K1 + 1.0) + } + + fn doc_weight_cache_key(&self) -> Option { + Some(u64::from(self.avg_doc_length().to_bits())) + } } pub struct IndexBM25Scorer<'a> { @@ -146,6 +189,14 @@ impl Scorer for IndexBM25Scorer<'_> { let doc_norm = K1 * (1.0 - B + B * doc_tokens / self.avg_doc_length); (K1 + 1.0) * freq / (freq + doc_norm) } + + fn doc_weight_upper_bound(&self) -> Option { + Some(K1 + 1.0) + } + + fn doc_weight_cache_key(&self) -> Option { + Some(u64::from(self.avg_doc_length.to_bits())) + } } #[inline] diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index 1785a23fedd..429ad965419 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use lance_core::{Error, Result}; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Deserializer, Serialize}; use std::{env, path::PathBuf}; #[cfg(feature = "tokenizer-jieba")] @@ -22,6 +22,10 @@ use crate::pbold; use crate::scalar::inverted::tokenizer::document_tokenizer::{ JsonTokenizer, LanceTokenizer, TextTokenizer, }; +use crate::scalar::inverted::{ + InvertedListFormatVersion, default_fts_format_version_for_block_size, + resolve_fts_format_version, validate_format_version_block_size, +}; pub use lance_tokenizer::Language; use lance_tokenizer::{ AsciiFoldingFilter, IcuTokenizer, LowerCaser, NgramTokenizer, RawTokenizer, RemoveLongFilter, @@ -29,6 +33,16 @@ use lance_tokenizer::{ WhitespaceTokenizer, }; +/// Posting block size for indexes whose metadata predates configurable block sizes. +/// +/// This must remain 128 so that legacy on-disk data is decoded correctly. +pub const LEGACY_BLOCK_SIZE: usize = 128; +/// Default posting block size for newly created indexes when none is configured. +/// +/// This intentionally matches [`LEGACY_BLOCK_SIZE`] today but may evolve independently. +pub const DEFAULT_BLOCK_SIZE: usize = 128; +pub const VALID_BLOCK_SIZES: [usize; 2] = [128, 256]; + /// Tokenizer configs #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct InvertedIndexParams { @@ -98,6 +112,17 @@ pub struct InvertedIndexParams { #[serde(default)] pub(crate) prefix_only: bool, + /// Number of documents in each compressed posting block. + /// + /// Missing serialized values come from indexes written before this + /// parameter existed and must read as 128 for backwards compatibility. New + /// indexes currently default to 128. + #[serde( + default = "legacy_block_size", + deserialize_with = "deserialize_block_size" + )] + pub(crate) block_size: usize, + /// Total memory limit in MiB for the build stage. /// /// This is split evenly across FTS workers at build time. By default Lance @@ -120,6 +145,21 @@ pub struct InvertedIndexParams { /// The effective worker count is clamped to `[1, num_cpus - 2]`. #[serde(rename = "num_workers", skip_serializing, default)] pub(crate) num_workers: Option, + + /// On-disk FTS format version to write when creating a new index. + /// + /// This is a build-time only parameter and is not persisted with the index. + /// If unset, Lance writes v2 for `block_size = 128` and v3 for + /// `block_size = 256`. + /// `format_version = 3` is experimental and is only valid with + /// `block_size = 256`. + #[serde( + rename = "format_version", + skip_serializing, + default, + deserialize_with = "deserialize_format_version" + )] + pub(crate) format_version: Option, } impl TryFrom<&InvertedIndexParams> for pbold::InvertedIndexDetails { @@ -138,6 +178,7 @@ impl TryFrom<&InvertedIndexParams> for pbold::InvertedIndexDetails { min_ngram_length: params.min_ngram_length, max_ngram_length: params.max_ngram_length, prefix_only: params.prefix_only, + block_size: Some(params.block_size as u32), }) } } @@ -165,8 +206,13 @@ impl TryFrom<&pbold::InvertedIndexDetails> for InvertedIndexParams { min_ngram_length: details.min_ngram_length, max_ngram_length: details.max_ngram_length, prefix_only: details.prefix_only, + block_size: match details.block_size { + Some(block_size) => validate_block_size(block_size as usize)?, + None => LEGACY_BLOCK_SIZE, + }, memory_limit_mb: defaults.memory_limit_mb, num_workers: defaults.num_workers, + format_version: defaults.format_version, }) } } @@ -183,6 +229,61 @@ fn default_max_ngram_length() -> u32 { 3 } +fn legacy_block_size() -> usize { + LEGACY_BLOCK_SIZE +} + +fn invalid_block_size_message(block_size: usize) -> String { + format!("FTS inverted index block_size must be one of 128 or 256, got {block_size}") +} + +pub fn validate_block_size(block_size: usize) -> Result { + if VALID_BLOCK_SIZES.contains(&block_size) { + Ok(block_size) + } else { + Err(Error::invalid_input(invalid_block_size_message(block_size))) + } +} + +fn deserialize_block_size<'de, D>(deserializer: D) -> std::result::Result +where + D: Deserializer<'de>, +{ + let block_size = usize::deserialize(deserializer)?; + validate_block_size(block_size).map_err(serde::de::Error::custom) +} + +fn deserialize_format_version<'de, D>( + deserializer: D, +) -> std::result::Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let value = Option::::deserialize(deserializer)?; + let Some(value) = value else { + return Ok(None); + }; + match value { + serde_json::Value::Null => Ok(None), + serde_json::Value::String(value) => resolve_fts_format_version(Some(&value)) + .map(Some) + .map_err(serde::de::Error::custom), + serde_json::Value::Number(value) => { + let Some(format_version) = value.as_u64() else { + return Err(serde::de::Error::custom(format!( + "FTS format_version must be 1, 2, or 3, got {value}" + ))); + }; + resolve_fts_format_version(Some(&format_version.to_string())) + .map(Some) + .map_err(serde::de::Error::custom) + } + other => Err(serde::de::Error::custom(format!( + "FTS format_version must be 1, 2, or 3, got {other}" + ))), + } +} + impl Default for InvertedIndexParams { fn default() -> Self { Self::new("simple".to_owned(), Language::English) @@ -220,8 +321,10 @@ impl InvertedIndexParams { min_ngram_length: default_min_ngram_length(), max_ngram_length: default_max_ngram_length(), prefix_only: false, + block_size: DEFAULT_BLOCK_SIZE, memory_limit_mb: None, num_workers: None, + format_version: None, } } @@ -310,6 +413,25 @@ impl InvertedIndexParams { self } + /// Set the compressed posting block size. + /// + /// Supported values are 128 and 256. Larger values reduce block-max metadata + /// and WAND skip granularity; smaller values preserve the legacy layout. + /// + /// `block_size = 256` is experimental and may introduce breaking changes. + /// Use `128` when stable compatibility with the legacy posting layout is required. + pub fn block_size(mut self, block_size: usize) -> Result { + self.block_size = validate_block_size(block_size)?; + Ok(self) + } + + /// Get the compressed posting block size. + /// + /// `256` is experimental and may introduce breaking changes. + pub fn posting_block_size(&self) -> usize { + self.block_size + } + pub fn memory_limit_mb(mut self, memory_limit_mb: u64) -> Self { self.memory_limit_mb = Some(memory_limit_mb); self @@ -324,6 +446,33 @@ impl InvertedIndexParams { self } + /// Set the on-disk FTS format version to use when creating a new index. + /// + /// If unset, Lance writes v2 for `block_size = 128` and v3 for + /// `block_size = 256`. Existing indexes keep their own on-disk format + /// during update and optimize operations. + /// `format_version = 3` is experimental and is only valid with + /// `block_size = 256`. + pub fn format_version(mut self, format_version: InvertedListFormatVersion) -> Self { + self.format_version = Some(format_version); + self + } + + /// Resolve the requested FTS format version, falling back to the default for + /// the configured block size. + pub fn resolved_format_version(&self) -> InvertedListFormatVersion { + self.format_version.unwrap_or_else(|| { + default_fts_format_version_for_block_size(self.block_size) + .expect("InvertedIndexParams block_size must be validated before use") + }) + } + + /// Validate that the requested FTS format version can safely encode the + /// configured posting block size. + pub fn validate_format_version(&self) -> Result<()> { + validate_format_version_block_size(self.resolved_format_version(), self.block_size) + } + /// Serialize params for the build/training path, including build-only fields. pub fn to_training_json(&self) -> serde_json::Result { let mut value = serde_json::to_value(self)?; @@ -342,9 +491,34 @@ impl InvertedIndexParams { serde_json::Value::from(num_workers), ); } + if let Some(format_version) = self.format_version { + object.insert( + "format_version".to_string(), + serde_json::Value::from(format_version.index_version()), + ); + } Ok(value) } + /// Deserialize params for new index training, using current creation defaults + /// for omitted fields. + pub(crate) fn from_training_json(params: &str) -> Result { + let supplied = serde_json::from_str::(params)?; + let mut value = serde_json::to_value(Self::default())?; + + let supplied = supplied.as_object().ok_or_else(|| { + Error::invalid_input("FTS inverted index params must be a JSON object".to_string()) + })?; + let object = value + .as_object_mut() + .expect("inverted index params should serialize to a JSON object"); + object.extend(supplied.clone()); + + let params: Self = serde_json::from_value(value)?; + params.validate_format_version()?; + Ok(params) + } + pub fn build(&self) -> Result> { let mut builder = self.build_base_tokenizer()?; if let Some(max_token_length) = self.max_token_length { @@ -450,18 +624,22 @@ pub fn language_model_home() -> Option { #[cfg(test)] mod tests { - use super::InvertedIndexParams; - use lance_tokenizer::TokenStream; + use crate::pbold; + + use super::{InvertedIndexParams, InvertedListFormatVersion}; + use lance_tokenizer::{Language, TokenStream}; use rstest::rstest; #[test] fn test_build_only_fields_are_not_serialized() { let params = InvertedIndexParams::default() .memory_limit_mb(4096) - .num_workers(7); + .num_workers(7) + .format_version(InvertedListFormatVersion::V1); let json = serde_json::to_value(¶ms).unwrap(); assert!(json.get("memory_limit").is_none()); assert!(json.get("num_workers").is_none()); + assert!(json.get("format_version").is_none()); } #[test] @@ -483,23 +661,162 @@ mod tests { let obj = json.as_object_mut().unwrap(); obj.insert("memory_limit".to_string(), serde_json::Value::from(4096)); obj.insert("num_workers".to_string(), serde_json::Value::from(3)); + obj.insert("format_version".to_string(), serde_json::Value::from("v1")); let params: InvertedIndexParams = serde_json::from_value(json).unwrap(); assert_eq!(params.memory_limit_mb, Some(4096)); assert_eq!(params.num_workers, Some(3)); + assert_eq!(params.format_version, Some(InvertedListFormatVersion::V1)); } #[test] fn test_training_json_serializes_build_only_fields() { let params = InvertedIndexParams::default() .memory_limit_mb(4096) - .num_workers(3); + .num_workers(3) + .format_version(InvertedListFormatVersion::V1); let json = params.to_training_json().unwrap(); assert_eq!( json.get("memory_limit"), Some(&serde_json::Value::from(4096)) ); assert_eq!(json.get("num_workers"), Some(&serde_json::Value::from(3))); + assert_eq!( + json.get("format_version"), + Some(&serde_json::Value::from(1)) + ); + } + + #[test] + fn test_default_format_version_resolves_to_v2() { + assert_eq!( + InvertedIndexParams::default().resolved_format_version(), + InvertedListFormatVersion::V2 + ); + } + + #[test] + fn test_block_size_256_defaults_to_v3() { + assert_eq!( + InvertedIndexParams::default() + .block_size(256) + .unwrap() + .resolved_format_version(), + InvertedListFormatVersion::V3 + ); + } + + #[test] + fn test_format_version_must_match_block_size() { + InvertedIndexParams::default() + .format_version(InvertedListFormatVersion::V2) + .validate_format_version() + .unwrap(); + InvertedIndexParams::default() + .block_size(256) + .unwrap() + .validate_format_version() + .unwrap(); + + let err = InvertedIndexParams::default() + .block_size(256) + .unwrap() + .format_version(InvertedListFormatVersion::V2) + .validate_format_version() + .unwrap_err(); + assert!(err.to_string().contains("block_size=256")); + + let err = InvertedIndexParams::default() + .format_version(InvertedListFormatVersion::V3) + .validate_format_version() + .unwrap_err(); + assert!(err.to_string().contains("format_version=3")); + } + + #[test] + fn test_training_json_rejects_incompatible_format_version_and_block_size() { + let err = + InvertedIndexParams::from_training_json(r#"{"block_size": 256, "format_version": 2}"#) + .unwrap_err(); + assert!(err.to_string().contains("block_size=256")); + } + + #[test] + fn test_training_json_invalid_numeric_format_version_includes_value() { + let err = InvertedIndexParams::from_training_json(r#"{"format_version": -1}"#).unwrap_err(); + assert!(matches!(&err, lance_core::Error::Arrow { .. })); + assert!(err.to_string().contains("got -1")); + } + + #[test] + fn test_block_size_default_serializes() { + let params = InvertedIndexParams::default(); + assert_eq!(params.block_size, 128); + let json = serde_json::to_value(¶ms).unwrap(); + assert_eq!(json.get("block_size"), Some(&serde_json::Value::from(128))); + } + + #[test] + fn test_block_size_missing_metadata_falls_back_to_128() { + let mut json = serde_json::to_value(InvertedIndexParams::default()).unwrap(); + json.as_object_mut().unwrap().remove("block_size"); + + let params: InvertedIndexParams = serde_json::from_value(json).unwrap(); + assert_eq!(params.block_size, 128); + } + + #[test] + fn test_block_size_details_conversion() { + let params = InvertedIndexParams::default().block_size(256).unwrap(); + let details = pbold::InvertedIndexDetails::try_from(¶ms).unwrap(); + assert_eq!(details.block_size, Some(256)); + + let old_details = pbold::InvertedIndexDetails { + base_tokenizer: Some("simple".to_string()), + language: serde_json::to_string(&Language::English).unwrap(), + with_position: false, + max_token_length: Some(40), + lower_case: true, + stem: true, + remove_stop_words: true, + ascii_folding: true, + min_ngram_length: 3, + max_ngram_length: 3, + prefix_only: false, + block_size: None, + }; + let params = InvertedIndexParams::try_from(&old_details).unwrap(); + assert_eq!(params.block_size, 128); + } + + #[rstest] + #[case::block_size_128(128)] + #[case::block_size_256(256)] + fn test_block_size_accepts_supported_values(#[case] block_size: usize) { + let params = InvertedIndexParams::default() + .block_size(block_size) + .unwrap(); + assert_eq!(params.block_size, block_size); + + let roundtrip: InvertedIndexParams = + serde_json::from_value(serde_json::to_value(¶ms).unwrap()).unwrap(); + assert_eq!(roundtrip.block_size, block_size); + } + + #[test] + fn test_block_size_rejects_invalid_values() { + let err = InvertedIndexParams::default().block_size(129).unwrap_err(); + assert!(err.to_string().contains("block_size")); + + let err = InvertedIndexParams::default().block_size(512).unwrap_err(); + assert!(err.to_string().contains("128 or 256")); + + let mut json = serde_json::to_value(InvertedIndexParams::default()).unwrap(); + json.as_object_mut() + .unwrap() + .insert("block_size".to_string(), serde_json::Value::from(1024)); + let err = serde_json::from_value::(json).unwrap_err(); + assert!(err.to_string().contains("128 or 256")); } #[test] @@ -582,4 +899,75 @@ mod tests { } assert_eq!(tokens, vec!["lance".to_string(), "data".to_string()]); } + + // Common English pronouns/function words such as `you`/`my`/`your`/`we` + // must be removed by the ICU `all()` stop-word path. These are among the + // highest-frequency tokens, so leaking them builds pathologically large + // single-term posting lists (and previously overflowed the u32 posting-list + // size counter, panicking the whole index build). The leak is independent + // of stemming, so we assert it for both stem=false and stem=true. + #[rstest] + #[case::icu_no_stem("icu", false)] + #[case::icu_stem("icu", true)] + #[case::icu_split_no_stem("icu/split", false)] + #[case::icu_split_stem("icu/split", true)] + // `simple` is the recommended tokenizer for monolingual English corpora and + // uses StopWordFilter::new(English) rather than the ICU all() path, so it + // must be covered too. + #[case::simple_no_stem("simple", false)] + #[case::simple_stem("simple", true)] + fn test_icu_common_english_stop_words_do_not_leak( + #[case] base_tokenizer: &str, + #[case] stem: bool, + ) { + let mut tokenizer = InvertedIndexParams::default() + .base_tokenizer(base_tokenizer.to_string()) + .stem(stem) + .remove_stop_words(true) + .build() + .unwrap(); + let mut stream = tokenizer.token_stream_for_search("you my your we lance data"); + let tokens: Vec = std::iter::from_fn(|| stream.next().map(|t| t.text.clone())) + .filter(|t| matches!(t.as_str(), "you" | "my" | "your" | "we")) + .collect(); + assert!( + tokens.is_empty(), + "common English stop words leaked through the icu pipeline (stem={stem}): {tokens:?}" + ); + } + + // Common Chinese function words/particles (了 是 在 的 和 有 我) are the + // highest-frequency Chinese tokens; like the English pronouns they must be + // removed by the ICU `all()` stop-word path so they don't build huge + // posting lists. Real content words (英语 = "English", 数据 = "data") must + // survive. ICU dictionary segmentation splits the input into words, so this + // exercises the CJK stop-word path end to end. + #[rstest] + #[case::icu("icu")] + #[case::icu_split("icu/split")] + fn test_icu_common_chinese_stop_words_do_not_leak(#[case] base_tokenizer: &str) { + let mut tokenizer = InvertedIndexParams::default() + .base_tokenizer(base_tokenizer.to_string()) + .stem(true) + .remove_stop_words(true) + .build() + .unwrap(); + let mut stream = tokenizer.token_stream_for_search("我 在 有 了 是 的 和 英语 数据"); + let tokens: Vec = + std::iter::from_fn(|| stream.next().map(|t| t.text.clone())).collect(); + let stop = ["我", "在", "有", "了", "是", "的", "和"]; + let leaked: Vec<&String> = tokens + .iter() + .filter(|t| stop.contains(&t.as_str())) + .collect(); + assert!( + leaked.is_empty(), + "common Chinese stop words leaked through the icu pipeline: {leaked:?} (all tokens: {tokens:?})" + ); + // The real content words must still be indexed. + assert!( + tokens.iter().any(|t| t == "英语"), + "content word 英语 was dropped: {tokens:?}" + ); + } } diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 0fc8b95cddb..f0dc416b1b4 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -22,6 +22,7 @@ use crate::metrics::MetricsCollector; use super::{ CompressedPositionStorage, + impact::{IMPACT_LEVEL1_BLOCKS, ImpactScoreCache, ImpactSkipData}, query::Operator, scorer::{K1, idf}, }; @@ -29,8 +30,8 @@ use super::{ CompressedPostingList, DocSet, PostingList, RawDocInfo, builder::ScoredDoc, encoding::{ - decode_position_stream_block, decompress_positions, decompress_posting_block, - decompress_posting_remainder, + MAX_POSTING_BLOCK_SIZE, decode_position_stream_block, decompress_positions, + decompress_posting_block, decompress_posting_remainder, }, query::FtsSearchParams, scorer::Scorer, @@ -38,6 +39,7 @@ use super::{ use super::{DocInfo, builder::BLOCK_SIZE}; const TERMINATED_DOC_ID: u64 = u64::MAX; +const LINEAR_BLOCK_SKIP_LIMIT: usize = 8; pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { std::env::var("LANCE_FLAT_SEARCH_PERCENT_THRESHOLD") .unwrap_or_else(|_| "10".to_string()) @@ -45,6 +47,29 @@ pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { .unwrap_or(10) }); +#[inline] +fn conservative_bm25_upper_bound(query_weight: f32) -> f32 { + if query_weight <= 0.0 { + 0.0 + } else { + query_weight * (K1 + 1.0) + } +} + +#[inline] +fn scorer_upper_bound(query_weight: f32, scorer: &S) -> f32 { + if query_weight.is_nan() { + return f32::INFINITY; + } + if query_weight <= 0.0 { + return 0.0; + } + match scorer.doc_weight_upper_bound() { + Some(bound) if bound.is_finite() && bound >= 0.0 => query_weight * bound, + _ => f32::INFINITY, + } +} + pub struct PostingIterator { token: String, token_id: u32, @@ -55,6 +80,7 @@ pub struct PostingIterator { index: usize, // the index of current block, this can be changed by `next() and shallow_next()` block_idx: usize, + current_doc: Option, approximate_upper_bound: f32, // for compressed posting list @@ -66,24 +92,26 @@ struct CompressedState { block_idx: usize, doc_ids: Vec, freqs: Vec, - buffer: Box<[u32; BLOCK_SIZE]>, + buffer: Box<[u32; MAX_POSTING_BLOCK_SIZE]>, position_block_idx: Option, position_values: Vec, position_offsets: Vec, block_max_window: BlockMaxWindow, + current_block_max_score: Option<(usize, f32)>, } impl CompressedState { - fn new() -> Self { + fn new(block_size: usize) -> Self { Self { block_idx: 0, - doc_ids: Vec::with_capacity(BLOCK_SIZE), - freqs: Vec::with_capacity(BLOCK_SIZE), - buffer: Box::new([0; BLOCK_SIZE]), + doc_ids: Vec::with_capacity(block_size), + freqs: Vec::with_capacity(block_size), + buffer: Box::new([0; MAX_POSTING_BLOCK_SIZE]), position_block_idx: None, position_values: Vec::new(), position_offsets: Vec::new(), block_max_window: BlockMaxWindow::new(), + current_block_max_score: None, } } @@ -95,21 +123,29 @@ impl CompressedState { num_blocks: usize, length: u32, tail_codec: super::PostingTailCodec, + block_size: usize, ) { self.doc_ids.clear(); self.freqs.clear(); - let remainder = length as usize % BLOCK_SIZE; + let remainder = length as usize % block_size; if block_idx + 1 == num_blocks && remainder != 0 { decompress_posting_remainder( block, remainder, tail_codec, + block_size, &mut self.doc_ids, &mut self.freqs, ); } else { - decompress_posting_block(block, &mut self.buffer, &mut self.doc_ids, &mut self.freqs); + decompress_posting_block( + block, + &mut self.buffer[..], + &mut self.doc_ids, + &mut self.freqs, + block_size, + ); } self.block_idx = block_idx; self.position_block_idx = None; @@ -125,6 +161,7 @@ struct BlockMaxWindow { start_block_idx: usize, next_block_idx: usize, max_scores: VecDeque<(usize, f32)>, + impact_score_cache: ImpactScoreCache, } struct BlockMaxScore { @@ -138,6 +175,7 @@ impl BlockMaxWindow { start_block_idx: 0, next_block_idx: 0, max_scores: VecDeque::new(), + impact_score_cache: ImpactScoreCache::default(), } } @@ -147,12 +185,28 @@ impl BlockMaxWindow { self.max_scores.clear(); } - fn max_score_up_to( + fn max_score_up_to( &mut self, list: &CompressedPostingList, start_block_idx: usize, up_to: u64, + query_weight: f32, + scorer: &S, ) -> BlockMaxScore { + if let Some(impacts) = &list.impacts { + let score = impacts.max_score_up_to_cached( + start_block_idx, + up_to, + query_weight, + scorer, + &mut self.impact_score_cache, + ); + return BlockMaxScore { + score: score.score, + blocks_scanned: score.entries_scanned, + }; + } + if start_block_idx >= list.blocks.len() { self.reset(start_block_idx); return BlockMaxScore { @@ -177,6 +231,17 @@ impl BlockMaxWindow { }; } + // V3 postings score quantized doc lengths, which can be shorter than + // the exact lengths used to bake a legacy max score. Without impacts, + // use the score-independent BM25 ceiling instead of that stale bound. + if list.block_size == MAX_POSTING_BLOCK_SIZE { + self.reset(start_block_idx); + return BlockMaxScore { + score: scorer_upper_bound(query_weight, scorer), + blocks_scanned: 0, + }; + } + self.next_block_idx = self.next_block_idx.max(start_block_idx); let mut blocks_scanned = 0; while self.next_block_idx < list.blocks.len() @@ -256,6 +321,80 @@ impl Ord for PostingIterator { } impl PostingIterator { + fn block_idx_for_doc( + &self, + list: &CompressedPostingList, + mut block_idx: usize, + least_id: u32, + ) -> usize { + let mut linear_skips = 0; + while block_idx + 1 < list.blocks.len() && linear_skips < LINEAR_BLOCK_SKIP_LIMIT { + if list.block_least_doc_id(block_idx + 1) > least_id { + return block_idx; + } + block_idx += 1; + linear_skips += 1; + } + + if block_idx + 1 >= list.blocks.len() { + return block_idx; + } + + if let Some(impacts) = list.impacts.as_ref() + && let Some(block_idx) = + self.block_idx_for_doc_with_impacts(list, impacts, block_idx, least_id) + { + return block_idx; + } + + self.block_idx_for_doc_by_least_doc_id(list, block_idx, least_id, list.blocks.len()) + } + + fn block_idx_for_doc_with_impacts( + &self, + list: &CompressedPostingList, + impacts: &ImpactSkipData, + mut block_idx: usize, + least_id: u32, + ) -> Option { + while block_idx + 1 < list.blocks.len() { + let group_idx = (block_idx + 1) / IMPACT_LEVEL1_BLOCKS; + let group_end = ((group_idx + 1) * IMPACT_LEVEL1_BLOCKS).min(list.blocks.len()); + let group_doc_up_to = impacts.level1_doc_up_to(group_idx)?; + if group_doc_up_to < least_id { + block_idx = group_end - 1; + continue; + } + if group_doc_up_to == least_id { + return Some(group_end - 1); + } + return Some( + self.block_idx_for_doc_by_least_doc_id(list, block_idx, least_id, group_end), + ); + } + Some(block_idx) + } + + fn block_idx_for_doc_by_least_doc_id( + &self, + list: &CompressedPostingList, + block_idx: usize, + least_id: u32, + right: usize, + ) -> usize { + let mut left = block_idx + 1; + let mut right = right; + while left < right { + let mid = left + (right - left) / 2; + if list.block_least_doc_id(mid) <= least_id { + left = mid + 1; + } else { + right = mid; + } + } + left - 1 + } + #[inline] fn compressed_state_ptr(&self) -> *mut CompressedState { debug_assert!(self.compressed.is_some()); @@ -279,6 +418,7 @@ impl PostingIterator { list.blocks.len(), list.length, list.posting_tail_codec, + list.block_size, ); } compressed as *mut CompressedState @@ -303,14 +443,24 @@ impl PostingIterator { list: PostingList, num_doc: usize, ) -> Self { - let approximate_upper_bound = match list.max_score() { - Some(max_score) => max_score, - None => idf(list.len(), num_doc) * (K1 + 1.0), + let approximate_upper_bound = match &list { + PostingList::Compressed(posting) if posting.impacts.is_some() => f32::INFINITY, + PostingList::Compressed(posting) if posting.block_size == MAX_POSTING_BLOCK_SIZE => { + conservative_bm25_upper_bound(query_weight) + } + _ => match list.max_score() { + Some(max_score) => max_score, + None => idf(list.len(), num_doc) * (K1 + 1.0), + }, + }; + let compressed = match &list { + PostingList::Compressed(list) => { + Some(UnsafeCell::new(CompressedState::new(list.block_size))) + } + PostingList::Plain(_) => None, }; - let is_compressed = matches!(list, PostingList::Compressed(_)); - - Self { + let mut posting = Self { token, token_id, position, @@ -318,9 +468,12 @@ impl PostingIterator { list, index: 0, block_idx: 0, + current_doc: None, approximate_upper_bound, - compressed: is_compressed.then(|| UnsafeCell::new(CompressedState::new())), - } + compressed, + }; + posting.refresh_current_doc(); + posting } #[inline] @@ -338,8 +491,37 @@ impl PostingIterator { self.approximate_upper_bound } + /// Tightest known list-wide score bound. Impact lists answer from the + /// baked doc-weight slab (the data-driven equivalent of the max_score the + /// non-impact format bakes at build time); everything else falls back to + /// `approximate_upper_bound`. A finite, tight global bound is what lets + /// lagging iterators park in the WAND tail instead of being force-advanced + /// on every candidate. #[inline] - fn score(&self, scorer: &S, freq: u32, doc_length: u32) -> f32 { + fn global_upper_bound(&self, scorer: &S) -> f32 { + if self.query_weight <= 0.0 { + return 0.0; + } + if let PostingList::Compressed(ref list) = self.list + && let Some(impacts) = list.impacts.as_ref() + { + let compressed = unsafe { &mut *self.compressed_state_ptr() }; + return self.query_weight + * impacts.global_max_doc_weight_cached( + scorer, + &mut compressed.block_max_window.impact_score_cache, + ); + } + if let PostingList::Compressed(ref list) = self.list + && list.block_size == MAX_POSTING_BLOCK_SIZE + { + return scorer_upper_bound(self.query_weight, scorer); + } + self.approximate_upper_bound + } + + #[inline] + fn score(&self, scorer: &S, freq: u32, doc_length: u32) -> f32 { self.query_weight * scorer.doc_weight(freq, doc_length) } @@ -355,14 +537,19 @@ impl PostingIterator { #[inline] fn doc(&self) -> Option { + self.current_doc + } + + fn refresh_current_doc(&mut self) { if self.empty() { - return None; + self.current_doc = None; + return; } - match self.list { + let current_doc = match self.list { PostingList::Compressed(ref list) => { - let block_idx = self.index / BLOCK_SIZE; - let block_offset = self.index % BLOCK_SIZE; + let block_idx = self.index >> list.block_shift(); + let block_offset = self.index & list.block_mask(); let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; // Read from the decompressed block @@ -372,7 +559,8 @@ impl PostingIterator { Some(doc) } PostingList::Plain(ref list) => Some(DocInfo::Located(list.doc(self.index))), - } + }; + self.current_doc = current_doc; } fn position_cursor(&self) -> Option> { @@ -400,8 +588,8 @@ impl PostingIterator { )) } CompressedPositionStorage::SharedStream(stream) => { - let block_idx = self.index / BLOCK_SIZE; - let block_offset = self.index % BLOCK_SIZE; + let block_idx = self.index >> list.block_shift(); + let block_offset = self.index & list.block_mask(); let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; if compressed.position_block_idx != Some(block_idx) { @@ -441,36 +629,41 @@ impl PostingIterator { PostingList::Compressed(ref list) => { debug_assert!(least_id <= u32::MAX as u64); let least_id = least_id as u32; - let mut block_idx = self.index / BLOCK_SIZE; - while block_idx + 1 < list.blocks.len() - && list.block_least_doc_id(block_idx + 1) <= least_id - { - block_idx += 1; - } - self.index = self.index.max(block_idx * BLOCK_SIZE); + let shift = list.block_shift(); + let block_idx = self.block_idx_for_doc(list, self.index >> shift, least_id); + self.index = self.index.max(block_idx << shift); let length = list.length as usize; while self.index < length { - let block_idx = self.index / BLOCK_SIZE; - let block_offset = self.index % BLOCK_SIZE; + let block_idx = self.index >> shift; + let block_offset = self.index & list.block_mask(); let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; let in_block = &compressed.doc_ids[block_offset..]; let offset_in_block = in_block.partition_point(|&doc_id| doc_id < least_id); let new_offset = block_offset + offset_in_block; if new_offset < compressed.doc_ids.len() { - self.index = block_idx * BLOCK_SIZE + new_offset; - break; + self.index = (block_idx << shift) + new_offset; + self.block_idx = block_idx; + self.current_doc = Some(DocInfo::Raw(RawDocInfo { + doc_id: compressed.doc_ids[new_offset], + frequency: compressed.freqs[new_offset], + })); + return; } if block_idx + 1 >= list.blocks.len() { self.index = length; + self.block_idx = self.index >> shift; + self.current_doc = None; break; } - self.index = (block_idx + 1) * BLOCK_SIZE; + self.index = (block_idx + 1) << shift; } - self.block_idx = self.index / BLOCK_SIZE; + self.block_idx = self.index >> shift; + self.current_doc = None; } PostingList::Plain(ref list) => { self.index += list.row_ids[self.index..].partition_point(|&id| id < least_id); + self.current_doc = (!self.empty()).then(|| DocInfo::Located(list.doc(self.index))); } } } @@ -480,11 +673,7 @@ impl PostingIterator { PostingList::Compressed(ref list) => { debug_assert!(least_id <= u32::MAX as u64); let least_id = least_id as u32; - while self.block_idx + 1 < list.blocks.len() - && list.block_least_doc_id(self.block_idx + 1) <= least_id - { - self.block_idx += 1; - } + self.block_idx = self.block_idx_for_doc(list, self.block_idx, least_id); } PostingList::Plain(_) => { // we don't have block max score for legacy index, @@ -494,21 +683,51 @@ impl PostingIterator { } #[inline] - fn block_max_score(&self) -> f32 { + fn block_max_score(&self, scorer: &S) -> f32 { match self.list { - PostingList::Compressed(ref list) => list.block_max_score(self.block_idx), + PostingList::Compressed(ref list) => { + if let Some(impacts) = list.impacts.as_ref() { + let compressed = unsafe { &mut *self.compressed_state_ptr() }; + if let Some((block_idx, score)) = compressed.current_block_max_score + && block_idx == self.block_idx + { + return score; + } + + let score = impacts.level0_score_cached( + self.block_idx, + self.query_weight, + scorer, + &mut compressed.block_max_window.impact_score_cache, + ); + compressed.current_block_max_score = Some((self.block_idx, score)); + return score; + } + if list.block_size == MAX_POSTING_BLOCK_SIZE { + return scorer_upper_bound(self.query_weight, scorer); + } + list.block_max_score(self.block_idx) + } PostingList::Plain(_) => self.approximate_upper_bound, } } #[inline] - fn block_max_score_up_to_with_stats(&mut self, up_to: u64) -> BlockMaxScore { + fn block_max_score_up_to_with_stats( + &mut self, + up_to: u64, + scorer: &S, + ) -> BlockMaxScore { match self.list { PostingList::Compressed(ref list) => { let compressed = unsafe { &mut *self.compressed_state_ptr() }; - compressed - .block_max_window - .max_score_up_to(list, self.block_idx, up_to) + compressed.block_max_window.max_score_up_to( + list, + self.block_idx, + up_to, + self.query_weight, + scorer, + ) } PostingList::Plain(_) => BlockMaxScore { score: self.approximate_upper_bound, @@ -903,7 +1122,7 @@ impl<'a, S: Scorer> Wand<'a, S> { } let doc_length = match &doc { - DocInfo::Raw(doc) => self.docs.num_tokens(doc.doc_id), + DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), }; @@ -1077,7 +1296,7 @@ impl<'a, S: Scorer> Wand<'a, S> { // score the doc let doc_length = match is_compressed { - true => self.docs.num_tokens(doc_id as u32), + true => self.docs.scoring_num_tokens(doc_id as u32), false => self.docs.num_tokens_by_row_id(row_id), }; if self.operator == Operator::Or && !self.refine_or_candidate(doc_id, doc_length) { @@ -1174,7 +1393,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let remaining_upper_bound = remaining .iter() - .map(|posting| posting.block_max_score()) + .map(|posting| posting.block_max_score(&self.scorer)) .sum::(); first.score(&self.scorer, doc.frequency(), doc_length) + remaining_upper_bound <= self.threshold @@ -1227,7 +1446,7 @@ impl<'a, S: Scorer> Wand<'a, S> { continue; }; let doc_length = match &first_doc { - DocInfo::Raw(doc) => self.docs.num_tokens(doc.doc_id), + DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), }; let mut lead_score = 0.0; @@ -1309,7 +1528,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let lead_doc = self.lead.first().and_then(|posting| posting.doc())?; let doc_length = match &lead_doc { - DocInfo::Raw(doc) => self.docs.num_tokens(doc.doc_id), + DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), }; if self.and_candidate_cannot_beat_threshold(doc_length) { @@ -1361,7 +1580,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let narrow_max_score = self .lead .iter() - .map(|posting| posting.block_max_score()) + .map(|posting| posting.block_max_score(&self.scorer)) .sum::(); if narrow_max_score >= self.threshold { @@ -1384,7 +1603,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut wide_max_score = 0.0; let mut range_blocks_scanned = 0; for posting in &mut self.lead { - let block_max = posting.block_max_score_up_to_with_stats(lead_up_to); + let block_max = posting.block_max_score_up_to_with_stats(lead_up_to, &self.scorer); wide_max_score += block_max.score; range_blocks_scanned += block_max.blocks_scanned; } @@ -1446,9 +1665,19 @@ impl<'a, S: Scorer> Wand<'a, S> { // Move all head iterators that are already known to be behind `target` // into `tail`, possibly overflowing low-value entries back into `head`. fn move_head_before_target_to_tail(&mut self, target: u64) { + if self.threshold <= 0.0 { + while matches!(self.head_doc(), Some(doc_id) if doc_id < target) { + if let Some(mut posting) = self.head.pop().map(|posting| posting.posting) { + posting.next(target); + self.push_head(posting); + } + } + return; + } + while matches!(self.head_doc(), Some(doc_id) if doc_id < target) { if let Some(posting) = self.head.pop() { - let upper_bound = posting.posting.approximate_upper_bound(); + let upper_bound = posting.posting.global_upper_bound(&self.scorer); if let Some(mut evicted) = self.insert_tail_with_overflow(posting.posting, upper_bound) { @@ -1467,12 +1696,12 @@ impl<'a, S: Scorer> Wand<'a, S> { let lead: f32 = self .lead .iter() - .map(|posting| posting.block_max_score()) + .map(|posting| posting.block_max_score(&self.scorer)) .sum(); let head: f32 = self .head .iter() - .map(|posting| posting.posting.block_max_score()) + .map(|posting| posting.posting.block_max_score(&self.scorer)) .sum(); lead + head + self.tail_max_score } @@ -1485,12 +1714,12 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut sum = self .lead .iter() - .map(|posting| posting.block_max_score()) + .map(|posting| posting.block_max_score(&self.scorer)) .sum::(); let mut possible_matches = self.lead.len(); for posting in &self.tail { if matches!(posting.posting.block_first_doc(), Some(block_doc) if block_doc <= target) { - sum += posting.posting.block_max_score(); + sum += posting.posting.block_max_score(&self.scorer); possible_matches += 1; } } @@ -1556,7 +1785,7 @@ impl<'a, S: Scorer> Wand<'a, S> { Some(block_doc) if block_doc <= target => { tail_posting .posting - .block_max_score_up_to_with_stats(up_to) + .block_max_score_up_to_with_stats(up_to, &self.scorer) .score } _ => 0.0, @@ -1693,9 +1922,9 @@ impl<'a, S: Scorer> Wand<'a, S> { && posting.is_compressed() && self.up_to.is_some_and(|up_to| target <= up_to) { - posting.block_max_score() + posting.block_max_score(&self.scorer) } else { - posting.approximate_upper_bound() + posting.global_upper_bound(&self.scorer) } } @@ -1730,6 +1959,14 @@ impl<'a, S: Scorer> Wand<'a, S> { // into lagging iterators. Entries that do not stay in `tail` are // advanced to `target` and returned to `head`. // pop() drains in place, keeping self.lead's capacity for reuse. + if self.threshold <= 0.0 { + while let Some(mut posting) = self.lead.pop() { + posting.next(target); + self.push_head(posting); + } + return; + } + while let Some(posting) = self.lead.pop() { let upper_bound = self.lead_to_tail_upper_bound(&posting, target); if let Some(mut evicted) = self.insert_tail_with_overflow(posting, upper_bound) { @@ -1977,13 +2214,17 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; + use super::super::impact::build_impact_skip_data; use super::*; - use crate::scalar::inverted::scorer::IndexBM25Scorer; + use crate::scalar::inverted::scorer::{IndexBM25Scorer, MemBM25Scorer}; use crate::{ metrics::{MetricsCollector, NoOpMetricsCollector}, scalar::inverted::{ - CompressedPostingList, PlainPostingList, PostingListBuilder, builder::PositionRecorder, - encoding::compress_posting_list, + CompressedPostingList, PlainPostingList, PostingListBuilder, + builder::PositionRecorder, + encoding::{ + compress_posting_list, compress_posting_list_with_tail_codec_and_block_size, + }, }, }; @@ -2146,6 +2387,8 @@ mod tests { max_score, doc_ids.len() as u32, crate::scalar::inverted::PostingTailCodec::VarintDelta, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, + None, None, )) } else { @@ -2158,6 +2401,75 @@ mod tests { } } + fn generate_impact_posting_list_with_freqs( + doc_ids: Vec, + freqs: Vec, + doc_lengths: Vec, + ) -> PostingList { + generate_impact_posting_list_with_freqs_and_block_size( + doc_ids, + freqs, + doc_lengths, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, + ) + } + + fn generate_impact_posting_list_with_freqs_and_block_size( + doc_ids: Vec, + freqs: Vec, + doc_lengths: Vec, + block_size: usize, + ) -> PostingList { + assert_eq!(doc_ids.len(), freqs.len()); + assert_eq!(doc_ids.len(), doc_lengths.len()); + let block_max_scores = vec![0.0; doc_ids.len().div_ceil(block_size)]; + let blocks = compress_posting_list_with_tail_codec_and_block_size( + doc_ids.len(), + doc_ids.iter(), + freqs.iter(), + block_max_scores.into_iter(), + crate::scalar::inverted::PostingTailCodec::VarintDelta, + block_size, + ) + .unwrap(); + let impact_blocks = doc_ids + .chunks(block_size) + .zip(freqs.chunks(block_size)) + .zip(doc_lengths.chunks(block_size)) + .map(|((doc_ids, freqs), doc_lengths)| { + doc_ids + .iter() + .copied() + .zip(freqs.iter().copied()) + .zip(doc_lengths.iter().copied()) + .map(|((doc_id, freq), doc_length)| (doc_id, freq, doc_length)) + .collect::>() + }) + .collect::>(); + let impacts = build_impact_skip_data(impact_blocks.as_slice()).unwrap(); + PostingList::Compressed(CompressedPostingList::new( + blocks, + 0.0, + doc_ids.len() as u32, + crate::scalar::inverted::PostingTailCodec::VarintDelta, + block_size, + None, + Some(impacts), + )) + } + + fn generate_contiguous_impact_posting_list_with_block_size( + total: usize, + block_size: usize, + ) -> PostingList { + generate_impact_posting_list_with_freqs_and_block_size( + (0..total as u32).collect(), + vec![1; total], + vec![1; total], + block_size, + ) + } + fn generate_posting_list_with_positions( doc_ids: Vec, positions_by_doc: Vec>, @@ -2652,6 +2964,56 @@ mod tests { assert_eq!(candidate.0.doc_id(), BLOCK_SIZE as u64); } + #[test] + fn test_non_positive_threshold_advances_without_impact_bound_scoring() { + let mut docs = DocSet::default(); + for doc_id in 0..3 { + docs.append(doc_id, 1); + } + let make_posting = || { + PostingIterator::with_query_weight( + String::from("term"), + 0, + 0, + 1.0, + generate_impact_posting_list_with_freqs( + vec![0, 1, 2], + vec![1, 1, 1], + vec![1, 1, 1], + ), + docs.len(), + ) + }; + let scored = Arc::new(AtomicUsize::new(0)); + + let mut wand = Wand::new( + Operator::Or, + std::iter::once(make_posting()), + &docs, + CountingScorer { + scored: scored.clone(), + }, + ); + wand.move_head_before_target_to_tail(1); + assert_eq!(wand.head_doc(), Some(1)); + assert!(wand.tail.is_empty()); + assert_eq!(scored.load(Ordering::Relaxed), 0); + + let mut wand = Wand::new( + Operator::Or, + std::iter::once(make_posting()), + &docs, + CountingScorer { + scored: scored.clone(), + }, + ); + wand.move_head_doc_to_lead(0); + wand.push_back_leads(1); + assert_eq!(wand.head_doc(), Some(1)); + assert!(wand.tail.is_empty()); + assert_eq!(scored.load(Ordering::Relaxed), 0); + } + #[test] fn test_or_plain_tail_does_not_advance_headless_window() { let mut docs = DocSet::default(); @@ -2982,7 +3344,7 @@ mod tests { posting.shallow_next(0); assert_eq!( posting - .block_max_score_up_to_with_stats((3 * BLOCK_SIZE - 1) as u64) + .block_max_score_up_to_with_stats((3 * BLOCK_SIZE - 1) as u64, &UnitScorer) .score, 4.0 ); @@ -2990,7 +3352,7 @@ mod tests { posting.shallow_next((2 * BLOCK_SIZE) as u64); assert_eq!( posting - .block_max_score_up_to_with_stats((4 * BLOCK_SIZE - 1) as u64) + .block_max_score_up_to_with_stats((4 * BLOCK_SIZE - 1) as u64, &UnitScorer) .score, 5.0 ); @@ -2998,12 +3360,147 @@ mod tests { posting.shallow_next((4 * BLOCK_SIZE) as u64); assert_eq!( posting - .block_max_score_up_to_with_stats((5 * BLOCK_SIZE - 1) as u64) + .block_max_score_up_to_with_stats((5 * BLOCK_SIZE - 1) as u64, &UnitScorer) .score, 3.0 ); } + #[test] + fn test_impact_level1_skip_keeps_boundary_equality_in_group() { + for block_size in [crate::scalar::inverted::LEGACY_BLOCK_SIZE, 256] { + let total = (IMPACT_LEVEL1_BLOCKS + 1) * block_size; + let mut posting = PostingIterator::new( + String::from("term"), + 0, + 0, + generate_contiguous_impact_posting_list_with_block_size(total, block_size), + total, + ); + let target = (IMPACT_LEVEL1_BLOCKS * block_size - 1) as u64; + + posting.shallow_next(target); + assert_eq!(posting.block_idx, IMPACT_LEVEL1_BLOCKS - 1); + + posting.next(target); + assert_eq!(posting.block_idx, IMPACT_LEVEL1_BLOCKS - 1); + assert_eq!(posting.doc().map(|doc| doc.doc_id()), Some(target)); + } + } + + #[test] + fn test_impact_level1_skip_handles_partial_final_group() { + for block_size in [crate::scalar::inverted::LEGACY_BLOCK_SIZE, 256] { + let total = (IMPACT_LEVEL1_BLOCKS + 3) * block_size + 17; + let mut posting = PostingIterator::new( + String::from("term"), + 0, + 0, + generate_contiguous_impact_posting_list_with_block_size(total, block_size), + total, + ); + let target = (total - 1) as u64; + let expected_block = total.div_ceil(block_size) - 1; + + posting.shallow_next(target); + assert_eq!(posting.block_idx, expected_block); + + posting.next(target); + assert_eq!(posting.block_idx, expected_block); + assert_eq!(posting.doc().map(|doc| doc.doc_id()), Some(target)); + } + } + + #[test] + fn test_impact_level1_skip_reaches_far_target_doc() { + for block_size in [crate::scalar::inverted::LEGACY_BLOCK_SIZE, 256] { + let total = (IMPACT_LEVEL1_BLOCKS * 3 + 5) * block_size; + let target_block = IMPACT_LEVEL1_BLOCKS * 2 + 2; + let target = (target_block * block_size + 17) as u64; + let mut posting = PostingIterator::new( + String::from("term"), + 0, + 0, + generate_contiguous_impact_posting_list_with_block_size(total, block_size), + total, + ); + + posting.shallow_next(target); + assert_eq!(posting.block_idx, target_block); + + posting.next(target); + assert_eq!(posting.block_idx, target_block); + assert_eq!(posting.doc().map(|doc| doc.doc_id()), Some(target)); + } + } + + #[test] + fn test_compressed_impact_block_max_score_memoizes_current_block() { + let total = 2 * BLOCK_SIZE as u32; + let doc_ids = (0..total).collect::>(); + let freqs = doc_ids + .iter() + .map(|doc_id| if *doc_id < BLOCK_SIZE as u32 { 1 } else { 2 }) + .collect::>(); + let doc_lengths = vec![1; total as usize]; + let posting_list = generate_impact_posting_list_with_freqs(doc_ids, freqs, doc_lengths); + let mut posting = + PostingIterator::new(String::from("term"), 0, 0, posting_list, total as usize); + let scored = Arc::new(AtomicUsize::new(0)); + let scorer = CountingScorer { + scored: scored.clone(), + }; + + let first_score = posting.block_max_score(&scorer); + assert_eq!(first_score, 1.0); + // Baking the query-local doc-weight bounds visits every frontier pair + // once (two level0 entries plus one level1 entry for this list). + let baked = scored.load(Ordering::Relaxed); + assert!(baked >= 2); + { + let compressed = unsafe { &mut *posting.compressed_state_ptr() }; + assert_eq!(compressed.current_block_max_score, Some((0, first_score))); + } + + let second_score = posting.block_max_score(&scorer); + assert_eq!(second_score, first_score); + assert_eq!( + scored.load(Ordering::Relaxed), + baked, + "repeated block max scores must not recompute doc weights" + ); + + posting.shallow_next(BLOCK_SIZE as u64); + let next_block_score = posting.block_max_score(&scorer); + assert_eq!(next_block_score, 2.0); + assert_eq!( + scored.load(Ordering::Relaxed), + baked, + "other blocks answer from the baked bounds without rescoring" + ); + } + + #[rstest] + #[case(0.0)] + #[case(-1.0)] + fn test_non_positive_query_weight_skips_global_impact_bound(#[case] query_weight: f32) { + let posting = PostingIterator::with_query_weight( + String::from("term"), + 0, + 0, + query_weight, + generate_impact_posting_list_with_freqs(vec![0], vec![1], vec![1]), + 1, + ); + let scored = Arc::new(AtomicUsize::new(0)); + let scorer = CountingScorer { + scored: scored.clone(), + }; + + assert_eq!(posting.global_upper_bound(&scorer), 0.0); + assert_eq!(scored.load(Ordering::Relaxed), 0); + } + #[test] fn test_and_candidate_prune_scores_first_term_before_full_score() { let total_docs = 2 * BLOCK_SIZE as u32 + 1; @@ -3340,13 +3837,95 @@ mod tests { let posting = PostingIterator::new(String::from("test"), 0, 0, posting_list, 1); - let actual = posting.block_max_score(); + let actual = posting.block_max_score(&UnitScorer); assert!( (actual - expected).abs() < 1e-6, "block max score should match stored value" ); } + #[test] + fn test_v3_without_impacts_uses_conservative_quantized_score_bound() { + let exact_doc_length = 300; + let quantized_doc_length = super::super::index::dequantize_doc_length( + super::super::index::quantize_doc_length(exact_doc_length), + ); + assert!(quantized_doc_length < exact_doc_length); + + let scorer = Arc::new(MemBM25Scorer::new(100, 1, Default::default())); + let stored_exact_score = scorer.doc_weight(1, exact_doc_length); + let quantized_score = scorer.doc_weight(1, quantized_doc_length); + assert!(quantized_score > stored_exact_score); + + let doc_ids = [0_u32]; + let frequencies = [1_u32]; + let blocks = compress_posting_list_with_tail_codec_and_block_size( + doc_ids.len(), + doc_ids.iter(), + frequencies.iter(), + std::iter::once(stored_exact_score), + crate::scalar::inverted::PostingTailCodec::VarintDelta, + MAX_POSTING_BLOCK_SIZE, + ) + .unwrap(); + let posting_list = PostingList::Compressed(CompressedPostingList::new( + blocks, + stored_exact_score, + doc_ids.len() as u32, + crate::scalar::inverted::PostingTailCodec::VarintDelta, + MAX_POSTING_BLOCK_SIZE, + None, + None, + )); + let mut posting = + PostingIterator::new(String::from("term"), 0, 0, posting_list, doc_ids.len()); + let expected_bound = K1 + 1.0; + + assert_eq!(posting.approximate_upper_bound(), expected_bound); + assert_eq!(posting.global_upper_bound(&scorer), expected_bound); + assert_eq!(posting.block_max_score(&scorer), expected_bound); + assert_eq!( + posting.block_max_score_up_to_with_stats(0, &scorer).score, + expected_bound + ); + assert!(expected_bound >= quantized_score); + } + + #[test] + fn test_v3_without_impacts_unknown_scorer_uses_infinite_bound() { + let doc_ids = [0_u32]; + let frequencies = [10_u32]; + let blocks = compress_posting_list_with_tail_codec_and_block_size( + doc_ids.len(), + doc_ids.iter(), + frequencies.iter(), + std::iter::once(10.0), + crate::scalar::inverted::PostingTailCodec::VarintDelta, + MAX_POSTING_BLOCK_SIZE, + ) + .unwrap(); + let posting_list = PostingList::Compressed(CompressedPostingList::new( + blocks, + 10.0, + doc_ids.len() as u32, + crate::scalar::inverted::PostingTailCodec::VarintDelta, + MAX_POSTING_BLOCK_SIZE, + None, + None, + )); + let mut posting = + PostingIterator::new(String::from("term"), 0, 0, posting_list, doc_ids.len()); + + assert!(posting.global_upper_bound(&UnitScorer).is_infinite()); + assert!(posting.block_max_score(&UnitScorer).is_infinite()); + assert!( + posting + .block_max_score_up_to_with_stats(0, &UnitScorer) + .score + .is_infinite() + ); + } + #[rstest] fn test_exact_phrase_with_repeated_terms(#[values(false, true)] is_compressed: bool) { let mut docs = DocSet::default(); diff --git a/rust/lance-index/src/scalar/json.rs b/rust/lance-index/src/scalar/json.rs index b20eb10ff20..09678d6787e 100644 --- a/rust/lance-index/src/scalar/json.rs +++ b/rust/lance-index/src/scalar/json.rs @@ -1,8 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::{ - collections::HashMap, ops::Bound, sync::{Arc, Mutex}, }; @@ -33,13 +33,15 @@ use lance_core::{Error, ROW_ID, Result, cache::LanceCache, error::LanceOptionExt use crate::{ Index, IndexType, - frag_reuse::FragReuseIndex, metrics::MetricsCollector, registry::IndexPluginRegistry, scalar::{ - AnyQuery, CreatedIndex, IndexStore, ScalarIndex, SearchResult, UpdateCriteria, + AnyQuery, CreatedIndex, IndexStore, RowIdRemapper, ScalarIndex, SearchResult, + UpdateCriteria, expression::{IndexedExpression, ScalarIndexExpr, ScalarIndexSearch, ScalarQueryParser}, - registry::{ScalarIndexPlugin, TrainingCriteria, TrainingRequest, VALUE_COLUMN_NAME}, + registry::{ + BasicTrainer, ScalarIndexPlugin, TrainingCriteria, TrainingRequest, VALUE_COLUMN_NAME, + }, }, }; @@ -114,7 +116,7 @@ impl ScalarIndex for JsonIndex { async fn remap( &self, - mapping: &HashMap>, + mapping: &RowAddrRemap, dest_store: &dyn IndexStore, ) -> Result { let target_created = self.target_index.remap(mapping, dest_store).await?; @@ -666,11 +668,7 @@ impl JsonIndexPlugin { } #[async_trait] -impl ScalarIndexPlugin for JsonIndexPlugin { - fn name(&self) -> &str { - "Json" - } - +impl BasicTrainer for JsonIndexPlugin { fn new_training_request( &self, params: &str, @@ -688,7 +686,12 @@ impl ScalarIndexPlugin for JsonIndexPlugin { let params = serde_json::from_str::(params)?; let registry = self.registry()?; let target_plugin = registry.get_plugin_by_name(¶ms.target_index_type)?; - let target_request = target_plugin.new_training_request( + let target_trainer = target_plugin.basic_trainer().ok_or_else(|| { + Error::invalid_input_source( + format!("The '{}' index type does not support basic training, please refer to the index's documentation for more details on how to create this index.", params.target_index_type).into(), + ) + })?; + let target_request = target_trainer.new_training_request( params.target_index_parameters.as_deref().unwrap_or("{}"), &Field::new("", target_type, true), )?; @@ -696,39 +699,6 @@ impl ScalarIndexPlugin for JsonIndexPlugin { Ok(Box::new(JsonTrainingRequest::new(params, target_request))) } - fn provides_exact_answer(&self) -> bool { - // TODO: Need to lookup target plugin via details to figure this out correctly - true - } - - fn attach_registry(&self, registry: Arc) { - let mut reg_ref = self.registry.lock().unwrap(); - *reg_ref = Some(registry); - } - - fn version(&self) -> u32 { - JSON_INDEX_VERSION - } - - fn new_query_parser( - &self, - index_name: String, - index_details: &prost_types::Any, - ) -> Option> { - // TODO: Allow return Result here - let registry = self.registry().unwrap(); - let json_details = - crate::pb::JsonIndexDetails::decode(index_details.value.as_slice()).unwrap(); - let target_details = json_details.target_details.as_ref().expect_ok().unwrap(); - let target_plugin = registry.get_plugin_by_details(target_details).unwrap(); - // TODO: Use something like ${index_name}_${path} for the index name? Don't have access to path here tho - let target_parser = target_plugin.new_query_parser(index_name, index_details)?; - Some(Box::new(JsonQueryParser::new( - json_details.path.clone(), - target_parser, - )) as Box) - } - async fn train_index( &self, data: SendableRecordBatchStream, @@ -755,7 +725,12 @@ impl ScalarIndexPlugin for JsonIndexPlugin { let target_plugin = registry.get_plugin_by_name(&request.parameters.target_index_type)?; // Create a new training request with the inferred type - let target_request = target_plugin.new_training_request( + let target_trainer = target_plugin.basic_trainer().ok_or_else(|| { + Error::invalid_input_source( + format!("The '{}' index type does not support basic training, please refer to the index's documentation for more details on how to create this index.", request.parameters.target_index_type).into(), + ) + })?; + let target_request = target_trainer.new_training_request( request .parameters .target_index_parameters @@ -764,7 +739,7 @@ impl ScalarIndexPlugin for JsonIndexPlugin { &Field::new("", inferred_type, true), )?; - let target_index = target_plugin + let target_index = target_trainer .train_index( converted_stream, index_store, @@ -784,12 +759,56 @@ impl ScalarIndexPlugin for JsonIndexPlugin { files: target_index.files, }) } +} + +#[async_trait] +impl ScalarIndexPlugin for JsonIndexPlugin { + fn basic_trainer(&self) -> Option<&dyn BasicTrainer> { + Some(self) + } + + fn name(&self) -> &str { + "Json" + } + + fn provides_exact_answer(&self) -> bool { + // TODO: Need to lookup target plugin via details to figure this out correctly + true + } + + fn attach_registry(&self, registry: Arc) { + let mut reg_ref = self.registry.lock().unwrap(); + *reg_ref = Some(registry); + } + + fn version(&self) -> u32 { + JSON_INDEX_VERSION + } + + fn new_query_parser( + &self, + index_name: String, + index_details: &prost_types::Any, + ) -> Option> { + // TODO: Allow return Result here + let registry = self.registry().unwrap(); + let json_details = + crate::pb::JsonIndexDetails::decode(index_details.value.as_slice()).unwrap(); + let target_details = json_details.target_details.as_ref().expect_ok().unwrap(); + let target_plugin = registry.get_plugin_by_details(target_details).unwrap(); + // TODO: Use something like ${index_name}_${path} for the index name? Don't have access to path here tho + let target_parser = target_plugin.new_query_parser(index_name, index_details)?; + Some(Box::new(JsonQueryParser::new( + json_details.path.clone(), + target_parser, + )) as Box) + } async fn load_index( &self, index_store: Arc, index_details: &prost_types::Any, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result> { let registry = self.registry().unwrap(); diff --git a/rust/lance-index/src/scalar/label_list.rs b/rust/lance-index/src/scalar/label_list.rs index 77232952419..b96a3ba0346 100644 --- a/rust/lance-index/src/scalar/label_list.rs +++ b/rust/lance-index/src/scalar/label_list.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::{ any::Any, collections::HashMap, @@ -28,18 +29,20 @@ use lance_select::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; use roaring::RoaringBitmap; use tracing::instrument; -use super::{AnyQuery, IndexFile, IndexStore, LabelListQuery, ScalarIndex, bitmap::BitmapIndex}; +use super::{ + AnyQuery, IndexFile, IndexStore, LabelListQuery, OldIndexDataFilter, ScalarIndex, + bitmap::BitmapIndex, +}; use super::{BuiltinIndexType, SargableQuery, ScalarIndexParams}; use super::{MetricsCollector, SearchResult}; -use crate::frag_reuse::FragReuseIndex; use crate::pbold; use crate::scalar::bitmap::{BitmapIndexPlugin, BitmapIndexState}; use crate::scalar::expression::{LabelListQueryParser, ScalarQueryParser}; use crate::scalar::registry::{ - DefaultTrainingRequest, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, - VALUE_COLUMN_NAME, + BasicTrainer, DefaultTrainingRequest, ScalarIndexLoad, ScalarIndexPlugin, TrainingCriteria, + TrainingOrdering, TrainingRequest, VALUE_COLUMN_NAME, single_flight_open, }; -use crate::scalar::{CreatedIndex, UpdateCriteria}; +use crate::scalar::{CreatedIndex, RowIdRemapper, UpdateCriteria}; use crate::{Index, IndexType}; pub const BITMAP_LOOKUP_NAME: &str = "bitmap_page_lookup.lance"; @@ -90,7 +93,7 @@ impl LabelListIndex { async fn load( store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, index_cache: &LanceCache, ) -> Result> { let values_index = @@ -210,7 +213,7 @@ impl ScalarIndex for LabelListIndex { /// Remap the row ids, creating a new remapped version of this index in `dest_store` async fn remap( &self, - mapping: &HashMap>, + mapping: &RowAddrRemap, dest_store: &dyn IndexStore, ) -> Result { let state = self.values_index.load_bitmap_index_state().await?; @@ -218,10 +221,7 @@ impl ScalarIndex for LabelListIndex { let remapped_nulls = RowAddrTreeMap::from_iter(self.list_nulls.row_addrs().unwrap().filter_map(|addr| { let addr_as_u64 = u64::from(addr); - mapping - .get(&addr_as_u64) - .copied() - .unwrap_or(Some(addr_as_u64)) + mapping.get(addr_as_u64).unwrap_or(Some(addr_as_u64)) })); let file = write_label_list_bitmap_index( remapped_state, @@ -440,7 +440,7 @@ fn unnest_chunks( async fn read_list_nulls( store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, ) -> Result { let reader = store.open_index_file(BITMAP_LOOKUP_NAME).await?; if let Some(buffer_idx_str) = reader.schema().metadata.get(LABEL_LIST_NULLS_METADATA_KEY) { @@ -486,6 +486,87 @@ async fn write_label_list_bitmap_index( .await } +/// Merge multiple LabelList index segments into a single index. +/// +/// A [`LabelListIndex`] is a [`BitmapIndex`] over the unnested list values plus a +/// separate `list_nulls` row set. Because distributed segments cover disjoint rows +/// (distinct fragments), merging is a cheap union of the underlying bitmap states +/// and of the `list_nulls` sets — no re-scan of source data is required. This +/// mirrors [`crate::scalar::bitmap::merge_bitmap_indices`] but also carries the +/// per-segment `list_nulls`. When `old_data_filter` is provided, rows from +/// retired fragments are removed from both the value bitmaps and `list_nulls`. +pub async fn merge_label_list_indices( + source_indices: &[Arc], + dest_store: &dyn IndexStore, + old_data_filter: Option, + progress: Arc, +) -> Result { + if source_indices.is_empty() { + return Err(Error::invalid_input( + "LabelList segment merge requires at least one source segment".to_string(), + )); + } + + let value_type = source_indices[0].values_index.value_type().clone(); + let mut merged_state = HashMap::::new(); + let mut merged_nulls = RowAddrTreeMap::new(); + + progress + .stage_start( + "merge_label_list_segments", + Some(source_indices.len() as u64), + "segments", + ) + .await?; + for (idx, source_index) in source_indices.iter().enumerate() { + if source_index.values_index.value_type() != &value_type { + return Err(Error::invalid_input(format!( + "LabelList segment has value type {:?}, expected {:?}", + source_index.values_index.value_type(), + value_type + ))); + } + + let state = source_index.values_index.load_bitmap_index_state().await?; + for (key, mut bitmap) in state { + if let Some(filter) = old_data_filter.as_ref() { + filter.retain_old_rows(&mut bitmap); + } + if bitmap.is_empty() { + continue; + } + merged_state + .entry(key) + .and_modify(|existing| *existing |= &bitmap) + .or_insert(bitmap); + } + let mut list_nulls = source_index.list_nulls.as_ref().clone(); + if let Some(filter) = old_data_filter.as_ref() { + filter.retain_old_rows(&mut list_nulls); + } + merged_nulls |= &list_nulls; + progress + .stage_progress("merge_label_list_segments", (idx + 1) as u64) + .await?; + } + progress.stage_complete("merge_label_list_segments").await?; + + progress + .stage_start("write_label_list_index", Some(1), "files") + .await?; + let file = + write_label_list_bitmap_index(merged_state, dest_store, &value_type, &merged_nulls).await?; + progress.stage_progress("write_label_list_index", 1).await?; + progress.stage_complete("write_label_list_index").await?; + + Ok(CreatedIndex { + index_details: prost_types::Any::from_msg(&pbold::LabelListIndexDetails::default()) + .unwrap(), + index_version: LABEL_LIST_INDEX_VERSION, + files: vec![file], + }) +} + /// The serializable state of a [`LabelListIndex`]. /// /// `LabelListIndex` is a thin wrapper around a [`BitmapIndex`] plus a separate @@ -513,11 +594,23 @@ impl LabelListIndexState { }) } + fn from_scalar_index(index: &dyn ScalarIndex) -> Result { + let label_list = index + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::internal( + "LabelListIndexState::from_scalar_index called with a non-label-list index", + ) + })?; + Self::from_index(label_list) + } + fn into_label_list_index( self, store: Arc, index_cache: &LanceCache, - frag_reuse_index: Option>, + frag_reuse_index: Option>, ) -> Result> { let bitmap = self .bitmap_state @@ -580,11 +673,7 @@ impl CacheKey for LabelListIndexStateKey { pub struct LabelListIndexPlugin; #[async_trait] -impl ScalarIndexPlugin for LabelListIndexPlugin { - fn name(&self) -> &str { - "LabelList" - } - +impl BasicTrainer for LabelListIndexPlugin { fn new_training_request( &self, _params: &str, @@ -606,25 +695,6 @@ impl ScalarIndexPlugin for LabelListIndexPlugin { ))) } - fn provides_exact_answer(&self) -> bool { - true - } - - fn version(&self) -> u32 { - LABEL_LIST_INDEX_VERSION - } - - fn new_query_parser( - &self, - index_name: String, - _index_details: &prost_types::Any, - ) -> Option> { - Some(Box::new(LabelListQueryParser::new( - index_name, - self.name().to_string(), - ))) - } - /// Train a new index /// /// The provided data must fulfill all the criteria returned by `training_criteria` @@ -634,15 +704,13 @@ impl ScalarIndexPlugin for LabelListIndexPlugin { data: SendableRecordBatchStream, index_store: &dyn IndexStore, _request: Box, - fragment_ids: Option>, + // Training over a fragment subset is supported for distributed builds: the + // provided `data` stream is already scoped to those fragments, so a partial + // index covering exactly those rows is produced. Segments are recombined by + // `merge_label_list_indices`. + _fragment_ids: Option>, _progress: Arc, ) -> Result { - if fragment_ids.is_some() { - return Err(Error::invalid_input_source( - "LabelList index does not support fragment training".into(), - )); - } - let schema = data.schema(); let field = schema .column_with_name(VALUE_COLUMN_NAME) @@ -681,13 +749,43 @@ impl ScalarIndexPlugin for LabelListIndexPlugin { files: vec![file], }) } +} + +#[async_trait] +impl ScalarIndexPlugin for LabelListIndexPlugin { + fn basic_trainer(&self) -> Option<&dyn BasicTrainer> { + Some(self) + } + + fn name(&self) -> &str { + "LabelList" + } + + fn provides_exact_answer(&self) -> bool { + true + } + + fn version(&self) -> u32 { + LABEL_LIST_INDEX_VERSION + } + + fn new_query_parser( + &self, + index_name: String, + _index_details: &prost_types::Any, + ) -> Option> { + Some(Box::new(LabelListQueryParser::new( + index_name, + self.name().to_string(), + ))) + } /// Load an index from storage async fn load_index( &self, index_store: Arc, _index_details: &prost_types::Any, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result> { Ok( @@ -699,7 +797,7 @@ impl ScalarIndexPlugin for LabelListIndexPlugin { async fn get_from_cache( &self, index_store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result>> { let Some(state) = cache.get_with_key(&LabelListIndexStateKey).await else { @@ -711,20 +809,34 @@ impl ScalarIndexPlugin for LabelListIndexPlugin { } async fn put_in_cache(&self, cache: &LanceCache, index: Arc) -> Result<()> { - let label_list = index - .as_any() - .downcast_ref::() - .ok_or_else(|| { - Error::internal( - "LabelListIndexPlugin::put_in_cache called with a non-label-list index", - ) - })?; - let state = LabelListIndexState::from_index(label_list)?; + let state = LabelListIndexState::from_scalar_index(index.as_ref())?; cache .insert_with_key(&LabelListIndexStateKey, Arc::new(state)) .await; Ok(()) } + + async fn get_or_insert_in_cache( + &self, + index_store: Arc, + frag_reuse_index: Option>, + cache: &LanceCache, + load: ScalarIndexLoad<'_>, + ) -> Result> { + single_flight_open( + cache, + LabelListIndexStateKey, + load, + LabelListIndexState::from_scalar_index, + move |state| { + Ok((*state) + .clone() + .into_label_list_index(index_store, cache, frag_reuse_index)? + as Arc) + }, + ) + .await + } } #[cfg(test)] diff --git a/rust/lance-index/src/scalar/lance_format.rs b/rust/lance-index/src/scalar/lance_format.rs index 76bcfcead57..0be5707d859 100644 --- a/rust/lance-index/src/scalar/lance_format.rs +++ b/rust/lance-index/src/scalar/lance_format.rs @@ -51,9 +51,12 @@ pub struct LanceIndexStore { impl DeepSizeOf for LanceIndexStore { fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + // Exclude the shared, session-scoped `metadata_cache` (accounted once by + // `Session::deep_size_of_children`): it is not this store's own footprint, and + // sizing it here makes every opened index that holds a store report the whole + // cache as its own bytes — inflating and N-times double-counting it. self.object_store.deep_size_of_children(context) + self.index_dir.as_ref().deep_size_of_children(context) - + self.metadata_cache.deep_size_of_children(context) } } @@ -136,10 +139,10 @@ impl IndexWriter for PreviousFileWrit } async fn finish(&mut self) -> Result { - Self::finish(self).await?; + let summary = Self::finish(self).await?; Ok(IndexFile { path: String::new(), - size_bytes: self.tell().await? as u64, + size_bytes: summary.size_bytes, }) } @@ -147,10 +150,10 @@ impl IndexWriter for PreviousFileWrit &mut self, metadata: HashMap, ) -> Result { - Self::finish_with_metadata(self, &metadata).await?; + let summary = Self::finish_with_metadata(self, &metadata).await?; Ok(IndexFile { path: String::new(), - size_bytes: self.tell().await? as u64, + size_bytes: summary.size_bytes, }) } } @@ -569,7 +572,7 @@ mod tests { use crate::scalar::bitmap::BitmapIndexPlugin; use crate::scalar::btree::{BTreeIndexPlugin, BTreeParameters}; use crate::scalar::label_list::LabelListIndexPlugin; - use crate::scalar::registry::{ScalarIndexPlugin, VALUE_COLUMN_NAME}; + use crate::scalar::registry::{BasicTrainer, ScalarIndexPlugin, VALUE_COLUMN_NAME}; use crate::scalar::{ LabelListQuery, SargableQuery, ScalarIndex, SearchResult, bitmap::BitmapIndex, @@ -590,6 +593,7 @@ mod tests { use datafusion_common::ScalarValue; use futures::FutureExt; use lance_core::ROW_ID; + use lance_core::utils::row_addr_remap::RowAddrRemap; use lance_core::utils::tempfile::TempDir; use lance_datagen::{ArrayGeneratorExt, BatchCount, ByteCount, RowCount, array, gen_batch}; use lance_select::{RowAddrTreeMap, RowSetOps}; @@ -606,6 +610,46 @@ mod tests { Arc::new(LanceIndexStore::new(object_store, test_path, cache)) } + #[tokio::test] + async fn test_store_deep_size_excludes_metadata_cache() { + // The metadata cache is session-scoped and shared; a store must not count + // it as its own footprint, so growing the cache must not change store size. + struct BlobKey; + impl lance_core::cache::CacheKey for BlobKey { + type ValueType = Vec; + fn key(&self) -> std::borrow::Cow<'_, str> { + std::borrow::Cow::Borrowed("blob") + } + fn type_name() -> &'static str { + "Vec" + } + } + + let index_dir = TempDir::default(); + let test_path = index_dir.obj_path(); + let (object_store, test_path) = ObjectStore::from_uri(test_path.as_ref()) + .now_or_never() + .unwrap() + .unwrap(); + let cache = Arc::new(lance_core::cache::LanceCache::with_capacity( + 128 * 1024 * 1024, + )); + let store = LanceIndexStore::new(object_store, test_path, cache.clone()); + + let before = store.deep_size_of(); + cache + .insert_with_key(&BlobKey, Arc::new(vec![0u8; 4 * 1024 * 1024])) + .await; + // Force moka to commit the write so a cache-inclusive size would grow. + let _ = cache.size_bytes().await; + let after = store.deep_size_of(); + + assert_eq!( + before, after, + "store deep size must exclude the shared metadata cache" + ); + } + async fn train_index( index_store: &Arc, data: impl RecordBatchReader + Send + Sync + 'static, @@ -1666,7 +1710,7 @@ mod tests { let remapped_dir = TempDir::default(); let remapped_store = test_store(&remapped_dir); index - .remap(&mapping, remapped_store.as_ref()) + .remap(&RowAddrRemap::direct(mapping), remapped_store.as_ref()) .await .unwrap(); let remapped_index = BitmapIndex::load(remapped_store, None, &LanceCache::no_cache()) diff --git a/rust/lance-index/src/scalar/ngram.rs b/rust/lance-index/src/scalar/ngram.rs index 582c7aab157..ade22cfc333 100644 --- a/rust/lance-index/src/scalar/ngram.rs +++ b/rust/lance-index/src/scalar/ngram.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::any::Any; use std::collections::BTreeMap; use std::iter::once; @@ -15,17 +16,17 @@ use super::{ AnyQuery, BuiltinIndexType, IndexFile, IndexReader, IndexStore, IndexWriter, MetricsCollector, ScalarIndex, ScalarIndexParams, SearchResult, TextQuery, }; -use crate::frag_reuse::FragReuseIndex; use crate::metrics::NoOpMetricsCollector; use crate::pbold; use crate::scalar::expression::{ScalarQueryParser, TextQueryParser}; use crate::scalar::registry::{ - DefaultTrainingRequest, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, - VALUE_COLUMN_NAME, + BasicTrainer, DefaultTrainingRequest, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, + TrainingRequest, VALUE_COLUMN_NAME, }; -use crate::scalar::{CreatedIndex, UpdateCriteria}; +use crate::scalar::{CreatedIndex, RowIdRemapper, UpdateCriteria}; use crate::{Index, IndexType}; use arrow::array::{AsArray, UInt32Builder}; +use arrow::buffer::{Buffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{UInt32Type, UInt64Type}; use arrow_array::{BinaryArray, RecordBatch, UInt32Array}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; @@ -59,6 +60,12 @@ const POSTING_LIST_COL: &str = "posting_list"; const POSTINGS_FILENAME: &str = "ngram_postings.lance"; const NGRAM_INDEX_VERSION: u32 = 0; +/// An i32-offset Binary array can hold at most i32::MAX bytes of values in total, +/// so a spill state whose serialized posting lists exceed that must be written as +/// multiple record batches (same approach as the bitmap index). Leave headroom. +const MAX_POSTING_LIST_BATCH_BYTES: usize = i32::MAX as usize - 1024 * 1024; +const POSTING_LIST_STREAM_BATCH_ROWS: usize = 64; + use std::sync::LazyLock; pub static TOKENS_FIELD: LazyLock = @@ -186,7 +193,7 @@ impl CacheKey for NGramPostingListKey { impl NGramPostingList { fn try_from_batch( batch: RecordBatch, - frag_reuse_index: Option>, + frag_reuse_index: Option>, ) -> Result { let bitmap_bytes = batch.column(0).as_binary::().value(0); let mut bitmap = RoaringTreemap::deserialize_from(bitmap_bytes) @@ -213,7 +220,7 @@ impl NGramPostingList { /// Reads on-demand ngram posting lists from storage (and stores them in a cache) struct NGramPostingListReader { reader: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, index_cache: WeakLanceCache, } @@ -298,7 +305,7 @@ impl DeepSizeOf for NGramIndex { impl NGramIndex { async fn from_store( store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, index_cache: &LanceCache, ) -> Result { let tokens = store.open_index_file(POSTINGS_FILENAME).await?; @@ -332,49 +339,35 @@ impl NGramIndex { }) } - fn remap_batch( + fn remap_state( &self, - batch: RecordBatch, - mapping: &HashMap>, - ) -> Result { - let posting_lists_array = batch - .column_by_name(POSTING_LIST_COL) - .expect_ok()? - .as_binary::(); - - let new_posting_lists = posting_lists_array - .iter() + state: NGramIndexSpillState, + mapping: &RowAddrRemap, + ) -> Result> { + let bitmaps = state + .bitmaps + .into_iter() .map(|posting_list| { - let posting_list = posting_list.unwrap(); - let posting_list = RoaringTreemap::deserialize_from(posting_list)?; - let new_posting_list = - RoaringTreemap::from_iter(posting_list.into_iter().filter_map(|row_id| { - match mapping.get(&row_id) { - Some(Some(new_row_id)) => Some(*new_row_id), - Some(None) => None, - None => Some(row_id), - } - })); - let mut buf = Vec::with_capacity(new_posting_list.serialized_size()); - new_posting_list.serialize_into(&mut buf)?; - Ok(buf) + RoaringTreemap::from_iter(posting_list.into_iter().filter_map(|row_id| { + match mapping.get(row_id) { + Some(Some(new_row_id)) => Some(new_row_id), + Some(None) => None, + None => Some(row_id), + } + })) }) - .collect::>>()?; + .collect(); - let new_posting_lists_array = BinaryArray::from_iter_values(new_posting_lists); - - Ok(RecordBatch::try_new( - POSTINGS_SCHEMA.clone(), - vec![ - batch.column_by_name(TOKENS_COL).expect_ok()?.clone(), - Arc::new(new_posting_lists_array), - ], - )?) + NGramIndexSpillState { + tokens: state.tokens, + bitmaps, + } + .try_into_batches() } async fn load( store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, index_cache: &LanceCache, ) -> Result> where @@ -523,7 +516,7 @@ impl ScalarIndex for NGramIndex { async fn remap( &self, - mapping: &HashMap>, + mapping: &RowAddrRemap, dest_store: &dyn IndexStore, ) -> Result { let reader = self.store.open_index_file(POSTINGS_FILENAME).await?; @@ -531,15 +524,12 @@ impl ScalarIndex for NGramIndex { .new_index_file(POSTINGS_FILENAME, POSTINGS_SCHEMA.clone()) .await?; - let mut offset = 0; - let num_rows = reader.num_rows(); - const BATCH_SIZE: usize = 64; - while offset < num_rows { - let batch_size = BATCH_SIZE.min(num_rows - offset); - let batch = reader.read_range(offset..offset + batch_size, None).await?; - let batch = self.remap_batch(batch, mapping)?; - writer.write_record_batch(batch).await?; - offset += BATCH_SIZE; + let mut spill_stream = + NGramIndexBuilder::stream_spill_reader(reader, MAX_POSTING_LIST_BATCH_BYTES)?; + while let Some(state) = spill_stream.try_next().await? { + for batch in self.remap_state(state, mapping)? { + writer.write_record_batch(batch).await?; + } } let file = writer.finish().await?; @@ -630,6 +620,85 @@ struct NGramIndexSpillState { bitmaps: Vec, } +struct NGramIndexSpillStateBuilder { + tokens: UInt32Builder, + bitmaps: Vec, + serialized_bytes: usize, +} + +impl NGramIndexSpillStateBuilder { + fn new() -> Self { + Self { + tokens: UInt32Builder::with_capacity(0), + bitmaps: Vec::new(), + serialized_bytes: 0, + } + } + + fn is_empty(&self) -> bool { + self.bitmaps.is_empty() + } + + fn len(&self) -> usize { + self.bitmaps.len() + } + + fn push( + &mut self, + token: u32, + bitmap: RoaringTreemap, + max_batch_bytes: usize, + ) -> Result> { + let posting_size = bitmap.serialized_size(); + if posting_size > max_batch_bytes { + return Err(Error::invalid_input(format!( + "posting list for ngram token {} serializes to {} bytes, which exceeds the {} bytes that fit in a single binary array", + token, posting_size, max_batch_bytes, + ))); + } + + let new_size = self + .serialized_bytes + .checked_add(posting_size) + .ok_or_else(|| { + Error::invalid_input(format!( + "posting list byte size overflowed while adding ngram token {}", + token + )) + })?; + let full_state = if !self.is_empty() && new_size > max_batch_bytes { + Some(self.finish()) + } else { + None + }; + + self.tokens.append_value(token); + self.bitmaps.push(bitmap); + self.serialized_bytes = posting_size + .checked_add(if full_state.is_some() { + 0 + } else { + self.serialized_bytes + }) + .ok_or_else(|| { + Error::invalid_input(format!( + "posting list byte size overflowed while adding ngram token {}", + token + )) + })?; + + Ok(full_state) + } + + fn finish(&mut self) -> NGramIndexSpillState { + self.serialized_bytes = 0; + NGramIndexSpillState { + tokens: std::mem::replace(&mut self.tokens, UInt32Builder::with_capacity(0)).finish(), + bitmaps: std::mem::take(&mut self.bitmaps), + } + } +} + impl NGramIndexSpillState { fn try_from_batch(batch: RecordBatch) -> Result { let tokens = batch @@ -653,16 +722,63 @@ impl NGramIndexSpillState { Ok(Self { tokens, bitmaps }) } - fn try_into_batch(self) -> Result { - let bitmap_array = BinaryArray::from_iter_values(self.bitmaps.into_iter().map(|bitmap| { - let mut buf = Vec::with_capacity(bitmap.serialized_size()); - bitmap.serialize_into(&mut buf).unwrap(); - buf - })); - Ok(RecordBatch::try_new( - POSTINGS_SCHEMA.clone(), - vec![Arc::new(self.tokens), Arc::new(bitmap_array)], - )?) + fn try_into_batches(self) -> Result> { + self.try_into_batches_impl(MAX_POSTING_LIST_BATCH_BYTES) + } + + // Split into multiple batches so that the cumulative serialized posting bytes + // of each batch stay under `max_batch_bytes`, avoiding i32 offset overflow in + // the Binary posting array. Postings are serialized straight into each batch's + // values buffer to avoid a second contiguous copy of multi-GiB payloads. + fn try_into_batches_impl(self, max_batch_bytes: usize) -> Result> { + debug_assert_eq!(self.tokens.len(), self.bitmaps.len()); + debug_assert!(max_batch_bytes <= i32::MAX as usize); + let make_batch = + |tokens: UInt32Array, values: Vec, offsets: Vec| -> Result { + let posting_array = BinaryArray::new( + OffsetBuffer::new(ScalarBuffer::from(offsets)), + Buffer::from_vec(values), + None, + ); + Ok(RecordBatch::try_new( + POSTINGS_SCHEMA.clone(), + vec![Arc::new(tokens), Arc::new(posting_array)], + )?) + }; + + let mut batches = Vec::new(); + let mut values: Vec = Vec::new(); + let mut offsets: Vec = vec![0]; + let mut batch_start = 0; + for (idx, bitmap) in self.bitmaps.into_iter().enumerate() { + let posting_size = bitmap.serialized_size(); + if posting_size > max_batch_bytes { + return Err(Error::invalid_input(format!( + "posting list for ngram token {} serializes to {} bytes, which exceeds the {} bytes that fit in a single binary array", + self.tokens.value(idx), + posting_size, + max_batch_bytes, + ))); + } + if values.len() + posting_size > max_batch_bytes { + batches.push(make_batch( + self.tokens.slice(batch_start, idx - batch_start), + std::mem::take(&mut values), + std::mem::replace(&mut offsets, vec![0]), + )?); + batch_start = idx; + } + bitmap.serialize_into(&mut values)?; + offsets.push(values.len() as i32); + } + if offsets.len() > 1 || batches.is_empty() { + batches.push(make_batch( + self.tokens.slice(batch_start, offsets.len() - 1), + values, + offsets, + )?); + } + Ok(batches) } } @@ -961,33 +1077,67 @@ impl NGramIndexBuilder { mut writer: Box, state: NGramIndexSpillState, ) -> Result<()> { - writer.write_record_batch(state.try_into_batch()?).await?; + Self::write_state(writer.as_mut(), state).await?; writer.finish().await?; Ok(()) } - async fn stream_spill_reader( + async fn write_state(writer: &mut dyn IndexWriter, state: NGramIndexSpillState) -> Result<()> { + for batch in state.try_into_batches()? { + writer.write_record_batch(batch).await?; + } + Ok(()) + } + + fn stream_spill_reader( reader: Arc, + max_batch_bytes: usize, ) -> Result>> { let num_rows = reader.num_rows(); - Ok(stream::try_unfold(0, move |offset| { - let reader = reader.clone(); - async move { - // These are small batches but, in the worst case scenario, each row could - // be massive (up to 128MB per row at 1B rows) and we end up breaking memory - let batch_size = std::cmp::min(num_rows - offset, 64); - if batch_size == 0 { - return Ok(None); + Ok(stream::try_unfold( + (0, NGramIndexSpillStateBuilder::new()), + move |(mut offset, mut builder)| { + let reader = reader.clone(); + async move { + while offset < num_rows { + // A single posting list is already bounded by + // MAX_POSTING_LIST_BATCH_BYTES. Reading one row at a time avoids + // materializing several large postings into the same BinaryArray + // before the byte-bounded writer can split them again. + let batch = reader.read_range(offset..offset + 1, None).await?; + offset += 1; + + let state = NGramIndexSpillState::try_from_batch(batch)?; + if state.tokens.len() != 1 || state.bitmaps.len() != 1 { + return Err(Error::internal(format!( + "expected one ngram posting row at offset {}, got {} tokens and {} posting lists", + offset - 1, + state.tokens.len(), + state.bitmaps.len(), + ))); + } + let token = state.tokens.value(0); + let mut bitmaps = state.bitmaps.into_iter(); + let bitmap = bitmaps.next().expect_ok()?; + if let Some(state) = builder.push(token, bitmap, max_batch_bytes)? { + return Ok(Some((state, (offset, builder)))); + } + if builder.len() >= POSTING_LIST_STREAM_BATCH_ROWS { + return Ok(Some((builder.finish(), (offset, builder)))); + } + } + + if builder.is_empty() { + Ok(None) + } else { + Ok(Some((builder.finish(), (offset, builder)))) + } } - let batch = reader.read_range(offset..offset + batch_size, None).await?; - let state = NGramIndexSpillState::try_from_batch(batch)?; - let new_offset = offset + batch_size; - Ok(Some((state, new_offset))) - } - .boxed() - })) + .boxed() + }, + )) } async fn stream_spill( @@ -997,7 +1147,7 @@ impl NGramIndexBuilder { let reader = spill_store .open_index_file(&Self::spill_filename(id)) .await?; - Self::stream_spill_reader(reader).await + Self::stream_spill_reader(reader, MAX_POSTING_LIST_BATCH_BYTES) } fn merge_spill_states( @@ -1093,21 +1243,21 @@ impl NGramIndexBuilder { if left_state.is_none() { // Left is done, full drain right let state = right_state.take().expect_ok()?; - writer.write_record_batch(state.try_into_batch()?).await?; + Self::write_state(writer, state).await?; while let Some(state) = right_stream.try_next().await? { - writer.write_record_batch(state.try_into_batch()?).await?; + Self::write_state(writer, state).await?; } } else if right_state.is_none() { // Right is done, full drain left let state = left_state.take().expect_ok()?; - writer.write_record_batch(state.try_into_batch()?).await?; + Self::write_state(writer, state).await?; while let Some(state) = left_stream.try_next().await? { - writer.write_record_batch(state.try_into_batch()?).await?; + Self::write_state(writer, state).await?; } } else { // There is a batch from both left and right. Need to merge them let merged = Self::merge_spill_states(&mut left_state, &mut right_state); - writer.write_record_batch(merged.try_into_batch()?).await?; + Self::write_state(writer, merged).await?; if left_state.is_none() { left_state = left_stream.try_next().await?; } @@ -1203,7 +1353,7 @@ impl NGramIndexBuilder { let left_stream = Self::stream_spill(self.spill_store.clone(), new_data_num).await?; let old_reader = old_index.open_index_file(POSTINGS_FILENAME).await?; - let right_stream = Self::stream_spill_reader(old_reader).await?; + let right_stream = Self::stream_spill_reader(old_reader, MAX_POSTING_LIST_BATCH_BYTES)?; Self::merge_spill_streams(left_stream, right_stream, writer.as_mut()).await?; @@ -1248,14 +1398,9 @@ impl NGramIndexBuilder { .open_index_file(&Self::spill_filename(index_to_copy)) .await?; - let num_rows = reader.num_rows(); - let mut offset = 0; - - while offset < num_rows { - let batch_size = std::cmp::min(num_rows - offset, 64); - let batch = reader.read_range(offset..offset + batch_size, None).await?; - writer.write_record_batch(batch).await?; - offset += batch_size; + let mut spill_stream = Self::stream_spill_reader(reader, MAX_POSTING_LIST_BATCH_BYTES)?; + while let Some(state) = spill_stream.try_next().await? { + Self::write_state(writer.as_mut(), state).await?; } writer.finish().await @@ -1279,11 +1424,7 @@ impl NGramIndexPlugin { } #[async_trait] -impl ScalarIndexPlugin for NGramIndexPlugin { - fn name(&self) -> &str { - "NGram" - } - +impl BasicTrainer for NGramIndexPlugin { fn new_training_request( &self, _params: &str, @@ -1301,29 +1442,6 @@ impl ScalarIndexPlugin for NGramIndexPlugin { ))) } - fn provides_exact_answer(&self) -> bool { - false - } - - fn version(&self) -> u32 { - NGRAM_INDEX_VERSION - } - - fn new_query_parser( - &self, - index_name: String, - _index_details: &prost_types::Any, - ) -> Option> { - Some(Box::new(TextQueryParser::new( - index_name, - self.name().to_string(), - // needs_recheck: ngram results are an inexact candidate superset. - true, - // supports_regex: the ngram index can answer regex queries. - true, - ))) - } - async fn train_index( &self, data: SendableRecordBatchStream, @@ -1346,12 +1464,46 @@ impl ScalarIndexPlugin for NGramIndexPlugin { files: vec![file], }) } +} + +#[async_trait] +impl ScalarIndexPlugin for NGramIndexPlugin { + fn basic_trainer(&self) -> Option<&dyn BasicTrainer> { + Some(self) + } + + fn name(&self) -> &str { + "NGram" + } + + fn provides_exact_answer(&self) -> bool { + false + } + + fn version(&self) -> u32 { + NGRAM_INDEX_VERSION + } + + fn new_query_parser( + &self, + index_name: String, + _index_details: &prost_types::Any, + ) -> Option> { + Some(Box::new(TextQueryParser::new( + index_name, + self.name().to_string(), + // needs_recheck: ngram results are an inexact candidate superset. + true, + // supports_regex: the ngram index can answer regex queries. + true, + ))) + } async fn load_index( &self, index_store: Arc, _index_details: &prost_types::Any, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result> { Ok(NGramIndex::load(index_store, frag_reuse_index, cache).await? as Arc) @@ -1360,34 +1512,85 @@ impl ScalarIndexPlugin for NGramIndexPlugin { #[cfg(test)] mod tests { + use lance_core::utils::row_addr_remap::RowAddrRemap; + use rstest::rstest; use std::{ collections::{HashMap, HashSet}, sync::Arc, }; - use arrow::datatypes::UInt64Type; - use arrow_array::{Array, RecordBatch, StringArray, UInt64Array}; + use arrow::array::AsArray; + use arrow::datatypes::{UInt32Type, UInt64Type}; + use arrow_array::{Array, RecordBatch, StringArray, UInt32Array, UInt64Array}; use arrow_schema::{DataType, Field, Schema}; + use async_trait::async_trait; use datafusion::{ execution::SendableRecordBatchStream, physical_plan::stream::RecordBatchStreamAdapter, }; use datafusion_common::DataFusionError; use futures::{TryStreamExt, stream}; use itertools::Itertools; - use lance_core::{ROW_ID, cache::LanceCache, utils::tempfile::TempDir}; + use lance_core::{Error, ROW_ID, Result, cache::LanceCache, utils::tempfile::TempDir}; use lance_datagen::{BatchCount, ByteCount, RowCount}; use lance_io::object_store::ObjectStore; use lance_select::RowAddrTreeMap; use lance_tokenizer::TextAnalyzer; + use roaring::RoaringTreemap; use crate::scalar::{ - ScalarIndex, SearchResult, TextQuery, + IndexReader, IndexStore, ScalarIndex, SearchResult, TextQuery, lance_format::LanceIndexStore, ngram::{NGramIndex, NGramIndexBuilder, NGramIndexBuilderOptions}, }; use crate::{metrics::NoOpMetricsCollector, scalar::registry::VALUE_COLUMN_NAME}; - use super::{NGRAM_TOKENIZER, ngram_to_token, tokenize_visitor}; + use super::{ + NGRAM_TOKENIZER, NGramIndexSpillState, POSTINGS_FILENAME, POSTINGS_SCHEMA, ngram_to_token, + tokenize_visitor, + }; + + struct MaxReadRangeReader { + inner: Arc, + max_rows: usize, + } + + #[async_trait] + impl IndexReader for MaxReadRangeReader { + async fn read_record_batch(&self, n: u64, batch_size: u64) -> Result { + self.inner.read_record_batch(n, batch_size).await + } + + async fn read_range( + &self, + range: std::ops::Range, + projection: Option<&[&str]>, + ) -> Result { + let rows = range.end - range.start; + if rows > self.max_rows { + return Err(Error::invalid_input(format!( + "read_range requested {} rows, max is {}", + rows, self.max_rows, + ))); + } + self.inner.read_range(range, projection).await + } + + async fn num_batches(&self, batch_size: u64) -> u32 { + self.inner.num_batches(batch_size).await + } + + fn num_rows(&self) -> usize { + self.inner.num_rows() + } + + fn schema(&self) -> &lance_core::datatypes::Schema { + self.inner.schema() + } + + fn file_size_bytes(&self) -> Option { + self.inner.file_size_bytes() + } + } fn collect_tokens(analyzer: &TextAnalyzer, text: &str) -> Vec { let mut tokens = Vec::with_capacity(text.len() * 3); @@ -1791,7 +1994,10 @@ mod tests { )); let remapping = HashMap::from([(2, Some(100)), (3, None), (4, Some(101))]); - index.remap(&remapping, test_store.as_ref()).await.unwrap(); + index + .remap(&RowAddrRemap::direct(remapping), test_store.as_ref()) + .await + .unwrap(); let index = NGramIndex::from_store(test_store, None, &LanceCache::no_cache()) .await @@ -1803,6 +2009,61 @@ mod tests { assert_eq!(null_posting_list, vec![100]); } + // Like `test_ngram_index_remap` but covering both RowAddrRemap modes: rows + // 0..4 of frag 0 are rewritten into frag 10; row 4 is deleted. + fn ngram_remap_compact() -> RowAddrRemap { + use lance_core::utils::row_addr_remap::GroupInput; + use roaring::RoaringTreemap; + RowAddrRemap::compact([GroupInput { + rewritten_old_row_addrs: RoaringTreemap::from_iter(0u64..4), + old_frag_ids: vec![0], + new_frags: vec![(10, 4)], + }]) + .unwrap() + } + + fn ngram_remap_explicit() -> RowAddrRemap { + RowAddrRemap::direct( + (0u64..4) + .map(|i| (i, Some((10u64 << 32) | i))) + .chain(std::iter::once((4u64, None))) + .collect(), + ) + } + + #[rstest] + #[case(ngram_remap_compact())] + #[case(ngram_remap_explicit())] + #[tokio::test] + async fn test_ngram_index_remap_compact(#[case] remap: RowAddrRemap) { + let data = simple_data_with_nulls(); + let builder = NGramIndexBuilder::try_new(NGramIndexBuilderOptions::default()).unwrap(); + let (index, _tmpdir) = do_train(builder, data).await; + + let row_ids = row_ids_in_index(&index).await; + assert_eq!(row_ids, vec![0, 1, 2, 3, 4]); + + let new_tmpdir = Arc::new(TempDir::default()); + let test_store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + new_tmpdir.obj_path(), + Arc::new(LanceCache::no_cache()), + )); + + index.remap(&remap, test_store.as_ref()).await.unwrap(); + + let index = NGramIndex::from_store(test_store, None, &LanceCache::no_cache()) + .await + .unwrap(); + let addr = |offset: u64| (10u64 << 32) | offset; + let row_ids = row_ids_in_index(&index).await; + assert_eq!(row_ids, vec![addr(0), addr(1), addr(2), addr(3)]); + + // rows 2 and 3 are the null docs; both are rewritten into frag 10. + let null_posting_list = get_null_posting_list(&index).await; + assert_eq!(null_posting_list, vec![addr(2), addr(3)]); + } + #[test_log::test(tokio::test)] async fn test_ngram_index_merge() { let data = simple_data_with_nulls(); @@ -1874,4 +2135,145 @@ mod tests { assert_eq!(index.tokens.len(), 29012); } + + #[test] + fn test_spill_state_chunks_by_byte_size() { + let bitmaps = (0..8u64) + .map(|i| RoaringTreemap::from_iter(0..(i + 1) * 100)) + .collect::>(); + let tokens = UInt32Array::from_iter_values(0..8); + let state = NGramIndexSpillState { + tokens: tokens.clone(), + bitmaps: bitmaps.clone(), + }; + + // Small enough that several splits are required, large enough that some + // batches hold more than one posting + let max_batch_bytes = bitmaps.iter().map(|b| b.serialized_size()).max().unwrap() * 2; + let batches = state.try_into_batches_impl(max_batch_bytes).unwrap(); + assert!(batches.len() > 1); + + // Token order and posting contents survive the chunking + let mut row = 0; + for batch in &batches { + let batch_tokens = batch["tokens"].as_primitive::(); + let batch_postings = batch["posting_list"].as_binary::(); + let mut batch_bytes = 0; + for i in 0..batch.num_rows() { + assert_eq!(batch_tokens.value(i), tokens.value(row)); + let posting = batch_postings.value(i); + batch_bytes += posting.len(); + assert_eq!( + RoaringTreemap::deserialize_from(posting).unwrap(), + bitmaps[row] + ); + row += 1; + } + assert!(batch_bytes <= max_batch_bytes || batch.num_rows() == 1); + } + assert_eq!(row, 8); + } + + #[test_log::test(tokio::test)] + async fn test_spill_reader_does_not_materialize_multirow_posting_batches() { + let bitmaps = (0..8u64) + .map(|i| RoaringTreemap::from_iter(0..(i + 1) * 100)) + .collect::>(); + let tokens = UInt32Array::from_iter_values(0..8); + let state = NGramIndexSpillState { + tokens: tokens.clone(), + bitmaps: bitmaps.clone(), + }; + let max_batch_bytes = bitmaps.iter().map(|b| b.serialized_size()).max().unwrap() * 2; + + let tmpdir = Arc::new(TempDir::default()); + let store = LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.obj_path(), + Arc::new(LanceCache::no_cache()), + ); + let mut writer = store + .new_index_file(POSTINGS_FILENAME, POSTINGS_SCHEMA.clone()) + .await + .unwrap(); + for batch in state.try_into_batches().unwrap() { + writer.write_record_batch(batch).await.unwrap(); + } + writer.finish().await.unwrap(); + + let reader = store.open_index_file(POSTINGS_FILENAME).await.unwrap(); + let reader = Arc::new(MaxReadRangeReader { + inner: reader, + max_rows: 1, + }); + let states = NGramIndexBuilder::stream_spill_reader(reader, max_batch_bytes) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!(states.len() > 1); + + let mut row = 0; + for state in states { + let batch_bytes = state + .bitmaps + .iter() + .map(RoaringTreemap::serialized_size) + .sum::(); + assert!(batch_bytes <= max_batch_bytes || state.bitmaps.len() == 1); + for (token, bitmap) in state.tokens.values().iter().zip(state.bitmaps) { + assert_eq!(*token, tokens.value(row)); + assert_eq!(bitmap, bitmaps[row]); + row += 1; + } + } + assert_eq!(row, 8); + } + + #[test] + fn test_spill_state_rejects_oversized_posting() { + let bitmap = RoaringTreemap::from_iter(0..1000u64); + let too_small = bitmap.serialized_size() - 1; + let state = NGramIndexSpillState { + tokens: UInt32Array::from_iter_values([42]), + bitmaps: vec![bitmap], + }; + let err = state.try_into_batches_impl(too_small).unwrap_err(); + assert!(err.to_string().contains("token 42"), "{}", err); + } + + #[test] + fn test_empty_spill_state_yields_one_empty_batch() { + let state = NGramIndexSpillState { + tokens: UInt32Array::from_iter_values([]), + bitmaps: vec![], + }; + let batches = state.try_into_batches().unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 0); + } + + // Reproduces https://linear.app/lancedb/issue/ENT-874: serialized posting lists + // totalling more than i32::MAX bytes used to panic with "byte array offset + // overflow" when packed into a single Binary array. + #[test] + #[ignore = "needs ~8 GiB of RAM and a couple of minutes; run manually"] + fn test_spill_state_over_i32_max_bytes() { + // Every 16th value keeps each container an array container (4096 entries, + // 2 bytes per value, immune to run compression), so the treemap serializes + // to ~450 MiB. Six copies exceed i32::MAX total bytes. + let bitmap = RoaringTreemap::from_sorted_iter((0..225_000_000u64).map(|v| v * 16)).unwrap(); + assert!(bitmap.serialized_size() > 400 * 1024 * 1024); + let bitmaps = vec![bitmap; 6]; + let tokens = UInt32Array::from_iter_values(0..6); + let state = NGramIndexSpillState { tokens, bitmaps }; + + let batches = state.try_into_batches().unwrap(); + assert!(batches.len() > 1); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 6); + for batch in &batches { + let postings = batch["posting_list"].as_binary::(); + assert!(postings.value_data().len() <= i32::MAX as usize); + } + } } diff --git a/rust/lance-index/src/scalar/registry.rs b/rust/lance-index/src/scalar/registry.rs index 0add98d8ab3..39dfff2c2c8 100644 --- a/rust/lance-index/src/scalar/registry.rs +++ b/rust/lance-index/src/scalar/registry.rs @@ -7,17 +7,17 @@ use std::sync::Arc; use arrow_schema::Field; use async_trait::async_trait; use datafusion::execution::SendableRecordBatchStream; +use futures::future::BoxFuture; use lance_core::{ Result, - cache::{LanceCache, UnsizedCacheKey}, + cache::{CacheKey, LanceCache, UnsizedCacheKey}, + deepsize::DeepSizeOf, }; use crate::progress::IndexBuildProgress; use crate::registry::IndexPluginRegistry; -use crate::{ - frag_reuse::FragReuseIndex, - scalar::{CreatedIndex, IndexStore, ScalarIndex, expression::ScalarQueryParser}, -}; +use crate::scalar::RowIdRemapper; +use crate::scalar::{CreatedIndex, IndexStore, ScalarIndex, expression::ScalarQueryParser}; pub const VALUE_COLUMN_NAME: &str = "value"; @@ -58,15 +58,12 @@ impl TrainingCriteria { } } -/// A trait that describes what criteria is needed to train an index -/// -/// The training process has two steps. First, the parameters are given to the -/// plugin and it creates a TrainingRequest. Then, the caller prepares the training -/// data and calls train_index. +/// A trait object for plugin-specific training parameters and data requirements. /// -/// The call to train_index will include the training request. This allows the plugin -/// to stash any deserialized parameter info in the request and fetch it later during -/// training by downcasting to the appropriate type. +/// Returned by [`BasicTrainer::new_training_request`]. The caller uses +/// [`criteria`](Self::criteria) to prepare the training data stream, then passes +/// the request back to [`BasicTrainer::train_index`], which may downcast +/// it to the plugin-specific concrete type to recover parsed parameters. pub trait TrainingRequest: std::any::Any + Send + Sync { fn as_any(&self) -> &dyn std::any::Any; fn criteria(&self) -> &TrainingCriteria; @@ -93,26 +90,35 @@ impl TrainingRequest for DefaultTrainingRequest { } } -/// A trait for scalar index plugins +/// Implemented by indexes that can train on a stream of column data. +/// +/// The training process has two stages. In the first stage, the caller provides +/// index parameters and receives a [`TrainingRequest`] that describes what criteria +/// the training data must satisfy (e.g. sort order, row-ID availability). In the +/// second stage, the caller prepares the data accordingly and calls +/// [`train_index`](Self::train_index). +/// +/// Any scalar index plugin that builds from a column data stream should implement +/// this trait. #[async_trait] -pub trait ScalarIndexPlugin: Send + Sync + std::fmt::Debug { - /// Creates a new training request from the given parameters +pub trait BasicTrainer: Send + Sync { + /// Creates a new training request from the given parameters. /// - /// This training request specifies the criteria that the data must satisfy to train the index. - /// For example, does the index require the input data to be sorted? + /// The returned request specifies the criteria the training data must satisfy. + /// It is the caller's responsibility to prepare data that meets those criteria + /// before calling [`train_index`](Self::train_index). fn new_training_request(&self, params: &str, field: &Field) -> Result>; - /// Train a new index + /// Train a new index from a prepared data stream. /// - /// The provided data must fulfill all the criteria returned by `training_criteria`. - /// It is the caller's responsibility to ensure this. + /// The provided data must fulfill all the criteria returned by + /// [`new_training_request`](Self::new_training_request). It is the caller's + /// responsibility to ensure this. /// - /// Returns index details that describe the index. These details can potentially be - /// useful for planning (although this will currently require inside information on - /// the index type) and they will need to be provided when loading the index. - /// - /// It is the caller's responsibility to store these details somewhere. + /// Returns index details describing the index. These details may be useful for + /// planning and must be provided when loading the index. It is the caller's + /// responsibility to store them. async fn train_index( &self, data: SendableRecordBatchStream, @@ -121,6 +127,36 @@ pub trait ScalarIndexPlugin: Send + Sync + std::fmt::Debug { fragment_ids: Option>, progress: Arc, ) -> Result; +} + +/// A trait for scalar index plugins +#[async_trait] +pub trait ScalarIndexPlugin: Send + Sync + std::fmt::Debug { + /// Returns this plugin's [`BasicTrainer`] implementation, if any. + /// + /// Training an index can be a complex process. For example, a btree index might + /// be trained using a shuffler from a distributed OLAP system such as + /// Spark or Ray. A vector index can be trained by sampling the column to create + /// a kmeans model and then streaming the vectors to assign partitions. Encapsulating + /// the entire set of possible approaches is beyond what this trait can model. + /// This is especially true because this is a low-level crate with no concept of a table + /// or a dataset. + /// + /// However, in many cases, an index can be trained on a (potentially sorted) stream + /// of column data. There is also significant utility in being able to provide users + /// with a simple generic "create an index" API. + /// + /// This method is a compromise. Indexes that support training on a stream of column + /// data should override this to return `Some(self)`. Indexes that need their own + /// individualized training approaches should return `None` and provide their own + /// methods for training. + /// + /// An index can take both approaches. Providing a simple (but maybe less + /// efficient) stream-based trainer while also providing more specialized index + /// creation methods elsewhere. + fn basic_trainer(&self) -> Option<&dyn BasicTrainer> { + None + } /// A short name for the index /// @@ -158,7 +194,7 @@ pub trait ScalarIndexPlugin: Send + Sync + std::fmt::Debug { &self, index_store: Arc, index_details: &prost_types::Any, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result>; @@ -170,14 +206,14 @@ pub trait ScalarIndexPlugin: Send + Sync + std::fmt::Debug { /// The default implementation reads an in-memory `Arc` entry. /// Plugins whose index has a serializable representation should override this /// (together with [`put_in_cache`](Self::put_in_cache)) to store that - /// representation under a sized [`CacheKey`](lance_core::cache::CacheKey) with + /// representation under a sized [`CacheKey`] with /// a codec, and reconstruct the index here. `index_store` and /// `frag_reuse_index` are provided so the override can rebuild the index /// without re-reading metadata. async fn get_from_cache( &self, _index_store: Arc, - _frag_reuse_index: Option>, + _frag_reuse_index: Option>, cache: &LanceCache, ) -> Result>> { Ok(cache.get_unsized_with_key(&ScalarIndexCacheKey).await) @@ -196,6 +232,31 @@ pub trait ScalarIndexPlugin: Send + Sync + std::fmt::Debug { Ok(()) } + /// Open an index through the cache, awaiting `load` only on a miss. + /// + /// The default read-through / write-back over + /// [`get_from_cache`](Self::get_from_cache) and + /// [`put_in_cache`](Self::put_in_cache) does not coalesce concurrent cold + /// opens; plugins with a serializable form should override with + /// [`single_flight_open`] so one shared load populates the sized state key. + async fn get_or_insert_in_cache( + &self, + index_store: Arc, + frag_reuse_index: Option>, + cache: &LanceCache, + load: ScalarIndexLoad<'_>, + ) -> Result> { + if let Some(index) = self + .get_from_cache(index_store, frag_reuse_index, cache) + .await? + { + return Ok(index); + } + let index = load.await?; + self.put_in_cache(cache, index.clone()).await?; + Ok(index) + } + /// Optional hook allowing a plugin to provide statistics without loading the index. async fn load_statistics( &self, @@ -219,6 +280,48 @@ pub trait ScalarIndexPlugin: Send + Sync + std::fmt::Debug { } } +/// A boxed, `Send` future performing the storage-level load of a scalar index +/// (compat checks, `load_index`, metrics). +/// +/// Passed to [`ScalarIndexPlugin::get_or_insert_in_cache`], which awaits it at +/// most once on a cache miss and drops it un-awaited on a warm hit. +pub type ScalarIndexLoad<'a> = BoxFuture<'a, Result>>; + +/// Single-flight open helper for plugins with a serializable form. +/// +/// Concurrent cold opens of the same index coalesce onto one `load`: it runs +/// once, `to_state` converts the opened index to its sized state, and the state +/// is cached under `state_key` (persisted via the key's +/// [`CacheCodec`](lance_core::cache::CacheCodec)). Every caller — warm hits +/// included — then rebuilds the index from the shared state via `from_state`, an +/// IO-free reconstruct. +/// +/// `to_state` / `from_state` mirror the plugin's +/// [`put_in_cache`](ScalarIndexPlugin::put_in_cache) / +/// [`get_from_cache`](ScalarIndexPlugin::get_from_cache), keeping the state +/// representation defined in one place. +pub async fn single_flight_open( + cache: &LanceCache, + state_key: K, + load: ScalarIndexLoad<'_>, + to_state: ToState, + from_state: FromState, +) -> Result> +where + K: CacheKey + Send, + K::ValueType: DeepSizeOf + Send + Sync + 'static, + ToState: FnOnce(&dyn ScalarIndex) -> Result + Send, + FromState: FnOnce(Arc) -> Result> + Send, +{ + let state = cache + .get_or_insert_with_key(state_key, move || async move { + let index = load.await?; + to_state(index.as_ref()) + }) + .await?; + from_state(state) +} + /// In-memory cache key for a whole `Arc`. /// /// Used by the default [`ScalarIndexPlugin::get_from_cache`] / diff --git a/rust/lance-index/src/scalar/rtree.rs b/rust/lance-index/src/scalar/rtree.rs index e26177b4b4a..0d6b54a02ca 100644 --- a/rust/lance-index/src/scalar/rtree.rs +++ b/rust/lance-index/src/scalar/rtree.rs @@ -1,17 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use crate::frag_reuse::FragReuseIndex; use crate::metrics::{MetricsCollector, NoOpMetricsCollector}; use crate::scalar::expression::{GeoQueryParser, ScalarQueryParser}; use crate::scalar::lance_format::LanceIndexStore; use crate::scalar::registry::{ - ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, + BasicTrainer, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, }; use crate::scalar::rtree::sort::Sorter; use crate::scalar::{ AnyQuery, BuiltinIndexType, CreatedIndex, GeoQuery, IndexFile, IndexReader, IndexStore, - IndexWriter, ScalarIndex, ScalarIndexParams, SearchResult, UpdateCriteria, + IndexWriter, RowIdRemapper, ScalarIndex, ScalarIndexParams, SearchResult, UpdateCriteria, }; use crate::{Index, IndexType, pb}; use arrow_array::UInt32Array; @@ -33,6 +32,7 @@ use lance_arrow::RecordBatchExt; use lance_core::cache::{CacheKey, LanceCache, WeakLanceCache}; use lance_core::deepsize::DeepSizeOf; use lance_core::utils::address::RowAddress; +use lance_core::utils::row_addr_remap::RowAddrRemap; use lance_core::utils::tempfile::TempDir; use lance_core::{Error, ROW_ID, Result}; use lance_datafusion::chunker::chunk_concat_stream; @@ -258,7 +258,7 @@ impl CacheKey for RTreeCacheKey { pub struct RTreeIndex { pub(crate) metadata: Arc, store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, index_cache: WeakLanceCache, pages_reader: Arc, nulls_reader: Arc, @@ -276,7 +276,7 @@ impl std::fmt::Debug for RTreeIndex { impl RTreeIndex { pub async fn load( store: Arc, - frag_reuse_index: Option>, + frag_reuse_index: Option>, index_cache: &LanceCache, ) -> Result> { let pages_reader = store.open_index_file(RTREE_PAGES_NAME).await?; @@ -544,7 +544,7 @@ impl ScalarIndex for RTreeIndex { async fn remap( &self, - _mapping: &HashMap>, + _mapping: &RowAddrRemap, _dest_store: &dyn IndexStore, ) -> Result { Err(Error::invalid_input_source( @@ -908,11 +908,7 @@ impl RTreeIndexPlugin { } #[async_trait] -impl ScalarIndexPlugin for RTreeIndexPlugin { - fn name(&self) -> &str { - "RTree" - } - +impl BasicTrainer for RTreeIndexPlugin { fn new_training_request( &self, params: &str, @@ -966,6 +962,17 @@ impl ScalarIndexPlugin for RTreeIndexPlugin { files, }) } +} + +#[async_trait] +impl ScalarIndexPlugin for RTreeIndexPlugin { + fn basic_trainer(&self) -> Option<&dyn BasicTrainer> { + Some(self) + } + + fn name(&self) -> &str { + "RTree" + } fn provides_exact_answer(&self) -> bool { false @@ -990,7 +997,7 @@ impl ScalarIndexPlugin for RTreeIndexPlugin { &self, index_store: Arc, _index_details: &prost_types::Any, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result> { Ok(RTreeIndex::load(index_store, frag_reuse_index, cache).await? as Arc) diff --git a/rust/lance-index/src/scalar/zoned.rs b/rust/lance-index/src/scalar/zoned.rs index 7ceed851bae..9c090f5e1d8 100644 --- a/rust/lance-index/src/scalar/zoned.rs +++ b/rust/lance-index/src/scalar/zoned.rs @@ -92,13 +92,14 @@ where pub async fn train( mut self, stream: SendableRecordBatchStream, - ) -> Result> { + ) -> Result<(Vec, RowAddrTreeMap)> { let zone_size = usize::try_from(self.zone_capacity).map_err(|_| { Error::invalid_input("zone capacity does not fit into usize on this platform") })?; let mut batches = chunk_concat_stream(stream, zone_size); let mut zones = Vec::new(); + let mut null_rows = RowAddrTreeMap::new(); let mut current_fragment_id: Option = None; let mut current_zone_len: usize = 0; let mut zone_start_offset: Option = None; @@ -155,6 +156,16 @@ where self.processor .process_chunk(&values.slice(batch_offset, take))?; + // Record exact row addresses for null values in this chunk. + let chunk = values.slice(batch_offset, take); + if chunk.null_count() > 0 { + for i in 0..take { + if values.is_null(batch_offset + i) { + null_rows.insert(row_addr_col.value(batch_offset + i)); + } + } + } + // Track the first and last row offsets to handle non-contiguous offsets // after deletions. Zone length (offset span) is computed as (last - first + 1), // not the actual row count. @@ -200,7 +211,7 @@ where } } - Ok(zones) + Ok((zones, null_rows)) } /// Flushes a non-empty zone and resets the processor state. @@ -266,21 +277,21 @@ where } /// Helper that retrains zones from `stream` and appends them to the existing -/// statistics. Useful for index update paths that need to merge new fragments -/// into an existing zone list. +/// statistics. Returns the combined zone list and the null-row bitmap for the +/// new data only — callers are responsible for merging with any existing bitmap. pub async fn rebuild_zones

( existing: &[P::ZoneStatistics], trainer: ZoneTrainer

, stream: SendableRecordBatchStream, -) -> Result> +) -> Result<(Vec, RowAddrTreeMap)> where P: ZoneProcessor, P::ZoneStatistics: Clone, { let mut combined = existing.to_vec(); - let mut new_zones = trainer.train(stream).await?; + let (mut new_zones, null_rows) = trainer.train(stream).await?; combined.append(&mut new_zones); - Ok(combined) + Ok((combined, null_rows)) } #[cfg(test)] @@ -362,7 +373,7 @@ mod tests { let processor = MockProcessor::new(); let trainer = ZoneTrainer::new(processor, 4).unwrap(); - let stats = trainer.train(stream).await.unwrap(); + let (stats, _) = trainer.train(stream).await.unwrap(); // Three zones: offsets [0..=3], [4..=7], [8..=9] assert_eq!(stats.len(), 3); @@ -393,7 +404,7 @@ mod tests { let processor = MockProcessor::new(); let trainer = ZoneTrainer::new(processor, 10).unwrap(); - let stats = trainer.train(stream).await.unwrap(); + let (stats, _) = trainer.train(stream).await.unwrap(); // Two zones, one per fragment (capacity=10 is large enough) assert_eq!(stats.len(), 2); @@ -447,7 +458,7 @@ mod tests { let processor = MockProcessor::new(); let trainer = ZoneTrainer::new(processor, 10).unwrap(); - let stats = trainer.train(stream).await.unwrap(); + let (stats, _) = trainer.train(stream).await.unwrap(); // One zone containing the 3 valid rows (empty batches skipped) assert_eq!(stats.len(), 1); @@ -469,7 +480,7 @@ mod tests { let processor = MockProcessor::new(); let trainer = ZoneTrainer::new(processor, 1).unwrap(); - let stats = trainer.train(stream).await.unwrap(); + let (stats, _) = trainer.train(stream).await.unwrap(); // Three zones, one per row (capacity=1) assert_eq!(stats.len(), 3); @@ -494,7 +505,7 @@ mod tests { let processor = MockProcessor::new(); let trainer = ZoneTrainer::new(processor, 10000).unwrap(); - let stats = trainer.train(stream).await.unwrap(); + let (stats, _) = trainer.train(stream).await.unwrap(); // One zone containing all 100 rows (capacity is large enough) assert_eq!(stats.len(), 1); @@ -530,7 +541,7 @@ mod tests { let processor = MockProcessor::new(); let trainer = ZoneTrainer::new(processor, 4).unwrap(); - let stats = trainer.train(stream).await.unwrap(); + let (stats, _) = trainer.train(stream).await.unwrap(); // Two zones: first 4 rows, then remaining 2 rows assert_eq!(stats.len(), 2); @@ -561,7 +572,7 @@ mod tests { let processor = MockProcessor::new(); let trainer = ZoneTrainer::new(processor, 3).unwrap(); - let stats = trainer.train(stream).await.unwrap(); + let (stats, _) = trainer.train(stream).await.unwrap(); // Three zones: frag 0 full zone, frag 0 partial (flushed at boundary), frag 1 assert_eq!(stats.len(), 3); @@ -602,7 +613,7 @@ mod tests { let processor = MockProcessor::new(); let trainer = ZoneTrainer::new(processor, 4).unwrap(); - let stats = trainer.train(stream).await.unwrap(); + let (stats, _) = trainer.train(stream).await.unwrap(); // Should create 2 zones (capacity=4): // Zone 0: rows at offsets [0, 1, 5, 7] (4 rows) @@ -637,7 +648,7 @@ mod tests { let processor = MockProcessor::new(); let trainer = ZoneTrainer::new(processor, 10).unwrap(); - let stats = trainer.train(stream).await.unwrap(); + let (stats, _) = trainer.train(stream).await.unwrap(); // One zone with 3 rows, but offset span [0..=200] so length=201 due to large gaps assert_eq!(stats.len(), 1); @@ -663,7 +674,7 @@ mod tests { let processor = MockProcessor::new(); let trainer = ZoneTrainer::new(processor, 10).unwrap(); - let stats = trainer.train(stream).await.unwrap(); + let (stats, _) = trainer.train(stream).await.unwrap(); // Should create 3 zones (one per fragment) assert_eq!(stats.len(), 3); @@ -810,7 +821,7 @@ mod tests { )); let trainer = ZoneTrainer::new(MockProcessor::new(), 2).unwrap(); - let rebuilt = rebuild_zones(&existing, trainer, stream).await.unwrap(); + let (rebuilt, _) = rebuild_zones(&existing, trainer, stream).await.unwrap(); // Existing zone should remain unchanged and new stats appended afterwards assert_eq!(rebuilt.len(), 2); assert_eq!(rebuilt[0].sum, 50); @@ -840,7 +851,7 @@ mod tests { )); let trainer = ZoneTrainer::new(MockProcessor::new(), 2).unwrap(); - let rebuilt = rebuild_zones(&existing, trainer, stream).await.unwrap(); + let (rebuilt, _) = rebuild_zones(&existing, trainer, stream).await.unwrap(); // Existing zone plus two new fragments should yield three total zones assert_eq!(rebuilt.len(), 3); assert_eq!(rebuilt[0].bound.fragment_id, 0); diff --git a/rust/lance-index/src/scalar/zonemap.rs b/rust/lance-index/src/scalar/zonemap.rs index af5380cce30..766a563593d 100644 --- a/rust/lance-index/src/scalar/zonemap.rs +++ b/rust/lance-index/src/scalar/zonemap.rs @@ -16,7 +16,7 @@ use crate::Any; use crate::pbold; use crate::scalar::expression::{SargableQueryParser, ScalarQueryParser}; use crate::scalar::registry::{ - ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, + BasicTrainer, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, }; use crate::scalar::{ BuiltinIndexType, CreatedIndex, IndexFile, SargableQuery, ScalarIndexParams, UpdateCriteria, @@ -24,6 +24,7 @@ use crate::scalar::{ }; use lance_arrow_stats::StatisticsAccumulator; use lance_core::cache::{LanceCache, WeakLanceCache}; +use lance_core::utils::row_addr_remap::RowAddrRemap; use serde::{Deserialize, Serialize}; use std::sync::LazyLock; @@ -33,10 +34,11 @@ use arrow_array::{ use arrow_schema::{DataType, Field}; use datafusion::execution::SendableRecordBatchStream; use datafusion_common::ScalarValue; +use lance_select::RowAddrTreeMap; use std::{collections::HashMap, sync::Arc}; use super::{AnyQuery, IndexStore, MetricsCollector, ScalarIndex, SearchResult}; -use crate::scalar::FragReuseIndex; +use crate::scalar::RowIdRemapper; use crate::{Index, IndexType}; use async_trait::async_trait; use lance_core::Error; @@ -49,6 +51,7 @@ const ROWS_PER_ZONE_DEFAULT: u64 = 8192; // 1 zone every two batches const ZONEMAP_FILENAME: &str = "zonemap.lance"; const ZONEMAP_SIZE_META_KEY: &str = "rows_per_zone"; +const NULL_BITMAP_META_KEY: &str = "null_bitmap"; const ZONEMAP_INDEX_VERSION: u32 = 0; /// Basic stats about zonemap index @@ -107,8 +110,11 @@ pub struct ZoneMapIndex { // The maximum rows per zone provided by user rows_per_zone: u64, store: Arc, - fri: Option>, + fri: Option>, index_cache: WeakLanceCache, + // Exact set of null row addresses across all zones; None when loaded from an + // older index that did not persist this bitmap. + null_rows: Option, } impl std::fmt::Debug for ZoneMapIndex { @@ -126,7 +132,7 @@ impl std::fmt::Debug for ZoneMapIndex { impl DeepSizeOf for ZoneMapIndex { fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - self.zones.deep_size_of_children(context) + self.zones.deep_size_of_children(context) + self.null_rows.deep_size_of_children(context) } } @@ -150,6 +156,48 @@ impl ZoneMapIndex { Self::zone_has_finite_min(zone) && !(zone.max.is_null() || Self::scalar_is_nan(&zone.max)) } + /// Global `[min, max]` folded across one or more ZoneMap segments (the + /// disjoint per-column segments of a multi-segment index), without a scan. + /// + /// `None` when no zone has a finite bound, or when any zone's `max` is NaN: + /// `ScalarValue`'s total order ranks NaN above every finite value, so a + /// NaN-bearing zone hides its true finite max and no sound upper bound exists + /// without a scan — folding only the finite maxes would yield a *subset* that + /// prunes live rows. Folding raw zones (not each segment's `value_range`) + /// keeps this exact across segments. + /// + /// Otherwise the range is a superset of the segments' live values, + /// conservative under deletion vectors: safe to prune with, not guaranteed + /// tight. The caller must ensure the segments jointly cover every live + /// fragment. + pub fn value_range_over<'a>( + segments: impl IntoIterator, + ) -> Option<(ScalarValue, ScalarValue)> { + let mut min: Option<&ScalarValue> = None; + let mut max: Option<&ScalarValue> = None; + for zone in segments.into_iter().flat_map(|seg| seg.zones.iter()) { + if Self::scalar_is_nan(&zone.max) { + return None; + } + if Self::scalar_is_finite_bound(&zone.min) + && min.is_none_or(|cur| zone.min.partial_cmp(cur).is_some_and(|o| o.is_lt())) + { + min = Some(&zone.min); + } + if Self::scalar_is_finite_bound(&zone.max) + && max.is_none_or(|cur| zone.max.partial_cmp(cur).is_some_and(|o| o.is_gt())) + { + max = Some(&zone.max); + } + } + Some((min?.clone(), max?.clone())) + } + + /// A scalar usable as a global-range bound: non-null and, for floats, non-NaN. + fn scalar_is_finite_bound(v: &ScalarValue) -> bool { + !v.is_null() && !Self::scalar_is_nan(v) + } + /// Evaluates whether a zone could potentially contain values matching the query. /// /// NaN query values use the explicit `nan_count`. For finite query values, @@ -217,10 +265,8 @@ impl ZoneMapIndex { return Ok(zone.nan_count > 0); } } - ScalarValue::Float64(Some(f)) => { - if f.is_nan() { - return Ok(zone.nan_count > 0); - } + ScalarValue::Float64(Some(f)) if f.is_nan() => { + return Ok(zone.nan_count > 0); } _ => {} } @@ -247,10 +293,8 @@ impl ZoneMapIndex { return Ok(false); // Nothing is greater than NaN } } - ScalarValue::Float64(Some(f)) => { - if f.is_nan() { - return Ok(false); // Nothing is greater than NaN - } + ScalarValue::Float64(Some(f)) if f.is_nan() => { + return Ok(false); // Nothing is greater than NaN } _ => {} } @@ -274,10 +318,8 @@ impl ZoneMapIndex { return Ok(zone.nan_count > 0 || zone_min <= e); } } - ScalarValue::Float64(Some(f)) => { - if f.is_nan() { - return Ok(zone.nan_count > 0 || zone_min <= e); - } + ScalarValue::Float64(Some(f)) if f.is_nan() => { + return Ok(zone.nan_count > 0 || zone_min <= e); } _ => {} } @@ -297,10 +339,8 @@ impl ZoneMapIndex { return Ok(true); } } - ScalarValue::Float64(Some(f)) => { - if f.is_nan() { - return Ok(true); - } + ScalarValue::Float64(Some(f)) if f.is_nan() => { + return Ok(true); } _ => {} } @@ -409,7 +449,7 @@ impl ZoneMapIndex { /// Load the scalar index from storage async fn load( store: Arc, - fri: Option>, + fri: Option>, index_cache: &LanceCache, ) -> Result> where @@ -426,21 +466,34 @@ impl ZoneMapIndex { .get(ZONEMAP_SIZE_META_KEY) .and_then(|bs| bs.parse().ok()) .unwrap_or(ROWS_PER_ZONE_DEFAULT); + + let null_rows = if let Some(idx_str) = file_schema.metadata.get(NULL_BITMAP_META_KEY) { + let idx = idx_str.parse::().map_err(|e| { + Error::invalid_input(format!("invalid null bitmap buffer index: {e}")) + })?; + let bytes = index_file.read_global_buffer(idx).await?; + Some(RowAddrTreeMap::deserialize_from(bytes.as_ref())?) + } else { + None + }; + Ok(Arc::new(Self::try_from_serialized( zone_maps, store, fri, index_cache, rows_per_zone, + null_rows, )?)) } fn try_from_serialized( data: RecordBatch, store: Arc, - fri: Option>, + fri: Option>, index_cache: &LanceCache, rows_per_zone: u64, + null_rows: Option, ) -> Result { // The RecordBatch should have columns: min, max, null_count let min_col = data @@ -502,6 +555,7 @@ impl ZoneMapIndex { store, fri, index_cache: WeakLanceCache::from(index_cache), + null_rows, }); } @@ -533,6 +587,7 @@ impl ZoneMapIndex { store, fri, index_cache: WeakLanceCache::from(index_cache), + null_rows, }) } } @@ -583,11 +638,21 @@ impl ScalarIndex for ZoneMapIndex { metrics: &dyn MetricsCollector, ) -> Result { let query = query.as_any().downcast_ref::().unwrap(); + if let SargableQuery::IsNull() = query + && let Some(null_rows) = &self.null_rows + { + return Ok(SearchResult::exact(null_rows.clone())); + } + search_zones(&self.zones, metrics, |zone| { self.evaluate_zone_against_query(zone, query) }) } + fn results_are_row_addresses(&self) -> bool { + true + } + fn can_remap(&self) -> bool { false } @@ -595,7 +660,7 @@ impl ScalarIndex for ZoneMapIndex { /// Remap the row ids, creating a new remapped version of this index in `dest_store` async fn remap( &self, - _mapping: &HashMap>, + _mapping: &RowAddrRemap, _dest_store: &dyn IndexStore, ) -> Result { Err(Error::invalid_input_source( @@ -617,19 +682,30 @@ impl ScalarIndex for ZoneMapIndex { let options = ZoneMapIndexBuilderParams::new(self.rows_per_zone); let processor = ZoneMapProcessor::new(value_type.clone())?; let trainer = ZoneTrainer::new(processor, self.rows_per_zone)?; - let updated_zones = rebuild_zones(&self.zones, trainer, new_data).await?; + let (updated_zones, new_null_rows) = rebuild_zones(&self.zones, trainer, new_data).await?; + + // Merge existing and new null rows. If the existing index had no null bitmap + // (legacy format — null positions unknown), preserve that None: updating cannot + // recover the missing information, and claiming the result has zero nulls would + // be a false negative. Only a full retrain produces a fresh, complete bitmap. + let merged_null_rows = self.null_rows.as_ref().map(|existing| { + let mut merged = existing.clone(); + merged |= &new_null_rows; + merged + }); // Serialize the combined zones back into the index file let mut builder = ZoneMapIndexBuilder::try_new(options, self.data_type.clone())?; builder.options.rows_per_zone = self.rows_per_zone; builder.maps = updated_zones; - let file = builder.write_index(dest_store).await?; + builder.null_rows = merged_null_rows; + let files = builder.write_index(dest_store).await?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&pbold::ZoneMapIndexDetails::default()) .unwrap(), index_version: ZONEMAP_INDEX_VERSION, - files: vec![file], + files, }) } @@ -643,6 +719,12 @@ impl ScalarIndex for ZoneMapIndex { let params = serde_json::to_value(ZoneMapIndexBuilderParams::new(self.rows_per_zone))?; Ok(ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap).with_params(¶ms)) } + + /// Single-segment `[min, max]` folded from this index's zones; see + /// [`value_range_over`](Self::value_range_over) for the full contract. + fn value_range(&self) -> Option<(ScalarValue, ScalarValue)> { + Self::value_range_over([self]) + } } /// Merge caller-selected ZoneMap segments into one self-contained segment. @@ -658,6 +740,8 @@ pub async fn merge_zonemap_indices( let data_type = first.data_type.clone(); let mut zones = Vec::new(); + let mut merged_null_rows = RowAddrTreeMap::new(); + let mut any_missing_bitmap = false; for source in source_indices { if source.rows_per_zone != rows_per_zone { return Err(Error::invalid_input(format!( @@ -681,18 +765,29 @@ pub async fn merge_zonemap_indices( }) .cloned(), ); + match &source.null_rows { + Some(null_rows) => { + let mut filtered = null_rows.clone(); + filtered.retain_fragments(fragment_filter.iter()); + merged_null_rows |= &filtered; + } + None => any_missing_bitmap = true, + } } zones.sort_by_key(|zone| (zone.bound.fragment_id, zone.bound.start)); let mut builder = ZoneMapIndexBuilder::try_new(ZoneMapIndexBuilderParams::new(rows_per_zone), data_type)?; builder.maps = zones; - builder.write_index(dest_store).await?; + if !any_missing_bitmap { + builder.null_rows = Some(merged_null_rows); + } + let files = builder.write_index(dest_store).await?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&pbold::ZoneMapIndexDetails::default()).unwrap(), index_version: ZONEMAP_INDEX_VERSION, - files: dest_store.list_files_with_sizes().await?, + files, }) } @@ -737,6 +832,10 @@ pub struct ZoneMapIndexBuilder { items_type: DataType, maps: Vec, + // None means "legacy index — null positions unknown"; Some means a complete bitmap. + // write_index omits the null-bitmap global buffer when this is None, preserving the + // legacy format so that downstream searches remain conservative. + null_rows: Option, } impl ZoneMapIndexBuilder { @@ -745,6 +844,7 @@ impl ZoneMapIndexBuilder { options, items_type, maps: Vec::new(), + null_rows: None, }) } @@ -754,7 +854,9 @@ impl ZoneMapIndexBuilder { pub async fn train(&mut self, batches_source: SendableRecordBatchStream) -> Result<()> { let processor = ZoneMapProcessor::new(self.items_type.clone())?; let trainer = ZoneTrainer::new(processor, self.options.rows_per_zone)?; - self.maps = trainer.train(batches_source).await?; + let (maps, null_rows) = trainer.train(batches_source).await?; + self.maps = maps; + self.null_rows = Some(null_rows); Ok(()) } @@ -807,7 +909,7 @@ impl ZoneMapIndexBuilder { Ok(RecordBatch::try_new(schema, columns)?) } - pub async fn write_index(self, index_store: &dyn IndexStore) -> Result { + pub async fn write_index(self, index_store: &dyn IndexStore) -> Result> { let record_batch = self.zonemap_stats_as_batch()?; let mut file_schema = record_batch.schema().as_ref().clone(); @@ -820,7 +922,24 @@ impl ZoneMapIndexBuilder { .new_index_file(ZONEMAP_FILENAME, Arc::new(file_schema)) .await?; index_file.write_record_batch(record_batch).await?; - index_file.finish().await + + let zonemap_file = if let Some(null_rows) = self.null_rows { + let mut null_bitmap_bytes = Vec::with_capacity(null_rows.serialized_size()); + null_rows.serialize_into(&mut null_bitmap_bytes)?; + let null_bitmap_idx = index_file + .add_global_buffer(bytes::Bytes::from(null_bitmap_bytes)) + .await?; + index_file + .finish_with_metadata(HashMap::from([( + NULL_BITMAP_META_KEY.to_string(), + null_bitmap_idx.to_string(), + )])) + .await? + } else { + index_file.finish_with_metadata(HashMap::new()).await? + }; + + Ok(vec![zonemap_file]) } } @@ -925,8 +1044,7 @@ impl ZoneMapIndexPlugin { batches_source: SendableRecordBatchStream, index_store: &dyn IndexStore, options: Option, - ) -> Result { - // train_zonemap_index: calling scan_aligned_chunks + ) -> Result> { let value_type = batches_source.schema().field(0).data_type().clone(); let mut builder = ZoneMapIndexBuilder::try_new(options.unwrap_or_default(), value_type)?; @@ -961,11 +1079,7 @@ impl TrainingRequest for ZoneMapIndexTrainingRequest { } #[async_trait] -impl ScalarIndexPlugin for ZoneMapIndexPlugin { - fn name(&self) -> &str { - "ZoneMap" - } - +impl BasicTrainer for ZoneMapIndexPlugin { fn new_training_request( &self, params: &str, @@ -982,26 +1096,6 @@ impl ScalarIndexPlugin for ZoneMapIndexPlugin { Ok(Box::new(ZoneMapIndexTrainingRequest::new(params))) } - fn provides_exact_answer(&self) -> bool { - false - } - - fn version(&self) -> u32 { - ZONEMAP_INDEX_VERSION - } - - fn new_query_parser( - &self, - index_name: String, - _index_details: &prost_types::Any, - ) -> Option> { - Some(Box::new(SargableQueryParser::new( - index_name, - self.name().to_string(), - true, - ))) - } - async fn train_index( &self, data: SendableRecordBatchStream, @@ -1017,20 +1111,51 @@ impl ScalarIndexPlugin for ZoneMapIndexPlugin { "must provide training request created by new_training_request".into(), ) })?; - let file = Self::train_zonemap_index(data, index_store, Some(request.params)).await?; + let files = Self::train_zonemap_index(data, index_store, Some(request.params)).await?; Ok(CreatedIndex { index_details: prost_types::Any::from_msg(&pbold::ZoneMapIndexDetails::default()) .unwrap(), index_version: ZONEMAP_INDEX_VERSION, - files: vec![file], + files, }) } +} + +#[async_trait] +impl ScalarIndexPlugin for ZoneMapIndexPlugin { + fn basic_trainer(&self) -> Option<&dyn BasicTrainer> { + Some(self) + } + + fn name(&self) -> &str { + "ZoneMap" + } + + fn provides_exact_answer(&self) -> bool { + false + } + + fn version(&self) -> u32 { + ZONEMAP_INDEX_VERSION + } + + fn new_query_parser( + &self, + index_name: String, + _index_details: &prost_types::Any, + ) -> Option> { + Some(Box::new(SargableQueryParser::new( + index_name, + self.name().to_string(), + true, + ))) + } async fn load_index( &self, index_store: Arc, _index_details: &prost_types::Any, - frag_reuse_index: Option>, + frag_reuse_index: Option>, cache: &LanceCache, ) -> Result> { Ok(ZoneMapIndex::load(index_store, frag_reuse_index, cache).await? as Arc) @@ -1045,8 +1170,8 @@ mod tests { use crate::scalar::zoned::ZoneBound; use crate::scalar::zonemap::{ZoneMapIndexPlugin, ZoneMapStatistics}; - use arrow::datatypes::Float32Type; - use arrow_array::{Array, RecordBatch, UInt64Array, record_batch}; + use arrow::datatypes::{ArrowPrimitiveType, Float32Type, Int64Type}; + use arrow_array::{Array, PrimitiveArray, RecordBatch, UInt64Array, record_batch}; use arrow_schema::{DataType, Field, Schema}; use datafusion::execution::SendableRecordBatchStream; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; @@ -1061,13 +1186,14 @@ mod tests { use lance_datagen::ArrayGeneratorExt; use lance_datagen::{BatchCount, RowCount, array}; use lance_io::object_store::ObjectStore; - use lance_select::{NullableRowAddrSet, RowAddrTreeMap}; + use lance_select::RowAddrTreeMap; use crate::scalar::{ SargableQuery, ScalarIndex, SearchResult, lance_format::LanceIndexStore, zonemap::{ ZONEMAP_FILENAME, ZONEMAP_SIZE_META_KEY, ZoneMapIndex, ZoneMapIndexBuilderParams, + merge_zonemap_indices, }, }; @@ -1098,6 +1224,124 @@ mod tests { Box::pin(RecordBatchStreamAdapter::new(schema, stream)) } + /// Build a single-column ZoneMap of primitive type `T` from `fragments` + /// (one batch -> one fragment), with small zones, then load it back. + async fn train_and_load( + fragments: Vec>>, + ) -> Arc + where + PrimitiveArray: From>>, + { + let tmpdir = TempObjDir::default(); + let test_store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + let schema = Arc::new(Schema::new(vec![Field::new( + VALUE_COLUMN_NAME, + T::DATA_TYPE, + true, + )])); + let batches: Vec = fragments + .into_iter() + .map(|vals| { + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(PrimitiveArray::::from(vals))], + ) + .unwrap() + }) + .collect(); + let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + schema.clone(), + stream::iter(batches.into_iter().map(Ok)), + )); + let stream = add_row_addr(stream); + + ZoneMapIndexPlugin::train_zonemap_index( + stream, + test_store.as_ref(), + Some(ZoneMapIndexBuilderParams::new(2)), + ) + .await + .unwrap(); + + ZoneMapIndex::load(test_store.clone(), None, &LanceCache::no_cache()) + .await + .expect("Failed to load ZoneMapIndex") + } + + #[tokio::test] + async fn test_value_range_spans_fragments() { + // Two fragments, multiple zones each; global min/max straddle both. + let index = train_and_load::(vec![ + vec![Some(10), Some(50), Some(30)], + vec![Some(5), Some(99), Some(42)], + ]) + .await; + assert_eq!( + index.value_range(), + Some((ScalarValue::Int64(Some(5)), ScalarValue::Int64(Some(99)))) + ); + } + + #[tokio::test] + async fn test_value_range_all_null_is_none() { + let index = train_and_load::(vec![vec![None, None, None]]).await; + assert_eq!(index.value_range(), None); + } + + #[tokio::test] + async fn test_value_range_nan_max_is_none() { + // Zones of size 2: [1.0, 2.0] then [100.0, NaN]. The NaN-bearing zone hides + // its finite max (100.0), so the only sound answer is None. + let index = train_and_load::(vec![vec![ + Some(1.0), + Some(2.0), + Some(100.0), + Some(f32::NAN), + ]]) + .await; + assert_eq!(index.value_range(), None); + } + + #[tokio::test] + async fn test_value_range_over_folds_segments() { + // Two disjoint segments of one logical index; the global range straddles + // both (min and max from `b`), proving the fold spans segments. + let a = train_and_load::(vec![vec![Some(5), Some(9)]]).await; + let b = train_and_load::(vec![vec![Some(1), Some(20)]]).await; + assert_eq!( + ZoneMapIndex::value_range_over([a.as_ref(), b.as_ref()]), + Some((ScalarValue::Int64(Some(1)), ScalarValue::Int64(Some(20)))) + ); + } + + #[tokio::test] + async fn test_value_range_over_nan_in_any_segment_is_none() { + // NaN in one segment hides that segment's finite max; the cross-segment + // fold must bail to None just as the single-segment path does. + let a = train_and_load::(vec![vec![Some(1.0), Some(2.0)]]).await; + let b = train_and_load::(vec![vec![Some(100.0), Some(f32::NAN)]]).await; + assert_eq!( + ZoneMapIndex::value_range_over([a.as_ref(), b.as_ref()]), + None + ); + } + + #[tokio::test] + async fn test_value_range_over_skips_all_null_segment() { + // An all-null segment yields no finite zone; folding it with a finite + // segment returns the finite segment's range (null contributes nothing). + let a = train_and_load::(vec![vec![None, None]]).await; + let b = train_and_load::(vec![vec![Some(3), Some(7)]]).await; + assert_eq!( + ZoneMapIndex::value_range_over([a.as_ref(), b.as_ref()]), + Some((ScalarValue::Int64(Some(3)), ScalarValue::Int64(Some(7)))) + ); + } + #[tokio::test] async fn test_empty_zonemap_index() { let tmpdir = TempObjDir::default(); @@ -1502,7 +1746,7 @@ mod tests { // Test IsNull query (should match nothing since there are no null values) let query = SargableQuery::IsNull(); let result = index.search(&query, &NoOpMetricsCollector).await.unwrap(); - assert_eq!(result, SearchResult::AtMost(NullableRowAddrSet::empty())); + assert_eq!(result, SearchResult::exact(RowAddrTreeMap::new())); // Test range queries with NaN bounds // Range with NaN as start bound (included) @@ -1545,7 +1789,7 @@ mod tests { ); let result = index.search(&query, &NoOpMetricsCollector).await.unwrap(); // Should match nothing since nothing is greater than NaN - assert_eq!(result, SearchResult::AtMost(NullableRowAddrSet::empty())); + assert_eq!(result, SearchResult::at_most(RowAddrTreeMap::new())); // Test IsIn query with mixed float types (Float16, Float32, Float64) let query = SargableQuery::IsIn(vec![ @@ -1729,7 +1973,7 @@ mod tests { // 8. IsNull query (no nulls in data, should match nothing) let query = SargableQuery::IsNull(); let result = index.search(&query, &NoOpMetricsCollector).await.unwrap(); - assert_eq!(result, SearchResult::at_most(RowAddrTreeMap::new())); + assert_eq!(result, SearchResult::exact(RowAddrTreeMap::new())); // 9. IsIn query: [0, 100, 101, 50] let query = SargableQuery::IsIn(vec![ ScalarValue::Int32(Some(0)), @@ -2438,6 +2682,7 @@ mod tests { store: test_store, fri: None, index_cache: WeakLanceCache::from(&LanceCache::no_cache()), + null_rows: None, }; // Test LikePrefix query for "foo" @@ -2510,6 +2755,7 @@ mod tests { store: test_store, fri: None, index_cache: WeakLanceCache::from(&LanceCache::no_cache()), + null_rows: None, }; // Test LikePrefix "test" @@ -2577,6 +2823,7 @@ mod tests { store: test_store, fri: None, index_cache: WeakLanceCache::from(&LanceCache::no_cache()), + null_rows: None, }; // Test LikePrefix with LargeUtf8 @@ -2629,4 +2876,204 @@ mod tests { // All max characters assert_eq!(compute_next_prefix("\u{10FFFF}\u{10FFFF}"), None); } + + // When merging zone map segments, if ANY source segment has null_rows = None + // (legacy — null positions unknown), the merged result must also be None. + // The bug: any_null_bitmap is set to true as soon as one source has Some(...), + // and the None sources are silently skipped. The merged index then has + // null_rows = Some(partial_bitmap), so an IsNull search returns exact results + // that only cover the modern segment's nulls — a false negative for the legacy + // segment whose null positions were never tracked. + #[tokio::test] + async fn test_merge_with_legacy_none_segment_not_treated_as_no_nulls() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + use arrow_array::{Int32Array, UInt32Array}; + + // Index A: fragment 0, modern — has a complete null bitmap with 2 known null rows. + let schema_a = Arc::new(Schema::new(vec![ + Field::new("min", DataType::Int32, true), + Field::new("max", DataType::Int32, true), + Field::new("null_count", DataType::UInt32, false), + Field::new("nan_count", DataType::UInt32, false), + Field::new("fragment_id", DataType::UInt64, false), + Field::new("zone_start", DataType::UInt64, false), + Field::new("zone_length", DataType::UInt64, false), + ])); + let batch_a = RecordBatch::try_new( + schema_a, + vec![ + Arc::new(Int32Array::from(vec![Some(1i32)])) as _, + Arc::new(Int32Array::from(vec![Some(5i32)])) as _, + Arc::new(UInt32Array::from(vec![2u32])) as _, + Arc::new(UInt32Array::from(vec![0u32])) as _, + Arc::new(UInt64Array::from(vec![0u64])) as _, + Arc::new(UInt64Array::from(vec![0u64])) as _, + Arc::new(UInt64Array::from(vec![10u64])) as _, + ], + ) + .unwrap(); + let mut modern_null_rows = RowAddrTreeMap::new(); + modern_null_rows.insert(3); // frag 0 row 3 + modern_null_rows.insert(7); // frag 0 row 7 + let cache = LanceCache::no_cache(); + let index_a = Arc::new( + ZoneMapIndex::try_from_serialized( + batch_a, + store.clone(), + None, + &cache, + 10, + Some(modern_null_rows), // modern: complete bitmap + ) + .unwrap(), + ); + + // Index B: fragment 1, legacy — null_rows = None despite null_count = 3. + let schema_b = Arc::new(Schema::new(vec![ + Field::new("min", DataType::Int32, true), + Field::new("max", DataType::Int32, true), + Field::new("null_count", DataType::UInt32, false), + Field::new("nan_count", DataType::UInt32, false), + Field::new("fragment_id", DataType::UInt64, false), + Field::new("zone_start", DataType::UInt64, false), + Field::new("zone_length", DataType::UInt64, false), + ])); + let batch_b = RecordBatch::try_new( + schema_b, + vec![ + Arc::new(Int32Array::from(vec![Some(10i32)])) as _, + Arc::new(Int32Array::from(vec![Some(20i32)])) as _, + Arc::new(UInt32Array::from(vec![3u32])) as _, + Arc::new(UInt32Array::from(vec![0u32])) as _, + Arc::new(UInt64Array::from(vec![1u64])) as _, + Arc::new(UInt64Array::from(vec![0u64])) as _, + Arc::new(UInt64Array::from(vec![10u64])) as _, + ], + ) + .unwrap(); + let index_b = Arc::new( + ZoneMapIndex::try_from_serialized( + batch_b, + store.clone(), + None, + &cache, + 10, + None, // legacy: null positions unknown + ) + .unwrap(), + ); + + let dest_tmpdir = TempObjDir::default(); + let dest_store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + dest_tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + let all_frags = RoaringBitmap::from_iter([0u32, 1]); + merge_zonemap_indices( + &[index_a.as_ref(), index_b.as_ref()], + dest_store.as_ref(), + &all_frags, + ) + .await + .unwrap(); + + let merged = ZoneMapIndex::load(dest_store.clone(), None, &LanceCache::no_cache()) + .await + .unwrap(); + + // Index B had null_rows = None, so the merged index cannot know all null positions. + // IsNull must NOT return exact — that would be a false negative for fragment 1's nulls. + let result = merged + .search(&SargableQuery::IsNull(), &NoOpMetricsCollector) + .await + .unwrap(); + + // With the bug: any_null_bitmap=true (from A) → null_rows=Some(A's bitmap only) + // → IsNull returns exact, missing B's unknown nulls ← FALSE NEGATIVE + // With the fix: any_null_bitmap=false (because B is None) → null_rows=None + // → IsNull falls through to zone scan → AtMost + assert!( + !result.is_exact(), + "IsNull on a merged index where one source had null_rows=None must not return \ + exact; the legacy segment had null_count=3 so its nulls exist at unknown positions" + ); + } + + // Writes a zonemap file in the legacy format (no null bitmap global buffer), + // simulating an index created before the null bitmap feature was added. + async fn write_legacy_zonemap(store: &dyn IndexStore, null_count: u32) { + use arrow_array::{Int32Array, UInt32Array}; + let schema = Arc::new(Schema::new(vec![ + Field::new("min", DataType::Int32, true), + Field::new("max", DataType::Int32, true), + Field::new("null_count", DataType::UInt32, false), + Field::new("nan_count", DataType::UInt32, false), + Field::new("fragment_id", DataType::UInt64, false), + Field::new("zone_start", DataType::UInt64, false), + Field::new("zone_length", DataType::UInt64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![Some(0)])) as _, + Arc::new(Int32Array::from(vec![Some(99)])) as _, + Arc::new(UInt32Array::from(vec![null_count])) as _, + Arc::new(UInt32Array::from(vec![0u32])) as _, + Arc::new(UInt64Array::from(vec![0u64])) as _, + Arc::new(UInt64Array::from(vec![0u64])) as _, + Arc::new(UInt64Array::from(vec![100u64])) as _, + ], + ) + .unwrap(); + let mut file_schema = schema.as_ref().clone(); + file_schema + .metadata + .insert(ZONEMAP_SIZE_META_KEY.to_string(), "8192".to_string()); + let mut writer = store + .new_index_file(ZONEMAP_FILENAME, Arc::new(file_schema)) + .await + .unwrap(); + writer.write_record_batch(batch).await.unwrap(); + writer.finish().await.unwrap(); + } + + #[tokio::test] + async fn test_legacy_zonemap_no_null_bitmap() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + // Write a legacy index with one zone that has nulls but no null bitmap. + write_legacy_zonemap(store.as_ref(), 10).await; + + let index = ZoneMapIndex::load(store, None, &LanceCache::no_cache()) + .await + .expect("failed to load legacy zonemap"); + + assert!( + index.null_rows.is_none(), + "legacy index should have no null bitmap" + ); + + // IS NULL should fall back to the zone-scan path and return AtMost, not Exact. + let result = index + .search(&SargableQuery::IsNull(), &NoOpMetricsCollector) + .await + .unwrap(); + assert!( + !result.is_exact(), + "IS NULL on a legacy index should not be exact" + ); + } } diff --git a/rust/lance-index/src/vector.rs b/rust/lance-index/src/vector.rs index 3c5a6601a8a..76b9ccc7a51 100644 --- a/rust/lance-index/src/vector.rs +++ b/rust/lance-index/src/vector.rs @@ -4,9 +4,10 @@ //! Vector Index //! +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::any::Any; use std::fmt::Debug; -use std::{collections::HashMap, sync::Arc}; +use std::sync::Arc; use arrow_array::{ArrayRef, Float32Array, RecordBatch, UInt32Array}; use arrow_schema::Field; @@ -408,7 +409,7 @@ pub trait VectorIndex: Send + Sync + std::fmt::Debug + Index { /// /// If an old row id is not in the mapping then it should be /// left alone. - async fn remap(&mut self, mapping: &HashMap>) -> Result<()>; + async fn remap(&mut self, mapping: &RowAddrRemap) -> Result<()>; /// The metric type of this vector index. fn metric_type(&self) -> DistanceType; diff --git a/rust/lance-index/src/vector/bq/storage.rs b/rust/lance-index/src/vector/bq/storage.rs index 2f4fe69792a..72fa8d4b056 100644 --- a/rust/lance-index/src/vector/bq/storage.rs +++ b/rust/lance-index/src/vector/bq/storage.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::borrow::Cow; use std::collections::{BinaryHeap, HashMap}; use std::ops::Sub; @@ -2534,7 +2535,7 @@ impl QuantizerStorage for RabitQuantizationStorage { }; match build_frag_reuse_mapping(fri.as_deref(), &storage.row_ids) { - Some(mapping) => storage.remap(&mapping), + Some(mapping) => storage.remap(&RowAddrRemap::direct(mapping)), None => Ok(storage), } } @@ -2555,7 +2556,7 @@ impl QuantizerStorage for RabitQuantizationStorage { Self::try_from_batch(batch, metadata, distance_type, frag_reuse_index) } - fn remap(&self, mapping: &HashMap>) -> Result { + fn remap(&self, mapping: &RowAddrRemap) -> Result { let num_vectors = self.codes.len(); let num_code_bytes = self.codes.value_length() as usize; let codes = self.codes.values().as_primitive::().values(); @@ -2565,10 +2566,10 @@ impl QuantizerStorage for RabitQuantizationStorage { let row_ids = self.row_ids.values(); for (i, row_id) in row_ids.iter().enumerate() { - match mapping.get(row_id) { + match mapping.get(*row_id) { Some(Some(new_id)) => { indices.push(i as u32); - new_row_ids.push(*new_id); + new_row_ids.push(new_id); new_codes.extend(get_rq_code(codes, i, num_vectors, num_code_bytes)); } Some(None) => {} @@ -2743,6 +2744,7 @@ fn get_rq_code( #[cfg(test)] mod tests { use super::*; + use rstest::rstest; use std::collections::{BinaryHeap, HashMap}; use arrow_array::{ArrayRef, Float32Array, Float64Array, UInt64Array}; @@ -2750,7 +2752,6 @@ mod tests { use lance_linalg::distance::DistanceType; use rand::rngs::SmallRng; use rand::{Rng, SeedableRng}; - use rstest::rstest; use crate::vector::bq::{RQRotationType, builder::RabitQuantizer}; use crate::vector::quantizer::{Quantization, QuantizerStorage}; @@ -3929,7 +3930,7 @@ mod tests { .into_iter() .map(|(id, dist)| (id as u64, dist)) .collect::>(); - expected.sort_by(|left, right| left.0.cmp(&right.0)); + expected.sort_by_key(|left| left.0); let mut heap = BinaryHeap::with_capacity(k); let mut distances = Vec::new(); @@ -3951,7 +3952,7 @@ mod tests { .into_iter() .map(|node| (node.id, node.dist.0)) .collect::>(); - actual.sort_by(|left, right| left.0.cmp(&right.0)); + actual.sort_by_key(|left| left.0); assert_eq!(actual.len(), expected.len()); for ((actual_id, actual_dist), (expected_id, expected_dist)) in @@ -4425,7 +4426,7 @@ mod tests { mapping.insert(3, None); mapping.insert(4, Some(104)); - let remapped = storage.remap(&mapping).unwrap(); + let remapped = storage.remap(&RowAddrRemap::direct(mapping)).unwrap(); assert!(remapped.metadata().packed); let remapped_batch = remapped.to_batches().unwrap().next().unwrap(); @@ -4470,7 +4471,7 @@ mod tests { mapping.insert(3, None); mapping.insert(4, Some(104)); - let remapped = storage.remap(&mapping).unwrap(); + let remapped = storage.remap(&RowAddrRemap::direct(mapping)).unwrap(); let remapped_batch = remapped.to_batches().unwrap().next().unwrap(); let remapped_row_ids = remapped_batch[ROW_ID].as_primitive::().values(); let expected_row_ids = UInt64Array::from_iter_values( diff --git a/rust/lance-index/src/vector/flat/index.rs b/rust/lance-index/src/vector/flat/index.rs index 2c099eb5435..89c28a6ffd7 100644 --- a/rust/lance-index/src/vector/flat/index.rs +++ b/rust/lance-index/src/vector/flat/index.rs @@ -4,7 +4,8 @@ //! Flat Vector Index. //! -use std::collections::{BinaryHeap, HashMap}; +use lance_core::utils::row_addr_remap::RowAddrRemap; +use std::collections::BinaryHeap; use std::sync::Arc; use arrow::array::AsArray; @@ -312,7 +313,7 @@ impl IvfSubIndex for FlatIndex { Ok(Self {}) } - fn remap(&self, _: &HashMap>, _: &impl VectorStore) -> Result { + fn remap(&self, _: &RowAddrRemap, _: &impl VectorStore) -> Result { Ok(self.clone()) } @@ -569,7 +570,7 @@ mod tests { .zip(dists.values().iter()) .map(|(row_id, dist)| (*row_id, *dist)) .collect::>(); - results.sort_by(|left, right| left.0.cmp(&right.0)); + results.sort_by_key(|left| left.0); results } @@ -578,7 +579,7 @@ mod tests { .into_iter() .map(|node| (node.id, node.dist.0)) .collect::>(); - results.sort_by(|left, right| left.0.cmp(&right.0)); + results.sort_by_key(|left| left.0); results } diff --git a/rust/lance-index/src/vector/flat/storage.rs b/rust/lance-index/src/vector/flat/storage.rs index c3ec30d5086..ff8756cbe41 100644 --- a/rust/lance-index/src/vector/flat/storage.rs +++ b/rust/lance-index/src/vector/flat/storage.rs @@ -134,7 +134,7 @@ impl VectorStore for FlatFloatStorage { fn append_batch(&self, batch: RecordBatch, _vector_column: &str) -> Result { // TODO: use chunked storage - let new_batch = concat_batches(&batch.schema(), vec![&self.batch, &batch].into_iter())?; + let new_batch = concat_batches(&batch.schema(), vec![&self.batch, &batch])?; let mut storage = self.clone(); storage.row_ids = Arc::new( new_batch @@ -296,7 +296,7 @@ impl VectorStore for FlatBinStorage { fn append_batch(&self, batch: RecordBatch, _vector_column: &str) -> Result { // TODO: use chunked storage - let new_batch = concat_batches(&batch.schema(), vec![&self.batch, &batch].into_iter())?; + let new_batch = concat_batches(&batch.schema(), vec![&self.batch, &batch])?; let mut storage = self.clone(); storage.row_ids = Arc::new( new_batch diff --git a/rust/lance-index/src/vector/hnsw/builder.rs b/rust/lance-index/src/vector/hnsw/builder.rs index 214750dfafa..cc1ac1abf81 100644 --- a/rust/lance-index/src/vector/hnsw/builder.rs +++ b/rust/lance-index/src/vector/hnsw/builder.rs @@ -10,6 +10,7 @@ use arrow_array::{ArrayRef, Float32Array, ListArray, RecordBatch, UInt64Array}; use crossbeam_queue::ArrayQueue; use itertools::Itertools; use lance_core::deepsize::DeepSizeOf; +use lance_core::utils::row_addr_remap::RowAddrRemap; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_linalg::distance::DistanceType; @@ -1222,7 +1223,7 @@ impl IvfSubIndex for HNSW { fn remap( &self, - _mapping: &HashMap>, // we don't need the mapping here because we rebuild the graph from remapped storage + _mapping: &RowAddrRemap, // we don't need the mapping here because we rebuild the graph from remapped storage store: &impl VectorStore, ) -> Result { // We can't simply remap the row ids in the graph because the vectors are changed, @@ -1239,17 +1240,23 @@ impl IvfSubIndex for HNSW { // batch verbatim, re-stamped with up-to-date HNSW metadata. // The IVF partition cache re-serializes loaded indices through // here (`ivf/partition_serde.rs`), so this must round-trip. + // + // Merge into (not replace) the existing schema metadata: a + // disk-loaded batch inherits other keys from the index file + // schema (e.g. `INDEX_METADATA_SCHEMA_KEY`, `IVF_METADATA_KEY`), + // and `RecordBatch::with_schema` requires the new metadata to be + // a superset of the current one. Dropping those keys here makes + // the new schema a non-superset and fails the round-trip with + // "target schema is not superset of current schema". let metadata = serde_json::to_string(&self.metadata())?; - let schema = - graph - .batch - .schema() - .as_ref() - .clone() - .with_metadata(HashMap::from_iter(vec![( - HNSW_METADATA_KEY.to_string(), - metadata, - )])); + let mut schema_metadata = graph.batch.schema_ref().metadata().clone(); + schema_metadata.insert(HNSW_METADATA_KEY.to_string(), metadata); + let schema = graph + .batch + .schema() + .as_ref() + .clone() + .with_metadata(schema_metadata); return Ok(graph.batch.clone().with_schema(Arc::new(schema))?); } }; @@ -1756,6 +1763,69 @@ mod tests { assert_eq!(a, b); } + /// Regression for the IVF partition-cache round-trip: a disk-loaded batch + /// inherits extra schema metadata keys from the index file (e.g. + /// `INDEX_METADATA_SCHEMA_KEY`, `IVF_METADATA_KEY`). `to_batch()` on the + /// loaded graph must *merge* the HNSW key into that metadata rather than + /// replacing it -- otherwise the new schema is not a superset of the + /// current one and `RecordBatch::with_schema` fails with "target schema is + /// not superset of current schema". + #[tokio::test] + async fn test_to_batch_loaded_preserves_extra_schema_metadata() { + use super::HNSW_METADATA_KEY; + + const DIM: usize = 24; + const TOTAL: usize = 512; + let fsl = + FixedSizeListArray::try_new_from_values(generate_random_array(TOTAL * DIM), DIM as i32) + .unwrap(); + let store = Arc::new(FlatFloatStorage::new(fsl, DistanceType::L2)); + let builder = HNSW::index_vectors( + store.as_ref(), + HnswBuildParams::default().num_edges(16).ef_construction(50), + ) + .unwrap(); + + // Simulate the disk-load path (`ivf/v2.rs::load_partition_entry`): + // the batch reaching `HNSW::load` carries the index file's schema + // metadata in addition to the HNSW key. + let built_batch = builder.to_batch().unwrap(); + let mut metadata = built_batch.schema_ref().metadata().clone(); + metadata.insert( + "lance:index_metadata".to_string(), + "{\"distance_type\":\"l2\"}".to_string(), + ); + metadata.insert("lance:ivf".to_string(), "42".to_string()); + let schema = built_batch + .schema() + .as_ref() + .clone() + .with_metadata(metadata); + let batch_with_extra = + RecordBatch::try_new(Arc::new(schema), built_batch.columns().to_vec()).unwrap(); + + let loaded = HNSW::load(batch_with_extra).unwrap(); + assert!(matches!(loaded.inner.graph, HnswGraph::Loaded(_))); + + // Before the fix this fails: the re-stamped schema dropped the extra + // keys, so `with_schema`'s superset check rejected the round-trip. + let out = loaded.to_batch().unwrap(); + let out_metadata = out.schema_ref().metadata(); + assert!(out_metadata.contains_key(HNSW_METADATA_KEY)); + assert_eq!( + out_metadata.get("lance:index_metadata").map(String::as_str), + Some("{\"distance_type\":\"l2\"}"), + ); + assert_eq!( + out_metadata.get("lance:ivf").map(String::as_str), + Some("42") + ); + + // The HNSW key must still decode to valid metadata after the merge. + let reloaded = HNSW::load(out).unwrap(); + assert_eq!(reloaded.len(), loaded.len()); + } + /// The loaded graph shares the Arrow batch and reconstructs no per-node /// `Vec` / `Vec`, so it is strictly /// lighter than the in-memory build representation. diff --git a/rust/lance-index/src/vector/hnsw/index.rs b/rust/lance-index/src/vector/hnsw/index.rs index c8c9e5164fe..dca9ea78955 100644 --- a/rust/lance-index/src/vector/hnsw/index.rs +++ b/rust/lance-index/src/vector/hnsw/index.rs @@ -1,9 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::{ any::Any, - collections::HashMap, fmt::{Debug, Formatter}, sync::Arc, }; @@ -307,7 +307,7 @@ impl VectorIndex for HNSWIndex { Box::new(self.storage.as_ref().unwrap().row_ids()) } - async fn remap(&mut self, _mapping: &HashMap>) -> Result<()> { + async fn remap(&mut self, _mapping: &RowAddrRemap) -> Result<()> { Err(Error::index( "Remapping HNSW in this way not supported".to_string(), )) diff --git a/rust/lance-index/src/vector/pq/storage.rs b/rust/lance-index/src/vector/pq/storage.rs index de5a7ac28bd..4e98df9906d 100644 --- a/rust/lance-index/src/vector/pq/storage.rs +++ b/rust/lance-index/src/vector/pq/storage.rs @@ -5,9 +5,9 @@ //! //! Used as storage backend for Graph based algorithms. +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::{ cmp::min, - collections::HashMap, sync::{Arc, OnceLock}, }; @@ -550,16 +550,16 @@ impl QuantizerStorage for ProductQuantizationStorage { // we can't use the default implementation of remap, // because PQ Storage transposed the PQ codes - fn remap(&self, mapping: &HashMap>) -> Result { + fn remap(&self, mapping: &RowAddrRemap) -> Result { let transposed_codes = self.pq_code.values(); let mut new_row_ids = Vec::with_capacity(self.len()); let mut new_codes = Vec::with_capacity(self.len() * self.metadata.num_sub_vectors); let row_ids = self.row_ids.values(); for (i, row_id) in row_ids.iter().enumerate() { - match mapping.get(row_id) { + match mapping.get(*row_id) { Some(Some(new_id)) => { - new_row_ids.push(*new_id); + new_row_ids.push(new_id); new_codes.extend(get_pq_code( transposed_codes, self.metadata.nbits, @@ -1154,6 +1154,7 @@ mod tests { use lance_arrow::FixedSizeListArrayExt; use lance_core::ROW_ID_FIELD; use rand::Rng; + use rstest::rstest; const DIM: usize = 32; const TOTAL: usize = 512; @@ -1300,21 +1301,40 @@ mod tests { assert!((storage.dist_between(u, v) - expected).abs() < 1e-4); } + // The first half of the rows is rewritten in order into frag 1; the second + // half is deleted. remap must behave the same in either RowAddrRemap mode. + fn pq_remap_compact() -> RowAddrRemap { + use lance_core::utils::row_addr_remap::GroupInput; + use roaring::RoaringTreemap; + RowAddrRemap::compact([GroupInput { + rewritten_old_row_addrs: RoaringTreemap::from_iter((0..TOTAL / 2).map(|i| i as u64)), + old_frag_ids: vec![0], + new_frags: vec![(1, (TOTAL / 2) as u32)], + }]) + .unwrap() + } + + fn pq_remap_explicit() -> RowAddrRemap { + RowAddrRemap::direct( + (0..TOTAL / 2) + .map(|i| (i as u64, Some((1u64 << 32) | i as u64))) + .chain((TOTAL / 2..TOTAL).map(|i| (i as u64, None))) + .collect(), + ) + } + + #[rstest] + #[case(pq_remap_compact())] + #[case(pq_remap_explicit())] #[tokio::test] - async fn test_remap_with_extra_column() { + async fn test_remap_with_extra_column(#[case] remap: RowAddrRemap) { let storage = create_pq_storage_with_extra_column().await; - let mut mapping = HashMap::new(); - for i in 0..TOTAL / 2 { - mapping.insert(i as u64, Some((TOTAL + i) as u64)); - } - for i in TOTAL / 2..TOTAL { - mapping.insert(i as u64, None); - } - let new_storage = storage.remap(&mapping).unwrap(); + let new_storage = storage.remap(&remap).unwrap(); assert_eq!(new_storage.len(), TOTAL / 2); assert_eq!(new_storage.row_ids.len(), TOTAL / 2); for (i, row_id) in new_storage.row_ids().enumerate() { - assert_eq!(*row_id, (TOTAL + i) as u64); + // Rewritten row i lands at offset i of frag 1. + assert_eq!(*row_id, (1u64 << 32) | i as u64); } assert_eq!(new_storage.batch.num_columns(), 2); assert!(new_storage.batch.column_by_name(ROW_ID).is_some()); diff --git a/rust/lance-index/src/vector/quantizer.rs b/rust/lance-index/src/vector/quantizer.rs index 8ee64669f32..6f4b191098b 100644 --- a/rust/lance-index/src/vector/quantizer.rs +++ b/rust/lance-index/src/vector/quantizer.rs @@ -2,9 +2,10 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use core::fmt; +use lance_core::utils::row_addr_remap::RowAddrRemap; +use std::fmt::Debug; use std::str::FromStr; use std::sync::Arc; -use std::{collections::HashMap, fmt::Debug}; use arrow::{array::AsArray, compute::concat_batches, datatypes::UInt64Type}; use arrow_array::{Array, ArrayRef, FixedSizeListArray, RecordBatch, UInt32Array, UInt64Array}; @@ -240,7 +241,7 @@ pub trait QuantizerStorage: Clone + Sized + DeepSizeOf + VectorStore { fn metadata(&self) -> &Self::Metadata; - fn remap(&self, mapping: &HashMap>) -> Result { + fn remap(&self, mapping: &RowAddrRemap) -> Result { let batches = self .to_batches()? .map(|b| { @@ -249,10 +250,10 @@ pub trait QuantizerStorage: Clone + Sized + DeepSizeOf + VectorStore { let row_ids = b.column(0).as_primitive::().values(); for (i, row_id) in row_ids.iter().enumerate() { - match mapping.get(row_id) { + match mapping.get(*row_id) { Some(Some(new_id)) => { indices.push(i as u32); - new_row_ids.push(*new_id); + new_row_ids.push(new_id); } Some(None) => {} None => { diff --git a/rust/lance-index/src/vector/sq/storage.rs b/rust/lance-index/src/vector/sq/storage.rs index 1e5eebda0d9..785aadc289d 100644 --- a/rust/lance-index/src/vector/sq/storage.rs +++ b/rust/lance-index/src/vector/sq/storage.rs @@ -12,12 +12,14 @@ use arrow_array::{ }; use arrow_schema::{DataType, SchemaRef}; use async_trait::async_trait; +use lance_arrow::ArrowFloatType; use lance_core::deepsize::DeepSizeOf; use lance_core::{Error, ROW_ID, Result}; use lance_file::previous::reader::FileReader as PreviousFileReader; use lance_io::object_store::ObjectStore; -use lance_linalg::distance::{DistanceType, dot_distance, l2_u8::l2_u8}; +use lance_linalg::distance::{DistanceType, dot_u8::dot_u8, l2_u8::l2_u8}; use lance_table::format::SelfDescribingFileReader; +use num_traits::AsPrimitive; use object_store::path::Path; use serde::{Deserialize, Serialize}; use std::sync::Arc; @@ -29,7 +31,7 @@ use crate::{ vector::{ SQ_CODE_COLUMN, quantizer::{QuantizerMetadata, QuantizerStorage}, - storage::{DistCalculator, VectorStore}, + storage::{DistCalculator, DistanceCalculatorOptions, QueryResidual, VectorStore}, transform::Transformer, }, }; @@ -370,30 +372,118 @@ impl VectorStore for ScalarQuantizationStorage { SQDistCalculator::new(query, self, self.quantizer.bounds()) } + fn dist_calculator_with_scratch<'a>( + &'a self, + query: ArrayRef, + _dist_q_c: f32, + _residual: Option>, + f32_scratch: &'a mut Vec, + _options: DistanceCalculatorOptions, + ) -> Self::DistanceCalculator<'a> { + SQDistCalculator::new_with_scratch(query, self, self.quantizer.bounds(), f32_scratch) + } + fn dist_calculator_from_id(&self, id: u32) -> Self::DistanceCalculator<'_> { let (offset, chunk) = self.chunk(id); let query_sq_code = chunk.sq_code_slice(id - offset); let bounds = self.quantizer.bounds(); + let lower_bound = bounds.start as f32; + let value_scale = sq_value_scale(&bounds); + let query_dot = match self.distance_type { + DistanceType::Dot => Some(SQDotQuery::from_sq_code(query_sq_code)), + _ => None, + }; SQDistCalculator { query_sq_code: SQQueryCode::Borrowed(query_sq_code), + query_dot, scale: sq_distance_scale(&bounds), + lower_bound, + value_scale, storage: self, } } } +#[inline] +fn sq_value_scale(bounds: &Range) -> f32 { + (bounds.end - bounds.start) as f32 / 255.0_f32 +} + #[inline] fn sq_distance_scale(bounds: &Range) -> f32 { - let range = (bounds.end - bounds.start) as f32; - (range * range) / (255.0_f32 * 255.0_f32) + let scale = sq_value_scale(bounds); + scale * scale } pub struct SQDistCalculator<'a> { query_sq_code: SQQueryCode<'a>, + query_dot: Option>, scale: f32, + lower_bound: f32, + value_scale: f32, storage: &'a ScalarQuantizationStorage, } +enum SQDotQuery<'a> { + Values { values: SQFloatQuery<'a>, sum: f32 }, + SqCode { code: &'a [u8], sum: f32 }, +} + +enum SQFloatQuery<'a> { + Borrowed(&'a [f32]), + Owned(Vec), +} + +impl SQFloatQuery<'_> { + fn as_slice(&self) -> &[f32] { + match self { + Self::Borrowed(values) => values, + Self::Owned(values) => values, + } + } +} + +impl<'a> SQDotQuery<'a> { + fn from_values(values: &[T::Native]) -> Self + where + T::Native: AsPrimitive, + { + let values: Vec<_> = values.iter().map(|v| v.as_()).collect(); + let sum = values.iter().sum(); + Self::Values { + values: SQFloatQuery::Owned(values), + sum, + } + } + + fn from_values_with_scratch( + values: &[T::Native], + scratch: &'a mut Vec, + ) -> Self + where + T::Native: AsPrimitive, + { + scratch.clear(); + scratch.extend(values.iter().map(|v| v.as_())); + let sum = scratch.iter().sum(); + Self::Values { + values: SQFloatQuery::Borrowed(scratch.as_slice()), + sum, + } + } + + fn from_sq_code(sq_code: &'a [u8]) -> Self { + Self::SqCode { + code: sq_code, + sum: sq_code_sum(sq_code), + } + } +} + +fn sq_code_sum(sq_code: &[u8]) -> f32 { + sq_code.iter().map(|code| *code as u32).sum::() as f32 +} + enum SQQueryCode<'a> { Borrowed(&'a [u8]), Owned(Vec), @@ -415,26 +505,127 @@ impl<'a> SQDistCalculator<'a> { // since we search 10s-100s of partitions, we can afford the overhead // this could be annoying at indexing time for HNSW, which requires constructing the // dist calculator frequently. However, HNSW isn't first-class citizen in Lance yet. so be it. - let query_sq_code = match query.data_type() { - DataType::Float16 => { - scale_to_u8::(query.as_primitive::().values(), &bounds) - } - DataType::Float32 => { - scale_to_u8::(query.as_primitive::().values(), &bounds) + let (query_sq_code, query_dot) = match storage.distance_type { + DistanceType::Dot => { + let query_dot = match query.data_type() { + DataType::Float16 => SQDotQuery::from_values::( + query.as_primitive::().values(), + ), + DataType::Float32 => SQDotQuery::from_values::( + query.as_primitive::().values(), + ), + DataType::Float64 => SQDotQuery::from_values::( + query.as_primitive::().values(), + ), + _ => { + panic!("Unsupported data type for ScalarQuantizationStorage"); + } + }; + (SQQueryCode::Owned(Vec::new()), Some(query_dot)) } - DataType::Float64 => { - scale_to_u8::(query.as_primitive::().values(), &bounds) + DistanceType::L2 | DistanceType::Cosine => { + let query_sq_code = match query.data_type() { + DataType::Float16 => scale_to_u8::( + query.as_primitive::().values(), + &bounds, + ), + DataType::Float32 => scale_to_u8::( + query.as_primitive::().values(), + &bounds, + ), + DataType::Float64 => scale_to_u8::( + query.as_primitive::().values(), + &bounds, + ), + _ => { + panic!("Unsupported data type for ScalarQuantizationStorage"); + } + }; + (SQQueryCode::Owned(query_sq_code), None) } + _ => panic!("We should not reach here: sq distance can only be L2 or Dot"), + }; + let lower_bound = bounds.start as f32; + let value_scale = sq_value_scale(&bounds); + Self { + query_sq_code, + query_dot, + scale: sq_distance_scale(&bounds), + lower_bound, + value_scale, + storage, + } + } + + fn new_with_scratch( + query: ArrayRef, + storage: &'a ScalarQuantizationStorage, + bounds: Range, + f32_scratch: &'a mut Vec, + ) -> Self { + if storage.distance_type != DistanceType::Dot { + return Self::new(query, storage, bounds); + } + + let query_dot = match query.data_type() { + DataType::Float16 => SQDotQuery::from_values_with_scratch::( + query.as_primitive::().values(), + f32_scratch, + ), + DataType::Float32 => SQDotQuery::from_values_with_scratch::( + query.as_primitive::().values(), + f32_scratch, + ), + DataType::Float64 => SQDotQuery::from_values_with_scratch::( + query.as_primitive::().values(), + f32_scratch, + ), _ => { panic!("Unsupported data type for ScalarQuantizationStorage"); } }; + let lower_bound = bounds.start as f32; + let value_scale = sq_value_scale(&bounds); Self { - query_sq_code: SQQueryCode::Owned(query_sq_code), + query_sq_code: SQQueryCode::Owned(Vec::new()), + query_dot: Some(query_dot), scale: sq_distance_scale(&bounds), + lower_bound, + value_scale, storage, } } + + fn dot_distance(&self, sq_code: &[u8]) -> f32 { + let query = self + .query_dot + .as_ref() + .expect("SQ dot distance requires a dot query"); + let dot = match query { + SQDotQuery::Values { values, sum } => { + let values = values.as_slice(); + self.lower_bound * *sum + + self.value_scale + * sq_code + .iter() + .zip(values.iter()) + .map(|(code, query_value)| *code as f32 * *query_value) + .sum::() + } + SQDotQuery::SqCode { + code: query_sq_code, + sum: query_code_sum, + } => { + let dim = sq_code.len() as f32; + let code_dot = dot_u8(sq_code, query_sq_code) as f32; + let code_sum = sq_code_sum(sq_code); + dim * self.lower_bound * self.lower_bound + + self.lower_bound * self.value_scale * (code_sum + *query_code_sum) + + self.scale * code_dot + } + }; + 1.0 - dot + } } impl DistCalculator for SQDistCalculator<'_> { @@ -442,12 +633,13 @@ impl DistCalculator for SQDistCalculator<'_> { let (offset, chunk) = self.storage.chunk(id); let sq_code = chunk.sq_code_slice(id - offset); let query_sq_code = self.query_sq_code.as_slice(); - let dist = match self.storage.distance_type { - DistanceType::L2 | DistanceType::Cosine => l2_u8(sq_code, query_sq_code) as f32, - DistanceType::Dot => dot_distance(sq_code, query_sq_code), + match self.storage.distance_type { + DistanceType::L2 | DistanceType::Cosine => { + l2_u8(sq_code, query_sq_code) as f32 * self.scale + } + DistanceType::Dot => self.dot_distance(sq_code), _ => panic!("We should not reach here: sq distance can only be L2 or Dot"), - }; - dist * self.scale + } } fn distance_all(&self, _k_hint: usize) -> Vec { @@ -473,9 +665,8 @@ impl DistCalculator for SQDistCalculator<'_> { c.sq_codes .values() .chunks_exact(c.dim()) - .map(|sq_codes| dot_distance(sq_codes, query_sq_code)) + .map(|sq_codes| self.dot_distance(sq_codes)) }) - .map(|dist| dist * self.scale) .collect(), _ => panic!("We should not reach here: sq distance can only be L2 or Dot"), } @@ -511,7 +702,7 @@ mod tests { use std::iter::repeat_with; use std::sync::Arc; - use arrow_array::FixedSizeListArray; + use arrow_array::{FixedSizeListArray, Float32Array}; use arrow_schema::{DataType, Field, Schema}; use lance_arrow::FixedSizeListArrayExt; use lance_testing::datagen::generate_random_array; @@ -541,6 +732,31 @@ mod tests { RecordBatch::try_new(schema, vec![Arc::new(row_ids), Arc::new(code_arr)]).unwrap() } + fn create_record_batch_with_sq_codes( + row_ids: Vec, + sq_codes: Vec, + dim: usize, + ) -> RecordBatch { + assert_eq!(sq_codes.len(), row_ids.len() * dim); + + let row_ids = UInt64Array::from_iter_values(row_ids); + let sq_code = UInt8Array::from_iter_values(sq_codes); + let code_arr = FixedSizeListArray::try_new_from_values(sq_code, dim as i32).unwrap(); + + let schema = Arc::new(Schema::new(vec![ + Field::new(ROW_ID, DataType::UInt64, false), + Field::new( + SQ_CODE_COLUMN, + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::UInt8, true)), + dim as i32, + ), + false, + ), + ])); + RecordBatch::try_new(schema, vec![Arc::new(row_ids), Arc::new(code_arr)]).unwrap() + } + #[test] fn test_get_chunks() { const DIM: usize = 64; @@ -592,4 +808,53 @@ mod tests { assert_eq!(offset, 400); assert_eq!(chunk.row_id(5), 105); } + + #[test] + fn test_dot_distance_accounts_for_sq_offset() { + const DIM: usize = 2; + + let storage = ScalarQuantizationStorage::try_new( + 8, + DistanceType::Dot, + -1.0..1.0, + [create_record_batch_with_sq_codes( + vec![0, 1], + vec![ + 255, 255, // [1.0, 1.0] + 0, 191, // [-1.0, 0.498] + ], + DIM, + )], + None, + ) + .unwrap(); + + let query = Arc::new(Float32Array::from(vec![-1.0, 1.0])) as ArrayRef; + let calculator = storage.dist_calculator(query.clone(), 0.0); + let distances = calculator.distance_all(2); + + assert!( + distances[1] < distances[0], + "expected [-1.0, 0.498] to rank before [1.0, 1.0], got {distances:?}" + ); + assert!((calculator.distance(0) - 1.0).abs() < 1e-6); + assert!((calculator.distance(1) - -0.49803925).abs() < 1e-6); + + let mut scratch = Vec::new(); + let scratch_distances = { + let scratch_calculator = storage.dist_calculator_with_scratch( + query, + 0.0, + None, + &mut scratch, + DistanceCalculatorOptions::default(), + ); + scratch_calculator.distance_all(2) + }; + assert_eq!(scratch_distances, distances); + assert_eq!(scratch.len(), DIM); + + let stored_query_calculator = storage.dist_calculator_from_id(0); + assert!((stored_query_calculator.distance(1) - 1.5019608).abs() < 1e-6); + } } diff --git a/rust/lance-index/src/vector/utils.rs b/rust/lance-index/src/vector/utils.rs index fb4f9004c57..587e7dc3f15 100644 --- a/rust/lance-index/src/vector/utils.rs +++ b/rust/lance-index/src/vector/utils.rs @@ -277,6 +277,7 @@ mod tests { use half::f16; use lance_arrow::FixedSizeListArrayExt; use num_traits::identities::Zero; + use rayon::ThreadPoolBuilder; use arrow::compute::cast; use rstest::rstest; @@ -307,7 +308,8 @@ mod tests { (0..100).flat_map(|i| std::iter::repeat_n(i as f32, 16)).collect::>(), )) as ArrayRef, 42.0f32)] fn test_simple_index_nearest_centroid(#[case] centroids: ArrayRef, #[case] query_val: f32) { - let index = build_index(centroids, 16); + let thread_pool = ThreadPoolBuilder::new().num_threads(1).build().unwrap(); + let index = thread_pool.install(|| build_index(centroids, 16)); let query: ArrayRef = Arc::new(Float32Array::from(vec![query_val; 16])); let (id, dist) = index.search(query).unwrap(); assert_eq!(id, 42); diff --git a/rust/lance-index/src/vector/v3/subindex.rs b/rust/lance-index/src/vector/v3/subindex.rs index 9a49bc95f1d..9e82921347f 100644 --- a/rust/lance-index/src/vector/v3/subindex.rs +++ b/rust/lance-index/src/vector/v3/subindex.rs @@ -1,7 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::collections::{BinaryHeap, HashMap}; +use lance_core::utils::row_addr_remap::RowAddrRemap; +use std::collections::BinaryHeap; use std::fmt::Debug; use std::sync::Arc; @@ -106,7 +107,7 @@ pub trait IvfSubIndex: Send + Sync + Debug + DeepSizeOf { where Self: Sized; - fn remap(&self, mapping: &HashMap>, store: &impl VectorStore) -> Result + fn remap(&self, mapping: &RowAddrRemap, store: &impl VectorStore) -> Result where Self: Sized; diff --git a/rust/lance-io/Cargo.toml b/rust/lance-io/Cargo.toml index 6cee04d5fd9..41368a84a7f 100644 --- a/rust/lance-io/Cargo.toml +++ b/rust/lance-io/Cargo.toml @@ -37,6 +37,7 @@ chrono.workspace = true futures.workspace = true http.workspace = true log.workspace = true +metrics = { workspace = true, optional = true } moka.workspace = true pin-project.workspace = true prost.workspace = true @@ -60,6 +61,7 @@ rstest.workspace = true mock_instant.workspace = true tokio = { workspace = true, features = ["test-util"] } tracing-mock = { workspace = true } +metrics-util = { workspace = true } [[bench]] name = "scheduler" @@ -67,6 +69,7 @@ harness = false [features] default = ["aws", "azure", "gcp"] +metrics = ["dep:metrics"] gcs-test = [] goosefs-test = [] gcp = ["object_store/gcp", "dep:opendal", "opendal/services-gcs", "dep:object_store_opendal"] diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs index 31d2e8f997e..15905ff4e69 100644 --- a/rust/lance-io/src/object_store.rs +++ b/rust/lance-io/src/object_store.rs @@ -3,6 +3,7 @@ //! Extend [object_store::ObjectStore] functionalities +use std::borrow::Cow; use std::collections::HashMap; use std::ops::Range; use std::pin::Pin; @@ -39,6 +40,8 @@ pub(crate) mod dynamic_credentials; #[cfg(any(feature = "oss", feature = "huggingface", feature = "tos"))] pub(crate) mod dynamic_opendal; mod list_retry; +#[cfg(feature = "metrics")] +pub mod metrics; pub mod providers; pub mod storage_options; #[cfg(test)] @@ -69,6 +72,7 @@ const DEFAULT_LOCAL_BLOCK_SIZE: usize = 4 * 1024; // 4KB block size feature = "tencent", feature = "huggingface", feature = "tos", + feature = "goosefs", ))] const DEFAULT_CLOUD_BLOCK_SIZE: usize = 64 * 1024; // 64KB block size @@ -82,8 +86,10 @@ pub const DEFAULT_DOWNLOAD_RETRY_COUNT: usize = 3; pub use providers::{ObjectStoreProvider, ObjectStoreRegistry}; pub use storage_options::{ - EXPIRES_AT_MILLIS_KEY, LanceNamespaceStorageOptionsProvider, REFRESH_OFFSET_MILLIS_KEY, - StorageOptionsAccessor, StorageOptionsProvider, + BASE_SCOPED_OPTION_PREFIX, BaseScopedStorageOptionsProvider, EXPIRES_AT_MILLIS_KEY, + LanceNamespaceStorageOptionsProvider, REFRESH_OFFSET_MILLIS_KEY, StorageOptionsAccessor, + StorageOptionsProvider, has_base_scoped_options, parse_base_scoped_key, + resolve_base_scoped_options, }; #[async_trait] @@ -257,6 +263,27 @@ impl ObjectStoreParams { .as_ref() .and_then(|a| a.initial_storage_options()) } + + /// Resolve these params for a single base path scope. + /// + /// Storage options may carry base-scoped entries (`base_.`) that + /// apply only to one registered base path; see + /// [`StorageOptionsAccessor::scoped_to_base`]. Returns the params unchanged + /// when the storage options contain no base-scoped entries. + pub fn scoped_to_base(&self, base_id: Option) -> Cow<'_, Self> { + let Some(accessor) = &self.storage_options_accessor else { + return Cow::Borrowed(self); + }; + let scoped = accessor.scoped_to_base(base_id); + if Arc::ptr_eq(&scoped, accessor) { + Cow::Borrowed(self) + } else { + Cow::Owned(Self { + storage_options_accessor: Some(scoped), + ..self.clone() + }) + } + } } // We implement hash for caching diff --git a/rust/lance-io/src/object_store/metrics.rs b/rust/lance-io/src/object_store/metrics.rs new file mode 100644 index 00000000000..d11ae76bd22 --- /dev/null +++ b/rust/lance-io/src/object_store/metrics.rs @@ -0,0 +1,1602 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Publishes object store metrics via the [`metrics`] crate. +//! +//! Two layers cooperate: +//! +//! * [`MeteredObjectStore`] wraps any [`object_store::ObjectStore`] and records +//! per-operation request counts, transferred bytes, latency, errors, and the +//! number of requests currently in flight. It works for every store +//! regardless of backend. +//! * [`MeteringHttpConnector`] wraps the HTTP client used by the native cloud +//! stores (S3 / GCS / Azure) and records throttle / retryable responses per +//! attempt. Because `object_store`'s retry loop re-issues each request +//! through the [`HttpService`](object_store::client::HttpService), this sees +//! every retried response, which a store-level wrapper cannot observe. +//! +//! The two layers have different coverage: every store gets the request-level +//! metrics from [`MeteredObjectStore`], but only the native cloud stores get +//! the HTTP-level throttle metrics. Opendal-backed stores (tos, oss, etc.) +//! bypass `object_store`'s HTTP client, so there is no place to install the +//! connector for them. +//! +//! Metrics carry a `base` label identifying the store. Its cardinality is +//! controlled by the `LANCE_OBJECT_STORE_METRICS_LABEL` environment variable +//! ([`BASE_LABEL_ENV_VAR`]): +//! +//! * `scheme` (default) — scheme only, e.g. `s3`; low, bounded cardinality. +//! * `full` — the full store prefix, e.g. `s3$bucket` or `az$container@account`, +//! so multiple buckets on the same cloud can be told apart. +//! * `off` — omit the `base` label entirely. +//! +//! The metric name constants ([`METRIC_REQUESTS`] etc.) and the recording +//! helpers ([`record_request`], [`record_count`], [`record_error`], +//! [`InFlightGuard`]) are public so custom object stores can emit the same +//! metrics. + +use std::ops::Range; +use std::pin::Pin; +use std::sync::{Arc, OnceLock}; +use std::task::{Context, Poll}; +use std::time::Instant; + +use bytes::Bytes; +use futures::stream::BoxStream; +use futures::{FutureExt, Stream, StreamExt}; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta, + PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as OSResult, + UploadPart, +}; + +/// Total number of object store requests, labelled by `operation` and `base`. +pub const METRIC_REQUESTS: &str = "lance_object_store_requests_total"; +/// Total bytes transferred by object store requests, labelled by `operation` and `base`. +pub const METRIC_BYTES: &str = "lance_object_store_request_bytes_total"; +/// Object store request latency in seconds, labelled by `operation` and `base`. +pub const METRIC_DURATION: &str = "lance_object_store_request_duration_seconds"; +/// Total number of failed object store requests, labelled by `operation` and `base`. +pub const METRIC_ERRORS: &str = "lance_object_store_errors_total"; +/// Total number of throttle responses (HTTP 429 / 503) seen at the HTTP layer, +/// labelled by `status` and `base`. Counts every attempt, including retries. +pub const METRIC_THROTTLE: &str = "lance_object_store_throttle_total"; +/// Total number of retryable responses (HTTP 5xx / 429 / 408) seen at the HTTP +/// layer, labelled by `status` and `base`. Counts every attempt, including +/// retries. This is a superset of [`METRIC_THROTTLE`]; 409 (conflict) is +/// deliberately excluded so commit conflicts are not counted as retries. +pub const METRIC_RETRYABLE: &str = "lance_object_store_retryable_responses_total"; +/// Number of object store requests currently in flight, labelled by `operation` +/// and `base`. +pub const METRIC_IN_FLIGHT: &str = "lance_object_store_in_flight_requests"; + +/// Environment variable controlling the cardinality of the `base` label. +pub const BASE_LABEL_ENV_VAR: &str = "LANCE_OBJECT_STORE_METRICS_LABEL"; + +/// Controls how much of a store's identity the `base` label carries, traded off +/// against metric cardinality. Selected via [`BASE_LABEL_ENV_VAR`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BaseLabelMode { + /// Full store prefix, e.g. `s3$bucket` or `az$container@account`. Highest + /// cardinality: one series family per bucket/container. + Full, + /// Scheme only, e.g. `s3`. The default: low, bounded cardinality. + Scheme, + /// Omit the `base` label entirely. + Off, +} + +fn parse_base_label_mode(value: Option<&str>) -> BaseLabelMode { + match value { + Some("full") => BaseLabelMode::Full, + Some("off") | Some("none") => BaseLabelMode::Off, + Some("scheme") | None => BaseLabelMode::Scheme, + Some(other) => { + tracing::warn!( + "Unrecognized {BASE_LABEL_ENV_VAR}={other:?}; \ + expected one of full, scheme, off. Defaulting to scheme." + ); + BaseLabelMode::Scheme + } + } +} + +/// The label mode is read once from the environment and cached for the process. +fn base_label_mode() -> BaseLabelMode { + static MODE: OnceLock = OnceLock::new(); + *MODE.get_or_init(|| parse_base_label_mode(std::env::var(BASE_LABEL_ENV_VAR).ok().as_deref())) +} + +/// Reduce a full store prefix (`scheme$authority`, or just `scheme` for stores +/// without buckets) to the configured `base` label value, or `None` when the +/// label should be omitted. +fn scoped_base(mode: BaseLabelMode, base: &str) -> Option { + match mode { + BaseLabelMode::Full => Some(base.to_owned()), + BaseLabelMode::Scheme => Some(base.split('$').next().unwrap_or(base).to_owned()), + BaseLabelMode::Off => None, + } +} + +/// Build the `operation` (+ optional `base`) label set shared by all +/// store-level metrics, honoring the configured label mode. +fn operation_labels(base: &str, operation: &'static str) -> Vec { + let mut labels = vec![metrics::Label::new("operation", operation)]; + if let Some(base) = scoped_base(base_label_mode(), base) { + labels.push(metrics::Label::new("base", base)); + } + labels +} + +/// Recommended histogram bucket boundaries for [`METRIC_DURATION`], in seconds. +/// +/// Object store requests can take anywhere from a few milliseconds to the +/// client timeout (commonly ~120s), so the boundaries are dense below 10s and +/// keep useful resolution through the timeout band out to 5 minutes. Exporters +/// that aggregate into fixed buckets (e.g. the OpenTelemetry bridge in the +/// Python bindings) use these. +pub const REQUEST_DURATION_BOUNDS: &[f64] = &[ + 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, // sub-10s + 10.0, 20.0, 30.0, 45.0, 60.0, 90.0, 120.0, 150.0, 180.0, 240.0, 300.0, // 10s–5min +]; + +/// Register descriptions (units and help text) for the object store metrics. +/// +/// This routes through whatever [`metrics::Recorder`] is currently installed, +/// so it must be called *after* the recorder is set. Exporters that build a +/// catalog of available metrics (such as the OpenTelemetry bridge) rely on +/// these descriptions to discover metric names, kinds, and units up front. +pub fn describe_metrics() { + metrics::describe_counter!( + METRIC_REQUESTS, + metrics::Unit::Count, + "Total number of object store requests, by operation and scheme." + ); + metrics::describe_counter!( + METRIC_BYTES, + metrics::Unit::Bytes, + "Total bytes transferred by object store requests, by operation and scheme." + ); + metrics::describe_histogram!( + METRIC_DURATION, + metrics::Unit::Seconds, + "Object store request latency in seconds, by operation and scheme." + ); + metrics::describe_counter!( + METRIC_ERRORS, + metrics::Unit::Count, + "Total number of failed object store requests, by operation and scheme." + ); + metrics::describe_counter!( + METRIC_THROTTLE, + metrics::Unit::Count, + "Total number of throttle responses (HTTP 429 / 503) seen at the HTTP layer, by status and scheme." + ); + metrics::describe_counter!( + METRIC_RETRYABLE, + metrics::Unit::Count, + "Total number of retryable responses (HTTP 5xx / 429 / 408) seen at the HTTP layer, by status and scheme." + ); + metrics::describe_gauge!( + METRIC_IN_FLIGHT, + metrics::Unit::Count, + "Number of object store requests currently in flight, by operation and scheme." + ); +} + +/// Recommended fixed bucket boundaries for the histogram metrics defined here, +/// as `(metric_name, boundaries)` pairs. Exporters that aggregate histograms +/// into fixed buckets read this to configure each histogram. +pub fn histogram_bounds() -> &'static [(&'static str, &'static [f64])] { + &[(METRIC_DURATION, REQUEST_DURATION_BOUNDS)] +} + +/// Record the outcome of a unary request: count, latency, bytes (on success), and errors. +pub fn record_request( + base: &str, + operation: &'static str, + start: Instant, + bytes: u64, + result: &OSResult, +) { + record_outcome(base, operation, start, bytes, result.is_err()); +} + +/// Record count, latency, and either transferred bytes or an error for a +/// completed request. Used both for unary requests and for streamed GETs whose +/// bytes are only known once the body finishes. +pub fn record_outcome( + base: &str, + operation: &'static str, + start: Instant, + bytes: u64, + is_error: bool, +) { + let elapsed = start.elapsed().as_secs_f64(); + let labels = operation_labels(base, operation); + metrics::counter!(METRIC_REQUESTS, labels.clone()).increment(1); + metrics::histogram!(METRIC_DURATION, labels.clone()).record(elapsed); + if is_error { + metrics::counter!(METRIC_ERRORS, labels).increment(1); + } else if bytes > 0 { + metrics::counter!(METRIC_BYTES, labels).increment(bytes); + } +} + +/// Record a single request count without latency, used for streaming operations +/// (list / delete) whose work happens lazily as the stream is polled. +pub fn record_count(base: &str, operation: &'static str) { + metrics::counter!(METRIC_REQUESTS, operation_labels(base, operation)).increment(1); +} + +/// Record a single error for an operation. +pub fn record_error(base: &str, operation: &'static str) { + metrics::counter!(METRIC_ERRORS, operation_labels(base, operation)).increment(1); +} + +/// Raises the in-flight gauge for an operation on creation and lowers it on +/// drop, so the count stays balanced even if the request future or stream is +/// cancelled or dropped before completing. +pub struct InFlightGuard { + labels: Vec, +} + +impl InFlightGuard { + pub fn new(base: &str, operation: &'static str) -> Self { + let labels = operation_labels(base, operation); + metrics::gauge!(METRIC_IN_FLIGHT, labels.clone()).increment(1.0); + Self { labels } + } +} + +impl Drop for InFlightGuard { + fn drop(&mut self) { + metrics::gauge!(METRIC_IN_FLIGHT, self.labels.clone()).decrement(1.0); + } +} + +#[derive(Debug)] +pub struct MeteredObjectStore { + target: Arc, + base: String, +} + +impl std::fmt::Display for MeteredObjectStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "MeteredObjectStore({})", self.target) + } +} + +#[async_trait::async_trait] +#[deny(clippy::missing_trait_methods)] +impl object_store::ObjectStore for MeteredObjectStore { + async fn put_opts( + &self, + location: &Path, + bytes: PutPayload, + opts: PutOptions, + ) -> OSResult { + let size = bytes.content_length() as u64; + let _in_flight = InFlightGuard::new(&self.base, "put"); + let start = Instant::now(); + let result = self.target.put_opts(location, bytes, opts).await; + record_request(&self.base, "put", start, size, &result); + result + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult> { + let upload = self.target.put_multipart_opts(location, opts).await?; + Ok(Box::new(MeteredMultipartUpload { + target: upload, + base: self.base.clone(), + })) + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { + // `head()` is implemented as a `get_opts` call with `head = true`, so we + // distinguish it here to keep HEAD and GET as separate operations. + let is_head = options.head; + let operation = if is_head { "head" } else { "get" }; + let in_flight = InFlightGuard::new(&self.base, operation); + let start = Instant::now(); + let result = self.target.get_opts(location, options).await; + + // A HEAD transfers only metadata, and errors carry no payload, so both + // are recorded immediately. `get_opts` only resolves once the response + // headers arrive; the body is streamed afterwards, so for a successful + // GET we defer recording until the body has been drained (see below). + if is_head || result.is_err() { + record_request(&self.base, operation, start, 0, &result); + return result; + } + + let result = result.expect("checked to be Ok above"); + Ok(meter_get_result( + result, + self.base.clone(), + start, + in_flight, + )) + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OSResult> { + let _in_flight = InFlightGuard::new(&self.base, "get"); + let start = Instant::now(); + let result = self.target.get_ranges(location, ranges).await; + let bytes = match &result { + Ok(parts) => parts.iter().map(|b| b.len() as u64).sum(), + Err(_) => 0, + }; + record_request(&self.base, "get", start, bytes, &result); + result + } + + fn delete_stream( + &self, + locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + let base = self.base.clone(); + // Count one logical delete request per call, matching `list`: a single + // `delete_stream` maps to one batched request on stores that support it + // (e.g. S3's `DeleteObjects`), so counting per yielded path would + // over-count. Errors are still recorded per failing path. + record_count(&self.base, "delete"); + let in_flight = InFlightGuard::new(&self.base, "delete"); + self.target + .delete_stream(locations) + .map(move |result| { + // Reference `in_flight` so this `move` closure captures (owns) + // the guard, keeping the gauge raised until the stream is + // dropped (a move closure only captures the variables it uses). + let _in_flight = &in_flight; + if result.is_err() { + record_error(&base, "delete"); + } + result + }) + .boxed() + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + record_count(&self.base, "list"); + meter_list_stream( + self.target.list(prefix), + self.base.clone(), + InFlightGuard::new(&self.base, "list"), + ) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OSResult> { + record_count(&self.base, "list"); + meter_list_stream( + self.target.list_with_offset(prefix, offset), + self.base.clone(), + InFlightGuard::new(&self.base, "list"), + ) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { + let _in_flight = InFlightGuard::new(&self.base, "list"); + let start = Instant::now(); + let result = self.target.list_with_delimiter(prefix).await; + record_request(&self.base, "list", start, 0, &result); + result + } + + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> { + let _in_flight = InFlightGuard::new(&self.base, "copy"); + let start = Instant::now(); + let result = self.target.copy_opts(from, to, opts).await; + record_request(&self.base, "copy", start, 0, &result); + result + } + + async fn rename_opts(&self, from: &Path, to: &Path, opts: RenameOptions) -> OSResult<()> { + let _in_flight = InFlightGuard::new(&self.base, "rename"); + let start = Instant::now(); + let result = self.target.rename_opts(from, to, opts).await; + record_request(&self.base, "rename", start, 0, &result); + result + } +} + +/// Count errors yielded while draining a list stream. The request itself is +/// counted once when the stream is created (a single LIST may return many items). +fn meter_list_stream( + stream: BoxStream<'static, OSResult>, + base: String, + in_flight: InFlightGuard, +) -> BoxStream<'static, OSResult> { + stream + .map(move |result| { + // Reference `in_flight` so this `move` closure captures (owns) the + // guard: a move closure only captures the variables it uses, and + // holding it here keeps the gauge raised until the stream is dropped. + let _in_flight = &in_flight; + if result.is_err() { + record_error(&base, "list"); + } + result + }) + .boxed() +} + +/// Wrap a successful GET so the request is recorded once its body has been +/// fully read, capturing the true transfer duration and byte count rather than +/// the time-to-first-byte and declared range. For payloads without a body +/// stream (e.g. a local file handle) the request is recorded immediately. +fn meter_get_result( + mut result: GetResult, + base: String, + start: Instant, + in_flight: InFlightGuard, +) -> GetResult { + match result.payload { + GetResultPayload::Stream(stream) => { + result.payload = GetResultPayload::Stream( + MeteredGetStream { + inner: stream, + base, + start, + bytes: 0, + errored: false, + recorded: false, + _in_flight: in_flight, + } + .boxed(), + ); + result + } + // No body stream to observe (e.g. a local file), so record now. + other => { + let bytes = result.range.end - result.range.start; + record_outcome(&base, "get", start, bytes, false); + result.payload = other; + result + } + } +} + +/// Stream wrapper over a GET body that records the request (count, duration, +/// bytes, errors) once the body is fully drained or the stream is dropped. +struct MeteredGetStream { + inner: BoxStream<'static, OSResult>, + base: String, + start: Instant, + bytes: u64, + errored: bool, + recorded: bool, + _in_flight: InFlightGuard, +} + +impl MeteredGetStream { + fn record(&mut self) { + if self.recorded { + return; + } + self.recorded = true; + record_outcome(&self.base, "get", self.start, self.bytes, self.errored); + } +} + +impl Stream for MeteredGetStream { + type Item = OSResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.inner.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(chunk))) => { + self.bytes += chunk.len() as u64; + Poll::Ready(Some(Ok(chunk))) + } + Poll::Ready(Some(Err(e))) => { + self.errored = true; + Poll::Ready(Some(Err(e))) + } + Poll::Ready(None) => { + self.record(); + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + } + } +} + +impl Drop for MeteredGetStream { + fn drop(&mut self) { + // Records the partial transfer if the body was dropped before it drained. + self.record(); + } +} + +#[derive(Debug)] +struct MeteredMultipartUpload { + target: Box, + base: String, +} + +#[async_trait::async_trait] +impl MultipartUpload for MeteredMultipartUpload { + fn put_part(&mut self, data: PutPayload) -> UploadPart { + // Each part upload is a distinct request, recorded under the `put_part` + // operation with the same count / bytes / latency / error set as a + // unary put. + let base = self.base.clone(); + let size = data.content_length() as u64; + let inner = self.target.put_part(data); + async move { + let _in_flight = InFlightGuard::new(&base, "put_part"); + let start = Instant::now(); + let result = inner.await; + record_request(&base, "put_part", start, size, &result); + result + } + .boxed() + } + + async fn complete(&mut self) -> OSResult { + // Completing a multipart upload issues its own request that can throttle + // or fail, so it is metered like any other operation. + let _in_flight = InFlightGuard::new(&self.base, "complete_multipart"); + let start = Instant::now(); + let result = self.target.complete().await; + record_request(&self.base, "complete_multipart", start, 0, &result); + result + } + + async fn abort(&mut self) -> OSResult<()> { + let _in_flight = InFlightGuard::new(&self.base, "abort_multipart"); + let start = Instant::now(); + let result = self.target.abort().await; + record_request(&self.base, "abort_multipart", start, 0, &result); + result + } +} + +pub trait ObjectStoreMetricsExt { + /// Wrap this store so its operations publish metrics under the given `base` label. + fn metered(self, base: String) -> Arc; +} + +impl ObjectStoreMetricsExt for Arc { + fn metered(self, base: String) -> Arc { + Arc::new(MeteredObjectStore { target: self, base }) + } +} + +// --- Layer 2: HTTP-level throttle metrics for native cloud stores --- + +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +mod http { + use super::*; + use object_store::client::{ + ClientOptions, HttpClient, HttpConnector, HttpError, HttpRequest, HttpResponse, + HttpService, ReqwestConnector, + }; + + /// An [`HttpConnector`] that records throttle and retryable responses + /// observed by the underlying HTTP client. Install it on the S3 / GCS / + /// Azure builders via `with_http_connector`. + #[derive(Debug)] + pub struct MeteringHttpConnector { + base: String, + inner: ReqwestConnector, + } + + impl MeteringHttpConnector { + pub fn new(base: String) -> Self { + Self { + base, + inner: ReqwestConnector::default(), + } + } + } + + impl HttpConnector for MeteringHttpConnector { + fn connect(&self, options: &ClientOptions) -> object_store::Result { + let client = self.inner.connect(options)?; + Ok(HttpClient::new(MeteringHttpService { + base: self.base.clone(), + inner: client, + })) + } + } + + #[derive(Debug)] + struct MeteringHttpService { + base: String, + inner: HttpClient, + } + + #[async_trait::async_trait] + impl HttpService for MeteringHttpService { + async fn call(&self, req: HttpRequest) -> Result { + let response = self.inner.execute(req).await?; + let status = response.status().as_u16(); + // Each attempt that object_store may retry is recorded with its + // numeric status. Throttles (429 / 503) are a distinct, narrower + // signal than the broader set of retryable responses, so they get + // their own counter. 409 (conflict) is intentionally excluded from + // the retryable set so commit conflicts are not counted as retries. + let is_throttle = status == 429 || status == 503; + let is_retryable = status == 429 || status == 408 || (500..600).contains(&status); + if is_throttle { + metrics::counter!(METRIC_THROTTLE, status_labels(&self.base, status)).increment(1); + } + if is_retryable { + metrics::counter!(METRIC_RETRYABLE, status_labels(&self.base, status)).increment(1); + } + Ok(response) + } + } + + /// Build the `status` (+ optional `base`) label set for HTTP-layer metrics, + /// honoring the configured label mode. + fn status_labels(base: &str, status: u16) -> Vec { + let mut labels = vec![metrics::Label::new("status", status.to_string())]; + if let Some(base) = scoped_base(base_label_mode(), base) { + labels.push(metrics::Label::new("base", base)); + } + labels + } + + #[cfg(test)] + mod tests { + use super::*; + use metrics_util::debugging::{DebugValue, DebuggingRecorder}; + use object_store::client::{HttpRequestBody, HttpResponseBody}; + + /// A mock [`HttpService`] that always responds with a fixed status code. + #[derive(Debug)] + struct StaticStatusService { + status: u16, + } + + #[async_trait::async_trait] + impl HttpService for StaticStatusService { + async fn call(&self, _req: HttpRequest) -> Result { + Ok(::http::Response::builder() + .status(self.status) + .body(HttpResponseBody::from(Bytes::new())) + .unwrap()) + } + } + + fn request() -> HttpRequest { + ::http::Request::builder() + .method("GET") + .uri("http://example.com/obj") + .body(HttpRequestBody::empty()) + .unwrap() + } + + fn metric_count( + metrics: &[(metrics::Key, DebugValue)], + name: &str, + base: &str, + status: &str, + ) -> u64 { + for (key, value) in metrics { + if key.name() != name { + continue; + } + let labels: std::collections::HashMap<&str, &str> = + key.labels().map(|l| (l.key(), l.value())).collect(); + if labels.get("base") == Some(&base) + && labels.get("status") == Some(&status) + && let DebugValue::Counter(v) = value + { + return *v; + } + } + 0 + } + + #[test] + fn test_throttle_and_retryable_responses_counted_by_status() { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + // Each attempt that object_store retries flows through `call` + // again; here we simulate that by issuing several responses. + // The base is baked into the connector, so it labels the + // metric. Bases here have no `$`, so they are unaffected by + // the label mode and this test isolates status handling. + for (base, status) in [ + ("s3", 429u16), + ("s3", 503), + ("s3", 503), + ("s3", 500), + ("s3", 408), + ("s3", 409), + ("s3", 200), + ("s3", 404), + ("gs", 429), + ] { + let service = MeteringHttpService { + base: base.into(), + inner: HttpClient::new(StaticStatusService { status }), + }; + service.call(request()).await.unwrap(); + } + }); + }); + + let recorded: Vec<_> = snapshotter + .snapshot() + .into_vec() + .into_iter() + .map(|(ck, _unit, _desc, value)| (ck.key().clone(), value)) + .collect(); + + let throttle = |base, status| metric_count(&recorded, METRIC_THROTTLE, base, status); + let retryable = |base, status| metric_count(&recorded, METRIC_RETRYABLE, base, status); + + // Throttles are only 429 and 503. + assert_eq!(throttle("s3", "429"), 1); + assert_eq!(throttle("s3", "503"), 2); + assert_eq!(throttle("s3", "500"), 0); + assert_eq!(throttle("s3", "408"), 0); + + // Retryable is the broader set: 5xx, 429, 408 (but not 409). + assert_eq!(retryable("s3", "429"), 1); + assert_eq!(retryable("s3", "503"), 2); + assert_eq!(retryable("s3", "500"), 1); + assert_eq!(retryable("s3", "408"), 1); + // 409 conflict is excluded so commit conflicts are not counted as retries. + assert_eq!(retryable("s3", "409"), 0); + + // Success and non-retryable client errors count as neither. + assert_eq!(throttle("s3", "200"), 0); + assert_eq!(retryable("s3", "404"), 0); + + // The base label is taken from the connector, not shared across stores. + assert_eq!(throttle("gs", "429"), 1); + assert_eq!(retryable("gs", "429"), 1); + assert_eq!(throttle("gs", "503"), 0); + } + } +} + +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +pub use http::MeteringHttpConnector; + +#[cfg(test)] +mod tests { + use super::*; + + use metrics_util::debugging::{DebugValue, DebuggingRecorder, Snapshotter}; + use object_store::memory::InMemory; + use object_store::{ObjectStoreExt, PutPayload}; + + fn payload(data: &[u8]) -> PutPayload { + PutPayload::from_bytes(Bytes::copy_from_slice(data)) + } + + fn metered_store() -> Arc { + (Arc::new(InMemory::new()) as Arc).metered("memory".into()) + } + + /// A single materialized snapshot of recorded metrics. It must be taken + /// only once: the snapshotter *drains* histogram samples on every + /// `snapshot()` call, so a second snapshot would see empty histograms. + type Metrics = Vec<(metrics::Key, DebugValue)>; + + /// Materialize the current recorder state. Histogram samples are *drained* + /// on each call, so a metric must be read from a single snapshot. + fn snapshot(snapshotter: &Snapshotter) -> Metrics { + snapshotter + .snapshot() + .into_vec() + .into_iter() + .map(|(ck, _unit, _desc, value)| (ck.key().clone(), value)) + .collect() + } + + /// Run an async closure with a thread-local metrics recorder installed and + /// return the resulting metrics. Uses a current-thread runtime so all polls + /// happen on the thread that holds the recorder guard. + fn capture_metrics(f: F) -> Metrics + where + F: FnOnce() -> Fut, + Fut: std::future::Future, + { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(f()); + }); + snapshot(&snapshotter) + } + + fn key_matches(key: &metrics::Key, name: &str, labels: &[(&str, &str)]) -> bool { + if key.name() != name { + return false; + } + let actual: std::collections::HashSet<(&str, &str)> = + key.labels().map(|l| (l.key(), l.value())).collect(); + labels.len() == actual.len() && labels.iter().all(|l| actual.contains(l)) + } + + fn counter_value(metrics: &Metrics, name: &str, labels: &[(&str, &str)]) -> u64 { + for (key, value) in metrics { + if key_matches(key, name, labels) + && let DebugValue::Counter(v) = value + { + return *v; + } + } + 0 + } + + fn histogram_count(metrics: &Metrics, name: &str, labels: &[(&str, &str)]) -> usize { + for (key, value) in metrics { + if key_matches(key, name, labels) + && let DebugValue::Histogram(samples) = value + { + return samples.len(); + } + } + 0 + } + + fn gauge_value(metrics: &Metrics, name: &str, labels: &[(&str, &str)]) -> f64 { + for (key, value) in metrics { + if key_matches(key, name, labels) + && let DebugValue::Gauge(v) = value + { + return v.0; + } + } + 0.0 + } + + fn has_metric(metrics: &Metrics, name: &str, labels: &[(&str, &str)]) -> bool { + metrics + .iter() + .any(|(key, _)| key_matches(key, name, labels)) + } + + #[test] + fn test_parse_base_label_mode() { + assert_eq!(parse_base_label_mode(None), BaseLabelMode::Scheme); + assert_eq!(parse_base_label_mode(Some("scheme")), BaseLabelMode::Scheme); + assert_eq!(parse_base_label_mode(Some("full")), BaseLabelMode::Full); + assert_eq!(parse_base_label_mode(Some("off")), BaseLabelMode::Off); + assert_eq!(parse_base_label_mode(Some("none")), BaseLabelMode::Off); + // Unrecognized values fall back to the conservative default. + assert_eq!(parse_base_label_mode(Some("bogus")), BaseLabelMode::Scheme); + } + + #[test] + fn test_scoped_base() { + assert_eq!( + scoped_base(BaseLabelMode::Full, "s3$bucket").as_deref(), + Some("s3$bucket") + ); + assert_eq!( + scoped_base(BaseLabelMode::Scheme, "s3$bucket").as_deref(), + Some("s3") + ); + // Azure keeps only the scheme even though its prefix carries the account. + assert_eq!( + scoped_base(BaseLabelMode::Scheme, "az$container@account").as_deref(), + Some("az") + ); + // A prefix without `$` (e.g. memory/file) is unchanged by scheme mode. + assert_eq!( + scoped_base(BaseLabelMode::Scheme, "memory").as_deref(), + Some("memory") + ); + assert_eq!(scoped_base(BaseLabelMode::Off, "s3$bucket"), None); + } + + #[test] + fn test_base_label_defaults_to_scheme() { + // No env var is set in the test process, so the default `scheme` mode + // applies: the full prefix collapses to just the scheme. + let recorded = capture_metrics(|| async { + let store = (Arc::new(InMemory::new()) as Arc) + .metered("s3$my-bucket".into()); + store.put(&Path::from("a"), payload(b"x")).await.unwrap(); + }); + + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "put"), ("base", "s3")] + ), + 1 + ); + // The full prefix is not emitted as the label under the default mode. + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "put"), ("base", "s3$my-bucket")] + ), + 0 + ); + } + + #[test] + fn test_put_records_count_bytes_and_latency() { + let data = b"hello world"; + let recorded = capture_metrics(|| async { + let store = metered_store(); + store + .put(&Path::from("a/b.bin"), payload(data)) + .await + .unwrap(); + }); + + let labels = [("operation", "put"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!( + counter_value(&recorded, METRIC_BYTES, &labels), + data.len() as u64 + ); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + } + + #[test] + fn test_get_records_count_and_bytes() { + let data = b"hello world"; + let recorded = capture_metrics(|| async { + let store = metered_store(); + let path = Path::from("a/b.bin"); + store.put(&path, payload(data)).await.unwrap(); + // The GET is only recorded once its body has been fully drained. + store.get(&path).await.unwrap().bytes().await.unwrap(); + }); + + let labels = [("operation", "get"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!( + counter_value(&recorded, METRIC_BYTES, &labels), + data.len() as u64 + ); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + } + + #[test] + fn test_get_not_recorded_until_body_drained() { + let data = b"hello world"; + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + let labels = [("operation", "get"), ("base", "memory")]; + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + let store = metered_store(); + let path = Path::from("a/b.bin"); + store.put(&path, payload(data)).await.unwrap(); + + // Holding the result without reading the body records nothing yet. + let result = store.get(&path).await.unwrap(); + assert_eq!( + counter_value(&snapshot(&snapshotter), METRIC_REQUESTS, &labels), + 0 + ); + + // Draining the body records the request with the true byte count. + let bytes = result.bytes().await.unwrap(); + assert_eq!(bytes.len(), data.len()); + let recorded = snapshot(&snapshotter); + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!( + counter_value(&recorded, METRIC_BYTES, &labels), + data.len() as u64 + ); + }); + }); + } + + #[test] + fn test_head_is_a_separate_operation() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + let path = Path::from("a/b.bin"); + store.put(&path, payload(b"hello world")).await.unwrap(); + store.head(&path).await.unwrap(); + }); + + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "head"), ("base", "memory")] + ), + 1 + ); + // The head call must not be counted as a get. + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "get"), ("base", "memory")] + ), + 0 + ); + // A HEAD transfers only metadata, so it records no payload bytes. + assert_eq!( + counter_value( + &recorded, + METRIC_BYTES, + &[("operation", "head"), ("base", "memory")] + ), + 0 + ); + } + + #[test] + fn test_delete_records_one_request_per_call() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + for i in 0..3 { + store + .put(&Path::from(format!("a/{i}.bin")), payload(b"x")) + .await + .unwrap(); + } + // `delete` drives `delete_stream`; deleting three paths is still one + // logical delete request (a single batched request on real stores). + let paths = + futures::stream::iter((0..3).map(|i| Ok(Path::from(format!("a/{i}.bin"))))).boxed(); + let _: Vec<_> = store.delete_stream(paths).collect().await; + }); + + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "delete"), ("base", "memory")] + ), + 1 + ); + } + + #[test] + fn test_list_counts_one_request_not_per_item() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + for i in 0..3 { + store + .put(&Path::from(format!("a/{i}.bin")), payload(b"x")) + .await + .unwrap(); + } + let _: Vec<_> = store.list(Some(&Path::from("a"))).collect().await; + }); + + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "list"), ("base", "memory")] + ), + 1 + ); + } + + #[test] + fn test_error_is_counted() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + // Getting a missing object errors. + let _ = store.get(&Path::from("does/not/exist")).await; + }); + + let labels = [("operation", "get"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &labels), 1); + // A failed request is still counted as a request, with latency recorded. + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + // No bytes are transferred on a failed get. + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 0); + } + + #[test] + fn test_get_ranges_sums_part_bytes_and_labels_get() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + let path = Path::from("a/b.bin"); + store.put(&path, payload(b"hello world")).await.unwrap(); + // Two disjoint ranges of 3 bytes each. + store.get_ranges(&path, &[2..5, 6..9]).await.unwrap(); + }); + + let labels = [("operation", "get"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 6); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + } + + #[test] + fn test_copy_and_rename_record_zero_bytes() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + store + .put(&Path::from("a/src"), payload(b"x")) + .await + .unwrap(); + store + .copy(&Path::from("a/src"), &Path::from("a/copy")) + .await + .unwrap(); + store + .rename(&Path::from("a/copy"), &Path::from("a/moved")) + .await + .unwrap(); + }); + + for operation in ["copy", "rename"] { + let labels = [("operation", operation), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 0); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + } + } + + #[test] + fn test_list_with_delimiter_records_latency() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + store.put(&Path::from("a/b"), payload(b"x")).await.unwrap(); + store + .list_with_delimiter(Some(&Path::from("a"))) + .await + .unwrap(); + }); + + let labels = [("operation", "list"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 0); + } + + #[test] + fn test_list_with_offset_counts_one_request() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + for i in 0..3 { + store + .put(&Path::from(format!("a/{i}")), payload(b"x")) + .await + .unwrap(); + } + let _: Vec<_> = store + .list_with_offset(Some(&Path::from("a")), &Path::from("a/0")) + .collect() + .await; + }); + + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "list"), ("base", "memory")] + ), + 1 + ); + } + + #[test] + fn test_multipart_records_each_part_and_complete() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + let mut upload = store.put_multipart(&Path::from("a/big")).await.unwrap(); + upload.put_part(payload(b"hello")).await.unwrap(); // 5 bytes + upload.put_part(payload(b"world!!")).await.unwrap(); // 7 bytes + upload.complete().await.unwrap(); + }); + + let part_labels = [("operation", "put_part"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &part_labels), 2); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &part_labels), 12); + // Each part records its own latency sample, like a unary put. + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &part_labels), 2); + // A successful part upload records no error. + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &part_labels), 0); + + // Completing the upload is its own metered request. + let complete_labels = [("operation", "complete_multipart"), ("base", "memory")]; + assert_eq!( + counter_value(&recorded, METRIC_REQUESTS, &complete_labels), + 1 + ); + assert_eq!( + histogram_count(&recorded, METRIC_DURATION, &complete_labels), + 1 + ); + } + + #[test] + fn test_multipart_abort_is_recorded() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + let mut upload = store.put_multipart(&Path::from("a/big")).await.unwrap(); + upload.put_part(payload(b"hello")).await.unwrap(); + upload.abort().await.unwrap(); + }); + + let labels = [("operation", "abort_multipart"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + } + + #[test] + fn test_multipart_part_error_is_counted() { + let recorded = capture_metrics(|| async { + let store = (Arc::new(FailingStreamStore) as Arc) + .metered("memory".into()); + let mut upload = store.put_multipart(&Path::from("a/big")).await.unwrap(); + let _ = upload.put_part(payload(b"data")).await; + }); + + let labels = [("operation", "put_part"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &labels), 1); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + // A failed part transfers no counted bytes. + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 0); + } + + #[test] + fn test_in_flight_guard_tracks_and_releases() { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + let labels = [("operation", "get"), ("base", "memory")]; + metrics::with_local_recorder(&recorder, || { + let g1 = InFlightGuard::new("memory", "get"); + let g2 = InFlightGuard::new("memory", "get"); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 2.0 + ); + drop(g1); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 1.0 + ); + drop(g2); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 0.0 + ); + }); + } + + #[test] + fn test_in_flight_gauge_is_wired_and_balances() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + let path = Path::from("a/b.bin"); + store.put(&path, payload(b"hello")).await.unwrap(); + store.get(&path).await.unwrap(); + }); + + // The gauge is emitted for each operation (guard is wired in) and, once + // the operation completes, balances back to zero. + for operation in ["put", "get"] { + let labels = [("operation", operation), ("base", "memory")]; + assert!(has_metric(&recorded, METRIC_IN_FLIGHT, &labels)); + assert_eq!(gauge_value(&recorded, METRIC_IN_FLIGHT, &labels), 0.0); + } + } + + #[test] + fn test_list_stream_holds_in_flight_until_dropped() { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + let labels = [("operation", "list"), ("base", "memory")]; + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + let store = metered_store(); + store.put(&Path::from("a/x"), payload(b"x")).await.unwrap(); + + // Creating the stream raises the gauge; it stays raised until the + // stream is dropped, even before any items are drained. + let stream = store.list(Some(&Path::from("a"))); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 1.0 + ); + drop(stream); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 0.0 + ); + }); + }); + } + + #[test] + fn test_delete_stream_holds_in_flight_until_dropped() { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + let labels = [("operation", "delete"), ("base", "memory")]; + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + let store = metered_store(); + let locations = futures::stream::iter(vec![Ok(Path::from("a/b"))]).boxed(); + + // Like list, creating the delete stream raises the gauge and holds + // it until the stream is dropped, before any items are drained. + let stream = store.delete_stream(locations); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 1.0 + ); + drop(stream); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 0.0 + ); + }); + }); + } + + #[test] + fn test_in_flight_released_when_operation_future_dropped() { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + let labels = [("operation", "get"), ("base", "memory")]; + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + let started = Arc::new(tokio::sync::Notify::new()); + // Never signalled: the request stays blocked mid-flight. + let release = Arc::new(tokio::sync::Notify::new()); + let store = (Arc::new(BlockingStore { + started: started.clone(), + release, + }) as Arc) + .metered("memory".into()); + + let path = Path::from("a/b"); + let mut fut = Box::pin(store.get(&path)); + // Drive the request until it is blocked inside the inner store. + tokio::select! { + _ = &mut fut => unreachable!("the blocking store never returns"), + _ = started.notified() => {} + } + + // The gauge is raised while the request is outstanding, and + // dropping the future before it completes releases it. + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 1.0 + ); + drop(fut); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 0.0 + ); + }); + }); + } + + #[test] + fn test_streaming_errors_are_counted() { + let recorded = capture_metrics(|| async { + let delete_store = (Arc::new(FailingStreamStore) as Arc) + .metered("memory".into()); + let _ = delete_store.delete(&Path::from("a/b")).await; + + let list_store = (Arc::new(FailingStreamStore) as Arc) + .metered("memory".into()); + let _: Vec<_> = list_store.list(None).collect().await; + }); + + // delete_stream counts the item and records an error when it fails. + let delete_labels = [("operation", "delete"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &delete_labels), 1); + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &delete_labels), 1); + + // A list request is counted once; a failure while draining records an error. + let list_labels = [("operation", "list"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &list_labels), 1); + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &list_labels), 1); + } + + /// A store whose stream-producing operations always yield an error, used to + /// exercise the error branches of the streaming wrappers. + #[derive(Debug)] + struct FailingStreamStore; + + impl std::fmt::Display for FailingStreamStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "FailingStreamStore") + } + } + + fn test_error() -> object_store::Error { + object_store::Error::Generic { + store: "FailingStreamStore", + source: "injected failure".into(), + } + } + + #[async_trait::async_trait] + impl object_store::ObjectStore for FailingStreamStore { + async fn put_opts( + &self, + _location: &Path, + _bytes: PutPayload, + _opts: PutOptions, + ) -> OSResult { + unimplemented!() + } + + async fn put_multipart_opts( + &self, + _location: &Path, + _opts: PutMultipartOptions, + ) -> OSResult> { + Ok(Box::new(FailingUpload)) + } + + async fn get_opts(&self, _location: &Path, _options: GetOptions) -> OSResult { + unimplemented!() + } + + fn delete_stream( + &self, + _locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + futures::stream::once(async { Err(test_error()) }).boxed() + } + + fn list(&self, _prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + futures::stream::once(async { Err(test_error()) }).boxed() + } + + fn list_with_offset( + &self, + _prefix: Option<&Path>, + _offset: &Path, + ) -> BoxStream<'static, OSResult> { + unimplemented!() + } + + async fn list_with_delimiter(&self, _prefix: Option<&Path>) -> OSResult { + unimplemented!() + } + + async fn copy_opts(&self, _from: &Path, _to: &Path, _opts: CopyOptions) -> OSResult<()> { + unimplemented!() + } + + async fn rename_opts( + &self, + _from: &Path, + _to: &Path, + _opts: RenameOptions, + ) -> OSResult<()> { + unimplemented!() + } + } + + /// A [`MultipartUpload`] whose part uploads always fail, used to exercise the + /// error branch of the metered `put_part`. + #[derive(Debug)] + struct FailingUpload; + + #[async_trait::async_trait] + impl MultipartUpload for FailingUpload { + fn put_part(&mut self, _data: PutPayload) -> UploadPart { + async { Err(test_error()) }.boxed() + } + + async fn complete(&mut self) -> OSResult { + unimplemented!() + } + + async fn abort(&mut self) -> OSResult<()> { + unimplemented!() + } + } + + /// A store whose `get_opts` blocks after signalling `started`, so a request + /// can be observed mid-flight and then dropped before it completes. + #[derive(Debug)] + struct BlockingStore { + started: Arc, + release: Arc, + } + + impl std::fmt::Display for BlockingStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "BlockingStore") + } + } + + #[async_trait::async_trait] + impl object_store::ObjectStore for BlockingStore { + async fn put_opts( + &self, + _location: &Path, + _bytes: PutPayload, + _opts: PutOptions, + ) -> OSResult { + unimplemented!() + } + + async fn put_multipart_opts( + &self, + _location: &Path, + _opts: PutMultipartOptions, + ) -> OSResult> { + unimplemented!() + } + + async fn get_opts(&self, _location: &Path, _options: GetOptions) -> OSResult { + self.started.notify_one(); + self.release.notified().await; + unreachable!("release is never signalled in the test") + } + + fn delete_stream( + &self, + _locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + unimplemented!() + } + + fn list(&self, _prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + unimplemented!() + } + + fn list_with_offset( + &self, + _prefix: Option<&Path>, + _offset: &Path, + ) -> BoxStream<'static, OSResult> { + unimplemented!() + } + + async fn list_with_delimiter(&self, _prefix: Option<&Path>) -> OSResult { + unimplemented!() + } + + async fn copy_opts(&self, _from: &Path, _to: &Path, _opts: CopyOptions) -> OSResult<()> { + unimplemented!() + } + + async fn rename_opts( + &self, + _from: &Path, + _to: &Path, + _opts: RenameOptions, + ) -> OSResult<()> { + unimplemented!() + } + } +} diff --git a/rust/lance-io/src/object_store/providers.rs b/rust/lance-io/src/object_store/providers.rs index c0b3a156502..5e7963e5e97 100644 --- a/rust/lance-io/src/object_store/providers.rs +++ b/rust/lance-io/src/object_store/providers.rs @@ -215,6 +215,12 @@ impl ObjectStoreRegistry { base_path: Url, params: &ObjectStoreParams, ) -> Result> { + // Base-scoped storage options (`base_.`) are directives for + // other registered base paths; resolve them away before building or + // caching a store for this location. Params already resolved for a + // base contain no scoped entries, so this is a no-op for them. + let params = params.scoped_to_base(None); + let params = params.as_ref(); let scheme = base_path.scheme(); let Some(provider) = self.get_provider(scheme) else { return Err(self.scheme_not_found_error(scheme)); @@ -259,6 +265,14 @@ impl ObjectStoreRegistry { store.inner = store.inner.traced(); + #[cfg(feature = "metrics")] + { + // Label metrics by the store's unique prefix (e.g. `s3$bucket`, + // `az$container@account`) so multiple stores on one cloud differ. + use crate::object_store::metrics::ObjectStoreMetricsExt; + store.inner = store.inner.metered(cache_path.clone()); + } + if let Some(wrapper) = ¶ms.object_store_wrapper { store.inner = wrapper.wrap(&cache_path, store.inner); } @@ -394,6 +408,36 @@ mod tests { ); } + #[tokio::test] + async fn test_get_store_resolves_base_scoped_options() { + use crate::object_store::StorageOptionsAccessor; + + let registry = ObjectStoreRegistry::default(); + let url = Url::parse("memory://test").unwrap(); + + let with_scoped = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("shared".to_string(), "value".to_string()), + ("base_1.account_key".to_string(), "base1-key".to_string()), + ]), + ))), + ..Default::default() + }; + let without_scoped = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([("shared".to_string(), "value".to_string())]), + ))), + ..Default::default() + }; + + // Base-scoped entries are resolved away before the store is built and + // cached, so params with and without them yield the same cached store. + let store_scoped = registry.get_store(url.clone(), &with_scoped).await.unwrap(); + let store_plain = registry.get_store(url, &without_scoped).await.unwrap(); + assert!(Arc::ptr_eq(&store_scoped, &store_plain)); + } + #[test] fn test_calculate_object_store_scheme_not_found() { let registry = ObjectStoreRegistry::empty(); diff --git a/rust/lance-io/src/object_store/providers/aws.rs b/rust/lance-io/src/object_store/providers/aws.rs index 9aad637bce2..8cad196f087 100644 --- a/rust/lance-io/src/object_store/providers/aws.rs +++ b/rust/lance-io/src/object_store/providers/aws.rs @@ -73,6 +73,12 @@ impl AwsStoreProvider { s3_storage_options.insert(AmazonS3ConfigKey::S3Express, true.to_string()); } + // Compute the metrics label before rewriting the url below, so it + // matches the prefix the registry uses to key this store. + #[cfg(feature = "metrics")] + let store_prefix = + self.calculate_object_store_prefix(base_path, Some(&storage_options.0))?; + // before creating the OSObjectStore we need to rewrite the url to drop ddb related parts base_path.set_scheme("s3").unwrap(); base_path.set_query(None); @@ -89,6 +95,13 @@ impl AwsStoreProvider { .with_retry(retry_config) .with_region(region); + #[cfg(feature = "metrics")] + { + builder = builder.with_http_connector( + crate::object_store::metrics::MeteringHttpConnector::new(store_prefix), + ); + } + Ok(Arc::new(builder.build()?) as Arc) } diff --git a/rust/lance-io/src/object_store/providers/azure.rs b/rust/lance-io/src/object_store/providers/azure.rs index e61f3f3b364..9de866bff91 100644 --- a/rust/lance-io/src/object_store/providers/azure.rs +++ b/rust/lance-io/src/object_store/providers/azure.rs @@ -192,6 +192,15 @@ impl AzureBlobStoreProvider { builder = builder.with_credentials(credentials); } + #[cfg(feature = "metrics")] + { + builder = builder.with_http_connector( + crate::object_store::metrics::MeteringHttpConnector::new( + self.calculate_object_store_prefix(base_path, Some(&storage_options.0))?, + ), + ); + } + Ok(Arc::new(builder.build()?) as Arc) } diff --git a/rust/lance-io/src/object_store/providers/gcp.rs b/rust/lance-io/src/object_store/providers/gcp.rs index f7f0a7672ff..35f79b66804 100644 --- a/rust/lance-io/src/object_store/providers/gcp.rs +++ b/rust/lance-io/src/object_store/providers/gcp.rs @@ -89,6 +89,15 @@ impl GcsStoreProvider { builder = builder.with_credentials(credential_provider); } + #[cfg(feature = "metrics")] + { + builder = builder.with_http_connector( + crate::object_store::metrics::MeteringHttpConnector::new( + self.calculate_object_store_prefix(base_path, Some(&storage_options.0))?, + ), + ); + } + Ok(Arc::new(builder.build()?) as Arc) } } diff --git a/rust/lance-io/src/object_store/providers/goosefs.rs b/rust/lance-io/src/object_store/providers/goosefs.rs index d6173571551..5d2a648522b 100644 --- a/rust/lance-io/src/object_store/providers/goosefs.rs +++ b/rust/lance-io/src/object_store/providers/goosefs.rs @@ -4,7 +4,6 @@ use std::collections::HashMap; use std::sync::Arc; -use object_store::path::Path; use object_store_opendal::OpendalStore; use opendal::{Operator, services::GooseFs}; use url::Url; @@ -27,7 +26,34 @@ const DEFAULT_GOOSEFS_PORT: u16 = 9200; /// - `host:port` is the GooseFS Master address (default port: 9200) /// - `/path` is the filesystem path within GooseFS /// -/// Configuration priority: storage_options > environment variables > URL authority > defaults +/// Path handling model (S3-style): +/// - The OpenDAL `root` is fixed to `/` (or a user-supplied cluster-wide base) +/// so that a single `Operator` can serve every dataset under the same +/// master. This keeps the `ObjectStoreRegistry` cache correct: two URLs +/// like `goosefs://host:9200/a.lance` and `goosefs://host:9200/b.lance` +/// share one store and each request carries its own object key. +/// - Path extraction relies on the default [`ObjectStoreProvider::extract_path`] +/// implementation, which returns the URL path (percent-decoded) as the key +/// passed to `ObjectStore::get`, `put`, etc. — mirroring how `s3://bucket/k` +/// yields key `k`. +/// +/// Supported configuration keys (via `storage_options` or environment variables, +/// resolved with priority: `storage_options` > env var > URL authority > default): +/// +/// | storage_options key | env var | purpose | +/// |---------------------------|-------------------------|-----------------------------------------------------------------------------------------------| +/// | `goosefs_master_addr` | `GOOSEFS_MASTER_ADDR` | Master gRPC address, e.g. `host:9200`. Supports HA: `addr1:port,addr2:port`. | +/// | `goosefs_root` | `GOOSEFS_ROOT` | Cluster-wide OpenDAL root shared by all datasets under the same master. Defaults to `/`. | +/// | `goosefs_write_type` | `GOOSEFS_WRITE_TYPE` | GooseFS write type (e.g. `MUST_CACHE`, `CACHE_THROUGH`, `THROUGH`, `ASYNC_THROUGH`). | +/// | `goosefs_block_size` | `GOOSEFS_BLOCK_SIZE` | GooseFS block size (bytes). Distinct from Lance's own `block_size`. | +/// | `goosefs_chunk_size` | `GOOSEFS_CHUNK_SIZE` | GooseFS chunk size (bytes) used by the client. | +/// | `goosefs_auth_type` | `GOOSEFS_AUTH_TYPE` | Authentication mode: `nosasl` or `simple`. | +/// | `goosefs_auth_username` | `GOOSEFS_AUTH_USERNAME` | Username for `simple` auth mode. | +/// +/// Note on `goosefs_root`: it is deliberately cluster-wide (not per-URL) so +/// that many datasets under the same master share a single cached `Operator`. +/// A custom root also participates in the `ObjectStoreRegistry` cache prefix, +/// so stores rooted at different subtrees do not collide. #[derive(Default, Debug)] pub struct GooseFsStoreProvider; @@ -79,6 +105,13 @@ impl GooseFsStoreProvider { .or_else(|| std::env::var(env_key).ok()) .filter(|v| !v.is_empty()) } + + /// Resolve the OpenDAL `root` for this Operator. See the file-level docs on + /// [`GooseFsStoreProvider`] for the semantics of `goosefs_root`. + fn resolve_root(storage_options: &StorageOptions) -> String { + Self::resolve_option(storage_options, "goosefs_root", "GOOSEFS_ROOT") + .unwrap_or_else(|| "/".to_string()) + } } #[async_trait::async_trait] @@ -90,16 +123,15 @@ impl ObjectStoreProvider for GooseFsStoreProvider { // Resolve master address let master_addr = Self::resolve_master_addr(&base_path, &storage_options)?; - // Extract root path from URL - let root = base_path.path().to_string(); + // Resolve a stable cluster-wide root. The URL path is *not* used here + // because it varies per dataset; per-request keys are supplied by + // `extract_path` instead. + let root = Self::resolve_root(&storage_options); // Build OpenDAL config map let mut config_map: HashMap = HashMap::new(); config_map.insert("master_addr".to_string(), master_addr); - - if !root.is_empty() && root != "/" { - config_map.insert("root".to_string(), root); - } + config_map.insert("root".to_string(), root); // Optional: write_type if let Some(wt) = @@ -163,26 +195,33 @@ impl ObjectStoreProvider for GooseFsStoreProvider { }) } - /// Extract the path relative to the root of the GooseFS filesystem. - /// - /// For GooseFS, the entire URL path is set as the OpenDAL `root` in `new_store`, - /// so the relative path returned here must be empty to avoid path duplication. - /// - /// `goosefs://host:port/data/file.lance` → root="/data/file.lance", extract_path="" - fn extract_path(&self, _url: &Url) -> Result { - Ok(Path::from("")) - } + // `extract_path` uses the default `ObjectStoreProvider` trait implementation: + // it percent-decodes the URL path and returns it as the object key, exactly + // like S3 does for `s3://bucket/key`. Overriding it here would only + // duplicate that behavior. See the file-level doc comment above for the + // full path-handling model. - /// Calculate the object store prefix for caching. + /// Calculate the object store prefix used as the registry cache key. /// - /// Format: `goosefs$host:port` - /// This ensures different GooseFS clusters get separate caches. + /// Format: `goosefs$host:port`. Because the OpenDAL root is now cluster- + /// wide (not per-URL), all datasets under the same master intentionally + /// share the same cached [`ObjectStore`]; the URL path is disambiguated + /// by [`Self::extract_path`] on each request. This is analogous to how + /// two `s3://bucket/a` and `s3://bucket/b` URLs share one store. fn calculate_object_store_prefix( &self, url: &Url, - _storage_options: Option<&HashMap>, + storage_options: Option<&HashMap>, ) -> Result { - Ok(format!("{}${}", url.scheme(), url.authority())) + // If a custom `goosefs_root` is provided, include it in the prefix so + // that stores built with different roots don't accidentally collide. + let opts = StorageOptions(storage_options.cloned().unwrap_or_default()); + let root = Self::resolve_root(&opts); + if root == "/" { + Ok(format!("{}${}", url.scheme(), url.authority())) + } else { + Ok(format!("{}${}#{}", url.scheme(), url.authority(), root)) + } } } @@ -191,38 +230,42 @@ mod tests { use super::*; #[test] - fn test_goosefs_store_path() { + fn test_goosefs_extract_path_basic() { let provider = GooseFsStoreProvider; - let url = Url::parse("goosefs://10.0.0.1:9200/data/embeddings.lance").unwrap(); let path = provider.extract_path(&url).unwrap(); - // extract_path returns empty because the full path is used as OpenDAL root - assert_eq!(path.to_string(), ""); + assert_eq!(path.to_string(), "data/embeddings.lance"); } #[test] - fn test_goosefs_store_root_path() { + fn test_goosefs_extract_path_root() { let provider = GooseFsStoreProvider; - let url = Url::parse("goosefs://10.0.0.1:9200/").unwrap(); let path = provider.extract_path(&url).unwrap(); assert_eq!(path.to_string(), ""); } #[test] - fn test_goosefs_store_deep_path() { + fn test_goosefs_extract_path_deep() { let provider = GooseFsStoreProvider; - let url = Url::parse("goosefs://master:9200/a/b/c/d.lance").unwrap(); let path = provider.extract_path(&url).unwrap(); - // All path components are in the OpenDAL root, extract_path is empty - assert_eq!(path.to_string(), ""); + assert_eq!(path.to_string(), "a/b/c/d.lance"); } #[test] - fn test_calculate_object_store_prefix() { + fn test_goosefs_extract_path_percent_decoded() { + // The URL contains a percent-encoded space; extract_path must decode + // it once so the ObjectStore layer does not double-encode later. let provider = GooseFsStoreProvider; + let url = Url::parse("goosefs://master:9200/dir/with%20space/f.lance").unwrap(); + let path = provider.extract_path(&url).unwrap(); + assert_eq!(path.to_string(), "dir/with space/f.lance"); + } + #[test] + fn test_calculate_object_store_prefix_default_root() { + let provider = GooseFsStoreProvider; let url = Url::parse("goosefs://10.0.0.1:9200/data").unwrap(); let prefix = provider.calculate_object_store_prefix(&url, None).unwrap(); assert_eq!(prefix, "goosefs$10.0.0.1:9200"); @@ -231,12 +274,69 @@ mod tests { #[test] fn test_calculate_object_store_prefix_with_hostname() { let provider = GooseFsStoreProvider; - let url = Url::parse("goosefs://myhost:9200/data").unwrap(); let prefix = provider.calculate_object_store_prefix(&url, None).unwrap(); assert_eq!(prefix, "goosefs$myhost:9200"); } + /// Regression test: two URLs pointing at different datasets under the + /// same master must produce the *same* cache prefix so they share one + /// Operator, and correctness must come from `extract_path` returning + /// distinct keys — never from a per-URL root baked into the prefix. + #[test] + fn test_prefix_shared_across_datasets_same_master() { + let provider = GooseFsStoreProvider; + let url_a = Url::parse("goosefs://10.0.0.1:9200/repro/a.lance").unwrap(); + let url_b = Url::parse("goosefs://10.0.0.1:9200/repro/b.lance").unwrap(); + + let pa = provider + .calculate_object_store_prefix(&url_a, None) + .unwrap(); + let pb = provider + .calculate_object_store_prefix(&url_b, None) + .unwrap(); + assert_eq!(pa, pb, "same master must share one cache prefix"); + + // Extracted keys must differ so the shared Operator can route + // requests to the correct dataset. + assert_ne!( + provider.extract_path(&url_a).unwrap(), + provider.extract_path(&url_b).unwrap(), + "distinct URLs must yield distinct object keys", + ); + } + + /// Different masters must never share a cache entry. + #[test] + fn test_prefix_isolated_across_masters() { + let provider = GooseFsStoreProvider; + let u1 = Url::parse("goosefs://host-a:9200/x.lance").unwrap(); + let u2 = Url::parse("goosefs://host-b:9200/x.lance").unwrap(); + assert_ne!( + provider.calculate_object_store_prefix(&u1, None).unwrap(), + provider.calculate_object_store_prefix(&u2, None).unwrap(), + ); + } + + /// A user-supplied `goosefs_root` participates in the cache prefix so + /// stores rooted at different subtrees don't collide. + #[test] + fn test_prefix_includes_custom_root() { + let provider = GooseFsStoreProvider; + let url = Url::parse("goosefs://host:9200/x.lance").unwrap(); + + let default_prefix = provider.calculate_object_store_prefix(&url, None).unwrap(); + let custom_opts: HashMap = + HashMap::from([("goosefs_root".to_string(), "/tenant-a".to_string())]); + let custom_prefix = provider + .calculate_object_store_prefix(&url, Some(&custom_opts)) + .unwrap(); + + assert_eq!(default_prefix, "goosefs$host:9200"); + assert_eq!(custom_prefix, "goosefs$host:9200#/tenant-a"); + assert_ne!(default_prefix, custom_prefix); + } + #[test] fn test_resolve_master_addr_from_url() { let url = Url::parse("goosefs://10.0.0.1:9200/data").unwrap(); @@ -263,4 +363,19 @@ mod tests { let addr = GooseFsStoreProvider::resolve_master_addr(&url, &storage_options).unwrap(); assert_eq!(addr, "10.0.0.2:9200,10.0.0.3:9200"); } + + #[test] + fn test_resolve_root_defaults_to_slash() { + let opts = StorageOptions(HashMap::new()); + assert_eq!(GooseFsStoreProvider::resolve_root(&opts), "/"); + } + + #[test] + fn test_resolve_root_from_storage_options() { + let opts = StorageOptions(HashMap::from([( + "goosefs_root".to_string(), + "/tenant-a".to_string(), + )])); + assert_eq!(GooseFsStoreProvider::resolve_root(&opts), "/tenant-a"); + } } diff --git a/rust/lance-io/src/object_store/storage_options.rs b/rust/lance-io/src/object_store/storage_options.rs index 640e0a26b1e..fa5f4995afd 100644 --- a/rust/lance-io/src/object_store/storage_options.rs +++ b/rust/lance-io/src/object_store/storage_options.rs @@ -79,6 +79,15 @@ pub trait StorageOptionsProvider: Send + Sync + fmt::Debug { /// If [`EXPIRES_AT_MILLIS_KEY`] is not provided, the options are considered to never expire. async fn fetch_storage_options(&self) -> Result>>; + /// Fetch fresh storage options, bypassing caches along the chain. + /// + /// Providers that serve from an upstream cache (e.g. base-scoped wrappers + /// reading through a parent accessor) override this to force the upstream + /// to re-fetch. Defaults to [`Self::fetch_storage_options`]. + async fn force_fetch_storage_options(&self) -> Result>> { + self.fetch_storage_options().await + } + /// Return a human-readable unique identifier for this provider instance /// /// This is used for equality comparison and hashing in the object store registry. @@ -155,6 +164,104 @@ impl StorageOptionsProvider for LanceNamespaceStorageOptionsProvider { } } +/// Prefix marking a storage option as scoped to a single registered base path. +/// +/// A key of the form `base_.` applies `` only to the base path +/// with manifest id ``, overriding the shared (unscoped) options for that +/// base. For example `base_1.account_key = abc` gives the base with id 1 the +/// option `account_key = abc` while it inherits every unscoped option. +pub const BASE_SCOPED_OPTION_PREFIX: &str = "base_"; + +/// Parse a base-scoped storage option key of the form `base_.`. +/// +/// Returns `Some((base_id, key))` only for keys that match the convention +/// exactly: the `base_` prefix, an all-digit u32 base id, a `.` separator, and +/// a non-empty remainder. Any other key (e.g. `base_url`, `base_x.key`, +/// `base_1.`) is not base-scoped. +pub fn parse_base_scoped_key(key: &str) -> Option<(u32, &str)> { + let rest = key.strip_prefix(BASE_SCOPED_OPTION_PREFIX)?; + let (id_str, scoped_key) = rest.split_once('.')?; + if scoped_key.is_empty() || id_str.is_empty() || !id_str.bytes().all(|b| b.is_ascii_digit()) { + return None; + } + let id = id_str.parse::().ok()?; + Some((id, scoped_key)) +} + +/// Returns true if any key in `options` is base-scoped (`base_.`). +pub fn has_base_scoped_options(options: &HashMap) -> bool { + options + .keys() + .any(|key| parse_base_scoped_key(key).is_some()) +} + +/// Resolve the effective storage options for one base path scope. +/// +/// All base-scoped keys are removed from the result. When `base_id` is +/// `Some(id)`, entries scoped to that id are overlaid on the unscoped options, +/// adding or overriding keys. `None` resolves the default scope (the primary +/// dataset base), which simply drops every base-scoped entry. +pub fn resolve_base_scoped_options( + options: &HashMap, + base_id: Option, +) -> HashMap { + let mut resolved = HashMap::with_capacity(options.len()); + let mut overrides = Vec::new(); + for (key, value) in options { + match parse_base_scoped_key(key) { + Some((id, scoped_key)) => { + if Some(id) == base_id { + overrides.push((scoped_key.to_string(), value.clone())); + } + } + None => { + resolved.insert(key.clone(), value.clone()); + } + } + } + resolved.extend(overrides); + resolved +} + +/// A [`StorageOptionsProvider`] that resolves another accessor's options for a +/// single base path scope. +/// +/// Fetching through this provider first refreshes the parent accessor when its +/// options have expired, then resolves the refreshed options for the scope. As +/// a result, dynamically vended per-base credentials (e.g. a namespace server +/// returning `base_.` entries in one flat map) stay fresh per base. +#[derive(Debug)] +pub struct BaseScopedStorageOptionsProvider { + inner: Arc, + base_id: Option, +} + +impl BaseScopedStorageOptionsProvider { + pub fn new(inner: Arc, base_id: Option) -> Self { + Self { inner, base_id } + } +} + +#[async_trait] +impl StorageOptionsProvider for BaseScopedStorageOptionsProvider { + async fn fetch_storage_options(&self) -> Result>> { + let options = self.inner.get_storage_options().await?; + Ok(Some(resolve_base_scoped_options(&options.0, self.base_id))) + } + + async fn force_fetch_storage_options(&self) -> Result>> { + let options = self.inner.refresh_storage_options().await?; + Ok(Some(resolve_base_scoped_options(&options.0, self.base_id))) + } + + fn provider_id(&self) -> String { + match self.base_id { + Some(id) => format!("base-scoped[base_id={}]({})", id, self.inner.accessor_id()), + None => format!("base-scoped[default]({})", self.inner.accessor_id()), + } + } +} + /// Unified access to storage options with automatic caching and refresh /// /// This struct bundles static storage options with an optional dynamic provider, @@ -184,6 +291,11 @@ pub struct StorageOptionsAccessor { /// Duration before expiry to trigger refresh refresh_offset: Duration, + + /// True when this accessor was produced by [`Self::scoped_to_base`]; its + /// options are already resolved for one base path scope, so scoping again + /// is a no-op. + scope_resolved: bool, } impl fmt::Debug for StorageOptionsAccessor { @@ -212,14 +324,33 @@ impl StorageOptionsAccessor { .unwrap_or(Duration::from_millis(DEFAULT_REFRESH_OFFSET_MILLIS)) } + /// Effective expiration of a raw options map: the minimum of the unscoped + /// `expires_at_millis` and every `base_.expires_at_millis` entry. + /// + /// A flat map may vend per-base credentials that expire before the shared + /// ones. Refreshing when the earliest credential is due keeps base-scoped + /// accessors (which refresh through this accessor) from re-resolving stale + /// per-base credentials out of a still-"valid" cache. + fn effective_expires_at_millis(options: &HashMap) -> Option { + options + .iter() + .filter(|(key, _)| { + key.as_str() == EXPIRES_AT_MILLIS_KEY + || matches!( + parse_base_scoped_key(key), + Some((_, scoped_key)) if scoped_key == EXPIRES_AT_MILLIS_KEY + ) + }) + .filter_map(|(_, value)| value.parse::().ok()) + .min() + } + /// Create an accessor with only static options (no refresh capability) /// /// The returned accessor will always return the provided options. /// This is useful when credentials don't expire or are managed externally. pub fn with_static_options(options: HashMap) -> Self { - let expires_at_millis = options - .get(EXPIRES_AT_MILLIS_KEY) - .and_then(|s| s.parse::().ok()); + let expires_at_millis = Self::effective_expires_at_millis(&options); let refresh_offset = Self::extract_refresh_offset(&options); Self { @@ -230,6 +361,7 @@ impl StorageOptionsAccessor { expires_at_millis, }))), refresh_offset, + scope_resolved: false, } } @@ -247,6 +379,7 @@ impl StorageOptionsAccessor { provider: Some(provider), cache: Arc::new(RwLock::new(None)), refresh_offset: Duration::from_millis(DEFAULT_REFRESH_OFFSET_MILLIS), + scope_resolved: false, } } @@ -263,9 +396,7 @@ impl StorageOptionsAccessor { initial_options: HashMap, provider: Arc, ) -> Self { - let expires_at_millis = initial_options - .get(EXPIRES_AT_MILLIS_KEY) - .and_then(|s| s.parse::().ok()); + let expires_at_millis = Self::effective_expires_at_millis(&initial_options); let refresh_offset = Self::extract_refresh_offset(&initial_options); Self { @@ -276,6 +407,7 @@ impl StorageOptionsAccessor { expires_at_millis, }))), refresh_offset, + scope_resolved: false, } } @@ -306,8 +438,9 @@ impl StorageOptionsAccessor { /// Fetch fresh options from the provider and update the cache. /// /// This bypasses the cache for callers that need to validate provider-vended - /// credentials even when initial metadata has no expiration. - #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] + /// credentials even when initial metadata has no expiration. The force + /// propagates through provider chains (e.g. base-scoped wrappers), so the + /// origin provider is re-queried even when intermediate caches are valid. pub(crate) async fn refresh_storage_options(&self) -> Result { let Some(provider) = &self.provider else { return self.get_storage_options().await; @@ -318,7 +451,7 @@ impl StorageOptionsAccessor { provider.provider_id() ); - let storage_options_map = provider.fetch_storage_options().await.map_err(|e| { + let storage_options_map = provider.force_fetch_storage_options().await.map_err(|e| { Error::io_source(Box::new(std::io::Error::other(format!( "Failed to fetch storage options: {}", e @@ -336,9 +469,7 @@ impl StorageOptionsAccessor { return Ok(super::StorageOptions(HashMap::new())); }; - let expires_at_millis = options - .get(EXPIRES_AT_MILLIS_KEY) - .and_then(|s| s.parse::().ok()); + let expires_at_millis = Self::effective_expires_at_millis(&options); let mut cache = self.cache.write().await; *cache = Some(CachedStorageOptions { @@ -409,9 +540,7 @@ impl StorageOptionsAccessor { return Ok(Some(super::StorageOptions(HashMap::new()))); }; - let expires_at_millis = options - .get(EXPIRES_AT_MILLIS_KEY) - .and_then(|s| s.parse::().ok()); + let expires_at_millis = Self::effective_expires_at_millis(&options); if let Some(expires_at) = expires_at_millis { let now_ms = SystemTime::now() @@ -493,6 +622,50 @@ impl StorageOptionsAccessor { } } + /// Resolve this accessor for a single base path scope. + /// + /// Storage options may carry base-scoped entries (`base_.`) that + /// apply only to one registered base path. The returned accessor resolves + /// options for `base_id`: entries scoped to that base overlay the unscoped + /// defaults, and all other scoped entries are dropped. `None` resolves the + /// default scope used for the primary dataset base. + /// + /// A static accessor whose options contain no base-scoped entries is + /// returned unchanged, preserving accessor identity (and thus object store + /// registry cache keys). A provider-backed accessor is always wrapped + /// through [`BaseScopedStorageOptionsProvider`] — fetched options may vend + /// base-scoped entries at any refresh, even when the initial options carry + /// none — so refreshed options are re-resolved for the scope on every + /// fetch. Accessors already produced by this method are returned unchanged. + pub fn scoped_to_base(self: &Arc, base_id: Option) -> Arc { + if self.scope_resolved { + return self.clone(); + } + if self.has_provider() { + let provider = Arc::new(BaseScopedStorageOptionsProvider::new(self.clone(), base_id)); + let mut scoped = match self.initial_storage_options() { + Some(initial) => Self::with_initial_and_provider( + resolve_base_scoped_options(initial, base_id), + provider, + ), + None => Self::with_provider(provider), + }; + scoped.scope_resolved = true; + Arc::new(scoped) + } else { + match self.initial_storage_options() { + Some(initial) if has_base_scoped_options(initial) => { + let mut scoped = + Self::with_static_options(resolve_base_scoped_options(initial, base_id)); + scoped.scope_resolved = true; + Arc::new(scoped) + } + // Static options never change, so there is nothing to scope. + _ => self.clone(), + } + } + } + /// Check if this accessor has a dynamic provider pub fn has_provider(&self) -> bool { self.provider.is_some() @@ -748,4 +921,322 @@ mod tests { accessor.get_storage_options().await.unwrap(); assert_eq!(mock_provider.get_call_count().await, 1); } + + #[test] + fn test_parse_base_scoped_key() { + assert_eq!( + parse_base_scoped_key("base_1.account_key"), + Some((1, "account_key")) + ); + assert_eq!( + parse_base_scoped_key("base_12.headers.x-ms-version"), + Some((12, "headers.x-ms-version")) + ); + assert_eq!(parse_base_scoped_key("base_0.region"), Some((0, "region"))); + + // Not base-scoped keys + assert_eq!(parse_base_scoped_key("account_key"), None); + assert_eq!(parse_base_scoped_key("base_url"), None); + assert_eq!(parse_base_scoped_key("base_hot.account_key"), None); + assert_eq!(parse_base_scoped_key("base_1x.account_key"), None); + assert_eq!(parse_base_scoped_key("base_+1.account_key"), None); + assert_eq!(parse_base_scoped_key("base_.account_key"), None); + assert_eq!(parse_base_scoped_key("base_1."), None); + assert_eq!(parse_base_scoped_key("base_1"), None); + // Overflows u32 + assert_eq!(parse_base_scoped_key("base_4294967296.key"), None); + } + + #[test] + fn test_resolve_base_scoped_options() { + let options = HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("account_key".to_string(), "shared-key".to_string()), + ("base_1.account_key".to_string(), "base1-key".to_string()), + ("base_2.account_key".to_string(), "base2-key".to_string()), + ("base_2.endpoint".to_string(), "http://b2".to_string()), + ]); + assert!(has_base_scoped_options(&options)); + + let base1 = resolve_base_scoped_options(&options, Some(1)); + assert_eq!( + base1, + HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("account_key".to_string(), "base1-key".to_string()), + ]) + ); + + let base2 = resolve_base_scoped_options(&options, Some(2)); + assert_eq!( + base2, + HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("account_key".to_string(), "base2-key".to_string()), + ("endpoint".to_string(), "http://b2".to_string()), + ]) + ); + + // A base without scoped entries inherits only the unscoped options + let base3 = resolve_base_scoped_options(&options, Some(3)); + assert_eq!( + base3, + HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("account_key".to_string(), "shared-key".to_string()), + ]) + ); + + // The default scope drops every scoped entry + let default = resolve_base_scoped_options(&options, None); + assert_eq!(default, base3); + + assert!(!has_base_scoped_options(&HashMap::from([( + "account_key".to_string(), + "shared-key".to_string() + )]))); + } + + #[tokio::test] + async fn test_scoped_to_base_identity_and_idempotency() { + // Static accessors without scoped keys are returned unchanged. + let accessor = Arc::new(StorageOptionsAccessor::with_static_options(HashMap::from( + [("account_key".to_string(), "shared-key".to_string())], + ))); + assert!(Arc::ptr_eq(&accessor.scoped_to_base(Some(1)), &accessor)); + assert!(Arc::ptr_eq(&accessor.scoped_to_base(None), &accessor)); + + // Scoping an already-scoped accessor is a no-op (the registry choke + // point re-applies the default scope to every params it sees). + let scoped = Arc::new(StorageOptionsAccessor::with_static_options(HashMap::from( + [ + ("account_key".to_string(), "shared-key".to_string()), + ("base_1.account_key".to_string(), "base1-key".to_string()), + ], + ))) + .scoped_to_base(Some(1)); + assert!(Arc::ptr_eq(&scoped.scoped_to_base(None), &scoped)); + + let provider_scoped = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + MockStorageOptionsProvider::new(None), + ))) + .scoped_to_base(Some(1)); + assert!(Arc::ptr_eq( + &provider_scoped.scoped_to_base(None), + &provider_scoped + )); + } + + #[tokio::test] + async fn test_scoped_to_base_provider_only_resolves_vended_options() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + // No initial options: scoped entries arrive only through the provider. + let provider = Arc::new(MockBaseScopedVendingProvider { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis: 600_000, + }); + let parent = Arc::new(StorageOptionsAccessor::with_provider(provider.clone())); + + let base1 = parent.scoped_to_base(Some(1)); + assert!(!Arc::ptr_eq(&base1, &parent)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_1"); + assert!(!result.0.contains_key("base_1.account_key")); + + let default = parent.scoped_to_base(None); + let result = default.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "SHARED_1"); + assert!(!result.0.contains_key("base_1.account_key")); + + // Both scopes were served from one parent fetch. + assert_eq!(*provider.call_count.read().await, 1); + } + + #[tokio::test] + async fn test_scoped_earlier_base_expiry_refreshes_parent() { + MockClock::set_system_time(Duration::from_secs(100_000)); + let now_ms = MockClock::system_time().as_millis() as u64; + + // Base 1 credentials expire before the shared ones; the parent must + // refresh when the earliest credential is due, or the scoped accessor + // would keep re-resolving stale base-1 credentials from a still- + // "valid" parent cache. + let provider = Arc::new(MockBaseScopedVendingProvider { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis: 600_000, + }); + let initial = HashMap::from([ + ("account_key".to_string(), "SHARED_0".to_string()), + ("base_1.account_key".to_string(), "BASE1_0".to_string()), + ( + EXPIRES_AT_MILLIS_KEY.to_string(), + (now_ms + 600_000).to_string(), + ), + ( + format!("base_1.{}", EXPIRES_AT_MILLIS_KEY), + (now_ms + 120_000).to_string(), + ), + ]); + let parent = Arc::new(StorageOptionsAccessor::with_initial_and_provider( + initial, + provider.clone(), + )); + + let base1 = parent.scoped_to_base(Some(1)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_0"); + assert_eq!(*provider.call_count.read().await, 0); + + // Past the base-1 expiry but before the shared expiry: the parent's + // effective expiry is the earlier one, so the refresh chain fetches + // fresh credentials instead of re-serving BASE1_0. + MockClock::set_system_time(Duration::from_secs(100_000 + 121)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_1"); + assert_eq!(*provider.call_count.read().await, 1); + } + + #[tokio::test] + async fn test_scoped_to_base_static() { + let accessor = Arc::new(StorageOptionsAccessor::with_static_options(HashMap::from( + [ + ("account_key".to_string(), "shared-key".to_string()), + ("base_1.account_key".to_string(), "base1-key".to_string()), + ], + ))); + + let base1 = accessor.scoped_to_base(Some(1)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!( + result.0, + HashMap::from([("account_key".to_string(), "base1-key".to_string())]) + ); + assert!(!base1.has_provider()); + + let default = accessor.scoped_to_base(None); + let result = default.get_storage_options().await.unwrap(); + assert_eq!( + result.0, + HashMap::from([("account_key".to_string(), "shared-key".to_string())]) + ); + + // Scoped accessor ids are stable across derivations and distinct per scope + assert_eq!( + accessor.scoped_to_base(Some(1)).accessor_id(), + base1.accessor_id() + ); + assert_ne!(base1.accessor_id(), default.accessor_id()); + assert_ne!(base1.accessor_id(), accessor.accessor_id()); + } + + #[derive(Debug)] + struct MockBaseScopedVendingProvider { + call_count: Arc>, + expires_in_millis: u64, + } + + #[async_trait] + impl StorageOptionsProvider for MockBaseScopedVendingProvider { + async fn fetch_storage_options(&self) -> Result>> { + let count = { + let mut c = self.call_count.write().await; + *c += 1; + *c + }; + let now_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() as u64; + Ok(Some(HashMap::from([ + ("account_key".to_string(), format!("SHARED_{}", count)), + ("base_1.account_key".to_string(), format!("BASE1_{}", count)), + ( + EXPIRES_AT_MILLIS_KEY.to_string(), + (now_ms + self.expires_in_millis).to_string(), + ), + ]))) + } + + fn provider_id(&self) -> String { + "MockBaseScopedVendingProvider".to_string() + } + } + + #[tokio::test] + async fn test_scoped_to_base_refreshes_through_parent() { + MockClock::set_system_time(Duration::from_secs(100_000)); + let now_ms = MockClock::system_time().as_millis() as u64; + + let provider = Arc::new(MockBaseScopedVendingProvider { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis: 600_000, + }); + let initial = HashMap::from([ + ("account_key".to_string(), "SHARED_0".to_string()), + ("base_1.account_key".to_string(), "BASE1_0".to_string()), + ( + EXPIRES_AT_MILLIS_KEY.to_string(), + (now_ms + 600_000).to_string(), + ), + ]); + let parent = Arc::new(StorageOptionsAccessor::with_initial_and_provider( + initial, + provider.clone(), + )); + + let base1 = parent.scoped_to_base(Some(1)); + let default = parent.scoped_to_base(None); + assert!(base1.has_provider()); + + // Initial options are resolved per scope without fetching + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_0"); + assert!(!result.0.contains_key("base_1.account_key")); + let result = default.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "SHARED_0"); + assert_eq!(*provider.call_count.read().await, 0); + + // Expire the vended options; the scoped accessor refreshes through the + // parent and re-resolves the refreshed options for its scope. + MockClock::set_system_time(Duration::from_secs(100_000 + 601)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_1"); + assert_eq!(*provider.call_count.read().await, 1); + + // The parent refresh is shared: other scopes see it without refetching + let result = default.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "SHARED_1"); + assert_eq!(*provider.call_count.read().await, 1); + } + + #[tokio::test] + async fn test_scoped_forced_refresh_reaches_origin_provider() { + MockClock::set_system_time(Duration::from_secs(100_000)); + let now_ms = MockClock::system_time().as_millis() as u64; + + let provider = Arc::new(MockBaseScopedVendingProvider { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis: 600_000, + }); + let initial = HashMap::from([ + ("account_key".to_string(), "SHARED_0".to_string()), + ("base_1.account_key".to_string(), "BASE1_0".to_string()), + ( + EXPIRES_AT_MILLIS_KEY.to_string(), + (now_ms + 600_000).to_string(), + ), + ]); + let parent = Arc::new(StorageOptionsAccessor::with_initial_and_provider( + initial, + provider.clone(), + )); + let base1 = parent.scoped_to_base(Some(1)); + + // A forced refresh must reach the origin provider even though both the + // scoped and the parent caches are still valid. + let result = base1.refresh_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_1"); + assert_eq!(*provider.call_count.read().await, 1); + } } diff --git a/rust/lance-io/src/object_writer.rs b/rust/lance-io/src/object_writer.rs index 0fd0a30f9e7..6dd6c4d992f 100644 --- a/rust/lance-io/src/object_writer.rs +++ b/rust/lance-io/src/object_writer.rs @@ -392,25 +392,23 @@ impl AsyncWrite for ObjectWriter { let fut = Box::pin(async move { store.put_multipart(path.as_ref()).await }); self.state = UploadState::CreatingUpload(fut); } + // TODO: Make max concurrency configurable from storage options. UploadState::InProgress { upload, part_idx, futures, .. - } => { - // TODO: Make max concurrency configurable from storage options. - if futures.len() < max_upload_parallelism() { - let data = Self::next_part_buffer( - &mut mut_self.buffer, - *part_idx, - mut_self.use_constant_size_upload_parts, - ); - futures.spawn( - Self::put_part(upload.as_mut(), data, *part_idx, None) - .instrument(tracing::Span::current()), - ); - *part_idx += 1; - } + } if futures.len() < max_upload_parallelism() => { + let data = Self::next_part_buffer( + &mut mut_self.buffer, + *part_idx, + mut_self.use_constant_size_upload_parts, + ); + futures.spawn( + Self::put_part(upload.as_mut(), data, *part_idx, None) + .instrument(tracing::Span::current()), + ); + *part_idx += 1; } _ => {} } diff --git a/rust/lance-io/src/scheduler.rs b/rust/lance-io/src/scheduler.rs index 07edcfbf70d..c8b3a09e17c 100644 --- a/rust/lance-io/src/scheduler.rs +++ b/rust/lance-io/src/scheduler.rs @@ -29,6 +29,7 @@ mod lite; const BACKPRESSURE_MIN: u64 = 5; // Don't log backpressure warnings more than once / minute const BACKPRESSURE_DEBOUNCE: u64 = 60; +const SCHEDULER_STATE_EVENT_TARGET: &str = "lance_io::scheduler::state"; // Global counter of how many IOPS we have issued static IOPS_COUNTER: AtomicU64 = AtomicU64::new(0); @@ -66,6 +67,10 @@ impl PrioritiesInFlight { self.in_flight.first().copied().unwrap_or(u128::MAX) } + fn contains(&self, prio: u128) -> bool { + self.in_flight.binary_search(&prio).is_ok() + } + fn push(&mut self, prio: u128) { let pos = match self.in_flight.binary_search(&prio) { Ok(pos) => pos, @@ -79,11 +84,23 @@ impl PrioritiesInFlight { self.in_flight.remove(pos); } } + + fn len(&self) -> usize { + self.in_flight.len() + } + + fn is_empty(&self) -> bool { + self.in_flight.is_empty() + } } struct IoQueueState { + // The configured number of IOPS that can be issued concurrently. + io_capacity: u32, // Number of IOPS we can issue concurrently before pausing I/O iops_avail: u32, + // The configured byte budget for unread I/O. + io_buffer_size: u64, // Number of bytes we are allowed to buffer in memory before pausing I/O // // This can dip below 0 due to I/O prioritization @@ -106,7 +123,9 @@ struct IoQueueState { impl IoQueueState { fn new(io_capacity: u32, io_buffer_size: u64) -> Self { Self { + io_capacity, iops_avail: io_capacity, + io_buffer_size, bytes_avail: io_buffer_size as i64, pending_requests: BinaryHeap::new(), priorities_in_flight: PrioritiesInFlight::new(io_capacity), @@ -117,6 +136,65 @@ impl IoQueueState { } } + fn scheduler_state_event(&self) -> Option { + if !tracing::enabled!(target: SCHEDULER_STATE_EVENT_TARGET, tracing::Level::TRACE) { + return None; + } + + let pending_bytes = self + .pending_requests + .iter() + .map(IoTask::num_bytes) + .sum::(); + let head_task = self.pending_requests.peek(); + let min_in_flight_priority = if self.priorities_in_flight.is_empty() { + None + } else { + Some(self.priorities_in_flight.min_in_flight()) + }; + let head_task_priority_bypass = head_task.map(|task| { + self.no_backpressure + || task.bypass_backpressure + || task.priority <= self.priorities_in_flight.min_in_flight() + }); + let head_task_blocked_by_iops = head_task.map(|_| self.iops_avail == 0); + let head_task_blocked_by_bytes = head_task.map(|task| { + let bypasses_bytes = self.no_backpressure + || task.bypass_backpressure + || task.priority <= self.priorities_in_flight.min_in_flight(); + !bypasses_bytes && task.num_bytes() as i64 > self.bytes_avail + }); + let head_task_can_deliver = head_task.map(|task| self.can_deliver_without_warning(task)); + let head_task_bytes = head_task.map(IoTask::num_bytes); + let (head_task_priority_high, head_task_priority_low) = + split_priority(head_task.map(|task| task.priority)); + let (min_in_flight_priority_high, min_in_flight_priority_low) = + split_priority(min_in_flight_priority); + + Some(SchedulerStateEvent { + queue_kind: "standard", + io_capacity: u64::from(self.io_capacity), + iops_available: u64::from(self.iops_avail), + active_iops: u64::from(self.io_capacity.saturating_sub(self.iops_avail)), + pending_iops: self.pending_requests.len() as u64, + pending_bytes, + bytes_available: self.bytes_avail, + bytes_reserved: self.io_buffer_size as i64 - self.bytes_avail, + io_buffer_size_bytes: self.io_buffer_size, + priorities_in_flight: self.priorities_in_flight.len() as u64, + no_backpressure: self.no_backpressure, + head_task_bytes, + head_task_priority_high, + head_task_priority_low, + min_in_flight_priority_high, + min_in_flight_priority_low, + head_task_can_deliver, + head_task_priority_bypass, + head_task_blocked_by_iops, + head_task_blocked_by_bytes, + }) + } + fn warn_if_needed(&self) { let seconds_elapsed = self.start.elapsed().as_secs(); let last_warn = self.last_warn.load(Ordering::Acquire); @@ -136,18 +214,33 @@ impl IoQueueState { } fn can_deliver(&self, task: &IoTask) -> bool { + let can_deliver = self.can_deliver_without_warning(task); + if !can_deliver + && self.iops_avail > 0 + && !(self.no_backpressure + || task.bypass_backpressure + || task.priority <= self.priorities_in_flight.min_in_flight()) + && task.num_bytes() as i64 > self.bytes_avail + { + self.warn_if_needed(); + } + can_deliver + } + + fn can_deliver_without_warning(&self, task: &IoTask) -> bool { if self.iops_avail == 0 { false } else if self.no_backpressure || task.bypass_backpressure || task.priority <= self.priorities_in_flight.min_in_flight() + // Chunks from an admitted logical request must keep moving. A + // higher-priority request may be scheduled later and remain + // unconsumed while the caller awaits this request. + || self.priorities_in_flight.contains(task.priority) { true - } else if task.num_bytes() as i64 > self.bytes_avail { - self.warn_if_needed(); - false } else { - true + task.num_bytes() as i64 <= self.bytes_avail } } @@ -183,13 +276,15 @@ struct IoQueue { state: Mutex, // Used to signal new I/O requests have arrived that might potentially be runnable notify: Notify, + stats: IoStats, } impl IoQueue { - fn new(io_capacity: u32, io_buffer_size: u64) -> Self { + fn new(io_capacity: u32, io_buffer_size: u64, stats: IoStats) -> Self { Self { state: Mutex::new(IoQueueState::new(io_capacity, io_buffer_size)), notify: Notify::new(), + stats, } } @@ -200,9 +295,12 @@ impl IoQueue { task.priority >> 64, task.priority & 0xFFFFFFFFFFFFFFFF ); - let mut state = self.state.lock().unwrap(); - state.pending_requests.push(task); - drop(state); + let event = { + let mut state = self.state.lock().unwrap(); + state.pending_requests.push(task); + state.scheduler_state_event() + }; + emit_scheduler_state_event(event, &self.stats); self.notify.notify_one(); } @@ -212,6 +310,9 @@ impl IoQueue { { let mut state = self.state.lock().unwrap(); if let Some(task) = state.next_task() { + let event = state.scheduler_state_event(); + drop(state); + emit_scheduler_state_event(event, &self.stats); return Some(task); } @@ -225,29 +326,39 @@ impl IoQueue { } fn on_iop_complete(&self) { - let mut state = self.state.lock().unwrap(); - state.iops_avail += 1; - drop(state); + let event = { + let mut state = self.state.lock().unwrap(); + state.iops_avail += 1; + state.scheduler_state_event() + }; + emit_scheduler_state_event(event, &self.stats); self.notify.notify_one(); } fn on_bytes_consumed(&self, bytes: u64, priority: u128, num_reqs: usize) { - let mut state = self.state.lock().unwrap(); - state.bytes_avail += bytes as i64; - for _ in 0..num_reqs { - state.priorities_in_flight.remove(priority); - } - drop(state); + let event = { + let mut state = self.state.lock().unwrap(); + state.bytes_avail += bytes as i64; + for _ in 0..num_reqs { + state.priorities_in_flight.remove(priority); + } + state.scheduler_state_event() + }; + emit_scheduler_state_event(event, &self.stats); self.notify.notify_one(); } fn close(&self) { - let mut state = self.state.lock().unwrap(); - state.done_scheduling = true; - let pending_requests = std::mem::take(&mut state.pending_requests); - drop(state); + let (pending_requests, event) = { + let mut state = self.state.lock().unwrap(); + state.done_scheduling = true; + let pending_requests = std::mem::take(&mut state.pending_requests); + let event = state.scheduler_state_event(); + (pending_requests, event) + }; + emit_scheduler_state_event(event, &self.stats); for request in pending_requests { request.cancel(); } @@ -269,6 +380,9 @@ struct MutableBatch { err: Option>, // When true, report 0 bytes consumed so the backpressure budget is unaffected bypass_backpressure: bool, + // Queue the batch's backpressure reservation is refunded to once its response + // is delivered or discarded (see `Response`'s `Drop`). + io_queue: Arc, } impl MutableBatch { @@ -278,6 +392,7 @@ impl MutableBatch { priority: u128, num_reqs: usize, bypass_backpressure: bool, + io_queue: Arc, ) -> Self { Self { when_done: Some(when_done), @@ -287,6 +402,7 @@ impl MutableBatch { num_reqs, err: None, bypass_backpressure, + io_queue, } } } @@ -308,7 +424,8 @@ impl Drop for MutableBatch { // We don't really care if no one is around to receive it, just let // the result go out of scope and get cleaned up let response = Response { - data: result, + data: Some(result), + io_queue: self.io_queue.clone(), // Report 0 bytes for bypass tasks so the backpressure budget is unaffected num_bytes: if self.bypass_backpressure { 0 @@ -511,6 +628,84 @@ impl ScanStats { } } +fn split_priority(priority: Option) -> (Option, Option) { + priority + .map(|priority| ((priority >> 64) as u64, priority as u64)) + .unzip() +} + +#[derive(Debug, Clone, Copy)] +pub(super) struct SchedulerStateEvent { + pub(super) queue_kind: &'static str, + pub(super) io_capacity: u64, + pub(super) iops_available: u64, + pub(super) active_iops: u64, + pub(super) pending_iops: u64, + pub(super) pending_bytes: u64, + pub(super) bytes_available: i64, + pub(super) bytes_reserved: i64, + pub(super) io_buffer_size_bytes: u64, + pub(super) priorities_in_flight: u64, + pub(super) no_backpressure: bool, + pub(super) head_task_bytes: Option, + pub(super) head_task_priority_high: Option, + pub(super) head_task_priority_low: Option, + pub(super) min_in_flight_priority_high: Option, + pub(super) min_in_flight_priority_low: Option, + pub(super) head_task_can_deliver: Option, + pub(super) head_task_priority_bypass: Option, + pub(super) head_task_blocked_by_iops: Option, + pub(super) head_task_blocked_by_bytes: Option, +} + +impl SchedulerStateEvent { + fn trace(self, stats: ScanStats) { + tracing::event!( + target: SCHEDULER_STATE_EVENT_TARGET, + tracing::Level::TRACE, + queue_kind = self.queue_kind, + scheduler_iops = stats.iops, + scheduler_requests = stats.requests, + scheduler_bytes_read = stats.bytes_read, + io_capacity = self.io_capacity, + iops_available = self.iops_available, + active_iops = self.active_iops, + pending_iops = self.pending_iops, + pending_bytes = self.pending_bytes, + bytes_available = self.bytes_available, + bytes_reserved = self.bytes_reserved, + io_buffer_size_bytes = self.io_buffer_size_bytes, + priorities_in_flight = self.priorities_in_flight, + no_backpressure = self.no_backpressure, + head_task_bytes_present = self.head_task_bytes.is_some(), + head_task_bytes = self.head_task_bytes.unwrap_or_default(), + head_task_priority_high_present = self.head_task_priority_high.is_some(), + head_task_priority_high = self.head_task_priority_high.unwrap_or_default(), + head_task_priority_low_present = self.head_task_priority_low.is_some(), + head_task_priority_low = self.head_task_priority_low.unwrap_or_default(), + min_in_flight_priority_high_present = self.min_in_flight_priority_high.is_some(), + min_in_flight_priority_high = self.min_in_flight_priority_high.unwrap_or_default(), + min_in_flight_priority_low_present = self.min_in_flight_priority_low.is_some(), + min_in_flight_priority_low = self.min_in_flight_priority_low.unwrap_or_default(), + head_task_can_deliver_present = self.head_task_can_deliver.is_some(), + head_task_can_deliver = self.head_task_can_deliver.unwrap_or(false), + head_task_priority_bypass_present = self.head_task_priority_bypass.is_some(), + head_task_priority_bypass = self.head_task_priority_bypass.unwrap_or(false), + head_task_blocked_by_iops_present = self.head_task_blocked_by_iops.is_some(), + head_task_blocked_by_iops = self.head_task_blocked_by_iops.unwrap_or(false), + head_task_blocked_by_bytes_present = self.head_task_blocked_by_bytes.is_some(), + head_task_blocked_by_bytes = self.head_task_blocked_by_bytes.unwrap_or(false), + "Scheduler state" + ); + } +} + +pub(super) fn emit_scheduler_state_event(event: Option, stats: &IoStats) { + if let Some(event) = event { + event.trace(stats.snapshot()); + } +} + /// A shareable, cloneable handle to a set of cumulative I/O counters. /// /// All clones share the same underlying counters. This serves two purposes: @@ -590,12 +785,25 @@ impl Debug for ScanScheduler { } struct Response { - data: Result>, + // `Option` so the caller can take the data out while the response (and its + // backpressure refund on drop) stays intact. + data: Option>>, + io_queue: Arc, priority: u128, num_reqs: usize, num_bytes: u64, } +// Refund the batch's backpressure reservation when the response is dropped, be +// that on delivery or when a cancelled request's undelivered response is +// discarded. This releases the budget even if the caller drops the future early. +impl Drop for Response { + fn drop(&mut self) { + self.io_queue + .on_bytes_consumed(self.num_bytes, self.priority, self.num_reqs); + } +} + #[derive(Debug, Clone, Copy)] pub struct SchedulerConfig { /// the # of bytes that can be buffered but not yet requested. @@ -651,6 +859,7 @@ impl ScanScheduler { /// * config - configuration settings for the scheduler pub fn new(object_store: Arc, config: SchedulerConfig) -> Arc { let io_capacity = object_store.io_parallelism(); + let stats = IoStats::new(); let use_lite = config .use_lite_scheduler .unwrap_or_else(|| object_store.prefers_lite_scheduler()); @@ -658,12 +867,14 @@ impl ScanScheduler { let io_queue = Arc::new(lite::IoQueue::new( io_capacity as u64, config.io_buffer_size_bytes, + stats.clone(), )); IoQueueType::Lite(io_queue) } else { let io_queue = Arc::new(IoQueue::new( io_capacity as u32, config.io_buffer_size_bytes, + stats.clone(), )); let io_queue_clone = io_queue.clone(); // Best we can do here is fire and forget. If the I/O loop is still running when the scheduler is @@ -675,7 +886,7 @@ impl ScanScheduler { Arc::new(Self { object_store, io_queue, - stats: IoStats::new(), + stats, }) } @@ -774,6 +985,7 @@ impl ScanScheduler { priority, request.len(), bypass_backpressure, + io_queue.clone(), )))); for (task_idx, iop) in request.into_iter().enumerate() { @@ -812,14 +1024,11 @@ impl ScanScheduler { self.do_submit_request(reader, request, tx, priority, io_queue, bypass_backpressure); - let io_queue_clone = io_queue.clone(); - - rx.map(move |wrapped_rsp| { - // Right now, it isn't possible for I/O to be cancelled so a cancel error should - // not occur - let rsp = wrapped_rsp.unwrap(); - io_queue_clone.on_bytes_consumed(rsp.num_bytes, rsp.priority, rsp.num_reqs); - rsp.data + rx.map(|wrapped_rsp| { + // A cancel error can't occur: the sender always sends before dropping. + // The reservation is refunded on `Response` drop, so just take the data. + let mut rsp = wrapped_rsp.unwrap(); + rsp.data.take().unwrap() }) } @@ -1151,6 +1360,47 @@ mod tests { } } + #[test] + fn test_scheduler_state_event_fields() { + use tracing_mock::{expect, subscriber}; + + let event = expect::event() + .with_target(SCHEDULER_STATE_EVENT_TARGET) + .at_level(tracing::Level::TRACE) + .with_fields( + expect::field("queue_kind") + .with_value(&"standard") + .and(expect::field("scheduler_iops").with_value(&7u64)) + .and(expect::field("scheduler_requests").with_value(&3u64)) + .and(expect::field("scheduler_bytes_read").with_value(&4096u64)) + .and(expect::field("io_capacity").with_value(&4u64)) + .and(expect::field("pending_iops").with_value(&1u64)) + .and(expect::field("bytes_available").with_value(&128i64)) + .and(expect::field("head_task_bytes_present").with_value(&true)) + .and(expect::field("head_task_bytes").with_value(&1u64)) + .and(expect::field("head_task_can_deliver_present").with_value(&true)) + .and(expect::field("head_task_can_deliver").with_value(&true)), + ); + let (subscriber, handle) = subscriber::mock().event(event).run_with_handle(); + + let stats = IoStats::new(); + stats.add_scan_stats(&ScanStats { + iops: 7, + requests: 3, + bytes_read: 4096, + }); + let mut state = IoQueueState::new(4, 192); + state.iops_avail = 2; + state.bytes_avail = 128; + state.pending_requests.push(make_task(1, false)); + + tracing::subscriber::with_default(subscriber, || { + emit_scheduler_state_event(state.scheduler_state_event(), &stats); + }); + + handle.assert_finished(); + } + #[test] fn test_iotask_ordering() { // Bypass tasks must come out of the heap before non-bypass tasks. @@ -1466,6 +1716,136 @@ mod tests { assert!(second_fut.await.unwrap().unwrap().len() == 20); } + #[tokio::test] + async fn test_standard_scheduler_state_tracks_queue_state() { + let some_path = Path::parse("foo").unwrap(); + let base_store = Arc::new(InMemory::new()); + base_store + .put(&some_path, vec![0; 1000].into()) + .await + .unwrap(); + + let semaphore = Arc::new(tokio::sync::Semaphore::new(0)); + let mut obj_store = MockObjectStore::default(); + let semaphore_copy = semaphore.clone(); + obj_store + .expect_get_opts() + .returning(move |location, options| { + let semaphore = semaphore.clone(); + let base_store = base_store.clone(); + let location = location.clone(); + async move { + semaphore.acquire().await.unwrap().forget(); + base_store.get_opts(&location, options).await + } + .boxed() + }); + let obj_store = Arc::new(ObjectStore::new( + Arc::new(obj_store), + Url::parse("mem://").unwrap(), + Some(500), + None, + false, + false, + 1, + DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + )); + + let scheduler = ScanScheduler::new( + obj_store, + SchedulerConfig { + io_buffer_size_bytes: 1024 * 1024, + use_lite_scheduler: Some(false), + }, + ); + let file_scheduler = scheduler + .open_file(&Path::parse("foo").unwrap(), &CachedFileSize::new(1000)) + .await + .unwrap(); + + let first_fut = timeout( + Duration::from_secs(10), + file_scheduler.submit_single(0..10, 0), + ) + .boxed(); + let second_fut = timeout( + Duration::from_secs(10), + file_scheduler.submit_single(0..20, 100), + ) + .boxed(); + let third_fut = timeout( + Duration::from_secs(10), + file_scheduler.submit_single(0..30, 0), + ) + .boxed(); + + let io_queue = match &scheduler.io_queue { + IoQueueType::Standard(io_queue) => io_queue.clone(), + IoQueueType::Lite(_) => unreachable!("test forces the standard scheduler"), + }; + let ( + io_capacity, + iops_available, + pending_bytes, + bytes_reserved, + priorities_in_flight, + head_task_bytes, + head_task_blocked_by_iops, + head_task_blocked_by_bytes, + ) = timeout(Duration::from_secs(5), async { + loop { + let observed = { + let state = io_queue.state.lock().unwrap(); + let active_iops = state.io_capacity.saturating_sub(state.iops_avail); + if active_iops == 1 && state.pending_requests.len() == 2 { + let pending_bytes = state + .pending_requests + .iter() + .map(IoTask::num_bytes) + .sum::(); + let head_task = state.pending_requests.peek().unwrap(); + let bypasses_bytes = state.no_backpressure + || head_task.bypass_backpressure + || head_task.priority <= state.priorities_in_flight.min_in_flight(); + Some(( + state.io_capacity, + state.iops_avail, + pending_bytes, + state.io_buffer_size as i64 - state.bytes_avail, + state.priorities_in_flight.len(), + head_task.num_bytes(), + state.iops_avail == 0, + !bypasses_bytes && head_task.num_bytes() as i64 > state.bytes_avail, + )) + } else { + None + } + }; + if let Some(observed) = observed { + break observed; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + assert_eq!(io_capacity, 1); + assert_eq!(iops_available, 0); + assert_eq!(pending_bytes, 50); + assert_eq!(bytes_reserved, 10); + assert_eq!(priorities_in_flight, 1); + assert_eq!(head_task_bytes, 30); + assert!(head_task_blocked_by_iops); + assert!(!head_task_blocked_by_bytes); + + semaphore_copy.add_permits(3); + assert_eq!(first_fut.await.unwrap().unwrap().len(), 10); + assert_eq!(third_fut.await.unwrap().unwrap().len(), 30); + assert_eq!(second_fut.await.unwrap().unwrap().len(), 20); + } + #[tokio::test(flavor = "multi_thread")] async fn test_backpressure() { let some_path = Path::parse("foo").unwrap(); @@ -1598,6 +1978,102 @@ mod tests { assert_eq!(second_fut.await.unwrap().len(), 10); } + #[derive(Debug)] + struct BlockingReader { + semaphore: Arc, + get_range_count: Arc, + path: Path, + } + + impl lance_core::deepsize::DeepSizeOf for BlockingReader { + fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize { + 0 + } + } + + impl Reader for BlockingReader { + fn path(&self) -> &Path { + &self.path + } + + fn block_size(&self) -> usize { + 4096 + } + + fn io_parallelism(&self) -> usize { + 1 + } + + fn size(&self) -> futures::future::BoxFuture<'_, object_store::Result> { + Box::pin(async { Ok(1_000_000) }) + } + + fn get_range( + &self, + range: Range, + ) -> futures::future::BoxFuture<'static, object_store::Result> { + self.get_range_count.fetch_add(1, Ordering::Release); + let semaphore = self.semaphore.clone(); + let num_bytes = range.end - range.start; + Box::pin(async move { + semaphore.acquire().await.unwrap().forget(); + Ok(Bytes::from(vec![0u8; num_bytes])) + }) + } + + fn get_all(&self) -> futures::future::BoxFuture<'_, object_store::Result> { + Box::pin(async { Ok(Bytes::from(vec![0u8; 1_000_000])) }) + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_same_priority_chunks_continue_after_higher_priority_request() { + let obj_store = Arc::new(ObjectStore::new( + Arc::new(InMemory::new()), + Url::parse("mem://").unwrap(), + Some(4096), + None, + false, + false, + 1, + DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + )); + let scheduler = ScanScheduler::new( + obj_store, + SchedulerConfig { + io_buffer_size_bytes: 10, + use_lite_scheduler: Some(false), + }, + ); + let semaphore = Arc::new(tokio::sync::Semaphore::new(0)); + let reader: Arc = Arc::new(BlockingReader { + semaphore: semaphore.clone(), + get_range_count: Arc::new(AtomicU64::new(0)), + path: Path::parse("test").unwrap(), + }); + + let low_priority = + scheduler.submit_request(reader.clone(), vec![0..6, 100..106], 10, false); + let high_priority = scheduler.submit_request(reader, vec![200..204], 0, false); + + semaphore.add_permits(3); + let low_priority = timeout(Duration::from_secs(5), low_priority) + .await + .unwrap() + .unwrap(); + assert_eq!( + low_priority.iter().map(|bytes| bytes.len()).sum::(), + 12 + ); + + let high_priority = timeout(Duration::from_secs(5), high_priority) + .await + .unwrap() + .unwrap(); + assert_eq!(high_priority[0].len(), 4); + } + /// A Reader that tracks how many times get_range has been called. #[derive(Debug)] struct TrackingReader { @@ -1880,4 +2356,105 @@ mod tests { .unwrap(); assert_eq!(bytes_dispatched.load(Ordering::Acquire), 30); } + + // Against a 100-byte budget: submit fut1 (50 bytes, priority 0), drop it while + // its read is still blocked in get_range, then submit fut2 (60 bytes, priority 1). + // fut2's priority can't win the priority-bypass, so it needs 60 of the budget -- + // available only if fut1's dropped reservation was refunded. Returns whether fut2 + // completed within 2s (false = the reservation leaked and fut2 deadlocked). + async fn run_caller_drop_scenario(use_lite_scheduler: bool) -> (bool, Duration) { + let obj_store = Arc::new(ObjectStore::new( + Arc::new(InMemory::new()), + Url::parse("mem://").unwrap(), + Some(4096), + None, + false, + false, + 1, + DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + )); + let scheduler = ScanScheduler::new( + obj_store, + SchedulerConfig { + io_buffer_size_bytes: 100, + use_lite_scheduler: Some(use_lite_scheduler), + }, + ); + + let semaphore = Arc::new(tokio::sync::Semaphore::new(0)); + let get_range_count = Arc::new(AtomicU64::new(0)); + let reader: Arc = Arc::new(BlockingReader { + semaphore: semaphore.clone(), + get_range_count: get_range_count.clone(), + path: Path::parse("test").unwrap(), + }); + + // Step 1: reserve 50 of the 100 budget bytes with a read we never consume. + // Spawn it so we can cancel the caller-side future while it is still parked + // waiting for the (blocked) read to finish. + let fut1 = scheduler.submit_request(reader.clone(), vec![0..50], 0, false); + let handle = tokio::spawn(async move { + let _ = fut1.await; + }); + + // Wait until the read is genuinely in flight (blocked on the semaphore). + // This guarantees the 50-byte reservation has been taken before we drop + // the caller, closing the race between the I/O loop and the abort. + while get_range_count.load(Ordering::Acquire) == 0 { + tokio::time::sleep(Duration::from_millis(1)).await; + } + + // Step 2: drop the caller-side future while its `rx` is still pending. + handle.abort(); + let _ = handle.await; + + // Step 3: let the in-flight read finish. The reservation should be refunded + // now that the request is done, whether or not the caller is still around. + semaphore.add_permits(1); + // Give the read time to run to completion so the refund would already have + // happened. + tokio::time::sleep(Duration::from_millis(50)).await; + + // Step 4: submit the follow-up. Add a permit up front so that, if it *is* + // admitted, its own read can complete rather than block on the semaphore. + semaphore.add_permits(1); + let fut2 = scheduler.submit_request(reader, vec![100..160], 1, false); + + let start = std::time::Instant::now(); + let outcome = timeout(Duration::from_secs(2), fut2).await; + let elapsed = start.elapsed(); + match outcome { + Ok(res) => { + assert_eq!(res.unwrap().iter().map(|b| b.len()).sum::(), 60); + (true, elapsed) + } + Err(_) => (false, elapsed), + } + } + + /// Dropping a standard-scheduler request future while its read is in flight must + /// still refund the backpressure reservation, so a later request that needs the + /// budget does not deadlock. + #[tokio::test(flavor = "multi_thread")] + async fn standard_scheduler_refunds_reservation_on_caller_drop() { + let (completed, elapsed) = run_caller_drop_scenario(false).await; + assert!( + completed, + "standard scheduler deadlocked the follow-up request (elapsed {elapsed:?}); \ + the dropped request's reservation was not refunded" + ); + } + + /// Same guarantee for the lite scheduler: dropping a request future mid-read + /// releases its reservation via the `TaskHandle` drop path. + #[tokio::test(flavor = "multi_thread")] + async fn lite_scheduler_refunds_reservation_on_caller_drop() { + let (completed, elapsed) = run_caller_drop_scenario(true).await; + assert!( + completed, + "lite scheduler deadlocked the follow-up request (elapsed {elapsed:?}); \ + the dropped request's reservation was not refunded" + ); + } } diff --git a/rust/lance-io/src/scheduler/lite.rs b/rust/lance-io/src/scheduler/lite.rs index 338a4e2a46f..b7a06e11bb7 100644 --- a/rust/lance-io/src/scheduler/lite.rs +++ b/rust/lance-io/src/scheduler/lite.rs @@ -33,7 +33,10 @@ use std::{ use bytes::Bytes; use lance_core::{Error, Result}; -use super::{BACKPRESSURE_DEBOUNCE, BACKPRESSURE_MIN}; +use super::{ + BACKPRESSURE_DEBOUNCE, BACKPRESSURE_MIN, IoStats, SCHEDULER_STATE_EVENT_TARGET, + SchedulerStateEvent, emit_scheduler_state_event, +}; type RunFn = Box Pin> + Send>> + Send>; @@ -67,6 +70,26 @@ enum TaskState { }, } +impl TaskState { + fn backpressure_reservation(&self) -> Option { + match self { + Self::Reserved { + backpressure_reservation, + .. + } + | Self::Running { + backpressure_reservation, + .. + } + | Self::Finished { + backpressure_reservation, + .. + } => Some(*backpressure_reservation), + Self::Initial { .. } | Self::Broken => None, + } + } +} + /// A custom error type that might have a backpressure reservation /// /// This is used instead of Lance's standard error type so we can ensure @@ -85,25 +108,14 @@ impl BrokenTaskError { // This will capture any backpressure reservation the task has and put it into the // error so we make sure to release it when returning the error. fn new(task_state: TaskState, message: String) -> Self { - match task_state { - TaskState::Reserved { - backpressure_reservation, - .. - } - | TaskState::Running { - backpressure_reservation, - .. - } - | TaskState::Finished { - backpressure_reservation, - .. - } => Self { + match task_state.backpressure_reservation() { + None => Self { message, - backpressure_reservation: Some(backpressure_reservation), + backpressure_reservation: None, }, - TaskState::Broken | TaskState::Initial { .. } => Self { + Some(reservation) => Self { message, - backpressure_reservation: None, + backpressure_reservation: Some(reservation), }, } } @@ -237,6 +249,7 @@ trait BackpressureThrottle: Send { /// Unconditionally acquire a zero-cost reservation, tracking only the priority. /// Used for bypass tasks that must never be blocked by backpressure. fn force_acquire(&mut self, priority: u128) -> BackpressureReservation; + fn state(&self) -> BackpressureState; } // We want to allow requests that have a lower priority than any @@ -262,6 +275,10 @@ impl PrioritiesInFlight { self.in_flight.first().copied().unwrap_or(u128::MAX) } + fn contains(&self, prio: u128) -> bool { + self.in_flight.binary_search(&prio).is_ok() + } + fn push(&mut self, prio: u128) { let pos = match self.in_flight.binary_search(&prio) { Ok(pos) => pos, @@ -275,9 +292,22 @@ impl PrioritiesInFlight { self.in_flight.remove(pos); } } + + fn len(&self) -> usize { + self.in_flight.len() + } +} + +#[derive(Debug, Clone, Copy)] +struct BackpressureState { + max_bytes: u64, + bytes_available: i64, + priorities_in_flight: u64, + no_backpressure: bool, } struct SimpleBackpressureThrottle { + max_bytes: u64, start: Instant, last_warn: AtomicU64, bytes_available: i64, @@ -293,6 +323,7 @@ impl SimpleBackpressureThrottle { panic!("Max bytes must be less than {}", i64::MAX); } Self { + max_bytes, start: Instant::now(), last_warn: AtomicU64::new(0), bytes_available: max_bytes as i64, @@ -325,6 +356,10 @@ impl BackpressureThrottle for SimpleBackpressureThrottle { if self.no_backpressure || self.bytes_available >= num_bytes as i64 || self.priorities_in_flight.min_in_flight() >= priority + // Chunks from an admitted logical request must keep moving. A + // higher-priority request may be scheduled later and remain + // unconsumed while the caller awaits this request. + || self.priorities_in_flight.contains(priority) { self.bytes_available -= num_bytes as i64; self.priorities_in_flight.push(priority); @@ -350,6 +385,15 @@ impl BackpressureThrottle for SimpleBackpressureThrottle { priority, } } + + fn state(&self) -> BackpressureState { + BackpressureState { + max_bytes: self.max_bytes, + bytes_available: self.bytes_available, + priorities_in_flight: self.priorities_in_flight.len() as u64, + no_backpressure: self.no_backpressure, + } + } } struct TaskEntry { @@ -419,6 +463,48 @@ impl IoQueueState { Ok(()) } } + + fn scheduler_state_event(&self) -> Option { + if !tracing::enabled!(target: SCHEDULER_STATE_EVENT_TARGET, tracing::Level::TRACE) { + return None; + } + + let backpressure = self.backpressure_throttle.state(); + let pending_bytes = self + .pending_tasks + .iter() + .filter_map(|entry| self.tasks.get(&entry.task_id)) + .map(|task| task.num_bytes) + .sum::(); + let active_iops = self + .tasks + .values() + .filter(|task| matches!(task.state, TaskState::Running { .. })) + .count() as u64; + + Some(SchedulerStateEvent { + queue_kind: "lite", + io_capacity: 0, + iops_available: 0, + active_iops, + pending_iops: self.pending_tasks.len() as u64, + pending_bytes, + bytes_available: backpressure.bytes_available, + bytes_reserved: backpressure.max_bytes as i64 - backpressure.bytes_available, + io_buffer_size_bytes: backpressure.max_bytes, + priorities_in_flight: backpressure.priorities_in_flight, + no_backpressure: backpressure.no_backpressure, + head_task_bytes: None, + head_task_priority_high: None, + head_task_priority_low: None, + min_in_flight_priority_high: None, + min_in_flight_priority_low: None, + head_task_can_deliver: None, + head_task_priority_bypass: None, + head_task_blocked_by_iops: None, + head_task_blocked_by_bytes: None, + }) + } } /// A queue of I/O tasks to be shared between the I/O scheduler and the I/O decoder. @@ -441,12 +527,14 @@ impl IoQueueState { /// day as well) pub(super) struct IoQueue { state: Arc>, + stats: IoStats, } impl IoQueue { - pub fn new(max_concurrency: u64, max_bytes: u64) -> Self { + pub fn new(max_concurrency: u64, max_bytes: u64, stats: IoStats) -> Self { Self { state: Arc::new(Mutex::new(IoQueueState::new(max_concurrency, max_bytes))), + stats, } } @@ -463,6 +551,9 @@ impl IoQueue { state.handle_result(task.reserve(reservation))?; state.handle_result(task.start())?; state.tasks.insert(task_id, task); + let event = state.scheduler_state_event(); + drop(state); + emit_scheduler_state_event(event, &self.stats); return Ok(()); } @@ -472,6 +563,9 @@ impl IoQueue { reserved: task.is_reserved(), }); state.tasks.insert(task_id, task); + let event = state.scheduler_state_event(); + drop(state); + emit_scheduler_state_event(event, &self.stats); Ok(()) } @@ -510,34 +604,42 @@ impl IoQueue { // When a task completes we should check to see if any other tasks are now runnable fn on_task_complete(&self, mut state: MutexGuard) -> Result<()> { - let state_ref = &mut *state; - let mut task_result = TaskResult::Ok(()); - while !state_ref.pending_tasks.is_empty() { - // Unwrap safe here since we just checked the queue is not empty - let next_task = state_ref.pending_tasks.peek().unwrap(); - let Some(task) = state_ref.tasks.get_mut(&next_task.task_id) else { - log::warn!("Task with id {} was lost", next_task.task_id); - continue; - }; - if !task.is_reserved() { - let Some(reservation) = state_ref - .backpressure_throttle - .try_acquire(task.num_bytes, task.priority) - else { - break; + let result = { + let state_ref = &mut *state; + let mut task_result = TaskResult::Ok(()); + while !state_ref.pending_tasks.is_empty() { + // Unwrap safe here since we just checked the queue is not empty + let task_id = state_ref.pending_tasks.peek().unwrap().task_id; + let Some(task) = state_ref.tasks.get_mut(&task_id) else { + // The caller dropped this task's handle (see `abandon`); discard the + // stale queue entry instead of spinning on it. + state_ref.pending_tasks.pop(); + continue; }; - if let Err(e) = task.reserve(reservation) { + if !task.is_reserved() { + let Some(reservation) = state_ref + .backpressure_throttle + .try_acquire(task.num_bytes, task.priority) + else { + break; + }; + if let Err(e) = task.reserve(reservation) { + task_result = Err(e); + break; + } + } + state_ref.pending_tasks.pop(); + if let Err(e) = task.start() { task_result = Err(e); break; } } - state_ref.pending_tasks.pop(); - if let Err(e) = task.start() { - task_result = Err(e); - break; - } - } - state_ref.handle_result(task_result) + state_ref.handle_result(task_result) + }; + let event = state.scheduler_state_event(); + drop(state); + emit_scheduler_state_event(event, &self.stats); + result } fn poll(&self, task_id: u64, cx: &mut Context<'_>) -> Poll> { @@ -565,10 +667,34 @@ impl IoQueue { } pub(super) fn close(&self) { + let event = { + let mut state = self.state.lock().unwrap(); + for task in std::mem::take(&mut state.tasks).values_mut() { + task.cancel(); + } + state.scheduler_state_event() + }; + emit_scheduler_state_event(event, &self.stats); + } + + // Called when a caller drops a task's handle before the task finishes. Removes + // the task and returns any backpressure reservation it holds to the budget, then + // re-checks the queue so newly-affordable tasks can start. Unlike the standard + // release path (`poll`), this runs without the task being polled to completion, + // so a cancelled read does not leak its reservation. + fn abandon(&self, task_id: u64) { let mut state = self.state.lock().unwrap(); - for task in std::mem::take(&mut state.tasks).values_mut() { - task.cancel(); + let Some(task) = state.tasks.remove(&task_id) else { + // Already consumed by `poll`; nothing to release. + return; + }; + + if let Some(reservation) = task.state.backpressure_reservation() { + state.backpressure_throttle.release(reservation); } + // Freed budget may make queued tasks runnable; there is no caller to surface + // an error to here. + let _ = self.on_task_complete(state); } } @@ -584,6 +710,12 @@ impl Future for TaskHandle { } } +impl Drop for TaskHandle { + fn drop(&mut self) { + self.queue.abandon(self.task_id); + } +} + #[cfg(test)] mod tests { use super::*; @@ -592,7 +724,7 @@ mod tests { #[tokio::test] async fn test_priority_ordering() { // Backpressure budget of 10 bytes: only one 10-byte task runs at a time. - let queue = Arc::new(IoQueue::new(128, 10)); + let queue = Arc::new(IoQueue::new(128, 10, IoStats::default())); // Records the priority of each task when its run_fn is invoked (i.e. when // the task transitions to Running). @@ -700,7 +832,7 @@ mod tests { async fn test_zero_buffer_bypasses_backpressure() { // Budget = 0 sets no_backpressure = true, so all tasks start immediately // regardless of how many bytes are "outstanding". - let queue = Arc::new(IoQueue::new(128, 0)); + let queue = Arc::new(IoQueue::new(128, 0, IoStats::default())); let start_order: Arc>> = Arc::new(Mutex::new(Vec::new())); let make_run_fn = @@ -742,7 +874,7 @@ mod tests { async fn test_bypass_flag_proceeds_past_exhausted_budget() { // Budget of 10 bytes. A blocker task fills it. A task with bypass=true starts // immediately despite the exhausted budget; a normal task stays queued. - let queue = Arc::new(IoQueue::new(128, 10)); + let queue = Arc::new(IoQueue::new(128, 10, IoStats::default())); let start_order: Arc>> = Arc::new(Mutex::new(Vec::new())); let make_run_fn = @@ -802,4 +934,22 @@ mod tests { normal_tx.send(Bytes::from_static(b"done")).unwrap(); normal.await.unwrap(); } + + #[test] + fn test_same_priority_reservation_continues_after_higher_priority() { + let mut throttle = SimpleBackpressureThrottle::new(10, 128); + + let low_priority_first = throttle.try_acquire(6, 10).unwrap(); + let high_priority = throttle.try_acquire(4, 0).unwrap(); + let low_priority_next = throttle.try_acquire(6, 10); + + assert!( + low_priority_next.is_some(), + "chunks from an already admitted logical request should continue" + ); + + throttle.release(low_priority_first); + throttle.release(high_priority); + throttle.release(low_priority_next.unwrap()); + } } diff --git a/rust/lance-io/src/utils.rs b/rust/lance-io/src/utils.rs index b36dff75133..a36aadec9d9 100644 --- a/rust/lance-io/src/utils.rs +++ b/rust/lance-io/src/utils.rs @@ -84,7 +84,11 @@ pub async fn read_fixed_stride_array( /// followed by the message itself. pub async fn read_message(reader: &dyn Reader, pos: usize) -> Result { let file_size = reader.size().await?; - if pos > file_size { + // A message is a u32 length prefix followed by its body; both must lie before + // the end. A `pos` too close to the end means the reader size is too small + // (e.g. a stale cached size). Reject it rather than slice a short buffer and + // panic. + if pos + 4 > file_size { return Err(Error::io("file size is too small".to_string())); } @@ -96,7 +100,9 @@ pub async fn read_message(reader: &dyn Reader, pos: usize) let remaining_range = range.end..min(4 + pos + msg_len, file_size); let remaining_bytes = reader.get_range(remaining_range).await?; let buf = [buf, remaining_bytes].concat(); - assert!(buf.len() >= msg_len + 4); + if buf.len() < msg_len + 4 { + return Err(Error::io("file size is too small".to_string())); + } Ok(M::decode(&buf[4..4 + msg_len])?) } else { Ok(M::decode(&buf[4..4 + msg_len])?) diff --git a/rust/lance-linalg/Cargo.toml b/rust/lance-linalg/Cargo.toml index 6a188ec3c62..d94efa7b2fa 100644 --- a/rust/lance-linalg/Cargo.toml +++ b/rust/lance-linalg/Cargo.toml @@ -25,6 +25,7 @@ approx = { workspace = true } criterion = { workspace = true } lance-testing = { path = "../lance-testing" } proptest.workspace = true +rstest.workspace = true [build-dependencies] cc = "1.0.83" diff --git a/rust/lance-linalg/proptest-regressions/distance/cosine.txt b/rust/lance-linalg/proptest-regressions/distance/cosine.txt new file mode 100644 index 00000000000..04e6dcb967f --- /dev/null +++ b/rust/lance-linalg/proptest-regressions/distance/cosine.txt @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 679c764b02e2726ec2d943b3abdc7d7d3565e168d4081b155860c1e69b616fce # shrinks to (x, y) = ([-1.2933521e-15, -1.1194652e-8, 7.447341e-32, -0.0, -14247490.0, -8.55e-43, -4576.872, 0.009181444], [4.926e-42, -9.267065e-15, -0.0, 7e-45, -0.00036999694, -0.0, -5.7810876e-8, 1.9273644e-21]) diff --git a/rust/lance-linalg/src/distance.rs b/rust/lance-linalg/src/distance.rs index 23d1cae2d63..d5f0dcba7aa 100644 --- a/rust/lance-linalg/src/distance.rs +++ b/rust/lance-linalg/src/distance.rs @@ -25,6 +25,96 @@ pub mod l2; pub mod l2_u8; pub mod norm_l2; +/// What a per-batch distance kernel yields. +/// +/// The two batch paths produce different shapes: the build-baseline path maps +/// lazily over the batch and allocates nothing, while a `#[target_feature]` +/// kernel must collect eagerly because it cannot be inlined into a lazy +/// closure. A concrete enum keeps both statically dispatched. +/// +/// Only sub-AVX2 builds need this. On an AVX2-baseline build the batch methods +/// return the bare `Map` instead, because any wrapper — trait object or enum — +/// loses `TrustedLen` (so `.collect()` stops preallocating) and loses +/// `Map::fold`'s inlined loop. Benchmarks showed that costing 2.5x on the dim-8 +/// batch, far more than the per-vector dispatch it was meant to remove. +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] +pub(crate) enum BatchIter { + /// Lazy per-vector map. No allocation. + Lazy(L), + /// Eagerly collected by a `#[target_feature]` kernel. + Eager(std::vec::IntoIter), +} + +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] +impl> Iterator for BatchIter { + type Item = f32; + + #[inline] + fn next(&mut self) -> Option { + match self { + Self::Lazy(iter) => iter.next(), + Self::Eager(iter) => iter.next(), + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + match self { + Self::Lazy(iter) => iter.size_hint(), + Self::Eager(iter) => iter.size_hint(), + } + } + + /// Delegated, not defaulted. `Map` overrides `fold` to drive the underlying + /// `ChunksExact` in one inlined, auto-vectorized loop; the default `fold` + /// would instead call `next()` per element, paying an enum branch and + /// losing that loop. On the dim-8 batch that costs ~2.5x. + #[inline] + fn fold(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + match self { + Self::Lazy(iter) => iter.fold(init, f), + Self::Eager(iter) => iter.fold(init, f), + } + } + + /// `for_each`, `sum` and `collect` all route through `fold`, so delegating + /// it covers them too. (`try_fold` cannot be overridden on stable: its + /// `Try` bound is unstable.) + #[inline] + fn for_each(self, f: F) + where + F: FnMut(Self::Item), + { + match self { + Self::Lazy(iter) => iter.for_each(f), + Self::Eager(iter) => iter.for_each(f), + } + } +} + +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] +impl> ExactSizeIterator for BatchIter { + #[inline] + fn len(&self) -> usize { + match self { + Self::Lazy(iter) => iter.len(), + Self::Eager(iter) => iter.len(), + } + } +} + pub use cosine::*; pub use dot::*; pub use hamming::{ diff --git a/rust/lance-linalg/src/distance/cosine.rs b/rust/lance-linalg/src/distance/cosine.rs index 995191b77eb..457c0c0dd6a 100644 --- a/rust/lance-linalg/src/distance/cosine.rs +++ b/rust/lance-linalg/src/distance/cosine.rs @@ -17,12 +17,12 @@ use arrow_array::{ use arrow_schema::DataType; use half::{bf16, f16}; use lance_arrow::{ArrowFloatType, FixedSizeListArrayExt, FloatArray}; -use lance_core::utils::cpu::SIMD_SUPPORT; -#[cfg(feature = "fp16kernels")] -use lance_core::utils::cpu::SimdSupport; +#[allow(unused_imports)] +use lance_core::utils::cpu::{SIMD_SUPPORT, SimdSupport}; use super::{Dot, norm_l2::norm_l2}; use super::{Normalize, dot::dot}; +#[allow(unused_imports)] use crate::simd::{ FloatSimd, SIMD, f32::{f32x8, f32x16}, @@ -127,6 +127,10 @@ impl Cosine for bf16 { SimdSupport::Lsx => unsafe { bf16_kernel::cosine_bf16_lsx(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the bf16 C kernels in `bf16_kernel::*` are compiled with + // `-march=haswell` minimum (which requires AVX2), so they cannot + // run on AVX-only or AVX+FMA hosts. Scalar is the correct route. _ => cosine_scalar(x, x_norm, y), } } @@ -179,88 +183,348 @@ impl Cosine for f16 { SimdSupport::Lsx => unsafe { kernel::cosine_f16_lsx(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the f16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => cosine_scalar(x, x_norm, y), } } } -/// f32 kernels for Cosine +/// f32 single-vector cosine helpers used by `cosine_batch` for fixed +/// dimensions 8 and 16. +/// +/// These were previously a single generic `cosine_once` but the +/// monomorphizations have to dispatch on `SIMD_SUPPORT` for the SIMD path +/// to stay correct under any compile baseline. Splitting them into two +/// concrete entry points keeps the dispatch site flat and lets each width +/// route to a `#[target_feature]` AVX2 inner function. mod f32 { use super::*; - // TODO: how can we explicitly infer N? #[inline] - pub(super) fn cosine_once, const N: usize>( + pub(super) fn cosine_once_8(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + cosine_once_x86::cosine_once_8_avx512(x, x_norm, y) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { + cosine_once_x86::cosine_once_8_avx_fma(x, x_norm, y) + }, + SimdSupport::Avx => unsafe { cosine_once_x86::cosine_once_8_avx(x, x_norm, y) }, + _ => cosine_once_8_scalar(x, x_norm, y), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + cosine_once_8_other(x, x_norm, y) + } + } + + #[inline] + pub(super) fn cosine_once_16(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + cosine_once_x86::cosine_once_16_avx512(x, x_norm, y) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { + cosine_once_x86::cosine_once_16_avx_fma(x, x_norm, y) + }, + SimdSupport::Avx => unsafe { cosine_once_x86::cosine_once_16_avx(x, x_norm, y) }, + _ => cosine_once_16_scalar(x, x_norm, y), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + cosine_once_16_other(x, x_norm, y) + } + } + + /// Portable scalar `cosine_once` for length-8 vectors. Matches the SIMD + /// path modulo summation order. + #[cfg(target_arch = "x86_64")] + #[inline] + pub(super) fn cosine_once_8_scalar(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let mut xy = 0.0f32; + let mut y2 = 0.0f32; + for i in 0..8 { + xy += x[i] * y[i]; + y2 += y[i] * y[i]; + } + 1.0 - xy / x_norm / y2.sqrt() + } + + #[cfg(target_arch = "x86_64")] + #[inline] + pub(super) fn cosine_once_16_scalar(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let mut xy = 0.0f32; + let mut y2 = 0.0f32; + for i in 0..16 { + xy += x[i] * y[i]; + y2 += y[i] * y[i]; + } + 1.0 - xy / x_norm / y2.sqrt() + } + + #[cfg(target_arch = "x86_64")] + pub(super) mod cosine_once_x86 { + use std::arch::x86_64::*; + + use super::{f32x8, f32x16}; + use crate::simd::SIMD; + + /// AVX + FMA path for 8-lane cosine. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn cosine_once_8_avx_fma(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = f32x8::load_unaligned(x.as_ptr()); + let yv = f32x8::load_unaligned(y.as_ptr()); + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + + /// AVX + FMA path for 16-lane cosine. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn cosine_once_16_avx_fma(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = f32x16::load_unaligned(x.as_ptr()); + let yv = f32x16::load_unaligned(y.as_ptr()); + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + + /// AVX-only path for 8-lane cosine (no FMA): body unchanged from AVX2 path; gated on Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn cosine_once_8_avx(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = f32x8::load_unaligned(x.as_ptr()); + let yv = f32x8::load_unaligned(y.as_ptr()); + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + + /// AVX-only path for 16-lane cosine (no FMA): body unchanged from AVX2 path; gated on Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn cosine_once_16_avx(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = f32x16::load_unaligned(x.as_ptr()); + let yv = f32x16::load_unaligned(y.as_ptr()); + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + + /// AVX-512 path for 8-lane cosine: masked load into a `__m512` lower half, reduce. + #[target_feature(enable = "avx512f")] + pub unsafe fn cosine_once_8_avx512(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + // mask 0x00FF: load the lower 8 f32 lanes, zero the upper 8. + let mask: __mmask16 = 0x00FF; + let xv = _mm512_maskz_loadu_ps(mask, x.as_ptr()); + let yv = _mm512_maskz_loadu_ps(mask, y.as_ptr()); + let xy = _mm512_mul_ps(xv, yv); + let y2 = _mm512_mul_ps(yv, yv); + let xy_sum = _mm512_reduce_add_ps(xy); + let y2_sum = _mm512_reduce_add_ps(y2); + 1.0 - xy_sum / x_norm / y2_sum.sqrt() + } + + /// AVX-512 path for 16-lane cosine: single full-width `__m512` load (16 f32 fits one `zmm`). + #[target_feature(enable = "avx512f")] + pub unsafe fn cosine_once_16_avx512(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = _mm512_loadu_ps(x.as_ptr()); + let yv = _mm512_loadu_ps(y.as_ptr()); + let xy = _mm512_mul_ps(xv, yv); + let y2 = _mm512_mul_ps(yv, yv); + let xy_sum = _mm512_reduce_add_ps(xy); + let y2_sum = _mm512_reduce_add_ps(y2); + 1.0 - xy_sum / x_norm / y2_sum.sqrt() + } + } + + #[cfg(not(target_arch = "x86_64"))] + #[inline] + fn cosine_once_8_other(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = unsafe { f32x8::load_unaligned(x.as_ptr()) }; + let yv = unsafe { f32x8::load_unaligned(y.as_ptr()) }; + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + + #[cfg(not(target_arch = "x86_64"))] + #[inline] + fn cosine_once_16_other(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = unsafe { f32x16::load_unaligned(x.as_ptr()) }; + let yv = unsafe { f32x16::load_unaligned(y.as_ptr()) }; + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + + /// Batch-level SIMD dispatch: the tier is chosen once by the caller, and the + /// whole `chunks_exact` loop runs inside one `#[target_feature]` context so + /// the per-vector `cosine_once_*` / `cosine_fast` kernels inline (no + /// per-vector `*SIMD_SUPPORT` branch, no per-vector call boundary). Used for + /// the AVX-512 path on the default wheel and for all tiers on sub-AVX2 builds. + #[cfg(target_arch = "x86_64")] + #[target_feature(enable = "avx,fma")] + pub(super) unsafe fn cosine_batch_avx_fma( x: &[f32], x_norm: f32, - y: &[f32], - ) -> f32 { - let x = unsafe { S::load_unaligned(x.as_ptr()) }; - let y = unsafe { S::load_unaligned(y.as_ptr()) }; - let y2 = y * y; - let xy = x * y; - 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + batch: &[f32], + dimension: usize, + ) -> Vec { + match dimension { + 8 => batch + .chunks_exact(8) + .map(|y| unsafe { cosine_once_x86::cosine_once_8_avx_fma(x, x_norm, y) }) + .collect(), + 16 => batch + .chunks_exact(16) + .map(|y| unsafe { cosine_once_x86::cosine_once_16_avx_fma(x, x_norm, y) }) + .collect(), + _ => batch + .chunks_exact(dimension) + .map(|y| unsafe { super::f32_x86::cosine_fast_avx_fma(x, x_norm, y) }) + .collect(), + } + } + + #[cfg(target_arch = "x86_64")] + #[target_feature(enable = "avx512f")] + pub(super) unsafe fn cosine_batch_avx512( + x: &[f32], + x_norm: f32, + batch: &[f32], + dimension: usize, + ) -> Vec { + match dimension { + 8 => batch + .chunks_exact(8) + .map(|y| unsafe { cosine_once_x86::cosine_once_8_avx512(x, x_norm, y) }) + .collect(), + 16 => batch + .chunks_exact(16) + .map(|y| unsafe { cosine_once_x86::cosine_once_16_avx512(x, x_norm, y) }) + .collect(), + _ => batch + .chunks_exact(dimension) + .map(|y| unsafe { super::f32_x86::cosine_fast_avx512(x, x_norm, y) }) + .collect(), + } + } + + #[cfg(target_arch = "x86_64")] + #[target_feature(enable = "avx")] + pub(super) unsafe fn cosine_batch_avx( + x: &[f32], + x_norm: f32, + batch: &[f32], + dimension: usize, + ) -> Vec { + match dimension { + 8 => batch + .chunks_exact(8) + .map(|y| unsafe { cosine_once_x86::cosine_once_8_avx(x, x_norm, y) }) + .collect(), + 16 => batch + .chunks_exact(16) + .map(|y| unsafe { cosine_once_x86::cosine_once_16_avx(x, x_norm, y) }) + .collect(), + _ => batch + .chunks_exact(dimension) + .map(|y| unsafe { super::f32_x86::cosine_fast_avx(x, x_norm, y) }) + .collect(), + } } } -impl Cosine for f32 { +/// Inlined f32 cosine kernels for builds whose baseline already guarantees AVX2 +/// (the default `haswell` wheel). No `#[target_feature]`, no runtime dispatch: +/// under `target-feature=+avx2,+fma` these compile to AVX2 and inline into the +/// batch loop exactly like the pre-PR code, so the modern path is not taxed by +/// the runtime-dispatch machinery (only needed below the AVX2 baseline). +#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] +mod f32_baseline { + use super::{dot, f32x8, f32x16, norm_l2}; + use crate::simd::{FloatSimd, SIMD}; + #[inline] - fn cosine_fast(x: &[Self], x_norm: Self, other: &[Self]) -> f32 { - let dim = x.len(); - let unrolled_len = dim / 16 * 16; - let mut y_norm16 = f32x16::zeros(); - let mut xy16 = f32x16::zeros(); - for i in (0..unrolled_len).step_by(16) { - unsafe { - let x = f32x16::load_unaligned(x.as_ptr().add(i)); - let y = f32x16::load_unaligned(other.as_ptr().add(i)); - xy16.multiply_add(x, y); - y_norm16.multiply_add(y, y); - } + pub fn cosine_once_8(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + unsafe { + let xv = f32x8::load_unaligned(x.as_ptr()); + let yv = f32x8::load_unaligned(y.as_ptr()); + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() } - let aligned_len = dim / 8 * 8; - let mut y_norm8 = f32x8::zeros(); - let mut xy8 = f32x8::zeros(); - for i in (unrolled_len..aligned_len).step_by(8) { - unsafe { - let x = f32x8::load_unaligned(x.as_ptr().add(i)); - let y = f32x8::load_unaligned(other.as_ptr().add(i)); - xy8.multiply_add(x, y); - y_norm8.multiply_add(y, y); - } + } + + #[inline] + pub fn cosine_once_16(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + unsafe { + let xv = f32x16::load_unaligned(x.as_ptr()); + let yv = f32x16::load_unaligned(y.as_ptr()); + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() } - let y_norm = - y_norm16.reduce_sum() + y_norm8.reduce_sum() + norm_l2(&other[aligned_len..]).powi(2); - let xy = - xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &other[aligned_len..]); - 1.0 - xy / x_norm / y_norm.sqrt() } #[inline] - fn cosine_with_norms(x: &[Self], x_norm: Self, y_norm: Self, y: &[Self]) -> Self { - let dim = x.len(); - let unrolled_len = dim / 16 * 16; - let mut xy16 = f32x16::zeros(); - for i in (0..unrolled_len).step_by(16) { - unsafe { - let x = f32x16::load_unaligned(x.as_ptr().add(i)); - let y = f32x16::load_unaligned(y.as_ptr().add(i)); - xy16.multiply_add(x, y); + pub fn cosine_fast(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { + unsafe { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + let mut y_norm16 = f32x16::zeros(); + let mut xy16 = f32x16::zeros(); + for i in (0..unrolled_len).step_by(16) { + let xv = f32x16::load_unaligned(x.as_ptr().add(i)); + let yv = f32x16::load_unaligned(other.as_ptr().add(i)); + xy16.multiply_add(xv, yv); + y_norm16.multiply_add(yv, yv); } - } - let aligned_len = dim / 8 * 8; - let mut xy8 = f32x8::zeros(); - for i in (unrolled_len..aligned_len).step_by(8) { - unsafe { - let x = f32x8::load_unaligned(x.as_ptr().add(i)); - let y = f32x8::load_unaligned(y.as_ptr().add(i)); - xy8.multiply_add(x, y); + let aligned_len = dim / 8 * 8; + let mut y_norm8 = f32x8::zeros(); + let mut xy8 = f32x8::zeros(); + for i in (unrolled_len..aligned_len).step_by(8) { + let xv = f32x8::load_unaligned(x.as_ptr().add(i)); + let yv = f32x8::load_unaligned(other.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + y_norm8.multiply_add(yv, yv); } + let y_norm = y_norm16.reduce_sum() + + y_norm8.reduce_sum() + + norm_l2(&other[aligned_len..]).powi(2); + let xy = xy16.reduce_sum() + + xy8.reduce_sum() + + dot(&x[aligned_len..], &other[aligned_len..]); + 1.0 - xy / x_norm / y_norm.sqrt() } - let xy = xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &y[aligned_len..]); - 1.0 - xy / x_norm / y_norm + } +} + +impl Cosine for f32 { + #[inline] + fn cosine_fast(x: &[Self], x_norm: Self, other: &[Self]) -> f32 { + // Trait methods cannot carry `#[target_feature]` attributes, so the body + // lives in a free function that runtime-dispatches via `*SIMD_SUPPORT` + // to an AVX2 inner kernel on capable hosts, or a portable scalar fallback. + cosine_fast_f32_dispatched(x, x_norm, other) } + #[inline] + fn cosine_with_norms(x: &[Self], x_norm: Self, y_norm: Self, y: &[Self]) -> Self { + // Trait methods cannot carry `#[target_feature]` attributes, so the body + // lives in a free function that runtime-dispatches via `*SIMD_SUPPORT` + // to an AVX2 inner kernel on capable hosts, or a portable scalar fallback. + cosine_with_norms_f32_dispatched(x, x_norm, y_norm, y) + } + + #[allow(unreachable_code)] fn cosine_batch<'a>( x: &'a [Self], batch: &'a [Self], @@ -268,16 +532,83 @@ impl Cosine for f32 { ) -> Box + 'a> { let x_norm = norm_l2(x); + // On a build whose baseline already guarantees AVX2 (the default + // `haswell` wheel), avoid the per-vector runtime dispatch + `#[target_feature]` + // wrapping that taxes the modern path. Dispatch ONCE per batch: AVX-512 + // hosts get the wide kernel; everyone else uses the inlined AVX2 baseline + // path (base-equivalent). The runtime-dispatch path below is only + // compiled/reached when the baseline is below AVX2 (pre-Haswell builds). + #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] + { + // dim 8/16 always use the inlined AVX2 baseline: AVX-512 gives no + // benefit for such tiny vectors (a masked 512-bit load is slower than + // a plain AVX2 load) and only adds dispatch + eager-collect overhead. + // Only the larger-dim path routes to AVX-512 on capable hosts — that's + // where the wider lanes actually pay off. + return match dimension { + 8 => Box::new( + batch + .chunks_exact(8) + .map(move |y| f32_baseline::cosine_once_8(x, x_norm, y)), + ), + 16 => Box::new( + batch + .chunks_exact(16) + .map(move |y| f32_baseline::cosine_once_16(x, x_norm, y)), + ), + _ => { + if matches!(*SIMD_SUPPORT, SimdSupport::Avx512 | SimdSupport::Avx512FP16) { + Box::new( + unsafe { f32::cosine_batch_avx512(x, x_norm, batch, dimension) } + .into_iter(), + ) + } else { + Box::new( + batch + .chunks_exact(dimension) + .map(move |y| f32_baseline::cosine_fast(x, x_norm, y)), + ) + } + } + }; + } + + // Sub-AVX2 / non-x86 build: hoisted per-batch runtime dispatch. + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => { + return Box::new( + unsafe { f32::cosine_batch_avx512(x, x_norm, batch, dimension) } + .into_iter(), + ); + } + SimdSupport::Avx2 | SimdSupport::AvxFma => { + return Box::new( + unsafe { f32::cosine_batch_avx_fma(x, x_norm, batch, dimension) } + .into_iter(), + ); + } + SimdSupport::Avx => { + return Box::new( + unsafe { f32::cosine_batch_avx(x, x_norm, batch, dimension) }.into_iter(), + ); + } + _ => {} + } + } + + // Scalar / non-x86 fallback. match dimension { 8 => Box::new( batch .chunks_exact(dimension) - .map(move |y| f32::cosine_once::(x, x_norm, y)), + .map(move |y| f32::cosine_once_8(x, x_norm, y)), ), 16 => Box::new( batch .chunks_exact(dimension) - .map(move |y| f32::cosine_once::(x, x_norm, y)), + .map(move |y| f32::cosine_once_16(x, x_norm, y)), ), _ => Box::new( batch @@ -291,34 +622,102 @@ impl Cosine for f32 { impl Cosine for f64 { #[inline] fn cosine_fast(x: &[Self], x_norm: f32, y: &[Self]) -> f32 { - use crate::simd::f64::{f64x4, f64x8}; - use crate::simd::{FloatSimd, SIMD}; + // Trait methods cannot carry `#[target_feature]` attributes, so the body + // lives in a free function that runtime-dispatches via `*SIMD_SUPPORT` + // to an AVX2 inner kernel on capable hosts, or a portable scalar fallback. + cosine_fast_f64_dispatched(x, x_norm, y) + } +} +/// Fast cosine for f64, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the SIMD +/// primitives in `crate::simd::f64`. +#[inline] +fn cosine_fast_f64_dispatched(x: &[f64], x_norm: f32, y: &[f64]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + f64_x86::cosine_fast_avx512(x, x_norm, y) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { + f64_x86::cosine_fast_avx_fma(x, x_norm, y) + }, + SimdSupport::Avx => unsafe { f64_x86::cosine_fast_avx(x, x_norm, y) }, + _ => cosine_scalar(x, x_norm, y), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + cosine_fast_f64_simd_other(x, x_norm, y) + } +} + +/// AVX2 + FMA implementation of the f64 cosine_fast kernel. +/// +/// Lives in a `#[target_feature]`-annotated function so the SIMD primitives +/// in `crate::simd::f64` (which use raw AVX intrinsics) inline correctly +/// even when the compile baseline does not have AVX2 enabled. Caller must +/// ensure the host supports AVX2 + FMA. +#[cfg(target_arch = "x86_64")] +mod f64_x86 { + use std::arch::x86_64::*; + + use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::x86::hsum256_pd; + use crate::simd::{FloatSimd, SIMD}; + + /// AVX-512 path for f64 fast cosine: 8-wide `__m512d` xy/yy with `vfmadd231pd` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn cosine_fast_avx512(x: &[f64], x_norm: f32, y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc_xy = _mm512_setzero_pd(); + let mut acc_yy = _mm512_setzero_pd(); + for i in (0..unrolled_len).step_by(8) { + let xv = _mm512_loadu_pd(x.as_ptr().add(i)); + let yv = _mm512_loadu_pd(y.as_ptr().add(i)); + acc_xy = _mm512_fmadd_pd(xv, yv, acc_xy); + acc_yy = _mm512_fmadd_pd(yv, yv, acc_yy); + } + + let mut xy = _mm512_reduce_add_pd(acc_xy); + let mut yy = _mm512_reduce_add_pd(acc_yy); + for i in unrolled_len..dim { + xy += x[i] * y[i]; + yy += y[i] * y[i]; + } + + let y_norm_sq = yy as f32; + let xy_f32 = xy as f32; + 1.0 - xy_f32 / x_norm / y_norm_sq.sqrt() + } + + /// AVX + FMA path for f64 fast cosine. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn cosine_fast_avx_fma(x: &[f64], x_norm: f32, y: &[f64]) -> f32 { let dim = x.len(); let unrolled_len = dim / 8 * 8; let mut y_norm8 = f64x8::zeros(); let mut xy8 = f64x8::zeros(); for i in (0..unrolled_len).step_by(8) { - unsafe { - let xv = f64x8::load_unaligned(x.as_ptr().add(i)); - let yv = f64x8::load_unaligned(y.as_ptr().add(i)); - xy8.multiply_add(xv, yv); - y_norm8.multiply_add(yv, yv); - } + let xv = f64x8::load_unaligned(x.as_ptr().add(i)); + let yv = f64x8::load_unaligned(y.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + y_norm8.multiply_add(yv, yv); } let aligned_len = dim / 4 * 4; let mut y_norm4 = f64x4::zeros(); let mut xy4 = f64x4::zeros(); for i in (unrolled_len..aligned_len).step_by(4) { - unsafe { - let xv = f64x4::load_unaligned(x.as_ptr().add(i)); - let yv = f64x4::load_unaligned(y.as_ptr().add(i)); - xy4.multiply_add(xv, yv); - y_norm4.multiply_add(yv, yv); - } + let xv = f64x4::load_unaligned(x.as_ptr().add(i)); + let yv = f64x4::load_unaligned(y.as_ptr().add(i)); + xy4.multiply_add(xv, yv); + y_norm4.multiply_add(yv, yv); } - let tail_y_norm: Self = y[aligned_len..].iter().map(|&v| v * v).sum(); - let tail_xy: Self = x[aligned_len..] + let tail_y_norm: f64 = y[aligned_len..].iter().map(|&v| v * v).sum(); + let tail_xy: f64 = x[aligned_len..] .iter() .zip(y[aligned_len..].iter()) .map(|(&a, &b)| a * b) @@ -328,6 +727,336 @@ impl Cosine for f64 { let xy = (xy8.reduce_sum() + xy4.reduce_sum() + tail_xy) as f32; 1.0 - xy / x_norm / y_norm_sq.sqrt() } + + /// AVX-only path for f64 fast cosine (no FMA): `_mm256_mul_pd` + `_mm256_add_pd` per iteration; tail handled inline. + #[target_feature(enable = "avx")] + pub unsafe fn cosine_fast_avx(x: &[f64], x_norm: f32, y: &[f64]) -> f32 { + let dim = x.len(); + let aligned_len = dim / 4 * 4; + + let mut acc_xy = _mm256_setzero_pd(); + let mut acc_yy = _mm256_setzero_pd(); + for i in (0..aligned_len).step_by(4) { + let xv = _mm256_loadu_pd(x.as_ptr().add(i)); + let yv = _mm256_loadu_pd(y.as_ptr().add(i)); + acc_xy = _mm256_add_pd(acc_xy, _mm256_mul_pd(xv, yv)); + acc_yy = _mm256_add_pd(acc_yy, _mm256_mul_pd(yv, yv)); + } + + let xy_main = hsum256_pd(acc_xy); + let yy_main = hsum256_pd(acc_yy); + + let tail_y_norm: f64 = y[aligned_len..].iter().map(|&v| v * v).sum(); + let tail_xy: f64 = x[aligned_len..] + .iter() + .zip(y[aligned_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + let y_norm_sq = (yy_main + tail_y_norm) as f32; + let xy = (xy_main + tail_xy) as f32; + 1.0 - xy / x_norm / y_norm_sq.sqrt() + } +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn cosine_fast_f64_simd_other(x: &[f64], x_norm: f32, y: &[f64]) -> f32 { + use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::{FloatSimd, SIMD}; + + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + let mut y_norm8 = f64x8::zeros(); + let mut xy8 = f64x8::zeros(); + for i in (0..unrolled_len).step_by(8) { + unsafe { + let xv = f64x8::load_unaligned(x.as_ptr().add(i)); + let yv = f64x8::load_unaligned(y.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + y_norm8.multiply_add(yv, yv); + } + } + let aligned_len = dim / 4 * 4; + let mut y_norm4 = f64x4::zeros(); + let mut xy4 = f64x4::zeros(); + for i in (unrolled_len..aligned_len).step_by(4) { + unsafe { + let xv = f64x4::load_unaligned(x.as_ptr().add(i)); + let yv = f64x4::load_unaligned(y.as_ptr().add(i)); + xy4.multiply_add(xv, yv); + y_norm4.multiply_add(yv, yv); + } + } + let tail_y_norm: f64 = y[aligned_len..].iter().map(|&v| v * v).sum(); + let tail_xy: f64 = x[aligned_len..] + .iter() + .zip(y[aligned_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + let y_norm_sq = (y_norm8.reduce_sum() + y_norm4.reduce_sum() + tail_y_norm) as f32; + let xy = (xy8.reduce_sum() + xy4.reduce_sum() + tail_xy) as f32; + 1.0 - xy / x_norm / y_norm_sq.sqrt() +} + +/// Cosine for f32 with known norms, runtime-dispatched via `SIMD_SUPPORT` +/// on x86_64 (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses +/// the auto-vectorised scalar loop. +#[inline] +fn cosine_with_norms_f32_dispatched(x: &[f32], x_norm: f32, y_norm: f32, y: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + f32_x86::cosine_with_norms_avx512(x, x_norm, y_norm, y) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { + f32_x86::cosine_with_norms_avx_fma(x, x_norm, y_norm, y) + }, + SimdSupport::Avx => unsafe { f32_x86::cosine_with_norms_avx(x, x_norm, y_norm, y) }, + _ => cosine_scalar_fast(x, x_norm, y, y_norm), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + cosine_with_norms_f32_simd_other(x, x_norm, y_norm, y) + } +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn cosine_with_norms_f32_simd_other(x: &[f32], x_norm: f32, y_norm: f32, y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + let mut xy16 = f32x16::zeros(); + for i in (0..unrolled_len).step_by(16) { + unsafe { + let xv = f32x16::load_unaligned(x.as_ptr().add(i)); + let yv = f32x16::load_unaligned(y.as_ptr().add(i)); + xy16.multiply_add(xv, yv); + } + } + let aligned_len = dim / 8 * 8; + let mut xy8 = f32x8::zeros(); + for i in (unrolled_len..aligned_len).step_by(8) { + unsafe { + let xv = f32x8::load_unaligned(x.as_ptr().add(i)); + let yv = f32x8::load_unaligned(y.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + } + } + let xy = xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &y[aligned_len..]); + 1.0 - xy / x_norm / y_norm +} + +/// Fast cosine for f32, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the +/// `simd::f32` primitives, unconditionally backed by NEON / LSX. +#[inline] +fn cosine_fast_f32_dispatched(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + f32_x86::cosine_fast_avx512(x, x_norm, other) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { + f32_x86::cosine_fast_avx_fma(x, x_norm, other) + }, + SimdSupport::Avx => unsafe { f32_x86::cosine_fast_avx(x, x_norm, other) }, + _ => cosine_scalar(x, x_norm, other), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + cosine_fast_f32_simd_other(x, x_norm, other) + } +} + +/// AVX2 + FMA implementation of the f32 fast cosine kernel. +/// +/// Lives in a `#[target_feature]`-annotated function so the SIMD primitives +/// in `crate::simd::f32` (which use raw AVX intrinsics) inline correctly +/// even when the compile baseline does not have AVX2 enabled. Caller must +/// ensure the host supports AVX2 + FMA. +#[cfg(target_arch = "x86_64")] +mod f32_x86 { + use std::arch::x86_64::*; + + use super::{dot, f32x8, f32x16, norm_l2}; + use crate::simd::x86::hsum256_ps; + use crate::simd::{FloatSimd, SIMD}; + + /// AVX + FMA path for f32 fast cosine. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn cosine_fast_avx_fma(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + let mut y_norm16 = f32x16::zeros(); + let mut xy16 = f32x16::zeros(); + for i in (0..unrolled_len).step_by(16) { + let xv = f32x16::load_unaligned(x.as_ptr().add(i)); + let yv = f32x16::load_unaligned(other.as_ptr().add(i)); + xy16.multiply_add(xv, yv); + y_norm16.multiply_add(yv, yv); + } + let aligned_len = dim / 8 * 8; + let mut y_norm8 = f32x8::zeros(); + let mut xy8 = f32x8::zeros(); + for i in (unrolled_len..aligned_len).step_by(8) { + let xv = f32x8::load_unaligned(x.as_ptr().add(i)); + let yv = f32x8::load_unaligned(other.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + y_norm8.multiply_add(yv, yv); + } + let y_norm = + y_norm16.reduce_sum() + y_norm8.reduce_sum() + norm_l2(&other[aligned_len..]).powi(2); + let xy = + xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &other[aligned_len..]); + 1.0 - xy / x_norm / y_norm.sqrt() + } + + /// AVX-only path for f32 fast cosine (no FMA): `_mm256_mul_ps` + `_mm256_add_ps` per iteration; tail via trait-routed `dot`/`norm_l2`. + #[target_feature(enable = "avx")] + pub unsafe fn cosine_fast_avx(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { + let dim = x.len(); + let aligned_len = dim / 8 * 8; + + let mut acc_xy = _mm256_setzero_ps(); + let mut acc_yy = _mm256_setzero_ps(); + for i in (0..aligned_len).step_by(8) { + let xv = _mm256_loadu_ps(x.as_ptr().add(i)); + let yv = _mm256_loadu_ps(other.as_ptr().add(i)); + acc_xy = _mm256_add_ps(acc_xy, _mm256_mul_ps(xv, yv)); + acc_yy = _mm256_add_ps(acc_yy, _mm256_mul_ps(yv, yv)); + } + + let xy_main = hsum256_ps(acc_xy); + let yy_main = hsum256_ps(acc_yy); + + let y_norm = yy_main + norm_l2(&other[aligned_len..]).powi(2); + let xy = xy_main + dot(&x[aligned_len..], &other[aligned_len..]); + 1.0 - xy / x_norm / y_norm.sqrt() + } + + /// AVX-512 path for f32 fast cosine: 16-wide `__m512` xy/yy with `vfmadd231ps` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn cosine_fast_avx512(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + + let mut acc_xy = _mm512_setzero_ps(); + let mut acc_yy = _mm512_setzero_ps(); + for i in (0..unrolled_len).step_by(16) { + let xv = _mm512_loadu_ps(x.as_ptr().add(i)); + let yv = _mm512_loadu_ps(other.as_ptr().add(i)); + acc_xy = _mm512_fmadd_ps(xv, yv, acc_xy); + acc_yy = _mm512_fmadd_ps(yv, yv, acc_yy); + } + + let mut xy = _mm512_reduce_add_ps(acc_xy); + let mut yy = _mm512_reduce_add_ps(acc_yy); + for i in unrolled_len..dim { + xy += x[i] * other[i]; + yy += other[i] * other[i]; + } + + 1.0 - xy / x_norm / yy.sqrt() + } + + /// AVX-512 path for f32 cosine with known norms: 16-wide `__m512` with `vfmadd231ps` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn cosine_with_norms_avx512(x: &[f32], x_norm: f32, y_norm: f32, y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + + let mut acc = _mm512_setzero_ps(); + for i in (0..unrolled_len).step_by(16) { + let xv = _mm512_loadu_ps(x.as_ptr().add(i)); + let yv = _mm512_loadu_ps(y.as_ptr().add(i)); + acc = _mm512_fmadd_ps(xv, yv, acc); + } + + let mut xy = _mm512_reduce_add_ps(acc); + for i in unrolled_len..dim { + xy += x[i] * y[i]; + } + + 1.0 - xy / x_norm / y_norm + } + + /// AVX + FMA path for f32 cosine with known norms. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn cosine_with_norms_avx_fma(x: &[f32], x_norm: f32, y_norm: f32, y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + let mut xy16 = f32x16::zeros(); + for i in (0..unrolled_len).step_by(16) { + let xv = f32x16::load_unaligned(x.as_ptr().add(i)); + let yv = f32x16::load_unaligned(y.as_ptr().add(i)); + xy16.multiply_add(xv, yv); + } + let aligned_len = dim / 8 * 8; + let mut xy8 = f32x8::zeros(); + for i in (unrolled_len..aligned_len).step_by(8) { + let xv = f32x8::load_unaligned(x.as_ptr().add(i)); + let yv = f32x8::load_unaligned(y.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + } + let xy = xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &y[aligned_len..]); + 1.0 - xy / x_norm / y_norm + } + + /// AVX-only path for f32 cosine with known norms (no FMA): `_mm256_mul_ps` + `_mm256_add_ps` per iteration; tail via trait-routed `dot`. + #[target_feature(enable = "avx")] + pub unsafe fn cosine_with_norms_avx(x: &[f32], x_norm: f32, y_norm: f32, y: &[f32]) -> f32 { + let dim = x.len(); + let aligned_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..aligned_len).step_by(8) { + let xv = _mm256_loadu_ps(x.as_ptr().add(i)); + let yv = _mm256_loadu_ps(y.as_ptr().add(i)); + acc = _mm256_add_ps(acc, _mm256_mul_ps(xv, yv)); + } + + let xy_main = hsum256_ps(acc); + let xy = xy_main + dot(&x[aligned_len..], &y[aligned_len..]); + 1.0 - xy / x_norm / y_norm + } +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn cosine_fast_f32_simd_other(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + let mut y_norm16 = f32x16::zeros(); + let mut xy16 = f32x16::zeros(); + for i in (0..unrolled_len).step_by(16) { + unsafe { + let xv = f32x16::load_unaligned(x.as_ptr().add(i)); + let yv = f32x16::load_unaligned(other.as_ptr().add(i)); + xy16.multiply_add(xv, yv); + y_norm16.multiply_add(yv, yv); + } + } + let aligned_len = dim / 8 * 8; + let mut y_norm8 = f32x8::zeros(); + let mut xy8 = f32x8::zeros(); + for i in (unrolled_len..aligned_len).step_by(8) { + unsafe { + let xv = f32x8::load_unaligned(x.as_ptr().add(i)); + let yv = f32x8::load_unaligned(other.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + y_norm8.multiply_add(yv, yv); + } + } + let y_norm = + y_norm16.reduce_sum() + y_norm8.reduce_sum() + norm_l2(&other[aligned_len..]).powi(2); + let xy = xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &other[aligned_len..]); + 1.0 - xy / x_norm / y_norm.sqrt() } /// Fallback non-SIMD implementation @@ -437,6 +1166,25 @@ pub fn cosine_distance_arrow_batch( } } +/// Portable scalar reference cosine over f64 inputs. Used by parity tests +/// to compare against every dispatched per-tier inner kernel. Computes +/// `1 - xy / (x_norm * y_norm_sq.sqrt())` in f64 then casts to f32, matching +/// the reduction order of the dispatched kernels. +#[cfg(test)] +fn cosine_fast_scalar(x: &[f64], x_norm: f32, y: &[f64]) -> f32 { + let xy: f64 = x.iter().zip(y.iter()).map(|(&a, &b)| a * b).sum(); + let y_norm_sq: f64 = y.iter().map(|&v| v * v).sum(); + 1.0 - (xy as f32) / x_norm / (y_norm_sq as f32).sqrt() +} + +/// Portable scalar reference cosine when both norms are known. Mirrors +/// `cosine_with_norms_f32_dispatched` for parity testing. +#[cfg(test)] +fn cosine_with_norms_scalar(x: &[f64], x_norm: f32, y_norm: f32, y: &[f64]) -> f32 { + let xy: f64 = x.iter().zip(y.iter()).map(|(&a, &b)| a * b).sum(); + 1.0 - (xy as f32) / x_norm / y_norm +} + #[cfg(test)] mod tests { use super::*; @@ -559,5 +1307,447 @@ mod tests { prop_assume!(norm_l2(&y) > 1e-20); do_cosine_test(&x, &y)?; } + + /// Cross-backend parity for the f32 cosine_fast kernel. Exercises the + /// scalar fallback (`cosine_scalar`) against the dispatched SIMD path + /// so the runtime fallback is exercised even on AVX2-capable CI hosts. + #[test] + fn test_cosine_fast_f32_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let scalar = cosine_fast_scalar(&x_f64, x_norm, &y_f64); + let simd = ::cosine_fast(&x, x_norm, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3)); + } + + /// AVX-512-direct parity for the f32 cosine_fast kernel. Early-returns + /// on hosts without AVX-512F so the test stays portable. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_fast_f32_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = cosine_scalar(&x, x_norm, &y); + let avx512 = unsafe { f32_x86::cosine_fast_avx512(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-5)); + } + + /// AVX + FMA-direct parity for the f32 cosine_fast kernel. Covers + /// the AMD Piledriver / Steamroller / FX-7500 tier. Early-returns + /// on hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_fast_f32_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = cosine_scalar(&x, x_norm, &y); + let avx_fma = unsafe { f32_x86::cosine_fast_avx_fma(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-5)); + } + + /// AVX-only-direct parity for the f32 cosine_fast kernel. Covers + /// the Intel Sandy Bridge / Ivy Bridge tier. Early-returns on + /// hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_fast_f32_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = cosine_scalar(&x, x_norm, &y); + let avx = unsafe { f32_x86::cosine_fast_avx(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-5)); + } + + /// Cross-backend parity for the f32 cosine_with_norms kernel. + /// Exercises the scalar fallback (`cosine_scalar_fast`) against the + /// dispatched SIMD path so the runtime fallback is exercised even on + /// AVX2-capable CI hosts. + #[test] + fn test_cosine_with_norms_f32_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let y_norm = norm_l2(&y); + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let scalar = cosine_with_norms_scalar(&x_f64, x_norm, y_norm, &y_f64); + let simd = ::cosine_with_norms(&x, x_norm, y_norm, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3)); + } + + /// AVX-512-direct parity for the f32 cosine_with_norms kernel. + /// Early-returns on hosts without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_with_norms_f32_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let y_norm = norm_l2(&y); + let scalar = cosine_scalar_fast(&x, x_norm, &y, y_norm); + let avx512 = unsafe { f32_x86::cosine_with_norms_avx512(&x, x_norm, y_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-5)); + } + + /// AVX + FMA-direct parity for the f32 cosine_with_norms kernel. + /// Covers the AMD Piledriver / Steamroller / FX-7500 tier. + /// Early-returns on hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_with_norms_f32_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let y_norm = norm_l2(&y); + let scalar = cosine_scalar_fast(&x, x_norm, &y, y_norm); + let avx_fma = unsafe { f32_x86::cosine_with_norms_avx_fma(&x, x_norm, y_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-5)); + } + + /// AVX-only-direct parity for the f32 cosine_with_norms kernel. + /// Covers the Intel Sandy Bridge / Ivy Bridge tier. Early-returns + /// on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_with_norms_f32_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let y_norm = norm_l2(&y); + let scalar = cosine_scalar_fast(&x, x_norm, &y, y_norm); + let avx = unsafe { f32_x86::cosine_with_norms_avx(&x, x_norm, y_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-5)); + } + + /// Cross-backend parity for the f64 cosine_fast kernel. Uses the + /// hand-rolled `cosine_fast_scalar` (not the trait-routed + /// `cosine_scalar`, which would itself dispatch through `dot::`) + /// so the reference stays free of any AVX path on AVX2-capable hosts. + #[test] + fn test_cosine_fast_f64_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + prop_assume!(norm_l2(&x) > 1e-20); + prop_assume!(norm_l2(&y) > 1e-20); + let x_norm = norm_l2(&x); + let scalar = cosine_fast_scalar(&x, x_norm, &y); + let simd = ::cosine_fast(&x, x_norm, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3)); + } + + /// AVX-512-direct parity for the f64 cosine_fast kernel. Early-returns + /// on hosts without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_fast_f64_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-20); + prop_assume!(norm_l2(&y) > 1e-20); + let x_norm = norm_l2(&x); + let scalar = cosine_fast_scalar(&x, x_norm, &y); + let avx512 = unsafe { f64_x86::cosine_fast_avx512(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-5)); + } + + /// AVX + FMA-direct parity for the f64 cosine_fast kernel. Covers + /// the AMD Piledriver / Steamroller / FX-7500 tier. Early-returns + /// on hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_fast_f64_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-20); + prop_assume!(norm_l2(&y) > 1e-20); + let x_norm = norm_l2(&x); + let scalar = cosine_fast_scalar(&x, x_norm, &y); + let avx_fma = unsafe { f64_x86::cosine_fast_avx_fma(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-5)); + } + + /// AVX-only-direct parity for the f64 cosine_fast kernel. Covers + /// the Intel Sandy Bridge / Ivy Bridge tier. Early-returns on + /// hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_fast_f64_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-20); + prop_assume!(norm_l2(&y) > 1e-20); + let x_norm = norm_l2(&x); + let scalar = cosine_fast_scalar(&x, x_norm, &y); + let avx = unsafe { f64_x86::cosine_fast_avx(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-5)); + } + + /// Parity check for `cosine_once_8` (despecialised 8-lane width). + /// + /// The `epsilon = 1e-6` clause handles the case where the proptest + /// generator produces inputs with extreme dynamic range (e.g., mixing + /// `1e-43` with `1e7` in the same vector). When the dot product is + /// dominated by one large term and the cosine result is near zero, + /// the f32-precision SIMD path and the f64-precision scalar reference + /// can legitimately differ by more than `max_relative = 1e-3` of the + /// (near-zero) result. The absolute epsilon catches these without + /// masking real bugs (where the absolute error would be > 1e-6). + #[test] + fn test_cosine_once_8_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 8..9) + ) { + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let scalar = cosine_fast_scalar(&x_f64, x_norm, &y_f64); + let simd = f32::cosine_once_8(&x, x_norm, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3, epsilon = 1e-6)); + } + + /// Parity check for `cosine_once_16` (despecialised 16-lane width). + /// See `test_cosine_once_8_scalar_simd_parity` for `epsilon` rationale. + #[test] + fn test_cosine_once_16_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 16..17) + ) { + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let scalar = cosine_fast_scalar(&x_f64, x_norm, &y_f64); + let simd = f32::cosine_once_16(&x, x_norm, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3, epsilon = 1e-6)); + } + + /// AVX-512-direct parity for the 8-lane cosine_once kernel. Verifies + /// the masked-load (mask 0x00FF) AVX-512 implementation produces + /// the same result as the scalar reference. Early-returns on hosts + /// without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_once_8_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 8..9) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = super::f32::cosine_once_8_scalar(&x, x_norm, &y); + let avx512 = + unsafe { super::f32::cosine_once_x86::cosine_once_8_avx512(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-5)); + } + + /// AVX-512-direct parity for the 16-lane cosine_once kernel. Verifies + /// the full-width `__m512` load implementation produces the same + /// result as the scalar reference. Early-returns on hosts without + /// AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_once_16_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 16..17) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = super::f32::cosine_once_16_scalar(&x, x_norm, &y); + let avx512 = + unsafe { super::f32::cosine_once_x86::cosine_once_16_avx512(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-5)); + } + + /// AVX + FMA-direct parity for the 8-lane cosine_once kernel. + /// Covers the AMD Piledriver / Steamroller / FX-7500 tier. + /// Early-returns on hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_once_8_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 8..9) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = super::f32::cosine_once_8_scalar(&x, x_norm, &y); + let avx_fma = + unsafe { super::f32::cosine_once_x86::cosine_once_8_avx_fma(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-5)); + } + + /// AVX + FMA-direct parity for the 16-lane cosine_once kernel. + /// Covers the AMD Piledriver / Steamroller / FX-7500 tier. + /// Early-returns on hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_once_16_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 16..17) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = super::f32::cosine_once_16_scalar(&x, x_norm, &y); + let avx_fma = + unsafe { super::f32::cosine_once_x86::cosine_once_16_avx_fma(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-5)); + } + + /// AVX-only-direct parity for the 8-lane cosine_once kernel. + /// Covers the Intel Sandy Bridge / Ivy Bridge tier. Early-returns + /// on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_once_8_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 8..9) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = super::f32::cosine_once_8_scalar(&x, x_norm, &y); + let avx = unsafe { super::f32::cosine_once_x86::cosine_once_8_avx(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-5)); + } + + /// AVX-only-direct parity for the 16-lane cosine_once kernel. + /// Covers the Intel Sandy Bridge / Ivy Bridge tier. Early-returns + /// on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_once_16_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 16..17) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = super::f32::cosine_once_16_scalar(&x, x_norm, &y); + let avx = unsafe { super::f32::cosine_once_x86::cosine_once_16_avx(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-5)); + } + } + + /// Asserts a batch-level f32 cosine SIMD kernel matches the scalar + /// `cosine_fast` reference for every vector in a multi-vector batch. Runs + /// each of the kernel's three internal dimension arms (8, 16, and the + /// general `chunks_exact` path). The batch kernels only run at runtime on + /// sub-AVX2 builds, so a direct call is the only way they get covered. + #[cfg(target_arch = "x86_64")] + fn check_cosine_batch_kernel(kernel: unsafe fn(&[f32], f32, &[f32], usize) -> Vec) { + for dimension in [8_usize, 16, 40] { + let x: Vec = (0..dimension).map(|i| (i as f32) * 0.5 + 1.0).collect(); + let x_norm = norm_l2(&x); + let num_vectors = 3; + let batch: Vec = (0..dimension * num_vectors) + .map(|i| ((i % 7) as f32) + 1.0) + .collect(); + + let got = unsafe { kernel(&x, x_norm, &batch, dimension) }; + assert_eq!(got.len(), num_vectors); + + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + for (chunk, &g) in batch.chunks_exact(dimension).zip(got.iter()) { + let y_f64: Vec = chunk.iter().map(|&v| v as f64).collect(); + let expected = cosine_fast_scalar(&x_f64, x_norm, &y_f64); + assert_relative_eq!(g, expected, max_relative = 1e-3, epsilon = 1e-6); + } + } + } + + /// AVX + FMA batch kernel parity (AVX2 / AVX+FMA tiers). Runs on any + /// Haswell-or-newer host; early-returns without AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_batch_avx_fma_matches_scalar() { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return; + } + check_cosine_batch_kernel(super::f32::cosine_batch_avx_fma); + } + + /// AVX-only batch kernel parity (Sandy Bridge / Ivy Bridge tier). + /// Early-returns on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_batch_avx_matches_scalar() { + if !std::is_x86_feature_detected!("avx") { + return; + } + check_cosine_batch_kernel(super::f32::cosine_batch_avx); + } + + /// AVX-512 batch kernel parity. Early-returns on hosts without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_batch_avx512_matches_scalar() { + if !std::is_x86_feature_detected!("avx512f") { + return; + } + check_cosine_batch_kernel(super::f32::cosine_batch_avx512); } } diff --git a/rust/lance-linalg/src/distance/cosine_u8.rs b/rust/lance-linalg/src/distance/cosine_u8.rs index b2d06d35c31..59833e020a7 100644 --- a/rust/lance-linalg/src/distance/cosine_u8.rs +++ b/rust/lance-linalg/src/distance/cosine_u8.rs @@ -199,6 +199,10 @@ fn select_backend() -> CosineU8AccumFn { if is_x86_feature_detected!("avx2") { return |a, b| unsafe { x86::cosine_u8_accum_avx2(a, b) }; } + // AvxFma and Avx hosts (AMD Piledriver / Steamroller, Intel Sandy + // Bridge / Ivy Bridge) fall through to scalar: the AVX2 inner uses + // `vpmaddubsw` / `vpmaddwd` integer ops which neither AVX nor + // AVX+FMA provides. } cosine_u8_accum_scalar diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs index 5903d24e0e5..3411d7e8dce 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/rust/lance-linalg/src/distance/dot.rs @@ -14,12 +14,16 @@ use arrow_schema::DataType; use half::{bf16, f16}; use lance_arrow::{ArrowFloatType, FixedSizeListArrayExt, FloatArray}; use lance_core::assume_eq; -use lance_core::utils::cpu::SIMD_SUPPORT; -#[cfg(feature = "fp16kernels")] -use lance_core::utils::cpu::SimdSupport; +#[allow(unused_imports)] +use lance_core::utils::cpu::{SIMD_SUPPORT, SimdSupport}; use num_traits::{AsPrimitive, Num, real::Real}; use crate::Result; +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] +use crate::distance::BatchIter; /// Default implementation of dot product. /// @@ -111,6 +115,24 @@ pub fn dot_distance(from: &[T], to: &[T]) -> f32 { pub trait Dot: Num { /// Dot product. fn dot(x: &[Self], y: &[Self]) -> f32; + + /// Dot product of `x` against each `dimension`-sized vector in `batch`. + /// + /// The default calls [`Dot::dot`] per vector. `f32` overrides it so the + /// SIMD tier is chosen once for the whole batch instead of once per + /// vector — on a build whose baseline already implies AVX2, per-vector + /// dispatch costs more than the kernel it selects. + /// + /// Returns `impl Iterator` rather than a trait object: hot consumers drive + /// this one element at a time, so a `Box` would cost a + /// virtual call per element and an allocation per batch. + fn dot_batch<'a>( + x: &'a [Self], + batch: &'a [Self], + dimension: usize, + ) -> impl Iterator + 'a { + batch.chunks_exact(dimension).map(move |y| Self::dot(x, y)) + } } #[cfg(feature = "fp16kernels")] @@ -161,6 +183,9 @@ impl Dot for bf16 { SimdSupport::Lsx => unsafe { bf16_kernel::dot_bf16_lsx(x.as_ptr(), y.as_ptr(), x.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the bf16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => dot_scalar::(x, y), } } @@ -214,6 +239,9 @@ impl Dot for f16 { SimdSupport::Lsx => unsafe { kernel::dot_f16_lsx(x.as_ptr(), y.as_ptr(), x.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the f16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => dot_scalar::(x, y), } } @@ -222,10 +250,128 @@ impl Dot for f16 { impl Dot for f32 { #[inline] fn dot(x: &[Self], y: &[Self]) -> f32 { - dot_scalar::(x, y) + // Trait methods cannot carry `#[target_feature]` attributes, so the body + // lives in a free function that runtime-dispatches via `*SIMD_SUPPORT` + // to an AVX2 or AVX-512 inner kernel on capable hosts, or a portable + // scalar fallback. Same shape as the f64 sibling and the existing + // u8 distance kernels in `dot_u8.rs`. + dot_f32_dispatched(x, y) + } + + fn dot_batch<'a>( + x: &'a [Self], + batch: &'a [Self], + dimension: usize, + ) -> impl Iterator + 'a { + // Exactly one arm compiles. Keeping each a tail expression (rather than + // an early `return` guarded by `cfg`) mirrors `dot_f32_dispatched` and + // avoids an unreachable tail on AVX2-baseline builds. + // AVX2-baseline build (the default `haswell` wheel). Hoist the tier + // choice out of the loop, but keep the SIMD kernel: the baseline already + // guarantees avx2+fma, so call the AVX+FMA kernel directly rather than + // re-checking per vector. Falling back to the scalar kernel here would + // lose ~4x at small dimensions, which is where batch calls live (PQ + // sub-vectors are 8 wide). + // + // The iterator is a bare `Map`: `Map` is `TrustedLen`, + // so `.collect()` preallocates, and `Map::fold` drives `ChunksExact` in + // one inlined loop. Any wrapper — trait object or enum — loses both. + #[cfg(all( + target_arch = "x86_64", + target_feature = "avx2", + target_feature = "fma" + ))] + { + // See `L2::l2_batch` for f32: below 16 lanes `dot_scalar`'s chunking + // degenerates to a scalar remainder loop, so the explicit AVX kernel + // wins big; above it the autovectorizer is already good and the + // 8-wide kernel can lose, so keep the pre-dispatch kernel exactly. + // + // SAFETY: avx2+fma are enabled for the whole crate by the build + // baseline, so the kernel's `#[target_feature]` contract holds + // statically. + let narrow = dimension <= 16; + batch.chunks_exact(dimension).map(move |y| { + if narrow { + unsafe { x86::dot_f32_avx_fma(x, y) } + } else { + dot_f32_scalar(x, y) + } + }) + } + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] + { + dot_batch_f32_runtime_dispatch(x, batch, dimension) + } + #[cfg(not(target_arch = "x86_64"))] + { + batch.chunks_exact(dimension).map(move |y| Self::dot(x, y)) + } + } +} + +/// Sub-AVX2 builds: the scalar kernel cannot reach the wide registers, so pick +/// a `#[target_feature]` kernel — once for the batch, not once per vector. +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] +#[inline] +fn dot_batch_f32_runtime_dispatch<'a>( + x: &'a [f32], + batch: &'a [f32], + dimension: usize, +) -> impl Iterator + 'a { + // SAFETY: each kernel is entered only under its matching runtime detection. + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => { + BatchIter::Eager(unsafe { x86::dot_batch_f32_avx512(x, batch, dimension) }.into_iter()) + } + SimdSupport::Avx2 | SimdSupport::AvxFma => { + BatchIter::Eager(unsafe { x86::dot_batch_f32_avx_fma(x, batch, dimension) }.into_iter()) + } + SimdSupport::Avx => { + BatchIter::Eager(unsafe { x86::dot_batch_f32_avx(x, batch, dimension) }.into_iter()) + } + _ => BatchIter::Lazy( + batch + .chunks_exact(dimension) + .map(move |y| dot_f32_scalar(x, y)), + ), + } +} + +/// Dot product for f32, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the +/// auto-vectorised scalar loop. +#[inline] +fn dot_f32_dispatched(x: &[f32], y: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { x86::dot_f32_avx512(x, y) }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { x86::dot_f32_avx_fma(x, y) }, + SimdSupport::Avx => unsafe { x86::dot_f32_avx(x, y) }, + _ => dot_f32_scalar(x, y), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + dot_f32_scalar(x, y) } } +/// Portable scalar dot product for f32. Used as the x86_64 fallback when no +/// AVX2 is detected, and as the only path on non-x86 architectures. The +/// `LANES = 16` chunking matches the explicit-SIMD inner kernels above. +#[inline] +fn dot_f32_scalar(x: &[f32], y: &[f32]) -> f32 { + dot_scalar::(x, y) +} + impl Dot for f64 { #[inline] fn dot(x: &[Self], y: &[Self]) -> f32 { @@ -233,9 +379,246 @@ impl Dot for f64 { } } -/// Explicit SIMD dot product for f64. +/// Dot product for f64, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the SIMD +/// primitives in `crate::simd::f64`, unconditionally backed by NEON / LSX-LASX. #[inline] fn dot_f64_simd(x: &[f64], y: &[f64]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { x86::dot_f64_avx512(x, y) }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { x86::dot_f64_avx_fma(x, y) }, + SimdSupport::Avx => unsafe { x86::dot_f64_avx(x, y) }, + _ => dot_f64_scalar(x, y), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + dot_f64_simd_other(x, y) + } +} + +/// Portable scalar dot product for f64. Used as the x86_64 fallback when no +/// AVX2 is detected, and exposed for cross-backend parity testing. +#[cfg(target_arch = "x86_64")] +#[inline] +fn dot_f64_scalar(x: &[f64], y: &[f64]) -> f32 { + x.iter().zip(y.iter()).map(|(&a, &b)| a * b).sum::() as f32 +} + +#[cfg(target_arch = "x86_64")] +mod x86 { + use std::arch::x86_64::*; + + use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::x86::hsum256_ps; + use crate::simd::{FloatSimd, SIMD}; + + /// Dot product of `x` against every `dimension`-sized vector in `batch`, + /// entering the AVX-512 tier once for the whole batch rather than once per + /// vector. + /// + /// # Safety + /// The host must support AVX-512F. + /// + /// Only compiled for builds whose baseline is below avx2+fma; at or above + /// that baseline `dot_batch` inlines the kernel directly and never runtime- + /// dispatches, so this wrapper would be dead code (see `dot_batch`). + #[cfg(not(all(target_feature = "avx2", target_feature = "fma")))] + #[target_feature(enable = "avx512f")] + pub(super) unsafe fn dot_batch_f32_avx512( + x: &[f32], + batch: &[f32], + dimension: usize, + ) -> Vec { + batch + .chunks_exact(dimension) + .map(|y| unsafe { dot_f32_avx512(x, y) }) + .collect() + } + + /// As [`dot_batch_f32_avx512`], for the AVX2 and AVX+FMA tiers. + /// + /// # Safety + /// The host must support AVX and FMA. + #[cfg(not(all(target_feature = "avx2", target_feature = "fma")))] + #[target_feature(enable = "avx,fma")] + pub(super) unsafe fn dot_batch_f32_avx_fma( + x: &[f32], + batch: &[f32], + dimension: usize, + ) -> Vec { + batch + .chunks_exact(dimension) + .map(|y| unsafe { dot_f32_avx_fma(x, y) }) + .collect() + } + + /// As [`dot_batch_f32_avx512`], for the AVX-without-FMA tier. + /// + /// # Safety + /// The host must support AVX. + #[cfg(not(all(target_feature = "avx2", target_feature = "fma")))] + #[target_feature(enable = "avx")] + pub(super) unsafe fn dot_batch_f32_avx(x: &[f32], batch: &[f32], dimension: usize) -> Vec { + batch + .chunks_exact(dimension) + .map(|y| unsafe { dot_f32_avx(x, y) }) + .collect() + } + + /// AVX-512 path for f64: 8-wide `__m512d` with `vfmadd231pd` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn dot_f64_avx512(x: &[f64], y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm512_setzero_pd(); + for i in (0..unrolled_len).step_by(8) { + let a = _mm512_loadu_pd(x.as_ptr().add(i)); + let b = _mm512_loadu_pd(y.as_ptr().add(i)); + acc = _mm512_fmadd_pd(a, b, acc); + } + + let tail: f64 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + (_mm512_reduce_add_pd(acc) + tail) as f32 + } + + /// AVX + FMA path for f64. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn dot_f64_avx_fma(x: &[f64], y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc8 = f64x8::zeros(); + for i in (0..unrolled_len).step_by(8) { + let a = f64x8::load_unaligned(x.as_ptr().add(i)); + let b = f64x8::load_unaligned(y.as_ptr().add(i)); + acc8.multiply_add(a, b); + } + + let aligned_len = dim / 4 * 4; + let mut acc4 = f64x4::zeros(); + for i in (unrolled_len..aligned_len).step_by(4) { + let a = f64x4::load_unaligned(x.as_ptr().add(i)); + let b = f64x4::load_unaligned(y.as_ptr().add(i)); + acc4.multiply_add(a, b); + } + + let tail: f64 = x[aligned_len..] + .iter() + .zip(y[aligned_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + (acc8.reduce_sum() + acc4.reduce_sum() + tail) as f32 + } + + /// AVX-only path for f64 (no FMA): `_mm256_mul_pd` + `_mm256_add_pd` per iteration for Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn dot_f64_avx(x: &[f64], y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 4 * 4; + + let mut acc = _mm256_setzero_pd(); + for i in (0..unrolled_len).step_by(4) { + let a = _mm256_loadu_pd(x.as_ptr().add(i)); + let b = _mm256_loadu_pd(y.as_ptr().add(i)); + acc = _mm256_add_pd(acc, _mm256_mul_pd(a, b)); + } + + // Horizontal sum of __m256d -> f64. Two pairwise adds across lanes. + let lo = _mm256_castpd256_pd128(acc); + let hi = _mm256_extractf128_pd(acc, 1); + let sum128 = _mm_add_pd(lo, hi); + let sum64 = _mm_add_pd(sum128, _mm_unpackhi_pd(sum128, sum128)); + let acc_sum = _mm_cvtsd_f64(sum64); + + let tail: f64 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + (acc_sum + tail) as f32 + } + + /// AVX-512 path for f32: 16-wide `__m512` with `vfmadd231ps` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn dot_f32_avx512(x: &[f32], y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + + let mut acc = _mm512_setzero_ps(); + for i in (0..unrolled_len).step_by(16) { + let a = _mm512_loadu_ps(x.as_ptr().add(i)); + let b = _mm512_loadu_ps(y.as_ptr().add(i)); + acc = _mm512_fmadd_ps(a, b, acc); + } + + let tail: f32 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + _mm512_reduce_add_ps(acc) + tail + } + + /// AVX + FMA path for f32. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn dot_f32_avx_fma(x: &[f32], y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..unrolled_len).step_by(8) { + let a = _mm256_loadu_ps(x.as_ptr().add(i)); + let b = _mm256_loadu_ps(y.as_ptr().add(i)); + acc = _mm256_fmadd_ps(a, b, acc); + } + + let tail: f32 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + hsum256_ps(acc) + tail + } + + /// AVX-only path for f32 (no FMA): `_mm256_mul_ps` + `_mm256_add_ps` per iteration for Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn dot_f32_avx(x: &[f32], y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..unrolled_len).step_by(8) { + let a = _mm256_loadu_ps(x.as_ptr().add(i)); + let b = _mm256_loadu_ps(y.as_ptr().add(i)); + acc = _mm256_add_ps(acc, _mm256_mul_ps(a, b)); + } + + let tail: f32 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + hsum256_ps(acc) + tail + } +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn dot_f64_simd_other(x: &[f64], y: &[f64]) -> f32 { use crate::simd::f64::{f64x4, f64x8}; use crate::simd::{FloatSimd, SIMD}; @@ -285,7 +668,7 @@ pub fn dot_distance_batch<'a, T: Dot>( ) -> Box + 'a> { assume_eq!(from.len(), dimension); assume_eq!(to.len() % dimension, 0); - Box::new(to.chunks_exact(dimension).map(|v| dot_distance(from, v))) + Box::new(T::dot_batch(from, to, dimension).map(|d| 1.0 - d)) } fn do_dot_distance_arrow_batch( @@ -309,10 +692,9 @@ where to.value_type() )))?; - let dists = to_values - .as_slice() - .chunks_exact(dimension) - .map(|v| dot_distance(from.as_slice(), v)); + // Route through `dot_distance_batch` rather than mapping `dot_distance` per + // vector, so this entry point gets the same hoisted dispatch. + let dists = dot_distance_batch(from.as_slice(), to_values.as_slice(), dimension); Ok(Arc::new(Float32Array::new( dists.collect(), @@ -480,5 +862,264 @@ mod tests { fn test_dot_f64((x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048)){ do_dot_test(&x, &y)?; } + + /// Cross-backend parity: scalar fallback must match the dispatched + /// SIMD path within numerical tolerance. Exercises `dot_f64_scalar` + /// directly so the runtime fallback is exercised even on AVX2-capable + /// CI hosts. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f64_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + let scalar = dot_f64_scalar(&x, &y); + let simd = dot_f64_simd(&x, &y); + let max_error = max_error::(&x, &y); + prop_assert!(approx::relative_eq!(scalar, simd, epsilon = max_error)); + } + + /// Parity check for `dot_f32_dispatched` (Branch B exclusive: the + /// auto-vectorised scalar dot path). The dispatched kernel must + /// agree with a portable f64-precision scalar reference within + /// numerical tolerance. The reference is hand-rolled here to keep + /// this test architecture-agnostic (the x86_64-only `dot_f64_scalar` + /// helper is gated above). + #[test] + fn test_dot_f32_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let scalar = x_f64 + .iter() + .zip(y_f64.iter()) + .map(|(&a, &b)| a * b) + .sum::() as f32; + let simd = ::dot(&x, &y); + let max_error = max_error::(&x_f64, &y_f64); + prop_assert!(approx::relative_eq!(scalar, simd, epsilon = max_error)); + } + + /// AVX-512-direct parity for f32: explicitly compares the scalar + /// fallback against the native f32 AVX-512 inner kernel on + /// AVX-512F-capable hosts. Early-returns on hosts without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f32_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + let scalar = dot_f32_scalar(&x, &y); + let avx512 = unsafe { x86::dot_f32_avx512(&x, &y) }; + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let max_error = max_error::(&x_f64, &y_f64); + prop_assert!(approx::relative_eq!(scalar, avx512, epsilon = max_error)); + } + + /// AVX + FMA-direct parity for the f32 dot kernel. Covers the AMD + /// Piledriver / Steamroller / FX-7500 tier. Early-returns on hosts + /// without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f32_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + let scalar = dot_f32_scalar(&x, &y); + let avx_fma = unsafe { x86::dot_f32_avx_fma(&x, &y) }; + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let max_error = max_error::(&x_f64, &y_f64); + prop_assert!(approx::relative_eq!(scalar, avx_fma, epsilon = max_error)); + } + + /// AVX-only-direct parity for the f32 dot kernel. Covers the Intel + /// Sandy Bridge / Ivy Bridge tier. Early-returns on hosts without + /// AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f32_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + let scalar = dot_f32_scalar(&x, &y); + let avx = unsafe { x86::dot_f32_avx(&x, &y) }; + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let max_error = max_error::(&x_f64, &y_f64); + prop_assert!(approx::relative_eq!(scalar, avx, epsilon = max_error)); + } + + /// AVX-512-direct parity: explicitly compares the scalar fallback + /// against the native AVX-512 inner kernel on AVX-512F-capable hosts + /// (Skylake-X+, Ice Lake, Sapphire Rapids, Zen 4). Early-returns on + /// hosts without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f64_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + let scalar = dot_f64_scalar(&x, &y); + let avx512 = unsafe { x86::dot_f64_avx512(&x, &y) }; + let max_error = max_error::(&x, &y); + prop_assert!(approx::relative_eq!(scalar, avx512, epsilon = max_error)); + } + + /// AVX + FMA-direct parity for the f64 dot kernel. Covers the AMD + /// Piledriver / Steamroller / FX-7500 tier. Early-returns on hosts + /// without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f64_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + let scalar = dot_f64_scalar(&x, &y); + let avx_fma = unsafe { x86::dot_f64_avx_fma(&x, &y) }; + let max_error = max_error::(&x, &y); + prop_assert!(approx::relative_eq!(scalar, avx_fma, epsilon = max_error)); + } + + /// AVX-only-direct parity for the f64 dot kernel. Covers the Intel + /// Sandy Bridge / Ivy Bridge tier (AVX without FMA). Early-returns + /// on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f64_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + let scalar = dot_f64_scalar(&x, &y); + let avx = unsafe { x86::dot_f64_avx(&x, &y) }; + let max_error = max_error::(&x, &y); + prop_assert!(approx::relative_eq!(scalar, avx, epsilon = max_error)); + } + } + + /// `dot_batch` must agree with the per-vector `dot` it replaced, on every + /// build: AVX2-baseline, hoisted-dispatch, and portable fallback all + /// funnel through here. + #[rstest::rstest] + #[case::dim_8(8)] + #[case::dim_16(16)] + #[case::dim_32(32)] + #[case::dim_1024(1024)] + fn test_dot_batch_f32_matches_per_vector_dot(#[case] dimension: usize) { + let num_vectors = 5; + let x: Vec = (0..dimension) + .map(|i| ((i % 13) as f32) * 0.25 + 1.0) + .collect(); + let batch: Vec = (0..dimension * num_vectors) + .map(|i| ((i % 11) as f32) * 0.5 - 2.0) + .collect(); + + let got: Vec = f32::dot_batch(&x, &batch, dimension).collect(); + let want: Vec = batch + .chunks_exact(dimension) + .map(|y| f32::dot(&x, y)) + .collect(); + + assert_eq!(got.len(), num_vectors); + for (g, w) in got.iter().zip(want.iter()) { + assert!( + approx::relative_eq!(g, w, epsilon = 1e-4), + "dim {dimension}: batch {g} != per-vector {w}" + ); + } + } + + /// `dot_distance_batch` still yields `1.0 - dot`, unchanged by the hoist. + #[test] + fn test_dot_distance_batch_preserves_distance_semantics() { + let dimension = 32; + let x: Vec = (0..dimension).map(|i| (i as f32) * 0.1).collect(); + let batch: Vec = (0..dimension * 3).map(|i| (i as f32) * 0.05).collect(); + + let got: Vec = dot_distance_batch(&x, &batch, dimension).collect(); + for (chunk, &g) in batch.chunks_exact(dimension).zip(got.iter()) { + assert!(approx::relative_eq!( + g, + 1.0 - f32::dot(&x, chunk), + epsilon = 1e-5 + )); + } + } + + /// The per-batch `#[target_feature]` kernels are only reached on sub-AVX2 + /// builds or AVX-512 hosts, so call them directly to cover them. + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] + fn check_dot_batch_kernel(kernel: unsafe fn(&[f32], &[f32], usize) -> Vec) { + for dimension in [8_usize, 16, 40] { + let num_vectors = 3; + let x: Vec = (0..dimension).map(|i| (i as f32) * 0.5 + 1.0).collect(); + let batch: Vec = (0..dimension * num_vectors) + .map(|i| ((i % 7) as f32) + 1.0) + .collect(); + + let got = unsafe { kernel(&x, &batch, dimension) }; + assert_eq!(got.len(), num_vectors); + for (chunk, &g) in batch.chunks_exact(dimension).zip(got.iter()) { + let want = dot_scalar::(&x, chunk); + assert!( + approx::relative_eq!(g, want, epsilon = 1e-4), + "dim {dimension}: kernel {g} != scalar {want}" + ); + } + } + } + + // The runtime-dispatch batch kernels only exist in sub-avx2+fma builds + // (see `x86::dot_batch_f32_avx512`), so gate their tests the same way. + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] + #[test] + fn test_dot_batch_avx_fma_matches_scalar() { + if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { + return; + } + check_dot_batch_kernel(x86::dot_batch_f32_avx_fma); + } + + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] + #[test] + fn test_dot_batch_avx_matches_scalar() { + if !std::is_x86_feature_detected!("avx") { + return; + } + check_dot_batch_kernel(x86::dot_batch_f32_avx); + } + + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] + #[test] + fn test_dot_batch_avx512_matches_scalar() { + if !std::is_x86_feature_detected!("avx512f") { + return; + } + check_dot_batch_kernel(x86::dot_batch_f32_avx512); } } diff --git a/rust/lance-linalg/src/distance/dot_u8.rs b/rust/lance-linalg/src/distance/dot_u8.rs index de5522cddfe..7b2e335f094 100644 --- a/rust/lance-linalg/src/distance/dot_u8.rs +++ b/rust/lance-linalg/src/distance/dot_u8.rs @@ -134,6 +134,10 @@ fn select_backend() -> DotU8Fn { if is_x86_feature_detected!("avx2") { return |a, b| unsafe { x86::dot_u8_avx2(a, b) }; } + // AvxFma and Avx hosts (AMD Piledriver / Steamroller, Intel Sandy + // Bridge / Ivy Bridge) fall through to scalar: the AVX2 inner uses + // `vpmaddubsw` / `vpmaddwd` integer ops which neither AVX nor + // AVX+FMA provides. } dot_u8_scalar diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs index c47aedd749f..7ee2f4bc537 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/rust/lance-linalg/src/distance/l2.rs @@ -20,18 +20,41 @@ use lance_arrow::{ArrowFloatType, FixedSizeListArrayExt, FloatArray}; use lance_core::assume_eq; use lance_core::deepsize::DeepSizeOf; use lance_core::utils::cpu::SIMD_SUPPORT; -#[cfg(feature = "fp16kernels")] +// Named tiers are only matched on x86_64, or by the fp16 kernels on the other +// architectures; without either, nothing below names a `SimdSupport` variant. +#[cfg(any(feature = "fp16kernels", target_arch = "x86_64"))] use lance_core::utils::cpu::SimdSupport; use num_traits::{AsPrimitive, Num}; +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] +use crate::distance::BatchIter; + /// Calculate the L2 distance between two vectors. /// pub trait L2: Num { /// Calculate the L2 distance between two vectors. fn l2(x: &[Self], y: &[Self]) -> f32; - fn l2_batch(x: &[Self], y: &[Self], dimension: usize) -> impl Iterator { - y.chunks_exact(dimension).map(|v| Self::l2(x, v)) + /// L2 distance from `x` to each `dimension`-sized vector in `y`. + /// + /// The default calls [`L2::l2`] per vector. `f32` overrides it so the SIMD + /// tier is chosen once for the whole batch instead of once per vector — + /// on a build whose baseline already implies AVX2, per-vector dispatch + /// costs more than the kernel it selects. + /// + /// Returns `impl Iterator` rather than a trait object: the k-means + /// assignment loop drives this one element at a time, so a + /// `Box` would cost a virtual call per element and an + /// allocation per batch. + fn l2_batch<'a>( + x: &'a [Self], + y: &'a [Self], + dimension: usize, + ) -> impl Iterator + 'a { + y.chunks_exact(dimension).map(move |v| Self::l2(x, v)) } } @@ -52,7 +75,6 @@ pub fn l2(from: &[T], to: &[T]) -> f32 { pub fn l2_f32(x: &[f32], y: &[f32]) -> f32 { #[cfg(target_arch = "x86_64")] { - use lance_core::utils::cpu::SimdSupport; if matches!(*SIMD_SUPPORT, SimdSupport::Avx512 | SimdSupport::Avx512FP16) { // SAFETY: guarded by the runtime AVX-512 detection above. return unsafe { l2_f32_avx512(x, y) }; @@ -191,6 +213,9 @@ impl L2 for bf16 { SimdSupport::Lsx => unsafe { bf16_kernel::l2_bf16_lsx(x.as_ptr(), y.as_ptr(), x.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the bf16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => l2_scalar::(x, y), } } @@ -244,6 +269,9 @@ impl L2 for f16 { SimdSupport::Lsx => unsafe { kernel::l2_f16_lsx(x.as_ptr(), y.as_ptr(), x.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the f16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => l2_scalar::(x, y), } } @@ -252,12 +280,114 @@ impl L2 for f16 { impl L2 for f32 { #[inline] fn l2(x: &[Self], y: &[Self]) -> f32 { - // 16 = 512 (avx512) / 8 bits / 4 (sizeof(f32)) - // See https://github.com/lance-format/lance/pull/2450. - l2_scalar::(x, y) + // Trait methods cannot carry `#[target_feature]` attributes, so the body + // lives in a free function that runtime-dispatches via `*SIMD_SUPPORT` + // to an AVX2 or AVX-512 inner kernel on capable hosts, or a portable + // scalar fallback. + l2_f32_dispatched(x, y) + } + + fn l2_batch<'a>( + x: &'a [Self], + y: &'a [Self], + dimension: usize, + ) -> impl Iterator + 'a { + // Exactly one arm compiles; see `Dot::dot_batch` for f32. + // See `Dot::dot_batch` for f32. + #[cfg(all( + target_arch = "x86_64", + target_feature = "avx2", + target_feature = "fma" + ))] + { + // `l2_scalar::<_, _, 16>` chunks the vector by 16 lanes. At or below + // that width the chunking degenerates to its scalar remainder loop + // and vectorizes nothing, so the explicit AVX kernel is worth ~40%. + // Above it the autovectorizer already does well and the 8-wide + // kernel can lose, so keep the exact kernel the pre-dispatch code + // used and stay non-regressing by construction. + // + // SAFETY: the build baseline enables avx2+fma, which imply avx+fma, + // so the kernel's `#[target_feature]` contract is met statically. + let narrow = dimension <= 16; + y.chunks_exact(dimension).map(move |v| { + if narrow { + unsafe { x86::l2_f32_avx_fma(x, v) } + } else { + l2_f32_scalar(x, v) + } + }) + } + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] + { + l2_batch_f32_runtime_dispatch(x, y, dimension) + } + #[cfg(not(target_arch = "x86_64"))] + { + y.chunks_exact(dimension).map(move |v| Self::l2(x, v)) + } + } +} + +/// Sub-AVX2 builds: pick a `#[target_feature]` kernel once for the batch. +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] +#[inline] +fn l2_batch_f32_runtime_dispatch<'a>( + x: &'a [f32], + y: &'a [f32], + dimension: usize, +) -> impl Iterator + 'a { + // SAFETY: each kernel is entered only under its matching runtime detection. + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => { + BatchIter::Eager(unsafe { x86::l2_batch_f32_avx512(x, y, dimension) }.into_iter()) + } + SimdSupport::Avx2 | SimdSupport::AvxFma => { + BatchIter::Eager(unsafe { x86::l2_batch_f32_avx_fma(x, y, dimension) }.into_iter()) + } + SimdSupport::Avx => { + BatchIter::Eager(unsafe { x86::l2_batch_f32_avx(x, y, dimension) }.into_iter()) + } + _ => BatchIter::Lazy(y.chunks_exact(dimension).map(move |v| l2_f32_scalar(x, v))), } } +/// L2 distance for f32, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the +/// auto-vectorised scalar loop. +#[inline] +fn l2_f32_dispatched(x: &[f32], y: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { x86::l2_f32_avx512(x, y) }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { x86::l2_f32_avx_fma(x, y) }, + SimdSupport::Avx => unsafe { x86::l2_f32_avx(x, y) }, + _ => l2_f32_scalar(x, y), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + l2_f32_scalar(x, y) + } +} + +/// Portable scalar L2 distance for f32. Used as the x86_64 fallback when no +/// AVX2 is detected, and as the only path on non-x86 architectures. The +/// `LANES = 16` chunking matches the explicit-SIMD inner kernels above. +#[inline] +fn l2_f32_scalar(x: &[f32], y: &[f32]) -> f32 { + // 16 = 512 (avx512) / 8 bits / 4 (sizeof(f32)) + // See https://github.com/lance-format/lance/pull/2450. + l2_scalar::(x, y) +} + impl L2 for f64 { #[inline] fn l2(x: &[Self], y: &[Self]) -> f32 { @@ -265,9 +395,277 @@ impl L2 for f64 { } } -/// Explicit SIMD L2 distance for f64. +/// L2 distance for f64, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the SIMD +/// primitives in `crate::simd::f64`, unconditionally backed by NEON / LSX-LASX. #[inline] fn l2_f64_simd(x: &[f64], y: &[f64]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { x86::l2_f64_avx512(x, y) }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { x86::l2_f64_avx_fma(x, y) }, + SimdSupport::Avx => unsafe { x86::l2_f64_avx(x, y) }, + _ => l2_f64_scalar(x, y), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + l2_f64_simd_other(x, y) + } +} + +/// Portable scalar L2 distance for f64. Used as the x86_64 fallback when no +/// AVX2 is detected, and exposed for cross-backend parity testing. +#[cfg(target_arch = "x86_64")] +#[inline] +fn l2_f64_scalar(x: &[f64], y: &[f64]) -> f32 { + x.iter() + .zip(y.iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum::() as f32 +} + +#[cfg(target_arch = "x86_64")] +mod x86 { + use std::arch::x86_64::*; + + use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::x86::hsum256_ps; + use crate::simd::{FloatSimd, SIMD}; + + /// L2 distance from `x` to every `dimension`-sized vector in `batch`, with + /// the AVX-512 tier entered once for the whole batch rather than once per + /// vector. + /// + /// # Safety + /// The host must support AVX-512F. + /// + /// Only compiled for builds whose baseline is below avx2+fma; at or above + /// that baseline `l2_batch` inlines the kernel directly and never runtime- + /// dispatches, so this wrapper would be dead code (see `l2_batch`). + #[cfg(not(all(target_feature = "avx2", target_feature = "fma")))] + #[target_feature(enable = "avx512f")] + pub(super) unsafe fn l2_batch_f32_avx512( + x: &[f32], + batch: &[f32], + dimension: usize, + ) -> Vec { + batch + .chunks_exact(dimension) + .map(|y| unsafe { l2_f32_avx512(x, y) }) + .collect() + } + + /// As [`l2_batch_f32_avx512`], for the AVX+FMA and AVX2 tiers. + /// + /// # Safety + /// The host must support AVX and FMA. + #[cfg(not(all(target_feature = "avx2", target_feature = "fma")))] + #[target_feature(enable = "avx,fma")] + pub(super) unsafe fn l2_batch_f32_avx_fma( + x: &[f32], + batch: &[f32], + dimension: usize, + ) -> Vec { + batch + .chunks_exact(dimension) + .map(|y| unsafe { l2_f32_avx_fma(x, y) }) + .collect() + } + + /// As [`l2_batch_f32_avx512`], for the AVX-without-FMA tier. + /// + /// # Safety + /// The host must support AVX. + #[cfg(not(all(target_feature = "avx2", target_feature = "fma")))] + #[target_feature(enable = "avx")] + pub(super) unsafe fn l2_batch_f32_avx(x: &[f32], batch: &[f32], dimension: usize) -> Vec { + batch + .chunks_exact(dimension) + .map(|y| unsafe { l2_f32_avx(x, y) }) + .collect() + } + + /// AVX-512 path for f64: 8-wide `__m512d` with `vsubpd` + `vfmadd231pd` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn l2_f64_avx512(x: &[f64], y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm512_setzero_pd(); + for i in (0..unrolled_len).step_by(8) { + let a = _mm512_loadu_pd(x.as_ptr().add(i)); + let b = _mm512_loadu_pd(y.as_ptr().add(i)); + let diff = _mm512_sub_pd(a, b); + acc = _mm512_fmadd_pd(diff, diff, acc); + } + + let tail: f64 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum(); + + (_mm512_reduce_add_pd(acc) + tail) as f32 + } + + /// AVX + FMA path for f64. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn l2_f64_avx_fma(x: &[f64], y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc8 = f64x8::zeros(); + for i in (0..unrolled_len).step_by(8) { + let a = f64x8::load_unaligned(x.as_ptr().add(i)); + let b = f64x8::load_unaligned(y.as_ptr().add(i)); + let diff = a - b; + acc8.multiply_add(diff, diff); + } + + let aligned_len = dim / 4 * 4; + let mut acc4 = f64x4::zeros(); + for i in (unrolled_len..aligned_len).step_by(4) { + let a = f64x4::load_unaligned(x.as_ptr().add(i)); + let b = f64x4::load_unaligned(y.as_ptr().add(i)); + let diff = a - b; + acc4.multiply_add(diff, diff); + } + + let tail: f64 = x[aligned_len..] + .iter() + .zip(y[aligned_len..].iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum(); + + (acc8.reduce_sum() + acc4.reduce_sum() + tail) as f32 + } + + /// AVX-only path for f64 (no FMA): squared diff via `_mm256_mul_pd` + `_mm256_add_pd` for Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn l2_f64_avx(x: &[f64], y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 4 * 4; + + let mut acc = _mm256_setzero_pd(); + for i in (0..unrolled_len).step_by(4) { + let a = _mm256_loadu_pd(x.as_ptr().add(i)); + let b = _mm256_loadu_pd(y.as_ptr().add(i)); + let diff = _mm256_sub_pd(a, b); + acc = _mm256_add_pd(acc, _mm256_mul_pd(diff, diff)); + } + + // Horizontal sum of __m256d -> f64. + let lo = _mm256_castpd256_pd128(acc); + let hi = _mm256_extractf128_pd(acc, 1); + let sum128 = _mm_add_pd(lo, hi); + let sum64 = _mm_add_pd(sum128, _mm_unpackhi_pd(sum128, sum128)); + let acc_sum = _mm_cvtsd_f64(sum64); + + let tail: f64 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum(); + + (acc_sum + tail) as f32 + } + + /// AVX-512 path for f32: 16-wide `__m512` with `vsubps` + `vfmadd231ps` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn l2_f32_avx512(x: &[f32], y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + + let mut acc = _mm512_setzero_ps(); + for i in (0..unrolled_len).step_by(16) { + let a = _mm512_loadu_ps(x.as_ptr().add(i)); + let b = _mm512_loadu_ps(y.as_ptr().add(i)); + let diff = _mm512_sub_ps(a, b); + acc = _mm512_fmadd_ps(diff, diff, acc); + } + + let tail: f32 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum(); + + _mm512_reduce_add_ps(acc) + tail + } + + /// AVX + FMA path for f32. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn l2_f32_avx_fma(x: &[f32], y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..unrolled_len).step_by(8) { + let a = _mm256_loadu_ps(x.as_ptr().add(i)); + let b = _mm256_loadu_ps(y.as_ptr().add(i)); + let diff = _mm256_sub_ps(a, b); + acc = _mm256_fmadd_ps(diff, diff, acc); + } + + let tail: f32 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum(); + + hsum256_ps(acc) + tail + } + + /// AVX-only path for f32 (no FMA): squared diff via `_mm256_mul_ps` + `_mm256_add_ps` for Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn l2_f32_avx(x: &[f32], y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..unrolled_len).step_by(8) { + let a = _mm256_loadu_ps(x.as_ptr().add(i)); + let b = _mm256_loadu_ps(y.as_ptr().add(i)); + let diff = _mm256_sub_ps(a, b); + acc = _mm256_add_ps(acc, _mm256_mul_ps(diff, diff)); + } + + let tail: f32 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum(); + + hsum256_ps(acc) + tail + } +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn l2_f64_simd_other(x: &[f64], y: &[f64]) -> f32 { use crate::simd::f64::{f64x4, f64x8}; use crate::simd::{FloatSimd, SIMD}; @@ -671,6 +1069,136 @@ mod tests { fn test_l2_distance_f64((x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048)){ do_l2_test(&x, &y)?; } + + /// Cross-backend parity: scalar fallback must match the dispatched + /// SIMD path within numerical tolerance. Exercises `l2_f64_scalar` + /// directly so the runtime fallback is exercised even on AVX2-capable + /// CI hosts. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f64_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + let scalar = l2_f64_scalar(&x, &y); + let simd = l2_f64_simd(&x, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-6)); + } + + /// Parity check for `l2_f32_dispatched` (Branch B exclusive: the + /// auto-vectorised scalar L2 path). The dispatched kernel must + /// agree with a portable f64-precision scalar reference within + /// numerical tolerance. The reference is hand-rolled here to keep + /// this test architecture-agnostic (the x86_64-only `l2_f64_scalar` + /// helper is gated above). + #[test] + fn test_l2_f32_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + let scalar = x + .iter() + .zip(y.iter()) + .map(|(&a, &b)| ((a as f64) - (b as f64)).powi(2)) + .sum::() as f32; + let simd = ::l2(&x, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3)); + } + + /// AVX-512-direct parity: explicitly compares the scalar fallback + /// against the native AVX-512 inner kernel on AVX-512F-capable hosts + /// (Skylake-X+, Ice Lake, Sapphire Rapids, Zen 4). Early-returns on + /// hosts without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f64_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + let scalar = l2_f64_scalar(&x, &y); + let avx512 = unsafe { x86::l2_f64_avx512(&x, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-6)); + } + + /// AVX + FMA-direct parity for the f64 L2 kernel. Covers the AMD + /// Piledriver / Steamroller / FX-7500 tier. Early-returns on hosts + /// without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f64_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + let scalar = l2_f64_scalar(&x, &y); + let avx_fma = unsafe { x86::l2_f64_avx_fma(&x, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-6)); + } + + /// AVX-only-direct parity for the f64 L2 kernel. Covers the Intel + /// Sandy Bridge / Ivy Bridge tier. Early-returns on hosts without + /// AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f64_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + let scalar = l2_f64_scalar(&x, &y); + let avx = unsafe { x86::l2_f64_avx(&x, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-6)); + } + + /// AVX-512-direct parity for f32: explicitly compares the scalar + /// fallback against the native f32 AVX-512 inner kernel on + /// AVX-512F-capable hosts. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f32_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + let scalar = l2_f32_scalar(&x, &y); + let avx512 = unsafe { x86::l2_f32_avx512(&x, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-3)); + } + + /// AVX + FMA-direct parity for the f32 L2 kernel. Covers the AMD + /// Piledriver / Steamroller / FX-7500 tier. Early-returns on hosts + /// without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f32_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + let scalar = l2_f32_scalar(&x, &y); + let avx_fma = unsafe { x86::l2_f32_avx_fma(&x, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-3)); + } + + /// AVX-only-direct parity for the f32 L2 kernel. Covers the Intel + /// Sandy Bridge / Ivy Bridge tier. Early-returns on hosts without + /// AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f32_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + let scalar = l2_f32_scalar(&x, &y); + let avx = unsafe { x86::l2_f32_avx(&x, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-3)); + } } #[test] @@ -787,4 +1315,100 @@ mod tests { assert_relative_eq!(d1[0], 0.0); // q1 == target[0] assert_relative_eq!(d2[1], 0.0); // q2 == target[1] } + + /// `l2_batch` must agree with the per-vector `l2` it replaced, on every + /// build: the AVX2-baseline path, the hoisted-dispatch path, and the + /// portable fallback all funnel through here. + #[rstest::rstest] + #[case::dim_8(8)] + #[case::dim_16(16)] + #[case::dim_32(32)] + #[case::dim_1024(1024)] + fn test_l2_batch_f32_matches_per_vector_l2(#[case] dimension: usize) { + let num_vectors = 5; + let x: Vec = (0..dimension) + .map(|i| ((i % 13) as f32) * 0.25 + 1.0) + .collect(); + let batch: Vec = (0..dimension * num_vectors) + .map(|i| ((i % 11) as f32) * 0.5 - 2.0) + .collect(); + + let got: Vec = f32::l2_batch(&x, &batch, dimension).collect(); + let want: Vec = batch + .chunks_exact(dimension) + .map(|y| f32::l2(&x, y)) + .collect(); + + assert_eq!(got.len(), num_vectors); + for (g, w) in got.iter().zip(want.iter()) { + assert!( + approx::relative_eq!(g, w, epsilon = 1e-4), + "dim {dimension}: batch {g} != per-vector {w}" + ); + } + } + + /// The per-batch `#[target_feature]` kernels are only reached on sub-AVX2 + /// builds or AVX-512 hosts, so call them directly to cover them. + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] + fn check_l2_batch_kernel(kernel: unsafe fn(&[f32], &[f32], usize) -> Vec) { + for dimension in [8_usize, 16, 40] { + let num_vectors = 3; + let x: Vec = (0..dimension).map(|i| (i as f32) * 0.5 + 1.0).collect(); + let batch: Vec = (0..dimension * num_vectors) + .map(|i| ((i % 7) as f32) + 1.0) + .collect(); + + let got = unsafe { kernel(&x, &batch, dimension) }; + assert_eq!(got.len(), num_vectors); + for (chunk, &g) in batch.chunks_exact(dimension).zip(got.iter()) { + let want = l2_scalar::(&x, chunk); + assert!( + approx::relative_eq!(g, want, epsilon = 1e-4), + "dim {dimension}: kernel {g} != scalar {want}" + ); + } + } + } + + // The runtime-dispatch batch kernels only exist in sub-avx2+fma builds + // (see `x86::l2_batch_f32_avx512`), so gate their tests the same way. + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] + #[test] + fn test_l2_batch_avx_fma_matches_scalar() { + if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { + return; + } + check_l2_batch_kernel(x86::l2_batch_f32_avx_fma); + } + + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] + #[test] + fn test_l2_batch_avx_matches_scalar() { + if !std::is_x86_feature_detected!("avx") { + return; + } + check_l2_batch_kernel(x86::l2_batch_f32_avx); + } + + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] + #[test] + fn test_l2_batch_avx512_matches_scalar() { + if !std::is_x86_feature_detected!("avx512f") { + return; + } + check_l2_batch_kernel(x86::l2_batch_f32_avx512); + } } diff --git a/rust/lance-linalg/src/distance/l2_u8.rs b/rust/lance-linalg/src/distance/l2_u8.rs index 1b111f91338..5fbf8a3af55 100644 --- a/rust/lance-linalg/src/distance/l2_u8.rs +++ b/rust/lance-linalg/src/distance/l2_u8.rs @@ -143,6 +143,10 @@ fn select_backend() -> L2U8Fn { if is_x86_feature_detected!("avx2") { return |a, b| unsafe { x86::l2_u8_avx2(a, b) }; } + // AvxFma and Avx hosts (AMD Piledriver / Steamroller, Intel Sandy + // Bridge / Ivy Bridge) fall through to scalar: the AVX2 inner uses + // AVX2 integer ops (`vpsubusb` / `vpmaddwd`) which neither AVX nor + // AVX+FMA provides. } l2_u8_scalar diff --git a/rust/lance-linalg/src/distance/norm_l2.rs b/rust/lance-linalg/src/distance/norm_l2.rs index b1daf85ab3b..8d126d6ed57 100644 --- a/rust/lance-linalg/src/distance/norm_l2.rs +++ b/rust/lance-linalg/src/distance/norm_l2.rs @@ -9,9 +9,7 @@ use arrow_array::types::{Float16Type, Float32Type, Float64Type}; use arrow_schema::DataType; use half::{bf16, f16}; #[allow(unused_imports)] -use lance_core::utils::cpu::SIMD_SUPPORT; -#[cfg(feature = "fp16kernels")] -use lance_core::utils::cpu::SimdSupport; +use lance_core::utils::cpu::{SIMD_SUPPORT, SimdSupport}; use num_traits::{AsPrimitive, Float, Num}; /// L2 normalization @@ -75,6 +73,9 @@ impl Normalize for f16 { SimdSupport::Lsx => unsafe { kernel::norm_l2_f16_lsx(vector.as_ptr(), vector.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the f16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => norm_l2_impl::(vector), } } @@ -126,6 +127,9 @@ impl Normalize for bf16 { SimdSupport::Lsx => unsafe { bf16_kernel::norm_l2_bf16_lsx(vector.as_ptr(), vector.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the bf16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => norm_l2_impl::(vector), } } @@ -134,10 +138,40 @@ impl Normalize for bf16 { impl Normalize for f32 { #[inline] fn norm_l2(vector: &[Self]) -> f32 { - norm_l2_impl::(vector) + norm_l2_f32_dispatched(vector) + } +} + +/// L2 norm for f32, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the +/// auto-vectorised scalar loop. +#[inline] +fn norm_l2_f32_dispatched(vector: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + x86::norm_l2_f32_avx512(vector) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { x86::norm_l2_f32_avx_fma(vector) }, + SimdSupport::Avx => unsafe { x86::norm_l2_f32_avx(vector) }, + _ => norm_l2_f32_scalar(vector), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + norm_l2_f32_scalar(vector) } } +/// Portable scalar L2 norm for f32. Used as the x86_64 fallback when no +/// AVX2 is detected, and as the only path on non-x86 architectures. The +/// `LANES = 16` chunking matches the explicit-SIMD inner kernels above. +#[inline] +fn norm_l2_f32_scalar(vector: &[f32]) -> f32 { + norm_l2_impl::(vector) +} + impl Normalize for f64 { #[inline] fn norm_l2(vector: &[Self]) -> f32 { @@ -145,11 +179,166 @@ impl Normalize for f64 { } } -/// Explicit SIMD implementation of L2 norm for f64. +/// L2 norm for f64. Runtime-dispatched to the best available backend. /// -/// Two-level unrolling: f64x8 main loop, f64x4 remainder, scalar tail. +/// On x86_64, dispatches via `SIMD_SUPPORT` to a native AVX-512 inner kernel +/// (Skylake-X+, Ice Lake, Sapphire Rapids, Zen 4), an AVX2 + FMA kernel +/// (Haswell+), an AVX + FMA kernel (AMD Piledriver / Steamroller), an +/// AVX-only kernel (Intel Sandy Bridge / Ivy Bridge), or a portable scalar +/// fallback. The per-tier inner functions each carry their own +/// `#[target_feature]` so they stay correct under any compile baseline. +/// On aarch64 and loongarch64, the SIMD primitives in `crate::simd::f64` +/// are unconditionally backed by NEON / LSX-LASX respectively, so no +/// runtime gate is required. #[inline] pub fn norm_l2_f64_simd(vector: &[f64]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + x86::norm_l2_f64_avx512(vector) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { x86::norm_l2_f64_avx_fma(vector) }, + SimdSupport::Avx => unsafe { x86::norm_l2_f64_avx(vector) }, + _ => norm_l2_f64_scalar(vector), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + norm_l2_f64_simd_other(vector) + } +} + +/// Portable scalar L2 norm. Used as the x86_64 fallback when no AVX2 is +/// detected, and exposed for cross-backend parity testing. +#[cfg(target_arch = "x86_64")] +#[inline] +fn norm_l2_f64_scalar(vector: &[f64]) -> f32 { + vector.iter().map(|v| v * v).sum::().sqrt() as f32 +} + +#[cfg(target_arch = "x86_64")] +mod x86 { + use std::arch::x86_64::*; + + use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::x86::hsum256_ps; + use crate::simd::{FloatSimd, SIMD}; + + /// AVX-512 path for f64: 8-wide `__m512d` with `vfmadd231pd` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn norm_l2_f64_avx512(vector: &[f64]) -> f32 { + let dim = vector.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm512_setzero_pd(); + for i in (0..unrolled_len).step_by(8) { + let v = _mm512_loadu_pd(vector.as_ptr().add(i)); + acc = _mm512_fmadd_pd(v, v, acc); + } + + let tail: f64 = vector[unrolled_len..].iter().map(|&v| v * v).sum(); + (_mm512_reduce_add_pd(acc) + tail).sqrt() as f32 + } + + /// AVX + FMA path for f64. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn norm_l2_f64_avx_fma(vector: &[f64]) -> f32 { + let dim = vector.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc8 = f64x8::zeros(); + for i in (0..unrolled_len).step_by(8) { + let v = f64x8::load_unaligned(vector.as_ptr().add(i)); + acc8.multiply_add(v, v); + } + + let aligned_len = dim / 4 * 4; + let mut acc4 = f64x4::zeros(); + for i in (unrolled_len..aligned_len).step_by(4) { + let v = f64x4::load_unaligned(vector.as_ptr().add(i)); + acc4.multiply_add(v, v); + } + + let tail: f64 = vector[aligned_len..].iter().map(|&v| v * v).sum(); + (acc8.reduce_sum() + acc4.reduce_sum() + tail).sqrt() as f32 + } + + /// AVX-only path for f64 (no FMA): `_mm256_mul_pd` + `_mm256_add_pd` per iteration for Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn norm_l2_f64_avx(vector: &[f64]) -> f32 { + let dim = vector.len(); + let unrolled_len = dim / 4 * 4; + + let mut acc = _mm256_setzero_pd(); + for i in (0..unrolled_len).step_by(4) { + let v = _mm256_loadu_pd(vector.as_ptr().add(i)); + acc = _mm256_add_pd(acc, _mm256_mul_pd(v, v)); + } + + // Horizontal sum of __m256d -> f64. Two pairwise adds across lanes. + let lo = _mm256_castpd256_pd128(acc); + let hi = _mm256_extractf128_pd(acc, 1); + let sum128 = _mm_add_pd(lo, hi); + let sum64 = _mm_add_pd(sum128, _mm_unpackhi_pd(sum128, sum128)); + let acc_sum = _mm_cvtsd_f64(sum64); + + let tail: f64 = vector[unrolled_len..].iter().map(|&v| v * v).sum(); + (acc_sum + tail).sqrt() as f32 + } + + /// AVX-512 path for f32: 16-wide `__m512` with `vfmadd231ps` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn norm_l2_f32_avx512(vector: &[f32]) -> f32 { + let dim = vector.len(); + let unrolled_len = dim / 16 * 16; + + let mut acc = _mm512_setzero_ps(); + for i in (0..unrolled_len).step_by(16) { + let v = _mm512_loadu_ps(vector.as_ptr().add(i)); + acc = _mm512_fmadd_ps(v, v, acc); + } + + let tail: f32 = vector[unrolled_len..].iter().map(|&v| v * v).sum(); + (_mm512_reduce_add_ps(acc) + tail).sqrt() + } + + /// AVX + FMA path for f32. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn norm_l2_f32_avx_fma(vector: &[f32]) -> f32 { + let dim = vector.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..unrolled_len).step_by(8) { + let v = _mm256_loadu_ps(vector.as_ptr().add(i)); + acc = _mm256_fmadd_ps(v, v, acc); + } + + let tail: f32 = vector[unrolled_len..].iter().map(|&v| v * v).sum(); + (hsum256_ps(acc) + tail).sqrt() + } + + /// AVX-only path for f32 (no FMA): `_mm256_mul_ps` + `_mm256_add_ps` per iteration for Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn norm_l2_f32_avx(vector: &[f32]) -> f32 { + let dim = vector.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..unrolled_len).step_by(8) { + let v = _mm256_loadu_ps(vector.as_ptr().add(i)); + acc = _mm256_add_ps(acc, _mm256_mul_ps(v, v)); + } + + let tail: f32 = vector[unrolled_len..].iter().map(|&v| v * v).sum(); + (hsum256_ps(acc) + tail).sqrt() + } +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn norm_l2_f64_simd_other(vector: &[f64]) -> f32 { use crate::simd::f64::{f64x4, f64x8}; use crate::simd::{FloatSimd, SIMD}; @@ -291,5 +480,133 @@ mod tests { fn test_l2_norm_f64(data in prop::collection::vec(arbitrary_f64(), 4..4048)){ do_norm_l2_test(&data)?; } + + /// Cross-backend parity: scalar fallback must match the dispatched + /// SIMD path within numerical tolerance. Exercises `norm_l2_f64_scalar` + /// directly so the runtime fallback is exercised even on AVX2-capable + /// CI hosts. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_norm_f64_scalar_simd_parity( + data in prop::collection::vec(arbitrary_f64(), 4..4048) + ) { + let scalar = norm_l2_f64_scalar(&data); + let simd = norm_l2_f64_simd(&data); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-6)); + } + + /// Parity check for `norm_l2_f32_dispatched` (Branch B exclusive: the + /// auto-vectorised scalar L2-norm path). The dispatched kernel must + /// agree with a portable f64-precision scalar reference within + /// numerical tolerance. The reference is hand-rolled here to keep this + /// test architecture-agnostic (the x86_64-only `norm_l2_f64_scalar` + /// helper is gated above). + #[test] + fn test_l2_norm_f32_scalar_simd_parity( + data in prop::collection::vec(arbitrary_f32(), 4..4048) + ) { + let scalar = data.iter().map(|&v| (v as f64).powi(2)).sum::().sqrt() as f32; + let simd = ::norm_l2(&data); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3)); + } + + /// AVX-512-direct parity: explicitly compares the scalar fallback + /// against the native AVX-512 inner kernel on AVX-512F-capable hosts + /// (Skylake-X+, Ice Lake, Sapphire Rapids, Zen 4). Early-returns on + /// hosts without AVX-512F so the test stays portable; CI runners with + /// AVX-512F exercise the `_mm512_*` path. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_norm_f64_scalar_vs_avx512_parity( + data in prop::collection::vec(arbitrary_f64(), 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + let scalar = norm_l2_f64_scalar(&data); + let avx512 = unsafe { x86::norm_l2_f64_avx512(&data) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-6)); + } + + /// AVX + FMA-direct parity for the f64 L2-norm kernel. Covers the + /// AMD Piledriver / Steamroller / FX-7500 tier. Early-returns on + /// hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_norm_f64_scalar_vs_avx_fma_parity( + data in prop::collection::vec(arbitrary_f64(), 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + let scalar = norm_l2_f64_scalar(&data); + let avx_fma = unsafe { x86::norm_l2_f64_avx_fma(&data) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-6)); + } + + /// AVX-only-direct parity for the f64 L2-norm kernel. Covers the + /// Intel Sandy Bridge / Ivy Bridge tier (AVX without FMA). + /// Early-returns on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_norm_f64_scalar_vs_avx_parity( + data in prop::collection::vec(arbitrary_f64(), 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + let scalar = norm_l2_f64_scalar(&data); + let avx = unsafe { x86::norm_l2_f64_avx(&data) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-6)); + } + + /// AVX-512-direct parity for the f32 L2-norm kernel. Explicitly + /// compares the scalar fallback against the native AVX-512 inner + /// kernel on AVX-512F-capable hosts. Early-returns on hosts without + /// AVX-512F; CI runners with AVX-512F exercise the `_mm512_*` path. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_norm_l2_f32_scalar_vs_avx512_parity( + data in prop::collection::vec(arbitrary_f32(), 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + let scalar = norm_l2_f32_scalar(&data); + let avx512 = unsafe { x86::norm_l2_f32_avx512(&data) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-3)); + } + + /// AVX + FMA-direct parity for the f32 L2-norm kernel. Covers the + /// AMD Piledriver / Steamroller / FX-7500 tier. Early-returns on + /// hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_norm_l2_f32_scalar_vs_avx_fma_parity( + data in prop::collection::vec(arbitrary_f32(), 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + let scalar = norm_l2_f32_scalar(&data); + let avx_fma = unsafe { x86::norm_l2_f32_avx_fma(&data) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-3)); + } + + /// AVX-only-direct parity for the f32 L2-norm kernel. Covers the + /// Intel Sandy Bridge / Ivy Bridge tier (AVX without FMA). + /// Early-returns on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_norm_l2_f32_scalar_vs_avx_parity( + data in prop::collection::vec(arbitrary_f32(), 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + let scalar = norm_l2_f32_scalar(&data); + let avx = unsafe { x86::norm_l2_f32_avx(&data) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-3)); + } } } diff --git a/rust/lance-linalg/src/simd.rs b/rust/lance-linalg/src/simd.rs index 91dc1c6959d..6722eafd768 100644 --- a/rust/lance-linalg/src/simd.rs +++ b/rust/lance-linalg/src/simd.rs @@ -19,6 +19,8 @@ pub mod f32; pub mod f64; pub mod i32; pub mod u8; +#[cfg(target_arch = "x86_64")] +pub(crate) mod x86; use num_traits::{Float, Num}; use u8::u8x16; diff --git a/rust/lance-linalg/src/simd/dist_table.rs b/rust/lance-linalg/src/simd/dist_table.rs index 626c1581b15..e337953ec10 100644 --- a/rust/lance-linalg/src/simd/dist_table.rs +++ b/rust/lance-linalg/src/simd/dist_table.rs @@ -73,6 +73,10 @@ pub fn sum_4bit_dist_table( ) } }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the AVX2 inner uses `_mm256_shuffle_epi8` / `_mm256_and_si256` / + // `_mm256_srli_epi16` / `_mm256_add_epi16` integer ops which + // neither AVX nor AVX+FMA provides. Scalar is the correct route. _ => sum_4bit_dist_table_scalar(code_len, codes, dist_table, dists), } } diff --git a/rust/lance-linalg/src/simd/f32.rs b/rust/lance-linalg/src/simd/f32.rs index 78042997121..4ce7f64706d 100644 --- a/rust/lance-linalg/src/simd/f32.rs +++ b/rust/lance-linalg/src/simd/f32.rs @@ -46,14 +46,42 @@ impl std::fmt::Debug for f32x8 { } impl f32x8 { + /// Gather 8 f32 values from `slice` at the offsets in `indices`. + /// + /// On x86_64 this uses the AVX2 `vgatherdps` instruction when the host + /// supports it (gated at runtime via `is_x86_feature_detected!`). On + /// other architectures (and on x86_64 hosts without AVX2) the function + /// falls back to a per-index scalar load followed by `Self::from(&out)`, + /// which goes through `load_unaligned` (NEON / LASX / `_mm256_loadu_ps` + /// depending on platform). Per-tier macro stamping (e.g., the + /// `multiversion` crate) was considered but doesn't fit here: the function + /// returns `Self` and `_mm256_i32gather_ps::<4>` requires the const-generic + /// stride to be a compile-time literal — neither composes with the macro. + /// + /// # Panics + /// + /// If any index is negative or lands outside `slice`. #[inline] pub fn gather(slice: &[f32], indices: &[i32; 8]) -> Self { - #[cfg(target_arch = "x86_64")] - unsafe { - use super::i32::i32x8; + // Every backend below reads without bounds checking: `vgatherdps` does + // none, and the NEON / LASX arms offset a raw pointer. Check once here + // so an out-of-range index panics on every host rather than reading out + // of bounds on some and panicking on others. + for &i in indices { + assert!( + (i as usize) < slice.len(), + "gather index {i} is out of bounds for a slice of length {}", + slice.len() + ); + } - let idx = i32x8::from(indices); - Self(_mm256_i32gather_ps::<4>(slice.as_ptr(), idx.0)) + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx2") { + unsafe { gather_avx2(slice, indices) } + } else { + gather_scalar_x86(slice, indices) + } } #[cfg(target_arch = "aarch64")] @@ -94,6 +122,30 @@ impl f32x8 { } } +/// AVX2 gather. Caller must ensure the host supports AVX2 (gated by +/// the `is_x86_feature_detected!("avx2")` check in `f32x8::gather`). +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +unsafe fn gather_avx2(slice: &[f32], indices: &[i32; 8]) -> f32x8 { + use super::i32::i32x8; + + let idx = i32x8::from(indices); + f32x8(_mm256_i32gather_ps::<4>(slice.as_ptr(), idx.0)) +} + +/// Portable scalar gather for x86_64 hosts without AVX2. +/// +/// Indexes the slice rather than offsetting a raw pointer: this is the slow +/// path already, so an out-of-range index should panic instead of reading out +/// of bounds. +#[cfg(target_arch = "x86_64")] +#[inline] +fn gather_scalar_x86(slice: &[f32], indices: &[i32; 8]) -> f32x8 { + let values = indices.map(|i| slice[i as usize]); + // SAFETY: `values` is eight contiguous, initialized `f32`. + unsafe { f32x8::load_unaligned(values.as_ptr()) } +} + impl From<&[f32]> for f32x8 { fn from(value: &[f32]) -> Self { unsafe { Self::load_unaligned(value.as_ptr()) } @@ -439,15 +491,18 @@ impl Mul for f32x8 { } } -/// 16 of 32-bit `f32` values. Use 512-bit SIMD if possible. +/// 16 of 32-bit `f32` values. Stored as a pair of 256-bit AVX vectors on +/// x86_64. Originally there was a sibling AVX-512 variant gated on +/// `target_feature = "avx512f"`, but no project CI configuration enables +/// `+avx512f` globally (and one of the avx512 arms still contained `todo!()`), +/// so the variant was dead code. Removed in the runtime-SIMD-dispatch +/// retrofit; per-tier dispatch happens in the kernel functions in +/// `crate::distance::*` via `match *SIMD_SUPPORT` + per-tier +/// `#[target_feature(enable = "...")]` inner functions. #[allow(non_camel_case_types)] -#[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] +#[cfg(target_arch = "x86_64")] #[derive(Clone, Copy)] pub struct f32x16(__m256, __m256); -#[allow(non_camel_case_types)] -#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] -#[derive(Clone, Copy)] -pub struct f32x16(__m512); /// 16 of 32-bit `f32` values. Use 512-bit SIMD if possible. #[allow(non_camel_case_types)] @@ -486,11 +541,7 @@ impl<'a> From<&'a [f32; 16]> for f32x16 { impl SIMD for f32x16 { #[inline] fn splat(val: f32) -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_set1_ps(val)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_set1_ps(val), _mm256_set1_ps(val)) } @@ -514,11 +565,7 @@ impl SIMD for f32x16 { #[inline] fn zeros() -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_setzero_ps()) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_setzero_ps(), _mm256_setzero_ps()) } @@ -534,14 +581,10 @@ impl SIMD for f32x16 { #[inline] unsafe fn load(ptr: *const f32) -> Self { - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_load_ps(ptr), _mm256_load_ps(ptr.add(8))) } - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_load_ps(ptr)) - } #[cfg(target_arch = "aarch64")] { Self::load_unaligned(ptr) @@ -557,14 +600,10 @@ impl SIMD for f32x16 { #[inline] unsafe fn load_unaligned(ptr: *const f32) -> Self { - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_loadu_ps(ptr), _mm256_loadu_ps(ptr.add(8))) } - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_loadu_ps(ptr)) - } #[cfg(target_arch = "aarch64")] { Self(vld1q_f32_x4(ptr)) @@ -580,11 +619,7 @@ impl SIMD for f32x16 { #[inline] unsafe fn store(&self, ptr: *mut f32) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_store_ps(ptr, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { _mm256_store_ps(ptr, self.0); _mm256_store_ps(ptr.add(8), self.1); @@ -602,11 +637,7 @@ impl SIMD for f32x16 { #[inline] unsafe fn store_unaligned(&self, ptr: *mut f32) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_storeu_ps(ptr, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { _mm256_storeu_ps(ptr, self.0); _mm256_storeu_ps(ptr.add(8), self.1); @@ -622,12 +653,9 @@ impl SIMD for f32x16 { } } + #[inline] fn reduce_sum(&self) -> f32 { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_mask_reduce_add_ps(0xFFFF, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { let mut sum = _mm256_add_ps(self.0, self.1); // Shift and add vector, until only 1 value left. @@ -657,11 +685,7 @@ impl SIMD for f32x16 { #[inline] fn reduce_min(&self) -> f32 { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_mask_reduce_min_ps(0xFFFF, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { let mut m1 = _mm256_min_ps(self.0, self.1); let mut m2 = _mm256_permute2f128_ps(m1, m1, 1); @@ -695,11 +719,7 @@ impl SIMD for f32x16 { #[inline] fn min(&self, rhs: &Self) -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_min_ps(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_min_ps(self.0, rhs.0), _mm256_min_ps(self.1, rhs.1)) } @@ -718,21 +738,15 @@ impl SIMD for f32x16 { } } + #[inline] fn find(&self, val: f32) -> Option { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - // let tgt = _mm512_set1_ps(val); - // let mask = _mm512_cmpeq_ps_mask(self.0, tgt); - // if mask != 0 { - // return Some(mask.trailing_zeros() as i32); - // } - todo!() - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { - // _mm256_cmpeq_ps_mask requires "avx512l". + // _mm256_cmpeq_ps_mask requires AVX-512 (avx512f); use a scalar scan here + // since we only require AVX2. + let arr = self.as_array(); for i in 0..16 { - if self.as_array().get_unchecked(i) == &val { + if arr.get_unchecked(i) == &val { return Some(i as i32); } } @@ -774,11 +788,7 @@ impl SIMD for f32x16 { impl FloatSimd for f32x16 { #[inline] fn multiply_add(&mut self, a: Self, b: Self) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - self.0 = _mm512_fmadd_ps(a.0, b.0, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { self.0 = _mm256_fmadd_ps(a.0, b.0, self.0); self.1 = _mm256_fmadd_ps(a.1, b.1, self.1); @@ -803,11 +813,7 @@ impl Add for f32x16 { #[inline] fn add(self, rhs: Self) -> Self::Output { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_add_ps(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_add_ps(self.0, rhs.0), _mm256_add_ps(self.1, rhs.1)) } @@ -830,11 +836,7 @@ impl Add for f32x16 { impl AddAssign for f32x16 { #[inline] fn add_assign(&mut self, rhs: Self) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - self.0 = _mm512_add_ps(self.0, rhs.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { self.0 = _mm256_add_ps(self.0, rhs.0); self.1 = _mm256_add_ps(self.1, rhs.1); @@ -859,11 +861,7 @@ impl Mul for f32x16 { #[inline] fn mul(self, rhs: Self) -> Self::Output { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_mul_ps(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_mul_ps(self.0, rhs.0), _mm256_mul_ps(self.1, rhs.1)) } @@ -888,11 +886,7 @@ impl Sub for f32x16 { #[inline] fn sub(self, rhs: Self) -> Self::Output { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_sub_ps(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_sub_ps(self.0, rhs.0), _mm256_sub_ps(self.1, rhs.1)) } @@ -915,11 +909,7 @@ impl Sub for f32x16 { impl SubAssign for f32x16 { #[inline] fn sub_assign(&mut self, rhs: Self) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - self.0 = _mm512_sub_ps(self.0, rhs.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { self.0 = _mm256_sub_ps(self.0, rhs.0); self.1 = _mm256_sub_ps(self.1, rhs.1); @@ -943,9 +933,17 @@ impl SubAssign for f32x16 { mod tests { use super::*; + use rstest::rstest; #[test] fn test_basic_ops() { + // Load / store / arithmetic on `f32x8` lower to AVX intrinsics, and + // `multiply_add` lowers to `_mm256_fmadd_ps`, which needs FMA. Both + // are present from the AvxFma tier up. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { + return; + } let a = (0..8).map(|f| f as f32).collect::>(); let b = (10..18).map(|f| f as f32).collect::>(); @@ -983,6 +981,11 @@ mod tests { #[test] fn test_f32x8_cmp_ops() { + // `min` / `reduce_min` are AVX intrinsics; `find` is a scalar scan. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx") { + return; + } let a = [1.0_f32, 2.0, 5.0, 6.0, 7.0, 3.0, 2.0, 1.0]; let b = [2.0_f32, 1.0, 4.0, 5.0, 9.0, 5.0, 6.0, 2.0]; let c = [2.0_f32, 1.0, 4.0, 5.0, 7.0, 3.0, 2.0, 1.0]; @@ -1007,6 +1010,11 @@ mod tests { #[test] fn test_basic_f32x16_ops() { + // `f32x16` is a pair of `__m256`; `multiply_add` needs FMA. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { + return; + } let a = (0..16).map(|f| f as f32).collect::>(); let b = (10..26).map(|f| f as f32).collect::>(); @@ -1041,6 +1049,11 @@ mod tests { #[test] fn test_f32x16_cmp_ops() { + // `min` / `reduce_min` are AVX intrinsics; `find` is a scalar scan. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx") { + return; + } let a = [ 1.0_f32, 2.0, 5.0, 6.0, 7.0, 3.0, 2.0, 1.0, -0.5, 5.0, 6.0, 7.0, 8.0, 9.0, 1.0, 2.0, ]; @@ -1074,9 +1087,59 @@ mod tests { #[test] fn test_f32x8_gather() { + // `f32x8::gather` does its own runtime AVX2 detection and falls back + // to a scalar gather, so this test only needs whatever reading the + // `__m256`-backed result costs: AVX, for `reduce_sum`. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx") { + return; + } let a = (0..256).map(|f| f as f32).collect::>(); let idx = [0_i32, 4, 8, 12, 16, 20, 24, 29]; let v = f32x8::gather(&a, &idx); assert_eq!(v.reduce_sum(), 113.0); } + + /// Directly exercises `gather_scalar_x86`, the per-index scalar fallback + /// `f32x8::gather` takes on x86_64 hosts without AVX2. Runtime AVX2 hosts + /// route through `gather_avx2` instead, so the fallback is otherwise never + /// hit under coverage. Reading the `__m256`-backed result needs AVX, so + /// skip on hosts without it. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_gather_scalar_x86() { + if !std::is_x86_feature_detected!("avx") { + return; + } + let a = (0..256).map(|f| f as f32).collect::>(); + let idx = [0_i32, 4, 8, 12, 16, 20, 24, 29]; + let v = gather_scalar_x86(&a, &idx); + let expected = idx.map(|i| a[i as usize]); + assert_eq!(v.as_array(), expected); + } + + /// An index past the end of the slice panics rather than reading out of + /// bounds. The bounds check fires before any AVX instruction, so this + /// case runs on every x86_64 host. + #[cfg(target_arch = "x86_64")] + #[test] + #[should_panic(expected = "index out of bounds")] + fn test_gather_scalar_x86_rejects_out_of_range_index() { + let a = (0..8).map(|f| f as f32).collect::>(); + let idx = [0_i32, 1, 2, 3, 4, 5, 6, 99]; + let _ = gather_scalar_x86(&a, &idx); + } + + /// `gather` validates before dispatching, so every backend — `vgatherdps`, + /// the x86 scalar fallback, and the NEON / LASX raw-pointer arms — rejects + /// a bad index identically instead of reading out of bounds. + #[rstest] + #[case::past_end(99)] + #[case::negative(-1)] + #[should_panic(expected = "out of bounds")] + fn test_gather_rejects_invalid_index(#[case] bad_index: i32) { + let a = (0..8).map(|f| f as f32).collect::>(); + let idx = [0_i32, 1, 2, 3, 4, 5, 6, bad_index]; + let _ = f32x8::gather(&a, &idx); + } } diff --git a/rust/lance-linalg/src/simd/f64.rs b/rust/lance-linalg/src/simd/f64.rs index 32c0d389e5b..1276f54da56 100644 --- a/rust/lance-linalg/src/simd/f64.rs +++ b/rust/lance-linalg/src/simd/f64.rs @@ -355,14 +355,15 @@ impl Mul for f64x4 { // f64x8: 8 × f64 values (512-bit SIMD or 2 × 256-bit) // --------------------------------------------------------------------------- -/// 8 of 64-bit `f64` values. Uses 512-bit SIMD if possible. +/// 8 of 64-bit `f64` values. Stored as a pair of 256-bit AVX vectors on +/// x86_64. Originally there was a sibling AVX-512 variant gated on +/// `target_feature = "avx512f"`, but no project CI configuration enables +/// `+avx512f` globally, so the variant was dead code. Removed in the +/// runtime-SIMD-dispatch retrofit; per-tier dispatch happens in the kernel +/// functions in `crate::distance::*` via `match *SIMD_SUPPORT` + per-tier +/// `#[target_feature(enable = "...")]` inner functions. #[allow(non_camel_case_types)] -#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] -#[derive(Clone, Copy)] -pub struct f64x8(__m512d); - -#[allow(non_camel_case_types)] -#[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] +#[cfg(target_arch = "x86_64")] #[derive(Clone, Copy)] pub struct f64x8(__m256d, __m256d); @@ -401,11 +402,7 @@ impl<'a> From<&'a [f64; 8]> for f64x8 { impl SIMD for f64x8 { #[inline] fn splat(val: f64) -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_set1_pd(val)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_set1_pd(val), _mm256_set1_pd(val)) } @@ -423,11 +420,7 @@ impl SIMD for f64x8 { #[inline] fn zeros() -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_setzero_pd()) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_setzero_pd(), _mm256_setzero_pd()) } @@ -443,11 +436,7 @@ impl SIMD for f64x8 { #[inline] unsafe fn load(ptr: *const f64) -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_load_pd(ptr)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_load_pd(ptr), _mm256_load_pd(ptr.add(4))) } @@ -466,11 +455,7 @@ impl SIMD for f64x8 { #[inline] unsafe fn load_unaligned(ptr: *const f64) -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_loadu_pd(ptr)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_loadu_pd(ptr), _mm256_loadu_pd(ptr.add(4))) } @@ -489,11 +474,7 @@ impl SIMD for f64x8 { #[inline] unsafe fn store(&self, ptr: *mut f64) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_store_pd(ptr, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { _mm256_store_pd(ptr, self.0); _mm256_store_pd(ptr.add(4), self.1); @@ -512,11 +493,7 @@ impl SIMD for f64x8 { #[inline] unsafe fn store_unaligned(&self, ptr: *mut f64) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_storeu_pd(ptr, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { _mm256_storeu_pd(ptr, self.0); _mm256_storeu_pd(ptr.add(4), self.1); @@ -533,12 +510,9 @@ impl SIMD for f64x8 { } } + #[inline] fn reduce_sum(&self) -> f64 { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_mask_reduce_add_pd(0xFF, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { let sum = _mm256_add_pd(self.0, self.1); let hi = _mm256_permute2f128_pd(sum, sum, 1); @@ -561,11 +535,7 @@ impl SIMD for f64x8 { #[inline] fn reduce_min(&self) -> f64 { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_mask_reduce_min_pd(0xFF, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { let m = _mm256_min_pd(self.0, self.1); let hi = _mm256_permute2f128_pd(m, m, 1); @@ -592,11 +562,7 @@ impl SIMD for f64x8 { #[inline] fn min(&self, rhs: &Self) -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_min_pd(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_min_pd(self.0, rhs.0), _mm256_min_pd(self.1, rhs.1)) } @@ -613,6 +579,7 @@ impl SIMD for f64x8 { } } + #[inline] fn find(&self, val: f64) -> Option { unsafe { for i in 0..8 { @@ -628,11 +595,7 @@ impl SIMD for f64x8 { impl FloatSimd for f64x8 { #[inline] fn multiply_add(&mut self, a: Self, b: Self) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - self.0 = _mm512_fmadd_pd(a.0, b.0, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { self.0 = _mm256_fmadd_pd(a.0, b.0, self.0); self.1 = _mm256_fmadd_pd(a.1, b.1, self.1); @@ -657,11 +620,7 @@ impl Add for f64x8 { #[inline] fn add(self, rhs: Self) -> Self::Output { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_add_pd(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_add_pd(self.0, rhs.0), _mm256_add_pd(self.1, rhs.1)) } @@ -682,11 +641,7 @@ impl Add for f64x8 { impl AddAssign for f64x8 { #[inline] fn add_assign(&mut self, rhs: Self) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - self.0 = _mm512_add_pd(self.0, rhs.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { self.0 = _mm256_add_pd(self.0, rhs.0); self.1 = _mm256_add_pd(self.1, rhs.1); @@ -711,11 +666,7 @@ impl Mul for f64x8 { #[inline] fn mul(self, rhs: Self) -> Self::Output { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_mul_pd(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_mul_pd(self.0, rhs.0), _mm256_mul_pd(self.1, rhs.1)) } @@ -738,11 +689,7 @@ impl Sub for f64x8 { #[inline] fn sub(self, rhs: Self) -> Self::Output { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_sub_pd(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_sub_pd(self.0, rhs.0), _mm256_sub_pd(self.1, rhs.1)) } @@ -763,11 +710,7 @@ impl Sub for f64x8 { impl SubAssign for f64x8 { #[inline] fn sub_assign(&mut self, rhs: Self) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - self.0 = _mm512_sub_pd(self.0, rhs.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { self.0 = _mm256_sub_pd(self.0, rhs.0); self.1 = _mm256_sub_pd(self.1, rhs.1); @@ -793,6 +736,12 @@ mod tests { #[test] fn test_f64x4_basic_ops() { + // The `f64x4` constructor / load / store / arithmetic paths all lower + // to AVX intrinsics on x86_64; none of them need AVX2. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx") { + return; + } let a = [1.0_f64, 2.0, 3.0, 4.0]; let b = [5.0_f64, 6.0, 7.0, 8.0]; @@ -814,6 +763,11 @@ mod tests { #[test] fn test_f64x4_fma() { + // `multiply_add` lowers to `_mm256_fmadd_pd`, which needs FMA. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { + return; + } let a = [1.0_f64, 2.0, 3.0, 4.0]; let b = [2.0_f64, 3.0, 4.0, 5.0]; @@ -826,6 +780,11 @@ mod tests { #[test] fn test_f64x4_min() { + // `min` / `reduce_min` are AVX intrinsics. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx") { + return; + } let a = [1.0_f64, 5.0, 2.0, 8.0]; let b = [3.0_f64, 2.0, 4.0, 1.0]; let simd_a: f64x4 = (&a).into(); @@ -838,6 +797,11 @@ mod tests { #[test] fn test_f64x8_basic_ops() { + // `f64x8` is a pair of `__m256d`; add / reduce are AVX intrinsics. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx") { + return; + } let a: [f64; 8] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; let b: [f64; 8] = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]; @@ -856,6 +820,11 @@ mod tests { #[test] fn test_f64x8_fma() { + // `multiply_add` lowers to `_mm256_fmadd_pd`, which needs FMA. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { + return; + } let a: [f64; 8] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; let b: [f64; 8] = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]; @@ -869,6 +838,11 @@ mod tests { #[test] fn test_f64x8_min() { + // `min` / `reduce_min` are AVX intrinsics. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx") { + return; + } let a: [f64; 8] = [5.0, 1.0, 8.0, 3.0, 9.0, 2.0, 7.0, 4.0]; let b: [f64; 8] = [2.0, 6.0, 3.0, 7.0, 1.0, 8.0, 4.0, 9.0]; let simd_a: f64x8 = (&a).into(); diff --git a/rust/lance-linalg/src/simd/x86.rs b/rust/lance-linalg/src/simd/x86.rs new file mode 100644 index 00000000000..437018669e2 --- /dev/null +++ b/rust/lance-linalg/src/simd/x86.rs @@ -0,0 +1,89 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Reduction helpers shared by the x86_64 AVX kernels in [`crate::distance`]. +//! +//! Each distance kernel accumulates into a 256-bit register and folds it down +//! to a scalar once, after the main loop. The fold is identical across +//! `cosine`, `dot`, `l2` and `norm_l2`, so it lives here instead of being +//! copied into each kernel's private `mod x86`. +//! +//! The module itself is `pub(crate)`, which is what keeps these helpers off the +//! public API; the items are `pub` rather than `pub(crate)` only because +//! `clippy::redundant_pub_crate` fires on the narrower visibility. + +use std::arch::x86_64::*; + +/// Horizontal sum of the eight `f32` lanes of an `__m256`. +/// +/// Folds the upper 128-bit lane into the lower one, then reduces the +/// remaining four lanes pairwise. Uses `movehl`/`shuffle` plus scalar adds +/// rather than two `vhaddps`, which is one fewer uop on most cores. +/// +/// # Safety +/// +/// The host must support AVX. Callers are `#[target_feature]`-annotated +/// kernels that the runtime dispatcher only selects after checking. +#[inline] +#[target_feature(enable = "avx")] +pub unsafe fn hsum256_ps(v: __m256) -> f32 { + let lo = _mm256_castps256_ps128(v); + let hi = _mm256_extractf128_ps(v, 1); + let sum128 = _mm_add_ps(lo, hi); + let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); + // 0x55 broadcasts lane 1 into lane 0, so the scalar add below lands the + // last of the four partial sums. + let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0x55)); + _mm_cvtss_f32(sum32) +} + +/// Horizontal sum of the four `f64` lanes of an `__m256d`. +/// +/// Folds the upper 128-bit lane into the lower one, then adds the remaining +/// pair. +/// +/// # Safety +/// +/// The host must support AVX. Callers are `#[target_feature]`-annotated +/// kernels that the runtime dispatcher only selects after checking. +#[inline] +#[target_feature(enable = "avx")] +pub unsafe fn hsum256_pd(v: __m256d) -> f64 { + let lo = _mm256_castpd256_pd128(v); + let hi = _mm256_extractf128_pd(v, 1); + let sum128 = _mm_add_pd(lo, hi); + let sum64 = _mm_add_pd(sum128, _mm_unpackhi_pd(sum128, sum128)); + _mm_cvtsd_f64(sum64) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::ascending([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], 36.0)] + #[case::negative_lanes([-1.5, 2.0, -3.0, 4.5, 0.0, -0.5, 1.0, 2.5], 5.0)] + #[case::zeros([0.0; 8], 0.0)] + #[case::cancelling([1.0, -1.0, 2.0, -2.0, 3.0, -3.0, 4.0, -4.0], 0.0)] + fn hsum256_ps_sums_every_lane(#[case] lanes: [f32; 8], #[case] expected: f32) { + if !std::is_x86_feature_detected!("avx") { + return; + } + let sum = unsafe { hsum256_ps(_mm256_loadu_ps(lanes.as_ptr())) }; + assert_eq!(sum, expected); + } + + #[rstest] + #[case::ascending([1.0, 2.0, 3.0, 4.0], 10.0)] + #[case::negative_lanes([-1.5, 2.0, -3.0, 4.5], 2.0)] + #[case::zeros([0.0; 4], 0.0)] + #[case::cancelling([1.0, -1.0, 2.0, -2.0], 0.0)] + fn hsum256_pd_sums_every_lane(#[case] lanes: [f64; 4], #[case] expected: f64) { + if !std::is_x86_feature_detected!("avx") { + return; + } + let sum = unsafe { hsum256_pd(_mm256_loadu_pd(lanes.as_ptr())) }; + assert_eq!(sum, expected); + } +} diff --git a/rust/lance-namespace-impls/Cargo.toml b/rust/lance-namespace-impls/Cargo.toml index 27b9a4bc0e2..c5d8e120740 100644 --- a/rust/lance-namespace-impls/Cargo.toml +++ b/rust/lance-namespace-impls/Cargo.toml @@ -91,7 +91,7 @@ rustls-pki-types = { version = "1", optional = true } # Azure credential vending dependencies (optional, enabled by "credential-vendor-azure" feature) chrono = { workspace = true, optional = true } hmac = { version = "0.12", optional = true } -quick-xml = { version = "0.38", optional = true } +quick-xml = { version = "0.40", optional = true } [dev-dependencies] opendal = { workspace = true, features = ["services-goosefs"] } diff --git a/rust/lance-namespace-impls/src/credentials.rs b/rust/lance-namespace-impls/src/credentials.rs index e841ac620f6..23f9346cfb3 100644 --- a/rust/lance-namespace-impls/src/credentials.rs +++ b/rust/lance-namespace-impls/src/credentials.rs @@ -223,6 +223,14 @@ pub mod aws_props { /// AWS credential duration in milliseconds. /// Default: 3600000 (1 hour). Range: 900000 (15 min) to 43200000 (12 hours). pub const DURATION_MILLIS: &str = "aws_duration_millis"; + + /// When "true", the scoped assume is performed via `AssumeRoleWithWebIdentity` + /// using the pod's projected service-account OIDC token + pub const ASSUME_VIA_POD_WEB_IDENTITY: &str = "aws_assume_via_pod_web_identity"; + + /// Explicit path to the pod's projected SA OIDC token file. Overrides + /// `AWS_WEB_IDENTITY_TOKEN_FILE` when set. + pub const POD_WEB_IDENTITY_TOKEN_FILE: &str = "aws_pod_web_identity_token_file"; } /// GCP-specific property keys (short form, without prefix) @@ -503,6 +511,43 @@ async fn create_aws_vendor( config = config.with_role_session_name(session_name); } + // Direct (non-chained) web-identity assume for the pod, when enabled. An + // explicit token-file path wins; otherwise, if opted in, resolve the + // EKS-injected `AWS_WEB_IDENTITY_TOKEN_FILE`. Falling back to the chained + // AssumeRole path when neither is present keeps existing behavior. + let assume_via_pod = properties + .get(aws_props::ASSUME_VIA_POD_WEB_IDENTITY) + .map(|v| v.eq_ignore_ascii_case("true")) + .unwrap_or(false); + let pod_token_file = properties + .get(aws_props::POD_WEB_IDENTITY_TOKEN_FILE) + .cloned() + .or_else(|| { + assume_via_pod + .then(|| std::env::var("AWS_WEB_IDENTITY_TOKEN_FILE").ok()) + .flatten() + }); + // Log the resolved assume path once at vendor init so a deployment can + // confirm at runtime which branch `assume_scoped` will take. + match &pod_token_file { + Some(path) => log::info!( + "AWS credential vendor (role {role_arn}): direct AssumeRoleWithWebIdentity \ + via pod token file '{path}'" + ), + None if assume_via_pod => log::warn!( + "AWS credential vendor (role {role_arn}): aws_assume_via_pod_web_identity=true \ + but no token file resolved (aws_pod_web_identity_token_file unset and \ + AWS_WEB_IDENTITY_TOKEN_FILE not in env); falling back to chained AssumeRole" + ), + None => log::info!( + "AWS credential vendor (role {role_arn}): chained AssumeRole \ + (pod web-identity not enabled)" + ), + } + if let Some(path) = pod_token_file { + config = config.with_pod_web_identity_token_file(path); + } + let vendor = AwsCredentialVendor::new(config).await?; Ok(Some(Box::new(vendor))) } diff --git a/rust/lance-namespace-impls/src/credentials/aws.rs b/rust/lance-namespace-impls/src/credentials/aws.rs index 7dedaa6e108..56dc9a54c2c 100644 --- a/rust/lance-namespace-impls/src/credentials/aws.rs +++ b/rust/lance-namespace-impls/src/credentials/aws.rs @@ -24,6 +24,18 @@ use super::{ redact_credential, }; +/// Render an error together with its full `source()` chain. +fn full_error_chain(err: &dyn std::error::Error) -> String { + let mut out = err.to_string(); + let mut source = err.source(); + while let Some(cause) = source { + out.push_str(": "); + out.push_str(&cause.to_string()); + source = cause.source(); + } + out +} + /// Configuration for AWS credential vending. #[derive(Debug, Clone)] pub struct AwsCredentialVendorConfig { @@ -60,6 +72,17 @@ pub struct AwsCredentialVendorConfig { /// When an API key is provided, its hash is looked up in this map. /// If found, the mapped permission is used instead of the default permission. pub api_key_hash_permissions: HashMap, + + /// Optional path to the pod's projected service-account OIDC token file (typically the + /// EKS-injected `AWS_WEB_IDENTITY_TOKEN_FILE`). + /// This is the recommended method when running in Kubernetes. + /// When set, the scoped assume is performed with `AssumeRoleWithWebIdentity` + /// using this token instead of a role-chained `AssumeRole` from the process's + /// ambient credentials. A web-identity assumption is *not* role chaining, so + /// STS honors the role's `MaxSessionDuration` (up to 12h) and the returned + /// expiration is accurate. When `None`, the `AssumeRole` path + /// is used and `duration_millis` is subject to the source role duration and may be invalid + pub pod_web_identity_token_file: Option, } impl AwsCredentialVendorConfig { @@ -74,6 +97,7 @@ impl AwsCredentialVendorConfig { permission: VendedPermission::default(), api_key_salt: None, api_key_hash_permissions: HashMap::new(), + pod_web_identity_token_file: None, } } @@ -101,6 +125,14 @@ impl AwsCredentialVendorConfig { self } + /// Set the pod web-identity token file, enabling direct (non-chained) + /// `AssumeRoleWithWebIdentity` for the scoped assume. See + /// [`AwsCredentialVendorConfig::pod_web_identity_token_file`]. + pub fn with_pod_web_identity_token_file(mut self, path: impl Into) -> Self { + self.pod_web_identity_token_file = Some(path.into()); + self + } + /// Set the permission level for vended credentials. pub fn with_permission(mut self, permission: VendedPermission) -> Self { self.permission = permission; @@ -407,7 +439,8 @@ impl AwsCredentialVendor { lance_core::Error::from(NamespaceError::Internal { message: format!( "AssumeRoleWithWebIdentity failed for role '{}': {}", - self.config.role_arn, e + self.config.role_arn, + full_error_chain(&e) ), }) })?; @@ -421,6 +454,95 @@ impl AwsCredentialVendor { } /// Vend credentials using AssumeRole with API key validation. + /// Perform the scoped assume for `(bucket, prefix, permission)`, attaching the + /// per-table session policy. + /// + /// When [`AwsCredentialVendorConfig::pod_web_identity_token_file`] is set this + /// uses `AssumeRoleWithWebIdentity` with the pod's projected SA OIDC token -- a + /// *direct*, non-chained assumption that honors the role's `MaxSessionDuration` + /// (so `expires_at_millis` is accurate). Otherwise it falls back to the legacy + /// role-chained `AssumeRole` from ambient credentials (STS-capped at 1h). + /// + /// `external_id` is only applied on the chained `AssumeRole` path; + /// `AssumeRoleWithWebIdentity` has no external-id parameter (the OIDC + /// `sub`/`aud` trust condition is the binding instead). + async fn assume_scoped( + &self, + bucket: &str, + prefix: &str, + permission: VendedPermission, + session_name: &str, + external_id: Option<&str>, + ) -> Result { + let policy = Self::build_policy(bucket, prefix, permission); + let duration_secs = self.config.duration_millis.div_ceil(1000).clamp(900, 43200) as i32; + + let credentials = if let Some(token_file) = &self.config.pod_web_identity_token_file { + // DIRECT (non-chained): AssumeRoleWithWebIdentity with the pod's SA + // OIDC token. Re-read the file every vend -- kubelet rotates it. + let token = tokio::fs::read_to_string(token_file).await.map_err(|e| { + lance_core::Error::from(NamespaceError::Internal { + message: format!( + "failed to read pod web identity token '{}': {}", + token_file, e + ), + }) + })?; + debug!( + "AWS AssumeRoleWithWebIdentity (pod): role={}, session={}, permission={}", + self.config.role_arn, session_name, permission + ); + let response = self + .sts_client + .assume_role_with_web_identity() + .role_arn(&self.config.role_arn) + .web_identity_token(token.trim()) + .role_session_name(session_name) + .policy(&policy) + .duration_seconds(duration_secs) + .send() + .await + .map_err(|e| { + lance_core::Error::from(NamespaceError::Internal { + message: format!( + "AssumeRoleWithWebIdentity (pod) failed for role '{}': {}", + self.config.role_arn, + full_error_chain(&e) + ), + }) + })?; + response.credentials().cloned() + } else { + // LEGACY chained path: AssumeRole from ambient credentials (1h cap). + debug!( + "AWS AssumeRole (chained): role={}, session={}, permission={}", + self.config.role_arn, session_name, permission + ); + let mut request = self + .sts_client + .assume_role() + .role_arn(&self.config.role_arn) + .role_session_name(session_name) + .policy(&policy) + .duration_seconds(duration_secs); + if let Some(external_id) = external_id { + request = request.external_id(external_id); + } + let response = request.send().await.map_err(|e| { + lance_core::Error::from(NamespaceError::Internal { + message: format!( + "AssumeRole failed for role '{}': {}", + self.config.role_arn, + full_error_chain(&e) + ), + }) + })?; + response.credentials().cloned() + }; + + self.extract_credentials(credentials.as_ref(), bucket, prefix, permission) + } + async fn vend_with_api_key( &self, bucket: &str, @@ -452,42 +574,19 @@ impl AwsCredentialVendor { }) })?; - let policy = Self::build_policy(bucket, prefix, permission); + // The api_key authorizes the client and picks the permission; the AWS + // assume itself goes through `assume_scoped` (pod web-identity when + // configured, else chained AssumeRole with the key hash as external_id). let session_name = Self::cap_session_name(&format!("lance-api-{}", &key_hash[..16])); - let duration_secs = self.config.duration_millis.div_ceil(1000).clamp(900, 43200) as i32; - - debug!( - "AWS AssumeRole with API key: role={}, session={}, permission={}", - self.config.role_arn, session_name, permission - ); - - let request = self - .sts_client - .assume_role() - .role_arn(&self.config.role_arn) - .role_session_name(&session_name) - .policy(&policy) - .duration_seconds(duration_secs) - .external_id(&key_hash); // Use hash as external_id - - let response = request.send().await.map_err(|e| { - lance_core::Error::from(NamespaceError::Internal { - message: format!( - "AssumeRole with API key failed for role '{}': {}", - self.config.role_arn, e - ), - }) - })?; - - self.extract_credentials(response.credentials(), bucket, prefix, permission) + self.assume_scoped(bucket, prefix, permission, &session_name, Some(&key_hash)) + .await } - /// Vend credentials using AssumeRole with static configuration. + /// Vend credentials using the vendor's static (default) permission. async fn vend_with_static_config( &self, bucket: &str, prefix: &str, - policy: &str, ) -> Result { let role_session_name = self .config @@ -496,40 +595,14 @@ impl AwsCredentialVendor { .unwrap_or_else(|| "lance-credential-vending".to_string()); let role_session_name = Self::cap_session_name(&role_session_name); - let duration_secs = self.config.duration_millis.div_ceil(1000).clamp(900, 43200) as i32; - - debug!( - "AWS AssumeRole (static): role={}, session={}, permission={}", - self.config.role_arn, role_session_name, self.config.permission - ); - - let mut request = self - .sts_client - .assume_role() - .role_arn(&self.config.role_arn) - .role_session_name(&role_session_name) - .policy(policy) - .duration_seconds(duration_secs); - - if let Some(ref external_id) = self.config.external_id { - request = request.external_id(external_id); - } - - let response = request.send().await.map_err(|e| { - lance_core::Error::from(NamespaceError::Internal { - message: format!( - "AssumeRole failed for role '{}': {}", - self.config.role_arn, e - ), - }) - })?; - - self.extract_credentials( - response.credentials(), + self.assume_scoped( bucket, prefix, self.config.permission, + &role_session_name, + self.config.external_id.as_deref(), ) + .await } } @@ -575,10 +648,8 @@ impl CredentialVendor for AwsCredentialVendor { .into()) } None => { - // Use AssumeRole with static configuration - let policy = Self::build_policy(&bucket, &prefix, self.config.permission); - self.vend_with_static_config(&bucket, &prefix, &policy) - .await + // Use the vendor's static (default) permission + self.vend_with_static_config(&bucket, &prefix).await } } } @@ -743,6 +814,41 @@ mod tests { assert_eq!(config.duration_millis, 7200000); assert_eq!(config.role_session_name, Some("my-session".to_string())); assert_eq!(config.region, Some("us-west-2".to_string())); + // Defaults to the legacy chained AssumeRole path. + assert_eq!(config.pod_web_identity_token_file, None); + + let pod_config = AwsCredentialVendorConfig::new("arn:aws:iam::123456789012:role/MyRole") + .with_pod_web_identity_token_file("/var/run/secrets/.../token"); + assert_eq!( + pod_config.pod_web_identity_token_file, + Some("/var/run/secrets/.../token".to_string()) + ); + } + + #[tokio::test] + async fn test_pod_web_identity_path_reads_token_file() { + // When pod_web_identity_token_file is set, the scoped assume takes the + // AssumeRoleWithWebIdentity branch and reads the token file first. Point + // it at a missing file so we deterministically hit the read error without + // needing a live STS -- this proves the branch selection + file read. + let sdk_config = aws_config::SdkConfig::builder() + .behavior_version(aws_config::BehaviorVersion::latest()) + .region(aws_config::Region::new("us-east-2")) + .build(); + let sts_client = StsClient::new(&sdk_config); + let config = AwsCredentialVendorConfig::new("arn:aws:iam::123456789012:role/MyRole") + .with_pod_web_identity_token_file("/nonexistent/pod/web-identity-token"); + let vendor = AwsCredentialVendor::with_sts_client(config, sts_client); + + let err = vendor + .vend_credentials("s3://bucket/prefix", None) + .await + .expect_err("missing token file must fail before any STS call"); + assert!( + err.to_string() + .contains("failed to read pod web identity token"), + "unexpected error: {err}" + ); } // ============================================================================ diff --git a/rust/lance-namespace-impls/src/dir.rs b/rust/lance-namespace-impls/src/dir.rs index 5305612b13a..814da6435c5 100644 --- a/rust/lance-namespace-impls/src/dir.rs +++ b/rust/lance-namespace-impls/src/dir.rs @@ -21,8 +21,8 @@ use lance::dataset::scanner::Scanner; use lance::dataset::statistics::DatasetStatisticsExt; use lance::dataset::transaction::{Operation, Transaction}; use lance::dataset::{ - Dataset, MergeInsertBuilder, WhenMatched, WhenNotMatched, WhenNotMatchedBySource, WriteMode, - WriteParams, + Dataset, MergeInsertBuilder, UpdateBuilder, WhenMatched, WhenNotMatched, + WhenNotMatchedBySource, WriteMode, WriteParams, }; use lance::index::{DatasetIndexExt, IndexParams, vector::VectorIndexParams}; use lance::session::Session; @@ -43,25 +43,28 @@ use object_store::{Error as ObjectStoreError, ObjectStore as OSObjectStore, PutM use std::collections::HashMap; use std::io::Cursor; use std::sync::{Arc, Mutex}; +use tokio::sync::OnceCell; use crate::context::DynamicContextProvider; use lance_namespace::models::{ AlterTableAddColumnsRequest, AlterTableAddColumnsResponse, AlterTableAlterColumnsRequest, AlterTableAlterColumnsResponse, AlterTableDropColumnsRequest, AlterTableDropColumnsResponse, - AnalyzeTableQueryPlanRequest, BatchDeleteTableVersionsRequest, - BatchDeleteTableVersionsResponse, BranchContents as ModelBranchContents, CountTableRowsRequest, - CreateNamespaceRequest, CreateNamespaceResponse, CreateTableBranchRequest, - CreateTableBranchResponse, CreateTableIndexRequest, CreateTableIndexResponse, - CreateTableRequest, CreateTableResponse, CreateTableScalarIndexResponse, CreateTableTagRequest, - CreateTableTagResponse, CreateTableVersionRequest, CreateTableVersionResponse, - DeclareTableRequest, DeclareTableResponse, DeleteTableBranchRequest, DeleteTableBranchResponse, - DeleteTableTagRequest, DeleteTableTagResponse, DescribeNamespaceRequest, - DescribeNamespaceResponse, DescribeTableIndexStatsRequest, DescribeTableIndexStatsResponse, - DescribeTableRequest, DescribeTableResponse, DescribeTableVersionRequest, - DescribeTableVersionResponse, DescribeTransactionRequest, DescribeTransactionResponse, - DropNamespaceRequest, DropNamespaceResponse, DropTableIndexRequest, DropTableIndexResponse, - DropTableRequest, DropTableResponse, ExplainTableQueryPlanRequest, FragmentStats, - FragmentSummary, GetTableStatsRequest, GetTableStatsResponse, GetTableTagVersionRequest, + AlterTransactionRequest, AlterTransactionResponse, AnalyzeTableQueryPlanRequest, + BatchDeleteTableVersionsRequest, BatchDeleteTableVersionsResponse, + BranchContents as ModelBranchContents, CountTableRowsRequest, CreateNamespaceRequest, + CreateNamespaceResponse, CreateTableBranchRequest, CreateTableBranchResponse, + CreateTableIndexRequest, CreateTableIndexResponse, CreateTableRequest, CreateTableResponse, + CreateTableScalarIndexResponse, CreateTableTagRequest, CreateTableTagResponse, + CreateTableVersionRequest, CreateTableVersionResponse, DeclareTableRequest, + DeclareTableResponse, DeleteFromTableRequest, DeleteFromTableResponse, + DeleteTableBranchRequest, DeleteTableBranchResponse, DeleteTableTagRequest, + DeleteTableTagResponse, DescribeNamespaceRequest, DescribeNamespaceResponse, + DescribeTableIndexStatsRequest, DescribeTableIndexStatsResponse, DescribeTableRequest, + DescribeTableResponse, DescribeTableVersionRequest, DescribeTableVersionResponse, + DescribeTransactionRequest, DescribeTransactionResponse, DropNamespaceRequest, + DropNamespaceResponse, DropTableIndexRequest, DropTableIndexResponse, DropTableRequest, + DropTableResponse, ExplainTableQueryPlanRequest, FragmentStats, FragmentSummary, + GetTableStatsRequest, GetTableStatsResponse, GetTableTagVersionRequest, GetTableTagVersionResponse, Identity, IndexContent, InsertIntoTableRequest, InsertIntoTableResponse, ListNamespacesRequest, ListNamespacesResponse, ListTableBranchesRequest, ListTableBranchesResponse, ListTableIndicesRequest, @@ -70,8 +73,8 @@ use lance_namespace::models::{ MergeInsertIntoTableRequest, MergeInsertIntoTableResponse, NamespaceExistsRequest, QueryTableRequest, QueryTableRequestColumns, QueryTableRequestVector, RestoreTableRequest, RestoreTableResponse, TableExistsRequest, TableVersion, TagContents as ModelTagContents, - UpdateTableSchemaMetadataRequest, UpdateTableSchemaMetadataResponse, UpdateTableTagRequest, - UpdateTableTagResponse, + UpdateTableRequest, UpdateTableResponse, UpdateTableSchemaMetadataRequest, + UpdateTableSchemaMetadataResponse, UpdateTableTagRequest, UpdateTableTagResponse, }; use lance_core::{Error, Result, box_error}; @@ -736,7 +739,7 @@ impl DirectoryNamespaceBuilder { Self::initialize_object_store(&self.root, &self.storage_options, &self.session).await?; let manifest_ns = if self.manifest_enabled { - match manifest::ManifestNamespace::from_directory( + match manifest::ManifestNamespace::open_from_directory( self.root.clone(), self.storage_options.clone(), self.session.clone(), @@ -755,18 +758,19 @@ impl DirectoryNamespaceBuilder { // degrading to a directory-listing view that ignores it. return Err(e); } - Err(e) => { - // Failed to initialize manifest namespace, fall back to directory listing only - log::warn!( - "Failed to initialize manifest namespace, falling back to directory listing only: {}", - e - ); + Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => { + log::debug!("Manifest namespace does not exist yet: {}", e); None } + Err(e) => return Err(e), } } else { None }; + let manifest_cell = OnceCell::new(); + if let Some(manifest_ns) = manifest_ns { + let _ = manifest_cell.set(manifest_ns); + } // Create credential vendor once during initialization if enabled let credential_vendor = if has_credential_vendor_config(&self.credential_vendor_properties) @@ -790,8 +794,12 @@ impl DirectoryNamespaceBuilder { session: self.session, object_store, base_path, - manifest_ns, + manifest_ns: manifest_cell, + write_manifest_ns: OnceCell::new(), + manifest_enabled: self.manifest_enabled, dir_listing_enabled: self.dir_listing_enabled, + inline_optimization_enabled: self.inline_optimization_enabled, + commit_retries: self.commit_retries, dir_listing_to_manifest_migration_enabled: self .dir_listing_to_manifest_migration_enabled, table_version_tracking_enabled: self.table_version_tracking_enabled, @@ -868,8 +876,12 @@ pub struct DirectoryNamespace { session: Option>, object_store: Arc, base_path: Path, - manifest_ns: Option>, + manifest_ns: OnceCell>, + write_manifest_ns: OnceCell>, + manifest_enabled: bool, dir_listing_enabled: bool, + inline_optimization_enabled: bool, + commit_retries: Option, /// When true, root-level table operations check the manifest first before /// falling back to directory listing. When false, root-level tables skip /// the manifest check and use directory listing directly. @@ -912,7 +924,121 @@ struct TableDeleteEntry { ranges: Vec<(i64, i64)>, } +/// Persistent record of `alter_transaction` outcomes for a single transaction. +/// +/// Lance's transaction file is immutable once written, so we record any +/// modifications (status transitions, extra properties, tombstoned properties) +/// in a namespace-owned sidecar file. The sidecar is then merged into the +/// response of subsequent `describe_transaction` / `alter_transaction` calls. +/// +/// Serialization is implemented manually via `serde_json::Value` to avoid +/// pulling in `serde`'s `derive` feature for this crate. +#[derive(Debug, Clone, Default)] +struct TransactionAlteration { + /// The most recently applied status, if any. + status: Option, + /// User-defined properties layered on top of the immutable transaction + /// properties. Values here take precedence over the transaction's own + /// properties when both are present. + properties: HashMap, + /// Names of transaction properties that have been tombstoned via + /// `unset_property_action`. A tombstoned key is hidden from the response + /// even when the immutable transaction still carries it. + removed_properties: std::collections::HashSet, +} + +impl TransactionAlteration { + /// JSON field names used for the sidecar on-disk representation. + const F_STATUS: &'static str = "status"; + const F_PROPERTIES: &'static str = "properties"; + const F_REMOVED_PROPERTIES: &'static str = "removed_properties"; + + /// Serialize this alteration to a JSON byte vector. + /// + /// Uses the same pattern as `dir/manifest.rs`: rely on the built-in + /// `Serialize` impls for `Option`, `HashMap` and + /// `HashSet` provided by the `serde` crate (transitively pulled in + /// by `serde_json`), so no `serde` derive nor extra dependency is needed. + fn to_json_bytes(&self) -> serde_json::Result> { + serde_json::to_vec(&serde_json::json!({ + Self::F_STATUS: self.status, + Self::F_PROPERTIES: self.properties, + Self::F_REMOVED_PROPERTIES: self.removed_properties, + })) + } + + /// Deserialize an alteration from JSON bytes, mirroring the + /// `serde_json::from_slice::>(...)` idiom already + /// used in `dir/manifest.rs`. Missing / null fields fall back to defaults + /// so that the sidecar format stays forward-compatible. + fn from_json_slice(bytes: &[u8]) -> serde_json::Result { + let mut obj: serde_json::Map = serde_json::from_slice(bytes)?; + Ok(Self { + status: serde_json::from_value( + obj.remove(Self::F_STATUS) + .unwrap_or(serde_json::Value::Null), + )?, + properties: serde_json::from_value( + obj.remove(Self::F_PROPERTIES) + .unwrap_or(serde_json::Value::Null), + ) + .unwrap_or_default(), + removed_properties: serde_json::from_value( + obj.remove(Self::F_REMOVED_PROPERTIES) + .unwrap_or(serde_json::Value::Null), + ) + .unwrap_or_default(), + }) + } +} + impl DirectoryNamespace { + fn manifest_ns_for_read(&self) -> Option<&Arc> { + self.write_manifest_ns + .get() + .or_else(|| self.manifest_ns.get()) + } + + async fn manifest_ns_for_write(&self) -> Result>> { + if !self.manifest_enabled { + return Ok(None); + } + + let manifest_ns = self + .write_manifest_ns + .get_or_try_init(|| async { + manifest::ManifestNamespace::from_directory( + self.root.clone(), + self.storage_options.clone(), + self.session.clone(), + self.object_store.clone(), + self.base_path.clone(), + self.dir_listing_enabled, + self.inline_optimization_enabled, + self.commit_retries, + ) + .await + .map(Arc::new) + }) + .await?; + Ok(Some(manifest_ns.clone())) + } + + fn child_namespace_requires_manifest_error(&self) -> Error { + if self.manifest_enabled { + NamespaceError::NamespaceNotFound { + message: "Child namespace reads require an existing __manifest dataset".to_string(), + } + .into() + } else { + NamespaceError::Unsupported { + message: "Child namespaces are only supported when manifest mode is enabled" + .to_string(), + } + .into() + } + } + /// Apply pagination to a list of table names /// /// Sorts the list alphabetically and applies pagination using page_token (start_after) and limit. @@ -1135,6 +1261,85 @@ impl DirectoryNamespace { } } + /// Map a Lance error from a table mutation (update / delete / merge-insert) into the most + /// specific `NamespaceError` we can determine from the underlying variant. + /// + /// Collapsing every failure into `InvalidInput`/`Internal` hides the real cause from callers; + /// mapping per variant lets them branch on a meaningful error code (e.g. retry on + /// `ConcurrentModification`, surface `TableNotFound` to the user). + /// + /// Commit-conflict variants are mapped consistently with `convert_lance_commit_error` in + /// `manifest.rs`: `CommitConflict` (retries exhausted, safe to retry) -> `Throttling`, while + /// semantic conflicts (`TooMuchWriteContention` / `RetryableCommitConflict` / + /// `IncompatibleTransaction` / `VersionConflict`) -> `ConcurrentModification`. + fn map_mutation_error( + err: lance_core::Error, + operation: &str, + table_uri: &str, + ) -> lance_core::Error { + let detail = err.to_string(); + let ns_err = match &err { + lance_core::Error::InvalidInput { .. } + | lance_core::Error::Unprocessable { .. } + | lance_core::Error::InvalidRef { .. } => NamespaceError::InvalidInput { + message: format!( + "Invalid input for {} on table at '{}': {}", + operation, table_uri, detail + ), + }, + lance_core::Error::NotFound { .. } | lance_core::Error::DatasetNotFound { .. } => { + NamespaceError::TableNotFound { + message: format!( + "Table at '{}' not found while running {}: {}", + table_uri, operation, detail + ), + } + } + lance_core::Error::SchemaMismatch { .. } | lance_core::Error::Schema { .. } => { + NamespaceError::TableSchemaValidationError { + message: format!( + "Schema validation failed for {} on table at '{}': {}", + operation, table_uri, detail + ), + } + } + // `CommitConflict` means the version-collision retries were exhausted; the operation + // is safe to retry as-is, so surface it as `Throttling` (kept aligned with + // `convert_lance_commit_error` in manifest.rs). + lance_core::Error::CommitConflict { .. } => NamespaceError::Throttling { + message: format!( + "Too many concurrent writes for {} on table at '{}', please retry later: {}", + operation, table_uri, detail + ), + }, + // Semantic conflicts: a concurrent change is incompatible with this one and retrying + // as-is would not help, so surface them as `ConcurrentModification` (kept aligned with + // `convert_lance_commit_error` in manifest.rs). + lance_core::Error::TooMuchWriteContention { .. } + | lance_core::Error::RetryableCommitConflict { .. } + | lance_core::Error::IncompatibleTransaction { .. } + | lance_core::Error::VersionConflict { .. } => NamespaceError::ConcurrentModification { + message: format!( + "Concurrent modification detected for {} on table at '{}': {}", + operation, table_uri, detail + ), + }, + lance_core::Error::NotSupported { .. } => NamespaceError::Unsupported { + message: format!( + "{} is not supported on table at '{}': {}", + operation, table_uri, detail + ), + }, + _ => NamespaceError::Internal { + message: format!( + "Failed to run {} on table at '{}': {}", + operation, table_uri, detail + ), + }, + }; + ns_err.into() + } + async fn table_has_actual_manifests(&self, table_name: &str) -> Result { manifest::ManifestNamespace::path_has_actual_manifests( &self.object_store, @@ -1464,9 +1669,9 @@ impl DirectoryNamespace { if needs_sort { if descending { - table_versions.sort_by(|a, b| b.version.cmp(&a.version)); + table_versions.sort_by_key(|v| std::cmp::Reverse(v.version)); } else { - table_versions.sort_by(|a, b| a.version.cmp(&b.version)); + table_versions.sort_by_key(|v| v.version); } } @@ -1485,10 +1690,11 @@ impl DirectoryNamespace { request: DescribeTableRequest, ) -> Result { let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1); + let is_child_table = request.id.as_ref().is_some_and(|id| id.len() > 1); let skip_manifest_for_root = self.dir_listing_enabled && is_root_level && !self.dir_listing_to_manifest_migration_enabled; - if let Some(ref manifest_ns) = self.manifest_ns + if let Some(manifest_ns) = self.manifest_ns_for_read() && !skip_manifest_for_root { match manifest_ns.describe_table(request.clone()).await { @@ -1518,9 +1724,16 @@ impl DirectoryNamespace { Err(e) => return Err(e), } } + if is_child_table { + return Err(self.child_namespace_requires_manifest_error()); + } let table_name = Self::table_name_from_id(&request.id)?; let table_id = Self::format_table_id_from_request(&request.id); + if !self.dir_listing_enabled { + return Err(NamespaceError::TableNotFound { message: table_id }.into()); + } + let table_uri = self.table_full_uri(&table_name); // Atomically check table existence and deregistration status @@ -1984,12 +2197,29 @@ impl DirectoryNamespace { fn transaction_response( version: u64, transaction: &Transaction, + alteration: Option, ) -> DescribeTransactionResponse { let mut properties = transaction .transaction_properties .as_ref() .map(|properties| (**properties).clone()) .unwrap_or_default(); + + // Apply persisted alterations on top of the immutable transaction + // properties so callers see the current effective state. + let mut effective_status = "SUCCEEDED".to_string(); + if let Some(alteration) = alteration { + for key in &alteration.removed_properties { + properties.remove(key); + } + for (key, value) in alteration.properties { + properties.insert(key, value); + } + if let Some(status) = alteration.status { + effective_status = status; + } + } + properties.insert("uuid".to_string(), transaction.uuid.clone()); properties.insert("version".to_string(), version.to_string()); properties.insert( @@ -2005,7 +2235,7 @@ impl DirectoryNamespace { } DescribeTransactionResponse { - status: "SUCCEEDED".to_string(), + status: effective_status, properties: Some(properties), } } @@ -2094,8 +2324,86 @@ impl DirectoryNamespace { .into()) } + /// Relative directory (under a table's Lance root) used to persist + /// alter_transaction outcomes. The Lance transaction file itself is + /// immutable, so we keep alterations in a namespace-owned sidecar. + const TRANSACTION_ALTERATIONS_DIR: &'static str = "_alter_transactions"; + + fn transaction_alteration_path(&self, table_uri: &str, txn_uuid: &str) -> Result { + let table_path = self.object_store_path_from_uri(table_uri)?; + Ok(table_path + .join(Self::TRANSACTION_ALTERATIONS_DIR) + .join(format!("{}.json", txn_uuid).as_str())) + } + + async fn load_transaction_alteration( + &self, + table_uri: &str, + txn_uuid: &str, + ) -> Result> { + let path = self.transaction_alteration_path(table_uri, txn_uuid)?; + match self.object_store.inner.get(&path).await { + Ok(get_result) => { + let bytes = get_result.bytes().await.map_err(|e| { + lance_core::Error::from(NamespaceError::Internal { + message: format!( + "Failed to read alter_transaction sidecar for '{}': {}", + txn_uuid, e + ), + }) + })?; + let alteration = TransactionAlteration::from_json_slice(&bytes).map_err(|e| { + lance_core::Error::from(NamespaceError::Internal { + message: format!( + "Failed to parse alter_transaction sidecar for '{}': {}", + txn_uuid, e + ), + }) + })?; + Ok(Some(alteration)) + } + Err(ObjectStoreError::NotFound { .. }) => Ok(None), + Err(e) => Err(lance_core::Error::from(NamespaceError::Internal { + message: format!( + "Failed to load alter_transaction sidecar for '{}': {}", + txn_uuid, e + ), + })), + } + } + + async fn save_transaction_alteration( + &self, + table_uri: &str, + txn_uuid: &str, + alteration: &TransactionAlteration, + ) -> Result<()> { + let path = self.transaction_alteration_path(table_uri, txn_uuid)?; + let bytes = alteration.to_json_bytes().map_err(|e| { + lance_core::Error::from(NamespaceError::Internal { + message: format!( + "Failed to serialize alter_transaction sidecar for '{}': {}", + txn_uuid, e + ), + }) + })?; + self.object_store + .inner + .put(&path, bytes.into()) + .await + .map_err(|e| { + lance_core::Error::from(NamespaceError::Internal { + message: format!( + "Failed to persist alter_transaction sidecar for '{}': {}", + txn_uuid, e + ), + }) + })?; + Ok(()) + } + fn table_full_uri(&self, table_name: &str) -> String { - format!("{}/{}.lance", &self.root, table_name) + format!("{}/{}.lance", self.root, table_name) } /// Get the object store path for a table (relative to base_path) @@ -2281,7 +2589,7 @@ impl DirectoryNamespace { /// - Manifest registration fails pub async fn migrate(&self) -> Result { // We only care about tables in the root namespace - let Some(ref manifest_ns) = self.manifest_ns else { + let Some(manifest_ns) = self.manifest_ns_for_write().await? else { return Ok(0); // No manifest, nothing to migrate }; @@ -2553,10 +2861,13 @@ impl LanceNamespace for DirectoryNamespace { request: ListNamespacesRequest, ) -> Result { self.record_op("list_namespaces"); - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_read() { return manifest_ns.list_namespaces(request).await; } + if request.id.as_ref().is_some_and(|id| !id.is_empty()) { + return Err(self.child_namespace_requires_manifest_error()); + } Self::validate_root_namespace_id(&request.id)?; Ok(ListNamespacesResponse::new(vec![])) } @@ -2566,10 +2877,13 @@ impl LanceNamespace for DirectoryNamespace { request: DescribeNamespaceRequest, ) -> Result { self.record_op("describe_namespace"); - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_read() { return manifest_ns.describe_namespace(request).await; } + if request.id.as_ref().is_some_and(|id| !id.is_empty()) { + return Err(self.child_namespace_requires_manifest_error()); + } Self::validate_root_namespace_id(&request.id)?; #[allow(clippy::needless_update)] Ok(DescribeNamespaceResponse { @@ -2583,7 +2897,7 @@ impl LanceNamespace for DirectoryNamespace { request: CreateNamespaceRequest, ) -> Result { self.record_op("create_namespace"); - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_write().await? { return manifest_ns.create_namespace(request).await; } @@ -2603,7 +2917,7 @@ impl LanceNamespace for DirectoryNamespace { async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result { self.record_op("drop_namespace"); - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_write().await? { return manifest_ns.drop_namespace(request).await; } @@ -2623,7 +2937,7 @@ impl LanceNamespace for DirectoryNamespace { async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> { self.record_op("namespace_exists"); - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_read() { return manifest_ns.namespace_exists(request).await; } @@ -2631,11 +2945,7 @@ impl LanceNamespace for DirectoryNamespace { return Ok(()); } - Err(NamespaceError::NamespaceNotFound { - message: "Child namespaces are only supported when manifest mode is enabled" - .to_string(), - } - .into()) + Err(self.child_namespace_requires_manifest_error()) } async fn list_tables(&self, request: ListTablesRequest) -> Result { @@ -2649,31 +2959,30 @@ impl LanceNamespace for DirectoryNamespace { // For child namespaces, always delegate to manifest (if enabled) if !namespace_id.is_empty() { - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_read() { return manifest_ns.list_tables(request).await; } - return Err(NamespaceError::Unsupported { - message: "Child namespaces are only supported when manifest mode is enabled" - .to_string(), - } - .into()); + return Err(self.child_namespace_requires_manifest_error()); } // When only manifest is enabled (no directory listing), delegate directly to manifest - if let Some(ref manifest_ns) = self.manifest_ns + if let Some(manifest_ns) = self.manifest_ns_for_read() && !self.dir_listing_enabled { return manifest_ns.list_tables(request).await; } + if !self.dir_listing_enabled { + return Ok(ListTablesResponse::new(vec![])); + } // When both manifest and directory listing are enabled with migration mode, // we need to merge and deduplicate - let mut tables = if self.manifest_ns.is_some() + let mut tables = if self.manifest_ns_for_read().is_some() && self.dir_listing_enabled && self.dir_listing_to_manifest_migration_enabled { // Get all manifest table locations (for deduplication) - let manifest_locations = if let Some(ref manifest_ns) = self.manifest_ns { + let manifest_locations = if let Some(manifest_ns) = self.manifest_ns_for_read() { manifest_ns.list_manifest_table_locations().await? } else { std::collections::HashSet::new() @@ -2683,7 +2992,7 @@ impl LanceNamespace for DirectoryNamespace { let mut manifest_request = request.clone(); manifest_request.limit = None; manifest_request.page_token = None; - let manifest_tables = if let Some(ref manifest_ns) = self.manifest_ns { + let manifest_tables = if let Some(manifest_ns) = self.manifest_ns_for_read() { let manifest_response = manifest_ns.list_tables(manifest_request).await?; manifest_response.tables } else { @@ -2731,10 +3040,11 @@ impl LanceNamespace for DirectoryNamespace { async fn table_exists(&self, request: TableExistsRequest) -> Result<()> { self.record_op("table_exists"); let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1); + let is_child_table = request.id.as_ref().is_some_and(|id| id.len() > 1); let skip_manifest_for_root = self.dir_listing_enabled && is_root_level && !self.dir_listing_to_manifest_migration_enabled; - if let Some(ref manifest_ns) = self.manifest_ns + if let Some(manifest_ns) = self.manifest_ns_for_read() && !skip_manifest_for_root { match manifest_ns.table_exists(request.clone()).await { @@ -2750,9 +3060,15 @@ impl LanceNamespace for DirectoryNamespace { Err(e) => return Err(e), } } + if is_child_table { + return Err(self.child_namespace_requires_manifest_error()); + } let table_name = Self::table_name_from_id(&request.id)?; let table_id = Self::format_table_id_from_request(&request.id); + if !self.dir_listing_enabled { + return Err(NamespaceError::TableNotFound { message: table_id }.into()); + } // Atomically check table existence and deregistration status let status = self.check_table_status(&table_name).await; @@ -2776,7 +3092,7 @@ impl LanceNamespace for DirectoryNamespace { async fn drop_table(&self, request: DropTableRequest) -> Result { self.record_op("drop_table"); - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_write().await? { return manifest_ns.drop_table(request).await; } @@ -2806,7 +3122,7 @@ impl LanceNamespace for DirectoryNamespace { request_data: Bytes, ) -> Result { self.record_op("create_table"); - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_write().await? { return manifest_ns.create_table(request, request_data).await; } @@ -2853,7 +3169,7 @@ impl LanceNamespace for DirectoryNamespace { async fn declare_table(&self, request: DeclareTableRequest) -> Result { self.record_op("declare_table"); - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_write().await? { let mut response = manifest_ns.declare_table(request.clone()).await?; if let Some(ref location) = response.location { // For backwards compatibility, only skip vending credentials when explicitly set to false @@ -2944,7 +3260,7 @@ impl LanceNamespace for DirectoryNamespace { ) -> Result { self.record_op("register_table"); // If manifest is enabled, delegate to manifest namespace - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_write().await? { return LanceNamespace::register_table(manifest_ns.as_ref(), request).await; } @@ -2961,7 +3277,7 @@ impl LanceNamespace for DirectoryNamespace { ) -> Result { self.record_op("deregister_table"); // If manifest is enabled, delegate to manifest namespace - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_write().await? { return LanceNamespace::deregister_table(manifest_ns.as_ref(), request).await; } @@ -3019,7 +3335,7 @@ impl LanceNamespace for DirectoryNamespace { &self, request: AlterTableAddColumnsRequest, ) -> Result { - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_write().await? { return manifest_ns.alter_table_add_columns(request).await; } @@ -3077,7 +3393,7 @@ impl LanceNamespace for DirectoryNamespace { &self, request: AlterTableAlterColumnsRequest, ) -> Result { - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_write().await? { return manifest_ns.alter_table_alter_columns(request).await; } @@ -3127,7 +3443,7 @@ impl LanceNamespace for DirectoryNamespace { &self, request: AlterTableDropColumnsRequest, ) -> Result { - if let Some(ref manifest_ns) = self.manifest_ns { + if let Some(manifest_ns) = self.manifest_ns_for_write().await? { return manifest_ns.alter_table_drop_columns(request).await; } @@ -3681,7 +3997,212 @@ impl LanceNamespace for DirectoryNamespace { .await?; let (version, transaction) = self.find_transaction(&dataset, &id).await?; - Ok(Self::transaction_response(version, &transaction)) + // Merge any persisted alter_transaction changes stored in the sidecar + // so that describe_transaction reflects the latest altered state. + let sidecar = self + .load_transaction_alteration(&table_uri, &transaction.uuid) + .await?; + + Ok(Self::transaction_response(version, &transaction, sidecar)) + } + + async fn alter_transaction( + &self, + request: AlterTransactionRequest, + ) -> Result { + self.record_op("alter_transaction"); + + // Parse the request ID: must include table id and transaction identifier + let mut request_id = request.id.ok_or_else(|| { + lance_core::Error::from(NamespaceError::InvalidInput { + message: "Transaction id must include table id and transaction identifier" + .to_string(), + }) + })?; + if request_id.len() < 2 { + return Err(NamespaceError::InvalidInput { + message: format!( + "Transaction request id must include table id and transaction identifier, got {:?}", + request_id + ), + } + .into()); + } + + let txn_id = request_id.pop().expect("request_id len checked above"); + let table_id = Some(request_id); + let table_uri = self.resolve_table_location(&table_id).await?; + let dataset = self + .load_dataset(&table_uri, None, "alter_transaction") + .await?; + let (version, transaction) = self.find_transaction(&dataset, &txn_id).await?; + + // Reserved keys are derived from the immutable Transaction metadata and + // must not be modified via alter_transaction. They are only surfaced in + // the response for the caller's convenience. + const RESERVED_KEYS: &[&str] = &["uuid", "version", "read_version", "operation", "tag"]; + let is_reserved = |key: &str| RESERVED_KEYS.contains(&key); + + // Load the existing sidecar (if any) so alterations accumulate across + // successive alter_transaction calls. + let mut sidecar = self + .load_transaction_alteration(&table_uri, &transaction.uuid) + .await? + .unwrap_or_default(); + + for action in &request.actions { + if let Some(ref set_status) = action.set_status_action + && let Some(ref status) = set_status.status + { + // Validate the status value (case-insensitive) + let normalized = status.to_lowercase().replace('_', ""); + match normalized.as_str() { + "queued" | "running" | "succeeded" | "failed" | "canceled" => { + sidecar.status = Some(status.clone()); + } + _ => { + return Err(NamespaceError::InvalidInput { + message: format!( + "Invalid transaction status '{}'. Valid values are: Queued, Running, Succeeded, Failed, Canceled", + status + ), + } + .into()); + } + } + } + + if let Some(ref set_property) = action.set_property_action + && let (Some(key), Some(value)) = (&set_property.key, &set_property.value) + { + if is_reserved(key) { + return Err(NamespaceError::InvalidInput { + message: format!("Property '{}' is reserved and cannot be modified", key), + } + .into()); + } + let mode = set_property + .mode + .as_deref() + .unwrap_or("Overwrite") + .to_lowercase(); + match mode.as_str() { + "overwrite" => { + sidecar.properties.insert(key.clone(), value.clone()); + } + "fail" => { + // Consider both the immutable transaction properties + // and any values previously written to the sidecar. + let exists = sidecar.properties.contains_key(key) + || transaction + .transaction_properties + .as_ref() + .is_some_and(|props| props.contains_key(key)); + if exists { + return Err(NamespaceError::ConcurrentModification { + message: format!( + "Property '{}' already exists and mode is 'Fail'", + key + ), + } + .into()); + } + sidecar.properties.insert(key.clone(), value.clone()); + } + "skip" => { + let exists = sidecar.properties.contains_key(key) + || transaction + .transaction_properties + .as_ref() + .is_some_and(|props| props.contains_key(key)); + if !exists { + sidecar.properties.insert(key.clone(), value.clone()); + } + } + _ => { + return Err(NamespaceError::InvalidInput { + message: format!( + "Invalid set_property mode '{}'. Valid values are: Overwrite, Fail, Skip", + mode + ), + } + .into()); + } + } + } + + if let Some(ref unset_property) = action.unset_property_action + && let Some(ref key) = unset_property.key + { + if is_reserved(key) { + return Err(NamespaceError::InvalidInput { + message: format!("Property '{}' is reserved and cannot be modified", key), + } + .into()); + } + let mode = unset_property + .mode + .as_deref() + .unwrap_or("Skip") + .to_lowercase(); + let exists_in_transaction = transaction + .transaction_properties + .as_ref() + .is_some_and(|props| props.contains_key(key)); + match mode.as_str() { + "skip" => { + sidecar.properties.remove(key); + if exists_in_transaction { + // Track a tombstone so describe_transaction can + // hide the immutable property from the response. + sidecar.removed_properties.insert(key.clone()); + } + } + "fail" => { + if !sidecar.properties.contains_key(key) && !exists_in_transaction { + return Err(NamespaceError::InvalidInput { + message: format!( + "Property '{}' does not exist and mode is 'Fail'", + key + ), + } + .into()); + } + sidecar.properties.remove(key); + if exists_in_transaction { + sidecar.removed_properties.insert(key.clone()); + } + } + _ => { + return Err(NamespaceError::InvalidInput { + message: format!( + "Invalid unset_property mode '{}'. Valid values are: Skip, Fail", + mode + ), + } + .into()); + } + } + } + } + + // Persist the accumulated alterations so subsequent calls observe + // them. The transaction file itself is immutable in Lance, so we + // record alter_transaction outcomes in a namespace-owned sidecar. + self.save_transaction_alteration(&table_uri, &transaction.uuid, &sidecar) + .await?; + + // Assemble the response by merging the immutable transaction metadata + // with the persisted alterations. + let final_status = sidecar + .status + .clone() + .unwrap_or_else(|| "SUCCEEDED".to_string()); + let response = Self::transaction_response(version, &transaction, Some(sidecar)); + Ok(AlterTransactionResponse { + status: final_status, + properties: response.properties, + }) } async fn create_table_scalar_index( @@ -4194,12 +4715,7 @@ impl LanceNamespace for DirectoryNamespace { })? .execute_reader(reader) .await - .map_err(|e| NamespaceError::Internal { - message: format!( - "Failed to merge_insert_into_table at '{}': {}", - table_uri, e - ), - })?; + .map_err(|e| Self::map_mutation_error(e, "merge_insert_into_table", &table_uri))?; Ok(MergeInsertIntoTableResponse { transaction_id: None, @@ -4210,30 +4726,143 @@ impl LanceNamespace for DirectoryNamespace { }) } - async fn query_table(&self, request: QueryTableRequest) -> Result { - use arrow::ipc::writer::FileWriter; - - self.record_op("query_table"); - let table_uri = self.resolve_table_location(&request.id).await?; - let dataset = self - .load_dataset(&table_uri, request.version, "query_table") - .await?; + async fn update_table(&self, request: UpdateTableRequest) -> Result { + self.record_op("update_table"); - // Build scanner - let mut scanner = dataset.scan(); + if request.updates.is_empty() { + return Err(NamespaceError::InvalidInput { + message: "update_table requires at least one [column, expression] pair".to_string(), + } + .into()); + } - // Check if this is a vector search query - // vector is Box, not Option - let has_vector_query = request - .vector - .single_vector - .as_ref() - .map(|sv| !sv.is_empty()) - .unwrap_or(false) - || request - .vector - .multi_vector - .as_ref() + // Validate every update pair shape and detect duplicate columns up front so we + // surface a clean error instead of failing deep inside the planner. + let mut seen_columns: HashMap = HashMap::with_capacity(request.updates.len()); + for (idx, pair) in request.updates.iter().enumerate() { + if pair.len() != 2 { + return Err(NamespaceError::InvalidInput { + message: format!( + "update_table updates[{}] must be a [column, expression] pair, got {} elements", + idx, + pair.len() + ), + } + .into()); + } + let column = &pair[0]; + if column.trim().is_empty() { + return Err(NamespaceError::InvalidInput { + message: format!("update_table updates[{}] has an empty column name", idx), + } + .into()); + } + if seen_columns.insert(column.clone(), ()).is_some() { + return Err(NamespaceError::InvalidInput { + message: format!( + "update_table cannot update column '{}' more than once", + column + ), + } + .into()); + } + } + + let table_uri = self.resolve_table_location(&request.id).await?; + let dataset = Arc::new(self.load_dataset(&table_uri, None, "update_table").await?); + + let mut builder = UpdateBuilder::new(dataset); + for pair in &request.updates { + // Indexing by 0/1 is safe due to the length check above. + builder = builder.set(&pair[0], &pair[1]).map_err(|e| { + lance_core::Error::from(NamespaceError::InvalidInput { + message: format!("Invalid update expression for column '{}': {}", pair[0], e), + }) + })?; + } + if let Some(predicate) = request.predicate.as_deref() + && !predicate.trim().is_empty() + { + builder = builder.update_where(predicate).map_err(|e| { + lance_core::Error::from(NamespaceError::InvalidInput { + message: format!("Invalid update_table predicate '{}': {}", predicate, e), + }) + })?; + } + + let job = builder.build().map_err(|e| { + lance_core::Error::from(NamespaceError::InvalidInput { + message: format!("Failed to build update_table job: {}", e), + }) + })?; + + let result = job + .execute() + .await + .map_err(|e| Self::map_mutation_error(e, "update_table", &table_uri))?; + + let version = result.new_dataset.version().version as i64; + Ok(UpdateTableResponse { + transaction_id: None, + updated_rows: result.rows_updated as i64, + version, + properties: None, + }) + } + + async fn delete_from_table( + &self, + request: DeleteFromTableRequest, + ) -> Result { + self.record_op("delete_from_table"); + + if request.predicate.trim().is_empty() { + return Err(NamespaceError::InvalidInput { + message: "delete_from_table requires a non-empty predicate".to_string(), + } + .into()); + } + + let table_uri = self.resolve_table_location(&request.id).await?; + let mut dataset = self + .load_dataset(&table_uri, None, "delete_from_table") + .await?; + + let result = dataset + .delete(&request.predicate) + .await + .map_err(|e| Self::map_mutation_error(e, "delete_from_table", &table_uri))?; + + Ok(DeleteFromTableResponse { + transaction_id: None, + version: Some(result.new_dataset.version().version as i64), + }) + } + + async fn query_table(&self, request: QueryTableRequest) -> Result { + use arrow::ipc::writer::FileWriter; + + self.record_op("query_table"); + let table_uri = self.resolve_table_location(&request.id).await?; + let dataset = self + .load_dataset(&table_uri, request.version, "query_table") + .await?; + + // Build scanner + let mut scanner = dataset.scan(); + + // Check if this is a vector search query + // vector is Box, not Option + let has_vector_query = request + .vector + .single_vector + .as_ref() + .map(|sv| !sv.is_empty()) + .unwrap_or(false) + || request + .vector + .multi_vector + .as_ref() .map(|mv| !mv.is_empty()) .unwrap_or(false); @@ -4796,6 +5425,52 @@ mod tests { } } + fn mutation_error_code(err: lance_core::Error) -> ErrorCode { + match err { + lance_core::Error::Namespace { source, .. } => source + .downcast_ref::() + .expect("mutation error should wrap a NamespaceError") + .code(), + other => panic!("expected Namespace error, got: {other:?}"), + } + } + + /// `map_mutation_error` must classify commit-conflict variants the same way as + /// `convert_lance_commit_error` in `manifest.rs`: `CommitConflict` is a retries-exhausted + /// version collision that is safe to retry (`Throttling`), while the semantic-conflict variants + /// map to `ConcurrentModification`. + #[test] + fn test_map_mutation_error_commit_conflict_alignment() { + let boxed = || -> Box { + Box::::from("inner conflict") + }; + + let throttling_cases = vec![lance_core::Error::commit_conflict_source(1, boxed())]; + for err in throttling_cases { + let code = mutation_error_code(DirectoryNamespace::map_mutation_error( + err, + "update", + "memory://t", + )); + assert_eq!(code, ErrorCode::Throttling); + } + + let concurrent_cases = vec![ + lance_core::Error::too_much_write_contention("contention"), + lance_core::Error::retryable_commit_conflict_source(1, boxed()), + lance_core::Error::incompatible_transaction_source(boxed()), + lance_core::Error::version_conflict("conflict", 0, 3), + ]; + for err in concurrent_cases { + let code = mutation_error_code(DirectoryNamespace::map_mutation_error( + err, + "update", + "memory://t", + )); + assert_eq!(code, ErrorCode::ConcurrentModification); + } + } + /// Helper to create a test DirectoryNamespace with a temporary directory async fn create_test_namespace() -> (DirectoryNamespace, TempStdDir) { let temp_dir = TempStdDir::default(); @@ -4971,6 +5646,43 @@ mod tests { create_ipc_data_from_batches(schema, vec![batch]) } + async fn create_legacy_manifest_without_primary_key_metadata(root: &str) { + use arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; + use arrow::record_batch::{RecordBatch, RecordBatchIterator}; + + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("object_id", DataType::Utf8, false), + Field::new("object_type", DataType::Utf8, false), + Field::new("location", DataType::Utf8, true), + Field::new("metadata", DataType::Utf8, true), + Field::new( + "base_objects", + DataType::List(Arc::new(Field::new("object_id", DataType::Utf8, true))), + true, + ), + ])); + let batch = RecordBatch::new_empty(schema.clone()); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + Dataset::write(Box::new(reader), &format!("{}/__manifest", root), None) + .await + .unwrap(); + } + + async fn manifest_has_primary_key_metadata(root: &str) -> bool { + let dataset = Dataset::open(&format!("{}/__manifest", root)) + .await + .unwrap(); + dataset + .schema() + .field("object_id") + .map(|field| { + field + .metadata + .contains_key(lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION) + }) + .unwrap_or(false) + } + fn create_vector_table_ipc_data() -> Vec { use arrow::array::{FixedSizeListArray, Float32Array, Int32Array}; use arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; @@ -11218,6 +11930,204 @@ mod tests { assert!(total_rows > 0); assert!(total_rows < 3); } + + // ---------------------- update_table / delete_from_table ---------------------- + + #[tokio::test] + async fn test_update_full_table() { + let (namespace, _temp_dir, table_id) = create_ns_with_table().await; + + // Capture base version so we can assert the update bumped it. + let base_version = open_dataset(&namespace, &table_id[0]) + .await + .version() + .version; + + let request = UpdateTableRequest { + id: Some(table_id.clone()), + updates: vec![vec!["name".to_string(), "'updated'".to_string()]], + predicate: None, + ..Default::default() + }; + + let response = namespace.update_table(request).await.unwrap(); + assert_eq!(response.updated_rows, 3); + assert!(response.version as u64 > base_version); + + // Validate that all rows now carry the new value. + let count_req = CountTableRowsRequest { + id: Some(table_id), + version: None, + predicate: Some("name = 'updated'".to_string()), + ..Default::default() + }; + assert_eq!(namespace.count_table_rows(count_req).await.unwrap(), 3); + } + + #[tokio::test] + async fn test_update_with_predicate() { + let (namespace, _temp_dir, table_id) = create_ns_with_table().await; + + let request = UpdateTableRequest { + id: Some(table_id.clone()), + updates: vec![vec!["name".to_string(), "'matched'".to_string()]], + predicate: Some("id > 1".to_string()), + ..Default::default() + }; + + let response = namespace.update_table(request).await.unwrap(); + assert_eq!(response.updated_rows, 2); + + // Rows that did not match the predicate must remain unchanged. + let untouched = CountTableRowsRequest { + id: Some(table_id.clone()), + version: None, + predicate: Some("name = 'Alice'".to_string()), + ..Default::default() + }; + assert_eq!(namespace.count_table_rows(untouched).await.unwrap(), 1); + + let touched = CountTableRowsRequest { + id: Some(table_id), + version: None, + predicate: Some("name = 'matched'".to_string()), + ..Default::default() + }; + assert_eq!(namespace.count_table_rows(touched).await.unwrap(), 2); + } + + #[tokio::test] + async fn test_update_invalid_expression_returns_invalid_input() { + let (namespace, _temp_dir, table_id) = create_ns_with_table().await; + + let request = UpdateTableRequest { + id: Some(table_id), + // Reference an unknown column on the right-hand side. + updates: vec![vec!["name".to_string(), "no_such_column + 1".to_string()]], + predicate: None, + ..Default::default() + }; + + let err = namespace.update_table(request).await.unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("Invalid input"), + "expected Invalid input error, got: {}", + msg + ); + } + + #[tokio::test] + async fn test_update_rejects_duplicate_columns() { + let (namespace, _temp_dir, table_id) = create_ns_with_table().await; + + let request = UpdateTableRequest { + id: Some(table_id), + updates: vec![ + vec!["name".to_string(), "'a'".to_string()], + vec!["name".to_string(), "'b'".to_string()], + ], + predicate: None, + ..Default::default() + }; + + let err = namespace.update_table(request).await.unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("Invalid input") && msg.contains("more than once"), + "expected duplicate column InvalidInput error, got: {}", + msg + ); + } + + #[tokio::test] + async fn test_delete_with_predicate() { + let (namespace, _temp_dir, table_id) = create_ns_with_table().await; + + let request = DeleteFromTableRequest { + id: Some(table_id.clone()), + predicate: "id > 1".to_string(), + ..Default::default() + }; + + let response = namespace.delete_from_table(request).await.unwrap(); + assert!(response.version.is_some()); + + let count_req = CountTableRowsRequest { + id: Some(table_id), + version: None, + predicate: None, + ..Default::default() + }; + // Original rows = 3; after deleting `id > 1` only row id=1 remains. + assert_eq!(namespace.count_table_rows(count_req).await.unwrap(), 1); + } + + #[tokio::test] + async fn test_delete_empty_predicate_returns_invalid_input() { + let (namespace, _temp_dir, table_id) = create_ns_with_table().await; + + let request = DeleteFromTableRequest { + id: Some(table_id), + predicate: " ".to_string(), + ..Default::default() + }; + + let err = namespace.delete_from_table(request).await.unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("Invalid input") && msg.contains("non-empty predicate"), + "expected non-empty predicate InvalidInput error, got: {}", + msg + ); + } + + #[tokio::test] + async fn test_delete_table_not_found() { + let (namespace, _temp_dir) = create_test_namespace().await; + + let request = DeleteFromTableRequest { + id: Some(vec!["does_not_exist".to_string()]), + predicate: "id = 1".to_string(), + ..Default::default() + }; + + let err = namespace.delete_from_table(request).await.unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("Table not found"), + "expected TableNotFound for missing table, got: {}", + msg + ); + } + + #[tokio::test] + async fn test_delete_invalid_predicate_returns_invalid_input() { + let (namespace, _temp_dir, table_id) = create_ns_with_table().await; + + // A predicate referencing a column that does not exist reaches `Dataset::delete` + // and surfaces as `Error::InvalidInput`, which must map to `InvalidInput` rather + // than a generic `Internal`. + let request = DeleteFromTableRequest { + id: Some(table_id), + predicate: "no_such_column = 1".to_string(), + ..Default::default() + }; + + let err = namespace.delete_from_table(request).await.unwrap_err(); + let lance_core::Error::Namespace { source, .. } = &err else { + panic!("expected a Namespace error, got: {}", err); + }; + let ns_err = source + .downcast_ref::() + .expect("expected a NamespaceError source"); + assert_eq!( + ns_err.code(), + lance_namespace::ErrorCode::InvalidInput, + "expected InvalidInput for an invalid delete predicate, got: {}", + err + ); + } } #[tokio::test] @@ -11634,6 +12544,109 @@ mod tests { ); } + #[tokio::test] + async fn test_build_and_root_reads_do_not_create_manifest() { + let temp_dir = TempStdDir::default(); + let temp_path = temp_dir.to_str().unwrap(); + let manifest_path = std::path::Path::new(temp_path).join("__manifest"); + + let dir_only_ns = DirectoryNamespaceBuilder::new(temp_path) + .manifest_enabled(false) + .dir_listing_enabled(true) + .build() + .await + .unwrap(); + create_scalar_table(&dir_only_ns, "catalog").await; + assert!(!manifest_path.exists()); + + let namespace = DirectoryNamespaceBuilder::new(temp_path) + .manifest_enabled(true) + .dir_listing_enabled(true) + .build() + .await + .unwrap(); + assert!(!manifest_path.exists()); + + let mut exists_req = TableExistsRequest::new(); + exists_req.id = Some(vec!["catalog".to_string()]); + namespace.table_exists(exists_req).await.unwrap(); + assert!(!manifest_path.exists()); + + let mut describe_req = DescribeTableRequest::new(); + describe_req.id = Some(vec!["catalog".to_string()]); + namespace.describe_table(describe_req).await.unwrap(); + assert!(!manifest_path.exists()); + + let list_response = namespace + .list_tables(ListTablesRequest { + id: Some(vec![]), + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(list_response.tables, vec!["catalog".to_string()]); + assert!(!manifest_path.exists()); + + let mut list_namespaces_req = ListNamespacesRequest::new(); + list_namespaces_req.id = Some(vec!["workspace".to_string()]); + let err = namespace + .list_namespaces(list_namespaces_req) + .await + .unwrap_err(); + assert!(err.to_string().contains("__manifest")); + assert!(!manifest_path.exists()); + + let err = namespace + .list_tables(ListTablesRequest { + id: Some(vec!["workspace".to_string()]), + ..Default::default() + }) + .await + .unwrap_err(); + assert!(err.to_string().contains("__manifest")); + assert!(!manifest_path.exists()); + + let mut child_describe_req = DescribeTableRequest::new(); + child_describe_req.id = Some(vec!["workspace".to_string(), "catalog".to_string()]); + let err = namespace + .describe_table(child_describe_req) + .await + .unwrap_err(); + assert!(err.to_string().contains("__manifest")); + assert!(!manifest_path.exists()); + + let mut child_exists_req = TableExistsRequest::new(); + child_exists_req.id = Some(vec!["workspace".to_string(), "catalog".to_string()]); + let err = namespace.table_exists(child_exists_req).await.unwrap_err(); + assert!(err.to_string().contains("__manifest")); + assert!(!manifest_path.exists()); + + let mut create_ns_req = CreateNamespaceRequest::new(); + create_ns_req.id = Some(vec!["workspace".to_string()]); + namespace.create_namespace(create_ns_req).await.unwrap(); + assert!(manifest_path.exists()); + } + + #[tokio::test] + async fn test_migrate_updates_read_opened_legacy_manifest() { + let temp_dir = TempStdDir::default(); + let temp_path = temp_dir.to_str().unwrap(); + create_legacy_manifest_without_primary_key_metadata(temp_path).await; + assert!(!manifest_has_primary_key_metadata(temp_path).await); + + let namespace = DirectoryNamespaceBuilder::new(temp_path) + .manifest_enabled(true) + .dir_listing_enabled(true) + .build() + .await + .unwrap(); + assert!(!manifest_has_primary_key_metadata(temp_path).await); + + let migrated = namespace.migrate().await.unwrap(); + assert_eq!(migrated, 0); + assert!(manifest_has_primary_key_metadata(temp_path).await); + } + #[tokio::test] async fn test_describe_declared_table_checks_versions_only_when_requested() { let temp_dir = TempStdDir::default(); @@ -12536,4 +13549,341 @@ mod tests { .expect("reopen branch failed"); assert_eq!(scan_id_column(&reopened).await, vec![1, 2]); } + + #[tokio::test] + async fn test_alter_transaction_set_status() { + use lance_namespace::models::{ + AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus, + DescribeTransactionRequest, + }; + + let (namespace, _temp_dir) = create_test_namespace().await; + create_scalar_table(&namespace, "users").await; + let txn_id = create_scalar_index(&namespace, "users", "users_id_idx") + .await + .expect("create_scalar_index should return a transaction id"); + + // First verify the transaction exists + let describe_resp = namespace + .describe_transaction(DescribeTransactionRequest { + id: Some(vec!["users".to_string(), txn_id.clone()]), + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(describe_resp.status, "SUCCEEDED"); + + // Alter the transaction status + let response = namespace + .alter_transaction(AlterTransactionRequest { + id: Some(vec!["users".to_string(), txn_id.clone()]), + actions: vec![AlterTransactionAction { + set_status_action: Some(Box::new(AlterTransactionSetStatus { + status: Some("Canceled".to_string()), + })), + set_property_action: None, + unset_property_action: None, + }], + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(response.status, "Canceled"); + assert!(response.properties.is_some()); + let props = response.properties.unwrap(); + assert_eq!(props.get("uuid"), Some(&txn_id)); + assert_eq!(props.get("operation"), Some(&"CreateIndex".to_string())); + } + + #[tokio::test] + async fn test_alter_transaction_set_property() { + use lance_namespace::models::{ + AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty, + }; + + let (namespace, _temp_dir) = create_test_namespace().await; + create_scalar_table(&namespace, "users").await; + let txn_id = create_scalar_index(&namespace, "users", "users_id_idx") + .await + .expect("create_scalar_index should return a transaction id"); + + let response = namespace + .alter_transaction(AlterTransactionRequest { + id: Some(vec!["users".to_string(), txn_id.clone()]), + actions: vec![AlterTransactionAction { + set_status_action: None, + set_property_action: Some(Box::new(AlterTransactionSetProperty { + key: Some("custom_key".to_string()), + value: Some("custom_value".to_string()), + mode: None, + })), + unset_property_action: None, + }], + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(response.status, "SUCCEEDED"); + let props = response.properties.unwrap(); + assert_eq!(props.get("custom_key"), Some(&"custom_value".to_string())); + } + + #[tokio::test] + async fn test_alter_transaction_set_property_fail_mode() { + use lance_namespace::models::{ + AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty, + }; + + let (namespace, _temp_dir) = create_test_namespace().await; + create_scalar_table(&namespace, "users").await; + let txn_id = create_scalar_index(&namespace, "users", "users_id_idx") + .await + .expect("create_scalar_index should return a transaction id"); + + // First, set a non-reserved property so it exists in the sidecar. + namespace + .alter_transaction(AlterTransactionRequest { + id: Some(vec!["users".to_string(), txn_id.clone()]), + actions: vec![AlterTransactionAction { + set_status_action: None, + set_property_action: Some(Box::new(AlterTransactionSetProperty { + key: Some("custom_key".to_string()), + value: Some("initial_value".to_string()), + mode: None, + })), + unset_property_action: None, + }], + ..Default::default() + }) + .await + .unwrap(); + + // Now try to set the same property again with Fail mode, which must + // exercise the mode='Fail' branch (not the reserved-key guard). + let result = namespace + .alter_transaction(AlterTransactionRequest { + id: Some(vec!["users".to_string(), txn_id.clone()]), + actions: vec![AlterTransactionAction { + set_status_action: None, + set_property_action: Some(Box::new(AlterTransactionSetProperty { + key: Some("custom_key".to_string()), + value: Some("new_value".to_string()), + mode: Some("Fail".to_string()), + })), + unset_property_action: None, + }], + ..Default::default() + }) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_alter_transaction_unset_property() { + use lance_namespace::models::{ + AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty, + AlterTransactionUnsetProperty, + }; + + let (namespace, _temp_dir) = create_test_namespace().await; + create_scalar_table(&namespace, "users").await; + let txn_id = create_scalar_index(&namespace, "users", "users_id_idx") + .await + .expect("create_scalar_index should return a transaction id"); + + // First set a custom property, then unset it + let response = namespace + .alter_transaction(AlterTransactionRequest { + id: Some(vec!["users".to_string(), txn_id.clone()]), + actions: vec![ + AlterTransactionAction { + set_status_action: None, + set_property_action: Some(Box::new(AlterTransactionSetProperty { + key: Some("temp_key".to_string()), + value: Some("temp_value".to_string()), + mode: None, + })), + unset_property_action: None, + }, + AlterTransactionAction { + set_status_action: None, + set_property_action: None, + unset_property_action: Some(Box::new(AlterTransactionUnsetProperty { + key: Some("temp_key".to_string()), + mode: None, + })), + }, + ], + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(response.status, "SUCCEEDED"); + let props = response.properties.unwrap(); + assert!(!props.contains_key("temp_key")); + } + + #[tokio::test] + async fn test_alter_transaction_invalid_status() { + use lance_namespace::models::{ + AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus, + }; + + let (namespace, _temp_dir) = create_test_namespace().await; + create_scalar_table(&namespace, "users").await; + let txn_id = create_scalar_index(&namespace, "users", "users_id_idx") + .await + .expect("create_scalar_index should return a transaction id"); + + let result = namespace + .alter_transaction(AlterTransactionRequest { + id: Some(vec!["users".to_string(), txn_id.clone()]), + actions: vec![AlterTransactionAction { + set_status_action: Some(Box::new(AlterTransactionSetStatus { + status: Some("InvalidStatus".to_string()), + })), + set_property_action: None, + unset_property_action: None, + }], + ..Default::default() + }) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_alter_transaction_not_found() { + use lance_namespace::models::{ + AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus, + }; + + let (namespace, _temp_dir) = create_test_namespace().await; + create_scalar_table(&namespace, "users").await; + + // Try to alter a non-existent transaction + let result = namespace + .alter_transaction(AlterTransactionRequest { + id: Some(vec!["users".to_string(), "non_existent_txn".to_string()]), + actions: vec![AlterTransactionAction { + set_status_action: Some(Box::new(AlterTransactionSetStatus { + status: Some("Canceled".to_string()), + })), + set_property_action: None, + unset_property_action: None, + }], + ..Default::default() + }) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_alter_transaction_missing_id() { + use lance_namespace::models::{ + AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus, + }; + + let (namespace, _temp_dir) = create_test_namespace().await; + + // Try with missing id + let result = namespace + .alter_transaction(AlterTransactionRequest { + id: None, + actions: vec![AlterTransactionAction { + set_status_action: Some(Box::new(AlterTransactionSetStatus { + status: Some("Canceled".to_string()), + })), + set_property_action: None, + unset_property_action: None, + }], + ..Default::default() + }) + .await; + assert!(result.is_err()); + + // Try with insufficient id parts + let result = namespace + .alter_transaction(AlterTransactionRequest { + id: Some(vec!["users".to_string()]), + actions: vec![AlterTransactionAction { + set_status_action: Some(Box::new(AlterTransactionSetStatus { + status: Some("Canceled".to_string()), + })), + set_property_action: None, + unset_property_action: None, + }], + ..Default::default() + }) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_alter_transaction_persists_changes() { + use lance_namespace::models::{ + AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty, + AlterTransactionSetStatus, DescribeTransactionRequest, + }; + + let (namespace, _temp_dir) = create_test_namespace().await; + create_scalar_table(&namespace, "users").await; + let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await; + + let txn_id = transaction_id.expect("scalar index should produce a transaction id"); + + // Alter status and set a custom property. + namespace + .alter_transaction(AlterTransactionRequest { + id: Some(vec!["users".to_string(), txn_id.clone()]), + actions: vec![ + AlterTransactionAction { + set_status_action: Some(Box::new(AlterTransactionSetStatus { + status: Some("Canceled".to_string()), + })), + set_property_action: None, + unset_property_action: None, + }, + AlterTransactionAction { + set_status_action: None, + set_property_action: Some(Box::new(AlterTransactionSetProperty { + key: Some("owner".to_string()), + value: Some("alice".to_string()), + mode: None, + })), + unset_property_action: None, + }, + ], + ..Default::default() + }) + .await + .unwrap(); + + // The changes must survive across a fresh describe_transaction call, + // proving the alteration was persisted to the transaction file. + let describe_resp = namespace + .describe_transaction(DescribeTransactionRequest { + id: Some(vec!["users".to_string(), txn_id.clone()]), + ..Default::default() + }) + .await + .unwrap(); + let props = describe_resp.properties.expect("properties should be set"); + assert_eq!(props.get("owner"), Some(&"alice".to_string())); + // The internal `_status` marker should not leak into the response but + // must be present on disk so subsequent alter_transaction calls can + // observe the previously set status. + assert!(!props.contains_key("_status")); + + let follow_up = namespace + .alter_transaction(AlterTransactionRequest { + id: Some(vec!["users".to_string(), txn_id.clone()]), + actions: vec![], + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(follow_up.status, "Canceled"); + let follow_up_props = follow_up.properties.unwrap(); + assert_eq!(follow_up_props.get("owner"), Some(&"alice".to_string())); + } } diff --git a/rust/lance-namespace-impls/src/dir/manifest.rs b/rust/lance-namespace-impls/src/dir/manifest.rs index c05db91ad56..ab916821d29 100644 --- a/rust/lance-namespace-impls/src/dir/manifest.rs +++ b/rust/lance-namespace-impls/src/dir/manifest.rs @@ -852,7 +852,60 @@ impl ManifestNamespace { Self::ensure_manifest_table_up_to_date(&root, &storage_options, session.clone()) .await?; - Ok(Self { + Ok(Self::new( + root, + storage_options, + session, + object_store, + base_path, + manifest_dataset, + dir_listing_enabled, + inline_optimization_enabled, + commit_retries, + )) + } + + /// Open an existing manifest dataset without creating or migrating it. + #[allow(clippy::too_many_arguments)] + pub async fn open_from_directory( + root: String, + storage_options: Option>, + session: Option>, + object_store: Arc, + base_path: Path, + dir_listing_enabled: bool, + inline_optimization_enabled: bool, + commit_retries: Option, + ) -> Result { + let manifest_dataset = + Self::open_manifest_table(&root, &storage_options, session.clone()).await?; + + Ok(Self::new( + root, + storage_options, + session, + object_store, + base_path, + manifest_dataset, + dir_listing_enabled, + inline_optimization_enabled, + commit_retries, + )) + } + + #[allow(clippy::too_many_arguments)] + fn new( + root: String, + storage_options: Option>, + session: Option>, + object_store: Arc, + base_path: Path, + manifest_dataset: DatasetConsistencyWrapper, + dir_listing_enabled: bool, + inline_optimization_enabled: bool, + commit_retries: Option, + ) -> Self { + Self { root, storage_options, session, @@ -863,7 +916,7 @@ impl ManifestNamespace { inline_optimization_enabled, commit_retries, manifest_mutation_lock: Arc::new(Mutex::new(())), - }) + } } /// Build object ID from namespace path and name @@ -1179,10 +1232,21 @@ impl ManifestNamespace { index_uuid: Uuid, ) -> Result { let index_store = LanceIndexStore::from_dataset_for_new(dataset, &index_uuid)?; - let plugin = registry.get_plugin_by_name(&input.params.index_type)?; - let training_request = plugin + let trainer = registry + .get_plugin_by_name(&input.params.index_type)? + .basic_trainer() + .ok_or_else(|| { + lance_core::Error::invalid_input_source( + format!( + "The '{}' index type does not support basic training, please refer to the index's documentation for more details on how to create this index.", + input.params.index_type + ) + .into(), + ) + })?; + let training_request = trainer .new_training_request(input.params.params.as_deref().unwrap_or("{}"), &input.field)?; - let created_index = plugin + let created_index = trainer .train_index( input.stream, &index_store, @@ -2400,6 +2464,37 @@ impl ManifestNamespace { Ok(found_result) } + /// Load an existing manifest dataset without creating or migrating it. + async fn open_manifest_table( + root: &str, + storage_options: &Option>, + session: Option>, + ) -> Result { + let manifest_path = format!("{}/{}", root, MANIFEST_TABLE_NAME); + log::debug!("Attempting to load manifest from {}", manifest_path); + let store_options = ObjectStoreParams { + storage_options_accessor: storage_options.as_ref().map(|opts| { + Arc::new( + lance_io::object_store::StorageOptionsAccessor::with_static_options( + opts.clone(), + ), + ) + }), + ..Default::default() + }; + let read_params = ReadParams { + session, + store_options: Some(store_options), + ..Default::default() + }; + let dataset = DatasetBuilder::from_uri(&manifest_path) + .with_read_params(read_params) + .load() + .await?; + ensure_readable(dataset.metadata())?; + Ok(DatasetConsistencyWrapper::new(dataset)) + } + /// Create or load the manifest dataset, ensuring it has the latest schema setup. /// /// This function will: @@ -2432,129 +2527,135 @@ impl ManifestNamespace { .with_read_params(read_params) .load() .await; - if let Ok(mut dataset) = dataset_result { - // Reject a manifest written with a reader feature flag this build - // does not understand before touching it. - ensure_readable(dataset.metadata())?; - - // Check if the object_id field has primary key metadata, migrate if not - let needs_pk_migration = dataset - .schema() - .field("object_id") - .map(|f| { - !f.metadata - .contains_key(LANCE_UNENFORCED_PRIMARY_KEY_POSITION) - }) - .unwrap_or(false); - - if needs_pk_migration { - // This legacy migration writes to the manifest, so confirm this - // build is allowed to write the current format first. - ensure_writable(dataset.metadata())?; - log::info!("Migrating __manifest table to add primary key metadata on object_id"); - dataset - .update_field_metadata() - .update("object_id", [(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "0")]) - .map_err(|e| { - lance_core::Error::from(NamespaceError::Internal { - message: format!( - "Failed to find object_id field for migration: {:?}", - e - ), - }) - })? - .await - .map_err(|e| { - lance_core::Error::from(NamespaceError::Internal { - message: format!("Failed to migrate primary key metadata: {:?}", e), - }) - })?; - } - - Ok(DatasetConsistencyWrapper::new(dataset)) - } else { - log::info!("Creating new manifest table at {}", manifest_path); - let schema = Self::manifest_schema(); - let empty_batch = RecordBatch::new_empty(schema.clone()); - let reader = RecordBatchIterator::new(vec![Ok(empty_batch)], schema.clone()); - - let store_params = ObjectStoreParams { - storage_options_accessor: storage_options.as_ref().map(|opts| { - Arc::new( - lance_io::object_store::StorageOptionsAccessor::with_static_options( - opts.clone(), - ), - ) - }), - ..Default::default() - }; - let write_params = WriteParams { - session: session.clone(), - store_params: Some(store_params), - ..Default::default() - }; - - let dataset = - Dataset::write(Box::new(reader), &manifest_path, Some(write_params)).await; + match dataset_result { + Ok(mut dataset) => { + // Reject a manifest written with a reader feature flag this build + // does not understand before touching it. + ensure_readable(dataset.metadata())?; + + // Check if the object_id field has primary key metadata, migrate if not + let needs_pk_migration = dataset + .schema() + .field("object_id") + .map(|f| { + !f.metadata + .contains_key(LANCE_UNENFORCED_PRIMARY_KEY_POSITION) + }) + .unwrap_or(false); - // Handle race condition where another process created the manifest concurrently - match dataset { - Ok(dataset) => { + if needs_pk_migration { + // This legacy migration writes to the manifest, so confirm this + // build is allowed to write the current format first. + ensure_writable(dataset.metadata())?; log::info!( - "Successfully created manifest table at {}, version={}, uri={}", - manifest_path, - dataset.version().version, - dataset.uri() + "Migrating __manifest table to add primary key metadata on object_id" ); - Ok(DatasetConsistencyWrapper::new(dataset)) - } - Err(ref e) - if matches!( - e, - LanceError::DatasetAlreadyExists { .. } - | LanceError::CommitConflict { .. } - | LanceError::IncompatibleTransaction { .. } - | LanceError::RetryableCommitConflict { .. } - ) => - { - // Another process created the manifest concurrently, try to load it - log::info!( - "Manifest table was created by another process, loading it: {}", - manifest_path - ); - let recovery_store_options = ObjectStoreParams { - storage_options_accessor: storage_options.as_ref().map(|opts| { - Arc::new( - lance_io::object_store::StorageOptionsAccessor::with_static_options( - opts.clone(), - ), - ) - }), - ..Default::default() - }; - let recovery_read_params = ReadParams { - session, - store_options: Some(recovery_store_options), - ..Default::default() - }; - let dataset = DatasetBuilder::from_uri(&manifest_path) - .with_read_params(recovery_read_params) - .load() - .await + dataset + .update_field_metadata() + .update("object_id", [(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "0")]) .map_err(|e| { lance_core::Error::from(NamespaceError::Internal { message: format!( - "Failed to load manifest dataset after creation conflict: {}", + "Failed to find object_id field for migration: {:?}", e ), }) + })? + .await + .map_err(|e| { + lance_core::Error::from(NamespaceError::Internal { + message: format!("Failed to migrate primary key metadata: {:?}", e), + }) })?; - Ok(DatasetConsistencyWrapper::new(dataset)) } - Err(e) => Err(lance_core::Error::from(NamespaceError::Internal { - message: format!("Failed to create manifest dataset: {:?}", e), - })), + + Ok(DatasetConsistencyWrapper::new(dataset)) + } + Err(err) if Self::is_not_found_load_error(&err) => { + log::info!("Creating new manifest table at {}", manifest_path); + let schema = Self::manifest_schema(); + let empty_batch = RecordBatch::new_empty(schema.clone()); + let reader = RecordBatchIterator::new(vec![Ok(empty_batch)], schema.clone()); + + let store_params = ObjectStoreParams { + storage_options_accessor: storage_options.as_ref().map(|opts| { + Arc::new( + lance_io::object_store::StorageOptionsAccessor::with_static_options( + opts.clone(), + ), + ) + }), + ..Default::default() + }; + let write_params = WriteParams { + session: session.clone(), + store_params: Some(store_params), + ..Default::default() + }; + + let dataset = + Dataset::write(Box::new(reader), &manifest_path, Some(write_params)).await; + + // Handle race condition where another process created the manifest concurrently + match dataset { + Ok(dataset) => { + log::info!( + "Successfully created manifest table at {}, version={}, uri={}", + manifest_path, + dataset.version().version, + dataset.uri() + ); + Ok(DatasetConsistencyWrapper::new(dataset)) + } + Err(ref e) + if matches!( + e, + LanceError::DatasetAlreadyExists { .. } + | LanceError::CommitConflict { .. } + | LanceError::IncompatibleTransaction { .. } + | LanceError::RetryableCommitConflict { .. } + ) => + { + // Another process created the manifest concurrently, try to load it + log::info!( + "Manifest table was created by another process, loading it: {}", + manifest_path + ); + let recovery_store_options = ObjectStoreParams { + storage_options_accessor: storage_options.as_ref().map(|opts| { + Arc::new( + lance_io::object_store::StorageOptionsAccessor::with_static_options( + opts.clone(), + ), + ) + }), + ..Default::default() + }; + let recovery_read_params = ReadParams { + session, + store_options: Some(recovery_store_options), + ..Default::default() + }; + let dataset = DatasetBuilder::from_uri(&manifest_path) + .with_read_params(recovery_read_params) + .load() + .await + .map_err(|e| { + lance_core::Error::from(NamespaceError::Internal { + message: format!( + "Failed to load manifest dataset after creation conflict: {}", + e + ), + }) + })?; + Ok(DatasetConsistencyWrapper::new(dataset)) + } + Err(e) => Err(lance_core::Error::from(NamespaceError::Internal { + message: format!("Failed to create manifest dataset: {:?}", e), + })), + } } + Err(err) => Err(err), } } diff --git a/rust/lance-table/benches/manifest_intern.rs b/rust/lance-table/benches/manifest_intern.rs index 78b7e352207..81bd57c1a22 100644 --- a/rust/lance-table/benches/manifest_intern.rs +++ b/rust/lance-table/benches/manifest_intern.rs @@ -59,6 +59,7 @@ fn make_uniform_pb_fragments(n: u64, num_fields: usize) -> Vec file_size_bytes: 0, base_id: None, }], + overlays: vec![], deletion_file: None, row_id_sequence: None, physical_rows: 1000, @@ -135,6 +136,7 @@ fn make_diverse_pb_fragments( file_size_bytes: 0, base_id: None, }], + overlays: vec![], deletion_file: None, row_id_sequence: None, physical_rows: 1000, diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index 096f0da79e5..41b8e415f8e 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -20,8 +20,22 @@ pub const FLAG_TABLE_CONFIG: u64 = 8; pub const FLAG_BASE_PATHS: u64 = 16; /// Disable writing transaction file under _transaction/, this flag is set when we only want to write inline transaction in manifest pub const FLAG_DISABLE_TRANSACTION_FILE: u64 = 32; +/// Fragments contain data overlay files, which supply new values for a subset of +/// cells without rewriting base data files. A reader that does not understand +/// overlays must refuse the dataset, since ignoring an overlay would silently +/// return stale base values. +/// +/// Data overlay files are not yet a released feature: in release builds this flag +/// is treated as unknown (so a release reader/writer refuses an overlay dataset) +/// unless [`ENABLE_UNSTABLE_DATA_OVERLAY_FILES_ENV`] is set, which lets benchmarks opt in. +/// Debug builds always understand it so tests exercise the path. +pub const FLAG_UNSTABLE_DATA_OVERLAY_FILES: u64 = 64; /// The first bit that is unknown as a feature flag -pub const FLAG_UNKNOWN: u64 = 64; +pub const FLAG_UNKNOWN: u64 = 128; + +/// Environment variable that opts a release build into reading and writing data +/// overlay files before the feature is generally released. +pub const ENABLE_UNSTABLE_DATA_OVERLAY_FILES_ENV: &str = "LANCE_ENABLE_UNSTABLE_DATA_OVERLAY_FILES"; /// Set the reader and writer feature flags in the manifest based on the contents of the manifest. pub fn apply_feature_flags( @@ -71,18 +85,62 @@ pub fn apply_feature_flags( manifest.writer_feature_flags |= FLAG_BASE_PATHS; } + // Overlay files change cell values on read, so a reader that ignores them + // would return stale base values. Both readers and writers must understand + // them. + let has_overlays = manifest + .fragments + .iter() + .any(|frag| !frag.overlays.is_empty()); + if has_overlays { + manifest.reader_feature_flags |= FLAG_UNSTABLE_DATA_OVERLAY_FILES; + manifest.writer_feature_flags |= FLAG_UNSTABLE_DATA_OVERLAY_FILES; + } + if disable_transaction_file { manifest.writer_feature_flags |= FLAG_DISABLE_TRANSACTION_FILE; } Ok(()) } +/// Whether this build understands data overlay files: always in debug builds, +/// and in release builds only when [`ENABLE_UNSTABLE_DATA_OVERLAY_FILES_ENV`] is set. +fn data_overlay_files_enabled() -> bool { + cfg!(debug_assertions) || std::env::var_os(ENABLE_UNSTABLE_DATA_OVERLAY_FILES_ENV).is_some() +} + +/// Clear `flag` from `flags` when its gating feature is not enabled in this +/// build; leave it set otherwise. One call per unstable flag, so support for +/// several unstable features chains cleanly. +fn mark_supported(flags: &mut u64, flag: u64, feature_enabled: bool) { + if !feature_enabled { + *flags &= !flag; + } +} + +/// The feature-flag bits this build understands, given whether overlay support +/// is enabled. Split out from [`supported_flags`] so the policy is testable +/// without toggling the build profile or environment. +fn supported_flags_when(overlay_enabled: bool) -> u64 { + let mut supported = FLAG_UNKNOWN - 1; + mark_supported( + &mut supported, + FLAG_UNSTABLE_DATA_OVERLAY_FILES, + overlay_enabled, + ); + supported +} + +fn supported_flags() -> u64 { + supported_flags_when(data_overlay_files_enabled()) +} + pub fn can_read_dataset(reader_flags: u64) -> bool { - reader_flags < FLAG_UNKNOWN + reader_flags & !supported_flags() == 0 } pub fn can_write_dataset(writer_flags: u64) -> bool { - writer_flags < FLAG_UNKNOWN + writer_flags & !supported_flags() == 0 } pub fn has_deprecated_v2_feature_flag(writer_flags: u64) -> bool { @@ -103,6 +161,13 @@ mod tests { assert!(can_read_dataset(super::FLAG_TABLE_CONFIG)); assert!(can_read_dataset(super::FLAG_BASE_PATHS)); assert!(can_read_dataset(super::FLAG_DISABLE_TRANSACTION_FILE)); + // Overlay support is gated on the build profile / env opt-in, so the + // flag is readable exactly when overlays are enabled (see + // test_data_overlay_flag_release_gating for the full policy). + assert_eq!( + can_read_dataset(super::FLAG_UNSTABLE_DATA_OVERLAY_FILES), + data_overlay_files_enabled() + ); assert!(can_read_dataset( super::FLAG_DELETION_FILES | super::FLAG_STABLE_ROW_IDS @@ -111,6 +176,58 @@ mod tests { assert!(!can_read_dataset(super::FLAG_UNKNOWN)); } + #[test] + fn test_data_overlay_flag_release_gating() { + // Release default (overlays disabled): the overlay flag is treated as + // unknown so the dataset is refused, while other known flags still pass. + let supported = supported_flags_when(false); + assert_eq!(supported & FLAG_UNSTABLE_DATA_OVERLAY_FILES, 0); + assert_eq!(FLAG_DELETION_FILES & !supported, 0); + assert_ne!(FLAG_UNSTABLE_DATA_OVERLAY_FILES & !supported, 0); + // Enabled (debug or env opt-in): the overlay flag is understood. + let supported = supported_flags_when(true); + assert_eq!(FLAG_UNSTABLE_DATA_OVERLAY_FILES & !supported, 0); + } + + #[test] + fn test_apply_feature_flags_sets_overlay_flag() { + use crate::format::overlay::{DataOverlayFile, OverlayCoverage}; + use crate::format::{DataFile, DataStorageFormat, Fragment}; + use arrow_schema::{Field as ArrowField, Schema as ArrowSchema}; + use lance_core::datatypes::Schema; + use roaring::RoaringBitmap; + use std::collections::HashMap; + use std::sync::Arc; + + let arrow_schema = ArrowSchema::new(vec![ArrowField::new( + "id", + arrow_schema::DataType::Int64, + false, + )]); + let schema = Schema::try_from(&arrow_schema).unwrap(); + let mut fragment = Fragment::new(0); + fragment.overlays = vec![DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![0], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 1, + }]; + let mut manifest = Manifest::new( + schema, + Arc::new(vec![fragment]), + DataStorageFormat::default(), + HashMap::new(), + ); + apply_feature_flags(&mut manifest, false, false).unwrap(); + assert_ne!( + manifest.reader_feature_flags & FLAG_UNSTABLE_DATA_OVERLAY_FILES, + 0 + ); + assert_ne!( + manifest.writer_feature_flags & FLAG_UNSTABLE_DATA_OVERLAY_FILES, + 0 + ); + } + #[test] fn test_write_check() { assert!(can_write_dataset(0)); @@ -120,6 +237,13 @@ mod tests { assert!(can_write_dataset(super::FLAG_TABLE_CONFIG)); assert!(can_write_dataset(super::FLAG_BASE_PATHS)); assert!(can_write_dataset(super::FLAG_DISABLE_TRANSACTION_FILE)); + // Overlay support is gated on the build profile / env opt-in, so the + // flag is writable exactly when overlays are enabled (see + // test_data_overlay_flag_release_gating for the full policy). + assert_eq!( + can_write_dataset(super::FLAG_UNSTABLE_DATA_OVERLAY_FILES), + data_overlay_files_enabled() + ); assert!(can_write_dataset( super::FLAG_DELETION_FILES | super::FLAG_STABLE_ROW_IDS diff --git a/rust/lance-table/src/format.rs b/rust/lance-table/src/format.rs index 842c76f1e58..5a5db7919d3 100644 --- a/rust/lance-table/src/format.rs +++ b/rust/lance-table/src/format.rs @@ -7,6 +7,7 @@ use uuid::Uuid; mod fragment; mod index; mod manifest; +pub mod overlay; mod transaction; pub use crate::rowids::version::{ diff --git a/rust/lance-table/src/format/fragment.rs b/rust/lance-table/src/format/fragment.rs index 431e466dbd4..4929f54b7ff 100644 --- a/rust/lance-table/src/format/fragment.rs +++ b/rust/lance-table/src/format/fragment.rs @@ -13,6 +13,7 @@ use lance_io::utils::CachedFileSize; use object_store::path::Path; use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use super::overlay::{DataOverlayFile, sort_overlays_newest_last}; use crate::format::pb; use crate::rowids::version::{ @@ -375,6 +376,15 @@ impl DataFileFieldInterner { .into_iter() .map(|f| self.intern_data_file(f)) .collect::>()?, + overlays: { + let mut overlays = p + .overlays + .into_iter() + .map(DataOverlayFile::try_from) + .collect::>>()?; + sort_overlays_newest_last(&mut overlays); + overlays + }, deletion_file: p.deletion_file.map(DeletionFile::try_from).transpose()?, row_id_meta: p.row_id_sequence.map(RowIdMeta::try_from).transpose()?, physical_rows, @@ -483,6 +493,12 @@ pub struct Fragment { /// Files within the fragment. pub files: Vec, + /// Overlay files supplying new values for a subset of cells without + /// rewriting the base data files. Order is significant: a later entry is + /// newer than an earlier one. See [`DataOverlayFile`] for resolution rules. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub overlays: Vec, + /// Optional file with deleted local row offsets. #[serde(skip_serializing_if = "Option::is_none")] pub deletion_file: Option, @@ -510,6 +526,7 @@ impl Fragment { Self { id, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -549,6 +566,7 @@ impl Fragment { Self { id, files: vec![DataFile::new_legacy(path, schema, None, None)], + overlays: vec![], deletion_file: None, physical_rows, row_id_meta: None, @@ -669,6 +687,15 @@ impl TryFrom for Fragment { .into_iter() .map(DataFile::try_from) .collect::>()?, + overlays: { + let mut overlays = p + .overlays + .into_iter() + .map(DataOverlayFile::try_from) + .collect::>>()?; + sort_overlays_newest_last(&mut overlays); + overlays + }, deletion_file: p.deletion_file.map(DeletionFile::try_from).transpose()?, row_id_meta: p.row_id_sequence.map(RowIdMeta::try_from).transpose()?, physical_rows, @@ -716,6 +743,7 @@ impl From<&Fragment> for pb::DataFragment { Self { id: f.id, files: f.files.iter().map(pb::DataFile::from).collect(), + overlays: f.overlays.iter().map(pb::DataOverlayFile::from).collect(), deletion_file, row_id_sequence, physical_rows: f.physical_rows.unwrap_or_default() as u64, @@ -728,12 +756,108 @@ impl From<&Fragment> for pb::DataFragment { #[cfg(test)] mod tests { use super::*; + use crate::format::overlay::OverlayCoverage; use arrow_schema::{ DataType, Field as ArrowField, Fields as ArrowFields, Schema as ArrowSchema, }; use object_store::path::Path; + use roaring::RoaringBitmap; use serde_json::{Value, json}; + #[test] + fn test_data_overlay_roundtrip() { + // A fragment carrying a dense overlay round-trips through protobuf and + // back, and the parsed coverage bitmap is recovered per field. + let mut bitmap = RoaringBitmap::new(); + bitmap.insert(1); + bitmap.insert(3); + + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay-0.lance", vec![3], None), + coverage: OverlayCoverage::dense(bitmap.clone()), + committed_version: 7, + }; + let mut fragment = Fragment::new(0); + fragment.files = vec![DataFile::new_legacy_from_fields( + "base.lance", + vec![1, 3], + None, + )]; + fragment.overlays = vec![overlay]; + + let proto = pb::DataFragment::from(&fragment); + assert_eq!(proto.overlays.len(), 1); + let round_tripped = Fragment::try_from(proto).unwrap(); + assert_eq!(round_tripped, fragment); + + // Dense coverage applies to every field. + let recovered = round_tripped.overlays[0].coverage_for_field(0).unwrap(); + assert_eq!(*recovered, bitmap); + assert_eq!( + *round_tripped.overlays[0].coverage_for_field(5).unwrap(), + bitmap + ); + } + + #[test] + fn test_data_overlay_sparse_per_field_coverage() { + // A sparse overlay carries one bitmap per field, recovered by position. + let name_coverage = RoaringBitmap::from_iter([2u32, 3]); + let embedding_coverage = RoaringBitmap::from_iter([1u32]); + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay-1.lance", vec![2, 4], None), + coverage: OverlayCoverage::sparse(vec![ + name_coverage.clone(), + embedding_coverage.clone(), + ]), + committed_version: 3, + }; + let mut fragment = Fragment::new(1); + fragment.overlays = vec![overlay]; + + let round_tripped = Fragment::try_from(pb::DataFragment::from(&fragment)).unwrap(); + assert_eq!( + *round_tripped.overlays[0].coverage_for_field(0).unwrap(), + name_coverage + ); + assert_eq!( + *round_tripped.overlays[0].coverage_for_field(1).unwrap(), + embedding_coverage + ); + } + + #[test] + fn test_overlays_sorted_newest_last_on_load() { + // Overlays load stable-sorted by committed_version (newest last), with + // list position preserved as the tiebreak for equal versions. + let mk = |version: u64, field: i32| DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![field], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: version, + }; + let mut fragment = Fragment::new(0); + // Written out of order: v5, v2, v2 (second), v3. + fragment.overlays = vec![mk(5, 1), mk(2, 2), mk(2, 3), mk(3, 4)]; + + let loaded = Fragment::try_from(pb::DataFragment::from(&fragment)).unwrap(); + let versions: Vec = loaded + .overlays + .iter() + .map(|o| o.committed_version) + .collect(); + assert_eq!(versions, vec![2, 2, 3, 5]); + // Stable: the two v2 overlays keep their original relative order (field 2 + // before field 3). + assert_eq!( + loaded.overlays[0].data_file.fields.as_ref(), + [2i32].as_slice() + ); + assert_eq!( + loaded.overlays[1].data_file.fields.as_ref(), + [3i32].as_slice() + ); + } + #[test] fn test_new_fragment() { let path = "foobar.lance"; diff --git a/rust/lance-table/src/format/manifest.rs b/rust/lance-table/src/format/manifest.rs index 9845061b7e4..cd0a403621f 100644 --- a/rust/lance-table/src/format/manifest.rs +++ b/rust/lance-table/src/format/manifest.rs @@ -1316,6 +1316,7 @@ mod tests { vec![0, 1, 2], None, )], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1328,6 +1329,7 @@ mod tests { DataFile::new_legacy_from_fields("path2", vec![0, 1, 43], None), DataFile::new_legacy_from_fields("path3", vec![2], None), ], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, diff --git a/rust/lance-table/src/format/overlay.rs b/rust/lance-table/src/format/overlay.rs new file mode 100644 index 00000000000..5e729411502 --- /dev/null +++ b/rust/lance-table/src/format/overlay.rs @@ -0,0 +1,443 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Data overlay files. +//! +//! An overlay file supplies new values for a subset of `(physical offset, field)` +//! cells within a fragment, without rewriting the fragment's base data files. See +//! the Data Overlay Files specification for the full rules; the invariants this +//! module relies on are: +//! +//! - **Physical-offset coverage.** Coverage bitmaps index *physical* row offsets +//! (positions in the base data files, counting deleted rows), so they are stable +//! across deletions, like deletion vectors. +//! - **Rank-based values.** The overlay's `data_file` stores one value column per +//! field, with no row-offset key column. Within a value column, a covered +//! offset's value sits at its **rank** — the 0-based count of set bits below it +//! in that field's coverage bitmap. +//! - **Dense vs. sparse coverage.** A dense overlay shares one bitmap across every +//! field ([`OverlayCoverage::Shared`]); a sparse overlay carries one bitmap per +//! field ([`OverlayCoverage::PerField`]). +//! - **Parse once.** Bitmaps are parsed from their 32-bit Roaring encoding a single +//! time when the fragment loads and held behind an `Arc`, so cloning a fragment +//! is cheap. +//! - **Newest-last ordering.** A fragment's overlays are stored newest-last and +//! stable-sorted by `committed_version` on load (see [`sort_overlays_newest_last`]), +//! with list position breaking ties for equal versions. When two overlays cover +//! the same `(offset, field)`, the higher `committed_version` wins. +//! - **Field tombstones.** When new base values are written for a field (a +//! DataReplacement, or an in-place column rewrite), any overlay value for that +//! field is stale and must stop shadowing the fresh base. The field is marked +//! obsolete in the overlay's `data_file.fields` with [`TOMBSTONE_FIELD_ID`] +//! (the same sentinel used for obsolete base columns) rather than physically +//! removed, so the overlay's other fields — and its coverage positions — stay +//! intact (see [`tombstone_overlay_fields`]). + +use std::sync::Arc; + +use lance_core::Error; +use lance_core::deepsize::DeepSizeOf; +use lance_core::error::Result; +use roaring::RoaringBitmap; +use serde::{Deserialize, Serialize}; + +use object_store::path::Path; + +use super::DataFile; +use crate::format::pb; + +/// Field-id sentinel marking a tombstoned (obsolete) field within an overlay's +/// `data_file.fields`. Matches the tombstone convention for obsolete columns in +/// base data files; a tombstoned field's values are ignored on read. +pub const TOMBSTONE_FIELD_ID: i32 = -2; + +/// Which `(physical offset, field)` cells a [`DataOverlayFile`] provides values +/// for. +/// +/// Bitmaps are parsed from their 32-bit Roaring encoding once when the fragment +/// is loaded and held behind an `Arc` so cloning a fragment is cheap; use +/// [`DataOverlayFile::coverage_for_field`] to obtain the one that applies to a +/// given field. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(into = "OverlayCoverageBytes", try_from = "OverlayCoverageBytes")] +pub enum OverlayCoverage { + /// A single bitmap that applies to every field in the overlay's + /// `data_file.fields` (a dense / rectangular overlay): every covered offset + /// has a value for every field. + Shared(Arc), + /// One bitmap per field, in the same order as the overlay's + /// `data_file.fields` (a sparse overlay): different fields may cover + /// different offset sets. + PerField(Vec>), +} + +/// Serialized form of [`OverlayCoverage`] — each bitmap as its 32-bit Roaring +/// byte encoding. The in-memory form parses these once at load. +#[derive(Debug, Clone, Serialize, Deserialize)] +enum OverlayCoverageBytes { + Shared(Vec), + PerField(Vec>), +} + +// The bytes come from a persisted overlay (the protobuf manifest or a +// serialized fragment), so a decode failure is on-disk corruption, not caller +// input. `path` locates the overlay's data file when known (empty on the serde +// path, which deserializes coverage in isolation). +fn deserialize_roaring(bytes: &[u8], path: &Path) -> Result { + RoaringBitmap::deserialize_from(bytes).map_err(|e| { + Error::corrupt_file( + path.clone(), + format!("failed to deserialize overlay coverage bitmap: {e}"), + ) + }) +} + +fn serialize_roaring(bitmap: &RoaringBitmap) -> Vec { + let mut bytes = Vec::with_capacity(bitmap.serialized_size()); + // Writing to a Vec is infallible. + bitmap.serialize_into(&mut bytes).unwrap(); + bytes +} + +impl From for OverlayCoverageBytes { + fn from(coverage: OverlayCoverage) -> Self { + match coverage { + OverlayCoverage::Shared(bitmap) => Self::Shared(serialize_roaring(&bitmap)), + OverlayCoverage::PerField(bitmaps) => { + Self::PerField(bitmaps.iter().map(|b| serialize_roaring(b)).collect()) + } + } + } +} + +impl TryFrom for OverlayCoverage { + type Error = Error; + + fn try_from(bytes: OverlayCoverageBytes) -> Result { + // Serde deserializes the coverage in isolation, so the owning data + // file's path is not available here. + let path = Path::default(); + Ok(match bytes { + OverlayCoverageBytes::Shared(b) => { + Self::Shared(Arc::new(deserialize_roaring(&b, &path)?)) + } + OverlayCoverageBytes::PerField(bs) => Self::PerField( + bs.iter() + .map(|b| deserialize_roaring(b, &path).map(Arc::new)) + .collect::>()?, + ), + }) + } +} + +impl DeepSizeOf for OverlayCoverage { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + // The same `Arc` is shared across every clone of a + // fragment, so mark each Arc's pointer and count its heap only the first + // time it is seen — otherwise walking many fragments double-counts the + // shared bitmaps. RoaringBitmap does not expose its allocation size; its + // serialized size is a cheap, close proxy for the heap it holds. + let bitmap_heap = |bitmap: &Arc, + context: &mut lance_core::deepsize::Context| { + if context.mark_seen(Arc::as_ptr(bitmap) as usize) { + std::mem::size_of::() + bitmap.serialized_size() + } else { + 0 + } + }; + match self { + Self::Shared(bitmap) => bitmap_heap(bitmap, context), + Self::PerField(bitmaps) => { + bitmaps.capacity() * std::mem::size_of::>() + + bitmaps + .iter() + .map(|b| bitmap_heap(b, context)) + .sum::() + } + } + } +} + +impl OverlayCoverage { + /// Build a dense coverage from a single bitmap shared across every field. + pub fn dense(bitmap: RoaringBitmap) -> Self { + Self::Shared(Arc::new(bitmap)) + } + + /// Build a sparse coverage from one bitmap per field. + pub fn sparse(bitmaps: Vec) -> Self { + Self::PerField(bitmaps.into_iter().map(Arc::new).collect()) + } +} + +/// An overlay file supplies new values for a subset of `(physical offset, field)` +/// cells within a fragment, without rewriting the fragment's base data files. See +/// the [module documentation](self) for the coverage, rank, and versioning rules. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)] +pub struct DataOverlayFile { + /// The data file storing the overlay's new cell values. + pub data_file: DataFile, + /// Which cells this overlay provides values for. + pub coverage: OverlayCoverage, + /// The dataset version at which this overlay became effective (the version of + /// the commit that introduced it, stamped at commit time and re-stamped on + /// retry). Higher wins when two overlays cover the same `(offset, field)`. + pub committed_version: u64, +} + +impl DataOverlayFile { + /// The parsed coverage bitmap that applies to the field stored at + /// `field_pos` within `data_file.fields`. + /// + /// For a dense overlay the same shared bitmap is returned for every field; + /// for a sparse overlay the per-field bitmap at `field_pos` is returned. The + /// bitmap is already parsed, so this is a cheap `Arc` clone. + pub fn coverage_for_field(&self, field_pos: usize) -> Result> { + match &self.coverage { + OverlayCoverage::Shared(bitmap) => Ok(bitmap.clone()), + OverlayCoverage::PerField(bitmaps) => { + bitmaps.get(field_pos).cloned().ok_or_else(|| { + Error::invalid_input(format!( + "overlay per-field coverage has {} bitmaps but field position {} was requested", + bitmaps.len(), + field_pos + )) + }) + } + } + } +} + +/// Stable-sort a fragment's overlays newest-last by `committed_version`. The +/// stable sort preserves list position as the tiebreak for equal versions, so +/// resolution can rely on the ordering without re-checking. See the [module +/// documentation](self) for the ordering invariant. +pub fn sort_overlays_newest_last(overlays: &mut [DataOverlayFile]) { + overlays.sort_by_key(|overlay| overlay.committed_version); +} + +/// Verify a fragment's overlays are stored newest-last (non-decreasing +/// `committed_version`), the ordering invariant readers rely on for +/// resolution. Returns an error identifying the first out-of-order pair. +/// +/// [`sort_overlays_newest_last`] normalizes on load; this is the write-side +/// guard that rejects any commit path that assembled overlays out of order. See +/// the [module documentation](self) for the ordering invariant. +pub fn verify_overlays_newest_last(overlays: &[DataOverlayFile]) -> Result<()> { + for pair in overlays.windows(2) { + if pair[0].committed_version > pair[1].committed_version { + return Err(Error::invalid_input(format!( + "overlay files must be stored newest-last, but committed_version {} precedes {}", + pair[0].committed_version, pair[1].committed_version + ))); + } + } + Ok(()) +} + +/// Tombstone `fields` across a fragment's `overlays`, dropping any overlay left +/// with no live fields. +/// +/// Called when new base values are written for those fields (a DataReplacement, +/// or an in-place column rewrite): the stale overlay values must stop shadowing +/// the fresh base. Each matching field id is replaced with [`TOMBSTONE_FIELD_ID`] +/// in place, preserving the overlay's remaining fields and its coverage positions +/// (a per-field coverage bitmap stays aligned with `data_file.fields`). An overlay +/// whose fields are now all tombstoned is removed entirely. See the [module +/// documentation](self) for the tombstone invariant. +pub fn tombstone_overlay_fields(overlays: &mut Vec, fields: &[u32]) { + for overlay in overlays.iter_mut() { + let tombstoned: Vec = overlay + .data_file + .fields + .iter() + .map(|&field| { + if field >= 0 && fields.contains(&(field as u32)) { + TOMBSTONE_FIELD_ID + } else { + field + } + }) + .collect(); + overlay.data_file.fields = tombstoned.into(); + } + overlays.retain(|overlay| { + overlay + .data_file + .fields + .iter() + .any(|&field| field != TOMBSTONE_FIELD_ID) + }); +} + +impl From<&DataOverlayFile> for pb::DataOverlayFile { + fn from(overlay: &DataOverlayFile) -> Self { + let coverage = match &overlay.coverage { + OverlayCoverage::Shared(bitmap) => { + pb::data_overlay_file::Coverage::SharedOffsetBitmap(serialize_roaring(bitmap)) + } + OverlayCoverage::PerField(bitmaps) => { + pb::data_overlay_file::Coverage::FieldCoverage(pb::FieldCoverage { + offset_bitmaps: bitmaps.iter().map(|b| serialize_roaring(b)).collect(), + }) + } + }; + Self { + data_file: Some(pb::DataFile::from(&overlay.data_file)), + coverage: Some(coverage), + committed_version: overlay.committed_version, + } + } +} + +impl TryFrom for DataOverlayFile { + type Error = Error; + + fn try_from(proto: pb::DataOverlayFile) -> Result { + let data_file = proto + .data_file + .ok_or_else(|| Error::invalid_input("DataOverlayFile is missing its data_file"))?; + let path = Path::from(data_file.path.as_str()); + let coverage = match proto.coverage { + Some(pb::data_overlay_file::Coverage::SharedOffsetBitmap(bytes)) => { + OverlayCoverage::Shared(Arc::new(deserialize_roaring(&bytes, &path)?)) + } + Some(pb::data_overlay_file::Coverage::FieldCoverage(fc)) => OverlayCoverage::PerField( + fc.offset_bitmaps + .iter() + .map(|b| deserialize_roaring(b, &path).map(Arc::new)) + .collect::>()?, + ), + None => { + return Err(Error::invalid_input( + "DataOverlayFile is missing its coverage", + )); + } + }; + Ok(Self { + data_file: DataFile::try_from(data_file)?, + coverage, + committed_version: proto.committed_version, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_data_overlay_missing_fields_error() { + // A DataOverlayFile proto missing its coverage or data_file is rejected. + let no_coverage = pb::DataOverlayFile { + data_file: Some(pb::DataFile::from(&DataFile::new_legacy_from_fields( + "overlay.lance", + vec![3], + None, + ))), + coverage: None, + committed_version: 1, + }; + let err = DataOverlayFile::try_from(no_coverage).unwrap_err(); + assert!(err.to_string().contains("missing its coverage"), "{err}"); + + let no_data_file = pb::DataOverlayFile { + data_file: None, + coverage: Some(pb::data_overlay_file::Coverage::SharedOffsetBitmap( + serialize_roaring(&RoaringBitmap::from_iter([0u32])), + )), + committed_version: 1, + }; + let err = DataOverlayFile::try_from(no_data_file).unwrap_err(); + assert!(err.to_string().contains("missing its data_file"), "{err}"); + } + + #[test] + fn test_overlay_coverage_serde_json_roundtrip() { + // The custom serde impl round-trips through JSON for dense/sparse, + // including empty bitmaps and a zero-bitmap sparse coverage. + for coverage in [ + OverlayCoverage::dense(RoaringBitmap::from_iter([1u32, 5, 100])), + OverlayCoverage::dense(RoaringBitmap::new()), + OverlayCoverage::sparse(vec![ + RoaringBitmap::from_iter([2u32, 3]), + RoaringBitmap::new(), + ]), + OverlayCoverage::sparse(vec![]), + ] { + let json = serde_json::to_string(&coverage).unwrap(); + let back: OverlayCoverage = serde_json::from_str(&json).unwrap(); + assert_eq!(back, coverage); + } + } + + #[test] + fn test_tombstone_overlay_fields() { + // An overlay covering fields [3, 5]: replacing field 5 tombstones just + // field 5's slot and keeps field 3. An overlay covering only field 5 is + // dropped entirely. An overlay touching no replaced field is untouched. + let mut overlays = vec![ + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("a.lance", vec![3, 5], None), + coverage: OverlayCoverage::sparse(vec![ + RoaringBitmap::from_iter([0u32]), + RoaringBitmap::from_iter([1u32]), + ]), + committed_version: 1, + }, + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("b.lance", vec![5], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 1, + }, + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("c.lance", vec![7], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 1, + }, + ]; + + tombstone_overlay_fields(&mut overlays, &[5]); + + // The single-field overlay on field 5 is gone; the others remain. + assert_eq!(overlays.len(), 2); + // Field 3 preserved, field 5 tombstoned in place (coverage stays aligned). + assert_eq!( + overlays[0].data_file.fields.as_ref(), + &[3, TOMBSTONE_FIELD_ID] + ); + // The untouched overlay keeps its field. + assert_eq!(overlays[1].data_file.fields.as_ref(), &[7]); + } + + #[test] + fn test_verify_overlays_newest_last() { + let mk = |version: u64| DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![3], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: version, + }; + // Non-decreasing (including equal versions) is accepted. + assert!(verify_overlays_newest_last(&[]).is_ok()); + assert!(verify_overlays_newest_last(&[mk(1), mk(2), mk(2), mk(5)]).is_ok()); + // A newer version before an older one is rejected. + let err = verify_overlays_newest_last(&[mk(2), mk(1)]).unwrap_err(); + assert!(err.to_string().contains("newest-last"), "{err}"); + } + + #[test] + fn test_coverage_for_field_out_of_bounds() { + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![2, 4], None), + coverage: OverlayCoverage::sparse(vec![ + RoaringBitmap::from_iter([1u32]), + RoaringBitmap::from_iter([2u32]), + ]), + committed_version: 1, + }; + assert!(overlay.coverage_for_field(0).is_ok()); + assert!(overlay.coverage_for_field(1).is_ok()); + let err = overlay.coverage_for_field(5).unwrap_err(); + assert!(err.to_string().contains("field position"), "{err}"); + } +} diff --git a/rust/lance-table/src/io/manifest.rs b/rust/lance-table/src/io/manifest.rs index 4b918177176..6a6bfd2724c 100644 --- a/rust/lance-table/src/io/manifest.rs +++ b/rust/lance-table/src/io/manifest.rs @@ -115,14 +115,18 @@ pub async fn read_manifest_indexes( manifest: &Manifest, ) -> Result> { if let Some(pos) = manifest.index_section.as_ref() { - let reader = if let Some(size) = location.size { - object_store - .open_with_size(&location.path, size as usize) - .await? - } else { - object_store.open(&location.path).await? + let result = read_index_section(object_store, &location.path, location.size, *pos).await; + // A stale cached size makes the index offset fall outside the sized view, + // so the read fails as "file size is too small". Retry once with the true + // size; surface any other error unchanged. + let section = match result { + Err(e) + if location.size.is_some() && e.to_string().contains("file size is too small") => + { + read_index_section(object_store, &location.path, None, *pos).await? + } + other => other?, }; - let section: pb::IndexSection = read_message(reader.as_ref(), *pos).await?; let indices = section .indices @@ -135,6 +139,22 @@ pub async fn read_manifest_indexes( } } +/// Read the index section message at `pos`, opening the manifest with a known +/// size when one is provided. +async fn read_index_section( + object_store: &ObjectStore, + path: &Path, + size: Option, + pos: usize, +) -> Result { + let reader = if let Some(size) = size { + object_store.open_with_size(path, size as usize).await? + } else { + object_store.open(path).await? + }; + read_message(reader.as_ref(), pos).await +} + async fn do_write_manifest( writer: &mut dyn Writer, manifest: &mut Manifest, diff --git a/rust/lance-table/src/rowids.rs b/rust/lance-table/src/rowids.rs index ab5dac72b48..6cdc907cda5 100644 --- a/rust/lance-table/src/rowids.rs +++ b/rust/lance-table/src/rowids.rs @@ -355,11 +355,7 @@ impl RowIdSequence { while (index - rows_passed) >= cur_seg_len { rows_passed += cur_seg_len; cur_seg = seg_iter.next(); - if let Some(cur_seg) = cur_seg { - cur_seg_len = cur_seg.len(); - } else { - return None; - } + cur_seg_len = cur_seg?.len(); } Some(cur_seg.unwrap().get(index - rows_passed).unwrap()) diff --git a/rust/lance-table/src/rowids/index.rs b/rust/lance-table/src/rowids/index.rs index 66720ed1f25..4fed0e651c5 100644 --- a/rust/lance-table/src/rowids/index.rs +++ b/rust/lance-table/src/rowids/index.rs @@ -5,10 +5,10 @@ use std::ops::RangeInclusive; use std::sync::Arc; use super::{RowIdSequence, U64Segment}; -use lance_core::Result; use lance_core::deepsize::DeepSizeOf; use lance_core::utils::address::RowAddress; use lance_core::utils::deletion::DeletionVector; +use lance_core::{Error, Result}; use rangemap::RangeInclusiveMap; /// An index of row ids @@ -46,18 +46,10 @@ impl RowIdIndex { RawIndexChunk::NonOverlapping(chunk) => { final_chunks.push(chunk); } - RawIndexChunk::Overlapping(range, overlapping_chunks) => { - debug_assert_eq!( - range.end() - range.start() + 1, - overlapping_chunks - .iter() - .map(|(_, (seq, _))| seq.len() as u64) - .sum::(), - "Wrong range for {:?}, chunks: {:?}", - range, - overlapping_chunks, - ); - // Merge overlapping chunks. + RawIndexChunk::Overlapping(_range, overlapping_chunks) => { + // Intersecting row-id ranges don't imply intersecting id sets; + // sparse ids and deletion holes leave the union short of the span. + // The real invariant (no id in two fragments) is checked in the merge. let merged_chunk = merge_overlapping_chunks(overlapping_chunks)?; final_chunks.push(merged_chunk); } @@ -357,6 +349,14 @@ fn merge_overlapping_chunks(overlapping_chunks: Vec) -> Result a hole in the span. + deletion_vector: Arc::new(DeletionVector::from_iter(vec![2])), + }, + ]; + + let index = RowIdIndex::new(&fragment_indices).unwrap(); + + assert_eq!(index.get(0), Some(RowAddress::new_from_parts(20, 0))); + assert_eq!(index.get(1), Some(RowAddress::new_from_parts(10, 0))); + assert_eq!(index.get(2), Some(RowAddress::new_from_parts(20, 1))); + assert_eq!(index.get(3), Some(RowAddress::new_from_parts(10, 1))); + assert_eq!(index.get(4), None); + // Surviving ids keep their original offsets (the hole is not compacted). + assert_eq!(index.get(6), Some(RowAddress::new_from_parts(20, 3))); + assert_eq!(index.get(8), Some(RowAddress::new_from_parts(20, 4))); + assert_eq!(index.get(9), Some(RowAddress::new_from_parts(10, 4))); + } + #[test] fn test_index_with_deletion_vector() { let deletion_vector = DeletionVector::from_iter(vec![2, 3]); diff --git a/rust/lance-table/src/system_index/mem_wal.rs b/rust/lance-table/src/system_index/mem_wal.rs index 3bf279df062..1d82fd9e44f 100644 --- a/rust/lance-table/src/system_index/mem_wal.rs +++ b/rust/lance-table/src/system_index/mem_wal.rs @@ -141,6 +141,39 @@ impl TryFrom for IndexCatchupProgress { } } +/// Lifecycle status of a WAL shard, persisted in [`ShardManifest`]. +/// +/// `Sealed` is the durable in-doubt record for drop-table two-phase +/// commit: a sealed shard refuses new writer claims (enforced in +/// `claim_epoch`) but is reversible back to `Active` on rollback. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +pub enum ShardStatus { + /// Normal: the shard accepts writer claims. + #[default] + Active, + /// A drop is in flight: claims are refused. Reversible. + Sealed, +} + +impl ShardStatus { + /// Map to the protobuf enum discriminant (`pb::ShardStatus`). + fn to_i32(self) -> i32 { + match self { + Self::Active => 0, + Self::Sealed => 1, + } + } + + /// Map from the protobuf enum discriminant; unknown values decode as + /// `Active` (forward-compatible default). + fn from_i32(v: i32) -> Self { + match v { + 1 => Self::Sealed, + _ => Self::Active, + } + } +} + /// Shard manifest containing epoch-based fencing and WAL state. /// Each shard has exactly one active writer at any time. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -166,6 +199,9 @@ pub struct ShardManifest { pub wal_entry_position_last_seen: u64, pub current_generation: u64, pub flushed_generations: Vec, + /// Lifecycle status (drop-table 2PC). Defaults to `Active`; preserved + /// across claims via `..base` so only fresh constructions set it. + pub status: ShardStatus, } impl DeepSizeOf for ShardManifest { @@ -194,6 +230,7 @@ impl From<&ShardManifest> for pb::ShardManifest { wal_entry_position_last_seen: rm.wal_entry_position_last_seen, current_generation: rm.current_generation, flushed_generations: rm.flushed_generations.iter().map(|fg| fg.into()).collect(), + status: rm.status.to_i32(), } } } @@ -226,6 +263,7 @@ impl TryFrom for ShardManifest { .into_iter() .map(FlushedGeneration::from) .collect(), + status: ShardStatus::from_i32(rm.status), }) } } diff --git a/rust/lance-tokenizer/src/stop_word_filter.rs b/rust/lance-tokenizer/src/stop_word_filter.rs index 2acf0b3dbd5..9a690b0ec06 100644 --- a/rust/lance-tokenizer/src/stop_word_filter.rs +++ b/rust/lance-tokenizer/src/stop_word_filter.rs @@ -17,7 +17,12 @@ fn all_stop_words() -> impl Iterator { stop_words::get("ar"), stopwords::DANISH, stopwords::DUTCH, - stopwords::ENGLISH, + // Use the fuller `stop-words` crate English list (~198 words) rather + // than the local Tantivy-style list (~33 words), which omits extremely + // common pronouns/function words (you, my, your, we, she, what, ...). + // Those omissions let the highest-frequency English tokens through the + // ICU stop-word path and build pathologically large posting lists. + stop_words::get("en"), stopwords::FINNISH, stopwords::FRENCH, stopwords::GERMAN, @@ -51,7 +56,11 @@ impl StopWordFilter { Language::Arabic => stop_words::get("ar"), Language::Danish => stopwords::DANISH, Language::Dutch => stopwords::DUTCH, - Language::English => stopwords::ENGLISH, + // Use the fuller `stop-words` crate English list (~198 words); the + // local Tantivy-style list (~33 words) omits common pronouns/function + // words (you, my, your, we, ...) that would otherwise leak through + // stop-word removal and build pathologically large posting lists. + Language::English => stop_words::get("en"), Language::Finnish => stopwords::FINNISH, Language::French => stopwords::FRENCH, Language::German => stopwords::GERMAN, diff --git a/rust/lance-tokenizer/src/stop_word_filter/stopwords.rs b/rust/lance-tokenizer/src/stop_word_filter/stopwords.rs index 227556ba527..2ac3f4a28aa 100644 --- a/rust/lance-tokenizer/src/stop_word_filter/stopwords.rs +++ b/rust/lance-tokenizer/src/stop_word_filter/stopwords.rs @@ -37,12 +37,6 @@ ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. */ -pub const ENGLISH: &[&str] = &[ - "a", "an", "and", "are", "as", "at", "be", "but", "by", "for", "if", "in", "into", "is", "it", - "no", "not", "of", "on", "or", "such", "that", "the", "their", "then", "there", "these", - "they", "this", "to", "was", "will", "with", -]; - pub const DANISH: &[&str] = &[ "og", "i", "jeg", "det", "at", "en", "den", "til", "er", "som", "på", "de", "med", "han", "af", "for", "ikke", "der", "var", "mig", "sig", "men", "et", "har", "om", "vi", "min", "havde", diff --git a/rust/lance/Cargo.toml b/rust/lance/Cargo.toml index 100aa42ea20..762ca513a19 100644 --- a/rust/lance/Cargo.toml +++ b/rust/lance/Cargo.toml @@ -32,6 +32,7 @@ lance-namespace = { workspace = true } lance-select = { workspace = true } lance-tokenizer = { workspace = true } lance-table = { workspace = true } +lance-bitpacking = { workspace = true } arrow-arith = { workspace = true } arrow-array = { workspace = true } arrow-buffer = { workspace = true } @@ -43,7 +44,6 @@ arrow-schema = { workspace = true } arrow-select = { workspace = true } async-recursion.workspace = true async-trait.workspace = true -bitpacking = { workspace = true } byteorder.workspace = true bytes.workspace = true chrono.workspace = true @@ -139,6 +139,7 @@ reqwest = { version = "0.12", default-features = false, features = ["rustls-tls" [features] default = ["aws", "azure", "gcp", "oss", "huggingface", "tencent", "tos", "goosefs", "geo"] +backtrace = ["lance-core/backtrace"] fp16kernels = ["lance-linalg/fp16kernels"] # Prevent dynamic linking of lzma, which comes from datafusion cli = ["dep:clap", "lzma-sys/static"] @@ -160,6 +161,8 @@ tencent = ["lance-io/tencent"] goosefs = ["lance-io/goosefs"] tos = ["lance-io/tos"] huggingface = ["lance-io/huggingface"] +# Publish object store metrics via the `metrics` crate. +metrics = ["lance-io/metrics"] geo = ["lance-datafusion/geo", "lance-index/geo"] # Enable slow integration tests (disabled by default in CI) slow_tests = [] @@ -171,6 +174,14 @@ slow_tests = [] name = "lq" required-features = ["cli"] +[[bin]] +name = "fm_contains_bench" +required-features = ["cli"] + +[[bin]] +name = "fm_index_tool" +required-features = ["cli"] + [[bench]] name = "scalar_index" harness = false @@ -187,6 +198,10 @@ harness = false name = "scan" harness = false +[[bench]] +name = "s3_file_reader_diagnostics" +harness = false + [[bench]] name = "count_pushdown" harness = false diff --git a/rust/lance/benches/mem_wal/fts/mem_wal_fineweb_fts.rs b/rust/lance/benches/mem_wal/fts/mem_wal_fineweb_fts.rs index 53bfc61fe23..63f9368c3c5 100644 --- a/rust/lance/benches/mem_wal/fts/mem_wal_fineweb_fts.rs +++ b/rust/lance/benches/mem_wal/fts/mem_wal_fineweb_fts.rs @@ -706,10 +706,14 @@ async fn run_read(args: &Args, uri: &str, corpus: &[String]) -> Result = stream.try_collect().await?; diff --git a/rust/lance/benches/mem_wal/fts/mem_wal_fts_read_bench.rs b/rust/lance/benches/mem_wal/fts/mem_wal_fts_read_bench.rs index 578eb0bc0ba..08fdb76b95f 100644 --- a/rust/lance/benches/mem_wal/fts/mem_wal_fts_read_bench.rs +++ b/rust/lance/benches/mem_wal/fts/mem_wal_fts_read_bench.rs @@ -524,7 +524,7 @@ async fn run_local(planner: &LsmFtsSearchPlanner, queries: &[String], k: usize) for q in queries { let t0 = Instant::now(); let plan = planner - .plan_search(TEXT_COL, FullTextSearchQuery::new(q.clone()), k, None) + .plan_search(TEXT_COL, FullTextSearchQuery::new(q.clone()), Some(k), None) .await?; let stream = plan.execute(0, ctx.task_ctx())?; let batches: Vec = stream.try_collect().await?; @@ -684,6 +684,8 @@ async fn run_search(args: &Args) -> Result { let batches_per_gen = (args.max_memtable_rows / args.batch_rows).max(1); let config = ShardWriterConfig { shard_id, + max_wal_persist_retries: 3, + wal_persist_retry_base_delay: std::time::Duration::from_millis(50), shard_spec_id: 0, durable_write: false, sync_indexed_write: false, diff --git a/rust/lance/benches/mem_wal/kv/mem_wal_kv_point_lookup.rs b/rust/lance/benches/mem_wal/kv/mem_wal_kv_point_lookup.rs index d99ea97a70a..757c2539642 100644 --- a/rust/lance/benches/mem_wal/kv/mem_wal_kv_point_lookup.rs +++ b/rust/lance/benches/mem_wal/kv/mem_wal_kv_point_lookup.rs @@ -788,6 +788,8 @@ async fn run_lance( let big = args.rows.saturating_mul(args.value_size + 256).max(1 << 30); let config = ShardWriterConfig { shard_id, + max_wal_persist_retries: 3, + wal_persist_retry_base_delay: std::time::Duration::from_millis(50), shard_spec_id: 0, durable_write: true, sync_indexed_write: true, @@ -817,8 +819,12 @@ async fn run_lance( }; while lo < part_end { let hi = (lo + args.batch_rows).min(part_end); - let batch = - make_batch(schema.clone(), &insert_order[lo..hi], args.value_size, key_type); + let batch = make_batch( + schema.clone(), + &insert_order[lo..hi], + args.value_size, + key_type, + ); writer.put(vec![batch]).await?; lo = hi; } @@ -1519,7 +1525,14 @@ async fn run_lance_flushed( let peak_rss_mb = sampler.stop(); println!( "[lance] read p50={:.2}us p95={:.2}us p99={:.2}us mean={:.2}us qps_1t={:.0} qps_{}t={:.0} (hits={hits} miss={misses_resolved}) peak_rss={:.0}MB", - stats.p50_us, stats.p95_us, stats.p99_us, stats.mean_us, read_qps_1t, args.threads, read_qps_nt, peak_rss_mb + stats.p50_us, + stats.p95_us, + stats.p99_us, + stats.mean_us, + read_qps_1t, + args.threads, + read_qps_nt, + peak_rss_mb ); Ok(EngineResult { diff --git a/rust/lance/benches/mem_wal/point_lookup/mem_wal_point_lookup_bench.rs b/rust/lance/benches/mem_wal/point_lookup/mem_wal_point_lookup_bench.rs index ca57b475560..cb9e6413ac1 100644 --- a/rust/lance/benches/mem_wal/point_lookup/mem_wal_point_lookup_bench.rs +++ b/rust/lance/benches/mem_wal/point_lookup/mem_wal_point_lookup_bench.rs @@ -334,6 +334,8 @@ async fn run_lookup(args: &Args) -> Result { let max_memtable_rows = args.max_memtable_rows; let config = ShardWriterConfig { shard_id, + max_wal_persist_retries: 3, + wal_persist_retry_base_delay: std::time::Duration::from_millis(50), shard_spec_id: 0, durable_write: false, sync_indexed_write: false, diff --git a/rust/lance/benches/mem_wal/read/memtable_read.rs b/rust/lance/benches/mem_wal/read/memtable_read.rs index 5d24943e4cc..27d21e95763 100644 --- a/rust/lance/benches/mem_wal/read/memtable_read.rs +++ b/rust/lance/benches/mem_wal/read/memtable_read.rs @@ -773,8 +773,14 @@ fn bench_fts(c: &mut Criterion) { async move { let mut total_found = 0usize; for (mut scanner, term) in scanners.into_iter().zip(terms.iter()) { + scanner + .full_text_search( + FullTextSearchQuery::new(term.to_string()) + .with_column("text".to_string()) + .unwrap(), + ) + .unwrap(); let batches: Vec = scanner - .full_text_search("text", term) .try_into_stream() .await .unwrap() @@ -947,11 +953,14 @@ fn bench_vector_search(c: &mut Criterion) { let query_array: Arc = Arc::new(Float32Array::from(q.clone())); b.to_async(&rt).iter(|| { let query_array = query_array.clone(); - async { + let memtable = &memtable; + async move { let mut scanner = memtable.scan(); + scanner + .nearest("vector", query_array.as_ref(), k) + .unwrap() + .nprobes(8); let batches: Vec = scanner - .nearest("vector", query_array, k) - .nprobes(8) .try_into_stream() .await .unwrap() diff --git a/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs b/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs index 1b2824bcc83..7b9f35b73c9 100644 --- a/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs +++ b/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs @@ -385,6 +385,8 @@ async fn run_checkpoint( } let writer_config = ShardWriterConfig { shard_id, + max_wal_persist_retries: 3, + wal_persist_retry_base_delay: std::time::Duration::from_millis(50), shard_spec_id: 0, durable_write: false, sync_indexed_write: true, @@ -665,7 +667,7 @@ async fn run_checkpoint( active.schema.clone(), ); let q_arr: ArrayRef = Arc::new(q_fsl); - scanner.nearest(VECTOR_COL, q_arr, k); + scanner.nearest(VECTOR_COL, q_arr.as_ref(), k)?; let h_t = Instant::now(); let stream = scanner.try_into_stream().await?; let batches: Vec = stream.try_collect().await?; diff --git a/rust/lance/benches/mem_wal/vector/mem_wal_index_micro.rs b/rust/lance/benches/mem_wal/vector/mem_wal_index_micro.rs index 57658637b33..8bff64dfcca 100644 --- a/rust/lance/benches/mem_wal/vector/mem_wal_index_micro.rs +++ b/rust/lance/benches/mem_wal/vector/mem_wal_index_micro.rs @@ -208,6 +208,8 @@ async fn main() -> lance_core::Result<()> { let total_batches_max = max_rows.div_ceil(batch_size); let writer_config = ShardWriterConfig { shard_id, + max_wal_persist_retries: 3, + wal_persist_retry_base_delay: std::time::Duration::from_millis(50), shard_spec_id: 0, durable_write, sync_indexed_write: true, @@ -298,7 +300,7 @@ async fn main() -> lance_core::Result<()> { active.schema.clone(), ); let q_arr: ArrayRef = Arc::new(q_fsl); - scanner.nearest(VECTOR_COL, q_arr, 10); + scanner.nearest(VECTOR_COL, q_arr.as_ref(), 10)?; let t = Instant::now(); let stream = scanner.try_into_stream().await?; let _: Vec = stream.try_collect().await?; diff --git a/rust/lance/benches/mem_wal/vector/mem_wal_vector_bench.rs b/rust/lance/benches/mem_wal/vector/mem_wal_vector_bench.rs index 07c58a72bc4..632bd37626c 100644 --- a/rust/lance/benches/mem_wal/vector/mem_wal_vector_bench.rs +++ b/rust/lance/benches/mem_wal/vector/mem_wal_vector_bench.rs @@ -516,6 +516,8 @@ async fn run_search(args: &Args) -> Result { let row_bytes = VECTOR_DIM * 4 + 8; // embedding + id let config = ShardWriterConfig { shard_id, + max_wal_persist_retries: 3, + wal_persist_retry_base_delay: std::time::Duration::from_millis(50), shard_spec_id: 0, durable_write: false, sync_indexed_write: false, diff --git a/rust/lance/benches/mem_wal/write/mem_wal_write.rs b/rust/lance/benches/mem_wal/write/mem_wal_write.rs index 9a5fc71ab17..6d3a31c4011 100644 --- a/rust/lance/benches/mem_wal/write/mem_wal_write.rs +++ b/rust/lance/benches/mem_wal/write/mem_wal_write.rs @@ -630,6 +630,9 @@ fn bench_lance_memwal_write(c: &mut Criterion) { let config = ShardWriterConfig { shard_id, shard_spec_id: 0, + max_wal_persist_retries: 3, + wal_persist_retry_base_delay: + std::time::Duration::from_millis(50), durable_write: durable, sync_indexed_write: indexed, max_wal_buffer_size: max_wal_buffer_size diff --git a/rust/lance/benches/s3_file_reader_diagnostics.rs b/rust/lance/benches/s3_file_reader_diagnostics.rs new file mode 100644 index 00000000000..5c661f447e3 --- /dev/null +++ b/rust/lance/benches/s3_file_reader_diagnostics.rs @@ -0,0 +1,2356 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +#![allow(clippy::print_stdout)] +#![recursion_limit = "256"] + +use std::collections::BTreeMap; +use std::env; +use std::fs; +use std::hint::black_box; +use std::ops::Range; +use std::sync::{ + Arc, Mutex, + atomic::{AtomicU64, Ordering}, +}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use arrow_array::RecordBatch; +use futures::future::BoxFuture; +use futures::stream::{FuturesOrdered, FuturesUnordered}; +use futures::{FutureExt, StreamExt, TryStreamExt}; +use lance::dataset::ProjectionRequest; +use lance::dataset::builder::DatasetBuilder; +use lance::dataset::fragment::{FileFragment, FragReadConfig}; +use lance::dataset::scanner::{ExecutionStatsCallback, ExecutionSummaryCounts}; +use lance_core::datatypes::Schema; +use lance_encoding::decoder::PageEncoding; +use lance_encoding::format::pb21; +use lance_file::reader::{ + DEFAULT_READ_CHUNK_SIZE, FileReader as LanceFileReader, FileReaderOptions, +}; +use lance_io::object_store::ObjectStore as LanceObjectStore; +use lance_io::scheduler::{FileScheduler, ScanScheduler, ScanStats, SchedulerConfig}; +use lance_io::utils::CachedFileSize; +use serde_json::{Value, json}; +use tracing::field::{Field, Visit}; +use tracing::subscriber::Interest; +use tracing::{Event, Metadata, Subscriber}; +use tracing_subscriber::layer::{Context, Layer}; +use tracing_subscriber::prelude::*; + +type Error = Box; +type Result = std::result::Result; + +const GIB: u64 = 1024 * 1024 * 1024; +const SCHEDULER_STATE_EVENT_TARGET: &str = "lance_io::scheduler::state"; + +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +enum SchedulerQueueKind { + Standard, + Lite, +} + +#[derive(Debug, Clone, Copy)] +struct SchedulerDiagnostics { + kind: SchedulerQueueKind, + stats: ScanStats, + io_capacity: u64, + iops_available: u64, + active_iops: u64, + pending_iops: u64, + pending_bytes: u64, + bytes_available: i64, + bytes_reserved: i64, + io_buffer_size_bytes: u64, + priorities_in_flight: u64, + no_backpressure: bool, + head_task_bytes: Option, + head_task_priority_high: Option, + head_task_priority_low: Option, + min_in_flight_priority_high: Option, + min_in_flight_priority_low: Option, + head_task_can_deliver: Option, + head_task_priority_bypass: Option, + head_task_blocked_by_iops: Option, + head_task_blocked_by_bytes: Option, +} + +#[derive(Debug, Clone)] +struct Config { + backend: Backend, + uri: String, + dataset_version: u64, + columns: Option>, + limit_rows: u64, + target_bytes: Option, + raw_range_size_bytes: u64, + raw_range_mode: RawRangeMode, + raw_column_indices: Option>, + raw_submit_mode: RawSubmitMode, + raw_completion_mode: RawCompletionMode, + take_repetitions: u64, + io_buffer_gib: Vec>, + batch_size: u32, + batch_size_bytes: Option, + skip_batch_byte_accounting: bool, + read_chunk_size: Option, + fragment_concurrency: usize, + batch_concurrency: usize, + sample_ms: u64, + out_dir: String, + case_name: String, + describe_layout: bool, + detach_fragment_streams: bool, + drop_read_tasks: bool, +} + +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +enum Backend { + FileReader, + Scanner, + SchedulerRaw, + DatasetTake, +} + +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +enum RawSubmitMode { + Single, + SplitNoConcat, + SplitConcat, +} + +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +enum RawRangeMode { + FileSequential, + MetadataPages, + MetadataPagesRoundRobin, +} + +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +enum RawCompletionMode { + Unordered, + Ordered, +} + +impl RawRangeMode { + fn name(self) -> &'static str { + match self { + Self::FileSequential => "file-sequential", + Self::MetadataPages => "metadata-pages", + Self::MetadataPagesRoundRobin => "metadata-pages-round-robin", + } + } +} + +impl RawSubmitMode { + fn name(self) -> &'static str { + match self { + Self::Single => "single", + Self::SplitNoConcat => "split-no-concat", + Self::SplitConcat => "split-concat", + } + } +} + +impl RawCompletionMode { + fn name(self) -> &'static str { + match self { + Self::Unordered => "unordered", + Self::Ordered => "ordered", + } + } +} + +impl Backend { + fn name(self) -> &'static str { + match self { + Self::FileReader => "lance-file-reader", + Self::Scanner => "lance-scanner", + Self::SchedulerRaw => "lance-scheduler-raw", + Self::DatasetTake => "lance-dataset-take", + } + } + + fn layer(self) -> &'static str { + match self { + Self::FileReader => "file-reader", + Self::Scanner => "scanner", + Self::SchedulerRaw => "scheduler", + Self::DatasetTake => "dataset-take", + } + } +} + +#[derive(Debug, Clone, Copy)] +struct CpuSample { + idle: u64, + total: u64, +} + +#[derive(Debug, Default)] +struct SharedCounters { + fragments_started: AtomicU64, + fragments_completed: AtomicU64, + batch_futures_emitted: AtomicU64, + batch_futures_received: AtomicU64, + batches_completed: AtomicU64, + rows_completed: AtomicU64, + arrow_bytes: AtomicU64, + open_reader_ns: AtomicU64, + read_stream_create_ns: AtomicU64, + next_batch_poll_ns: AtomicU64, + channel_send_wait_ns: AtomicU64, + decode_ns: AtomicU64, + raw_reassemble_ns: AtomicU64, +} + +#[derive(Debug)] +struct CaseStats { + rows: u64, + batches: u64, + arrow_bytes: u64, + planned_fragments: usize, + planned_rows: u64, + elapsed: Duration, + producer_finished_at: Option, + peak_decode_in_flight: usize, + cpu_avg: Option, + scheduler_diagnostics: SchedulerDiagnostics, + counters: Arc, + samples: Vec, +} + +#[derive(Debug)] +struct LastSample { + elapsed: Duration, + scheduler_stats: ScanStats, + rows: u64, + batches: u64, + arrow_bytes: u64, +} + +fn usage() -> &'static str { + "usage: s3_file_reader_diagnostics --uri \ + [--backend ] \ + [--dataset-version ] [--columns ] \ + [--limit-rows ] [--target-bytes ] [--raw-range-size-bytes ] \ + [--raw-range-mode ] \ + [--raw-column-indices ] \ + [--raw-submit-mode ] \ + [--raw-completion-mode ] \ + [--take-repetitions ] \ + [--io-buffer-gib ] \ + [--batch-size ] [--batch-size-bytes ] [--read-chunk-size ] \ + [--skip-batch-byte-accounting] \ + [--fragment-concurrency ] [--batch-concurrency ] \ + [--sample-ms ] [--out-dir ] [--case ] \ + [--detach-fragment-streams] [--drop-read-tasks] [--describe-layout]" +} + +fn parse_args() -> Result { + let mut backend = Backend::FileReader; + let mut uri = None; + let mut dataset_version = 1u64; + let mut columns = Some(vec!["vector".to_string()]); + let mut limit_rows = 67_108_864u64; + let mut target_bytes = None; + let mut raw_range_size_bytes = 16 * 1024 * 1024; + let mut raw_range_mode = RawRangeMode::FileSequential; + let mut raw_column_indices = None; + let mut raw_submit_mode = RawSubmitMode::Single; + let mut raw_completion_mode = RawCompletionMode::Unordered; + let mut take_repetitions = 100u64; + let mut io_buffer_gib = vec![Some(8)]; + let mut batch_size = 8192u32; + let mut batch_size_bytes = None; + let mut skip_batch_byte_accounting = false; + let mut read_chunk_size = None; + let mut fragment_concurrency = 256usize; + let mut batch_concurrency = 256usize; + let mut sample_ms = 1000u64; + let mut out_dir = "/tmp/lance-s3-bottleneck-results".to_string(); + let mut case_name = "lance-file-reader-diagnostics".to_string(); + let mut describe_layout = false; + let mut detach_fragment_streams = false; + let mut drop_read_tasks = false; + + let mut args = env::args().skip(1); + while let Some(arg) = args.next() { + match arg.as_str() { + "--backend" => { + let value = args + .next() + .ok_or_else(|| format!("missing value for --backend. {}", usage()))?; + backend = parse_backend(&value)?; + } + "--uri" => uri = args.next(), + "--dataset-version" => { + dataset_version = parse_required_value(&mut args, "--dataset-version")?; + } + "--columns" => { + let value = args + .next() + .ok_or_else(|| format!("missing value for --columns. {}", usage()))?; + columns = parse_columns(&value)?; + } + "--limit-rows" => { + limit_rows = parse_required_value(&mut args, "--limit-rows")?; + } + "--target-bytes" => { + target_bytes = Some(parse_required_value(&mut args, "--target-bytes")?); + } + "--raw-range-size-bytes" => { + raw_range_size_bytes = parse_required_value(&mut args, "--raw-range-size-bytes")?; + } + "--raw-range-mode" => { + let value = args + .next() + .ok_or_else(|| format!("missing value for --raw-range-mode. {}", usage()))?; + raw_range_mode = parse_raw_range_mode(&value)?; + } + "--raw-column-indices" => { + let value = args.next().ok_or_else(|| { + format!("missing value for --raw-column-indices. {}", usage()) + })?; + raw_column_indices = parse_raw_column_indices(&value)?; + } + "--raw-submit-mode" => { + let value = args + .next() + .ok_or_else(|| format!("missing value for --raw-submit-mode. {}", usage()))?; + raw_submit_mode = parse_raw_submit_mode(&value)?; + } + "--raw-completion-mode" => { + let value = args.next().ok_or_else(|| { + format!("missing value for --raw-completion-mode. {}", usage()) + })?; + raw_completion_mode = parse_raw_completion_mode(&value)?; + } + "--take-repetitions" => { + take_repetitions = parse_required_value(&mut args, "--take-repetitions")?; + } + "--io-buffer-gib" => { + let value = args + .next() + .ok_or_else(|| format!("missing value for --io-buffer-gib. {}", usage()))?; + io_buffer_gib = parse_io_buffer_gib(&value)?; + } + "--batch-size" => { + batch_size = parse_required_value(&mut args, "--batch-size")?; + } + "--batch-size-bytes" => { + batch_size_bytes = Some(parse_required_value(&mut args, "--batch-size-bytes")?); + } + "--skip-batch-byte-accounting" => { + skip_batch_byte_accounting = true; + } + "--read-chunk-size" => { + read_chunk_size = Some(parse_required_value(&mut args, "--read-chunk-size")?); + } + "--fragment-concurrency" => { + fragment_concurrency = parse_required_value(&mut args, "--fragment-concurrency")?; + } + "--batch-concurrency" => { + batch_concurrency = parse_required_value(&mut args, "--batch-concurrency")?; + } + "--sample-ms" => { + sample_ms = parse_required_value(&mut args, "--sample-ms")?; + } + "--out-dir" => { + out_dir = args + .next() + .ok_or_else(|| format!("missing value for --out-dir. {}", usage()))?; + } + "--case" => { + case_name = args + .next() + .ok_or_else(|| format!("missing value for --case. {}", usage()))?; + } + "--describe-layout" => { + describe_layout = true; + } + "--detach-fragment-streams" => { + detach_fragment_streams = true; + } + "--drop-read-tasks" => { + drop_read_tasks = true; + } + "--help" | "-h" => { + println!("{}", usage()); + std::process::exit(0); + } + "--bench" => { + // Cargo appends this flag when running harness-free benches. + } + other => { + return Err(format!("unknown argument {other}. {}", usage()).into()); + } + } + } + + let uri = uri.ok_or_else(|| format!("missing required --uri. {}", usage()))?; + if limit_rows == 0 { + return Err("--limit-rows must be greater than zero".into()); + } + if matches!(target_bytes, Some(0)) { + return Err("--target-bytes must be greater than zero".into()); + } + if raw_range_size_bytes == 0 { + return Err("--raw-range-size-bytes must be greater than zero".into()); + } + if take_repetitions == 0 { + return Err("--take-repetitions must be greater than zero".into()); + } + if io_buffer_gib.is_empty() { + return Err("--io-buffer-gib must not be empty".into()); + } + if batch_size == 0 { + return Err("--batch-size must be greater than zero".into()); + } + if matches!(batch_size_bytes, Some(0)) { + return Err("--batch-size-bytes must be greater than zero".into()); + } + if matches!(read_chunk_size, Some(0)) { + return Err("--read-chunk-size must be greater than zero".into()); + } + if fragment_concurrency == 0 && !matches!(backend, Backend::Scanner) { + return Err("--fragment-concurrency must be greater than zero".into()); + } + if batch_concurrency == 0 && !matches!(backend, Backend::Scanner) { + return Err("--batch-concurrency must be greater than zero".into()); + } + if sample_ms == 0 { + return Err("--sample-ms must be greater than zero".into()); + } + + Ok(Config { + backend, + uri, + dataset_version, + columns, + limit_rows, + target_bytes, + raw_range_size_bytes, + raw_range_mode, + raw_column_indices, + raw_submit_mode, + raw_completion_mode, + take_repetitions, + io_buffer_gib, + batch_size, + batch_size_bytes, + skip_batch_byte_accounting, + read_chunk_size, + fragment_concurrency, + batch_concurrency, + sample_ms, + out_dir, + case_name, + describe_layout, + detach_fragment_streams, + drop_read_tasks, + }) +} + +fn parse_backend(value: &str) -> Result { + match value { + "file-reader" | "lance-file-reader" => Ok(Backend::FileReader), + "scanner" | "lance-scanner" => Ok(Backend::Scanner), + "scheduler-raw" | "lance-scheduler-raw" => Ok(Backend::SchedulerRaw), + "dataset-take" | "take" | "lance-dataset-take" => Ok(Backend::DatasetTake), + other => Err(format!( + "invalid --backend value {other}; expected file-reader, scanner, scheduler-raw, or dataset-take" + ) + .into()), + } +} + +fn parse_raw_submit_mode(value: &str) -> Result { + match value { + "single" => Ok(RawSubmitMode::Single), + "split-no-concat" => Ok(RawSubmitMode::SplitNoConcat), + "split-concat" => Ok(RawSubmitMode::SplitConcat), + other => Err(format!( + "invalid --raw-submit-mode value {other}; expected single, split-no-concat, or split-concat" + ) + .into()), + } +} + +fn parse_raw_range_mode(value: &str) -> Result { + match value { + "file-sequential" => Ok(RawRangeMode::FileSequential), + "metadata-pages" => Ok(RawRangeMode::MetadataPages), + "metadata-pages-round-robin" => Ok(RawRangeMode::MetadataPagesRoundRobin), + other => Err(format!( + "invalid --raw-range-mode value {other}; expected file-sequential, metadata-pages, or metadata-pages-round-robin" + ) + .into()), + } +} + +fn parse_raw_column_indices(value: &str) -> Result>> { + if value == "all" { + return Ok(None); + } + + let indices = value + .split(',') + .map(str::trim) + .filter(|part| !part.is_empty()) + .map(|part| { + part.parse::() + .map_err(|err| format!("invalid raw column index {part}: {err}").into()) + }) + .collect::>>()?; + if indices.is_empty() { + return Err("--raw-column-indices must specify at least one column index or all".into()); + } + Ok(Some(indices)) +} + +fn parse_raw_completion_mode(value: &str) -> Result { + match value { + "unordered" => Ok(RawCompletionMode::Unordered), + "ordered" => Ok(RawCompletionMode::Ordered), + other => Err(format!( + "invalid --raw-completion-mode value {other}; expected unordered or ordered" + ) + .into()), + } +} + +fn parse_required_value(args: &mut impl Iterator, name: &str) -> Result +where + T: std::str::FromStr, + T::Err: std::fmt::Display + Send + Sync + 'static, +{ + let value = args + .next() + .ok_or_else(|| format!("missing value for {name}. {}", usage()))?; + value + .parse() + .map_err(|err| format!("invalid {name} value {value}: {err}").into()) +} + +fn parse_columns(value: &str) -> Result>> { + match value { + "all" => Ok(None), + "empty" => Err("FileReader benchmark requires at least one data column".into()), + _ => Ok(Some( + value + .split(',') + .map(str::trim) + .filter(|column| !column.is_empty()) + .map(ToString::to_string) + .collect(), + )), + } +} + +fn parse_io_buffer_gib(value: &str) -> Result>> { + value + .split(',') + .map(|part| { + let part = part.trim(); + if part == "auto" { + Ok(None) + } else { + part.parse::() + .map(Some) + .map_err(|err| format!("invalid --io-buffer-gib value {part}: {err}").into()) + } + }) + .collect() +} + +fn projection_name(columns: &Option>) -> String { + match columns { + None => "all".to_string(), + Some(columns) => columns.join(","), + } +} + +fn page_layout_kind(encoding: &PageEncoding) -> &'static str { + match encoding { + PageEncoding::Legacy(_) => "legacy", + PageEncoding::Structural(layout) => match layout.layout.as_ref() { + Some(pb21::page_layout::Layout::MiniBlockLayout(_)) => "miniblock", + Some(pb21::page_layout::Layout::ConstantLayout(_)) => "constant", + Some(pb21::page_layout::Layout::FullZipLayout(_)) => "fullzip", + Some(pb21::page_layout::Layout::BlobLayout(_)) => "blob", + None => "missing", + }, + } +} + +fn summarize_u64(values: &[u64]) -> Value { + if values.is_empty() { + return json!({ + "count": 0, + "min": null, + "p50": null, + "p90": null, + "max": null, + "sum": 0, + }); + } + let mut sorted = values.to_vec(); + sorted.sort_unstable(); + let percentile = |p: f64| { + let idx = ((sorted.len() - 1) as f64 * p).round() as usize; + sorted[idx] + }; + json!({ + "count": sorted.len(), + "min": sorted[0], + "p50": percentile(0.5), + "p90": percentile(0.9), + "max": *sorted.last().unwrap(), + "sum": values.iter().sum::(), + }) +} + +async fn describe_layout(config: &Config) -> Result<()> { + let dataset = Arc::new( + DatasetBuilder::from_uri(&config.uri) + .with_version(config.dataset_version) + .load() + .await?, + ); + let fragment = dataset + .fragments() + .first() + .ok_or("dataset has no fragments")?; + let data_file = fragment + .files + .first() + .ok_or("first fragment has no data files")?; + if data_file.base_id.is_some() { + return Err("layout diagnostics do not support external base data files yet".into()); + } + + let data_path = dataset.data_dir().join(data_file.path.as_str()); + let (object_store, _) = LanceObjectStore::from_uri(&config.uri).await?; + let scheduler = ScanScheduler::new(object_store, SchedulerConfig::new(8 * GIB)); + let file_scheduler = scheduler + .open_file(&data_path, &CachedFileSize::unknown()) + .await?; + let metadata = LanceFileReader::read_all_metadata(&file_scheduler).await?; + + let columns = metadata + .column_infos + .iter() + .map(|column| { + let mut layout_counts = BTreeMap::new(); + let mut page_rows = Vec::with_capacity(column.page_infos.len()); + let mut page_bytes = Vec::with_capacity(column.page_infos.len()); + for page in column.page_infos.iter() { + *layout_counts + .entry(page_layout_kind(&page.encoding)) + .or_insert(0usize) += 1; + page_rows.push(page.num_rows); + page_bytes.push( + page.buffer_offsets_and_sizes + .iter() + .map(|(_, size)| *size) + .sum::(), + ); + } + json!({ + "column_index": column.index, + "num_pages": column.page_infos.len(), + "layout_counts": layout_counts, + "page_rows": summarize_u64(&page_rows), + "page_bytes": summarize_u64(&page_bytes), + }) + }) + .collect::>(); + + println!( + "{}", + serde_json::to_string_pretty(&json!({ + "dataset_uri": config.uri, + "dataset_version": config.dataset_version, + "fragment_id": fragment.id, + "data_file_path": data_file.path, + "resolved_data_path": data_path.to_string(), + "file_version": metadata.version().to_string(), + "num_rows": metadata.num_rows, + "num_data_bytes": metadata.num_data_bytes, + "columns": columns, + }))? + ); + Ok(()) +} + +fn projected_schema(dataset_schema: &Schema, columns: &Option>) -> Result { + Ok(match columns { + None => dataset_schema.clone(), + Some(columns) => dataset_schema.project(columns)?, + }) +} + +fn file_reader_options(config: &Config) -> Option { + if config.batch_size_bytes.is_none() && config.read_chunk_size.is_none() { + return None; + } + Some(FileReaderOptions { + batch_size_bytes: config.batch_size_bytes, + read_chunk_size: config.read_chunk_size.unwrap_or(DEFAULT_READ_CHUNK_SIZE), + ..Default::default() + }) +} + +fn add_duration(counter: &AtomicU64, duration: Duration) { + let nanos = duration.as_nanos().min(u128::from(u64::MAX)) as u64; + counter.fetch_add(nanos, Ordering::Relaxed); +} + +fn ns_to_seconds(ns: u64) -> f64 { + ns as f64 / 1_000_000_000.0 +} + +fn diff_u64(current: u64, previous: u64) -> u64 { + current.saturating_sub(previous) +} + +fn scheduler_kind_name(kind: SchedulerQueueKind) -> &'static str { + match kind { + SchedulerQueueKind::Standard => "standard", + SchedulerQueueKind::Lite => "lite", + } +} + +fn diagnostics_json(diagnostics: SchedulerDiagnostics) -> Value { + json!({ + "queue_kind": scheduler_kind_name(diagnostics.kind), + "scheduler_iops": diagnostics.stats.iops, + "scheduler_requests": diagnostics.stats.requests, + "scheduler_bytes_read": diagnostics.stats.bytes_read, + "io_capacity": diagnostics.io_capacity, + "iops_available": diagnostics.iops_available, + "active_iops": diagnostics.active_iops, + "pending_iops": diagnostics.pending_iops, + "pending_bytes": diagnostics.pending_bytes, + "bytes_available": diagnostics.bytes_available, + "bytes_reserved": diagnostics.bytes_reserved, + "io_buffer_size_bytes": diagnostics.io_buffer_size_bytes, + "priorities_in_flight": diagnostics.priorities_in_flight, + "no_backpressure": diagnostics.no_backpressure, + "head_task_bytes": diagnostics.head_task_bytes, + "head_task_priority_high": diagnostics.head_task_priority_high, + "head_task_priority_low": diagnostics.head_task_priority_low, + "min_in_flight_priority_high": diagnostics.min_in_flight_priority_high, + "min_in_flight_priority_low": diagnostics.min_in_flight_priority_low, + "head_task_can_deliver": diagnostics.head_task_can_deliver, + "head_task_priority_bypass": diagnostics.head_task_priority_bypass, + "head_task_blocked_by_iops": diagnostics.head_task_blocked_by_iops, + "head_task_blocked_by_bytes": diagnostics.head_task_blocked_by_bytes, + }) +} + +#[derive(Debug, Default)] +struct ExecutionStatsHolder { + collected_stats: Arc>>, +} + +impl ExecutionStatsHolder { + fn get_setter(&self) -> ExecutionStatsCallback { + let collected_stats = self.collected_stats.clone(); + Arc::new(move |stats| { + *collected_stats.lock().unwrap() = Some(stats.clone()); + }) + } + + fn consume(self) -> Option { + self.collected_stats.lock().unwrap().take() + } +} + +#[derive(Debug, Clone, Default)] +struct SchedulerDiagnosticsCollector { + latest: Arc>>, +} + +impl SchedulerDiagnosticsCollector { + fn clear(&self) { + *self.latest.lock().unwrap() = None; + } + + fn observe(&self, diagnostics: SchedulerDiagnostics) { + *self.latest.lock().unwrap() = Some(diagnostics); + } + + fn snapshot(&self, io_buffer_gib: Option) -> SchedulerDiagnostics { + self.latest + .lock() + .unwrap() + .as_ref() + .copied() + .unwrap_or_else(|| diagnostics_from_scan_stats(ScanStats::default(), io_buffer_gib)) + } +} + +#[derive(Debug, Clone)] +struct SchedulerDiagnosticsLayer { + collector: SchedulerDiagnosticsCollector, +} + +impl SchedulerDiagnosticsLayer { + fn new(collector: SchedulerDiagnosticsCollector) -> Self { + Self { collector } + } +} + +fn is_scheduler_state_metadata(metadata: &Metadata<'_>) -> bool { + // The scheduler uses `tracing::enabled!` before constructing the event; + // that guard registers a HINT callsite, not an EVENT callsite. + metadata.target() == SCHEDULER_STATE_EVENT_TARGET && *metadata.level() == tracing::Level::TRACE +} + +impl Layer for SchedulerDiagnosticsLayer +where + S: Subscriber, +{ + fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest { + if is_scheduler_state_metadata(metadata) { + Interest::always() + } else { + Interest::never() + } + } + + fn enabled(&self, metadata: &Metadata<'_>, _ctx: Context<'_, S>) -> bool { + is_scheduler_state_metadata(metadata) + } + + fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) { + if !is_scheduler_state_metadata(event.metadata()) { + return; + } + let mut visitor = SchedulerDiagnosticsVisitor::default(); + event.record(&mut visitor); + if let Some(diagnostics) = visitor.into_diagnostics() { + self.collector.observe(diagnostics); + } + } +} + +#[derive(Debug, Default)] +struct SchedulerDiagnosticsVisitor { + kind: Option, + scheduler_iops: Option, + scheduler_requests: Option, + scheduler_bytes_read: Option, + io_capacity: Option, + iops_available: Option, + active_iops: Option, + pending_iops: Option, + pending_bytes: Option, + bytes_available: Option, + bytes_reserved: Option, + io_buffer_size_bytes: Option, + priorities_in_flight: Option, + no_backpressure: Option, + head_task_bytes_present: bool, + head_task_bytes: Option, + head_task_priority_high_present: bool, + head_task_priority_high: Option, + head_task_priority_low_present: bool, + head_task_priority_low: Option, + min_in_flight_priority_high_present: bool, + min_in_flight_priority_high: Option, + min_in_flight_priority_low_present: bool, + min_in_flight_priority_low: Option, + head_task_can_deliver_present: bool, + head_task_can_deliver: Option, + head_task_priority_bypass_present: bool, + head_task_priority_bypass: Option, + head_task_blocked_by_iops_present: bool, + head_task_blocked_by_iops: Option, + head_task_blocked_by_bytes_present: bool, + head_task_blocked_by_bytes: Option, +} + +impl SchedulerDiagnosticsVisitor { + fn into_diagnostics(self) -> Option { + Some(SchedulerDiagnostics { + kind: self.kind?, + stats: ScanStats { + iops: self.scheduler_iops.unwrap_or_default(), + requests: self.scheduler_requests.unwrap_or_default(), + bytes_read: self.scheduler_bytes_read.unwrap_or_default(), + }, + io_capacity: self.io_capacity.unwrap_or_default(), + iops_available: self.iops_available.unwrap_or_default(), + active_iops: self.active_iops.unwrap_or_default(), + pending_iops: self.pending_iops.unwrap_or_default(), + pending_bytes: self.pending_bytes.unwrap_or_default(), + bytes_available: self.bytes_available.unwrap_or_default(), + bytes_reserved: self.bytes_reserved.unwrap_or_default(), + io_buffer_size_bytes: self.io_buffer_size_bytes.unwrap_or_default(), + priorities_in_flight: self.priorities_in_flight.unwrap_or_default(), + no_backpressure: self.no_backpressure.unwrap_or(false), + head_task_bytes: optional_u64(self.head_task_bytes_present, self.head_task_bytes), + head_task_priority_high: optional_u64( + self.head_task_priority_high_present, + self.head_task_priority_high, + ), + head_task_priority_low: optional_u64( + self.head_task_priority_low_present, + self.head_task_priority_low, + ), + min_in_flight_priority_high: optional_u64( + self.min_in_flight_priority_high_present, + self.min_in_flight_priority_high, + ), + min_in_flight_priority_low: optional_u64( + self.min_in_flight_priority_low_present, + self.min_in_flight_priority_low, + ), + head_task_can_deliver: optional_bool( + self.head_task_can_deliver_present, + self.head_task_can_deliver, + ), + head_task_priority_bypass: optional_bool( + self.head_task_priority_bypass_present, + self.head_task_priority_bypass, + ), + head_task_blocked_by_iops: optional_bool( + self.head_task_blocked_by_iops_present, + self.head_task_blocked_by_iops, + ), + head_task_blocked_by_bytes: optional_bool( + self.head_task_blocked_by_bytes_present, + self.head_task_blocked_by_bytes, + ), + }) + } +} + +impl Visit for SchedulerDiagnosticsVisitor { + fn record_bool(&mut self, field: &Field, value: bool) { + match field.name() { + "no_backpressure" => self.no_backpressure = Some(value), + "head_task_bytes_present" => self.head_task_bytes_present = value, + "head_task_priority_high_present" => self.head_task_priority_high_present = value, + "head_task_priority_low_present" => self.head_task_priority_low_present = value, + "min_in_flight_priority_high_present" => { + self.min_in_flight_priority_high_present = value; + } + "min_in_flight_priority_low_present" => { + self.min_in_flight_priority_low_present = value; + } + "head_task_can_deliver_present" => self.head_task_can_deliver_present = value, + "head_task_can_deliver" => self.head_task_can_deliver = Some(value), + "head_task_priority_bypass_present" => { + self.head_task_priority_bypass_present = value; + } + "head_task_priority_bypass" => self.head_task_priority_bypass = Some(value), + "head_task_blocked_by_iops_present" => { + self.head_task_blocked_by_iops_present = value; + } + "head_task_blocked_by_iops" => self.head_task_blocked_by_iops = Some(value), + "head_task_blocked_by_bytes_present" => { + self.head_task_blocked_by_bytes_present = value; + } + "head_task_blocked_by_bytes" => self.head_task_blocked_by_bytes = Some(value), + _ => {} + } + } + + fn record_i64(&mut self, field: &Field, value: i64) { + match field.name() { + "bytes_available" => self.bytes_available = Some(value), + "bytes_reserved" => self.bytes_reserved = Some(value), + _ => {} + } + } + + fn record_u64(&mut self, field: &Field, value: u64) { + match field.name() { + "scheduler_iops" => self.scheduler_iops = Some(value), + "scheduler_requests" => self.scheduler_requests = Some(value), + "scheduler_bytes_read" => self.scheduler_bytes_read = Some(value), + "io_capacity" => self.io_capacity = Some(value), + "iops_available" => self.iops_available = Some(value), + "active_iops" => self.active_iops = Some(value), + "pending_iops" => self.pending_iops = Some(value), + "pending_bytes" => self.pending_bytes = Some(value), + "io_buffer_size_bytes" => self.io_buffer_size_bytes = Some(value), + "priorities_in_flight" => self.priorities_in_flight = Some(value), + "head_task_bytes" => self.head_task_bytes = Some(value), + "head_task_priority_high" => self.head_task_priority_high = Some(value), + "head_task_priority_low" => self.head_task_priority_low = Some(value), + "min_in_flight_priority_high" => self.min_in_flight_priority_high = Some(value), + "min_in_flight_priority_low" => self.min_in_flight_priority_low = Some(value), + _ => {} + } + } + + fn record_str(&mut self, field: &Field, value: &str) { + if field.name() == "queue_kind" { + self.kind = match value { + "standard" => Some(SchedulerQueueKind::Standard), + "lite" => Some(SchedulerQueueKind::Lite), + _ => None, + }; + } + } + + fn record_debug(&mut self, _field: &Field, _value: &dyn std::fmt::Debug) {} +} + +fn optional_u64(present: bool, value: Option) -> Option { + present.then(|| value.unwrap_or_default()) +} + +fn optional_bool(present: bool, value: Option) -> Option { + present.then(|| value.unwrap_or(false)) +} + +fn diagnostics_from_scan_stats( + stats: ScanStats, + io_buffer_gib: Option, +) -> SchedulerDiagnostics { + SchedulerDiagnostics { + kind: SchedulerQueueKind::Standard, + stats, + io_capacity: 0, + iops_available: 0, + active_iops: 0, + pending_iops: 0, + pending_bytes: 0, + bytes_available: 0, + bytes_reserved: 0, + io_buffer_size_bytes: io_buffer_gib.map(|value| value * GIB).unwrap_or_default(), + priorities_in_flight: 0, + no_backpressure: false, + head_task_bytes: None, + head_task_priority_high: None, + head_task_priority_low: None, + min_in_flight_priority_high: None, + min_in_flight_priority_low: None, + head_task_can_deliver: None, + head_task_priority_bypass: None, + head_task_blocked_by_iops: None, + head_task_blocked_by_bytes: None, + } +} + +fn scan_stats_from_execution_summary(summary: &ExecutionSummaryCounts) -> ScanStats { + ScanStats { + iops: summary.iops as u64, + requests: summary.requests as u64, + bytes_read: summary.bytes_read as u64, + } +} + +fn sample_json( + started: Instant, + counters: &SharedCounters, + diagnostics: SchedulerDiagnostics, + decode_in_flight: usize, + channel_buffered: usize, + last: &mut LastSample, +) -> Value { + let elapsed = started.elapsed(); + let interval = elapsed.saturating_sub(last.elapsed); + let interval_secs = interval.as_secs_f64(); + let rows = counters.rows_completed.load(Ordering::Relaxed); + let batches = counters.batches_completed.load(Ordering::Relaxed); + let arrow_bytes = counters.arrow_bytes.load(Ordering::Relaxed); + let scheduler_stats = diagnostics.stats; + let delta_scheduler_bytes = + diff_u64(scheduler_stats.bytes_read, last.scheduler_stats.bytes_read); + let delta_rows = diff_u64(rows, last.rows); + let delta_arrow_bytes = diff_u64(arrow_bytes, last.arrow_bytes); + let physical_gbps = if interval_secs > 0.0 { + delta_scheduler_bytes as f64 * 8.0 / interval_secs / 1_000_000_000.0 + } else { + 0.0 + }; + let logical_gbps = if interval_secs > 0.0 { + delta_arrow_bytes as f64 * 8.0 / interval_secs / 1_000_000_000.0 + } else { + 0.0 + }; + let rows_per_second = if interval_secs > 0.0 { + delta_rows as f64 / interval_secs + } else { + 0.0 + }; + + last.elapsed = elapsed; + last.scheduler_stats = scheduler_stats; + last.rows = rows; + last.batches = batches; + last.arrow_bytes = arrow_bytes; + + json!({ + "elapsed_seconds": elapsed.as_secs_f64(), + "interval_seconds": interval_secs, + "physical_gbps": physical_gbps, + "logical_gbps": logical_gbps, + "rows_per_second": rows_per_second, + "rows": rows, + "batches": batches, + "arrow_bytes": arrow_bytes, + "delta_rows": delta_rows, + "delta_arrow_bytes": delta_arrow_bytes, + "delta_scheduler_bytes": delta_scheduler_bytes, + "fragments_started": counters.fragments_started.load(Ordering::Relaxed), + "fragments_completed": counters.fragments_completed.load(Ordering::Relaxed), + "batch_futures_emitted": counters.batch_futures_emitted.load(Ordering::Relaxed), + "batch_futures_received": counters.batch_futures_received.load(Ordering::Relaxed), + "batches_completed": counters.batches_completed.load(Ordering::Relaxed), + "decode_in_flight": decode_in_flight, + "channel_buffered": channel_buffered, + "open_reader_seconds_total": ns_to_seconds(counters.open_reader_ns.load(Ordering::Relaxed)), + "read_stream_create_seconds_total": ns_to_seconds(counters.read_stream_create_ns.load(Ordering::Relaxed)), + "next_batch_poll_seconds_total": ns_to_seconds(counters.next_batch_poll_ns.load(Ordering::Relaxed)), + "channel_send_wait_seconds_total": ns_to_seconds(counters.channel_send_wait_ns.load(Ordering::Relaxed)), + "decode_seconds_total": ns_to_seconds(counters.decode_ns.load(Ordering::Relaxed)), + "raw_reassemble_seconds_total": ns_to_seconds(counters.raw_reassemble_ns.load(Ordering::Relaxed)), + "scheduler": diagnostics_json(diagnostics), + }) +} + +async fn run_scanner_case( + config: &Config, + io_buffer_gib: Option, + scheduler_diagnostics: &SchedulerDiagnosticsCollector, +) -> Result { + scheduler_diagnostics.clear(); + let dataset = Arc::new( + DatasetBuilder::from_uri(&config.uri) + .with_version(config.dataset_version) + .load() + .await?, + ); + + let mut remaining_rows = config.limit_rows; + let mut planned_fragments = 0usize; + let mut planned_rows = 0u64; + for fragment in dataset.fragments().iter() { + if remaining_rows == 0 { + break; + } + let fragment_rows = fragment + .num_rows() + .ok_or_else(|| format!("fragment {} is missing num_rows", fragment.id))? + as u64; + let rows = fragment_rows.min(remaining_rows); + planned_fragments += 1; + planned_rows += rows; + remaining_rows -= rows; + } + if planned_fragments == 0 { + return Err("no fragments selected".into()); + } + + let counters = Arc::new(SharedCounters::default()); + let stats_holder = ExecutionStatsHolder::default(); + let cpu_before = read_cpu_sample(); + let started = Instant::now(); + + let mut scanner = dataset.scan(); + if let Some(columns) = config.columns.as_ref() { + scanner.project(columns)?; + } + scanner + .batch_size(config.batch_size as usize) + .scan_in_order(false) + .scan_stats_callback(stats_holder.get_setter()); + if config.batch_concurrency > 0 { + scanner + .batch_readahead(config.batch_concurrency) + .target_parallelism(config.batch_concurrency); + } + if config.fragment_concurrency > 0 { + scanner.fragment_readahead(config.fragment_concurrency); + } + if let Some(file_reader_options) = file_reader_options(config) { + scanner.with_file_reader_options(file_reader_options); + } + if let Some(batch_size_bytes) = config.batch_size_bytes { + scanner.batch_size_bytes(batch_size_bytes); + } + if let Some(io_buffer_gib) = io_buffer_gib { + scanner.io_buffer_size(io_buffer_gib * GIB); + } + let limit_rows = i64::try_from(config.limit_rows) + .map_err(|_| "--limit-rows is too large for scanner limit")?; + scanner.limit(Some(limit_rows), None)?; + + let mut stream = scanner.try_into_stream().await?; + let mut rows = 0u64; + let mut batches = 0u64; + let mut arrow_bytes = 0u64; + let mut samples = Vec::new(); + let mut sample_interval = tokio::time::interval(Duration::from_millis(config.sample_ms)); + sample_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut last_sample = LastSample { + elapsed: Duration::default(), + scheduler_stats: ScanStats::default(), + rows: 0, + batches: 0, + arrow_bytes: 0, + }; + loop { + tokio::select! { + maybe_batch = stream.next() => { + let Some(batch) = maybe_batch else { + break; + }; + let batch = batch?; + let batch_bytes = if config.skip_batch_byte_accounting { + 0 + } else { + batch.get_array_memory_size() as u64 + }; + rows += batch.num_rows() as u64; + batches += 1; + arrow_bytes += batch_bytes; + counters.batches_completed.fetch_add(1, Ordering::Relaxed); + counters + .rows_completed + .fetch_add(batch.num_rows() as u64, Ordering::Relaxed); + counters.arrow_bytes.fetch_add(batch_bytes, Ordering::Relaxed); + } + _ = sample_interval.tick() => { + samples.push(sample_json( + started, + counters.as_ref(), + scheduler_diagnostics.snapshot(io_buffer_gib), + 0, + 0, + &mut last_sample, + )); + } + } + } + drop(stream); + + let elapsed = started.elapsed(); + let summary = stats_holder + .consume() + .ok_or("scanner execution stats callback did not run")?; + let scheduler_stats = scan_stats_from_execution_summary(&summary); + let mut final_diagnostics = scheduler_diagnostics.snapshot(io_buffer_gib); + final_diagnostics.stats = scheduler_stats; + let cpu_after = read_cpu_sample(); + samples.push(sample_json( + started, + counters.as_ref(), + final_diagnostics, + 0, + 0, + &mut last_sample, + )); + + Ok(CaseStats { + rows, + batches, + arrow_bytes, + planned_fragments, + planned_rows, + elapsed, + producer_finished_at: Some(elapsed), + peak_decode_in_flight: 0, + cpu_avg: cpu_before.zip(cpu_after).and_then(|(before, after)| { + let total = after.total.checked_sub(before.total)?; + let idle = after.idle.checked_sub(before.idle)?; + if total == 0 { + return None; + } + Some((total - idle) as f64 / total as f64 * 100.0) + }), + scheduler_diagnostics: final_diagnostics, + counters, + samples, + }) +} + +async fn run_dataset_take_case( + config: &Config, + scheduler_diagnostics: &SchedulerDiagnosticsCollector, +) -> Result { + scheduler_diagnostics.clear(); + let dataset = DatasetBuilder::from_uri(&config.uri) + .with_version(config.dataset_version) + .load() + .await?; + let projection = Arc::new(projected_schema(dataset.schema(), &config.columns)?); + let total_rows = dataset + .fragments() + .iter() + .map(|fragment| { + fragment + .num_rows() + .map(|rows| rows as u64) + .ok_or_else(|| format!("fragment {} is missing num_rows", fragment.id)) + }) + .collect::, _>>()? + .into_iter() + .sum::(); + if total_rows == 0 { + return Err("dataset has no rows".into()); + } + + let counters = Arc::new(SharedCounters::default()); + let cpu_before = read_cpu_sample(); + let started = Instant::now(); + let mut rows = 0u64; + let mut batches = 0u64; + let mut arrow_bytes = 0u64; + const STRIDE: u64 = 104_729; + + for repetition in 0..config.take_repetitions { + let row_ids = (0..config.limit_rows) + .map(|offset| { + repetition + .wrapping_mul(STRIDE) + .wrapping_add(offset.wrapping_mul(STRIDE)) + % total_rows + }) + .collect::>(); + let batch = dataset + .take(&row_ids, ProjectionRequest::Schema(projection.clone())) + .await?; + if batch.num_rows() as u64 != config.limit_rows { + return Err(format!( + "take_rows returned {} rows, expected {}", + batch.num_rows(), + config.limit_rows + ) + .into()); + } + black_box(&batch); + rows += batch.num_rows() as u64; + batches += 1; + let batch_bytes = if config.skip_batch_byte_accounting { + 0 + } else { + batch.get_array_memory_size() as u64 + }; + arrow_bytes += batch_bytes; + counters.batches_completed.fetch_add(1, Ordering::Relaxed); + counters + .rows_completed + .fetch_add(batch.num_rows() as u64, Ordering::Relaxed); + counters + .arrow_bytes + .fetch_add(batch_bytes, Ordering::Relaxed); + } + + let elapsed = started.elapsed(); + let cpu_after = read_cpu_sample(); + Ok(CaseStats { + rows, + batches, + arrow_bytes, + planned_fragments: dataset.fragments().len(), + planned_rows: total_rows, + elapsed, + producer_finished_at: Some(elapsed), + peak_decode_in_flight: 0, + cpu_avg: cpu_before.zip(cpu_after).and_then(|(before, after)| { + let total = after.total.checked_sub(before.total)?; + let idle = after.idle.checked_sub(before.idle)?; + if total == 0 { + return None; + } + Some((total - idle) as f64 / total as f64 * 100.0) + }), + scheduler_diagnostics: scheduler_diagnostics.snapshot(None), + counters, + samples: Vec::new(), + }) +} + +enum RawInFlight { + Unordered(FuturesUnordered>>), + Ordered(FuturesOrdered>>), +} + +impl RawInFlight { + fn new(mode: RawCompletionMode) -> Self { + match mode { + RawCompletionMode::Unordered => Self::Unordered(FuturesUnordered::new()), + RawCompletionMode::Ordered => Self::Ordered(FuturesOrdered::new()), + } + } + + fn len(&self) -> usize { + match self { + Self::Unordered(in_flight) => in_flight.len(), + Self::Ordered(in_flight) => in_flight.len(), + } + } + + fn is_empty(&self) -> bool { + match self { + Self::Unordered(in_flight) => in_flight.is_empty(), + Self::Ordered(in_flight) => in_flight.is_empty(), + } + } + + fn push(&mut self, future: BoxFuture<'static, lance_core::Result>) { + match self { + Self::Unordered(in_flight) => in_flight.push(future), + Self::Ordered(in_flight) => in_flight.push_back(future), + } + } + + async fn next(&mut self) -> Option> { + match self { + Self::Unordered(in_flight) => in_flight.next().await, + Self::Ordered(in_flight) => in_flight.next().await, + } + } +} + +fn raw_read_future( + file_scheduler: FileScheduler, + range: Range, + priority: u64, + raw_submit_mode: RawSubmitMode, + read_chunk_size: u64, + counters: Arc, +) -> BoxFuture<'static, lance_core::Result> { + async move { + match raw_submit_mode { + RawSubmitMode::Single => { + let bytes = file_scheduler.submit_single(range, priority).await?; + Ok(bytes.len()) + } + RawSubmitMode::SplitNoConcat => { + let ranges = split_range_by_size(range, read_chunk_size); + let bytes = file_scheduler.submit_request(ranges, priority).await?; + Ok(bytes.iter().map(bytes::Bytes::len).sum()) + } + RawSubmitMode::SplitConcat => { + let ranges = split_range_by_size(range, read_chunk_size); + let bytes = file_scheduler.submit_request(ranges, priority).await?; + let reassemble_started = Instant::now(); + let total_size = bytes.iter().map(bytes::Bytes::len).sum(); + let mut combined = Vec::with_capacity(total_size); + for chunk in bytes { + combined.extend_from_slice(&chunk); + } + add_duration(&counters.raw_reassemble_ns, reassemble_started.elapsed()); + let len = combined.len(); + black_box(&combined); + Ok(len) + } + } + } + .boxed() +} + +fn split_range_by_size(range: Range, chunk_size: u64) -> Vec> { + let range_size = range.end - range.start; + if range_size <= chunk_size { + return vec![range]; + } + + let num_chunks = range_size.div_ceil(chunk_size); + let per_chunk = range_size / num_chunks; + let mut ranges = Vec::with_capacity(num_chunks as usize); + for idx in 0..num_chunks { + let start = range.start + idx * per_chunk; + let end = if idx == num_chunks - 1 { + range.end + } else { + start + per_chunk + }; + ranges.push(start..end); + } + ranges +} + +fn push_split_planned_ranges( + planned: &mut Vec<(usize, Range)>, + file_idx: usize, + range: Range, + chunk_size: u64, + remaining: &mut u64, +) { + let mut start = range.start; + while start < range.end && *remaining > 0 { + let bytes_to_read = chunk_size.min(range.end - start).min(*remaining); + if bytes_to_read == 0 { + break; + } + let end = start + bytes_to_read; + planned.push((file_idx, start..end)); + *remaining -= bytes_to_read; + start = end; + } +} + +fn push_split_ranges(ranges: &mut Vec>, range: Range, chunk_size: u64) { + let mut start = range.start; + while start < range.end { + let bytes_to_read = chunk_size.min(range.end - start); + if bytes_to_read == 0 { + break; + } + let end = start + bytes_to_read; + ranges.push(start..end); + start = end; + } +} + +async fn run_scheduler_raw_case( + config: &Config, + io_buffer_gib: Option, + scheduler_diagnostics: &SchedulerDiagnosticsCollector, +) -> Result { + scheduler_diagnostics.clear(); + let target_bytes = config + .target_bytes + .ok_or("--target-bytes is required for --backend scheduler-raw")?; + let dataset = Arc::new( + DatasetBuilder::from_uri(&config.uri) + .with_version(config.dataset_version) + .load() + .await?, + ); + let (object_store, _) = LanceObjectStore::from_uri(&config.uri).await?; + let scheduler_config = io_buffer_gib + .map(|gib| SchedulerConfig::new(gib * GIB)) + .unwrap_or_else(|| SchedulerConfig::max_bandwidth(object_store.as_ref())); + let scheduler = ScanScheduler::new(object_store, scheduler_config); + + let (selected_files, planned) = match config.raw_range_mode { + RawRangeMode::FileSequential => { + let mut selected_files = Vec::new(); + let mut selected_file_bytes = 0u64; + for fragment in dataset.fragments().iter() { + for data_file in &fragment.files { + if data_file.base_id.is_some() { + continue; + } + let Some(file_size) = data_file.file_size_bytes.get() else { + continue; + }; + let path = dataset.data_dir().join(data_file.path.as_str()); + let file_scheduler = scheduler + .open_file_with_priority(&path, 0, &data_file.file_size_bytes) + .await?; + selected_file_bytes += file_size.get(); + selected_files.push((file_scheduler, file_size.get())); + if selected_file_bytes >= target_bytes { + break; + } + } + if selected_file_bytes >= target_bytes { + break; + } + } + if selected_files.is_empty() { + return Err("scheduler-raw found no data files with known sizes".into()); + } + + let mut offsets = vec![0u64; selected_files.len()]; + let mut planned = Vec::new(); + let mut remaining = target_bytes; + let mut file_idx = 0usize; + while remaining > 0 { + let idx = file_idx % selected_files.len(); + let file_size = selected_files[idx].1; + if offsets[idx] >= file_size { + offsets[idx] = 0; + } + let available = file_size - offsets[idx]; + let bytes_to_read = config.raw_range_size_bytes.min(available).min(remaining); + let start = offsets[idx]; + let end = start + bytes_to_read; + planned.push((idx, start..end)); + offsets[idx] = end; + remaining -= bytes_to_read; + file_idx += 1; + } + (selected_files, planned) + } + RawRangeMode::MetadataPages | RawRangeMode::MetadataPagesRoundRobin => { + let mut selected_files = Vec::new(); + let mut per_file_ranges = Vec::>>::new(); + let mut candidate_bytes = 0u64; + + 'fragments: for fragment in dataset.fragments().iter() { + for data_file in &fragment.files { + if data_file.base_id.is_some() { + continue; + } + if data_file.file_size_bytes.get().is_none() { + continue; + } + let path = dataset.data_dir().join(data_file.path.as_str()); + let file_scheduler = scheduler + .open_file_with_priority(&path, 0, &data_file.file_size_bytes) + .await?; + let metadata = LanceFileReader::read_all_metadata(&file_scheduler).await?; + let mut file_ranges = Vec::new(); + + let raw_column_indices = config + .raw_column_indices + .clone() + .unwrap_or_else(|| (0..metadata.column_infos.len() as u32).collect()); + for column_index in raw_column_indices { + let column_info = metadata + .column_infos + .get(column_index as usize) + .ok_or_else(|| { + format!( + "raw metadata-pages requested column index {column_index} but file has {} columns", + metadata.column_infos.len() + ) + })?; + for page in column_info.page_infos.iter() { + for (offset, size) in page.buffer_offsets_and_sizes.iter() { + if *size == 0 { + continue; + } + push_split_ranges( + &mut file_ranges, + *offset..(*offset + *size), + config.raw_range_size_bytes, + ); + } + } + } + + if !file_ranges.is_empty() { + candidate_bytes += file_ranges + .iter() + .map(|range| range.end - range.start) + .sum::(); + selected_files.push(( + file_scheduler, + data_file.file_size_bytes.get().unwrap().get(), + )); + per_file_ranges.push(file_ranges); + if candidate_bytes >= target_bytes { + break 'fragments; + } + } + } + } + if selected_files.is_empty() || per_file_ranges.is_empty() { + return Err("scheduler-raw metadata-pages found no readable page buffers".into()); + } + + let mut planned = Vec::new(); + let mut remaining = target_bytes; + match config.raw_range_mode { + RawRangeMode::MetadataPages => { + 'ranges: for (file_idx, ranges) in per_file_ranges.iter().enumerate() { + for range in ranges { + push_split_planned_ranges( + &mut planned, + file_idx, + range.clone(), + config.raw_range_size_bytes, + &mut remaining, + ); + if remaining == 0 { + break 'ranges; + } + } + } + } + RawRangeMode::MetadataPagesRoundRobin => { + let mut positions = vec![0usize; per_file_ranges.len()]; + while remaining > 0 { + let mut made_progress = false; + for (file_idx, ranges) in per_file_ranges.iter().enumerate() { + if positions[file_idx] >= ranges.len() { + continue; + } + let range = ranges[positions[file_idx]].clone(); + positions[file_idx] += 1; + made_progress = true; + push_split_planned_ranges( + &mut planned, + file_idx, + range, + config.raw_range_size_bytes, + &mut remaining, + ); + if remaining == 0 { + break; + } + } + if !made_progress { + break; + } + } + } + RawRangeMode::FileSequential => unreachable!(), + } + + if remaining > 0 { + return Err(format!( + "scheduler-raw metadata-pages planned {} bytes but target is {target_bytes}", + target_bytes - remaining + ) + .into()); + } + (selected_files, planned) + } + }; + let planned_bytes = planned + .iter() + .map(|(_, range)| range.end - range.start) + .sum::(); + + let counters = Arc::new(SharedCounters::default()); + let cpu_before = read_cpu_sample(); + let started = Instant::now(); + let mut samples = Vec::new(); + let mut sample_interval = tokio::time::interval(Duration::from_millis(config.sample_ms)); + sample_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut last_sample = LastSample { + elapsed: Duration::default(), + scheduler_stats: ScanStats::default(), + rows: 0, + batches: 0, + arrow_bytes: 0, + }; + let mut in_flight = RawInFlight::new(config.raw_completion_mode); + let mut next_range = 0usize; + let read_chunk_size = config.read_chunk_size.unwrap_or(DEFAULT_READ_CHUNK_SIZE); + while next_range < planned.len() && in_flight.len() < config.batch_concurrency { + let (idx, range) = planned[next_range].clone(); + in_flight.push(raw_read_future( + selected_files[idx].0.clone(), + range, + next_range as u64, + config.raw_submit_mode, + read_chunk_size, + counters.clone(), + )); + next_range += 1; + } + + let mut bytes_read = 0u64; + let mut requests_completed = 0u64; + while !in_flight.is_empty() { + tokio::select! { + maybe_bytes = in_flight.next() => { + let bytes = maybe_bytes.expect("raw read future disappeared")?; + let bytes = bytes as u64; + bytes_read += bytes; + requests_completed += 1; + counters.batches_completed.fetch_add(1, Ordering::Relaxed); + counters.arrow_bytes.fetch_add(bytes, Ordering::Relaxed); + if next_range < planned.len() { + let (idx, range) = planned[next_range].clone(); + in_flight.push(raw_read_future( + selected_files[idx].0.clone(), + range, + next_range as u64, + config.raw_submit_mode, + read_chunk_size, + counters.clone(), + )); + next_range += 1; + } + } + _ = sample_interval.tick() => { + samples.push(sample_json( + started, + counters.as_ref(), + scheduler_diagnostics.snapshot(io_buffer_gib), + in_flight.len(), + planned.len().saturating_sub(next_range), + &mut last_sample, + )); + } + } + } + counters.arrow_bytes.store(bytes_read, Ordering::Relaxed); + let mut final_diagnostics = scheduler_diagnostics.snapshot(io_buffer_gib); + final_diagnostics.stats = scheduler.stats(); + samples.push(sample_json( + started, + counters.as_ref(), + final_diagnostics, + in_flight.len(), + 0, + &mut last_sample, + )); + let elapsed = started.elapsed(); + let cpu_after = read_cpu_sample(); + + Ok(CaseStats { + rows: 0, + batches: requests_completed, + arrow_bytes: bytes_read, + planned_fragments: selected_files.len(), + planned_rows: planned_bytes, + elapsed, + producer_finished_at: Some(elapsed), + peak_decode_in_flight: config.batch_concurrency, + cpu_avg: cpu_before.zip(cpu_after).and_then(|(before, after)| { + let total = after.total.checked_sub(before.total)?; + let idle = after.idle.checked_sub(before.idle)?; + if total == 0 { + return None; + } + Some((total - idle) as f64 / total as f64 * 100.0) + }), + scheduler_diagnostics: final_diagnostics, + counters, + samples, + }) +} + +async fn run_case( + config: &Config, + io_buffer_gib: Option, + scheduler_diagnostics: &SchedulerDiagnosticsCollector, +) -> Result { + scheduler_diagnostics.clear(); + let dataset = Arc::new( + DatasetBuilder::from_uri(&config.uri) + .with_version(config.dataset_version) + .load() + .await?, + ); + let projection = Arc::new(projected_schema(dataset.schema(), &config.columns)?); + let (object_store, _) = LanceObjectStore::from_uri(&config.uri).await?; + let scheduler_config = io_buffer_gib + .map(|gib| SchedulerConfig::new(gib * GIB)) + .unwrap_or_else(|| SchedulerConfig::max_bandwidth(object_store.as_ref())); + let scheduler = ScanScheduler::new(object_store, scheduler_config); + + let mut planned = Vec::new(); + let mut remaining_rows = config.limit_rows; + for fragment in dataset.fragments().iter() { + if remaining_rows == 0 { + break; + } + let fragment_rows = fragment + .num_rows() + .ok_or_else(|| format!("fragment {} is missing num_rows", fragment.id))? + as u64; + let rows = fragment_rows.min(remaining_rows); + planned.push((fragment.clone(), rows)); + remaining_rows -= rows; + } + if planned.is_empty() { + return Err("no fragments selected".into()); + } + let planned_rows: u64 = planned.iter().map(|(_, rows)| *rows).sum(); + let planned_fragments = planned.len(); + + let counters = Arc::new(SharedCounters::default()); + let cpu_before = read_cpu_sample(); + let started = Instant::now(); + + let (tx, mut rx) = tokio::sync::mpsc::channel::< + BoxFuture<'static, lance_core::Result>, + >(config.batch_concurrency * 2); + let producer = if config.detach_fragment_streams { + tokio::spawn({ + let dataset = dataset.clone(); + let projection = projection.clone(); + let scheduler = scheduler.clone(); + let counters = counters.clone(); + let batch_size = config.batch_size; + let file_reader_options = file_reader_options(config); + let fragment_concurrency = config.fragment_concurrency; + async move { + let drainers = futures::stream::iter(planned.into_iter().enumerate()) + .map({ + move |(priority, (fragment, rows))| { + let dataset = dataset.clone(); + let projection = projection.clone(); + let scheduler = scheduler.clone(); + let tx = tx.clone(); + let counters = counters.clone(); + let file_reader_options = file_reader_options.clone(); + async move { + counters.fragments_started.fetch_add(1, Ordering::Relaxed); + let file_fragment = FileFragment::new(dataset, fragment); + let read_config = FragReadConfig::default() + .with_scan_scheduler(scheduler) + .with_reader_priority(priority as u32); + let read_config = if let Some(file_reader_options) = + file_reader_options.clone() + { + read_config.with_file_reader_options(file_reader_options) + } else { + read_config + }; + + let open_started = Instant::now(); + let reader = + file_fragment.open(projection.as_ref(), read_config).await?; + add_duration(&counters.open_reader_ns, open_started.elapsed()); + + let create_stream_started = Instant::now(); + let mut read_stream = + reader.read_ranges(vec![0..rows].into(), batch_size).await?; + add_duration( + &counters.read_stream_create_ns, + create_stream_started.elapsed(), + ); + + let drainer = tokio::spawn(async move { + loop { + let next_started = Instant::now(); + let maybe_batch_fut = read_stream.next().await; + add_duration( + &counters.next_batch_poll_ns, + next_started.elapsed(), + ); + let Some(batch_fut) = maybe_batch_fut else { + break; + }; + counters + .batch_futures_emitted + .fetch_add(1, Ordering::Relaxed); + let send_started = Instant::now(); + tx.send(batch_fut) + .await + .map_err(|_| "batch consumer dropped")?; + add_duration( + &counters.channel_send_wait_ns, + send_started.elapsed(), + ); + } + counters.fragments_completed.fetch_add(1, Ordering::Relaxed); + Ok::<_, Error>(()) + }); + Ok::<_, Error>(drainer) + } + } + }) + .buffer_unordered(fragment_concurrency) + .try_collect::>() + .await?; + + for drainer in drainers { + drainer.await.map_err(Error::from)??; + } + Ok::<_, Error>(()) + } + }) + } else { + tokio::spawn({ + let dataset = dataset.clone(); + let projection = projection.clone(); + let scheduler = scheduler.clone(); + let counters = counters.clone(); + let batch_size = config.batch_size; + let file_reader_options = file_reader_options(config); + let fragment_concurrency = config.fragment_concurrency; + async move { + futures::stream::iter(planned.into_iter().enumerate()) + .map({ + move |(priority, (fragment, rows))| { + let dataset = dataset.clone(); + let projection = projection.clone(); + let scheduler = scheduler.clone(); + let tx = tx.clone(); + let counters = counters.clone(); + let file_reader_options = file_reader_options.clone(); + async move { + counters.fragments_started.fetch_add(1, Ordering::Relaxed); + let file_fragment = FileFragment::new(dataset, fragment); + let read_config = FragReadConfig::default() + .with_scan_scheduler(scheduler) + .with_reader_priority(priority as u32); + let read_config = if let Some(file_reader_options) = + file_reader_options.clone() + { + read_config.with_file_reader_options(file_reader_options) + } else { + read_config + }; + + let open_started = Instant::now(); + let reader = + file_fragment.open(projection.as_ref(), read_config).await?; + add_duration(&counters.open_reader_ns, open_started.elapsed()); + + let create_stream_started = Instant::now(); + let mut read_stream = + reader.read_ranges(vec![0..rows].into(), batch_size).await?; + add_duration( + &counters.read_stream_create_ns, + create_stream_started.elapsed(), + ); + + loop { + let next_started = Instant::now(); + let maybe_batch_fut = read_stream.next().await; + add_duration( + &counters.next_batch_poll_ns, + next_started.elapsed(), + ); + let Some(batch_fut) = maybe_batch_fut else { + break; + }; + counters + .batch_futures_emitted + .fetch_add(1, Ordering::Relaxed); + let send_started = Instant::now(); + tx.send(batch_fut) + .await + .map_err(|_| "batch consumer dropped")?; + add_duration( + &counters.channel_send_wait_ns, + send_started.elapsed(), + ); + } + counters.fragments_completed.fetch_add(1, Ordering::Relaxed); + Ok::<_, Error>(()) + } + } + }) + .buffer_unordered(fragment_concurrency) + .try_collect::>() + .await?; + Ok::<_, Error>(()) + } + }) + }; + + let mut in_flight = FuturesUnordered::new(); + let skip_batch_byte_accounting = config.skip_batch_byte_accounting; + let drop_read_tasks = config.drop_read_tasks; + let mut producer_done = false; + let mut producer_finished_at = None; + let mut rows = 0u64; + let mut batches = 0u64; + let mut arrow_bytes = 0u64; + let mut peak_decode_in_flight = 0usize; + let mut samples = Vec::new(); + let mut sample_interval = tokio::time::interval(Duration::from_millis(config.sample_ms)); + sample_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut last_sample = LastSample { + elapsed: Duration::default(), + scheduler_stats: ScanStats::default(), + rows: 0, + batches: 0, + arrow_bytes: 0, + }; + + loop { + if producer_done && in_flight.is_empty() { + break; + } + tokio::select! { + maybe_batch_fut = rx.recv(), if !producer_done && in_flight.len() < config.batch_concurrency => { + if let Some(batch_fut) = maybe_batch_fut { + counters.batch_futures_received.fetch_add(1, Ordering::Relaxed); + if drop_read_tasks { + drop(batch_fut); + counters.batches_completed.fetch_add(1, Ordering::Relaxed); + batches += 1; + continue; + } + let counters_for_task = counters.clone(); + in_flight.push(async move { + let decode_started = Instant::now(); + let batch = batch_fut.await?; + let batch_bytes = if skip_batch_byte_accounting { + 0 + } else { + batch.get_array_memory_size() as u64 + }; + add_duration(&counters_for_task.decode_ns, decode_started.elapsed()); + counters_for_task.batches_completed.fetch_add(1, Ordering::Relaxed); + counters_for_task.rows_completed.fetch_add(batch.num_rows() as u64, Ordering::Relaxed); + counters_for_task.arrow_bytes.fetch_add(batch_bytes, Ordering::Relaxed); + Ok::<_, lance_core::Error>((batch, batch_bytes)) + }); + peak_decode_in_flight = peak_decode_in_flight.max(in_flight.len()); + } else { + producer_done = true; + producer_finished_at = Some(started.elapsed()); + } + } + maybe_batch = in_flight.next(), if !in_flight.is_empty() => { + let (batch, batch_bytes) = maybe_batch.expect("in-flight batch future disappeared")?; + rows += batch.num_rows() as u64; + batches += 1; + arrow_bytes += batch_bytes; + } + _ = sample_interval.tick() => { + samples.push(sample_json( + started, + counters.as_ref(), + scheduler_diagnostics.snapshot(io_buffer_gib), + in_flight.len(), + rx.len(), + &mut last_sample, + )); + } + } + } + producer.await??; + + let mut final_diagnostics = scheduler_diagnostics.snapshot(io_buffer_gib); + final_diagnostics.stats = scheduler.stats(); + samples.push(sample_json( + started, + counters.as_ref(), + final_diagnostics, + in_flight.len(), + rx.len(), + &mut last_sample, + )); + + let elapsed = started.elapsed(); + let cpu_after = read_cpu_sample(); + + Ok(CaseStats { + rows, + batches, + arrow_bytes, + planned_fragments, + planned_rows, + elapsed, + producer_finished_at, + peak_decode_in_flight, + cpu_avg: cpu_before.zip(cpu_after).and_then(|(before, after)| { + let total = after.total.checked_sub(before.total)?; + let idle = after.idle.checked_sub(before.idle)?; + if total == 0 { + return None; + } + Some((total - idle) as f64 / total as f64 * 100.0) + }), + scheduler_diagnostics: final_diagnostics, + counters, + samples, + }) +} + +fn read_cpu_sample() -> Option { + let contents = fs::read_to_string("/proc/stat").ok()?; + let line = contents.lines().next()?; + let values = line + .split_whitespace() + .skip(1) + .map(|value| value.parse::()) + .collect::, _>>() + .ok()?; + if values.len() < 5 { + return None; + } + + let idle = values[3] + values[4]; + let total = values.iter().sum(); + Some(CpuSample { idle, total }) +} + +fn now_unix_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs()) + .unwrap_or_default() +} + +fn current_commit() -> String { + option_env!("LANCE_BENCH_COMMIT") + .or_else(|| option_env!("GIT_COMMIT")) + .unwrap_or("unknown") + .to_string() +} + +fn env_var(name: &str) -> Option { + env::var(name).ok() +} + +#[tokio::main] +async fn main() -> Result<()> { + let config = parse_args()?; + if config.describe_layout { + return describe_layout(&config).await; + } + let scheduler_diagnostics = SchedulerDiagnosticsCollector::default(); + let subscriber = tracing_subscriber::registry().with(SchedulerDiagnosticsLayer::new( + scheduler_diagnostics.clone(), + )); + tracing::subscriber::set_global_default(subscriber).map_err(|error| { + std::io::Error::other(format!( + "failed to install scheduler diagnostics subscriber: {error}" + )) + })?; + + fs::create_dir_all(&config.out_dir)?; + let output_path = format!( + "{}/s3_file_reader_diagnostics_{}.jsonl", + config.out_dir, + now_unix_secs() + ); + let mut jsonl = String::new(); + let commit = current_commit(); + let instance = env_var("EC2_INSTANCE_TYPE").unwrap_or_else(|| "unknown".to_string()); + let region = env_var("AWS_REGION") + .or_else(|| env_var("AWS_DEFAULT_REGION")) + .unwrap_or_else(|| "unknown".to_string()); + let projection = projection_name(&config.columns); + + for io_buffer_gib in &config.io_buffer_gib { + println!( + "running case={} backend={} projection={} limit_rows={} io_buffer_gib={} batch_size={} fragment_concurrency={} batch_concurrency={} sample_ms={}", + config.case_name, + config.backend.name(), + projection, + config.limit_rows, + io_buffer_gib + .map(|value| value.to_string()) + .unwrap_or_else(|| "auto".to_string()), + config.batch_size, + config.fragment_concurrency, + config.batch_concurrency, + config.sample_ms + ); + let stats = match config.backend { + Backend::FileReader => { + run_case(&config, *io_buffer_gib, &scheduler_diagnostics).await? + } + Backend::Scanner => { + run_scanner_case(&config, *io_buffer_gib, &scheduler_diagnostics).await? + } + Backend::SchedulerRaw => { + run_scheduler_raw_case(&config, *io_buffer_gib, &scheduler_diagnostics).await? + } + Backend::DatasetTake => run_dataset_take_case(&config, &scheduler_diagnostics).await?, + }; + let elapsed_secs = stats.elapsed.as_secs_f64(); + let scheduler_stats = stats.scheduler_diagnostics.stats; + let logical_gbps = if elapsed_secs > 0.0 { + stats.arrow_bytes as f64 * 8.0 / elapsed_secs / 1_000_000_000.0 + } else { + 0.0 + }; + let physical_gbps = if elapsed_secs > 0.0 { + scheduler_stats.bytes_read as f64 * 8.0 / elapsed_secs / 1_000_000_000.0 + } else { + 0.0 + }; + let rows_per_second = if elapsed_secs > 0.0 { + stats.rows as f64 / elapsed_secs + } else { + 0.0 + }; + let bytes_per_row = stats.arrow_bytes.checked_div(stats.rows).unwrap_or(0); + let avg_bytes_per_scheduler_request = scheduler_stats + .bytes_read + .checked_div(scheduler_stats.requests) + .unwrap_or(0); + let avg_bytes_per_scheduler_iop = scheduler_stats + .bytes_read + .checked_div(scheduler_stats.iops) + .unwrap_or(0); + let counters = stats.counters.as_ref(); + let record = json!({ + "case": config.case_name, + "instance": instance, + "region": region, + "layer": config.backend.layer(), + "backend": config.backend.name(), + "dataset_uri": config.uri, + "dataset_version": config.dataset_version, + "lance_commit": commit, + "projection": projection, + "limit_rows": config.limit_rows, + "target_bytes": config.target_bytes, + "raw_range_size_bytes": config.raw_range_size_bytes, + "raw_range_mode": config.raw_range_mode.name(), + "raw_column_indices": config.raw_column_indices.clone(), + "raw_submit_mode": config.raw_submit_mode.name(), + "raw_completion_mode": config.raw_completion_mode.name(), + "take_repetitions": config.take_repetitions, + "raw_read_chunk_size_bytes": config.read_chunk_size.unwrap_or(DEFAULT_READ_CHUNK_SIZE), + "planned_rows": stats.planned_rows, + "planned_fragments": stats.planned_fragments, + "rows": stats.rows, + "batches": stats.batches, + "batch_size": config.batch_size, + "batch_size_bytes": config.batch_size_bytes, + "skip_batch_byte_accounting": config.skip_batch_byte_accounting, + "read_chunk_size": config.read_chunk_size, + "fragment_concurrency": config.fragment_concurrency, + "batch_concurrency": config.batch_concurrency, + "detach_fragment_streams": config.detach_fragment_streams, + "drop_read_tasks": config.drop_read_tasks, + "sample_ms": config.sample_ms, + "io_buffer_bytes": io_buffer_gib.map(|value| value * GIB), + "io_buffer_mode": if io_buffer_gib.is_some() { "explicit" } else { "auto" }, + "lance_io_threads": env_var("LANCE_IO_THREADS").or_else(|| env_var("IO_THREADS")), + "lance_default_io_buffer_size": env_var("LANCE_DEFAULT_IO_BUFFER_SIZE"), + "lance_max_iop_size": env_var("LANCE_MAX_IOP_SIZE"), + "lance_use_lite_scheduler": env_var("LANCE_USE_LITE_SCHEDULER"), + "lance_inline_scheduling_threshold": env_var("LANCE_INLINE_SCHEDULING_THRESHOLD"), + "elapsed_seconds": elapsed_secs, + "producer_finished_seconds": stats.producer_finished_at.map(|duration| duration.as_secs_f64()), + "logical_gbps": logical_gbps, + "physical_gbps": physical_gbps, + "rows_per_second": rows_per_second, + "arrow_bytes": stats.arrow_bytes, + "bytes_per_row": bytes_per_row, + "avg_bytes_per_scheduler_request": avg_bytes_per_scheduler_request, + "avg_bytes_per_scheduler_iop": avg_bytes_per_scheduler_iop, + "scheduler_iops": scheduler_stats.iops, + "scheduler_requests": scheduler_stats.requests, + "scheduler_bytes_read": scheduler_stats.bytes_read, + "scheduler_diagnostics": diagnostics_json(stats.scheduler_diagnostics), + "fragments_started": counters.fragments_started.load(Ordering::Relaxed), + "fragments_completed": counters.fragments_completed.load(Ordering::Relaxed), + "batch_futures_emitted": counters.batch_futures_emitted.load(Ordering::Relaxed), + "batch_futures_received": counters.batch_futures_received.load(Ordering::Relaxed), + "batches_completed": counters.batches_completed.load(Ordering::Relaxed), + "peak_decode_in_flight": stats.peak_decode_in_flight, + "open_reader_seconds_total": ns_to_seconds(counters.open_reader_ns.load(Ordering::Relaxed)), + "read_stream_create_seconds_total": ns_to_seconds(counters.read_stream_create_ns.load(Ordering::Relaxed)), + "next_batch_poll_seconds_total": ns_to_seconds(counters.next_batch_poll_ns.load(Ordering::Relaxed)), + "channel_send_wait_seconds_total": ns_to_seconds(counters.channel_send_wait_ns.load(Ordering::Relaxed)), + "decode_seconds_total": ns_to_seconds(counters.decode_ns.load(Ordering::Relaxed)), + "raw_reassemble_seconds_total": ns_to_seconds(counters.raw_reassemble_ns.load(Ordering::Relaxed)), + "cpu_avg": stats.cpu_avg, + "samples": stats.samples, + }); + println!( + "case={} backend={} projection={} io_buffer_gib={} elapsed={:.3}s logical_gbps={:.2} physical_gbps={:.2} rows={} batches={} scheduler_bytes_read={} scheduler_iops={} scheduler_requests={} active_iops={} pending_iops={} cpu_avg={}", + config.case_name, + config.backend.name(), + projection, + io_buffer_gib + .map(|value| value.to_string()) + .unwrap_or_else(|| "auto".to_string()), + elapsed_secs, + logical_gbps, + physical_gbps, + stats.rows, + stats.batches, + scheduler_stats.bytes_read, + scheduler_stats.iops, + scheduler_stats.requests, + stats.scheduler_diagnostics.active_iops, + stats.scheduler_diagnostics.pending_iops, + stats + .cpu_avg + .map(|value| format!("{value:.1}%")) + .unwrap_or_else(|| "unknown".to_string()) + ); + jsonl.push_str(&serde_json::to_string(&record)?); + jsonl.push('\n'); + fs::write(&output_path, &jsonl)?; + } + + println!("wrote {output_path}"); + Ok(()) +} diff --git a/rust/lance/src/bin/fm_contains_bench.rs b/rust/lance/src/bin/fm_contains_bench.rs new file mode 100644 index 00000000000..b62ad8699f2 --- /dev/null +++ b/rust/lance/src/bin/fm_contains_bench.rs @@ -0,0 +1,610 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +#![allow(clippy::print_stdout)] + +use std::collections::HashMap; +use std::fs::File; +use std::io::{BufWriter, Write}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use arrow_array::{Array, LargeStringArray, RecordBatch, StringArray, UInt64Array}; +use clap::Parser; +use futures::future::join_all; +use lance::dataset::builder::DatasetBuilder; +use lance::dataset::{Dataset, ReadParams}; +use lance::index::DatasetIndexExt; +use lance::{Error, Result}; +use lance_io::object_store::{ObjectStoreParams, StorageOptionsAccessor}; +use serde::Serialize; +use tokio::sync::Semaphore; +use tokio::time::timeout; + +const ROW_ID_COLUMN: &str = "_rowid"; + +#[derive(Parser, Debug, Clone)] +#[command(author, version, about, long_about = None)] +struct Args { + #[arg( + long, + default_value = "az://datasets/mmlb/mmlb_100m_fts_en_fm_20260626.lance" + )] + uri: String, + + #[arg(long, default_value = "full_content_idx")] + index_name: String, + + #[arg(long, default_value = "full_content")] + text_column: String, + + #[arg(long, default_value = "summary_in_image")] + pattern_source_column: String, + + #[arg(long, default_value = "oailancepub")] + storage_account: String, + + #[arg(long, value_parser = parse_storage_option)] + storage_option: Vec<(String, String)>, + + #[arg(long, default_value_t = 1_099_511_627_776)] + index_cache_bytes: usize, + + #[arg(long, default_value_t = 8_589_934_592)] + metadata_cache_bytes: usize, + + #[arg(long, default_value_t = 5_000)] + sample_rows: usize, + + #[arg(long, default_value_t = 5)] + query_term_count: usize, + + #[arg(long, default_value_t = 100)] + k: i64, + + #[arg(long, default_value_t = 0)] + warmup: usize, + + #[arg(long, default_value_t = 4)] + queries: usize, + + #[arg(long, default_value = "1 2 4 8")] + threads: String, + + #[arg(long, default_value_t = 60.0)] + query_timeout_secs: f64, + + #[arg(long, default_value_t = false)] + skip_prewarm: bool, + + #[arg(long, default_value_t = true)] + explain: bool, + + #[arg(long)] + output_json: String, + + #[arg(long)] + output_csv: String, +} + +#[derive(Debug, Clone, Serialize)] +struct QueryPattern { + query_id: String, + pattern: String, + filter_sql: String, +} + +#[derive(Debug, Serialize)] +struct QueryResult { + query_id: String, + pattern: String, + filter_sql: String, + success: bool, + error: Option, + elapsed_ms: f64, + result_count: usize, + checksum: u64, +} + +#[derive(Debug, Serialize)] +struct BatchSummary { + round_type: String, + threads: usize, + queries: usize, + successes: usize, + success_rate: f64, + wall_s: f64, + qps: f64, + mean_ms: f64, + p50_ms: f64, + p95_ms: f64, + p99_ms: f64, + avg_result_count: f64, + max_result_count: usize, + checksum: u64, + results: Vec, +} + +#[derive(Debug, Serialize)] +struct BenchmarkOutput { + uri: String, + index_name: String, + text_column: String, + pattern_source_column: String, + sample_rows: usize, + query_term_count: usize, + k: i64, + warmup: usize, + queries: usize, + threads: Vec, + query_timeout_secs: f64, + index_cache_bytes: usize, + metadata_cache_bytes: usize, + prewarm_ms: Option, + explain_plan: Option, + summaries: Vec, +} + +fn parse_threads(value: &str) -> std::result::Result, String> { + let mut threads = Vec::new(); + for part in value.replace(',', " ").split_whitespace() { + let thread_count = part + .parse::() + .map_err(|err| format!("invalid thread count {part:?}: {err}"))?; + if thread_count == 0 { + return Err("thread counts must be greater than zero".to_string()); + } + threads.push(thread_count); + } + if threads.is_empty() { + return Err("at least one thread count is required".to_string()); + } + Ok(threads) +} + +fn parse_storage_option(value: &str) -> std::result::Result<(String, String), String> { + let Some((key, val)) = value.split_once('=') else { + return Err("storage options must be key=value".to_string()); + }; + if key.is_empty() { + return Err("storage option key cannot be empty".to_string()); + } + Ok((key.to_string(), val.to_string())) +} + +fn sql_quote(value: &str) -> String { + format!("'{}'", value.replace('\'', "''")) +} + +fn query_terms_from_text(text: &str, term_count: usize) -> Result { + let terms = text.split_whitespace().take(term_count).collect::>(); + if terms.len() < term_count { + return Err(Error::invalid_input(format!( + "sampled query text has {} terms, but requested {}", + terms.len(), + term_count + ))); + } + Ok(terms.join(" ")) +} + +fn string_value(batch: &RecordBatch, column: &str, row: usize) -> Result> { + let array = batch + .column_by_name(column) + .ok_or_else(|| Error::invalid_input(format!("column {column:?} not found")))?; + if array.is_null(row) { + return Ok(None); + } + if let Some(strings) = array.as_any().downcast_ref::() { + return Ok(Some(strings.value(row).to_string())); + } + if let Some(strings) = array.as_any().downcast_ref::() { + return Ok(Some(strings.value(row).to_string())); + } + Err(Error::invalid_input(format!( + "column {column:?} must be Utf8 or LargeUtf8, got {:?}", + array.data_type() + ))) +} + +async fn open_dataset(args: &Args) -> Result { + let mut options = HashMap::from([("account_name".to_string(), args.storage_account.clone())]); + for (key, val) in &args.storage_option { + options.insert(key.clone(), val.clone()); + } + let read_params = ReadParams { + index_cache_size_bytes: args.index_cache_bytes, + metadata_cache_size_bytes: args.metadata_cache_bytes, + store_options: Some(ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + options, + ))), + ..Default::default() + }), + ..Default::default() + }; + + DatasetBuilder::from_uri(&args.uri) + .with_read_params(read_params) + .load() + .await +} + +async fn sample_patterns(dataset: &Dataset, args: &Args) -> Result> { + let total = args.warmup + args.queries; + let mut scanner = dataset.scan(); + scanner + .project(&[args.pattern_source_column.as_str()])? + .limit(Some(args.sample_rows as i64), None)?; + let batch = scanner.try_into_batch().await?; + let mut patterns = Vec::with_capacity(total); + let mut seen = std::collections::HashSet::new(); + + for row in 0..batch.num_rows() { + let Some(text) = string_value(&batch, &args.pattern_source_column, row)? else { + continue; + }; + let pattern = query_terms_from_text(&text, args.query_term_count)?; + if seen.insert(pattern.clone()) { + let filter_sql = format!( + "contains({}, {})", + args.text_column, + sql_quote(pattern.as_str()) + ); + patterns.push(QueryPattern { + query_id: format!("Q{:06}", patterns.len() + 1), + pattern, + filter_sql, + }); + if patterns.len() == total { + return Ok(patterns); + } + } + } + + Err(Error::invalid_input(format!( + "sampled only {} unique patterns, need {}; increase --sample-rows", + patterns.len(), + total + ))) +} + +async fn explain_first_query( + dataset: &Dataset, + pattern: &QueryPattern, + args: &Args, +) -> Result { + let mut scanner = dataset.scan(); + scanner + .with_row_id() + .project::<&str>(&[])? + .filter(&pattern.filter_sql)? + .limit(Some(args.k), None)?; + scanner.explain_plan(false).await +} + +async fn execute_query( + dataset: Arc, + pattern: &QueryPattern, + args: &Args, +) -> Result { + let start = Instant::now(); + let mut scanner = dataset.scan(); + scanner + .with_row_id() + .project::<&str>(&[])? + .filter(&pattern.filter_sql)? + .limit(Some(args.k), None)?; + let batch = scanner.try_into_batch().await?; + let elapsed_ms = start.elapsed().as_secs_f64() * 1000.0; + let row_ids = batch + .column_by_name(ROW_ID_COLUMN) + .ok_or_else(|| Error::invalid_input(format!("{ROW_ID_COLUMN} column missing")))?; + let row_ids = row_ids + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::invalid_input(format!("{ROW_ID_COLUMN} must be UInt64")))?; + let mut checksum = 0_u64; + for row in 0..row_ids.len() { + if !row_ids.is_null(row) { + checksum ^= row_ids.value(row); + } + } + Ok(QueryResult { + query_id: pattern.query_id.clone(), + pattern: pattern.pattern.clone(), + filter_sql: pattern.filter_sql.clone(), + success: true, + error: None, + elapsed_ms, + result_count: batch.num_rows(), + checksum, + }) +} + +async fn run_one(dataset: Arc, pattern: QueryPattern, args: Arc) -> QueryResult { + let start = Instant::now(); + let result = if args.query_timeout_secs > 0.0 { + timeout( + Duration::from_secs_f64(args.query_timeout_secs), + execute_query(dataset, &pattern, &args), + ) + .await + .map_err(|_| { + Error::io(format!( + "query timed out after {}s", + args.query_timeout_secs + )) + }) + .and_then(|inner| inner) + } else { + execute_query(dataset, &pattern, &args).await + }; + + match result { + Ok(result) => result, + Err(err) => QueryResult { + query_id: pattern.query_id, + pattern: pattern.pattern, + filter_sql: pattern.filter_sql, + success: false, + error: Some(err.to_string()), + elapsed_ms: start.elapsed().as_secs_f64() * 1000.0, + result_count: 0, + checksum: 0, + }, + } +} + +async fn run_batch( + dataset: Arc, + args: Arc, + patterns: &[QueryPattern], + thread_count: usize, + round_type: &str, +) -> BatchSummary { + let wall_start = Instant::now(); + let semaphore = Arc::new(Semaphore::new(thread_count)); + let tasks = patterns + .iter() + .map(|pattern| { + let pattern = pattern.clone(); + let dataset = Arc::clone(&dataset); + let args = Arc::clone(&args); + let semaphore = Arc::clone(&semaphore); + tokio::spawn(async move { + let _permit = semaphore.acquire_owned().await.expect("semaphore closed"); + run_one(dataset, pattern, args).await + }) + }) + .collect::>(); + + let results = join_all(tasks) + .await + .into_iter() + .map(|result| match result { + Ok(query_result) => query_result, + Err(err) => QueryResult { + query_id: "join_error".to_string(), + pattern: String::new(), + filter_sql: String::new(), + success: false, + error: Some(err.to_string()), + elapsed_ms: 0.0, + result_count: 0, + checksum: 0, + }, + }) + .collect::>(); + summarize(round_type, thread_count, results, wall_start.elapsed()) +} + +fn percentile(values: &[f64], pct: usize) -> f64 { + if values.is_empty() { + return 0.0; + } + let mut ordered = values.to_vec(); + ordered.sort_by(|left, right| left.total_cmp(right)); + let idx = ((ordered.len() * pct).div_ceil(100)).saturating_sub(1); + ordered[idx.min(ordered.len() - 1)] +} + +fn summarize( + round_type: &str, + threads: usize, + results: Vec, + wall: Duration, +) -> BatchSummary { + let successes = results.iter().filter(|result| result.success).count(); + let latencies = results + .iter() + .filter(|result| result.success) + .map(|result| result.elapsed_ms) + .collect::>(); + let result_counts = results + .iter() + .filter(|result| result.success) + .map(|result| result.result_count) + .collect::>(); + let wall_s = wall.as_secs_f64(); + BatchSummary { + round_type: round_type.to_string(), + threads, + queries: results.len(), + successes, + success_rate: if results.is_empty() { + 0.0 + } else { + successes as f64 / results.len() as f64 + }, + wall_s, + qps: if wall_s == 0.0 { + 0.0 + } else { + successes as f64 / wall_s + }, + mean_ms: if latencies.is_empty() { + 0.0 + } else { + latencies.iter().sum::() / latencies.len() as f64 + }, + p50_ms: percentile(&latencies, 50), + p95_ms: percentile(&latencies, 95), + p99_ms: percentile(&latencies, 99), + avg_result_count: if result_counts.is_empty() { + 0.0 + } else { + result_counts.iter().sum::() as f64 / result_counts.len() as f64 + }, + max_result_count: result_counts.into_iter().max().unwrap_or(0), + checksum: results + .iter() + .filter(|result| result.success) + .map(|result| result.checksum) + .fold(0, |acc, checksum| acc ^ checksum), + results, + } +} + +fn write_csv(path: &str, summaries: &[BatchSummary]) -> Result<()> { + let file = File::create(path)?; + let mut writer = BufWriter::new(file); + writeln!( + writer, + "round_type,threads,queries,successes,success_rate,wall_s,qps,mean_ms,p50_ms,p95_ms,p99_ms,avg_result_count,max_result_count,checksum" + )?; + for summary in summaries { + writeln!( + writer, + "{},{},{},{},{:.6},{:.6},{:.6},{:.6},{:.6},{:.6},{:.6},{:.6},{},{}", + csv_escape(&summary.round_type), + summary.threads, + summary.queries, + summary.successes, + summary.success_rate, + summary.wall_s, + summary.qps, + summary.mean_ms, + summary.p50_ms, + summary.p95_ms, + summary.p99_ms, + summary.avg_result_count, + summary.max_result_count, + summary.checksum + )?; + } + Ok(()) +} + +fn csv_escape(value: &str) -> String { + if value.contains([',', '"', '\n']) { + format!("\"{}\"", value.replace('"', "\"\"")) + } else { + value.to_string() + } +} + +#[tokio::main] +async fn main() -> Result<()> { + let args = Arc::new(Args::parse()); + let thread_counts = parse_threads(&args.threads).map_err(Error::invalid_input)?; + if args.k <= 0 { + return Err(Error::invalid_input("--k must be greater than zero")); + } + if args.query_term_count == 0 { + return Err(Error::invalid_input( + "--query-term-count must be greater than zero", + )); + } + if args.query_timeout_secs < 0.0 { + return Err(Error::invalid_input( + "--query-timeout-secs must be non-negative", + )); + } + + println!("opening dataset {}", args.uri); + let dataset = Arc::new(open_dataset(&args).await?); + println!("dataset version {}", dataset.version().version); + + let patterns = sample_patterns(&dataset, &args).await?; + println!("sampled {} patterns", patterns.len()); + + let explain_plan = if args.explain { + Some(explain_first_query(&dataset, &patterns[0], &args).await?) + } else { + None + }; + + let prewarm_ms = if args.skip_prewarm { + None + } else { + println!("prewarming index {}", args.index_name); + let start = Instant::now(); + dataset.prewarm_index(&args.index_name).await?; + let elapsed_ms = start.elapsed().as_secs_f64() * 1000.0; + println!("prewarm finished in {:.3} ms", elapsed_ms); + Some(elapsed_ms) + }; + + let warmup_patterns = &patterns[..args.warmup]; + let query_patterns = &patterns[args.warmup..]; + let mut summaries = Vec::new(); + for thread_count in thread_counts.iter().copied() { + if !warmup_patterns.is_empty() { + summaries.push( + run_batch( + Arc::clone(&dataset), + Arc::clone(&args), + warmup_patterns, + thread_count, + "warmup", + ) + .await, + ); + } + let summary = run_batch( + Arc::clone(&dataset), + Arc::clone(&args), + query_patterns, + thread_count, + "test", + ) + .await; + println!( + "test threads={} queries={} successes={} wall_s={:.3} qps={:.3} mean_ms={:.3} p95_ms={:.3}", + summary.threads, + summary.queries, + summary.successes, + summary.wall_s, + summary.qps, + summary.mean_ms, + summary.p95_ms + ); + summaries.push(summary); + } + + let output = BenchmarkOutput { + uri: args.uri.clone(), + index_name: args.index_name.clone(), + text_column: args.text_column.clone(), + pattern_source_column: args.pattern_source_column.clone(), + sample_rows: args.sample_rows, + query_term_count: args.query_term_count, + k: args.k, + warmup: args.warmup, + queries: args.queries, + threads: thread_counts, + query_timeout_secs: args.query_timeout_secs, + index_cache_bytes: args.index_cache_bytes, + metadata_cache_bytes: args.metadata_cache_bytes, + prewarm_ms, + explain_plan, + summaries, + }; + + let mut json = serde_json::to_string_pretty(&output)?; + json.push('\n'); + std::fs::write(&args.output_json, json)?; + write_csv(&args.output_csv, &output.summaries)?; + Ok(()) +} diff --git a/rust/lance/src/bin/fm_index_tool.rs b/rust/lance/src/bin/fm_index_tool.rs new file mode 100644 index 00000000000..1eb03b78b26 --- /dev/null +++ b/rust/lance/src/bin/fm_index_tool.rs @@ -0,0 +1,161 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +#![allow(clippy::print_stdout)] + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Instant; + +use clap::{Parser, ValueEnum}; +use lance::dataset::ReadParams; +use lance::dataset::builder::DatasetBuilder; +use lance::index::{CreateIndexBuilder, DatasetIndexExt}; +use lance::{Error, Result}; +use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams}; +use lance_io::object_store::{ObjectStoreParams, StorageOptionsAccessor}; +use uuid::Uuid; + +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +struct Args { + #[arg(value_enum)] + action: Action, + + #[arg(long)] + uri: String, + + #[arg(long, default_value = "full_content_idx")] + index_name: String, + + #[arg(long, default_value = "full_content")] + column: String, + + #[arg(long, default_value = "oailancepub")] + storage_account: String, + + #[arg(long, value_parser = parse_storage_option)] + storage_option: Vec<(String, String)>, + + #[arg(long, default_value_t = 1)] + num_segments: u32, + + #[arg(long)] + index_uuid: Option, +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq, ValueEnum)] +enum Action { + List, + Drop, + Create, +} + +fn parse_storage_option(value: &str) -> std::result::Result<(String, String), String> { + let Some((key, val)) = value.split_once('=') else { + return Err("storage options must be key=value".to_string()); + }; + if key.is_empty() { + return Err("storage option key cannot be empty".to_string()); + } + Ok((key.to_string(), val.to_string())) +} + +async fn open_dataset(args: &Args) -> Result { + let mut options = HashMap::from([("account_name".to_string(), args.storage_account.clone())]); + for (key, val) in &args.storage_option { + options.insert(key.clone(), val.clone()); + } + let read_params = ReadParams { + store_options: Some(ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + options, + ))), + ..Default::default() + }), + ..Default::default() + }; + + DatasetBuilder::from_uri(&args.uri) + .with_read_params(read_params) + .load() + .await +} + +fn fm_params(num_segments: u32) -> ScalarIndexParams { + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::Fm); + if num_segments == 1 { + params + } else { + params.with_params(&serde_json::json!({ "num_segments": num_segments })) + } +} + +#[tokio::main] +async fn main() -> Result<()> { + let args = Args::parse(); + if args.num_segments == 0 { + return Err(Error::invalid_input( + "--num-segments must be greater than 0", + )); + } + + println!("opening dataset {}", args.uri); + let mut dataset = open_dataset(&args).await?; + println!("dataset version {}", dataset.version().version); + + match args.action { + Action::List => { + let indices = dataset.load_indices().await?; + println!("indices={}", indices.len()); + for index in indices.iter() { + println!( + "name={} uuid={} version={} fields={:?} fragments={}", + index.name, + index.uuid, + index.index_version, + index.fields, + index + .fragment_bitmap + .as_ref() + .map(|bitmap| bitmap.len()) + .unwrap_or(0) + ); + } + } + Action::Drop => { + let start = Instant::now(); + dataset.drop_index(&args.index_name).await?; + println!( + "dropped index {} in {:.3}s; new version {}", + args.index_name, + start.elapsed().as_secs_f64(), + dataset.version().version + ); + } + Action::Create => { + let params = fm_params(args.num_segments); + let start = Instant::now(); + let mut builder = CreateIndexBuilder::new( + &mut dataset, + &[args.column.as_str()], + lance_index::IndexType::Fm, + ¶ms, + ) + .name(args.index_name.clone()) + .replace(true); + if let Some(index_uuid) = args.index_uuid { + builder = builder.index_uuid(index_uuid); + } + let metadata = builder.await?; + println!( + "created index {} uuid={} in {:.3}s; new version {}", + metadata.name, + metadata.uuid, + start.elapsed().as_secs_f64(), + dataset.version().version + ); + } + } + Ok(()) +} diff --git a/rust/lance/src/blob.rs b/rust/lance/src/blob.rs index 58df42b5cd3..b112f011419 100644 --- a/rust/lance/src/blob.rs +++ b/rust/lance/src/blob.rs @@ -1,22 +1,42 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -//! Convenience builders for Lance blob v2 input columns. +//! Builders and file-level writer helpers for Lance blob v2 columns. //! -//! Blob v2 expects a column shaped as `Struct` and -//! tagged with `ARROW:extension:name = "lance.blob.v2"`. This module offers a -//! type-safe builder to construct that struct without manually wiring metadata +//! Logical blob input uses `Struct`. File-level blob +//! descriptors use a physical writer-side struct with `kind`, `blob_id`, and range fields. +use std::collections::{HashMap, HashSet}; use std::num::NonZeroUsize; -use std::sync::Arc; +use std::ops::Range; +use std::sync::{ + Arc, Mutex, + atomic::{AtomicU32, Ordering}, +}; -use arrow_array::{ArrayRef, StructArray, builder::LargeBinaryBuilder, builder::StringBuilder}; +use arrow_array::{ + Array, ArrayRef, StructArray, + builder::{LargeBinaryBuilder, PrimitiveBuilder, StringBuilder}, + cast::AsArray, + types::{UInt8Type, UInt32Type, UInt64Type}, +}; use arrow_buffer::NullBufferBuilder; -use arrow_schema::{DataType, Field}; +use arrow_schema::{DataType, Field, Fields}; +use bytes::Bytes; use lance_arrow::{ ARROW_EXT_NAME_KEY, BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, - BLOB_INLINE_SIZE_THRESHOLD_META_KEY, BLOB_V2_EXT_NAME, + BLOB_INLINE_SIZE_THRESHOLD_META_KEY, BLOB_V2_EXT_NAME, FieldExt, +}; +use lance_core::{ + datatypes::{BlobKind, Field as LanceField, Schema as LanceSchema}, + utils::blob::blob_path, +}; +use lance_io::{ + object_store::ObjectStore, + traits::{Reader, WriteExt, Writer}, }; +use object_store::path::Path; +use tokio::io::AsyncWriteExt; use crate::{Error, Result}; @@ -104,6 +124,788 @@ pub fn blob_field_with_options(name: &str, nullable: bool, options: BlobFieldOpt .with_metadata(metadata) } +fn prepared_blob_child_fields() -> Fields { + Fields::from(vec![ + Field::new("kind", DataType::UInt8, true), + Field::new("data", DataType::LargeBinary, true), + Field::new("uri", DataType::Utf8, true), + Field::new("blob_id", DataType::UInt32, true), + Field::new("blob_size", DataType::UInt64, true), + Field::new("position", DataType::UInt64, true), + ]) +} + +fn prepared_blob_field_with_metadata( + name: &str, + nullable: bool, + metadata: HashMap, +) -> Field { + Field::new( + name, + DataType::Struct(prepared_blob_child_fields()), + nullable, + ) + .with_metadata(metadata) +} + +fn logical_blob_lance_children() -> Result> { + [ + Field::new("data", DataType::LargeBinary, true), + Field::new("uri", DataType::Utf8, true), + ] + .iter() + .map(LanceField::try_from) + .collect() +} + +fn field_matches(field: &Field, name: &str, data_type: &DataType, nullable: bool) -> bool { + field.name() == name && field.data_type() == data_type && field.is_nullable() == nullable +} + +fn blob_v2_shape_error(field: &Field) -> Error { + Error::invalid_input(format!( + "Blob v2 field '{}' must use either logical struct \ + with optional position/size UInt64 fields or prepared struct", + field.name() + )) +} + +/// Returns true when `field` is the writer-side prepared blob v2 struct. +pub(crate) fn is_prepared_blob_v2_field(field: &Field) -> bool { + if !field.is_blob_v2() { + return false; + } + let DataType::Struct(fields) = field.data_type() else { + return false; + }; + let expected = prepared_blob_child_fields(); + fields.len() == expected.len() + && fields + .iter() + .zip(expected.iter()) + .all(|(actual, expected)| actual.as_ref() == expected.as_ref()) +} + +/// Returns true when `field` is the logical blob v2 input struct. +pub(crate) fn is_logical_blob_v2_field(field: &Field) -> bool { + if !field.is_blob_v2() { + return false; + } + let DataType::Struct(fields) = field.data_type() else { + return false; + }; + match fields.len() { + 2 => { + field_matches(fields[0].as_ref(), "data", &DataType::LargeBinary, true) + && field_matches(fields[1].as_ref(), "uri", &DataType::Utf8, true) + } + 4 => { + field_matches(fields[0].as_ref(), "data", &DataType::LargeBinary, true) + && field_matches(fields[1].as_ref(), "uri", &DataType::Utf8, true) + && fields[2].name() == "position" + && fields[2].data_type() == &DataType::UInt64 + && fields[3].name() == "size" + && fields[3].data_type() == &DataType::UInt64 + } + _ => false, + } +} + +fn normalize_prepared_blob_lance_field(field: &LanceField) -> Result { + if field.is_blob_v2() { + let arrow_field = Field::from(field); + if is_prepared_blob_v2_field(&arrow_field) { + let mut normalized = field.clone(); + let mut logical_children = logical_blob_lance_children()?; + for (logical_child, prepared_child) in + logical_children.iter_mut().zip(field.children.iter()) + { + logical_child.id = prepared_child.id; + logical_child.parent_id = field.id; + } + normalized.children = logical_children; + return Ok(normalized); + } + if is_logical_blob_v2_field(&arrow_field) { + return Ok(field.clone()); + } + return Err(blob_v2_shape_error(&arrow_field)); + } + + if field.children.is_empty() { + return Ok(field.clone()); + } + + let normalized_children = field + .children + .iter() + .map(normalize_prepared_blob_lance_field) + .collect::>>()?; + + Ok(LanceField { + children: normalized_children, + ..field.clone() + }) +} + +pub(crate) fn normalize_prepared_blob_schema(schema: &LanceSchema) -> Result { + let fields = schema + .fields + .iter() + .map(normalize_prepared_blob_lance_field) + .collect::>>()?; + Ok(LanceSchema { + fields, + metadata: schema.metadata.clone(), + }) +} + +#[derive(Clone, Debug)] +pub(crate) struct BlobIdAllocator { + inner: Arc, +} + +#[derive(Debug)] +struct BlobIdAllocatorInner { + next: AtomicU32, + used: Mutex>, +} + +impl BlobIdAllocator { + pub(crate) fn new(start: u32) -> Self { + Self { + inner: Arc::new(BlobIdAllocatorInner { + next: AtomicU32::new(start), + used: Mutex::new(HashSet::new()), + }), + } + } + + pub(crate) fn next(&self) -> Result { + loop { + let id = self.inner.next.load(Ordering::Relaxed); + if id == u32::MAX { + return Err(Error::invalid_input( + "Blob id allocator exhausted u32 id space", + )); + } + if self + .inner + .next + .compare_exchange(id, id + 1, Ordering::Relaxed, Ordering::Relaxed) + .is_err() + { + continue; + } + let mut used = + self.inner.used.lock().map_err(|_| { + Error::internal("Blob id allocator mutex was poisoned".to_string()) + })?; + if used.insert(id) { + return Ok(id); + } + } + } +} + +fn validate_blob_id(blob_id: u32) -> Result<()> { + if blob_id == 0 { + return Err(Error::invalid_input("Blob id 0 is reserved")); + } + Ok(()) +} + +fn validate_range(offset: u64, size: u64, object_size: u64, label: &str) -> Result<()> { + let end = offset.checked_add(size).ok_or_else(|| { + Error::invalid_input(format!( + "{label} range overflows u64: offset={offset}, size={size}" + )) + })?; + if end > object_size { + return Err(Error::invalid_input(format!( + "{label} range [{offset}, {end}) exceeds blob object size {object_size}" + ))); + } + Ok(()) +} + +fn validate_prepared_blob_value_array(field: &Field, array: &ArrayRef) -> Result<()> { + if !is_prepared_blob_v2_field(field) { + return Err(blob_v2_shape_error(field)); + } + + let struct_arr = array + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::invalid_input("Prepared blob column was not a struct array"))?; + let kind_col = struct_arr + .column_by_name("kind") + .ok_or_else(|| Error::invalid_input("Prepared blob struct missing `kind` field"))? + .as_primitive::(); + let data_col = struct_arr + .column_by_name("data") + .ok_or_else(|| Error::invalid_input("Prepared blob struct missing `data` field"))? + .as_binary::(); + let uri_col = struct_arr + .column_by_name("uri") + .ok_or_else(|| Error::invalid_input("Prepared blob struct missing `uri` field"))? + .as_string::(); + let blob_id_col = struct_arr + .column_by_name("blob_id") + .ok_or_else(|| Error::invalid_input("Prepared blob struct missing `blob_id` field"))? + .as_primitive::(); + let blob_size_col = struct_arr + .column_by_name("blob_size") + .ok_or_else(|| Error::invalid_input("Prepared blob struct missing `blob_size` field"))? + .as_primitive::(); + let position_col = struct_arr + .column_by_name("position") + .ok_or_else(|| Error::invalid_input("Prepared blob struct missing `position` field"))? + .as_primitive::(); + + for row in 0..struct_arr.len() { + if struct_arr.is_null(row) { + continue; + } + if kind_col.is_null(row) { + return Err(Error::invalid_input(format!( + "Prepared blob row {row} is non-null but `kind` is null" + ))); + } + + match BlobKind::try_from(kind_col.value(row))? { + BlobKind::Inline => { + if data_col.is_null(row) { + return Err(Error::invalid_input(format!( + "Prepared inline blob row {row} must set `data`" + ))); + } + } + BlobKind::Packed => { + if blob_id_col.is_null(row) + || blob_size_col.is_null(row) + || position_col.is_null(row) + { + return Err(Error::invalid_input(format!( + "Prepared packed blob row {row} must set `blob_id`, `blob_size`, and `position`" + ))); + } + validate_blob_id(blob_id_col.value(row))?; + let offset = position_col.value(row); + let size = blob_size_col.value(row); + offset.checked_add(size).ok_or_else(|| { + Error::invalid_input(format!( + "Prepared packed blob row {row} range overflows u64: offset={offset}, size={size}" + )) + })?; + } + BlobKind::Dedicated => { + if blob_id_col.is_null(row) || blob_size_col.is_null(row) { + return Err(Error::invalid_input(format!( + "Prepared dedicated blob row {row} must set `blob_id` and `blob_size`" + ))); + } + validate_blob_id(blob_id_col.value(row))?; + } + BlobKind::External => { + if uri_col.is_null(row) || uri_col.value(row).is_empty() { + return Err(Error::invalid_input(format!( + "Prepared external blob row {row} must set a non-empty `uri`" + ))); + } + let offset = if position_col.is_null(row) { + 0 + } else { + position_col.value(row) + }; + let size = if blob_size_col.is_null(row) { + 0 + } else { + blob_size_col.value(row) + }; + offset.checked_add(size).ok_or_else(|| { + Error::invalid_input(format!( + "Prepared external blob row {row} range overflows u64: offset={offset}, size={size}" + )) + })?; + } + } + } + + Ok(()) +} + +/// Validate a writer-side prepared blob v2 array before it reaches the encoder. +pub(crate) fn validate_prepared_blob_array(field: &Field, array: &ArrayRef) -> Result<()> { + validate_prepared_blob_value_array(field, array) +} + +fn sidecar_path_for_data_file(data_file_path: &Path, blob_id: u32) -> Result { + validate_blob_id(blob_id)?; + let file_name = data_file_path.filename().ok_or_else(|| { + Error::invalid_input("Data file path must include a file name".to_string()) + })?; + let data_file_key = file_name.strip_suffix(".lance").ok_or_else(|| { + Error::invalid_input(format!( + "Data file path '{}' must end with '.lance'", + data_file_path + )) + })?; + let data_dir = data_file_path.parent().unwrap_or_default(); + Ok(blob_path(&data_dir, data_file_key, blob_id)) +} + +/// Byte range inside a packed or external blob object. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct BlobRange { + /// Byte offset relative to the beginning of the blob object. + pub offset: u64, + /// Number of bytes in this range. + pub size: u64, +} + +/// A physical blob descriptor row. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum BlobDescriptor { + /// A null blob row. + Null, + /// Payload bytes embedded into the data file by the blob encoder. + Inline { data: Bytes }, + /// Payload bytes stored as a range in a packed sidecar blob. + Packed { + blob_id: u32, + offset: u64, + size: u64, + }, + /// Payload bytes stored as the full contents of a dedicated sidecar blob. + Dedicated { blob_id: u32, size: u64 }, + /// Payload bytes referenced from an external object or registered base. + External { + base_id: u32, + uri: String, + offset: u64, + size: u64, + }, +} + +/// A physical blob descriptor column ready to be included in a [`RecordBatch`](arrow_array::RecordBatch). +pub struct BlobDescriptorColumn { + field: Field, + array: ArrayRef, +} + +impl BlobDescriptorColumn { + /// Return the Arrow field for the descriptor column. + pub fn field(&self) -> &Field { + &self.field + } + + /// Return the Arrow array for the descriptor column. + pub fn array(&self) -> &ArrayRef { + &self.array + } + + /// Consume this column into `(field, array)` parts. + pub fn into_parts(self) -> (Field, ArrayRef) { + (self.field, self.array) + } +} + +/// Builds physical blob descriptors for one blob v2 column. +/// +/// This builder only produces the writer-side descriptor struct array. It does not allocate blob ids, +/// choose sidecar paths, write blob objects, or commit data files. +pub struct BlobDescriptorArrayBuilder { + field: Field, + values: Vec, +} + +impl BlobDescriptorArrayBuilder { + /// Create a descriptor array builder for one blob column. + pub fn new(column: impl Into) -> Self { + let mut metadata = HashMap::with_capacity(1); + metadata.insert(ARROW_EXT_NAME_KEY.to_string(), BLOB_V2_EXT_NAME.to_string()); + Self::new_with_metadata(column, true, metadata) + } + + pub(crate) fn new_with_metadata( + column: impl Into, + nullable: bool, + metadata: HashMap, + ) -> Self { + Self { + field: prepared_blob_field_with_metadata(&column.into(), nullable, metadata), + values: Vec::new(), + } + } + + /// Append one packed blob descriptor. + pub fn push_packed(&mut self, blob_id: u32, range: BlobRange) -> Result<()> { + self.push(BlobDescriptor::Packed { + blob_id, + offset: range.offset, + size: range.size, + }) + } + + /// Append multiple packed blob descriptors for the same blob object. + pub fn extend_packed( + &mut self, + blob_id: u32, + ranges: impl IntoIterator, + ) -> Result<()> { + for range in ranges { + self.push_packed(blob_id, range)?; + } + Ok(()) + } + + /// Append one dedicated blob descriptor. + pub fn push_dedicated(&mut self, blob_id: u32, size: u64) -> Result<()> { + self.push(BlobDescriptor::Dedicated { blob_id, size }) + } + + /// Append one blob descriptor to this column. + pub fn push(&mut self, value: BlobDescriptor) -> Result<()> { + validate_blob_descriptor(&value)?; + self.values.push(value); + Ok(()) + } + + /// Append multiple blob descriptors to this column. + pub fn extend(&mut self, values: impl IntoIterator) -> Result<()> { + for value in values { + self.push(value)?; + } + Ok(()) + } + + /// Append an inline blob value. + pub fn push_inline(&mut self, data: impl Into) -> Result<()> { + self.push(BlobDescriptor::Inline { data: data.into() }) + } + + /// Append an external blob reference. + pub fn push_external( + &mut self, + uri: impl Into, + range: Option, + ) -> Result<()> { + let range = range.unwrap_or(BlobRange { offset: 0, size: 0 }); + self.push(BlobDescriptor::External { + base_id: 0, + uri: uri.into(), + offset: range.offset, + size: range.size, + }) + } + + /// Append a null blob row. + pub fn push_null(&mut self) -> Result<()> { + self.push(BlobDescriptor::Null) + } + + /// Return the descriptor Arrow field for this blob column. + pub fn field(&self) -> &Field { + &self.field + } + + /// Finish this column and return the writer-side descriptor struct array. + pub fn finish(self) -> Result { + let mut kind_builder = PrimitiveBuilder::::with_capacity(self.values.len()); + let mut data_builder = LargeBinaryBuilder::with_capacity(self.values.len(), 0); + let mut uri_builder = StringBuilder::with_capacity(self.values.len(), 0); + let mut blob_id_builder = PrimitiveBuilder::::with_capacity(self.values.len()); + let mut blob_size_builder = + PrimitiveBuilder::::with_capacity(self.values.len()); + let mut position_builder = PrimitiveBuilder::::with_capacity(self.values.len()); + let mut validity = NullBufferBuilder::new(self.values.len()); + + for value in self.values { + match value { + BlobDescriptor::Null => { + validity.append_null(); + kind_builder.append_null(); + data_builder.append_null(); + uri_builder.append_null(); + blob_id_builder.append_null(); + blob_size_builder.append_null(); + position_builder.append_null(); + } + BlobDescriptor::Inline { data } => { + validity.append_non_null(); + kind_builder.append_value(BlobKind::Inline as u8); + data_builder.append_value(data.as_ref()); + uri_builder.append_null(); + blob_id_builder.append_null(); + blob_size_builder.append_null(); + position_builder.append_null(); + } + BlobDescriptor::Packed { + blob_id, + offset, + size, + } => { + validity.append_non_null(); + kind_builder.append_value(BlobKind::Packed as u8); + data_builder.append_null(); + uri_builder.append_null(); + blob_id_builder.append_value(blob_id); + blob_size_builder.append_value(size); + position_builder.append_value(offset); + } + BlobDescriptor::Dedicated { blob_id, size } => { + validity.append_non_null(); + kind_builder.append_value(BlobKind::Dedicated as u8); + data_builder.append_null(); + uri_builder.append_null(); + blob_id_builder.append_value(blob_id); + blob_size_builder.append_value(size); + position_builder.append_null(); + } + BlobDescriptor::External { + base_id, + uri, + offset, + size, + } => { + validity.append_non_null(); + kind_builder.append_value(BlobKind::External as u8); + data_builder.append_null(); + uri_builder.append_value(uri); + blob_id_builder.append_value(base_id); + blob_size_builder.append_value(size); + position_builder.append_value(offset); + } + } + } + + let array = Arc::new(StructArray::try_new( + prepared_blob_child_fields(), + vec![ + Arc::new(kind_builder.finish()), + Arc::new(data_builder.finish()), + Arc::new(uri_builder.finish()), + Arc::new(blob_id_builder.finish()), + Arc::new(blob_size_builder.finish()), + Arc::new(position_builder.finish()), + ], + validity.finish(), + )?) as ArrayRef; + validate_prepared_blob_array(&self.field, &array)?; + + Ok(BlobDescriptorColumn { + field: self.field, + array, + }) + } +} + +fn validate_blob_descriptor(value: &BlobDescriptor) -> Result<()> { + match value { + BlobDescriptor::Null => Ok(()), + BlobDescriptor::Inline { .. } => Ok(()), + BlobDescriptor::Packed { + blob_id, + offset, + size, + } => { + validate_blob_id(*blob_id)?; + offset.checked_add(*size).ok_or_else(|| { + Error::invalid_input(format!( + "Packed blob range overflows u64: offset={offset}, size={size}" + )) + })?; + Ok(()) + } + BlobDescriptor::Dedicated { blob_id, .. } => validate_blob_id(*blob_id), + BlobDescriptor::External { + uri, offset, size, .. + } => { + if uri.is_empty() { + return Err(Error::invalid_input("External blob URI cannot be empty")); + } + offset.checked_add(*size).ok_or_else(|| { + Error::invalid_input(format!( + "External blob range overflows u64: offset={offset}, size={size}" + )) + })?; + Ok(()) + } + } +} + +/// Writes a Lance-owned packed sidecar blob for one data file and returns descriptors. +pub struct PackedBlobWriter { + object_store: ObjectStore, + path: Path, + blob_id: u32, + writer: Box, + offset: u64, + values: Vec, +} + +impl PackedBlobWriter { + /// Create a packed blob writer for `data_file_path` and `blob_id`. + pub async fn try_new( + object_store: ObjectStore, + data_file_path: Path, + blob_id: u32, + ) -> Result { + let path = sidecar_path_for_data_file(&data_file_path, blob_id)?; + let writer = object_store.create(&path).await?; + Ok(Self { + object_store, + path, + blob_id, + writer, + offset: 0, + values: Vec::new(), + }) + } + + /// Return the blob id for this packed sidecar. + pub fn blob_id(&self) -> u32 { + self.blob_id + } + + /// Return the derived sidecar path for this packed blob. + pub fn path(&self) -> &Path { + &self.path + } + + /// Append one logical blob payload to the packed sidecar. + pub async fn write_blob(&mut self, bytes: impl AsRef<[u8]>) -> Result<()> { + let bytes = bytes.as_ref(); + self.write_blob_bytes(bytes).await?; + Ok(()) + } + + pub(crate) async fn write_blob_bytes(&mut self, bytes: &[u8]) -> Result { + let size = bytes.len() as u64; + let offset = self.offset; + self.writer.write_all(bytes).await?; + self.record_written_blob(offset, size) + } + + pub(crate) async fn write_blob_from_reader( + &mut self, + reader: &dyn Reader, + range: Range, + ) -> Result { + let size = range.len() as u64; + let offset = self.offset; + self.writer.copy_range_from_reader(reader, range).await?; + self.record_written_blob(offset, size) + } + + fn record_written_blob(&mut self, offset: u64, size: u64) -> Result { + self.offset = self.offset.checked_add(size).ok_or_else(|| { + Error::invalid_input(format!( + "Packed blob writer offset overflowed: offset={offset}, size={size}" + )) + })?; + let value = BlobDescriptor::Packed { + blob_id: self.blob_id, + offset, + size, + }; + self.values.push(value.clone()); + Ok(value) + } + + /// Finish the packed sidecar and return descriptors in write order. + pub async fn finish(mut self) -> Result> { + Writer::shutdown(self.writer.as_mut()).await?; + let object_size = self.object_store.size(&self.path).await?; + validate_range(0, self.offset, object_size, "Packed blob")?; + Ok(self.values) + } +} + +/// Writes a Lance-owned dedicated sidecar blob for one data file and returns its descriptor. +pub struct DedicatedBlobWriter { + object_store: ObjectStore, + path: Path, + blob_id: u32, + writer: Box, + size: u64, +} + +impl DedicatedBlobWriter { + /// Create a dedicated blob writer for `data_file_path` and `blob_id`. + pub async fn try_new( + object_store: ObjectStore, + data_file_path: Path, + blob_id: u32, + ) -> Result { + let path = sidecar_path_for_data_file(&data_file_path, blob_id)?; + let writer = object_store.create(&path).await?; + Ok(Self { + object_store, + path, + blob_id, + writer, + size: 0, + }) + } + + /// Return the blob id for this dedicated sidecar. + pub fn blob_id(&self) -> u32 { + self.blob_id + } + + /// Return the derived sidecar path for this dedicated blob. + pub fn path(&self) -> &Path { + &self.path + } + + /// Append bytes to the dedicated sidecar. + pub async fn write(&mut self, bytes: impl AsRef<[u8]>) -> Result<()> { + let bytes = bytes.as_ref(); + let size = bytes.len() as u64; + self.writer.write_all(bytes).await?; + self.record_written_bytes(size) + } + + pub(crate) async fn write_from_reader( + &mut self, + reader: &dyn Reader, + range: Range, + ) -> Result<()> { + let size = range.len() as u64; + self.writer.copy_range_from_reader(reader, range).await?; + self.record_written_bytes(size) + } + + fn record_written_bytes(&mut self, size: u64) -> Result<()> { + self.size = self.size.checked_add(size).ok_or_else(|| { + Error::invalid_input(format!( + "Dedicated blob writer size overflowed: current={}, append={size}", + self.size + )) + })?; + Ok(()) + } + + /// Finish the dedicated sidecar and return its descriptor. + pub async fn finish(mut self) -> Result { + Writer::shutdown(self.writer.as_mut()).await?; + let object_size = self.object_store.size(&self.path).await?; + if object_size != self.size { + return Err(Error::io(format!( + "Dedicated blob sidecar '{}' has size {}, expected {}", + self.path, object_size, self.size + ))); + } + Ok(BlobDescriptor::Dedicated { + blob_id: self.blob_id, + size: self.size, + }) + } +} + /// Builder for blob v2 input struct columns. /// /// The builder enforces that each row contains exactly one of `data` or `uri` (or is null). @@ -211,8 +1013,10 @@ mod tests { use std::num::NonZeroUsize; use super::*; - use arrow_array::Array; use arrow_array::cast::AsArray; + use arrow_array::{Array, StringArray}; + use arrow_schema::Schema as ArrowSchema; + use lance_core::utils::tempfile::TempDir; #[test] fn test_field_metadata() { @@ -287,4 +1091,178 @@ mod tests { let err = b.push_uri("").unwrap_err(); assert!(err.to_string().contains("URI cannot be empty")); } + + #[test] + fn test_prepared_blob_column_finish() { + let mut writer = BlobDescriptorArrayBuilder::new("blob"); + writer.push_inline(Bytes::from_static(b"hello")).unwrap(); + writer + .push_packed(7, BlobRange { offset: 3, size: 5 }) + .unwrap(); + writer.push_dedicated(8, 9).unwrap(); + writer.push_null().unwrap(); + + let column = writer.finish().unwrap(); + assert!(is_prepared_blob_v2_field(column.field())); + let struct_arr = column.array().as_struct(); + let kinds = struct_arr + .column_by_name("kind") + .unwrap() + .as_primitive::(); + let blob_ids = struct_arr + .column_by_name("blob_id") + .unwrap() + .as_primitive::(); + let sizes = struct_arr + .column_by_name("blob_size") + .unwrap() + .as_primitive::(); + let positions = struct_arr + .column_by_name("position") + .unwrap() + .as_primitive::(); + + assert_eq!(kinds.value(0), BlobKind::Inline as u8); + assert_eq!(kinds.value(1), BlobKind::Packed as u8); + assert_eq!(blob_ids.value(1), 7); + assert_eq!(positions.value(1), 3); + assert_eq!(sizes.value(1), 5); + assert_eq!(kinds.value(2), BlobKind::Dedicated as u8); + assert_eq!(blob_ids.value(2), 8); + assert_eq!(sizes.value(2), 9); + assert!(struct_arr.is_null(3)); + } + + #[test] + fn test_prepared_blob_array_rejects_range_overflow() { + let mut metadata = HashMap::new(); + metadata.insert(ARROW_EXT_NAME_KEY.to_string(), BLOB_V2_EXT_NAME.to_string()); + let field = prepared_blob_field_with_metadata("blob", true, metadata); + + for (kind, uri) in [ + (BlobKind::Packed, None), + (BlobKind::External, Some("file:///external.bin")), + ] { + let array = Arc::new( + StructArray::try_new( + prepared_blob_child_fields(), + vec![ + Arc::new(arrow_array::UInt8Array::from(vec![kind as u8])) as ArrayRef, + Arc::new(arrow_array::LargeBinaryArray::from_iter([None::<&[u8]>])), + Arc::new(arrow_array::StringArray::from_iter([uri])), + Arc::new(arrow_array::UInt32Array::from(vec![1])), + Arc::new(arrow_array::UInt64Array::from(vec![4])), + Arc::new(arrow_array::UInt64Array::from(vec![u64::MAX - 1])), + ], + None, + ) + .unwrap(), + ) as ArrayRef; + + let err = validate_prepared_blob_array(&field, &array).unwrap_err(); + assert!(err.to_string().contains("range overflows u64")); + } + } + + #[test] + fn test_normalize_prepared_blob_schema_preserves_non_blob_fields() { + let mut metadata = HashMap::new(); + metadata.insert(ARROW_EXT_NAME_KEY.to_string(), BLOB_V2_EXT_NAME.to_string()); + let prepared_field = prepared_blob_field_with_metadata("blob", true, metadata); + let dict_field = Field::new( + "dict", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true, + ); + let mut schema = + LanceSchema::try_from(&ArrowSchema::new(vec![dict_field, prepared_field])).unwrap(); + schema.fields[0].id = 42; + schema.fields[1].id = 7; + + let dictionary_values = Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef; + schema.fields[0].set_dictionary_values(&dictionary_values); + + let normalized = normalize_prepared_blob_schema(&schema).unwrap(); + + assert_eq!(normalized.fields[0].id, 42); + assert_eq!( + normalized.fields[0] + .dictionary + .as_ref() + .and_then(|dict| dict.values.as_ref()) + .map(|values| values.len()), + Some(2) + ); + assert_eq!(normalized.fields[1].id, 7); + assert_eq!(normalized.fields[1].children.len(), 2); + assert_eq!(normalized.fields[1].children[0].name, "data"); + assert_eq!(normalized.fields[1].children[1].name, "uri"); + assert!(normalized.fields[1].children[0].id >= 0); + assert!(normalized.fields[1].children[1].id >= 0); + } + + #[tokio::test] + async fn test_sidecar_writers_return_prepared_values() { + let temp_dir = TempDir::default(); + let data_dir = Path::from_absolute_path(temp_dir.std_path().join("data")).unwrap(); + let data_file_key = "data-file".to_string(); + let data_file_path = data_dir.clone().join(format!("{data_file_key}.lance")); + let object_store = ObjectStore::local(); + + let packed_id = 7; + let packed_path = blob_path(&data_dir, &data_file_key, packed_id); + let mut packed = + PackedBlobWriter::try_new(object_store.clone(), data_file_path.clone(), packed_id) + .await + .unwrap(); + assert_eq!(packed.path(), &packed_path); + packed.write_blob(b"abc").await.unwrap(); + packed.write_blob(b"de").await.unwrap(); + let packed_values = packed.finish().await.unwrap(); + assert_eq!( + packed_values, + vec![ + BlobDescriptor::Packed { + blob_id: packed_id, + offset: 0, + size: 3, + }, + BlobDescriptor::Packed { + blob_id: packed_id, + offset: 3, + size: 2, + }, + ] + ); + + let dedicated_id = 8; + let dedicated_path = blob_path(&data_dir, &data_file_key, dedicated_id); + let mut dedicated = + DedicatedBlobWriter::try_new(object_store.clone(), data_file_path, dedicated_id) + .await + .unwrap(); + assert_eq!(dedicated.path(), &dedicated_path); + dedicated.write(b"abcdef").await.unwrap(); + assert_eq!( + dedicated.finish().await.unwrap(), + BlobDescriptor::Dedicated { + blob_id: dedicated_id, + size: 6, + } + ); + + let mut builder = BlobDescriptorArrayBuilder::new("blob"); + builder + .extend_packed( + 42, + vec![ + BlobRange { offset: 1, size: 2 }, + BlobRange { offset: 4, size: 1 }, + ], + ) + .unwrap(); + builder.push_dedicated(43, 3).unwrap(); + let column = builder.finish().unwrap(); + assert_eq!(column.array().len(), 3); + } } diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 3e0d77704da..72b13e903b5 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -42,7 +42,8 @@ use lance_io::utils::{ }; use lance_namespace::LanceNamespace; use lance_table::format::{ - DataFile, DataStorageFormat, DeletionFile, Fragment, IndexMetadata, Manifest, RowIdMeta, pb, + DataFile, DataStorageFormat, DeletionFile, Fragment, IndexMetadata, MAGIC, Manifest, RowIdMeta, + pb, }; use lance_table::io::commit::{ CommitConfig, CommitError, CommitHandler, CommitLock, ManifestLocation, ManifestNamingScheme, @@ -100,6 +101,7 @@ use self::cleanup::RemovalStats; use self::fragment::FileFragment; use self::refs::Refs; use self::scanner::{DatasetRecordBatchStream, Scanner}; +use self::statistics::DatasetStatistics; use self::transaction::{Operation, Transaction, TransactionBuilder, UpdateMapEntry}; use self::write::{cleanup_data_fragments, write_fragments_internal}; use crate::dataset::branch_location::BranchLocation; @@ -453,6 +455,12 @@ impl Dataset { self.refs.tags() } + /// A handle for cheap, index-derived statistics about this dataset (e.g. a + /// column's global value range) that never scan data. + pub fn statistics(&self) -> DatasetStatistics<'_> { + DatasetStatistics::new(self) + } + pub fn branches(&self) -> Branches<'_> { self.refs.branches() } @@ -591,7 +599,7 @@ impl Dataset { return Ok(self.clone()); } - let manifest = Self::load_manifest( + let manifest = Self::get_manifest( self.object_store.as_ref(), &manifest_location, &new_location.uri, @@ -617,7 +625,7 @@ impl Dataset { self.object_store.clone(), new_location.path, new_location.uri, - Arc::new(manifest), + manifest, manifest_location, self.session.clone(), self.commit_handler.clone(), @@ -654,6 +662,24 @@ impl Dataset { } _ => Error::io_source(err.into()), })?; + + // A stale cached size yields a bogus footer offset. Detect it (the block + // lacks the trailing magic) and retry with the true size, like + // read_manifest. + if manifest_location.size.is_some() && !last_block.ends_with(MAGIC) { + let manifest_location = ManifestLocation { + size: None, + ..manifest_location.clone() + }; + return Box::pin(Self::load_manifest( + object_store, + &manifest_location, + uri, + session, + )) + .await; + } + let offset = read_metadata_offset(&last_block)?; // If manifest is in the last block, we can decode directly from memory. @@ -731,6 +757,35 @@ impl Dataset { Ok(manifest) } + /// Fetch the manifest for `manifest_location` from the session metadata + /// cache, loading and caching it on a miss. + pub(crate) async fn get_manifest( + object_store: &ObjectStore, + manifest_location: &ManifestLocation, + uri: &str, + session: &Session, + ) -> Result> { + if manifest_location.size.is_none() { + return Ok(Arc::new( + Self::load_manifest(object_store, manifest_location, uri, session).await?, + )); + } + let metadata_cache = session.metadata_cache.for_dataset(uri); + let manifest_key = ManifestKey { + version: manifest_location.version, + e_tag: manifest_location.e_tag.as_deref(), + }; + if let Some(cached) = metadata_cache.get_with_key(&manifest_key).await { + return Ok(cached); + } + let loaded = + Arc::new(Self::load_manifest(object_store, manifest_location, uri, session).await?); + metadata_cache + .insert_with_key(&manifest_key, loaded.clone()) + .await; + Ok(loaded) + } + #[allow(clippy::too_many_arguments)] fn checkout_manifest( object_store: Arc, @@ -756,6 +811,11 @@ impl Dataset { let metadata_cache = Arc::new(session.metadata_cache.for_dataset(&uri)); let index_cache = Arc::new(session.index_cache.for_dataset(&uri)); let fragment_bitmap = Arc::new(manifest.fragments.iter().map(|f| f.id as u32).collect()); + write::log_unregistered_base_scoped_options( + store_params.as_ref(), + &manifest.base_paths, + log::Level::Debug, + ); Ok(Self { object_store, base: base_path, @@ -1836,21 +1896,26 @@ impl Dataset { store_params } - fn store_params_for_base( + pub(crate) fn store_params_for_base( &self, base_path: Option<&lance_table::format::BasePath>, ) -> ObjectStoreParams { // Base-specific bindings are exact ObjectStoreParams keyed by - // `BasePath.path`. If a base has no explicit binding then reads fall back - // to the dataset-level default store params. - base_path - .and_then(|base_path| { - self.base_store_params - .as_ref() - .and_then(|params| params.get(&base_path.path)) - }) - .cloned() - .unwrap_or_else(|| self.store_params.as_deref().cloned().unwrap_or_default()) + // `BasePath.path` and are used as-is. Otherwise the dataset-level + // default params are resolved for the base scope: `base_.` + // storage options overlay the shared defaults for that base. + if let Some(params) = base_path.and_then(|base_path| { + self.base_store_params + .as_ref() + .and_then(|params| params.get(&base_path.path)) + }) { + return params.clone(); + } + let default_params = self.store_params.as_deref().cloned().unwrap_or_default(); + match default_params.scoped_to_base(base_path.map(|base_path| base_path.id)) { + Cow::Owned(scoped_params) => scoped_params, + Cow::Borrowed(_) => default_params, + } } /// Returns the initial storage options used when opening this dataset, if any. @@ -1974,58 +2039,215 @@ impl Dataset { file_metadata.minor_version as u32, )?; - // Get top-level column names from file schema in file order - let column_names: Vec<&str> = file_metadata - .file_schema - .fields - .iter() - .map(|f| f.name.as_str()) - .collect(); + let is_structural = file_version >= LanceFileVersion::V2_1; + let physical_columns = file_metadata.column_metadatas.len(); + let has_footer_orphans = file_metadata.file_schema.fields.len() > physical_columns; + let dataset_schema = self.schema(); + let mut represented_columns = 0usize; + let mut column_names = Vec::new(); + let mut consumed_top_level_fields = 0usize; + + fn physical_column_count( + field: &lance_core::datatypes::Field, + is_structural: bool, + ) -> usize { + if !is_structural { + return 1 + field + .children + .iter() + .map(|child| physical_column_count(child, is_structural)) + .sum::(); + } - // Project dataset schema by file column names to get dataset field IDs - let projected_ds_schema = self.schema().project(&column_names)?; + if field.children.is_empty() || field.is_blob() || field.is_packed_struct() { + 1 + } else { + field + .children + .iter() + .map(|child| physical_column_count(child, is_structural)) + .sum() + } + } - // Walk both schemas in parallel to build fields and column_indices - let is_structural = file_version >= LanceFileVersion::V2_1; - let ds_fields: Vec<_> = projected_ds_schema.fields_pre_order().collect(); - let file_fields: Vec<_> = file_metadata.file_schema.fields_pre_order().collect(); + fn field_contains_blob(field: &lance_core::datatypes::Field) -> bool { + field.is_blob() || field.children.iter().any(field_contains_blob) + } - if ds_fields.len() != file_fields.len() { - return Err(Error::invalid_input(format!( - "Schema mismatch: dataset projection has {} fields but file has {} fields", - ds_fields.len(), - file_fields.len() - ))); + fn field_names_match( + fields: &[lance_core::datatypes::Field], + start: usize, + names: &[&str], + ) -> bool { + fields + .get(start..start + names.len()) + .is_some_and(|candidate| { + candidate + .iter() + .zip(names) + .all(|(field, name)| field.name == *name) + }) } - let mut fields = Vec::new(); - let mut column_indices = Vec::new(); - let mut curr_column_idx: i32 = 0; - let mut packed_struct_fields_num: usize = 0; + fn blob_descriptor_orphan_len( + fields: &[lance_core::datatypes::Field], + start: usize, + ) -> usize { + const BLOB_V2_DESCRIPTOR_FIELDS: &[&str] = + &["kind", "position", "size", "blob_id", "blob_uri"]; + const BLOB_V1_DESCRIPTOR_FIELDS: &[&str] = &["position", "size"]; + + if field_names_match(fields, start, BLOB_V2_DESCRIPTOR_FIELDS) { + BLOB_V2_DESCRIPTOR_FIELDS.len() + } else if field_names_match(fields, start, BLOB_V1_DESCRIPTOR_FIELDS) { + BLOB_V1_DESCRIPTOR_FIELDS.len() + } else { + 0 + } + } - for (ds_field, file_field) in ds_fields.iter().zip(file_fields.iter()) { - if ds_field.name != file_field.name { + fn collect_columns( + field: &lance_core::datatypes::Field, + is_structural: bool, + fields: &mut Vec, + column_indices: &mut Vec, + curr_column_idx: &mut i32, + ) { + let contributes = !is_structural + || field.children.is_empty() + || field.is_blob() + || field.is_packed_struct(); + let recurse = !is_structural || (!field.is_blob() && !field.is_packed_struct()); + + if contributes { + fields.push(field.id); + column_indices.push(*curr_column_idx); + *curr_column_idx += 1; + } + + if recurse { + for child in &field.children { + collect_columns( + child, + is_structural, + fields, + column_indices, + curr_column_idx, + ); + } + } + } + + fn validate_file_field_matches_dataset( + dataset_field: &lance_core::datatypes::Field, + file_field: &lance_core::datatypes::Field, + path: &str, + ) -> Result<()> { + if dataset_field.name != file_field.name { return Err(Error::invalid_input(format!( "Schema mismatch: expected field '{}' but file has '{}'", - ds_field.name, file_field.name + path, file_field.name ))); } - if packed_struct_fields_num > 0 { - packed_struct_fields_num -= 1; - continue; + if dataset_field.is_blob() && file_field.is_blob() { + return Ok(()); } - if file_field.is_packed_struct() { - fields.push(ds_field.id); - column_indices.push(curr_column_idx); - curr_column_idx += 1; - packed_struct_fields_num = file_field.children.len(); - } else if file_field.children.is_empty() || !is_structural { - fields.push(ds_field.id); - column_indices.push(curr_column_idx); - curr_column_idx += 1; + if dataset_field.children.len() != file_field.children.len() { + return Err(Error::invalid_input(format!( + "Schema mismatch: field '{}' has {} children in dataset schema but {} children in file schema", + path, + dataset_field.children.len(), + file_field.children.len() + ))); + } + + for (dataset_child, file_child) in + dataset_field.children.iter().zip(&file_field.children) + { + let child_path = format!("{}.{}", path, dataset_child.name); + validate_file_field_matches_dataset(dataset_child, file_child, &child_path)?; } + + Ok(()) + } + + let file_schema_fields = &file_metadata.file_schema.fields; + let mut idx = 0usize; + while represented_columns < physical_columns { + let Some(field) = file_schema_fields.get(idx) else { + return Err(Error::invalid_input(format!( + "Schema mismatch: file schema ended after representing {} physical columns but file has {} columns", + represented_columns, physical_columns + ))); + }; + + let Some(dataset_field) = dataset_schema.field(&field.name) else { + return Err(Error::invalid_input(format!( + "Schema mismatch: file has extra field '{}'", + field.name + ))); + }; + validate_file_field_matches_dataset(dataset_field, field, &field.name)?; + + represented_columns += physical_column_count(field, is_structural); + column_names.push(field.name.as_str()); + consumed_top_level_fields = idx + 1; + idx += 1; + + if has_footer_orphans && field_contains_blob(field) { + loop { + let skipped = blob_descriptor_orphan_len(file_schema_fields, idx); + if skipped == 0 { + break; + } + consumed_top_level_fields = idx + skipped; + idx += skipped; + } + } + } + + if represented_columns != physical_columns { + return Err(Error::invalid_input(format!( + "Schema mismatch: file schema represents {} physical columns but file has {} columns", + represented_columns, physical_columns + ))); + } + + if let Some(field) = file_schema_fields.get(consumed_top_level_fields) { + return Err(Error::invalid_input(format!( + "Schema mismatch: file has extra field '{}'", + field.name + ))); + } + + let projected_ds_schema = self.schema().project(&column_names)?; + + let mut fields = Vec::new(); + let mut column_indices = Vec::new(); + let mut curr_column_idx: i32 = 0; + for field in &projected_ds_schema.fields { + collect_columns( + field, + is_structural, + &mut fields, + &mut column_indices, + &mut curr_column_idx, + ); + } + + if curr_column_idx as usize != physical_columns { + return Err(Error::invalid_input(format!( + "Schema mismatch: dataset projection maps to {} physical columns but file has {} columns", + curr_column_idx, physical_columns + ))); + } + + if fields.is_empty() && physical_columns > 0 { + return Err(Error::invalid_input( + "Schema mismatch: file has columns but none matched the dataset schema", + )); } let file_size_nz = NonZero::new(file_size); @@ -2043,7 +2265,7 @@ impl Dataset { /// Resolve the data directory for a given base_id. /// /// If `base_id` is `None`, returns the default data directory. - fn data_file_dir_for_base(&self, base_id: Option) -> Result { + pub(crate) fn data_file_dir_for_base(&self, base_id: Option) -> Result { match base_id { Some(base_id) => { let base_path = self.manifest.base_paths.get(&base_id).ok_or_else(|| { @@ -2557,7 +2779,7 @@ impl Dataset { self.manifest_location.path.clone(), format!( "Duplicate index id {} found in dataset {:?}", - &index.uuid, self.base + index.uuid, self.base ), )); } @@ -2904,30 +3126,13 @@ pub(crate) fn load_new_transactions(dataset: &Dataset) -> NewTransactionResult<' .map_ok(move |location| { let latest_tx = latest_tx.take(); async move { - let manifest_key = ManifestKey { - version: location.version, - e_tag: location.e_tag.as_deref(), - }; - let manifest = if let Some(cached) = - dataset.metadata_cache.get_with_key(&manifest_key).await - { - cached - } else { - let loaded = Arc::new( - Dataset::load_manifest( - dataset.object_store.as_ref(), - &location, - &dataset.uri, - dataset.session.as_ref(), - ) - .await?, - ); - dataset - .metadata_cache - .insert_with_key(&manifest_key, loaded.clone()) - .await; - loaded - }; + let manifest = Dataset::get_manifest( + dataset.object_store.as_ref(), + &location, + &dataset.uri, + dataset.session.as_ref(), + ) + .await?; if let Some(latest_tx) = latest_tx { // We ignore the error, since we don't care if the receiver is dropped. diff --git a/rust/lance/src/dataset/blob.rs b/rust/lance/src/dataset/blob.rs index 8cdde543e4e..c0559753330 100644 --- a/rust/lance/src/dataset/blob.rs +++ b/rust/lance/src/dataset/blob.rs @@ -12,33 +12,38 @@ use std::{ use arrow::array::AsArray; use arrow::datatypes::{UInt8Type, UInt32Type, UInt64Type}; -use arrow_array::RecordBatch; -use arrow_array::builder::{LargeBinaryBuilder, PrimitiveBuilder, StringBuilder}; -use arrow_array::{Array, ArrayRef}; -use arrow_schema::{DataType as ArrowDataType, Field as ArrowField}; +use arrow_array::{ + Array, ArrayRef, GenericListArray, OffsetSizeTrait, RecordBatch, builder::LargeBinaryBuilder, +}; +use arrow_buffer::{OffsetBuffer, ScalarBuffer}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; use bytes::Bytes; use futures::future::BoxFuture; use futures::stream::BoxStream; use futures::{FutureExt, StreamExt, TryStreamExt, stream}; use lance_arrow::{ - BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, BLOB_INLINE_SIZE_THRESHOLD_META_KEY, FieldExt, - r#struct::StructArrayExt, + ARROW_EXT_NAME_KEY, BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, + BLOB_INLINE_SIZE_THRESHOLD_META_KEY, BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY, FieldExt, + list::ListArrayExt, r#struct::StructArrayExt, }; use lance_io::object_store::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry}; use lance_io::scheduler::{FileScheduler, ScanScheduler, SchedulerConfig}; use object_store::path::Path; -use tokio::io::AsyncWriteExt; use tokio::sync::{Mutex, OnceCell, oneshot}; use url::Url; use super::take::TakeBuilder; use super::write::ExternalBlobMode; use super::{Dataset, ProjectionRequest}; +use crate::blob::{ + BlobDescriptor, BlobDescriptorArrayBuilder, BlobIdAllocator, BlobRange, PackedBlobWriter, + is_logical_blob_v2_field, is_prepared_blob_v2_field, validate_prepared_blob_array, +}; use arrow_array::StructArray; -use lance_core::datatypes::{BlobKind, BlobVersion, parse_field_path}; +use lance_core::datatypes::{BlobKind, BlobVersion, Field as LanceField, Schema, parse_field_path}; use lance_core::utils::blob::blob_path; -use lance_core::{Error, Result, utils::address::RowAddress}; -use lance_io::traits::{Reader, WriteExt, Writer}; +use lance_core::{Error, ROW_ADDR, Result, utils::address::RowAddress}; +use lance_io::traits::Reader; use lance_io::utils::CachedFileSize; const INLINE_MAX: usize = 64 * 1024; // 64KB inline cutoff @@ -71,6 +76,19 @@ pub(super) fn blob_dedicated_threshold_from_metadata( ) } +pub(super) fn blob_pack_file_threshold_from_metadata( + metadata: &HashMap, + field_name: &str, +) -> Result { + blob_threshold_from_metadata( + metadata, + field_name, + BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY, + PACK_FILE_MAX_SIZE, + false, + ) +} + fn blob_threshold_from_metadata( metadata: &HashMap, field_name: &str, @@ -168,84 +186,86 @@ impl ExternalBaseResolver { } } -// Maintains rolling `.blob` sidecar files for packed blobs. -// Layout: data/{data_file_key}/{obfuscated_blob_id:032b}.blob where each file is an -// unframed concatenation of blob payloads; descriptors store (blob_id, -// position, size) to locate each slice. A dedicated struct keeps path state -// and rolling size separate from the per-batch preprocessor logic, so we can -// reuse the same writer across rows and close/roll files cleanly on finish. -struct PackWriter { - object_store: ObjectStore, - data_dir: Path, - data_file_key: String, - max_pack_size: usize, - current_blob_id: Option, - writer: Option>, +struct RollingPackedBlobWriter { + current: Option, current_size: usize, + current_max_pack_size: Option, } -impl PackWriter { - fn new(object_store: ObjectStore, data_dir: Path, data_file_key: String) -> Self { +impl RollingPackedBlobWriter { + fn new() -> Self { Self { - object_store, - data_dir, - data_file_key, - max_pack_size: PACK_FILE_MAX_SIZE, - current_blob_id: None, - writer: None, + current: None, current_size: 0, + current_max_pack_size: None, } } - async fn start_new_pack(&mut self, blob_id: u32) -> Result<()> { - let path = blob_path(&self.data_dir, &self.data_file_key, blob_id); - let writer = self.object_store.create(&path).await?; - self.writer = Some(writer); - self.current_blob_id = Some(blob_id); + async fn start_new_pack( + &mut self, + object_store: ObjectStore, + data_dir: Path, + data_file_key: String, + blob_id_allocator: BlobIdAllocator, + max_pack_size: usize, + ) -> Result<()> { + self.finish().await?; + let blob_id = blob_id_allocator.next()?; + let data_file_path = data_dir.join(format!("{data_file_key}.lance")); + self.current = + Some(PackedBlobWriter::try_new(object_store, data_file_path, blob_id).await?); self.current_size = 0; + self.current_max_pack_size = Some(max_pack_size); Ok(()) } - /// Append `data` to the current `.blob` file, rolling to a new file when - /// `max_pack_size` would be exceeded. - /// - /// alloc_blob_id: called only when a new pack file is opened; returns the - /// blob_id used as the file name. - /// - /// Returns `(blob_id, position)` where - /// position is the start offset of this payload in that pack file. - async fn write_with_allocator( + async fn write( &mut self, - alloc_blob_id: &mut F, + object_store: ObjectStore, + data_dir: Path, + data_file_key: String, + blob_id_allocator: BlobIdAllocator, + max_pack_size: usize, source: BlobWriteSource<'_>, - ) -> Result<(u32, u64)> - where - F: FnMut() -> u32, - { + ) -> Result { let len = source.size(); - if self - .current_blob_id - .map(|_| self.current_size + len > self.max_pack_size) - .unwrap_or(true) - { - let blob_id = alloc_blob_id(); - self.finish().await?; - self.start_new_pack(blob_id).await?; + let needs_new_pack = match (self.current.as_ref(), self.current_max_pack_size) { + (Some(_), Some(current_max_pack_size)) => { + current_max_pack_size != max_pack_size + || self.current_size + len > current_max_pack_size + } + _ => true, + }; + if needs_new_pack { + self.start_new_pack( + object_store, + data_dir, + data_file_key, + blob_id_allocator, + max_pack_size, + ) + .await?; } - let writer = self.writer.as_mut().expect("pack writer is initialized"); - let position = self.current_size as u64; - source.write_to(writer.as_mut()).await?; + let writer = self.current.as_mut().expect("pack writer is initialized"); + let value = match source { + BlobWriteSource::Bytes(data) => writer.write_blob_bytes(data).await?, + BlobWriteSource::External(source) => { + writer + .write_blob_from_reader(source.reader.as_ref(), source.reader_range()?) + .await? + } + }; self.current_size += len; - Ok((self.current_blob_id.expect("pack blob id"), position)) + Ok(value) } async fn finish(&mut self) -> Result<()> { - if let Some(mut writer) = self.writer.take() { - Writer::shutdown(writer.as_mut()).await?; + if let Some(writer) = self.current.take() { + writer.finish().await?; } - self.current_blob_id = None; self.current_size = 0; + self.current_max_pack_size = None; Ok(()) } } @@ -259,8 +279,12 @@ pub struct BlobPreprocessor { object_store: ObjectStore, data_dir: Path, data_file_key: String, - local_counter: u32, - pack_writer: PackWriter, + blob_id_allocator: BlobIdAllocator, + pack_writer: RollingPackedBlobWriter, + /// Write-param override for the pack-file roll size. When set, it takes + /// precedence over each field's `blob-pack-file-size-threshold` metadata for + /// this write job only; it is not persisted into the dataset schema. + pack_file_size_override: Option, field_processors: Vec, external_base_resolver: Option>, allow_external_blob_outside_bases: bool, @@ -296,17 +320,35 @@ enum BlobPreprocessFieldKind { BlobV2 { inline_threshold: usize, dedicated_threshold: usize, + pack_file_threshold: usize, writer_metadata: HashMap, }, Struct { children: Vec, }, + List { + child: Box, + }, Passthrough, } impl BlobPreprocessField { fn new(field: &ArrowField) -> Result { if field.is_blob_v2() { + if is_prepared_blob_v2_field(field) { + return Ok(Self { + kind: BlobPreprocessFieldKind::Passthrough, + }); + } + if !is_logical_blob_v2_field(field) { + return Err(Error::invalid_input(format!( + "Blob v2 field '{}' must use either logical struct with optional position/size UInt64 fields or prepared \ + struct", + field.name() + ))); + } return Ok(Self { kind: BlobPreprocessFieldKind::BlobV2 { inline_threshold: blob_inline_threshold_from_metadata( @@ -317,6 +359,10 @@ impl BlobPreprocessField { field.metadata(), field.name(), )?, + pack_file_threshold: blob_pack_file_threshold_from_metadata( + field.metadata(), + field.name(), + )?, writer_metadata: field.metadata().clone(), }, }); @@ -334,6 +380,17 @@ impl BlobPreprocessField { } } + if let ArrowDataType::List(child) | ArrowDataType::LargeList(child) = field.data_type() { + let child = Self::new(child.as_ref())?; + if child.requires_preprocessing() { + return Ok(Self { + kind: BlobPreprocessFieldKind::List { + child: Box::new(child), + }, + }); + } + } + Ok(Self { kind: BlobPreprocessFieldKind::Passthrough, }) @@ -378,15 +435,6 @@ impl ExternalBlobSource { let range = self.reader_range()?; self.reader.get_range(range).await.map_err(Into::into) } - - /// Stream the slice into a writer for packed or dedicated blob storage. - async fn copy_to_writer(&self, writer: &mut dyn Writer) -> Result<()> { - let range = self.reader_range()?; - writer - .copy_range_from_reader(self.reader.as_ref(), range) - .await?; - Ok(()) - } } impl BlobWriteSource<'_> { @@ -394,22 +442,12 @@ impl BlobWriteSource<'_> { fn size(&self) -> usize { match self { Self::Bytes(data) => data.len(), - Self::External(source) => usize::try_from(source.size()) + Self::External(source) => source + .size() + .try_into() .expect("packed and inline external blobs must fit into usize"), } } - - /// Write the payload into Lance-managed storage without forcing callers to branch on source - /// type. - async fn write_to(&self, writer: &mut dyn Writer) -> Result<()> { - match self { - Self::Bytes(data) => { - writer.write_all(data).await?; - Ok(()) - } - Self::External(source) => source.copy_to_writer(writer).await, - } - } } impl BlobPreprocessor { @@ -424,16 +462,9 @@ impl BlobPreprocessor { external_blob_mode: ExternalBlobMode, source_store_registry: Arc, source_store_params: ObjectStoreParams, - pack_file_size_threshold: Option, + pack_file_size_override: Option, ) -> Result { - let mut pack_writer = PackWriter::new( - object_store.clone(), - data_dir.clone(), - data_file_key.clone(), - ); - if let Some(max_bytes) = pack_file_size_threshold { - pack_writer.max_pack_size = max_bytes; - } + let pack_writer = RollingPackedBlobWriter::new(); let arrow_schema = arrow_schema::Schema::from(schema); let field_processors = arrow_schema .fields() @@ -444,9 +475,9 @@ impl BlobPreprocessor { object_store, data_dir, data_file_key, - // Start at 1 to avoid a potential all-zero blob_id value. - local_counter: 1, + blob_id_allocator: BlobIdAllocator::new(1), pack_writer, + pack_file_size_override, field_processors, external_base_resolver, allow_external_blob_outside_bases, @@ -456,29 +487,52 @@ impl BlobPreprocessor { }) } - fn next_blob_id(&mut self) -> u32 { - let id = self.local_counter; - self.local_counter += 1; - id + fn blob_writer_with_metadata( + &self, + field: &ArrowField, + metadata: HashMap, + ) -> BlobDescriptorArrayBuilder { + BlobDescriptorArrayBuilder::new_with_metadata(field.name(), field.is_nullable(), metadata) } - async fn write_dedicated(&mut self, blob_id: u32, source: BlobWriteSource<'_>) -> Result { - let path = blob_path(&self.data_dir, &self.data_file_key, blob_id); - let mut writer = self.object_store.create(&path).await?; - source.write_to(writer.as_mut()).await?; - Writer::shutdown(&mut writer).await?; - Ok(path) + async fn write_dedicated( + object_store: ObjectStore, + data_dir: Path, + data_file_key: String, + blob_id_allocator: BlobIdAllocator, + source: BlobWriteSource<'_>, + ) -> Result { + let blob_id = blob_id_allocator.next()?; + let data_file_path = data_dir.join(format!("{data_file_key}.lance")); + let mut writer = + crate::blob::DedicatedBlobWriter::try_new(object_store, data_file_path, blob_id) + .await?; + match source { + BlobWriteSource::Bytes(data) => writer.write(data).await?, + BlobWriteSource::External(source) => { + writer + .write_from_reader(source.reader.as_ref(), source.reader_range()?) + .await?; + } + } + writer.finish().await } - async fn write_packed(&mut self, source: BlobWriteSource<'_>) -> Result<(u32, u64)> { - let (counter, pack_writer) = (&mut self.local_counter, &mut self.pack_writer); - pack_writer - .write_with_allocator( - &mut || { - let id = *counter; - *counter += 1; - id - }, + async fn write_packed( + &mut self, + field_pack_file_threshold: usize, + source: BlobWriteSource<'_>, + ) -> Result { + let max_pack_size = self + .pack_file_size_override + .unwrap_or(field_pack_file_threshold); + self.pack_writer + .write( + self.object_store.clone(), + self.data_dir.clone(), + self.data_file_key.clone(), + self.blob_id_allocator.clone(), + max_pack_size, source, ) .await @@ -595,11 +649,17 @@ impl BlobPreprocessor { field: &'a Arc, ) -> BoxFuture<'a, Result<(ArrayRef, Arc)>> { async move { + if is_prepared_blob_v2_field(field.as_ref()) { + validate_prepared_blob_array(field.as_ref(), &array)?; + return Ok((array, field.clone())); + } + match &processor.kind { BlobPreprocessFieldKind::Passthrough => Ok((array, field.clone())), BlobPreprocessFieldKind::BlobV2 { inline_threshold, dedicated_threshold, + pack_file_threshold, writer_metadata, } => { self.preprocess_blob_array( @@ -607,6 +667,7 @@ impl BlobPreprocessor { field.as_ref(), *inline_threshold, *dedicated_threshold, + *pack_file_threshold, writer_metadata, ) .await @@ -615,6 +676,20 @@ impl BlobPreprocessor { self.preprocess_struct_array(array, field.as_ref(), children) .await } + BlobPreprocessFieldKind::List { child } => match field.data_type() { + ArrowDataType::List(_) => { + self.preprocess_list_array::(array, field.as_ref(), child) + .await + } + ArrowDataType::LargeList(_) => { + self.preprocess_list_array::(array, field.as_ref(), child) + .await + } + _ => Err(Error::internal(format!( + "Blob list preprocessor received non-list field '{}'", + field.name() + ))), + }, } } .boxed() @@ -647,10 +722,8 @@ impl BlobPreprocessor { let mut new_columns = Vec::with_capacity(children.len()); let mut new_fields = Vec::with_capacity(children.len()); - for ((child_processor, child_array), child_field) in children - .iter() - .zip(child_columns.into_iter()) - .zip(child_fields.iter()) + for ((child_processor, child_array), child_field) in + children.iter().zip(child_columns).zip(child_fields.iter()) { let (new_column, new_field) = self .preprocess_field(child_processor, child_array, child_field) @@ -672,12 +745,86 @@ impl BlobPreprocessor { Ok((Arc::new(struct_array), field)) } + async fn preprocess_list_array( + &mut self, + array: ArrayRef, + field: &ArrowField, + child: &BlobPreprocessField, + ) -> Result<(ArrayRef, Arc)> { + let list_arr = array.as_list::(); + let list_arr = if list_arr.null_count() > 0 { + list_arr.filter_garbage_nulls() + } else { + list_arr.clone() + }; + + let first_offset = *list_arr + .offsets() + .first() + .ok_or_else(|| Error::invalid_input("List offsets cannot be empty"))?; + let last_offset = *list_arr + .offsets() + .last() + .ok_or_else(|| Error::invalid_input("List offsets cannot be empty"))?; + let values_len = list_arr.values().len(); + let needs_trim = first_offset != O::zero() + || last_offset.to_usize().ok_or_else(|| { + Error::invalid_input(format!( + "List field '{}' offset does not fit into usize", + field.name() + )) + })? != values_len; + + let (offsets, values) = if needs_trim { + let values = list_arr.trimmed_values(); + let offsets = list_arr + .offsets() + .iter() + .map(|offset| *offset - first_offset) + .collect::>(); + (OffsetBuffer::new(ScalarBuffer::from(offsets)), values) + } else { + (list_arr.offsets().clone(), list_arr.values().clone()) + }; + + let child_field = match field.data_type() { + ArrowDataType::List(child_field) | ArrowDataType::LargeList(child_field) => { + child_field.clone() + } + other => { + return Err(Error::invalid_input(format!( + "Blob list preprocessor expected list field '{}', got {other}", + field.name() + ))); + } + }; + let (new_values, new_child_field) = + self.preprocess_field(child, values, &child_field).await?; + + let list_array = GenericListArray::::try_new( + new_child_field, + offsets, + new_values, + list_arr.nulls().cloned(), + )?; + let field = Arc::new( + ArrowField::new( + field.name(), + list_array.data_type().clone(), + field.is_nullable(), + ) + .with_metadata(field.metadata().clone()), + ); + Ok((Arc::new(list_array), field)) + } + async fn preprocess_blob_array( &mut self, array: ArrayRef, field: &ArrowField, inline_threshold: usize, dedicated_threshold: usize, + pack_file_threshold: usize, writer_metadata: &HashMap, ) -> Result<(ArrayRef, Arc)> { let struct_arr = array @@ -700,26 +847,11 @@ impl BlobPreprocessor { .column_by_name("size") .map(|col| col.as_primitive::()); - let mut data_builder = LargeBinaryBuilder::with_capacity(struct_arr.len(), 0); - let mut uri_builder = StringBuilder::with_capacity(struct_arr.len(), 0); - let mut blob_id_builder = - PrimitiveBuilder::::with_capacity(struct_arr.len()); - let mut blob_size_builder = - PrimitiveBuilder::::with_capacity(struct_arr.len()); - let mut kind_builder = PrimitiveBuilder::::with_capacity(struct_arr.len()); - let mut position_builder = - PrimitiveBuilder::::with_capacity(struct_arr.len()); - - let struct_nulls = struct_arr.nulls(); + let mut blob_writer = self.blob_writer_with_metadata(field, writer_metadata.clone()); for i in 0..struct_arr.len() { if struct_arr.is_null(i) { - data_builder.append_null(); - uri_builder.append_null(); - blob_id_builder.append_null(); - blob_size_builder.append_null(); - kind_builder.append_null(); - position_builder.append_null(); + blob_writer.push_null()?; continue; } @@ -736,30 +868,26 @@ impl BlobPreprocessor { let data_len = if has_data { data_col.value(i).len() } else { 0 }; if has_data && data_len > dedicated_threshold { - let blob_id = self.next_blob_id(); - self.write_dedicated(blob_id, BlobWriteSource::Bytes(data_col.value(i))) - .await?; - - kind_builder.append_value(BlobKind::Dedicated as u8); - data_builder.append_null(); - uri_builder.append_null(); - blob_id_builder.append_value(blob_id); - blob_size_builder.append_value(data_len as u64); - position_builder.append_null(); + let value = Self::write_dedicated( + self.object_store.clone(), + self.data_dir.clone(), + self.data_file_key.clone(), + self.blob_id_allocator.clone(), + BlobWriteSource::Bytes(data_col.value(i)), + ) + .await?; + blob_writer.push(value)?; continue; } if has_data && data_len > inline_threshold { - let (pack_blob_id, position) = self - .write_packed(BlobWriteSource::Bytes(data_col.value(i))) + let value = self + .write_packed( + pack_file_threshold, + BlobWriteSource::Bytes(data_col.value(i)), + ) .await?; - - kind_builder.append_value(BlobKind::Packed as u8); - data_builder.append_null(); - uri_builder.append_null(); - blob_id_builder.append_value(pack_blob_id); - blob_size_builder.append_value(data_len as u64); - position_builder.append_value(position); + blob_writer.push(value)?; continue; } @@ -785,114 +913,63 @@ impl BlobPreprocessor { let data_len = source.size(); if data_len > dedicated_threshold as u64 { - let blob_id = self.next_blob_id(); - self.write_dedicated(blob_id, BlobWriteSource::External(&source)) - .await?; - - kind_builder.append_value(BlobKind::Dedicated as u8); - data_builder.append_null(); - uri_builder.append_null(); - blob_id_builder.append_value(blob_id); - blob_size_builder.append_value(data_len); - position_builder.append_null(); + let value = Self::write_dedicated( + self.object_store.clone(), + self.data_dir.clone(), + self.data_file_key.clone(), + self.blob_id_allocator.clone(), + BlobWriteSource::External(&source), + ) + .await?; + blob_writer.push(value)?; continue; } if data_len > inline_threshold as u64 { - let (pack_blob_id, position) = self - .write_packed(BlobWriteSource::External(&source)) + let value = self + .write_packed(pack_file_threshold, BlobWriteSource::External(&source)) .await?; - - kind_builder.append_value(BlobKind::Packed as u8); - data_builder.append_null(); - uri_builder.append_null(); - blob_id_builder.append_value(pack_blob_id); - blob_size_builder.append_value(data_len); - position_builder.append_value(position); + blob_writer.push(value)?; continue; } let data = source.read_all().await?; - - kind_builder.append_value(BlobKind::Inline as u8); - data_builder.append_value(data.as_ref()); - uri_builder.append_null(); - blob_id_builder.append_null(); - blob_size_builder.append_null(); - position_builder.append_null(); + blob_writer.push_inline(data)?; continue; } let (external_base_id, external_uri_or_path) = self.resolve_external_reference(uri_val).await?; - kind_builder.append_value(BlobKind::External as u8); - data_builder.append_null(); - uri_builder.append_value(external_uri_or_path); - blob_id_builder.append_value(external_base_id); - if has_position && has_size { - let position = position_col - .as_ref() - .expect("position column must exist") - .value(i); - let size = size_col.as_ref().expect("size column must exist").value(i); - blob_size_builder.append_value(size); - position_builder.append_value(position); + let range = if has_position && has_size { + BlobRange { + offset: position_col + .as_ref() + .expect("position column must exist") + .value(i), + size: size_col.as_ref().expect("size column must exist").value(i), + } } else { - blob_size_builder.append_null(); - position_builder.append_null(); - } + BlobRange { offset: 0, size: 0 } + }; + blob_writer.push(BlobDescriptor::External { + base_id: external_base_id, + uri: external_uri_or_path, + offset: range.offset, + size: range.size, + })?; continue; } if has_data { - kind_builder.append_value(BlobKind::Inline as u8); - let value = data_col.value(i); - data_builder.append_value(value); - uri_builder.append_null(); - blob_id_builder.append_null(); - blob_size_builder.append_null(); - position_builder.append_null(); + blob_writer.push_inline(Bytes::copy_from_slice(data_col.value(i)))?; } else { - data_builder.append_null(); - uri_builder.append_null(); - blob_id_builder.append_null(); - blob_size_builder.append_null(); - kind_builder.append_null(); - position_builder.append_null(); + blob_writer.push_null()?; } } - let child_fields = vec![ - ArrowField::new("kind", ArrowDataType::UInt8, true), - ArrowField::new("data", ArrowDataType::LargeBinary, true), - ArrowField::new("uri", ArrowDataType::Utf8, true), - ArrowField::new("blob_id", ArrowDataType::UInt32, true), - ArrowField::new("blob_size", ArrowDataType::UInt64, true), - ArrowField::new("position", ArrowDataType::UInt64, true), - ]; - - let struct_array = StructArray::try_new( - child_fields.clone().into(), - vec![ - Arc::new(kind_builder.finish()), - Arc::new(data_builder.finish()), - Arc::new(uri_builder.finish()), - Arc::new(blob_id_builder.finish()), - Arc::new(blob_size_builder.finish()), - Arc::new(position_builder.finish()), - ], - struct_nulls.cloned(), - )?; - - let field = Arc::new( - ArrowField::new( - field.name(), - ArrowDataType::Struct(child_fields.into()), - field.is_nullable(), - ) - .with_metadata(writer_metadata.clone()), - ); - Ok((Arc::new(struct_array), field)) + let column = blob_writer.finish()?; + let (field, array) = column.into_parts(); + Ok((array, Arc::new(field))) } pub(crate) async fn finish(&mut self) -> Result<()> { @@ -1827,6 +1904,27 @@ async fn execute_blob_read_plan( .collect()) } +async fn execute_blob_entries( + entries: Vec, + io_parallelism: usize, + io_buffer_size_bytes: Option, +) -> Result> { + let plans = plan_blob_read_plans(entries); + if plans.is_empty() { + return Ok(Vec::new()); + } + + let execution = Arc::new(ReadBlobsExecution::new(io_buffer_size_bytes)); + let batches = stream::iter(plans.into_iter().map(move |plan| { + let execution = execution.clone(); + execute_blob_read_plan(plan, execution) + })) + .buffer_unordered(io_parallelism.max(1)) + .try_collect::>() + .await?; + Ok(batches.into_iter().flatten().collect()) +} + pub(super) async fn take_blobs( dataset: &Arc, row_ids: &[u64], @@ -1960,16 +2058,57 @@ async fn collect_blob_entries_for_selection( /// descriptor `StructArray`, descending through nested struct children for /// dotted paths. fn leaf_descriptor_struct<'a>(batch: &'a RecordBatch, column: &str) -> Result<&'a StructArray> { + let current = leaf_descriptor_array(batch, column)?; + current + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Blob column '{}' expected descriptor struct but got {}", + column, + current.data_type() + ) + .into(), + ) + }) +} + +fn leaf_descriptor_array<'a>(batch: &'a RecordBatch, column: &str) -> Result<&'a dyn Array> { let path = parse_field_path(column)?; - let mut current = batch + let mut current: &dyn Array = batch .column_by_name(&path[0]) - .expect("validate_blob_column ensured column exists") - .as_struct(); + .ok_or_else(|| { + Error::invalid_input(format!( + "Blob column '{}' was not found in descriptor batch", + column + )) + })? + .as_ref(); for segment in &path[1..] { - current = current + let struct_array = current + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Blob column path '{}' expected struct before segment '{}' but got {}", + column, + segment, + current.data_type() + ) + .into(), + ) + })?; + current = struct_array .column_by_name(segment) - .expect("validate_blob_column ensured all path segments exist") - .as_struct(); + .ok_or_else(|| { + Error::invalid_input(format!( + "Blob column path '{}' missing segment '{}'", + column, segment + )) + })? + .as_ref(); } Ok(current) } @@ -1995,6 +2134,38 @@ fn blob_version_from_descriptions(descriptions: &StructArray) -> Result { + descriptions: &'a StructArray, + kinds: &'a arrow::array::PrimitiveArray, + positions: &'a arrow::array::PrimitiveArray, + sizes: &'a arrow::array::PrimitiveArray, + blob_ids: &'a arrow::array::PrimitiveArray, + blob_uris: &'a arrow::array::GenericStringArray, +} + +impl<'a> BlobV2DescriptorColumns<'a> { + fn new(descriptions: &'a StructArray) -> Self { + Self { + descriptions, + kinds: descriptions.column(0).as_primitive::(), + positions: descriptions.column(1).as_primitive::(), + sizes: descriptions.column(2).as_primitive::(), + blob_ids: descriptions.column(3).as_primitive::(), + blob_uris: descriptions.column(4).as_string::(), + } + } + + fn is_null_blob(&self, idx: usize) -> Result { + if self.descriptions.is_null(idx) || self.kinds.is_null(idx) { + return Ok(true); + } + let kind = BlobKind::try_from(self.kinds.value(idx))?; + Ok(matches!(kind, BlobKind::Inline) + && self.positions.value(idx) == 0 + && self.sizes.value(idx) == 0) + } +} + /// Convert blob v1 descriptors into logical blob entries. fn collect_blob_entries_v1( dataset: &Arc, @@ -2057,193 +2228,609 @@ async fn collect_blob_entries_v2( descriptions: &StructArray, row_addrs: &arrow::array::PrimitiveArray, ) -> Result> { - let kinds = descriptions.column(0).as_primitive::(); - let positions = descriptions.column(1).as_primitive::(); - let sizes = descriptions.column(2).as_primitive::(); - let blob_ids = descriptions.column(3).as_primitive::(); - let blob_uris = descriptions.column(4).as_string::(); - + let columns = BlobV2DescriptorColumns::new(descriptions); let mut files = Vec::with_capacity(row_addrs.len()); - let mut fragment_cache = HashMap::::new(); - let mut store_cache = HashMap::>::new(); - let mut external_base_path_cache = HashMap::::new(); - let mut source_cache = HashMap::>::new(); + let mut read_context = BlobV2ReadContext::new(dataset, blob_field_id); for (selection_index, row_addr) in row_addrs.values().iter().enumerate() { - let idx = selection_index; - let kind = BlobKind::try_from(kinds.value(idx))?; - - // Struct is non-nullable; null rows are encoded as inline with zero position/size and empty uri - if matches!(kind, BlobKind::Inline) && positions.value(idx) == 0 && sizes.value(idx) == 0 { - continue; - } - - match kind { - BlobKind::Inline => { - let position = positions.value(idx); - let size = sizes.value(idx); - let location = resolve_blob_read_location( - dataset, - blob_field_id, - *row_addr, - &mut fragment_cache, - &mut store_cache, - ) - .await?; - let source = shared_blob_source( - &mut source_cache, - location.object_store, - &location.data_file_path, - ); - files.push(BlobEntry { - selection_index, - row_address: *row_addr, - file: BlobFile::with_source(source, position, size, BlobKind::Inline, None), - }); - } - BlobKind::Dedicated => { - let blob_id = blob_ids.value(idx); - let size = sizes.value(idx); - let location = resolve_blob_read_location( - dataset, - blob_field_id, - *row_addr, - &mut fragment_cache, - &mut store_cache, - ) - .await?; - let path = blob_path(&location.data_file_dir, &location.data_file_key, blob_id); - let source = shared_blob_source(&mut source_cache, location.object_store, &path); - files.push(BlobEntry { - selection_index, - row_address: *row_addr, - file: BlobFile::with_source(source, 0, size, BlobKind::Dedicated, None), - }); - } - BlobKind::Packed => { - let blob_id = blob_ids.value(idx); - let size = sizes.value(idx); - let position = positions.value(idx); - let location = resolve_blob_read_location( - dataset, - blob_field_id, - *row_addr, - &mut fragment_cache, - &mut store_cache, - ) - .await?; - let path = blob_path(&location.data_file_dir, &location.data_file_key, blob_id); - let source = shared_blob_source(&mut source_cache, location.object_store, &path); - files.push(BlobEntry { - selection_index, - row_address: *row_addr, - file: BlobFile::with_source(source, position, size, BlobKind::Packed, None), - }); - } - BlobKind::External => { - let uri_or_path = blob_uris.value(idx).to_string(); - let position = positions.value(idx); - let size = sizes.value(idx); - let base_id = blob_ids.value(idx); - let (object_store, path) = if base_id == 0 { - let registry = dataset.session.store_registry(); - let params = dataset - .store_params - .as_ref() - .map(|p| Arc::new((**p).clone())) - .unwrap_or_else(|| Arc::new(ObjectStoreParams::default())); - ObjectStore::from_uri_and_params(registry, &uri_or_path, ¶ms).await? - } else { - let object_store = if let Some(store) = store_cache.get(&base_id) { - store.clone() - } else { - let store = dataset.object_store(Some(base_id)).await?; - store_cache.insert(base_id, store.clone()); - store - }; - let base_root = if let Some(path) = external_base_path_cache.get(&base_id) { - path.clone() - } else { - let base = dataset.manifest.base_paths.get(&base_id).ok_or_else(|| { - Error::invalid_input(format!( - "External blob references unknown base_id {}", - base_id - )) - })?; - let path = base.extract_path(dataset.session.store_registry())?; - external_base_path_cache.insert(base_id, path.clone()); - path - }; - let path = join_base_and_relative_path(&base_root, &uri_or_path)?; - (object_store, path) - }; - let size = if size > 0 { - size - } else { - object_store.size(&path).await? - }; - let source = shared_blob_source(&mut source_cache, object_store, &path); - files.push(BlobEntry { - selection_index, - row_address: *row_addr, - file: BlobFile::with_source( - source, - position, - size, - BlobKind::External, - Some(uri_or_path), - ), - }); - } + if let Some(entry) = read_context + .collect_entry(&columns, selection_index, selection_index, *row_addr) + .await? + { + files.push(entry); } } Ok(files) } -fn normalize_external_absolute_uri(uri: &str) -> Result { - let url = Url::parse(uri).map_err(|_| { - Error::invalid_input(format!( - "External URI '{}' is outside registered external bases and is not a valid absolute URI", - uri - )) - })?; - Ok(url.to_string()) +fn is_blob_v2_binary_view(field: &LanceField) -> bool { + field.is_blob_v2() && matches!(field.data_type(), ArrowDataType::LargeBinary) } -fn join_base_and_relative_path(base: &Path, relative_path: &str) -> Result { - let relative = Path::parse(relative_path).map_err(|e| { - Error::invalid_input(format!( - "Invalid relative external blob path '{}': {}", - relative_path, e - )) - })?; - Ok(Path::from_iter(base.parts().chain(relative.parts()))) +fn public_blob_v2_binary_output_field(mut field: LanceField) -> LanceField { + if is_blob_v2_binary_view(&field) { + field.metadata.remove(ARROW_EXT_NAME_KEY); + } + field.children = field + .children + .into_iter() + .map(public_blob_v2_binary_output_field) + .collect(); + field } -/// Resolve the physical read location for a blob row in a base-aware way. +/// Return the public Arrow-facing schema for a blob v2 binary scan. /// -/// Given a `row_addr`, this helper locates the owning fragment and the blob field's -/// data file, then returns the concrete object store and paths needed to read blob -/// bytes correctly under multi-base datasets. +/// Scan planning uses a blob v2 extension marker on `LargeBinary` leaves to +/// identify payloads that need descriptor-based materialization. This helper +/// removes that internal marker before the schema is exposed to callers. +pub fn public_blob_v2_binary_output_schema(schema: &Schema) -> Schema { + Schema { + fields: schema + .fields + .iter() + .cloned() + .map(public_blob_v2_binary_output_field) + .collect(), + metadata: schema.metadata.clone(), + } +} + +fn field_has_blob_v2_binary_view(field: &LanceField) -> bool { + is_blob_v2_binary_view(field) || field.children.iter().any(field_has_blob_v2_binary_view) +} + +/// Return true if the schema contains a blob v2 leaf in binary payload view. /// -/// It uses two caller-provided caches: -/// - `fragment_cache` memoizes per-fragment path metadata (`data_file_dir`, -/// `data_file_path`, and `data_file_key`) plus the resolved store. -/// - `store_cache` memoizes `base_id -> ObjectStore` so multiple fragments that -/// share the same base do not repeat async store resolution. -async fn resolve_blob_read_location( - dataset: &Arc, - blob_field_id: u32, - row_addr: u64, - fragment_cache: &mut HashMap, - store_cache: &mut HashMap>, -) -> Result { - let frag_id = RowAddress::from(row_addr).fragment_id(); - if let Some(location) = fragment_cache.get(&frag_id) { - return Ok(location.clone()); +/// This detects the internal `LargeBinary` view created by +/// [`BlobHandling::AllBinary`](lance_core::datatypes::BlobHandling::AllBinary) +/// or selective binary blob handling. +pub fn schema_has_blob_v2_binary_view(schema: &Schema) -> bool { + schema.fields.iter().any(field_has_blob_v2_binary_view) +} + +fn blob_v2_descriptor_field(mut field: LanceField) -> LanceField { + if is_blob_v2_binary_view(&field) { + field.unloaded_mut(); + return field; } - let frag = dataset + field.children = field + .children + .into_iter() + .map(blob_v2_descriptor_field) + .collect(); + field +} + +/// Convert blob v2 binary-view leaves back to descriptor-view leaves. +/// +/// Readers use this schema to fetch stored blob descriptors first. The scan +/// layer then materializes those descriptors into the caller's binary payload +/// view after row addresses are available. +pub fn blob_v2_descriptor_schema(schema: &Schema) -> Schema { + Schema { + fields: schema + .fields + .iter() + .cloned() + .map(blob_v2_descriptor_field) + .collect(), + metadata: schema.metadata.clone(), + } +} + +/// Materialize blob v2 descriptor arrays in a decoded batch into binary arrays. +/// +/// The input batch must include `_rowaddr`, which is used to resolve packed, +/// dedicated, inline, and external blob payload locations. `output_schema` +/// defines the exact returned columns, including requested system columns, with +/// blob v2 binary leaves exposed as plain `LargeBinary` fields. +pub async fn materialize_blob_v2_binary_batch( + dataset: &Arc, + output_schema: &Schema, + batch: RecordBatch, +) -> Result { + let row_addr_idx = batch + .schema() + .column_with_name(ROW_ADDR) + .ok_or_else(|| { + Error::internal(format!( + "_rowaddr column missing from blob v2 binary scan batch, columns: {:?}", + batch + .schema() + .fields() + .iter() + .map(|field| field.name()) + .collect::>() + )) + })? + .0; + let row_addrs = batch + .column(row_addr_idx) + .as_primitive::() + .values() + .iter() + .copied() + .collect::>(); + let row_addrs: Arc<[u64]> = row_addrs.into(); + + let mut columns = Vec::with_capacity(output_schema.fields.len()); + let mut fields = Vec::with_capacity(output_schema.fields.len()); + + for field in &output_schema.fields { + let input = batch + .column_by_name(&field.name) + .ok_or_else(|| { + Error::internal(format!( + "blob v2 binary scan batch missing projected column '{}'", + field.name + )) + })? + .clone(); + let materialized = + materialize_blob_v2_binary_array(dataset, field, input, row_addrs.clone()).await?; + columns.push(materialized); + let output_field = public_blob_v2_binary_output_field(field.clone()); + fields.push(ArrowField::from(&output_field)); + } + + Ok(RecordBatch::try_new( + Arc::new(ArrowSchema::new_with_metadata( + fields, + batch.schema().metadata().clone(), + )), + columns, + )?) +} + +fn materialize_blob_v2_binary_array<'a>( + dataset: &'a Arc, + field: &'a LanceField, + array: ArrayRef, + row_addrs: Arc<[u64]>, +) -> BoxFuture<'a, Result> { + async move { + if is_blob_v2_binary_view(field) { + let descriptions = array.as_struct(); + return materialize_blob_v2_descriptors( + dataset, + field.id as u32, + descriptions, + row_addrs.as_ref(), + ) + .await; + } + + match field.data_type() { + ArrowDataType::Struct(_) => { + let struct_array = array.as_struct(); + let mut children = Vec::with_capacity(field.children.len()); + for (child_field, child_array) in + field.children.iter().zip(struct_array.columns().iter()) + { + children.push( + materialize_blob_v2_binary_array( + dataset, + child_field, + child_array.clone(), + row_addrs.clone(), + ) + .await?, + ); + } + let public_field = public_blob_v2_binary_output_field(field.clone()); + let ArrowDataType::Struct(fields) = public_field.data_type() else { + unreachable!("public output field preserved struct type") + }; + Ok(Arc::new(StructArray::try_new( + fields, + children, + struct_array.nulls().cloned(), + )?) as ArrayRef) + } + ArrowDataType::List(_) => { + let list_array = array.as_list::(); + materialize_blob_v2_list_array::(dataset, field, list_array, row_addrs).await + } + ArrowDataType::LargeList(_) => { + let list_array = array.as_list::(); + materialize_blob_v2_list_array::(dataset, field, list_array, row_addrs).await + } + _ => Ok(array), + } + } + .boxed() +} + +async fn materialize_blob_v2_list_array( + dataset: &Arc, + field: &LanceField, + list_array: &GenericListArray, + row_addrs: Arc<[u64]>, +) -> Result { + let offsets = list_array.value_offsets(); + let values_start = offsets[0].as_usize(); + let values_end = offsets[list_array.len()].as_usize(); + if values_end < values_start { + return Err(Error::internal(format!( + "List field '{}' has invalid offsets while materializing blob v2 binary scan", + field.name + ))); + } + + let values_len = values_end - values_start; + let mut normalized_offsets = Vec::with_capacity(list_array.len() + 1); + normalized_offsets.push(O::usize_as(0)); + let mut child_row_addrs = Vec::with_capacity(values_len); + for row_idx in 0..list_array.len() { + let start = offsets[row_idx].as_usize(); + let end = offsets[row_idx + 1].as_usize(); + if end < start { + return Err(Error::internal(format!( + "List field '{}' has decreasing offsets while materializing blob v2 binary scan", + field.name + ))); + } + let row_addr = row_addrs.get(row_idx).copied().ok_or_else(|| { + Error::internal(format!( + "List field '{}' row address count {} did not match row count {}", + field.name, + row_addrs.len(), + list_array.len() + )) + })?; + for _ in start..end { + child_row_addrs.push(row_addr); + } + normalized_offsets.push(O::usize_as(end - values_start)); + } + let child_row_addrs: Arc<[u64]> = child_row_addrs.into(); + let child = field.children.first().ok_or_else(|| { + Error::internal(format!( + "List field '{}' missing child while materializing blob v2 binary scan", + field.name + )) + })?; + let values = list_array.values().slice(values_start, values_len); + let values = materialize_blob_v2_binary_array(dataset, child, values, child_row_addrs).await?; + let child_field = public_blob_v2_binary_output_field(child.clone()); + let list_array = GenericListArray::::try_new( + Arc::new(ArrowField::from(&child_field)), + OffsetBuffer::new(ScalarBuffer::from(normalized_offsets)), + values, + list_array.nulls().cloned(), + )?; + Ok(Arc::new(list_array)) +} + +async fn materialize_blob_v2_descriptors( + dataset: &Arc, + blob_field_id: u32, + descriptions: &StructArray, + row_addrs: &[u64], +) -> Result { + if descriptions.len() != row_addrs.len() { + return Err(Error::internal(format!( + "blob v2 descriptor count {} did not match row address count {}", + descriptions.len(), + row_addrs.len() + ))); + } + match blob_version_from_descriptions(descriptions)? { + BlobVersion::V1 => { + return Err(Error::not_supported( + "Blob v2 binary materialization received a legacy blob descriptor".to_string(), + )); + } + BlobVersion::V2 => {} + } + + let columns = BlobV2DescriptorColumns::new(descriptions); + let mut read_context = BlobV2ReadContext::new(dataset, blob_field_id); + let mut entries = Vec::with_capacity(descriptions.len()); + let mut payloads = vec![None; descriptions.len()]; + + for (idx, row_addr) in row_addrs.iter().copied().enumerate() { + if descriptions.is_null(idx) || columns.kinds.is_null(idx) { + continue; + } + + let kind = BlobKind::try_from(columns.kinds.value(idx))?; + if matches!(kind, BlobKind::Inline) + && columns.positions.value(idx) == 0 + && columns.sizes.value(idx) == 0 + { + payloads[idx] = Some(Bytes::new()); + continue; + } + + let entry = read_context + .collect_entry(&columns, idx, idx, row_addr) + .await? + .ok_or_else(|| { + Error::internal(format!( + "blob v2 descriptor at index {idx} unexpectedly resolved to null" + )) + })?; + entries.push(entry); + } + + let blobs = execute_blob_entries(entries, dataset.object_store.io_parallelism(), None).await?; + for blob in blobs { + let payload = payloads.get_mut(blob.selection_index).ok_or_else(|| { + Error::internal(format!( + "blob result selection index {} exceeded descriptor count {}", + blob.selection_index, + descriptions.len() + )) + })?; + if payload.replace(blob.data).is_some() { + return Err(Error::internal(format!( + "blob result selection index {} was produced more than once", + blob.selection_index + ))); + } + } + + let payload_capacity = payloads.iter().flatten().map(Bytes::len).sum::(); + let mut builder = LargeBinaryBuilder::with_capacity(descriptions.len(), payload_capacity); + for (idx, payload) in payloads.into_iter().enumerate() { + if descriptions.is_null(idx) || columns.kinds.is_null(idx) { + builder.append_null(); + } else { + let payload = payload.ok_or_else(|| { + Error::internal(format!( + "blob v2 descriptor at index {idx} did not produce a payload" + )) + })?; + builder.append_value(payload); + } + } + + Ok(Arc::new(builder.finish())) +} + +struct BlobV2ReadContext<'a> { + dataset: &'a Arc, + blob_field_id: u32, + fragment_cache: HashMap, + store_cache: HashMap>, + external_base_path_cache: HashMap, + source_cache: HashMap>, +} + +impl<'a> BlobV2ReadContext<'a> { + fn new(dataset: &'a Arc, blob_field_id: u32) -> Self { + Self { + dataset, + blob_field_id, + fragment_cache: HashMap::new(), + store_cache: HashMap::new(), + external_base_path_cache: HashMap::new(), + source_cache: HashMap::new(), + } + } + + async fn collect_entry( + &mut self, + columns: &BlobV2DescriptorColumns<'_>, + idx: usize, + selection_index: usize, + row_addr: u64, + ) -> Result> { + if columns.is_null_blob(idx)? { + return Ok(None); + } + + let kind = BlobKind::try_from(columns.kinds.value(idx))?; + let entry = match kind { + BlobKind::Inline => { + self.collect_inline(columns, idx, selection_index, row_addr) + .await? + } + BlobKind::Dedicated => { + self.collect_dedicated(columns, idx, selection_index, row_addr) + .await? + } + BlobKind::Packed => { + self.collect_packed(columns, idx, selection_index, row_addr) + .await? + } + BlobKind::External => { + self.collect_external(columns, idx, selection_index, row_addr) + .await? + } + }; + + Ok(Some(entry)) + } + + async fn blob_read_location(&mut self, row_addr: u64) -> Result { + resolve_blob_read_location( + self.dataset, + self.blob_field_id, + row_addr, + &mut self.fragment_cache, + &mut self.store_cache, + ) + .await + } + + async fn collect_inline( + &mut self, + columns: &BlobV2DescriptorColumns<'_>, + idx: usize, + selection_index: usize, + row_addr: u64, + ) -> Result { + let position = columns.positions.value(idx); + let size = columns.sizes.value(idx); + let location = self.blob_read_location(row_addr).await?; + let source = shared_blob_source( + &mut self.source_cache, + location.object_store, + &location.data_file_path, + ); + Ok(BlobEntry { + selection_index, + row_address: row_addr, + file: BlobFile::with_source(source, position, size, BlobKind::Inline, None), + }) + } + + async fn collect_dedicated( + &mut self, + columns: &BlobV2DescriptorColumns<'_>, + idx: usize, + selection_index: usize, + row_addr: u64, + ) -> Result { + let blob_id = columns.blob_ids.value(idx); + let size = columns.sizes.value(idx); + let location = self.blob_read_location(row_addr).await?; + let path = blob_path(&location.data_file_dir, &location.data_file_key, blob_id); + let source = shared_blob_source(&mut self.source_cache, location.object_store, &path); + Ok(BlobEntry { + selection_index, + row_address: row_addr, + file: BlobFile::with_source(source, 0, size, BlobKind::Dedicated, None), + }) + } + + async fn collect_packed( + &mut self, + columns: &BlobV2DescriptorColumns<'_>, + idx: usize, + selection_index: usize, + row_addr: u64, + ) -> Result { + let blob_id = columns.blob_ids.value(idx); + let size = columns.sizes.value(idx); + let position = columns.positions.value(idx); + let location = self.blob_read_location(row_addr).await?; + let path = blob_path(&location.data_file_dir, &location.data_file_key, blob_id); + let source = shared_blob_source(&mut self.source_cache, location.object_store, &path); + Ok(BlobEntry { + selection_index, + row_address: row_addr, + file: BlobFile::with_source(source, position, size, BlobKind::Packed, None), + }) + } + + async fn collect_external( + &mut self, + columns: &BlobV2DescriptorColumns<'_>, + idx: usize, + selection_index: usize, + row_addr: u64, + ) -> Result { + let uri_or_path = columns.blob_uris.value(idx).to_string(); + let position = columns.positions.value(idx); + let size = columns.sizes.value(idx); + let base_id = columns.blob_ids.value(idx); + let (object_store, path) = if base_id == 0 { + let registry = self.dataset.session.store_registry(); + let params = self + .dataset + .store_params + .as_ref() + .map(|p| Arc::new((**p).clone())) + .unwrap_or_else(|| Arc::new(ObjectStoreParams::default())); + ObjectStore::from_uri_and_params(registry, &uri_or_path, ¶ms).await? + } else { + let object_store = if let Some(store) = self.store_cache.get(&base_id) { + store.clone() + } else { + let store = self.dataset.object_store(Some(base_id)).await?; + self.store_cache.insert(base_id, store.clone()); + store + }; + let base_root = if let Some(path) = self.external_base_path_cache.get(&base_id) { + path.clone() + } else { + let base = self + .dataset + .manifest + .base_paths + .get(&base_id) + .ok_or_else(|| { + Error::invalid_input(format!( + "External blob references unknown base_id {}", + base_id + )) + })?; + let path = base.extract_path(self.dataset.session.store_registry())?; + self.external_base_path_cache.insert(base_id, path.clone()); + path + }; + let path = join_base_and_relative_path(&base_root, &uri_or_path)?; + (object_store, path) + }; + let size = if size > 0 { + size + } else { + object_store.size(&path).await? + }; + let source = shared_blob_source(&mut self.source_cache, object_store, &path); + Ok(BlobEntry { + selection_index, + row_address: row_addr, + file: BlobFile::with_source( + source, + position, + size, + BlobKind::External, + Some(uri_or_path), + ), + }) + } +} + +fn normalize_external_absolute_uri(uri: &str) -> Result { + let url = Url::parse(uri).map_err(|_| { + Error::invalid_input(format!( + "External URI '{}' is outside registered external bases and is not a valid absolute URI", + uri + )) + })?; + Ok(url.to_string()) +} + +fn join_base_and_relative_path(base: &Path, relative_path: &str) -> Result { + let relative = Path::parse(relative_path).map_err(|e| { + Error::invalid_input(format!( + "Invalid relative external blob path '{}': {}", + relative_path, e + )) + })?; + Ok(Path::from_iter(base.parts().chain(relative.parts()))) +} + +/// Resolve the physical read location for a blob row in a base-aware way. +/// +/// Given a `row_addr`, this helper locates the owning fragment and the blob field's +/// data file, then returns the concrete object store and paths needed to read blob +/// bytes correctly under multi-base datasets. +/// +/// It uses two caller-provided caches: +/// - `fragment_cache` memoizes per-fragment path metadata (`data_file_dir`, +/// `data_file_path`, and `data_file_key`) plus the resolved store. +/// - `store_cache` memoizes `base_id -> ObjectStore` so multiple fragments that +/// share the same base do not repeat async store resolution. +async fn resolve_blob_read_location( + dataset: &Arc, + blob_field_id: u32, + row_addr: u64, + fragment_cache: &mut HashMap, + store_cache: &mut HashMap>, +) -> Result { + let frag_id = RowAddress::from(row_addr).fragment_id(); + if let Some(location) = fragment_cache.get(&frag_id) { + return Ok(location.clone()); + } + + let frag = dataset .get_fragment(frag_id as usize) .ok_or_else(|| Error::internal("Fragment not found".to_string()))?; let data_file = frag @@ -2293,18 +2880,23 @@ mod tests { }; use arrow_array::RecordBatch; use arrow_array::{ - Array, ArrayRef, RecordBatchIterator, StringArray, StructArray, UInt32Array, UInt64Array, + Array, ArrayRef, Int32Array, LargeBinaryArray, RecordBatchIterator, StringArray, + StructArray, UInt8Array, UInt32Array, UInt64Array, }; use arrow_schema::{DataType, Field, Schema}; use async_trait::async_trait; use bytes::Bytes; use chrono::Utc; - use futures::{StreamExt, TryStreamExt, future::try_join_all}; + use futures::{StreamExt, TryStreamExt}; use lance_arrow::{ ARROW_EXT_NAME_KEY, BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, - BLOB_INLINE_SIZE_THRESHOLD_META_KEY, BLOB_V2_EXT_NAME, DataTypeExt, + BLOB_INLINE_SIZE_THRESHOLD_META_KEY, BLOB_META_KEY, BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY, + BLOB_V2_EXT_NAME, DataTypeExt, + }; + use lance_core::{ + datatypes::{BlobHandling, BlobKind}, + utils::blob::blob_path, }; - use lance_core::datatypes::BlobKind; use lance_io::object_store::{ ObjectStore, ObjectStoreParams, ObjectStoreRegistry, StorageOptionsAccessor, }; @@ -2319,21 +2911,29 @@ mod tests { use url::Url; use lance_core::{ - Error, Result, + Error, ROW_ADDR, ROW_CREATED_AT_VERSION, ROW_ID, ROW_LAST_UPDATED_AT_VERSION, Result, utils::tempfile::{TempDir, TempStrDir}, }; use lance_datagen::{BatchCount, RowCount, array}; - use lance_file::version::LanceFileVersion; + use lance_file::{ + version::LanceFileVersion, + writer::{FileWriter, FileWriterOptions}, + }; + use uuid::Uuid; use super::{ BlobEntry, BlobFile, BlobSource, ExternalBaseCandidate, ExternalBaseResolver, - ReadBlobsExecution, collect_blob_entries_v1, data_file_key_from_path, + ReadBlobsExecution, collect_blob_entries_v1, data_file_key_from_path, execute_blob_entries, execute_blob_read_plan, plan_blob_read_plans, }; use crate::{ Dataset, - blob::{BlobArrayBuilder, blob_field}, - dataset::{ExternalBlobMode, WriteMode, WriteParams}, + blob::{BlobArrayBuilder, BlobDescriptorArrayBuilder, PackedBlobWriter, blob_field}, + dataset::{ + CommitBuilder, ExternalBlobMode, WriteMode, WriteParams, + scanner::MaterializationStyle, + transaction::{DataReplacementGroup, Operation, Transaction}, + }, utils::test::TestDatasetGenerator, }; @@ -3196,84 +3796,867 @@ mod tests { .try_into_batch() .await .unwrap(); - assert_eq!(bytes.column(0).data_type(), &DataType::LargeBinary); - let blobs = bytes.column(0).as_binary::(); - assert_eq!(blobs.value(0), b"foo"); - assert_eq!(blobs.value(1), b"bar"); - assert_eq!(blobs.value(2), b"baz"); + assert_eq!(bytes.column(0).data_type(), &DataType::LargeBinary); + let blobs = bytes.column(0).as_binary::(); + assert_eq!(blobs.value(0), b"foo"); + assert_eq!(blobs.value(1), b"bar"); + assert_eq!(blobs.value(2), b"baz"); + + // Pass 3: back to descriptor view to verify both directions are safe. + let mut scanner = dataset.scan(); + scanner.blob_handling(BlobHandling::BlobsDescriptions); + let descriptors = scanner + .project(&["blobs"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert!(descriptors.column(0).data_type().is_struct()); + } + + #[tokio::test] + async fn test_v2_0_legacy_blob_descriptor_projection_and_reads() { + let test_dir = TempStrDir::default(); + let blob_meta = HashMap::from([(BLOB_META_KEY.to_string(), "true".to_string())]); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("blob", DataType::LargeBinary, true).with_metadata(blob_meta), + ])); + let payloads = [ + b"abc".as_slice(), + b"defgh".as_slice(), + b"ijklmnop".as_slice(), + ]; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef, + Arc::new(LargeBinaryArray::from_iter_values(payloads)) as ArrayRef, + ], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_0), + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + for blob_handling in [ + None, + Some(BlobHandling::BlobsDescriptions), + Some(BlobHandling::AllDescriptions), + ] { + let mut scanner = dataset.scan(); + if let Some(blob_handling) = blob_handling { + scanner.blob_handling(blob_handling); + } + let descriptors = scanner + .project(&["blob"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let descriptor = descriptors.column(0).as_struct(); + assert_eq!(descriptor.fields().len(), 2); + assert_eq!(descriptor.fields()[0].name(), "position"); + assert_eq!(descriptor.fields()[1].name(), "size"); + let sizes = descriptor + .column_by_name("size") + .unwrap() + .as_primitive::(); + assert_eq!(sizes.values(), &[3, 5, 8]); + } + + let mut scanner = dataset.scan(); + scanner.blob_handling(BlobHandling::AllBinary); + let bytes = scanner + .project(&["blob"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let bytes = bytes.column(0).as_binary::(); + for (idx, expected) in payloads.iter().enumerate() { + assert_eq!(bytes.value(idx), *expected); + } + + let blob_files = dataset + .take_blobs_by_indices(&[0, 1, 2], "blob") + .await + .unwrap(); + assert_eq!(blob_files.len(), payloads.len()); + for (blob_file, expected) in blob_files.iter().zip(payloads) { + assert_eq!(blob_file.read().await.unwrap().as_ref(), expected); + } + + let read_blobs = dataset + .read_blobs("blob") + .unwrap() + .with_row_indices(vec![0, 1, 2]) + .execute() + .await + .unwrap(); + assert_eq!(read_blobs.len(), payloads.len()); + for (read_blob, expected) in read_blobs.iter().zip(payloads) { + assert_eq!(read_blob.data.as_ref(), expected); + } + } + + #[test] + fn test_data_file_key_from_path() { + assert_eq!(data_file_key_from_path("data/abc.lance"), "abc"); + assert_eq!(data_file_key_from_path("abc.lance"), "abc"); + assert_eq!(data_file_key_from_path("nested/path/xyz"), "xyz"); + } + + #[tokio::test] + async fn test_write_and_take_blobs_with_blob_descriptor_array_builder() { + let test_dir = TempStrDir::default(); + + // Build a blob column with the new BlobArrayBuilder + let mut blob_builder = BlobArrayBuilder::new(2); + blob_builder.push_bytes(b"hello").unwrap(); + blob_builder.push_bytes(b"world").unwrap(); + let blob_array: arrow_array::ArrayRef = blob_builder.finish().unwrap(); + + let id_array: arrow_array::ArrayRef = Arc::new(UInt32Array::from(vec![0, 1])); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + blob_field("blob", true), + ])); + + let batch = RecordBatch::try_new(schema.clone(), vec![id_array, blob_array]).unwrap(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); + + let params = WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }; + let dataset = Arc::new( + Dataset::write(reader, &test_dir, Some(params)) + .await + .unwrap(), + ); + + let blobs = dataset + .take_blobs_by_indices(&[0, 1], "blob") + .await + .unwrap(); + + assert_eq!(blobs.len(), 2); + let first = blobs[0].read().await.unwrap(); + let second = blobs[1].read().await.unwrap(); + assert_eq!(first.as_ref(), b"hello"); + assert_eq!(second.as_ref(), b"world"); + } + + #[tokio::test] + async fn test_write_prepared_blob_column_normalizes_schema() { + let test_dir = TempStrDir::default(); + let mut prepared_writer = BlobDescriptorArrayBuilder::new("blob"); + prepared_writer + .push_inline(Bytes::from_static(b"prepared-inline")) + .unwrap(); + let (prepared_field, prepared_array) = prepared_writer.finish().unwrap().into_parts(); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + prepared_field, + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt32Array::from(vec![0])) as ArrayRef, + prepared_array, + ], + ) + .unwrap(); + + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + let dataset_schema = Schema::from(dataset.schema()); + let blob_field = dataset_schema.field_with_name("blob").unwrap(); + let DataType::Struct(fields) = blob_field.data_type() else { + panic!("expected blob field to be a struct"); + }; + assert_eq!(fields.len(), 2); + assert_eq!(fields[0].name(), "data"); + assert_eq!(fields[1].name(), "uri"); + + let blobs = dataset.take_blobs_by_indices(&[0], "blob").await.unwrap(); + assert_eq!(blobs.len(), 1); + assert_eq!(blobs[0].read().await.unwrap().as_ref(), b"prepared-inline"); + } + + #[tokio::test] + async fn test_data_replacement_uses_blob_descriptor_array_builder_prepared_packed_blob_column() + { + let test_dir = TempStrDir::default(); + let logical_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + blob_field("blob", true), + ])); + let mut initial_builder = BlobArrayBuilder::new(1); + initial_builder.push_bytes(b"initial").unwrap(); + let initial_batch = RecordBatch::try_new( + logical_schema.clone(), + vec![ + Arc::new(UInt32Array::from(vec![0])) as ArrayRef, + initial_builder.finish().unwrap(), + ], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(initial_batch)], logical_schema.clone()), + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + let file_id = Uuid::new_v4().to_string(); + let data_file_name = format!("{file_id}.lance"); + let data_file_path = dataset.data_dir().join(data_file_name.as_str()); + let blob_id = 1; + let packed_path = blob_path(&dataset.data_dir(), &file_id, blob_id); + let mut blob_writer = BlobDescriptorArrayBuilder::new("blob"); + let mut packed = PackedBlobWriter::try_new( + dataset.object_store.as_ref().clone(), + data_file_path.clone(), + blob_id, + ) + .await + .unwrap(); + assert_eq!(packed.path(), &packed_path); + packed.write_blob(b"prepared-packed").await.unwrap(); + blob_writer.extend(packed.finish().await.unwrap()).unwrap(); + let prepared_field = blob_writer.field().clone(); + let prepared_array = blob_writer.finish().unwrap().into_parts().1; + let append_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + prepared_field, + ])); + let replacement_batch = RecordBatch::try_new( + append_schema.clone(), + vec![ + Arc::new(UInt32Array::from(vec![1])) as ArrayRef, + prepared_array, + ], + ) + .unwrap(); + + let object_writer = dataset.object_store.create(&data_file_path).await.unwrap(); + let mut file_writer = FileWriter::try_new( + object_writer, + crate::datatypes::Schema::try_from(append_schema.as_ref()).unwrap(), + FileWriterOptions { + format_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }, + ) + .unwrap(); + file_writer.write_batch(&replacement_batch).await.unwrap(); + file_writer.finish().await.unwrap(); + + let data_file = dataset + .create_data_file(&data_file_name, None) + .await + .unwrap(); + let transaction = Transaction { + read_version: dataset.manifest.version, + uuid: Uuid::new_v4().hyphenated().to_string(), + operation: Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file)], + }, + tag: None, + transaction_properties: None, + }; + let dataset = Arc::new( + CommitBuilder::new(dataset) + .execute(transaction) + .await + .unwrap(), + ); + + let blobs = dataset.take_blobs_by_indices(&[0], "blob").await.unwrap(); + assert_eq!(blobs.len(), 1); + assert_eq!(blobs[0].read().await.unwrap().as_ref(), b"prepared-packed"); + assert_eq!(blobs[0].kind(), BlobKind::Packed); + } + + #[tokio::test] + async fn test_append_uses_blob_descriptor_array_builder_prepared_blob_column() { + let test_dir = TempStrDir::default(); + let logical_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + blob_field("blob", true), + ])); + let mut initial_builder = BlobArrayBuilder::new(1); + initial_builder.push_bytes(b"initial").unwrap(); + let initial_batch = RecordBatch::try_new( + logical_schema.clone(), + vec![ + Arc::new(UInt32Array::from(vec![0])) as ArrayRef, + initial_builder.finish().unwrap(), + ], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(initial_batch)], logical_schema.clone()), + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + let mut blob_writer = BlobDescriptorArrayBuilder::new("blob"); + blob_writer + .push_inline(Bytes::from_static(b"append")) + .unwrap(); + let prepared_column = blob_writer.finish().unwrap(); + let (prepared_field, prepared_array) = prepared_column.into_parts(); + let append_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + prepared_field, + ])); + let append_batch = RecordBatch::try_new( + append_schema.clone(), + vec![ + Arc::new(UInt32Array::from(vec![1])) as ArrayRef, + prepared_array, + ], + ) + .unwrap(); + + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(append_batch)], append_schema), + dataset, + Some(WriteParams { + mode: WriteMode::Append, + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + let ids = dataset + .scan() + .project(&["id"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let id_values = ids.column(0).as_primitive::(); + assert_eq!(id_values.values(), &[0, 1]); + + let blobs = dataset + .take_blobs_by_indices(&[0, 1], "blob") + .await + .unwrap(); + assert_eq!(blobs.len(), 2); + assert_eq!(blobs[0].read().await.unwrap().as_ref(), b"initial"); + assert_eq!(blobs[1].read().await.unwrap().as_ref(), b"append"); + } + + #[tokio::test] + async fn test_data_replacement_uses_blob_descriptor_array_builder_prepared_nested_blob_column() + { + let test_dir = TempStrDir::default(); + let mut initial_builder = BlobArrayBuilder::new(1); + initial_builder.push_bytes(b"initial").unwrap(); + let (logical_schema, initial_batch) = + nested_blob_v2_batch(initial_builder.finish().unwrap()); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(initial_batch)], logical_schema.clone()), + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + let file_id = Uuid::new_v4().to_string(); + let data_file_name = format!("{file_id}.lance"); + let data_file_path = dataset.data_dir().join(data_file_name.as_str()); + let blob_id = 1; + let mut blob_writer = BlobDescriptorArrayBuilder::new("blob"); + let mut packed = PackedBlobWriter::try_new( + dataset.object_store.as_ref().clone(), + data_file_path.clone(), + blob_id, + ) + .await + .unwrap(); + packed.write_blob(b"nested-replacement").await.unwrap(); + blob_writer.extend(packed.finish().await.unwrap()).unwrap(); + let prepared_column = blob_writer.finish().unwrap(); + let (prepared_field, prepared_array) = prepared_column.into_parts(); + + let info_fields = vec![Field::new("name", DataType::Utf8, false), prepared_field]; + let info_array = Arc::new( + StructArray::try_new( + info_fields.clone().into(), + vec![ + Arc::new(StringArray::from(vec!["replacement"])) as ArrayRef, + prepared_array, + ], + None, + ) + .unwrap(), + ) as ArrayRef; + let replacement_schema = Arc::new(Schema::new(vec![Field::new( + "info", + DataType::Struct(info_fields.into()), + true, + )])); + let replacement_batch = + RecordBatch::try_new(replacement_schema.clone(), vec![info_array]).unwrap(); + + let object_writer = dataset.object_store.create(&data_file_path).await.unwrap(); + let mut file_writer = FileWriter::try_new( + object_writer, + crate::datatypes::Schema::try_from(replacement_schema.as_ref()).unwrap(), + FileWriterOptions { + format_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }, + ) + .unwrap(); + file_writer.write_batch(&replacement_batch).await.unwrap(); + file_writer.finish().await.unwrap(); + + let data_file = dataset + .create_data_file(&data_file_name, None) + .await + .unwrap(); + let transaction = Transaction { + read_version: dataset.manifest.version, + uuid: Uuid::new_v4().hyphenated().to_string(), + operation: Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file)], + }, + tag: None, + transaction_properties: None, + }; + let dataset = Arc::new( + CommitBuilder::new(dataset) + .execute(transaction) + .await + .unwrap(), + ); + + let blobs = dataset + .take_blobs_by_indices(&[0], "info.blob") + .await + .unwrap(); + assert_eq!(blobs.len(), 1); + assert_eq!( + blobs[0].read().await.unwrap().as_ref(), + b"nested-replacement" + ); + assert_eq!(blobs[0].kind(), BlobKind::Packed); + } + + #[tokio::test] + async fn test_create_data_file_skips_blob_orphans_with_top_level_name_collision() { + let test_dir = TempStrDir::default(); + let mut initial_builder = BlobArrayBuilder::new(1); + initial_builder.push_bytes(b"initial").unwrap(); + let (_, initial_info_batch) = nested_blob_v2_batch(initial_builder.finish().unwrap()); + let info_field = initial_info_batch.schema().field(0).as_ref().clone(); + let logical_schema = Arc::new(Schema::new(vec![ + info_field, + Field::new("kind", DataType::UInt8, false), + ])); + let initial_batch = RecordBatch::try_new( + logical_schema.clone(), + vec![ + initial_info_batch.column(0).clone(), + Arc::new(UInt8Array::from(vec![7])) as ArrayRef, + ], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(initial_batch)], logical_schema.clone()), + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + let file_id = Uuid::new_v4().to_string(); + let data_file_name = format!("{file_id}.lance"); + let data_file_path = dataset.data_dir().join(data_file_name.as_str()); + let mut blob_writer = BlobDescriptorArrayBuilder::new("blob"); + blob_writer + .push_inline(Bytes::from_static(b"replacement")) + .unwrap(); + let prepared_column = blob_writer.finish().unwrap(); + let (prepared_field, prepared_array) = prepared_column.into_parts(); + + let info_fields = vec![Field::new("name", DataType::Utf8, false), prepared_field]; + let info_array = Arc::new( + StructArray::try_new( + info_fields.clone().into(), + vec![ + Arc::new(StringArray::from(vec!["replacement"])) as ArrayRef, + prepared_array, + ], + None, + ) + .unwrap(), + ) as ArrayRef; + let replacement_schema = Arc::new(Schema::new(vec![Field::new( + "info", + DataType::Struct(info_fields.into()), + true, + )])); + let replacement_batch = + RecordBatch::try_new(replacement_schema.clone(), vec![info_array]).unwrap(); + + let object_writer = dataset.object_store.create(&data_file_path).await.unwrap(); + let mut file_writer = FileWriter::try_new( + object_writer, + crate::datatypes::Schema::try_from(replacement_schema.as_ref()).unwrap(), + FileWriterOptions { + format_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }, + ) + .unwrap(); + file_writer.write_batch(&replacement_batch).await.unwrap(); + file_writer.finish().await.unwrap(); + + let data_file = dataset + .create_data_file(&data_file_name, None) + .await + .unwrap(); + + assert_eq!(data_file.fields.len(), 2); + assert_eq!(data_file.column_indices.as_ref(), &[0, 1]); + } + + #[tokio::test] + async fn test_write_and_take_nested_blob_v2() { + let test_dir = TempStrDir::default(); + let packed_payload = vec![0x4A; super::INLINE_MAX + 1024]; + + let mut blob_builder = BlobArrayBuilder::new(3); + blob_builder.push_bytes(b"hello").unwrap(); + blob_builder.push_bytes(&packed_payload).unwrap(); + blob_builder.push_null().unwrap(); + let blob_array: ArrayRef = blob_builder.finish().unwrap(); + + let (schema, batch) = nested_blob_v2_batch(blob_array); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + + let dataset = Arc::new( + Dataset::write( + reader, + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + let info_batch = dataset + .scan() + .project(&["info"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let blob_desc = info_batch + .column(0) + .as_struct() + .column_by_name("blob") + .unwrap() + .as_struct(); + assert_eq!( + blob_desc + .column_by_name("kind") + .unwrap() + .as_primitive::() + .value(0), + BlobKind::Inline as u8 + ); + assert_eq!( + blob_desc + .column_by_name("kind") + .unwrap() + .as_primitive::() + .value(1), + BlobKind::Packed as u8 + ); + + let blobs = dataset + .take_blobs_by_indices(&[0, 1], "info.blob") + .await + .unwrap(); + assert_eq!(blobs.len(), 2); + assert_eq!(blobs[0].read().await.unwrap().as_ref(), b"hello"); + assert_eq!( + blobs[1].read().await.unwrap().as_ref(), + packed_payload.as_slice() + ); + + let null_blobs = dataset + .take_blobs_by_indices(&[2], "info.blob") + .await + .unwrap(); + assert!(null_blobs.is_empty()); + + let filtered = dataset + .scan() + .project(&["info"]) + .unwrap() + .filter("info.blob IS NOT NULL") + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(filtered.num_rows(), 2); + } + + #[tokio::test] + async fn test_write_and_scan_list_blob_v2_descriptions() { + let test_dir = TempStrDir::default(); + let packed_payload = vec![0x4B; super::INLINE_MAX + 1024]; + + let mut blob_builder = BlobArrayBuilder::new(4); + blob_builder.push_bytes(b"hello").unwrap(); + blob_builder.push_null().unwrap(); + blob_builder.push_bytes(&packed_payload).unwrap(); + blob_builder.push_bytes(b"tail").unwrap(); + let blob_values = blob_builder.finish().unwrap(); + + let item_field = Arc::new(blob_field("item", true)); + let list_array: ArrayRef = Arc::new( + arrow_array::ListArray::try_new( + item_field.clone(), + arrow_buffer::OffsetBuffer::new(arrow_buffer::ScalarBuffer::from(vec![ + 0i32, 3, 3, 3, 4, + ])), + blob_values, + Some(arrow_buffer::NullBuffer::from(vec![ + true, true, false, true, + ])), + ) + .unwrap(), + ); + + let schema = Arc::new(Schema::new(vec![ + Field::new("blobs", DataType::List(item_field), true), + Field::new("id", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![list_array, Arc::new(Int32Array::from(vec![0, 1, 2, 3]))], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + + let dataset = Arc::new( + Dataset::write( + reader, + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + let descriptions = dataset + .scan() + .project(&["blobs"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let lists = descriptions.column(0).as_list::(); + assert_eq!(lists.offsets().inner().as_ref(), &[0, 3, 3, 3, 4]); + assert!(lists.is_valid(0)); + assert!(lists.is_valid(1)); + assert!(lists.is_null(2)); + assert!(lists.is_valid(3)); + + let DataType::List(descriptor_field) = lists.data_type() else { + panic!("unexpected list type: {}", lists.data_type()); + }; + assert!(matches!(descriptor_field.data_type(), DataType::Struct(_))); + assert!(!descriptor_field.metadata().contains_key(ARROW_EXT_NAME_KEY)); + let descriptors = lists.values().as_struct(); + assert_eq!(descriptors.fields().len(), 5); + assert_eq!(descriptors.fields()[0].name(), "kind"); + assert!(descriptors.is_valid(0)); + assert!(descriptors.is_null(1)); + assert!(descriptors.is_valid(2)); + assert!(descriptors.is_valid(3)); + let kinds = descriptors + .column_by_name("kind") + .unwrap() + .as_primitive::(); + assert_eq!(kinds.value(0), BlobKind::Inline as u8); + assert_eq!(kinds.value(2), BlobKind::Packed as u8); + assert_eq!(kinds.value(3), BlobKind::Inline as u8); + + let filtered = dataset + .scan() + .project(&["blobs"]) + .unwrap() + .filter("blobs IS NOT NULL") + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(filtered.num_rows(), 3); - // Pass 3: back to descriptor view to verify both directions are safe. let mut scanner = dataset.scan(); - scanner.blob_handling(BlobHandling::BlobsDescriptions); - let descriptors = scanner + scanner.blob_handling(BlobHandling::AllBinary); + let bytes = scanner .project(&["blobs"]) .unwrap() .try_into_batch() .await .unwrap(); - assert!(descriptors.column(0).data_type().is_struct()); - } - - #[test] - fn test_data_file_key_from_path() { - assert_eq!(data_file_key_from_path("data/abc.lance"), "abc"); - assert_eq!(data_file_key_from_path("abc.lance"), "abc"); - assert_eq!(data_file_key_from_path("nested/path/xyz"), "xyz"); - } - - #[tokio::test] - async fn test_write_and_take_blobs_with_blob_array_builder() { - let test_dir = TempStrDir::default(); - - // Build a blob column with the new BlobArrayBuilder - let mut blob_builder = BlobArrayBuilder::new(2); - blob_builder.push_bytes(b"hello").unwrap(); - blob_builder.push_bytes(b"world").unwrap(); - let blob_array: arrow_array::ArrayRef = blob_builder.finish().unwrap(); - - let id_array: arrow_array::ArrayRef = Arc::new(UInt32Array::from(vec![0, 1])); - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::UInt32, false), - blob_field("blob", true), - ])); - - let batch = RecordBatch::try_new(schema.clone(), vec![id_array, blob_array]).unwrap(); - let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); - - let params = WriteParams { - data_storage_version: Some(LanceFileVersion::V2_2), - ..Default::default() + let lists = bytes.column(0).as_list::(); + assert_eq!(lists.offsets().inner().as_ref(), &[0, 3, 3, 3, 4]); + assert!(lists.is_valid(0)); + assert!(lists.is_valid(1)); + assert!(lists.is_null(2)); + assert!(lists.is_valid(3)); + let DataType::List(value_field) = lists.data_type() else { + panic!("unexpected list type: {}", lists.data_type()); }; - let dataset = Arc::new( - Dataset::write(reader, &test_dir, Some(params)) - .await - .unwrap(), - ); - - let blobs = dataset - .take_blobs_by_indices(&[0, 1], "blob") - .await - .unwrap(); + assert_eq!(value_field.data_type(), &DataType::LargeBinary); + assert!(!value_field.metadata().contains_key(ARROW_EXT_NAME_KEY)); + let values = lists.values().as_binary::(); + assert_eq!(values.value(0), b"hello"); + assert!(values.is_null(1)); + assert_eq!(values.value(2), packed_payload.as_slice()); + assert_eq!(values.value(3), b"tail"); + + for (filter, materialization_style) in [ + (None, MaterializationStyle::Heuristic), + (Some("id >= 2"), MaterializationStyle::Heuristic), + (Some("id >= 2"), MaterializationStyle::AllEarly), + ] { + let mut scanner = dataset.scan(); + scanner.blob_handling(BlobHandling::AllBinary); + scanner.materialization_style(materialization_style); + scanner + .project(&["blobs", ROW_LAST_UPDATED_AT_VERSION, ROW_CREATED_AT_VERSION]) + .unwrap() + .with_row_id() + .with_row_address(); + if let Some(filter) = filter { + scanner.filter(filter).unwrap(); + } - assert_eq!(blobs.len(), 2); - let first = blobs[0].read().await.unwrap(); - let second = blobs[1].read().await.unwrap(); - assert_eq!(first.as_ref(), b"hello"); - assert_eq!(second.as_ref(), b"world"); + let expected_schema = scanner.schema().await.unwrap(); + let batch = scanner.try_into_batch().await.unwrap(); + assert_eq!(batch.schema().as_ref(), expected_schema.as_ref()); + assert_eq!(batch.num_rows(), if filter.is_some() { 2 } else { 4 }); + for column in [ + ROW_ID, + ROW_ADDR, + ROW_LAST_UPDATED_AT_VERSION, + ROW_CREATED_AT_VERSION, + ] { + assert!( + batch.column_by_name(column).is_some(), + "requested system column {column} was missing" + ); + } + } } #[tokio::test] - async fn test_write_and_take_nested_blob_v2() { + async fn test_write_and_scan_struct_nested_list_blob_v2() { let test_dir = TempStrDir::default(); - let packed_payload = vec![0x4A; super::INLINE_MAX + 1024]; - let mut blob_builder = BlobArrayBuilder::new(3); - blob_builder.push_bytes(b"hello").unwrap(); - blob_builder.push_bytes(&packed_payload).unwrap(); + let mut blob_builder = BlobArrayBuilder::new(2); + blob_builder.push_bytes(b"nested").unwrap(); blob_builder.push_null().unwrap(); - let blob_array: ArrayRef = blob_builder.finish().unwrap(); + let blob_values = blob_builder.finish().unwrap(); + + let item_field = Arc::new(blob_field("item", true)); + let list_field = Field::new("blobs", DataType::List(item_field.clone()), true); + let list_array: ArrayRef = Arc::new( + arrow_array::ListArray::try_new( + item_field, + arrow_buffer::OffsetBuffer::new(arrow_buffer::ScalarBuffer::from(vec![0i32, 2, 2])), + blob_values, + None, + ) + .unwrap(), + ); + let info_fields = vec![Field::new("name", DataType::Utf8, false), list_field]; + let info_array: ArrayRef = Arc::new( + StructArray::try_new( + info_fields.clone().into(), + vec![ + Arc::new(StringArray::from(vec!["row-0", "row-1"])) as ArrayRef, + list_array, + ], + None, + ) + .unwrap(), + ); - let (schema, batch) = nested_blob_v2_batch(blob_array); + let schema = Arc::new(Schema::new(vec![Field::new( + "info", + DataType::Struct(info_fields.into()), + true, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![info_array]).unwrap(); let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); let dataset = Arc::new( @@ -3289,52 +4672,52 @@ mod tests { .unwrap(), ); - let info_batch = dataset + let descriptions = dataset .scan() .project(&["info"]) .unwrap() .try_into_batch() .await .unwrap(); - let blob_desc = info_batch - .column(0) - .as_struct() - .column_by_name("blob") - .unwrap() - .as_struct(); + let info = descriptions.column(0).as_struct(); assert_eq!( - blob_desc - .column_by_name("kind") + info.column_by_name("name") .unwrap() - .as_primitive::() + .as_string::() .value(0), - BlobKind::Inline as u8 - ); - assert_eq!( - blob_desc - .column_by_name("kind") - .unwrap() - .as_primitive::() - .value(1), - BlobKind::Packed as u8 - ); - - let blobs = dataset - .take_blobs_by_indices(&[0, 1], "info.blob") - .await - .unwrap(); - assert_eq!(blobs.len(), 2); - assert_eq!(blobs[0].read().await.unwrap().as_ref(), b"hello"); - assert_eq!( - blobs[1].read().await.unwrap().as_ref(), - packed_payload.as_slice() + "row-0" ); + let lists = info.column_by_name("blobs").unwrap().as_list::(); + assert_eq!(lists.offsets().inner().as_ref(), &[0, 2, 2]); + let DataType::List(descriptor_field) = lists.data_type() else { + panic!("unexpected nested list type: {}", lists.data_type()); + }; + assert!(matches!(descriptor_field.data_type(), DataType::Struct(_))); + assert!(!descriptor_field.metadata().contains_key(ARROW_EXT_NAME_KEY)); + let descriptors = lists.values().as_struct(); + assert_eq!(descriptors.fields().len(), 5); + assert!(descriptors.is_valid(0)); + assert!(descriptors.is_null(1)); - let null_blobs = dataset - .take_blobs_by_indices(&[2], "info.blob") + let mut scanner = dataset.scan(); + scanner.blob_handling(BlobHandling::AllBinary); + let bytes = scanner + .project(&["info"]) + .unwrap() + .try_into_batch() .await .unwrap(); - assert!(null_blobs.is_empty()); + let info = bytes.column(0).as_struct(); + let lists = info.column_by_name("blobs").unwrap().as_list::(); + assert_eq!(lists.offsets().inner().as_ref(), &[0, 2, 2]); + let DataType::List(value_field) = lists.data_type() else { + panic!("unexpected nested list type: {}", lists.data_type()); + }; + assert_eq!(value_field.data_type(), &DataType::LargeBinary); + assert!(!value_field.metadata().contains_key(ARROW_EXT_NAME_KEY)); + let values = lists.values().as_binary::(); + assert_eq!(values.value(0), b"nested"); + assert!(values.is_null(1)); } #[tokio::test] @@ -3445,7 +4828,7 @@ mod tests { } #[tokio::test] - async fn test_read_blobs_plan_preserves_order_and_coalesces() { + async fn test_execute_blob_entries_preserves_order_and_coalesces() { let (store, inner) = recording_range_store(Bytes::from_static(b"abcdefghij")); let source = Arc::new(BlobSource::new(store, Path::from("blobs/test.bin"))); let entries = vec![ @@ -3460,15 +4843,7 @@ mod tests { file: BlobFile::with_source(source, 1, 3, BlobKind::Packed, None), }, ]; - let execution = Arc::new(ReadBlobsExecution::new(None)); - let blobs = try_join_all( - plan_blob_read_plans(entries) - .into_iter() - .map(|plan| execute_blob_read_plan(plan, execution.clone())), - ) - .await - .unwrap(); - let mut blobs = blobs.into_iter().flatten().collect::>(); + let mut blobs = execute_blob_entries(entries, 2, None).await.unwrap(); blobs.sort_by_key(|blob| blob.selection_index); assert_eq!(blobs.len(), 2); @@ -4082,10 +5457,12 @@ mod tests { ); } - async fn try_preprocess_kind_with_blob_metadata( + async fn try_preprocess_blobs_with_blob_metadata( metadata_entries: Vec<(&'static str, String)>, - data_len: usize, - ) -> Result { + pack_file_size_override: Option, + blob_len: usize, + num_blobs: usize, + ) -> Result { let (object_store, base_path) = ObjectStore::from_uri_and_params( Arc::new(ObjectStoreRegistry::default()), "memory://blob_preprocessor", @@ -4116,11 +5493,13 @@ mod tests { ExternalBlobMode::Reference, Arc::new(ObjectStoreRegistry::default()), ObjectStoreParams::default(), - None, + pack_file_size_override, )?; - let mut blob_builder = BlobArrayBuilder::new(1); - blob_builder.push_bytes(vec![0u8; data_len]).unwrap(); + let mut blob_builder = BlobArrayBuilder::new(num_blobs); + for _ in 0..num_blobs { + blob_builder.push_bytes(vec![0u8; blob_len]).unwrap(); + } let blob_array: arrow_array::ArrayRef = blob_builder.finish().unwrap(); let field_without_metadata = @@ -4129,11 +5508,20 @@ mod tests { let batch = RecordBatch::try_new(batch_schema, vec![blob_array]).unwrap(); let out = preprocessor.preprocess_batch(&batch).await?; - let struct_arr = out + Ok(out .column(0) .as_any() .downcast_ref::() - .unwrap(); + .unwrap() + .clone()) + } + + async fn try_preprocess_kind_with_blob_metadata( + metadata_entries: Vec<(&'static str, String)>, + data_len: usize, + ) -> Result { + let struct_arr = + try_preprocess_blobs_with_blob_metadata(metadata_entries, None, data_len, 1).await?; Ok(struct_arr .column_by_name("kind") .unwrap() @@ -4150,6 +5538,27 @@ mod tests { .unwrap() } + async fn packed_blobs_with_blob_metadata( + metadata_entries: Vec<(&'static str, String)>, + pack_file_size_override: Option, + blob_len: usize, + num_blobs: usize, + ) -> Vec { + let struct_arr = try_preprocess_blobs_with_blob_metadata( + metadata_entries, + pack_file_size_override, + blob_len, + num_blobs, + ) + .await + .unwrap(); + let blob_ids = struct_arr + .column_by_name("blob_id") + .unwrap() + .as_primitive::(); + (0..struct_arr.len()).map(|i| blob_ids.value(i)).collect() + } + #[tokio::test] async fn test_blob_v2_dedicated_threshold_rejects_non_positive_metadata() { let err = try_preprocess_kind_with_blob_metadata( @@ -4240,6 +5649,61 @@ mod tests { assert_eq!(kind, lance_core::datatypes::BlobKind::Packed as u8); } + #[tokio::test] + async fn test_blob_v2_pack_file_threshold_rolls_at_metadata_value() { + // Blobs must exceed the inline cutoff to be packed at all. With a pack-file + // threshold equal to a single blob's size, each of the three blobs rolls to its + // own pack file (a distinct blob_id). + let blob_len = super::INLINE_MAX + 1024; + let blob_ids = packed_blobs_with_blob_metadata( + vec![(BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY, blob_len.to_string())], + None, + blob_len, + 3, + ) + .await; + let distinct: std::collections::HashSet = blob_ids.iter().copied().collect(); + assert_eq!( + distinct.len(), + 3, + "expected one pack file per blob: {blob_ids:?}" + ); + } + + #[tokio::test] + async fn test_blob_v2_pack_file_threshold_packs_within_metadata_value() { + // A pack-file threshold large enough for all three blobs keeps them in a single + // pack file (one shared blob_id). + let blob_len = super::INLINE_MAX + 1024; + let blob_ids = packed_blobs_with_blob_metadata( + vec![( + BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY, + (blob_len * 3).to_string(), + )], + None, + blob_len, + 3, + ) + .await; + let distinct: std::collections::HashSet = blob_ids.iter().copied().collect(); + assert_eq!( + distinct.len(), + 1, + "expected a single shared pack file: {blob_ids:?}" + ); + } + + #[tokio::test] + async fn test_blob_v2_pack_file_threshold_rejects_non_positive_metadata() { + let err = try_preprocess_kind_with_blob_metadata( + vec![(BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY, "0".to_string())], + 1024, + ) + .await + .expect_err("zero pack-file threshold should be rejected"); + assert!(err.to_string().contains("expected a positive integer")); + } + #[tokio::test] async fn test_blob_v2_inline_threshold_respects_smaller_metadata() { let kind = preprocess_kind_with_blob_metadata( @@ -4390,6 +5854,91 @@ mod tests { assert_eq!(packed_kind, lance_core::datatypes::BlobKind::Packed as u8); } + #[tokio::test] + async fn test_blob_v2_pack_file_threshold_starts_new_packed_blob() { + let (object_store, base_path) = ObjectStore::from_uri_and_params( + Arc::new(ObjectStoreRegistry::default()), + "memory://blob_preprocessor", + &ObjectStoreParams::default(), + ) + .await + .unwrap(); + let object_store = object_store.as_ref().clone(); + let data_dir = base_path.clone().join("data"); + + let mut field = blob_field("blob", true); + let mut metadata = field.metadata().clone(); + metadata.insert( + BLOB_INLINE_SIZE_THRESHOLD_META_KEY.to_string(), + "0".to_string(), + ); + metadata.insert( + BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY.to_string(), + "1024".to_string(), + ); + field = field.with_metadata(metadata); + + let writer_arrow_schema = Schema::new(vec![field.clone()]); + let writer_schema = lance_core::datatypes::Schema::try_from(&writer_arrow_schema).unwrap(); + + let mut preprocessor = super::BlobPreprocessor::new( + object_store.clone(), + data_dir, + "data_file_key".to_string(), + &writer_schema, + None, + false, + ExternalBlobMode::Reference, + Arc::new(ObjectStoreRegistry::default()), + ObjectStoreParams::default(), + Some(3), + ) + .unwrap(); + + let mut blob_builder = BlobArrayBuilder::new(3); + blob_builder.push_bytes(vec![0x11; 2]).unwrap(); + blob_builder.push_bytes(vec![0x22; 2]).unwrap(); + blob_builder.push_bytes(vec![0x33; 1]).unwrap(); + let blob_array: arrow_array::ArrayRef = blob_builder.finish().unwrap(); + + let batch_schema = Arc::new(Schema::new(vec![Field::new( + "blob", + field.data_type().clone(), + field.is_nullable(), + )])); + let batch = RecordBatch::try_new(batch_schema, vec![blob_array]).unwrap(); + + let out = preprocessor.preprocess_batch(&batch).await.unwrap(); + preprocessor.finish().await.unwrap(); + + let struct_arr = out + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let kinds = struct_arr + .column_by_name("kind") + .unwrap() + .as_primitive::(); + let blob_ids = struct_arr + .column_by_name("blob_id") + .unwrap() + .as_primitive::(); + let sizes = struct_arr + .column_by_name("blob_size") + .unwrap() + .as_primitive::(); + let positions = struct_arr + .column_by_name("position") + .unwrap() + .as_primitive::(); + + assert_eq!(kinds.values(), &[BlobKind::Packed as u8; 3]); + assert_eq!(blob_ids.values(), &[1, 2, 2]); + assert_eq!(sizes.values(), &[2, 2, 1]); + assert_eq!(positions.values(), &[0, 0, 2]); + } + #[tokio::test] async fn test_blob_v2_append_rejects_explicit_inline_threshold_mismatch() { let dataset_dir = TempDir::default(); @@ -4566,4 +6115,25 @@ mod tests { .unwrap(); assert_eq!(dataset.count_rows(None).await.unwrap(), 2); } + + #[tokio::test] + async fn test_blob_v2_pack_file_threshold_write_param_overrides_metadata() { + let blob_len = super::INLINE_MAX + 1024; + let blob_ids = packed_blobs_with_blob_metadata( + vec![( + BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY, + (blob_len * 3).to_string(), + )], + Some(blob_len), + blob_len, + 3, + ) + .await; + let distinct: std::collections::HashSet = blob_ids.iter().copied().collect(); + assert_eq!( + distinct.len(), + 3, + "write-param override should force one pack file per blob: {blob_ids:?}" + ); + } } diff --git a/rust/lance/src/dataset/builder.rs b/rust/lance/src/dataset/builder.rs index 23db7254b15..aecbf92f0d8 100644 --- a/rust/lance/src/dataset/builder.rs +++ b/rust/lance/src/dataset/builder.rs @@ -305,6 +305,14 @@ impl DatasetBuilder { /// - [Azure options](https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html#variants) /// - [S3 options](https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html#variants) /// - [Google options](https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html#variants) + /// + /// For datasets with additional registered base paths, a key of the form + /// `base_.` applies `` only to the base path with that + /// manifest id, overriding the unscoped options that every base inherits. + /// For example `base_1.account_key = abc` makes base 1 use + /// `account_key = abc` while all other options are shared. Exact per-base + /// bindings set via [`Self::with_base_store_params`] take precedence over + /// base-scoped keys. pub fn with_storage_options(mut self, storage_options: HashMap) -> Self { // Merge with existing options if accessor exists, otherwise create new static accessor if let Some(existing) = self.options.storage_options_accessor.take() { @@ -459,8 +467,9 @@ impl DatasetBuilder { /// /// These params are not persisted in the manifest. They are used as-is /// whenever the dataset resolves an object store for the given - /// `BasePath.path`. Dataset-level store params remain the fallback for bases - /// without an explicit binding. + /// `BasePath.path`, taking precedence over `base_.` storage + /// options (see [`Self::with_storage_options`]). Dataset-level store params + /// remain the fallback for bases without an explicit binding. pub fn with_base_store_params( mut self, base_path: impl AsRef, @@ -833,7 +842,7 @@ impl DatasetBuilder { let reader = object_store.open(&location.path).await?; populate_schema_dictionary(&mut manifest.schema, reader.as_ref()).await?; } - (manifest, location) + (Arc::new(manifest), location) } else { let manifest_location = match version_number { Some(version) => { @@ -866,7 +875,8 @@ impl DatasetBuilder { _ => e, })?, }; - let manifest = Dataset::load_manifest( + + let manifest = Dataset::get_manifest( &object_store, &manifest_location, &table_uri, @@ -880,7 +890,7 @@ impl DatasetBuilder { object_store, base_path, table_uri, - Arc::new(manifest), + manifest, location, session, commit_handler, diff --git a/rust/lance/src/dataset/cleanup.rs b/rust/lance/src/dataset/cleanup.rs index 65928038cea..d8f69e7d36a 100644 --- a/rust/lance/src/dataset/cleanup.rs +++ b/rust/lance/src/dataset/cleanup.rs @@ -58,8 +58,8 @@ use lance_table::{ manifest::{read_manifest, read_manifest_indexes}, }, }; +use object_store::ObjectMeta; use object_store::path::Path; -use object_store::{Error as ObjectStoreError, ObjectMeta}; use std::fmt::Debug; use std::{ collections::{HashMap, HashSet}, @@ -625,16 +625,7 @@ impl<'a> CleanupTask<'a> { let verification_threshold = utc_now() - TimeDelta::try_days(UNVERIFIED_THRESHOLD_DAYS).expect("TimeDelta::try_days"); - let is_not_found_err = |e: &Error| { - matches!( - e, - Error::IO { source,.. } - if source - .downcast_ref::() - .map(|os_err| matches!(os_err, ObjectStoreError::NotFound {.. })) - .unwrap_or(false) - ) - }; + let is_not_found_err = |e: &Error| matches!(e, Error::NotFound { .. }); // Build stream for a managed subtree let build_listing_stream = |dir: Path| { let inspection_ref = &inspection; @@ -2987,10 +2978,8 @@ mod tests { async fn cleanup_before_ts_and_retain_n_recent_versions() { let fixture = MockDatasetFixture::try_new().unwrap(); fixture.create_some_data().await.unwrap(); - let mut time = 1i64; - for _ in 0..4 { + for time in (1i64..).take(4) { MockClock::set_system_time(TimeDelta::try_days(time).unwrap().to_std().unwrap()); - time += 1i64; fixture.overwrite_some_data().await.unwrap(); } diff --git a/rust/lance/src/dataset/files.rs b/rust/lance/src/dataset/files.rs index 848add7e4a8..2214822ed90 100644 --- a/rust/lance/src/dataset/files.rs +++ b/rust/lance/src/dataset/files.rs @@ -1036,6 +1036,7 @@ mod tests { // No base_id -> falls back to the dataset base_uri. mk_file("c.lance", None), ], + overlays: vec![], // Deletion files also carry a base_id when they originate from a // shallow clone, and must resolve against base_paths too. deletion_file: Some(DeletionFile { diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 4ea1b51f073..4df89ab71b8 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -22,6 +22,7 @@ use datafusion::logical_expr::Expr; use datafusion::scalar::ScalarValue; use futures::future::{BoxFuture, try_join_all}; use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt, join, stream}; +use lance_arrow::json::{convert_json_columns, has_json_fields, is_arrow_json_field}; use lance_arrow::{RecordBatchExt, SchemaExt}; use lance_core::datatypes::{OnMissing, OnTypeMismatch, SchemaCompareOptions}; use lance_core::utils::address::RowAddress; @@ -910,6 +911,16 @@ impl FileFragment { projection: &Schema, read_config: FragReadConfig, ) -> Result { + // Overlay files supply newer cell values that must be merged on read. + // Until the scan/take merge path lands (the rest of OSS-1322 / OSS-1324), + // reading a fragment that has overlays would silently return stale base + // values, so we refuse rather than serve incorrect data. + if !self.metadata.overlays.is_empty() { + return Err(Error::not_supported( + "reading fragments with data overlay files is not yet supported \ + (overlay merge is in progress)", + )); + } let open_files = self.open_readers(projection, &read_config); let deletion_vec_load = self.get_deletion_vector(); @@ -1572,8 +1583,22 @@ impl FileFragment { }; // Then call take rows - self.take_rows(&row_ids, projection, false, false, false, false) - .await + let batch = self + .take_rows(&row_ids, projection, false, false, false, false) + .await?; + + // Convert Lance JSON columns (LargeBinary/JSONB) back to Arrow JSON (Utf8) + // for user-facing output. + if batch + .schema() + .fields() + .iter() + .any(|f| lance_arrow::json::is_json_field(f) || lance_arrow::json::has_json_fields(f)) + { + Ok(lance_arrow::json::convert_lance_json_to_arrow(&batch)?) + } else { + Ok(batch) + } } /// Get the deletion vector for this fragment, using the cache if available. @@ -1892,6 +1917,17 @@ impl FileFragment { ) .await?; // Hash join: rows matched on the right-hand stream rewrite columns; track physical offsets via `_rowaddr`. + // Convert Arrow JSON columns (Utf8) to Lance JSON (LargeBinary) in the right stream + // so they match the physical storage format read from the fragment's left batch. + let right_stream: Box = if right_schema + .fields() + .iter() + .any(|f| is_arrow_json_field(f) || has_json_fields(f)) + { + Box::new(JsonConvertingReader::new(right_stream)) + } else { + right_stream + }; let joiner = Arc::new(HashJoiner::try_new(right_stream, right_on).await?); let mut matched_offsets = RoaringBitmap::new(); let frag_id_u32 = u32::try_from(self.metadata.id).map_err(|_| { @@ -2932,6 +2968,59 @@ impl FragmentReader { } } +/// A wrapper around a `RecordBatchReader` that converts Arrow JSON columns +/// (Utf8/LargeUtf8 with `arrow.json` extension) to Lance JSON columns +/// (LargeBinary with `lance.json` extension / JSONB format). +/// +/// This is needed when user-provided data contains Arrow JSON fields but the +/// dataset stores them in Lance's JSONB binary format. +struct JsonConvertingReader { + inner: Box, + schema: arrow_schema::SchemaRef, +} + +impl JsonConvertingReader { + fn new(inner: Box) -> Self { + use lance_arrow::json::arrow_json_to_lance_json; + + // Build the converted schema (Arrow JSON fields → Lance JSON fields) + let orig_schema = inner.schema(); + let new_fields: Vec = orig_schema + .fields() + .iter() + .map(|f| { + if is_arrow_json_field(f) || has_json_fields(f) { + Arc::new(arrow_json_to_lance_json(f)) + } else { + Arc::clone(f) + } + }) + .collect(); + let schema = Arc::new(arrow_schema::Schema::new_with_metadata( + new_fields, + orig_schema.metadata().clone(), + )); + + Self { inner, schema } + } +} + +impl Iterator for JsonConvertingReader { + type Item = std::result::Result; + + fn next(&mut self) -> Option { + self.inner + .next() + .map(|result| result.and_then(|batch| convert_json_columns(&batch))) + } +} + +impl RecordBatchReader for JsonConvertingReader { + fn schema(&self) -> arrow_schema::SchemaRef { + self.schema.clone() + } +} + #[cfg(test)] mod tests { use arrow_arith::numeric::mul; @@ -4406,4 +4495,80 @@ mod tests { assert_io_eq!(stats, read_iops, 1); assert_io_lt!(stats, read_bytes, 4096); } + + #[tokio::test] + async fn test_update_columns_with_json_extension_type() { + use arrow_array::UInt64Array; + use lance_arrow::ARROW_EXT_NAME_KEY; + use lance_arrow::json::ARROW_JSON_EXT_NAME; + use lance_core::ROW_ID; + use std::collections::HashMap; + + // Create a dataset with an Arrow JSON extension column + let test_dir = TempStrDir::default(); + let mut json_metadata = HashMap::new(); + json_metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int64, false), + ArrowField::new("name", DataType::Utf8, true), + ArrowField::new("meta", DataType::Utf8, true).with_metadata(json_metadata.clone()), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5])), + Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"])), + Arc::new(StringArray::from(vec![ + r#"{"x":1}"#, + r#"{"x":2}"#, + r#"{"x":3}"#, + r#"{"x":4}"#, + r#"{"x":5}"#, + ])), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write(reader, test_dir.as_ref(), None) + .await + .unwrap(); + + // Build the right stream with Arrow JSON column (Utf8 + arrow.json extension) + // Only update rows with row_id 1 and 3 + let update_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new(ROW_ID, DataType::UInt64, false), + ArrowField::new("meta", DataType::Utf8, true).with_metadata(json_metadata), + ])); + let update_batch = RecordBatch::try_new( + update_schema.clone(), + vec![ + Arc::new(UInt64Array::from(vec![1, 3])), + Arc::new(StringArray::from(vec![ + r#"{"updated":true,"id":2}"#, + r#"{"updated":true,"id":4}"#, + ])), + ], + ) + .unwrap(); + let right_stream: Box = Box::new(RecordBatchIterator::new( + vec![Ok(update_batch)], + update_schema, + )); + + // Perform update_columns - this should NOT fail with type mismatch + // Previously this would error with: + // "It is not possible to interleave arrays of different data types (Utf8 and LargeBinary)" + let mut fragment = dataset.get_fragment(0).unwrap(); + let (updated_fragment, fields_modified) = fragment + .update_columns(right_stream, ROW_ID, ROW_ID) + .await + .unwrap(); + + // Verify the operation produced valid results + assert!(!fields_modified.is_empty()); + assert!(!updated_fragment.files.is_empty()); + } } diff --git a/rust/lance/src/dataset/fragment/session.rs b/rust/lance/src/dataset/fragment/session.rs index 64fa427580a..de50255bb72 100644 --- a/rust/lance/src/dataset/fragment/session.rs +++ b/rust/lance/src/dataset/fragment/session.rs @@ -54,7 +54,20 @@ impl FragmentSession { }; // Then call take rows - self.take_rows(&row_ids).await + let batch = self.take_rows(&row_ids).await?; + + // Convert Lance JSON columns (LargeBinary/JSONB) back to Arrow JSON (Utf8) + // for user-facing output, mirroring FileFragment::take. + if batch + .schema() + .fields() + .iter() + .any(|f| lance_arrow::json::is_json_field(f) || lance_arrow::json::has_json_fields(f)) + { + Ok(lance_arrow::json::convert_lance_json_to_arrow(&batch)?) + } else { + Ok(batch) + } } pub(crate) async fn take_rows(&self, row_offsets: &[u32]) -> Result { @@ -204,6 +217,78 @@ mod tests { ); } + #[tokio::test] + async fn test_fragment_session_take_json() { + use arrow_array::LargeBinaryArray; + use lance_arrow::ARROW_EXT_NAME_KEY; + use lance_arrow::json::{ARROW_JSON_EXT_NAME, JsonArray, json_field}; + + let test_dir = TempStrDir::default(); + let test_uri = &test_dir; + + // Build a schema with a Lance JSON column (LargeBinary + lance.json metadata). + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("i", DataType::Int32, true), + json_field("j", true), + ])); + + let json_strings = (0..10) + .map(|v| Some(format!("{{\"v\":{}}}", v))) + .collect::>(); + let jsonb = JsonArray::try_from_iter(json_strings).unwrap().into_inner(); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..10)), + Arc::new(jsonb), + ], + ) + .unwrap(); + + let write_params = WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }; + let batches = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write(batches, test_uri, Some(write_params)) + .await + .unwrap(); + + let fragment = dataset.get_fragments().into_iter().next().unwrap(); + let take_session = fragment + .open_session(dataset.schema(), false) + .await + .unwrap(); + let result = take_session.take(&[0, 3, 7]).await.unwrap(); + + // The JSON column must be returned as Arrow JSON (Utf8 + arrow.json), not raw JSONB. + let field = result.schema().field_with_name("j").unwrap().clone(); + assert_eq!(field.data_type(), &DataType::Utf8); + assert_eq!( + field.metadata().get(ARROW_EXT_NAME_KEY).map(|s| s.as_str()), + Some(ARROW_JSON_EXT_NAME) + ); + assert!( + result + .column_by_name("j") + .unwrap() + .as_any() + .downcast_ref::() + .is_none(), + "JSON column should not be raw LargeBinary JSONB" + ); + let json_col = result + .column_by_name("j") + .unwrap() + .as_any() + .downcast_ref::() + .expect("JSON column should be Utf8 strings"); + assert_eq!(json_col.value(0), "{\"v\":0}"); + assert_eq!(json_col.value(1), "{\"v\":3}"); + assert_eq!(json_col.value(2), "{\"v\":7}"); + } + async fn create_dataset(test_uri: &str, data_storage_version: LanceFileVersion) -> Dataset { let schema = Arc::new(ArrowSchema::new(vec![ ArrowField::new("i", DataType::Int32, true), diff --git a/rust/lance/src/dataset/fragment/write.rs b/rust/lance/src/dataset/fragment/write.rs index 9731be0c0eb..42b05051124 100644 --- a/rust/lance/src/dataset/fragment/write.rs +++ b/rust/lance/src/dataset/fragment/write.rs @@ -21,7 +21,8 @@ use uuid::Uuid; use crate::Result; use crate::dataset::builder::DatasetBuilder; -use crate::dataset::write::{do_write_fragments, validate_and_resolve_target_bases}; +use crate::dataset::utils::SchemaAdapter; +use crate::dataset::write::{do_write_fragments, validate_and_resolve_target_bases_with_primary}; use crate::dataset::{DATA_DIR, Dataset, ReadParams, WriteMode, WriteParams}; /// Generates a filename optimized for S3 throughput using a UUID-based approach. @@ -106,6 +107,12 @@ impl<'a> FragmentCreateBuilder<'a> { id: Option, ) -> Result { let (stream, schema) = self.get_stream_and_schema(Box::new(source)).await?; + // Convert Arrow JSON columns (`arrow.json`, stored as Utf8) into Lance JSON + // (`lance.json`, stored as JSONB-encoded LargeBinary) before writing. The + // multi-fragment and dataset write paths perform this through `do_write_fragments`; + // the single-fragment create path must do the same or the raw UTF-8 string bytes + // would be written into a column whose schema declares JSONB, corrupting reads. + let stream = SchemaAdapter::new(stream.schema()).to_physical_stream(stream); self.write_impl(stream, schema, id).await } @@ -206,6 +213,7 @@ impl<'a> FragmentCreateBuilder<'a> { let version = params.data_storage_version.unwrap_or_default(); let needs_existing_dataset = params.target_base_names_or_paths.is_some() || params.target_bases.is_some() + || params.target_all_bases.is_some() || params.initial_bases.is_some(); let existing_dataset = if needs_existing_dataset { self.existing_dataset(¶ms).await? @@ -215,17 +223,24 @@ impl<'a> FragmentCreateBuilder<'a> { let existing_base_paths = existing_dataset .as_ref() .map(|dataset| &dataset.manifest.base_paths); - let target_bases_info = if needs_existing_dataset { - validate_and_resolve_target_bases(&mut params, existing_base_paths).await? - } else { - None - }; let (object_store, base_path) = ObjectStore::from_uri_and_params( params.store_registry(), self.dataset_uri, ¶ms.store_params.clone().unwrap_or_default(), ) .await?; + let target_bases_info = if needs_existing_dataset { + validate_and_resolve_target_bases_with_primary( + &mut params, + existing_base_paths, + &object_store, + &base_path, + self.dataset_uri, + ) + .await? + } else { + None + }; do_write_fragments( existing_dataset.as_ref(), object_store, @@ -287,7 +302,7 @@ impl<'a> FragmentCreateBuilder<'a> { return Err(Error::invalid_input("Input data was empty.")); } - fragment.physical_rows = Some(writer.finish().await?); + fragment.physical_rows = Some(writer.finish().await?.num_rows as usize); progress.complete(&fragment).await?; @@ -407,7 +422,7 @@ mod tests { matches!(result.as_ref().unwrap_err(), Error::InvalidInput { source, .. } if source.to_string().contains("Cannot write with an empty schema.")), "{:?}", - &result + result ); // Writing empty reader produces an error @@ -421,7 +436,7 @@ mod tests { matches!(result.as_ref().unwrap_err(), Error::InvalidInput { source, .. } if source.to_string().contains("Input data was empty.")), "{:?}", - &result + result ); // Writing with incorrect schema produces an error. @@ -439,7 +454,7 @@ mod tests { matches!(result.as_ref().unwrap_err(), Error::SchemaMismatch { difference, .. } if difference.contains("fields did not match")), "{:?}", - &result + result ); } @@ -498,7 +513,7 @@ mod tests { matches!(result.as_ref().unwrap_err(), Error::InvalidInput { source, .. } if source.to_string().contains("Cannot write with an empty schema.")), "{:?}", - &result + result ); // Writing empty reader produces an error @@ -525,7 +540,7 @@ mod tests { matches!(result.as_ref().unwrap_err(), Error::SchemaMismatch { difference, .. } if difference.contains("fields did not match")), "{:?}", - &result + result ); } @@ -606,6 +621,51 @@ mod tests { assert_eq!(fragments[0].files[0].base_id, Some(2)); } + #[tokio::test] + async fn test_write_fragments_with_target_all_bases() { + let primary = TempStrDir::default(); + let base1 = TempStrDir::default(); + let base2 = TempStrDir::default(); + let create_params = WriteParams::default().with_initial_bases(vec![ + BasePath::new(0, base1.to_string(), Some("base1".to_string()), false), + BasePath::new(0, base2.to_string(), Some("base2".to_string()), false), + ]); + + let dataset = InsertBuilder::new(primary.as_str()) + .with_params(&create_params) + .execute_stream(test_data()) + .await + .unwrap(); + + // Without primary, the first slot is the lowest registered base id. + let append_params = WriteParams { + mode: WriteMode::Append, + ..Default::default() + } + .with_target_all_bases(false); + let fragments = FragmentCreateBuilder::new(dataset.uri.as_str()) + .write_params(&append_params) + .write_fragments(test_data()) + .await + .unwrap(); + assert_eq!(fragments.len(), 1); + assert_eq!(fragments[0].files[0].base_id, Some(1)); + + // With primary included, the first slot is primary storage. + let append_params = WriteParams { + mode: WriteMode::Append, + ..Default::default() + } + .with_target_all_bases(true); + let fragments = FragmentCreateBuilder::new(dataset.uri.as_str()) + .write_params(&append_params) + .write_fragments(test_data()) + .await + .unwrap(); + assert_eq!(fragments.len(), 1); + assert_eq!(fragments[0].files[0].base_id, None); + } + #[rstest] #[tokio::test] async fn test_write_with_format_version( diff --git a/rust/lance/src/dataset/index.rs b/rust/lance/src/dataset/index.rs index 770c68b89e9..7cac7815125 100644 --- a/rust/lance/src/dataset/index.rs +++ b/rust/lance/src/dataset/index.rs @@ -3,7 +3,8 @@ pub mod frag_reuse; -use std::collections::{HashMap, HashSet}; +use lance_core::utils::row_addr_remap::RowAddrRemap; +use std::collections::HashSet; use std::sync::Arc; use crate::Dataset; @@ -47,7 +48,7 @@ impl DatasetIndexRemapper { async fn remap_index( &self, index: &IndexMetadata, - mapping: &HashMap>, + mapping: &RowAddrRemap, ) -> Result { remap_index(&self.dataset, &index.uuid, mapping).await } @@ -57,7 +58,7 @@ impl DatasetIndexRemapper { impl IndexRemapper for DatasetIndexRemapper { async fn remap_indices( &self, - mapping: HashMap>, + mapping: RowAddrRemap, affected_fragment_ids: &[u64], ) -> Result> { let affected_frag_ids = HashSet::::from_iter(affected_fragment_ids.iter().copied()); @@ -175,8 +176,6 @@ impl LanceIndexStoreExt for LanceIndexStore { #[cfg(test)] mod tests { - use std::collections::HashMap; - use super::*; use crate::dataset::WriteParams; use crate::index::DatasetIndexExt; @@ -299,7 +298,7 @@ mod tests { .create_remapper(&dataset) .unwrap(); let remapped = remapper - .remap_indices(HashMap::new(), &[target_fragments[0].id() as u64]) + .remap_indices(RowAddrRemap::empty(), &[target_fragments[0].id() as u64]) .await .unwrap(); diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index 9e8b6a4ff0e..f5b89d06ff4 100644 --- a/rust/lance/src/dataset/mem_wal.rs +++ b/rust/lance/src/dataset/mem_wal.rs @@ -39,6 +39,8 @@ mod manifest; pub mod memtable; pub mod scanner; pub mod sharding; +#[cfg(test)] +pub(crate) mod test_util; pub mod util; mod wal; pub mod write; diff --git a/rust/lance/src/dataset/mem_wal/api.rs b/rust/lance/src/dataset/mem_wal/api.rs index 79184c13ec8..597c65dce83 100644 --- a/rust/lance/src/dataset/mem_wal/api.rs +++ b/rust/lance/src/dataset/mem_wal/api.rs @@ -400,6 +400,14 @@ fn writer_config_to_defaults(config: &ShardWriterConfig) -> HashMap Result { + match index_meta.index_version { + // Legacy Arrow FTS indexes did not use the v1/v2 metadata values, but + // the maintained-index path can only write the modern format. + 0 | 1 => Ok(InvertedListFormatVersion::V1), + 2 => Ok(InvertedListFormatVersion::V2), + 3 => Ok(InvertedListFormatVersion::V3), + version => Err(Error::invalid_input(format!( + "FTS index '{}' has unsupported index_version {}; expected 0, 1, 2, or 3", + index_meta.name, version + ))), + } + } + /// Extract field ID and column name from index metadata. fn extract_field_info( index_meta: &IndexMetadata, @@ -240,6 +258,11 @@ pub struct IndexStore { /// visible to scanners. Advanced unconditionally after a WAL append /// succeeds; not gated on whether any indexes are configured. max_visible_batch_position: AtomicUsize, + /// Conservative flag set once this memtable has observed any primary-key + /// rewrite while maintaining a search index. Search planners can push top-k + /// into HNSW/FTS for append-only PK data, but must switch to + /// newest-before-top-k search after an overwrite. + pk_has_overrides: AtomicBool, } impl Default for IndexStore { @@ -250,6 +273,7 @@ impl Default for IndexStore { fts_indexes: HashMap::new(), pk_index: None, max_visible_batch_position: AtomicUsize::new(0), + pk_has_overrides: AtomicBool::new(false), } } } @@ -280,6 +304,10 @@ impl std::fmt::Debug for IndexStore { "max_visible_batch_position", &self.max_visible_batch_position.load(Ordering::Acquire), ) + .field( + "pk_has_overrides", + &self.pk_has_overrides.load(Ordering::Acquire), + ) .finish() } } @@ -325,8 +353,11 @@ impl IndexStore { registry.hnsw_indexes.insert(c.name.clone(), index); } MemIndexConfig::Fts(c) => { - let index = - FtsMemIndex::with_params(c.field_id, c.column.clone(), c.params.clone()); + let index = FtsMemIndex::try_with_params( + c.field_id, + c.column.clone(), + c.params.clone(), + )?; registry.fts_indexes.insert(c.name.clone(), index); } } @@ -343,6 +374,10 @@ impl IndexStore { } /// Add an HNSW vector index with default build parameters. + /// + /// HNSW indexes must be configured before rows are inserted into a + /// PK-indexed memtable. The vector planner's append-only fast path relies + /// on `pk_has_overrides` being maintained for every row visible to HNSW. pub fn add_hnsw( &mut self, name: String, @@ -352,6 +387,10 @@ impl IndexStore { capacity: usize, max_batches: usize, ) { + assert!( + self.pk_index.is_none() || self.pk_is_empty(), + "HNSW indexes must be configured before inserting rows into a PK memtable" + ); self.hnsw_indexes.insert( name, HnswMemIndex::with_capacity( @@ -366,6 +405,8 @@ impl IndexStore { } /// Add an HNSW vector index with explicit build parameters. + /// + /// See [`Self::add_hnsw`] for the PK-indexed memtable lifecycle invariant. #[allow(clippy::too_many_arguments)] pub fn add_hnsw_with_params( &mut self, @@ -377,6 +418,10 @@ impl IndexStore { capacity: usize, max_batches: usize, ) { + assert!( + self.pk_index.is_none() || self.pk_is_empty(), + "HNSW indexes must be configured before inserting rows into a PK memtable" + ); self.hnsw_indexes.insert( name, HnswMemIndex::with_capacity( @@ -391,7 +436,15 @@ impl IndexStore { } /// Add an FTS index with default tokenizer parameters. + /// + /// FTS indexes must be configured before rows are inserted into a + /// PK-indexed memtable. FTS top-k pushdown relies on `pk_has_overrides` + /// being maintained for every row visible to the index. pub fn add_fts(&mut self, name: String, field_id: i32, column: String) { + assert!( + self.pk_index.is_none() || self.pk_is_empty(), + "FTS indexes must be configured before inserting rows into a PK memtable" + ); self.fts_indexes .insert(name, FtsMemIndex::new(field_id, column)); } @@ -403,9 +456,16 @@ impl IndexStore { field_id: i32, column: String, params: InvertedIndexParams, - ) { - self.fts_indexes - .insert(name, FtsMemIndex::with_params(field_id, column, params)); + ) -> Result<()> { + assert!( + self.pk_index.is_none() || self.pk_is_empty(), + "FTS indexes must be configured before inserting rows into a PK memtable" + ); + self.fts_indexes.insert( + name, + FtsMemIndex::try_with_params(field_id, column, params)?, + ); + Ok(()) } /// Maintain a primary-key index so the memtable can answer "newest visible @@ -417,8 +477,16 @@ impl IndexStore { /// order-preserving encoded tuple (synthetic `PK_KEY_COLUMN`), maintained /// explicitly in the insert paths. Call once at construction, after /// [`Self::from_configs`] and before any inserts; a no-op when `pk_columns` - /// is empty. + /// is empty. Search indexes (HNSW/FTS) must also still be empty so every + /// search-visible row participates in PK override tracking. pub fn enable_pk_index(&mut self, pk_columns: &[(String, i32)]) { + if !pk_columns.is_empty() { + assert!( + self.hnsw_indexes.values().all(|idx| idx.is_empty()) + && self.fts_indexes.values().all(|idx| idx.is_empty()), + "Primary-key indexes must be configured before inserting rows into a search-indexed memtable" + ); + } self.pk_index = match pk_columns { [] => None, [(column, field_id)] => { @@ -483,7 +551,12 @@ impl IndexStore { /// Maintain the composite PK index for `batch` (no-op for single/no PK): /// encode the PK columns into the synthetic `PK_KEY_COLUMN` `Binary` column /// and feed that to the keyed `BTreeMemIndex`. - fn insert_composite_pk(&self, batch: &RecordBatch, row_offset: u64) -> Result<()> { + fn insert_composite_pk( + &self, + batch: &RecordBatch, + row_offset: u64, + report_existing: bool, + ) -> Result { if let Some(PkIndex::Composite { index, columns }) = &self.pk_index { let pk_indices = Self::pk_batch_indices(batch, columns)?; let encoded = encode_pk_batch(batch, &pk_indices)?; @@ -494,9 +567,12 @@ impl IndexStore { )])); let key_batch = RecordBatch::try_new(schema, vec![Arc::new(encoded)]) .map_err(|e| Error::invalid_input(e.to_string()))?; + if report_existing { + return index.insert_and_report_existing(&key_batch, row_offset); + } index.insert(&key_batch, row_offset)?; } - Ok(()) + Ok(false) } /// The newest row position of the primary-key tuple `values` (in PK order) @@ -560,6 +636,32 @@ impl IndexStore { } } + /// Whether this memtable has observed at least one PK rewrite. + /// + /// This is intentionally conservative: once true, it never resets for the + /// lifetime of the memtable. That is enough for query planning because a + /// memtable is flushed as a unit, and any rewrite means search-index top-k + /// pushdown can be polluted by stale entries that must be removed before + /// top-k. Scalar-only PK tables skip tracking because no search index uses + /// the flag. + pub fn pk_has_overrides(&self) -> bool { + self.pk_has_overrides.load(Ordering::Acquire) + } + + fn should_track_pk_overrides(&self) -> bool { + (!self.hnsw_indexes.is_empty() || !self.fts_indexes.is_empty()) && !self.pk_has_overrides() + } + + fn is_single_pk_btree(&self, index: &Arc) -> bool { + matches!(&self.pk_index, Some(PkIndex::Single(pk)) if Arc::ptr_eq(pk, index)) + } + + fn mark_pk_overrides_if_needed(&self, had_existing_pk: bool) { + if had_existing_pk { + self.pk_has_overrides.store(true, Ordering::Release); + } + } + /// Insert a batch into all indexes. pub fn insert(&self, batch: &RecordBatch, row_offset: u64) -> Result<()> { self.insert_with_batch_position(batch, row_offset, None) @@ -573,8 +675,14 @@ impl IndexStore { row_offset: u64, batch_position: Option, ) -> Result<()> { + let track_pk_overrides = self.should_track_pk_overrides(); for index in self.btree_indexes.values() { - index.insert(batch, row_offset)?; + if track_pk_overrides && self.is_single_pk_btree(index) { + let had_existing = index.insert_and_report_existing(batch, row_offset)?; + self.mark_pk_overrides_if_needed(had_existing); + } else { + index.insert(batch, row_offset)?; + } } for index in self.hnsw_indexes.values() { index.insert(batch, row_offset)?; @@ -584,7 +692,8 @@ impl IndexStore { } // Single-column PK aliases a `btree_indexes` entry (maintained above); // a composite PK has its own index, maintained here. - self.insert_composite_pk(batch, row_offset)?; + let had_existing = self.insert_composite_pk(batch, row_offset, track_pk_overrides)?; + self.mark_pk_overrides_if_needed(had_existing); // Update global watermark after all indexes have been updated if let Some(bp) = batch_position { @@ -596,11 +705,11 @@ impl IndexStore { /// Advance the visibility watermark to at least `batch_pos`. /// - /// The watermark only ever moves forward (idempotent max). Public so the - /// WAL flush handler can advance it after a successful WAL append, which - /// is the authoritative durability signal regardless of whether any - /// indexes are configured. - pub fn advance_max_visible_batch_position(&self, batch_pos: usize) { + /// The watermark only ever moves forward (idempotent max). The vector + /// planner relies on the insert paths setting `pk_has_overrides` before + /// calling this method, so any snapshot that can see a PK rewrite also + /// observes `pk_has_overrides == true`. + pub(crate) fn advance_max_visible_batch_position(&self, batch_pos: usize) { let mut current = self.max_visible_batch_position.load(Ordering::Acquire); while batch_pos > current { match self.max_visible_batch_position.compare_exchange_weak( @@ -622,11 +731,20 @@ impl IndexStore { return Ok(()); } + let track_pk_overrides = self.should_track_pk_overrides(); // BTree indexes: iterate batches (no cross-batch optimization benefit) for index in self.btree_indexes.values() { + let track_this_index = track_pk_overrides && self.is_single_pk_btree(index); + let mut had_existing = false; for stored in batches { - index.insert(&stored.data, stored.row_offset)?; + if track_this_index { + had_existing |= + index.insert_and_report_existing(&stored.data, stored.row_offset)?; + } else { + index.insert(&stored.data, stored.row_offset)?; + } } + self.mark_pk_overrides_if_needed(had_existing); } // HNSW indexes: use batched insert @@ -643,9 +761,12 @@ impl IndexStore { // Single-column PK aliases a `btree_indexes` entry (maintained above); // a composite PK has its own index, maintained here. + let mut had_existing = false; for stored in batches { - self.insert_composite_pk(&stored.data, stored.row_offset)?; + had_existing |= + self.insert_composite_pk(&stored.data, stored.row_offset, track_pk_overrides)?; } + self.mark_pk_overrides_if_needed(had_existing); // Update global watermark to the max batch position let max_bp = batches.iter().map(|b| b.batch_position).max().unwrap(); @@ -675,24 +796,32 @@ impl IndexStore { return Ok(std::collections::HashMap::new()); } + let track_pk_overrides = self.should_track_pk_overrides(); // Use std::thread::scope for parallel CPU-bound work std::thread::scope(|scope| { // Each handle returns (index_name, index_type, duration, Result) let mut handles: Vec<( &str, &str, - std::thread::ScopedJoinHandle<'_, (std::time::Duration, Result<()>)>, + std::thread::ScopedJoinHandle<'_, (std::time::Duration, Result)>, )> = Vec::new(); // Spawn a thread for each BTree index for (name, index) in &self.btree_indexes { - let handle = scope.spawn(move || -> (std::time::Duration, Result<()>) { + let track_this_index = track_pk_overrides && self.is_single_pk_btree(index); + let handle = scope.spawn(move || -> (std::time::Duration, Result) { let start = Instant::now(); let result = (|| { + let mut had_existing = false; for stored in batches { - index.insert(&stored.data, stored.row_offset)?; + if track_this_index { + had_existing |= index + .insert_and_report_existing(&stored.data, stored.row_offset)?; + } else { + index.insert(&stored.data, stored.row_offset)?; + } } - Ok(()) + Ok(had_existing) })(); (start.elapsed(), result) }); @@ -701,9 +830,9 @@ impl IndexStore { // Spawn a thread for each HNSW index for (name, index) in &self.hnsw_indexes { - let handle = scope.spawn(move || -> (std::time::Duration, Result<()>) { + let handle = scope.spawn(move || -> (std::time::Duration, Result) { let start = Instant::now(); - let result = index.insert_batches(batches); + let result = index.insert_batches(batches).map(|_| false); (start.elapsed(), result) }); handles.push((name.as_str(), "hnsw", handle)); @@ -711,13 +840,13 @@ impl IndexStore { // Spawn a thread for each FTS index for (name, index) in &self.fts_indexes { - let handle = scope.spawn(move || -> (std::time::Duration, Result<()>) { + let handle = scope.spawn(move || -> (std::time::Duration, Result) { let start = Instant::now(); let result = (|| { for stored in batches { index.insert(&stored.data, stored.row_offset)?; } - Ok(()) + Ok(false) })(); (start.elapsed(), result) }); @@ -729,11 +858,13 @@ impl IndexStore { // BTree updates) are preserved instead of getting truncated to 0. let mut first_error: Option = None; let mut timings: Vec<(&str, &str, std::time::Duration)> = Vec::new(); + let mut had_existing_pk = false; for (name, idx_type, handle) in handles { match handle.join() { - Ok((duration, Ok(()))) => { + Ok((duration, Ok(had_existing))) => { timings.push((name, idx_type, duration)); + had_existing_pk |= had_existing; } Ok((duration, Err(e))) => { timings.push((name, idx_type, duration)); @@ -753,6 +884,7 @@ impl IndexStore { if let Some(e) = first_error { return Err(e); } + self.mark_pk_overrides_if_needed(had_existing_pk); let duration_map: std::collections::HashMap = timings .into_iter() @@ -763,9 +895,12 @@ impl IndexStore { // already maintained it (and joined). A composite PK has its own // index; maintain it here before the watermark advances so the // visible prefix is fully indexed. + let mut had_existing = false; for stored in batches { - self.insert_composite_pk(&stored.data, stored.row_offset)?; + had_existing |= + self.insert_composite_pk(&stored.data, stored.row_offset, track_pk_overrides)?; } + self.mark_pk_overrides_if_needed(had_existing); // Update global watermark to the max batch position let max_bp = batches.iter().map(|b| b.batch_position).max().unwrap(); @@ -869,6 +1004,7 @@ mod tests { use arrow_schema::{DataType, Field, Schema as ArrowSchema}; use log::warn; use std::sync::Arc; + use uuid::Uuid; /// Check if an index type is supported and log warning if not. fn check_index_type_supported(index_type: &str) -> bool { @@ -911,6 +1047,39 @@ mod tests { .unwrap() } + fn fts_index_metadata(index_version: i32) -> IndexMetadata { + let details = + pbold::InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap(); + fts_index_metadata_with_details(index_version, Some(details)) + } + + fn fts_index_metadata_with_details( + index_version: i32, + details: Option, + ) -> IndexMetadata { + let index_details = details.map(|details| { + let mut value = Vec::new(); + details.encode(&mut value).unwrap(); + Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.index.InvertedIndexDetails".to_string(), + value, + }) + }); + + IndexMetadata { + uuid: Uuid::new_v4(), + fields: vec![2], + name: "desc_idx".to_string(), + dataset_version: 1, + fragment_bitmap: None, + index_details, + index_version, + created_at: None, + base_id: None, + files: None, + } + } + /// Single-column `id` batch for primary-key lookup tests. fn id_batch(ids: &[i32]) -> RecordBatch { RecordBatch::try_new( @@ -924,6 +1093,66 @@ mod tests { .unwrap() } + fn id_vector_batch(ids: &[i32]) -> RecordBatch { + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; + + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2), + false, + ), + ])); + let mut vectors = FixedSizeListBuilder::new(Float32Builder::new(), 2); + for id in ids { + vectors.values().append_value(*id as f32); + vectors.values().append_value(*id as f32 + 0.5); + vectors.append(true); + } + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids.to_vec())), + Arc::new(vectors.finish()), + ], + ) + .unwrap() + } + + fn id_name_vector_batch(rows: &[(i32, &str)]) -> RecordBatch { + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; + + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2), + false, + ), + ])); + let mut ids = Vec::with_capacity(rows.len()); + let mut names = Vec::with_capacity(rows.len()); + let mut vectors = FixedSizeListBuilder::new(Float32Builder::new(), 2); + for (id, name) in rows { + ids.push(*id); + names.push(*name); + vectors.values().append_value(*id as f32); + vectors.values().append_value(name.len() as f32); + vectors.append(true); + } + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(names)), + Arc::new(vectors.finish()), + ], + ) + .unwrap() + } + #[test] fn pk_newest_visible_single_column() { let mut store = IndexStore::new(); @@ -942,6 +1171,114 @@ mod tests { assert!(!store.pk_contains_key(&ScalarValue::Int32(Some(9)), 5)); } + #[test] + fn pk_has_overrides_tracks_single_column_rewrites() { + let mut store = IndexStore::new(); + store.add_hnsw( + "vector_hnsw".to_string(), + 1, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + 64, + 8, + ); + store.enable_pk_index(&[("id".to_string(), 0)]); + + store.insert(&id_vector_batch(&[1, 2]), 0).unwrap(); + assert!( + !store.pk_has_overrides(), + "append-only PK inserts should keep HNSW eligible" + ); + + store.insert(&id_vector_batch(&[3, 3]), 2).unwrap(); + assert!( + store.pk_has_overrides(), + "duplicate PKs within one insert must disable HNSW" + ); + } + + #[test] + #[should_panic( + expected = "Primary-key indexes must be configured before inserting rows into a search-indexed memtable" + )] + fn enable_pk_index_after_search_rows_panics() { + let mut store = IndexStore::new(); + store.add_hnsw( + "vector_hnsw".to_string(), + 1, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + 64, + 8, + ); + store.insert(&id_vector_batch(&[1, 2]), 0).unwrap(); + + store.enable_pk_index(&[("id".to_string(), 0)]); + } + + #[test] + fn pk_has_overrides_tracks_single_column_rewrites_across_inserts() { + let mut store = IndexStore::new(); + store.add_hnsw( + "vector_hnsw".to_string(), + 1, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + 64, + 8, + ); + store.enable_pk_index(&[("id".to_string(), 0)]); + + store.insert(&id_vector_batch(&[1, 2]), 0).unwrap(); + assert!( + !store.pk_has_overrides(), + "append-only PK inserts should keep HNSW eligible" + ); + + store.insert(&id_vector_batch(&[1]), 2).unwrap(); + assert!( + store.pk_has_overrides(), + "single-column PK rewrites across inserts must disable HNSW" + ); + } + + #[test] + fn pk_has_overrides_skips_scalar_only_tables() { + let mut store = IndexStore::new(); + store.enable_pk_index(&[("id".to_string(), 0)]); + + store.insert(&id_batch(&[1, 1]), 0).unwrap(); + assert!( + !store.pk_has_overrides(), + "scalar-only PK tables should not pay override tracking cost" + ); + } + + #[test] + fn pk_has_overrides_tracks_fts_rewrites() { + let mut store = IndexStore::new(); + store.enable_pk_index(&[("id".to_string(), 0)]); + store.add_fts("text_fts".to_string(), 1, "text".to_string()); + + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(StringArray::from(vec!["alpha", "beta"])), + ], + ) + .unwrap(); + store.insert(&batch, 0).unwrap(); + assert!( + store.pk_has_overrides(), + "FTS PK rewrites must disable index-level FTS limit/WAND pushdown" + ); + } + #[test] fn pk_newest_visible_composite_seeks_encoded_tuple() { let mut store = IndexStore::new(); @@ -978,6 +1315,30 @@ mod tests { assert!(!store.pk_contains_key(&key_2a, 5)); } + #[test] + fn pk_has_overrides_tracks_composite_rewrites() { + let mut store = IndexStore::new(); + store.add_hnsw( + "vector_hnsw".to_string(), + 2, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + 64, + 8, + ); + store.enable_pk_index(&[("id".to_string(), 0), ("name".to_string(), 1)]); + let first = id_name_vector_batch(&[(1, "a"), (1, "b")]); + store.insert(&first, 0).unwrap(); + assert!(!store.pk_has_overrides()); + + let rewrite = id_name_vector_batch(&[(1, "a")]); + store.insert(&rewrite, 2).unwrap(); + assert!( + store.pk_has_overrides(), + "repeated composite PK must disable HNSW" + ); + } + #[test] fn test_index_registry() { let schema = create_test_schema(); @@ -1011,6 +1372,102 @@ mod tests { assert!(!check_index_type_supported("unknown")); } + #[test] + fn fts_from_metadata_preserves_format_version() { + let arrow_schema = create_test_schema(); + let schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + for (index_version, expected_format_version) in [ + (0, InvertedListFormatVersion::V1), + (1, InvertedListFormatVersion::V1), + (2, InvertedListFormatVersion::V2), + ] { + let config = + MemIndexConfig::fts_from_metadata(&fts_index_metadata(index_version), &schema) + .unwrap(); + + match config { + MemIndexConfig::Fts(config) => { + assert_eq!( + config.params.resolved_format_version(), + expected_format_version + ); + } + _ => unreachable!("fts metadata should create an FTS config"), + } + } + } + + #[test] + fn fts_from_metadata_rejects_unsupported_format_version() { + let arrow_schema = create_test_schema(); + let schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + let err = MemIndexConfig::fts_from_metadata(&fts_index_metadata(4), &schema).unwrap_err(); + assert!( + err.to_string().contains("unsupported index_version 4"), + "{err}" + ); + } + + #[test] + fn fts_from_metadata_rejects_v3_without_256_block_size() { + let arrow_schema = create_test_schema(); + let schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + let missing_details = + MemIndexConfig::fts_from_metadata(&fts_index_metadata_with_details(3, None), &schema) + .unwrap_err(); + assert!( + missing_details + .to_string() + .contains("requires block_size=256"), + "{missing_details}" + ); + assert!( + missing_details.to_string().contains("got 128"), + "{missing_details}" + ); + + let default_details = + MemIndexConfig::fts_from_metadata(&fts_index_metadata(3), &schema).unwrap_err(); + assert!( + default_details + .to_string() + .contains("requires block_size=256"), + "{default_details}" + ); + assert!( + default_details.to_string().contains("got 128"), + "{default_details}" + ); + } + + #[test] + fn fts_from_metadata_accepts_v3_with_256_block_size() { + let arrow_schema = create_test_schema(); + let schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + let params = InvertedIndexParams::default().block_size(256).unwrap(); + let details = pbold::InvertedIndexDetails::try_from(¶ms).unwrap(); + + let config = MemIndexConfig::fts_from_metadata( + &fts_index_metadata_with_details(3, Some(details)), + &schema, + ) + .unwrap(); + + match config { + MemIndexConfig::Fts(config) => { + assert_eq!( + config.params.resolved_format_version(), + InvertedListFormatVersion::V3 + ); + assert_eq!(config.params.posting_block_size(), 256); + } + _ => unreachable!("fts metadata should create an FTS config"), + } + } + #[test] fn test_from_configs() { let configs = vec![ diff --git a/rust/lance/src/dataset/mem_wal/index/arena_skiplist.rs b/rust/lance/src/dataset/mem_wal/index/arena_skiplist.rs index b600f6296de..6b7361e9f1b 100644 --- a/rust/lance/src/dataset/mem_wal/index/arena_skiplist.rs +++ b/rust/lance/src/dataset/mem_wal/index/arena_skiplist.rs @@ -250,9 +250,22 @@ impl SkipListWriter { height } - /// Insert `key`. Keys must be unique (the MemTable key carries a row - /// position, making every entry distinct); equal keys are not expected. + #[cfg(test)] pub fn insert(&mut self, key: K) { + self.insert_and_check_neighbors(key, |_, _| false); + } + + /// Insert `key` and let the caller inspect its immediate level-0 neighbors. + /// + /// Keys must be unique (the MemTable key carries a row position, making + /// every entry distinct). Backends that index `(value, row_position)` use + /// the neighbor check to detect whether the same value already existed, + /// reusing the insertion traversal instead of issuing a second lookup. + pub fn insert_and_check_neighbors( + &mut self, + key: K, + check: impl FnOnce(Option<&K>, Option<&K>) -> bool, + ) -> bool { let cur_height = self.core.height.load(Ordering::Relaxed); // Find the predecessor at every level. For levels at/above the current @@ -274,6 +287,15 @@ impl SkipListWriter { preds[level] = pred; } + let succ = self.core.next_slot(preds[0], 0).load(Ordering::Acquire); + // SAFETY: `preds[0]` and `succ` are either null or point to nodes owned + // by this skiplist's arena. The single writer never removes nodes, so + // any non-null neighbor remains alive for the duration of this call. + let had_neighbor_match = + check(unsafe { preds[0].as_ref().map(|node| &node.key) }, unsafe { + succ.as_ref().map(|node| &node.key) + }); + let height = self.random_height(); // Allocate the node in one block and initialize it fully (key + tower) @@ -309,6 +331,7 @@ impl SkipListWriter { } self.core.len.fetch_add(1, Ordering::Release); + had_neighbor_match } } diff --git a/rust/lance/src/dataset/mem_wal/index/btree.rs b/rust/lance/src/dataset/mem_wal/index/btree.rs index c2f89b9932d..ca4dd178548 100644 --- a/rust/lance/src/dataset/mem_wal/index/btree.rs +++ b/rust/lance/src/dataset/mem_wal/index/btree.rs @@ -305,7 +305,8 @@ impl FixedIntBackend { } } - fn insert_array(&self, array: &dyn Array, row_offset: u64) -> Result<()> { + fn insert_array_and_report_existing(&self, array: &dyn Array, row_offset: u64) -> Result { + let mut had_existing = false; macro_rules! insert_int { ($array_type:ty, $to_i64:expr) => {{ let typed = array @@ -314,14 +315,26 @@ impl FixedIntBackend { .unwrap(); let mut writer = self.writer.lock().unwrap(); let mut nulls: Vec = Vec::new(); + let had_existing_nulls = !self.null_positions.lock().unwrap().is_empty(); + let mut saw_null = false; for (row_idx, value) in typed.iter().enumerate() { let position = row_offset + row_idx as u64; match value { - Some(v) => writer.insert(FixedKey { - enc: $to_i64(v), - position, - }), - None => nulls.push(position), + Some(v) => { + let enc = $to_i64(v); + let key = FixedKey { enc, position }; + had_existing |= writer.insert_and_check_neighbors(key, |prev, next| { + prev.is_some_and(|key| key.enc == enc) + || next.is_some_and(|key| key.enc == enc) + }); + } + None => { + if had_existing_nulls || saw_null { + had_existing = true; + } + saw_null = true; + nulls.push(position); + } } } drop(writer); @@ -348,7 +361,7 @@ impl FixedIntBackend { ))); } } - Ok(()) + Ok(had_existing) } fn get_newest_visible( @@ -458,7 +471,8 @@ impl BytesBackend { } } - fn insert_array(&self, array: &dyn Array, row_offset: u64) -> Result<()> { + fn insert_array_and_report_existing(&self, array: &dyn Array, row_offset: u64) -> Result { + let mut had_existing = false; // Append (position, key bytes) for each row; nulls go to the side list. // `$v => $to_bytes` extracts the key bytes from each non-null value // inline (no closure, so the borrow ties directly to the row value). @@ -467,16 +481,26 @@ impl BytesBackend { let typed = array.as_any().downcast_ref::<$array_type>().unwrap(); let mut writer = self.writer.lock().unwrap(); let mut nulls: Vec = Vec::new(); + let had_existing_nulls = !self.null_positions.lock().unwrap().is_empty(); + let mut saw_null = false; for row_idx in 0..typed.len() { let position = row_offset + row_idx as u64; if typed.is_null(row_idx) { + if had_existing_nulls || saw_null { + had_existing = true; + } + saw_null = true; nulls.push(position); } else { let $v = typed.value(row_idx); let bytes: &[u8] = $to_bytes; - writer.insert(BytesKey { + let key = BytesKey { bytes: InlineBytes::new(bytes), position, + }; + had_existing |= writer.insert_and_check_neighbors(key, |prev, next| { + prev.is_some_and(|key| key.bytes.as_slice() == bytes) + || next.is_some_and(|key| key.bytes.as_slice() == bytes) }); } } @@ -501,7 +525,7 @@ impl BytesBackend { ))); } } - Ok(()) + Ok(had_existing) } fn get_newest_visible( @@ -609,14 +633,22 @@ impl ScalarBackend { } } - fn add(&self, value: OrderableScalarValue, row_position: RowPosition) { - self.writer.lock().unwrap().insert(IndexKey { - value, - row_position, - }); + fn add(&self, value: OrderableScalarValue, row_position: RowPosition) -> bool { + let probe = value.clone(); + self.writer.lock().unwrap().insert_and_check_neighbors( + IndexKey { + value, + row_position, + }, + |prev, next| { + prev.is_some_and(|key| key.value == probe) + || next.is_some_and(|key| key.value == probe) + }, + ) } - fn insert_array(&self, array: &dyn Array, row_offset: u64) -> Result<()> { + fn insert_array_and_report_existing(&self, array: &dyn Array, row_offset: u64) -> Result { + let mut had_existing = false; macro_rules! insert_primitive { ($array_type:ty, $scalar_variant:ident) => {{ let typed_array = array @@ -625,7 +657,7 @@ impl ScalarBackend { .unwrap(); for (row_idx, value) in typed_array.iter().enumerate() { let row_position = row_offset + row_idx as u64; - self.add( + had_existing |= self.add( OrderableScalarValue(ScalarValue::$scalar_variant(value)), row_position, ); @@ -653,7 +685,7 @@ impl ScalarBackend { .unwrap(); for (row_idx, value) in typed_array.iter().enumerate() { let row_position = row_offset + row_idx as u64; - self.add( + had_existing |= self.add( OrderableScalarValue(ScalarValue::Utf8(value.map(|s| s.to_string()))), row_position, ); @@ -666,7 +698,7 @@ impl ScalarBackend { .unwrap(); for (row_idx, value) in typed_array.iter().enumerate() { let row_position = row_offset + row_idx as u64; - self.add( + had_existing |= self.add( OrderableScalarValue(ScalarValue::LargeUtf8(value.map(|s| s.to_string()))), row_position, ); @@ -679,7 +711,7 @@ impl ScalarBackend { .unwrap(); for (row_idx, value) in typed_array.iter().enumerate() { let row_position = row_offset + row_idx as u64; - self.add( + had_existing |= self.add( OrderableScalarValue(ScalarValue::Boolean(value)), row_position, ); @@ -690,11 +722,11 @@ impl ScalarBackend { for row_idx in 0..array.len() { let value = ScalarValue::try_from_array(array, row_idx)?; let row_position = row_offset + row_idx as u64; - self.add(OrderableScalarValue(value), row_position); + had_existing |= self.add(OrderableScalarValue(value), row_position); } } } - Ok(()) + Ok(had_existing) } fn get_newest_visible( @@ -769,11 +801,11 @@ impl Backend { } } - fn insert_array(&self, array: &dyn Array, row_offset: u64) -> Result<()> { + fn insert_array_and_report_existing(&self, array: &dyn Array, row_offset: u64) -> Result { match self { - Self::FixedInt(b) => b.insert_array(array, row_offset), - Self::Bytes(b) => b.insert_array(array, row_offset), - Self::Scalar(b) => b.insert_array(array, row_offset), + Self::FixedInt(b) => b.insert_array_and_report_existing(array, row_offset), + Self::Bytes(b) => b.insert_array_and_report_existing(array, row_offset), + Self::Scalar(b) => b.insert_array_and_report_existing(array, row_offset), } } @@ -873,6 +905,13 @@ impl BTreeMemIndex { /// Insert rows from a batch into the index. pub fn insert(&self, batch: &RecordBatch, row_offset: u64) -> Result<()> { + self.insert_and_report_existing(batch, row_offset) + .map(|_| ()) + } + + /// Insert rows and report whether any inserted key already existed in the + /// index or repeated earlier in this batch. + pub fn insert_and_report_existing(&self, batch: &RecordBatch, row_offset: u64) -> Result { let col_idx = batch .schema() .column_with_name(&self.column_name) @@ -885,7 +924,7 @@ impl BTreeMemIndex { let backend = self .backend .get_or_init(|| Backend::for_type(column.data_type())); - backend.insert_array(column.as_ref(), row_offset) + backend.insert_array_and_report_existing(column.as_ref(), row_offset) } /// Look up row positions for an exact value. diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 8c6c9d85de7..6e0d12a276e 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -54,9 +54,9 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use arc_swap::ArcSwap; use arrow_array::{Array, LargeStringArray, RecordBatch, StringArray, StringViewArray}; use arrow_schema::DataType; -use bitpacking::{BitPacker, BitPacker4x}; use crossbeam_skiplist::SkipMap; use fst::{Map, Streamer}; +use lance_bitpacking::{BitPacker, BitPacker4x}; use lance_core::{Error, Result}; use lance_index::scalar::InvertedIndexParams; use lance_index::scalar::inverted::query::Operator; @@ -88,6 +88,8 @@ pub enum FtsQueryExpr { Match { /// The search query string. query: String, + /// The operator used to combine tokenized query terms. + operator: Operator, /// Boost factor applied to the score (default 1.0). boost: f32, }, @@ -107,6 +109,8 @@ pub enum FtsQueryExpr { /// Maximum edit distance (Levenshtein distance). /// `None` means auto-fuzziness based on token length. fuzziness: Option, + /// Number of initial characters that must match exactly. + prefix_length: u32, /// Maximum number of terms to expand to. max_expansions: usize, /// Boost factor applied to the score (default 1.0). @@ -204,8 +208,13 @@ impl SearchOptions { impl FtsQueryExpr { pub fn match_query(query: impl Into) -> Self { + Self::match_query_with_operator(query, Operator::Or) + } + + pub fn match_query_with_operator(query: impl Into, operator: Operator) -> Self { Self::Match { query: query.into(), + operator, boost: 1.0, } } @@ -230,6 +239,7 @@ impl FtsQueryExpr { Self::Fuzzy { query: query.into(), fuzziness: None, + prefix_length: 0, max_expansions: DEFAULT_MAX_EXPANSIONS, boost: 1.0, } @@ -239,6 +249,7 @@ impl FtsQueryExpr { Self::Fuzzy { query: query.into(), fuzziness: Some(fuzziness), + prefix_length: 0, max_expansions: DEFAULT_MAX_EXPANSIONS, boost: 1.0, } @@ -247,11 +258,13 @@ impl FtsQueryExpr { pub fn fuzzy_with_options( query: impl Into, fuzziness: Option, + prefix_length: u32, max_expansions: usize, ) -> Self { Self::Fuzzy { query: query.into(), fuzziness, + prefix_length, max_expansions, boost: 1.0, } @@ -279,16 +292,24 @@ impl FtsQueryExpr { pub fn with_boost(self, boost: f32) -> Self { match self { - Self::Match { query, .. } => Self::Match { query, boost }, + Self::Match { + query, operator, .. + } => Self::Match { + query, + operator, + boost, + }, Self::Phrase { query, slop, .. } => Self::Phrase { query, slop, boost }, Self::Fuzzy { query, fuzziness, + prefix_length, max_expansions, .. } => Self::Fuzzy { query, fuzziness, + prefix_length, max_expansions, boost, }, @@ -340,6 +361,16 @@ pub fn levenshtein_distance(a: &str, b: &str) -> u32 { prev_row[n] } +fn char_prefix(term: &str, prefix_length: u32) -> &str { + if prefix_length == 0 { + return ""; + } + term.char_indices() + .nth(prefix_length as usize) + .map(|(idx, _)| &term[..idx]) + .unwrap_or(term) +} + /// Builder for constructing Boolean queries. #[derive(Debug, Clone, Default)] pub struct BooleanQueryBuilder { @@ -929,10 +960,20 @@ impl FtsMemIndex { /// Create a new FTS index with custom tokenizer parameters. pub fn with_params(field_id: i32, column_name: String, params: InvertedIndexParams) -> Self { - let pool = TokenizerPool::new(¶ms, Self::DEFAULT_TOKENIZER_POOL_CAP) - .expect("Failed to build tokenizer"); + Self::try_with_params(field_id, column_name, params) + .expect("invalid MemWAL FTS index parameters") + } + + /// Try to create a new FTS index with custom tokenizer parameters. + pub fn try_with_params( + field_id: i32, + column_name: String, + params: InvertedIndexParams, + ) -> Result { + params.validate_format_version()?; + let pool = TokenizerPool::new(¶ms, Self::DEFAULT_TOKENIZER_POOL_CAP)?; let writer_tokenizer = pool.template.box_clone(); - Self { + Ok(Self { field_id, column_name, params, @@ -941,7 +982,7 @@ impl FtsMemIndex { state: ArcSwap::from(IndexState::empty()), freeze_threshold_rows: Self::DEFAULT_FREEZE_THRESHOLD_ROWS, merge: Arc::new(Mutex::new(None)), - } + }) } /// Override the tail freeze threshold (docs) — the analogue of Lucene's @@ -1237,7 +1278,7 @@ impl FtsMemIndex { pub fn search(&self, term: &str) -> Vec { let st = self.state.load_full(); let tokens = self.tokenize_for_search(term); - self.search_match(&st, &tokens, None, true, true) + self.search_match(&st, &tokens, Operator::Or, None, true, true) } /// Search for documents containing an exact phrase, optionally allowing @@ -1269,7 +1310,7 @@ impl FtsMemIndex { max_expansions: usize, ) -> Vec<(String, u32)> { let st = self.state.load_full(); - self.expand_fuzzy_term(&st, term, max_distance, max_expansions, true) + self.expand_fuzzy_term(&st, term, max_distance, 0, max_expansions, true) } /// Search for documents using fuzzy matching on each query token. @@ -1281,7 +1322,7 @@ impl FtsMemIndex { ) -> Vec { let st = self.state.load_full(); let tokens = self.tokenize_for_search(query); - self.search_fuzzy_tokens(&st, &tokens, fuzziness, max_expansions, true) + self.search_fuzzy_tokens(&st, &tokens, fuzziness, 0, max_expansions, true) } /// BM25 OR-search over the query tokens, scored with one corpus-wide @@ -1294,6 +1335,7 @@ impl FtsMemIndex { &self, st: &IndexState, tokens: &[String], + operator: Operator, limit: Option, include_tail: bool, tail_skip: bool, @@ -1309,8 +1351,8 @@ impl FtsMemIndex { if scorer.num_docs() == 0 { return Vec::new(); } - match limit { - Some(k) if k > 0 => { + match (operator, limit) { + (Operator::Or, Some(k)) if k > 0 => { let mut topk = TopK::new(k); // Scan the block-max partitions first to warm the shared // threshold, then the (un-skippable) tail last — so the tail @@ -1325,7 +1367,14 @@ impl FtsMemIndex { } else { f32::NEG_INFINITY }; - for e in score_terms(&tail_snap, &st.tail.terms, tokens, &scorer, theta) { + for e in score_terms( + &tail_snap, + &st.tail.terms, + tokens, + Operator::Or, + &scorer, + theta, + ) { topk.offer(e.score, e.row_position); } } @@ -1334,13 +1383,14 @@ impl FtsMemIndex { _ => { let mut results = Vec::new(); for p in st.partitions.iter() { - results.extend(p.search_match(tokens, Operator::Or, &scorer)); + results.extend(p.search_match(tokens, operator, &scorer)); } if scan_tail { results.extend(score_terms( &tail_snap, &st.tail.terms, tokens, + operator, &scorer, f32::NEG_INFINITY, )); @@ -1362,7 +1412,7 @@ impl FtsMemIndex { } if tokens.len() == 1 { // A single-token phrase reduces to a regular term search. - return self.search_match(st, tokens, None, include_tail, true); + return self.search_match(st, tokens, Operator::Or, None, include_tail, true); } // A multi-token phrase needs token positions; without them (the index // was built `with_position = false`) phrase search is unsupported, as @@ -1397,6 +1447,7 @@ impl FtsMemIndex { st: &IndexState, tokens: &[String], fuzziness: Option, + prefix_length: u32, max_expansions: usize, include_tail: bool, ) -> Vec { @@ -1407,9 +1458,14 @@ impl FtsMemIndex { let mut seen: HashSet = HashSet::new(); for tok in tokens { let max_dist = fuzziness.unwrap_or_else(|| auto_fuzziness(tok)); - for (matched, _) in - self.expand_fuzzy_term(st, tok, max_dist, max_expansions, include_tail) - { + for (matched, _) in self.expand_fuzzy_term( + st, + tok, + max_dist, + prefix_length, + max_expansions, + include_tail, + ) { if seen.insert(matched.clone()) { expanded.push(matched); } @@ -1418,7 +1474,7 @@ impl FtsMemIndex { if expanded.is_empty() { return Vec::new(); } - self.search_match(st, &expanded, None, include_tail, true) + self.search_match(st, &expanded, Operator::Or, None, include_tail, true) } /// Expand `term` against the term dictionaries of every partition (and the @@ -1428,6 +1484,7 @@ impl FtsMemIndex { st: &IndexState, term: &str, max_distance: u32, + prefix_length: u32, max_expansions: usize, include_tail: bool, ) -> Vec<(String, u32)> { @@ -1449,11 +1506,15 @@ impl FtsMemIndex { } let mut matches: Vec<(String, u32)> = Vec::new(); let mut seen: HashSet = HashSet::new(); + let prefix = char_prefix(term, prefix_length); for entry in st.tail.terms.iter() { if !include_tail || !has_visible_chunk(&entry.value().load(), tail_snap.visible_count) { continue; } let key: &Arc = entry.key(); + if !key.starts_with(prefix) { + continue; + } let dist = levenshtein_distance(term, key); if dist <= max_distance && seen.insert(key.to_string()) { matches.push((key.to_string(), dist)); @@ -1461,6 +1522,9 @@ impl FtsMemIndex { } for p in st.partitions.iter() { for key in p.collect_terms() { + if !key.starts_with(prefix) { + continue; + } let dist = levenshtein_distance(term, key.as_ref()); if dist <= max_distance && seen.insert(key.to_string()) { matches.push((key.to_string(), dist)); @@ -1497,9 +1561,14 @@ impl FtsMemIndex { tail_skip: bool, ) -> Vec { match query { - FtsQueryExpr::Match { query, boost } => { + FtsQueryExpr::Match { + query, + operator, + boost, + } => { let tokens = self.tokenize_for_search(query); - let mut results = self.search_match(st, &tokens, limit, include_tail, tail_skip); + let mut results = + self.search_match(st, &tokens, *operator, limit, include_tail, tail_skip); apply_boost(&mut results, *boost); results } @@ -1512,6 +1581,7 @@ impl FtsMemIndex { FtsQueryExpr::Fuzzy { query, fuzziness, + prefix_length, max_expansions, boost, } => { @@ -1520,6 +1590,7 @@ impl FtsMemIndex { st, &tokens, *fuzziness, + *prefix_length, *max_expansions, include_tail, ); @@ -1551,6 +1622,10 @@ impl FtsMemIndex { query: &FtsQueryExpr, options: SearchOptions, ) -> Vec { + if options.limit == Some(0) { + return Vec::new(); + } + let st = self.state.load_full(); let mut results = self.search_query_with_state( query, @@ -1699,6 +1774,9 @@ impl FtsMemIndex { let st = self.state.load_full(); let with_position = self.params.has_positions(); + let block_size = self.params.posting_block_size(); + let format_version = self.params.resolved_format_version(); + let posting_tail_codec = format_version.posting_tail_codec(); let total_rows_u64 = total_rows as u64; // Step 1: collect (original_pos, num_tokens) for every doc across all @@ -1716,10 +1794,12 @@ impl FtsMemIndex { } } if all_docs.is_empty() { - return Ok(InnerBuilder::new( + return Ok(InnerBuilder::new_with_format_version_and_block_size( partition_id, with_position, Default::default(), + format_version, + block_size, )); } @@ -1805,7 +1885,13 @@ impl FtsMemIndex { docs_for_term.sort_by_key(|(doc_id, _, _)| *doc_id); let token_id = tokens.add(token) as usize; debug_assert_eq!(token_id, posting_lists.len()); - posting_lists.push(PostingListBuilder::new(with_position)); + posting_lists.push( + PostingListBuilder::new_with_posting_tail_codec_and_block_size( + with_position, + posting_tail_codec, + block_size, + ), + ); let plb = &mut posting_lists[token_id]; for (doc_id, freq, pos) in docs_for_term { let recorder = if with_position { @@ -1817,7 +1903,13 @@ impl FtsMemIndex { } } - let mut builder = InnerBuilder::new(partition_id, with_position, Default::default()); + let mut builder = InnerBuilder::new_with_format_version_and_block_size( + partition_id, + with_position, + Default::default(), + format_version, + block_size, + ); builder.set_tokens(tokens); builder.set_docs(docs); builder.set_posting_lists(posting_lists); @@ -2021,12 +2113,13 @@ fn tail_token_df( } } -/// OR-score `tokens` against the visible tail, summing each token's BM25 +/// Score `tokens` against the visible tail, summing each token's BM25 /// contribution per document. Uses the shared corpus-wide `scorer`. fn score_terms( snap: &Snapshot, terms: &SkipMap, Arc>>, tokens: &[String], + operator: Operator, scorer: &MemBM25Scorer, theta: f32, ) -> Vec { @@ -2038,6 +2131,9 @@ fn score_terms( let mut tail_terms: Vec<(f32, Arc)> = Vec::with_capacity(tokens.len()); for token in tokens { let Some(entry) = terms.get(token.as_str()) else { + if operator == Operator::And { + return Vec::new(); + } continue; }; let qw = scorer.query_weight(token); @@ -2058,6 +2154,8 @@ fn score_terms( return Vec::new(); } let mut doc_scores: HashMap = HashMap::new(); + let mut doc_hits: Option> = + (operator == Operator::And).then(HashMap::new); for (qw, slice) in tail_terms { for chunk in slice.chunks() { if chunk.batch_position >= snap.visible_count { @@ -2070,11 +2168,21 @@ fn score_terms( let dl = meta.dl(row_position).unwrap_or(1); let score = qw * scorer.doc_weight(chunk.frequencies[i], dl); *doc_scores.entry(row_position).or_default() += score; + if let Some(doc_hits) = &mut doc_hits { + *doc_hits.entry(row_position).or_default() += 1; + } } } } doc_scores .into_iter() + .filter(|(row_position, _)| { + operator == Operator::Or + || doc_hits + .as_ref() + .and_then(|doc_hits| doc_hits.get(row_position)) + .is_some_and(|hits| *hits >= tokens.len() as u32) + }) .map(|(row_position, score)| FtsEntry { row_position, score, @@ -2240,12 +2348,23 @@ impl FtsIndexConfig { column: String, params: InvertedIndexParams, ) -> Self { - Self { + Self::try_with_params(name, field_id, column, params) + .expect("invalid MemWAL FTS index config parameters") + } + + pub fn try_with_params( + name: String, + field_id: i32, + column: String, + params: InvertedIndexParams, + ) -> Result { + params.validate_format_version()?; + Ok(Self { name, field_id, column, params, - } + }) } } @@ -4057,6 +4176,47 @@ mod tests { assert!(!entries.is_empty()); } + #[test] + fn test_search_query_fuzzy_respects_prefix_length() { + let schema = create_test_schema(); + let index = FtsMemIndex::new(1, "description".to_string()); + + let batch = create_fuzzy_test_batch(&schema); + index.insert(&batch, 0).unwrap(); + + let query = FtsQueryExpr::fuzzy_with_options("alpha", Some(5), 2, 50); + let entries = index.search_query(&query); + assert!(!entries.is_empty()); + assert!( + entries.iter().all(|entry| entry.row_position != 3), + "the omega row is within edit distance but does not share the required prefix" + ); + } + + #[test] + fn test_search_query_fuzzy_prefix_length_uses_char_boundaries() { + let schema = create_test_schema(); + let index = FtsMemIndex::new(1, "description".to_string()); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![0, 1])), + Arc::new(StringArray::from(vec!["éclair", "xclair"])), + ], + ) + .unwrap(); + index.insert(&batch, 0).unwrap(); + + let query = FtsQueryExpr::fuzzy_with_options("éclair", Some(2), 1, 50); + let entries = index.search_query(&query); + assert!(!entries.is_empty()); + assert!( + entries.iter().all(|entry| entry.row_position != 1), + "the ASCII-prefixed row is within edit distance but not the character prefix" + ); + } + #[test] fn test_search_query_fuzzy_with_boost() { let schema = create_test_schema(); @@ -4332,6 +4492,20 @@ mod tests { } } + #[test] + fn test_search_with_options_zero_limit_with_wand_factor() { + let schema = create_test_schema(); + let index = FtsMemIndex::new(1, "description".to_string()); + + let batch = create_wand_test_batch(&schema); + index.insert(&batch, 0).unwrap(); + + let query = FtsQueryExpr::match_query("alpha"); + let options = SearchOptions::new().with_limit(0).with_wand_factor(0.5); + let results = index.search_with_options(&query, options); + assert!(results.is_empty()); + } + #[test] fn test_search_with_options_empty_results() { let schema = create_test_schema(); @@ -4519,6 +4693,18 @@ mod tests { assert!(builder.id() > 0 || builder.id() == 42); } + #[test] + fn test_to_index_builder_supports_block_size_256() { + let schema = create_test_schema(); + let params = InvertedIndexParams::default().block_size(256).unwrap(); + let index = FtsMemIndex::try_with_params(1, "description".to_string(), params).unwrap(); + let batch = create_test_batch(&schema); + index.insert(&batch, 0).unwrap(); + + let builder = index.to_index_builder(42, 3).unwrap(); + assert_eq!(builder.id(), 42); + } + #[test] fn test_unsupported_column_type_errors() { let schema = Arc::new(ArrowSchema::new(vec![ diff --git a/rust/lance/src/dataset/mem_wal/index/hnsw.rs b/rust/lance/src/dataset/mem_wal/index/hnsw.rs index 7ac162432ee..7c7993bb298 100644 --- a/rust/lance/src/dataset/mem_wal/index/hnsw.rs +++ b/rust/lance/src/dataset/mem_wal/index/hnsw.rs @@ -527,13 +527,11 @@ mod tests { let results = index.search(&query, 5, Some(32), u64::MAX).unwrap(); assert!(!results.is_empty()); - let (best_dist, best_pos) = results[0]; assert!( - best_dist < 1e-4, - "expected near-zero distance, got {}", - best_dist + results.iter().any(|&(dist, pos)| pos == 82 && dist < 1e-4), + "expected exact row position 82 in top-5 candidates, got {:?}", + results ); - assert_eq!(best_pos, 82); } #[test] diff --git a/rust/lance/src/dataset/mem_wal/manifest.rs b/rust/lance/src/dataset/mem_wal/manifest.rs index 48df22dca88..9c2a3aa2163 100644 --- a/rust/lance/src/dataset/mem_wal/manifest.rs +++ b/rust/lance/src/dataset/mem_wal/manifest.rs @@ -35,7 +35,7 @@ use bytes::Bytes; use futures::StreamExt; use futures::stream::FuturesUnordered; use lance_core::{Error, Result}; -use lance_index::mem_wal::ShardManifest; +use lance_index::mem_wal::{ShardManifest, ShardStatus}; use lance_io::object_store::ObjectStore; use lance_table::format::pb; use log::{info, warn}; @@ -147,6 +147,7 @@ impl ShardManifestStore { wal_entry_position_last_seen: 0, current_generation: 1, flushed_generations: vec![], + status: ShardStatus::Active, }; match self.write(&manifest).await { @@ -425,6 +426,21 @@ impl ShardManifestStore { for _ in 0..MAX_CLAIM_RETRIES { let current = self.read_latest().await?; + // A sealed shard is mid-drop (drop-table 2PC). Refuse the claim + // with a distinguishable error rather than minting a new epoch, + // so a caller that skips its own status check still cannot + // resurrect a shard being dropped. Sophon's reconcile keys on + // the "sealed" marker in this message to tell it apart from an + // ordinary epoch fence. + if let Some(m) = ¤t + && m.status == ShardStatus::Sealed + { + return Err(Error::invalid_input(format!( + "shard {} is sealed; refusing claim (drop in flight)", + self.shard_id + ))); + } + let (next_version, next_epoch, base_manifest) = match current { Some(m) => (m.version + 1, m.writer_epoch + 1, Some(m)), None => (1, 1, None), @@ -447,6 +463,7 @@ impl ShardManifestStore { wal_entry_position_last_seen: 0, current_generation: 1, flushed_generations: vec![], + status: ShardStatus::Active, } }; @@ -502,8 +519,8 @@ impl ShardManifestStore { shard_id: Uuid, ) -> Result<()> { match manifest { - Some(m) if m.writer_epoch > local_epoch => Err(Error::io(format!( - "Writer fenced: local epoch {} < stored epoch {} for shard {}", + Some(m) if m.writer_epoch > local_epoch => Err(Error::fenced_by_peer(format!( + "local epoch {} < stored epoch {} for shard {}", local_epoch, m.writer_epoch, shard_id ))), _ => Ok(()), @@ -603,6 +620,7 @@ mod tests { wal_entry_position_last_seen: 0, current_generation: 1, flushed_generations: vec![], + status: ShardStatus::Active, } } @@ -752,6 +770,51 @@ mod tests { assert_eq!(second.wal_entry_position_last_seen, 42); } + #[tokio::test] + async fn test_claim_epoch_refuses_sealed_manifest() { + // A `Sealed` manifest is the drop-table 2PC in-doubt marker: + // `claim_epoch` must refuse it with a distinguishable error rather + // than mint a new epoch, so a sealed shard can't be resurrected even + // by a caller that skips its own status check. Rolling the status + // back to `Active` makes the shard claimable again (reversible). + let (store, base_path, _temp_dir) = create_local_store().await; + let shard_id = Uuid::new_v4(); + let manifest_store = ShardManifestStore::new(store, &base_path, shard_id, 2); + + let (epoch, claimed) = manifest_store.claim_epoch(0).await.unwrap(); + assert_eq!(claimed.status, ShardStatus::Active); + + // Seal it (drop-table prepare). + let sealed = ShardManifest { + version: claimed.version + 1, + status: ShardStatus::Sealed, + ..claimed + }; + manifest_store.write(&sealed).await.unwrap(); + + // The claim is refused with the distinguishable "sealed" error — + // and the manifest is left untouched (no new epoch minted). + let err = manifest_store.claim_epoch(0).await.unwrap_err(); + assert!( + err.to_string().contains("sealed"), + "expected a distinguishable sealed-refusal error, got: {err}" + ); + let after = manifest_store.read_latest().await.unwrap().unwrap(); + assert_eq!(after.writer_epoch, sealed.writer_epoch, "no epoch minted"); + assert_eq!(after.status, ShardStatus::Sealed); + + // Roll back to Active (drop-table abort) → re-claimable. + let active = ShardManifest { + version: sealed.version + 1, + status: ShardStatus::Active, + ..sealed + }; + manifest_store.write(&active).await.unwrap(); + let (next_epoch, reclaimed) = manifest_store.claim_epoch(0).await.unwrap(); + assert!(next_epoch > epoch, "rolled-back shard mints the next epoch"); + assert_eq!(reclaimed.status, ShardStatus::Active); + } + #[tokio::test] async fn test_initialize_shard_rejects_conflict_with_mismatch() { let (store, base_path, _temp_dir) = create_local_store().await; diff --git a/rust/lance/src/dataset/mem_wal/memtable.rs b/rust/lance/src/dataset/mem_wal/memtable.rs index cb95a4ab531..77611fd7c47 100644 --- a/rust/lance/src/dataset/mem_wal/memtable.rs +++ b/rust/lance/src/dataset/mem_wal/memtable.rs @@ -22,6 +22,7 @@ use uuid::Uuid; use super::index::IndexStore; use super::util::{WatchableOnceCell, WatchableOnceCellReader}; +use super::wal::WalFlushFailure; use super::write::{DurabilityResult, WalFlushResult}; use crate::Dataset; use batch_store::BatchStore; @@ -103,9 +104,10 @@ pub struct MemTable { /// Set when the memtable is frozen and a WAL flush request is sent. /// The reader can be awaited to know when WAL flush is complete. /// Uses Mutex for interior mutability since the MemTable is wrapped in Arc when frozen. - /// Uses Result since lance_core::Error doesn't implement Clone. + /// Uses `WalFlushFailure` (not `Error`) since `lance_core::Error` doesn't + /// implement Clone; the carrier preserves the fence reason for the waiter. wal_flush_completion: std::sync::Mutex< - Option>>, + Option>>, >, /// Cell for memtable flush completion notification. @@ -266,7 +268,7 @@ impl MemTable { /// the WAL flush is complete. pub fn set_wal_flush_completion( &self, - reader: WatchableOnceCellReader>, + reader: WatchableOnceCellReader>, ) { *self.wal_flush_completion.lock().unwrap() = Some(reader); } @@ -278,7 +280,7 @@ impl MemTable { /// Thread-safe via interior mutability. pub fn take_wal_flush_completion( &self, - ) -> Option>> { + ) -> Option>> { self.wal_flush_completion.lock().unwrap().take() } diff --git a/rust/lance/src/dataset/mem_wal/memtable/batch_store.rs b/rust/lance/src/dataset/mem_wal/memtable/batch_store.rs index 054d9b1630e..30e47e1a9bb 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/batch_store.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/batch_store.rs @@ -43,7 +43,9 @@ use std::cell::UnsafeCell; use std::mem::MaybeUninit; use std::sync::atomic::{AtomicUsize, Ordering}; +use arrow::array::ArrayData; use arrow_array::RecordBatch; +use arrow_schema::DataType; /// A batch stored in the lock-free store. #[derive(Clone)] @@ -75,14 +77,56 @@ impl StoredBatch { } /// Estimate the memory size of a RecordBatch. + /// + /// Sums each column's slice-aware buffer size (see + /// [`Self::estimate_array_size`]) plus the struct overhead, so a column that + /// is a zero-copy slice of a larger parent contributes only its own window + /// rather than the whole shared buffer. fn estimate_batch_size(batch: &RecordBatch) -> usize { batch .columns() .iter() - .map(|col| col.get_array_memory_size()) + .map(|col| Self::estimate_array_size(&col.to_data())) .sum::() + std::mem::size_of::() } + + /// Slice-aware buffer size of a single array. + /// + /// [`ArrayData::get_slice_memory_size`] reports each buffer's own window + /// (not the whole shared buffer), but omits the variadic data buffers of + /// `Utf8View`/`BinaryView` (values > 12 bytes) while still returning `Ok`, so + /// [`Self::view_data_buffers_size`] adds them. Those buffers are shared across + /// zero-copy slices and are counted at full capacity for each slice — an + /// over-count in the safe direction. + fn estimate_array_size(data: &ArrayData) -> usize { + match data.get_slice_memory_size() { + Ok(size) => size + Self::view_data_buffers_size(data), + // Fall back to the full-buffer sum for layouts the slice-aware call + // cannot handle. + Err(_) => data.get_array_memory_size(), + } + } + + /// Capacity of the variadic `Utf8View`/`BinaryView` data buffers that + /// [`ArrayData::get_slice_memory_size`] omits, summed recursively over children. + fn view_data_buffers_size(data: &ArrayData) -> usize { + let mut size = 0; + if matches!(data.data_type(), DataType::Utf8View | DataType::BinaryView) { + // buffers()[0] is the 16-byte view array that get_slice_memory_size + // already counts; [1..] are the data buffers it skips. + size += data + .buffers() + .iter() + .skip(1) + .map(|b| b.capacity()) + .sum::(); + } + for child in data.child_data() { + size += Self::view_data_buffers_size(child); + } + size + } } /// Snapshot of the active batches that have not yet been flushed to WAL. @@ -972,6 +1016,122 @@ mod tests { assert_eq!(cap, 16); // minimum } + #[test] + fn test_estimated_size_is_slice_aware() { + // A batch that is a zero-copy slice of a larger parent must contribute + // only its own window to the estimate, not the whole shared buffer. + // `get_array_memory_size` counts every buffer's full capacity regardless + // of offset/length, so N slices tiling one parent each report the + // parent's size and inflate the memtable estimate ~N×, tripping the + // flush threshold far below the configured size. + let chunk = 1_000; + let num_slices = 100; + let parent = create_test_batch(chunk * num_slices); + + // One window vs an equivalently-sized owned batch should track each + // other; the buggy per-slice estimate would be ~num_slices× larger. + let slice_est = StoredBatch::estimate_batch_size(&parent.slice(0, chunk)); + let owned_est = StoredBatch::estimate_batch_size(&create_test_batch(chunk)); + assert!( + slice_est <= owned_est * 2, + "slice estimate {slice_est} should track its own window (~{owned_est}), not the parent" + ); + + // End-to-end: tiling the parent with zero-copy slices must not multiply + // the store's running estimate. Track what the old full-buffer behavior + // would have summed to for contrast. + let store = BatchStore::with_capacity(num_slices); + let mut over_counting_sum = 0usize; + for k in 0..num_slices { + let s = parent.slice(k * chunk, chunk); + over_counting_sum += s + .columns() + .iter() + .map(|col| col.get_array_memory_size()) + .sum::() + + std::mem::size_of::(); + store.append(s).unwrap(); + } + + // Two non-nullable Int32 columns → exactly 4 bytes/row/col of payload. + let payload_bytes = num_slices * chunk * 2 * std::mem::size_of::(); + let estimated = store.estimated_bytes(); + assert!( + estimated >= payload_bytes, + "estimate {estimated} should cover the actual payload {payload_bytes}" + ); + // The old behavior over-counts by ~num_slices×; the fix must be far + // below it (generous 10× margin against struct/alignment overhead). + assert!( + estimated * 10 < over_counting_sum, + "estimate {estimated} should be far below the over-counting sum {over_counting_sum}" + ); + } + + #[test] + fn test_estimated_size_counts_view_data_buffers() { + // Long Utf8View/BinaryView values live in variadic data buffers that + // `get_slice_memory_size` ignores (returning ~16 * rows). The estimate + // must include them, both for a top-level view column and for a view + // array nested in a container, which is only reached via child_data + // recursion. + use arrow_array::{Array, ArrayRef, StringViewArray, StructArray}; + + let num_rows = 1_000; + // Each value exceeds the 12-byte inline limit, so it spills to a data buffer. + let long_value = "x".repeat(64); + let payload_bytes = num_rows * long_value.len(); + // What the slice-aware call alone reports: just the 16-byte view entries. + let view_entries_only = num_rows * 16; + + let make_views = || { + StringViewArray::from( + (0..num_rows) + .map(|_| Some(long_value.as_str())) + .collect::>(), + ) + }; + let assert_covers = |batch: &RecordBatch| { + let estimated = StoredBatch::estimate_batch_size(batch); + assert!( + estimated >= payload_bytes, + "estimate {estimated} should cover the view data-buffer payload {payload_bytes}" + ); + assert!( + estimated > view_entries_only * 2, + "estimate {estimated} must exceed the ~{view_entries_only}-byte view-entry-only undercount" + ); + }; + + // Top-level view column. + let flat = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![Field::new( + "s", + DataType::Utf8View, + false, + )])), + vec![Arc::new(make_views())], + ) + .unwrap(); + assert_covers(&flat); + + // View nested inside a struct — reachable only through child_data recursion. + let nested = StructArray::from(vec![( + Arc::new(Field::new("s", DataType::Utf8View, false)), + Arc::new(make_views()) as ArrayRef, + )]); + let nested = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![Field::new( + "st", + nested.data_type().clone(), + false, + )])), + vec![Arc::new(nested)], + ) + .unwrap(); + assert_covers(&nested); + } + #[test] fn test_to_vec() { let store = BatchStore::with_capacity(10); diff --git a/rust/lance/src/dataset/mem_wal/memtable/flush.rs b/rust/lance/src/dataset/mem_wal/memtable/flush.rs index 122fcc90cbb..50c21a3eaf5 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/flush.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/flush.rs @@ -650,7 +650,6 @@ impl MemTableFlusher { total_rows: usize, ) -> Result> { use lance_index::pbold; - use lance_index::scalar::inverted::current_fts_format_version; use lance_index::scalar::lance_format::LanceIndexStore; let fts_configs: Vec<_> = index_configs @@ -713,6 +712,7 @@ impl MemTableFlusher { })?; let fragment_ids: roaring::RoaringBitmap = dataset.fragment_bitmap.as_ref().clone(); + let format_version = fts_cfg.params.resolved_format_version(); let index_meta = IndexMetadata { uuid: index_uuid, @@ -721,7 +721,7 @@ impl MemTableFlusher { dataset_version: dataset.version().version, fragment_bitmap: Some(fragment_ids), index_details: Some(Arc::new(index_details)), - index_version: current_fts_format_version().index_version() as i32, + index_version: format_version.index_version() as i32, created_at: None, base_id: None, files: None, @@ -748,22 +748,50 @@ impl MemTableFlusher { use arrow_schema::{DataType, Field, Schema}; use std::sync::Arc; - use lance_index::scalar::inverted::TokenSetFormat; + use lance_index::scalar::inverted::{ + FTS_FORMAT_VERSION_KEY, POSITIONS_CODEC_KEY, POSITIONS_CODEC_PACKED_DELTA_V1, + POSITIONS_LAYOUT_KEY, POSITIONS_LAYOUT_SHARED_STREAM_V2, POSTING_BLOCK_SIZE_KEY, + POSTING_TAIL_CODEC_KEY, TokenSetFormat, + }; // Create metadata with params and partitions in schema metadata (this is what InvertedIndex expects) let params_json = serde_json::to_string(&config.params)?; let partitions_json = serde_json::to_string(&[partition_id])?; let token_set_format = TokenSetFormat::default().to_string(); + let format_version = config.params.resolved_format_version(); + let mut metadata = [ + ("params".to_string(), params_json), + ("partitions".to_string(), partitions_json), + ("token_set_format".to_string(), token_set_format), + ( + POSTING_TAIL_CODEC_KEY.to_string(), + format_version.posting_tail_codec().as_str().to_string(), + ), + ( + FTS_FORMAT_VERSION_KEY.to_string(), + format_version.index_version().to_string(), + ), + ( + POSTING_BLOCK_SIZE_KEY.to_string(), + config.params.posting_block_size().to_string(), + ), + ] + .into_iter() + .collect::>(); + if config.params.has_positions() && format_version.uses_shared_position_stream() { + metadata.insert( + POSITIONS_LAYOUT_KEY.to_string(), + POSITIONS_LAYOUT_SHARED_STREAM_V2.to_string(), + ); + metadata.insert( + POSITIONS_CODEC_KEY.to_string(), + POSITIONS_CODEC_PACKED_DELTA_V1.to_string(), + ); + } let schema = Arc::new( - Schema::new(vec![Field::new("_placeholder", DataType::Utf8, true)]).with_metadata( - [ - ("params".to_string(), params_json), - ("partitions".to_string(), partitions_json), - ("token_set_format".to_string(), token_set_format), - ] - .into(), - ), + Schema::new(vec![Field::new("_placeholder", DataType::Utf8, true)]) + .with_metadata(metadata), ); // Create a minimal batch (schema metadata is what matters) diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs index 17fa9c76a65..e1f0d6e689c 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs @@ -7,21 +7,24 @@ use std::sync::Arc; use arrow_array::{Array, RecordBatch}; use arrow_schema::{DataType, Field, SchemaRef}; -use datafusion::common::{ScalarValue, ToDFSchema}; +use datafusion::common::ScalarValue; use datafusion::physical_plan::limit::GlobalLimitExec; use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream}; use datafusion::prelude::{Expr, SessionContext}; +use datafusion_physical_expr::PhysicalExprRef; use futures::TryStreamExt; use lance_core::{Error, ROW_ID, Result}; use lance_datafusion::expr::safe_coerce_scalar; use lance_datafusion::planner::Planner; +use lance_index::scalar::FullTextSearchQuery; +use lance_index::scalar::inverted::query::{FtsQuery as IndexFtsQuery, Operator}; use lance_linalg::distance::DistanceType; use super::exec::{ BTreeIndexExec, FtsIndexExec, MemTableBruteForceVectorExec, MemTableDedupScanExec, - MemTableScanExec, VectorIndexExec, + MemTableScanExec, SCORE_COLUMN, VectorIndexExec, }; -use crate::dataset::mem_wal::scanner::exec::validate_pk_types; +use crate::dataset::mem_wal::scanner::{exec::validate_pk_types, parse_filter_expr}; use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; /// Vector search query parameters. @@ -58,6 +61,10 @@ pub enum FtsQueryType { Match { /// The search query string. query: String, + /// The operator used to combine tokenized query terms. + operator: Operator, + /// Boost factor applied to the score. + boost: f32, }, /// Phrase query with slop. Phrase { @@ -82,8 +89,12 @@ pub enum FtsQueryType { /// Maximum edit distance (Levenshtein distance). /// None means auto-fuzziness based on token length. fuzziness: Option, + /// Number of initial characters that must match exactly. + prefix_length: u32, /// Maximum number of terms to expand to. max_expansions: usize, + /// Boost factor applied to the score. + boost: f32, }, } @@ -97,6 +108,8 @@ pub struct FtsQuery { /// WAND factor for early termination (0.0 to 1.0). /// 1.0 = full recall (default), <1.0 = faster but may miss low-scoring results. pub wand_factor: f32, + /// Query-level result limit. + pub limit: Option, /// Whether to also search the mutable tail (rows written since the last /// freeze). `true` (default) = read-your-writes; `false` = search only the /// immutable frozen partitions (the Lucene model), trading read-recency for @@ -113,12 +126,23 @@ pub const DEFAULT_WAND_FACTOR: f32 = 1.0; impl FtsQuery { /// Create a simple term match query. pub fn match_query(column: impl Into, query: impl Into) -> Self { + Self::match_query_with_operator(column, query, Operator::Or) + } + + pub fn match_query_with_operator( + column: impl Into, + query: impl Into, + operator: Operator, + ) -> Self { Self { column: column.into(), query_type: FtsQueryType::Match { query: query.into(), + operator, + boost: 1.0, }, wand_factor: DEFAULT_WAND_FACTOR, + limit: None, include_tail: true, } } @@ -132,6 +156,7 @@ impl FtsQuery { slop, }, wand_factor: DEFAULT_WAND_FACTOR, + limit: None, include_tail: true, } } @@ -151,6 +176,7 @@ impl FtsQuery { must_not, }, wand_factor: DEFAULT_WAND_FACTOR, + limit: None, include_tail: true, } } @@ -167,9 +193,12 @@ impl FtsQuery { query_type: FtsQueryType::Fuzzy { query: query.into(), fuzziness: None, + prefix_length: 0, max_expansions: DEFAULT_MAX_EXPANSIONS, + boost: 1.0, }, wand_factor: DEFAULT_WAND_FACTOR, + limit: None, include_tail: true, } } @@ -185,9 +214,12 @@ impl FtsQuery { query_type: FtsQueryType::Fuzzy { query: query.into(), fuzziness: Some(fuzziness), + prefix_length: 0, max_expansions: DEFAULT_MAX_EXPANSIONS, + boost: 1.0, }, wand_factor: DEFAULT_WAND_FACTOR, + limit: None, include_tail: true, } } @@ -197,6 +229,7 @@ impl FtsQuery { column: impl Into, query: impl Into, fuzziness: Option, + prefix_length: u32, max_expansions: usize, ) -> Self { Self { @@ -204,9 +237,12 @@ impl FtsQuery { query_type: FtsQueryType::Fuzzy { query: query.into(), fuzziness, + prefix_length, max_expansions, + boost: 1.0, }, wand_factor: DEFAULT_WAND_FACTOR, + limit: None, include_tail: true, } } @@ -221,12 +257,89 @@ impl FtsQuery { self } + pub fn with_limit(mut self, limit: Option) -> Self { + self.limit = limit; + self + } + /// Set whether to search the mutable tail (read-your-writes) or only the /// immutable frozen partitions (the Lucene model). Default `true`. pub fn with_include_tail(mut self, include_tail: bool) -> Self { self.include_tail = include_tail; self } + + fn with_boost(mut self, boost: f32) -> Self { + match &mut self.query_type { + FtsQueryType::Match { boost: b, .. } | FtsQueryType::Fuzzy { boost: b, .. } => { + *b = boost; + } + FtsQueryType::Phrase { .. } | FtsQueryType::Boolean { .. } => {} + } + self + } +} + +/// Convert an index-level [`FullTextSearchQuery`] into the MemTable's local +/// [`FtsQuery`], so the MemTable scanner shares the dataset `Scanner`'s FTS +/// entry type. Supports match (exact `fuzziness == Some(0)` and fuzzy) and +/// phrase leaf queries; the column must be bound on the query. Compound queries +/// (boolean / boost / multi-match) cannot be modeled by the MemTable path and +/// return a `not_supported` error rather than failing deep in planning. +fn local_fts_query(query: FullTextSearchQuery) -> Result { + let wand_factor = query.wand_factor.unwrap_or(DEFAULT_WAND_FACTOR); + let limit = query + .limit + .map(|limit| { + if limit < 0 { + Err(Error::invalid_input( + "full-text search limit must be non-negative".to_string(), + )) + } else { + Ok(limit as usize) + } + }) + .transpose()?; + let require_column = |column: Option| { + column.ok_or_else(|| { + Error::invalid_input( + "full-text search requires a column; set it with \ + `FullTextSearchQuery::with_column`" + .to_string(), + ) + }) + }; + let local = match query.query { + IndexFtsQuery::Match(m) => { + let column = require_column(m.column)?; + match m.fuzziness { + // Some(0) is an exact match in the index model. + Some(0) => FtsQuery::match_query_with_operator(column, m.terms, m.operator) + .with_boost(m.boost), + _ if m.operator != Operator::Or => { + return Err(Error::not_supported( + "MemTable fuzzy full-text search only supports OR match operators" + .to_string(), + )); + } + fuzziness => FtsQuery::fuzzy_with_options( + column, + m.terms, + fuzziness, + m.prefix_length, + m.max_expansions, + ) + .with_boost(m.boost), + } + } + IndexFtsQuery::Phrase(p) => FtsQuery::phrase(require_column(p.column)?, p.terms, p.slop), + other => { + return Err(Error::not_supported(format!( + "MemTable full-text search supports match and phrase queries, got: {other}" + ))); + } + }; + Ok(local.with_wand_factor(wand_factor).with_limit(limit)) } /// Scalar predicate for BTree index queries. @@ -271,11 +384,15 @@ impl ScalarPredicate { /// /// # Example /// +/// The builder methods take `&mut self` (mirroring +/// [`crate::dataset::scanner::Scanner`]), so configure the scanner with +/// statements rather than a fluent chain: +/// /// ```ignore -/// let scanner = MemTableScanner::new(batch_store, indexes, schema) -/// .project(&["id", "name"])? -/// .filter("id > 10")? -/// .limit(100, None)?; +/// let mut scanner = MemTableScanner::new(batch_store, indexes, schema); +/// scanner.project(&["id", "name"])?; +/// scanner.filter("id > 10")?; +/// scanner.limit(Some(100), None)?; /// /// let stream = scanner.try_into_stream().await?; /// ``` @@ -300,6 +417,11 @@ pub struct MemTableScanner { /// Whether to include _rowaddr column in output. /// Same value as _rowid but named for compatibility with LSM scanner. with_row_address: bool, + /// Primary-key columns, supplied by the LSM planner. When set, a filtered + /// vector/FTS search evaluates the predicate against the newest version of + /// each PK only, so an in-memtable update whose current version fails the + /// predicate is excluded rather than leaking a stale older match. + pk_columns: Option>, } impl MemTableScanner { @@ -334,18 +456,34 @@ impl MemTableScanner { batch_size: None, with_row_id: false, with_row_address: false, + pk_columns: None, } } - /// Project only the specified columns. + /// Provide the primary-key columns. When set, a filtered vector/FTS search + /// evaluates the predicate against the newest version of each PK only, + /// preventing a stale older match from leaking past an in-memtable update + /// whose current version fails the predicate. + pub fn with_pk_columns(&mut self, pk_columns: Vec) -> &mut Self { + self.pk_columns = if pk_columns.is_empty() { + None + } else { + Some(pk_columns) + }; + self + } + + /// Project only the specified columns. Mirrors + /// [`crate::dataset::scanner::Scanner::project`]. /// /// Special columns: /// - `_rowid`: Returns the row position (global row offset in MemTable) - pub fn project(&mut self, columns: &[&str]) -> &mut Self { + pub fn project>(&mut self, columns: &[T]) -> Result<&mut Self> { // Check if _rowid is requested in projection let mut filtered_columns = Vec::new(); for col in columns { - if *col == ROW_ID { + let col = col.as_ref(); + if col == ROW_ID { self.with_row_id = true; } else { filtered_columns.push(col.to_string()); @@ -355,7 +493,7 @@ impl MemTableScanner { if !filtered_columns.is_empty() || self.with_row_id { self.projection = Some(filtered_columns); } - self + Ok(self) } /// Include the _rowid column in output. @@ -385,15 +523,7 @@ impl MemTableScanner { /// Set a filter expression using SQL-like syntax. pub fn filter(&mut self, filter_expr: &str) -> Result<&mut Self> { - let ctx = SessionContext::new(); - let df_schema = self - .schema - .clone() - .to_dfschema() - .map_err(|e| Error::invalid_input(format!("Failed to create DFSchema: {}", e)))?; - let expr = ctx.parse_sql_expr(filter_expr, &df_schema).map_err(|e| { - Error::invalid_input(format!("Failed to parse filter expression: {}", e)) - })?; + let expr = parse_filter_expr(self.schema.as_ref(), filter_expr)?; self.filter = Some(expr); Ok(self) } @@ -404,24 +534,50 @@ impl MemTableScanner { self } - /// Limit the number of results. - pub fn limit(&mut self, limit: usize, offset: Option) -> &mut Self { - self.limit = Some(limit); - self.offset = offset; - self + /// Limit the number of results, with an optional offset. Mirrors + /// [`crate::dataset::scanner::Scanner::limit`]: both bounds are `Option` + /// and must be non-negative. + pub fn limit(&mut self, limit: Option, offset: Option) -> Result<&mut Self> { + if let Some(value) = limit + && value < 0 + { + return Err(Error::invalid_input( + "limit must be non-negative".to_string(), + )); + } + if let Some(value) = offset + && value < 0 + { + return Err(Error::invalid_input( + "offset must be non-negative".to_string(), + )); + } + self.limit = limit.map(|value| value as usize); + self.offset = offset.map(|value| value as usize); + Ok(self) } - /// Set up a vector similarity search. + /// Set up a vector similarity search. Mirrors + /// [`crate::dataset::scanner::Scanner::nearest`] — the query vector is passed + /// by reference. /// /// # Arguments /// /// * `column` - The name of the vector column to search. /// * `query` - The query vector. /// * `k` - Number of nearest neighbors to return. - pub fn nearest(&mut self, column: &str, query: Arc, k: usize) -> &mut Self { + pub fn nearest(&mut self, column: &str, query: &dyn Array, k: usize) -> Result<&mut Self> { + if k == 0 { + return Err(Error::invalid_input("k must be positive".to_string())); + } + if query.is_empty() { + return Err(Error::invalid_input( + "query vector must have non-zero length".to_string(), + )); + } self.nearest = Some(VectorQuery { column: column.to_string(), - query_vector: query, + query_vector: query.slice(0, query.len()), k, nprobes: 1, maximum_nprobes: None, @@ -431,7 +587,7 @@ impl MemTableScanner { distance_lower_bound: None, distance_upper_bound: None, }); - self + Ok(self) } /// Set the number of probes for IVF search. @@ -525,10 +681,15 @@ impl MemTableScanner { self } - /// Set up a full-text search with simple term matching. - pub fn full_text_search(&mut self, column: &str, query: &str) -> &mut Self { - self.full_text_query = Some(FtsQuery::match_query(column, query)); - self + /// Set up a full-text search. Mirrors + /// [`crate::dataset::scanner::Scanner::full_text_search`], taking a + /// [`FullTextSearchQuery`] whose column is set via + /// `FullTextSearchQuery::with_column`. Match (exact/fuzzy) and phrase leaf + /// queries are supported; compound queries (boolean/boost/multi-match) are + /// not yet supported by the MemTable path and return an error. + pub fn full_text_search(&mut self, query: FullTextSearchQuery) -> Result<&mut Self> { + self.full_text_query = Some(local_fts_query(query)?); + Ok(self) } /// Set up a full-text phrase search. @@ -615,6 +776,7 @@ impl MemTableScanner { column, query, fuzziness, + 0, max_expansions, )); self @@ -764,6 +926,12 @@ impl MemTableScanner { /// Create the execution plan based on the query configuration. pub async fn create_plan(&self) -> Result> { + if self.nearest.is_some() && self.full_text_query.is_some() { + return Err(Error::invalid_input( + "MemTableScanner cannot combine vector and full-text search".to_string(), + )); + } + // Determine which type of plan to create if let Some(ref vector_query) = self.nearest { return self.plan_vector_search(vector_query).await; @@ -815,12 +983,12 @@ impl MemTableScanner { let mut plan: Arc = Arc::new(scan); - // Apply limit if present - if let Some(limit) = self.limit { + // Apply limit / offset if present. + if self.limit.is_some() || self.offset.unwrap_or(0) > 0 { plan = Arc::new(GlobalLimitExec::new( plan, self.offset.unwrap_or(0), - Some(limit), + self.limit, )); } @@ -913,12 +1081,44 @@ impl MemTableScanner { /// hasn't reached this writer yet (cold-start, or rows written between an /// index commit and the next memtable rotation), KNN must still produce /// correct, distance-bearing results so the LSM-level merge stays sound. + /// Compile the optional logical `filter` into a physical predicate against + /// the memtable schema. Shared by the vector and FTS search arms; mirrors the + /// compilation in [`Self::plan_full_scan`] (`optimize_expr` before + /// `create_physical_expr` for literal type coercion). + fn filter_predicate(&self) -> Result> { + let Some(ref filter) = self.filter else { + return Ok(None); + }; + let planner = Planner::new(self.schema.clone()); + let optimized = planner.optimize_expr(filter.clone())?; + Ok(Some(planner.create_physical_expr(&optimized)?)) + } + async fn plan_vector_search(&self, query: &VectorQuery) -> Result> { let max_visible = self.max_visible_batch_position; let projection_indices = self.compute_projection_indices()?; let base_schema = self.base_output_schema(); + let filter_predicate = self.filter_predicate()?; + if let Some(pk_columns) = &self.pk_columns { + validate_pk_types(&self.schema, pk_columns)?; + } - let exec: Arc = if self.has_vector_index(&query.column) { + // With a prefilter we use brute force rather than HNSW because graph + // traversal cannot honor an arbitrary predicate. With PK rewrites, we + // also need exact newest-before-top-k semantics: a stale near vector + // must not consume an HNSW top-k slot and hide the next live row. Pure + // append-only PK data can still use HNSW safely. This relies on + // `IndexStore` marking PK overrides before advancing the visible batch + // watermark, so any snapshot that sees a rewrite also sees the flag. + let hnsw_safe_with_pk = self + .pk_columns + .as_ref() + .map(|_| self.indexes.has_pk_index() && !self.indexes.pk_has_overrides()) + .unwrap_or(true); + let exec: Arc = if filter_predicate.is_none() + && hnsw_safe_with_pk + && self.has_vector_index(&query.column) + { Arc::new(VectorIndexExec::new( self.batch_store.clone(), self.indexes.clone(), @@ -929,14 +1129,18 @@ impl MemTableScanner { self.with_row_id, )?) } else { - Arc::new(MemTableBruteForceVectorExec::new( - self.batch_store.clone(), - query.clone(), - max_visible, - projection_indices, - base_schema, - self.with_row_id, - )?) + Arc::new( + MemTableBruteForceVectorExec::new( + self.batch_store.clone(), + query.clone(), + max_visible, + projection_indices, + base_schema, + self.with_row_id, + )? + .with_filter(filter_predicate) + .with_pk_columns(self.pk_columns.clone()), + ) }; self.apply_post_index_ops(exec).await } @@ -944,14 +1148,18 @@ impl MemTableScanner { /// Plan a full-text search. /// /// Uses the effective visibility (min of max_visible and max_indexed) to ensure - /// queries only see indexed data. Falls back to full scan if no index exists. + /// queries only see indexed data. async fn plan_fts_search(&self, query: &FtsQuery) -> Result> { if !self.has_fts_index(&query.column) { - return self.plan_full_scan().await; + return self.empty_fts_plan(); } let max_visible = self.max_visible_batch_position; let projection_indices = self.compute_projection_indices()?; + let filter_predicate = self.filter_predicate()?; + if let Some(pk_columns) = &self.pk_columns { + validate_pk_types(&self.schema, pk_columns)?; + } let index_exec = FtsIndexExec::new( self.batch_store.clone(), @@ -961,10 +1169,29 @@ impl MemTableScanner { projection_indices, self.base_output_schema(), self.with_row_id, - )?; + )? + .with_filter(filter_predicate) + .with_pk_columns(self.pk_columns.clone()); self.apply_post_index_ops(Arc::new(index_exec)).await } + fn empty_fts_plan(&self) -> Result> { + use datafusion::physical_plan::empty::EmptyExec; + + let mut fields: Vec = self + .base_output_schema() + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + fields.push(Field::new(SCORE_COLUMN, DataType::Float32, true)); + if self.with_row_id { + fields.push(Field::new(ROW_ID, DataType::UInt64, true)); + } + let schema = Arc::new(arrow_schema::Schema::new(fields)); + Ok(Arc::new(EmptyExec::new(schema))) + } + /// Apply limit and other post-processing operations. async fn apply_post_index_ops( &self, @@ -972,11 +1199,11 @@ impl MemTableScanner { ) -> Result> { let mut result = plan; - if let Some(limit) = self.limit { + if self.limit.is_some() || self.offset.unwrap_or(0) > 0 { result = Arc::new(GlobalLimitExec::new( result, self.offset.unwrap_or(0), - Some(limit), + self.limit, )); } @@ -1026,16 +1253,14 @@ impl MemTableScanner { value: coerced_lit, }); } - datafusion::logical_expr::Operator::Lt - | datafusion::logical_expr::Operator::LtEq => { + datafusion::logical_expr::Operator::Lt => { return Some(ScalarPredicate::Range { column: col.name.clone(), lower: None, upper: Some(coerced_lit), }); } - datafusion::logical_expr::Operator::Gt - | datafusion::logical_expr::Operator::GtEq => { + datafusion::logical_expr::Operator::GtEq => { return Some(ScalarPredicate::Range { column: col.name.clone(), lower: Some(coerced_lit), @@ -1046,7 +1271,7 @@ impl MemTableScanner { } } } - Expr::InList(in_list) => { + Expr::InList(in_list) if !in_list.negated => { if let Expr::Column(col) = in_list.expr.as_ref() { let values: Vec = in_list .list @@ -1108,7 +1333,7 @@ impl MemTableScanner { #[cfg(test)] mod tests { use super::*; - use arrow_array::{Int32Array, StringArray}; + use arrow_array::{BooleanArray, Int32Array, StringArray}; use arrow_schema::{DataType, Field, Schema}; fn create_test_schema() -> SchemaRef { @@ -1213,7 +1438,7 @@ mod tests { let indexes = create_index_store_with_batches(&batch_store, &schema, &[(0, 10)]); let mut scanner = MemTableScanner::new(batch_store, indexes, schema.clone()); - scanner.project(&["id"]); + scanner.project(&["id"]).unwrap(); let result = scanner.try_into_batch().await.unwrap(); assert_eq!(result.num_columns(), 1); @@ -1228,12 +1453,559 @@ mod tests { let indexes = create_index_store_with_batches(&batch_store, &schema, &[(0, 100)]); let mut scanner = MemTableScanner::new(batch_store, indexes, schema.clone()); - scanner.limit(10, None); + scanner.limit(Some(10), None).unwrap(); let result = scanner.try_into_batch().await.unwrap(); assert_eq!(result.num_rows(), 10); } + #[tokio::test] + async fn test_scanner_offset_without_limit() { + let schema = create_test_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(100)); + + let indexes = create_index_store_with_batches(&batch_store, &schema, &[(0, 10)]); + + let mut scanner = MemTableScanner::new(batch_store, indexes, schema.clone()); + scanner.limit(Some(3), None).unwrap(); + scanner.limit(None, Some(2)).unwrap(); + + let result = scanner.try_into_batch().await.unwrap(); + let ids = result + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!(ids, vec![2, 3, 4, 5, 6, 7, 8, 9]); + } + + #[tokio::test] + async fn btree_filter_fallback_preserves_non_representable_predicates() { + let schema = create_test_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(100)); + let indexes = create_index_store_with_batches(&batch_store, &schema, &[(0, 10)]); + + async fn ids_for( + batch_store: Arc, + indexes: Arc, + schema: SchemaRef, + filter: &str, + ) -> Vec { + let mut scanner = MemTableScanner::new(batch_store, indexes, schema); + scanner.filter(filter).unwrap(); + scanner + .try_into_batch() + .await + .unwrap() + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + } + + assert_eq!( + ids_for( + batch_store.clone(), + indexes.clone(), + schema.clone(), + "id NOT IN (1, 2)" + ) + .await, + vec![0, 3, 4, 5, 6, 7, 8, 9] + ); + assert_eq!( + ids_for( + batch_store.clone(), + indexes.clone(), + schema.clone(), + "id <= 5" + ) + .await, + vec![0, 1, 2, 3, 4, 5] + ); + assert_eq!( + ids_for(batch_store, indexes, schema, "id > 5").await, + vec![6, 7, 8, 9] + ); + } + + /// `full_text_search` now takes a structured `FullTextSearchQuery` (matching + /// the dataset `Scanner`); `local_fts_query` maps the supported leaf shapes + /// and rejects compound queries and missing columns. + #[test] + fn local_fts_query_maps_leaf_shapes_and_rejects_the_rest() { + use lance_index::scalar::inverted::query::{ + BooleanQuery, MatchQuery, Occur, Operator, PhraseQuery, + }; + + // Exact match (default fuzziness Some(0)) -> local Match, preserving the + // old `full_text_search(col, terms)` behavior. + let q = FullTextSearchQuery::new("hello".to_string()) + .with_column("text".to_string()) + .unwrap(); + let local = local_fts_query(q).unwrap(); + assert_eq!(local.column, "text"); + assert!( + matches!(local.query_type, FtsQueryType::Match { query, operator, .. } + if query == "hello" && operator == Operator::Or) + ); + + let exact_and = FullTextSearchQuery::new_query(IndexFtsQuery::Match( + MatchQuery::new("hello world".to_string()) + .with_operator(Operator::And) + .with_boost(3.0) + .with_column(Some("text".to_string())), + )); + let local = local_fts_query(exact_and).unwrap(); + assert!( + matches!(local.query_type, FtsQueryType::Match { query, operator, boost } + if query == "hello world" && operator == Operator::And && boost == 3.0) + ); + + // Fuzzy match -> local Fuzzy carrying edit distance, prefix length, and boost. + let fuzzy = FullTextSearchQuery::new_query(IndexFtsQuery::Match( + MatchQuery::new("lance".to_string()) + .with_fuzziness(Some(2)) + .with_prefix_length(2) + .with_boost(2.5) + .with_column(Some("text".to_string())), + )); + let local = local_fts_query(fuzzy).unwrap(); + assert!( + matches!(local.query_type, FtsQueryType::Fuzzy { fuzziness, prefix_length, boost, .. } + if fuzziness == Some(2) && prefix_length == 2 && boost == 2.5) + ); + + let fuzzy_and = FullTextSearchQuery::new_query(IndexFtsQuery::Match( + MatchQuery::new("lance memwal".to_string()) + .with_operator(Operator::And) + .with_fuzziness(Some(1)) + .with_column(Some("text".to_string())), + )); + assert!( + local_fts_query(fuzzy_and).is_err(), + "fuzzy AND cannot be represented by the local memtable query" + ); + + // Phrase -> local Phrase. + let phrase = FullTextSearchQuery::new_query(IndexFtsQuery::Phrase( + PhraseQuery::new("quick fox".to_string()).with_column(Some("text".to_string())), + )); + let local = local_fts_query(phrase).unwrap(); + assert!(matches!(local.query_type, FtsQueryType::Phrase { .. })); + + // Compound (boolean) -> not supported. + let boolean = + FullTextSearchQuery::new_query(IndexFtsQuery::Boolean(BooleanQuery::new(vec![( + Occur::Must, + IndexFtsQuery::Match( + MatchQuery::new("x".to_string()).with_column(Some("text".to_string())), + ), + )]))); + assert!( + local_fts_query(boolean).is_err(), + "boolean must be rejected" + ); + + // Missing column -> error. + let no_col = FullTextSearchQuery::new("hi".to_string()); + assert!( + local_fts_query(no_col).is_err(), + "missing column must error" + ); + } + + #[tokio::test] + async fn full_text_search_honors_query_limit() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8, true), + ])); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec![ + "lance", + "lance filler", + "lance filler filler", + ])), + ], + ) + .unwrap(); + let mut indexes = IndexStore::new(); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + + let mut scanner = MemTableScanner::new(batch_store, Arc::new(indexes), schema); + scanner + .full_text_search( + FullTextSearchQuery::new("lance".to_string()) + .with_column("text".to_string()) + .unwrap() + .limit(Some(1)), + ) + .unwrap(); + + let result = scanner.try_into_batch().await.unwrap(); + assert_eq!( + result.num_rows(), + 1, + "query-level FTS limit must cap direct MemTableScanner results" + ); + } + + #[tokio::test] + async fn full_text_search_without_index_returns_empty_score_schema() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8, true), + ])); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec!["needle", "needle"])), + ], + ) + .unwrap(); + batch_store.append(batch).unwrap(); + + let mut scanner = MemTableScanner::new(batch_store, Arc::new(IndexStore::new()), schema); + scanner + .full_text_search( + FullTextSearchQuery::new("needle".to_string()) + .with_column("text".to_string()) + .unwrap(), + ) + .unwrap(); + + let result = scanner.try_into_batch().await.unwrap(); + assert_eq!(result.num_rows(), 0); + assert!( + result.schema().field_with_name("_score").is_ok(), + "missing FTS indexes should produce an empty FTS-shaped result" + ); + } + + #[tokio::test] + async fn full_text_search_prefilter_null_predicate_excludes_rows() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8, true), + Field::new("active", DataType::Boolean, true), + ])); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["needle", "needle", "needle"])), + Arc::new(BooleanArray::from(vec![None, Some(true), Some(false)])), + ], + ) + .unwrap(); + let mut indexes = IndexStore::new(); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + + let mut scanner = MemTableScanner::new(batch_store, Arc::new(indexes), schema); + scanner.filter("active = true").unwrap(); + scanner + .full_text_search( + FullTextSearchQuery::new("needle".to_string()) + .with_column("text".to_string()) + .unwrap(), + ) + .unwrap(); + + let result = scanner.try_into_batch().await.unwrap(); + let ids = result + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + ids, + vec![2], + "NULL predicate results must be excluded from FTS prefilter candidates" + ); + } + + #[tokio::test] + async fn full_text_search_prefilter_disables_wand_pruning() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8, true), + Field::new("active", DataType::Boolean, true), + ])); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec!["alpha beta gamma delta", "alpha"])), + Arc::new(BooleanArray::from(vec![Some(false), Some(true)])), + ], + ) + .unwrap(); + let mut indexes = IndexStore::new(); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + + let mut scanner = MemTableScanner::new(batch_store, Arc::new(indexes), schema); + scanner.filter("active = true").unwrap(); + scanner + .full_text_search( + FullTextSearchQuery::new("alpha beta gamma delta".to_string()) + .with_column("text".to_string()) + .unwrap() + .wand_factor(Some(0.99)), + ) + .unwrap(); + + let result = scanner.try_into_batch().await.unwrap(); + let ids = result + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + ids, + vec![2], + "filtered FTS must not let WAND prune rows before the prefilter is applied" + ); + } + + #[tokio::test] + async fn full_text_search_append_only_pk_keeps_wand_pruning() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8, true), + ])); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec!["alpha beta gamma delta", "alpha"])), + ], + ) + .unwrap(); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + + let mut scanner = MemTableScanner::new(batch_store, Arc::new(indexes), schema); + scanner.with_pk_columns(vec!["id".to_string()]); + scanner + .full_text_search( + FullTextSearchQuery::new("alpha beta gamma delta".to_string()) + .with_column("text".to_string()) + .unwrap() + .wand_factor(Some(0.99)), + ) + .unwrap(); + + let result = scanner.try_into_batch().await.unwrap(); + let ids = result + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + ids, + vec![1], + "append-only PK data should keep index WAND pruning enabled" + ); + } + + #[tokio::test] + async fn full_text_search_with_pk_rewrite_disables_index_limit_pushdown() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8, true), + ])); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 3])), + Arc::new(StringArray::from(vec![ + "alpha beta gamma delta epsilon", + "other", + "alpha beta gamma delta", + "alpha", + ])), + ], + ) + .unwrap(); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + + let mut scanner = MemTableScanner::new(batch_store, Arc::new(indexes), schema); + scanner.with_pk_columns(vec!["id".to_string()]); + scanner + .full_text_search( + FullTextSearchQuery::new("alpha beta gamma delta epsilon".to_string()) + .with_column("text".to_string()) + .unwrap() + .limit(Some(2)), + ) + .unwrap(); + + let result = scanner.try_into_batch().await.unwrap(); + let ids = result + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + ids, + vec![2, 3], + "FTS-only PK rewrites must disable index limit pushdown so live lower-scoring PKs can backfill" + ); + } + + #[tokio::test] + async fn full_text_search_with_pk_columns_drops_stale_filtered_hits() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8, true), + Field::new("active", DataType::Boolean, false), + ])); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(StringArray::from(vec!["needle", "needle"])), + Arc::new(BooleanArray::from(vec![true, false])), + ], + ) + .unwrap(); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + + let mut scanner = MemTableScanner::new(batch_store, Arc::new(indexes), schema); + scanner.with_pk_columns(vec!["id".to_string()]); + scanner.filter("active = true").unwrap(); + scanner + .full_text_search( + FullTextSearchQuery::new("needle".to_string()) + .with_column("text".to_string()) + .unwrap(), + ) + .unwrap(); + + let result = scanner.try_into_batch().await.unwrap(); + assert_eq!( + result.num_rows(), + 0, + "the older matching version must not leak when the newest PK fails the filter" + ); + } + + #[tokio::test] + async fn full_text_search_with_pk_columns_falls_back_without_pk_index() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8, true), + ])); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 2])), + Arc::new(StringArray::from(vec![ + "needle stale", + "other", + "other", + "needle fresh", + ])), + ], + ) + .unwrap(); + let mut indexes = IndexStore::new(); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + + let mut scanner = MemTableScanner::new(batch_store, Arc::new(indexes), schema); + scanner.with_pk_columns(vec!["id".to_string()]); + scanner + .full_text_search( + FullTextSearchQuery::new("needle".to_string()) + .with_column("text".to_string()) + .unwrap(), + ) + .unwrap(); + + let result = scanner + .try_into_batch() + .await + .expect("FTS PK recency should fall back without a PK index"); + let ids = result + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!( + ids, + vec![2], + "without a PK index the batch-scan fallback must drop stale id=1 \ + but keep id=2 whose newest version still matches" + ); + } + #[tokio::test] async fn test_scanner_count_rows() { let schema = create_test_schema(); @@ -1290,7 +2062,7 @@ mod tests { let mut scanner = MemTableScanner::new(batch_store, indexes, schema.clone()); // Project only "id" and "_rowid" - scanner.project(&["id", "_rowid"]); + scanner.project(&["id", "_rowid"]).unwrap(); // Verify output schema let output_schema = scanner.output_schema(); @@ -1360,7 +2132,7 @@ mod tests { let mut scanner = MemTableScanner::new(batch_store, indexes, schema); // Project with _rowid should set with_row_id flag - scanner.project(&["id", "_rowid"]); + scanner.project(&["id", "_rowid"]).unwrap(); // with_row_id should be true now assert!(scanner.with_row_id); @@ -1408,7 +2180,7 @@ mod tests { let indexes = create_index_store_with_batches(&batch_store, &schema, &[(0, 10)]); let mut scanner = MemTableScanner::new(batch_store, indexes, schema.clone()); - scanner.project(&["id", "_rowid"]); + scanner.project(&["id", "_rowid"]).unwrap(); let plan = scanner.create_plan().await.unwrap(); @@ -1587,7 +2359,7 @@ mod tests { let mut scanner = MemTableScanner::new(batch_store, indexes, schema.clone()); let query: Arc = Arc::new(arrow_array::Float32Array::from(vec![0.0_f32, 0.0_f32])); - scanner.nearest("vector", query, 5); + scanner.nearest("vector", query.as_ref(), 5).unwrap(); let plan = scanner .create_plan() @@ -1600,4 +2372,106 @@ mod tests { out_schema ); } + + #[tokio::test] + async fn test_nearest_rejects_invalid_query_shape() { + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2), + true, + ), + ])); + let batch_store = Arc::new(BatchStore::with_capacity(4)); + let indexes = Arc::new(IndexStore::new()); + + let mut scanner = + MemTableScanner::new(batch_store.clone(), indexes.clone(), schema.clone()); + let query: Arc = + Arc::new(arrow_array::Float32Array::from(vec![0.0_f32, 0.0_f32])); + let Err(err) = scanner.nearest("vector", query.as_ref(), 0) else { + panic!("zero-k vector search should fail"); + }; + assert!( + err.to_string().contains("k must be positive"), + "unexpected zero-k error: {err}" + ); + + let mut scanner = MemTableScanner::new(batch_store, indexes, schema); + let empty_query: Arc = + Arc::new(arrow_array::Float32Array::from(Vec::::new())); + let Err(err) = scanner.nearest("vector", empty_query.as_ref(), 5) else { + panic!("empty vector search should fail"); + }; + assert!( + err.to_string().contains("non-zero length"), + "unexpected empty-query error: {err}" + ); + } + + #[tokio::test] + async fn test_create_plan_rejects_vector_and_fts_combination() { + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8, true), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2), + true, + ), + ])); + let batch_store = Arc::new(BatchStore::with_capacity(4)); + let indexes = Arc::new(IndexStore::new()); + + let mut scanner = MemTableScanner::new(batch_store, indexes, schema); + let query: Arc = + Arc::new(arrow_array::Float32Array::from(vec![0.0_f32, 0.0_f32])); + scanner.nearest("vector", query.as_ref(), 5).unwrap(); + scanner + .full_text_search( + FullTextSearchQuery::new("needle".to_string()) + .with_column("text".to_string()) + .unwrap(), + ) + .unwrap(); + + let err = scanner + .create_plan() + .await + .expect_err("vector and FTS search must not be silently combined"); + assert!( + err.to_string().contains("vector and full-text search"), + "unexpected combined-search error: {err}" + ); + } + + #[tokio::test] + async fn test_plan_vector_search_validates_pk_types() { + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Float64, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2), + true, + ), + ])); + let batch_store = Arc::new(BatchStore::with_capacity(4)); + let indexes = Arc::new(IndexStore::new()); + + let mut scanner = MemTableScanner::new(batch_store, indexes, schema); + scanner.with_pk_columns(vec!["id".to_string()]); + let query: Arc = + Arc::new(arrow_array::Float32Array::from(vec![0.0_f32, 0.0_f32])); + scanner.nearest("vector", query.as_ref(), 5).unwrap(); + + let err = scanner + .create_plan() + .await + .expect_err("unsupported vector PK type must be rejected"); + assert!( + err.to_string().contains("unsupported type Float64"), + "unexpected error: {err}" + ); + } } diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec.rs index e6d79ba6b48..2a593a0f215 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec.rs @@ -10,6 +10,12 @@ //! - `MemTableBruteForceVectorExec` - KNN over the active memtable without an HNSW //! - `FtsIndexExec` - Full-text search +use std::collections::{HashMap, HashSet}; + +use arrow_array::RecordBatch; +use datafusion::common::ScalarValue; +use datafusion::error::Result as DataFusionResult; + mod brute_force_vector; mod btree; mod dedup_scan; @@ -17,9 +23,54 @@ mod fts; mod scan; mod vector; +use crate::dataset::mem_wal::scanner::exec::resolve_pk_indices; +use crate::dataset::mem_wal::write::BatchStore; + pub use brute_force_vector::MemTableBruteForceVectorExec; pub use btree::BTreeIndexExec; pub use dedup_scan::MemTableDedupScanExec; -pub use fts::FtsIndexExec; +pub use fts::{FtsIndexExec, SCORE_COLUMN}; pub use scan::{MemTableScanExec, ROW_ADDRESS_COLUMN}; pub use vector::VectorIndexExec; + +pub(super) fn newest_pk_positions( + batch_store: &BatchStore, + pk_columns: &[String], + max_visible_batch_position: usize, + max_visible_row: u64, +) -> DataFusionResult> { + let mut newest: HashMap, u64> = HashMap::new(); + let mut current_row: u64 = 0; + for (batch_position, stored_batch) in batch_store.iter().enumerate() { + let n = stored_batch.num_rows; + if n == 0 { + continue; + } + if batch_position > max_visible_batch_position { + current_row += n as u64; + continue; + } + let pk_indices = resolve_pk_indices(&stored_batch.data, pk_columns)?; + for row in 0..n { + let pos = current_row + row as u64; + if pos > max_visible_row { + break; + } + let key = pk_key(&stored_batch.data, &pk_indices, row)?; + newest.insert(key, pos); + } + current_row += n as u64; + } + Ok(newest.into_values().collect()) +} + +fn pk_key( + batch: &RecordBatch, + pk_indices: &[usize], + row: usize, +) -> DataFusionResult> { + pk_indices + .iter() + .map(|&col_idx| ScalarValue::try_from_array(batch.column(col_idx).as_ref(), row)) + .collect() +} diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/brute_force_vector.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/brute_force_vector.rs index 43169f92d2d..8a239605b50 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/brute_force_vector.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/brute_force_vector.rs @@ -15,7 +15,7 @@ use std::any::Any; use std::fmt::{Debug, Formatter}; use std::sync::Arc; -use arrow_array::{Array, Float32Array, RecordBatch, UInt64Array, cast::AsArray}; +use arrow_array::{Array, BooleanArray, Float32Array, RecordBatch, UInt64Array, cast::AsArray}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion::common::stats::Precision; use datafusion::error::Result as DataFusionResult; @@ -27,12 +27,13 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, }; -use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExprRef}; use futures::stream::{self, StreamExt}; use lance_core::{Error, Result}; use lance_linalg::distance::DistanceType; use super::super::builder::VectorQuery; +use super::newest_pk_positions; use super::vector::DISTANCE_COLUMN; use crate::dataset::mem_wal::write::BatchStore; @@ -53,6 +54,14 @@ pub struct MemTableBruteForceVectorExec { properties: Arc, metrics: ExecutionPlanMetricsSet, with_row_id: bool, + /// Optional prefilter predicate, compiled against the memtable schema. + /// Applied per row before the top-k cut so the KNN only ranks matching + /// rows (true prefilter, not a lossy post-filter on the top-k). + filter: Option, + /// Primary-key columns. When set, only the newest version of each PK is + /// eligible for top-k. With a filter, this evaluates the predicate against + /// the current PK version instead of falling back to a stale older version. + pk_columns: Option>, } impl Debug for MemTableBruteForceVectorExec { @@ -108,9 +117,49 @@ impl MemTableBruteForceVectorExec { properties, metrics: ExecutionPlanMetricsSet::new(), with_row_id, + filter: None, + pk_columns: None, }) } + /// Attach an optional prefilter predicate (compiled against the memtable + /// schema). Rows that fail the predicate are excluded before the top-k cut. + pub fn with_filter(mut self, filter: Option) -> Self { + self.filter = filter; + self + } + + /// Provide the primary-key columns so search keeps only the newest version + /// of each PK (see `pk_columns`). + pub fn with_pk_columns(mut self, pk_columns: Option>) -> Self { + self.pk_columns = pk_columns.filter(|columns| !columns.is_empty()); + self + } + + /// Evaluate the prefilter predicate against a memtable batch, returning a + /// keep-mask (`true` = retain). `Ok(None)` when no filter is configured. + fn filter_mask(&self, batch: &RecordBatch) -> Result> { + let Some(ref predicate) = self.filter else { + return Ok(None); + }; + let values = predicate + .evaluate(batch) + .and_then(|v| v.into_array(batch.num_rows())) + .map_err(|e| { + Error::invalid_input(format!("vector prefilter evaluation failed: {}", e)) + })?; + let mask = values + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::invalid_input( + "vector prefilter predicate did not evaluate to boolean".to_string(), + ) + })? + .clone(); + Ok(Some(mask)) + } + /// Last row position visible under `max_visible_batch_position`, or `None` /// if no batches are visible. Identical to `VectorIndexExec`'s helper so /// both arms cut at the same MVCC boundary. @@ -167,6 +216,24 @@ impl MemTableBruteForceVectorExec { let distance_type = self.query.distance_type.unwrap_or(DEFAULT_DISTANCE_TYPE); let batch_func = distance_type.arrow_batch_func(); + // When PK columns are configured, only the newest version of each PK is + // eligible. This keeps top-k slots from being consumed by superseded + // rows and makes filtered search evaluate the predicate against the + // current version of the PK. + let newest_positions = if let Some(pk_columns) = &self.pk_columns { + Some( + newest_pk_positions( + &self.batch_store, + pk_columns, + self.max_visible_batch_position, + max_visible_row, + ) + .map_err(|e| Error::invalid_input(e.to_string()))?, + ) + } else { + None + }; + // Walk batches in append order. `current_row` is the global row offset // of the *next* row about to be visited; rows past `max_visible_row` // are dropped before they reach the heap. @@ -207,11 +274,27 @@ impl MemTableBruteForceVectorExec { )) })?; + // Prefilter: drop rows that fail the predicate before they reach the + // top-k heap (a NULL predicate result excludes the row, matching SQL). + let filter_mask = self.filter_mask(&stored_batch.data)?; + for row in 0..n { let pos = current_row + row as u64; if pos > max_visible_row { break; } + // Skip superseded versions: only the newest version of each PK is + // eligible, so a newer non-matching version excludes the PK. + if let Some(ref newest) = newest_positions + && !newest.contains(&pos) + { + continue; + } + if let Some(ref mask) = filter_mask + && (!mask.is_valid(row) || !mask.value(row)) + { + continue; + } if distances.is_null(row) { continue; } @@ -407,10 +490,11 @@ impl ExecutionPlan for MemTableBruteForceVectorExec { #[cfg(test)] mod tests { use super::*; - use arrow_array::{FixedSizeListArray, Float32Array, Int32Array}; + use arrow_array::{BooleanArray, FixedSizeListArray, Float32Array, Int32Array}; use arrow_schema::{DataType, Field, Schema}; use datafusion::physical_plan::common::collect; - use datafusion::prelude::SessionContext; + use datafusion::prelude::{Expr, SessionContext, col, lit}; + use lance_datafusion::planner::Planner; fn make_schema() -> SchemaRef { Arc::new(Schema::new(vec![ @@ -423,6 +507,18 @@ mod tests { ])) } + fn make_schema_with_active() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2), + true, + ), + Field::new("active", DataType::Boolean, true), + ])) + } + fn make_batch(schema: SchemaRef, ids: &[i32], vectors: &[[f32; 2]]) -> RecordBatch { let id_array = Arc::new(Int32Array::from(ids.to_vec())) as Arc; let values: Vec = vectors.iter().flat_map(|v| v.iter().copied()).collect(); @@ -434,6 +530,23 @@ mod tests { RecordBatch::try_new(schema, vec![id_array, vec_array]).expect("build batch") } + fn make_batch_with_active( + schema: SchemaRef, + ids: &[i32], + vectors: &[[f32; 2]], + active: &[Option], + ) -> RecordBatch { + let id_array = Arc::new(Int32Array::from(ids.to_vec())) as Arc; + let values: Vec = vectors.iter().flat_map(|v| v.iter().copied()).collect(); + let inner = Arc::new(Float32Array::from(values)); + let field = Arc::new(Field::new("item", DataType::Float32, true)); + let vec_array = + Arc::new(FixedSizeListArray::try_new(field, 2, inner, None).expect("build fsl")) + as Arc; + let active_array = Arc::new(BooleanArray::from(active.to_vec())) as Arc; + RecordBatch::try_new(schema, vec![id_array, vec_array, active_array]).expect("build batch") + } + fn store_with_batches(batches: Vec) -> Arc { let store = Arc::new(BatchStore::with_capacity(batches.len().max(1))); for batch in batches { @@ -458,12 +571,34 @@ mod tests { } } + fn physical_filter(schema: SchemaRef, expr: Expr) -> PhysicalExprRef { + let planner = Planner::new(schema); + let optimized = planner.optimize_expr(expr).expect("optimize filter"); + planner + .create_physical_expr(&optimized) + .expect("create physical filter") + } + async fn execute_to_batches(exec: Arc) -> Vec { let ctx = SessionContext::new(); let stream = exec.execute(0, ctx.task_ctx()).expect("execute"); collect(stream).await.expect("collect") } + fn ids_from_batches(batches: &[RecordBatch]) -> Vec { + let mut ids = Vec::new(); + for batch in batches { + let id_arr = batch + .column_by_name("id") + .unwrap() + .as_primitive::(); + for row in 0..batch.num_rows() { + ids.push(id_arr.value(row)); + } + } + ids + } + #[tokio::test] async fn top_k_by_distance() { // Five rows; query at (0,0); L2 distances are id² each — expect ids in @@ -637,4 +772,121 @@ mod tests { pairs.sort_by_key(|(id, _)| *id); assert_eq!(pairs, vec![(10, 0), (11, 1), (12, 2)]); } + + #[tokio::test] + async fn prefilter_null_predicate_excludes_rows() { + let schema = make_schema_with_active(); + let batch = make_batch_with_active( + schema.clone(), + &[1, 2, 3], + &[[0.0, 0.0], [3.0, 0.0], [1.0, 0.0]], + &[None, Some(true), Some(false)], + ); + let store = store_with_batches(vec![batch]); + let query = query_for([0.0, 0.0], 3); + let filter = physical_filter(schema.clone(), col("active").eq(lit(true))); + let exec = Arc::new( + MemTableBruteForceVectorExec::new(store, query, usize::MAX, None, schema, false) + .expect("ctor") + .with_filter(Some(filter)), + ); + + let out = execute_to_batches(exec).await; + assert_eq!( + ids_from_batches(&out), + vec![2], + "NULL predicate results must be excluded from vector prefilter candidates" + ); + } + + #[tokio::test] + async fn prefilter_with_pk_columns_drops_stale_matching_version() { + let schema = make_schema_with_active(); + let batch = make_batch_with_active( + schema.clone(), + &[5, 5], + &[[0.0, 0.0], [10.0, 0.0]], + &[Some(true), Some(false)], + ); + let store = store_with_batches(vec![batch]); + let query = query_for([0.0, 0.0], 10); + let filter = physical_filter(schema.clone(), col("active").eq(lit(true))); + let exec = Arc::new( + MemTableBruteForceVectorExec::new(store, query, usize::MAX, None, schema, false) + .expect("ctor") + .with_filter(Some(filter)) + .with_pk_columns(Some(vec!["id".to_string()])), + ); + + let out = execute_to_batches(exec).await; + assert!( + out.iter().all(|batch| batch.num_rows() == 0), + "the older matching vector version must not leak when the newest PK fails the filter" + ); + } + + #[tokio::test] + async fn pk_columns_keep_newest_version_without_filter() { + let schema = make_schema(); + let batch = make_batch(schema.clone(), &[5, 5], &[[0.0, 0.0], [10.0, 0.0]]); + let store = store_with_batches(vec![batch]); + let query = query_for([0.0, 0.0], 10); + let exec = Arc::new( + MemTableBruteForceVectorExec::new( + store, + query, + usize::MAX, + None, + schema, + /* with_row_id = */ true, + ) + .expect("ctor") + .with_pk_columns(Some(vec!["id".to_string()])), + ); + + let out = execute_to_batches(exec).await; + let row_ids: Vec = out + .iter() + .flat_map(|batch| { + let row_ids = batch + .column_by_name(lance_core::ROW_ID) + .unwrap() + .as_primitive::(); + (0..batch.num_rows()) + .map(|row| row_ids.value(row)) + .collect::>() + }) + .collect(); + assert_eq!( + row_ids, + vec![1], + "brute-force vector PK recency must keep only the newest duplicate" + ); + } + + #[tokio::test] + async fn empty_pk_columns_do_not_collapse_results() { + let schema = make_schema(); + let batch = make_batch( + schema.clone(), + &[1, 2, 3], + &[[0.0, 0.0], [1.0, 0.0], [2.0, 0.0]], + ); + let store = store_with_batches(vec![batch]); + let query = query_for([0.0, 0.0], 3); + let exec = Arc::new( + MemTableBruteForceVectorExec::new(store, query, usize::MAX, None, schema, false) + .expect("ctor") + .with_pk_columns(Some(vec![])), + ); + + let out = execute_to_batches(exec).await; + let mut ids = ids_from_batches(&out); + ids.sort_unstable(); + assert_eq!( + ids, + vec![1, 2, 3], + "empty PK columns should behave like no PK columns, not one empty tuple key" + ); + } } diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/fts.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/fts.rs index 44b1a097d1a..364261c3276 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/fts.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/fts.rs @@ -7,8 +7,9 @@ use std::any::Any; use std::fmt::{Debug, Formatter}; use std::sync::Arc; -use arrow_array::{Float32Array, RecordBatch, UInt32Array, UInt64Array}; +use arrow_array::{BooleanArray, Float32Array, RecordBatch, UInt32Array, UInt64Array}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use datafusion::common::ScalarValue; use datafusion::common::stats::Precision; use datafusion::error::Result as DataFusionResult; use datafusion::execution::TaskContext; @@ -19,12 +20,14 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, }; -use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExprRef}; use futures::stream::{self, StreamExt}; use lance_core::{Error, Result}; use super::super::builder::{FtsQuery, FtsQueryType}; +use super::newest_pk_positions; use crate::dataset::mem_wal::index::{FtsQueryExpr, SearchOptions}; +use crate::dataset::mem_wal::scanner::exec::resolve_pk_indices; use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; /// Score column name in output. @@ -38,6 +41,8 @@ struct BatchRange { batch_id: usize, } +type MaterializedFtsRows = (Vec>, Vec, Vec); + /// ExecutionPlan node that queries FTS index with MVCC visibility. pub struct FtsIndexExec { batch_store: Arc, @@ -54,6 +59,13 @@ pub struct FtsIndexExec { max_visible_row: Option, /// Whether to include _rowid column (row position) in output. with_row_id: bool, + /// Optional prefilter predicate, compiled against the memtable schema. + /// Applied to the materialized full-schema hits before projection so the + /// FTS arm only returns rows matching the predicate. + filter: Option, + /// Primary-key columns. When set, materialized hits are kept only if their + /// row position is the newest visible version of that PK. + pk_columns: Option>, } impl Debug for FtsIndexExec { @@ -162,9 +174,24 @@ impl FtsIndexExec { batch_ranges, max_visible_row, with_row_id, + filter: None, + pk_columns: None, }) } + /// Attach an optional prefilter predicate (compiled against the memtable + /// schema). Hits that fail the predicate are dropped before projection. + pub fn with_filter(mut self, filter: Option) -> Self { + self.filter = filter; + self + } + + /// Provide primary-key columns for newest-version filtering. + pub fn with_pk_columns(mut self, pk_columns: Option>) -> Self { + self.pk_columns = pk_columns; + self + } + /// Find batch for a row position using binary search. O(log n). #[inline] fn find_batch(&self, row_pos: usize) -> Option<&BatchRange> { @@ -183,7 +210,11 @@ impl FtsIndexExec { // Convert FtsQueryType to FtsQueryExpr let query_expr = match &self.query.query_type { - FtsQueryType::Match { query } => FtsQueryExpr::match_query(query), + FtsQueryType::Match { + query, + operator, + boost, + } => FtsQueryExpr::match_query_with_operator(query, *operator).with_boost(*boost), FtsQueryType::Phrase { query, slop } => FtsQueryExpr::phrase_with_slop(query, *slop), FtsQueryType::Boolean { must, @@ -205,15 +236,33 @@ impl FtsIndexExec { FtsQueryType::Fuzzy { query, fuzziness, + prefix_length, max_expansions, - } => FtsQueryExpr::fuzzy_with_options(query, *fuzziness, *max_expansions), + boost, + } => { + FtsQueryExpr::fuzzy_with_options(query, *fuzziness, *prefix_length, *max_expansions) + .with_boost(*boost) + } }; - // Search the index using the query expression. `include_tail` selects - // read-your-writes vs immutable-only; `wand_factor` adds pruning. - let options = SearchOptions::new() - .with_wand_factor(self.query.wand_factor) - .with_include_tail(self.query.include_tail); + let all_rows_visible = self.batch_ranges.last().is_none_or(|last| { + self.max_visible_row + .map(|max_visible| max_visible + 1 >= last.end as u64) + .unwrap_or(last.end == 0) + }); + let pk_recency_is_noop = self.pk_columns.is_none() + || (self.indexes.has_pk_index() && !self.indexes.pk_has_overrides()); + let can_prune_in_index = self.filter.is_none() && pk_recency_is_noop && all_rows_visible; + + // Search the index using the query expression. WAND pruning is only + // safe when the index search itself sees the final candidate set. + let mut options = SearchOptions::new().with_include_tail(self.query.include_tail); + if can_prune_in_index { + options = options.with_wand_factor(self.query.wand_factor); + if let Some(limit) = self.query.limit { + options = options.with_limit(limit); + } + } let entries = index.search_with_options(&query_expr, options); // Convert to (row_position, score) pairs @@ -298,6 +347,66 @@ impl FtsIndexExec { } } + // Prefilter: evaluate the predicate against the full-schema hits and drop + // non-matching rows before applying the query limit and projection (a + // NULL result excludes the row, matching SQL). When a predicate exists, + // query_index deliberately avoids pushing the limit into the index so + // this remains an exact prefilter, not a lossy post-filter. + let (final_columns, all_scores, all_row_positions) = + if let Some(ref predicate) = self.filter { + let Some(first) = self.batch_store.get(0) else { + return Ok(vec![]); + }; + let data_batch = RecordBatch::try_new(first.data.schema(), final_columns)?; + let mask = predicate + .evaluate(&data_batch)? + .into_array(data_batch.num_rows())?; + let mask = mask + .as_any() + .downcast_ref::() + .ok_or_else(|| { + datafusion::error::DataFusionError::Internal( + "FTS prefilter predicate did not evaluate to boolean".to_string(), + ) + })?; + let filtered_columns = data_batch + .columns() + .iter() + .map(|c| arrow_select::filter::filter(c.as_ref(), mask)) + .collect::, _>>()?; + let filtered_scores: Vec = all_scores + .iter() + .zip(mask.iter()) + .filter_map(|(s, keep)| keep.unwrap_or(false).then_some(*s)) + .collect(); + let filtered_positions: Vec = all_row_positions + .iter() + .zip(mask.iter()) + .filter_map(|(p, keep)| keep.unwrap_or(false).then_some(*p)) + .collect(); + (filtered_columns, filtered_scores, filtered_positions) + } else { + (final_columns, all_scores, all_row_positions) + }; + + let (mut final_columns, mut all_scores, mut all_row_positions) = + self.filter_to_newest_pk(final_columns, all_scores, all_row_positions)?; + + if all_scores.is_empty() { + return Ok(vec![]); + } + + if let Some(limit) = self.query.limit + && all_scores.len() > limit + { + final_columns = final_columns + .into_iter() + .map(|column| column.slice(0, limit)) + .collect(); + all_scores.truncate(limit); + all_row_positions.truncate(limit); + } + // Add score column final_columns.push(Arc::new(Float32Array::from(all_scores))); @@ -322,6 +431,76 @@ impl FtsIndexExec { let batch = RecordBatch::try_new(self.output_schema.clone(), projected_columns)?; Ok(vec![batch]) } + + fn filter_to_newest_pk( + &self, + final_columns: Vec>, + all_scores: Vec, + all_row_positions: Vec, + ) -> DataFusionResult { + let Some(pk_columns) = &self.pk_columns else { + return Ok((final_columns, all_scores, all_row_positions)); + }; + if pk_columns.is_empty() || all_scores.is_empty() { + return Ok((final_columns, all_scores, all_row_positions)); + } + let Some(max_visible_row) = self.max_visible_row else { + return Ok((final_columns, all_scores, all_row_positions)); + }; + if self.indexes.has_pk_index() && !self.indexes.pk_has_overrides() { + return Ok((final_columns, all_scores, all_row_positions)); + } + let Some(first) = self.batch_store.get(0) else { + return Ok((final_columns, all_scores, all_row_positions)); + }; + let newest_positions = if self.indexes.has_pk_index() { + None + } else { + Some(newest_pk_positions( + &self.batch_store, + pk_columns, + self.max_visible_batch_position, + max_visible_row, + )?) + }; + + let data_batch = RecordBatch::try_new(first.data.schema(), final_columns)?; + let pk_indices = resolve_pk_indices(&data_batch, pk_columns)?; + let keep = (0..data_batch.num_rows()) + .map(|row| { + Ok(match &newest_positions { + Some(newest) => newest.contains(&all_row_positions[row]), + None => { + let values: Vec = pk_indices + .iter() + .map(|&col| ScalarValue::try_from_array(data_batch.column(col), row)) + .collect::>()?; + self.indexes + .pk_is_newest(&values, all_row_positions[row], max_visible_row) + } + }) + }) + .collect::>>()?; + + let mask = BooleanArray::from_iter(keep.iter().copied()); + let filtered_columns = data_batch + .columns() + .iter() + .map(|c| arrow_select::filter::filter(c.as_ref(), &mask)) + .collect::, _>>()?; + let filtered_scores = all_scores + .into_iter() + .zip(keep.iter()) + .filter_map(|(s, keep)| keep.then_some(s)) + .collect(); + let filtered_positions = all_row_positions + .into_iter() + .zip(keep.iter()) + .filter_map(|(p, keep)| keep.then_some(p)) + .collect(); + + Ok((filtered_columns, filtered_scores, filtered_positions)) + } } impl DisplayAs for FtsIndexExec { diff --git a/rust/lance/src/dataset/mem_wal/scanner.rs b/rust/lance/src/dataset/mem_wal/scanner.rs index fe14bd82dd8..f1d84611e04 100644 --- a/rust/lance/src/dataset/mem_wal/scanner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner.rs @@ -24,13 +24,18 @@ //! use lance::dataset::mem_wal::scanner::LsmScanner; //! //! let scanner = LsmScanner::new(base_table, shard_snapshots, vec!["pk".to_string()]) -//! .project(&["id", "name"]) +//! .project(&["id", "name"])? //! .filter("id > 10")? -//! .limit(100, None); +//! .limit(Some(100), None)?; //! //! let stream = scanner.try_into_stream().await?; //! ``` +use arrow_schema::Schema as ArrowSchema; +use datafusion::common::ToDFSchema; +use datafusion::prelude::{Expr, SessionContext}; +use lance_core::{Error, Result}; + mod block_list; mod builder; mod collector; @@ -56,3 +61,24 @@ pub use fts_search::{LsmFtsSearchPlanner, SCORE_COLUMN}; pub use point_lookup::LsmPointLookupPlanner; pub use projection::DISTANCE_COLUMN; pub use vector_search::LsmVectorSearchPlanner; + +/// Parse a SQL filter expression against a MemWAL source schema. +pub fn parse_filter_expr(schema: &ArrowSchema, filter_expr: &str) -> Result { + let ctx = SessionContext::new(); + let df_schema = schema + .clone() + .to_dfschema() + .map_err(|e| Error::invalid_input(format!("Failed to create DFSchema: {}", e)))?; + ctx.parse_sql_expr(filter_expr, &df_schema) + .map_err(|e| Error::invalid_input(format!("Failed to parse filter expression: {}", e))) +} + +pub(crate) fn validate_overfetch_factor(factor: f64) -> Result { + if !factor.is_finite() || factor < 1.0 { + return Err(Error::invalid_input(format!( + "overfetch_factor must be finite and at least 1.0, got {}", + factor + ))); + } + Ok(factor) +} diff --git a/rust/lance/src/dataset/mem_wal/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/scanner/builder.rs index a006257493b..2947ef1464f 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/builder.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/builder.rs @@ -7,16 +7,20 @@ use std::collections::HashMap; use std::collections::hash_map::Entry; use std::sync::Arc; -use arrow_array::RecordBatch; -use arrow_schema::SchemaRef; +use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; +use arrow_array::cast::AsArray; +use arrow_array::types::Float32Type; +use arrow_array::{Array, FixedSizeListArray, RecordBatch}; +use arrow_schema::{DataType, SchemaRef}; use datafusion::common::ScalarValue; -use datafusion::common::ToDFSchema; use datafusion::logical_expr::Operator; use datafusion::physical_plan::limit::GlobalLimitExec; use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream}; use datafusion::prelude::{Expr, SessionContext}; use futures::TryStreamExt; use lance_core::{Error, Result, is_system_column}; +use lance_index::scalar::FullTextSearchQuery; +use lance_linalg::distance::DistanceType; use uuid::Uuid; use super::collector::{InMemoryMemTableRef, InMemoryMemTables, LsmDataSourceCollector}; @@ -24,9 +28,29 @@ use super::data_source::{FreshTierWatermark, ShardSnapshot}; use super::flushed_cache::{DatasetCache, GenerationWarmer}; use super::planner::LsmScanPlanner; use super::point_lookup::LsmPointLookupPlanner; +use super::projection::validate_projection_names; use crate::dataset::Dataset; use crate::session::Session; +/// Vector (KNN) search state, set by [`LsmScanner::nearest`] and friends. Mirrors +/// the subset of `lance::dataset::scanner::Query` the LSM vector planner honors. +#[derive(Debug, Clone)] +struct LsmVectorQuery { + /// Vector column to search. + column: String, + /// Flat query vector (one vector; validated against the column dim in `create_plan`). + key: Arc, + /// Number of nearest neighbors to fetch per source before the global merge. + k: usize, + /// Number of IVF partitions to probe on the base arm. + nprobes: usize, + /// Re-rank base candidates with exact distances when set (refine factor is + /// treated as a boolean; the LSM merge needs exact base distances). + refine: bool, + /// Distance metric; `None` defaults to L2 (matching the unindexed memtable arm). + metric_type: Option, +} + /// If `filter` is a point-lookup shape on `pk_col` — `pk = lit` (either /// operand order) or `pk IN (lit, …)` — return the literal key values. Any /// other shape returns `None`, so the scanner falls through to the general @@ -72,6 +96,60 @@ enum BaseSource { PathOnly(String), } +/// The fixed-size-list dimension of `column` in `schema`. Errors if the column +/// is missing or is not a fixed-size-list vector column. +fn vector_dim(schema: &arrow_schema::Schema, column: &str) -> Result { + let field = schema + .field_with_name(column) + .map_err(|_| Error::invalid_input(format!("vector column '{}' not found", column)))?; + match field.data_type() { + DataType::FixedSizeList(_, dim) => Ok(*dim), + other => Err(Error::invalid_input(format!( + "column '{}' is not a fixed-size-list vector (got {:?})", + column, other + ))), + } +} + +/// Wrap a flat `Float32` query vector into a single-row `FixedSizeListArray` of +/// `dim`, as the vector planner expects. A pre-built single-row FSL passes +/// through unchanged. The query length must match the column dimension. +fn key_to_fsl(key: &dyn Array, dim: i32) -> Result { + if let Some(fsl) = key.as_any().downcast_ref::() { + // The LSM cross-source merge assumes a single query vector; a multi-row + // FSL would otherwise flow through and silently produce wrong results. + if fsl.len() != 1 { + return Err(Error::invalid_input(format!( + "LSM vector search supports a single query vector, got {} rows", + fsl.len() + ))); + } + if fsl.value_length() != dim { + return Err(Error::invalid_input(format!( + "query vector dimension {} does not match column dimension {}", + fsl.value_length(), + dim + ))); + } + return Ok(fsl.clone()); + } + let values = key + .as_primitive_opt::() + .ok_or_else(|| Error::invalid_input("query vector must be Float32".to_string()))?; + if values.len() != dim as usize { + return Err(Error::invalid_input(format!( + "query vector dimension {} does not match column dimension {}", + values.len(), + dim + ))); + } + let mut builder = + FixedSizeListBuilder::with_capacity(Float32Builder::with_capacity(dim as usize), dim, 1); + builder.values().append_slice(values.values()); + builder.append(true); + Ok(builder.finish()) +} + /// Scanner for LSM tree data spanning base table, flushed MemTables, and active MemTable. /// /// This scanner provides a unified interface for querying data across multiple @@ -87,12 +165,18 @@ enum BaseSource { /// /// ```ignore /// let scanner = LsmScanner::new(base_table, shard_snapshots, vec!["pk".to_string()]) -/// .project(&["id", "name"]) +/// .project(&["id", "name"])? /// .filter("id > 10")? -/// .limit(100, None); +/// .limit(Some(100), None)?; /// /// let results = scanner.try_into_batch().await?; /// ``` +/// +/// The query-building methods mirror [`crate::dataset::scanner::Scanner`]: +/// [`Self::nearest`] (+ [`Self::nprobes`] / [`Self::refine`] / +/// [`Self::distance_metric`]) for vector search and [`Self::full_text_search`] +/// for FTS are state setters, and [`Self::create_plan`] dispatches to the right +/// planner — so an LSM read reads like a normal scan. pub struct LsmScanner { // Data sources base: BaseSource, @@ -110,6 +194,10 @@ pub struct LsmScanner { filter: Option, limit: Option, offset: Option, + /// Vector (KNN) search; when set, `create_plan` routes to the vector planner. + nearest: Option, + /// Full-text search; when set, `create_plan` routes to the FTS planner. + full_text_query: Option, // Internal columns with_row_address: bool, @@ -160,6 +248,8 @@ impl LsmScanner { filter: None, limit: None, offset: None, + nearest: None, + full_text_query: None, with_row_address: false, with_memtable_gen: false, pk_columns, @@ -200,6 +290,8 @@ impl LsmScanner { filter: None, limit: None, offset: None, + nearest: None, + full_text_query: None, with_row_address: false, with_memtable_gen: false, pk_columns, @@ -272,7 +364,8 @@ impl LsmScanner { /// Set the over-fetch multiple block-listed sources use in search plans /// so they still yield `k` live rows after cross-generation dedup. - /// Threaded into [`super::LsmFtsSearchPlanner`]; clamped to `>= 1.0`. + /// Threaded into the LSM search planners; values below `1.0` are rejected + /// when the plan is created. pub fn with_overfetch_factor(mut self, factor: f64) -> Self { self.overfetch_factor = Some(factor); self @@ -282,25 +375,16 @@ impl LsmScanner { /// /// If not called, all columns from the base schema are included. /// Primary key columns are always included for deduplication. - pub fn project(mut self, columns: &[&str]) -> Self { - self.projection = Some(columns.iter().map(|s| s.to_string()).collect()); - self + pub fn project>(mut self, columns: &[T]) -> Result { + self.projection = Some(columns.iter().map(|s| s.as_ref().to_string()).collect()); + Ok(self) } /// Set filter expression using SQL-like syntax. /// /// The filter is pushed down to each data source when possible. pub fn filter(mut self, filter_expr: &str) -> Result { - let ctx = SessionContext::new(); - let df_schema = self - .schema - .as_ref() - .clone() - .to_dfschema() - .map_err(|e| Error::invalid_input(format!("Failed to create DFSchema: {}", e)))?; - let expr = ctx.parse_sql_expr(filter_expr, &df_schema).map_err(|e| { - Error::invalid_input(format!("Failed to parse filter expression: {}", e)) - })?; + let expr = super::parse_filter_expr(self.schema.as_ref(), filter_expr)?; self.filter = Some(expr); Ok(self) } @@ -311,10 +395,80 @@ impl LsmScanner { self } - /// Limit the number of results. - pub fn limit(mut self, limit: usize, offset: Option) -> Self { - self.limit = Some(limit); - self.offset = offset; + /// Limit the number of results, with an optional offset. Matches + /// [`crate::dataset::scanner::Scanner::limit`]: both bounds are `Option` + /// and must be non-negative. + pub fn limit(mut self, limit: Option, offset: Option) -> Result { + if let Some(value) = limit + && value < 0 + { + return Err(Error::invalid_input( + "limit must be non-negative".to_string(), + )); + } + if let Some(value) = offset + && value < 0 + { + return Err(Error::invalid_input( + "offset must be non-negative".to_string(), + )); + } + self.limit = limit.map(|value| value as usize); + self.offset = offset.map(|value| value as usize); + Ok(self) + } + + /// Find the `k` nearest neighbors of `key` in `column`. Routes `create_plan` + /// through the LSM vector planner (base ∪ flushed ∪ in-memory). Mirrors + /// [`crate::dataset::scanner::Scanner::nearest`]; the LSM path supports a + /// single Float32 query vector. When combined with an offset, the LSM path + /// fetches `k + offset` per source before applying the final page. Tune with + /// [`Self::nprobes`], [`Self::refine`], and [`Self::distance_metric`]. + pub fn nearest(mut self, column: &str, key: &dyn Array, k: usize) -> Result { + if k == 0 { + return Err(Error::invalid_input("k must be positive".to_string())); + } + if key.is_empty() { + return Err(Error::invalid_input( + "query vector must have non-zero length".to_string(), + )); + } + self.nearest = Some(LsmVectorQuery { + column: column.to_string(), + key: key.slice(0, key.len()), + k, + nprobes: 1, + refine: false, + metric_type: None, + }); + Ok(self) + } + + /// Number of IVF partitions to probe on the base arm (default 1). No-op + /// unless [`Self::nearest`] was called. + pub fn nprobes(mut self, nprobes: usize) -> Self { + if let Some(q) = self.nearest.as_mut() { + q.nprobes = nprobes; + } + self + } + + /// Re-rank base-table candidates with exact distances so they are directly + /// comparable to the exact memtable distances in the cross-source merge. + /// The factor is treated as a boolean today. No-op unless `nearest` was set. + pub fn refine(mut self, refine_factor: u32) -> Self { + if let Some(q) = self.nearest.as_mut() { + q.refine = refine_factor > 0; + } + self + } + + /// Distance metric for the vector search. Defaults to L2. No-op unless + /// [`Self::nearest`] was called. + pub fn distance_metric(mut self, metric: DistanceType) -> Self { + if let Some(q) = self.nearest.as_mut() { + q.metric_type = Some(metric); + } self } @@ -354,8 +508,163 @@ impl LsmScanner { /// Create the execution plan. pub async fn create_plan(&self) -> Result> { + // Dispatch by builder state, mirroring `Scanner::create_plan` and + // `MemTableScanner::create_plan`: vector search, then full-text search, + // then the (point-lookup or union) plain scan. + // The LSM path has no hybrid mode, so reject combined search rather than + // silently dropping the full-text query under the vector dispatch. + if self.nearest.is_some() && self.full_text_query.is_some() { + return Err(Error::invalid_input( + "LSM scanner does not support combined vector and full-text search".to_string(), + )); + } + if self.nearest.is_some() { + return self.plan_vector().await; + } + if self.full_text_query.is_some() { + return self.plan_fts().await; + } + self.plan_scan().await + } + + /// Apply the builder's `limit`/`offset` on top of a search plan. The plain + /// scan applies them inside the planner; the vector/FTS arms wrap here. + fn apply_limit_offset(&self, plan: Arc) -> Arc { + let skip = self.offset.unwrap_or(0); + if skip == 0 && self.limit.is_none() { + return plan; + } + Arc::new(GlobalLimitExec::new(plan, skip, self.limit)) + } + + /// Vector (KNN) search across base ∪ flushed ∪ in-memory, via the LSM vector + /// planner. Honors the builder filter as a prefilter. + async fn plan_vector(&self) -> Result> { + let nearest = self + .nearest + .as_ref() + .expect("plan_vector requires a nearest query"); + let base_schema = self.schema(); + let dim = vector_dim(base_schema.as_ref(), &nearest.column)?; + let query_fsl = key_to_fsl(nearest.key.as_ref(), dim)?; + let distance_type = nearest.metric_type.unwrap_or(DistanceType::L2); + let collector = self.build_collector(); + let mut planner = super::LsmVectorSearchPlanner::new( + collector, + self.pk_columns.clone(), + base_schema, + nearest.column.clone(), + distance_type, + ) + .with_filter(self.filter.clone()); + if let BaseSource::Table(dataset) = &self.base { + planner = planner.with_dataset(dataset.clone()); + } + if let Some(session) = &self.session { + planner = planner.with_session(session.clone()); + } + if let Some(cache) = &self.flushed_cache { + planner = planner.with_flushed_cache(cache.clone()); + } + if let Some(warmer) = &self.warmer { + planner = planner.with_warmer(warmer.clone()); + } + // Over-fetch by `offset` so the per-source top-k still yields `k` live + // rows after `apply_limit_offset` skips the first `offset` (mirrors the + // FTS arm, which folds `offset` into its `k`). + let per_source_k = nearest.k.saturating_add(self.offset.unwrap_or(0)); + let overfetch_factor = self.overfetch_factor.unwrap_or(1.0); + let plan = planner + .plan_search( + &query_fsl, + per_source_k, + nearest.nprobes, + self.projection.as_deref(), + nearest.refine, + overfetch_factor, + ) + .await?; + Ok(self.apply_limit_offset(plan)) + } + + /// Full-text search across base ∪ flushed ∪ in-memory, via the LSM FTS + /// planner. Query/scanner limits bound per-source fetches when present; + /// otherwise the search remains unbounded and any offset is applied above. + async fn plan_fts(&self) -> Result> { + let query = self + .full_text_query + .as_ref() + .expect("plan_fts requires a full-text query"); + let columns: Vec = query.columns().into_iter().collect(); + if columns.len() > 1 { + return Err(Error::invalid_input( + "LSM full-text search supports a single column".to_string(), + )); + } + let column = columns.into_iter().next().ok_or_else(|| { + Error::invalid_input( + "full_text_search requires a column; set it with `FullTextSearchQuery::with_column`" + .to_string(), + ) + })?; let base_schema = self.schema(); + base_schema.field_with_name(&column).map_err(|_| { + Error::invalid_input(format!("Column '{}' not found in schema", column)) + })?; + let query_limit = query + .limit + .map(|limit| { + if limit < 0 { + Err(Error::invalid_input( + "full-text search limit must be non-negative".to_string(), + )) + } else { + Ok(limit as usize) + } + }) + .transpose()?; + // Match the dataset Scanner: a query-level FTS limit is the hard + // search cap, and scanner offset pages within that capped result set. + let source_limit = match query_limit { + Some(limit) => Some(limit), + None => self + .limit + .map(|limit| limit.saturating_add(self.offset.unwrap_or(0))), + }; + + let collector = self.build_collector(); + let mut planner = + super::LsmFtsSearchPlanner::new(collector, self.pk_columns.clone(), base_schema) + .with_filter(self.filter.clone()); + if let Some(session) = &self.session { + planner = planner.with_session(session.clone()); + } + if let Some(cache) = &self.flushed_cache { + planner = planner.with_flushed_cache(cache.clone()); + } + if let Some(warmer) = &self.warmer { + planner = planner.with_warmer(warmer.clone()); + } + if let Some(factor) = self.overfetch_factor { + planner = planner.with_overfetch_factor(factor); + } + let plan = planner + .plan_search( + &column, + query.clone(), + source_limit, + self.projection.as_deref(), + ) + .await?; + Ok(self.apply_limit_offset(plan)) + } + + /// Plain (filter / projection / limit) scan over base ∪ flushed ∪ in-memory. + async fn plan_scan(&self) -> Result> { + let collector = self.build_collector(); + let base_schema = self.schema(); + validate_projection_names(self.projection.as_deref(), &base_schema, &[])?; // Fast point-lookup routing: a `pk = lit` / `pk IN (..)` filter on the // single pk column — with no offset and no scan-only system columns — @@ -417,37 +726,17 @@ impl LsmScanner { .await } - /// Build a local-scoring FTS plan spanning base + flushed + active sources. - /// - /// Routes through [`super::LsmFtsSearchPlanner`]. Output schema is - /// `projection ∪ pk_columns + _score`; per-source local BM25 `_score` - /// is merged DESC and capped at `k`. `column` must be FTS-indexed on - /// the queried sources. - pub async fn full_text_search( - &self, - column: &str, - query: lance_index::scalar::FullTextSearchQuery, - k: usize, - ) -> Result> { - let collector = self.build_collector(); - let base_schema = self.schema(); - let mut planner = - super::LsmFtsSearchPlanner::new(collector, self.pk_columns.clone(), base_schema); - if let Some(session) = &self.session { - planner = planner.with_session(session.clone()); - } - if let Some(cache) = &self.flushed_cache { - planner = planner.with_flushed_cache(cache.clone()); - } - if let Some(warmer) = &self.warmer { - planner = planner.with_warmer(warmer.clone()); - } - if let Some(factor) = self.overfetch_factor { - planner = planner.with_overfetch_factor(factor); - } - planner - .plan_search(column, query, k, self.projection.as_deref()) - .await + /// Find rows matching a full-text query. Routes `create_plan` through the + /// LSM FTS planner (base ∪ flushed ∪ in-memory), local-scored by BM25 and + /// merged by `_score` DESC. Mirrors + /// [`crate::dataset::scanner::Scanner::full_text_search`]: the searched + /// column(s) come from the query (set via `FullTextSearchQuery::with_column`); + /// the query limit, when present, controls the FTS source/merge top-k; + /// scanner limit/offset still apply to the final merged result. The LSM + /// path supports a single FTS-indexed column. + pub fn full_text_search(mut self, query: FullTextSearchQuery) -> Result { + self.full_text_query = Some(query); + Ok(self) } /// Execute the scan and return a stream of record batches. @@ -462,14 +751,14 @@ impl LsmScanner { /// Execute the scan and collect all results into a single RecordBatch. pub async fn try_into_batch(&self) -> Result { let stream = self.try_into_stream().await?; + let output_schema = stream.schema(); let batches: Vec = stream .try_collect() .await .map_err(|e| Error::io(format!("Failed to collect batches: {}", e)))?; if batches.is_empty() { - let schema = self.schema(); - return Ok(RecordBatch::new_empty(schema)); + return Ok(RecordBatch::new_empty(output_schema)); } let schema = batches[0].schema(); @@ -683,6 +972,124 @@ mod tests { } } + #[tokio::test] + async fn invalid_overfetch_factor_is_rejected() { + let shard = Uuid::new_v4(); + let scanner = LsmScanner::without_base_table( + pk_schema(), + "memory://t", + vec![], + vec!["id".to_string()], + ) + .with_in_memory_memtables( + shard, + InMemoryMemTables { + active: mk_pk_memtable(&[1, 2], 2), + frozen: vec![], + }, + ) + .with_overfetch_factor(0.5); + + let Err(err) = scanner.try_into_stream().await else { + panic!("invalid overfetch factor should fail planning"); + }; + assert!( + err.to_string().contains("overfetch_factor"), + "unexpected error for invalid overfetch factor: {err}" + ); + + let empty_scanner = LsmScanner::without_base_table( + pk_schema(), + "memory://empty", + vec![], + vec!["id".to_string()], + ) + .with_overfetch_factor(0.5); + let Err(err) = empty_scanner.try_into_stream().await else { + panic!("invalid overfetch factor should fail even when there are no sources"); + }; + assert!( + err.to_string().contains("overfetch_factor"), + "unexpected error for invalid empty-source overfetch factor: {err}" + ); + } + + #[tokio::test] + async fn unknown_scan_projection_column_is_rejected() { + let scanner = LsmScanner::without_base_table( + pk_schema(), + "memory://t", + vec![], + vec!["id".to_string()], + ) + .project(&["missing"]) + .unwrap(); + + let Err(err) = scanner.try_into_stream().await else { + panic!("unknown projection column should fail planning"); + }; + assert!( + err.to_string().contains("missing"), + "unexpected missing-column projection error: {err}" + ); + } + + #[tokio::test] + async fn point_lookup_fast_route_rejects_missing_projection_column() { + let shard = Uuid::new_v4(); + let scanner = LsmScanner::without_base_table( + pk_schema(), + "memory://t", + vec![], + vec!["id".to_string()], + ) + .with_in_memory_memtables( + shard, + InMemoryMemTables { + active: mk_pk_memtable(&[1, 2], 2), + frozen: vec![], + }, + ) + .project(&["missing"]) + .unwrap() + .filter("id = 1") + .unwrap(); + + let Err(err) = scanner.try_into_stream().await else { + panic!("unknown projection column should fail on point-lookup fast route"); + }; + assert!( + err.to_string().contains("missing"), + "unexpected missing-column point lookup error: {err}" + ); + } + + #[tokio::test] + async fn full_text_search_missing_column_is_rejected() { + use arrow_schema::{DataType, Field}; + + let scanner = LsmScanner::without_base_table( + pk_schema_with(Field::new("text", DataType::Utf8, true)), + "memory://t", + vec![], + vec!["id".to_string()], + ) + .full_text_search( + FullTextSearchQuery::new("lance".to_string()) + .with_column("missing".to_string()) + .unwrap(), + ) + .unwrap(); + + let Err(err) = scanner.try_into_stream().await else { + panic!("unknown FTS column should fail planning"); + }; + assert!( + err.to_string().contains("missing"), + "unexpected missing FTS column error: {err}" + ); + } + #[tokio::test] async fn contains_pks_reports_fresh_tier_membership() { // Fresh-tier only: active gen 2 (pk=1,2) + frozen gen 1 (pk=3). @@ -758,6 +1165,646 @@ mod tests { assert_eq!(bounded, vec![true, false, false, true, false, false]); } + fn pk_schema_with(extra: arrow_schema::Field) -> SchemaRef { + use arrow_schema::{DataType, Field, Schema}; + let mut id_meta = HashMap::new(); + id_meta.insert( + "lance-schema:unenforced-primary-key".to_string(), + "true".to_string(), + ); + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(id_meta), + extra, + ])) + } + + /// `LsmScanner::nearest(..).create_plan()` must route through the vector + /// planner (exercising the Scanner-aligned facade, `key_to_fsl`, and + /// `vector_dim`) and surface the in-memory match. + #[tokio::test] + async fn nearest_dispatches_through_facade() { + use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; + use crate::dataset::{Dataset, WriteParams}; + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; + use arrow_array::{Float32Array, Int32Array, RecordBatchIterator}; + use arrow_schema::{DataType, Field}; + + let schema = pk_schema_with(Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), + false, + )); + let make_batch = |ids: &[i32]| -> RecordBatch { + let mut vb = FixedSizeListBuilder::new(Float32Builder::new(), 4); + for id in ids { + let base = *id as f32 * 0.1; + for d in 0..4 { + vb.values().append_value(base + d as f32 * 0.1); + } + vb.append(true); + } + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids.to_vec())), + Arc::new(vb.finish()), + ], + ) + .unwrap() + }; + + // Far, unindexed base rows contribute nothing under fast_search. + let tmp = tempfile::tempdir().unwrap(); + let uri = format!("{}/base", tmp.path().to_str().unwrap()); + let reader = RecordBatchIterator::new(vec![Ok(make_batch(&[100, 200]))], schema.clone()); + let base = Arc::new( + Dataset::write(reader, &uri, Some(WriteParams::default())) + .await + .unwrap(), + ); + + let store = Arc::new(BatchStore::with_capacity(16)); + let mut index = IndexStore::new(); + index.enable_pk_index(&[("id".to_string(), 0)]); + index.add_hnsw( + "vec_hnsw".to_string(), + 1, + "vector".to_string(), + DistanceType::L2, + 64, + 8, + ); + let batch = make_batch(&[1, 2, 3]); + store.append(batch.clone()).unwrap(); + index + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + + let scanner = LsmScanner::new(base, vec![], vec!["id".to_string()]) + .with_in_memory_memtables( + Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store: store, + index_store: Arc::new(index), + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ) + // A flat query vector (closest to id=1); the facade wraps it into an FSL. + .nearest( + "vector", + &Float32Array::from(vec![0.1f32, 0.2, 0.3, 0.4]), + 3, + ) + .unwrap(); + + let batches: Vec = scanner + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let ids: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!( + ids.first().copied(), + Some(1), + "nearest neighbor via the facade should be id=1; got {ids:?}" + ); + } + + /// Vector pagination: `limit(Some(2), Some(1))` over-fetches `k + offset` + /// per source, then skips `offset` and caps at `limit`. With ids ordered by + /// distance (id=1 nearest), the page must be the 2nd–3rd nearest (id=2, id=3), + /// not the 1st–2nd. Covers `plan_vector`'s `per_source_k` + `apply_limit_offset`. + #[tokio::test] + async fn nearest_pagination_skips_offset_and_caps_limit() { + use crate::dataset::{Dataset, WriteParams}; + use crate::index::DatasetIndexExt; + use crate::index::vector::VectorIndexParams; + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; + use arrow_array::{Float32Array, Int32Array, RecordBatchIterator}; + use arrow_schema::{DataType, Field}; + use lance_index::IndexType; + + let schema = pk_schema_with(Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), + false, + )); + let mut vb = FixedSizeListBuilder::new(Float32Builder::new(), 4); + let ids: Vec = (1..=6).collect(); + for id in &ids { + let base = *id as f32 * 0.1; + for d in 0..4 { + vb.values().append_value(base + d as f32 * 0.1); + } + vb.append(true); + } + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids.clone())), + Arc::new(vb.finish()), + ], + ) + .unwrap(); + + let tmp = tempfile::tempdir().unwrap(); + let uri = format!("{}/base", tmp.path().to_str().unwrap()); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + // `ivf_flat(1)` is exhaustive within its single partition, so ordering is exact. + let mut base = Dataset::write(reader, &uri, Some(WriteParams::default())) + .await + .unwrap(); + let ivf_flat = VectorIndexParams::ivf_flat(1, DistanceType::L2); + base.create_index(&["vector"], IndexType::Vector, None, &ivf_flat, true) + .await + .unwrap(); + let base = Arc::new(base); + + let scanner = LsmScanner::new(base, vec![], vec!["id".to_string()]) + .nearest( + "vector", + &Float32Array::from(vec![0.1f32, 0.2, 0.3, 0.4]), + 2, + ) + .unwrap() + .limit(Some(2), Some(1)) + .unwrap(); + + let batches: Vec = scanner + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let mut out: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + out.sort(); + assert_eq!( + out, + vec![2, 3], + "offset=1, limit=2 over k=2 must return the 2nd-3rd nearest (id=2, id=3); got {out:?}" + ); + } + + /// FTS pagination: `limit(Some(1), Some(1))` fetches `limit + offset` per + /// source then skips `offset`. With BM25 ranking id=1 > id=2 > id=3 (by doc + /// length), the page must be the 2nd-ranked hit (id=2). Covers `plan_fts`'s + /// `k = limit + offset` + `apply_limit_offset`. + #[tokio::test] + async fn full_text_search_pagination_skips_offset_and_caps_limit() { + use crate::dataset::{Dataset, WriteParams}; + use crate::index::DatasetIndexExt; + use arrow_array::{Int32Array, RecordBatchIterator, StringArray}; + use arrow_schema::{DataType, Field}; + use lance_index::IndexType; + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + + let schema = pk_schema_with(Field::new("text", DataType::Utf8, true)); + // Shorter docs score higher under BM25 length normalization, so the + // ranking is deterministically id=1 > id=2 > id=3. + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec![ + "lance", + "lance filler", + "lance filler filler", + ])), + ], + ) + .unwrap(); + + let tmp = tempfile::tempdir().unwrap(); + let uri = format!("{}/base", tmp.path().to_str().unwrap()); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let mut base = Dataset::write(reader, &uri, Some(WriteParams::default())) + .await + .unwrap(); + base.create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + let base = Arc::new(Dataset::open(&uri).await.unwrap()); + + let query_limited = LsmScanner::new(base.clone(), vec![], vec!["id".to_string()]) + .full_text_search( + FullTextSearchQuery::new("lance".to_string()) + .with_column("text".to_string()) + .unwrap() + .limit(Some(1)), + ) + .unwrap(); + let batches: Vec = query_limited + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let out: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!( + out, + vec![1], + "query-level FTS limit=1 must cap the unpaginated scanner result; got {out:?}" + ); + + let query_limited_with_offset = + LsmScanner::new(base.clone(), vec![], vec!["id".to_string()]) + .full_text_search( + FullTextSearchQuery::new("lance".to_string()) + .with_column("text".to_string()) + .unwrap() + .limit(Some(2)), + ) + .unwrap() + .limit(None, Some(1)) + .unwrap(); + let batches: Vec = query_limited_with_offset + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let out: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!( + out, + vec![2], + "query-level FTS limit=2 plus offset=1 must page within the top 2; got {out:?}" + ); + + let scanner = LsmScanner::new(base, vec![], vec!["id".to_string()]) + .full_text_search( + FullTextSearchQuery::new("lance".to_string()) + .with_column("text".to_string()) + .unwrap(), + ) + .unwrap() + .limit(Some(1), Some(1)) + .unwrap(); + + let batches: Vec = scanner + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let out: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!( + out, + vec![2], + "offset=1, limit=1 must return the 2nd-ranked hit (id=2); got {out:?}" + ); + } + + #[tokio::test] + async fn full_text_search_without_limit_returns_all_matches() { + use crate::dataset::{Dataset, WriteParams}; + use crate::index::DatasetIndexExt; + use arrow_array::{Int32Array, RecordBatchIterator, StringArray}; + use arrow_schema::{DataType, Field}; + use lance_index::IndexType; + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + + let schema = pk_schema_with(Field::new("text", DataType::Utf8, true)); + let ids: Vec = (0..12).collect(); + let texts: Vec<&str> = (0..12).map(|_| "lance").collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(texts)), + ], + ) + .unwrap(); + + let tmp = tempfile::tempdir().unwrap(); + let uri = format!("{}/base", tmp.path().to_str().unwrap()); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let mut base = Dataset::write(reader, &uri, Some(WriteParams::default())) + .await + .unwrap(); + base.create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + let base = Arc::new(Dataset::open(&uri).await.unwrap()); + + let scanner = LsmScanner::new(base, vec![], vec!["id".to_string()]) + .full_text_search( + FullTextSearchQuery::new("lance".to_string()) + .with_column("text".to_string()) + .unwrap(), + ) + .unwrap(); + let batches: Vec = scanner + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total, 12, + "unbounded LSM FTS must not apply the old default top-10 cap" + ); + } + + /// Setting both `nearest` and `full_text_search` must error rather than + /// silently running vector search and dropping the full-text query (the LSM + /// path has no hybrid mode). + #[tokio::test] + async fn combined_vector_and_fts_is_rejected() { + use crate::dataset::{Dataset, WriteParams}; + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; + use arrow_array::{Float32Array, Int32Array, RecordBatchIterator}; + use arrow_schema::{DataType, Field}; + + let schema = pk_schema_with(Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), + false, + )); + let mut vb = FixedSizeListBuilder::new(Float32Builder::new(), 4); + for d in 0..4 { + vb.values().append_value(d as f32); + } + vb.append(true); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1])), Arc::new(vb.finish())], + ) + .unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let uri = format!("{}/base", tmp.path().to_str().unwrap()); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let base = Arc::new( + Dataset::write(reader, &uri, Some(WriteParams::default())) + .await + .unwrap(), + ); + + let scanner = LsmScanner::new(base, vec![], vec!["id".to_string()]) + .nearest( + "vector", + &Float32Array::from(vec![0.0f32, 1.0, 2.0, 3.0]), + 1, + ) + .unwrap() + .full_text_search( + FullTextSearchQuery::new("lance".to_string()) + .with_column("vector".to_string()) + .unwrap(), + ) + .unwrap(); + let err = scanner.create_plan().await.unwrap_err(); + assert!( + err.to_string() + .contains("combined vector and full-text search"), + "expected combined-search rejection, got: {err}" + ); + } + + /// A multi-row query vector must be rejected, not silently flow through the + /// single-vector LSM merge as wrong results. + #[tokio::test] + async fn multi_row_query_vector_is_rejected() { + use crate::dataset::{Dataset, WriteParams}; + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; + use arrow_array::{Int32Array, RecordBatchIterator}; + use arrow_schema::{DataType, Field}; + + let schema = pk_schema_with(Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), + false, + )); + let mut col = FixedSizeListBuilder::new(Float32Builder::new(), 4); + for d in 0..4 { + col.values().append_value(d as f32); + } + col.append(true); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1])), Arc::new(col.finish())], + ) + .unwrap(); + let tmp = tempfile::tempdir().unwrap(); + let uri = format!("{}/base", tmp.path().to_str().unwrap()); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let base = Arc::new( + Dataset::write(reader, &uri, Some(WriteParams::default())) + .await + .unwrap(), + ); + + // A two-row FSL query vector. + let mut q = FixedSizeListBuilder::new(Float32Builder::new(), 4); + for _ in 0..2 { + for d in 0..4 { + q.values().append_value(d as f32); + } + q.append(true); + } + let query = q.finish(); + + let scanner = LsmScanner::new(base, vec![], vec!["id".to_string()]) + .nearest("vector", &query, 1) + .unwrap(); + let err = scanner.create_plan().await.unwrap_err(); + assert!( + err.to_string().contains("single query vector"), + "expected single-query-vector rejection, got: {err}" + ); + } + + /// `LsmScanner::full_text_search(query).create_plan()` must route through the + /// FTS planner and surface a memtable-only match. + #[tokio::test] + async fn full_text_search_dispatches_through_facade() { + use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; + use crate::dataset::{Dataset, WriteParams}; + use crate::index::DatasetIndexExt; + use arrow_array::{Int32Array, RecordBatchIterator, StringArray}; + use arrow_schema::{DataType, Field}; + use lance_index::IndexType; + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + + let schema = pk_schema_with(Field::new("text", DataType::Utf8, true)); + let make_batch = |ids: &[i32], texts: &[&str]| { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids.to_vec())), + Arc::new(StringArray::from(texts.to_vec())), + ], + ) + .unwrap() + }; + + let tmp = tempfile::tempdir().unwrap(); + let uri = format!("{}/base", tmp.path().to_str().unwrap()); + let reader = + RecordBatchIterator::new(vec![Ok(make_batch(&[1], &["alpha"]))], schema.clone()); + let mut base = Dataset::write(reader, &uri, Some(WriteParams::default())) + .await + .unwrap(); + base.create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + let base = Arc::new(Dataset::open(&uri).await.unwrap()); + + let store = Arc::new(BatchStore::with_capacity(16)); + let mut index = IndexStore::new(); + index.enable_pk_index(&[("id".to_string(), 0)]); + index.add_fts("text_fts".to_string(), 1, "text".to_string()); + let batch = make_batch(&[99], &["zebra"]); + store.append(batch.clone()).unwrap(); + index + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + + let scanner = LsmScanner::new(base, vec![], vec!["id".to_string()]) + .with_in_memory_memtables( + Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store: store, + index_store: Arc::new(index), + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ) + .full_text_search( + FullTextSearchQuery::new("zebra".to_string()) + .with_column("text".to_string()) + .unwrap(), + ) + .unwrap(); + + let batches: Vec = scanner + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + rows, 1, + "facade FTS should surface the memtable 'zebra' row" + ); + } + + /// Empty search results must still preserve the execution plan schema. + #[tokio::test] + async fn try_into_batch_empty_fts_keeps_score_schema() { + use arrow_schema::{DataType, Field}; + + let schema = pk_schema_with(Field::new("text", DataType::Utf8, true)); + let scanner = LsmScanner::without_base_table( + schema, + "memory://empty", + vec![], + vec!["id".to_string()], + ) + .full_text_search( + FullTextSearchQuery::new("missing".to_string()) + .with_column("text".to_string()) + .unwrap(), + ) + .unwrap(); + + let batch = scanner.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 0); + assert!( + batch.schema().field_with_name("_score").is_ok(), + "empty FTS batch must keep the planned _score column" + ); + } + /// One active memtable with a maintained BTree on `id`, all rows visible. fn mk_indexed_memtable(schema: &SchemaRef, ids: &[i32], names: &[&str]) -> InMemoryMemTableRef { use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; @@ -826,6 +1873,29 @@ mod tests { rows.iter().map(|b| b.num_rows()).sum::() } }; + let collect_ids = |plan: Arc| { + let ctx = ctx.clone(); + async move { + let rows: Vec = plan + .execute(0, ctx.task_ctx()) + .unwrap() + .try_collect() + .await + .unwrap(); + let mut ids = Vec::new(); + for batch in rows { + let id_array = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + ids.extend(id_array.values().iter().copied()); + } + ids.sort_unstable(); + ids + } + }; // `id = 2` routes to the direct point-lookup node (OneShotStream), not the // union/dedup scan, and returns the one matching row. @@ -865,5 +1935,21 @@ mod tests { "range filter must not route to the point-lookup node" ); assert_eq!(count(plan).await, 3); // 3,4,5 + + // Offset changes the semantics, so even a point-lookup-shaped filter + // must bypass the direct lookup route and use the general scan path. + let plan = scanner() + .filter_expr(col("id").in_list(vec![lit(1i32), lit(3i32), lit(5i32)], false)) + .limit(None, Some(1)) + .unwrap() + .create_plan() + .await + .unwrap(); + let disp = format!("{}", displayable(plan.as_ref()).indent(true)); + assert!( + !disp.contains("OneShotStream"), + "offset point-lookup filters must use the scan path: {disp}" + ); + assert_eq!(collect_ids(plan).await, vec![3, 5]); } } diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec.rs b/rust/lance/src/dataset/mem_wal/scanner/exec.rs index 115cffccc81..1498c5f60ea 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/exec.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/exec.rs @@ -10,19 +10,16 @@ //! - [`BloomFilterGuardExec`]: Guards child execution with bloom filter check //! - [`CoalesceFirstExec`]: Returns first non-empty result with short-circuit //! - [`PkBlockFilterExec`]: Drops rows whose PK was superseded by a newer generation (the cross-generation block-list) -//! - [`NewestPkFilterExec`]: Drops active-memtable hits that aren't the newest visible version of their PK (the within-source recency filter) mod bloom_guard; mod coalesce_first; mod generation_tag; -mod newest_pk_filter; mod pk; mod pk_block_filter; pub use bloom_guard::{BloomFilterGuardExec, compute_pk_hash_from_scalars}; pub use coalesce_first::CoalesceFirstExec; pub use generation_tag::{MEMTABLE_GEN_COLUMN, MemtableGenTagExec}; -pub use newest_pk_filter::NewestPkFilterExec; pub use pk::{ ROW_ADDRESS_COLUMN, compute_pk_hash, is_supported_pk_type, resolve_pk_indices, validate_pk_types, diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec/newest_pk_filter.rs b/rust/lance/src/dataset/mem_wal/scanner/exec/newest_pk_filter.rs deleted file mode 100644 index e1495cb0bb1..00000000000 --- a/rust/lance/src/dataset/mem_wal/scanner/exec/newest_pk_filter.rs +++ /dev/null @@ -1,393 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors - -//! Drop predicate-crossing stale rows from an active-memtable index search. -//! -//! The active memtable's HNSW / inverted index are append-only, so an updated -//! row's old entries stay live. When an update moves a row out of the query's -//! match set, the fresh version isn't in the index result, so a result-set -//! dedup (keep-newest among the returned rows) has nothing to suppress the -//! stale version against — and it leaks. -//! -//! This node closes that hole with a predicate-independent recency check: for -//! each hit it asks the memtable's maintained primary-key index -//! ([`IndexStore::pk_is_newest`]) whether the hit's own row position is the -//! newest version of its primary key visible at the query's `max_visible` -//! watermark, and keeps the hit **iff so**. A stale hit (some -//! newer version exists) is dropped even when that newer version never appears -//! in the result. This is exactly the seek point-lookup already does; the index -//! search arms simply didn't do it. - -use std::any::Any; -use std::fmt; -use std::pin::Pin; -use std::sync::Arc; -use std::task::{Context, Poll}; - -use arrow::compute::filter_record_batch; -use arrow_array::{Array, BooleanArray, RecordBatch, UInt64Array}; -use arrow_schema::SchemaRef; -use datafusion::common::ScalarValue; -use datafusion::error::{DataFusionError, Result as DFResult}; -use datafusion::execution::TaskContext; -use datafusion::physical_expr::EquivalenceProperties; -use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, - SendableRecordBatchStream, -}; -use futures::{Stream, StreamExt}; - -use super::pk::resolve_pk_indices; -use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; - -/// Keeps only the index hits that are the newest visible version of their PK. -/// -/// The input must expose all `pk_columns` and the `row_id_column` (`UInt64`, -/// the BatchStore row position). The output schema is unchanged. -pub struct NewestPkFilterExec { - input: Arc, - pk_columns: Vec, - row_id_column: String, - /// Holds the maintained primary-key index, queried per hit via - /// [`IndexStore::pk_is_newest`]. - index_store: Arc, - /// Resolves the `max_visible` row watermark from the visible batch prefix. - batch_store: Arc, - /// The MVCC batch-position snapshot the index search latched. Captured once - /// at plan time and shared with the search so the recency check keys on the - /// same snapshot the hits came from. - max_visible_batch_position: usize, - properties: Arc, -} - -impl fmt::Debug for NewestPkFilterExec { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - // `BatchStore` / `IndexStore` aren't `Debug`; show only the knobs. - f.debug_struct("NewestPkFilterExec") - .field("pk_columns", &self.pk_columns) - .field("row_id_column", &self.row_id_column) - .field( - "max_visible_batch_position", - &self.max_visible_batch_position, - ) - .finish() - } -} - -impl NewestPkFilterExec { - pub fn new( - input: Arc, - pk_columns: Vec, - row_id_column: impl Into, - index_store: Arc, - batch_store: Arc, - max_visible_batch_position: usize, - ) -> Self { - // A filter preserves the input schema and partitioning. - let properties = Arc::new(PlanProperties::new( - EquivalenceProperties::new(input.schema()), - input.output_partitioning().clone(), - input.pipeline_behavior(), - input.boundedness(), - )); - Self { - input, - pk_columns, - row_id_column: row_id_column.into(), - index_store, - batch_store, - max_visible_batch_position, - properties, - } - } - - /// The inclusive max visible row position for this snapshot, or `None` when - /// no rows are visible. - fn max_visible_row(&self) -> Option { - self.batch_store - .max_visible_row(self.max_visible_batch_position) - } -} - -impl DisplayAs for NewestPkFilterExec { - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - match t { - DisplayFormatType::Default - | DisplayFormatType::Verbose - | DisplayFormatType::TreeRender => { - write!( - f, - "NewestPkFilterExec: pk=[{}], row_id={}, max_visible_batch={}", - self.pk_columns.join(", "), - self.row_id_column, - self.max_visible_batch_position, - ) - } - } - } -} - -impl ExecutionPlan for NewestPkFilterExec { - fn name(&self) -> &str { - "NewestPkFilterExec" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - self.input.schema() - } - - fn properties(&self) -> &Arc { - &self.properties - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> DFResult> { - if children.len() != 1 { - return Err(DataFusionError::Internal( - "NewestPkFilterExec requires exactly one child".to_string(), - )); - } - Ok(Arc::new(Self::new( - children[0].clone(), - self.pk_columns.clone(), - self.row_id_column.clone(), - self.index_store.clone(), - self.batch_store.clone(), - self.max_visible_batch_position, - ))) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> DFResult { - let input_stream = self.input.execute(partition, context)?; - Ok(Box::pin(NewestPkFilterStream { - input: input_stream, - pk_columns: self.pk_columns.clone(), - row_id_column: self.row_id_column.clone(), - index_store: self.index_store.clone(), - max_visible_row: self.max_visible_row(), - schema: self.schema(), - })) - } -} - -struct NewestPkFilterStream { - input: SendableRecordBatchStream, - pk_columns: Vec, - row_id_column: String, - index_store: Arc, - /// Inclusive watermark snapshot; `None` when no rows are visible. - max_visible_row: Option, - schema: SchemaRef, -} - -impl NewestPkFilterStream { - fn filter_batch(&self, batch: RecordBatch) -> DFResult { - // No primary-key index (memtable without a primary key), no visible - // rows, or an empty batch: nothing to dedup against, so pass it through. - if !self.index_store.has_pk_index() { - return Ok(batch); - } - let Some(max_visible_row) = self.max_visible_row else { - return Ok(batch); - }; - if batch.num_rows() == 0 { - return Ok(batch); - } - - let pk_indices = resolve_pk_indices(&batch, &self.pk_columns)?; - let row_ids = batch - .column_by_name(&self.row_id_column) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "Row-id column '{}' not found in NewestPkFilterExec input", - self.row_id_column - )) - })? - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal(format!( - "Row-id column '{}' is not UInt64", - self.row_id_column - )) - })?; - - let mut keep = Vec::with_capacity(batch.num_rows()); - for row in 0..batch.num_rows() { - // A null row position can't be ordered; keep it rather than guess - // (callers always project a real position here). - if row_ids.is_null(row) { - keep.push(true); - continue; - } - let position = row_ids.value(row); - let values: Vec = pk_indices - .iter() - .map(|&col| ScalarValue::try_from_array(batch.column(col), row)) - .collect::>()?; - // Keep iff this hit is the newest visible version of its PK. - keep.push( - self.index_store - .pk_is_newest(&values, position, max_visible_row), - ); - } - filter_record_batch(&batch, &BooleanArray::from(keep)) - .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) - } -} - -impl Stream for NewestPkFilterStream { - type Item = DFResult; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - match self.input.poll_next_unpin(cx) { - Poll::Ready(Some(Ok(batch))) => Poll::Ready(Some(self.filter_batch(batch))), - other => other, - } - } -} - -impl datafusion::physical_plan::RecordBatchStream for NewestPkFilterStream { - fn schema(&self) -> SchemaRef { - self.schema.clone() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow_array::Int32Array; - use arrow_schema::{DataType, Field, Schema}; - use datafusion::prelude::SessionContext; - use datafusion_physical_plan::test::TestMemoryExec; - use futures::TryStreamExt; - - /// Single-column `id` PK batch, one per append so a caller can control - /// row-level visibility via `max_visible_batch_position`. - fn id_batch(id: i32) -> RecordBatch { - let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); - RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![id]))]).unwrap() - } - - /// Index-search "hits": `(id, _rowid)` pairs the filter evaluates. - fn hits(rows: &[(i32, u64)]) -> RecordBatch { - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new(lance_core::ROW_ID, DataType::UInt64, true), - ])); - let ids: Vec = rows.iter().map(|(id, _)| *id).collect(); - let rowids: Vec = rows.iter().map(|(_, p)| *p).collect(); - RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(ids)), - Arc::new(UInt64Array::from(rowids)), - ], - ) - .unwrap() - } - - /// Build an active memtable whose PK index + BatchStore hold one row per - /// `id` in `appended` (positions 0..n), all committed. - fn active(appended: &[i32]) -> (Arc, Arc) { - let batch_store = Arc::new(BatchStore::with_capacity(16)); - let mut index = IndexStore::new(); - index.enable_pk_index(&[("id".to_string(), 0)]); - for &id in appended { - let b = id_batch(id); - let (bp, off, _) = batch_store.append(b.clone()).unwrap(); - index.insert_with_batch_position(&b, off, Some(bp)).unwrap(); - } - (Arc::new(index), batch_store) - } - - async fn run( - index_store: Arc, - batch_store: Arc, - max_visible_batch_position: usize, - hits_batch: RecordBatch, - ) -> Vec<(i32, u64)> { - let input = - TestMemoryExec::try_new_exec(&[vec![hits_batch.clone()]], hits_batch.schema(), None) - .unwrap(); - let exec = NewestPkFilterExec::new( - input, - vec!["id".to_string()], - lance_core::ROW_ID, - index_store, - batch_store, - max_visible_batch_position, - ); - let ctx = SessionContext::new(); - let out: Vec = exec - .execute(0, ctx.task_ctx()) - .unwrap() - .try_collect() - .await - .unwrap(); - let mut rows = Vec::new(); - for b in &out { - let ids = b.column(0).as_any().downcast_ref::().unwrap(); - let pos = b.column(1).as_any().downcast_ref::().unwrap(); - for i in 0..b.num_rows() { - rows.push((ids.value(i), pos.value(i))); - } - } - rows - } - - #[tokio::test] - async fn keeps_only_the_newest_visible_position_per_pk() { - // id=1 written at positions 0 and 2 (an update), id=2 at position 1; all - // visible. A stale hit (id=1 @ 0) is dropped; the newest (id=1 @ 2) and - // the unrelated id=2 survive — even though all three were "returned" by - // the index search. - let (index, store) = active(&[1, 2, 1]); - let rows = run(index, store, 2, hits(&[(1, 0), (2, 1), (1, 2)])).await; - assert_eq!(rows, vec![(2, 1), (1, 2)]); - } - - #[tokio::test] - async fn does_not_vanish_a_visible_row_under_a_newer_invisible_write() { - // The store/index hold id=1 at positions 0 and 2, but the query latched - // `max_visible_batch_position = 0` (only position 0 visible) — i.e. the - // update at position 2 was committed *after* this query's snapshot. The - // visible older row (id=1 @ 0) must be KEPT (its newest *visible* version - // is itself), not dropped because of the not-yet-visible position 2. - let (index, store) = active(&[1, 2, 1]); - let kept = run(index.clone(), store.clone(), 0, hits(&[(1, 0)])).await; - assert_eq!(kept, vec![(1, 0)], "visible row must not vanish"); - - // And the not-yet-visible position is itself dropped (outside snapshot). - let dropped = run(index, store, 0, hits(&[(1, 2)])).await; - assert!( - dropped.is_empty(), - "row beyond the snapshot must be dropped" - ); - } - - #[tokio::test] - async fn passes_through_when_no_pk_index() { - // A memtable without a primary-key index can't be deduped here, so the - // filter is a pass-through rather than dropping everything. - let batch_store = Arc::new(BatchStore::with_capacity(16)); - batch_store.append(id_batch(1)).unwrap(); - let index = Arc::new(IndexStore::new()); // no enable_pk_index - let rows = run(index, batch_store, 0, hits(&[(1, 0), (1, 9)])).await; - assert_eq!(rows, vec![(1, 0), (1, 9)]); - } -} diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec/pk.rs b/rust/lance/src/dataset/mem_wal/scanner/exec/pk.rs index 0707eb5e8dd..af1e7696181 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/exec/pk.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/exec/pk.rs @@ -4,7 +4,7 @@ //! Shared primary-key helpers for the LSM scanner execution nodes. //! //! Centralizes PK column resolution and per-row hashing so that every -//! consumer (e.g. [`super::PkBlockFilterExec`], [`super::NewestPkFilterExec`]) +//! consumer (e.g. [`super::PkBlockFilterExec`]) //! resolves and hashes a primary key the same way. The row hash is kept //! consistent with the variants supported by [`super::compute_pk_hash_from_scalars`] //! so a single PK produces the same hash regardless of which exec consumes it. diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index e7c8d205d5d..d783540bf6b 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -44,17 +44,18 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::union::UnionExec; +use datafusion::prelude::Expr; use lance_core::{Error, Result, is_system_column}; use lance_index::scalar::FullTextSearchQuery; -use lance_index::scalar::inverted::query::FtsQuery as IndexFtsQuery; +use lance_index::scalar::inverted::query::{FtsQuery as IndexFtsQuery, Operator}; use tracing::instrument; use super::block_list::compute_source_block_lists; use super::collector::LsmDataSourceCollector; use super::data_source::LsmDataSource; -use super::exec::{NewestPkFilterExec, PkBlockFilterExec}; +use super::exec::PkBlockFilterExec; use super::flushed_cache::{DatasetCache, GenerationWarmer, open_flushed_dataset}; -use super::projection::project_to_canonical; +use super::projection::{project_to_canonical, validate_projection_names}; use crate::dataset::mem_wal::memtable::scanner::MemTableScanner; use crate::session::Session; @@ -68,6 +69,41 @@ pub const SCORE_COLUMN: &str = "_score"; /// so a blocked source still yields `k` live rows after the block-list filter. const DEFAULT_OVERFETCH_FACTOR: f64 = 1.0; +fn validate_lsm_fts_query(query: &FullTextSearchQuery) -> Result<()> { + match &query.query { + IndexFtsQuery::Match(m) => { + if m.fuzziness != Some(0) && m.operator != Operator::Or { + return Err(Error::not_supported( + "LSM fuzzy full-text search only supports OR match operators".to_string(), + )); + } + Ok(()) + } + IndexFtsQuery::Phrase(_) => Ok(()), + _ => Err(Error::not_supported( + "LSM full-text search only supports match and phrase leaf queries".to_string(), + )), + } +} + +fn active_source_can_execute_fts(source: &LsmDataSource, column: &str) -> bool { + match source { + LsmDataSource::ActiveMemTable { + batch_store, + index_store, + .. + } => { + index_store + .get_fts_by_column(column) + .is_some_and(|index| !index.is_empty()) + && batch_store + .max_visible_row(index_store.max_visible_batch_position()) + .is_some() + } + _ => false, + } +} + /// Plans local-scoring FTS queries over LSM data. pub struct LsmFtsSearchPlanner { collector: LsmDataSourceCollector, @@ -79,8 +115,12 @@ pub struct LsmFtsSearchPlanner { flushed_cache: Option>, /// Optional warmer fired on first open of a flushed generation. warmer: Option>, - /// Over-fetch multiple for blocked sources (clamped to `>= 1.0`). + /// Over-fetch multiple for blocked sources. overfetch_factor: f64, + /// Optional prefilter predicate applied to every source arm so FTS hits + /// failing the predicate are dropped. Base/flushed arms use the dataset + /// scanner's native filter; memtable arms filter the materialized hits. + filter: Option, } impl LsmFtsSearchPlanner { @@ -98,11 +138,21 @@ impl LsmFtsSearchPlanner { flushed_cache: None, warmer: None, overfetch_factor: DEFAULT_OVERFETCH_FACTOR, + filter: None, } } + /// Attach an optional prefilter predicate. Every source arm restricts its + /// FTS hits to rows matching the predicate, matching a normal filtered + /// full-text scan over base ∪ flushed ∪ in-memory data. + pub fn with_filter(mut self, filter: Option) -> Self { + self.filter = filter; + self + } + /// Set the over-fetch multiple for blocked sources so they still yield `k` - /// live rows after cross-generation block-list filtering. Clamped to `>= 1.0`. + /// live rows after cross-generation block-list filtering. Values below + /// `1.0` are rejected by [`Self::plan_search`]. pub fn with_overfetch_factor(mut self, factor: f64) -> Self { self.overfetch_factor = factor; self @@ -136,28 +186,36 @@ impl LsmFtsSearchPlanner { /// * `query` — the FTS query (match / phrase / boolean / fuzzy for /// base/flushed Lance sources; the active memtable currently /// supports `MatchQuery`). - /// * `k` — global top-k to return. + /// * `limit` — optional global top-k to return. /// * `projection` — user columns to project. PK columns are /// auto-included; `_score` is always appended. /// /// Each source is scored independently (local BM25), normalized to a - /// canonical schema, unioned, and merged by `_score` DESC with a - /// top-k cap pushed into each partition. + /// canonical schema, unioned, and merged by `_score` DESC. When a finite + /// limit is supplied, top-k caps are pushed into each partition. #[instrument( name = "lsm_fts_search", level = "info", skip_all, - fields(column = %column, k) + fields(column = %column, limit) )] pub async fn plan_search( &self, column: &str, query: FullTextSearchQuery, - k: usize, + limit: Option, projection: Option<&[String]>, ) -> Result> { let sources = self.collector.collect()?; + if sources + .iter() + .any(|source| active_source_can_execute_fts(source, column)) + { + validate_lsm_fts_query(&query)?; + } + validate_projection_names(projection, &self.base_schema, &[SCORE_COLUMN])?; let target_schema = self.canonical_fts_schema(projection); + let overfetch = super::validate_overfetch_factor(self.overfetch_factor)?; if sources.is_empty() { return self.empty_plan(&target_schema); @@ -173,7 +231,6 @@ impl LsmFtsSearchPlanner { self.flushed_cache.as_ref(), )) .await?; - let overfetch = self.overfetch_factor.max(1.0); // Stage the per-source over-fetch decisions, then build every source // plan concurrently — the builds are independent and a sequential loop @@ -183,48 +240,58 @@ impl LsmFtsSearchPlanner { .map(|source| { let is_active = matches!(source, LsmDataSource::ActiveMemTable { .. }); let blocked = block_lists.get(&(source.shard_id(), source.generation())); - // Over-fetch a blocked source so the post-filter still yields k live - // rows. The active arm returns all matches (no builder limit), so its - // within-source dedup needs no over-fetch hint. - let fetch_k = if blocked.is_some() { - ((k as f64) * overfetch).ceil() as usize + let active_needs_uncapped_recency = match source { + LsmDataSource::ActiveMemTable { index_store, .. } + if !self.pk_columns.is_empty() => + { + !index_store.has_pk_index() || index_store.pk_has_overrides() + } + _ => false, + }; + // Active PK arms only need to stay uncapped when recency + // filtering may drop hits. Append-only PK memtables can safely + // pass the limit through and let FtsIndexExec use WAND/top-k. + // Blocked non-active sources use heuristic over-fetch because + // their newer generation membership may also drop candidates. + let fetch_limit = if active_needs_uncapped_recency { + None + } else if blocked.is_some() && !self.pk_columns.is_empty() { + limit.map(|limit| ((limit as f64) * overfetch).ceil() as usize) } else { - k + limit }; - (source, is_active, blocked, fetch_k) + (source, is_active, blocked, fetch_limit) }) .collect(); let built = - futures::future::try_join_all(arm_inputs.iter().map(|(source, _, _, fetch_k)| { - Box::pin(self.build_source_plan(source, column, &query, *fetch_k, projection)) + futures::future::try_join_all(arm_inputs.iter().map(|(source, _, _, fetch_limit)| { + Box::pin(self.build_source_plan(source, column, &query, *fetch_limit, projection)) })) .await?; let mut per_source_plans: Vec> = Vec::with_capacity(sources.len()); - for ((_, is_active, blocked, _), plan) in arm_inputs.iter().zip(built) { - let is_active = *is_active; + for ((_, _, blocked, _), plan) in arm_inputs.iter().zip(built) { let blocked = *blocked; // Dedup, mirroring LsmVectorSearchPlanner: - // * active: already wrapped in `NewestPkFilterExec` inside - // `build_source_plan` (drops predicate-crossing stale hits, which a - // result-set dedup can't catch). - // * flushed/base: drop rows superseded by a newer generation via the - // block-list (within-gen is handled by the flushed deletion vector). - let deduped = if is_active { - plan - } else if let Some(set) = blocked { + // * each memtable: `FtsIndexExec` drops superseded PK versions + // within that memtable before the query limit whenever PK + // columns are present. + // * any source with a block-list: drop rows superseded by a newer + // generation, including frozen in-memory memtables. + let deduped = if let Some(set) = blocked + && !self.pk_columns.is_empty() + { Arc::new(PkBlockFilterExec::new( plan, self.pk_columns.clone(), set.clone(), - k, + limit.unwrap_or(usize::MAX), )) as Arc } else { plan }; - // Normalize to canonical. This also drops the active arm's _rowid, - // which the canonical FTS schema omits — it served only the dedup. + // Normalize to the canonical FTS schema before merging sources. let normalized = project_to_canonical(deduped, &target_schema)?; per_source_plans.push(normalized); } @@ -265,10 +332,10 @@ impl LsmFtsSearchPlanner { let per_partition_sorted: Arc = Arc::new( SortExec::new(lex_ordering.clone(), merged) .with_preserve_partitioning(true) - .with_fetch(Some(k)), + .with_fetch(limit), ); let merged_sorted: Arc = Arc::new( - SortPreservingMergeExec::new(lex_ordering, per_partition_sorted).with_fetch(Some(k)), + SortPreservingMergeExec::new(lex_ordering, per_partition_sorted).with_fetch(limit), ); Ok(merged_sorted) @@ -279,7 +346,7 @@ impl LsmFtsSearchPlanner { source: &LsmDataSource, column: &str, query: &FullTextSearchQuery, - k: usize, + limit: Option, projection: Option<&[String]>, ) -> Result> { match source { @@ -287,10 +354,19 @@ impl LsmFtsSearchPlanner { let mut scanner = dataset.scan(); let cols = self.fts_scanner_projection(projection); scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; - let bound_query = query - .clone() - .with_column(column.to_string())? - .limit(Some(k as i64)); + if let Some(ref filter) = self.filter { + // `prefilter(true)` is required: without it the scanner + // post-filters the unfiltered BM25 top-k, dropping matching + // rows that scored below non-matching ones. + scanner.filter_expr(filter.clone()); + scanner.prefilter(true); + } + let mut bound_query = query.clone().with_column(column.to_string())?; + if let Some(limit) = limit { + bound_query = bound_query.limit(Some(limit as i64)); + } else { + bound_query = bound_query.limit(None); + } scanner.full_text_search(bound_query)?; scanner.create_plan().await } @@ -305,10 +381,18 @@ impl LsmFtsSearchPlanner { let mut scanner = dataset.scan(); let cols = self.fts_scanner_projection(projection); scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; - let bound_query = query - .clone() - .with_column(column.to_string())? - .limit(Some(k as i64)); + if let Some(ref filter) = self.filter { + // See the base arm: `prefilter(true)` makes this a true + // prefilter rather than a lossy post-filter on the BM25 top-k. + scanner.filter_expr(filter.clone()); + scanner.prefilter(true); + } + let mut bound_query = query.clone().with_column(column.to_string())?; + if let Some(limit) = limit { + bound_query = bound_query.limit(Some(limit as i64)); + } else { + bound_query = bound_query.limit(None); + } scanner.full_text_search(bound_query)?; scanner.create_plan().await } @@ -318,47 +402,36 @@ impl LsmFtsSearchPlanner { schema, .. } => { + if !active_source_can_execute_fts(source, column) { + return self.empty_plan(&self.canonical_fts_schema(projection)); + } + validate_lsm_fts_query(query)?; let mut scanner = MemTableScanner::new(batch_store.clone(), index_store.clone(), schema.clone()); let cols = self.fts_scanner_projection(projection); - scanner.project(&cols.iter().map(|s| s.as_str()).collect::>()); - // Expose the row position so the recency filter can identify the - // newest visible version of each PK. The append-only inverted - // index keeps an updated row's old postings live, so a stale hit - // can match a query the fresh row no longer does; the filter - // drops it. `project_to_canonical` strips `_rowid` afterward. - scanner.with_row_id(); - // `MemTableScanner::full_text_search` takes a raw match - // string; richer query shapes (phrase/boolean/fuzzy) can - // be plumbed through once the MemTable scanner accepts a - // structured query. - let match_str = match &query.query { - IndexFtsQuery::Match(m) => m.terms.clone(), - other => { - return Err(Error::not_supported(format!( - "Active memtable FTS via LsmFtsSearchPlanner currently only \ - supports MatchQuery, got: {other:?}" - ))); - } - }; - let _ = scanner.full_text_search(column, &match_str); - // Active arm doesn't take a top-K hint via the builder - // today; the per-partition Sort+fetch above bounds the - // emitted rows. - let _ = k; - let plan = scanner.create_plan().await?; - // Drop predicate-crossing stale hits: keep a hit iff it is the - // newest visible version of its PK (collapses duplicate-PK - // appends too — supersedes the old WithinSourceDedupExec). - let filtered: Arc = Arc::new(NewestPkFilterExec::new( - plan, - self.pk_columns.clone(), - lance_core::ROW_ID, - index_store.clone(), - batch_store.clone(), - scanner.max_visible_batch_position(), - )); - Ok(filtered) + scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; + if let Some(ref filter) = self.filter { + // Honored inside `plan_fts_search`: the materialized hits are + // masked by the predicate before projection. + scanner.filter_expr(filter.clone()); + } + // The append-only inverted index keeps an updated row's old + // postings live, so the memtable FTS exec needs PK columns to + // drop stale hits before it applies the query limit. + if !self.pk_columns.is_empty() { + scanner.with_pk_columns(self.pk_columns.clone()); + } + // `MemTableScanner::full_text_search` now takes a structured + // `FullTextSearchQuery` (match/phrase); it rejects compound + // shapes the MemTable path can't model. + let mut bound_query = query.clone().with_column(column.to_string())?; + if let Some(limit) = limit { + bound_query = bound_query.limit(Some(limit as i64)); + } else { + bound_query = bound_query.limit(None); + } + scanner.full_text_search(bound_query)?; + scanner.create_plan().await } } } @@ -435,7 +508,7 @@ mod tests { use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; use crate::dataset::{Dataset, WriteParams}; - use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, StringArray}; + use arrow_array::{BooleanArray, Int32Array, RecordBatch, RecordBatchIterator, StringArray}; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; use futures::TryStreamExt; use std::collections::HashMap; @@ -453,6 +526,20 @@ mod tests { ])) } + fn fts_tombstone_schema() -> Arc { + let mut id_meta = HashMap::new(); + id_meta.insert( + "lance-schema:unenforced-primary-key".to_string(), + "true".to_string(), + ); + let id_field = Field::new("id", DataType::Int32, false).with_metadata(id_meta); + Arc::new(ArrowSchema::new(vec![ + id_field, + Field::new("text", DataType::Utf8, true), + Field::new(crate::dataset::mem_wal::TOMBSTONE, DataType::Boolean, false), + ])) + } + fn make_batch(schema: &ArrowSchema, ids: &[i32], texts: &[&str]) -> RecordBatch { RecordBatch::try_new( Arc::new(schema.clone()), @@ -464,12 +551,61 @@ mod tests { .unwrap() } + fn make_tombstone_batch( + schema: &ArrowSchema, + rows: &[(i32, Option<&str>, bool)], + ) -> RecordBatch { + let ids: Vec = rows.iter().map(|(id, _, _)| *id).collect(); + let texts: Vec> = rows.iter().map(|(_, text, _)| *text).collect(); + let tombstones: Vec = rows.iter().map(|(_, _, tombstone)| *tombstone).collect(); + RecordBatch::try_new( + Arc::new(schema.clone()), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(texts)), + Arc::new(BooleanArray::from(tombstones)), + ], + ) + .unwrap() + } + async fn write_dataset(uri: &str, batches: Vec) -> Dataset { let schema = batches[0].schema(); - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema); - Dataset::write(reader, uri, Some(WriteParams::default())) + let has_id = schema.column_with_name("id").is_some(); + let reader = RecordBatchIterator::new(batches.clone().into_iter().map(Ok), schema); + let dataset = Dataset::write(reader, uri, Some(WriteParams::default())) .await - .unwrap() + .unwrap(); + if has_id { + crate::dataset::mem_wal::scanner::block_list::write_pk_sidecar(uri, &batches, &["id"]) + .await + .unwrap(); + } + dataset + } + + #[tokio::test] + async fn rejects_missing_projection_column() { + let schema = fts_schema(); + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]); + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema); + + let projection = vec!["missing".to_string()]; + let err = planner + .plan_search( + "text", + FullTextSearchQuery::new("lance".to_string()), + Some(1), + Some(&projection), + ) + .await + .unwrap_err(); + assert!( + err.to_string().contains("missing"), + "unexpected missing-column projection error: {err}" + ); } #[tokio::test] @@ -543,7 +679,7 @@ mod tests { .plan_search( "text", FullTextSearchQuery::new("lance".to_string()), - 10, + Some(10), None, ) .await @@ -587,190 +723,1390 @@ mod tests { assert!(ids.contains(&3), "missing active hit id=3; got ids={ids:?}"); } + /// A prefilter on a full-text search must drop hits failing the predicate + /// from both the base arm (native scanner filter) and the active memtable + /// arm (materialized-hit mask), even though they match the query text. #[tokio::test] - async fn local_mode_active_memtable_only_returns_score_sorted_hits() { + async fn prefilter_drops_nonmatching_hits_across_base_and_active() { + use crate::index::DatasetIndexExt; + use datafusion::prelude::{col, lit}; + use lance_index::IndexType; + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + let schema = fts_schema(); + let tmp = tempfile::tempdir().unwrap(); + + // Base rows 1 and 2 both contain "lance". + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let mut base_ds = write_dataset( + &base_uri, + vec![make_batch(&schema, &[1, 2], &["lance one", "lance two"])], + ) + .await; + base_ds + .create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + let base_ds = Arc::new(Dataset::open(&base_uri).await.unwrap()); + + // Active memtable rows 0 and 3 both contain "lance". let batch_store = Arc::new(BatchStore::with_capacity(16)); let mut indexes = IndexStore::new(); - // text column has field_id 1 in fts_schema() + indexes.enable_pk_index(&[("id".to_string(), 0)]); indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); - let batch = make_batch( - &schema, - &[1, 2, 3, 4], - &[ - "lance is a columnar data format", - "memwal handles streaming writes", - "lance memwal lance lance", - "completely unrelated", - ], - ); - batch_store.append(batch.clone()).unwrap(); + let active_batch = make_batch(&schema, &[0, 3], &["lance zero", "lance three"]); + batch_store.append(active_batch.clone()).unwrap(); indexes - .insert_with_batch_position(&batch, 0, Some(0)) + .insert_with_batch_position(&active_batch, 0, Some(0)) .unwrap(); let indexes = Arc::new(indexes); - let tmp = tempfile::tempdir().unwrap(); - let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); - let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) - .with_in_memory_memtables( - uuid::Uuid::new_v4(), - InMemoryMemTables { - active: InMemoryMemTableRef { - batch_store, - index_store: indexes, - schema: schema.clone(), - generation: 1, - }, - frozen: vec![], + let collector = LsmDataSourceCollector::new(base_ds, vec![]).with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, }, - ); + frozen: vec![], + }, + ); - let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema); + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema) + // `id >= 2` keeps base id=2 and active id=3; drops base id=1 and active id=0. + .with_filter(Some(col("id").gt_eq(lit(2i32)))); let plan = planner .plan_search( "text", FullTextSearchQuery::new("lance".to_string()), - 10, + Some(10), None, ) .await - .expect("local mode planner should produce a plan"); + .expect("planner should produce a filtered base+active plan"); - // Plan executes and emits _score-sorted rows. let ctx = datafusion::prelude::SessionContext::new(); let stream = plan.execute(0, ctx.task_ctx()).unwrap(); let batches: Vec = stream.try_collect().await.unwrap(); - let total: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert!( - total >= 2, - "expected at least the 2 'lance' rows, got {total}" - ); - - // Schema must include _score and the PK id. - let out = batches[0].schema(); - assert!(out.field_with_name(SCORE_COLUMN).is_ok()); - assert!(out.field_with_name("id").is_ok()); - // _score must be non-ascending across the result. - let mut prev_score: Option = None; - for batch in &batches { - let score = batch - .column_by_name(SCORE_COLUMN) + let mut ids: Vec = Vec::new(); + for b in &batches { + let col = b + .column_by_name("id") .unwrap() .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(); - for i in 0..batch.num_rows() { - let s = score.value(i); - if let Some(p) = prev_score { - assert!(p >= s, "scores not sorted DESC: {p} then {s}"); - } - prev_score = Some(s); + for i in 0..b.num_rows() { + ids.push(col.value(i)); } } + ids.sort(); + assert_eq!( + ids, + vec![2, 3], + "prefilter must keep only id>=2 across base (id=2) and active (id=3), \ + dropping base id=1 and active id=0; got {ids:?}" + ); } + /// The flushed arm must apply the filter as a true FTS prefilter, and that + /// prefiltered candidate set must compose with cross-generation block-list + /// filtering plus over-fetch. Gen 1's best predicate-matching hit (id=3) is + /// superseded by gen 2; with over-fetch, gen 1 should still contribute id=4. #[tokio::test] - async fn local_mode_active_dedups_updated_pk_keeping_newest() { - // The active memtable is an append log and the FTS index is - // append-only, so a PK updated before flush is searchable as two - // row-positions. WithinSourceDedupExec(KeepMaxRowAddr) must collapse - // them to the newest insert. Without it the same PK would surface - // twice (criterion 2 violation). + async fn prefilter_on_flushed_composes_with_block_list() { + use crate::dataset::mem_wal::scanner::data_source::ShardSnapshot; + use crate::index::DatasetIndexExt; + use datafusion::prelude::{col, lit}; + use lance_index::IndexType; + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + let schema = fts_schema(); - let batch_store = Arc::new(BatchStore::with_capacity(16)); - let mut indexes = IndexStore::new(); - indexes.enable_pk_index(&[("id".to_string(), 0)]); - indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let shard_id = uuid::Uuid::new_v4(); - // First append (positions 0,1): id=1 is the stale version of the PK. - let batch_old = make_batch(&schema, &[1, 2], &["lance stale version", "other doc"]); - batch_store.append(batch_old.clone()).unwrap(); - indexes - .insert_with_batch_position(&batch_old, 0, Some(0)) - .unwrap(); + // Gen 1: id=1 matches strongly but fails the predicate. id=3 matches + // strongly but is stale (blocked by gen 2). id=4 is the next live + // predicate match that only survives if the flushed arm prefilters and + // over-fetches before the block-list drops id=3. + let gen1_uri = format!("{}/_mem_wal/{}/gen_1", base_uri, shard_id); + let mut gen1 = write_dataset( + &gen1_uri, + vec![make_batch( + &schema, + &[1, 3, 4], + &["lance lance lance lance", "lance lance lance", "lance"], + )], + ) + .await; + gen1.create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); - // Second append (position 2): id=1 updated — same PK, later row. - let batch_new = make_batch(&schema, &[1], &["lance fresh version"]); - batch_store.append(batch_new.clone()).unwrap(); - indexes - .insert_with_batch_position(&batch_new, 2, Some(1)) - .unwrap(); - let indexes = Arc::new(indexes); + // Gen 2: newer id=3 shadows gen 1's match but does not match the query. + let gen2_uri = format!("{}/_mem_wal/{}/gen_2", base_uri, shard_id); + let mut gen2 = + write_dataset(&gen2_uri, vec![make_batch(&schema, &[3], &["other text"])]).await; + gen2.create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); - let tmp = tempfile::tempdir().unwrap(); - let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); - let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) - .with_in_memory_memtables( - uuid::Uuid::new_v4(), - InMemoryMemTables { - active: InMemoryMemTableRef { - batch_store, - index_store: indexes, - schema: schema.clone(), - generation: 1, - }, - frozen: vec![], - }, - ); + let snapshot = ShardSnapshot::new(shard_id) + .with_current_generation(3) + .with_flushed_generation(1, "gen_1".to_string()) + .with_flushed_generation(2, "gen_2".to_string()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![snapshot]); - let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema); + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema) + .with_filter(Some(col("id").gt_eq(lit(3i32)))) + .with_overfetch_factor(2.0); let plan = planner .plan_search( "text", FullTextSearchQuery::new("lance".to_string()), - 10, + Some(1), None, ) .await - .expect("planner should produce an active-only plan"); + .expect("planner should produce a filtered flushed plan"); let ctx = datafusion::prelude::SessionContext::new(); let stream = plan.execute(0, ctx.task_ctx()).unwrap(); let batches: Vec = stream.try_collect().await.unwrap(); - let mut rows: Vec<(i32, String)> = Vec::new(); + let mut ids: Vec = Vec::new(); for b in &batches { - let ids = b + let col = b .column_by_name("id") .unwrap() .as_any() .downcast_ref::() .unwrap(); - let texts = b - .column_by_name("text") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); for i in 0..b.num_rows() { - rows.push((ids.value(i), texts.value(i).to_string())); + ids.push(col.value(i)); } } - - // id=1 must appear exactly once, and it must be the *newest* version. - let id1: Vec<&(i32, String)> = rows.iter().filter(|(id, _)| *id == 1).collect(); - assert_eq!( - id1.len(), - 1, - "updated PK id=1 must be deduped to one row; got {rows:?}" - ); assert_eq!( - id1[0].1, "lance fresh version", - "dedup must keep the newest (max row-position) version" + ids, + vec![4], + "flushed FTS prefilter should return live id=4 after stale id=3 is blocked; got {ids:?}" ); } #[tokio::test] - async fn active_stale_update_predicate_crossing_leaks() { - // A PK update that crosses out of the match set: pk=1 inserted as - // "alpha lance", then updated to "beta lance". The append-only inverted - // index keeps the old "alpha" posting live, so an "alpha" search still - // matches the STALE pk=1 row — and the fresh "beta lance" row isn't even - // a candidate, so a result-set dedup has nothing to suppress it against. - // `NewestPkFilterExec` drops it predicate-independently: pk=1's newest - // visible row is "beta lance", so the "alpha" hit is not the newest. + async fn active_tombstone_masks_base_fts_hit() { + use crate::index::DatasetIndexExt; + use lance_index::IndexType; + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + + let base_schema = fts_schema(); + let mem_schema = fts_tombstone_schema(); + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + + let mut base = write_dataset( + &base_uri, + vec![make_batch(&base_schema, &[1, 2], &["lance lance", "lance"])], + ) + .await; + base.create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + + let active_tombstone = make_tombstone_batch(&mem_schema, &[(1, None, true)]); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut index_store = IndexStore::new(); + index_store.enable_pk_index(&[("id".to_string(), 0)]); + index_store.add_fts("text_fts".to_string(), 1, "text".to_string()); + let (_, row_offset, batch_position) = batch_store.append(active_tombstone.clone()).unwrap(); + index_store + .insert_with_batch_position(&active_tombstone, row_offset, Some(batch_position)) + .unwrap(); + let index_store = Arc::new(index_store); + + let collector = LsmDataSourceCollector::new(Arc::new(base), vec![]) + .with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store, + schema: mem_schema, + generation: 1, + }, + frozen: vec![], + }, + ); + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], base_schema) + .with_overfetch_factor(2.0); + + let plan = planner + .plan_search( + "text", + FullTextSearchQuery::new("lance".to_string()), + Some(1), + None, + ) + .await + .unwrap(); + let stream = plan + .execute(0, datafusion::prelude::SessionContext::new().task_ctx()) + .unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut ids: Vec = Vec::new(); + for b in &batches { + let col = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + ids.push(col.value(i)); + } + } + assert_eq!( + ids, + vec![2], + "active tombstone for id=1 must block the older base FTS hit; got {ids:?}" + ); + } + + #[tokio::test] + async fn active_update_and_tombstone_mask_frozen_fts_hits() { + let base_schema = fts_schema(); + let mem_schema = fts_tombstone_schema(); + + let frozen_batch = make_tombstone_batch( + &mem_schema, + &[ + (1, Some("lance stale update"), false), + (2, Some("lance live"), false), + (3, Some("lance deleted"), false), + ], + ); + let frozen_batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut frozen_index_store = IndexStore::new(); + frozen_index_store.enable_pk_index(&[("id".to_string(), 0)]); + frozen_index_store.add_fts("text_fts".to_string(), 1, "text".to_string()); + let (_, frozen_row_offset, frozen_batch_position) = + frozen_batch_store.append(frozen_batch.clone()).unwrap(); + frozen_index_store + .insert_with_batch_position( + &frozen_batch, + frozen_row_offset, + Some(frozen_batch_position), + ) + .unwrap(); + + let active_batch = make_tombstone_batch( + &mem_schema, + &[(1, Some("fresh other text"), false), (3, None, true)], + ); + let active_batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut active_index_store = IndexStore::new(); + active_index_store.enable_pk_index(&[("id".to_string(), 0)]); + active_index_store.add_fts("text_fts".to_string(), 1, "text".to_string()); + let (_, active_row_offset, active_batch_position) = + active_batch_store.append(active_batch.clone()).unwrap(); + active_index_store + .insert_with_batch_position( + &active_batch, + active_row_offset, + Some(active_batch_position), + ) + .unwrap(); + + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let shard_id = uuid::Uuid::new_v4(); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) + .with_in_memory_memtables( + shard_id, + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store: active_batch_store, + index_store: Arc::new(active_index_store), + schema: mem_schema.clone(), + generation: 2, + }, + frozen: vec![InMemoryMemTableRef { + batch_store: frozen_batch_store, + index_store: Arc::new(frozen_index_store), + schema: mem_schema, + generation: 1, + }], + }, + ); + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], base_schema); + + let plan = planner + .plan_search( + "text", + FullTextSearchQuery::new("lance".to_string()), + Some(10), + None, + ) + .await + .unwrap(); + let stream = plan + .execute(0, datafusion::prelude::SessionContext::new().task_ctx()) + .unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut ids: Vec = Vec::new(); + for b in &batches { + let col = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + ids.push(col.value(i)); + } + } + ids.sort_unstable(); + assert_eq!( + ids, + vec![2], + "active update id=1 and tombstone id=3 must block stale frozen FTS hits; got {ids:?}" + ); + } + + #[tokio::test] + async fn active_filtered_search_without_pk_applies_small_limit_after_filter() { + use datafusion::prelude::{col, lit}; + + let schema = fts_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut indexes = IndexStore::new(); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + let active_batch = make_batch( + &schema, + &[1, 2, 3], + &["lance", "lance filler", "lance filler filler"], + ); + batch_store.append(active_batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&active_batch, 0, Some(0)) + .unwrap(); + let indexes = Arc::new(indexes); + + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) + .with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmFtsSearchPlanner::new(collector, vec![], schema) + .with_filter(Some(col("id").gt_eq(lit(1i32)))); + let plan = planner + .plan_search( + "text", + FullTextSearchQuery::new("lance".to_string()), + Some(2), + None, + ) + .await + .expect("planner should produce an active-only filtered plan"); + + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut ids: Vec = Vec::new(); + for b in &batches { + let col = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + ids.push(col.value(i)); + } + } + ids.sort_unstable(); + assert_eq!( + ids, + vec![1, 2], + "no-PK filtered active search must apply the limit after filtering \ + and keep the top-scoring matching hits; got ids={ids:?}" + ); + } + + #[tokio::test] + async fn fuzzy_and_query_is_rejected_when_active_memtable_is_present() { + use lance_index::scalar::inverted::query::{ + FtsQuery as IndexFtsQuery, MatchQuery, Operator, + }; + + let schema = fts_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut indexes = IndexStore::new(); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + let active_batch = make_batch(&schema, &[1], &["lance memwal"]); + let (_, row_offset, batch_position) = batch_store.append(active_batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&active_batch, row_offset, Some(batch_position)) + .unwrap(); + let indexes = Arc::new(indexes); + + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) + .with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + let planner = LsmFtsSearchPlanner::new(collector, vec![], schema); + let query = FullTextSearchQuery::new_query(IndexFtsQuery::Match( + MatchQuery::new("lance memwal".to_string()) + .with_operator(Operator::And) + .with_fuzziness(Some(1)), + )); + + let err = planner + .plan_search("text", query, Some(10), None) + .await + .expect_err("fuzzy AND should be rejected consistently"); + assert!( + err.to_string().contains("fuzzy full-text search"), + "unexpected error for fuzzy AND query: {err}" + ); + } + + #[tokio::test] + async fn base_only_boolean_query_uses_dataset_scanner_support() { + use crate::index::DatasetIndexExt; + use lance_index::IndexType; + use lance_index::scalar::inverted::query::{ + BooleanQuery, FtsQuery as IndexFtsQuery, MatchQuery, Occur, + }; + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + + let schema = fts_schema(); + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let mut base_ds = write_dataset( + &base_uri, + vec![make_batch( + &schema, + &[1, 2], + &["lance rocks", "unrelated text"], + )], + ) + .await; + base_ds + .create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + let base_ds = Arc::new(Dataset::open(&base_uri).await.unwrap()); + let collector = LsmDataSourceCollector::new(base_ds, vec![]); + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema); + + let query = FullTextSearchQuery::new_query(IndexFtsQuery::Boolean(BooleanQuery::new( + vec![(Occur::Must, MatchQuery::new("lance".to_string()).into())], + ))); + let plan = planner + .plan_search("text", query, Some(10), Some(&["id".to_string()])) + .await + .expect("base-only boolean query should be delegated to dataset scanner"); + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + let ids: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!(ids, vec![1]); + } + + #[tokio::test] + async fn boolean_query_ignores_active_memtable_without_relevant_fts_index() { + use crate::index::DatasetIndexExt; + use lance_index::IndexType; + use lance_index::scalar::inverted::query::{ + BooleanQuery, FtsQuery as IndexFtsQuery, MatchQuery, Occur, + }; + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + + let schema = fts_schema(); + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let mut base_ds = write_dataset( + &base_uri, + vec![make_batch( + &schema, + &[1, 2], + &["lance rocks", "unrelated text"], + )], + ) + .await; + base_ds + .create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + let base_ds = Arc::new(Dataset::open(&base_uri).await.unwrap()); + + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + let active_batch = make_batch(&schema, &[99], &["active text has no fts index"]); + let (_, row_offset, batch_position) = batch_store.append(active_batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&active_batch, row_offset, Some(batch_position)) + .unwrap(); + let indexes = Arc::new(indexes); + let collector = LsmDataSourceCollector::new(base_ds, vec![]).with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema); + + let query = FullTextSearchQuery::new_query(IndexFtsQuery::Boolean(BooleanQuery::new( + vec![(Occur::Must, MatchQuery::new("lance".to_string()).into())], + ))); + let plan = planner + .plan_search("text", query, Some(10), Some(&["id".to_string()])) + .await + .expect("irrelevant active memtable must not reject base-supported boolean query"); + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + let ids: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!(ids, vec![1]); + } + + /// The base arm must apply the filter as a true *prefilter*, not a + /// post-filter on the BM25 top-k. With `k = 1` and the higher-scoring base + /// doc failing the predicate, a post-filter would return zero rows; a + /// prefilter restricts BM25 to matching rows and returns the lower-scoring + /// match. Regression for a missing `scanner.prefilter(true)` on the base arm. + #[tokio::test] + async fn prefilter_on_base_is_not_a_lossy_postfilter() { + use crate::index::DatasetIndexExt; + use datafusion::prelude::{col, lit}; + use lance_index::IndexType; + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + + let schema = fts_schema(); + let tmp = tempfile::tempdir().unwrap(); + + // id=1 is a short doc ("lance") so it scores higher under BM25 length + // normalization; id=2 buries "lance" among filler so it scores lower. + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let mut base_ds = write_dataset( + &base_uri, + vec![make_batch( + &schema, + &[1, 2], + &["lance", "lance filler filler filler filler filler"], + )], + ) + .await; + base_ds + .create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + let base_ds = Arc::new(Dataset::open(&base_uri).await.unwrap()); + + // Base-only collector (no in-memory memtables): isolates the base arm. + let collector = LsmDataSourceCollector::new(base_ds, vec![]); + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema) + // Keeps only id=2, which is the *lower*-scoring match. A post-filter + // on the top-1 (id=1) would drop everything. + .with_filter(Some(col("id").gt_eq(lit(2i32)))); + let plan = planner + .plan_search( + "text", + FullTextSearchQuery::new("lance".to_string()), + Some(1), + None, + ) + .await + .expect("planner should produce a filtered base plan"); + + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut ids: Vec = Vec::new(); + for b in &batches { + let col = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + ids.push(col.value(i)); + } + } + assert_eq!( + ids, + vec![2], + "prefilter must return the lower-scoring match id=2, not post-filter \ + the top-1 (id=1) down to nothing; got {ids:?}" + ); + } + + /// The active memtable FTS arm must also apply the predicate before its + /// top-k cap. With `k = 1` and the higher-scoring active doc failing the + /// predicate, pushing the limit into the index would return zero rows. + #[tokio::test] + async fn prefilter_on_active_is_not_a_lossy_postfilter() { + use datafusion::prelude::{col, lit}; + + let mut id_meta = HashMap::new(); + id_meta.insert( + "lance-schema:unenforced-primary-key".to_string(), + "true".to_string(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(id_meta), + Field::new("text", DataType::Utf8, true), + Field::new("status", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec![ + "lance", + "lance filler filler filler filler filler", + ])), + Arc::new(StringArray::from(vec!["archived", "active"])), + ], + ) + .unwrap(); + + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + let indexes = Arc::new(indexes); + + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) + .with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema) + .with_filter(Some(col("status").eq(lit("active")))); + let plan = planner + .plan_search( + "text", + FullTextSearchQuery::new("lance".to_string()), + Some(1), + None, + ) + .await + .expect("planner should produce a filtered active plan"); + + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut ids: Vec = Vec::new(); + for b in &batches { + let col = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + ids.push(col.value(i)); + } + } + assert_eq!( + ids, + vec![2], + "active FTS prefilter must return the lower-scoring matching row, \ + not post-filter the top-1 down to nothing; got {ids:?}" + ); + } + + /// The active FTS arm must also avoid capping before newest-PK filtering. + /// A stale high-scoring hit can be removed by `FtsIndexExec`; a lower + /// scoring live hit must still be available for the final global top-k. + #[tokio::test] + async fn active_limit_applies_after_newest_pk_recency_filter() { + use datafusion::prelude::{col, lit}; + + let mut id_meta = HashMap::new(); + id_meta.insert( + "lance-schema:unenforced-primary-key".to_string(), + "true".to_string(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(id_meta), + Field::new("text", DataType::Utf8, true), + Field::new("status", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 1])), + Arc::new(StringArray::from(vec![ + "lance", + "lance filler filler filler filler filler", + "other text", + ])), + Arc::new(StringArray::from(vec!["active", "active", "active"])), + ], + ) + .unwrap(); + + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + let indexes = Arc::new(indexes); + + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) + .with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema) + .with_filter(Some(col("status").eq(lit("active")))); + let plan = planner + .plan_search( + "text", + FullTextSearchQuery::new("lance".to_string()), + Some(1), + None, + ) + .await + .expect("planner should produce a filtered active plan"); + + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut ids: Vec = Vec::new(); + for b in &batches { + let col = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + ids.push(col.value(i)); + } + } + assert_eq!( + ids, + vec![2], + "active FTS limit must apply after newest-PK filtering; got {ids:?}" + ); + } + + /// An in-memtable update whose *newest* version fails the prefilter must + /// exclude the PK, not leak the stale older hit that still passes. Both + /// versions of pk=5 match the query text "lance", but only the older one is + /// "active"; the current version is "archived" and must be dropped. + /// Regression for filter-before-dedup on the active FTS arm. + #[tokio::test] + async fn prefilter_excludes_pk_whose_newest_version_fails() { + use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; + use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; + use arrow_schema::{DataType, Field}; + use datafusion::prelude::{col, lit}; + + let mut id_meta = HashMap::new(); + id_meta.insert( + "lance-schema:unenforced-primary-key".to_string(), + "true".to_string(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(id_meta), + Field::new("text", DataType::Utf8, true), + Field::new("status", DataType::Utf8, false), + ])); + let make_row = |statuses: &[&str]| -> RecordBatch { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![5; statuses.len()])), + Arc::new(StringArray::from(vec!["lance text"; statuses.len()])), + Arc::new(StringArray::from(statuses.to_vec())), + ], + ) + .unwrap() + }; + + // Base unindexed → contributes nothing; isolate the active arm. + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let base_batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![999])), + Arc::new(StringArray::from(vec!["unrelated"])), + Arc::new(StringArray::from(vec!["active"])), + ], + ) + .unwrap(); + let base_ds = Arc::new(write_dataset(&base_uri, vec![base_batch]).await); + + // Active memtable: pk=5 appended twice (active then archived), both "lance". + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + let active_batch = make_row(&["active", "archived"]); + batch_store.append(active_batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&active_batch, 0, Some(0)) + .unwrap(); + let indexes = Arc::new(indexes); + + let collector = LsmDataSourceCollector::new(base_ds, vec![]).with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema) + .with_filter(Some(col("status").eq(lit("active")))); + let plan = planner + .plan_search( + "text", + FullTextSearchQuery::new("lance".to_string()), + Some(10), + None, + ) + .await + .expect("planner should produce a filtered active plan"); + + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total, 0, + "pk=5's current version is 'archived' and must be excluded; the stale \ + 'active' older hit must not leak (filter evaluated on newest version)" + ); + } + + /// Cross-arm stale hits must be blocked even if the newer active row fails + /// the prefilter. The base copy of pk=5 matches both text and status, but + /// the newer active copy is archived; pk=5 must not leak from the base arm. + #[tokio::test] + async fn prefilter_blocks_base_hit_when_active_newest_fails() { + use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; + use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; + use crate::index::DatasetIndexExt; + use datafusion::prelude::{col, lit}; + use lance_index::IndexType; + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + + let mut id_meta = HashMap::new(); + id_meta.insert( + "lance-schema:unenforced-primary-key".to_string(), + "true".to_string(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(id_meta), + Field::new("text", DataType::Utf8, true), + Field::new("status", DataType::Utf8, false), + ])); + let make_rows = |rows: &[(i32, &str, &str)]| -> RecordBatch { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from( + rows.iter().map(|(id, _, _)| *id).collect::>(), + )), + Arc::new(StringArray::from( + rows.iter().map(|(_, text, _)| *text).collect::>(), + )), + Arc::new(StringArray::from( + rows.iter() + .map(|(_, _, status)| *status) + .collect::>(), + )), + ], + ) + .unwrap() + }; + + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let mut base_ds = write_dataset( + &base_uri, + vec![make_rows(&[ + (5, "lance base stale", "active"), + (6, "lance base live", "active"), + ])], + ) + .await; + base_ds + .create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + let base_ds = Arc::new(Dataset::open(&base_uri).await.unwrap()); + + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + let active_batch = make_rows(&[(5, "lance active newest", "archived")]); + batch_store.append(active_batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&active_batch, 0, Some(0)) + .unwrap(); + let indexes = Arc::new(indexes); + + let collector = LsmDataSourceCollector::new(base_ds, vec![]).with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema) + .with_filter(Some(col("status").eq(lit("active")))); + let plan = planner + .plan_search( + "text", + FullTextSearchQuery::new("lance".to_string()), + Some(10), + None, + ) + .await + .expect("planner should produce a filtered base+active plan"); + + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + let mut ids = Vec::new(); + for batch in &batches { + let id_array = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + ids.push(id_array.value(row)); + } + } + ids.sort_unstable(); + assert_eq!( + ids, + vec![6], + "base pk=5 passes the filter but is superseded by active archived pk=5; got {ids:?}" + ); + } + + #[tokio::test] + async fn local_mode_active_memtable_only_returns_score_sorted_hits() { + let schema = fts_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + // text column has field_id 1 in fts_schema() + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + let batch = make_batch( + &schema, + &[1, 2, 3, 4], + &[ + "lance is a columnar data format", + "memwal handles streaming writes", + "lance memwal lance lance", + "completely unrelated", + ], + ); + batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + let indexes = Arc::new(indexes); + + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) + .with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema); + let plan = planner + .plan_search( + "text", + FullTextSearchQuery::new("lance".to_string()), + Some(10), + None, + ) + .await + .expect("local mode planner should produce a plan"); + + // Plan executes and emits _score-sorted rows. + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert!( + total >= 2, + "expected at least the 2 'lance' rows, got {total}" + ); + + // Schema must include _score and the PK id. + let out = batches[0].schema(); + assert!(out.field_with_name(SCORE_COLUMN).is_ok()); + assert!(out.field_with_name("id").is_ok()); + + // _score must be non-ascending across the result. + let mut prev_score: Option = None; + for batch in &batches { + let score = batch + .column_by_name(SCORE_COLUMN) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + let s = score.value(i); + if let Some(p) = prev_score { + assert!(p >= s, "scores not sorted DESC: {p} then {s}"); + } + prev_score = Some(s); + } + } + } + + #[tokio::test] + async fn active_match_query_preserves_and_operator() { + use lance_index::scalar::inverted::query::{ + FtsQuery as IndexFtsQuery, MatchQuery, Operator, + }; + + let schema = fts_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + let batch = make_batch( + &schema, + &[1, 2, 3], + &["lance only", "memwal only", "lance memwal"], + ); + batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + let indexes = Arc::new(indexes); + + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) + .with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema); + let query = FullTextSearchQuery::new_query(IndexFtsQuery::Match( + MatchQuery::new("lance memwal".to_string()) + .with_operator(Operator::And) + .with_column(Some("text".to_string())), + )); + let plan = planner + .plan_search("text", query, Some(10), None) + .await + .expect("planner should produce an active-only plan"); + + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut ids: Vec = Vec::new(); + for b in &batches { + let col = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + ids.push(col.value(i)); + } + } + ids.sort_unstable(); + assert_eq!( + ids, + vec![3], + "AND query must only return rows containing both terms; got ids={ids:?}" + ); + } + + #[tokio::test] + async fn local_mode_active_dedups_updated_pk_keeping_newest() { + // The active memtable is an append log and the FTS index is + // append-only, so a PK updated before flush is searchable as two + // row-positions. WithinSourceDedupExec(KeepMaxRowAddr) must collapse + // them to the newest insert. Without it the same PK would surface + // twice (criterion 2 violation). + let schema = fts_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut indexes = IndexStore::new(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + indexes.add_fts("text_fts".to_string(), 1, "text".to_string()); + + // First append (positions 0,1): id=1 is the stale version of the PK. + let batch_old = make_batch(&schema, &[1, 2], &["lance stale version", "other doc"]); + batch_store.append(batch_old.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch_old, 0, Some(0)) + .unwrap(); + + // Second append (position 2): id=1 updated — same PK, later row. + let batch_new = make_batch(&schema, &[1], &["lance fresh version"]); + batch_store.append(batch_new.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch_new, 2, Some(1)) + .unwrap(); + let indexes = Arc::new(indexes); + + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) + .with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store: indexes, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema); + let plan = planner + .plan_search( + "text", + FullTextSearchQuery::new("lance".to_string()), + Some(10), + None, + ) + .await + .expect("planner should produce an active-only plan"); + + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut rows: Vec<(i32, String)> = Vec::new(); + for b in &batches { + let ids = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let texts = b + .column_by_name("text") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + rows.push((ids.value(i), texts.value(i).to_string())); + } + } + + // id=1 must appear exactly once, and it must be the *newest* version. + let id1: Vec<&(i32, String)> = rows.iter().filter(|(id, _)| *id == 1).collect(); + assert_eq!( + id1.len(), + 1, + "updated PK id=1 must be deduped to one row; got {rows:?}" + ); + assert_eq!( + id1[0].1, "lance fresh version", + "dedup must keep the newest (max row-position) version" + ); + } + + #[tokio::test] + async fn active_stale_update_predicate_crossing_leaks() { + // A PK update that crosses out of the match set: pk=1 inserted as + // "alpha lance", then updated to "beta lance". The append-only inverted + // index keeps the old "alpha" posting live, so an "alpha" search still + // matches the STALE pk=1 row — and the fresh "beta lance" row isn't even + // a candidate, so a result-set dedup has nothing to suppress it against. + // `FtsIndexExec` drops it predicate-independently: pk=1's newest visible + // row is "beta lance", so the "alpha" hit is not the newest. let schema = fts_schema(); let batch_store = Arc::new(BatchStore::with_capacity(16)); let mut indexes = IndexStore::new(); @@ -813,7 +2149,7 @@ mod tests { .plan_search( "text", FullTextSearchQuery::new("alpha".to_string()), - 10, + Some(10), None, ) .await @@ -845,4 +2181,97 @@ mod tests { "live pk=2 ('alpha foo') must still match 'alpha'; got ids={ids:?}" ); } + + #[tokio::test] + async fn cross_gen_stale_update_blocked_by_newer_memtable() { + // The cross-generation analog of `active_stale_update_predicate_crossing_leaks`: + // pk=1's stale "alpha" version lives in a FROZEN memtable and its fresh + // "beta" version in the ACTIVE one. The frozen arm's recency filter is + // per-generation and can't see the newer gen, so only the cross-gen + // block-list can drop the stale "alpha" hit. The cluster constantly + // freezes memtables, so an insert and its later update/delete split + // across in-memory generations — this is the residual fuzz phantom. + let schema = fts_schema(); + + // Frozen gen=1: pk=1 "alpha lance" (matches), pk=2 "alpha foo" (live). + let frozen_store = Arc::new(BatchStore::with_capacity(16)); + let mut frozen_idx = IndexStore::new(); + frozen_idx.enable_pk_index(&[("id".to_string(), 0)]); + frozen_idx.add_fts("text_fts".to_string(), 1, "text".to_string()); + let fb = make_batch(&schema, &[1, 2], &["alpha lance", "alpha foo"]); + let (bp, off, _) = frozen_store.append(fb.clone()).unwrap(); + frozen_idx + .insert_with_batch_position(&fb, off, Some(bp)) + .unwrap(); + + // Active gen=2: pk=1 updated to "beta lance" (no longer matches "alpha"). + let active_store = Arc::new(BatchStore::with_capacity(16)); + let mut active_idx = IndexStore::new(); + active_idx.enable_pk_index(&[("id".to_string(), 0)]); + active_idx.add_fts("text_fts".to_string(), 1, "text".to_string()); + let ab = make_batch(&schema, &[1], &["beta lance"]); + let (bp, off, _) = active_store.append(ab.clone()).unwrap(); + active_idx + .insert_with_batch_position(&ab, off, Some(bp)) + .unwrap(); + + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) + .with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store: active_store, + index_store: Arc::new(active_idx), + schema: schema.clone(), + generation: 2, + }, + frozen: vec![InMemoryMemTableRef { + batch_store: frozen_store, + index_store: Arc::new(frozen_idx), + schema: schema.clone(), + generation: 1, + }], + }, + ); + + let planner = LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], schema); + let plan = planner + .plan_search( + "text", + FullTextSearchQuery::new("alpha".to_string()), + Some(10), + None, + ) + .await + .expect("planner should produce a plan"); + + let ctx = datafusion::prelude::SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut ids: Vec = Vec::new(); + for b in &batches { + let col = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + ids.push(col.value(i)); + } + } + + assert!( + !ids.contains(&1), + "stale frozen pk=1 ('alpha lance', now 'beta lance' in the active gen) \ + leaked on an 'alpha' search; got ids={ids:?}" + ); + assert!( + ids.contains(&2), + "live pk=2 ('alpha foo', only in the frozen gen) must still match; got ids={ids:?}" + ); + } } diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index f701b5b01e6..ec13da1df66 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -21,6 +21,7 @@ use super::exec::{MEMTABLE_GEN_COLUMN, MemtableGenTagExec, PkBlockFilterExec, RO use super::flushed_cache::{DatasetCache, GenerationWarmer, open_flushed_dataset}; use super::projection::{ build_scanner_projection, canonical_output_schema, null_columns, project_to_canonical, + validate_projection_names, }; use crate::session::Session; @@ -103,7 +104,8 @@ impl LsmScanPlanner { } /// Set the over-fetch multiple for the per-source limit pushdown - /// (see the field docs). Clamped to `>= 1.0` at use. + /// (see the field docs). Values below `1.0` are rejected by + /// [`Self::plan_scan`]. pub fn with_overfetch_factor(mut self, factor: f64) -> Self { self.overfetch_factor = factor; self @@ -148,9 +150,11 @@ impl LsmScanPlanner { .map(|p| p.iter().any(|c| c == ROW_ADDRESS_COLUMN)) .unwrap_or(false); let keep_row_address = keep_row_address || user_wants_rowaddr; + validate_projection_names(projection, &self.base_schema, &[])?; // 1. Collect all data sources let sources = self.collector.collect()?; + let overfetch = super::validate_overfetch_factor(self.overfetch_factor)?; if sources.is_empty() { // Return empty plan @@ -182,7 +186,6 @@ impl LsmScanPlanner { // is in-memory and within-gen append duplicates are resolved by its // own dedup, so it is never capped here. let n_needed = limit.map(|l| l.saturating_add(offset.unwrap_or(0))); - let overfetch = self.overfetch_factor.max(1.0); let mut source_plans = Vec::new(); for source in sources { @@ -192,7 +195,7 @@ impl LsmScanPlanner { .get(&(source.shard_id(), source.generation())) .cloned(); let fetch = match (n_needed, is_active) { - (Some(n), false) => Some(if blocked.is_some() { + (Some(n), false) => Some(if blocked.is_some() && !self.pk_columns.is_empty() { ((n as f64) * overfetch).ceil() as usize } else { n @@ -207,13 +210,14 @@ impl LsmScanPlanner { // With a limit, `k = n_needed` arms the under-fetch warning; with // no limit `k = 0` keeps it silent. let scan = match blocked { - Some(set) => Arc::new(PkBlockFilterExec::new( + Some(set) if !self.pk_columns.is_empty() => Arc::new(PkBlockFilterExec::new( scan, self.pk_columns.clone(), set, n_needed.unwrap_or(0), - )) as Arc, - None => scan, + )) + as Arc, + _ => scan, }; // Post-block-list cap: each source contributes at most `n_needed` @@ -261,9 +265,9 @@ impl LsmScanPlanner { &self.canonical_scan_schema(projection, with_memtable_gen, keep_row_address), )?; - // 6. Add limit if specified - if let Some(limit) = limit { - plan = Arc::new(GlobalLimitExec::new(plan, offset.unwrap_or(0), Some(limit))); + // 6. Add limit / offset if specified + if limit.is_some() || offset.unwrap_or(0) > 0 { + plan = Arc::new(GlobalLimitExec::new(plan, offset.unwrap_or(0), limit)); } Ok(plan) @@ -386,7 +390,7 @@ impl LsmScanPlanner { let cols = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); - scanner.project(&cols.iter().map(|s| s.as_str()).collect::>()); + scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; scanner.with_row_address(); // The dedup scan applies the filter post-dedup; pushing it @@ -939,8 +943,9 @@ mod integration_tests { setup_multi_level_lsm().await; // Create scanner with projection (only id column) - let mut scanner = - LsmScanner::new(base_dataset, shard_snapshots, pk_columns).project(&["id"]); + let mut scanner = LsmScanner::new(base_dataset, shard_snapshots, pk_columns) + .project(&["id"]) + .unwrap(); if let Some((shard_id, memtable)) = active_memtable { scanner = scanner.with_in_memory_memtables(shard_id, memtable); } @@ -970,7 +975,9 @@ mod integration_tests { setup_multi_level_lsm().await; // Create scanner with limit - let mut scanner = LsmScanner::new(base_dataset, shard_snapshots, pk_columns).limit(3, None); + let mut scanner = LsmScanner::new(base_dataset, shard_snapshots, pk_columns) + .limit(Some(3), None) + .unwrap(); if let Some((shard_id, memtable)) = active_memtable { scanner = scanner.with_in_memory_memtables(shard_id, memtable); } @@ -989,6 +996,35 @@ mod integration_tests { assert_eq!(total_rows, 3, "Should have 3 rows due to limit"); } + #[tokio::test] + async fn test_lsm_scan_with_offset_without_limit() { + let (base_dataset, shard_snapshots, active_memtable, pk_columns, _temp_path) = + setup_multi_level_lsm().await; + + let mut scanner = LsmScanner::new(base_dataset, shard_snapshots, pk_columns) + .limit(Some(3), None) + .unwrap() + .limit(None, Some(2)) + .unwrap(); + if let Some((shard_id, memtable)) = active_memtable { + scanner = scanner.with_in_memory_memtables(shard_id, memtable); + } + + let batches: Vec = scanner + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total_rows, 5, + "offset-only scan should skip 2 rows from the 7-row deduped result" + ); + } + #[tokio::test] async fn test_lsm_scan_base_only() { let (base_dataset, _, _, pk_columns, _temp_path) = setup_multi_level_lsm().await; @@ -1742,13 +1778,9 @@ mod integration_tests { let (base_dataset, shard_snapshots, active_memtable, pk_columns, _temp_path) = setup_multi_level_lsm().await; - let mut scanner = LsmScanner::new(base_dataset, shard_snapshots, pk_columns).project(&[ - "id", - "_rowoffset", - "name", - "_rowaddr", - "_rowid", - ]); + let mut scanner = LsmScanner::new(base_dataset, shard_snapshots, pk_columns) + .project(&["id", "_rowoffset", "name", "_rowaddr", "_rowid"]) + .unwrap(); if let Some((shard_id, memtable)) = active_memtable { scanner = scanner.with_in_memory_memtables(shard_id, memtable); } @@ -1838,7 +1870,8 @@ mod integration_tests { setup_multi_level_lsm().await; let mut scanner = LsmScanner::new(base_dataset, shard_snapshots, pk_columns) - .project(&["id", "_rowid", "name"]); + .project(&["id", "_rowid", "name"]) + .unwrap(); if let Some((shard_id, memtable)) = active_memtable { scanner = scanner.with_in_memory_memtables(shard_id, memtable); } @@ -1901,7 +1934,8 @@ mod integration_tests { let scanner = LsmScanner::without_base_table(schema, base_uri, vec![], vec!["id".to_string()]) - .project(&["id", "_rowaddr", "name", "_rowid"]); + .project(&["id", "_rowaddr", "name", "_rowid"]) + .unwrap(); let plan = scanner.create_plan().await.unwrap(); let names: Vec = plan @@ -2126,7 +2160,8 @@ mod integration_tests { vec![shard_snapshot], vec!["id".to_string()], ) - .limit(2, None); + .limit(Some(2), None) + .unwrap(); let batches: Vec = scanner .try_into_stream() .await diff --git a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs index f247ef12cb0..4f7c5f093f6 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs @@ -37,7 +37,7 @@ use super::exec::{BloomFilterGuardExec, CoalesceFirstExec, compute_pk_hash_from_ use super::flushed_cache::{DatasetCache, GenerationWarmer, open_flushed_dataset}; use super::projection::{ DISTANCE_COLUMN, build_scanner_projection, canonical_output_schema, null_columns, - project_to_canonical, wants_row_address, wants_row_id, + project_to_canonical, validate_projection_names, wants_row_address, wants_row_id, }; use crate::session::Session; @@ -189,6 +189,34 @@ impl LsmPointLookupPlanner { pk_values: &[ScalarValue], projection: Option<&[String]>, ) -> Result> { + match self.plan_lookup_coalesced(pk_values, projection).await? { + // Tombstones are dropped AFTER the coalesce, never per-source: the + // tombstone wins the newest-first coalesce (its source is the newest + // non-empty arm), so filtering it here yields "not found". Filtering + // per-source would empty the newest arm and let `CoalesceFirstExec` + // fall through to an older arm — resurrecting the deleted row. Then + // the carried `_tombstone` column is projected away. + Some(coalesced) => { + let canonical = + canonical_output_schema(projection, &self.base_schema, &self.pk_columns, false); + filter_tombstones_after_coalesce(coalesced, &canonical) + } + None => self.empty_plan(projection), + } + } + + /// Build the coalesced point-lookup plan: each source scanned newest-first, + /// unioned under `CoalesceFirstExec`, output in the *carry* schema (canonical + /// output + the `_tombstone` marker). `None` when there are no sources at + /// all. [`Self::plan_lookup`] drops the tombstone on top of this; + /// [`Self::lookup_keep_tombstone`] keeps it so partial-update merge can tell + /// a fresh-deleted key from an absent one. + async fn plan_lookup_coalesced( + &self, + pk_values: &[ScalarValue], + projection: Option<&[String]>, + ) -> Result>> { + validate_projection_names(projection, &self.base_schema, &[])?; if pk_values.len() != self.pk_columns.len() { return Err(lance_core::Error::invalid_input(format!( "Expected {} primary key values, got {}", @@ -202,7 +230,7 @@ impl LsmPointLookupPlanner { let sources = self.collector.collect()?; if sources.is_empty() { - return self.empty_plan(projection); + return Ok(None); } // Sort by generation DESC (newest first) @@ -242,17 +270,60 @@ impl LsmPointLookupPlanner { // needs (the in-memory mem_wal execs report empty column statistics, and // datafusion's projection statistics would index out of bounds without // this normalization). - let coalesced: Arc = Arc::new(CoalesceFirstExec::new(source_plans)); - - // Tombstones must be dropped AFTER the coalesce, never per-source: the - // tombstone wins the newest-first coalesce (its source is the newest - // non-empty arm), so filtering it here yields "not found". Filtering - // per-source would empty the newest arm and let `CoalesceFirstExec` - // fall through to an older arm — resurrecting the deleted row. Then the - // carried `_tombstone` column is projected away. + Ok(Some(Arc::new(CoalesceFirstExec::new(source_plans)))) + } + + /// Like [`Self::lookup`] but does NOT drop a tombstone: the returned 1-row + /// batch carries the `_tombstone` marker (`true` ⇒ the key's newest fresh + /// version is a delete); `None` still means the key is absent from every + /// source. Partial-update merge uses this to treat a fresh-deleted PK as + /// absent, so it never resurrects stale, not-yet-compacted base columns. + /// Always plans (no in-memory fast path) — used off the hot read path, on + /// small partial-update batches. + pub async fn lookup_keep_tombstone( + &self, + pk_values: &[ScalarValue], + projection: Option<&[String]>, + ) -> Result> { + let Some(plan) = self.plan_lookup_coalesced(pk_values, projection).await? else { + return Ok(None); + }; + let batches: Vec = plan + .execute(0, self.task_ctx.clone())? + .try_collect() + .await?; + for batch in batches { + if batch.num_rows() > 0 { + return Ok(Some(batch.slice(0, 1))); + } + } + Ok(None) + } + + /// Tombstone-preserving [`Self::lookup_many`] for single- or multi-column + /// keys: resolves each key with [`Self::lookup_keep_tombstone`] and + /// concatenates the hits in the carry schema (canonical output + + /// `_tombstone`); keys absent from every source are omitted. Per-key (no + /// batched fast path) — the partial-update batches that use it are small. + pub async fn lookup_many_keep_tombstone( + &self, + keys: &[Vec], + projection: Option<&[String]>, + ) -> Result { let canonical = canonical_output_schema(projection, &self.base_schema, &self.pk_columns, false); - filter_tombstones_after_coalesce(coalesced, &canonical) + let target = carry_schema(&canonical); + let mut out: Vec = Vec::with_capacity(keys.len()); + for key in keys { + if let Some(b) = self.lookup_keep_tombstone(key, projection).await? { + out.push(b); + } + } + match out.len() { + 0 => Ok(RecordBatch::new_empty(target)), + 1 => Ok(out.pop().unwrap()), + _ => Ok(arrow_select::concat::concat_batches(&target, &out)?), + } } /// Resolve a single-row point lookup, returning the newest matching row (a @@ -606,7 +677,7 @@ impl LsmPointLookupPlanner { // Carry `_tombstone` through so the post-coalesce filter can drop // a deleted key; it survives the sort below. let cols = cols_with_tombstone(&cols, schema.column_with_name(TOMBSTONE).is_some()); - scanner.project(&cols.iter().map(|s| s.as_str()).collect::>()); + scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; scanner.filter_expr(filter.clone()); // Expose `_rowid` (the BatchStore row offset, monotonic with // insert order) so we can pick the most recently inserted @@ -1269,6 +1340,27 @@ mod tests { ); } + #[tokio::test] + async fn test_point_lookup_rejects_missing_projection_column() { + let schema = create_pk_schema(); + let temp_dir = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap()); + + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]); + let planner = LsmPointLookupPlanner::new(collector, vec!["id".to_string()], schema); + + let projection = vec!["missing".to_string()]; + let pk_values = vec![ScalarValue::Int32(Some(2))]; + let err = planner + .plan_lookup(&pk_values, Some(&projection)) + .await + .expect_err("unknown projection column should fail planning"); + assert!( + err.to_string().contains("missing"), + "unexpected missing-column projection error: {err}" + ); + } + #[tokio::test] async fn test_point_lookup_active_memtable_returns_newest_duplicate() { // Regression: same primary key inserted twice into one active @@ -1984,6 +2076,85 @@ mod tests { LsmPointLookupPlanner::new(collector, vec!["id".to_string()], base_schema) } + /// Read the `_tombstone` marker from row 0 of a keep-tombstone result. + fn tombstone_at(b: &RecordBatch) -> bool { + let idx = b.schema().index_of(TOMBSTONE).expect("_tombstone column"); + b.column(idx) + .as_any() + .downcast_ref::() + .expect("_tombstone is Boolean") + .value(0) + } + + #[tokio::test] + async fn test_lookup_keep_tombstone_returns_deleted_row_with_marker() { + // The tombstone-preserving variant keeps a deleted key that the filtered + // `lookup` would drop: id=1 deleted, id=2 live, id=99 absent. + let schema = pk_ts_schema(); + let planner = active_ts_planner(&[ts_real(&schema, &[1, 2], "v"), ts_tomb(&schema, &[1])]); + + // Deleted key: present with `_tombstone = true` (vs `None` from `lookup`). + let deleted = planner + .lookup_keep_tombstone(&[ScalarValue::Int32(Some(1))], None) + .await + .unwrap() + .expect("deleted key kept by lookup_keep_tombstone"); + assert_eq!(id_at(&deleted), 1); + assert!( + tombstone_at(&deleted), + "deleted key carries _tombstone = true" + ); + + // Live key: present with `_tombstone = false`. + let live = planner + .lookup_keep_tombstone(&[ScalarValue::Int32(Some(2))], None) + .await + .unwrap() + .expect("live key found"); + assert_eq!(id_at(&live), 2); + assert!(!tombstone_at(&live), "live key carries _tombstone = false"); + + // Absent key: still `None` — no fresh entry at all to distinguish. + assert!( + planner + .lookup_keep_tombstone(&[ScalarValue::Int32(Some(99))], None) + .await + .unwrap() + .is_none(), + "absent key has no fresh entry" + ); + } + + #[tokio::test] + async fn test_lookup_many_keep_tombstone_includes_tombstoned_keys() { + // Batched variant: the tombstoned key is INCLUDED (carrying its marker), + // unlike `lookup_many` which omits it. id=2 deleted; 1 and 3 live. + let schema = pk_ts_schema(); + let planner = + active_ts_planner(&[ts_real(&schema, &[1, 2, 3], "v"), ts_tomb(&schema, &[2])]); + let keys = vec![ + vec![ScalarValue::Int32(Some(1))], + vec![ScalarValue::Int32(Some(2))], + vec![ScalarValue::Int32(Some(3))], + ]; + let batch = planner + .lookup_many_keep_tombstone(&keys, None) + .await + .unwrap(); + assert_eq!( + batch.num_rows(), + 3, + "all three keys kept, including the tombstone" + ); + // Exactly one row (id=2) is marked deleted. + let deleted: Vec = (0..batch.num_rows()) + .map(|r| batch.slice(r, 1)) + .filter(tombstone_at) + .map(|b| id_at(&b)) + .collect(); + assert_eq!(deleted, vec![2], "only id=2 is marked deleted"); + } + #[tokio::test] async fn test_lookup_tombstone_within_active_not_found() { // Fast path: id=1 written then tombstoned (newer); id=2 is a control. diff --git a/rust/lance/src/dataset/mem_wal/scanner/projection.rs b/rust/lance/src/dataset/mem_wal/scanner/projection.rs index 20d1b1a403d..0ec482aebf8 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/projection.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/projection.rs @@ -73,12 +73,37 @@ pub fn build_scanner_projection( cols } +/// Validate user-facing projection names before constructing a canonical schema. +pub fn validate_projection_names( + user_projection: Option<&[String]>, + base_schema: &SchemaRef, + extra_columns: &[&str], +) -> Result<()> { + let Some(user_projection) = user_projection else { + return Ok(()); + }; + for name in user_projection { + if is_system_column(name) + || extra_columns.contains(&name.as_str()) + || base_schema.field_with_name(name).is_ok() + { + continue; + } + return Err(lance_core::Error::invalid_input(format!( + "Column '{}' not found in schema", + name + ))); + } + Ok(()) +} + /// Canonical output schema honoring user column order. /// /// System cols → nullable `UInt64` at user position (filled by /// `project_to_canonical`). `_distance` (when `include_distance`) → /// nullable `Float32` at user position, appended if absent. PKs appended. -/// Unknown names are silently dropped. +/// Callers should run [`validate_projection_names`] before using this with a +/// user-supplied projection. pub fn canonical_output_schema( user_projection: Option<&[String]>, base_schema: &SchemaRef, diff --git a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs index 7f849f3d8bf..59f721aa08c 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs @@ -18,8 +18,9 @@ use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::union::UnionExec; -use lance_core::Result; +use datafusion::prelude::Expr; use lance_core::datatypes::OnMissing; +use lance_core::{Error, Result}; use tracing::instrument; use crate::dataset::Dataset; @@ -30,21 +31,20 @@ use super::data_source::LsmDataSource; use super::flushed_cache::{DatasetCache, GenerationWarmer, open_flushed_dataset}; use super::projection::{ DISTANCE_COLUMN, build_scanner_projection, canonical_output_schema, null_columns, - project_to_canonical, wants_row_id, + project_to_canonical, validate_projection_names, wants_row_id, }; use crate::session::Session; /// Plans vector search queries over LSM data. /// /// Each source is independently newest-per-PK before the union — the active -/// memtable via an over-fetched KNN + a newest-per-PK recency filter -/// ([`super::exec::NewestPkFilterExec`], which drops a hit that isn't the newest -/// visible version of its PK), flushed generations via their within-generation -/// deletion vector — and the cross-generation block-list -/// ([`super::exec::PkBlockFilterExec`]) drops any PK superseded by a newer -/// generation. So each PK reaches the union from exactly one source and a -/// distance-ordered merge yields the global top-k; no cross-source dedup is -/// needed. +/// memtable via exact brute-force KNN when PK rewrites or a filter require it +/// (append-only active data can still use HNSW), +/// flushed generations via their within-generation deletion vector — and the +/// cross-generation block-list ([`super::exec::PkBlockFilterExec`]) drops any +/// PK superseded by a newer generation. So each PK reaches the union from +/// exactly one source and a distance-ordered merge yields the global top-k; no +/// cross-source dedup is needed. /// /// # Query Plan Structure /// @@ -55,8 +55,7 @@ use crate::session::Session; /// UnionExec /// ProjectionExec (canonical output schema) /// SortExec(_distance, fetch=k) -/// NewestPkFilterExec: newest-per-PK recency (active) -/// KNNExec: active memtable, fetch=ceil(k*overfetch) +/// MemTableBruteForceVectorExec or VectorIndexExec: active memtable KNN /// ProjectionExec (canonical output schema) /// ProjectionExec (null_columns _rowid) /// PkBlockFilterExec: block-list (flushed) @@ -96,6 +95,11 @@ pub struct LsmVectorSearchPlanner { flushed_cache: Option>, /// Optional warmer fired on first open of a flushed generation. warmer: Option>, + /// Optional prefilter predicate applied to every source arm before its KNN + /// search, so rows failing the predicate never enter the top-k. Base and + /// flushed arms use the dataset scanner's native prefilter; memtable arms + /// route to a filtered brute-force scan. + filter: Option, } impl LsmVectorSearchPlanner { @@ -125,9 +129,18 @@ impl LsmVectorSearchPlanner { session: None, flushed_cache: None, warmer: None, + filter: None, } } + /// Attach an optional prefilter predicate. Every source arm restricts its + /// KNN to rows matching the predicate (true prefilter), so results match a + /// normal filtered vector scan over base ∪ flushed ∪ in-memory data. + pub fn with_filter(mut self, filter: Option) -> Self { + self.filter = filter; + self + } + /// Thread a session into flushed-generation opens so the first open /// populates the shared index / file-metadata caches. pub fn with_session(mut self, session: Arc) -> Self { @@ -171,17 +184,13 @@ impl LsmVectorSearchPlanner { /// * `refine_base_table` - When true, the base-table arm re-ranks its /// candidates with exact distances (refine factor 1). Useful when the base /// table uses an approximate index (IVF-PQ) so cross-source distance - /// comparison is exact. Memtable arms use exact HNSW search and never need - /// refine. Auto-enabled whenever stale filtering is on (see below). - /// * `overfetch_factor` - A single knob that controls **both** whether stale - /// rows are filtered and how aggressively sources over-fetch to backfill - /// the rows that filtering drops: + /// comparison is exact. Memtable arms already use exact distances and + /// never need refine. Auto-enabled whenever stale filtering is on (see + /// below). + /// * `overfetch_factor` - Controls how aggressively sources over-fetch to + /// backfill the rows dropped by stale-row filtering. Values below `1.0` + /// are rejected; stale filtering is always enabled. /// - /// - `factor < 1.0` (e.g. `0.0`): **stale filtering off.** The per-source - /// block-list / [`super::exec::PkBlockFilterExec`] is not built or applied, - /// so rows superseded by a newer generation can surface. The global PK - /// dedup still runs, so it still suppresses stale copies in the cases - /// where both the stale and the fresh row reach it. /// - `factor == 1.0`: **stale filtering on, no over-fetch.** Each source /// that has superseded rows fetches exactly `k` candidates, drops the /// stale ones, and may therefore return fewer than `k` live rows. @@ -189,15 +198,10 @@ impl LsmVectorSearchPlanner { /// fetches `ceil(k * factor)` candidates so that dropping the stale ones /// still leaves `k` live rows for the merge. /// - /// There is intentionally no separate on/off flag: over-fetch is only ever - /// meaningful while filtering, so the factor encodes both. A true KNN - /// prefilter would remove the need for over-fetch entirely. - /// /// # Returns /// /// An execution plan that returns the top-K nearest neighbors across all - /// LSM levels, with stale results filtered out (unless `overfetch_factor` - /// disables filtering). + /// LSM levels, with stale results filtered out. #[instrument(name = "lsm_vector_search", level = "info", skip_all, fields(k, nprobes, vector_column = %self.vector_column, distance_type = ?self.distance_type))] pub async fn plan_search( &self, @@ -208,18 +212,24 @@ impl LsmVectorSearchPlanner { refine_base_table: bool, overfetch_factor: f64, ) -> Result> { + if k == 0 { + return Err(Error::invalid_input("k must be positive".to_string())); + } + if nprobes == 0 { + return Err(Error::invalid_input("nprobes must be positive".to_string())); + } + let sources = self.collector.collect()?; + validate_projection_names(projection, &self.base_schema, &[DISTANCE_COLUMN])?; + // The block-list is the sole cross-generation dedup mechanism, so it + // runs unconditionally; `overfetch_factor` only tunes the over-fetch + // multiple for blocked sources. + let overfetch_factor = super::validate_overfetch_factor(overfetch_factor)?; if sources.is_empty() { return self.empty_plan(projection); } - // The block-list is the sole cross-generation dedup mechanism, so it - // runs unconditionally; `overfetch_factor` only tunes the over-fetch - // multiple and is clamped to >= 1.0 so blocked sources still yield k - // live candidates after the post-filter. - let overfetch_factor = overfetch_factor.max(1.0); - // Per-source PK block sets (`NEWER(G)`; base = union of all gens). // `Box::pin` keeps the future off `clippy::large_futures`. let block_lists = Box::pin(super::block_list::compute_source_block_lists( @@ -251,12 +261,12 @@ impl LsmVectorSearchPlanner { let generation = source.generation(); let is_base = matches!(source, LsmDataSource::BaseTable { .. }); let is_active = matches!(source, LsmDataSource::ActiveMemTable { .. }); - // Over-fetch when the post-source filter can drop candidates: a - // blocked source loses superseded rows; the active source's - // within-source dedup collapses duplicate-PK HNSW nodes. Block - // lookup is per shard — generations are per-shard. + // Over-fetch when the post-source block-list can drop + // candidates. Active memtable PK recency is handled inside the + // exact brute-force exec before top-k, so it does not need + // source over-fetch. let blocked = block_lists.get(&(source.shard_id(), generation)); - let fetch_k = if blocked.is_some() || is_active { + let fetch_k = if blocked.is_some() && !self.pk_columns.is_empty() { ((k as f64) * overfetch_factor).ceil() as usize } else { k @@ -279,53 +289,24 @@ impl LsmVectorSearchPlanner { .await?; let mut knn_plans = Vec::new(); - // `build_knn_plan` returns each active arm's max-visible snapshot - // alongside its plan; the active arm's NewestPkFilterExec needs both it - // and `source` (for the batch/index stores), so neither is discarded. - for ((source, is_base, is_active, blocked, _), (knn, active_max_visible)) in - arm_inputs.iter().zip(built) - { + for ((_, is_base, _, blocked, _), knn) in arm_inputs.iter().zip(built) { let is_base = *is_base; - let is_active = *is_active; let blocked = *blocked; // Make each source independently newest-per-PK before the union: - // * active: the append-only HNSW returns one node per inserted - // version *and* leaves stale versions of updated PKs live. The - // recency filter keeps only the hit that is the newest visible - // version of its PK (per the maintained MVCC PK-position index), - // closing the predicate-crossing stale read, then re-sort by - // distance. + // * active: append-only memtables can use HNSW directly; once a + // PK rewrite is observed, `MemTableBruteForceVectorExec` drops + // superseded versions before the top-k cut. // * flushed/base: drop cross-gen superseded rows via the // block-list (within-gen is handled by the flushed DV). - let knn = if is_active { - let (batch_store, index_store) = match source { - LsmDataSource::ActiveMemTable { - batch_store, - index_store, - .. - } => (batch_store.clone(), index_store.clone()), - _ => unreachable!("is_active implies ActiveMemTable"), - }; - let filtered: Arc = - Arc::new(super::exec::NewestPkFilterExec::new( - knn, - self.pk_columns.clone(), - lance_core::ROW_ID, - index_store, - batch_store, - active_max_visible.expect("active arm returns its max_visible snapshot"), - )); - sort_by_distance(filtered, k)? - } else { - match blocked { - Some(set) => Arc::new(super::exec::PkBlockFilterExec::new( - knn, - self.pk_columns.clone(), - set.clone(), - k, - )) as Arc, - None => knn, - } + let knn = match blocked { + Some(_) if self.pk_columns.is_empty() => knn, + Some(set) => Arc::new(super::exec::PkBlockFilterExec::new( + knn, + self.pk_columns.clone(), + set.clone(), + k, + )) as Arc, + None => knn, }; // Lance's `fast_search()` and the active scan both produce a // per-source `_rowid` that would collide with base row ids in the @@ -417,9 +398,6 @@ impl LsmVectorSearchPlanner { /// Build KNN plan for a single data source. /// - /// Returns the plan and, for the active memtable, the `max_visible_batch_position` - /// snapshot its scanner latched — threaded into the recency filter so it keys - /// on the same snapshot the search saw (`None` for base / flushed sources). async fn build_knn_plan( &self, source: &LsmDataSource, @@ -428,13 +406,22 @@ impl LsmVectorSearchPlanner { nprobes: usize, projection: Option<&[String]>, refine: bool, - ) -> Result<(Arc, Option)> { + ) -> Result> { match source { LsmDataSource::BaseTable { dataset } => { let mut scanner = dataset.scan(); let cols = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; + if let Some(ref filter) = self.filter { + // Native scanner prefilter: the ANN runs over rows matching + // the predicate, so the top-k holds only matching rows. + // `prefilter(true)` is required — without it the scanner + // post-filters the unfiltered top-k, dropping matching rows + // that ranked below non-matching ones. + scanner.filter_expr(filter.clone()); + scanner.prefilter(true); + } // Only the base produces a meaningful `_rowid`. `_rowaddr` // can't be combined with `fast_search()` — the IVF index // doesn't preserve it and `TakeExec` refuses to insert it @@ -453,7 +440,7 @@ impl LsmVectorSearchPlanner { if refine { scanner.refine(1); } - Ok((scanner.create_plan().await?, None)) + scanner.create_plan().await } LsmDataSource::FlushedMemTable { path, .. } => { let dataset = open_flushed_dataset( @@ -467,13 +454,19 @@ impl LsmVectorSearchPlanner { let cols = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; + if let Some(ref filter) = self.filter { + // See the base arm: `prefilter(true)` makes this a true + // prefilter rather than a lossy post-filter on the top-k. + scanner.filter_expr(filter.clone()); + scanner.prefilter(true); + } // No `with_row_id/address`: per-source IDs would collide with base. let query_arr = single_query_array(query_vector); scanner.nearest(&self.vector_column, query_arr.as_ref(), k)?; scanner.nprobes(nprobes); scanner.distance_metric(self.distance_type); scanner.fast_search(); - Ok((scanner.create_plan().await?, None)) + scanner.create_plan().await } LsmDataSource::ActiveMemTable { batch_store, @@ -482,29 +475,26 @@ impl LsmVectorSearchPlanner { .. } => { use crate::dataset::mem_wal::memtable::scanner::MemTableScanner; - use arrow_array::Array; let mut scanner = MemTableScanner::new(batch_store.clone(), index_store.clone(), schema.clone()); + // Supply PKs so the memtable scanner can choose HNSW for + // append-only data and exact newest-before-top-k search when + // PK rewrites or filters make stale suppression necessary. + scanner.with_pk_columns(self.pk_columns.clone()); // PK auto-included so the staleness filter retains its bloom hash key. let cols = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); - scanner.project(&cols.iter().map(|s| s.as_str()).collect::>()); - // Expose `_rowid` (BatchStore row offset, monotonic with - // insert order) so `NewestPkFilterExec` can compare each hit's - // position against the PK-position index. The value is - // per-source and NULL'd before reaching the canonical merge. - // (VectorIndexExec only plumbs `with_row_id`, not - // `with_row_address`, but the two yield identical values - // for an active memtable so either would work.) - scanner.with_row_id(); - let query_arr: Arc = Arc::new(query_vector.clone()); - scanner.nearest(&self.vector_column, query_arr, k); + scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; + if let Some(ref filter) = self.filter { + // Routed to filtered brute-force (see `plan_vector_search`): + // the predicate masks rows before the memtable top-k cut. + scanner.filter_expr(filter.clone()); + } + scanner.nearest(&self.vector_column, query_vector, k)?; scanner.nprobes(nprobes); scanner.distance_metric(self.distance_type); - let plan = scanner.create_plan().await?; - // Capture the scanner's own latched snapshot for the recency filter. - Ok((plan, Some(scanner.max_visible_batch_position()))) + scanner.create_plan().await } } } @@ -518,29 +508,6 @@ impl LsmVectorSearchPlanner { } } -/// Sort a single-partition plan by `_distance` ascending and cap at `k`. -/// -/// Used to re-order the active arm after its within-source dedup (which emits -/// rows unordered) so the cross-source distance merge sees a sorted stream. -fn sort_by_distance(plan: Arc, k: usize) -> Result> { - let idx = plan.schema().index_of(DISTANCE_COLUMN).map_err(|_| { - lance_core::Error::invalid_input(format!( - "Column '{}' not found in schema", - DISTANCE_COLUMN - )) - })?; - let sort_expr = vec![PhysicalSortExpr { - expr: Arc::new(Column::new(DISTANCE_COLUMN, idx)), - options: SortOptions { - descending: false, - nulls_first: false, - }, - }]; - let ordering = LexOrdering::new(sort_expr) - .ok_or_else(|| lance_core::Error::internal("Failed to create LexOrdering".to_string()))?; - Ok(Arc::new(SortExec::new(ordering, plan).with_fetch(Some(k)))) -} - /// Convert a (typically single-row) FixedSizeList query into the array shape /// `Scanner::nearest` expects: /// @@ -562,7 +529,7 @@ mod tests { use super::*; use crate::dataset::{Dataset, WriteParams}; use arrow_array::{ - Int32Array, RecordBatch, RecordBatchIterator, builder::FixedSizeListBuilder, + BooleanArray, Int32Array, RecordBatch, RecordBatchIterator, builder::FixedSizeListBuilder, }; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; use std::collections::HashMap; @@ -586,6 +553,25 @@ mod tests { ])) } + fn create_vector_tombstone_schema() -> Arc { + let mut id_metadata = HashMap::new(); + id_metadata.insert( + "lance-schema:unenforced-primary-key".to_string(), + "true".to_string(), + ); + let id_field = Field::new("id", DataType::Int32, false).with_metadata(id_metadata); + + Arc::new(ArrowSchema::new(vec![ + id_field, + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), + true, + ), + Field::new(crate::dataset::mem_wal::TOMBSTONE, DataType::Boolean, false), + ])) + } + fn create_query_vector() -> FixedSizeListArray { use arrow_array::builder::Float32Builder; @@ -622,6 +608,60 @@ mod tests { .unwrap() } + /// One row with an explicit `(id, vector)` — unlike [`create_test_batch`], + /// the vector is not derived from `id`, so the same PK can carry different + /// vectors across generations (an update / move-out-of-neighborhood). + fn create_id_vector_batch(schema: &ArrowSchema, id: i32, vector: [f32; 4]) -> RecordBatch { + use arrow_array::builder::Float32Builder; + let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 4); + for v in vector { + builder.values().append_value(v); + } + builder.append(true); + RecordBatch::try_new( + Arc::new(schema.clone()), + vec![ + Arc::new(Int32Array::from(vec![id])), + Arc::new(builder.finish()), + ], + ) + .unwrap() + } + + fn tombstone_vector_batch( + schema: &ArrowSchema, + rows: &[(i32, Option<[f32; 4]>, bool)], + ) -> RecordBatch { + use arrow_array::builder::Float32Builder; + + let mut vector_builder = FixedSizeListBuilder::new(Float32Builder::new(), 4); + for (_, vector, _) in rows { + if let Some(vector) = vector { + for value in vector { + vector_builder.values().append_value(*value); + } + vector_builder.append(true); + } else { + for _ in 0..4 { + vector_builder.values().append_null(); + } + vector_builder.append(false); + } + } + let ids: Vec = rows.iter().map(|(id, _, _)| *id).collect(); + let tombstones: Vec = rows.iter().map(|(_, _, tombstone)| *tombstone).collect(); + + RecordBatch::try_new( + Arc::new(schema.clone()), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(vector_builder.finish()), + Arc::new(BooleanArray::from(tombstones)), + ], + ) + .unwrap() + } + async fn create_dataset(uri: &str, batches: Vec) -> Dataset { let schema = batches[0].schema(); let has_id = schema.column_with_name("id").is_some(); @@ -665,6 +705,43 @@ mod tests { plan.expect("planner should produce a plan even when memtables are empty"); } + #[tokio::test] + async fn test_vector_search_validates_k_and_nprobes() { + let schema = create_vector_schema(); + let temp_dir = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap()); + let base_batch = create_test_batch(&schema, &[1, 2, 3]); + let base_dataset = Arc::new(create_dataset(&base_uri, vec![base_batch]).await); + let collector = LsmDataSourceCollector::new(base_dataset, vec![]); + + let planner = LsmVectorSearchPlanner::new( + collector, + vec!["id".to_string()], + schema, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + ); + + let query = create_query_vector(); + let err = planner + .plan_search(&query, 0, 1, None, false, 1.0) + .await + .unwrap_err(); + assert!( + err.to_string().contains("k must be positive"), + "expected k validation error, got {err}" + ); + + let err = planner + .plan_search(&query, 1, 0, None, false, 1.0) + .await + .unwrap_err(); + assert!( + err.to_string().contains("nprobes must be positive"), + "expected nprobes validation error, got {err}" + ); + } + #[tokio::test] async fn test_projection_includes_pk() { let schema = create_vector_schema(); @@ -673,26 +750,819 @@ mod tests { let base_batch = create_test_batch(&schema, &[1]); let base_dataset = Arc::new(create_dataset(&base_uri, vec![base_batch]).await); - let collector = LsmDataSourceCollector::new(base_dataset, vec![]); + let collector = LsmDataSourceCollector::new(base_dataset, vec![]); + + let _planner = LsmVectorSearchPlanner::new( + collector, + vec!["id".to_string()], + schema.clone(), + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + ); + + // Project only "vector" - should also include "id" for staleness detection + let cols = + build_scanner_projection(Some(&["vector".to_string()]), &schema, &["id".to_string()]); + + assert!(cols.contains(&"vector".to_string())); + assert!(cols.contains(&"id".to_string())); + } + + #[tokio::test] + async fn test_vector_search_base_plus_active_returns_distance() { + use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; + use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; + use datafusion::prelude::SessionContext; + use futures::TryStreamExt; + + let schema = create_vector_schema(); + let temp_dir = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap()); + let base_batch = create_test_batch(&schema, &[10, 20, 30]); + let base_dataset = Arc::new(create_dataset(&base_uri, vec![base_batch]).await); + + // Active memtable with HNSW index over the "vector" column. + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut index_store = IndexStore::new(); + index_store.enable_pk_index(&[("id".to_string(), 0)]); + index_store.add_hnsw( + "vector_hnsw".to_string(), + 1, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + 64, + 8, + ); + let batch = create_test_batch(&schema, &[1, 2, 3, 4]); + batch_store.append(batch.clone()).unwrap(); + index_store + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + let index_store = Arc::new(index_store); + + let shard_id = uuid::Uuid::new_v4(); + let collector = LsmDataSourceCollector::new(base_dataset, vec![]).with_in_memory_memtables( + shard_id, + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmVectorSearchPlanner::new( + collector, + vec!["id".to_string()], + schema, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + ); + + let query = create_query_vector(); + let plan = planner + .plan_search(&query, 3, 1, None, false, 1.0) + .await + .expect("planner should produce a plan"); + + let ctx = SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert!(total > 0, "expected at least one result row"); + + let out_schema = batches[0].schema(); + let out_cols: Vec = out_schema + .fields() + .iter() + .map(|f| f.name().clone()) + .collect(); + assert!( + out_schema.field_with_name(DISTANCE_COLUMN).is_ok(), + "output schema is missing `_distance` column. Got: {:?}", + out_cols + ); + // Internal columns must not leak: `_rowid` (added by Lance's fast_search + // in the base/flushed arms) and `_memtable_gen` (added by the LSM merge + // when bloom filters are present) are bookkeeping, not API. + assert!( + out_schema.field_with_name("_rowid").is_err(), + "`_rowid` leaked into output: {:?}", + out_cols + ); + assert!( + out_schema + .field_with_name(super::super::exec::MEMTABLE_GEN_COLUMN) + .is_err(), + "`_memtable_gen` leaked into output: {:?}", + out_cols + ); + + // The nearest neighbor to the query vector should be id=1 (exact match), + // and its `_distance` must be ~0 — verifying both the column is present + // and the values are correct. + let id_col = batches[0] + .column_by_name("id") + .expect("id column missing") + .as_any() + .downcast_ref::() + .expect("id column should be Int32"); + let dist_col = batches[0] + .column_by_name(DISTANCE_COLUMN) + .expect("_distance column missing") + .as_any() + .downcast_ref::() + .expect("_distance column should be Float32"); + assert_eq!(id_col.value(0), 1, "expected id=1 as nearest neighbor"); + assert!( + dist_col.value(0).abs() < 1e-3, + "expected near-zero distance for self-match, got {}", + dist_col.value(0) + ); + } + + #[tokio::test] + async fn test_vector_search_active_tombstone_masks_base_hit() { + use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; + use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; + use crate::index::DatasetIndexExt; + use crate::index::vector::VectorIndexParams; + use datafusion::prelude::SessionContext; + use futures::TryStreamExt; + use lance_index::IndexType; + + let base_schema = create_vector_schema(); + let mem_schema = create_vector_tombstone_schema(); + let temp_dir = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap()); + let q = [0.1, 0.2, 0.3, 0.4]; + let near = [0.12, 0.22, 0.32, 0.42]; + + let mut base_dataset = create_dataset( + &base_uri, + vec![batch_rows(&base_schema, &[(1, q), (2, near)])], + ) + .await; + let ivf_flat = VectorIndexParams::ivf_flat(1, lance_linalg::distance::DistanceType::L2); + base_dataset + .create_index(&["vector"], IndexType::Vector, None, &ivf_flat, true) + .await + .unwrap(); + let base_dataset = Arc::new(base_dataset); + + let active_tombstone = tombstone_vector_batch(&mem_schema, &[(1, None, true)]); + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut index_store = IndexStore::new(); + index_store.enable_pk_index(&[("id".to_string(), 0)]); + let (_, row_offset, batch_position) = batch_store.append(active_tombstone.clone()).unwrap(); + index_store + .insert_with_batch_position(&active_tombstone, row_offset, Some(batch_position)) + .unwrap(); + let index_store = Arc::new(index_store); + + let collector = LsmDataSourceCollector::new(base_dataset, vec![]).with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store, + schema: mem_schema, + generation: 1, + }, + frozen: vec![], + }, + ); + let planner = LsmVectorSearchPlanner::new( + collector, + vec!["id".to_string()], + base_schema, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + ); + + let query = create_query_vector(); + let plan = planner + .plan_search(&query, 1, 1, None, false, 2.0) + .await + .unwrap(); + let stream = plan.execute(0, SessionContext::new().task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + let rows = collect_id_dist(&batches); + + assert_eq!(rows.len(), 1, "expected one backfilled hit, got {rows:?}"); + assert_eq!( + rows[0].0, 2, + "active tombstone for id=1 must block the older base vector hit; got {rows:?}" + ); + } + + /// A prefilter on a vector search must restrict the KNN to rows matching the + /// predicate, even though the nearest (and second-nearest) rows fail it. The + /// active memtable has an HNSW index, but a filtered search routes to the + /// brute-force arm, which masks rows before the top-k cut. + #[tokio::test] + async fn test_vector_search_prefilter_restricts_to_matching_rows() { + use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; + use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; + use datafusion::prelude::{SessionContext, col, lit}; + use futures::TryStreamExt; + + let schema = create_vector_schema(); + let temp_dir = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap()); + // Base rows are far from the query and unindexed, so `fast_search` + // contributes nothing; the test isolates the memtable prefilter. + let base_dataset = Arc::new( + create_dataset(&base_uri, vec![create_test_batch(&schema, &[100, 200])]).await, + ); + + // Active memtable with ids 0..=3 (id=1 is the exact match, id=0 ties id=2). + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut index_store = IndexStore::new(); + index_store.enable_pk_index(&[("id".to_string(), 0)]); + index_store.add_hnsw( + "vector_hnsw".to_string(), + 1, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + 64, + 8, + ); + let batch = create_test_batch(&schema, &[0, 1, 2, 3]); + batch_store.append(batch.clone()).unwrap(); + index_store + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + let index_store = Arc::new(index_store); + + let collector = LsmDataSourceCollector::new(base_dataset, vec![]).with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmVectorSearchPlanner::new( + collector, + vec!["id".to_string()], + schema, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + ) + // `id >= 2` excludes the two nearest rows (id=1 exact, id=0 tie). + .with_filter(Some(col("id").gt_eq(lit(2i32)))); + + let query = create_query_vector(); + let plan = planner + .plan_search(&query, 10, 1, None, false, 1.0) + .await + .expect("planner should produce a filtered plan"); + + let ctx = SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut ids: Vec = Vec::new(); + for b in &batches { + let col = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + ids.push(col.value(i)); + } + } + + // Only id=2 and id=3 satisfy `id >= 2`; the nearer id=0/id=1 are excluded. + assert_eq!( + ids.first().copied(), + Some(2), + "nearest matching row is id=2" + ); + let mut sorted = ids.clone(); + sorted.sort(); + assert_eq!( + sorted, + vec![2, 3], + "prefilter must drop id=0 and id=1 (nearer but failing `id >= 2`)" + ); + } + + #[tokio::test] + async fn test_vector_search_filtered_active_without_pk_keeps_all_matching_rows() { + use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; + use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; + use datafusion::prelude::{SessionContext, col, lit}; + use futures::TryStreamExt; + + let schema = create_vector_schema(); + let temp_dir = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap()); + let base_dataset = Arc::new( + create_dataset(&base_uri, vec![create_test_batch(&schema, &[100, 200])]).await, + ); + + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let index_store = IndexStore::new(); + let batch = create_test_batch(&schema, &[1, 2, 3]); + batch_store.append(batch.clone()).unwrap(); + index_store + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + let index_store = Arc::new(index_store); + + let collector = LsmDataSourceCollector::new(base_dataset, vec![]).with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmVectorSearchPlanner::new( + collector, + vec![], + schema, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + ) + .with_filter(Some(col("id").gt_eq(lit(1i32)))); + + let query = create_query_vector(); + let plan = planner + .plan_search(&query, 3, 1, None, false, 1.0) + .await + .expect("planner should produce a filtered active plan"); + + let ctx = SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut ids: Vec = Vec::new(); + for b in &batches { + let col = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + ids.push(col.value(i)); + } + } + ids.sort_unstable(); + assert_eq!( + ids, + vec![1, 2, 3], + "no-PK filtered active vector search must not collapse all rows into one key" + ); + } + + /// The *base* arm must apply the filter as a true prefilter, not a + /// post-filter on the unfiltered top-k. The base is vector-indexed (so + /// `fast_search` uses the index); the two rows nearest the query (id=1, id=2) + /// fail the predicate while matching rows (id>=3) are farther. Without + /// `scanner.prefilter(true)` the base arm runs the ANN unfiltered, takes the + /// top-`k`, then post-filters — dropping every row. A true prefilter restricts + /// the ANN to matching rows. Regression for a missing base-arm `prefilter(true)`. + #[tokio::test] + async fn test_vector_search_base_prefilter_is_not_a_lossy_postfilter() { + use crate::index::DatasetIndexExt; + use crate::index::vector::VectorIndexParams; + use datafusion::prelude::{SessionContext, col, lit}; + use futures::TryStreamExt; + use lance_index::IndexType; + + let schema = create_vector_schema(); + let temp_dir = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap()); + + // L2 distance to the query grows monotonically with id (id=1 is the exact + // match), so the two nearest rows are id=1, id=2 — both excluded by `id>=3`. + let base_batch = create_test_batch(&schema, &[1, 2, 3, 4, 5, 6]); + let mut base_dataset = create_dataset(&base_uri, vec![base_batch]).await; + let ivf_flat = VectorIndexParams::ivf_flat(1, lance_linalg::distance::DistanceType::L2); + base_dataset + .create_index(&["vector"], IndexType::Vector, None, &ivf_flat, true) + .await + .unwrap(); + let base_dataset = Arc::new(base_dataset); + + let collector = LsmDataSourceCollector::new(base_dataset.clone(), vec![]); + let planner = LsmVectorSearchPlanner::new( + collector, + vec!["id".to_string()], + schema, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + ) + .with_dataset(base_dataset) + // Keeps only id>=3, the *farther* matches. A post-filter on the + // unfiltered top-2 (id=1, id=2) would drop everything. + .with_filter(Some(col("id").gt_eq(lit(3i32)))); + + let query = create_query_vector(); + let plan = planner + .plan_search(&query, 2, 1, None, false, 1.0) + .await + .expect("planner should produce a filtered base plan"); + + let ctx = SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut ids: Vec = Vec::new(); + for b in &batches { + let col = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + ids.push(col.value(i)); + } + } + assert_eq!( + ids.first().copied(), + Some(3), + "nearest matching base row is id=3; got {ids:?}" + ); + let mut sorted = ids.clone(); + sorted.sort(); + assert_eq!( + sorted, + vec![3, 4], + "base prefilter must return the two nearest matches (id=3, id=4), not \ + post-filter the unfiltered top-2 (id=1, id=2) down to nothing; got {ids:?}" + ); + } + + /// The flushed arm must also apply the filter as a true prefilter, and that + /// prefiltered candidate set must compose with cross-generation block-list + /// filtering plus over-fetch. Gen 1's closest predicate-matching row (id=3) + /// is superseded by gen 2; with over-fetch, gen 1 should still contribute + /// the next live predicate match (id=4). + #[tokio::test] + async fn test_vector_search_flushed_prefilter_composes_with_block_list() { + use crate::dataset::mem_wal::scanner::data_source::ShardSnapshot; + use crate::index::DatasetIndexExt; + use crate::index::vector::VectorIndexParams; + use datafusion::prelude::{SessionContext, col, lit}; + use futures::TryStreamExt; + use lance_index::IndexType; + + let schema = create_vector_schema(); + let temp_dir = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap()); + let shard_id = uuid::Uuid::new_v4(); + let ivf_flat = VectorIndexParams::ivf_flat(1, lance_linalg::distance::DistanceType::L2); + + let q = [0.1, 0.2, 0.3, 0.4]; + let near = [0.12, 0.22, 0.32, 0.42]; + let far = [9.0, 9.0, 9.0, 9.0]; + + // Gen 1: id=1 is nearest but fails the predicate. id=3 passes but is + // stale (blocked by gen 2). id=4 is the next live predicate match. + let gen1_uri = format!("{}/_mem_wal/{}/gen_1", base_uri, shard_id); + let mut gen1 = create_dataset( + &gen1_uri, + vec![batch_rows(&schema, &[(1, q), (3, q), (4, near)])], + ) + .await; + gen1.create_index(&["vector"], IndexType::Vector, None, &ivf_flat, true) + .await + .unwrap(); + + // Gen 2: newer id=3 shadows gen 1's close copy but is far from query. + let gen2_uri = format!("{}/_mem_wal/{}/gen_2", base_uri, shard_id); + let mut gen2 = create_dataset(&gen2_uri, vec![batch_rows(&schema, &[(3, far)])]).await; + gen2.create_index(&["vector"], IndexType::Vector, None, &ivf_flat, true) + .await + .unwrap(); + + let snapshot = ShardSnapshot::new(shard_id) + .with_current_generation(3) + .with_flushed_generation(1, "gen_1".to_string()) + .with_flushed_generation(2, "gen_2".to_string()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![snapshot]); + + let planner = LsmVectorSearchPlanner::new( + collector, + vec!["id".to_string()], + schema, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + ) + .with_filter(Some(col("id").gt_eq(lit(3i32)))); + + let query = create_query_vector(); + let plan = planner + .plan_search(&query, 1, 1, None, false, 2.0) + .await + .unwrap(); + let ctx = SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + let rows = collect_id_dist(&batches); + + assert_eq!(rows.len(), 1, "expected one result, got {:?}", rows); + assert_eq!( + rows[0].0, 4, + "flushed prefilter should return live id=4 after stale id=3 is blocked; got {:?}", + rows + ); + } + + /// An in-memtable update whose *newest* version fails the prefilter must + /// exclude the PK entirely, not leak the stale older version that still + /// passes. Regression for filter-before-dedup on the active arm: the filter + /// is evaluated against the newest version of each PK. + #[tokio::test] + async fn test_vector_search_prefilter_excludes_pk_whose_newest_version_fails() { + use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; + use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; + use arrow_array::StringArray; + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; + use arrow_schema::{DataType, Field}; + use datafusion::prelude::{SessionContext, col, lit}; + use futures::TryStreamExt; + + let mut id_meta = std::collections::HashMap::new(); + id_meta.insert( + "lance-schema:unenforced-primary-key".to_string(), + "true".to_string(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(id_meta), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), + false, + ), + Field::new("status", DataType::Utf8, false), + ])); + // Same PK id=5 appended twice (an in-memtable update): row0 "active" + // (passes), row1 "archived" (the current version, fails the filter). Both + // vectors are near the query. + let make_batch = |statuses: &[&str]| -> RecordBatch { + let mut vb = FixedSizeListBuilder::new(Float32Builder::new(), 4); + for _ in statuses { + for d in 0..4 { + vb.values().append_value(0.1 + d as f32 * 0.1); + } + vb.append(true); + } + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![5; statuses.len()])), + Arc::new(vb.finish()), + Arc::new(StringArray::from(statuses.to_vec())), + ], + ) + .unwrap() + }; + + // Base is unindexed, so `fast_search` contributes nothing; the test + // isolates the active-memtable dedup-before-filter behavior. Use a + // non-conflicting id to avoid any PK confusion. + let temp_dir = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap()); + let mut fb = FixedSizeListBuilder::new(Float32Builder::new(), 4); + for d in 0..4 { + fb.values().append_value(100.0 + d as f32); + } + fb.append(true); + let far = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![999])), + Arc::new(fb.finish()), + Arc::new(StringArray::from(vec!["active"])), + ], + ) + .unwrap(); + let base_dataset = Arc::new(create_dataset(&base_uri, vec![far]).await); + + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut index_store = IndexStore::new(); + index_store.enable_pk_index(&[("id".to_string(), 0)]); + index_store.add_hnsw( + "vector_hnsw".to_string(), + 1, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + 64, + 8, + ); + let batch = make_batch(&["active", "archived"]); + batch_store.append(batch.clone()).unwrap(); + index_store + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + let index_store = Arc::new(index_store); + + let collector = LsmDataSourceCollector::new(base_dataset, vec![]).with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmVectorSearchPlanner::new( + collector, + vec!["id".to_string()], + schema, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + ) + .with_filter(Some(col("status").eq(lit("active")))); + + let query = create_query_vector(); + let plan = planner + .plan_search(&query, 10, 1, None, false, 1.0) + .await + .expect("planner should produce a filtered plan"); + + let ctx = SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total, 0, + "pk=5's current version is 'archived' and must be excluded; a stale \ + 'active' older version must not leak (filter evaluated on newest)" + ); + } + + /// Cross-arm stale rows must be blocked even if the newer active row fails + /// the prefilter. The base copy of pk=5 matches `status = 'active'`, but the + /// newer active copy is `archived`; pk=5 must be absent from the result. + #[tokio::test] + async fn test_vector_search_prefilter_blocks_base_when_active_newest_fails() { + use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; + use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; + use crate::index::DatasetIndexExt; + use crate::index::vector::VectorIndexParams; + use arrow_array::StringArray; + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; + use arrow_schema::{DataType, Field}; + use datafusion::prelude::{SessionContext, col, lit}; + use futures::TryStreamExt; + use lance_index::IndexType; + + let mut id_meta = std::collections::HashMap::new(); + id_meta.insert( + "lance-schema:unenforced-primary-key".to_string(), + "true".to_string(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(id_meta), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), + false, + ), + Field::new("status", DataType::Utf8, false), + ])); + let make_batch = |rows: &[(i32, [f32; 4], &str)]| -> RecordBatch { + let mut vectors = FixedSizeListBuilder::new(Float32Builder::new(), 4); + for (_, vector, _) in rows { + for value in vector { + vectors.values().append_value(*value); + } + vectors.append(true); + } + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from( + rows.iter().map(|(id, _, _)| *id).collect::>(), + )), + Arc::new(vectors.finish()), + Arc::new(StringArray::from( + rows.iter() + .map(|(_, _, status)| *status) + .collect::>(), + )), + ], + ) + .unwrap() + }; + + let query = [0.1, 0.2, 0.3, 0.4]; + let fallback = [0.2, 0.3, 0.4, 0.5]; + + let temp_dir = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap()); + let base_batch = make_batch(&[(5, query, "active"), (6, fallback, "active")]); + let mut base_dataset = create_dataset(&base_uri, vec![base_batch]).await; + let ivf_flat = VectorIndexParams::ivf_flat(1, lance_linalg::distance::DistanceType::L2); + base_dataset + .create_index(&["vector"], IndexType::Vector, None, &ivf_flat, true) + .await + .unwrap(); + let base_dataset = Arc::new(base_dataset); + + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut index_store = IndexStore::new(); + index_store.enable_pk_index(&[("id".to_string(), 0)]); + index_store.add_hnsw( + "vector_hnsw".to_string(), + 1, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + 16, + 4, + ); + let active_batch = make_batch(&[(5, query, "archived")]); + batch_store.append(active_batch.clone()).unwrap(); + index_store + .insert_with_batch_position(&active_batch, 0, Some(0)) + .unwrap(); + let index_store = Arc::new(index_store); + + let collector = LsmDataSourceCollector::new(base_dataset.clone(), vec![]) + .with_in_memory_memtables( + uuid::Uuid::new_v4(), + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); - let _planner = LsmVectorSearchPlanner::new( + let planner = LsmVectorSearchPlanner::new( collector, vec!["id".to_string()], - schema.clone(), + schema, "vector".to_string(), lance_linalg::distance::DistanceType::L2, - ); + ) + .with_dataset(base_dataset) + .with_filter(Some(col("status").eq(lit("active")))); - // Project only "vector" - should also include "id" for staleness detection - let cols = - build_scanner_projection(Some(&["vector".to_string()]), &schema, &["id".to_string()]); + let query = create_query_vector(); + let plan = planner + .plan_search(&query, 10, 1, None, false, 1.0) + .await + .expect("planner should produce a filtered base+active plan"); - assert!(cols.contains(&"vector".to_string())); - assert!(cols.contains(&"id".to_string())); + let ctx = SessionContext::new(); + let stream = plan.execute(0, ctx.task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + let mut ids: Vec = collect_id_dist(&batches) + .into_iter() + .map(|(id, _)| id) + .collect(); + ids.sort_unstable(); + assert_eq!( + ids, + vec![6], + "base pk=5 passes the filter but is superseded by active archived pk=5; got {ids:?}" + ); } #[tokio::test] - async fn test_vector_search_base_plus_active_returns_distance() { + async fn cross_gen_stale_version_blocked_by_newer_memtable() { + // A PK updated ACROSS a freeze boundary: its old (near) version lives in + // a frozen memtable, its new (far) version in the active one. Both are + // in-memory `ActiveMemTable` sources, so each gets only a per-generation + // recency filter (`NewestPkFilterExec`) — which can't see the newer gen. + // The cross-generation block-list is the only thing that can drop the + // stale near copy; if the active arm discards it, the frozen near version + // leaks as a phantom near-neighbor. This is the residual wallop fuzz + // delete/update phantom: the cluster constantly freezes memtables + // (flush_interval ~250ms), so an insert and its later delete/update split + // across in-memory generations. use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; use datafusion::prelude::SessionContext; @@ -701,39 +1571,47 @@ mod tests { let schema = create_vector_schema(); let temp_dir = tempfile::tempdir().unwrap(); let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap()); - let base_batch = create_test_batch(&schema, &[10, 20, 30]); - let base_dataset = Arc::new(create_dataset(&base_uri, vec![base_batch]).await); + // Base seed far from the query; id=1 is not in the base. + let base_dataset = + Arc::new(create_dataset(&base_uri, vec![create_test_batch(&schema, &[99])]).await); + + // One in-memory memtable holding `id` at `vector`, with a PK + HNSW index. + let make_gen = |id: i32, vector: [f32; 4], generation: u64| { + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut index_store = IndexStore::new(); + index_store.enable_pk_index(&[("id".to_string(), 0)]); + index_store.add_hnsw( + "vector_hnsw".to_string(), + 1, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + 64, + 8, + ); + let batch = create_id_vector_batch(&schema, id, vector); + batch_store.append(batch.clone()).unwrap(); + index_store + .insert_with_batch_position(&batch, 0, Some(0)) + .unwrap(); + InMemoryMemTableRef { + batch_store, + index_store: Arc::new(index_store), + schema: schema.clone(), + generation, + } + }; - // Active memtable with HNSW index over the "vector" column. - let batch_store = Arc::new(BatchStore::with_capacity(16)); - let mut index_store = IndexStore::new(); - index_store.enable_pk_index(&[("id".to_string(), 0)]); - index_store.add_hnsw( - "vector_hnsw".to_string(), - 1, - "vector".to_string(), - lance_linalg::distance::DistanceType::L2, - 64, - 8, - ); - let batch = create_test_batch(&schema, &[1, 2, 3, 4]); - batch_store.append(batch.clone()).unwrap(); - index_store - .insert_with_batch_position(&batch, 0, Some(0)) - .unwrap(); - let index_store = Arc::new(index_store); + // Frozen gen=1: id=1 @ the query vector (distance ~0). + let frozen = make_gen(1, [0.1, 0.2, 0.3, 0.4], 1); + // Active gen=2: id=1 moved FAR — the newest version of the PK. + let active = make_gen(1, [9.0, 9.0, 9.0, 9.0], 2); let shard_id = uuid::Uuid::new_v4(); let collector = LsmDataSourceCollector::new(base_dataset, vec![]).with_in_memory_memtables( shard_id, InMemoryMemTables { - active: InMemoryMemTableRef { - batch_store, - index_store, - schema: schema.clone(), - generation: 1, - }, - frozen: vec![], + active, + frozen: vec![frozen], }, ); @@ -750,62 +1628,48 @@ mod tests { .plan_search(&query, 3, 1, None, false, 1.0) .await .expect("planner should produce a plan"); - let ctx = SessionContext::new(); - let stream = plan.execute(0, ctx.task_ctx()).unwrap(); - let batches: Vec = stream.try_collect().await.unwrap(); - - let total: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert!(total > 0, "expected at least one result row"); + let batches: Vec = plan + .execute(0, ctx.task_ctx()) + .unwrap() + .try_collect() + .await + .unwrap(); - let out_schema = batches[0].schema(); - let out_cols: Vec = out_schema - .fields() - .iter() - .map(|f| f.name().clone()) - .collect(); - assert!( - out_schema.field_with_name(DISTANCE_COLUMN).is_ok(), - "output schema is missing `_distance` column. Got: {:?}", - out_cols - ); - // Internal columns must not leak: `_rowid` (added by Lance's fast_search - // in the base/flushed arms) and `_memtable_gen` (added by the LSM merge - // when bloom filters are present) are bookkeeping, not API. - assert!( - out_schema.field_with_name("_rowid").is_err(), - "`_rowid` leaked into output: {:?}", - out_cols - ); - assert!( - out_schema - .field_with_name(super::super::exec::MEMTABLE_GEN_COLUMN) - .is_err(), - "`_memtable_gen` leaked into output: {:?}", - out_cols - ); + let mut id1_distances = Vec::new(); + for b in &batches { + let ids = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let dist = b + .column_by_name(DISTANCE_COLUMN) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..b.num_rows() { + if ids.value(i) == 1 { + id1_distances.push(dist.value(i)); + } + } + } - // The nearest neighbor to the query vector should be id=1 (exact match), - // and its `_distance` must be ~0 — verifying both the column is present - // and the values are correct. - let id_col = batches[0] - .column_by_name("id") - .expect("id column missing") - .as_any() - .downcast_ref::() - .expect("id column should be Int32"); - let dist_col = batches[0] - .column_by_name(DISTANCE_COLUMN) - .expect("_distance column missing") - .as_any() - .downcast_ref::() - .expect("_distance column should be Float32"); - assert_eq!(id_col.value(0), 1, "expected id=1 as nearest neighbor"); + // id=1 must surface at most once, and only at its NEWEST (far) distance — + // never the stale near ~0 copy the frozen generation still indexes. assert!( - dist_col.value(0).abs() < 1e-3, - "expected near-zero distance for self-match, got {}", - dist_col.value(0) + id1_distances.len() <= 1, + "id=1 leaked its stale frozen copy (appears {} times): {id1_distances:?}", + id1_distances.len() ); + if let Some(d) = id1_distances.first() { + assert!( + *d > 1.0, + "id=1 surfaced at the stale near distance {d}; its newest version is far" + ); + } } #[tokio::test] @@ -1064,9 +1928,10 @@ mod tests { .await .expect("planner should produce a plan"); - // Each arm is independently newest-per-PK (active within-source dedup, - // flushed DV) and the block-list handles cross-gen, merged by a - // distance SPM. No global PK dedup or source tag node is involved. + // Each arm is independently newest-per-PK (active append-only data can + // use HNSW directly; rewritten active data falls back to exact search) + // and the block-list handles cross-gen, merged by a distance SPM. No + // global PK dedup or source tag node is involved. let plan_str = format!( "{}", datafusion::physical_plan::displayable(plan.as_ref()).indent(true) @@ -1077,8 +1942,8 @@ mod tests { plan_str ); assert!( - plan_str.contains("NewestPkFilterExec") && plan_str.contains("SortPreservingMergeExec"), - "expected per-arm dedup + distance merge, got:\n{}", + plan_str.contains("VectorIndexExec") && plan_str.contains("SortPreservingMergeExec"), + "expected append-only active HNSW + distance merge, got:\n{}", plan_str ); @@ -1430,6 +2295,32 @@ mod tests { ); } + #[tokio::test] + async fn test_vector_search_rejects_missing_projection_column() { + let schema = create_vector_schema(); + let temp_dir = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]); + let planner = LsmVectorSearchPlanner::new( + collector, + vec!["id".to_string()], + schema, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + ); + + let projection = vec!["missing".to_string()]; + let query = create_query_vector(); + let err = planner + .plan_search(&query, 5, 1, Some(&projection), false, 1.0) + .await + .unwrap_err(); + assert!( + err.to_string().contains("missing"), + "unexpected missing-column projection error: {err}" + ); + } + #[tokio::test] async fn test_vector_search_without_base_table() { use futures::TryStreamExt; @@ -1507,9 +2398,8 @@ mod tests { #[tokio::test] async fn test_vector_search_dedup_within_active_memtable() { // Regression: same PK inserted twice into one active memtable with - // *different* vectors. HNSW indexes each as a distinct node, so without - // the recency filter a KNN can return both candidates for the same PK - // and pollute top-k. The newer insert must win. + // *different* vectors. The exact active search must consider only the + // newest version of each PK before top-k. The newer insert must win. use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; use datafusion::prelude::SessionContext; @@ -1575,25 +2465,23 @@ mod tests { lance_linalg::distance::DistanceType::L2, ); - // Query is exactly the *newer* vector for pk=1. If the older - // vector for pk=1 leaks through, it'd appear in top-k too because - // the older row's vector is far from the query but still a graph - // node. After dedup we should see pk=1 exactly once. + // Query is exactly the *newer* vector for pk=1. The exact active + // memtable arm should dedup before top-k so pk=1 appears exactly once. let query = create_query_vector(); let plan = planner .plan_search(&query, 5, 1, None, false, 1.0) .await .unwrap(); - // The active arm collapses duplicate-PK HNSW nodes itself via the - // recency filter — there is no cross-source dedup fallback. + // The active arm uses exact brute force when PK recency is configured; + // there is no cross-source dedup fallback. let plan_str = format!( "{}", datafusion::physical_plan::displayable(plan.as_ref()).indent(true) ); assert!( - plan_str.contains("NewestPkFilterExec"), - "active vector arm must self-dedup, got:\n{}", + plan_str.contains("MemTableBruteForceVectorExec"), + "active vector arm must use exact brute-force PK recency, got:\n{}", plan_str ); @@ -1626,19 +2514,9 @@ mod tests { // BUG REPRODUCTION (vector case: a PK update that moves out of the neighborhood). // // Within a *single* active memtable, pk=1 is first inserted ON the query - // (distance ~0), then updated to a FAR vector. The append-only HNSW keeps - // both nodes live. A result-set dedup only collapses duplicate PKs that - // are BOTH present in the over-fetched candidate set. - // - // Here the fresh (far) pk=1 is evicted from the candidate set — there are - // enough nearer filler rows that it ranks below the fetch cutoff — so the - // dedup never sees it and the STALE near pk=1 leaks as the nearest hit. - // This is the predicate-crossing hole: the row that *would* suppress the - // stale version isn't in the result set, so result-set dedup can't help. - // - // Desired (NewestPkFilterExec) behaviour: pk=1's newest row-position is - // the far one, computed predicate-independently over the whole memtable, - // so the stale near node is dropped and pk=1 must NOT surface at ~0. + // (distance ~0), then updated to a FAR vector. Exact brute force must + // compute newest-per-PK over the whole memtable before top-k, so the + // stale near row is dropped even though the fresh row is far away. use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; use datafusion::prelude::SessionContext; @@ -1711,8 +2589,9 @@ mod tests { lance_linalg::distance::DistanceType::L2, ); - // k=3, no over-fetch: the candidate set is {pk1@near, two nearest - // fillers}; fresh pk1@far ranks 7th and never enters the candidates. + // k=3, no over-fetch: exact active search still computes + // newest-per-PK across the whole memtable before top-k, so stale + // pk1@near cannot leak even though fresh pk1@far ranks 7th by distance. let query = create_query_vector(); let plan = planner .plan_search(&query, 3, 1, None, false, 1.0) @@ -1723,6 +2602,13 @@ mod tests { let batches: Vec = stream.try_collect().await.unwrap(); let rows = collect_id_dist(&batches); + assert_eq!( + rows.len(), + 3, + "active PK recency filtering must not underfill k after dropping stale candidates; \ + results={:?}", + rows + ); assert!( !rows.iter().any(|&(id, d)| id == 1 && d.abs() < 1e-3), "stale near pk=1 leaked: its live vector is far from the query, so it \ @@ -1865,27 +2751,15 @@ mod tests { rows ); - // The block-list is now unconditional: a sub-1.0 overfetch_factor is - // clamped to 1.0 and the stale base copy of pk=1 stays suppressed (the - // factor only tunes the over-fetch multiple, it cannot disable filtering). - let still_filtered = planner + // The block-list is unconditional and cannot be disabled via + // overfetch_factor; invalid sub-1.0 values are rejected instead. + let err = planner .plan_search(&query, 1, 1, None, false, 0.0) .await - .unwrap(); - let still_filtered_rows = { - let stream = still_filtered - .execute(0, SessionContext::new().task_ctx()) - .unwrap(); - let batches: Vec = stream.try_collect().await.unwrap(); - collect_id_dist(&batches) - }; + .unwrap_err(); assert!( - still_filtered_rows - .iter() - .all(|&(id, d)| !(id == 1 && d.abs() < 1e-3)), - "block-list is unconditional: stale pk=1 must stay suppressed even \ - with overfetch_factor < 1.0; got {:?}", - still_filtered_rows + err.to_string().contains("overfetch_factor"), + "unexpected error for invalid overfetch factor: {err}" ); } @@ -2253,6 +3127,131 @@ mod tests { ); } + #[tokio::test] + async fn test_vector_search_prefilter_excludes_composite_pk_newest_version() { + use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables}; + use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; + use arrow_array::StringArray; + use arrow_array::builder::Float32Builder; + use datafusion::prelude::{SessionContext, col, lit}; + use futures::TryStreamExt; + + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id1", DataType::Int32, false), + Field::new("id2", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), + false, + ), + Field::new("status", DataType::Utf8, false), + ])); + + fn batch(schema: &ArrowSchema, rows: &[((i32, i32), [f32; 4], &str)]) -> RecordBatch { + let mut vb = FixedSizeListBuilder::new(Float32Builder::new(), 4); + for (_, vector, _) in rows { + for value in vector { + vb.values().append_value(*value); + } + vb.append(true); + } + let id1: Vec = rows.iter().map(|((id1, _), _, _)| *id1).collect(); + let id2: Vec = rows.iter().map(|((_, id2), _, _)| *id2).collect(); + let status: Vec<&str> = rows.iter().map(|(_, _, status)| *status).collect(); + RecordBatch::try_new( + Arc::new(schema.clone()), + vec![ + Arc::new(Int32Array::from(id1)), + Arc::new(Int32Array::from(id2)), + Arc::new(vb.finish()), + Arc::new(StringArray::from(status)), + ], + ) + .unwrap() + } + + let q = [0.1, 0.2, 0.3, 0.4]; + let near = [0.12, 0.22, 0.32, 0.42]; + let far = [9.0, 9.0, 9.0, 9.0]; + let batch = batch( + &schema, + &[ + ((1, 1), q, "active"), + ((1, 1), far, "archived"), + ((2, 2), near, "active"), + ], + ); + + let batch_store = Arc::new(BatchStore::with_capacity(16)); + let mut index_store = IndexStore::new(); + index_store.enable_pk_index(&[("id1".to_string(), 0), ("id2".to_string(), 1)]); + let (_, row_offset, batch_position) = batch_store.append(batch.clone()).unwrap(); + index_store + .insert_with_batch_position(&batch, row_offset, Some(batch_position)) + .unwrap(); + let index_store = Arc::new(index_store); + + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let shard_id = uuid::Uuid::new_v4(); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![]) + .with_in_memory_memtables( + shard_id, + InMemoryMemTables { + active: InMemoryMemTableRef { + batch_store, + index_store, + schema: schema.clone(), + generation: 1, + }, + frozen: vec![], + }, + ); + + let planner = LsmVectorSearchPlanner::new( + collector, + vec!["id1".to_string(), "id2".to_string()], + schema, + "vector".to_string(), + lance_linalg::distance::DistanceType::L2, + ) + .with_filter(Some(col("status").eq(lit("active")))); + + let query = create_query_vector(); + let plan = planner + .plan_search(&query, 10, 1, None, false, 1.0) + .await + .unwrap(); + let stream = plan.execute(0, SessionContext::new().task_ctx()).unwrap(); + let batches: Vec = stream.try_collect().await.unwrap(); + + let mut ids = Vec::new(); + for batch in &batches { + let id1 = batch + .column_by_name("id1") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let id2 = batch + .column_by_name("id2") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + ids.push((id1.value(row), id2.value(row))); + } + } + + assert_eq!( + ids, + vec![(2, 2)], + "composite PK (1,1)'s newest version is archived and must exclude \ + the older active version; got {ids:?}" + ); + } + #[tokio::test] async fn test_vector_search_same_l0_override_newest_wins() { // Ported from the #6844 spec. The DANGEROUS within-memtable direction: diff --git a/rust/lance/src/dataset/mem_wal/test_util.rs b/rust/lance/src/dataset/mem_wal/test_util.rs new file mode 100644 index 00000000000..43a3861e686 --- /dev/null +++ b/rust/lance/src/dataset/mem_wal/test_util.rs @@ -0,0 +1,188 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Test-only object store that injects WAL-write failures, for exercising the +//! WAL persistence-failure fencing path. + +use std::fmt::{Debug, Display, Formatter}; +use std::ops::Range; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +use bytes::Bytes; +use futures::stream::BoxStream; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore as OSObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, + RenameOptions, Result as OSResult, +}; + +use lance_io::object_store::{ + ObjectStore, ObjectStoreParams, ObjectStoreRegistry, WrappingObjectStore, +}; + +/// Knobs controlling injected WAL-entry write failures, shared with the test. +#[derive(Debug, Default)] +pub struct FailControls { + /// Remaining WAL-entry `put_opts` calls to fail; saturates at 0. Set high to + /// "always fail", set to 0 to "recover". + wal_put_failures: AtomicUsize, + /// When failing, write to the inner store anyway before reporting the error, + /// simulating a lost acknowledgement (the PUT actually landed). + simulate_lost_ack: AtomicBool, + /// WAL-entry `put_opts` attempts observed, for assertions. + wal_put_attempts: AtomicUsize, +} + +impl FailControls { + pub fn fail_wal_puts(&self, n: usize) { + self.wal_put_failures.store(n, Ordering::SeqCst); + } + pub fn recover(&self) { + self.wal_put_failures.store(0, Ordering::SeqCst); + } + pub fn set_lost_ack(&self, v: bool) { + self.simulate_lost_ack.store(v, Ordering::SeqCst); + } + pub fn attempts(&self) -> usize { + self.wal_put_attempts.load(Ordering::SeqCst) + } +} + +/// Wraps the inner store with [`FailingObjectStore`] at construction. +#[derive(Debug)] +struct FailingWrapper { + controls: Arc, +} + +impl WrappingObjectStore for FailingWrapper { + fn wrap(&self, _prefix: &str, original: Arc) -> Arc { + Arc::new(FailingObjectStore { + inner: original, + controls: self.controls.clone(), + }) + } +} + +/// Delegates everything to `inner`, failing WAL-entry PUTs per [`FailControls`]. +#[derive(Debug)] +struct FailingObjectStore { + inner: Arc, + controls: Arc, +} + +impl FailingObjectStore { + /// WAL entries live under `.../wal/`; manifests (`.../manifest/`) are never + /// failed so `ShardWriter::open`/`claim_epoch` still work. + fn is_wal_entry(location: &Path) -> bool { + location.as_ref().contains("/wal/") + } + + fn injected_error() -> object_store::Error { + object_store::Error::Generic { + store: "failing-test-store", + source: "injected transient WAL put failure".to_string().into(), + } + } +} + +impl Display for FailingObjectStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "FailingObjectStore({})", self.inner) + } +} + +#[async_trait::async_trait] +impl OSObjectStore for FailingObjectStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> OSResult { + if Self::is_wal_entry(location) { + self.controls + .wal_put_attempts + .fetch_add(1, Ordering::SeqCst); + if self.controls.wal_put_failures.load(Ordering::SeqCst) > 0 { + self.controls + .wal_put_failures + .fetch_sub(1, Ordering::SeqCst); + if self.controls.simulate_lost_ack.load(Ordering::SeqCst) { + // The write lands but we still report failure (lost ack). + let _ = self.inner.put_opts(location, payload, opts).await; + } + return Err(Self::injected_error()); + } + } + self.inner.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { + self.inner.get_opts(location, options).await + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OSResult> { + self.inner.get_ranges(location, ranges).await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + self.inner.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + self.inner.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OSResult> { + self.inner.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> { + self.inner.copy_opts(from, to, opts).await + } + + async fn rename_opts(&self, from: &Path, to: &Path, opts: RenameOptions) -> OSResult<()> { + self.inner.rename_opts(from, to, opts).await + } +} + +/// Build an in-memory `ObjectStore` whose WAL-entry writes can be failed on +/// demand. Returns the store, its base path, and the shared controls. +pub async fn failing_memory_store() -> (Arc, Path, Arc) { + let controls = Arc::new(FailControls::default()); + let params = ObjectStoreParams { + object_store_wrapper: Some(Arc::new(FailingWrapper { + controls: controls.clone(), + })), + ..Default::default() + }; + let (store, base) = ObjectStore::from_uri_and_params( + Arc::new(ObjectStoreRegistry::default()), + "memory:///", + ¶ms, + ) + .await + .unwrap(); + (store, base, controls) +} diff --git a/rust/lance/src/dataset/mem_wal/wal.rs b/rust/lance/src/dataset/mem_wal/wal.rs index afccfbc0979..b4caff180e1 100644 --- a/rust/lance/src/dataset/mem_wal/wal.rs +++ b/rust/lance/src/dataset/mem_wal/wal.rs @@ -12,13 +12,15 @@ use std::sync::Mutex as StdMutex; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Instant; +use std::time::Duration; + use arrow_array::RecordBatch; use arrow_ipc::reader::StreamReader; use arrow_ipc::writer::StreamWriter; use arrow_schema::Schema as ArrowSchema; use bytes::Bytes; use futures::StreamExt; -use lance_core::{Error, Result}; +use lance_core::{Error, FenceReason, Result}; use lance_io::object_store::ObjectStore; use object_store::ObjectStoreExt; use object_store::path::Path; @@ -43,11 +45,46 @@ pub const WRITER_EPOCH_KEY: &str = "writer_epoch"; /// replay skips sentinels via their empty batch list). pub const FENCE_SENTINEL_KEY: &str = "fence_sentinel"; -/// True if `error` is the terminal fence emitted by `ManifestStore::check_fenced` -/// (a successor claimed a higher epoch). Matches the message it formats, since -/// fences surface as a plain `Error::io` rather than a typed variant. +/// True if `error` is a peer fence (a successor claimed a higher epoch). +#[cfg(test)] fn is_fence_error(error: &Error) -> bool { - error.to_string().contains("Writer fenced") + error.fence_reason() == Some(FenceReason::PeerClaimedEpoch) +} + +/// True if `error` is terminal for the writer (either fence kind): the WAL will +/// never advance, so durability waiters must be woken and later writes rejected. +fn is_terminal_failure(error: &Error) -> bool { + error.fence_reason().is_some() +} + +/// Cloneable carrier that ferries a flush error across the async completion +/// channels (`WalFlusher::terminal_error` and the per-flush `done` cell), since +/// `lance_core::Error` is not `Clone`. Preserves the [`FenceReason`] so the typed +/// [`Error::Fenced`] can be rebuilt rather than flattened to a string. +#[derive(Clone, Debug)] +pub struct WalFlushFailure { + /// The fence reason if terminal; `None` for an ordinary flush error. + pub fence_reason: Option, + /// The error message carrying details about the flush failure. + pub message: String, +} + +impl WalFlushFailure { + pub(crate) fn from_error(error: &Error) -> Self { + Self { + fence_reason: error.fence_reason(), + message: error.to_string(), + } + } + + /// Rebuild a typed `Error`, restoring the fence reason when present. + pub(crate) fn into_error(self) -> Error { + match self.fence_reason { + Some(FenceReason::PeerClaimedEpoch) => Error::fenced_by_peer(self.message), + Some(FenceReason::PersistenceFailure) => Error::writer_poisoned(self.message), + None => Error::io(self.message), + } + } } /// Watcher for batch durability using watermark-based tracking. @@ -60,10 +97,10 @@ pub struct BatchDurableWatcher { rx: watch::Receiver, /// Target batch ID to wait for. target_batch_position: usize, - /// Terminal flush failure (e.g. a fence) shared with the flusher. When - /// set, the watermark will never advance to the target, so `wait` - /// returns this error instead of blocking forever. - terminal_error: Arc>>, + /// Terminal flush failure shared with the flusher. When set, the watermark + /// can never reach the target, so `wait` returns this typed error instead of + /// blocking forever. + terminal_error: Arc>>, } impl BatchDurableWatcher { @@ -71,7 +108,7 @@ impl BatchDurableWatcher { pub fn new( rx: watch::Receiver, target_batch_position: usize, - terminal_error: Arc>>, + terminal_error: Arc>>, ) -> Self { Self { rx, @@ -87,8 +124,8 @@ impl BatchDurableWatcher { /// never reach the target. pub async fn wait(&mut self) -> Result<()> { loop { - if let Some(msg) = self.terminal_error.lock().unwrap().clone() { - return Err(Error::io(msg)); + if let Some(failure) = self.terminal_error.lock().unwrap().clone() { + return Err(failure.into_error()); } let current = *self.rx.borrow(); if current >= self.target_batch_position { @@ -186,8 +223,10 @@ pub struct TriggerWalFlush { /// this flush. Use `usize::MAX` to flush all pending batches. pub end_batch_position: usize, /// Optional cell to write completion result. - /// Uses Result since Error doesn't implement Clone. - pub done: Option>>, + /// Uses `WalFlushFailure` (not `Error`) since `Error` doesn't implement + /// `Clone`; the carrier preserves the fence reason so callers waiting on + /// this cell still get a typed error. + pub done: Option>>, } impl std::fmt::Debug for TriggerWalFlush { @@ -338,11 +377,10 @@ pub struct WalFlusher { /// Created at construction and recreated after each flush. /// Used by backpressure to wait for WAL flushes. wal_flush_cell: std::sync::Mutex>>, - /// First terminal flush failure (a fence). Shared with every - /// `BatchDurableWatcher` so a fenced flush — which never advances the - /// watermark — wakes durability waiters with the error instead of - /// hanging them forever. - terminal_error: Arc>>, + /// First terminal flush failure, shared with every `BatchDurableWatcher`. It + /// wakes durability waiters (the watermark never advances) and is read by + /// `check_poisoned` so the write path fails fast. + terminal_error: Arc>>, } impl WalFlusher { @@ -392,15 +430,14 @@ impl WalFlusher { ) } - /// Record a terminal flush failure (a fence) and wake every pending - /// durability waiter. A fence is permanent — the watermark will never - /// advance — so waiters must observe the error rather than block forever. + /// Latch a terminal flush failure and wake every durability waiter (the + /// watermark never advances, so they must observe the error, not block). /// Idempotent: only the first failure is retained. fn mark_terminal_failure(&self, error: &Error) { { let mut slot = self.terminal_error.lock().unwrap(); if slot.is_none() { - *slot = Some(error.to_string()); + *slot = Some(WalFlushFailure::from_error(error)); } } // Wake `wait`ers without advancing the watermark; each re-checks @@ -408,6 +445,17 @@ impl WalFlusher { self.durable_watermark_tx.send_modify(|_| {}); } + /// Fail fast with the typed error if this writer has been fenced (by a peer + /// or its own persistence failure). The write path calls this before touching + /// the memtable so a poisoned writer can't diverge further. Recovery is to + /// reopen the shard (replay the WAL). + pub fn check_poisoned(&self) -> Result<()> { + if let Some(failure) = self.terminal_error.lock().unwrap().clone() { + return Err(failure.into_error()); + } + Ok(()) + } + /// Get the current durable watermark. pub fn durable_watermark(&self) -> usize { *self.durable_watermark_rx.borrow() @@ -452,7 +500,7 @@ impl WalFlusher { &self, source: WalFlushSource, end_batch_position: usize, - done: Option>>, + done: Option>>, ) -> Result<()> { if let Some(tx) = &self.flush_tx { tx.send(TriggerWalFlush { @@ -488,11 +536,10 @@ impl WalFlusher { } WalFlushSource::WalOnly { state } => self.flush_from_wal_only(state).await, }; - // A fence is terminal: the append will never succeed, so the - // durability watermark can never advance. Wake any waiter (e.g. a - // `durable_write` put) with the fence error instead of hanging it. + // A terminal failure means the watermark can never advance; latch the + // poison so waiters wake with the typed error and later writes fail fast. if let Err(e) = &result - && is_fence_error(e) + && is_terminal_failure(e) { self.mark_terminal_failure(e); } @@ -734,6 +781,42 @@ const MAX_APPEND_CREATE_CONFLICTS: usize = 1024; const APPEND_CONFLICT_REFRESH_INTERVAL: usize = 16; const MAX_CURSOR_PROBE: u64 = 4096; +/// Retry policy for transient WAL persistence failures before the writer +/// self-fences. On a non-conflict object-store error the appender retries the +/// *same* WAL position up to `max_retries` times with exponential backoff from +/// `base_delay`; exhausting the budget poisons the writer +/// ([`Error::writer_poisoned`]). +#[derive(Debug, Clone, Copy)] +pub struct WalRetryConfig { + /// Maximum number of retries for a transient WAL write failure before the + /// writer self-fences. + pub max_retries: usize, + /// Base duration for exponential backoff between retry attempts. + pub base_delay: Duration, +} + +impl Default for WalRetryConfig { + fn default() -> Self { + Self { + max_retries: 3, + base_delay: Duration::from_millis(50), + } + } +} + +impl WalRetryConfig { + /// Backoff before retry `attempt` (1-based), capped so a wedged store can't + /// block the flush task indefinitely. + fn backoff(&self, attempt: usize) -> Duration { + const MAX_BACKOFF: Duration = Duration::from_secs(5); + let shift = attempt.saturating_sub(1).min(16) as u32; + self.base_delay + .checked_mul(1u32 << shift) + .unwrap_or(MAX_BACKOFF) + .min(MAX_BACKOFF) + } +} + /// Result of appending a WAL entry. #[derive(Debug, Clone)] pub struct WalAppendResult { @@ -772,6 +855,8 @@ pub struct WalAppender { /// so reopened shards report the post-recovery cursor immediately; /// updated after each successful append. next_entry_position_hint: AtomicU64, + /// Retry budget for transient persistence failures before self-fencing. + retry: WalRetryConfig, } impl WalAppender { @@ -800,6 +885,7 @@ impl WalAppender { manifest_store, writer_epoch, position_hint, + WalRetryConfig::default(), )) } @@ -821,6 +907,7 @@ impl WalAppender { manifest_store: Arc, writer_epoch: u64, next_entry_position_hint_seed: u64, + retry: WalRetryConfig, ) -> Self { Self { object_store, @@ -830,6 +917,7 @@ impl WalAppender { writer_epoch, next_entry_position: Mutex::new(None), next_entry_position_hint: AtomicU64::new(next_entry_position_hint_seed), + retry, } } @@ -882,7 +970,12 @@ impl WalAppender { *next_pos = Some(self.discover_next_position().await?); } + // `conflicts` counts position races across the append; `io_attempts` is + // the per-position retry budget for transient PUT failures. The latter + // only grows while we sit on one position — the sole advancing branch is + // `io_attempts == 0`, so it never carries across positions. let mut conflicts = 0; + let mut io_attempts = 0; loop { let pos = next_pos.ok_or_else(|| { Error::internal(format!( @@ -913,7 +1006,20 @@ impl WalAppender { }); } Err(AtomicPutError::AlreadyExists) => { - self.check_fenced().await?; + self.check_fenced().await?; // surfaces a peer takeover as a typed fence + // A slot we already failed to PUT is now occupied — ambiguous + // (our lost-ack, or a peer). We can't advance-and-rewrite (would + // duplicate) nor blindly accept it, so poison and let replay + // reconcile on reopen. + if io_attempts > 0 { + return Err(Error::writer_poisoned(format!( + "WAL position {} for shard {} was taken after a failed PUT; \ + in-memory state may diverge from the durable WAL, reopen to replay", + pos, self.shard_id + ))); + } + // First touch of this slot: an ordinary position conflict + // (stale cursor, our own earlier entries, or contention). conflicts += 1; if conflicts >= MAX_APPEND_CREATE_CONFLICTS { return Err(Error::io(format!( @@ -930,8 +1036,18 @@ impl WalAppender { } } Err(AtomicPutError::Other(error)) => { + // A successor fence is terminal — don't waste the retry budget. self.check_fenced().await?; - return Err(error); + if io_attempts >= self.retry.max_retries { + return Err(Error::writer_poisoned(format!( + "WAL persistence failed for shard {} at position {} after {} retries; \ + in-memory state may diverge from the durable WAL, reopen to replay: {}", + self.shard_id, pos, self.retry.max_retries, error + ))); + } + io_attempts += 1; + tokio::time::sleep(self.retry.backoff(io_attempts)).await; + // Retry the same position (next_pos unchanged). } } } @@ -1378,6 +1494,7 @@ async fn best_effort_cursor_update(manifest_store: &ShardManifestStore, entry_po #[cfg(test)] mod tests { use super::*; + use crate::dataset::mem_wal::test_util::failing_memory_store; use arrow_array::{Int32Array, StringArray}; use arrow_schema::{DataType, Field, Schema}; use std::sync::Arc; @@ -1430,6 +1547,7 @@ mod tests { writer_epoch, // Tests start with no entries, so seed the hint at 0. 0, + WalRetryConfig::default(), )); WalFlusher::new(appender) } @@ -1698,18 +1816,13 @@ mod tests { second.append(vec![batch.clone()]).await.unwrap(); let err = first.check_fenced().await.unwrap_err(); - assert!( - err.to_string().contains("Writer fenced"), - "expected fence error, got: {err}" - ); + // A peer takeover is a typed peer fence, distinct from a self-poison. + assert_eq!(err.fence_reason(), Some(FenceReason::PeerClaimedEpoch)); // Fenced writer's cached next_pos still points at 2; the conflict path // must surface the fence error rather than silently advance. let err = first.append(vec![batch]).await.unwrap_err(); - assert!( - err.to_string().contains("Writer fenced"), - "expected fence error from append, got: {err}" - ); + assert_eq!(err.fence_reason(), Some(FenceReason::PeerClaimedEpoch)); } #[tokio::test] @@ -1868,4 +1981,137 @@ mod tests { // Three entries from a fresh shard land at 1, 2, 3, so next is 4. assert_eq!(tailer.next_position().await.unwrap(), 4); } + + // A transient PUT failure is retried at the same position and then succeeds. + #[tokio::test] + async fn test_append_retries_transient_failure_then_succeeds() { + let (store, base, controls) = failing_memory_store().await; + let shard_id = Uuid::new_v4(); + controls.fail_wal_puts(2); // 2 failures, under the default budget of 3 + let appender = WalAppender::open(store, base, shard_id, 0).await.unwrap(); + + let schema = create_test_schema(); + let res = appender + .append(vec![create_test_batch(&schema, 1)]) + .await + .unwrap(); + assert_eq!(res.entry_position, FIRST_WAL_ENTRY_POSITION); + assert_eq!(controls.attempts(), 3); // 2 failed + 1 succeeded + } + + // Exhausting the retry budget poisons the writer with a typed persistence + // failure (distinct from a peer fence). + #[tokio::test] + async fn test_append_poisons_after_exhausting_retries() { + let (store, base, controls) = failing_memory_store().await; + let shard_id = Uuid::new_v4(); + controls.fail_wal_puts(usize::MAX); + let manifest_store = Arc::new(ShardManifestStore::new(store.clone(), &base, shard_id, 2)); + let (epoch, _) = manifest_store.claim_epoch(0).await.unwrap(); + let appender = WalAppender::with_claimed_epoch( + store, + base, + shard_id, + manifest_store, + epoch, + 0, + WalRetryConfig { + max_retries: 2, + base_delay: Duration::from_millis(1), + }, + ); + + let schema = create_test_schema(); + let err = appender + .append(vec![create_test_batch(&schema, 1)]) + .await + .unwrap_err(); + assert_eq!(err.fence_reason(), Some(FenceReason::PersistenceFailure)); + assert_eq!(controls.attempts(), 3); // io_attempts 0,1,2 then poison + } + + // A lost acknowledgement (PUT errored but landed) poisons the writer without + // ever writing a duplicate entry at the next position. + #[tokio::test] + async fn test_append_lost_ack_poisons_without_duplicate() { + let (store, base, controls) = failing_memory_store().await; + let shard_id = Uuid::new_v4(); + controls.set_lost_ack(true); + controls.fail_wal_puts(1); + let appender = WalAppender::open(store.clone(), base.clone(), shard_id, 0) + .await + .unwrap(); + + let schema = create_test_schema(); + let err = appender + .append(vec![create_test_batch(&schema, 1)]) + .await + .unwrap_err(); + assert_eq!(err.fence_reason(), Some(FenceReason::PersistenceFailure)); + + // The entry landed exactly once; the next slot stays empty. + let tailer = WalTailer::new(store, base, shard_id); + assert!( + tailer + .read_entry(FIRST_WAL_ENTRY_POSITION) + .await + .unwrap() + .is_some() + ); + assert!( + tailer + .read_entry(FIRST_WAL_ENTRY_POSITION + 1) + .await + .unwrap() + .is_none() + ); + } + + // A persistence failure during flush latches the poison: the flush result, + // `check_poisoned`, and the durability watcher all report the typed error + // (rather than the watcher hanging on a watermark that never advances). + #[tokio::test] + async fn test_flush_persistence_failure_poisons_and_wakes_waiter() { + let (store, base, controls) = failing_memory_store().await; + let shard_id = Uuid::new_v4(); + controls.fail_wal_puts(usize::MAX); + let manifest_store = Arc::new(ShardManifestStore::new(store.clone(), &base, shard_id, 2)); + let (epoch, _) = manifest_store.claim_epoch(0).await.unwrap(); + let appender = Arc::new(WalAppender::with_claimed_epoch( + store, + base, + shard_id, + manifest_store, + epoch, + 0, + WalRetryConfig { + max_retries: 1, + base_delay: Duration::from_millis(1), + }, + )); + let flusher = WalFlusher::new(appender); + + let schema = create_test_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(10)); + batch_store.append(create_test_batch(&schema, 1)).unwrap(); + let mut watcher = flusher.track_batch(0); + + let source = batch_store_source(&batch_store); + let flush_err = flusher.flush(&source, batch_store.len()).await.unwrap_err(); + assert_eq!( + flush_err.fence_reason(), + Some(FenceReason::PersistenceFailure) + ); + + assert_eq!( + flusher.check_poisoned().unwrap_err().fence_reason(), + Some(FenceReason::PersistenceFailure) + ); + + let waited = tokio::time::timeout(Duration::from_secs(5), watcher.wait()) + .await + .expect("watcher hung after a poisoning flush") + .expect_err("watcher must surface the poison"); + assert_eq!(waited.fence_reason(), Some(FenceReason::PersistenceFailure)); + } } diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 8403452f16d..e006b7505a9 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -44,13 +44,13 @@ pub use super::memtable::batch_store::{BatchStore, StoreFull, StoredBatch}; pub use super::memtable::flush::MemTableFlusher; pub use super::memtable::scanner::MemTableScanner; pub use super::util::{WatchableOnceCell, WatchableOnceCellReader}; -pub use super::wal::{WalEntry, WalEntryData, WalFlushResult, WalFlusher}; +pub use super::wal::{WalEntry, WalEntryData, WalFlushFailure, WalFlushResult, WalFlusher}; use super::memtable::flush::TriggerMemTableFlush; use super::scanner::GenerationWarmer; use super::wal::{ - BatchDurableWatcher, TriggerWalFlush, WalAppender, WalFlushSource, WalOnlyState, WalTailer, - empty_flush_result, + BatchDurableWatcher, TriggerWalFlush, WalAppender, WalFlushSource, WalOnlyState, + WalRetryConfig, WalTailer, empty_flush_result, }; use super::{TOMBSTONE, schema_with_tombstone}; @@ -109,6 +109,16 @@ pub struct ShardWriterConfig { /// Default: 100ms pub max_wal_flush_interval: Option, + /// Times a failed WAL PUT is retried (same position, exponential backoff) + /// before the writer self-fences. On exhaustion the writer is poisoned + /// (`Error::writer_poisoned`) and must be reopened to replay. Absorbs + /// transient errors beyond `object_store`'s own retries. Default: 3. + pub max_wal_persist_retries: usize, + + /// Base backoff before the first WAL persistence retry; subsequent retries + /// back off exponentially (capped). Default: 50ms. + pub wal_persist_retry_base_delay: Duration, + /// Maximum MemTable size in bytes before triggering a flush to storage. /// /// MemTable size is checked every `max_wal_flush_interval` (during WAL flush ticks). @@ -250,9 +260,11 @@ impl Default for ShardWriterConfig { sync_indexed_write: true, max_wal_buffer_size: 10 * 1024 * 1024, // 10MB max_wal_flush_interval: Some(Duration::from_millis(100)), // 100ms - max_memtable_size: 256 * 1024 * 1024, // 256MB - max_memtable_rows: 100_000, // 100k rows - max_memtable_batches: 8_000, // 8k batches + max_wal_persist_retries: 3, + wal_persist_retry_base_delay: Duration::from_millis(50), + max_memtable_size: 256 * 1024 * 1024, // 256MB + max_memtable_rows: 100_000, // 100k rows + max_memtable_batches: 8_000, // 8k batches manifest_scan_batch_size: 2, max_unflushed_memtable_bytes: 1024 * 1024 * 1024, // 1GB backpressure_log_interval: Duration::from_secs(30), @@ -306,6 +318,20 @@ impl ShardWriterConfig { self } + /// Set the number of WAL persistence retries before the writer self-fences. + /// See [`ShardWriterConfig::max_wal_persist_retries`]. + pub fn with_max_wal_persist_retries(mut self, retries: usize) -> Self { + self.max_wal_persist_retries = retries; + self + } + + /// Set the base backoff before the first WAL persistence retry. + /// See [`ShardWriterConfig::wal_persist_retry_base_delay`]. + pub fn with_wal_persist_retry_base_delay(mut self, delay: Duration) -> Self { + self.wal_persist_retry_base_delay = delay; + self + } + /// Set maximum MemTable size. pub fn with_max_memtable_size(mut self, size: usize) -> Self { self.max_memtable_size = size; @@ -843,7 +869,7 @@ async fn replay_memtable_from_wal( None => break, Some(entry) => { if entry.writer_epoch > our_epoch { - return Err(Error::io(format!( + return Err(Error::fenced_by_peer(format!( "WAL replay aborted: entry at position {} has writer_epoch {} > our claimed epoch {} for shard {} (writer was fenced during open)", position, entry.writer_epoch, our_epoch, shard_id ))); @@ -1054,8 +1080,9 @@ impl SharedWriterState { let _memtable_flush_watcher = old_memtable.create_memtable_flush_completion(); if pending_wal_range.is_some() { - let completion_cell: WatchableOnceCell> = - WatchableOnceCell::new(); + let completion_cell: WatchableOnceCell< + std::result::Result, + > = WatchableOnceCell::new(); let completion_reader = completion_cell.reader(); old_memtable.set_wal_flush_completion(completion_reader); @@ -1349,6 +1376,10 @@ impl ShardWriter { manifest_store.clone(), epoch, position_hint_seed, + WalRetryConfig { + max_retries: config.max_wal_persist_retries, + base_delay: config.wal_persist_retry_base_delay, + }, )); // Fence the predecessor before replay (see `write_fence_sentinel`). @@ -1661,6 +1692,26 @@ impl ShardWriter { /// ``` #[instrument(name = "sw_delete", level = "info", skip_all, fields(batch_count = keys.len(), shard_id = %self.config.shard_id))] pub async fn delete(&self, keys: Vec) -> Result { + let (result, watcher) = self.delete_no_wait(keys).await?; + // Wait for durability if configured (mirrors `put` → `put_memtable`). + if let Some(mut watcher) = watcher { + watcher.wait().await?; + } + Ok(result) + } + + /// Like [`Self::delete`], but returns the durability watcher *without* + /// awaiting it — the tombstone lands in the in-memory tier the instant this + /// returns, but the index-driven LSM read only folds it once the watcher + /// resolves the flush (which advances the visibility watermark and updates + /// the PK index). The delete analog of [`Self::put_no_wait`], so a caller + /// can hold an external lock across only the in-memory insert and await + /// durability after releasing it. MemTable mode only. + #[instrument(name = "sw_delete_no_wait", level = "info", skip_all, fields(batch_count = keys.len(), shard_id = %self.config.shard_id))] + pub async fn delete_no_wait( + &self, + keys: Vec, + ) -> Result<(WriteResult, Option)> { if keys.is_empty() { return Err(Error::invalid_input("Cannot delete with empty key list")); } @@ -1687,7 +1738,7 @@ impl ShardWriter { build_tombstone_batch(&k, &writer_state.schema, &writer_state.pk_columns) }) .collect::>>()?; - self.put_memtable(tombstones, state, writer_state, backpressure) + self.put_memtable_no_wait(tombstones, state, writer_state, backpressure) .await } WriterMode::WalOnly { .. } => Err(Error::invalid_input( @@ -1697,7 +1748,9 @@ impl ShardWriter { } /// Like [`Self::put`], but returns the durability watcher *without* awaiting - /// it. The row is visible to reads on this writer the instant this returns; + /// it. The row lands in the in-memory tier the instant this returns, but the + /// index-driven LSM read only surfaces it once the watcher resolves the + /// flush (which advances the visibility watermark and updates the indexes); /// the caller awaits durability via the watcher (`None` when `durable_write` /// is off). /// @@ -1774,6 +1827,10 @@ impl ShardWriter { writer_state: &Arc, backpressure: &BackpressureController, ) -> Result<(WriteResult, Option)> { + // Reject writes on a fenced writer before mutating the memtable, so a + // poisoned writer can't drift further from the durable WAL. + self.wal_flusher.check_poisoned()?; + // Apply backpressure if needed (before acquiring main lock) backpressure .maybe_apply_backpressure(|| { @@ -1845,6 +1902,10 @@ impl ShardWriter { trigger: &StdRwLock, backpressure: &BackpressureController, ) -> Result { + // Reject writes on a fenced writer before enqueuing — see + // `put_memtable_no_wait`. + self.wal_flusher.check_poisoned()?; + // Apply backpressure against the pending queue before pushing. The // budget reuses `max_unflushed_memtable_bytes` since WAL-only mode // shares the same "in-memory bytes waiting for durable storage" @@ -1898,7 +1959,9 @@ impl ShardWriter { let mut reader = reader; match reader.await_value().await { Some(Ok(_)) => {} - Some(Err(msg)) => return Err(Error::io(msg)), + // Rebuild the typed error (peer fence vs. persistence-failure + // self-fence) so a WAL-only durable caller can tell them apart. + Some(Err(failure)) => return Err(failure.into_error()), None => { return Err(Error::io( "WAL flush handler exited before reporting durability", @@ -2113,6 +2176,7 @@ impl ShardWriter { .. } => { self.check_fenced().await?; + self.wal_flusher.check_poisoned()?; let mut state = state.write().await; if state.memtable.batch_count() == 0 { return Ok(()); @@ -2173,6 +2237,37 @@ impl ShardWriter { } } + /// Abort the writer without flushing. + /// + /// Shuts down the background flush tasks and leaves all buffered + /// memtable state to be dropped with the writer. Unlike + /// [`Self::close`], no WAL/MemTable flush is issued: pending in-memory + /// rows are discarded, not made durable, and no object-store IO is + /// performed. Used on drop-table, where the dataset directory is about + /// to be removed and a flush would only race fresh files back into a + /// doomed path. + /// + /// Caller-quiesce contract: `abort` takes `&self` (so it can be called + /// through the `Arc` callers hold) and therefore cannot + /// structurally bar a concurrent or subsequent `put` the way consuming + /// `close(self)` does. After abort the dispatchers are gone, so a later + /// `put` would buffer data that never flushes. Callers MUST stop + /// issuing writes before calling abort. + /// + /// Blocks until any flush already mid-`handle()` settles — + /// cancellation only fires between messages — so no flush task lingers + /// after abort returns. Idempotent: a second call re-cancels an + /// already-cancelled token and joins an already-emptied task set. + #[instrument(name = "sw_abort", level = "info", skip_all, fields(shard_id = %self.config.shard_id, epoch = self.epoch))] + pub async fn abort(&self) -> Result<()> { + info!( + "Aborting ShardWriter for shard {} (no flush)", + self.config.shard_id + ); + self.task_executor.shutdown_all().await?; + Ok(()) + } + /// Close the writer gracefully. /// /// Flushes pending data and shuts down background tasks. @@ -2359,9 +2454,10 @@ impl MessageHandler for WalFlushHandler { state.last_flushed_wal_entry_position.max(entry.position); } - // Notify completion if requested + // Notify completion if requested. Carry the typed fence reason through + // the cell (not just a string) so a waiter rebuilds the right error. if let Some(cell) = done { - cell.write(result.map_err(|e| e.to_string())); + cell.write(result.map_err(|e| WalFlushFailure::from_error(&e))); } Ok(()) @@ -2561,7 +2657,9 @@ impl MemTableFlushHandler { if let Some(mut completion_reader) = memtable.take_wal_flush_completion() { match completion_reader.await_value().await { Some(Ok(flush_result)) => flush_result.entry.map(|e| e.position), - Some(Err(e)) => return Err(Error::io(format!("WAL flush failed: {}", e))), + // Rebuild the typed error so a fence/poison reason + // propagates to the memtable-flush caller too. + Some(Err(e)) => return Err(e.into_error()), None => { return Err(Error::io( "WAL flush handler exited before reporting completion", @@ -2837,11 +2935,7 @@ impl WriteStatsSnapshot { /// Get average WAL flush size in bytes. pub fn avg_wal_flush_bytes(&self) -> Option { - if self.wal_flush_count > 0 { - Some(self.wal_flush_bytes / self.wal_flush_count) - } else { - None - } + self.wal_flush_bytes.checked_div(self.wal_flush_count) } /// Get WAL write throughput (bytes per second based on WAL flush time). @@ -2873,11 +2967,7 @@ impl WriteStatsSnapshot { /// Get average rows per index update. pub fn avg_index_update_rows(&self) -> Option { - if self.index_update_count > 0 { - Some(self.index_update_rows / self.index_update_count) - } else { - None - } + self.index_update_rows.checked_div(self.index_update_count) } /// Get average MemTable flush latency. @@ -2891,11 +2981,8 @@ impl WriteStatsSnapshot { /// Get average MemTable flush size in rows. pub fn avg_memtable_flush_rows(&self) -> Option { - if self.memtable_flush_count > 0 { - Some(self.memtable_flush_rows / self.memtable_flush_count) - } else { - None - } + self.memtable_flush_rows + .checked_div(self.memtable_flush_count) } /// Log stats summary using tracing (for structured telemetry). @@ -2950,8 +3037,10 @@ pub fn new_shared_stats() -> SharedWriteStats { #[cfg(test)] mod tests { use super::*; + use crate::dataset::mem_wal::test_util::failing_memory_store; use arrow_array::{Int32Array, StringArray}; use arrow_schema::{DataType, Field}; + use lance_core::FenceReason; use tempfile::TempDir; async fn create_local_store() -> (Arc, Path, String, TempDir) { @@ -3129,6 +3218,128 @@ mod tests { writer.close().await.unwrap(); } + /// `delete_no_wait` lands the tombstone in the in-memory tier (visible at + /// the batch-store level the instant it returns) and hands back the + /// durability watcher *without* awaiting it. Index-driven LSM read + /// visibility — folding the tombstone out — follows once the watcher + /// resolves the flush that advances the visibility watermark and updates + /// the PK index. The delete analog of + /// `test_put_no_wait_durable_visible_then_durable`. + #[tokio::test] + async fn test_shard_writer_delete_no_wait_durable_visible_after_watcher() { + use crate::dataset::mem_wal::scanner::LsmScanner; + use futures::TryStreamExt; + + let (store, base_path, base_uri, _temp) = create_local_store().await; + let schema = create_pk_test_schema(); + let config = ShardWriterConfig { + shard_id: Uuid::new_v4(), + durable_write: true, + ..Default::default() + }; + let shard_id = config.shard_id; + let writer = ShardWriter::open( + store, + base_path, + base_uri.clone(), + config, + schema.clone(), + vec![], + ) + .await + .unwrap(); + + writer + .put(vec![create_test_batch(&schema, 0, 5)]) + .await + .unwrap(); + + // Tombstone id=2 without blocking on durability. + let (_result, watcher) = writer + .delete_no_wait(vec![id_only_keys(&[2])]) + .await + .unwrap(); + + // The tombstone is in the in-memory tier immediately (5 rows + 1 + // tombstone), even though the index-driven read can't fold it until the + // flush behind the watcher lands. `durable_write` is on, so a watcher is + // returned to await. + assert_eq!(writer.memtable_stats().await.unwrap().row_count, 6); + let mut watcher = watcher.expect("durable_write returns a watcher"); + + // Awaiting the watcher waits for the flush, which advances the + // visibility watermark and updates the PK index — only then does the LSM + // read fold the delete. + watcher.wait().await.unwrap(); + + let refs = writer.in_memory_memtable_refs().await.unwrap(); + let scanner = LsmScanner::without_base_table( + schema.clone(), + base_uri, + vec![], + vec!["id".to_string()], + ) + .with_in_memory_memtables(shard_id, refs); + let batches: Vec = scanner + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let mut ids: Vec = Vec::new(); + for b in &batches { + let arr = b + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + ids.extend((0..arr.len()).map(|i| arr.value(i))); + } + ids.sort_unstable(); + assert_eq!( + ids, + vec![0, 1, 3, 4], + "delete folded by the LSM read once the watcher resolves" + ); + + writer.close().await.unwrap(); + } + + /// With `durable_write` off, `delete_no_wait` returns no watcher (nothing to + /// await), but the tombstone still lands in the in-memory tier. The delete + /// analog of `test_put_no_wait_non_durable_returns_no_watcher`. + #[tokio::test] + async fn test_shard_writer_delete_no_wait_non_durable_returns_no_watcher() { + let (store, base_path, base_uri, _temp) = create_local_store().await; + let schema = create_pk_test_schema(); + let config = ShardWriterConfig { + shard_id: Uuid::new_v4(), + durable_write: false, + ..Default::default() + }; + let writer = ShardWriter::open(store, base_path, base_uri, config, schema.clone(), vec![]) + .await + .unwrap(); + + writer + .put(vec![create_test_batch(&schema, 0, 5)]) + .await + .unwrap(); + + let (_result, watcher) = writer + .delete_no_wait(vec![id_only_keys(&[2])]) + .await + .unwrap(); + assert!(watcher.is_none(), "non-durable delete has nothing to await"); + + // Tombstone landed in the in-memory tier (5 rows + 1 tombstone). + assert_eq!(writer.memtable_stats().await.unwrap().row_count, 6); + + writer.close().await.unwrap(); + } + /// Read every surviving `id` through the full LSM read path (which folds /// `NOT _tombstone` per source) over the writer's *flushed* generations, /// with an optional filter. Mirrors how a query reads a WAL table after a @@ -4556,6 +4767,60 @@ mod tests { ])) } + // A durable write whose WAL PUT keeps failing poisons the writer with a + // typed persistence failure; the next write fails fast with the same reason; + // and once storage heals, reopening replays the WAL and writes resume. + #[tokio::test] + async fn test_writer_poisons_on_persistence_failure_and_recovers_on_reopen() { + let (store, base_path, controls) = failing_memory_store().await; + let base_uri = "memory:///"; + let shard_id = Uuid::new_v4(); + let schema = schema_with_pk(); + controls.fail_wal_puts(usize::MAX); + + let config = ShardWriterConfig { + max_wal_persist_retries: 1, + wal_persist_retry_base_delay: Duration::from_millis(1), + ..memtable_config_with_pk(shard_id) + }; + + let writer = ShardWriter::open( + store.clone(), + base_path.clone(), + base_uri, + config.clone(), + schema.clone(), + vec![], + ) + .await + .unwrap(); + + // The durable put waits on the flush, which poisons after its retries. + let err = writer + .put(vec![create_test_batch(&schema, 0, 1)]) + .await + .unwrap_err(); + assert_eq!(err.fence_reason(), Some(FenceReason::PersistenceFailure)); + + // A poisoned writer rejects further writes fast, same typed reason. + let err = writer + .put(vec![create_test_batch(&schema, 1, 1)]) + .await + .unwrap_err(); + assert_eq!(err.fence_reason(), Some(FenceReason::PersistenceFailure)); + drop(writer); + + // Storage heals: reopening replays the WAL and accepts writes again. + controls.recover(); + let writer = ShardWriter::open(store, base_path, base_uri, config, schema.clone(), vec![]) + .await + .unwrap(); + writer + .put(vec![create_test_batch(&schema, 2, 1)]) + .await + .unwrap(); + } + /// Replay-on-open recovers durable WAL entries that were never flushed /// to a Lance generation. Setup: writer A durably writes batches, drops /// without close (so MemTable freeze never runs); writer B reopens and @@ -4794,6 +5059,7 @@ mod tests { // hint seed irrelevant; the real position counter is discovered // lazily on the first append. 0, + WalRetryConfig::default(), ); high_epoch_appender .append(vec![create_test_batch(&schema, 999, 1)]) @@ -4815,6 +5081,15 @@ mod tests { let Err(err) = result else { panic!("expected open to fail with fence error during replay"); }; + // Assert the *typed* fence reason, not just the message: a regression + // reverting this to `Error::io` would still carry a message containing + // "fenced" and slip past a string check, but must not report a + // `FenceReason`. + assert_eq!( + err.fence_reason(), + Some(FenceReason::PeerClaimedEpoch), + "replay must abort with a typed peer-fence error, got: {err}" + ); let msg = err.to_string(); assert!( msg.contains("WAL replay aborted") && msg.contains("fenced"), @@ -5108,6 +5383,63 @@ mod tests { writer.close().await.unwrap(); } + /// `abort` tears down the background flush tasks WITHOUT flushing — + /// buffered memtable rows are discarded, not sealed into an L0 + /// generation the way `close` would. Idempotent on a second call. + #[tokio::test] + async fn test_abort_discards_without_flushing_and_is_idempotent() { + let (store, base_path, base_uri, _temp_dir) = create_local_store().await; + let schema = create_test_schema(); + + // Thresholds high enough that nothing auto-flushes; the rows stay + // in the active memtable until abort discards them. + let config = ShardWriterConfig { + shard_id: Uuid::new_v4(), + shard_spec_id: 0, + durable_write: false, + sync_indexed_write: false, + max_wal_buffer_size: 64 * 1024 * 1024, + max_wal_flush_interval: None, + max_memtable_size: 64 * 1024 * 1024, + manifest_scan_batch_size: 2, + ..Default::default() + }; + + let writer = ShardWriter::open(store, base_path, base_uri, config, schema.clone(), vec![]) + .await + .unwrap(); + + writer + .put(vec![create_test_batch(&schema, 0, 10)]) + .await + .unwrap(); + let flushed_before = writer + .manifest() + .await + .unwrap() + .map(|m| m.flushed_generations.len()) + .unwrap_or(0); + + writer.abort().await.unwrap(); + + // No generation was sealed — contrast with `close`, which flushes + // the 10 buffered rows into a new L0 generation. + let flushed_after = writer + .manifest() + .await + .unwrap() + .map(|m| m.flushed_generations.len()) + .unwrap_or(0); + assert_eq!( + flushed_after, flushed_before, + "abort must not flush a new L0 generation" + ); + + // Idempotent: re-cancels the already-cancelled token, joins an + // already-emptied task set. + writer.abort().await.unwrap(); + } + /// On a successful flush commit the sealed generation's rows land in the /// manifest immediately, but the in-memory handle is NOT dropped — it /// lingers for `frozen_memtable_grace` (so in-flight as-of reads keep @@ -5308,8 +5640,8 @@ mod shard_writer_tests { use arrow_schema::{DataType, Field, Schema as ArrowSchema}; use lance_arrow::FixedSizeListArrayExt; use lance_index::IndexType; - use lance_index::scalar::ScalarIndexParams; - use lance_index::scalar::inverted::InvertedIndexParams; + use lance_index::scalar::inverted::{InvertedIndexParams, InvertedListFormatVersion}; + use lance_index::scalar::{FullTextSearchQuery, ScalarIndexParams}; use lance_index::vector::ivf::IvfBuildParams; use lance_index::vector::pq::builder::PQBuildParams; use lance_linalg::distance::MetricType; @@ -5613,6 +5945,93 @@ mod shard_writer_tests { writer.close().await.unwrap(); } + #[tokio::test] + async fn test_mem_wal_maintained_fts_v1_flush_preserves_format() { + use tempfile::TempDir; + + let vector_dim = 32; + let schema = create_test_schema(vector_dim); + let temp_dir = TempDir::new().expect("Failed to create temp dir"); + let uri = format!("file://{}", temp_dir.path().display()); + + let initial = create_test_batch(&schema, 0, 16, vector_dim); + let batches = RecordBatchIterator::new([Ok(initial)], schema.clone()); + let mut dataset = Dataset::write(batches, &uri, Some(WriteParams::default())) + .await + .expect("Failed to create dataset"); + + let fts_params = + InvertedIndexParams::default().format_version(InvertedListFormatVersion::V1); + dataset + .create_index( + &["text"], + IndexType::Inverted, + Some("text_fts".to_string()), + &fts_params, + false, + ) + .await + .expect("Failed to create v1 FTS index"); + let base_indices = dataset.load_indices().await.unwrap(); + assert_eq!(base_indices.len(), 1); + assert_eq!(base_indices[0].index_version, 1); + + dataset + .initialize_mem_wal() + .maintained_indexes(["text_fts"]) + .execute() + .await + .expect("Failed to initialize MemWAL"); + + let shard_id = Uuid::new_v4(); + let config = ShardWriterConfig::new(shard_id) + .with_durable_write(true) + .with_sync_indexed_write(true); + let writer = dataset + .mem_wal_writer(shard_id, config) + .await + .expect("Failed to create MemWAL writer"); + writer + .put(vec![create_test_batch(&schema, 1_000, 3, vector_dim)]) + .await + .expect("Failed to write MemWAL batch"); + writer.close().await.expect("Failed to close writer"); + + let (store, base_path) = lance_io::object_store::ObjectStore::from_uri(&uri) + .await + .expect("Failed to open store"); + let manifest_store = + super::super::manifest::ShardManifestStore::new(store, &base_path, shard_id, 2); + let manifest = manifest_store + .read_latest() + .await + .expect("Failed to read manifest") + .expect("Manifest should exist"); + assert_eq!(manifest.flushed_generations.len(), 1); + + let flushed = &manifest.flushed_generations[0]; + let gen_uri = format!("{}/_mem_wal/{}/{}", uri, shard_id, flushed.path); + let flushed_dataset = Dataset::open(&gen_uri) + .await + .expect("Failed to open flushed generation"); + let flushed_indices = flushed_dataset.load_indices().await.unwrap(); + assert_eq!(flushed_indices.len(), 1); + assert_eq!(flushed_indices[0].name, "text_fts"); + assert_eq!( + flushed_indices[0].index_version, 1, + "maintained v1 FTS index must flush as v1" + ); + + let results = flushed_dataset + .scan() + .full_text_search(FullTextSearchQuery::new("Sample".to_owned())) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 3); + } + #[tokio::test] async fn test_writer_hnsw_params_override() { use lance_index::vector::hnsw::builder::HnswBuildParams; @@ -5980,7 +6399,7 @@ mod shard_writer_tests { (0..vector_dim as usize).map(|d| (target_id as f32 * 0.1 + d as f32 * 0.01).sin()), ); let mut scanner = writer.scan().await.unwrap(); - scanner.nearest("vector", Arc::new(query), 80); + scanner.nearest("vector", &query, 80).unwrap(); let result = scanner.try_into_batch().await.expect("Failed to scan"); assert!(result.num_rows() > 0, "vector query returned no rows"); diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index bc69269c93f..d352659bdcb 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -80,6 +80,7 @@ //! the successful tasks can be committed. You can also commit in batches if //! you wish. As long as the tasks don't rewrite any of the same fragments, //! they can be committed in any order. +use lance_core::utils::row_addr_remap::{GroupInput, RowAddrRemap}; use std::borrow::Cow; use std::collections::HashMap; use std::io::Cursor; @@ -114,6 +115,7 @@ use lance_core::datatypes::{BlobHandling, BlobKind}; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_core::utils::tracing::{DATASET_COMPACTING_EVENT, TRACE_DATASET_EVENTS}; use lance_index::frag_reuse::FragReuseGroup; +use lance_index::is_system_index; use lance_table::format::{Fragment, RowIdMeta}; use roaring::{RoaringBitmap, RoaringTreemap}; use serde::{Deserialize, Serialize}; @@ -154,6 +156,39 @@ impl TryFrom<&str> for CompactionMode { } } +/// Controls how the old-to-new row-address mapping is built when remapping +/// indices during compaction. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +pub enum IndexRemapMode { + /// Store a compact remap and compute row-address mappings during lookup. + /// + /// Best for large compactions where peak memory is the constraint. Uses + /// less memory, but each lookup does extra bitmap/range computation. + Compact, + /// Store the full row-address remap in memory for fast direct lookups. + /// + /// Best when the remap fits comfortably in memory and remap speed is the + /// priority. Uses more peak memory because every rewritten/deleted row has + /// a materialized mapping entry. + #[default] + Direct, +} + +impl TryFrom<&str> for IndexRemapMode { + type Error = Error; + + fn try_from(value: &str) -> std::result::Result { + match value.to_lowercase().as_str() { + "compact" => Ok(Self::Compact), + "direct" => Ok(Self::Direct), + _ => Err(Error::invalid_input(format!( + "Invalid index remap mode \"{}\". Valid values: \"compact\", \"direct\"", + value + ))), + } + } +} + /// Options to be passed to [compact_files]. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct CompactionOptions { @@ -202,6 +237,10 @@ pub struct CompactionOptions { /// not be remapped during this compaction operation. Instead, the fragment reuse index /// is updated and will be used to perform remapping later. pub defer_index_remap: bool, + /// How the old-to-new row-address mapping used to remap indices is built. + /// Defaults to [`IndexRemapMode::Direct`]. + #[serde(default)] + pub index_remap_mode: IndexRemapMode, /// The compaction mode to use. When set, this takes priority over the /// deprecated `enable_binary_copy` and `enable_binary_copy_force` fields. /// @@ -246,6 +285,7 @@ impl Default for CompactionOptions { batch_size: None, io_buffer_size: None, defer_index_remap: false, + index_remap_mode: IndexRemapMode::Direct, compaction_mode: None, enable_binary_copy: false, enable_binary_copy_force: false, @@ -271,6 +311,7 @@ impl CompactionOptions { /// - `lance.compaction.materialize_deletions` /// - `lance.compaction.materialize_deletions_threshold` /// - `lance.compaction.defer_index_remap` + /// - `lance.compaction.index_remap_mode` /// - `lance.compaction.batch_size` /// - `lance.compaction.io_buffer_size` /// - `lance.compaction.compaction_mode` @@ -348,6 +389,9 @@ impl CompactionOptions { } }; } + "index_remap_mode" => { + self.index_remap_mode = IndexRemapMode::try_from(value.as_str())?; + } "batch_size" => { self.batch_size = Some(value.parse().map_err(|_| { Error::invalid_input(format!( @@ -796,7 +840,7 @@ pub async fn compact_files_with_planner( let dataset_ref = &dataset.clone(); - let result_stream = futures::stream::iter(compaction_plan.tasks.into_iter()) + let result_stream = futures::stream::iter(compaction_plan.tasks) .map(|task| rewrite_files(Cow::Borrowed(dataset_ref), task, &compaction_plan.options)) .buffer_unordered( compaction_plan @@ -1410,7 +1454,11 @@ impl CandidateBin { async fn load_index_fragmaps(dataset: &Dataset) -> Result> { let indices = dataset.load_indices().await?; let mut index_fragmaps = Vec::with_capacity(indices.len()); - for index in indices.iter() { + // System indices (fragment-reuse, mem-wal) don't define data coverage and + // aren't remapped per rewrite group, so they must not constrain compaction + // bins -- otherwise deferred compaction's fragment-reuse index repeatedly + // splits the small-fragment run and they never coalesce. + for index in indices.iter().filter(|idx| !is_system_index(idx)) { if let Some(fragment_bitmap) = index.fragment_bitmap.as_ref() { index_fragmaps.push(fragment_bitmap.clone()); } else { @@ -1717,7 +1765,8 @@ async fn rewrite_files( let row_addrs = match row_addrs_result { Ok(v) => v, Err(e) => { - cleanup_data_fragments(&dataset.object_store, &dataset.base, &new_fragments).await; + cleanup_data_fragments(&dataset.object_store, &dataset.base, None, &new_fragments) + .await; return Err(e); } }; @@ -1892,8 +1941,8 @@ async fn recalc_versions_for_rewritten_fragments( // Set both version metadata on new fragments for ((fragment, last_updated_seq), created_at_seq) in new_fragments .iter_mut() - .zip(new_last_updated_sequences.into_iter()) - .zip(new_created_at_sequences.into_iter()) + .zip(new_last_updated_sequences) + .zip(new_created_at_sequences) { fragment.last_updated_at_version_meta = Some( lance_table::format::RowDatasetVersionMeta::from_sequence(&last_updated_seq).unwrap(), @@ -1960,7 +2009,8 @@ pub async fn commit_compaction( let mut rewrite_groups = Vec::with_capacity(completed_tasks.len()); let mut metrics = CompactionMetrics::default(); - let mut row_id_map: HashMap> = HashMap::default(); + let mut remap_group_inputs: Vec = Vec::new(); + let mut direct_row_id_map: HashMap> = HashMap::default(); let mut frag_reuse_groups: Vec = Vec::new(); let mut new_fragment_bitmap: RoaringBitmap = RoaringBitmap::new(); @@ -1975,12 +2025,41 @@ pub async fn commit_compaction( if let Some(row_addrs_bytes) = task.row_addrs { let row_addrs = RoaringTreemap::deserialize_from(&mut Cursor::new(&row_addrs_bytes))?; - let transposed = remapping::transpose_row_addrs( - row_addrs, - &task.original_fragments, - &task.new_fragments, - ); - row_id_map.extend(transposed); + match options.index_remap_mode { + IndexRemapMode::Direct => { + let transposed = remapping::transpose_row_addrs( + row_addrs, + &task.original_fragments, + &task.new_fragments, + ); + direct_row_id_map.extend(transposed); + } + IndexRemapMode::Compact => { + let new_frags = task + .new_fragments + .iter() + .map(|f| { + let physical_rows = f.physical_rows.ok_or_else(|| { + Error::invalid_input(format!( + "compacted fragment {} is missing physical_rows", + f.id + )) + })?; + Ok((f.id as u32, physical_rows as u32)) + }) + .collect::>>()?; + + remap_group_inputs.push(GroupInput { + rewritten_old_row_addrs: row_addrs, + old_frag_ids: task + .original_fragments + .iter() + .map(|f| f.id as u32) + .collect(), + new_frags, + }); + } + } } } else if options.defer_index_remap { let changed_row_addrs = task.row_addrs.ok_or_else(|| { @@ -2008,9 +2087,11 @@ pub async fn commit_compaction( .flat_map(|group| group.old_fragments.iter().map(|frag| frag.id)) .collect::>(); - let remapped_indices = index_remapper - .remap_indices(row_id_map, &affected_ids) - .await?; + let remap = match options.index_remap_mode { + IndexRemapMode::Direct => RowAddrRemap::direct(direct_row_id_map), + IndexRemapMode::Compact => RowAddrRemap::compact(remap_group_inputs)?, + }; + let remapped_indices = index_remapper.remap_indices(remap, &affected_ids).await?; remapped_indices .into_iter() .map(|rewritten| RewrittenIndex { @@ -2070,7 +2151,13 @@ pub async fn commit_compaction( .apply_commit(transaction, &Default::default(), &Default::default()) .await { - cleanup_data_fragments(&dataset.object_store, &dataset.base, &all_new_fragments).await; + cleanup_data_fragments( + &dataset.object_store, + &dataset.base, + None, + &all_new_fragments, + ) + .await; return Err(e); } @@ -2134,6 +2221,7 @@ mod tests { let fragment = Fragment { id: 0, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(0), @@ -2234,18 +2322,31 @@ mod tests { impl IndexRemapper for MockIndexRemapper { async fn remap_indices( &self, - index_map: HashMap>, + index_map: RowAddrRemap, _: &[u64], ) -> Result> { for expectation in &self.expectations { - if expectation.expected == index_map { + let matches = match &index_map { + RowAddrRemap::Direct(map) => map == &expectation.expected, + RowAddrRemap::Compact(_) => { + let expected_frags: RoaringBitmap = expectation + .expected + .keys() + .map(|addr| (addr >> 32) as u32) + .collect(); + index_map.affected_fragments() == expected_frags + && expectation + .expected + .iter() + .all(|(k, v)| index_map.get(*k) == Some(*v)) + } + }; + if matches { return Ok(expectation.answer.clone()); } } panic!( - "Unexpected index map (len={}): {}\n Options: {}", - index_map.len(), - Self::stringify_map(&index_map), + "Unexpected index map; expected one of:\n {}", self.expectations .iter() .map(|expectation| Self::stringify_map(&expectation.expected)) @@ -2781,11 +2882,7 @@ mod tests { #[async_trait] impl IndexRemapper for IgnoreRemap { - async fn remap_indices( - &self, - _: HashMap>, - _: &[u64], - ) -> Result> { + async fn remap_indices(&self, _: RowAddrRemap, _: &[u64]) -> Result> { Ok(Vec::new()) } } @@ -3402,6 +3499,69 @@ mod tests { } } + #[tokio::test] + async fn test_deferred_compaction_not_split_by_frag_reuse_index() { + // A deferred compaction creates a fragment-reuse index covering its + // output. Later small fragments must still compact together with that + // (FRI-covered) output: the FRI is a system index and must not split the + // compaction bin. Without the fix the FRI-covered fragment is isolated, + // so only the new fragments merge and the count never returns to one. + let data = sample_data(); + let test_dir = TempStrDir::default(); + let test_uri = &test_dir; + let options = CompactionOptions { + defer_index_remap: true, + ..Default::default() + }; + + // Two small fragments -> deferred compaction folds them into one, + // creating the fragment-reuse index. + let reader = RecordBatchIterator::new(vec![Ok(data.slice(0, 400))], data.schema()); + let mut dataset = Dataset::write( + reader, + test_uri, + Some(WriteParams { + max_rows_per_file: 200, + ..Default::default() + }), + ) + .await + .unwrap(); + compact_files(&mut dataset, options.clone(), None) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), 1); + assert!( + dataset + .load_index_by_name(FRAG_REUSE_INDEX_NAME) + .await + .unwrap() + .is_some() + ); + + // Append two more small fragments, then compact again. + let reader = RecordBatchIterator::new(vec![Ok(data.slice(400, 400))], data.schema()); + let mut dataset = Dataset::write( + reader, + test_uri, + Some(WriteParams { + max_rows_per_file: 200, + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), 3); + + compact_files(&mut dataset, options, None).await.unwrap(); + assert_eq!( + dataset.get_fragments().len(), + 1, + "FRI-covered fragment must compact together with the new fragments" + ); + } + #[tokio::test] async fn test_remap_index_after_compaction() { let mut data_gen = BatchGenerator::new() @@ -4217,6 +4377,356 @@ mod tests { ); } + #[rstest] + #[case(IndexRemapMode::Compact)] + #[case(IndexRemapMode::Direct)] + #[tokio::test] + async fn test_btree_index_remap_after_compaction(#[case] index_remap_mode: IndexRemapMode) { + let mut dataset = lance_datagen::gen_batch() + .col( + "vec", + lance_datagen::array::rand_vec::(Dimension::from(32)), + ) + .col("id", lance_datagen::array::step::()) + .into_ram_dataset(FragmentCount::from(6), FragmentRowCount::from(1000)) + .await + .unwrap(); + + // Delete rows scattered across fragments so the remap must drop some old + // addresses and shift the survivors. + dataset.delete("id % 10 == 0").await.unwrap(); + + dataset + .create_index( + &["id"], + IndexType::BTree, + Some("id_idx".into()), + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + + let count_low = dataset + .count_rows(Some("id < 1000".to_owned())) + .await + .unwrap(); + let count_mid = dataset + .count_rows(Some("id >= 2000 and id < 3000".to_owned())) + .await + .unwrap(); + let count_high = dataset + .count_rows(Some("id >= 5000".to_owned())) + .await + .unwrap(); + + let options = CompactionOptions { + target_rows_per_fragment: 50_000, + index_remap_mode, + ..Default::default() + }; + let metrics = compact_files(&mut dataset, options, None).await.unwrap(); + assert!(metrics.fragments_removed > 0); + assert!(metrics.fragments_added > 0); + + // The index was remapped inline and must still drive scans. + let mut scanner = dataset.scan(); + scanner.filter("id >= 2000 and id < 3000").unwrap(); + scanner.project::(&[]).unwrap().with_row_id(); + let plan = scanner.explain_plan(false).await.unwrap(); + assert!( + plan.contains("ScalarIndexQuery: query=[id >= 2000 && id < 3000]@id_idx(BTree)"), + "Expected scalar index query in plan: {}", + plan + ); + + // Counts resolved through the remapped index match the pre-compaction + // values in both remap modes. + assert_eq!( + dataset + .count_rows(Some("id < 1000".to_owned())) + .await + .unwrap(), + count_low + ); + assert_eq!( + dataset + .count_rows(Some("id >= 2000 and id < 3000".to_owned())) + .await + .unwrap(), + count_mid + ); + assert_eq!( + dataset + .count_rows(Some("id >= 5000".to_owned())) + .await + .unwrap(), + count_high + ); + } + + #[rstest] + #[case(IndexRemapMode::Compact)] + #[case(IndexRemapMode::Direct)] + #[tokio::test] + async fn test_ivf_pq_index_remap_after_compaction(#[case] index_remap_mode: IndexRemapMode) { + use arrow_array::cast::AsArray; + use lance_index::vector::pq::PQBuildParams; + + const DIM: u32 = 32; + let mut dataset = lance_datagen::gen_batch() + .col("id", lance_datagen::array::step::()) + .col( + "vec", + lance_datagen::array::rand_vec::(Dimension::from(DIM)), + ) + .into_ram_dataset(FragmentCount::from(6), FragmentRowCount::from(1000)) + .await + .unwrap(); + + let params = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + small_ivf(), + PQBuildParams { + max_iters: 2, + num_sub_vectors: 2, + ..Default::default() + }, + ); + dataset + .create_index( + &["vec"], + IndexType::Vector, + Some("vec_idx".into()), + ¶ms, + false, + ) + .await + .unwrap(); + let original_uuid = dataset + .load_index_by_name("vec_idx") + .await + .unwrap() + .unwrap() + .uuid; + + // Delete rows scattered across fragments so the remap must drop some old + // addresses and shift the survivors. + dataset.delete("id % 10 == 0").await.unwrap(); + + // Sample queries from surviving vectors and capture the pre-compaction + // KNN answer and the surviving id set. + let mut survivors: Vec<(i32, Vec)> = Vec::new(); + { + let mut scanner = dataset.scan(); + scanner.project(&["id", "vec"]).unwrap(); + let batches = scanner + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + for batch in &batches { + let ids = batch["id"].as_primitive::(); + let vecs = batch["vec"].as_fixed_size_list(); + for i in 0..batch.num_rows() { + let v = vecs.value(i); + survivors.push(( + ids.value(i), + v.as_primitive::().values().to_vec(), + )); + } + } + } + let surviving_ids: std::collections::HashSet = + survivors.iter().map(|(id, _)| *id).collect(); + let step = (survivors.len() / 16).max(1); + let queries: Vec> = survivors + .iter() + .step_by(step) + .map(|(_, v)| v.clone()) + .collect(); + let k = 10; + let mut baseline: Vec> = Vec::new(); + for q in &queries { + baseline.push(vector_knn_ids(&dataset, q, k).await); + } + + // Inline remap (defer_index_remap = false): compaction physically + // rebuilds the vector index through the configured remap mode. + let metrics = compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 50_000, + index_remap_mode, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + assert!(metrics.fragments_removed > 0); + assert!(metrics.fragments_added > 0); + + // The index was physically remapped inline, so its uuid must change. + assert_ne!( + dataset + .load_index_by_name("vec_idx") + .await + .unwrap() + .unwrap() + .uuid, + original_uuid, + "vector index must be physically remapped inline" + ); + + // The remap only relabels row addresses; it must not resurrect deleted + // rows, and KNN must stay close to the pre-compaction answer in both + // remap modes. + for (i, q) in queries.iter().enumerate() { + let after = vector_knn_ids(&dataset, q, k).await; + for id in &after { + assert!( + surviving_ids.contains(id), + "KNN returned id {id} that is not a surviving row (query #{i}, mode {index_remap_mode:?})" + ); + } + let overlap = after.iter().filter(|id| baseline[i].contains(id)).count(); + assert!( + overlap >= 8, + "KNN top-{k} diverged after compaction: overlap {overlap} < 8 (query #{i}, mode {index_remap_mode:?})" + ); + } + } + + #[rstest] + #[case(IndexRemapMode::Compact)] + #[case(IndexRemapMode::Direct)] + #[tokio::test] + async fn test_inverted_index_remap_after_compaction(#[case] index_remap_mode: IndexRemapMode) { + use arrow_array::cast::AsArray; + + let mut dataset = lance_datagen::gen_batch() + .col("id", lance_datagen::array::step::()) + .col("doc", lance_datagen::array::random_sentence(1, 100, false)) + .into_ram_dataset(FragmentCount::from(6), FragmentRowCount::from(1000)) + .await + .unwrap(); + + dataset + .create_index( + &["doc"], + IndexType::Inverted, + Some("doc_idx".into()), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + let original_uuid = dataset + .load_index_by_name("doc_idx") + .await + .unwrap() + .unwrap() + .uuid; + + // Sample a few words from a real document to drive full-text searches. + let words: Vec = { + let mut scanner = dataset.scan(); + scanner + .project(&["doc"]) + .unwrap() + .limit(Some(1), None) + .unwrap(); + let batches = scanner + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut words: Vec = batches[0]["doc"] + .as_string::() + .value(0) + .split_whitespace() + .map(|s| s.to_string()) + .collect(); + words.sort(); + words.dedup(); + words.truncate(3); + words + }; + assert!(!words.is_empty(), "sampled document must contain words"); + + // Delete rows scattered across fragments so the remap must drop some old + // addresses and shift the survivors. + dataset.delete("id % 10 == 0").await.unwrap(); + + // Capture the post-deletion full-text-search counts (resolved through the + // index + deletion vectors) before compaction physically remaps. + let mut before = Vec::new(); + for word in &words { + let mut scanner = dataset.scan(); + scanner + .full_text_search(FullTextSearchQuery::new(word.clone())) + .unwrap(); + scanner.project::(&[]).unwrap().with_row_id(); + before.push(scanner.count_rows().await.unwrap()); + } + + // Inline remap (defer_index_remap = false): compaction physically rebuilds + // the inverted index through the configured remap mode. + let options = CompactionOptions { + target_rows_per_fragment: 50_000, + index_remap_mode, + ..Default::default() + }; + let metrics = compact_files(&mut dataset, options, None).await.unwrap(); + assert!(metrics.fragments_removed > 0); + assert!(metrics.fragments_added > 0); + + // The index was physically remapped inline, so its uuid must change. + assert_ne!( + dataset + .load_index_by_name("doc_idx") + .await + .unwrap() + .unwrap() + .uuid, + original_uuid, + "inverted index must be physically remapped inline (mode {index_remap_mode:?})" + ); + + // The remapped index must still drive full-text search. + let mut scanner = dataset.scan(); + scanner + .full_text_search(FullTextSearchQuery::new(words[0].clone())) + .unwrap(); + scanner.project::(&[]).unwrap().with_row_id(); + let plan = scanner.explain_plan(true).await.unwrap(); + assert!( + plan.contains("MatchQuery"), + "Expected inverted index scan in plan: {}", + plan + ); + + // Counts resolved through the remapped index match the pre-compaction + // values in both remap modes. + for (word, expected) in words.iter().zip(before) { + let mut scanner = dataset.scan(); + scanner + .full_text_search(FullTextSearchQuery::new(word.clone())) + .unwrap(); + scanner.project::(&[]).unwrap().with_row_id(); + assert_eq!( + scanner.count_rows().await.unwrap(), + expected, + "full-text count for {word:?} changed after compaction (mode {index_remap_mode:?})" + ); + } + } + #[tokio::test] async fn test_read_inverted_index_with_defer_index_remap() { // Generate random words using lance-datagen @@ -5611,6 +6121,10 @@ mod tests { "lance.compaction.binary_copy_read_batch_bytes".to_string(), "8388608".to_string(), ), + ( + "lance.compaction.index_remap_mode".to_string(), + "compact".to_string(), + ), ]); let opts = CompactionOptions::from_dataset_config(&config).unwrap(); @@ -5624,6 +6138,8 @@ mod tests { assert_eq!(opts.io_buffer_size, Some(1_073_741_824)); assert_eq!(opts.compaction_mode, Some(CompactionMode::TryBinaryCopy)); assert_eq!(opts.binary_copy_read_batch_bytes, Some(8_388_608)); + // A non-default value proves the config string was actually parsed. + assert_eq!(opts.index_remap_mode, IndexRemapMode::Compact); } #[test] @@ -5643,6 +6159,8 @@ mod tests { defaults.materialize_deletions_threshold ); assert_eq!(opts.defer_index_remap, defaults.defer_index_remap); + assert_eq!(opts.index_remap_mode, defaults.index_remap_mode); + assert_eq!(opts.index_remap_mode, IndexRemapMode::Direct); assert_eq!(opts.batch_size, defaults.batch_size); assert_eq!(opts.compaction_mode, defaults.compaction_mode); assert_eq!( diff --git a/rust/lance/src/dataset/optimize/remapping.rs b/rust/lance/src/dataset/optimize/remapping.rs index 266ac977a69..aef2cd231fc 100644 --- a/rust/lance/src/dataset/optimize/remapping.rs +++ b/rust/lance/src/dataset/optimize/remapping.rs @@ -12,6 +12,7 @@ use crate::{Dataset, index}; use async_trait::async_trait; use lance_core::Error; use lance_core::utils::address::RowAddress; +use lance_core::utils::row_addr_remap::RowAddrRemap; use lance_index::frag_reuse::{FRAG_REUSE_INDEX_NAME, FragDigest}; use lance_table::format::{Fragment, IndexFile, IndexMetadata}; use lance_table::io::manifest::read_manifest_indexes; @@ -51,7 +52,7 @@ pub struct RemappedIndex { pub trait IndexRemapper: Send + Sync { async fn remap_indices( &self, - index_map: HashMap>, + index_map: RowAddrRemap, affected_fragment_ids: &[u64], ) -> Result>; } @@ -69,11 +70,7 @@ pub struct IgnoreRemap {} #[async_trait] impl IndexRemapper for IgnoreRemap { - async fn remap_indices( - &self, - _: HashMap>, - _: &[u64], - ) -> Result> { + async fn remap_indices(&self, _: RowAddrRemap, _: &[u64]) -> Result> { Ok(Vec::new()) } } @@ -316,7 +313,8 @@ async fn remap_index(dataset: &mut Dataset, index_id: &Uuid) -> Result<()> { .map(|old_addr| (old_addr, frag_reuse_index.remap_row_id(old_addr))) .collect(); - let remap_result = index::remap_index(dataset, index_id, &composed_row_id_map).await?; + let remapper = RowAddrRemap::direct(composed_row_id_map); + let remap_result = index::remap_index(dataset, index_id, &remapper).await?; let new_index_meta = match remap_result { // The composed remap emptied the index (every row deleted). Matching the @@ -411,6 +409,87 @@ pub async fn remap_column_index( mod tests { use super::*; + #[test] + fn test_compact_matches_transpose() { + use lance_core::utils::row_addr_remap::GroupInput; + // Ascending old fragments (compaction's scan order), with deletions. + let old = vec![ + FragDigest { + id: 0, + physical_rows: 5, + num_deleted_rows: 2, + }, + FragDigest { + id: 1, + physical_rows: 4, + num_deleted_rows: 1, + }, + FragDigest { + id: 3, + physical_rows: 3, + num_deleted_rows: 0, + }, + ]; + // 9 rewritten rows (offsets that survived in each old fragment). + let rewritten = [ + (0, 1), + (0, 2), + (0, 4), + (1, 0), + (1, 1), + (1, 3), + (3, 0), + (3, 1), + (3, 2), + ]; + let addrs = RoaringTreemap::from_iter( + rewritten + .iter() + .map(|(f, o)| u64::from(RowAddress::new_from_parts(*f, *o))), + ); + // 9 rewritten rows split across two new fragments. + let new = vec![ + FragDigest { + id: 10, + physical_rows: 4, + num_deleted_rows: 0, + }, + FragDigest { + id: 11, + physical_rows: 5, + num_deleted_rows: 0, + }, + ]; + + let expected = transpose_row_ids_from_digest(addrs.clone(), &old, &new); + let compact = RowAddrRemap::compact([GroupInput { + rewritten_old_row_addrs: addrs, + old_frag_ids: old.iter().map(|f| f.id as u32).collect(), + new_frags: new + .iter() + .map(|f| (f.id as u32, f.physical_rows as u32)) + .collect(), + }]) + .unwrap(); + + // Every real address in the old fragments must map identically. + for f in &old { + for o in 0..f.physical_rows as u32 { + let a = u64::from(RowAddress::new_from_parts(f.id as u32, o)); + assert_eq!( + compact.get(a), + expected.get(&a).copied(), + "mismatch at ({}, {})", + f.id, + o + ); + } + } + // A fragment outside the group is unaffected by both. + let outside = u64::from(RowAddress::new_from_parts(99, 0)); + assert_eq!(compact.get(outside), expected.get(&outside).copied()); + } + #[test] fn test_missing_indices() { // Sanity test to make sure MissingIds works. Does not test actual functionality so diff --git a/rust/lance/src/dataset/refs.rs b/rust/lance/src/dataset/refs.rs index 98b4f0cbc0a..0d3f65f7959 100644 --- a/rust/lance/src/dataset/refs.rs +++ b/rust/lance/src/dataset/refs.rs @@ -21,7 +21,6 @@ use std::cmp::Ordering; use std::collections::HashMap; use std::fmt; use std::fmt::Formatter; -use std::io::ErrorKind; use uuid::Uuid; pub const MAIN_BRANCH: &str = "main"; @@ -476,7 +475,7 @@ impl Branches<'_> { if !self.object_store().exists(&manifest_file.path).await? { return Err(Error::VersionNotFound { - message: format!("Manifest file {} does not exist", &manifest_file.path), + message: format!("Manifest file {} does not exist", manifest_file.path), }); }; @@ -611,16 +610,8 @@ impl Branches<'_> { && let Err(e) = self.refs.object_store.remove_dir_all(delete_path).await { match &e { - Error::IO { source, .. } => { - if let Some(io_err) = source.downcast_ref::() { - if io_err.kind() == ErrorKind::NotFound { - log::debug!("Branch directory already deleted: {}", io_err); - } else { - return Err(e); - } - } else { - return Err(e); - } + Error::NotFound { .. } => { + log::debug!("Branch directory already deleted"); } _ => return Err(e), } diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 7c7a082ad7e..6832fa28ed9 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -91,7 +91,9 @@ use crate::index::scalar_logical::scalar_index_fragment_bitmap; use crate::index::vector::utils::{ default_distance_type_for, get_vector_dim, get_vector_type, validate_distance_type_for, }; -use crate::io::exec::filtered_read::{FilteredReadExec, FilteredReadOptions}; +use crate::io::exec::filtered_read::{ + FilteredReadExec, FilteredReadOptions, FilteredReadThreadingMode, +}; use crate::io::exec::fts::{ BoostQueryExec, FlatMatchFilterExec, FlatMatchQueryExec, MatchQueryExec, PhraseQueryExec, }; @@ -1358,24 +1360,25 @@ impl Scanner { /// used by the scanner. If the buffer is full then the scanner will block until /// the buffer is processed. /// - /// Generally this should scale with the number of concurrent I/O threads. The - /// default is 2GiB which comfortably provides enough space for somewhere between - /// 32 and 256 concurrent I/O threads. + /// Generally this should scale with the number of concurrent I/O threads. If + /// unset, v2 scans choose a default based on the object store and + /// `LANCE_DEFAULT_IO_BUFFER_SIZE` can override that default. /// /// This value is not a hard cap on the amount of RAM the scanner will use. Some /// space is used for the compute (which can be controlled by the batch size) and /// Lance does not keep track of memory after it is returned to the user. /// - /// Currently, if there is a single batch of data which is larger than the io buffer - /// size then the scanner will deadlock. This is a known issue and will be fixed in - /// a future release. pub fn io_buffer_size(&mut self, size: u64) -> &mut Self { self.io_buffer_size = Some(size); self } - /// Set the prefetch size. - /// Ignored in v2 and newer format + /// Set the number of batches to decode concurrently. + /// + /// This bounds the decode fan-out of the scan: at most this many batch-decode + /// tasks run in flight at once. Defaults to `get_num_compute_intensive_cpus()`. + /// + /// `nbatches` must be greater than zero. pub fn batch_readahead(&mut self, nbatches: usize) -> &mut Self { self.batch_readahead = nbatches; self @@ -1745,7 +1748,7 @@ impl Scanner { .field(&column.column_name) .ok_or(Error::invalid_input(format!( "Column {} not found", - &column.column_name + column.column_name )))?; } } @@ -1902,6 +1905,38 @@ impl Scanner { } } + /// Ensure `input` exposes `column_name` as a top-level column. + /// + /// Nested FTS flat-search paths read the projected struct column from storage + /// but the FTS executor consumes a single document column by name. + fn ensure_column_alias( + &self, + input: Arc, + column_name: &str, + ) -> Result> { + let input_schema = input.schema(); + if input_schema.column_with_name(column_name).is_some() { + return Ok(input); + } + + let mut projection_exprs = Vec::with_capacity(input_schema.fields().len() + 1); + for field in input_schema.fields() { + projection_exprs.push(( + Arc::new(Column::new_with_schema( + field.name(), + input_schema.as_ref(), + )?) as Arc, + field.name().clone(), + )); + } + projection_exprs.push(( + Self::create_column_expr(column_name, self.dataset.as_ref(), input_schema.as_ref())?, + column_name.to_string(), + )); + + Ok(Arc::new(ProjectionExec::try_new(projection_exprs, input)?)) + } + /// Set whether to use statistics to optimize the scan (default: true) /// /// This is used for debugging or benchmarking purposes. @@ -2313,11 +2348,12 @@ impl Scanner { MaterializationStyle::AllLate => false, MaterializationStyle::AllEarlyExcept(ref cols) => !cols.contains(&(field.id as u32)), MaterializationStyle::Heuristic => { - if field.is_blob() { - // By default, blobs are loaded as descriptions, and so should be early - // - // TODO: Once we make blob handling configurable, we should use the blob - // handling setting here. + if field.is_blob() && self.blob_handling.returns_description(field) { + // A blob returned as a description (offset + size) is tiny, so it is + // cheaper to read eagerly. When blob_handling materializes the full + // binary value instead (e.g. `all_binary`), fall through to the + // width-based heuristic so a selective filter can late-materialize it + // rather than reading the whole column. return true; } @@ -2365,6 +2401,12 @@ impl Scanner { } fn validate_options(&self) -> Result<()> { + if self.batch_readahead == 0 { + return Err(Error::invalid_input_source( + "batch_readahead must be greater than 0, got 0".into(), + )); + } + if self.include_deleted_rows && !self.projection_plan.physical_projection.with_row_id { return Err(Error::invalid_input_source( "include_deleted_rows is set but with_row_id is false".into(), @@ -2877,6 +2919,11 @@ impl Scanner { read_options = read_options.with_batch_size(batch_size as u32); } + // Bound the decode fan-out by `batch_readahead`. + read_options = read_options.with_threading_mode( + FilteredReadThreadingMode::OnePartitionMultipleThreads(self.batch_readahead), + ); + if let Some(file_reader_options) = self.resolved_file_reader_options() { read_options = read_options.with_file_reader_options(file_reader_options); } @@ -3590,7 +3637,7 @@ impl Scanner { ))? .clone(); - let mut columns = vec![column]; + let mut columns = vec![column.clone()]; if let Some(refine_expr) = filter_plan.refine_expr.as_ref() { columns.extend(Planner::column_names_in_expr(refine_expr)); } @@ -3614,6 +3661,7 @@ impl Scanner { if let Some(refine_expr) = filter_plan.refine_expr.as_ref() { plan = Arc::new(LanceFilterExec::try_new(refine_expr.clone(), plan)?); } + plan = self.ensure_column_alias(plan, &column)?; let flat_match_plan = Arc::new(FlatMatchQueryExec::new( self.dataset.clone(), @@ -4363,7 +4411,8 @@ impl Scanner { .dataset .empty_projection() .union_column(&column, OnMissing::Error)?; - self.take(input, projection)? + let input = self.take(input, projection)?; + self.ensure_column_alias(input, &column)? } else { input }; @@ -4413,7 +4462,8 @@ impl Scanner { .dataset .empty_projection() .union_column(&column, OnMissing::Error)?; - self.take(input, projection)? + let input = self.take(input, projection)?; + self.ensure_column_alias(input, &column)? } else { input }; @@ -4898,24 +4948,25 @@ impl Scanner { } } +fn is_fts_indexable_field(field: &Field) -> bool { + match field.data_type() { + DataType::Utf8 | DataType::LargeUtf8 => true, + DataType::List(inner_field) | DataType::LargeList(inner_field) => { + matches!( + inner_field.data_type(), + DataType::Utf8 | DataType::LargeUtf8 + ) + } + _ => false, + } +} + // Search over all indexed fields including nested ones, collecting columns that have an // inverted index async fn fts_indexed_columns(dataset: Arc) -> Result> { let mut indexed_columns = Vec::new(); for field in dataset.schema().fields_pre_order() { - // Check if this field is a string type that could have an inverted index - let is_string_field = match field.data_type() { - DataType::Utf8 | DataType::LargeUtf8 => true, - DataType::List(inner_field) | DataType::LargeList(inner_field) => { - matches!( - inner_field.data_type(), - DataType::Utf8 | DataType::LargeUtf8 - ) - } - _ => false, - }; - - if is_string_field { + if is_fts_indexable_field(field) { // Build the full field path for nested fields let column_path = if let Some(ancestors) = dataset.schema().field_ancestry_by_id(field.id) { @@ -5199,7 +5250,13 @@ pub mod test_dataset { } pub async fn make_fts_index(&mut self) -> Result<()> { - let params = InvertedIndexParams::default().with_position(true); + // These scanner tests search for the token "s" (from the `s-{N}` + // column values) to exercise fragment/append coverage, and "s" is + // in the full English stop-word list. Keep the token searchable; + // stop-word behavior itself is covered by the tokenizer tests. + let params = InvertedIndexParams::default() + .with_position(true) + .remove_stop_words(false); self.dataset .create_index(&["s"], IndexType::Inverted, None, ¶ms, true) .await?; @@ -5631,6 +5688,149 @@ mod test { Ok(()) } + // Regression for #6580: a scan with `filter` + `project` of a + // `(Large)List` column used to panic in `merge_with_schema` + // (called from `TakeStream::map_batch`) because the filtered batch arrived + // as a sliced view of a larger batch and the cloned list offsets did not + // start at zero. The trigger requires (a) a `(Large)List` + // projection where the struct is split across `filtered_read` and + // `TakeExec` and (b) a sparse-tail selectivity pattern so the trailing + // filter result lands deep inside the values buffer of its source batch. + // Parametrized over `List`/`LargeList` since the fix touches both offset + // widths in `merge_with_schema`. + #[rstest] + #[tokio::test] + async fn test_filter_project_list_struct_sparse_tail( + // The panic is specific to v2.x storage; the legacy reader takes a + // different code path. V2_0 and V2_2 are the versions called out in + // the original report. + #[values( + LanceFileVersion::V2_0, + LanceFileVersion::Stable, + LanceFileVersion::V2_2 + )] + data_storage_version: LanceFileVersion, + #[values(false, true)] large_list: bool, + ) { + use arrow_array::{LargeListArray, ListArray, UInt16Array}; + use arrow_buffer::{OffsetBuffer, ScalarBuffer}; + + let struct_fields = Fields::from(vec![ + Arc::new(ArrowField::new("a", DataType::Int32, true)), + Arc::new(ArrowField::new("b", DataType::Int32, true)), + ]); + let item_field = Arc::new(ArrowField::new( + "item", + DataType::Struct(struct_fields.clone()), + true, + )); + let items_dtype = if large_list { + DataType::LargeList(item_field.clone()) + } else { + DataType::List(item_field.clone()) + }; + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("grp", DataType::UInt16, false), + ArrowField::new("items", items_dtype, false), + ])); + + let make_batch = |start: i32, n: usize, group: u16| -> RecordBatch { + let ids = Int32Array::from_iter_values(start..start + n as i32); + let groups = UInt16Array::from(vec![group; n]); + + let mut offsets = Vec::with_capacity(n + 1); + let mut a_vals: Vec = Vec::new(); + let mut b_vals: Vec = Vec::new(); + offsets.push(0i64); + for i in 0..n { + // Variable-length lists (1..=18) so offsets don't land on + // batch-row boundaries. + let len = 1 + (i % 18); + for j in 0..len { + a_vals.push(j as i32); + b_vals.push(-(j as i32)); + } + offsets.push(a_vals.len() as i64); + } + let struct_arr = Arc::new(StructArray::new( + struct_fields.clone(), + vec![ + Arc::new(Int32Array::from(a_vals)) as ArrayRef, + Arc::new(Int32Array::from(b_vals)) as ArrayRef, + ], + None, + )); + let items: ArrayRef = if large_list { + Arc::new(LargeListArray::new( + item_field.clone(), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + struct_arr, + None, + )) + } else { + let offsets_i32: Vec = offsets.iter().map(|&o| o as i32).collect(); + Arc::new(ListArray::new( + item_field.clone(), + OffsetBuffer::new(ScalarBuffer::from(offsets_i32)), + struct_arr, + None, + )) + }; + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(ids) as ArrayRef, + Arc::new(groups) as ArrayRef, + items, + ], + ) + .unwrap() + }; + + // Sparse-tail selectivity (matching the original report's shape at a + // smaller scale): a large leading block of matches, a large gap of + // non-matches, then a small trailing match. Single fragment. + let batches = vec![ + make_batch(0, 100_000, 7), + make_batch(100_000, 400_000, 1), + make_batch(500_000, 7_300, 7), + ]; + + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let dataset = Dataset::write( + reader, + "memory://", + Some(WriteParams { + max_rows_per_file: 1_000_000, + data_storage_version: Some(data_storage_version), + ..Default::default() + }), + ) + .await + .unwrap(); + + // Force a column split inside the `items` struct by marking `items.b` + // as a late-materialized field: `filtered_read` returns the batch with + // `items.a`, and `TakeExec` adds `items.b`. `merge_with_schema` then + // takes its `List` branch, which is where the panic was. + let items_b_field_id = dataset + .schema() + .field("items") + .unwrap() + .child("item") + .unwrap() + .child("b") + .unwrap() + .id as u32; + let mut scan = dataset.scan(); + scan.filter("grp = 7").unwrap(); + scan.project(&["id", "items"]).unwrap(); + scan.materialization_style(MaterializationStyle::AllEarlyExcept(vec![items_b_field_id])); + let result = scan.try_into_batch().await.unwrap(); + assert_eq!(result.num_rows(), 107_300); + } + #[tokio::test] async fn test_scan_regexp_match_and_non_empty_captions() { // Build a small dataset with three Utf8 columns and verify the full @@ -10082,6 +10282,160 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") assert_io_lt!(io_stats, read_bytes, index_scan_bytes); } + #[tokio::test] + async fn test_blob_all_binary_late_materialization() { + // A selective filter that projects a blob column with `blob_handling=all_binary` + // must late-materialize the blob (take only the matched rows) rather than eagerly + // reading the whole column. Blobs returned as descriptions stay eager (they are + // tiny), but full binary values should follow the width-based heuristic like any + // other wide column. + use lance_io::assert_io_lt; + use lance_table::io::commit::RenameCommitHandler; + + // 8KB stays under the 64KB inline threshold, so the blob is a normal column in + // the data file rather than a dedicated blob file. + let blob_meta = std::collections::HashMap::from([( + "lance-encoding:blob".to_string(), + "true".to_string(), + )]); + let blobs = array::rand_fixedbin(ByteCount::from(8 * 1024), true).with_metadata(blob_meta); + let data = gen_batch() + .col("filterme", array::step::()) + .col("blobs", blobs) + .into_reader_rows(RowCount::from(500), BatchCount::from(8)); + + let dataset = Dataset::write( + data, + "memory://test", + Some(WriteParams { + commit_handler: Some(Arc::new(RenameCommitHandler)), + data_storage_version: Some(LanceFileVersion::Stable), + ..Default::default() + }), + ) + .await + .unwrap(); + + // Baseline: read the blob column as binary for the whole table. + let _ = dataset.object_store.as_ref().io_stats_incremental(); // reset + dataset + .scan() + .project(&["blobs"]) + .unwrap() + .blob_handling(BlobHandling::AllBinary) + .try_into_batch() + .await + .unwrap(); + let full_scan_bytes = dataset + .object_store + .as_ref() + .io_stats_incremental() + .read_bytes; + + // A filter matching a single row out of 4000 should read far less than the whole + // column: only the filter leaf plus the one materialized blob. + dataset + .scan() + .project(&["blobs"]) + .unwrap() + .blob_handling(BlobHandling::AllBinary) + .filter("filterme = 100") + .unwrap() + .try_into_batch() + .await + .unwrap(); + let io_stats = dataset.object_store.as_ref().io_stats_incremental(); + assert_io_lt!(io_stats, read_bytes, full_scan_bytes); + } + + #[tokio::test] + async fn test_nested_blob_all_binary_late_materialization() { + // Same as above, but the blob is a leaf *inside* a struct and the filter is on a + // sibling leaf. Materialization is decided per leaf (fields_pre_order), so the + // nested blob must late-materialize under `all_binary` just like a top-level one. + use lance_io::assert_io_lt; + use lance_table::io::commit::RenameCommitHandler; + + let blob_meta = std::collections::HashMap::from([( + "lance-encoding:blob".to_string(), + "true".to_string(), + )]); + let a_field = ArrowField::new("a", DataType::Int32, false); + let blob_field = + ArrowField::new("blob", DataType::LargeBinary, false).with_metadata(blob_meta); + let struct_fields: Fields = vec![a_field, blob_field].into(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "s", + DataType::Struct(struct_fields.clone()), + false, + )])); + + let rows_per_batch = 500usize; + let batches: Vec = (0..8) + .map(|b| { + let base = (b * rows_per_batch) as i32; + let a = Arc::new(Int32Array::from_iter_values( + base..base + rows_per_batch as i32, + )); + // Vary the payload per row so it does not collapse under compression. + let blobs: Vec> = (0..rows_per_batch) + .map(|r| { + let seed = (base as usize + r).wrapping_mul(2654435761); + (0usize..8 * 1024) + .map(|i| (i.wrapping_mul(31).wrapping_add(seed) & 0xff) as u8) + .collect() + }) + .collect(); + let blob = Arc::new(arrow_array::LargeBinaryArray::from_iter_values( + blobs.iter().map(|v| v.as_slice()), + )); + let s = StructArray::new(struct_fields.clone(), vec![a, blob as ArrayRef], None); + RecordBatch::try_new(schema.clone(), vec![Arc::new(s)]).unwrap() + }) + .collect(); + + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let dataset = Dataset::write( + reader, + "memory://test", + Some(WriteParams { + commit_handler: Some(Arc::new(RenameCommitHandler)), + data_storage_version: Some(LanceFileVersion::Stable), + ..Default::default() + }), + ) + .await + .unwrap(); + + let _ = dataset.object_store.as_ref().io_stats_incremental(); // reset + dataset + .scan() + .project(&["s"]) + .unwrap() + .blob_handling(BlobHandling::AllBinary) + .try_into_batch() + .await + .unwrap(); + let full_scan_bytes = dataset + .object_store + .as_ref() + .io_stats_incremental() + .read_bytes; + + dataset + .scan() + .project(&["s"]) + .unwrap() + .blob_handling(BlobHandling::AllBinary) + .filter("s.a = 100") + .unwrap() + .try_into_batch() + .await + .unwrap(); + let io_stats = dataset.object_store.as_ref().io_stats_incremental(); + assert_io_lt!(io_stats, read_bytes, full_scan_bytes); + } + #[rstest] #[tokio::test] async fn test_project_nested( @@ -12374,6 +12728,47 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") assert_eq!(filtered.options().io_buffer_size_bytes, Some(7777)); } + #[tokio::test] + async fn test_batch_readahead_bounds_decode_concurrency() { + let data = lance_datagen::gen_batch() + .col("x", lance_datagen::array::step::()) + .into_reader_rows(RowCount::from(8), BatchCount::from(1)); + let dataset = Dataset::write(data, "memory://test_batch_readahead_concurrency", None) + .await + .unwrap(); + + // Default: threading mode falls back to get_num_compute_intensive_cpus(). + let plan = dataset.scan().create_plan().await.unwrap(); + let filtered = find_filtered_read(plan.as_ref()) + .expect("expected a FilteredReadExec in the scan plan"); + assert_eq!( + filtered.options().threading_mode, + FilteredReadThreadingMode::OnePartitionMultipleThreads(get_num_compute_intensive_cpus()), + ); + + // Explicit batch_readahead(N) bounds the decode fan-out to N. + let mut scanner = dataset.scan(); + scanner.batch_readahead(3); + let plan = scanner.create_plan().await.unwrap(); + let filtered = find_filtered_read(plan.as_ref()) + .expect("expected a FilteredReadExec in the scan plan"); + assert_eq!( + filtered.options().threading_mode, + FilteredReadThreadingMode::OnePartitionMultipleThreads(3), + ); + + let mut scanner = dataset.scan(); + scanner.batch_readahead(0); + let Err(Error::InvalidInput { source, .. }) = scanner.create_plan().await else { + panic!("expected batch_readahead=0 to be rejected"); + }; + assert!( + source + .to_string() + .contains("batch_readahead must be greater than 0") + ); + } + // The env var key scopes serial_test's lock so this test only blocks others // that touch LANCE_DEFAULT_IO_BUFFER_SIZE — unrelated tests still run in // parallel. diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 867801b37a5..40eac95e919 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -375,11 +375,15 @@ pub(super) async fn add_columns_to_fragments( NewColumnTransform::AllNulls(output_schema) => { check_names(output_schema.as_ref())?; - // Check that the schema is compatible considering all the new columns must be nullable - let schema = Schema::try_from(output_schema.as_ref())?; - if !schema.all_fields_nullable() { + // AllNulls is metadata-only; missing columns are synthesized as nulls at + // read time, so only each new top-level column needs to be nullable. + if let Some(field) = output_schema.fields().iter().find(|f| !f.is_nullable()) { return Err(Error::invalid_input_source( - "All-null columns must be nullable.".into(), + format!( + "All-null columns must be nullable, but field '{}' is not.", + field.name() + ) + .into(), )); } @@ -488,6 +492,7 @@ async fn cleanup_new_column_data_files(fragments: &[FileFragment], new_fragments cleanup_data_fragments( &first_fragment.dataset().object_store, &first_fragment.dataset().base, + None, &fragments_to_cleanup, ) .await; @@ -1952,6 +1957,7 @@ mod test { Ok(Some(Fragment { files: vec![], id: 0, + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(50), @@ -2116,8 +2122,8 @@ mod test { .await .unwrap_err(); assert!( - err.to_string() - .contains("All-null columns must be nullable.") + matches!(err, Error::InvalidInput { .. }), + "unexpected error: {err}" ); let data = dataset.scan().try_into_batch().await?; @@ -2131,6 +2137,107 @@ mod test { Ok(()) } + /// `AllNulls` accepts any nullable top-level column whatever its inner-field + /// nullability (Map/List/Struct with non-null children); non-nullable ones are rejected. + #[tokio::test] + async fn test_add_column_all_nulls_nested() -> Result<()> { + let num_rows = 100; + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..num_rows))], + )?; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + + let test_dir = TempStrDir::default(); + let mut dataset = Dataset::write( + reader, + &test_dir, + Some(WriteParams { + max_rows_per_file: 50, + max_rows_per_group: 25, + data_storage_version: Some(LanceFileVersion::Stable), + ..Default::default() + }), + ) + .await?; + + let map_with_non_null_entries = DataType::Map( + Arc::new(ArrowField::new( + "entries", + DataType::Struct(ArrowFields::from(vec![ + ArrowField::new("key", DataType::Utf8, false), + ArrowField::new("value", DataType::Float64, true), + ])), + false, + )), + false, + ); + let list_with_non_null_items = + DataType::List(Arc::new(ArrowField::new("item", DataType::Utf8, false))); + let struct_with_non_null_child = + DataType::Struct(ArrowFields::from(vec![ArrowField::new( + "a", + DataType::Int32, + false, + )])); + + dataset + .add_columns( + NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new(vec![ + ArrowField::new("cutoffs", map_with_non_null_entries.clone(), true), + ArrowField::new("tags", list_with_non_null_items.clone(), true), + ArrowField::new("info", struct_with_non_null_child.clone(), true), + ]))), + None, + None, + ) + .await?; + + let data = dataset.scan().try_into_batch().await?; + assert_eq!(data.num_rows(), num_rows as usize); + for (name, expected_type) in [ + ("cutoffs", &map_with_non_null_entries), + ("tags", &list_with_non_null_items), + ("info", &struct_with_non_null_child), + ] { + let column = data.column_by_name(name).unwrap(); + assert_eq!( + column.data_type(), + expected_type, + "type mismatch for {name}" + ); + assert_eq!( + column.null_count(), + num_rows as usize, + "column {name} should be all-null" + ); + } + + // A non-nullable top-level field is still rejected, and the error names it. + let err = + dataset + .add_columns( + NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new(vec![ + ArrowField::new("non_null_cutoffs", map_with_non_null_entries, false), + ]))), + None, + None, + ) + .await + .unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { .. }), + "unexpected error: {err}" + ); + + Ok(()) + } + #[tokio::test] async fn test_add_column_all_nulls_legacy() -> Result<()> { let num_rows = 100; diff --git a/rust/lance/src/dataset/statistics.rs b/rust/lance/src/dataset/statistics.rs index 4719a639068..627ccfc1081 100644 --- a/rust/lance/src/dataset/statistics.rs +++ b/rust/lance/src/dataset/statistics.rs @@ -5,11 +5,16 @@ use std::{collections::HashMap, future::Future, sync::Arc}; +use datafusion::scalar::ScalarValue; use futures::{StreamExt, TryStreamExt}; -use lance_core::Result; +use lance_core::{Error, Result}; +use lance_index::metrics::NoOpMetricsCollector; +use lance_index::scalar::zonemap::ZoneMapIndex; use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; +use roaring::RoaringBitmap; use super::{Dataset, fragment::FileFragment}; +use crate::index::{DatasetIndexExt, DatasetIndexInternalExt}; /// Statistics about a single field in the dataset pub struct FieldStatistics { @@ -83,6 +88,98 @@ impl DatasetStatisticsExt for Dataset { } } +/// A read-only handle for cheap, index-derived statistics about a [`Dataset`]. +/// +/// Obtained via [`Dataset::statistics`]. Groups statistics accessors behind one +/// handle instead of accreting one-off methods on [`Dataset`]. Every accessor is +/// served from index metadata and never scans data. +#[derive(Debug, Clone, Copy)] +pub struct DatasetStatistics<'a> { + dataset: &'a Dataset, +} + +impl<'a> DatasetStatistics<'a> { + pub(crate) fn new(dataset: &'a Dataset) -> Self { + Self { dataset } + } + + /// Global `[min, max]` for `column` from its min/max-capable scalar index + /// (currently ZoneMap), without a scan. + /// + /// `None` unless the column's index segments *jointly* cover every live + /// fragment and the column can be soundly bounded — fragments appended after + /// the index was built, or a NaN-bearing column, yield `None`. The disjoint + /// segments of a multi-segment index are folded together. + /// + /// When `Some`, the range is a superset of live values, conservative under + /// deletion vectors: safe to prune with. See [`ScalarIndex::value_range`]. + /// + /// [`ScalarIndex::value_range`]: lance_index::scalar::ScalarIndex::value_range + pub async fn column_value_range( + &self, + column: &str, + ) -> Result> { + let dataset = self.dataset; + let Some(field) = dataset.schema().field(column) else { + return Err(Error::invalid_input(format!( + "column_value_range: column '{column}' not found in dataset schema" + ))); + }; + let field_id = field.id; + let field_path = dataset.schema().field_path(field_id)?; + + // A multi-segment ZoneMap is several index entries over the same column, + // each covering a disjoint fragment subset. Match the field, then the + // details type (the column may also carry e.g. a BTree). + let indices = dataset.load_indices().await?; + let segments: Vec<_> = indices + .iter() + .filter(|idx| matches!(idx.fields.as_slice(), [only] if *only == field_id)) + .filter(|idx| { + idx.index_details + .as_ref() + .is_some_and(|d| d.type_url.ends_with("ZoneMapIndexDetails")) + }) + .collect(); + if segments.is_empty() { + return Ok(None); + } + + // Soundness: the segments must *jointly* cover every live fragment, else + // the fold sees only a subset and could prune live rows (e.g. fragments + // appended after the index was built). Extra dead fragments are harmless. + let mut covered = RoaringBitmap::new(); + for idx in &segments { + let Some(bitmap) = idx.fragment_bitmap.as_ref() else { + return Ok(None); + }; + covered |= bitmap.clone(); + } + if !dataset.fragment_bitmap.as_ref().is_subset(&covered) { + return Ok(None); + } + + // Keep the opened indices alive so the `ZoneMapIndex` refs we fold over + // stay borrowed. + let mut opened = Vec::with_capacity(segments.len()); + for idx in &segments { + opened.push( + dataset + .open_generic_index(&field_path, &idx.uuid, &NoOpMetricsCollector) + .await?, + ); + } + let Some(zonemaps) = opened + .iter() + .map(|index| index.as_any().downcast_ref::()) + .collect::>>() + else { + return Ok(None); + }; + Ok(ZoneMapIndex::value_range_over(zonemaps)) + } +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 267296c984b..86682004f04 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -34,8 +34,10 @@ use lance_core::utils::tempfile::TempStrDir; use lance_datagen::{BatchCount, Dimension, RowCount, array, gen_batch}; use lance_file::reader::{FileReader, FileReaderOptions}; use lance_file::version::LanceFileVersion; +use lance_index::optimize::OptimizeOptions; use lance_index::scalar::FullTextSearchQuery; use lance_index::scalar::inverted::{ + InvertedListFormatVersion, query::{BooleanQuery, MatchQuery, Occur, Operator, PhraseQuery}, tokenizer::InvertedIndexParams, }; @@ -861,6 +863,148 @@ async fn test_fts_on_multiple_columns() { assert_eq!(results.num_rows(), 1); } +fn nested_fts_batch( + ids: Vec, + a_values: Vec>, + b_values: Vec>, +) -> RecordBatch { + let a_values = Arc::new(StringArray::from(a_values)) as ArrayRef; + let b_values = Arc::new(StringArray::from(b_values)) as ArrayRef; + let struct_array = StructArray::from(vec![ + ( + Arc::new(Field::new("a", DataType::Utf8, true)), + a_values.clone(), + ), + ( + Arc::new(Field::new("b", DataType::Utf8, true)), + b_values.clone(), + ), + ]); + let struct_type = struct_array.data_type().clone(); + RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("s", struct_type, true), + ])), + vec![ + Arc::new(UInt64Array::from(ids)) as ArrayRef, + Arc::new(struct_array) as ArrayRef, + ], + ) + .unwrap() +} + +async fn nested_fts_result_ids(dataset: &Dataset, query: FullTextSearchQuery) -> Vec { + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let mut ids = batch["id"].as_primitive::().values().to_vec(); + ids.sort_unstable(); + ids +} + +#[tokio::test] +async fn test_fts_on_nested_fields() { + let batch = nested_fts_batch( + vec![0, 1, 2, 3], + vec![ + Some("lance nested alpha"), + Some("plain text"), + None, + Some("phrase target here"), + ], + vec![ + Some("metadata only"), + Some("database nested beta"), + Some("lance beta"), + Some("other"), + ], + ); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write(batches, &test_uri, None).await.unwrap(); + + dataset + .create_index( + &["s.a"], + IndexType::Inverted, + None, + &InvertedIndexParams::default().with_position(true), + true, + ) + .await + .unwrap(); + dataset + .create_index( + &["s.b"], + IndexType::Inverted, + None, + &InvertedIndexParams::default(), + true, + ) + .await + .unwrap(); + + let indices = dataset.load_indices().await.unwrap(); + let indexed_fields = indices + .iter() + .map(|index| dataset.schema().field_path(index.fields[0]).unwrap()) + .collect::>(); + assert_eq!( + indexed_fields, + HashSet::from(["s.a".to_string(), "s.b".to_string()]) + ); + + let query = FullTextSearchQuery::new_query(FtsQuery::Match( + MatchQuery::new("alpha".to_owned()).with_column(Some("s.a".to_owned())), + )); + assert_eq!(nested_fts_result_ids(&dataset, query).await, vec![0]); + + let query = FullTextSearchQuery::new_query(FtsQuery::Match( + MatchQuery::new("beta".to_owned()).with_column(Some("s.b".to_owned())), + )); + assert_eq!(nested_fts_result_ids(&dataset, query).await, vec![1, 2]); + + assert_eq!( + nested_fts_result_ids(&dataset, FullTextSearchQuery::new("lance".to_owned())).await, + vec![0, 2] + ); + + let query = FullTextSearchQuery::new_query(FtsQuery::MultiMatch(MultiMatchQuery { + match_queries: vec![ + MatchQuery::new("nested".to_owned()).with_column(Some("s.a".to_owned())), + MatchQuery::new("nested".to_owned()).with_column(Some("s.b".to_owned())), + ], + })); + assert_eq!(nested_fts_result_ids(&dataset, query).await, vec![0, 1]); + + let query = FullTextSearchQuery::new_query( + PhraseQuery::new("phrase target".to_owned()) + .with_column(Some("s.a".to_owned())) + .into(), + ); + assert_eq!(nested_fts_result_ids(&dataset, query).await, vec![3]); + + let append_batch = nested_fts_batch( + vec![4, 5], + vec![Some("fresh lance append"), Some("plain append")], + vec![Some("other"), Some("fresh beta append")], + ); + let schema = append_batch.schema(); + let batches = RecordBatchIterator::new(vec![append_batch].into_iter().map(Ok), schema); + dataset.append(batches, None).await.unwrap(); + + assert_eq!( + nested_fts_result_ids(&dataset, FullTextSearchQuery::new("fresh".to_owned())).await, + vec![4, 5] + ); +} + #[tokio::test] async fn test_fts_unindexed_data() { let params = InvertedIndexParams::default(); @@ -936,6 +1080,68 @@ async fn test_fts_unindexed_data() { assert_eq!(results.num_rows(), 1); } +#[tokio::test] +async fn test_fts_v1_remains_queryable_after_append_optimize() { + let params = InvertedIndexParams::default().format_version(InvertedListFormatVersion::V1); + let text_col = StringArray::from(vec!["alpha original", "beta original"]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![Field::new( + "text", + text_col.data_type().to_owned(), + false, + )]) + .into(), + vec![Arc::new(text_col) as ArrayRef], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(batches, "memory://test.lance", None) + .await + .unwrap(); + dataset + .create_index(&["text"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + assert_eq!(dataset.load_indices().await.unwrap()[0].index_version, 1); + + let appended = StringArray::from(vec!["alpha appended"]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![Field::new( + "text", + appended.data_type().to_owned(), + false, + )]) + .into(), + vec![Arc::new(appended) as ArrayRef], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + dataset.append(batches, None).await.unwrap(); + dataset + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap(); + + let results = dataset + .scan() + .full_text_search(FullTextSearchQuery::new("alpha".to_owned())) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 2); + assert!( + dataset + .load_indices() + .await + .unwrap() + .iter() + .all(|index| index.index_version == 1) + ); +} + #[tokio::test] async fn test_fts_unindexed_data_with_stop_words() { // When indexed data has avg_doc_length < 1.0 (e.g. single-word stop words @@ -1739,6 +1945,169 @@ async fn test_fts_index_with_large_string() { test_fts_index::(true).await; } +#[tokio::test] +async fn test_fts_list_index_uses_row_level_documents() { + let tempdir = TempStrDir::default(); + let uri = tempdir.to_owned(); + drop(tempdir); + + let mut list_col = GenericListBuilder::::new(GenericStringBuilder::::new()); + list_col.values().append_value("lance"); + list_col.values().append_value("lance database"); + list_col.append(true); + list_col.values().append_value("database"); + list_col.append(true); + list_col.append(true); + list_col.values().append_null(); + list_col.append(true); + list_col.append(false); + + let docs = Arc::new(list_col.finish()) as ArrayRef; + let ids = Arc::new(UInt64Array::from_iter_values(0..docs.len() as u64)) as ArrayRef; + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("doc", docs.data_type().clone(), true), + ArrowField::new("id", DataType::UInt64, false), + ])), + vec![docs, ids], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let mut dataset = Dataset::write(batches, &uri, None).await.unwrap(); + + dataset + .create_index( + &["doc"], + IndexType::Inverted, + None, + &InvertedIndexParams::default(), + true, + ) + .await + .unwrap(); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new("lance".to_owned()).limit(Some(10))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result["id"].as_primitive::().values(), &[0]); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new("database".to_owned()).limit(Some(10))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let mut ids = result["id"] + .as_primitive::() + .values() + .iter() + .copied() + .collect::>(); + ids.sort_unstable(); + assert_eq!(ids, vec![0, 1], "{:?}", result); +} + +#[tokio::test] +async fn test_fts_list_phrase_query_can_cross_elements() { + assert_fts_list_phrase_query_can_cross_elements::().await; +} + +#[tokio::test] +async fn test_fts_large_list_phrase_query_can_cross_elements() { + assert_fts_list_phrase_query_can_cross_elements::().await; +} + +async fn assert_fts_list_phrase_query_can_cross_elements() { + let tempdir = TempStrDir::default(); + let uri = tempdir.to_owned(); + drop(tempdir); + + let mut list_col = GenericListBuilder::::new(GenericStringBuilder::::new()); + let rows: &[&[&str]] = &[ + &["alpha", "beta"], + &["want the", "apple"], + &["want", "apple"], + ]; + for values in rows.iter().copied() { + for value in values { + list_col.values().append_value(value); + } + list_col.append(true); + } + + let docs = Arc::new(list_col.finish()) as ArrayRef; + let ids = Arc::new(UInt64Array::from(vec![0u64, 1, 2])) as ArrayRef; + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("doc", docs.data_type().clone(), true), + ArrowField::new("id", DataType::UInt64, false), + ])), + vec![docs, ids], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let mut dataset = Dataset::write(batches, &uri, None).await.unwrap(); + + let params = InvertedIndexParams::default() + .with_position(true) + .remove_stop_words(true); + dataset + .create_index(&["doc"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query(PhraseQuery::new("alpha beta".to_owned()).into()) + .limit(Some(10)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result["id"].as_primitive::().values(), &[0]); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query(PhraseQuery::new("want the apple".to_owned()).into()) + .limit(Some(10)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result["id"].as_primitive::().values(), &[1]); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query(PhraseQuery::new("want apple".to_owned()).into()) + .limit(Some(10)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result["id"].as_primitive::().values(), &[2]); +} + #[tokio::test] async fn test_fts_accented_chars() { let ds = create_fts_dataset::(false, false, InvertedIndexParams::default()).await; @@ -3679,3 +4048,68 @@ async fn test_legacy_dataset_uses_v2_0_for_indexes() { "Index files should never use legacy format, even for legacy datasets" ); } + +#[tokio::test] +async fn test_manifest_read_recovers_from_stale_size() { + // A cached `ManifestLocation.size` can lag the real object: a reader may pick + // up a size from a stale listing/hint while another writer is committing + // concurrently. Reading the manifest (or its index section) with that stale + // size must not fail with a spurious "file size is too small" error. The + // reader should drop the cached size, fetch the true size, and succeed. + use crate::session::Session; + use lance_table::io::commit::ManifestLocation; + use lance_table::io::manifest::read_manifest_indexes; + + let test_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from((0..100).collect::>()))], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + + let mut dataset = Dataset::write(reader, &test_uri, None).await.unwrap(); + dataset + .create_index( + &["id"], + IndexType::BTree, + Some("id_idx".to_string()), + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + + let real_location = dataset.manifest_location().clone(); + assert!(real_location.size.is_some()); + + // A deliberately-too-small size stands in for a stale cached size. Without the + // retry, both reads below decode a bogus footer offset and fail with + // "file size is too small". + let stale_location = ManifestLocation { + size: Some(1), + ..real_location.clone() + }; + + let session = Session::default(); + let manifest = Dataset::load_manifest( + dataset.object_store.as_ref(), + &stale_location, + test_uri.as_ref(), + &session, + ) + .await + .expect("load_manifest should recover from a stale manifest size"); + assert_eq!(manifest.version, real_location.version); + + let indices = read_manifest_indexes(dataset.object_store.as_ref(), &stale_location, &manifest) + .await + .expect("read_manifest_indexes should recover from a stale manifest size"); + assert_eq!(indices.len(), 1); + assert_eq!(indices[0].name, "id_idx"); +} diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index c0c90fc71c9..f9618914037 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -12,10 +12,11 @@ use crate::dataset::WriteMode::Overwrite; use crate::dataset::builder::DatasetBuilder; use crate::dataset::{ManifestWriteConfig, write_manifest_file}; use crate::session::Session; +use crate::session::caches::ManifestKey; use crate::{Dataset, Error, Result}; use lance_table::format::DataStorageFormat; -use crate::dataset::write::{WriteMode, WriteParams}; +use crate::dataset::write::{CommitBuilder, InsertBuilder, WriteMode, WriteParams}; use arrow::array::as_struct_array; use arrow::compute::concat_batches; use arrow_array::RecordBatch; @@ -33,7 +34,10 @@ use lance_arrow::bfloat16::{self, BFLOAT16_EXT_NAME}; use lance_arrow::{ARROW_EXT_META_KEY, ARROW_EXT_NAME_KEY}; use lance_core::utils::tempfile::{TempStdDir, TempStrDir}; use lance_datagen::{BatchCount, RowCount, array, gen_batch}; -use lance_file::version::LanceFileVersion; +use lance_file::{ + version::LanceFileVersion, + writer::{FileWriter, FileWriterOptions}, +}; use lance_io::assert_io_eq; use lance_table::feature_flags; use lance_table::format::BasePath; @@ -241,6 +245,62 @@ async fn test_with_object_store_wrappers_wraps_base_store_params() { assert!(request_tracker.incremental_stats().read_iops > 0); } +#[tokio::test] +async fn test_store_params_for_base_resolves_base_scoped_options() { + let test_dir = TempStdDir::default(); + create_file(&test_dir, WriteMode::Create, LanceFileVersion::Stable).await; + let uri = test_dir.to_str().unwrap(); + let dataset = Arc::new(Dataset::open(uri).await.unwrap()); + + let base_dir = tempfile::tempdir().unwrap(); + let base_uri = file_object_store_uri(base_dir.path()); + let base = BasePath::new(1, base_uri, Some("base".to_string()), true); + dataset.add_bases(vec![base.clone()], None).await.unwrap(); + + // Reopen with a single flat storage options map carrying base-scoped + // entries (`base_.`) next to shared defaults. + let dataset = DatasetBuilder::from_uri(uri) + .with_storage_options(HashMap::from([ + ("shared_option".to_string(), "shared".to_string()), + ( + "base_1.scoped_option".to_string(), + "base1-value".to_string(), + ), + ])) + .load() + .await + .unwrap(); + + // The registered base resolves the scoped entry on top of shared defaults. + let base_path = dataset.manifest.base_paths.get(&1).unwrap().clone(); + let params = dataset.store_params_for_base(Some(&base_path)); + assert_eq!( + params.storage_options().unwrap(), + &HashMap::from([ + ("shared_option".to_string(), "shared".to_string()), + ("scoped_option".to_string(), "base1-value".to_string()), + ]) + ); + + // The default scope keeps only the shared defaults. + let params = dataset.store_params_for_base(None); + assert_eq!( + params.storage_options().unwrap(), + &HashMap::from([("shared_option".to_string(), "shared".to_string())]) + ); + + // The base store resolved from scoped options is usable end to end. + let base_store = dataset.object_store(Some(1)).await.unwrap(); + let probe = base + .extract_path(dataset.session().store_registry()) + .unwrap() + .join("data") + .join("probe.lance"); + base_store.put(&probe, b"hello").await.unwrap(); + let read = base_store.inner.get_range(&probe, 0..5).await.unwrap(); + assert_eq!(read.as_ref(), b"hello"); +} + #[tokio::test] async fn test_with_object_store_wrappers_wraps_refs_store() { let test_dir = tempfile::tempdir().unwrap(); @@ -327,6 +387,118 @@ async fn test_create_data_file_uses_base_object_store() { assert!(tracker.incremental_stats().read_iops > 0); } +#[tokio::test] +async fn test_create_data_file_rejects_nested_schema_mismatch() { + let dataset_uri = format!( + "memory://test_create_data_file_rejects_nested_schema_mismatch/{}", + uuid::Uuid::new_v4() + ); + + let dataset_struct_fields = vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ]; + let dataset_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "s", + DataType::Struct(dataset_struct_fields.clone().into()), + true, + )])); + let dataset_struct = StructArray::try_new( + dataset_struct_fields.clone().into(), + vec![ + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + Arc::new(Int32Array::from(vec![2])) as ArrayRef, + ], + None, + ) + .unwrap(); + let dataset_batch = + RecordBatch::try_new(dataset_schema.clone(), vec![Arc::new(dataset_struct)]).unwrap(); + let dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(dataset_batch)], dataset_schema), + &dataset_uri, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(); + + async fn write_replacement_file( + dataset: &Dataset, + file_name: &str, + struct_fields: Vec, + ) { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "s", + DataType::Struct(struct_fields.clone().into()), + true, + )])); + let values = struct_fields + .iter() + .enumerate() + .map(|(idx, _)| Arc::new(Int32Array::from(vec![idx as i32])) as ArrayRef) + .collect::>(); + let struct_array = StructArray::try_new(struct_fields.into(), values, None).unwrap(); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(struct_array) as ArrayRef]).unwrap(); + + let object_writer = dataset + .object_store + .create(&dataset.data_dir().join(file_name)) + .await + .unwrap(); + let mut writer = FileWriter::try_new( + object_writer, + crate::datatypes::Schema::try_from(schema.as_ref()).unwrap(), + FileWriterOptions { + format_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }, + ) + .unwrap(); + writer.write_batch(&batch).await.unwrap(); + writer.finish().await.unwrap(); + } + + write_replacement_file( + &dataset, + "nested_reordered.lance", + vec![ + ArrowField::new("b", DataType::Int32, true), + ArrowField::new("a", DataType::Int32, true), + ], + ) + .await; + let err = dataset + .create_data_file("nested_reordered.lance", None) + .await + .unwrap_err(); + assert!( + err.to_string().contains("Schema mismatch"), + "unexpected error: {err}" + ); + + write_replacement_file( + &dataset, + "nested_unknown.lance", + vec![ + ArrowField::new("x", DataType::Int32, true), + ArrowField::new("y", DataType::Int32, true), + ], + ) + .await; + let err = dataset + .create_data_file("nested_unknown.lance", None) + .await + .unwrap_err(); + assert!( + err.to_string().contains("Schema mismatch"), + "unexpected error: {err}" + ); +} + #[tokio::test] async fn test_shallow_clone_base_artifacts_use_base_object_store() { let source_dir = tempfile::tempdir().unwrap(); @@ -692,12 +864,82 @@ async fn test_load_manifest_iops() { .await .unwrap(); - // There should be only two IOPS: - // 1. List _versions directory to get the latest manifest location - // 2. Read the manifest file. (The manifest is small enough to be read in one go. - // Larger manifests would result in more IOPS.) + // The write above committed on this same Session, so the manifest is already + // in the metadata cache. Opening therefore issues a single IOP: + // 1. List _versions directory to resolve the latest manifest location. + // The manifest body is served from the cache instead of being read from storage. let io_stats = _dataset.object_store.as_ref().io_stats_incremental(); - assert_io_eq!(io_stats, read_iops, 2); + assert_io_eq!(io_stats, read_iops, 1); +} + +#[tokio::test] +async fn test_checkout_removed_version_not_served_from_cache() { + let test_uri = TempStrDir::default(); + let session = Arc::new(Session::default()); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..10_i32))], + ) + .unwrap(); + let dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + &test_uri, + Some(WriteParams { + session: Some(session.clone()), + ..Default::default() + }), + ) + .await + .unwrap(); + + let version = dataset.manifest().version; + let location = dataset.manifest_location().clone(); + let cache = session.metadata_cache.for_dataset(&dataset.uri); + + assert!( + cache + .get_with_key(&ManifestKey { + version, + e_tag: location.e_tag.as_deref(), + }) + .await + .is_some(), + "manifest should be cached after the write" + ); + dataset.checkout_version(version).await.unwrap(); + + // Remove the version from storage, as cleanup (or a manual delete) would. + dataset.object_store.delete(&location.path).await.unwrap(); + + let resolved = dataset + .commit_handler + .resolve_version_location(&dataset.base, version, &dataset.object_store.inner) + .await + .unwrap(); + assert!( + resolved.size.is_none(), + "resolving a removed version must fall back to a size-less location, got {:?}", + resolved.size + ); + + cache + .insert_with_key( + &ManifestKey { + version, + e_tag: None, + }, + Arc::new(dataset.manifest().clone()), + ) + .await; + assert!( + dataset.checkout_version(version).await.is_err(), + "checkout of a version removed from storage must not be served from cache" + ); } #[rstest] @@ -890,6 +1132,164 @@ async fn test_write_manifest( assert!(matches!(write_result, Err(Error::NotSupported { .. }))); } +#[tokio::test] +async fn test_rle_v2_v23_write_and_append() { + let test_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![7; 1000]))], + ) + .unwrap(); + + let batches = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema.clone()); + let mut dataset = Dataset::write( + batches, + &test_uri, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_3), + ..Default::default() + }), + ) + .await + .unwrap(); + + let manifest = read_manifest( + dataset.object_store.as_ref(), + &dataset + .commit_handler + .resolve_latest_location(&dataset.base, dataset.object_store.as_ref()) + .await + .unwrap() + .path, + None, + ) + .await + .unwrap(); + assert_eq!( + manifest.data_storage_format.lance_file_version().unwrap(), + LanceFileVersion::V2_3 + ); + + let append_batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![9; 1000]))], + ) + .unwrap(); + let append_batches = + RecordBatchIterator::new(vec![Ok(append_batch)].into_iter(), schema.clone()); + dataset = Dataset::write( + append_batches, + &test_uri, + Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await + .unwrap(); + + assert_eq!( + dataset + .manifest + .data_storage_format + .lance_file_version() + .unwrap(), + LanceFileVersion::V2_3 + ); + + let actual = dataset.scan().try_into_batch().await.unwrap(); + let expected = RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from( + [vec![7; 1000], vec![9; 1000]].concat(), + ))], + ) + .unwrap(); + assert_eq!(actual, expected); +} + +#[tokio::test] +async fn test_rle_v2_uncommitted_create_commits_v23_storage() { + let test_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![7; 1000]))], + ) + .unwrap(); + + let transaction = InsertBuilder::new(test_uri.as_str()) + .with_params(&WriteParams { + data_storage_version: Some(LanceFileVersion::V2_3), + ..Default::default() + }) + .execute_uncommitted(vec![batch]) + .await + .unwrap(); + + let dataset = CommitBuilder::new(test_uri.as_str()) + .execute(transaction) + .await + .unwrap(); + assert_eq!( + dataset + .manifest + .data_storage_format + .lance_file_version() + .unwrap(), + LanceFileVersion::V2_3 + ); +} + +#[tokio::test] +async fn test_rle_v2_shallow_clone_preserves_v23_storage() { + let test_uri = TempStrDir::default(); + let clone_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![7; 1000]))], + ) + .unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema), + &test_uri, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_3), + ..Default::default() + }), + ) + .await + .unwrap(); + + let clone = dataset + .shallow_clone(clone_uri.as_str(), dataset.version().version, None) + .await + .unwrap(); + assert_eq!( + clone + .manifest + .data_storage_format + .lance_file_version() + .unwrap(), + LanceFileVersion::V2_3 + ); +} + #[rstest] #[tokio::test] async fn append_dataset( diff --git a/rust/lance/src/dataset/tests/dataset_merge_update.rs b/rust/lance/src/dataset/tests/dataset_merge_update.rs index 73804d7b4c0..4b81f62dc6c 100644 --- a/rust/lance/src/dataset/tests/dataset_merge_update.rs +++ b/rust/lance/src/dataset/tests/dataset_merge_update.rs @@ -1,7 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use std::collections::HashSet; use std::sync::Arc; +use std::time::Duration; use std::vec; use crate::dataset::ROW_ID; @@ -36,7 +38,7 @@ use lance_datagen::{BatchCount, RowCount, array, gen_batch}; use lance_file::version::LanceFileVersion; use lance_file::writer::FileWriter; use lance_io::utils::CachedFileSize; -use lance_table::format::DataFile; +use lance_table::format::{BasePath, DataFile, Fragment}; use crate::dataset::write::merge_insert::{WhenMatched, WhenNotMatched}; use futures::TryStreamExt; @@ -2240,6 +2242,184 @@ async fn test_data_replacement_invalidates_index_bitmap() { ); } +/// Run a predicate over `col` twice -- once index-served, once via a forced flat scan +/// (`use_scalar_index(false)`) -- assert the two agree, and return the matching `col` +/// values sorted. Equality is the index-consistency invariant: a divergence means the +/// index served rows that disagree with the underlying data. +async fn index_consistent_values(dataset: &Dataset, col: &str, predicate: &str) -> Vec { + let sorted = |batch: &RecordBatch| -> Vec { + let mut v: Vec = batch + .column_by_name(col) + .unwrap() + .as_primitive::() + .iter() + .flatten() + .collect(); + v.sort(); + v + }; + + let indexed = dataset + .scan() + .filter(predicate) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let flat = dataset + .scan() + .use_scalar_index(false) + .filter(predicate) + .unwrap() + .try_into_batch() + .await + .unwrap(); + + let indexed_vals = sorted(&indexed); + let flat_vals = sorted(&flat); + assert_eq!( + indexed_vals, flat_vals, + "index-served `{predicate}` disagrees with a flat scan" + ); + indexed_vals +} + +/// Build a Merge overlay fragment that rewrites a single `field_id` in place: tombstone +/// (-2) the field in `prev`'s existing data files and back it with `new_file` (a new +/// single-column file) instead. A file left with no live field is dropped. This is the +/// manifest shape an in-place column rewrite produces when it falls back from a +/// DataReplacement to a Merge. +fn build_overlay_frag(prev: &Fragment, field_id: i32, new_file: &str) -> Fragment { + let mut overlay = prev.clone(); + overlay.files = prev + .files + .iter() + .filter_map(|df| { + let masked: Vec = df + .fields + .iter() + .map(|&f| if f == field_id { -2 } else { f }) + .collect(); + if masked.iter().all(|&f| f == -2) { + return None; // file holds only the tombstoned field + } + let mut m = df.clone(); + m.fields = masked.into(); + Some(m) + }) + .collect(); + overlay.add_file( + new_file, + vec![field_id], + vec![0], + &LanceFileVersion::default(), + None, + ); + overlay +} + +/// A `Merge` that rewrites an indexed column's data in place must keep that column's +/// index consistent: a query the index serves must return the same rows as a flat scan +/// of the rewritten data. The overlay fragment tombstones (-2) the column's field id in +/// the existing data file and appends a new file for it, so the field stays in the +/// schema and its index is retained -- the rewritten fragment must be pruned from that +/// index. This is the shape produced when an in-place column rewrite cannot be expressed +/// as a DataReplacement (e.g. an `update` has merged the fragment's column files) and +/// falls back to a Merge overlay. +#[tokio::test] +async fn test_merge_rewriting_indexed_column_keeps_index_consistent() { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(Int32Array::from(vec![10, 20, 30, 40])), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let mut dataset = Dataset::write(reader, "memory://merge_index_rewrite", None) + .await + .unwrap(); + + dataset + .create_index( + &["a"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + + // Baseline: the index serves correct results before the rewrite. + assert_eq!( + index_consistent_values(&dataset, "a", "a >= 3").await, + vec![3, 4] + ); + + let a_field_id = 0i32; + + // Write a new single-column file holding `a`'s replacement values. + let a_only = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "a", + DataType::Int32, + true, + )])); + let new_a = RecordBatch::try_new( + a_only.clone(), + vec![Arc::new(Int32Array::from(vec![91, 92, 93, 94]))], + ) + .unwrap(); + let new_a_path = dataset.data_dir().join("merge_new_a.lance"); + let object_writer = dataset.object_store.create(&new_a_path).await.unwrap(); + let mut writer = FileWriter::try_new( + object_writer, + a_only.as_ref().try_into().unwrap(), + Default::default(), + ) + .unwrap(); + writer.write_batch(&new_a).await.unwrap(); + writer.finish().await.unwrap(); + + // Overlay that file onto fragment 0, rewriting `a` in place. + let prev = dataset.get_fragment(0).unwrap().metadata().clone(); + let overlay = build_overlay_frag(&prev, a_field_id, "merge_new_a.lance"); + + let read_version = dataset.version().version; + let dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::Merge { + fragments: vec![overlay], + schema: schema.as_ref().try_into().unwrap(), + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + + // `a` now holds [91, 92, 93, 94]; an index-served query must reflect that. + assert_eq!( + index_consistent_values(&dataset, "a", "a >= 90").await, + vec![91, 92, 93, 94], + "index-served `a >= 90` must return the rewritten values" + ); + assert!( + index_consistent_values(&dataset, "a", "a < 90") + .await + .is_empty(), + "no row satisfies `a < 90` after the rewrite" + ); +} + /// DataReplacement on an indexed column should remove the fragment from /// fragment_bitmap AND add it to invalidated_fragment_bitmap so that /// stale index entries are blocked at query time. @@ -3227,3 +3407,890 @@ async fn test_fts_unfiltered_after_compaction_returns_remapped_row_ids() { assert!(live.contains(id), "stale row_id {id}"); } } + +// --------------------------------------------------------------------------- +// Multi-base tests: merge insert on datasets whose data lives across multiple +// registered base paths, with and without routing new fragments to bases. +// --------------------------------------------------------------------------- + +/// Fixture: primary storage plus two external bases. base1 is a dataset-root +/// style base (files under `{base1}/data/`), base2 is a plain data directory +/// (files directly under `{base2}/`). Initial data: ids 0..6 in two fragments +/// in base1, ids 6..9 in one fragment in primary storage. +struct MultiBaseFixture { + _tmp: TempDir, + dataset: Dataset, + base1_dir: std::path::PathBuf, + base2_dir: std::path::PathBuf, +} + +fn multi_base_schema() -> Arc { + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("a", DataType::Int32, false), + ArrowField::new("b", DataType::Utf8, true), + ])) +} + +fn multi_base_batch(ids: &[i32], a_offset: i32, b_prefix: &str) -> RecordBatch { + RecordBatch::try_new( + multi_base_schema(), + vec![ + Arc::new(Int32Array::from(ids.to_vec())), + Arc::new(Int32Array::from( + ids.iter().map(|id| id + a_offset).collect::>(), + )), + Arc::new(StringArray::from( + ids.iter() + .map(|id| format!("{}{}", b_prefix, id)) + .collect::>(), + )), + ], + ) + .unwrap() +} + +async fn multi_base_fixture(indexed: bool) -> MultiBaseFixture { + let tmp = TempDir::default(); + let primary_dir = tmp.std_path().join("primary"); + let base1_dir = tmp.std_path().join("base1"); + let base2_dir = tmp.std_path().join("base2"); + std::fs::create_dir_all(&base1_dir).unwrap(); + std::fs::create_dir_all(&base2_dir).unwrap(); + let primary_uri = format!("file://{}", primary_dir.display()); + + let reader = RecordBatchIterator::new( + vec![Ok(multi_base_batch(&[0, 1, 2, 3, 4, 5], 100, "orig"))], + multi_base_schema(), + ); + let dataset = Dataset::write( + reader, + &primary_uri, + Some(WriteParams { + mode: WriteMode::Create, + max_rows_per_file: 3, + initial_bases: Some(vec![ + BasePath { + id: 1, + name: Some("base1".to_string()), + is_dataset_root: true, + path: format!("file://{}", base1_dir.display()), + }, + BasePath { + id: 2, + name: Some("base2".to_string()), + is_dataset_root: false, + path: format!("file://{}", base2_dir.display()), + }, + ]), + target_bases: Some(vec![1]), + ..Default::default() + }), + ) + .await + .unwrap(); + + let reader = RecordBatchIterator::new( + vec![Ok(multi_base_batch(&[6, 7, 8], 100, "orig"))], + multi_base_schema(), + ); + let mut dataset = Dataset::write( + reader, + Arc::new(dataset), + Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await + .unwrap(); + + if indexed { + dataset + .create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + } + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 3); + for fragment in &fragments[..2] { + for file in &fragment.metadata.files { + assert_eq!(file.base_id, Some(1)); + } + } + for file in &fragments[2].metadata.files { + assert_eq!(file.base_id, None); + } + + MultiBaseFixture { + _tmp: tmp, + dataset, + base1_dir, + base2_dir, + } +} + +/// Collect (id, a, b) rows sorted by id. +async fn collect_multi_base_rows(dataset: &Dataset) -> Vec<(i32, i32, Option)> { + let mut scan = dataset.scan(); + scan.project(&["id", "a", "b"]).unwrap(); + let batches = scan + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut rows = vec![]; + for batch in batches { + let ids = batch.column(0).as_primitive::(); + let a = batch.column(1).as_primitive::(); + let b = batch.column(2).as_string::(); + for i in 0..batch.num_rows() { + let b_val = if b.is_null(i) { + None + } else { + Some(b.value(i).to_string()) + }; + rows.push((ids.value(i), a.value(i), b_val)); + } + } + rows.sort_unstable(); + rows +} + +fn expected_row(id: i32, a_offset: i32, b_prefix: &str) -> (i32, i32, Option) { + (id, id + a_offset, Some(format!("{}{}", b_prefix, id))) +} + +/// Merge insert against a multi-base table without routing: every path must +/// read fragments from external bases correctly and write all new files to +/// primary storage with no base id. `indexed` toggles the v2 plan path +/// (false) vs the legacy indexed-scan path (true). +#[rstest] +#[tokio::test] +async fn test_merge_insert_on_multi_base_table(#[values(false, true)] indexed: bool) { + let fixture = multi_base_fixture(indexed).await; + let dataset = Arc::new(fixture.dataset); + + // Update one row in each existing fragment (two in base1, one in + // primary), insert two new rows. + let source = multi_base_batch(&[1, 4, 7, 10, 11], 1000, "new"); + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(source)], + multi_base_schema(), + )); + let (dataset, stats) = job.execute(reader_to_stream(reader)).await.unwrap(); + + assert_eq!(stats.num_updated_rows, 3); + assert_eq!(stats.num_inserted_rows, 2); + assert_eq!(dataset.count_rows(None).await.unwrap(), 11); + + for fragment in dataset.get_fragments() { + let metadata = &fragment.metadata; + if metadata.id >= 3 { + // Fragments written by the merge live in primary storage. + for file in &metadata.files { + assert_eq!(file.base_id, None); + } + } else { + // Pre-existing fragments keep their base and get local deletion + // files for the rewritten rows. + let expected_base = if metadata.id < 2 { Some(1) } else { None }; + for file in &metadata.files { + assert_eq!(file.base_id, expected_base); + } + let deletion = metadata.deletion_file.as_ref().unwrap(); + assert_eq!(deletion.base_id, None); + } + } + + let mut expected = vec![]; + for id in [0, 2, 3, 5, 6, 8] { + expected.push(expected_row(id, 100, "orig")); + } + for id in [1, 4, 7, 10, 11] { + expected.push(expected_row(id, 1000, "new")); + } + expected.sort_unstable(); + assert_eq!(collect_multi_base_rows(&dataset).await, expected); + + // Re-open from scratch to make sure the result is readable without any + // cached state. + let dataset = Dataset::open(dataset.uri()).await.unwrap(); + assert_eq!(collect_multi_base_rows(&dataset).await, expected); +} + +/// Merge insert routing new fragments to target bases, by id and by name, +/// covering both base layouts (dataset-root and plain data directory). +#[rstest] +#[tokio::test] +async fn test_merge_insert_route_to_target_bases(#[values(false, true)] indexed: bool) { + let fixture = multi_base_fixture(indexed).await; + let dataset = Arc::new(fixture.dataset); + + let source = multi_base_batch(&[1, 4, 10, 11], 1000, "new"); + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_bases(vec![2]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(source)], + multi_base_schema(), + )); + let (dataset, stats) = job.execute(reader_to_stream(reader)).await.unwrap(); + assert_eq!(stats.num_updated_rows, 2); + assert_eq!(stats.num_inserted_rows, 2); + + // New fragments land in base2, which is a plain data directory, so the + // files sit directly under it. + let mut merge_files = 0; + for fragment in dataset.get_fragments() { + let metadata = &fragment.metadata; + if metadata.id >= 3 { + for file in &metadata.files { + assert_eq!(file.base_id, Some(2)); + let on_disk = fixture.base2_dir.join(file.path.as_str()); + assert!(on_disk.exists(), "missing data file {:?}", on_disk); + merge_files += 1; + } + } + } + assert!(merge_files > 0); + + let max_fragment_id = dataset.manifest.max_fragment_id().unwrap(); + + // A second merge referencing a base by name: base1 is a dataset-root + // base, so files go under `{base1}/data/`. + let source = multi_base_batch(&[12, 13], 1000, "new"); + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_base_names_or_paths(vec!["base1".to_string()]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(source)], + multi_base_schema(), + )); + let (dataset, stats) = job.execute(reader_to_stream(reader)).await.unwrap(); + assert_eq!(stats.num_inserted_rows, 2); + + let mut merge_files = 0; + for fragment in dataset.get_fragments() { + let metadata = &fragment.metadata; + if metadata.id > max_fragment_id { + for file in &metadata.files { + assert_eq!(file.base_id, Some(1)); + let on_disk = fixture.base1_dir.join("data").join(file.path.as_str()); + assert!(on_disk.exists(), "missing data file {:?}", on_disk); + merge_files += 1; + } + } + } + assert!(merge_files > 0); + + let mut expected = vec![]; + for id in [0, 2, 3, 5, 6, 7, 8] { + expected.push(expected_row(id, 100, "orig")); + } + for id in [1, 4, 10, 11, 12, 13] { + expected.push(expected_row(id, 1000, "new")); + } + expected.sort_unstable(); + assert_eq!(collect_multi_base_rows(&dataset).await, expected); + + let dataset = Dataset::open(dataset.uri()).await.unwrap(); + assert_eq!(collect_multi_base_rows(&dataset).await, expected); +} + +/// Partial-schema merge insert on a multi-base table: column patches for +/// existing fragments stay in primary storage (mixing with data files in +/// external bases within the same fragment) while inserted rows route to the +/// requested base. Requires an index on the join key to reach the in-place +/// update path. +#[tokio::test] +async fn test_merge_insert_partial_schema_multi_base() { + let fixture = multi_base_fixture(true).await; + let dataset = Arc::new(fixture.dataset); + + // Update column `a` for all rows of fragment 0 (full column rewrite) and + // one row of fragment 1 (incremental update), insert ids 10 and 11. + let partial_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("a", DataType::Int32, false), + ])); + let ids = vec![0, 1, 2, 3, 10, 11]; + let source = RecordBatch::try_new( + partial_schema.clone(), + vec![ + Arc::new(Int32Array::from(ids.clone())), + Arc::new(Int32Array::from( + ids.iter().map(|id| id + 1000).collect::>(), + )), + ], + ) + .unwrap(); + + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_bases(vec![2]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new(vec![Ok(source)], partial_schema)); + let (dataset, stats) = job.execute(reader_to_stream(reader)).await.unwrap(); + assert_eq!(stats.num_updated_rows, 4); + assert_eq!(stats.num_inserted_rows, 2); + + for fragment in dataset.get_fragments() { + let metadata = &fragment.metadata; + match metadata.id { + 0 | 1 => { + // Patched fragments: the original file in base1 plus a column + // patch written to primary storage. + assert_eq!(metadata.files.len(), 2); + assert_eq!(metadata.files[0].base_id, Some(1)); + assert_eq!(metadata.files[1].base_id, None); + } + 2 => { + assert_eq!(metadata.files.len(), 1); + assert_eq!(metadata.files[0].base_id, None); + } + _ => { + // Inserted rows route to base2. + for file in &metadata.files { + assert_eq!(file.base_id, Some(2)); + let on_disk = fixture.base2_dir.join(file.path.as_str()); + assert!(on_disk.exists(), "missing data file {:?}", on_disk); + } + } + } + } + + // Updated rows keep their `b` values, inserted rows have no `b`. + let mut expected = vec![]; + for id in [4, 5, 6, 7, 8] { + expected.push(expected_row(id, 100, "orig")); + } + for id in [0, 1, 2, 3] { + expected.push((id, id + 1000, Some(format!("orig{}", id)))); + } + for id in [10, 11] { + expected.push((id, id + 1000, None)); + } + expected.sort_unstable(); + assert_eq!(collect_multi_base_rows(&dataset).await, expected); + + let dataset = Dataset::open(dataset.uri()).await.unwrap(); + assert_eq!(collect_multi_base_rows(&dataset).await, expected); +} + +/// Round-robin distribution across multiple target bases within a single +/// merge insert. New data files are cut at `max_rows_per_file` (the write +/// default of 1Mi rows), so inserting more rows than that produces multiple +/// files, which must alternate between the target bases. +#[tokio::test] +async fn test_merge_insert_round_robin_target_bases() { + let tmp = TempDir::default(); + let primary_dir = tmp.std_path().join("primary"); + let base1_dir = tmp.std_path().join("base1"); + let base2_dir = tmp.std_path().join("base2"); + std::fs::create_dir_all(&base1_dir).unwrap(); + std::fs::create_dir_all(&base2_dir).unwrap(); + let primary_uri = format!("file://{}", primary_dir.display()); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + let reader = RecordBatchIterator::new( + vec![Ok(RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..10))], + ) + .unwrap())], + schema.clone(), + ); + let dataset = Dataset::write( + reader, + &primary_uri, + Some(WriteParams { + mode: WriteMode::Create, + initial_bases: Some(vec![ + BasePath { + id: 1, + name: Some("base1".to_string()), + is_dataset_root: true, + path: format!("file://{}", base1_dir.display()), + }, + BasePath { + id: 2, + name: Some("base2".to_string()), + is_dataset_root: false, + path: format!("file://{}", base2_dir.display()), + }, + ]), + ..Default::default() + }), + ) + .await + .unwrap(); + + // 1.2Mi new rows -> two data files -> one per base. + const BATCH_ROWS: i32 = 100_000; + let batches: Vec<_> = (0..12) + .map(|i| { + let start = 1000 + i * BATCH_ROWS; + Ok(RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values( + start..start + BATCH_ROWS, + ))], + ) + .unwrap()) + }) + .collect(); + let job = MergeInsertBuilder::try_new(Arc::new(dataset), vec!["id".to_string()]) + .unwrap() + .target_bases(vec![1, 2]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new(batches, schema.clone())); + let (dataset, stats) = job.execute(reader_to_stream(reader)).await.unwrap(); + assert_eq!(stats.num_inserted_rows, 1_200_000); + assert_eq!(dataset.count_rows(None).await.unwrap(), 1_200_010); + + let merge_file_bases: Vec<_> = dataset + .get_fragments() + .iter() + .filter(|fragment| fragment.metadata.id >= 1) + .flat_map(|fragment| fragment.metadata.files.iter().map(|file| file.base_id)) + .collect(); + assert_eq!(merge_file_bases, vec![Some(1), Some(2)]); +} + +/// Target base validation across build and execution paths. +#[tokio::test] +async fn test_merge_insert_target_bases_validation() { + let fixture = multi_base_fixture(false).await; + let dataset = Arc::new(fixture.dataset); + + // Both selectors set fails at build time. + let err = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_bases(vec![1]) + .target_base_names_or_paths(vec!["base2".to_string()]) + .try_build() + .err() + .unwrap(); + assert!( + err.to_string() + .contains("Cannot specify both target_base_names_or_paths and target_bases"), + "unexpected error: {}", + err + ); + + // Unknown base id fails at execution. + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_bases(vec![99]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(multi_base_batch(&[1], 1000, "new"))], + multi_base_schema(), + )); + let err = job.execute(reader_to_stream(reader)).await.unwrap_err(); + assert!( + err.to_string() + .contains("Target base ID 99 not found in available bases"), + "unexpected error: {}", + err + ); + + // An empty target base list is rejected rather than silently ignored. + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_bases(vec![]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(multi_base_batch(&[1], 1000, "new"))], + multi_base_schema(), + )); + let err = job.execute(reader_to_stream(reader)).await.unwrap_err(); + assert!( + err.to_string().contains("target_bases cannot be empty"), + "unexpected error: {}", + err + ); + + // Unknown base name fails at execution. + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_base_names_or_paths(vec!["nonexistent".to_string()]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(multi_base_batch(&[1], 1000, "new"))], + multi_base_schema(), + )); + let err = job.execute(reader_to_stream(reader)).await.unwrap_err(); + assert!( + err.to_string() + .contains("Base reference 'nonexistent' not found in available bases"), + "unexpected error: {}", + err + ); + + // Delete-only merges write no data files but still validate target bases. + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::Delete) + .when_not_matched(WhenNotMatched::DoNothing) + .target_bases(vec![99]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(multi_base_batch(&[1], 1000, "new"))], + multi_base_schema(), + )); + let err = job.execute(reader_to_stream(reader)).await.unwrap_err(); + assert!( + err.to_string() + .contains("Target base ID 99 not found in available bases"), + "unexpected error: {}", + err + ); + + // Valid target bases on a delete-only merge are a no-op. + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::Delete) + .when_not_matched(WhenNotMatched::DoNothing) + .target_bases(vec![2]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(multi_base_batch(&[5], 1000, "new"))], + multi_base_schema(), + )); + let (dataset, stats) = job.execute(reader_to_stream(reader)).await.unwrap(); + assert_eq!(stats.num_deleted_rows, 1); + assert_eq!(dataset.count_rows(None).await.unwrap(), 8); + + // Datasets with no registered bases reject target bases. + let plain_dir = TempStrDir::default(); + let reader = RecordBatchIterator::new( + vec![Ok(multi_base_batch(&[0, 1], 100, "orig"))], + multi_base_schema(), + ); + let plain_dataset = Dataset::write(reader, plain_dir.as_str(), None) + .await + .unwrap(); + let job = MergeInsertBuilder::try_new(Arc::new(plain_dataset), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_bases(vec![1]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(multi_base_batch(&[1], 1000, "new"))], + multi_base_schema(), + )); + let err = job.execute(reader_to_stream(reader)).await.unwrap_err(); + assert!( + err.to_string() + .contains("Target base ID 1 not found in available bases"), + "unexpected error: {}", + err + ); +} + +/// Base id 0 and the dataset URI include primary storage in the merge insert +/// target rotation. +#[tokio::test] +async fn test_merge_insert_target_bases_include_primary() { + let fixture = multi_base_fixture(false).await; + let dataset = Arc::new(fixture.dataset); + let primary_uri = dataset.uri().to_string(); + + // Single new file: the first slot (primary) receives it. + let source = multi_base_batch(&[1, 10], 1000, "new"); + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_bases(vec![0, 2]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(source)], + multi_base_schema(), + )); + let (dataset, _) = job.execute(reader_to_stream(reader)).await.unwrap(); + let new_files: Vec<_> = dataset + .get_fragments() + .iter() + .filter(|f| f.metadata.id >= 3) + .flat_map(|f| f.metadata.files.iter().map(|file| file.base_id)) + .collect(); + assert_eq!(new_files, vec![None]); + + // Flipped order: the first slot is base 2. + let max_id = dataset.manifest.max_fragment_id().unwrap(); + let source = multi_base_batch(&[11], 1000, "new"); + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_bases(vec![2, 0]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(source)], + multi_base_schema(), + )); + let (dataset, _) = job.execute(reader_to_stream(reader)).await.unwrap(); + let new_files: Vec<_> = dataset + .get_fragments() + .iter() + .filter(|f| f.metadata.id > max_id) + .flat_map(|f| f.metadata.files.iter().map(|file| file.base_id)) + .collect(); + assert_eq!(new_files, vec![Some(2)]); + + // Names variant: the dataset's URI selects primary storage. + let max_id = dataset.manifest.max_fragment_id().unwrap(); + let source = multi_base_batch(&[12], 1000, "new"); + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_base_names_or_paths(vec![primary_uri]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(source)], + multi_base_schema(), + )); + let (dataset, _) = job.execute(reader_to_stream(reader)).await.unwrap(); + let new_files: Vec<_> = dataset + .get_fragments() + .iter() + .filter(|f| f.metadata.id > max_id) + .flat_map(|f| f.metadata.files.iter().map(|file| file.base_id)) + .collect(); + assert_eq!(new_files, vec![None]); + + let mut expected = vec![]; + for id in [0, 2, 3, 4, 5, 6, 7, 8] { + expected.push(expected_row(id, 100, "orig")); + } + for id in [1, 10, 11, 12] { + expected.push(expected_row(id, 1000, "new")); + } + expected.sort_unstable(); + assert_eq!(collect_multi_base_rows(&dataset).await, expected); + + let dataset = Dataset::open(dataset.uri()).await.unwrap(); + assert_eq!(collect_multi_base_rows(&dataset).await, expected); +} + +/// Merge insert attempts discarded by a retryable commit conflict must clean +/// up the data files they routed to target bases; after concurrent merges the +/// bases must contain only files referenced by the final manifest. +#[tokio::test] +async fn test_merge_insert_conflict_retry_cleans_routed_files() { + let fixture = multi_base_fixture(false).await; + let dataset = Arc::new(fixture.dataset); + let concurrency: u32 = 5; + + let barrier = Arc::new(tokio::sync::Barrier::new(concurrency as usize)); + let mut handles = Vec::new(); + for i in 0..concurrency { + // Every task starts from the same dataset version and updates the same + // row, so all but one attempt per round hit a retryable conflict. + let dataset = dataset.clone(); + let barrier = barrier.clone(); + handles.push(tokio::spawn(async move { + barrier.wait().await; + let source = multi_base_batch(&[1, 100 + i as i32], 1000 + i as i32, "new"); + let job = MergeInsertBuilder::try_new(dataset, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .conflict_retries(20) + .retry_timeout(Duration::from_secs(60)) + .target_bases(vec![1, 2]) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(source)], + multi_base_schema(), + )); + job.execute(reader_to_stream(reader)).await.unwrap() + })); + } + let mut total_attempts = 0; + for handle in handles { + let (_dataset, stats) = handle.await.unwrap(); + total_attempts += stats.num_attempts; + } + assert!( + total_attempts > concurrency, + "expected at least one conflicted attempt, got {} attempts", + total_attempts + ); + + let dataset = Dataset::open(dataset.uri()).await.unwrap(); + assert_eq!( + dataset.count_rows(None).await.unwrap(), + 9 + concurrency as usize + ); + + let mut referenced: HashSet<(Option, String)> = HashSet::new(); + for fragment in dataset.get_fragments() { + for file in &fragment.metadata.files { + referenced.insert((file.base_id, file.path.to_string())); + } + } + let list_files = |dir: &std::path::Path| -> Vec { + if !dir.exists() { + return vec![]; + } + std::fs::read_dir(dir) + .unwrap() + .filter_map(|entry| { + let path = entry.unwrap().path(); + if path.extension().is_some_and(|ext| ext == "lance") { + Some(path.file_name().unwrap().to_string_lossy().to_string()) + } else { + None + } + }) + .collect() + }; + for name in list_files(&fixture.base1_dir.join("data")) { + assert!( + referenced.contains(&(Some(1), name.clone())), + "orphaned file in base1: {}", + name + ); + } + for name in list_files(&fixture.base2_dir) { + assert!( + referenced.contains(&(Some(2), name.clone())), + "orphaned file in base2: {}", + name + ); + } +} + +/// `target_all_bases` on merge insert resolves to every registered base at +/// execution time, with primary storage first when included. +#[tokio::test] +async fn test_merge_insert_target_all_bases() { + let fixture = multi_base_fixture(false).await; + let dataset = Arc::new(fixture.dataset); + + // Single new file: with primary included it takes the first slot. + let source = multi_base_batch(&[20], 1000, "new"); + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_all_bases(true) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(source)], + multi_base_schema(), + )); + let (dataset, _) = job.execute(reader_to_stream(reader)).await.unwrap(); + let new_files: Vec<_> = dataset + .get_fragments() + .iter() + .filter(|f| f.metadata.id >= 3) + .flat_map(|f| f.metadata.files.iter().map(|file| file.base_id)) + .collect(); + assert_eq!(new_files, vec![None]); + + // Without primary the first slot is the lowest registered base id. + let max_id = dataset.manifest.max_fragment_id().unwrap(); + let source = multi_base_batch(&[21], 1000, "new"); + let job = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_all_bases(false) + .try_build() + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + vec![Ok(source)], + multi_base_schema(), + )); + let (dataset, _) = job.execute(reader_to_stream(reader)).await.unwrap(); + let new_files: Vec<_> = dataset + .get_fragments() + .iter() + .filter(|f| f.metadata.id > max_id) + .flat_map(|f| f.metadata.files.iter().map(|file| file.base_id)) + .collect(); + assert_eq!(new_files, vec![Some(1)]); + + // Cannot be combined with explicit target bases. + let err = MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .target_bases(vec![1]) + .target_all_bases(true) + .try_build() + .err() + .unwrap(); + assert!( + err.to_string() + .contains("Cannot specify target_all_bases together with"), + "unexpected error: {}", + err + ); + + let expected_new: Vec<_> = [20, 21] + .iter() + .map(|id| expected_row(*id, 1000, "new")) + .collect(); + let all_rows = collect_multi_base_rows(&dataset).await; + for row in expected_new { + assert!(all_rows.contains(&row), "missing row {:?}", row); + } +} diff --git a/rust/lance/src/dataset/tests/dataset_scanner.rs b/rust/lance/src/dataset/tests/dataset_scanner.rs index 4c44cb0795b..363698d9966 100644 --- a/rust/lance/src/dataset/tests/dataset_scanner.rs +++ b/rust/lance/src/dataset/tests/dataset_scanner.rs @@ -42,6 +42,82 @@ use lance_index::vector::pq::PQBuildParams; use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, Query}; use pretty_assertions::assert_eq; +#[tokio::test] +async fn test_scan_wide_fixed_size_list_at_batch_boundary() { + const DIM_A: usize = 140_000; + const DIM_B: usize = 4_096; + const SHORT_ROWS: usize = 68; + const LONG_ROWS: usize = 128; + + fn make_batch(schema: SchemaRef, rows: usize, base: usize) -> RecordBatch { + let values_a = Float32Array::from_iter_values( + (0..rows * DIM_A).map(|idx| ((idx + base) % 1009) as f32 / 1009.0), + ); + let values_b = Float32Array::from_iter_values( + (0..rows * DIM_B).map(|idx| ((idx + base) % 251) as f32 / 251.0), + ); + let arr_a = FixedSizeListArray::try_new_from_values(values_a, DIM_A as i32).unwrap(); + let arr_b = FixedSizeListArray::try_new_from_values(values_b, DIM_B as i32).unwrap(); + RecordBatch::try_new(schema, vec![Arc::new(arr_a), Arc::new(arr_b)]).unwrap() + } + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "a", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + DIM_A as i32, + ), + true, + ), + ArrowField::new( + "b", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + DIM_B as i32, + ), + true, + ), + ])); + + let batches = vec![ + make_batch(schema.clone(), SHORT_ROWS, 0), + make_batch(schema.clone(), LONG_ROWS, 17), + ]; + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let write_params = WriteParams { + data_storage_version: Some(LanceFileVersion::V2_1), + ..WriteParams::default() + }; + let dir = tempfile::tempdir().unwrap(); + let dataset = Dataset::write(reader, dir.path().to_str().unwrap(), Some(write_params)) + .await + .unwrap(); + + // The first column splits into 9 read chunks. The second column is a + // higher-priority request that can reserve the remaining buffer while the + // first column is still awaited. + let mut scanner = dataset.scan(); + scanner.io_buffer_size(70 * 1024 * 1024); + scanner + .limit(Some(LONG_ROWS as i64), Some(SHORT_ROWS as i64)) + .unwrap(); + let mut stream = tokio::time::timeout( + std::time::Duration::from_secs(20), + scanner.try_into_stream(), + ) + .await + .expect("stream creation timed out") + .unwrap(); + let batch = tokio::time::timeout(std::time::Duration::from_secs(20), stream.try_next()) + .await + .expect("first batch timed out") + .unwrap() + .unwrap(); + + assert_eq!(batch.num_rows(), LONG_ROWS); +} + #[tokio::test] async fn test_vector_filter_fts_search() { let dataset = prepare_query_filter_dataset().await; diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs index 4555cd7ee6c..372adee27db 100644 --- a/rust/lance/src/dataset/transaction.rs +++ b/rust/lance/src/dataset/transaction.rs @@ -32,7 +32,8 @@ use lance_table::rowids::read_row_ids; use lance_table::{ format::{ BasePath, DataFile, DataStorageFormat, Fragment, IndexFile, IndexMetadata, Manifest, - RowDatasetVersionMeta, RowDatasetVersionRun, RowDatasetVersionSequence, RowIdMeta, pb, + RowDatasetVersionMeta, RowDatasetVersionRun, RowDatasetVersionSequence, RowIdMeta, + overlay::DataOverlayFile, pb, }, io::{ commit::CommitHandler, @@ -258,6 +259,17 @@ pub struct Transaction { #[derive(Debug, Clone, DeepSizeOf, PartialEq)] pub struct DataReplacementGroup(pub u64, pub DataFile); +/// Overlay files to append to a single fragment, in order (the last entry is +/// newest). The overlays are appended to the fragment's existing `overlays` +/// list rather than replacing it, so overlays written by concurrent commits are +/// preserved. Each overlay's `committed_version` is stamped to the new dataset +/// version at commit time (re-stamped on retry). +#[derive(Debug, Clone, DeepSizeOf, PartialEq)] +pub struct DataOverlayGroup { + pub fragment_id: u64, + pub overlays: Vec, +} + /// An entry for a map update. If value is None, the key will be removed from the map. #[derive(Debug, Clone, DeepSizeOf, PartialEq)] pub struct UpdateMapEntry { @@ -367,6 +379,11 @@ pub enum Operation { DataReplacement { replacements: Vec, }, + /// Attach overlay files to fragments, supplying new values for a subset of + /// `(physical offset, field)` cells without rewriting the fragments' base + /// data files. See [`DataOverlayFile`] and the Data Overlay Files + /// specification for resolution, coverage, and versioning rules. + DataOverlay { groups: Vec }, /// Merge a new column in /// 'fragments' is the final fragments include all data files, the new fragments must align with old ones at rows. /// 'schema' is not forced to include existed columns, which means we could use Merge to drop column data @@ -499,6 +516,7 @@ impl std::fmt::Display for Operation { Self::Project { .. } => write!(f, "Project"), Self::UpdateConfig { .. } => write!(f, "UpdateConfig"), Self::DataReplacement { .. } => write!(f, "DataReplacement"), + Self::DataOverlay { .. } => write!(f, "DataOverlay"), Self::Clone { .. } => write!(f, "Clone"), Self::UpdateMemWalState { .. } => write!(f, "UpdateMemWalState"), Self::UpdateBases { .. } => write!(f, "UpdateBases"), @@ -1345,6 +1363,8 @@ impl PartialEq for Operation { (Self::Clone { .. }, Self::UpdateBases { .. }) => { std::mem::discriminant(self) == std::mem::discriminant(other) } + (Self::DataOverlay { groups: a }, Self::DataOverlay { groups: b }) => compare_vec(a, b), + (Self::DataOverlay { .. }, _) | (_, Self::DataOverlay { .. }) => false, } } } @@ -1521,6 +1541,7 @@ impl Operation { Self::Project { .. } => "Project", Self::UpdateConfig { .. } => "UpdateConfig", Self::DataReplacement { .. } => "DataReplacement", + Self::DataOverlay { .. } => "DataOverlay", Self::UpdateMemWalState { .. } => "UpdateMemWalState", Self::Clone { .. } => "Clone", Self::UpdateBases { .. } => "UpdateBases", @@ -1861,6 +1882,8 @@ impl Transaction { )) }); + let new_version = current_manifest.map_or(1, |m| m.version + 1); + match &self.operation { Operation::Clone { .. } => { return Err(Error::internal( @@ -1875,7 +1898,6 @@ impl Transaction { if let Some(next_row_id) = &mut next_row_id { Self::assign_row_ids(next_row_id, new_fragments.as_mut_slice())?; // Add version metadata for all new fragments - let new_version = current_manifest.map(|m| m.version + 1).unwrap_or(1); for fragment in new_fragments.iter_mut() { let version_meta = build_version_meta(fragment, new_version); fragment.last_updated_at_version_meta = version_meta.clone(); @@ -1923,7 +1945,20 @@ impl Transaction { return None; } if let Some(updated) = updated_fragments.iter().find(|uf| uf.id == f.id) { - Some(updated.clone()) + let mut updated = updated.clone(); + // Carry forward the fragment's current overlays (which + // may include ones added by a concurrent commit). An + // in-place column rewrite then tombstones the overlaid + // fields it rewrote, since the fresh base values + // supersede them. + updated.overlays = f.overlays.clone(); + if matches!(update_mode, Some(RewriteColumns)) { + lance_table::format::overlay::tombstone_overlay_fields( + &mut updated.overlays, + fields_modified, + ); + } + Some(updated) } else { Some(f.clone()) } @@ -1945,7 +1980,6 @@ impl Transaction { && let Some(UpdatedFragmentOffsets(off_map)) = updated_fragment_offsets && !off_map.is_empty() { - let new_version = current_manifest.map(|m| m.version + 1).unwrap_or(1); let prev_version = current_manifest.map(|m| m.version).unwrap_or(0); for fragment in final_fragments.iter_mut() { let Some(bitmap) = off_map.get(&fragment.id) else { @@ -1988,7 +2022,6 @@ impl Transaction { } if next_row_id.is_some() { - let new_version = current_manifest.map(|m| m.version + 1).unwrap_or(1); resolve_update_version_metadata( existing_fragments, new_fragments.as_mut_slice(), @@ -2033,7 +2066,7 @@ impl Transaction { if !merged_generations.is_empty() { update_mem_wal_index_merged_generations( &mut final_indices, - current_manifest.map_or(1, |m| m.version + 1), + new_version, merged_generations.clone(), )?; } @@ -2045,7 +2078,6 @@ impl Transaction { if let Some(next_row_id) = &mut next_row_id { Self::assign_row_ids(next_row_id, new_fragments.as_mut_slice())?; // Add version metadata for all new fragments - let new_version = current_manifest.map(|m| m.version + 1).unwrap_or(1); for fragment in new_fragments.iter_mut() { let version_meta = build_version_meta(fragment, new_version); fragment.last_updated_at_version_meta = version_meta.clone(); @@ -2111,13 +2143,11 @@ impl Transaction { final_fragments.extend(maybe_existing_fragments?.clone()); } Operation::Merge { fragments, .. } => { + let existing_fragments = maybe_existing_fragments?; let mut merged_fragments = fragments.clone(); if next_row_id.is_some() { - let new_version = current_manifest.map(|m| m.version + 1).unwrap_or(1); - let prev_by_id: HashMap = maybe_existing_fragments? - .iter() - .map(|f| (f.id, f)) - .collect(); + let prev_by_id: HashMap = + existing_fragments.iter().map(|f| (f.id, f)).collect(); for fragment in merged_fragments.iter_mut() { match prev_by_id.get(&fragment.id) { Some(prev) => { @@ -2144,6 +2174,15 @@ impl Transaction { } final_fragments.extend(merged_fragments); + // A Merge can rewrite a column's data file in place; the field stays + // in the schema, so the index is retained -- prune its now-stale + // entries for the rewritten fragments. + Self::prune_merge_rewritten_fields_from_indices( + &mut final_indices, + existing_fragments, + fragments, + ); + // Some fields that have indices may have been removed, so we should // remove those indices as well. Self::retain_relevant_indices(&mut final_indices, &schema, &final_fragments) @@ -2240,27 +2279,24 @@ impl Transaction { && file.file_major_version == new_file.file_major_version && file.file_minor_version == new_file.file_minor_version { - // assign the new file path / size to the fragment + // assign the new file path / size / base to the fragment file.path = new_file.path.clone(); file.file_size_bytes = new_file.file_size_bytes.clone(); + file.base_id = new_file.base_id; } columns_covered.extend(file.fields.iter()); } // SPECIAL CASE: if the column(s) being replaced are not covered by the fragment // Then it means it's a all-NULL column that is being replaced with real data - // just add it to the final fragments + // just add it to the final fragments. Push the DataFile as + // given so every field (including base_id) is preserved. if columns_covered.is_disjoint(&new_file.fields.iter().collect()) { - new_frag.add_file( - new_file.path.clone(), - new_file.fields.to_vec(), - new_file.column_indices.to_vec(), - &LanceFileVersion::try_from_major_minor( - new_file.file_major_version, - new_file.file_minor_version, - ) - .expect("Expected valid file version"), - new_file.file_size_bytes.get(), - ); + LanceFileVersion::try_from_major_minor( + new_file.file_major_version, + new_file.file_minor_version, + ) + .expect("Expected valid file version"); + new_frag.files.push(new_file.clone()); } // Nothing changed in the current fragment, which is not expected -- error out @@ -2269,6 +2305,15 @@ impl Transaction { "Expected to modify the fragment but no changes were made. This means the new data files does not align with any exiting datafiles. Please check if the schema of the new data files matches the schema of the old data files including the file major and minor versions", )); } + + // New base values for these fields supersede any overlay + // still shadowing them; tombstone the overlaid fields so the + // replacement is not silently masked. + lance_table::format::overlay::tombstone_overlay_fields( + &mut new_frag.overlays, + &replaced_fields, + ); + final_fragments.push(new_frag); } @@ -2300,10 +2345,57 @@ impl Transaction { &replaced_fields, ); } + Operation::DataOverlay { groups } => { + // Stamp each overlay with the version this commit is producing. + // build_manifest re-runs on every retry with an updated + // current_manifest, so this is naturally re-stamped on retry. + let new_version = current_manifest.map_or(1, |m| m.version + 1); + + let existing_fragments = maybe_existing_fragments?; + // Multiple groups may target the same fragment; merge them in + // order rather than letting a HashMap collapse drop all but the + // last group's overlays. + let mut overlays_by_fragment: HashMap> = HashMap::new(); + for group in groups { + overlays_by_fragment + .entry(group.fragment_id) + .or_default() + .extend(group.overlays.iter()); + } + + // Every group must target an existing fragment. Build a set of + // existing ids once so this is O(groups + fragments) rather than + // O(groups * fragments). + let existing_fragment_ids: HashSet = + existing_fragments.iter().map(|f| f.id).collect(); + for fragment_id in overlays_by_fragment.keys() { + if !existing_fragment_ids.contains(fragment_id) { + return Err(Error::invalid_input(format!( + "DataOverlay targets fragment {fragment_id}, which does not exist" + ))); + } + } + + for fragment in existing_fragments { + let mut fragment = fragment.clone(); + if let Some(new_overlays) = overlays_by_fragment.get(&fragment.id) { + // Appended (not replaced) so concurrently-written overlays + // survive; later entries are newer. + fragment + .overlays + .extend(new_overlays.iter().map(|&overlay| { + let mut overlay = overlay.clone(); + overlay.committed_version = new_version; + overlay + })); + } + final_fragments.push(fragment); + } + } Operation::UpdateMemWalState { merged_generations } => { update_mem_wal_index_merged_generations( &mut final_indices, - current_manifest.map_or(1, |m| m.version + 1), + new_version, merged_generations.clone(), )?; } @@ -2320,6 +2412,15 @@ impl Transaction { // Clean up data files that only contain tombstoned fields Self::remove_tombstoned_data_files(&mut final_fragments); + // Enforce the newest-last overlay ordering invariant at the write + // boundary. Load normalizes with a sort; this rejects any commit path + // that assembled a fragment's overlays out of order. + for fragment in &final_fragments { + if !fragment.overlays.is_empty() { + lance_table::format::overlay::verify_overlays_newest_last(&fragment.overlays)?; + } + } + let user_requested_version = match (&config.storage_format, config.use_legacy_format) { (Some(storage_format), _) => Some(storage_format.lance_file_version()?), (None, Some(true)) => Some(LanceFileVersion::Legacy), @@ -2623,6 +2724,60 @@ impl Transaction { } } + /// Map each (non-tombstoned) field id in a fragment to the path of the data + /// file that backs it. + fn fragment_field_paths(frag: &Fragment) -> HashMap { + let mut map = HashMap::new(); + for file in &frag.files { + for &field_id in file.fields.iter() { + if field_id >= 0 { + map.insert(field_id, file.path.as_str()); + } + } + } + map + } + + /// A `Merge` can rewrite a column's data *in place* -- the field stays in the + /// schema but its backing data file changes (the overlay fragment carries a new + /// file for the field and tombstones its old field id). `retain_relevant_indices` + /// only drops indices for *removed* fields, so without this the index keeps + /// covering the rewritten fragments with stale entries. Remove each such fragment + /// from any index covering a field whose backing data file changed. + fn prune_merge_rewritten_fields_from_indices( + indices: &mut [IndexMetadata], + prev_fragments: &[Fragment], + new_fragments: &[Fragment], + ) { + let prev_by_id: HashMap = + prev_fragments.iter().map(|f| (f.id, f)).collect(); + for new_frag in new_fragments { + let Some(prev) = prev_by_id.get(&new_frag.id) else { + continue; // brand-new fragment: nothing stale to prune + }; + let prev_paths = Self::fragment_field_paths(prev); + let new_paths = Self::fragment_field_paths(new_frag); + // Fields still present whose backing file path changed == rewritten data. + let changed: Vec = prev_paths + .iter() + .filter(|(field_id, prev_path)| { + new_paths + .get(*field_id) + .is_some_and(|new_path| new_path != *prev_path) + }) + .map(|(field_id, _)| *field_id as u32) + .collect(); + if changed.is_empty() { + continue; + } + Self::prune_updated_fields_from_indices( + indices, + std::slice::from_ref(new_frag), + &changed, + ); + } + } + fn is_vector_index(index: &IndexMetadata) -> bool { if let Some(details) = &index.index_details { details.type_url.ends_with("VectorIndexDetails") @@ -3007,6 +3162,34 @@ impl TryFrom for DataReplacementGroup { } } +impl From<&DataOverlayGroup> for pb::transaction::DataOverlayGroup { + fn from(group: &DataOverlayGroup) -> Self { + Self { + fragment_id: group.fragment_id, + overlays: group + .overlays + .iter() + .map(pb::DataOverlayFile::from) + .collect(), + } + } +} + +impl TryFrom for DataOverlayGroup { + type Error = Error; + + fn try_from(message: pb::transaction::DataOverlayGroup) -> Result { + Ok(Self { + fragment_id: message.fragment_id, + overlays: message + .overlays + .into_iter() + .map(DataOverlayFile::try_from) + .collect::>>()?, + }) + } +} + impl TryFrom for Transaction { type Error = Error; @@ -3289,6 +3472,14 @@ impl TryFrom for Transaction { })) => Operation::UpdateBases { new_bases: new_bases.into_iter().map(BasePath::from).collect(), }, + Some(pb::transaction::Operation::DataOverlay(pb::transaction::DataOverlay { + groups, + })) => Operation::DataOverlay { + groups: groups + .into_iter() + .map(DataOverlayGroup::try_from) + .collect::>>()?, + }, None => { return Err(Error::internal( "Transaction message did not contain an operation".to_string(), @@ -3559,6 +3750,14 @@ impl From<&Transaction> for pb::Transaction { .collect(), }) } + Operation::DataOverlay { groups } => { + pb::transaction::Operation::DataOverlay(pb::transaction::DataOverlay { + groups: groups + .iter() + .map(pb::transaction::DataOverlayGroup::from) + .collect(), + }) + } Operation::UpdateMemWalState { merged_generations } => { pb::transaction::Operation::UpdateMemWalState(pb::transaction::UpdateMemWalState { merged_generations: merged_generations @@ -3837,6 +4036,7 @@ mod tests { use lance_core::{ROW_ADDR, ROW_CREATED_AT_VERSION, ROW_LAST_UPDATED_AT_VERSION}; use lance_file::version::LanceFileVersion; use lance_io::utils::CachedFileSize; + use lance_table::format::overlay::OverlayCoverage; use lance_table::format::{ RowDatasetVersionMeta, RowDatasetVersionRun, RowDatasetVersionSequence, RowIdMeta, }; @@ -4148,6 +4348,7 @@ mod tests { physical_rows: Some(100), row_id_meta: None, files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4180,6 +4381,7 @@ mod tests { physical_rows: Some(50), row_id_meta: Some(RowIdMeta::Inline(serialized)), files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4212,6 +4414,7 @@ mod tests { physical_rows: Some(50), // More physical rows than existing row IDs row_id_meta: Some(RowIdMeta::Inline(serialized)), files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4247,6 +4450,7 @@ mod tests { physical_rows: Some(50), // Less physical rows than existing row IDs row_id_meta: Some(RowIdMeta::Inline(serialized)), files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4275,6 +4479,7 @@ mod tests { physical_rows: Some(30), // No existing row IDs row_id_meta: None, files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4284,6 +4489,7 @@ mod tests { physical_rows: Some(25), // Partial existing row IDs row_id_meta: Some(RowIdMeta::Inline(serialized)), files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4328,6 +4534,7 @@ mod tests { physical_rows: None, row_id_meta: None, files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4835,6 +5042,7 @@ mod tests { let fragment = Fragment { id: 1, files: vec![data_file], + overlays: vec![], deletion_file: None, row_id_meta, physical_rows: Some(5), @@ -5104,6 +5312,7 @@ mod tests { None, )], physical_rows: Some(10), + overlays: vec![], deletion_file: None, row_id_meta: None, last_updated_at_version_meta: None, @@ -5195,6 +5404,7 @@ mod tests { let prev_fragment = Fragment { id: 0, files: vec![mk_file("before.lance")], + overlays: vec![], deletion_file: None, row_id_meta, physical_rows: Some(5), @@ -5267,6 +5477,7 @@ mod tests { let prev_fragment = Fragment { id: 0, files: vec![data_file.clone()], + overlays: vec![], deletion_file: None, row_id_meta: row_id_meta.clone(), physical_rows: Some(5), @@ -5286,6 +5497,7 @@ mod tests { let merged_fragment = Fragment { id: 0, files: vec![data_file], + overlays: vec![], deletion_file: None, row_id_meta, physical_rows: Some(5), @@ -5335,6 +5547,7 @@ mod tests { let prev_fragment = Fragment { id: 0, files: vec![mk_file("before.lance")], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(5), @@ -5399,6 +5612,7 @@ mod tests { let existing_fragment = Fragment { id: 0, files: vec![mk_file("existing.lance")], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&row_ids_0))), physical_rows: Some(3), @@ -5421,6 +5635,7 @@ mod tests { let new_fragment = Fragment { id: 1, files: vec![mk_file("new.lance")], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&row_ids_1))), physical_rows: Some(4), @@ -5483,6 +5698,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(3), @@ -5496,6 +5712,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5538,6 +5755,7 @@ mod tests { Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&frag_a_seq))), physical_rows: Some(2), @@ -5549,6 +5767,7 @@ mod tests { Fragment { id: 2, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&frag_b_seq))), physical_rows: Some(3), @@ -5564,6 +5783,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5605,6 +5825,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5619,6 +5840,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5663,6 +5885,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5676,6 +5899,7 @@ mod tests { let new_fragment = Fragment { id: 20, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(4), @@ -5709,6 +5933,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5720,6 +5945,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(1), @@ -5748,6 +5974,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5758,6 +5985,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(3), @@ -5788,6 +6016,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5801,6 +6030,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(1), @@ -5842,6 +6072,7 @@ mod tests { let in_range_frag = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&in_range_seq))), physical_rows: Some(2), @@ -5862,6 +6093,7 @@ mod tests { let out_of_range_frag = Fragment { id: 2, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&out_of_range_seq))), physical_rows: Some(2), @@ -5876,6 +6108,7 @@ mod tests { let new_frag = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5914,6 +6147,7 @@ mod tests { let existing = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&seq))), physical_rows: Some(3), @@ -5926,6 +6160,7 @@ mod tests { let new_frag = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5974,6 +6209,7 @@ mod tests { let src_frag = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&src_seq))), physical_rows: Some(100), @@ -5988,6 +6224,7 @@ mod tests { let new_frag = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(100), @@ -6035,6 +6272,7 @@ mod tests { Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&seq_a))), physical_rows: Some(3), @@ -6046,6 +6284,7 @@ mod tests { Fragment { id: 2, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&seq_b))), physical_rows: Some(3), @@ -6061,6 +6300,7 @@ mod tests { let new_frag = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -6127,4 +6367,272 @@ mod tests { assert!(!left.modifies_same_metadata(&different_key)); assert!(left.modifies_same_metadata(&replace)); } + + #[test] + fn test_data_overlay_operation_roundtrips() { + // A DataOverlay operation survives the protobuf round-trip, preserving + // the target fragment, the overlay's coverage, and its committed_version. + let mut bitmap = roaring::RoaringBitmap::new(); + bitmap.insert(1); + bitmap.insert(4); + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay-0.lance", vec![3], None), + coverage: OverlayCoverage::dense(bitmap.clone()), + committed_version: 6, + }; + let pb_overlay = pb::DataOverlayFile::from(&overlay); + + let message = pb::Transaction { + read_version: 1, + uuid: Uuid::new_v4().to_string(), + operation: Some(pb::transaction::Operation::DataOverlay( + pb::transaction::DataOverlay { + groups: vec![pb::transaction::DataOverlayGroup { + fragment_id: 7, + overlays: vec![pb_overlay], + }], + }, + )), + ..Default::default() + }; + + let txn = Transaction::try_from(message).unwrap(); + match txn.operation { + Operation::DataOverlay { groups } => { + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].fragment_id, 7); + assert_eq!(groups[0].overlays.len(), 1); + assert_eq!(groups[0].overlays[0].committed_version, 6); + assert_eq!( + *groups[0].overlays[0].coverage_for_field(0).unwrap(), + bitmap + ); + } + other => panic!("expected DataOverlay, got {other:?}"), + } + } + + fn overlay_with_field(field: i32, committed_version: u64) -> DataOverlayFile { + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![field], None), + coverage: OverlayCoverage::dense(roaring::RoaringBitmap::from_iter([0u32])), + committed_version, + } + } + + #[test] + fn test_data_overlay_build_manifest_multi_fragment() { + // Overlays targeting two distinct fragments are each applied and stamped. + // A targeted fragment already carrying an overlay (committed at v3) gets + // the new overlay appended and stamped while its existing overlay is + // preserved, and a fragment the operation does not target is passed + // through with its existing overlays untouched. + let mut frag0 = Fragment::new(0); + frag0.overlays = vec![overlay_with_field(5, 3)]; // targeted, pre-existing at v3 + let frag1 = Fragment::new(1); + let mut frag2 = Fragment::new(2); + frag2.overlays = vec![overlay_with_field(9, 3)]; // untargeted, committed at v3 + let schema = ArrowSchema::new(vec![ArrowField::new("id", DataType::Int32, false)]); + let mut manifest = Manifest::new( + LanceSchema::try_from(&schema).unwrap(), + Arc::new(vec![frag0, frag1, frag2]), + lance_table::format::DataStorageFormat::new(LanceFileVersion::V2_0), + HashMap::new(), + ); + // The pre-existing overlays were committed at v3, so the current + // manifest must be at least that version; the new commit then stamps + // its overlay at v4, keeping the fragment's overlays newest-last. + manifest.version = 3; + + let txn = Transaction::new( + manifest.version, + Operation::DataOverlay { + groups: vec![ + DataOverlayGroup { + fragment_id: 0, + overlays: vec![overlay_with_field(1, 0)], + }, + DataOverlayGroup { + fragment_id: 1, + overlays: vec![overlay_with_field(2, 0)], + }, + ], + }, + None, + ); + + let (result, _) = txn + .build_manifest( + Some(&manifest), + vec![], + "txn", + &ManifestWriteConfig::default(), + ) + .unwrap(); + + let frag = |id: u64| { + result + .fragments + .iter() + .find(|f| f.id == id) + .unwrap_or_else(|| panic!("fragment {id} missing from result")) + }; + // The already-overlaid target keeps its v3 overlay and appends the new + // one, stamped to the new version. + assert_eq!(frag(0).overlays.len(), 2); + assert_eq!(frag(0).overlays[0].committed_version, 3); + assert_eq!(frag(0).overlays[1].committed_version, result.version); + // The fresh target gets its overlay, stamped to the new version. + assert_eq!(frag(1).overlays.len(), 1); + assert_eq!(frag(1).overlays[0].committed_version, result.version); + // The untargeted fragment is unchanged: same overlay, original version. + assert_eq!(frag(2).overlays.len(), 1); + assert_eq!(frag(2).overlays[0].committed_version, 3); + assert!(result.version > manifest.version); + } + + #[test] + fn test_data_replacement_tombstones_overlaid_fields() { + // A DataReplacement writing new base values for field 5 must stop any + // overlay from shadowing those cells: field 5 is tombstoned in place + // (preserving the overlay's field 3), and an overlay covering only field + // 5 is dropped entirely. + let mut fragment = Fragment::new(0); + fragment.files = vec![ + DataFile::new_legacy_from_fields("f3.lance", vec![3], None), + DataFile::new_legacy_from_fields("f5.lance", vec![5], None), + ]; + fragment.overlays = vec![ + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o35.lance", vec![3, 5], None), + coverage: OverlayCoverage::sparse(vec![ + roaring::RoaringBitmap::from_iter([0u32]), + roaring::RoaringBitmap::from_iter([0u32]), + ]), + committed_version: 3, + }, + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o5.lance", vec![5], None), + coverage: OverlayCoverage::dense(roaring::RoaringBitmap::from_iter([0u32])), + committed_version: 3, + }, + ]; + + let schema = ArrowSchema::new(vec![ArrowField::new("id", DataType::Int32, false)]); + let manifest = Manifest::new( + LanceSchema::try_from(&schema).unwrap(), + Arc::new(vec![fragment]), + lance_table::format::DataStorageFormat::new(LanceFileVersion::V2_0), + HashMap::new(), + ); + + let txn = Transaction::new( + manifest.version, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 0, + DataFile::new_legacy_from_fields("f5-new.lance", vec![5], None), + )], + }, + None, + ); + + let (result, _) = txn + .build_manifest( + Some(&manifest), + vec![], + "txn", + &ManifestWriteConfig::default(), + ) + .unwrap(); + + let frag = &result.fragments[0]; + // The base data file for field 5 was swapped in. + assert!(frag.files.iter().any(|f| f.path == "f5-new.lance")); + // The [3, 5] overlay keeps field 3 and tombstones field 5; the [5]-only + // overlay is dropped. + assert_eq!(frag.overlays.len(), 1); + assert_eq!(frag.overlays[0].data_file.fields.as_ref(), &[3, -2]); + } + + #[test] + fn test_data_overlay_build_manifest_merges_duplicate_groups() { + // Two groups targeting the same fragment must both survive (a HashMap + // collapse would have dropped the first). + let manifest = sample_manifest(); + let txn = Transaction::new( + manifest.version, + Operation::DataOverlay { + groups: vec![ + DataOverlayGroup { + fragment_id: 0, + overlays: vec![overlay_with_field(1, 0)], + }, + DataOverlayGroup { + fragment_id: 0, + overlays: vec![overlay_with_field(2, 0)], + }, + ], + }, + None, + ); + + let (result, _) = txn + .build_manifest( + Some(&manifest), + vec![], + "txn", + &ManifestWriteConfig::default(), + ) + .unwrap(); + + let overlays = &result.fragments[0].overlays; + assert_eq!(overlays.len(), 2); + assert_eq!(overlays[0].data_file.fields.as_ref(), [1i32].as_slice()); + assert_eq!(overlays[1].data_file.fields.as_ref(), [2i32].as_slice()); + } + + #[test] + fn test_data_overlay_build_manifest_rejects_unknown_fragment() { + let manifest = sample_manifest(); + let txn = Transaction::new( + manifest.version, + Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id: 99, + overlays: vec![overlay_with_field(1, 0)], + }], + }, + None, + ); + let err = txn + .build_manifest( + Some(&manifest), + vec![], + "txn", + &ManifestWriteConfig::default(), + ) + .unwrap_err(); + assert!(err.to_string().contains("does not exist"), "{err}"); + } + + #[test] + fn test_data_overlay_operation_eq() { + let overlay = |field: i32| Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id: 0, + overlays: vec![overlay_with_field(field, 1)], + }], + }; + // Reflexive and value-based (the arm previously returned false for self). + assert_eq!(overlay(1), overlay(1)); + assert_ne!(overlay(1), overlay(2)); + // Not equal to a different operation kind (previously returned true vs Rewrite). + let rewrite = Operation::Rewrite { + groups: vec![], + rewritten_indices: vec![], + frag_reuse_index: None, + }; + assert_ne!(overlay(1), rewrite); + } } diff --git a/rust/lance/src/dataset/updater.rs b/rust/lance/src/dataset/updater.rs index 90ef8df914b..752a479e5b0 100644 --- a/rust/lance/src/dataset/updater.rs +++ b/rust/lance/src/dataset/updater.rs @@ -238,6 +238,7 @@ impl Updater { cleanup_data_fragments( &self.dataset().object_store, &self.dataset().base, + None, &[fragment], ) .await; diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index ff0a119158c..b35cfeb3eed 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -8,7 +8,8 @@ use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use futures::{Stream, StreamExt, TryStreamExt}; use lance_arrow::{ ARROW_EXT_NAME_KEY, BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, - BLOB_INLINE_SIZE_THRESHOLD_META_KEY, BLOB_META_KEY, BLOB_V2_EXT_NAME, + BLOB_INLINE_SIZE_THRESHOLD_META_KEY, BLOB_META_KEY, BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY, + BLOB_V2_EXT_NAME, }; use lance_core::datatypes::{ NullabilityComparison, OnMissing, OnTypeMismatch, SchemaCompareOptions, @@ -25,22 +26,26 @@ use lance_file::previous::writer::{ }; use lance_file::version::LanceFileVersion; use lance_file::writer::{self as current_writer, FileWriterOptions}; -use lance_io::object_store::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry}; +use lance_io::object_store::{ + ObjectStore, ObjectStoreParams, ObjectStoreRegistry, parse_base_scoped_key, +}; use lance_table::format::{BasePath, DataFile, Fragment}; use lance_table::io::commit::{CommitHandler, commit_handler_from_url}; use lance_table::io::manifest::ManifestDescribing; use object_store::path::Path; -use std::collections::{HashMap, HashSet}; +use std::borrow::Cow; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::num::NonZero; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use tracing::{info, instrument}; use crate::Dataset; +use crate::blob::normalize_prepared_blob_schema; use crate::dataset::blob::{ BlobPreprocessor, ExternalBaseCandidate, ExternalBaseResolver, blob_dedicated_threshold_from_metadata, blob_inline_threshold_from_metadata, - preprocess_blob_batches, + blob_pack_file_threshold_from_metadata, preprocess_blob_batches, }; use crate::session::Session; @@ -183,66 +188,63 @@ fn validate_blob_threshold_metadata_for_append( let Some(dataset_field) = dataset_schema.field(&input_field.name) else { continue; }; - let input_is_blob_v2 = input_field - .metadata - .get(ARROW_EXT_NAME_KEY) - .is_some_and(|extension_name| extension_name == BLOB_V2_EXT_NAME); - let dataset_is_blob_v2 = dataset_field - .metadata - .get(ARROW_EXT_NAME_KEY) - .is_some_and(|extension_name| extension_name == BLOB_V2_EXT_NAME); - if !input_is_blob_v2 && !dataset_is_blob_v2 { - continue; - } + validate_blob_threshold_metadata_for_field_recursive(input_field, dataset_field)?; + } - let has_inline_threshold = input_field - .metadata - .contains_key(BLOB_INLINE_SIZE_THRESHOLD_META_KEY); - let has_dedicated_threshold = input_field - .metadata - .contains_key(BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY); - if !has_inline_threshold && !has_dedicated_threshold { - continue; - } + Ok(()) +} - if has_inline_threshold { - let input_inline_threshold = - blob_inline_threshold_from_metadata(&input_field.metadata, &input_field.name)?; - let dataset_inline_threshold = - blob_inline_threshold_from_metadata(&dataset_field.metadata, &dataset_field.name)?; - if input_inline_threshold != dataset_inline_threshold { - return Err(Error::invalid_input(format!( - "Cannot append data with blob threshold metadata {}={} for field '{}'; \ - the dataset schema has effective value {}. Blob thresholds for existing \ - columns are stored in the dataset schema.", - BLOB_INLINE_SIZE_THRESHOLD_META_KEY, - input_inline_threshold, - input_field.name, - dataset_inline_threshold, - ))); +fn validate_blob_threshold_metadata_for_field_recursive( + input_field: &lance_core::datatypes::Field, + dataset_field: &lance_core::datatypes::Field, +) -> Result<()> { + let input_is_blob_v2 = input_field + .metadata + .get(ARROW_EXT_NAME_KEY) + .is_some_and(|extension_name| extension_name == BLOB_V2_EXT_NAME); + let dataset_is_blob_v2 = dataset_field + .metadata + .get(ARROW_EXT_NAME_KEY) + .is_some_and(|extension_name| extension_name == BLOB_V2_EXT_NAME); + if input_is_blob_v2 || dataset_is_blob_v2 { + for (key, read_threshold) in [ + ( + BLOB_INLINE_SIZE_THRESHOLD_META_KEY, + blob_inline_threshold_from_metadata + as fn(&HashMap, &str) -> Result, + ), + ( + BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, + blob_dedicated_threshold_from_metadata, + ), + ( + BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY, + blob_pack_file_threshold_from_metadata, + ), + ] { + if !input_field.metadata.contains_key(key) { + continue; } - } - if has_dedicated_threshold { - let input_dedicated_threshold = - blob_dedicated_threshold_from_metadata(&input_field.metadata, &input_field.name)?; - let dataset_dedicated_threshold = blob_dedicated_threshold_from_metadata( - &dataset_field.metadata, - &dataset_field.name, - )?; - if input_dedicated_threshold != dataset_dedicated_threshold { + let input_value = read_threshold(&input_field.metadata, &input_field.name)?; + let dataset_value = read_threshold(&dataset_field.metadata, &dataset_field.name)?; + if input_value != dataset_value { return Err(Error::invalid_input(format!( - "Cannot append data with blob threshold metadata {}={} for field '{}'; \ - the dataset schema has effective value {}. Blob thresholds for existing \ - columns are stored in the dataset schema.", - BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, - input_dedicated_threshold, + "Cannot append data with blob threshold metadata {key}={input_value} for \ + field '{}'; the dataset schema has effective value {dataset_value}. Blob \ + thresholds for existing columns are stored in the dataset schema.", input_field.name, - dataset_dedicated_threshold, ))); } } } + for input_child in &input_field.children { + let Some(dataset_child) = dataset_field.child(&input_child.name) else { + continue; + }; + validate_blob_threshold_metadata_for_field_recursive(input_child, dataset_child)?; + } + Ok(()) } @@ -287,8 +289,16 @@ pub struct WriteParams { /// Write mode pub mode: WriteMode, + /// Default object store params for the write. + /// + /// Storage options may carry base-scoped entries (`base_.`) that + /// apply only to the registered base path with that id, overriding the + /// unscoped options that every base inherits. pub store_params: Option, + /// Exact object store params per base path URI, taking precedence over + /// `base_.` storage options in [`Self::store_params`]. See + /// [`Self::with_base_store_params`]. pub base_store_params: Option>, pub progress: Arc, @@ -313,6 +323,7 @@ pub struct WriteParams { /// /// Newer versions are more efficient but the data can only be read by more recent versions /// of lance. + /// Lance file version 2.3 enables RLE v2 run length widths by default. /// /// If not specified then the latest stable version will be used. pub data_storage_version: Option, @@ -371,13 +382,23 @@ pub struct WriteParams { /// The IDs must correspond to either: /// - IDs in initial_bases (for CREATE/OVERWRITE modes) /// - IDs already registered in the existing dataset manifest (for APPEND mode) + /// - [`PRIMARY_BASE_ID`] (0), which targets the dataset's primary storage + /// and participates in the round-robin like any other entry pub target_bases: Option>, /// Target base names or paths as strings (unresolved). /// These will be resolved to IDs when the write operation executes. /// Resolution happens at builder execution time when dataset context is available. + /// An entry equal to the dataset's URI targets the dataset's primary storage. pub target_base_names_or_paths: Option>, + /// Target every base registered in the dataset manifest, resolved when the + /// write executes. `Some(include_primary)`: when `include_primary` is true + /// the dataset's primary storage participates in the rotation as the first + /// slot. Cannot be combined with `target_bases` or + /// `target_base_names_or_paths`. + pub target_all_bases: Option, + /// Allow writing external blob URIs that cannot be mapped to any registered /// non-dataset-root base path. When disabled, such rows are rejected. pub allow_external_blob_outside_bases: bool, @@ -415,6 +436,7 @@ impl Default for WriteParams { initial_bases: None, target_bases: None, target_base_names_or_paths: None, + target_all_bases: None, allow_external_blob_outside_bases: false, external_blob_mode: ExternalBlobMode::Reference, blob_pack_file_size_threshold: None, @@ -445,8 +467,10 @@ impl WriteParams { /// Set exact runtime object store params for a registered base path. /// - /// These params are used as-is for that base. The write-level default - /// `store_params` remain the fallback for bases without an explicit binding. + /// These params are used as-is for that base, taking precedence over + /// `base_.` storage options in `store_params`. The write-level + /// default `store_params` remain the fallback for bases without an + /// explicit binding. pub fn with_base_store_params( mut self, base_path: impl AsRef, @@ -512,6 +536,19 @@ impl WriteParams { } } + /// Target every base registered in the dataset manifest, resolved when the + /// write executes. When `include_primary` is true the dataset's primary + /// storage participates in the rotation as the first slot. + /// + /// Cannot be combined with [`Self::with_target_bases`] or + /// [`Self::with_target_base_names_or_paths`]. + pub fn with_target_all_bases(self, include_primary: bool) -> Self { + Self { + target_all_bases: Some(include_primary), + ..self + } + } + /// Configure whether external blobs outside registered bases are allowed. pub fn with_allow_external_blob_outside_bases(self, allow: bool) -> Self { Self { @@ -596,6 +633,8 @@ pub async fn do_write_fragments( .unwrap_or_else(|| params.store_registry()); let source_store_params = params.store_params.clone().unwrap_or_default(); + // Keep a copy so failure paths can clean up files written to target bases. + let cleanup_bases = target_bases_info.clone(); let writer_generator = WriterGenerator::new( object_store.clone(), base_dir, @@ -676,7 +715,13 @@ pub async fn do_write_fragments( // Drop the writer so its in-progress file is cleaned up (LocalWriter // removes its temp file; ObjectWriter aborts the multipart upload). drop(writer.take()); - cleanup_data_fragments(&object_store, base_dir, &fragments).await; + cleanup_data_fragments( + &object_store, + base_dir, + cleanup_bases.as_deref(), + &fragments, + ) + .await; return Err(e); } @@ -701,7 +746,13 @@ pub async fn do_write_fragments( } Err(e) => { drop(writer); - cleanup_data_fragments(&object_store, base_dir, &fragments).await; + cleanup_data_fragments( + &object_store, + base_dir, + cleanup_bases.as_deref(), + &fragments, + ) + .await; return Err(e); } } @@ -715,53 +766,67 @@ pub async fn do_write_fragments( /// Contract: /// - Errors from individual `delete` calls are logged and swallowed, never returned — /// callers should propagate the original write error. -/// - Only files in the dataset's default storage (`base_id == None`) are deleted; -/// files in external bases are skipped because we don't have their object stores here. +/// - Files in the dataset's default storage (`base_id == None`) are deleted via +/// `object_store`; files whose `base_id` matches an entry in `target_bases` are +/// deleted via that base's object store. Files in bases not listed in +/// `target_bases` are skipped because we don't have their object stores here. /// - Safe to call with an empty slice. /// - Must be called before the fragments are committed, otherwise live data may be deleted. pub(crate) async fn cleanup_data_fragments( object_store: &ObjectStore, base_dir: &Path, + target_bases: Option<&[TargetBaseInfo]>, fragments: &[Fragment], ) { let data_dir = base_dir.clone().join(DATA_DIR); let mut skipped_external = 0usize; for fragment in fragments { for file in &fragment.files { - if file.base_id.is_none() { - let path = data_dir.clone().join(file.path.as_str()); - if let Err(e) = object_store.delete(&path).await { - log::warn!("Failed to clean up orphaned data file '{}': {}", path, e); + let (store, file_dir) = if let Some(base_id) = file.base_id { + match target_bases.and_then(|bases| bases.iter().find(|b| b.base_id == base_id)) { + Some(base_info) => { + let dir = if base_info.is_dataset_root { + base_info.base_dir.clone().join(DATA_DIR) + } else { + base_info.base_dir.clone() + }; + (base_info.object_store.as_ref(), dir) + } + None => { + skipped_external += 1; + continue; + } } + } else { + (object_store, data_dir.clone()) + }; - // Clean up any blob v2 sidecars that might exist for this data file. - // Blob v2 sidecars are written to `data/{data_file_key}/{blob_id}.blob`. - // The `data_file_key` is the file stem of the .lance file. - if let Some(stem) = std::path::Path::new(file.path.as_str()) - .file_stem() - .and_then(|s| s.to_str()) - { - let blob_dir = data_dir.clone().join(stem); - match object_store.remove_dir_all(blob_dir.clone()).await { - Err(e) if !matches!(e, Error::NotFound { .. }) => { - log::warn!( - "Failed to clean up orphaned blob dir '{}': {}", - blob_dir, - e - ); - } - _ => {} + let path = file_dir.clone().join(file.path.as_str()); + if let Err(e) = store.delete(&path).await { + log::warn!("Failed to clean up orphaned data file '{}': {}", path, e); + } + + // Clean up any blob v2 sidecars that might exist for this data file. + // Blob v2 sidecars are written to `data/{data_file_key}/{blob_id}.blob`. + // The `data_file_key` is the file stem of the .lance file. + if let Some(stem) = std::path::Path::new(file.path.as_str()) + .file_stem() + .and_then(|s| s.to_str()) + { + let blob_dir = file_dir.clone().join(stem); + match store.remove_dir_all(blob_dir.clone()).await { + Err(e) if !matches!(e, Error::NotFound { .. }) => { + log::warn!("Failed to clean up orphaned blob dir '{}': {}", blob_dir, e); } + _ => {} } - } else { - skipped_external += 1; } } } if skipped_external > 0 { log::warn!( "Skipped cleanup of {} orphaned data file(s) in external bases: \ - cleanup not supported for external bases", + their object stores are not available here", skipped_external ); } @@ -785,6 +850,12 @@ pub async fn validate_and_resolve_target_bases( )); } + if params.target_all_bases.is_some() { + return Err(Error::invalid_input( + "target_all_bases requires dataset context to resolve; use the write or merge insert APIs to apply it.", + )); + } + // Step 2: Assign IDs to initial_bases and add them to all_bases let mut all_bases: HashMap = existing_base_paths.cloned().unwrap_or_default(); if let Some(initial_bases) = &mut params.initial_bases { @@ -798,6 +869,11 @@ pub async fn validate_and_resolve_target_bases( all_bases.insert(base_path.id, base_path.clone()); } } + log_unregistered_base_scoped_options( + params.store_params.as_ref(), + &all_bases, + log::Level::Warn, + ); // Step 3: Resolve target_base_names_or_paths to IDs let target_base_ids = if let Some(ref names_or_paths) = params.target_base_names_or_paths { @@ -833,6 +909,13 @@ pub async fn validate_and_resolve_target_bases( .unwrap_or_default(); if let Some(target_bases) = &target_base_ids { + // An empty list would panic in round-robin selection; reject it + // instead of silently writing to primary storage. + if target_bases.is_empty() { + return Err(Error::invalid_input( + "target_bases cannot be empty. Omit the option to write to primary storage.", + )); + } let mut bases_info = Vec::new(); for &target_base_id in target_bases { @@ -843,7 +926,7 @@ pub async fn validate_and_resolve_target_bases( )) })?; - let store_params = write_store_params_for_base(params, &base_path.path); + let store_params = write_store_params_for_base(params, base_path); let (target_object_store, extracted_path) = ObjectStore::from_uri_and_params( store_registry.clone(), &base_path.path, @@ -865,6 +948,136 @@ pub async fn validate_and_resolve_target_bases( } } +/// Like [`validate_and_resolve_target_bases`], but also resolves references to +/// the dataset's primary storage: base id [`PRIMARY_BASE_ID`] in +/// `target_bases`, or an entry equal to `primary_uri` in +/// `target_base_names_or_paths`. Primary slots participate in the round-robin +/// like any other target base; files written through them carry no base id. +pub(crate) async fn validate_and_resolve_target_bases_with_primary( + params: &mut WriteParams, + existing_base_paths: Option<&HashMap>, + primary_object_store: &Arc, + primary_base_dir: &Path, + primary_uri: &str, +) -> Result>> { + // Expand an all-bases request into an explicit id list (primary first, + // then registered bases in ascending id order) and continue below. + if let Some(include_primary) = params.target_all_bases { + if params.target_bases.is_some() || params.target_base_names_or_paths.is_some() { + return Err(Error::invalid_input( + "Cannot specify target_all_bases together with target_bases or target_base_names_or_paths.", + )); + } + let mut ids: Vec = existing_base_paths + .map(|bases| bases.keys().copied().collect()) + .unwrap_or_default(); + // CREATE mode registers initial_bases in the same write; assign their + // ids here (the delegate keeps non-zero ids as-is) so they join the + // rotation. + if let Some(initial_bases) = &mut params.initial_bases { + let mut next_id = ids.iter().max().map(|id| id + 1).unwrap_or(1); + for base_path in initial_bases.iter_mut() { + if base_path.id == 0 { + base_path.id = next_id; + next_id += 1; + } + ids.push(base_path.id); + } + } + ids.sort_unstable(); + ids.dedup(); + if include_primary { + ids.insert(0, PRIMARY_BASE_ID); + } + if ids.is_empty() { + return Err(Error::invalid_input( + "target_all_bases found no registered bases and include_primary is false. \ + Register bases or include primary storage.", + )); + } + params.target_bases = Some(ids); + params.target_all_bases = None; + } + + let has_primary_ids = params + .target_bases + .as_ref() + .is_some_and(|ids| ids.contains(&PRIMARY_BASE_ID)); + let has_primary_refs = params + .target_base_names_or_paths + .as_ref() + .is_some_and(|refs| refs.iter().any(|r| r == primary_uri)); + if !has_primary_ids && !has_primary_refs { + return validate_and_resolve_target_bases(params, existing_base_paths).await; + } + + // The delegate below may be skipped when only primary slots remain, so + // validate mutual exclusion here as well. + if params.target_base_names_or_paths.is_some() && params.target_bases.is_some() { + return Err(Error::invalid_input( + "Cannot specify both target_base_names_or_paths and target_bases. Use one or the other.", + )); + } + + // Strip the primary slots, resolve the remaining references through the + // normal path, then splice the primary slots back into their original + // positions so the round-robin order matches what the caller asked for. + let is_primary_slot: Vec = if let Some(ids) = ¶ms.target_bases { + ids.iter().map(|id| *id == PRIMARY_BASE_ID).collect() + } else { + params + .target_base_names_or_paths + .as_ref() + .unwrap() + .iter() + .map(|r| r == primary_uri) + .collect() + }; + + let mut shim = params.clone(); + if let Some(ids) = ¶ms.target_bases { + let rest: Vec = ids + .iter() + .copied() + .filter(|id| *id != PRIMARY_BASE_ID) + .collect(); + shim.target_bases = if rest.is_empty() { None } else { Some(rest) }; + } else { + let rest: Vec = params + .target_base_names_or_paths + .as_ref() + .unwrap() + .iter() + .filter(|r| *r != primary_uri) + .cloned() + .collect(); + shim.target_base_names_or_paths = if rest.is_empty() { None } else { Some(rest) }; + } + + let resolved_rest = validate_and_resolve_target_bases(&mut shim, existing_base_paths).await?; + // The delegate assigns ids to initial_bases in place; propagate that side + // effect back so CREATE-mode transactions register properly assigned ids. + params.initial_bases = shim.initial_bases; + + let mut rest_iter = resolved_rest.unwrap_or_default().into_iter(); + let mut bases_info = Vec::with_capacity(is_primary_slot.len()); + for is_primary in is_primary_slot { + if is_primary { + bases_info.push(TargetBaseInfo { + base_id: PRIMARY_BASE_ID, + object_store: primary_object_store.clone(), + base_dir: primary_base_dir.clone(), + is_dataset_root: true, + }); + } else { + bases_info.push(rest_iter.next().ok_or_else(|| { + Error::internal("target base resolution returned fewer bases than requested") + })?); + } + } + Ok(Some(bases_info)) +} + fn append_external_base_candidate( base_path: &BasePath, store_prefix: String, @@ -886,22 +1099,53 @@ fn append_external_base_candidate( } } -fn write_store_params_for_base(params: &WriteParams, base_path: &str) -> ObjectStoreParams { - params - .base_store_params - .as_ref() - .and_then(|base_store_params| base_store_params.get(base_path)) - .cloned() - .unwrap_or_else(|| params.store_params.clone().unwrap_or_default()) +/// Log base-scoped storage options (`base_.`) whose id does not +/// match any registered base path. Unregistered entries are ignored during +/// resolution. The open path logs at debug, since options may legitimately be +/// vended for bases the loaded version does not register; the write path logs +/// at warn, since ids are already assigned there and an unmatched id is much +/// more likely a mistake. +pub(crate) fn log_unregistered_base_scoped_options( + store_params: Option<&ObjectStoreParams>, + base_paths: &HashMap, + level: log::Level, +) { + if !log::log_enabled!(level) { + return; + } + let Some(options) = store_params.and_then(|params| params.storage_options()) else { + return; + }; + let unregistered = options + .keys() + .filter_map(|key| parse_base_scoped_key(key).map(|(id, _)| id)) + .filter(|id| !base_paths.contains_key(id)) + .collect::>(); + if !unregistered.is_empty() { + log::log!( + level, + "Ignoring base-scoped storage options for unregistered base path ids: {:?}", + unregistered + ); + } } -fn dataset_store_params_for_base(dataset: &Dataset, base_path: &str) -> ObjectStoreParams { - dataset +fn write_store_params_for_base(params: &WriteParams, base_path: &BasePath) -> ObjectStoreParams { + // Exact per-URI bindings are used as-is. Otherwise the write-level default + // params are resolved for the base scope: `base_.` storage + // options overlay the shared defaults for that base. + if let Some(store_params) = params .base_store_params .as_ref() - .and_then(|base_store_params| base_store_params.get(base_path)) - .cloned() - .unwrap_or_else(|| dataset.store_params.as_deref().cloned().unwrap_or_default()) + .and_then(|base_store_params| base_store_params.get(&base_path.path)) + { + return store_params.clone(); + } + let default_params = params.store_params.clone().unwrap_or_default(); + match default_params.scoped_to_base(Some(base_path.id)) { + Cow::Owned(scoped_params) => scoped_params, + Cow::Borrowed(_) => default_params, + } } async fn append_external_initial_bases( @@ -913,7 +1157,7 @@ async fn append_external_initial_bases( ) -> Result<()> { if let Some(initial_bases) = initial_bases { for base_path in initial_bases { - let store_params = write_store_params_for_base(params, &base_path.path); + let store_params = write_store_params_for_base(params, base_path); let (store, extracted_path) = ObjectStore::from_uri_and_params( store_registry.clone(), &base_path.path, @@ -946,7 +1190,7 @@ async fn build_external_base_resolver( if let Some(dataset) = dataset { for base_path in dataset.manifest.base_paths.values() { - let store_params = dataset_store_params_for_base(dataset, &base_path.path); + let store_params = dataset.store_params_for_base(Some(base_path)); let (store, extracted_path) = ObjectStore::from_uri_and_params( store_registry.clone(), &base_path.path, @@ -1012,12 +1256,13 @@ pub async fn write_fragments_internal( // Make sure the max rows per group is not larger than the max rows per file params.max_rows_per_group = std::cmp::min(params.max_rows_per_group, params.max_rows_per_file); validate_external_blob_write_params(¶ms)?; + let normalized_converted_schema = normalize_prepared_blob_schema(&converted_schema)?; let (schema, storage_version) = if let Some(dataset) = dataset { match params.mode { WriteMode::Append | WriteMode::Create => { // Append mode, so we need to check compatibility - converted_schema.check_compatible( + normalized_converted_schema.check_compatible( dataset.schema(), &SchemaCompareOptions { // We don't care if the user claims their data is nullable / non-nullable. We will @@ -1029,9 +1274,12 @@ pub async fn write_fragments_internal( ..Default::default() }, )?; - validate_blob_threshold_metadata_for_append(&converted_schema, dataset.schema())?; + validate_blob_threshold_metadata_for_append( + &normalized_converted_schema, + dataset.schema(), + )?; let write_schema = dataset.schema().project_by_schema( - &converted_schema, + &normalized_converted_schema, OnMissing::Error, OnTypeMismatch::Error, )?; @@ -1052,13 +1300,16 @@ pub async fn write_fragments_internal( .data_storage_format .lance_file_version()?, ); - (converted_schema, data_storage_version) + (normalized_converted_schema, data_storage_version) } } } else { // Brand new dataset, use the schema from the data and the storage version // from the user or the default. - (converted_schema, params.storage_version_or_default()) + ( + normalized_converted_schema, + params.storage_version_or_default(), + ) }; if storage_version < LanceFileVersion::V2_2 && schema.fields_pre_order().any(|f| f.is_blob_v2()) @@ -1142,14 +1393,13 @@ where Ok(self.writer.tell().await? as u64) } async fn finish(&mut self) -> Result<(u32, DataFile)> { - let num_rows = self.writer.finish().await? as u32; - let size_bytes = self.writer.tell().await?; + let summary = self.writer.finish().await?; Ok(( - num_rows, + summary.num_rows as u32, DataFile::new_legacy( self.path.clone(), self.writer.schema(), - NonZero::new(size_bytes as u64), + NonZero::new(summary.size_bytes), self.base_id, ), )) @@ -1358,9 +1608,18 @@ async fn open_writer_with_options( Ok(writer) } +/// Reserved base id that refers to the dataset's primary storage in +/// [`WriteParams::target_bases`]. Real base ids are assigned starting from 1, +/// so 0 is never a registered base. Files written through a primary slot +/// carry no base id, exactly like a write without target bases. +pub const PRIMARY_BASE_ID: u32 = 0; + /// Information about a target base for writing. /// Contains the base ID, object store, directory path, and whether it's a dataset root. +#[derive(Clone)] pub struct TargetBaseInfo { + /// The registered base id, or [`PRIMARY_BASE_ID`] for the dataset's + /// primary storage. pub base_id: u32, pub object_store: Arc, /// The base directory path (without /data subdirectory) @@ -1444,7 +1703,9 @@ impl WriterGenerator { self.storage_version, WriterOptions { add_data_dir: base_info.is_dataset_root, - base_id: Some(base_info.base_id), + // Primary-storage slots stamp no base id, like a write + // without target bases. + base_id: (base_info.base_id != PRIMARY_BASE_ID).then_some(base_info.base_id), external_base_resolver: self.external_base_resolver.clone(), allow_external_blob_outside_bases: self.allow_external_blob_outside_bases, external_blob_mode: self.external_blob_mode, @@ -2104,7 +2365,28 @@ mod tests { .with_base_store_params("az://container/path-a", azure_store_params("account-a")) .with_base_store_params("az://container/path-b", azure_store_params("account-b")); - let existing_base_paths = HashMap::from([ + let existing_base_paths = azure_base_paths_a_b(); + + let target_bases = + validate_and_resolve_target_bases(&mut params, Some(&existing_base_paths)) + .await + .unwrap() + .unwrap(); + + assert_eq!(target_bases.len(), 2); + assert_eq!( + target_bases[0].object_store.store_prefix, + "az$container@account-a" + ); + assert_eq!( + target_bases[1].object_store.store_prefix, + "az$container@account-b" + ); + } + + #[cfg(feature = "azure")] + fn azure_base_paths_a_b() -> HashMap { + HashMap::from([ ( 1, BasePath::new( @@ -2123,7 +2405,32 @@ mod tests { false, ), ), - ]); + ]) + } + + #[cfg(feature = "azure")] + #[tokio::test] + async fn test_validate_and_resolve_target_bases_uses_base_scoped_storage_options() { + // A single flat storage options map carries per-base credentials via + // the `base_.` convention; unscoped keys are shared defaults. + let store_params = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("account_name".to_string(), "account-shared".to_string()), + ("account_key".to_string(), "dGVzdA==".to_string()), + ("base_1.account_name".to_string(), "account-a".to_string()), + ("base_2.account_name".to_string(), "account-b".to_string()), + ]), + ))), + ..Default::default() + }; + let mut params = WriteParams { + store_params: Some(store_params), + ..Default::default() + } + .with_target_bases(vec![1, 2]); + + let existing_base_paths = azure_base_paths_a_b(); let target_bases = validate_and_resolve_target_bases(&mut params, Some(&existing_base_paths)) @@ -2142,6 +2449,43 @@ mod tests { ); } + #[cfg(feature = "azure")] + #[tokio::test] + async fn test_base_store_params_take_precedence_over_base_scoped_options() { + let store_params = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("account_key".to_string(), "dGVzdA==".to_string()), + ( + "base_1.account_name".to_string(), + "account-scoped".to_string(), + ), + ]), + ))), + ..Default::default() + }; + let mut params = WriteParams { + store_params: Some(store_params), + ..Default::default() + } + .with_target_bases(vec![1]) + .with_base_store_params("az://container/path-a", azure_store_params("account-exact")); + + let existing_base_paths = azure_base_paths_a_b(); + + let target_bases = + validate_and_resolve_target_bases(&mut params, Some(&existing_base_paths)) + .await + .unwrap() + .unwrap(); + + assert_eq!(target_bases.len(), 1); + assert_eq!( + target_bases[0].object_store.store_prefix, + "az$container@account-exact" + ); + } + #[tokio::test] async fn test_explicit_data_file_bases_writer_generator() { use arrow::datatypes::{DataType, Field as ArrowField, Schema as ArrowSchema}; @@ -2539,6 +2883,80 @@ mod tests { ); } + #[tokio::test] + async fn test_multi_base_write_read_with_base_scoped_storage_options() { + use crate::dataset::builder::DatasetBuilder; + use lance_core::utils::tempfile::TempStrDir; + use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; + + let primary_dir = TempStrDir::default(); + let base1_dir = TempStrDir::default(); + + // Local stores ignore these options; this verifies base-scoped entries + // flow through the full write/read path without breaking anything. + let scoped_options = HashMap::from([ + ("shared_option".to_string(), "shared".to_string()), + ( + "base_1.scoped_option".to_string(), + "base1-value".to_string(), + ), + ]); + let store_params = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + scoped_options.clone(), + ))), + ..Default::default() + }; + + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + let dataset = Dataset::write( + data_gen.batch(5), + primary_dir.as_str(), + Some(WriteParams { + mode: WriteMode::Create, + store_params: Some(store_params), + initial_bases: Some(vec![BasePath { + id: 1, + name: Some("base1".to_string()), + path: base1_dir.as_str().to_string(), + is_dataset_root: true, + }]), + target_bases: Some(vec![1]), + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.count_rows(None).await.unwrap(), 5); + for fragment in dataset.get_fragments() { + assert!( + fragment + .metadata + .files + .iter() + .all(|file| file.base_id == Some(1)) + ); + } + + // Reopen with the same flat options and scan through the base store. + let dataset = DatasetBuilder::from_uri(primary_dir.as_str()) + .with_storage_options(scoped_options) + .load() + .await + .unwrap(); + let batches = dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let num_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(num_rows, 5); + } + #[tokio::test] async fn test_multi_base_overwrite() { use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; @@ -3562,6 +3980,7 @@ mod tests { let fragments = vec![Fragment { id: 0, files: vec![external_file, local_file], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(0), @@ -3569,7 +3988,7 @@ mod tests { last_updated_at_version_meta: None, }]; - cleanup_data_fragments(&object_store, &base_dir, &fragments).await; + cleanup_data_fragments(&object_store, &base_dir, None, &fragments).await; // The local file should be removed; the external file is skipped without // erroring (its base store isn't known here). @@ -3579,4 +3998,487 @@ mod tests { "Local data file should be deleted by cleanup" ); } + + /// Verifies the target-base branch in `cleanup_data_fragments`: files whose + /// `base_id` matches a provided [`TargetBaseInfo`] are deleted via that base's + /// object store (respecting `is_dataset_root` layout), while files in bases + /// without a provided store are still skipped. + #[tokio::test] + async fn test_cleanup_data_fragments_deletes_target_base_files() { + use lance_core::utils::tempfile::TempStrDir; + + let primary_dir = TempStrDir::default(); + let base1_dir = TempStrDir::default(); + let base2_dir = TempStrDir::default(); + + let (object_store, base_dir) = ObjectStore::from_uri_and_params( + Default::default(), + primary_dir.as_str(), + &Default::default(), + ) + .await + .unwrap(); + let (base1_store, base1_path) = ObjectStore::from_uri_and_params( + Default::default(), + base1_dir.as_str(), + &Default::default(), + ) + .await + .unwrap(); + let (base2_store, base2_path) = ObjectStore::from_uri_and_params( + Default::default(), + base2_dir.as_str(), + &Default::default(), + ) + .await + .unwrap(); + + // base2 is a plain data directory: files sit at its root, not under data/. + let count_plain_files = |dir: &str| { + std::fs::read_dir(dir) + .map(|entries| { + entries + .filter(|e| e.as_ref().unwrap().path().is_file()) + .count() + }) + .unwrap_or(0) + }; + + // base1 is a dataset root (files under data/), base2 is a plain data dir. + let base1_file_path = base1_path.clone().join(DATA_DIR).join("one.lance"); + base1_store.put(&base1_file_path, b"x").await.unwrap(); + let base2_file_path = base2_path.clone().join("two.lance"); + base2_store.put(&base2_file_path, b"x").await.unwrap(); + assert_eq!(count_data_files(base1_dir.as_str()), 1); + assert_eq!(count_plain_files(base2_dir.as_str()), 1); + + let mut base1_file = DataFile::new_unstarted("one.lance", 2, 1); + base1_file.base_id = Some(1); + let mut base2_file = DataFile::new_unstarted("two.lance", 2, 1); + base2_file.base_id = Some(2); + let mut unknown_file = DataFile::new_unstarted("unknown.lance", 2, 1); + unknown_file.base_id = Some(42); + let fragments = vec![Fragment { + id: 0, + files: vec![base1_file, base2_file, unknown_file], + overlays: vec![], + deletion_file: None, + row_id_meta: None, + physical_rows: Some(0), + created_at_version_meta: None, + last_updated_at_version_meta: None, + }]; + + let target_bases = vec![ + TargetBaseInfo { + base_id: 1, + object_store: base1_store, + base_dir: base1_path, + is_dataset_root: true, + }, + TargetBaseInfo { + base_id: 2, + object_store: base2_store, + base_dir: base2_path, + is_dataset_root: false, + }, + ]; + + cleanup_data_fragments(&object_store, &base_dir, Some(&target_bases), &fragments).await; + + assert_eq!( + count_data_files(base1_dir.as_str()), + 0, + "File in dataset-root target base should be deleted" + ); + assert_eq!( + count_plain_files(base2_dir.as_str()), + 0, + "File in plain-directory target base should be deleted" + ); + } + + #[tokio::test] + async fn test_cleanup_routed_data_files_on_failed_write() { + // Files already completed in target bases must be removed when the + // write later fails. + use lance_core::utils::tempfile::TempStrDir; + + let primary_dir = TempStrDir::default(); + let base1_dir = TempStrDir::default(); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + let schema = Schema::try_from(arrow_schema.as_ref()).unwrap(); + + let (object_store, base_dir) = ObjectStore::from_uri_and_params( + Default::default(), + primary_dir.as_str(), + &Default::default(), + ) + .await + .unwrap(); + let (base1_store, base1_path) = ObjectStore::from_uri_and_params( + Default::default(), + base1_dir.as_str(), + &Default::default(), + ) + .await + .unwrap(); + + let good_batch = RecordBatch::try_new( + arrow_schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + + // 3 rows per file: the first batch fills and completes a file in the + // target base, then the stream fails. + let items: Vec> = vec![ + Ok(good_batch.clone()), + Ok(good_batch.clone()), + Err(DataFusionError::External("injected failure".into())), + ]; + let stream = Box::pin(RecordBatchStreamAdapter::new( + arrow_schema.clone(), + futures::stream::iter(items), + )); + + let target_bases = vec![TargetBaseInfo { + base_id: 1, + object_store: base1_store, + base_dir: base1_path, + is_dataset_root: true, + }]; + + let result = do_write_fragments( + None, + object_store.clone(), + &base_dir, + &schema, + stream, + WriteParams { + max_rows_per_file: 3, + ..Default::default() + }, + LanceFileVersion::V2_1, + Some(target_bases), + ) + .await; + + assert!(result.is_err(), "Expected write to fail"); + assert_eq!( + count_data_files(base1_dir.as_str()), + 0, + "Data files routed to the target base should be cleaned up on failure" + ); + assert_eq!(count_data_files(primary_dir.as_str()), 0); + } + + /// PRIMARY_BASE_ID (0) and the dataset URI include primary storage in the + /// target rotation, alongside registered bases. + #[tokio::test] + async fn test_multi_base_target_primary_and_bases() { + use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; + + let test_uri = "memory://primary_slot_test"; + let primary_uri = format!("{}/primary", test_uri); + let base1_uri = format!("{}/base1", test_uri); + let base2_uri = format!("{}/base2", test_uri); + + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + + // CREATE mode targeting primary + a new base: also verifies the id + // assignment on initial_bases reaches the committed manifest. + let dataset = Dataset::write( + data_gen.batch(6), + &primary_uri, + Some(WriteParams { + mode: WriteMode::Create, + max_rows_per_file: 3, + initial_bases: Some(vec![ + BasePath { + id: 1, + name: Some("base1".to_string()), + is_dataset_root: true, + path: base1_uri.clone(), + }, + BasePath { + id: 2, + name: Some("base2".to_string()), + is_dataset_root: false, + path: base2_uri.clone(), + }, + ]), + target_bases: Some(vec![PRIMARY_BASE_ID, 1]), + ..Default::default() + }), + ) + .await + .unwrap(); + + assert_eq!(dataset.manifest.base_paths.len(), 2); + let file_bases: Vec<_> = dataset + .get_fragments() + .iter() + .flat_map(|f| f.metadata.files.iter().map(|file| file.base_id)) + .collect(); + assert_eq!(file_bases, vec![None, Some(1)]); + + // APPEND across primary + both bases, one file per slot in order. + let mut data_gen2 = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + let dataset = Dataset::write( + data_gen2.batch(9), + Arc::new(dataset), + Some(WriteParams { + mode: WriteMode::Append, + max_rows_per_file: 3, + target_bases: Some(vec![PRIMARY_BASE_ID, 1, 2]), + ..Default::default() + }), + ) + .await + .unwrap(); + let file_bases: Vec<_> = dataset + .get_fragments() + .iter() + .skip(2) + .flat_map(|f| f.metadata.files.iter().map(|file| file.base_id)) + .collect(); + assert_eq!(file_bases, vec![None, Some(1), Some(2)]); + + // Names variant: the dataset's own URI selects primary storage. + let mut data_gen3 = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + let dataset = Dataset::write( + data_gen3.batch(6), + Arc::new(dataset), + Some(WriteParams { + mode: WriteMode::Append, + max_rows_per_file: 3, + target_base_names_or_paths: Some(vec![primary_uri.clone(), "base2".to_string()]), + ..Default::default() + }), + ) + .await + .unwrap(); + let file_bases: Vec<_> = dataset + .get_fragments() + .iter() + .skip(5) + .flat_map(|f| f.metadata.files.iter().map(|file| file.base_id)) + .collect(); + assert_eq!(file_bases, vec![None, Some(2)]); + + assert_eq!(dataset.count_rows(None).await.unwrap(), 21); + } + + /// `target_all_bases` resolves to every registered base at execution + /// time, with primary storage as the first slot when included. + #[tokio::test] + async fn test_multi_base_target_all_bases() { + use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; + + let test_uri = "memory://all_bases_test"; + let primary_uri = format!("{}/primary", test_uri); + let base1_uri = format!("{}/base1", test_uri); + let base2_uri = format!("{}/base2", test_uri); + + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + let dataset = Dataset::write( + data_gen.batch(3), + &primary_uri, + Some(WriteParams { + mode: WriteMode::Create, + initial_bases: Some(vec![ + BasePath { + id: 1, + name: Some("base1".to_string()), + is_dataset_root: true, + path: base1_uri.clone(), + }, + BasePath { + id: 2, + name: Some("base2".to_string()), + is_dataset_root: false, + path: base2_uri.clone(), + }, + ]), + ..Default::default() + }), + ) + .await + .unwrap(); + + // All bases including primary: slots are [primary, base1, base2]. + let mut data_gen2 = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + let dataset = Dataset::write( + data_gen2.batch(9), + Arc::new(dataset), + Some( + WriteParams { + mode: WriteMode::Append, + max_rows_per_file: 3, + ..Default::default() + } + .with_target_all_bases(true), + ), + ) + .await + .unwrap(); + let file_bases: Vec<_> = dataset + .get_fragments() + .iter() + .skip(1) + .flat_map(|f| f.metadata.files.iter().map(|file| file.base_id)) + .collect(); + assert_eq!(file_bases, vec![None, Some(1), Some(2)]); + + // Without primary: slots are [base1, base2]. + let mut data_gen3 = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + let dataset = Dataset::write( + data_gen3.batch(6), + Arc::new(dataset), + Some( + WriteParams { + mode: WriteMode::Append, + max_rows_per_file: 3, + ..Default::default() + } + .with_target_all_bases(false), + ), + ) + .await + .unwrap(); + let file_bases: Vec<_> = dataset + .get_fragments() + .iter() + .skip(4) + .flat_map(|f| f.metadata.files.iter().map(|file| file.base_id)) + .collect(); + assert_eq!(file_bases, vec![Some(1), Some(2)]); + + // Cannot be combined with explicit target bases. + let mut data_gen4 = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + let result = Dataset::write( + data_gen4.batch(3), + Arc::new(dataset), + Some( + WriteParams { + mode: WriteMode::Append, + target_bases: Some(vec![1]), + ..Default::default() + } + .with_target_all_bases(true), + ), + ) + .await; + assert!( + result + .unwrap_err() + .to_string() + .contains("Cannot specify target_all_bases together with") + ); + + // On a dataset with no registered bases: include_primary=true is a + // no-op rotation over primary, false is rejected. + let plain_uri = "memory://all_bases_plain"; + let mut data_gen5 = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + let plain = Dataset::write(data_gen5.batch(3), plain_uri, None) + .await + .unwrap(); + let mut data_gen6 = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + let plain = Dataset::write( + data_gen6.batch(3), + Arc::new(plain), + Some( + WriteParams { + mode: WriteMode::Append, + ..Default::default() + } + .with_target_all_bases(true), + ), + ) + .await + .unwrap(); + assert!( + plain.get_fragments().iter().all(|f| f + .metadata + .files + .iter() + .all(|file| file.base_id.is_none())) + ); + let mut data_gen7 = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + let result = Dataset::write( + data_gen7.batch(3), + Arc::new(plain), + Some( + WriteParams { + mode: WriteMode::Append, + ..Default::default() + } + .with_target_all_bases(false), + ), + ) + .await; + assert!( + result + .unwrap_err() + .to_string() + .contains("target_all_bases found no registered bases") + ); + + // CREATE mode: initial_bases join the rotation before their ids are + // committed to a manifest. + let create_uri = "memory://all_bases_create"; + let mut data_gen8 = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + let dataset = Dataset::write( + data_gen8.batch(9), + create_uri, + Some( + WriteParams { + mode: WriteMode::Create, + max_rows_per_file: 3, + initial_bases: Some(vec![ + BasePath { + id: 0, + name: Some("base1".to_string()), + is_dataset_root: true, + path: format!("{}/base1", create_uri), + }, + BasePath { + id: 0, + name: Some("base2".to_string()), + is_dataset_root: false, + path: format!("{}/base2", create_uri), + }, + ]), + ..Default::default() + } + .with_target_all_bases(true), + ), + ) + .await + .unwrap(); + assert_eq!(dataset.manifest.base_paths.len(), 2); + let file_bases: Vec<_> = dataset + .get_fragments() + .iter() + .flat_map(|f| f.metadata.files.iter().map(|file| file.base_id)) + .collect(); + assert_eq!(file_bases, vec![None, Some(1), Some(2)]); + } } diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index baad71b3e39..e0e9b32e998 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -345,7 +345,6 @@ impl<'a> CommitBuilder<'a> { } else { self.use_stable_row_ids.unwrap_or(false) }; - // Validate storage format matches existing dataset if let Some(ds) = dest.dataset() && let Some(storage_format) = self.storage_format @@ -503,7 +502,6 @@ impl<'a> CommitBuilder<'a> { }, read_version, tag: None, - //TODO: handle batch transaction merges in the future transaction_properties: None, }; let dataset = self.execute(merged.clone()).await?; @@ -551,6 +549,7 @@ mod tests { file_size_bytes: CachedFileSize::new(100), base_id: None, }], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(10), @@ -630,10 +629,12 @@ mod tests { .unwrap(); assert_eq!(new_ds.manifest().version, 7); // Session should still be re-used - // However, the dataset needs to be loaded and the read version checked out, - // so an additional 4 IOPs are needed. + // However, the dataset needs to be loaded and the read version checked out. + // The read version's manifest body is served from the session cache (it + // was cached when v1 was first created), so the checkout only pays the + // version-resolution head, not a manifest read. let io_stats = dataset.object_store.as_ref().io_stats_incremental(); - assert_io_eq!(io_stats, read_iops, 5, "load dataset + check version"); + assert_io_eq!(io_stats, read_iops, 3, "load dataset + check version"); assert_io_eq!(io_stats, write_iops, 2, "write txn + manifest"); // Commit transaction with URI and new session. Re-use the store diff --git a/rust/lance/src/dataset/write/insert.rs b/rust/lance/src/dataset/write/insert.rs index bfd702c9c3b..6e1db342f9c 100644 --- a/rust/lance/src/dataset/write/insert.rs +++ b/rust/lance/src/dataset/write/insert.rs @@ -19,10 +19,13 @@ use lance_table::io::commit::CommitHandler; use object_store::path::Path; use crate::Dataset; +use crate::blob::normalize_prepared_blob_schema; use crate::dataset::ReadParams; use crate::dataset::builder::DatasetBuilder; use crate::dataset::transaction::{Operation, Transaction, TransactionBuilder}; -use crate::dataset::write::{validate_and_resolve_target_bases, write_fragments_internal}; +use crate::dataset::write::{ + validate_and_resolve_target_bases_with_primary, write_fragments_internal, +}; use crate::{Error, Result}; use tracing::info; @@ -201,8 +204,14 @@ impl<'a> InsertBuilder<'a> { self.validate_write(&mut context, &schema)?; let existing_base_paths = context.dest.dataset().map(|ds| &ds.manifest.base_paths); - let target_base_info = - validate_and_resolve_target_bases(&mut context.params, existing_base_paths).await?; + let target_base_info = validate_and_resolve_target_bases_with_primary( + &mut context.params, + existing_base_paths, + &context.object_store, + &context.base_path, + &context.dest.uri(), + ) + .await?; let (written_fragments, written_schema) = write_fragments_internal( context.dest.dataset(), @@ -315,7 +324,8 @@ impl<'a> InsertBuilder<'a> { ..Default::default() }; - data_schema.check_compatible(dataset.schema(), &schema_cmp_opts)?; + let normalized_data_schema = normalize_prepared_blob_schema(data_schema)?; + normalized_data_schema.check_compatible(dataset.schema(), &schema_cmp_opts)?; } // Make sure we aren't using any reserved column names diff --git a/rust/lance/src/dataset/write/merge_insert.rs b/rust/lance/src/dataset/write/merge_insert.rs index 3889bdc66d5..bc13a923613 100644 --- a/rust/lance/src/dataset/write/merge_insert.rs +++ b/rust/lance/src/dataset/write/merge_insert.rs @@ -43,7 +43,10 @@ use inserted_rows::KeyExistenceFilter; use super::cleanup_data_fragments; use super::retry::{RetryConfig, RetryExecutor, execute_with_retry}; -use super::{CommitBuilder, WriteParams, write_fragments_internal}; +use super::{ + CommitBuilder, TargetBaseInfo, WriteMode, WriteParams, + validate_and_resolve_target_bases_with_primary, write_fragments_internal, +}; use crate::dataset::rowids::get_row_id_index; use crate::dataset::transaction::UpdateMode::{RewriteColumns, RewriteRows}; use crate::dataset::utils::CapturedRowIds; @@ -97,6 +100,7 @@ use futures::{ Stream, StreamExt, TryStreamExt, stream::{self}, }; +use lance_arrow::json::{convert_json_columns, has_json_fields, is_arrow_json_field}; use lance_arrow::{RecordBatchExt, SchemaExt, interleave_batches}; use lance_core::datatypes::NullabilityComparison; use lance_core::utils::address::RowAddress; @@ -351,6 +355,14 @@ struct MergeInsertParams { source_dedupe_behavior: SourceDedupeBehavior, // Number of inner commit retries for manifest version conflicts. Default is 20. commit_retries: Option, + // Target base IDs for routing new fragments, mirroring WriteParams::target_bases. + target_bases: Option>, + // Target base names or path URIs (unresolved), mirroring + // WriteParams::target_base_names_or_paths. Resolved at execution time. + target_base_names_or_paths: Option>, + // Target all registered bases, mirroring WriteParams::target_all_bases. + // Some(include_primary); resolved at execution time. + target_all_bases: Option, } /// A MergeInsertJob inserts new rows, deletes old rows, and updates existing rows all as @@ -471,6 +483,9 @@ impl MergeInsertBuilder { use_index: true, source_dedupe_behavior: SourceDedupeBehavior::Fail, commit_retries: None, + target_bases: None, + target_base_names_or_paths: None, + target_all_bases: None, }, }) } @@ -569,6 +584,46 @@ impl MergeInsertBuilder { self } + /// Write new fragments produced by this merge insert to these base IDs. + /// + /// New data files are distributed across the target bases round-robin, + /// the same way a normal write with [`WriteParams::target_bases`] routes + /// them. The IDs must be registered in the dataset manifest, or + /// [`super::PRIMARY_BASE_ID`] (0) to include the dataset's primary + /// storage in the rotation (e.g. `vec![0, 1, 2]` spreads across primary + /// plus bases 1 and 2). Data files that patch existing fragments and + /// deletion files are always written to the dataset's primary storage. + /// + /// Cannot be combined with [`Self::target_base_names_or_paths`]. + pub fn target_bases(&mut self, base_ids: Vec) -> &mut Self { + self.params.target_bases = Some(base_ids); + self + } + + /// Like [`Self::target_bases`], but referencing bases by name or path URI. + /// + /// References are resolved against the base paths registered in the + /// dataset manifest when the merge insert executes. An entry equal to the + /// dataset's URI includes the dataset's primary storage in the rotation. + /// + /// Cannot be combined with [`Self::target_bases`]. + pub fn target_base_names_or_paths(&mut self, refs: Vec) -> &mut Self { + self.params.target_base_names_or_paths = Some(refs); + self + } + + /// Write new fragments produced by this merge insert to every base + /// registered in the dataset manifest, resolved when the merge executes. + /// When `include_primary` is true the dataset's primary storage + /// participates in the rotation as the first slot. + /// + /// Cannot be combined with [`Self::target_bases`] or + /// [`Self::target_base_names_or_paths`]. + pub fn target_all_bases(&mut self, include_primary: bool) -> &mut Self { + self.params.target_all_bases = Some(include_primary); + self + } + /// Crate a merge insert job pub fn try_build(&mut self) -> Result { if !self.params.insert_not_matched @@ -579,6 +634,19 @@ impl MergeInsertBuilder { "The merge insert job is not configured to change the data in any way", )); } + if self.params.target_bases.is_some() && self.params.target_base_names_or_paths.is_some() { + return Err(Error::invalid_input( + "Cannot specify both target_base_names_or_paths and target_bases. Use one or the other.", + )); + } + if self.params.target_all_bases.is_some() + && (self.params.target_bases.is_some() + || self.params.target_base_names_or_paths.is_some()) + { + return Err(Error::invalid_input( + "Cannot specify target_all_bases together with target_bases or target_base_names_or_paths.", + )); + } Ok(MergeInsertJob { dataset: self.dataset.clone(), params: self.params.clone(), @@ -586,6 +654,45 @@ impl MergeInsertBuilder { } } +/// Resolve the merge insert target bases against the base paths registered in +/// the dataset manifest. Returns `None` when no target bases were requested. +/// Base id [`super::PRIMARY_BASE_ID`] and the dataset's URI refer to the +/// dataset's primary storage. +/// +/// Resolution runs once per execution attempt so retries validate against the +/// manifest version they are writing to. +async fn resolve_target_bases( + dataset: &Dataset, + params: &MergeInsertParams, +) -> Result>> { + if params.target_bases.is_none() + && params.target_base_names_or_paths.is_none() + && params.target_all_bases.is_none() + { + return Ok(None); + } + // Reuse the normal write path resolution (validation, name/path lookup, + // and per-base credential handling) through a parameter shim. + let mut write_params = WriteParams { + mode: WriteMode::Append, + target_bases: params.target_bases.clone(), + target_base_names_or_paths: params.target_base_names_or_paths.clone(), + target_all_bases: params.target_all_bases, + session: Some(dataset.session.clone()), + store_params: dataset.store_params.as_deref().cloned(), + base_store_params: dataset.base_store_params.as_deref().cloned(), + ..Default::default() + }; + validate_and_resolve_target_bases_with_primary( + &mut write_params, + Some(&dataset.manifest.base_paths), + &dataset.object_store, + &dataset.base, + dataset.uri(), + ) + .await +} + enum SchemaComparison { FullCompatible, Subschema, @@ -981,7 +1088,11 @@ impl MergeInsertJob { dataset: Arc, source: SendableRecordBatchStream, current_version: u64, + target_bases_info: Option>, ) -> Result<(Vec, Vec, Vec)> { + // Shared across the per-group tasks spawned below; only new fragments + // are routed to target bases, column patches stay in primary storage. + let target_bases_info = Arc::new(target_bases_info); // Expected source schema: _rowaddr, updated_cols* use datafusion::logical_expr::{col, lit}; let session_ctx = get_session_context(&LanceExecutionOptions { @@ -1029,7 +1140,38 @@ impl MergeInsertJob { let reservation = MemoryConsumer::new("MergeInsert").register(session_ctx.task_ctx().memory_pool()); - while let Some((frag_id, batches)) = group_stream.next().await.transpose()? { + // Best-effort removal of uncommitted files after a mid-update failure. + // Aborts in-flight tasks first, then deletes the new fragments written + // so far (including ones routed to target bases). Column-patch files + // from completed tasks stay in primary storage where regular dataset + // cleanup can reclaim them. + async fn cleanup_on_failure( + dataset: &Dataset, + target_bases_info: &Option>, + new_fragments: &Mutex>, + tasks: &mut JoinSet>, + ) { + tasks.shutdown().await; + let written = new_fragments.lock().unwrap().clone(); + cleanup_data_fragments( + &dataset.object_store, + &dataset.base, + target_bases_info.as_deref(), + &written, + ) + .await; + } + + loop { + let (frag_id, batches) = match group_stream.next().await.transpose() { + Ok(Some(group)) => group, + Ok(None) => break, + Err(e) => { + cleanup_on_failure(&dataset, &target_bases_info, &new_fragments, &mut tasks) + .await; + return Err(e.into()); + } + }; async fn handle_fragment( dataset: Arc, fragment: FileFragment, @@ -1081,6 +1223,20 @@ impl MergeInsertJob { Err(e) => Err(e), })?; + // Convert Arrow JSON columns (Utf8) to Lance JSON (LargeBinary/JSONB) + // before writing. Without this, Utf8 data is written raw while the + // schema says LargeBinary, causing decoder panics on subsequent reads. + let needs_json_conversion = batches[0] + .schema() + .fields() + .iter() + .any(|f| is_arrow_json_field(f) || has_json_fields(f)); + if needs_json_conversion { + for batch in batches.iter_mut() { + *batch = convert_json_columns(batch).map_err(Error::from)?; + } + } + if data_storage_version == LanceFileVersion::Legacy { // Need to match the existing batch size exactly, otherwise // we'll get errors. @@ -1133,9 +1289,16 @@ impl MergeInsertJob { // will be the original source data, and all subsequent batches // will be updates. let mut source_batches = Vec::with_capacity(batches.len() + 1); - source_batches.push(batches[0].clone()); // placeholder for source data + // Convert Arrow JSON columns (Utf8) to Lance JSON (LargeBinary) so every + // batch is in physical format, matching what the updater reads from the + // fragment. `convert_json_columns` is a no-op clone when there is nothing + // to convert, so it can be applied unconditionally. The first entry is a + // placeholder for the source data (overwritten each iteration below); it + // must be converted too, otherwise its schema would diverge from the rest. + source_batches.push(convert_json_columns(&batches[0]).map_err(Error::from)?); for batch in &batches { - source_batches.push(batch.drop_column(ROW_ADDR)?); + let dropped = batch.drop_column(ROW_ADDR)?; + source_batches.push(convert_json_columns(&dropped).map_err(Error::from)?); } // This function is here to help rustc with lifetimes. @@ -1241,6 +1404,7 @@ impl MergeInsertJob { batches: Vec, new_fragments: Arc>>, reservation_size: usize, + target_bases_info: Arc>>, ) -> Result { // Batches still have _rowaddr (used elsewhere to merge with existing data) // We need to remove it before writing to Lance files. @@ -1272,7 +1436,7 @@ impl MergeInsertJob { write_schema, stream, Default::default(), // TODO: support write params. - None, // Merge insert doesn't use target_bases + (*target_bases_info).clone(), ) .await?; @@ -1300,15 +1464,26 @@ impl MergeInsertJob { } if let Some(res) = tasks.join_next().await { - let size = res??; - reservation.shrink(size); + match res.map_err(Error::from).and_then(|size| size) { + Ok(size) => reservation.shrink(size), + Err(e) => { + cleanup_on_failure( + &dataset, + &target_bases_info, + &new_fragments, + &mut tasks, + ) + .await; + return Err(e); + } + } } } match frag_id.first() { Some(ScalarValue::UInt64(Some(frag_id))) => { let frag_id = *frag_id; - let fragment = dataset.get_fragment(frag_id as usize).ok_or_else(|| { + let Some(fragment) = dataset.get_fragment(frag_id as usize) else { error!( fragment_id = frag_id, dataset_uri = %dataset.uri(), @@ -1317,15 +1492,22 @@ impl MergeInsertJob { branch = ?dataset.manifest().branch, "Non-existent fragment id returned from merge result", ); - Error::internal(format!( + cleanup_on_failure( + &dataset, + &target_bases_info, + &new_fragments, + &mut tasks, + ) + .await; + return Err(Error::internal(format!( "Got non-existent fragment id from merge result: {} (uri={}, version={}, manifest={}, branch={})", frag_id, dataset.uri(), dataset.manifest().version, dataset.manifest_location().path, dataset.manifest().branch.as_deref().unwrap_or("main"), - )) - })?; + ))); + }; let metadata = fragment.metadata.clone(); let fut = handle_fragment( @@ -1345,10 +1527,13 @@ impl MergeInsertJob { batches, new_fragments.clone(), memory_size, + target_bases_info.clone(), ); tasks.spawn(fut); } _ => { + cleanup_on_failure(&dataset, &target_bases_info, &new_fragments, &mut tasks) + .await; return Err(Error::internal(format!( "Got non-fragment id from merge result: {:?}", frag_id @@ -1358,8 +1543,14 @@ impl MergeInsertJob { } while let Some(res) = tasks.join_next().await { - let size = res??; - reservation.shrink(size); + match res.map_err(Error::from).and_then(|size| size) { + Ok(size) => reservation.shrink(size), + Err(e) => { + cleanup_on_failure(&dataset, &target_bases_info, &new_fragments, &mut tasks) + .await; + return Err(e); + } + } } let mut updated_fragments = Arc::try_unwrap(updated_fragments) .unwrap() @@ -1691,7 +1882,17 @@ impl MergeInsertJob { } } + // A partial-schema source that both deletes matched rows and inserts + // unmatched rows cannot be expressed by the indexed-scan delete path + // (the delete cannot be folded into a partial write). Keep it off the + // scalar-index route so it falls through to the v2 plan, which handles + // delete + insert directly. + let is_partial_delete_with_insert = is_subset_schema + && self.params.insert_not_matched + && matches!(self.params.when_matched, WhenMatched::Delete); + let would_use_scalar_index = if self.params.use_index + && !is_partial_delete_with_insert && matches!( self.params.delete_not_matched_by_source, WhenNotMatchedBySource::Keep @@ -1757,6 +1958,8 @@ impl MergeInsertJob { }); } + let target_bases_info = resolve_target_bases(&self.dataset, &self.params).await?; + let source_schema = source.schema(); let lance_schema = lance_core::datatypes::Schema::try_from(source_schema.as_ref())?; let full_schema = self.dataset.schema(); @@ -1785,7 +1988,66 @@ impl MergeInsertJob { .try_flatten(); let stream = RecordBatchStreamAdapter::new(merger_schema, stream); - let (operation, affected_rows) = if !is_full_schema { + // A partial-schema source can patch columns or delete matched rows, but + // not both in one commit: the writer rejects subschema rows up front and + // the delete cannot be folded into the write. Reject that combination. + if !is_full_schema + && matches!(self.params.when_matched, WhenMatched::Delete) + && self.params.insert_not_matched + { + return Err(Error::not_supported_source("Combining when_matched(Delete) with inserts from a partial-schema source is not supported; provide the full target schema in the source".into())); + } + + // The commit strategy follows what the merge does to matched rows. A + // pure delete (no inserts) writes nothing: the merger emits no batches + // and only records the matched row ids. This holds for any source schema + // width, so it is keyed on the operation rather than `is_full_schema`. + let is_delete_only = matches!(self.params.when_matched, WhenMatched::Delete) + && !self.params.insert_not_matched; + + let (operation, affected_rows) = if is_delete_only { + // Consume the stream so the merger records the matched row ids in + // `deleted_rows`; it produces no batches. + let drained: Vec = Box::pin(stream).try_collect().await?; + debug_assert!(drained.is_empty(), "delete-only merge must not emit rows"); + + let removed_row_ids = Arc::into_inner(deleted_rows).unwrap().into_inner().unwrap(); + let removed_row_addr_vec = + if let Some(row_id_index) = get_row_id_index(&self.dataset).await? { + removed_row_ids + .iter() + .filter_map(|id| row_id_index.get(*id).map(|address| address.into())) + .collect::>() + } else { + removed_row_ids + }; + let removed_row_addrs = RoaringTreemap::from_iter(removed_row_addr_vec); + + let (updated_fragments, removed_fragment_ids) = + Self::apply_deletions(&self.dataset, &removed_row_addrs).await?; + + let operation = Operation::Update { + removed_fragment_ids, + updated_fragments, + new_fragments: vec![], + fields_modified: vec![], + merged_generations: self.params.merged_generations.clone(), + fields_for_preserving_frag_bitmap: full_schema + .fields + .iter() + .map(|f| f.id as u32) + .collect(), + update_mode: Some(RewriteRows), + inserted_rows_filter: None, // not implemented for v1 + updated_fragment_offsets: None, + }; + + let affected_rows = Some(RowAddrTreeMap::from(removed_row_addrs)); + (operation, affected_rows) + } else if !is_full_schema { + // Non-delete partial-schema merge: patch the provided columns into + // existing fragments in place. (Delete is handled above; a wider + // full source takes the row-rewrite branch below.) if !matches!( self.params.delete_not_matched_by_source, WhenNotMatchedBySource::Keep @@ -1799,6 +2061,7 @@ impl MergeInsertJob { self.dataset.clone(), Box::pin(stream), self.dataset.manifest.version + 1, + target_bases_info, ) .await?; @@ -1817,6 +2080,7 @@ impl MergeInsertJob { // we can't use affected rows here. (operation, None) } else { + let cleanup_bases = target_bases_info.clone(); let (mut new_fragments, _) = write_fragments_internal( Some(&self.dataset), self.dataset.object_store.clone(), @@ -1824,48 +2088,66 @@ impl MergeInsertJob { self.dataset.schema().clone(), Box::pin(stream), WriteParams::default(), - None, // Merge insert doesn't use target_bases + target_bases_info, ) .await?; - if let Some(row_id_sequence) = updating_row_ids.lock().unwrap().row_id_sequence() { - let fragment_sizes = new_fragments - .iter() - .map(|f| f.physical_rows.unwrap() as u64); + // The new data files exist but are not committed yet; clean them up + // (including files routed to target bases) if any later step fails. + let post_write_result: Result = async { + if let Some(row_id_sequence) = updating_row_ids.lock().unwrap().row_id_sequence() { + let fragment_sizes = new_fragments + .iter() + .map(|f| f.physical_rows.unwrap() as u64); - let sequences = lance_table::rowids::rechunk_sequences( - [row_id_sequence.clone()], - fragment_sizes, - true, - ) - .map_err(|e| { - Error::internal(format!( - "Captured row ids not equal to number of rows written: {}", - e - )) - })?; + let sequences = lance_table::rowids::rechunk_sequences( + [row_id_sequence.clone()], + fragment_sizes, + true, + ) + .map_err(|e| { + Error::internal(format!( + "Captured row ids not equal to number of rows written: {}", + e + )) + })?; - for (fragment, sequence) in new_fragments.iter_mut().zip(sequences) { - let serialized = lance_table::rowids::write_row_ids(&sequence); - fragment.row_id_meta = Some(RowIdMeta::Inline(serialized)); + for (fragment, sequence) in new_fragments.iter_mut().zip(sequences) { + let serialized = lance_table::rowids::write_row_ids(&sequence); + fragment.row_id_meta = Some(RowIdMeta::Inline(serialized)); + } } - } - // Apply deletions - let removed_row_ids = Arc::into_inner(deleted_rows).unwrap().into_inner().unwrap(); + // Apply deletions + let removed_row_ids = Arc::into_inner(deleted_rows).unwrap().into_inner().unwrap(); - let removed_row_addr_vec = - if let Some(row_id_index) = get_row_id_index(&self.dataset).await? { - let addresses: Vec = removed_row_ids - .iter() - .filter_map(|id| row_id_index.get(*id).map(|address| address.into())) - .collect::>(); - addresses - } else { - removed_row_ids - }; + let removed_row_addr_vec = + if let Some(row_id_index) = get_row_id_index(&self.dataset).await? { + let addresses: Vec = removed_row_ids + .iter() + .filter_map(|id| row_id_index.get(*id).map(|address| address.into())) + .collect::>(); + addresses + } else { + removed_row_ids + }; - let removed_row_addrs = RoaringTreemap::from_iter(removed_row_addr_vec.into_iter()); + Ok(RoaringTreemap::from_iter(removed_row_addr_vec)) + } + .await; + let removed_row_addrs = match post_write_result { + Ok(removed_row_addrs) => removed_row_addrs, + Err(e) => { + cleanup_data_fragments( + &self.dataset.object_store, + &self.dataset.base, + cleanup_bases.as_deref(), + &new_fragments, + ) + .await; + return Err(e); + } + }; let deletions_result = Self::apply_deletions(&self.dataset, &removed_row_addrs).await; let (old_fragments, removed_fragment_ids) = match deletions_result { @@ -1874,6 +2156,7 @@ impl MergeInsertJob { cleanup_data_fragments( &self.dataset.object_store, &self.dataset.base, + cleanup_bases.as_deref(), &new_fragments, ) .await; @@ -2111,6 +2394,10 @@ impl RetryExecutor for MergeInsertJobWithIterator { // Update stats with the current attempt count data.stats.num_attempts = self.attempt_count.load(Ordering::SeqCst); + // The dataset argument is the refreshed per-attempt dataset (the same + // manifest execute_impl resolved against); keep a handle so conflict + // cleanup resolves bases added between attempts. + let cleanup_dataset = dataset.clone(); let mut commit_builder = CommitBuilder::new(dataset).with_skip_auto_cleanup(self.job.params.skip_auto_cleanup); if let Some(commit_retries) = self.job.params.commit_retries { @@ -2119,9 +2406,36 @@ impl RetryExecutor for MergeInsertJobWithIterator { if let Some(affected_rows) = data.affected_rows { commit_builder = commit_builder.with_affected_rows(affected_rows); } - let new_dataset = commit_builder.execute(data.transaction).await?; - Ok((Arc::new(new_dataset), data.stats)) + let new_fragments = match &data.transaction.operation { + Operation::Update { new_fragments, .. } => new_fragments.clone(), + _ => Vec::new(), + }; + match commit_builder.execute(data.transaction).await { + Ok(new_dataset) => Ok((Arc::new(new_dataset), data.stats)), + Err(e) => { + // A retryable conflict discards this attempt and re-executes it, + // so its data files are provably uncommitted; remove them + // (including files routed to target bases, which version cleanup + // never scans). Other commit errors may be ambiguous about + // whether the manifest was written, so leave the files alone. + if matches!(e, Error::RetryableCommitConflict { .. }) && !new_fragments.is_empty() { + let target_bases_info = + resolve_target_bases(&cleanup_dataset, &self.job.params) + .await + .ok() + .flatten(); + cleanup_data_fragments( + &cleanup_dataset.object_store, + &cleanup_dataset.base, + target_bases_info.as_deref(), + &new_fragments, + ) + .await; + } + Err(e) + } + } } fn update_dataset(&mut self, dataset: Arc) { @@ -2341,40 +2655,27 @@ impl Merger { // borrow checker (the stream needs to be `sync` since it crosses an await point) let mut deleted_row_ids = self.deleted_rows.lock().unwrap(); - if self.params.when_matched != WhenMatched::DoNothing { - let mut matched = arrow::compute::filter_record_batch(&batch, &in_both)?; - - if let Some(match_filter) = self.match_filter_expr { - let unzipped = unzip_batch(&matched, &self.schema); - let filtered = match_filter.evaluate(&unzipped)?; - match filtered { - ColumnarValue::Array(mask) => { - // Some rows matched, filter down and replace those rows - matched = arrow::compute::filter_record_batch(&matched, mask.as_boolean())?; - } - ColumnarValue::Scalar(scalar) => { - if let ScalarValue::Boolean(Some(true)) = scalar { - // All rows matched, go ahead and replace the whole batch - } else { - // Nothing matched, replace nothing - matched = RecordBatch::new_empty(matched.schema()); - } - } - } - } - - merge_statistics.num_updated_rows += matched.num_rows() as u64; - - // If the filter eliminated all rows then its important we don't try and write - // the batch at all. Writing an empty batch currently panics - if matched.num_rows() > 0 { + // Each `WhenMatched` variant handles the matched rows (`in_both`) + // differently. + let match_filter_expr = self.match_filter_expr; + match &self.params.when_matched { + WhenMatched::DoNothing => {} + WhenMatched::Delete => { + // Matched rows are removed, not rewritten: record their row ids + // for the commit to delete and emit no replacement batch. A + // source with duplicate keys matches the same target row more + // than once; apply the same `source_dedupe_behavior` policy as + // updates so a duplicate either aborts (`Fail`) or is skipped + // and counted once (`FirstSeen`) — the commit deletes the row a + // single time regardless. + let matched = arrow::compute::filter_record_batch(&batch, &in_both)?; let row_ids = matched.column(row_id_col).as_primitive::(); let mut processed_row_ids = self.processed_row_ids.lock().unwrap(); - let mut keep_indices: Vec = Vec::with_capacity(matched.num_rows()); for (row_idx, &row_id) in row_ids.values().iter().enumerate() { if processed_row_ids.insert(row_id) { - keep_indices.push(row_idx as u32); + merge_statistics.num_deleted_rows += 1; + deleted_row_ids.push(row_id); } else { match self.params.source_dedupe_behavior { SourceDedupeBehavior::Fail => { @@ -2385,54 +2686,115 @@ impl Merger { )); } SourceDedupeBehavior::FirstSeen => { - // Skip this duplicate row (don't add to keep_indices) + merge_statistics.num_skipped_duplicates += 1; } } } } - drop(processed_row_ids); - - // Filter out duplicate rows if any were skipped - let num_skipped = matched.num_rows() - keep_indices.len(); - if num_skipped > 0 { - merge_statistics.num_skipped_duplicates += num_skipped as u64; - merge_statistics.num_updated_rows -= num_skipped as u64; - - let indices = UInt32Array::from(keep_indices); - matched = take_record_batch(&matched, &indices)?; + } + WhenMatched::Fail => { + // Any matched row aborts the whole operation. + if let Some(row_idx) = (0..in_both.len()).find(|&i| in_both.value(i)) { + return Err(DataFusionError::Execution(format!( + "Merge insert failed: found matching row with key values: {}", + format_key_values_on_columns(&batch, row_idx, &self.params.on) + ))); + } + } + WhenMatched::UpdateAll | WhenMatched::UpdateIf(_) | WhenMatched::UpdateIfExpr(_) => { + let mut matched = arrow::compute::filter_record_batch(&batch, &in_both)?; + + if let Some(match_filter) = match_filter_expr { + let unzipped = unzip_batch(&matched, &self.schema); + let filtered = match_filter.evaluate(&unzipped)?; + match filtered { + ColumnarValue::Array(mask) => { + // Some rows matched, filter down and replace those rows + matched = + arrow::compute::filter_record_batch(&matched, mask.as_boolean())?; + } + ColumnarValue::Scalar(scalar) => { + if let ScalarValue::Boolean(Some(true)) = scalar { + // All rows matched, go ahead and replace the whole batch + } else { + // Nothing matched, replace nothing + matched = RecordBatch::new_empty(matched.schema()); + } + } + } } - // Only process and write if there are remaining rows after filtering duplicates + merge_statistics.num_updated_rows += matched.num_rows() as u64; + + // If the filter eliminated all rows then its important we don't try and write + // the batch at all. Writing an empty batch currently panics if matched.num_rows() > 0 { - // Get row_ids again after filtering (if any duplicates were removed) let row_ids = matched.column(row_id_col).as_primitive::(); - deleted_row_ids.extend(row_ids.values()); - if self.enable_stable_row_ids { - self.updating_row_ids - .lock() - .unwrap() - .capture(row_ids.values())?; + + let mut processed_row_ids = self.processed_row_ids.lock().unwrap(); + let mut keep_indices: Vec = Vec::with_capacity(matched.num_rows()); + for (row_idx, &row_id) in row_ids.values().iter().enumerate() { + if processed_row_ids.insert(row_id) { + keep_indices.push(row_idx as u32); + } else { + match self.params.source_dedupe_behavior { + SourceDedupeBehavior::Fail => { + return Err(create_duplicate_row_error( + &matched, + row_idx, + &self.params.on, + )); + } + SourceDedupeBehavior::FirstSeen => { + // Skip this duplicate row (don't add to keep_indices) + } + } + } } + drop(processed_row_ids); - let projection = if let Some(row_addr_col) = row_addr_col { - let mut cols = Vec::from_iter(left_cols.iter().cloned()); - cols.push(row_addr_col); - cols - } else { - #[allow(clippy::redundant_clone)] - left_cols.clone() - }; - let matched = matched.project(&projection)?; - // The payload columns of an outer join are always nullable. We need to restore - // non-nullable to columns that were originally non-nullable. This should be safe - // since the not_matched rows should all be valid on the right_cols - // - // Sadly we can't use with_schema because it doesn't let you toggle nullability - let matched = RecordBatch::try_new( - self.output_schema.clone(), - Vec::from_iter(matched.columns().iter().cloned()), - )?; - batches.push(Ok(matched)); + // Filter out duplicate rows if any were skipped + let num_skipped = matched.num_rows() - keep_indices.len(); + if num_skipped > 0 { + merge_statistics.num_skipped_duplicates += num_skipped as u64; + merge_statistics.num_updated_rows -= num_skipped as u64; + + let indices = UInt32Array::from(keep_indices); + matched = take_record_batch(&matched, &indices)?; + } + + // Only process and write if there are remaining rows after filtering duplicates + if matched.num_rows() > 0 { + // Get row_ids again after filtering (if any duplicates were removed) + let row_ids = matched.column(row_id_col).as_primitive::(); + deleted_row_ids.extend(row_ids.values()); + if self.enable_stable_row_ids { + self.updating_row_ids + .lock() + .unwrap() + .capture(row_ids.values())?; + } + + let projection = if let Some(row_addr_col) = row_addr_col { + let mut cols = Vec::from_iter(left_cols.iter().cloned()); + cols.push(row_addr_col); + cols + } else { + #[allow(clippy::redundant_clone)] + left_cols.clone() + }; + let matched = matched.project(&projection)?; + // The payload columns of an outer join are always nullable. We need to restore + // non-nullable to columns that were originally non-nullable. This should be safe + // since the not_matched rows should all be valid on the right_cols + // + // Sadly we can't use with_schema because it doesn't let you toggle nullability + let matched = RecordBatch::try_new( + self.output_schema.clone(), + Vec::from_iter(matched.columns().iter().cloned()), + )?; + batches.push(Ok(matched)); + } } } } @@ -2541,6 +2903,127 @@ mod tests { t } + // An update-style merge_insert leaves the source and new fragments with + // overlapping id ranges; a scattered delete punches holes in that range. A + // filtered `with_row_id` scan must still resolve every id (round-tripped via take). + #[tokio::test(flavor = "multi_thread")] + async fn merge_insert_then_delete_resolves_overlapping_row_ids() { + use arrow_array::ArrayRef; + use arrow_array::types::UInt64Type; + let dir = TempStrDir::default(); + let uri = dir.as_str(); + let schema = Arc::new(Schema::new(vec![ + Field::new("slug", DataType::Utf8, false), + Field::new("title", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let mk = |slugs: Vec, titles: Vec| { + let cats: Vec = (0..slugs.len()) + .map(|i| ["A", "B", "C", "D", "E"][i % 5].to_string()) + .collect(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(slugs)) as ArrayRef, + Arc::new(StringArray::from(titles)) as ArrayRef, + Arc::new(StringArray::from(cats)) as ArrayRef, + ], + ) + .unwrap() + }; + + // Empty dataset with stable row ids; write-seed 40 rows (ids 0..40). + let params = WriteParams { + mode: WriteMode::Create, + enable_stable_row_ids: true, + ..Default::default() + }; + let mut ds = Dataset::write( + RecordBatchIterator::new(vec![Ok(mk(vec![], vec![]))], schema.clone()), + uri, + Some(params), + ) + .await + .unwrap(); + ds.append( + RecordBatchIterator::new( + vec![Ok(mk( + (1..=40).map(|i| format!("t{i}")).collect(), + (1..=40).map(|i| format!("r{i}")).collect(), + ))], + schema.clone(), + ), + None, + ) + .await + .unwrap(); + + // Update every other row (so the new fragment's ids interleave with -- and + // its range overlaps -- the source's) plus a few inserts. + let mut slugs: Vec = (1..=40).step_by(2).map(|i| format!("t{i}")).collect(); + let mut titles: Vec = (1..=40).step_by(2).map(|i| format!("e{i}")).collect(); + for i in 41..=45 { + slugs.push(format!("t{i}")); + titles.push(format!("e{i}")); + } + let mut b = MergeInsertBuilder::try_new(Arc::new(ds), vec!["slug".into()]).unwrap(); + b.when_matched(WhenMatched::UpdateAll); + b.when_not_matched(WhenNotMatched::InsertAll); + let (ds, _) = b + .try_build() + .unwrap() + .execute_reader(RecordBatchIterator::new( + vec![Ok(mk(slugs, titles))], + schema.clone(), + )) + .await + .unwrap(); + + // Scattered delete -> interior holes in the overlapped id range. + let mut ds = (*ds).clone(); + let ds = (*ds.delete("category = 'A'").await.unwrap().new_dataset).clone(); + + // The `with_row_id` scan builds the RowIdIndex; every scanned id must + // round-trip through take_rows. + let batches: Vec = ds + .scan() + .with_row_id() + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let row_ids: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name(ROW_ID) + .unwrap() + .as_primitive::() + .values() + .to_vec() + }) + .collect(); + let scanned_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + let taken = ds.take_rows(&row_ids, ds.schema().clone()).await.unwrap(); + assert_eq!(taken.num_rows(), scanned_rows); + + // A point lookup on a surviving row resolves to exactly one row. + let filtered: Vec = ds + .scan() + .with_row_id() + .filter("slug = 't30'") + .unwrap() + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let n: usize = filtered.iter().map(|b| b.num_rows()).sum(); + assert_eq!(n, 1, "expected exactly one row for slug='t30'"); + } + async fn check_then_refresh_dataset( new_data: RecordBatch, mut job: MergeInsertJob, @@ -3905,72 +4388,658 @@ mod tests { ) .unwrap(); - let (updated_ds, stats) = - MergeInsertBuilder::try_new(Arc::new(ds), vec!["a".to_string(), "b".to_string()]) - .unwrap() - .when_matched(WhenMatched::UpdateAll) - .when_not_matched(WhenNotMatched::InsertAll) - .try_build() - .unwrap() - .execute_reader(Box::new(RecordBatchIterator::new( - vec![Ok(source.clone())], - source.schema(), - ))) - .await - .unwrap(); + let (updated_ds, stats) = + MergeInsertBuilder::try_new(Arc::new(ds), vec!["a".to_string(), "b".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new( + vec![Ok(source.clone())], + source.schema(), + ))) + .await + .unwrap(); + + assert_eq!(stats.num_updated_rows, 1); + assert_eq!(stats.num_inserted_rows, 0); + let count = updated_ds + .count_rows(Some("a = 1 AND b = 20 AND value = 999".to_string())) + .await + .unwrap(); + assert_eq!(count, 1); + } + + /// Composite-key merge_insert must use standard SQL NULL semantics + /// (NULL != NULL) on the post-filter hash join so its behavior is + /// identical to the full-scan path; otherwise enabling the indexed + /// path for multi-column joins would silently change semantics. + #[tokio::test] + async fn test_indexed_merge_insert_composite_key_null_semantics() { + let initial = record_batch!( + ("a", Int32, [Some(1)]), + ("b", Utf8, [Option::<&str>::None]), + ("value", Int32, [Some(10)]) + ) + .unwrap(); + let schema = initial.schema(); + + let mut ds = Dataset::write( + RecordBatchIterator::new(vec![Ok(initial.clone())], schema.clone()), + "memory://", + None, + ) + .await + .unwrap(); + + ds.create_index( + &["a"], + IndexType::Scalar, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + + let source = record_batch!( + ("a", Int32, [Some(1)]), + ("b", Utf8, [Option::<&str>::None]), + ("value", Int32, [Some(99)]) + ) + .unwrap(); + + let (updated_ds, stats) = + MergeInsertBuilder::try_new(Arc::new(ds), vec!["a".to_string(), "b".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new( + vec![Ok(source.clone())], + source.schema(), + ))) + .await + .unwrap(); + + assert_eq!(stats.num_inserted_rows, 1); + assert_eq!(stats.num_updated_rows, 0); + assert_eq!(updated_ds.count_rows(None).await.unwrap(), 2); + } + + /// Composite-key merge_insert where new (unindexed) fragments are + /// appended after the indices were built. The indexed take only sees + /// fragments covered by every chosen index, so the unindexed remainder + /// must be unioned in via a full scan — otherwise updates to rows + /// that live in those fragments are silently dropped. + #[tokio::test] + async fn test_indexed_merge_insert_composite_key_unindexed_fragments() { + let first = record_batch!( + ("a", Int32, [1, 2]), + ("b", Int32, [10, 20]), + ("value", Int32, [100, 200]) + ) + .unwrap(); + let schema = first.schema(); + + let mut ds = Dataset::write( + RecordBatchIterator::new(vec![Ok(first.clone())], schema.clone()), + "memory://", + Some(WriteParams { + max_rows_per_file: 64, + ..Default::default() + }), + ) + .await + .unwrap(); + + let params = ScalarIndexParams::default(); + ds.create_index(&["a"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + ds.create_index(&["b"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + + // Append a fragment AFTER both indices are built. The new (3, 30) + // row lives in a fragment neither index covers, so the indexed + // take alone would miss it. + let appended = record_batch!( + ("a", Int32, [3]), + ("b", Int32, [30]), + ("value", Int32, [300]) + ) + .unwrap(); + ds.append( + RecordBatchIterator::new(vec![Ok(appended.clone())], appended.schema()), + None, + ) + .await + .unwrap(); + + // Source updates one row in the indexed fragment AND one row in + // the appended (unindexed) fragment. + let source = record_batch!( + ("a", Int32, [1, 3]), + ("b", Int32, [10, 30]), + ("value", Int32, [999, 333]) + ) + .unwrap(); + + let (updated_ds, stats) = + MergeInsertBuilder::try_new(Arc::new(ds), vec!["a".to_string(), "b".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new( + vec![Ok(source.clone())], + source.schema(), + ))) + .await + .unwrap(); + + assert_eq!( + stats.num_updated_rows, 2, + "row in the unindexed fragment must also be updated" + ); + assert_eq!(stats.num_inserted_rows, 0); + assert_eq!(updated_ds.count_rows(None).await.unwrap(), 3); + } + + /// Composite-key delete-only merge_insert (`when_matched(Delete)`, + /// `when_not_matched_by_source(Keep)`) removes the matched rows for every + /// combination of which join columns carry a scalar index, including when + /// every column is indexed. + #[rstest::rstest] + #[case::index_on_both(true, true)] + #[case::index_on_first(true, false)] + #[case::index_on_second(false, true)] + #[case::no_index(false, false)] + #[tokio::test] + async fn test_indexed_merge_insert_composite_key_delete( + #[case] index_on_a: bool, + #[case] index_on_b: bool, + ) { + let initial = record_batch!( + ("a", Int32, [1, 1, 2, 2]), + ("b", Int32, [10, 20, 10, 20]), + ("value", Int32, [100, 200, 300, 400]) + ) + .unwrap(); + let schema = initial.schema(); + + let mut ds = Dataset::write( + RecordBatchIterator::new(vec![Ok(initial.clone())], schema.clone()), + "memory://", + None, + ) + .await + .unwrap(); + + let params = ScalarIndexParams::default(); + if index_on_a { + ds.create_index(&["a"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + } + if index_on_b { + ds.create_index(&["b"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + } + + // Delete (1, 10) by composite key. Only key columns in the source. + let source = record_batch!(("a", Int32, [1]), ("b", Int32, [10])).unwrap(); + + let (updated_ds, stats) = + MergeInsertBuilder::try_new(Arc::new(ds), vec!["a".to_string(), "b".to_string()]) + .unwrap() + .when_matched(WhenMatched::Delete) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new( + vec![Ok(source.clone())], + source.schema(), + ))) + .await + .unwrap(); + + assert_eq!(stats.num_deleted_rows, 1, "matched row must be deleted"); + assert_eq!(updated_ds.count_rows(None).await.unwrap(), 3); + assert_eq!( + updated_ds + .count_rows(Some("a = 1 AND b = 10".to_string())) + .await + .unwrap(), + 0, + "(1, 10) must be gone after the delete" + ); + // The sibling key (1, 20) must remain untouched. + assert_eq!( + updated_ds + .count_rows(Some("a = 1 AND b = 20 AND value = 200".to_string())) + .await + .unwrap(), + 1, + ); + } + + /// Delete-only merge_insert on a single indexed key removes the matched + /// rows (the indexed path is not exclusive to composite keys). + #[tokio::test] + async fn test_indexed_merge_insert_single_key_delete() { + let initial = record_batch!( + ("id", Int32, [1, 2, 3, 4]), + ("value", Int32, [10, 20, 30, 40]) + ) + .unwrap(); + let schema = initial.schema(); + + let mut ds = Dataset::write( + RecordBatchIterator::new(vec![Ok(initial.clone())], schema.clone()), + "memory://", + None, + ) + .await + .unwrap(); + + ds.create_index( + &["id"], + IndexType::Scalar, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + + let source = record_batch!(("id", Int32, [2, 4])).unwrap(); + + let (updated_ds, stats) = MergeInsertBuilder::try_new(Arc::new(ds), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::Delete) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new( + vec![Ok(source.clone())], + source.schema(), + ))) + .await + .unwrap(); + + assert_eq!(stats.num_deleted_rows, 2); + assert_eq!(updated_ds.count_rows(None).await.unwrap(), 2); + assert_eq!( + updated_ds + .count_rows(Some("id = 2 OR id = 4".to_string())) + .await + .unwrap(), + 0, + ); + } + + /// `when_matched(Fail)` on an indexed key aborts the operation when a + /// source row matches an existing key, and inserts cleanly when none do. + #[tokio::test] + async fn test_indexed_merge_insert_when_matched_fail() { + let initial = + record_batch!(("id", Int32, [1, 2, 3]), ("value", Int32, [10, 20, 30])).unwrap(); + let schema = initial.schema(); + + let mut ds = Dataset::write( + RecordBatchIterator::new(vec![Ok(initial.clone())], schema.clone()), + "memory://", + None, + ) + .await + .unwrap(); + + ds.create_index( + &["id"], + IndexType::Scalar, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let ds = Arc::new(ds); + + // A source row matching an existing key must fail the operation. + let matching = record_batch!(("id", Int32, [2]), ("value", Int32, [999])).unwrap(); + let err = MergeInsertBuilder::try_new(ds.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::Fail) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new( + vec![Ok(matching.clone())], + matching.schema(), + ))) + .await + .unwrap_err(); + let msg = err.to_string(); + assert!(msg.contains("Merge insert failed"), "got: {msg}"); + assert!(msg.contains("found matching row"), "got: {msg}"); + + // A source with no matching key inserts without failing. + let new_rows = record_batch!(("id", Int32, [4]), ("value", Int32, [40])).unwrap(); + let (updated_ds, stats) = MergeInsertBuilder::try_new(ds.clone(), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::Fail) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new( + vec![Ok(new_rows.clone())], + new_rows.schema(), + ))) + .await + .unwrap(); + assert_eq!(stats.num_inserted_rows, 1); + assert_eq!(stats.num_updated_rows, 0); + assert_eq!(updated_ds.count_rows(None).await.unwrap(), 4); + } + + /// Fully-indexed composite-key `when_matched(Delete)` combined with + /// `when_not_matched(InsertAll)` must both delete matched rows and write + /// the inserted rows. + #[tokio::test] + async fn test_indexed_merge_insert_composite_key_delete_with_insert() { + let initial = record_batch!( + ("a", Int32, [1, 1, 2, 2]), + ("b", Int32, [10, 20, 10, 20]), + ("value", Int32, [100, 200, 300, 400]) + ) + .unwrap(); + let schema = initial.schema(); + + let mut ds = Dataset::write( + RecordBatchIterator::new(vec![Ok(initial.clone())], schema.clone()), + "memory://", + None, + ) + .await + .unwrap(); + + let params = ScalarIndexParams::default(); + ds.create_index(&["a"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + ds.create_index(&["b"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + + // Source matches (1, 10) -> delete, and (3, 30) is new -> insert. + let source = record_batch!( + ("a", Int32, [1, 3]), + ("b", Int32, [10, 30]), + ("value", Int32, [999, 333]) + ) + .unwrap(); + + let (updated_ds, stats) = + MergeInsertBuilder::try_new(Arc::new(ds), vec!["a".to_string(), "b".to_string()]) + .unwrap() + .when_matched(WhenMatched::Delete) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new( + vec![Ok(source.clone())], + source.schema(), + ))) + .await + .unwrap(); + + assert_eq!(stats.num_deleted_rows, 1); + assert_eq!(stats.num_inserted_rows, 1); + // 4 - 1 deleted + 1 inserted = 4. + assert_eq!(updated_ds.count_rows(None).await.unwrap(), 4); + assert_eq!( + updated_ds + .count_rows(Some("a = 1 AND b = 10".to_string())) + .await + .unwrap(), + 0, + "matched row must be deleted, not updated" + ); + assert_eq!( + updated_ds + .count_rows(Some("a = 3 AND b = 30 AND value = 333".to_string())) + .await + .unwrap(), + 1, + "unmatched source row must be inserted" + ); + } + + /// A delete whose source contains duplicate keys matching the same target + /// row applies `source_dedupe_behavior` on the indexed-scan path, exactly + /// like an update: the default `Fail` aborts (naming the ambiguous key), + /// while `FirstSeen` removes and counts the row once and reports the extra + /// match as a skipped duplicate. + #[rstest::rstest] + #[case::fail(SourceDedupeBehavior::Fail)] + #[case::first_seen(SourceDedupeBehavior::FirstSeen)] + #[tokio::test] + async fn test_indexed_merge_insert_delete_source_duplicates( + #[case] behavior: SourceDedupeBehavior, + ) { + let initial = record_batch!( + ("a", Int32, [1, 1, 2, 2]), + ("b", Int32, [10, 20, 10, 20]), + ("value", Int32, [100, 200, 300, 400]) + ) + .unwrap(); + let schema = initial.schema(); + + let mut ds = Dataset::write( + RecordBatchIterator::new(vec![Ok(initial.clone())], schema.clone()), + "memory://", + None, + ) + .await + .unwrap(); + + // Index every join column so the merge takes the indexed-scan delete path. + let params = ScalarIndexParams::default(); + ds.create_index(&["a"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + ds.create_index(&["b"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + + // Two source rows collide on the same target key (1, 10). + let source = record_batch!(("a", Int32, [1, 1]), ("b", Int32, [10, 10])).unwrap(); + + let result = + MergeInsertBuilder::try_new(Arc::new(ds), vec!["a".to_string(), "b".to_string()]) + .unwrap() + .when_matched(WhenMatched::Delete) + .when_not_matched(WhenNotMatched::DoNothing) + .source_dedupe_behavior(behavior) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new( + vec![Ok(source.clone())], + source.schema(), + ))) + .await; + + if behavior == SourceDedupeBehavior::Fail { + let err = result.unwrap_err().to_string(); + assert!( + err.contains("Ambiguous merge inserts") && err.contains("a = 1"), + "Fail must abort naming the ambiguous key, got: {err}" + ); + return; + } + + let (updated_ds, stats) = result.unwrap(); + assert_eq!(stats.num_deleted_rows, 1); + assert_eq!(stats.num_skipped_duplicates, 1); + assert_eq!(updated_ds.count_rows(None).await.unwrap(), 3); + assert_eq!( + updated_ds + .count_rows(Some("a = 1 AND b = 10".to_string())) + .await + .unwrap(), + 0, + "the matched row must be removed exactly once" + ); + } + + /// The v2 plans apply the same `source_dedupe_behavior` to deletes when the + /// source has duplicate keys matching one target row — covering both + /// `FullSchemaMergeInsertExec` (`Delete + InsertAll`) and + /// `DeleteOnlyMergeInsertExec` (pure delete). No scalar index, so routing + /// stays on the v2 path. + #[rstest::rstest] + #[case::full_schema_fail(true, SourceDedupeBehavior::Fail)] + #[case::full_schema_first_seen(true, SourceDedupeBehavior::FirstSeen)] + #[case::delete_only_fail(false, SourceDedupeBehavior::Fail)] + #[case::delete_only_first_seen(false, SourceDedupeBehavior::FirstSeen)] + #[tokio::test] + async fn test_v2_merge_insert_delete_source_duplicates( + #[case] with_insert: bool, + #[case] behavior: SourceDedupeBehavior, + ) { + let initial = + record_batch!(("a", Int32, [1, 2, 3]), ("value", Int32, [10, 20, 30])).unwrap(); + let schema = initial.schema(); + + let ds = Dataset::write( + RecordBatchIterator::new(vec![Ok(initial.clone())], schema.clone()), + "memory://", + None, + ) + .await + .unwrap(); + + // Two source rows collide on target key a=1. With insert, a=4 is new. + let (source, when_not_matched, expected_inserted, expected_total) = if with_insert { + ( + record_batch!(("a", Int32, [1, 1, 4]), ("value", Int32, [99, 99, 40])).unwrap(), + WhenNotMatched::InsertAll, + 1, + 3, // 3 - 1 deleted + 1 inserted + ) + } else { + ( + record_batch!(("a", Int32, [1, 1])).unwrap(), + WhenNotMatched::DoNothing, + 0, + 2, // 3 - 1 deleted + ) + }; + + let result = MergeInsertBuilder::try_new(Arc::new(ds), vec!["a".to_string()]) + .unwrap() + .when_matched(WhenMatched::Delete) + .when_not_matched(when_not_matched) + .source_dedupe_behavior(behavior) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new( + vec![Ok(source.clone())], + source.schema(), + ))) + .await; + + if behavior == SourceDedupeBehavior::Fail { + let err = result.unwrap_err().to_string(); + assert!( + err.contains("Ambiguous merge inserts") && err.contains("a = 1"), + "Fail must abort naming the ambiguous key, got: {err}" + ); + return; + } - assert_eq!(stats.num_updated_rows, 1); - assert_eq!(stats.num_inserted_rows, 0); - let count = updated_ds - .count_rows(Some("a = 1 AND b = 20 AND value = 999".to_string())) - .await - .unwrap(); - assert_eq!(count, 1); + let (updated_ds, stats) = result.unwrap(); + assert_eq!(stats.num_deleted_rows, 1, "the matched row is removed once"); + assert_eq!(stats.num_skipped_duplicates, 1); + assert_eq!(stats.num_inserted_rows, expected_inserted); + assert_eq!(updated_ds.count_rows(None).await.unwrap(), expected_total); + assert_eq!( + updated_ds + .count_rows(Some("a = 1".to_string())) + .await + .unwrap(), + 0, + "the matched row must be removed exactly once" + ); } - /// Composite-key merge_insert must use standard SQL NULL semantics - /// (NULL != NULL) on the post-filter hash join so its behavior is - /// identical to the full-scan path; otherwise enabling the indexed - /// path for multi-column joins would silently change semantics. + /// A partial-schema source that combines `when_matched(Delete)` with + /// `when_not_matched(InsertAll)` must succeed even when every join key is + /// indexed. The indexed-scan delete path cannot fold a delete into a + /// partial write, so this case routes to the v2 plan (which fills omitted + /// nullable target columns) instead of being rejected. #[tokio::test] - async fn test_indexed_merge_insert_composite_key_null_semantics() { - let initial = record_batch!( - ("a", Int32, [Some(1)]), - ("b", Utf8, [Option::<&str>::None]), - ("value", Int32, [Some(10)]) + async fn test_indexed_merge_insert_partial_schema_delete_with_insert() { + // Target carries two nullable non-key columns; the source omits `note`. + let full_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("value", DataType::Int32, true), + Field::new("note", DataType::Utf8, true), + ])); + let full_batch = RecordBatch::try_new( + full_schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 2])), + Arc::new(Int32Array::from(vec![10, 20, 10, 20])), + Arc::new(Int32Array::from(vec![100, 200, 300, 400])), + Arc::new(StringArray::from(vec!["w", "x", "y", "z"])), + ], ) .unwrap(); - let schema = initial.schema(); let mut ds = Dataset::write( - RecordBatchIterator::new(vec![Ok(initial.clone())], schema.clone()), + RecordBatchIterator::new(vec![Ok(full_batch)], full_schema.clone()), "memory://", None, ) .await .unwrap(); - ds.create_index( - &["a"], - IndexType::Scalar, - None, - &ScalarIndexParams::default(), - false, - ) - .await - .unwrap(); + let params = ScalarIndexParams::default(); + ds.create_index(&["a"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + ds.create_index(&["b"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); - let source = record_batch!( - ("a", Int32, [Some(1)]), - ("b", Utf8, [Option::<&str>::None]), - ("value", Int32, [Some(99)]) + // Source deletes matched (1, 10) and inserts new (3, 30), omitting `note`. + let partial_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("value", DataType::Int32, true), + ])); + let source = RecordBatch::try_new( + partial_schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 3])), + Arc::new(Int32Array::from(vec![10, 30])), + Arc::new(Int32Array::from(vec![999, 333])), + ], ) .unwrap(); let (updated_ds, stats) = MergeInsertBuilder::try_new(Arc::new(ds), vec!["a".to_string(), "b".to_string()]) .unwrap() - .when_matched(WhenMatched::UpdateAll) + .when_matched(WhenMatched::Delete) .when_not_matched(WhenNotMatched::InsertAll) .try_build() .unwrap() @@ -3981,31 +5050,59 @@ mod tests { .await .unwrap(); + assert_eq!(stats.num_deleted_rows, 1); assert_eq!(stats.num_inserted_rows, 1); - assert_eq!(stats.num_updated_rows, 0); - assert_eq!(updated_ds.count_rows(None).await.unwrap(), 2); + // 4 - 1 deleted + 1 inserted = 4. + assert_eq!(updated_ds.count_rows(None).await.unwrap(), 4); + assert_eq!( + updated_ds + .count_rows(Some("a = 1 AND b = 10".to_string())) + .await + .unwrap(), + 0, + "matched row must be deleted, not updated" + ); + // Inserted row carries the omitted `note` column as NULL. + assert_eq!( + updated_ds + .count_rows(Some( + "a = 3 AND b = 30 AND value = 333 AND note IS NULL".to_string() + )) + .await + .unwrap(), + 1, + "unmatched source row must be inserted with omitted column NULL-filled" + ); } - /// Composite-key merge_insert where new (unindexed) fragments are - /// appended after the indices were built. The indexed take only sees - /// fragments covered by every chosen index, so the unindexed remainder - /// must be unioned in via a full scan — otherwise updates to rows - /// that live in those fragments are silently dropped. + /// Fully-indexed composite-key delete across multiple fragments, with + /// stable row ids on/off. Exercises the indexed-scan delete commit + /// path: matched row ids are resolved to addresses (via the row-id + /// index when stable) and removed without rewriting any fragments. + /// Also covers an appended fragment that neither index covers, so the + /// delete must reach rows via the unindexed-remainder union too. + #[rstest::rstest] + #[case(true)] + #[case(false)] #[tokio::test] - async fn test_indexed_merge_insert_composite_key_unindexed_fragments() { - let first = record_batch!( - ("a", Int32, [1, 2]), - ("b", Int32, [10, 20]), - ("value", Int32, [100, 200]) + async fn test_indexed_merge_insert_composite_key_delete_multi_fragment( + #[case] enable_stable_row_ids: bool, + ) { + let initial = record_batch!( + ("a", Int32, [1, 1, 2, 2]), + ("b", Int32, [10, 20, 10, 20]), + ("value", Int32, [100, 200, 300, 400]) ) .unwrap(); - let schema = first.schema(); + let schema = initial.schema(); + // One row per fragment so the delete spans multiple fragments. let mut ds = Dataset::write( - RecordBatchIterator::new(vec![Ok(first.clone())], schema.clone()), + RecordBatchIterator::new(vec![Ok(initial.clone())], schema.clone()), "memory://", Some(WriteParams { - max_rows_per_file: 64, + max_rows_per_file: 1, + enable_stable_row_ids, ..Default::default() }), ) @@ -4020,13 +5117,12 @@ mod tests { .await .unwrap(); - // Append a fragment AFTER both indices are built. The new (3, 30) - // row lives in a fragment neither index covers, so the indexed - // take alone would miss it. + // Append a row AFTER the indices are built so it lives in a fragment + // neither index covers. The delete must still reach it. let appended = record_batch!( ("a", Int32, [3]), ("b", Int32, [30]), - ("value", Int32, [300]) + ("value", Int32, [500]) ) .unwrap(); ds.append( @@ -4036,20 +5132,14 @@ mod tests { .await .unwrap(); - // Source updates one row in the indexed fragment AND one row in - // the appended (unindexed) fragment. - let source = record_batch!( - ("a", Int32, [1, 3]), - ("b", Int32, [10, 30]), - ("value", Int32, [999, 333]) - ) - .unwrap(); + // Delete an indexed row (1, 10) and the unindexed appended row (3, 30). + let source = record_batch!(("a", Int32, [1, 3]), ("b", Int32, [10, 30])).unwrap(); let (updated_ds, stats) = MergeInsertBuilder::try_new(Arc::new(ds), vec!["a".to_string(), "b".to_string()]) .unwrap() - .when_matched(WhenMatched::UpdateAll) - .when_not_matched(WhenNotMatched::InsertAll) + .when_matched(WhenMatched::Delete) + .when_not_matched(WhenNotMatched::DoNothing) .try_build() .unwrap() .execute_reader(Box::new(RecordBatchIterator::new( @@ -4059,12 +5149,24 @@ mod tests { .await .unwrap(); + assert_eq!(stats.num_deleted_rows, 2); + assert_eq!(updated_ds.count_rows(None).await.unwrap(), 3); assert_eq!( - stats.num_updated_rows, 2, - "row in the unindexed fragment must also be updated" + updated_ds + .count_rows(Some("(a = 1 AND b = 10) OR (a = 3 AND b = 30)".to_string())) + .await + .unwrap(), + 0, + "both matched rows must be gone" + ); + // Untouched rows survive with their original values. + assert_eq!( + updated_ds + .count_rows(Some("a = 2 AND b = 20 AND value = 400".to_string())) + .await + .unwrap(), + 1, ); - assert_eq!(stats.num_inserted_rows, 0); - assert_eq!(updated_ds.count_rows(None).await.unwrap(), 3); } /// Composite-key `MapIndexExec` formats its Display so plans expose @@ -10065,4 +11167,285 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n "Newly written merge-insert data files should be cleaned up on apply_deletions failure" ); } + + #[tokio::test] + async fn test_merge_insert_full_fragment_rewrite_with_json_columns() { + // This test verifies the "all rows updated" fast path in handle_fragment + // correctly converts Arrow JSON (Utf8) to Lance JSON (LargeBinary/JSONB) + // before writing. Without conversion, the file would have Utf8 data (i32 + // offsets) but schema says LargeBinary (i64 offsets), causing decoder panic + // on subsequent reads. + // + // To trigger the fast path we need: + // 1. Subschema update (not all columns) → forces v1 update_fragments path + // 2. ALL rows in a fragment updated → triggers the fast path + use lance_arrow::ARROW_EXT_NAME_KEY; + use lance_arrow::json::{ARROW_JSON_EXT_NAME, is_arrow_json_field}; + + let test_dir = TempStrDir::default(); + let mut json_metadata = HashMap::new(); + json_metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + Field::new("score", DataType::Int64, true), + Field::new("meta", DataType::Utf8, true).with_metadata(json_metadata.clone()), + ])); + // Small fragment so ALL rows will be updated + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + Arc::new(Int64Array::from(vec![10, 20, 30])), + Arc::new(StringArray::from(vec![ + r#"{"x":1}"#, + r#"{"x":2}"#, + r#"{"x":3}"#, + ])), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let write_params = WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }; + let dataset = Arc::new( + Dataset::write(reader, test_dir.as_ref(), Some(write_params)) + .await + .unwrap(), + ); + assert_eq!(dataset.get_fragments().len(), 1); + + // Subschema update: only provide [id, meta] (missing "name" and "score") + // This forces the v1 path (update_fragments) instead of v2 (create_plan). + // Update ALL rows → triggers the "all rows updated" fast path. + let update_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("meta", DataType::Utf8, true).with_metadata(json_metadata), + ])); + let update_batch = RecordBatch::try_new( + update_schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])), // all rows + Arc::new(StringArray::from(vec![ + r#"{"updated":true,"id":1}"#, + r#"{"updated":true,"id":2}"#, + r#"{"updated":true,"id":3}"#, + ])), + ], + ) + .unwrap(); + let update_reader: Box = Box::new(RecordBatchIterator::new( + vec![Ok(update_batch)], + update_schema, + )); + let stream = reader_to_stream(update_reader); + + // Execute merge_insert with subschema + let mut builder = + MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]).unwrap(); + builder.when_matched(WhenMatched::UpdateAll); + builder.when_not_matched(WhenNotMatched::DoNothing); + let job = builder.try_build().unwrap(); + let (updated_dataset, stats) = job.execute(stream).await.unwrap(); + + assert_eq!(stats.num_updated_rows, 3); + + // Critical: read the data back. Without the fix, this would PANIC with: + // "the offset of the new Buffer cannot exceed the existing Length: + // slice offset=0 Length=N selfLen=N/2" + let batches = updated_dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let result = concat_batches(&batches[0].schema(), &batches).unwrap(); + assert_eq!(result.num_rows(), 3); + + // Verify JSON column is in Arrow JSON format (Utf8) on read + let result_schema = result.schema(); + let meta_field = result_schema.field_with_name("meta").unwrap(); + assert!( + is_arrow_json_field(meta_field), + "Expected Arrow JSON (Utf8 + arrow.json), got {:?}", + meta_field + ); + + // Verify data correctness + let metas = result + .column_by_name("meta") + .unwrap() + .as_any() + .downcast_ref::() + .expect("meta should be StringArray after read conversion"); + for i in 0..3 { + let val = metas.value(i); + assert!( + val.contains("updated"), + "row {} should have updated meta, got: {}", + i, + val + ); + } + + // Verify non-updated columns are preserved + let scores = result + .column_by_name("score") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(scores.values(), &[10, 20, 30]); + + // Also verify via take (exercises the take read conversion path) + let take_result = updated_dataset + .take(&[0, 1, 2], updated_dataset.schema().clone()) + .await + .unwrap(); + let take_schema = take_result.schema(); + let take_meta_field = take_schema.field_with_name("meta").unwrap(); + assert!( + is_arrow_json_field(take_meta_field), + "take() should return Arrow JSON, got {:?}", + take_meta_field + ); + } + + #[tokio::test] + async fn test_merge_insert_subschema_with_json_columns() { + use lance_arrow::ARROW_EXT_NAME_KEY; + use lance_arrow::json::ARROW_JSON_EXT_NAME; + + // Create a dataset with an Arrow JSON extension column + let test_dir = TempStrDir::default(); + let mut json_metadata = HashMap::new(); + json_metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + Field::new("score", DataType::Int64, true), + Field::new("meta", DataType::Utf8, true).with_metadata(json_metadata.clone()), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5])), + Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"])), + Arc::new(Int64Array::from(vec![10, 20, 30, 40, 50])), + Arc::new(StringArray::from(vec![ + r#"{"x":1}"#, + r#"{"x":2}"#, + r#"{"x":3}"#, + r#"{"x":4}"#, + r#"{"x":5}"#, + ])), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let dataset = Arc::new( + Dataset::write(reader, test_dir.as_ref(), None) + .await + .unwrap(), + ); + + // Perform a subschema merge_insert: only update "meta" column (JSON type) + // This exercises the update_fragments path with interleave_batches + let update_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("meta", DataType::Utf8, true).with_metadata(json_metadata), + ])); + let update_batch = RecordBatch::try_new( + update_schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![2, 4])), + Arc::new(StringArray::from(vec![ + r#"{"updated":true,"id":2}"#, + r#"{"updated":true,"id":4}"#, + ])), + ], + ) + .unwrap(); + let update_reader: Box = Box::new(RecordBatchIterator::new( + vec![Ok(update_batch)], + update_schema, + )); + let stream = reader_to_stream(update_reader); + + // Execute merge_insert with subschema (only id + meta columns) + let mut builder = + MergeInsertBuilder::try_new(dataset.clone(), vec!["id".to_string()]).unwrap(); + builder.when_matched(WhenMatched::UpdateAll); + builder.when_not_matched(WhenNotMatched::DoNothing); + let job = builder.try_build().unwrap(); + let (updated_dataset, stats) = job.execute(stream).await.unwrap(); + + // Verify: the merge should not fail with type mismatch + assert_eq!(stats.num_updated_rows, 2); + + // Read back and verify the JSON column was updated correctly + let batches = updated_dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let result = concat_batches(&batches[0].schema(), &batches).unwrap(); + assert_eq!(result.num_rows(), 5); + + // Verify the "score" column (not in update) is preserved, and "meta" updated + let ids = result + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let scores = result + .column_by_name("score") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let metas = result + .column_by_name("meta") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..5 { + let id = ids.value(i); + let score = scores.value(i); + let meta = metas.value(i); + // score = id * 10, regardless of row order + assert_eq!(score, id * 10, "id={} score mismatch", id); + if id == 2 || id == 4 { + assert!( + meta.contains("updated"), + "id={} should have updated meta, got: {}", + id, + meta + ); + } else { + assert!( + meta.contains("\"x\""), + "id={} should have original meta, got: {}", + id, + meta + ); + } + } + } } diff --git a/rust/lance/src/dataset/write/merge_insert/exec/delete.rs b/rust/lance/src/dataset/write/merge_insert/exec/delete.rs index be7897e2441..07fad758902 100644 --- a/rust/lance/src/dataset/write/merge_insert/exec/delete.rs +++ b/rust/lance/src/dataset/write/merge_insert/exec/delete.rs @@ -22,7 +22,10 @@ use roaring::RoaringTreemap; use crate::Dataset; use crate::dataset::transaction::{Operation, Transaction}; use crate::dataset::write::merge_insert::assign_action::Action; -use crate::dataset::write::merge_insert::{MERGE_ACTION_COLUMN, MergeInsertParams, MergeStats}; +use crate::dataset::write::merge_insert::{ + MERGE_ACTION_COLUMN, MergeInsertParams, MergeStats, SourceDedupeBehavior, + create_duplicate_row_error, resolve_target_bases, +}; use super::{MergeInsertMetrics, apply_deletions}; @@ -101,6 +104,8 @@ impl DeleteOnlyMergeInsertExec { async fn collect_deletions( mut input_stream: SendableRecordBatchStream, metrics: MergeInsertMetrics, + source_dedupe_behavior: SourceDedupeBehavior, + on_columns: &[String], ) -> DFResult { let schema = input_stream.schema(); @@ -156,8 +161,25 @@ impl DeleteOnlyMergeInsertExec { if action == Action::Delete && !row_addr_array.is_null(row_idx) { let row_addr = row_addr_array.value(row_idx); - delete_row_addrs.insert(row_addr); - metrics.num_deleted_rows.add(1); + // The treemap dedupes addresses, so a repeat insert signals + // a duplicate source row matching the same target; apply the + // same dedupe policy as updates. (Delete-only never carries + // `delete_not_matched_by_source`, so every delete here is a + // source match.) + if delete_row_addrs.insert(row_addr) { + metrics.num_deleted_rows.add(1); + } else { + match source_dedupe_behavior { + SourceDedupeBehavior::Fail => { + return Err(create_duplicate_row_error( + &batch, row_idx, on_columns, + )); + } + SourceDedupeBehavior::FirstSeen => { + metrics.num_skipped_duplicates.add(1); + } + } + } } } } @@ -257,13 +279,24 @@ impl ExecutionPlan for DeleteOnlyMergeInsertExec { let input_stream = self.input.execute(partition, context)?; let dataset = self.dataset.clone(); + let params = self.params.clone(); let merge_stats_holder = self.merge_stats.clone(); let transaction_holder = self.transaction.clone(); let affected_rows_holder = self.affected_rows.clone(); let merged_generations = self.params.merged_generations.clone(); + let source_dedupe_behavior = self.params.source_dedupe_behavior; + let on_columns = self.params.on.clone(); let result_stream = futures::stream::once(async move { - let delete_row_addrs = Self::collect_deletions(input_stream, metrics).await?; + // Delete-only merges write no data files, but still validate any + // requested target bases so unknown bases fail on every path. + resolve_target_bases(&dataset, ¶ms).await?; + // `metrics` is moved into `collect_deletions`; keep a handle on the + // skipped-duplicate counter so it can be folded into the stats below. + let skipped_duplicates = metrics.num_skipped_duplicates.clone(); + let delete_row_addrs = + Self::collect_deletions(input_stream, metrics, source_dedupe_behavior, &on_columns) + .await?; let (updated_fragments, removed_fragment_ids) = apply_deletions(&dataset, &delete_row_addrs) @@ -297,7 +330,7 @@ impl ExecutionPlan for DeleteOnlyMergeInsertExec { bytes_written: 0, num_files_written: 0, num_attempts: 1, - num_skipped_duplicates: 0, + num_skipped_duplicates: skipped_duplicates.value() as u64, }; if let Ok(mut transaction_guard) = transaction_holder.lock() { diff --git a/rust/lance/src/dataset/write/merge_insert/exec/write.rs b/rust/lance/src/dataset/write/merge_insert/exec/write.rs index d80c3b084f6..d5b51b3d97f 100644 --- a/rust/lance/src/dataset/write/merge_insert/exec/write.rs +++ b/rust/lance/src/dataset/write/merge_insert/exec/write.rs @@ -30,14 +30,14 @@ use crate::dataset::write::merge_insert::inserted_rows::{ }; use crate::dataset::write::merge_insert::{ MERGE_SOURCE_SENTINEL, SourceDedupeBehavior, create_duplicate_row_error, - format_key_values_on_columns, + format_key_values_on_columns, resolve_target_bases, }; use crate::{ Dataset, dataset::{ transaction::{Operation, Transaction}, write::{ - WriteParams, + WriteParams, cleanup_data_fragments, merge_insert::{ MERGE_ACTION_COLUMN, MergeInsertParams, MergeStats, assign_action::Action, exec::MergeInsertMetrics, @@ -103,6 +103,29 @@ impl MergeState { // Delete action - only delete, don't write back if !row_addr_array.is_null(row_idx) { let row_addr = row_addr_array.value(row_idx); + let row_id = row_id_array.value(row_idx); + + // A source with duplicate keys matches the same target row + // more than once; apply the same dedupe policy as updates. + // (Target-only deletes from `delete_not_matched_by_source` + // also reach here but never duplicate, so they never trip + // `Fail`.) + if !self.processed_row_ids.insert(row_id) { + match self.source_dedupe_behavior { + SourceDedupeBehavior::Fail => { + return Err(create_duplicate_row_error( + batch, + row_idx, + &self.on_columns, + )); + } + SourceDedupeBehavior::FirstSeen => { + self.metrics.num_skipped_duplicates.add(1); + return Ok(None); // Skip this duplicate row + } + } + } + self.delete_row_addrs.insert(row_addr); self.metrics.num_deleted_rows.add(1); } @@ -866,6 +889,7 @@ impl ExecutionPlan for FullSchemaMergeInsertExec { // Use flat_map to handle the async write operation let dataset = self.dataset.clone(); + let params = self.params.clone(); let merge_stats_holder = self.merge_stats.clone(); let transaction_holder = self.transaction.clone(); let affected_rows_holder = self.affected_rows.clone(); @@ -879,6 +903,9 @@ impl ExecutionPlan for FullSchemaMergeInsertExec { let result_stream = stream::once(async move { // Step 2: Write new fragments using the filtered data (inserts + updates) + let target_bases_info = resolve_target_bases(&dataset, ¶ms).await?; + // Keep a copy so failures after the write can clean up routed files. + let cleanup_bases = target_bases_info.clone(); let (mut new_fragments, _) = write_fragments_internal( Some(&dataset), dataset.object_store.clone(), @@ -886,31 +913,44 @@ impl ExecutionPlan for FullSchemaMergeInsertExec { dataset.schema().clone(), write_data_stream, WriteParams::default(), - None, // Merge insert doesn't use target_bases + target_bases_info, ) .await?; - if let Some(row_id_sequence) = updating_row_ids.lock().unwrap().row_id_sequence() { - let fragment_sizes = new_fragments - .iter() - .map(|f| f.physical_rows.unwrap() as u64); - - let sequences = lance_table::rowids::rechunk_sequences( - [row_id_sequence.clone()], - fragment_sizes, - true, - ) - .map_err(|e| { - Error::internal(format!( - "Captured row ids not equal to number of rows written: {}", - e - )) - })?; - - for (fragment, sequence) in new_fragments.iter_mut().zip(sequences) { - let serialized = lance_table::rowids::write_row_ids(&sequence); - fragment.row_id_meta = Some(RowIdMeta::Inline(serialized)); + let row_id_result: lance_core::Result<()> = (|| { + if let Some(row_id_sequence) = updating_row_ids.lock().unwrap().row_id_sequence() { + let fragment_sizes = new_fragments + .iter() + .map(|f| f.physical_rows.unwrap() as u64); + + let sequences = lance_table::rowids::rechunk_sequences( + [row_id_sequence.clone()], + fragment_sizes, + true, + ) + .map_err(|e| { + Error::internal(format!( + "Captured row ids not equal to number of rows written: {}", + e + )) + })?; + + for (fragment, sequence) in new_fragments.iter_mut().zip(sequences) { + let serialized = lance_table::rowids::write_row_ids(&sequence); + fragment.row_id_meta = Some(RowIdMeta::Inline(serialized)); + } } + Ok(()) + })(); + if let Err(e) = row_id_result { + cleanup_data_fragments( + &dataset.object_store, + &dataset.base, + cleanup_bases.as_deref(), + &new_fragments, + ) + .await; + return Err(e.into()); } // Step 2.5: Calculate write metrics from new fragments @@ -932,7 +972,21 @@ impl ExecutionPlan for FullSchemaMergeInsertExec { }; let (updated_fragments, removed_fragment_ids) = - apply_deletions(&dataset, &delete_row_addrs_clone).await?; + match apply_deletions(&dataset, &delete_row_addrs_clone).await { + Ok(result) => result, + Err(e) => { + // The new data files are not committed; remove them (including + // ones routed to target bases) before surfacing the error. + cleanup_data_fragments( + &dataset.object_store, + &dataset.base, + cleanup_bases.as_deref(), + &new_fragments, + ) + .await; + return Err(e.into()); + } + }; // Step 4: Create the transaction operation let operation = Operation::Update { diff --git a/rust/lance/src/dataset/write/update.rs b/rust/lance/src/dataset/write/update.rs index f8e71b834a7..a87b3a4260b 100644 --- a/rust/lance/src/dataset/write/update.rs +++ b/rust/lance/src/dataset/write/update.rs @@ -25,6 +25,7 @@ use datafusion::prelude::Expr; use datafusion::scalar::ScalarValue; use futures::StreamExt; use lance_arrow::RecordBatchExt; +use lance_core::datatypes::BlobHandling; use lance_core::error::{InvalidInputSnafu, box_error}; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_core::{ROW_ADDR_FIELD, ROW_ID_FIELD, ROW_OFFSET_FIELD}; @@ -280,6 +281,7 @@ impl UpdateJob { async fn execute_impl(self) -> Result { let mut scanner = self.dataset.scan(); scanner.with_row_id(); + scanner.blob_handling(BlobHandling::AllBinary); if let Some(expr) = &self.condition { scanner.filter_expr(expr.clone()); @@ -369,6 +371,7 @@ impl UpdateJob { cleanup_data_fragments( &self.dataset.object_store, &self.dataset.base, + None, &new_fragments, ) .await; @@ -1701,4 +1704,74 @@ mod tests { "Rewritten data files should be cleaned up on apply_deletions failure" ); } + + #[tokio::test] + async fn test_update_with_blob() { + use arrow_array::LargeBinaryArray; + use arrow_schema::Field; + use lance_arrow::BLOB_META_KEY; + + let test_dir = TempStrDir::default(); + let blob_meta = HashMap::from([(BLOB_META_KEY.to_string(), "true".to_string())]); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("blobs", DataType::LargeBinary, true).with_metadata(blob_meta), + Field::new("id", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(LargeBinaryArray::from(vec![ + Some(b"foo".as_slice()), + Some(b"bar".as_slice()), + Some(b"baz".as_slice()), + ])), + Arc::new(Int64Array::from(vec![0, 1, 2])), + ], + ) + .unwrap(); + + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write( + reader, + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_1), + ..Default::default() + }), + ) + .await + .unwrap(); + + // Perform an update: update the "blobs" column where id = 1 + let dataset = Arc::new(dataset); + let updated_dataset = UpdateBuilder::new(dataset) + .update_where("id = 1") + .unwrap() + .set("blobs", "arrow_cast('updated_bar', 'LargeBinary')") + .unwrap() + .build() + .unwrap() + .execute() + .await + .unwrap() + .new_dataset; + + // Verify the updated value + let mut scanner = updated_dataset.scan(); + // Read as binary to assert actual value + scanner.blob_handling(BlobHandling::AllBinary); + let batches = scanner.try_into_batch().await.unwrap(); + let blobs = batches.column_by_name("blobs").unwrap().as_binary::(); + let ids = batches + .column_by_name("id") + .unwrap() + .as_primitive::(); + + // Find the index of id = 1 + let idx = ids.values().iter().position(|&x| x == 1).unwrap(); + assert_eq!(blobs.value(idx), b"updated_bar"); + + let idx_foo = ids.values().iter().position(|&x| x == 0).unwrap(); + assert_eq!(blobs.value(idx_foo), b"foo"); + } } diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index d98ab8701c1..b960405b3ee 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -4,8 +4,10 @@ //! Secondary Index //! +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::collections::{HashMap, HashSet}; use std::sync::{Arc, OnceLock}; +use std::time::Instant; use arrow_schema::DataType; use async_trait::async_trait; @@ -15,7 +17,6 @@ use itertools::Itertools; use lance_core::cache::CacheKey; use lance_core::datatypes::Field; use lance_core::datatypes::Schema as LanceSchema; -use lance_core::utils::address::RowAddress; use lance_core::utils::parse::parse_env_as_bool; use lance_core::utils::tracing::{ IO_TYPE_OPEN_FRAG_REUSE, IO_TYPE_OPEN_MEM_WAL, IO_TYPE_OPEN_VECTOR, TRACE_IO_EVENTS, @@ -60,7 +61,7 @@ use lance_table::io::manifest::read_manifest_indexes; use roaring::RoaringBitmap; use scalar::index_matches_criteria; use serde_json::json; -use tracing::{info, instrument}; +use tracing::{info, instrument, warn}; use uuid::Uuid; use vector::details::{ derive_vector_index_type, infer_missing_vector_details, vector_details_as_json, @@ -207,6 +208,227 @@ fn retain_committed_inverted_files(files: &mut Vec) { files.retain(|file| !file.path.starts_with("staging/")); } +async fn prewarm_opened_index( + index: Arc, + options: Option<&PrewarmOptions>, +) -> Result<()> { + match options { + None => index.prewarm().await, + Some(PrewarmOptions::Fts(fts_options)) => { + let inverted = index + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::invalid_input(format!( + "FTS prewarm options are only supported for inverted indices, got {:?}", + index.index_type() + )) + })?; + inverted.prewarm_with_options(fts_options).await + } + Some(_) => Err(Error::not_supported( + "unsupported prewarm options for this lance version".to_owned(), + )), + } +} + +fn total_index_segment_size_bytes(indices: &[IndexMetadata]) -> Option { + let mut total = 0u64; + for index_meta in indices { + total += index_meta.total_size_bytes()?; + } + Some(total) +} + +fn cache_size_delta(after: usize, before: usize) -> i64 { + after.saturating_sub(before) as i64 - before.saturating_sub(after) as i64 +} + +fn prewarm_options_fields(options: Option<&PrewarmOptions>) -> (&'static str, bool) { + match options { + None => ("default", false), + Some(PrewarmOptions::Fts(fts_options)) => ("fts", fts_options.with_position), + Some(_) => ("unsupported", false), + } +} + +async fn prewarm_index_segments_by_metadata( + dataset: &Dataset, + name: &str, + indices: Vec, + options: Option<&PrewarmOptions>, + available_segment_count: usize, + requested_segment_count: Option, +) -> Result<()> { + let request_started = Instant::now(); + let selected_segment_count = indices.len(); + let selected_size_bytes = total_index_segment_size_bytes(&indices); + let (prewarm_options, fts_with_position) = prewarm_options_fields(options); + let cache_stats_before = dataset.session.index_cache_stats().await; + info!( + index_name = name, + selected_segment_count, + available_segment_count, + requested_segment_count = requested_segment_count.unwrap_or(0), + segment_filter = requested_segment_count.is_some(), + selected_size_bytes = selected_size_bytes.unwrap_or(0), + selected_size_bytes_known = selected_size_bytes.is_some(), + index_cache_entries_before = cache_stats_before.num_entries, + index_cache_size_bytes_before = cache_stats_before.size_bytes, + prewarm_options, + fts_with_position, + "prewarm index segments started" + ); + + let result = futures::future::try_join_all(indices.into_iter().map(|index_meta| async move { + let index_uuid = index_meta.uuid; + let size_bytes = index_meta.total_size_bytes(); + let fragment_count = index_meta + .fragment_bitmap + .as_ref() + .map(|bitmap| bitmap.len()); + let segment_started = Instant::now(); + info!( + index_name = name, + %index_uuid, + index_version = index_meta.index_version, + dataset_version = index_meta.dataset_version, + fragment_count = fragment_count.unwrap_or(0), + fragment_count_known = fragment_count.is_some(), + size_bytes = size_bytes.unwrap_or(0), + size_bytes_known = size_bytes.is_some(), + "prewarm index segment started" + ); + + let index = match dataset + .open_generic_index(name, &index_uuid, &NoOpMetricsCollector) + .await + { + Ok(index) => { + info!( + index_name = name, + %index_uuid, + index_type = ?index.index_type(), + elapsed_ms = segment_started.elapsed().as_millis() as u64, + "opened index segment for prewarm" + ); + index + } + Err(err) => { + warn!( + index_name = name, + %index_uuid, + error = %err, + elapsed_ms = segment_started.elapsed().as_millis() as u64, + "failed to open index segment for prewarm" + ); + return Err(err); + } + }; + + if let Err(err) = prewarm_opened_index(index, options).await { + warn!( + index_name = name, + %index_uuid, + error = %err, + elapsed_ms = segment_started.elapsed().as_millis() as u64, + "prewarm index segment failed" + ); + return Err(err); + } + + info!( + index_name = name, + %index_uuid, + elapsed_ms = segment_started.elapsed().as_millis() as u64, + "prewarm index segment finished" + ); + Ok(()) + })) + .await; + + match result { + Ok(_) => { + let cache_stats_after = dataset.session.index_cache_stats().await; + info!( + index_name = name, + selected_segment_count, + index_cache_entries_after = cache_stats_after.num_entries, + index_cache_entries_delta = cache_size_delta( + cache_stats_after.num_entries, + cache_stats_before.num_entries, + ), + index_cache_size_bytes_after = cache_stats_after.size_bytes, + index_cache_size_bytes_delta = + cache_size_delta(cache_stats_after.size_bytes, cache_stats_before.size_bytes,), + elapsed_ms = request_started.elapsed().as_millis() as u64, + "prewarm index segments finished" + ); + Ok(()) + } + Err(err) => { + let cache_stats_after = dataset.session.index_cache_stats().await; + warn!( + index_name = name, + selected_segment_count, + index_cache_entries_after = cache_stats_after.num_entries, + index_cache_entries_delta = cache_size_delta( + cache_stats_after.num_entries, + cache_stats_before.num_entries, + ), + index_cache_size_bytes_after = cache_stats_after.size_bytes, + index_cache_size_bytes_delta = cache_size_delta( + cache_stats_after.size_bytes, + cache_stats_before.size_bytes, + ), + error = %err, + elapsed_ms = request_started.elapsed().as_millis() as u64, + "prewarm index segments failed" + ); + Err(err) + } + } +} + +fn filter_index_segments_by_ids( + name: &str, + indices: Vec, + segment_ids: &[Uuid], +) -> Result> { + if segment_ids.is_empty() { + return Ok(Vec::new()); + } + + let requested = segment_ids.iter().copied().collect::>(); + let mut matched = HashSet::new(); + let filtered = indices + .into_iter() + .filter(|index_meta| { + if requested.contains(&index_meta.uuid) { + matched.insert(index_meta.uuid); + true + } else { + false + } + }) + .collect::>(); + + if matched.len() != requested.len() { + let mut missing = requested + .difference(&matched) + .map(ToString::to_string) + .collect::>(); + missing.sort(); + return Err(Error::index_not_found(format!( + "name={}, segment_ids=[{}]", + name, + missing.join(", ") + ))); + } + + Ok(filtered) +} + fn validate_segment_index_details(index_name: &str, segments: &[IndexMetadata]) -> Result<()> { let mut type_url = None::<&str>; for segment in segments { @@ -292,6 +514,13 @@ fn segment_has_fmindex_details(segment: &IndexMetadata) -> bool { .is_some_and(|details| details.type_url.ends_with("FMIndexDetails")) } +fn segment_has_label_list_details(segment: &IndexMetadata) -> bool { + segment + .index_details + .as_ref() + .is_some_and(|details| details.type_url.ends_with("LabelListIndexDetails")) +} + // Cache keys for different index types #[derive(Debug, Clone)] pub(crate) struct LegacyVectorIndexCacheKey<'a> { @@ -477,7 +706,7 @@ pub trait IndexBuilder { pub(crate) async fn remap_index( dataset: &Dataset, index_id: &Uuid, - row_id_map: &HashMap>, + row_id_map: &RowAddrRemap, ) -> Result { // Load indices from the disk. let indices = dataset.load_indices().await?; @@ -492,20 +721,14 @@ pub(crate) async fn remap_index( )); } - if row_id_map.values().all(|v| v.is_none()) { - let deleted_bitmap = RoaringBitmap::from_iter( - row_id_map - .keys() - .map(|row_id| RowAddress::new_from_u64(*row_id)) - .map(|addr| addr.fragment_id()), - ); - if Some(deleted_bitmap) == matched.fragment_bitmap { - // If remap deleted all rows, we can just return the same index ID. - // This can happen if there is a bug where the index is covering empty - // fragment that haven't been cleaned up. They should be cleaned up - // outside of this function. - return Ok(RemapResult::Keep(*index_id)); - } + if let Some(deleted_bitmap) = row_id_map.fully_deleted_fragments() + && Some(deleted_bitmap) == matched.fragment_bitmap + { + // If remap deleted all rows, we can just return the same index ID. + // This can happen if there is a bug where the index is covering empty + // fragment that haven't been cleaned up. They should be cleaned up + // outside of this function. + return Ok(RemapResult::Keep(*index_id)); } let field_id = matched @@ -780,10 +1003,15 @@ impl IndexDescriptionImpl { let mut missing_fragment_refs = 0u64; for shard in &segments { - let fragment_bitmap = shard - .fragment_bitmap - .as_ref() - .ok_or_else(|| Error::index("Fragment bitmap is required for index description. This index must be retrained to support this method.".to_string()))?; + let Some(fragment_bitmap) = shard.fragment_bitmap.as_ref() else { + // A system index (e.g. __mem_wal) indexes no fragments, so a + // missing bitmap means zero indexed rows. For a data index it + // means unknown coverage — reject rather than fabricate a count. + if is_system_index(shard) { + continue; + } + return Err(Error::index("Fragment bitmap is required for index description. This index must be retrained to support this method.".to_string())); + }; indexed_fragment_refs += fragment_bitmap.len(); for fragment_id in fragment_bitmap.iter() { @@ -973,49 +1201,70 @@ impl DatasetIndexExt for Dataset { return Err(Error::index_not_found(format!("name={}", name))); } - for index_meta in indices { - let index = self - .open_generic_index(name, &index_meta.uuid, &NoOpMetricsCollector) - .await?; - index.prewarm().await?; + let available_segment_count = indices.len(); + prewarm_index_segments_by_metadata(self, name, indices, None, available_segment_count, None) + .await + } + + async fn prewarm_index_with_options(&self, name: &str, options: &PrewarmOptions) -> Result<()> { + let indices = self.load_indices_by_name(name).await?; + if indices.is_empty() { + return Err(Error::index_not_found(format!("name={}", name))); } - Ok(()) + let available_segment_count = indices.len(); + prewarm_index_segments_by_metadata( + self, + name, + indices, + Some(options), + available_segment_count, + None, + ) + .await } - async fn prewarm_index_with_options(&self, name: &str, options: &PrewarmOptions) -> Result<()> { + async fn prewarm_index_segments(&self, name: &str, segment_ids: &[Uuid]) -> Result<()> { let indices = self.load_indices_by_name(name).await?; if indices.is_empty() { return Err(Error::index_not_found(format!("name={}", name))); } + let available_segment_count = indices.len(); + let indices = filter_index_segments_by_ids(name, indices, segment_ids)?; - for index_meta in indices { - let index = self - .open_generic_index(name, &index_meta.uuid, &NoOpMetricsCollector) - .await?; + prewarm_index_segments_by_metadata( + self, + name, + indices, + None, + available_segment_count, + Some(segment_ids.len()), + ) + .await + } - match options { - PrewarmOptions::Fts(fts_options) => { - let inverted = index - .as_any() - .downcast_ref::() - .ok_or_else(|| { - Error::invalid_input(format!( - "FTS prewarm options are only supported for inverted indices, got {:?}", - index.index_type() - )) - })?; - inverted.prewarm_with_options(fts_options).await?; - } - _ => { - return Err(Error::not_supported( - "unsupported prewarm options for this lance version".to_owned(), - )); - } - } + async fn prewarm_index_segments_with_options( + &self, + name: &str, + segment_ids: &[Uuid], + options: &PrewarmOptions, + ) -> Result<()> { + let indices = self.load_indices_by_name(name).await?; + if indices.is_empty() { + return Err(Error::index_not_found(format!("name={}", name))); } + let available_segment_count = indices.len(); + let indices = filter_index_segments_by_ids(name, indices, segment_ids)?; - Ok(()) + prewarm_index_segments_by_metadata( + self, + name, + indices, + Some(options), + available_segment_count, + Some(segment_ids.len()), + ) + .await } async fn describe_indices<'a, 'b>( @@ -1149,7 +1398,14 @@ impl DatasetIndexExt for Dataset { let all_btree = source_segments.iter().all(segment_has_btree_details); let all_fmindex = source_segments.iter().all(segment_has_fmindex_details); let all_zonemap = source_segments.iter().all(segment_has_zonemap_details); - if !all_vector && !all_inverted && !all_bitmap && !all_btree && !all_fmindex && !all_zonemap + let all_label_list = source_segments.iter().all(segment_has_label_list_details); + if !all_vector + && !all_inverted + && !all_bitmap + && !all_btree + && !all_fmindex + && !all_zonemap + && !all_label_list { return Err(Error::invalid_input( "merge_existing_index_segments requires all segments to have the same supported index type" @@ -1170,6 +1426,8 @@ impl DatasetIndexExt for Dataset { crate::index::scalar::fmindex::merge_segments(self, source_segments).await? } else if all_bitmap { crate::index::scalar::bitmap::merge_segments(self, source_segments).await? + } else if all_label_list { + crate::index::scalar::label_list::merge_segments(self, source_segments).await? } else if all_zonemap { crate::index::scalar::zonemap::merge_segments(self, source_segments).await? } else { @@ -4071,7 +4329,9 @@ mod tests { assert_io_eq!(stats, read_bytes, 0); assert_eq!(indices.len(), 1); - session.index_cache.clear().await; // Clear the cache + // Clear the cache + session.index_cache.clear().await; + session.file_metadata_cache().clear().await; let dataset2 = DatasetBuilder::from_uri(test_uri) .with_session(session.clone()) @@ -4111,7 +4371,7 @@ mod tests { let remap_to_empty = (0..dataset.count_all_rows().await.unwrap()) .map(|i| (i as u64, None)) .collect::>(); - let new_uuid = remap_index(&dataset, &index_uuid, &remap_to_empty) + let new_uuid = remap_index(&dataset, &index_uuid, &RowAddrRemap::direct(remap_to_empty)) .await .unwrap(); assert_eq!(new_uuid, RemapResult::Keep(index_uuid)); diff --git a/rust/lance/src/index/api.rs b/rust/lance/src/index/api.rs index f856e9004f3..54c710f42ea 100644 --- a/rust/lance/src/index/api.rs +++ b/rust/lance/src/index/api.rs @@ -167,6 +167,25 @@ pub trait DatasetIndexExt { )) } + /// Prewarm selected physical segments of an index by name. + async fn prewarm_index_segments(&self, _name: &str, _segment_ids: &[Uuid]) -> Result<()> { + Err(Error::not_supported( + "segment-level prewarm is not supported by this dataset implementation".to_owned(), + )) + } + + /// Prewarm selected physical segments of an index by name with additional options. + async fn prewarm_index_segments_with_options( + &self, + _name: &str, + _segment_ids: &[Uuid], + _options: &PrewarmOptions, + ) -> Result<()> { + Err(Error::not_supported( + "prewarm options are not supported by this dataset implementation".to_owned(), + )) + } + /// Read all indices of this Dataset version. /// /// The indices are lazy loaded and cached in memory within the `Dataset` instance. diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index d3ecde030c1..edcd4357ae4 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -863,8 +863,10 @@ mod tests { use arrow::datatypes::{Float32Type, UInt32Type}; use arrow_array::cast::AsArray; use arrow_array::{ - FixedSizeListArray, Int32Array, RecordBatch, RecordBatchIterator, StringArray, UInt32Array, + ArrayRef, FixedSizeListArray, Int32Array, RecordBatch, RecordBatchIterator, StringArray, + UInt32Array, }; + use arrow_buffer::{BooleanBufferBuilder, NullBuffer}; use arrow_schema::{DataType, Field, Schema}; use futures::TryStreamExt; use lance_arrow::FixedSizeListArrayExt; @@ -1005,6 +1007,118 @@ mod tests { assert_eq!(num_rows, 2000); } + #[tokio::test] + async fn test_optimize_append_preserves_case_sensitive_nullable_vector_column() { + const DIM: usize = 64; + const ROWS: usize = 1000; + + fn make_vectors(rows: usize, dim: usize, include_null: bool) -> FixedSizeListArray { + if include_null { + let mut nulls_builder = BooleanBufferBuilder::new(rows); + for row_idx in 0..rows { + nulls_builder.append(row_idx != 0); + } + let nulls = NullBuffer::new(nulls_builder.finish()); + FixedSizeListArray::try_new( + Arc::new(Field::new("item", DataType::Float32, true)), + dim as i32, + Arc::new(generate_random_array(rows * dim)), + Some(nulls), + ) + .unwrap() + } else { + FixedSizeListArray::try_new_from_values( + generate_random_array(rows * dim), + dim as i32, + ) + .unwrap() + } + } + + fn make_batch( + schema: Arc, + start_id: u32, + vectors: Arc, + ) -> RecordBatch { + let columns: Vec = vec![ + Arc::new(UInt32Array::from_iter_values( + start_id..start_id + ROWS as u32, + )) as ArrayRef, + vectors as ArrayRef, + ]; + RecordBatch::try_new(schema, columns).unwrap() + } + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let vector_type = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIM as i32, + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("VECTOR", vector_type, true), + ])); + + let initial_vectors = Arc::new(make_vectors(ROWS, DIM, false)); + let initial_batch = make_batch(schema.clone(), 0, initial_vectors); + let batches = RecordBatchIterator::new(std::iter::once(Ok(initial_batch)), schema.clone()); + let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap(); + + let params = VectorIndexParams::with_ivf_pq_params( + MetricType::L2, + IvfBuildParams::new(2), + PQBuildParams { + num_sub_vectors: 2, + ..Default::default() + }, + ); + dataset + .create_index(&["VECTOR"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + let appended_vectors = Arc::new(make_vectors(ROWS, DIM, true)); + let query = appended_vectors.value(5); + let appended_batch = make_batch(schema.clone(), ROWS as u32, appended_vectors); + let batches = RecordBatchIterator::new(std::iter::once(Ok(appended_batch)), schema); + dataset.append(batches, None).await.unwrap(); + + let index_name = dataset.load_indices().await.unwrap()[0].name.clone(); + assert!( + !dataset + .unindexed_fragments(&index_name) + .await + .unwrap() + .is_empty() + ); + + dataset + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap(); + + let dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap(); + assert!( + dataset + .unindexed_fragments(&index_name) + .await + .unwrap() + .is_empty() + ); + + let mut scanner = dataset.scan(); + scanner + .nearest("VECTOR", query.as_primitive::(), 10) + .unwrap(); + let results = scanner.try_into_batch().await.unwrap(); + assert_eq!( + results.num_rows(), + 10, + "expected the requested k=10 nearest-neighbor results" + ); + } + /// Regression: a second `OptimizeOptions::append()` call on a steady-state /// vector index used to fall through to `optimize_vector_indices` and write /// a new UUID directory + manifest even though nothing had changed. The @@ -1213,7 +1327,23 @@ mod tests { assert_eq!(results.num_rows(), 2); let mut id_arr = results["id"].as_primitive::().values().to_vec(); id_arr.sort(); - assert_eq!(id_arr, vec![0, 1000]); + // For exact indexes (e.g. IvfFlat) the top-2 nearest neighbors of the + // query vector are deterministic (id 0 in the first delta and id 1000 + // in the second delta, since the same vector is duplicated across + // the two fragments). For approximate indexes (e.g. IvfPq, IvfHnswSq) + // the returned ids are not guaranteed to be exactly [0, 1000]; the key + // property this test verifies is that both delta indices are queried + // (i.e. one result comes from each delta), so we only assert that. + let is_approximate = !matches!(index_params.index_type(), IndexType::IvfFlat); + if is_approximate { + assert!( + id_arr[0] < TOTAL as u32 && id_arr[1] >= TOTAL as u32, + "expected one result from each delta index, got {:?}", + id_arr + ); + } else { + assert_eq!(id_arr, vec![0, 1000]); + } } #[tokio::test] diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index 19f79f8dd66..702a8c4e49f 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -847,9 +847,9 @@ mod tests { use crate::dataset::{WriteMode, WriteParams}; use crate::index::{DatasetIndexExt, IndexSegment}; use crate::utils::test::{DatagenExt, FragmentCount, FragmentRowCount}; - use arrow::datatypes::{Float32Type, Int32Type}; + use arrow::datatypes::{Float32Type, Int32Type, Int64Type}; use arrow_array::cast::AsArray; - use arrow_array::{FixedSizeListArray, RecordBatchIterator}; + use arrow_array::{Array, FixedSizeListArray, ListArray, RecordBatchIterator}; use arrow_array::{Int32Array, RecordBatch, StringArray}; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use datafusion::common::ScalarValue; @@ -2271,6 +2271,313 @@ mod tests { ); } + // Distributed LabelList build: one segment per fragment via + // `execute_uncommitted`, then `merge_existing_index_segments` consolidates them + // into a single canonical segment that answers `array_has_any` across all rows. + #[tokio::test] + async fn test_label_list_merge_existing_index_segments() { + use lance_index::scalar::{LabelListQuery, SearchResult}; + + // Open `segment` and count rows whose `labels` list contains `label`. + async fn count_has_any(dataset: &Dataset, segment: &IndexMetadata, label: i64) -> usize { + let field_path = dataset.schema().field_path(segment.fields[0]).unwrap(); + let index = crate::index::scalar::open_scalar_index( + dataset, + &field_path, + segment, + &NoOpMetricsCollector, + ) + .await + .unwrap(); + let query = LabelListQuery::HasAnyLabel(vec![ScalarValue::Int64(Some(label))]); + match index.search(&query, &NoOpMetricsCollector).await.unwrap() { + SearchResult::Exact(row_addrs) => { + row_addrs.true_rows().row_addrs().unwrap().count() + } + other => panic!("expected exact result, got {other:?}"), + } + } + + let tmpdir = TempStrDir::default(); + let dataset_uri = format!("file://{}", tmpdir.as_str()); + + // 4000 rows across two 2000-row fragments; each `labels` list cycles over 1..=5. + let mut dataset = gen_batch() + .col( + "labels", + lance_datagen::array::rand_list_any( + lance_datagen::array::cycle::(vec![1, 2, 3, 4, 5]), + false, + ), + ) + .into_dataset( + &dataset_uri, + FragmentCount::from(2), + FragmentRowCount::from(2000), + ) + .await + .unwrap(); + + // Ground truth via a full scan before any index exists. + let expected = dataset + .scan() + .project(&["labels"]) + .unwrap() + .filter("array_has_any(labels, [3])") + .unwrap() + .try_into_batch() + .await + .unwrap() + .num_rows(); + assert!( + expected > 0, + "test dataset must contain at least one row whose labels include 3" + ); + + // One LabelList segment per fragment, committed as a multi-segment index. + let params = + ScalarIndexParams::for_builtin(lance_index::scalar::BuiltinIndexType::LabelList); + let mut staged = Vec::new(); + for fragment in dataset.get_fragments() { + staged.push( + CreateIndexBuilder::new(&mut dataset, &["labels"], IndexType::LabelList, ¶ms) + .name("labels_idx".to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(), + ); + } + dataset + .commit_existing_index_segments("labels_idx", "labels", staged) + .await + .unwrap(); + + // Merge the two per-fragment segments into a single segment covering both. + let merged = dataset + .merge_existing_index_segments( + dataset.load_indices_by_name("labels_idx").await.unwrap(), + ) + .await + .unwrap(); + assert_eq!( + merged.fragment_bitmap.as_ref().unwrap(), + &roaring::RoaringBitmap::from_iter([0u32, 1]) + ); + assert!( + merged + .index_details + .as_ref() + .unwrap() + .type_url + .ends_with("LabelListIndexDetails") + ); + // The merged segment returns every row whose labels include 3 across both + // fragments — i.e. the per-fragment bitmaps and null sets were unioned. + assert_eq!(count_has_any(&dataset, &merged, 3).await, expected); + } + + fn label_list_batch(labels: Vec>>>) -> (Arc, RecordBatch) { + let labels = ListArray::from_iter_primitive::(labels); + let row_count = labels.len(); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("labels", labels.data_type().clone(), true), + ])); + let ids = Int32Array::from_iter_values((0..row_count).map(|id| id as i32)); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(ids), Arc::new(labels)]).unwrap(); + (schema, batch) + } + + #[tokio::test] + async fn test_label_list_merge_existing_index_segments_drops_retired_fragments() { + use lance_index::scalar::{LabelListQuery, SearchResult}; + + async fn search_counts( + dataset: &Dataset, + segment: &IndexMetadata, + label: i64, + ) -> (usize, usize) { + let field_path = dataset.schema().field_path(segment.fields[0]).unwrap(); + let index = crate::index::scalar::open_scalar_index( + dataset, + &field_path, + segment, + &NoOpMetricsCollector, + ) + .await + .unwrap(); + let query = LabelListQuery::HasAnyLabel(vec![ScalarValue::Int64(Some(label))]); + match index.search(&query, &NoOpMetricsCollector).await.unwrap() { + SearchResult::Exact(row_addrs) => ( + row_addrs.true_rows().row_addrs().unwrap().count(), + row_addrs.null_rows().row_addrs().unwrap().count(), + ), + other => panic!("expected exact result, got {other:?}"), + } + } + + let tmpdir = TempStrDir::default(); + let dataset_uri = format!("file://{}", tmpdir.as_str()); + let (schema, batch) = label_list_batch(vec![ + Some(vec![Some(1)]), + None, + Some(vec![Some(1)]), + Some(vec![Some(1)]), + Some(vec![Some(2)]), + Some(vec![Some(2)]), + Some(vec![Some(2)]), + Some(vec![Some(2)]), + ]); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write( + reader, + &dataset_uri, + Some(WriteParams { + max_rows_per_file: 4, + mode: WriteMode::Overwrite, + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + + let params = + ScalarIndexParams::for_builtin(lance_index::scalar::BuiltinIndexType::LabelList); + let mut staged = Vec::new(); + for fragment in dataset.get_fragments() { + staged.push( + CreateIndexBuilder::new(&mut dataset, &["labels"], IndexType::LabelList, ¶ms) + .name("labels_retired_idx".to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(), + ); + } + dataset + .commit_existing_index_segments("labels_retired_idx", "labels", staged) + .await + .unwrap(); + + dataset.delete("id < 4").await.unwrap(); + crate::dataset::optimize::compact_files( + &mut dataset, + crate::dataset::optimize::CompactionOptions { + target_rows_per_fragment: 4, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + assert!( + !dataset.fragment_bitmap.contains(0), + "compaction should retire fragment 0" + ); + + let merged = dataset + .merge_existing_index_segments( + dataset + .load_indices_by_name("labels_retired_idx") + .await + .unwrap(), + ) + .await + .unwrap(); + let coverage = merged.fragment_bitmap.as_ref().unwrap(); + assert!(!coverage.contains(0), "must drop retired frag 0"); + assert!(coverage.contains(1), "must keep live indexed frag 1"); + + let (retired_true, _) = search_counts(&dataset, &merged, 1).await; + assert_eq!( + retired_true, 0, + "must filter value bitmaps from retired fragments" + ); + + let (live_true, null_rows) = search_counts(&dataset, &merged, 2).await; + assert_eq!(live_true, 4, "must keep live fragment rows"); + assert_eq!( + null_rows, 0, + "must filter list_nulls from retired fragments" + ); + } + + #[tokio::test] + async fn test_label_list_merge_rejects_nullable_segment_missing_null_metadata() { + use crate::dataset::index::LanceIndexStoreExt; + use lance_index::scalar::IndexStore; + use lance_index::scalar::label_list::{ + BITMAP_LOOKUP_NAME, LABEL_LIST_NULLS_METADATA_KEY, LABEL_LIST_NULLS_MIN_VERSION, + }; + use lance_index::scalar::lance_format::LanceIndexStore; + + let tmpdir = TempStrDir::default(); + let dataset_uri = format!("file://{}", tmpdir.as_str()); + let (schema, batch) = label_list_batch(vec![ + Some(vec![Some(1)]), + None, + Some(vec![Some(2)]), + Some(vec![Some(3)]), + ]); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write( + reader, + &dataset_uri, + Some(WriteParams { + mode: WriteMode::Overwrite, + ..Default::default() + }), + ) + .await + .unwrap(); + + let params = + ScalarIndexParams::for_builtin(lance_index::scalar::BuiltinIndexType::LabelList); + let segment = + CreateIndexBuilder::new(&mut dataset, &["labels"], IndexType::LabelList, ¶ms) + .name("labels_legacy_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + + let source_store = LanceIndexStore::from_dataset_for_existing(&dataset, &segment) + .await + .unwrap(); + let reader = source_store + .open_index_file(BITMAP_LOOKUP_NAME) + .await + .unwrap(); + let batch = reader.read_range(0..reader.num_rows(), None).await.unwrap(); + let source_schema: ArrowSchema = reader.schema().into(); + let schema_without_metadata = Arc::new(ArrowSchema::new(source_schema.fields().clone())); + + let legacy_uuid = Uuid::new_v4(); + let legacy_store = LanceIndexStore::from_dataset_for_new(&dataset, &legacy_uuid).unwrap(); + let mut writer = legacy_store + .new_index_file(BITMAP_LOOKUP_NAME, schema_without_metadata) + .await + .unwrap(); + writer.write_record_batch(batch).await.unwrap(); + let legacy_file = writer.finish().await.unwrap(); + + let mut legacy_segment = segment.clone(); + legacy_segment.uuid = legacy_uuid; + legacy_segment.index_version = LABEL_LIST_NULLS_MIN_VERSION; + legacy_segment.files = Some(vec![legacy_file]); + + let err = dataset + .merge_existing_index_segments(vec![legacy_segment]) + .await + .unwrap_err(); + assert!( + err.to_string().contains(LABEL_LIST_NULLS_METADATA_KEY), + "unexpected error: {err}" + ); + } + #[tokio::test] async fn test_commit_existing_index_supports_local_hnsw_segments() { let tmpdir = TempStrDir::default(); diff --git a/rust/lance/src/index/mem_wal.rs b/rust/lance/src/index/mem_wal.rs index de2c70b62d2..84c9ae25d86 100644 --- a/rust/lance/src/index/mem_wal.rs +++ b/rust/lance/src/index/mem_wal.rs @@ -490,4 +490,77 @@ mod tests { assert!(indices.is_empty()); } + + /// Regression: a committed `__mem_wal` (legitimately `fragment_bitmap: + /// None`) must not break `describe_indices` — the path behind lancedb's + /// `list_indices`/`wait_for_index`. It's described as zero indexed rows, + /// like `__frag_reuse`. + #[tokio::test] + async fn test_describe_indices_includes_mem_wal_system_index() { + use crate::index::DatasetIndexExt; + use lance_index::IndexType; + use lance_index::scalar::ScalarIndexParams; + + let mut dataset = test_dataset().await; + + // A real user index that describe_indices must keep returning. + dataset + .create_index( + &["a"], + IndexType::Scalar, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + + // Commit a __mem_wal index, as WAL provisioning does in production. + let shard = Uuid::new_v4(); + let txn = Transaction::new( + dataset.manifest.version, + Operation::UpdateMemWalState { + merged_generations: vec![MergedGeneration::new(shard, 1)], + }, + None, + ); + let dataset = CommitBuilder::new(Arc::new(dataset)) + .execute(txn) + .await + .unwrap(); + + // The system index is present with no fragment_bitmap (by design). + let mem_wal = dataset + .load_indices() + .await + .unwrap() + .iter() + .find(|i| i.name == MEM_WAL_INDEX_NAME) + .unwrap() + .clone(); + assert!(mem_wal.fragment_bitmap.is_none()); + + // describe_indices describes the bitmap-less __mem_wal alongside the + // real index instead of erroring. + let descriptions = dataset.describe_indices(None).await.unwrap(); + let mem_wal_desc = descriptions + .iter() + .find(|d| d.name() == MEM_WAL_INDEX_NAME) + .expect("__mem_wal must be described, not skipped"); + assert_eq!( + mem_wal_desc.index_type(), + "MemWal", + "system index type must resolve via infer_system_index_type" + ); + assert_eq!( + mem_wal_desc.rows_indexed(), + 0, + "a bitmap-less system index indexes zero rows" + ); + assert_eq!( + descriptions.len(), + 2, + "both the real scalar index and __mem_wal must be listed" + ); + } } diff --git a/rust/lance/src/index/scalar.rs b/rust/lance/src/index/scalar.rs index ae2478589fb..79e6eed44c9 100644 --- a/rust/lance/src/index/scalar.rs +++ b/rust/lance/src/index/scalar.rs @@ -8,6 +8,7 @@ pub(crate) mod bitmap; pub(crate) mod btree; pub(crate) mod fmindex; pub(crate) mod inverted; +pub(crate) mod label_list; pub(crate) mod zonemap; pub use inverted::{load_segment_details, load_segments}; @@ -46,12 +47,12 @@ use lance_index::scalar::label_list::{ LABEL_LIST_NULLS_METADATA_KEY, LABEL_LIST_NULLS_MIN_VERSION, }; use lance_index::scalar::registry::{ - ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, VALUE_COLUMN_NAME, + ScalarIndexLoad, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, VALUE_COLUMN_NAME, }; use lance_index::scalar::{BuiltinIndexType, CreatedIndex, InvertedIndexParams}; use lance_index::scalar::{ - ScalarIndex, ScalarIndexParams, bitmap::BITMAP_LOOKUP_NAME, inverted::INVERT_LIST_FILE, - lance_format::LanceIndexStore, + RowIdRemapper, ScalarIndex, ScalarIndexParams, bitmap::BITMAP_LOOKUP_NAME, + inverted::INVERT_LIST_FILE, lance_format::LanceIndexStore, }; use lance_index::{IndexCriteria, IndexType}; use lance_table::format::{Fragment, IndexMetadata}; @@ -296,8 +297,13 @@ pub(super) async fn build_scalar_index( let index_store = LanceIndexStore::from_dataset_for_new(dataset, &uuid)?; let plugin = SCALAR_INDEX_PLUGIN_REGISTRY.get_plugin_by_name(¶ms.index_type)?; + let trainer = plugin.basic_trainer().ok_or_else(|| { + Error::invalid_input_source( + format!("The '{}' index type does not support basic training, please refer to the index's documentation for more details on how to create this index.", params.index_type).into(), + ) + })?; let training_request = - plugin.new_training_request(params.params.as_deref().unwrap_or("{}"), &field)?; + trainer.new_training_request(params.params.as_deref().unwrap_or("{}"), &field)?; progress.stage_start("load_data", None, "rows").await?; let training_data = match preprocessed_data { @@ -316,7 +322,7 @@ pub(super) async fn build_scalar_index( }; progress.stage_complete("load_data").await?; - let created_index = plugin + let created_index = trainer .train_index( training_data, &index_store, @@ -353,8 +359,13 @@ pub(super) async fn build_bitmap_index_segment( let params = ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap); let plugin = SCALAR_INDEX_PLUGIN_REGISTRY.get_plugin_by_name(¶ms.index_type)?; + let trainer = plugin.basic_trainer().ok_or_else(|| { + Error::invalid_input_source( + format!("The '{}' index type does not support basic training, please refer to the index's documentation for more details on how to create this index.", params.index_type).into(), + ) + })?; let training_request = - plugin.new_training_request(params.params.as_deref().unwrap_or("{}"), &field)?; + trainer.new_training_request(params.params.as_deref().unwrap_or("{}"), &field)?; let criteria = training_request.criteria(); progress.stage_start("load_data", None, "rows").await?; @@ -363,7 +374,7 @@ pub(super) async fn build_bitmap_index_segment( progress.stage_complete("load_data").await?; let index_store = LanceIndexStore::from_dataset_for_new(dataset, &uuid)?; - plugin + trainer .train_index( training_data, &index_store, @@ -448,29 +459,35 @@ pub async fn open_scalar_index( .index_cache .for_index(&index.uuid, frag_reuse_index.as_ref().map(|f| &f.uuid)); - if let Some(index) = plugin - .get_from_cache(index_store.clone(), frag_reuse_index.clone(), &index_cache) - .await? - { - // Compatibility check is only needed on first load; a cache hit means - // the index was already validated when it was originally opened in - // this session, so we can skip the extra `open_index_file` IOP. - return Ok(index); - } - - if index_details.type_url.ends_with("LabelListIndexDetails") { - validate_label_list_index_compatibility(dataset, column, index, &index_store).await?; - } + let frag_reuse_index: Option> = + frag_reuse_index.map(|f| f as Arc); + + // Runs only on a cold miss, and at most once even under concurrent opens + // (the plugin coalesces). The compat check lives here because a warm hit was + // already validated this session, saving the extra `open_index_file` IOP. + let load: ScalarIndexLoad = Box::pin({ + let index_store = index_store.clone(); + let frag_reuse_index = frag_reuse_index.clone(); + let index_cache = index_cache.clone(); + async move { + if index_details.type_url.ends_with("LabelListIndexDetails") { + validate_label_list_index_compatibility(dataset, column, index, &index_store) + .await?; + } - let index = plugin - .load_index(index_store, &index_details, frag_reuse_index, &index_cache) - .await?; + let index = plugin + .load_index(index_store, &index_details, frag_reuse_index, &index_cache) + .await?; - tracing::info!(target: TRACE_IO_EVENTS, index_uuid = %index_uuid, r#type = IO_TYPE_OPEN_SCALAR, index_type = index.index_type().to_string()); - metrics.record_index_load(); + tracing::info!(target: TRACE_IO_EVENTS, index_uuid = %index_uuid, r#type = IO_TYPE_OPEN_SCALAR, index_type = index.index_type().to_string()); + metrics.record_index_load(); + Ok(index) + } + }); - plugin.put_in_cache(&index_cache, index.clone()).await?; - Ok(index) + plugin + .get_or_insert_in_cache(index_store, frag_reuse_index, &index_cache, load) + .await } pub(crate) async fn infer_scalar_index_details( @@ -1038,6 +1055,51 @@ mod tests { } } + #[tokio::test] + async fn test_open_scalar_index_coalesces_concurrent_cold_opens() { + use crate::dataset::builder::DatasetBuilder; + use arrow::datatypes::Int64Type; + use futures::future::try_join_all; + use lance_index::metrics::LocalMetricsCollector; + use std::sync::atomic::Ordering; + + let dir = TempStrDir::default(); + let uri = dir.as_ref(); + let mut ds = lance_datagen::gen_batch() + .col("id", array::step::()) + .into_dataset(uri, FragmentCount::from(1), FragmentRowCount::from(200)) + .await + .unwrap(); + + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::BTree); + ds.create_index(&["id"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + + // Reopen so the index cache is cold for the concurrent opens below. + let ds = DatasetBuilder::from_uri(uri).load().await.unwrap(); + let id_field = ds.schema().field("id").unwrap().id; + let index_meta = ds + .load_indices() + .await + .unwrap() + .iter() + .find(|idx| idx.fields == [id_field]) + .cloned() + .expect("btree index on `id`"); + + // Eight concurrent cold opens; single-flight means the loader (which calls + // `record_index_load`) runs exactly once, not once per open. + let metrics = LocalMetricsCollector::default(); + let indices = + try_join_all((0..8).map(|_| open_scalar_index(&ds, "id", &index_meta, &metrics))) + .await + .unwrap(); + + assert_eq!(indices.len(), 8); + assert_eq!(metrics.index_loads.load(Ordering::Relaxed), 1); + } + #[tokio::test] async fn test_load_training_data_addr_sort() { // Create test data using lance_datagen @@ -1774,6 +1836,122 @@ mod tests { ); } + #[tokio::test] + async fn test_dataset_column_value_range() { + use crate::index::DatasetIndexExt; + use arrow::datatypes::Int64Type; + use datafusion::scalar::ScalarValue; + use lance_datagen::array; + use lance_index::IndexType; + use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams}; + + // 2 fragments x 5 rows: `id` and `other` both step 0..9. + let mut ds = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("other", array::step::()) + .into_ram_dataset(FragmentCount::from(2), FragmentRowCount::from(5)) + .await + .unwrap(); + + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap); + ds.create_index(&["id"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + + // Folded straight from the ZoneMap summaries: global [0, 9]. + assert_eq!( + ds.statistics().column_value_range("id").await.unwrap(), + Some((ScalarValue::Int64(Some(0)), ScalarValue::Int64(Some(9)))) + ); + // `other` has no ZoneMap index -> no precomputed range. + assert_eq!( + ds.statistics().column_value_range("other").await.unwrap(), + None + ); + } + + #[tokio::test] + async fn test_column_value_range_none_when_index_misses_fragments() { + use crate::index::DatasetIndexExt; + use arrow::datatypes::Int64Type; + use lance_datagen::{BatchCount, RowCount, array}; + use lance_index::IndexType; + use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams}; + + // Index covers the initial 2 fragments. + let mut ds = lance_datagen::gen_batch() + .col("id", array::step::()) + .into_ram_dataset(FragmentCount::from(2), FragmentRowCount::from(5)) + .await + .unwrap(); + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap); + ds.create_index(&["id"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + assert!( + ds.statistics() + .column_value_range("id") + .await + .unwrap() + .is_some() + ); + + // Append a fragment the index doesn't cover. Its rows may hold values + // outside the indexed range, so any non-None range would be a subset + // unsound to prune with -> must return None until the index is rebuilt. + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .into_reader_rows(RowCount::from(5), BatchCount::from(1)); + ds.append(reader, None).await.unwrap(); + + assert_eq!( + ds.statistics().column_value_range("id").await.unwrap(), + None + ); + } + + #[tokio::test] + async fn test_column_value_range_folds_multiple_segments() { + use crate::index::DatasetIndexExt; + use arrow::datatypes::Int64Type; + use datafusion::scalar::ScalarValue; + use lance_datagen::array; + use lance_index::IndexType; + use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams}; + + // 2 fragments x 5 rows (`id` steps 0..9). Build one ZoneMap segment per + // fragment, then commit them as a single multi-segment logical index. + let mut ds = lance_datagen::gen_batch() + .col("id", array::step::()) + .into_ram_dataset(FragmentCount::from(2), FragmentRowCount::from(5)) + .await + .unwrap(); + + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap); + let columns = ["id"]; + let mut segments = Vec::new(); + for fragment_id in [0_u32, 1_u32] { + let mut builder = ds + .create_index_builder(&columns, IndexType::Scalar, ¶ms) + .name("id_idx".to_string()) + .fragments(vec![fragment_id]); + segments.push(builder.execute_uncommitted().await.unwrap()); + } + // execute_uncommitted yields ready IndexMetadata segments (IntoIndexSegment); + // build_all is vector-only, so commit the per-fragment segments directly. + ds.commit_existing_index_segments("id_idx", "id", segments) + .await + .unwrap(); + + // Two disjoint segments jointly cover every live fragment -> the fold + // spans both and yields the global range, not None. + assert_eq!(ds.load_indices_by_name("id_idx").await.unwrap().len(), 2); + assert_eq!( + ds.statistics().column_value_range("id").await.unwrap(), + Some((ScalarValue::Int64(Some(0)), ScalarValue::Int64(Some(9)))) + ); + } + #[tokio::test] async fn test_zonemap_index_then_deletion() { // Tests the opposite scenario: create index FIRST, then perform deletions @@ -2126,4 +2304,70 @@ mod tests { "Should have 0 rows with value='banana' after deletion" ); } + + // End-to-end: create index → delete a whole fragment → search (via index) → + // update index → search again. No deleted rows should ever appear. + #[tokio::test] + async fn test_zonemap_search_with_deleted_fragment_before_and_after_update() { + use arrow::datatypes::Int32Type; + use lance_datagen::array; + use lance_index::IndexType; + use lance_index::optimize::OptimizeOptions; + use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams}; + + // 3 fragments × 10 rows: id 0-9 (frag 0), 10-19 (frag 1), 20-29 (frag 2). + let mut ds = lance_datagen::gen_batch() + .col("id", array::step::()) + .into_ram_dataset(FragmentCount::from(3), FragmentRowCount::from(10)) + .await + .unwrap(); + + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap); + ds.create_index(&["id"], IndexType::Scalar, None, ¶ms, false) + .await + .unwrap(); + + // Delete the middle fragment entirely. + ds.delete("id >= 10 AND id < 20").await.unwrap(); + + // Helper: run a filter scan and return the sorted id values. + async fn live_ids(ds: &crate::Dataset) -> Vec { + let batch = ds + .scan() + .filter("id >= 0") + .unwrap() + .try_into_batch() + .await + .unwrap(); + let mut ids: Vec = batch["id"] + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + ids.sort_unstable(); + ids + } + + // --- Before index update --- + let ids = live_ids(&ds).await; + assert_eq!(ids.len(), 20, "expected 20 live rows before index update"); + assert!( + ids.iter().all(|&id| !(10..20).contains(&id)), + "deleted fragment rows (id 10-19) must not appear before index update; got: {:?}", + ids + ); + + // Update the zone map index to reflect the deletion. + ds.optimize_indices(&OptimizeOptions::new()).await.unwrap(); + + // --- After index update --- + let ids = live_ids(&ds).await; + assert_eq!(ids.len(), 20, "expected 20 live rows after index update"); + assert!( + ids.iter().all(|&id| !(10..20).contains(&id)), + "deleted fragment rows (id 10-19) must not appear after index update; got: {:?}", + ids + ); + } } diff --git a/rust/lance/src/index/scalar/fmindex.rs b/rust/lance/src/index/scalar/fmindex.rs index 6c33498d929..32684ebf9ab 100644 --- a/rust/lance/src/index/scalar/fmindex.rs +++ b/rust/lance/src/index/scalar/fmindex.rs @@ -14,6 +14,12 @@ use crate::{Dataset, Error, Result}; /// to combine two BWT structures. This function re-reads text data from the /// dataset for all fragments covered by the source segments and builds a fresh /// FM-Index over the combined data. +/// +/// As an exception, a single source segment whose fragment coverage is still +/// fully live is returned as-is: rebuilding it would produce equivalent content +/// over the same fragments, which makes distributed builds (one uncommitted +/// segment per worker, merged 1:1 into final segments) pay for every segment +/// twice. pub(in crate::index) async fn merge_segments( dataset: &Dataset, segments: Vec, @@ -75,6 +81,10 @@ pub(in crate::index) async fn merge_segments( }); } + if segments.len() == 1 && segments[0].fragment_bitmap.as_ref() == Some(&fragment_bitmap) { + return Ok(segments.into_iter().next().unwrap()); + } + let fragment_ids: Vec = fragment_bitmap.iter().collect(); let new_uuid = Uuid::new_v4(); diff --git a/rust/lance/src/index/scalar/label_list.rs b/rust/lance/src/index/scalar/label_list.rs new file mode 100644 index 00000000000..27bc49643bb --- /dev/null +++ b/rust/lance/src/index/scalar/label_list.rs @@ -0,0 +1,148 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use lance_index::metrics::NoOpMetricsCollector; +use lance_index::scalar::IndexStore; +use lance_index::scalar::label_list::{ + BITMAP_LOOKUP_NAME, LABEL_LIST_NULLS_METADATA_KEY, LABEL_LIST_NULLS_MIN_VERSION, LabelListIndex, +}; +use lance_index::scalar::lance_format::LanceIndexStore; +use lance_table::format::IndexMetadata; +use roaring::RoaringBitmap; +use std::sync::Arc; +use uuid::Uuid; + +use crate::{Dataset, Error, Result, dataset::index::LanceIndexStoreExt}; + +async fn validate_nullable_segment_for_merge( + dataset: &Dataset, + field_id: i32, + segment: &IndexMetadata, +) -> Result<()> { + let field = dataset.schema().field_by_id(field_id).ok_or_else(|| { + Error::invalid_input(format!( + "merge_existing_index_segments: field id {} does not exist", + field_id + )) + })?; + + if !field.nullable { + return Ok(()); + } + + if segment.index_version < LABEL_LIST_NULLS_MIN_VERSION { + return Err(Error::invalid_input(format!( + "Cannot merge nullable LabelList segment {} because it was created before list-null metadata was required. Rebuild the segment instead.", + segment.uuid + ))); + } + + let index_store = LanceIndexStore::from_dataset_for_existing(dataset, segment).await?; + let reader = index_store.open_index_file(BITMAP_LOOKUP_NAME).await?; + if !reader + .schema() + .metadata + .contains_key(LABEL_LIST_NULLS_METADATA_KEY) + { + return Err(Error::invalid_input(format!( + "Cannot merge nullable LabelList segment {} because it is missing required metadata key {}. Rebuild the segment instead.", + segment.uuid, LABEL_LIST_NULLS_METADATA_KEY + ))); + } + + Ok(()) +} + +/// Merge one caller-defined group of source LabelList segments into a single segment. +/// +/// A LabelList index is a bitmap over the unnested list values plus a `list_nulls` +/// row set, so segments are recombined by unioning the underlying bitmap states and +/// null sets (see [`lance_index::scalar::label_list::merge_label_list_indices`]) +/// rather than rebuilding from source text. +pub(in crate::index) async fn merge_segments( + dataset: &Dataset, + segments: Vec, +) -> Result { + if segments.is_empty() { + return Err(Error::index("No segment metadata was provided".to_string())); + } + + let field_id = *segments[0].fields.first().ok_or_else(|| { + Error::invalid_input(format!( + "CreateIndex: segment {} is missing field ids", + segments[0].uuid + )) + })?; + let field_path = dataset.schema().field_path(field_id)?; + + let dataset_fragments = dataset.fragment_bitmap.as_ref(); + let mut effective_old_frags = RoaringBitmap::new(); + let mut deleted_old_frags = RoaringBitmap::new(); + for segment in &segments { + if segment.fragment_bitmap.is_none() { + return Err(Error::invalid_input(format!( + "CreateIndex: segment {} is missing fragment coverage", + segment.uuid + ))); + } + if let Some(effective) = segment.effective_fragment_bitmap(dataset_fragments) { + effective_old_frags |= effective; + } + if let Some(deleted) = segment.deleted_fragment_bitmap(dataset_fragments) { + deleted_old_frags |= deleted; + } + } + + let fragment_bitmap = effective_old_frags.clone(); + let old_data_filter = if deleted_old_frags.is_empty() { + None + } else { + crate::index::append::build_old_data_filter( + dataset, + &effective_old_frags, + &deleted_old_frags, + ) + .await? + }; + + let mut source_indices = Vec::with_capacity(segments.len()); + for segment in &segments { + validate_nullable_segment_for_merge(dataset, field_id, segment).await?; + let scalar_index = + super::open_scalar_index(dataset, &field_path, segment, &NoOpMetricsCollector).await?; + let label_list_index = scalar_index + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::index(format!( + "merge_existing_index_segments: expected label list segment {}, got {:?}", + segment.uuid, + scalar_index.index_type() + )) + })?; + source_indices.push(Arc::new(label_list_index.clone())); + } + + let new_uuid = Uuid::new_v4(); + let new_store = LanceIndexStore::from_dataset_for_new(dataset, &new_uuid)?; + let created_index = lance_index::scalar::label_list::merge_label_list_indices( + &source_indices, + &new_store, + old_data_filter, + lance_index::progress::noop_progress(), + ) + .await?; + + Ok(IndexMetadata { + uuid: new_uuid, + fields: vec![field_id], + dataset_version: dataset.manifest.version, + fragment_bitmap: Some(fragment_bitmap), + index_details: Some(Arc::new(created_index.index_details)), + index_version: created_index.index_version as i32, + created_at: Some(chrono::Utc::now()), + base_id: None, + files: Some(created_index.files), + ..segments[0].clone() + }) +} diff --git a/rust/lance/src/index/scalar_logical.rs b/rust/lance/src/index/scalar_logical.rs index 8ef86a6cb5f..d4a631fa2db 100644 --- a/rust/lance/src/index/scalar_logical.rs +++ b/rust/lance/src/index/scalar_logical.rs @@ -3,6 +3,7 @@ //! Query-time logical views over scalar index segments. +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::any::Any; use std::sync::Arc; @@ -135,13 +136,19 @@ impl ScalarIndex for LogicalScalarIndex { combine_search_results(results) } + fn results_are_row_addresses(&self) -> bool { + // All segments of a logical index share the same underlying index type, + // so they agree on the result domain. + self.segments[0].results_are_row_addresses() + } + fn can_remap(&self) -> bool { false } async fn remap( &self, - _mapping: &std::collections::HashMap>, + _mapping: &RowAddrRemap, _dest_store: &dyn lance_index::scalar::IndexStore, ) -> Result { Err(Error::invalid_input(format!( @@ -1460,4 +1467,188 @@ mod tests { "rows from live fragment should still be searchable" ); } + + #[tokio::test] + async fn test_fmindex_merge_single_segment_passthrough() { + let test_dir = TempStrDir::default(); + + let schema = Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new( + "text", + arrow_schema::DataType::Utf8, + false, + )])); + let write_params = crate::dataset::write::WriteParams { + max_rows_per_file: 4, + ..Default::default() + }; + let batches = vec![ + arrow_array::RecordBatch::try_new( + schema.clone(), + vec![Arc::new(arrow_array::StringArray::from(vec![ + "alpha beta gamma delta", + "beta gamma delta epsilon", + "gamma delta epsilon zeta", + "delta epsilon zeta eta", + "epsilon zeta eta theta", + "zeta eta theta iota", + "eta theta iota kappa", + "theta iota kappa lambda", + ]))], + ) + .unwrap(), + ]; + let reader = + arrow_array::RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let mut dataset = Dataset::write(reader, test_dir.as_str(), Some(write_params)) + .await + .unwrap(); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 2); + let fragment_ids: Vec = fragments.iter().map(|f| f.id() as u32).collect(); + + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::Fm); + let segment = CreateIndexBuilder::new(&mut dataset, &["text"], IndexType::Fm, ¶ms) + .name("text_fmindex_single".to_string()) + .fragments(fragment_ids.clone()) + .execute_uncommitted() + .await + .unwrap(); + let source_uuid = segment.uuid; + + // A single segment whose coverage is fully live is reused, not rebuilt. + let merged = dataset + .merge_existing_index_segments(vec![segment]) + .await + .unwrap(); + assert_eq!( + merged.uuid, source_uuid, + "single-segment merge with unchanged coverage should reuse the segment" + ); + assert_eq!( + merged + .fragment_bitmap + .as_ref() + .unwrap() + .iter() + .collect::>(), + fragment_ids + ); + + dataset + .commit_existing_index_segments("text_fmindex_single", "text", vec![merged]) + .await + .unwrap(); + + let logical = open_named_scalar_index( + &dataset, + "text", + "text_fmindex_single", + &NoOpMetricsCollector, + ) + .await + .unwrap(); + assert_eq!(logical.index_type(), IndexType::Fm); + + let query = lance_index::scalar::TextQuery::StringContains("delta".to_string()); + let result = logical.search(&query, &NoOpMetricsCollector).await.unwrap(); + let row_addrs = match result { + SearchResult::Exact(row_addrs) => row_addrs, + other => panic!("expected exact result from fmindex, got {:?}", other), + }; + assert_eq!(row_addrs.true_rows().row_addrs().unwrap().count(), 4); + } + + #[tokio::test] + async fn test_fmindex_merge_single_segment_rebuilds_when_coverage_shrinks() { + let test_dir = TempStrDir::default(); + + let schema = Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new( + "text", + arrow_schema::DataType::Utf8, + false, + )])); + let write_params = crate::dataset::write::WriteParams { + max_rows_per_file: 4, + enable_stable_row_ids: true, + ..Default::default() + }; + let batches = vec![ + arrow_array::RecordBatch::try_new( + schema.clone(), + vec![Arc::new(arrow_array::StringArray::from(vec![ + "alpha beta gamma", + "beta gamma delta", + "gamma delta epsilon", + "delta epsilon zeta", + "epsilon zeta eta", + "zeta eta theta", + "eta theta iota", + "theta iota kappa", + ]))], + ) + .unwrap(), + ]; + let reader = + arrow_array::RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let mut dataset = Dataset::write(reader, test_dir.as_str(), Some(write_params)) + .await + .unwrap(); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 2); + + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::Fm); + let segment = CreateIndexBuilder::new(&mut dataset, &["text"], IndexType::Fm, ¶ms) + .name("text_fmindex_shrink".to_string()) + .fragments(fragments.iter().map(|f| f.id() as u32).collect()) + .execute_uncommitted() + .await + .unwrap(); + let source_uuid = segment.uuid; + + // Retire fragment 0: delete its rows and compact it away. + dataset.delete("text = 'alpha beta gamma'").await.unwrap(); + dataset.delete("text = 'beta gamma delta'").await.unwrap(); + crate::dataset::optimize::compact_files( + &mut dataset, + crate::dataset::optimize::CompactionOptions { + target_rows_per_fragment: 4, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + + let live_frags: RoaringBitmap = dataset + .get_fragments() + .iter() + .map(|f| f.id() as u32) + .collect(); + assert!( + !live_frags.contains(0), + "compaction should retire fragment 0" + ); + assert!(live_frags.contains(1), "fragment 1 should stay live"); + + // Coverage shrank, so even a single segment must be rebuilt. + let merged = dataset + .merge_existing_index_segments(vec![segment]) + .await + .unwrap(); + assert_ne!( + merged.uuid, source_uuid, + "shrunk coverage must trigger a rebuild" + ); + assert_eq!( + merged + .fragment_bitmap + .as_ref() + .unwrap() + .iter() + .collect::>(), + vec![1] + ); + } } diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index af48bc94c41..c9920fb8adf 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -4,6 +4,7 @@ //! Vector Index for Fast Approximate Nearest Neighbor (ANN) Search //! +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::sync::Arc; use std::{any::Any, collections::HashMap}; @@ -1487,14 +1488,14 @@ pub(crate) async fn build_empty_vector_index( )) } -#[instrument(level = "debug", skip_all, fields(old_uuid = old_uuid.to_string(), new_uuid = new_uuid.to_string(), num_rows = mapping.len()))] +#[instrument(level = "debug", skip_all, fields(old_uuid = old_uuid.to_string(), new_uuid = new_uuid.to_string()))] pub(crate) async fn remap_vector_index( dataset: Arc, column: &str, old_uuid: &Uuid, new_uuid: &Uuid, old_metadata: &IndexMetadata, - mapping: &HashMap>, + mapping: &RowAddrRemap, ) -> Result> { let old_index = dataset .open_vector_index(column, old_uuid, &NoOpMetricsCollector) diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index 1e4fec8c762..bf968d9744f 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::collections::HashSet; use std::sync::{Arc, Mutex}; use std::{collections::HashMap, pin::Pin}; @@ -332,7 +333,7 @@ impl IvfIndexBuilder }) } - pub async fn remap(&mut self, mapping: &HashMap>) -> Result> { + pub async fn remap(&mut self, mapping: &RowAddrRemap) -> Result> { if self.existing_indices.is_empty() { return Err(Error::invalid_input( "No existing indices available for remapping", diff --git a/rust/lance/src/index/vector/fixture_test.rs b/rust/lance/src/index/vector/fixture_test.rs index 1b82a7f6941..a03cc313466 100644 --- a/rust/lance/src/index/vector/fixture_test.rs +++ b/rust/lance/src/index/vector/fixture_test.rs @@ -5,10 +5,10 @@ #[cfg(test)] mod test { + use lance_core::utils::row_addr_remap::RowAddrRemap; use std::{ any::Any, cell::OnceCell, - collections::HashMap, sync::{Arc, Mutex}, }; @@ -149,7 +149,7 @@ mod test { todo!("this method is for only IVF_HNSW_* index"); } - async fn remap(&mut self, _mapping: &HashMap>) -> Result<()> { + async fn remap(&mut self, _mapping: &RowAddrRemap) -> Result<()> { Ok(()) } diff --git a/rust/lance/src/index/vector/hamming.rs b/rust/lance/src/index/vector/hamming.rs index ba6ea98c42d..c5c8c7d6dc8 100644 --- a/rust/lance/src/index/vector/hamming.rs +++ b/rust/lance/src/index/vector/hamming.rs @@ -5,123 +5,215 @@ //! //! This module provides functionality to perform pairwise hamming distance //! computation and clustering on specific partitions of IVF_FLAT indices. - +//! +//! A logical IVF_FLAT index may consist of multiple physical segments (e.g. +//! delta segments created by `optimize_indices` in append mode, or segments +//! committed by distributed index builds). All segments of one logical index +//! are assumed to share the same global IVF centroids, so one partition id +//! refers to the same centroid region in every segment; this is validated and +//! an error is returned if the centroids differ. + +use std::sync::Arc; use std::time::Instant; -use arrow_array::RecordBatchReader; use arrow_array::cast::AsArray; use arrow_array::types::UInt64Type; +use arrow_array::{Array, FixedSizeListArray, RecordBatchReader}; use arrow_schema::DataType; use lance_core::{Error, Result}; use lance_index::metrics::NoOpMetricsCollector; use lance_index::vector::VectorIndex; use lance_index::vector::flat::index::{FlatBinQuantizer, FlatIndex}; use lance_index::vector::flat::storage::FLAT_COLUMN; +use lance_index::vector::ivf::storage::IvfModel; use lance_index::vector::storage::VectorStore; use lance_linalg::distance::{ ClusteringResult, cluster_pairwise_result, extract_hashes_from_fixed_list, pairwise_hamming_distance_parallel, }; +use lance_table::format::IndexMetadata; use rand::rng; use rand::seq::index::sample; +use uuid::Uuid; use crate::dataset::Dataset; -use crate::index::{DatasetIndexExt, DatasetIndexInternalExt}; +use crate::index::{DatasetIndexExt, DatasetIndexInternalExt, filter_index_segments_by_ids}; use super::ivf::v2::IVFIndex; -/// Perform pairwise hamming distance clustering on a partition of an IVF_FLAT index. -/// -/// This function loads a specific partition from an IVF_FLAT index on a hash column, -/// computes pairwise hamming distances between all hashes in the partition, -/// filters by threshold, and clusters the results using union-find. -/// -/// # Arguments -/// -/// * `dataset` - The Lance dataset -/// * `index_name` - Name of the IVF_FLAT index on the hash column -/// * `partition_id` - The partition ID within the IVF_FLAT index -/// * `hamming_threshold` - Maximum hamming distance to consider as similar -/// -/// # Returns -/// -/// A `RecordBatchReader` yielding batches with columns: -/// - `representative`: UInt64 - The representative row ID for each cluster -/// - `duplicates`: `List` - List of duplicate row IDs in each cluster -/// -/// # Errors +/// One opened physical segment of a logical IVF_FLAT binary index. +struct HashIndexSegment { + metadata: IndexMetadata, + index: Arc, +} + +impl HashIndexSegment { + fn ivf_flat_bin(&self) -> &IVFIndex { + self.index + .as_any() + .downcast_ref::>() + .expect("segment type validated in open_hash_index_segments") + } +} + +/// Validate that a column stores 64-bit hashes as `FixedSizeList`. +fn validate_hash_column(column: &str, data_type: &DataType) -> Result<()> { + match data_type { + DataType::FixedSizeList(inner, 8) if *inner.data_type() == DataType::UInt8 => Ok(()), + DataType::FixedSizeList(inner, 8) => Err(Error::invalid_input(format!( + "Column '{}' must be FixedSizeList, got FixedSizeList<{:?}, 8>", + column, + inner.data_type() + ))), + _ => Err(Error::invalid_input(format!( + "Column '{}' must be FixedSizeList, got {:?}", + column, data_type + ))), + } +} + +/// Validate that every segment of a logical IVF index shares the same global +/// centroids, so one partition id refers to the same centroid region in each +/// segment. Fails if any segment has no centroids or diverging centroids. +fn validate_shared_centroids<'a>( + index_name: &str, + models: impl IntoIterator, +) -> Result<()> { + struct Reference<'a> { + uuid: Uuid, + num_partitions: usize, + centroids: &'a FixedSizeListArray, + } + + let mut reference: Option> = None; + for (uuid, model) in models { + let centroids = model.centroids_array().ok_or_else(|| { + Error::invalid_input(format!( + "Index '{}' segment {} has no IVF centroids; hamming clustering requires \ + segments built from a shared global IVF model", + index_name, uuid + )) + })?; + match &reference { + None => { + reference = Some(Reference { + uuid, + num_partitions: model.num_partitions(), + centroids, + }); + } + Some(reference) => { + if centroids.to_data() != reference.centroids.to_data() { + return Err(Error::invalid_input(format!( + "Index '{}' segments do not share the same global IVF centroids: \ + segment {} ({} partitions) differs from segment {} ({} partitions); \ + retrain the index to merge segments before hamming clustering", + index_name, + uuid, + model.num_partitions(), + reference.uuid, + reference.num_partitions + ))); + } + } + } + } + Ok(()) +} + +/// Open the physical segments of a logical IVF_FLAT binary index. /// -/// Returns an error if: -/// - The index doesn't exist or is not an IVF_FLAT index -/// - The indexed column has wrong type (must be `FixedSizeList`) -/// - The partition ID is out of range -pub async fn hamming_clustering_for_ivf_partition( +/// When `segment_ids` is `None` all segments are opened; otherwise only the +/// requested segments are opened and every requested id must exist. Validates +/// that all selected segments index the same `FixedSizeList` column, +/// are IVF_FLAT indices for binary data, and share the same global centroids. +async fn open_hash_index_segments( dataset: &Dataset, index_name: &str, - partition_id: usize, - hamming_threshold: u32, -) -> Result> { - // Load indices and find the IVF_FLAT index - let indices = dataset.load_indices().await?; - let index_meta = indices - .iter() - .find(|idx| idx.name == index_name) - .ok_or_else(|| { - Error::invalid_input(format!("Index '{}' not found on dataset", index_name)) - })?; + segment_ids: Option<&[Uuid]>, +) -> Result> { + let metadatas = dataset.load_indices_by_name(index_name).await?; + if metadatas.is_empty() { + return Err(Error::invalid_input(format!( + "Index '{}' not found on dataset", + index_name + ))); + } - // Get the column name from the index metadata - let schema = dataset.schema(); - let field_id = index_meta - .fields + let metadatas = match segment_ids { + None => metadatas, + Some(segment_ids) => { + if segment_ids.is_empty() { + return Err(Error::invalid_input(format!( + "Segment selection for index '{}' must not be empty; \ + omit index_segments to use all segments", + index_name + ))); + } + filter_index_segments_by_ids(index_name, metadatas, segment_ids)? + } + }; + + let fields = metadatas[0].fields.clone(); + if let Some(mismatched) = metadatas.iter().find(|meta| meta.fields != fields) { + return Err(Error::invalid_input(format!( + "Index '{}' segments cover different fields: segment {} covers {:?} \ + while segment {} covers {:?}", + index_name, metadatas[0].uuid, fields, mismatched.uuid, mismatched.fields + ))); + } + + let field_id = fields .first() .ok_or_else(|| Error::invalid_input(format!("Index '{}' has no fields", index_name)))?; + let schema = dataset.schema(); let field = schema.field_by_id(*field_id).ok_or_else(|| { Error::invalid_input(format!( "Field with id {} not found in schema for index '{}'", field_id, index_name )) })?; - let column = &field.name; + validate_hash_column(&field.name, &field.data_type())?; - // Check column is FixedSizeList - let data_type = field.data_type(); - match data_type { - DataType::FixedSizeList(inner, 8) => { - if *inner.data_type() != DataType::UInt8 { - return Err(Error::invalid_input(format!( - "Column '{}' must be FixedSizeList, got FixedSizeList<{:?}, 8>", - column, - inner.data_type() - ))); - } - } - _ => { + let mut segments = Vec::with_capacity(metadatas.len()); + for metadata in metadatas { + let index = dataset + .open_vector_index(&field.name, &metadata.uuid, &NoOpMetricsCollector) + .await?; + if index + .as_any() + .downcast_ref::>() + .is_none() + { return Err(Error::invalid_input(format!( - "Column '{}' must be FixedSizeList, got {:?}", - column, data_type + "Index '{}' segment {} is not an IVF_FLAT index for binary data", + index_name, metadata.uuid ))); } + segments.push(HashIndexSegment { metadata, index }); } - // Open the vector index - let index = dataset - .open_vector_index(column, &index_meta.uuid, &NoOpMetricsCollector) - .await?; + validate_shared_centroids( + index_name, + segments + .iter() + .map(|segment| (segment.metadata.uuid, segment.ivf_flat_bin().ivf_model())), + )?; - // Try to downcast to IVFIndex (IVF_FLAT for binary data) - let ivf_index = index - .as_any() - .downcast_ref::>() - .ok_or_else(|| { - Error::invalid_input(format!( - "Index '{}' is not an IVF_FLAT index for binary data", - index_name - )) - })?; + Ok(segments) +} + +async fn hamming_clustering_for_ivf_partition_impl( + dataset: &Dataset, + index_name: &str, + segment_ids: Option<&[Uuid]>, + partition_id: usize, + hamming_threshold: u32, +) -> Result> { + let segments = open_hash_index_segments(dataset, index_name, segment_ids).await?; - // Check partition ID is valid - let num_partitions = ivf_index.ivf_model().num_partitions(); + // All segments share centroids, so the partition count is uniform. + let num_partitions = segments[0].ivf_flat_bin().ivf_model().num_partitions(); if partition_id >= num_partitions { return Err(Error::invalid_input(format!( "Partition ID {} is out of range (0..{})", @@ -129,45 +221,49 @@ pub async fn hamming_clustering_for_ivf_partition( ))); } - // Load the partition storage - let storage = ivf_index.load_partition_storage(partition_id, None).await?; - - // Get row IDs - let row_id_slice: Vec = storage.row_ids().copied().collect(); - - if row_id_slice.is_empty() { - let empty = ClusteringResult { - clusters: Vec::new(), - }; - return Ok(empty.into_reader(None)); + // Concatenate the partition's row ids and hashes across segments; identical + // hashes land in the same partition of every segment, so one pairwise pass + // over the union finds cross-segment duplicates. + let mut all_row_ids: Vec = Vec::new(); + let mut all_hashes: Vec = Vec::new(); + for segment in &segments { + let storage = segment + .ivf_flat_bin() + .load_partition_storage(partition_id, None) + .await?; + all_row_ids.extend(storage.row_ids().copied()); + for batch in storage.to_batches()? { + let vectors = batch + .column_by_name(FLAT_COLUMN) + .ok_or_else(|| { + Error::invalid_input(format!("Column '{}' not found in storage", FLAT_COLUMN)) + })? + .as_fixed_size_list(); + all_hashes.extend(extract_hashes_from_fixed_list(vectors)?); + } + if all_row_ids.len() != all_hashes.len() { + return Err(Error::internal(format!( + "Index '{}' segment {} partition {}: row id count {} does not match hash count {}", + index_name, + segment.metadata.uuid, + partition_id, + all_row_ids.len(), + all_hashes.len() + ))); + } } - // Get vectors from the storage batches - let batches: Vec<_> = storage.to_batches()?.collect(); - if batches.is_empty() { + if all_row_ids.is_empty() { let empty = ClusteringResult { clusters: Vec::new(), }; return Ok(empty.into_reader(None)); } - // Extract the hash vectors from the FLAT_COLUMN - let mut all_hashes = Vec::new(); - for batch in &batches { - let vectors = batch - .column_by_name(FLAT_COLUMN) - .ok_or_else(|| { - Error::invalid_input(format!("Column '{}' not found in storage", FLAT_COLUMN)) - })? - .as_fixed_size_list(); - let hashes = extract_hashes_from_fixed_list(vectors)?; - all_hashes.extend(hashes); - } - // Compute pairwise hamming distances with threshold filtering let pairwise_result = pairwise_hamming_distance_parallel( &all_hashes, - Some(&row_id_slice), + Some(&all_row_ids), Some(hamming_threshold), ); @@ -177,60 +273,124 @@ pub async fn hamming_clustering_for_ivf_partition( Ok(clustering.into_reader(None)) } -/// Get partition statistics for an IVF_FLAT index. -pub async fn get_ivf_partition_info( +/// Perform pairwise hamming distance clustering on a partition of an IVF_FLAT index. +/// +/// This function loads a specific partition from every segment of an IVF_FLAT +/// index on a hash column, computes pairwise hamming distances between all +/// hashes in the combined partition, filters by threshold, and clusters the +/// results using union-find. See [`hamming_clustering_for_ivf_partition_segments`] +/// to restrict the computation to selected segments. +/// +/// # Arguments +/// +/// * `dataset` - The Lance dataset +/// * `index_name` - Name of the IVF_FLAT index on the hash column +/// * `partition_id` - The partition ID within the IVF_FLAT index +/// * `hamming_threshold` - Maximum hamming distance to consider as similar +/// +/// # Returns +/// +/// A `RecordBatchReader` yielding batches with columns: +/// - `representative`: UInt64 - The representative row ID for each cluster +/// - `duplicates`: `List` - List of duplicate row IDs in each cluster +/// +/// # Errors +/// +/// Returns an error if: +/// - The index doesn't exist or is not an IVF_FLAT index +/// - The indexed column has wrong type (must be `FixedSizeList`) +/// - The index segments do not share the same global IVF centroids +/// - The partition ID is out of range +pub async fn hamming_clustering_for_ivf_partition( dataset: &Dataset, index_name: &str, -) -> Result> { - let indices = dataset.load_indices().await?; - let index_meta = indices - .iter() - .find(|idx| idx.name == index_name) - .ok_or_else(|| { - Error::invalid_input(format!("Index '{}' not found on dataset", index_name)) - })?; - - // Get the column name from the index metadata - let schema = dataset.schema(); - let field_id = index_meta - .fields - .first() - .ok_or_else(|| Error::invalid_input(format!("Index '{}' has no fields", index_name)))?; - let field = schema.field_by_id(*field_id).ok_or_else(|| { - Error::invalid_input(format!( - "Field with id {} not found in schema for index '{}'", - field_id, index_name - )) - })?; - let column = &field.name; - - let index = dataset - .open_vector_index(column, &index_meta.uuid, &NoOpMetricsCollector) - .await?; - - let ivf_index = index - .as_any() - .downcast_ref::>() - .ok_or_else(|| { - Error::invalid_input(format!( - "Index '{}' is not an IVF_FLAT index for binary data", - index_name - )) - })?; + partition_id: usize, + hamming_threshold: u32, +) -> Result> { + hamming_clustering_for_ivf_partition_impl( + dataset, + index_name, + None, + partition_id, + hamming_threshold, + ) + .await +} - let num_partitions = ivf_index.ivf_model().num_partitions(); - let mut partition_infos = Vec::with_capacity(num_partitions); +/// Perform pairwise hamming distance clustering on a partition of selected +/// segments of an IVF_FLAT index. +/// +/// Same as [`hamming_clustering_for_ivf_partition`] but only the requested +/// physical segments contribute rows. Segment ids are the index UUIDs reported +/// by index descriptions; every requested id must belong to the named index and +/// the selection must not be empty. +pub async fn hamming_clustering_for_ivf_partition_segments( + dataset: &Dataset, + index_name: &str, + segment_ids: &[Uuid], + partition_id: usize, + hamming_threshold: u32, +) -> Result> { + hamming_clustering_for_ivf_partition_impl( + dataset, + index_name, + Some(segment_ids), + partition_id, + hamming_threshold, + ) + .await +} - for i in 0..num_partitions { - partition_infos.push(PartitionInfo { - partition_id: i, - size: ivf_index.ivf_model().partition_size(i), - }); +async fn get_ivf_partition_info_impl( + dataset: &Dataset, + index_name: &str, + segment_ids: Option<&[Uuid]>, +) -> Result> { + let segments = open_hash_index_segments(dataset, index_name, segment_ids).await?; + + let num_partitions = segments[0].ivf_flat_bin().ivf_model().num_partitions(); + let mut partition_infos: Vec = (0..num_partitions) + .map(|partition_id| PartitionInfo { + partition_id, + size: 0, + }) + .collect(); + for segment in &segments { + // Sizes come from the partition storage; the IVF model of a v3 index + // file does not carry partition lengths. + let index = segment.ivf_flat_bin(); + for info in partition_infos.iter_mut() { + info.size += index.partition_size(info.partition_id); + } } Ok(partition_infos) } +/// Get partition statistics for an IVF_FLAT index. +/// +/// Partition sizes are aggregated across all segments of the logical index. +/// See [`get_ivf_partition_info_segments`] to restrict the statistics to +/// selected segments. +pub async fn get_ivf_partition_info( + dataset: &Dataset, + index_name: &str, +) -> Result> { + get_ivf_partition_info_impl(dataset, index_name, None).await +} + +/// Get partition statistics for selected segments of an IVF_FLAT index. +/// +/// Same as [`get_ivf_partition_info`] but only the requested physical segments +/// contribute to the partition sizes. +pub async fn get_ivf_partition_info_segments( + dataset: &Dataset, + index_name: &str, + segment_ids: &[Uuid], +) -> Result> { + get_ivf_partition_info_impl(dataset, index_name, Some(segment_ids)).await +} + /// Information about an IVF partition. #[derive(Debug, Clone)] pub struct PartitionInfo { @@ -267,26 +427,7 @@ pub async fn hamming_clustering_for_sample( let field = schema.field(column).ok_or_else(|| { Error::invalid_input(format!("Column '{}' not found in dataset schema", column)) })?; - - // Check column is FixedSizeList - let data_type = field.data_type(); - match data_type { - DataType::FixedSizeList(inner, 8) => { - if *inner.data_type() != DataType::UInt8 { - return Err(Error::invalid_input(format!( - "Column '{}' must be FixedSizeList, got FixedSizeList<{:?}, 8>", - column, - inner.data_type() - ))); - } - } - _ => { - return Err(Error::invalid_input(format!( - "Column '{}' must be FixedSizeList, got {:?}", - column, data_type - ))); - } - } + validate_hash_column(column, &field.data_type())?; // Get total row count let total_rows: usize = dataset @@ -411,26 +552,7 @@ pub async fn hamming_clustering_for_range( let field = schema.field(column).ok_or_else(|| { Error::invalid_input(format!("Column '{}' not found in dataset schema", column)) })?; - - // Check column is FixedSizeList - let data_type = field.data_type(); - match data_type { - DataType::FixedSizeList(inner, 8) => { - if *inner.data_type() != DataType::UInt8 { - return Err(Error::invalid_input(format!( - "Column '{}' must be FixedSizeList, got FixedSizeList<{:?}, 8>", - column, - inner.data_type() - ))); - } - } - _ => { - return Err(Error::invalid_input(format!( - "Column '{}' must be FixedSizeList, got {:?}", - column, data_type - ))); - } - } + validate_hash_column(column, &field.data_type())?; // Get the fragment let fragment = dataset.get_fragment(fragment_id).ok_or_else(|| { @@ -764,6 +886,195 @@ mod tests { assert!(err.to_string().contains("not found"), "Error: {}", err); } + #[test] + fn test_validate_shared_centroids() { + use arrow_array::UInt8Array; + use lance_arrow::FixedSizeListArrayExt; + + fn model_from_bytes(bytes: Vec) -> IvfModel { + let centroids = + FixedSizeListArray::try_new_from_values(UInt8Array::from(bytes), 8).unwrap(); + IvfModel::new(centroids, None) + } + + let uuid_a = Uuid::new_v4(); + let uuid_b = Uuid::new_v4(); + + let model_a = model_from_bytes(vec![0u8; 16]); + let model_b = model_from_bytes(vec![0u8; 16]); + validate_shared_centroids("idx", [(uuid_a, &model_a), (uuid_b, &model_b)]).unwrap(); + + let mut diverged = vec![0u8; 16]; + diverged[0] = 1; + let model_c = model_from_bytes(diverged); + let err = + validate_shared_centroids("idx", [(uuid_a, &model_a), (uuid_b, &model_c)]).unwrap_err(); + assert!( + err.to_string() + .contains("do not share the same global IVF centroids"), + "{}", + err + ); + assert!(err.to_string().contains(&uuid_b.to_string()), "{}", err); + + let model_d = model_from_bytes(vec![0u8; 24]); + let err = + validate_shared_centroids("idx", [(uuid_a, &model_a), (uuid_b, &model_d)]).unwrap_err(); + assert!(err.to_string().contains("2 partitions"), "{}", err); + assert!(err.to_string().contains("3 partitions"), "{}", err); + + let err = validate_shared_centroids("idx", [(uuid_a, &IvfModel::empty())]).unwrap_err(); + assert!(err.to_string().contains("has no IVF centroids"), "{}", err); + } + + #[tokio::test] + async fn test_hamming_clustering_for_ivf_partition_multi_segment() { + use arrow_array::{FixedSizeListArray, RecordBatchIterator, UInt8Array}; + use arrow_schema::{Field, Schema}; + use lance_arrow::FixedSizeListArrayExt; + use lance_index::optimize::OptimizeOptions; + use lance_index::vector::ivf::IvfBuildParams; + use std::sync::Arc; + use tempfile::tempdir; + + fn hash_batch(schema: Arc, values: &[u64]) -> arrow_array::RecordBatch { + let mut bytes = Vec::with_capacity(values.len() * 8); + for value in values { + bytes.extend_from_slice(&value.to_le_bytes()); + } + let array = + FixedSizeListArray::try_new_from_values(UInt8Array::from(bytes), 8).unwrap(); + arrow_array::RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap() + } + + let schema = Arc::new(Schema::new(vec![Field::new( + "hash", + arrow_schema::DataType::FixedSizeList( + Arc::new(Field::new("item", arrow_schema::DataType::UInt8, true)), + 8, + ), + false, + )])); + + // 25 distinct hash values, two copies each; the same batch is written to + // fragment 0 and appended as fragment 1, so every value has duplicates + // in both fragments. + let values: Vec = (0..50) + .map(|i| ((i / 2) as u64).wrapping_mul(0x9E3779B97F4A7C15)) + .collect(); + let num_values = 25; + + let temp_dir = tempdir().unwrap(); + let uri = temp_dir.path().to_str().unwrap(); + let reader = RecordBatchIterator::new( + vec![Ok(hash_batch(schema.clone(), &values))], + schema.clone(), + ); + let mut dataset = crate::Dataset::write(reader, uri, None).await.unwrap(); + + let params = crate::index::vector::VectorIndexParams::with_ivf_flat_params( + lance_linalg::distance::MetricType::Hamming, + IvfBuildParams::new(4), + ); + dataset + .create_index( + &["hash"], + crate::index::IndexType::Vector, + Some("hash_idx".into()), + ¶ms, + false, + ) + .await + .unwrap(); + + let reader = RecordBatchIterator::new( + vec![Ok(hash_batch(schema.clone(), &values))], + schema.clone(), + ); + dataset.append(reader, None).await.unwrap(); + dataset + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap(); + + let segments = dataset.load_indices_by_name("hash_idx").await.unwrap(); + assert_eq!( + segments.len(), + 2, + "expected a delta segment after optimize append" + ); + + // Partition sizes aggregate across both segments. + let infos = get_ivf_partition_info(&dataset, "hash_idx").await.unwrap(); + assert_eq!(infos.len(), 4); + assert_eq!(infos.iter().map(|info| info.size).sum::(), 100); + + // Clustering each partition with threshold 0 must group all four copies + // of every value, including the copies in the appended fragment. + const FRAG1_START: u64 = 1 << 32; + let mut clusters = Vec::new(); + for partition_id in 0..4 { + let reader = + hamming_clustering_for_ivf_partition(&dataset, "hash_idx", partition_id, 0) + .await + .unwrap(); + clusters.extend(collect_clusters(reader)); + } + assert_eq!(clusters.len(), num_values); + for (representative, duplicates) in &clusters { + assert_eq!(duplicates.len(), 3); + assert!(*representative < FRAG1_START); + assert!( + duplicates.iter().any(|row_id| *row_id >= FRAG1_START), + "cluster {} should contain rows from the appended fragment", + representative + ); + } + + // Selecting only the original segment reproduces the single-segment scope. + let first_segment = segments + .iter() + .find(|meta| meta.fragment_bitmap.as_ref().unwrap().contains(0)) + .unwrap(); + let mut old_clusters = Vec::new(); + for partition_id in 0..4 { + let reader = hamming_clustering_for_ivf_partition_segments( + &dataset, + "hash_idx", + &[first_segment.uuid], + partition_id, + 0, + ) + .await + .unwrap(); + old_clusters.extend(collect_clusters(reader)); + } + assert_eq!(old_clusters.len(), num_values); + for (representative, duplicates) in &old_clusters { + assert_eq!(duplicates.len(), 1); + assert!(*representative < FRAG1_START); + assert!(duplicates.iter().all(|row_id| *row_id < FRAG1_START)); + } + let infos = get_ivf_partition_info_segments(&dataset, "hash_idx", &[first_segment.uuid]) + .await + .unwrap(); + assert_eq!(infos.iter().map(|info| info.size).sum::(), 50); + + // Invalid segment selections are rejected. + let err = hamming_clustering_for_ivf_partition_segments(&dataset, "hash_idx", &[], 0, 0) + .await + .err() + .unwrap(); + assert!(err.to_string().contains("must not be empty"), "{}", err); + let missing = Uuid::new_v4(); + let err = + hamming_clustering_for_ivf_partition_segments(&dataset, "hash_idx", &[missing], 0, 0) + .await + .err() + .unwrap(); + assert!(err.to_string().contains(&missing.to_string()), "{}", err); + } + #[tokio::test] async fn test_hamming_clustering_for_sample_integration() { use arrow_array::{FixedSizeListArray, RecordBatchIterator, UInt8Array}; diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 5d477a2e8dd..6e6810c454a 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -44,6 +44,7 @@ use futures::{ use io::write_hnsw_quantization_index_partitions; use lance_arrow::*; use lance_core::deepsize::DeepSizeOf; +use lance_core::utils::row_addr_remap::RowAddrRemap; use lance_core::{ Error, ROW_ID_FIELD, Result, cache::{LanceCache, UnsizedCacheKey, WeakLanceCache}, @@ -958,13 +959,12 @@ async fn optimize_ivf_hnsw_indices( writer.add_metadata(IVF_PARTITION_KEY, &hnsw_metadata_json.to_string()); ivf_mut.write(&mut writer).await?; - let index_size = writer.tell().await? as u64; - writer.finish().await?; + // `finish` writes the footer and returns the authoritative on-disk size. + let index_size = writer.finish().await?.size_bytes; // Write the aux file aux_ivf.write(&mut aux_writer).await?; - let aux_size = aux_writer.tell().await? as u64; - aux_writer.finish().await?; + let aux_size = aux_writer.finish().await?.size_bytes; Ok(( existing_indices.len() - start_pos, @@ -1244,7 +1244,7 @@ impl VectorIndex for IVFIndex { todo!("this method is for only IVF_HNSW_* index"); } - async fn remap(&mut self, _mapping: &HashMap>) -> Result<()> { + async fn remap(&mut self, _mapping: &RowAddrRemap) -> Result<()> { // This will be needed if we want to clean up IVF to allow more than just // one layer (e.g. IVF -> IVF -> PQ). We need to pass on the call to // remap to the lower layers. @@ -1719,7 +1719,7 @@ impl RemapPageTask { mut self, reader: Arc, index: &IVFIndex, - mapping: &HashMap>, + mapping: &RowAddrRemap, ) -> Result { let mut page = index .sub_index @@ -1765,7 +1765,7 @@ pub(crate) async fn remap_index_file_v3( dataset: &Dataset, new_uuid: &Uuid, index: Arc, - mapping: &HashMap>, + mapping: &RowAddrRemap, column: String, ) -> Result> { let dataset = dataset.clone(); @@ -1864,7 +1864,7 @@ pub(crate) async fn remap_index_file( new_uuid: &Uuid, old_version: u64, index: &IVFIndex, - mapping: &HashMap>, + mapping: &RowAddrRemap, name: String, column: String, transforms: Vec, @@ -1892,7 +1892,7 @@ pub(crate) async fn remap_index_file( let tasks = generate_remap_tasks(&index.ivf.offsets, &index.ivf.lengths)?; - let mut task_stream = stream::iter(tasks.into_iter()) + let mut task_stream = stream::iter(tasks) .map(|task| task.load_and_remap(reader.clone(), index, mapping)) .buffered(object_store.io_parallelism()); @@ -4267,19 +4267,26 @@ async fn train_streaming_coreset_ivf_model( let coreset_len = coreset.len(); let (coreset_data, coreset_weights, coreset_losses) = coreset.into_fsl_parts(dimension)?; - let weighted_hierarchical_params = WeightedHierarchicalKMeansParams { - dimension, - target_k: num_partitions, - metric_type: DistanceType::L2, - max_iters: params.max_iters, - on_progress: on_progress.clone(), + // Scope `weighted_hierarchical_params` so the `on_progress` clone it holds + // (which owns a clone of `progress_tx`) is dropped as soon as training + // returns. Otherwise it would outlive the `progress_worker.await` below, + // keeping a channel sender alive so `progress_rx.recv()` never returns + // `None` and the progress worker — and thus this function — hangs forever. + let mut centroids = { + let weighted_hierarchical_params = WeightedHierarchicalKMeansParams { + dimension, + target_k: num_partitions, + metric_type: DistanceType::L2, + max_iters: params.max_iters, + on_progress: on_progress.clone(), + }; + train_weighted_hierarchical_f32_kmeans( + &coreset_data, + &coreset_weights, + &coreset_losses, + &weighted_hierarchical_params, + )? }; - let mut centroids = train_weighted_hierarchical_f32_kmeans( - &coreset_data, - &coreset_weights, - &coreset_losses, - &weighted_hierarchical_params, - )?; let refine_iters = 3; if refine_iters > 0 { let refined = refine_weighted_f32_kmeans( @@ -5143,7 +5150,7 @@ mod tests { &new_uuid, dataset_mut.version().version, ivf_index, - &mapping, + &RowAddrRemap::direct(mapping), INDEX_NAME.to_string(), WellKnownIvfPqData::COLUMN.to_string(), vec![], @@ -5661,6 +5668,94 @@ mod tests { ); } + /// Regression test for a hang in the streaming *coreset* trainer + /// (`train_streaming_coreset_ivf_model`, taken when `num_partitions > 256`). + /// + /// That function spawns a `progress_worker` task that loops on + /// `progress_rx.recv()` and only terminates once every clone of the mpsc + /// sender is dropped. The `on_progress` closure owns a sender clone, and + /// `WeightedHierarchicalKMeansParams` used to retain an `on_progress` clone + /// that outlived the `progress_worker.await` at the end of the function. + /// With a live sender remaining, `recv()` never returned `None`, the worker + /// never finished, and the trainer hung forever after all compute was done. + /// + /// The build is wrapped in a timeout so the regression fails fast rather + /// than hanging the test process indefinitely. + #[tokio::test(flavor = "multi_thread")] + async fn test_streaming_coreset_ivf_training_terminates() { + use lance_index::progress::IndexBuildProgress; + use std::sync::atomic::{AtomicU64, Ordering}; + use std::time::Duration; + + #[derive(Debug, Default)] + struct CountingProgress { + progress_calls: AtomicU64, + } + + #[async_trait::async_trait] + impl IndexBuildProgress for CountingProgress { + async fn stage_start(&self, _: &str, _: Option, _: &str) -> Result<()> { + Ok(()) + } + async fn stage_progress(&self, _: &str, _: u64) -> Result<()> { + self.progress_calls.fetch_add(1, Ordering::Relaxed); + Ok(()) + } + async fn stage_complete(&self, _: &str) -> Result<()> { + Ok(()) + } + } + + const SMALL_DIM: usize = 8; + + let test_dir = TempStrDir::default(); + let uri = format!("{}/ds", test_dir.as_str()); + let reader = gen_batch() + .col("id", array::step::()) + .col( + "vector", + array::rand_vec::((SMALL_DIM as u32).into()), + ) + .into_reader_rows(RowCount::from(2048), BatchCount::from(4)); + let dataset = Dataset::write(reader, &uri, None).await.unwrap(); + + // > 256 partitions routes through `train_streaming_coreset_ivf_model`. + let mut params = IvfBuildParams::new(257); + params.sample_rate = 8; + params.streaming_sample_rate = Some(4); + params.streaming_refine_passes = 1; + params.max_iters = 2; + + let progress = Arc::new(CountingProgress::default()); + + let ivf_model = tokio::time::timeout( + Duration::from_secs(120), + build_ivf_model( + &dataset, + "vector", + SMALL_DIM, + MetricType::L2, + ¶ms, + None, + progress.clone(), + ), + ) + .await + .expect( + "streaming coreset IVF training hung: progress worker never terminated after training", + ) + .unwrap(); + + assert_eq!(ivf_model.num_partitions(), 257); + assert_eq!(ivf_model.dimension(), SMALL_DIM); + // The progress worker must have processed reports and then joined + // cleanly (proven by `build_ivf_model` returning at all). + assert!( + progress.progress_calls.load(Ordering::Relaxed) > 0, + "expected the progress worker to receive at least one report" + ); + } + #[test] fn test_fixed_training_ranges_are_sorted_and_bounded() { let ranges = generate_fixed_training_ranges(10_000, 1_234, 1_024, 16); @@ -6564,4 +6659,80 @@ mod tests { let indices = dataset.load_indices().await.unwrap(); assert!(!indices.is_empty(), "should have at least one index"); } + + #[tokio::test] + async fn test_optimize_ivf_hnsw_records_actual_file_sizes() { + // Regression test: `optimize_ivf_hnsw_indices` must record the on-disk + // file sizes, which are only known after `finish()` writes the footer. + const DIM: usize = 16; + let data = gen_batch() + .col( + "vec", + array::rand_vec::(Dimension::from(DIM as u32)), + ) + .into_batch_rows(RowCount::from(1_000)) + .unwrap(); + let schema = data.schema(); + + let mut dataset = InsertBuilder::new("memory://") + .execute(vec![data]) + .await + .unwrap(); + + // Legacy file version keeps the index on the v1 path that goes through + // `optimize_ivf_hnsw_indices`. + let mut params = VectorIndexParams::with_ivf_hnsw_sq_params( + MetricType::L2, + IvfBuildParams::new(2), + HnswBuildParams::default().num_edges(50), + SQBuildParams::default(), + ); + params.version(IndexFileVersion::Legacy); + dataset + .create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Append unindexed data so `optimize_indices` has something to merge. + let more = gen_batch() + .col( + "vec", + array::rand_vec::(Dimension::from(DIM as u32)), + ) + .into_batch_rows(RowCount::from(500)) + .unwrap(); + let more = RecordBatch::try_new(schema.clone(), more.columns().to_vec()).unwrap(); + let mut dataset = InsertBuilder::new(Arc::new(dataset)) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute(vec![more]) + .await + .unwrap(); + + dataset.optimize_indices(&Default::default()).await.unwrap(); + + let indices = dataset.load_indices().await.unwrap(); + let index = indices.first().expect("should have an index"); + let files = index + .files + .as_ref() + .expect("optimized index should record file sizes"); + assert!(!files.is_empty(), "index should record at least one file"); + + let indices_dir = dataset.indices_dir().join(index.uuid.to_string()); + for file in files { + let actual = dataset + .object_store + .size(&indices_dir.clone().join(file.path.as_str())) + .await + .unwrap(); + assert_eq!( + file.size_bytes, actual, + "recorded size for {} must match the on-disk size", + file.path + ); + } + } } diff --git a/rust/lance/src/index/vector/ivf/io.rs b/rust/lance/src/index/vector/ivf/io.rs index 56d220aeed2..f8f713a07af 100644 --- a/rust/lance/src/index/vector/ivf/io.rs +++ b/rust/lance/src/index/vector/ivf/io.rs @@ -44,6 +44,7 @@ use lance_table::format::SelfDescribingFileReader; use lance_table::io::manifest::ManifestDescribing; use object_store::path::Path; use tokio::sync::Semaphore; +use tokio::task::JoinHandle; use crate::Result; @@ -278,164 +279,234 @@ pub(super) async fn write_hnsw_quantization_index_partitions( } let object_store = ObjectStore::local(); - let mut part_files = Vec::with_capacity(ivf.num_partitions()); - let mut aux_part_files = Vec::with_capacity(ivf.num_partitions()); - let tmp_part_dir = Path::from_filesystem_path(TempStdDir::default())?; - let mut tasks = Vec::with_capacity(ivf.num_partitions()); - let sem = Arc::new(Semaphore::new(*HNSW_PARTITIONS_BUILD_PARALLEL)); - for part_id in 0..ivf.num_partitions() { - part_files.push(tmp_part_dir.clone().join(format!("hnsw_part_{}", part_id))); - aux_part_files.push( - tmp_part_dir - .clone() - .join(format!("hnsw_part_aux_{}", part_id)), - ); + // Partitions are staged in this scratch dir, then merged into the final index. + // Share the guard with every task via `Arc` so its `Drop` removes the dir only + // after the last task finishes -- never while one is still writing. + let tmp_part_dir_guard = Arc::new(TempStdDir::default()); + let tmp_part_dir = Path::from_filesystem_path(&**tmp_part_dir_guard)?; + + // `Option` per handle so the consume loop can `take()` each one, leaving the + // not-yet-consumed handles for the error-path drain. + let mut tasks: Vec>>> = + Vec::with_capacity(ivf.num_partitions()); + + let build_result: Result<(Vec, IvfModel)> = async { + let mut part_files = Vec::with_capacity(ivf.num_partitions()); + let mut aux_part_files = Vec::with_capacity(ivf.num_partitions()); + let sem = Arc::new(Semaphore::new(*HNSW_PARTITIONS_BUILD_PARALLEL)); + for part_id in 0..ivf.num_partitions() { + part_files.push(tmp_part_dir.clone().join(format!("hnsw_part_{}", part_id))); + aux_part_files.push( + tmp_part_dir + .clone() + .join(format!("hnsw_part_aux_{}", part_id)), + ); - let mut code_array: Vec> = vec![]; - let mut row_id_array: Vec> = vec![]; + let mut code_array: Vec> = vec![]; + let mut row_id_array: Vec> = vec![]; - // We don't transform vectors to SQ codes while shuffling, - // so we won't merge SQ codes from the stream. + // We don't transform vectors to SQ codes while shuffling, + // so we won't merge SQ codes from the stream. - if let Some(&previous_indices) = existing_indices.as_ref() { - for &idx in previous_indices.iter() { - let sub_index = idx - .load_partition(part_id, true, &NoOpMetricsCollector) - .await?; - let row_ids = Arc::new(UInt64Array::from_iter_values(sub_index.row_ids().cloned())); - row_id_array.push(row_ids); + if let Some(&previous_indices) = existing_indices.as_ref() { + for &idx in previous_indices.iter() { + let sub_index = idx + .load_partition(part_id, true, &NoOpMetricsCollector) + .await?; + let row_ids = + Arc::new(UInt64Array::from_iter_values(sub_index.row_ids().cloned())); + row_id_array.push(row_ids); + } } - } - - let code_column = match &quantizer { - Quantizer::Product(pq) => Some(pq.column()), - _ => None, - }; - merge_streams( - &mut streams_heap, - &mut new_streams, - part_id as u32, - code_column, - &mut code_array, - &mut row_id_array, - ) - .await?; - if row_id_array.is_empty() { - tasks.push(tokio::spawn(async { Ok(0) })); - continue; - } - - let (part_file, aux_part_file) = (&part_files[part_id], &aux_part_files[part_id]); - let part_writer = PreviousFileWriter::::try_new( - &object_store, - part_file, - Schema::try_from(writer.schema())?, - &Default::default(), - ) - .await?; + let code_column = match &quantizer { + Quantizer::Product(pq) => Some(pq.column()), + _ => None, + }; + merge_streams( + &mut streams_heap, + &mut new_streams, + part_id as u32, + code_column, + &mut code_array, + &mut row_id_array, + ) + .await?; - let aux_part_writer = match auxiliary_writer.as_ref() { - Some(writer) => Some( - PreviousFileWriter::::try_new( - &object_store, - aux_part_file, - Schema::try_from(writer.schema())?, - &Default::default(), - ) - .await?, - ), - None => None, - }; + if row_id_array.is_empty() { + tasks.push(Some(tokio::spawn(async { Ok(0) }))); + continue; + } - let dataset = dataset.clone(); - let column = column.to_owned(); - let hnsw_params = hnsw_params.clone(); - let quantizer = quantizer.clone(); - let sem = sem.clone(); - tasks.push(tokio::spawn(async move { - let _permit = sem.acquire().await.expect("semaphore error"); - - log::debug!("Building HNSW partition {}", part_id); - let result = build_hnsw_quantization_partition( - dataset, - &column, - distance_type, - hnsw_params, - part_writer, - aux_part_writer, - quantizer, - row_id_array, - code_array, + let (part_file, aux_part_file) = (&part_files[part_id], &aux_part_files[part_id]); + let part_writer = PreviousFileWriter::::try_new( + &object_store, + part_file, + Schema::try_from(writer.schema())?, + &Default::default(), ) - .await; - log::debug!("Finished building HNSW partition {}", part_id); - result - })); - } - - let mut aux_ivf = IvfModel::empty(); - let mut hnsw_metadata = Vec::with_capacity(ivf.num_partitions()); - for (part_id, task) in tasks.into_iter().enumerate() { - let offset = writer.len(); - let num_rows = task.await??; + .await?; - if num_rows == 0 { - ivf.add_partition(0); - aux_ivf.add_partition(0); - hnsw_metadata.push(HnswMetadata::default()); - continue; + let aux_part_writer = match auxiliary_writer.as_ref() { + Some(writer) => Some( + PreviousFileWriter::::try_new( + &object_store, + aux_part_file, + Schema::try_from(writer.schema())?, + &Default::default(), + ) + .await?, + ), + None => None, + }; + + let dataset = dataset.clone(); + let column = column.to_owned(); + let hnsw_params = hnsw_params.clone(); + let quantizer = quantizer.clone(); + let sem = sem.clone(); + let tmp_part_dir_guard = tmp_part_dir_guard.clone(); + tasks.push(Some(tokio::spawn(async move { + // Hold a guard clone so the scratch dir stays alive while this task writes. + let _tmp_part_dir_guard = tmp_part_dir_guard; + let _permit = sem.acquire().await.map_err(|err| { + Error::io(format!( + "failed to acquire HNSW partition build permit: {err}" + )) + })?; + + log::debug!("Building HNSW partition {}", part_id); + let result = build_hnsw_quantization_partition( + dataset, + &column, + distance_type, + hnsw_params, + part_writer, + aux_part_writer, + quantizer, + row_id_array, + code_array, + ) + .await; + log::debug!("Finished building HNSW partition {}", part_id); + result + }))); } - let (part_file, aux_part_file) = (&part_files[part_id], &aux_part_files[part_id]); - let part_reader = - PreviousFileReader::try_new_self_described(&object_store, part_file, None).await?; + let mut aux_ivf = IvfModel::empty(); + let mut hnsw_metadata = Vec::with_capacity(ivf.num_partitions()); + for (part_id, task) in tasks.iter_mut().enumerate() { + let task = task + .take() + .expect("each partition task is consumed exactly once"); + let offset = writer.len(); + let num_rows = task.await??; + + if num_rows == 0 { + ivf.add_partition(0); + aux_ivf.add_partition(0); + hnsw_metadata.push(HnswMetadata::default()); + continue; + } - let batches = futures::stream::iter(0..part_reader.num_batches()) - .map(|batch_id| { - part_reader.read_batch( - batch_id as i32, - ReadBatchParams::RangeFull, - part_reader.schema(), - ) - }) - .buffered(object_store.io_parallelism()) - .try_collect::>() - .await?; - writer.write(&batches).await?; - - ivf.add_partition((writer.len() - offset) as u32); - hnsw_metadata.push(serde_json::from_str( - part_reader.schema().metadata[HNSW::metadata_key()].as_str(), - )?); - std::mem::drop(part_reader); - object_store.delete(part_file).await?; - - if let Some(aux_writer) = auxiliary_writer.as_mut() { - let aux_part_reader = - PreviousFileReader::try_new_self_described(&object_store, aux_part_file, None) - .await?; + let (part_file, aux_part_file) = (&part_files[part_id], &aux_part_files[part_id]); + let part_reader = + PreviousFileReader::try_new_self_described(&object_store, part_file, None).await?; - let batches = futures::stream::iter(0..aux_part_reader.num_batches()) + let batches = futures::stream::iter(0..part_reader.num_batches()) .map(|batch_id| { - aux_part_reader.read_batch( + part_reader.read_batch( batch_id as i32, ReadBatchParams::RangeFull, - aux_part_reader.schema(), + part_reader.schema(), ) }) .buffered(object_store.io_parallelism()) .try_collect::>() .await?; - std::mem::drop(aux_part_reader); - object_store.delete(aux_part_file).await?; + writer.write(&batches).await?; + + ivf.add_partition((writer.len() - offset) as u32); + hnsw_metadata.push(serde_json::from_str( + part_reader.schema().metadata[HNSW::metadata_key()].as_str(), + )?); + std::mem::drop(part_reader); + object_store.delete(part_file).await?; + + if let Some(aux_writer) = auxiliary_writer.as_mut() { + let aux_part_reader = + PreviousFileReader::try_new_self_described(&object_store, aux_part_file, None) + .await?; + + let batches = futures::stream::iter(0..aux_part_reader.num_batches()) + .map(|batch_id| { + aux_part_reader.read_batch( + batch_id as i32, + ReadBatchParams::RangeFull, + aux_part_reader.schema(), + ) + }) + .buffered(object_store.io_parallelism()) + .try_collect::>() + .await?; + std::mem::drop(aux_part_reader); + object_store.delete(aux_part_file).await?; + + aux_writer.write(&batches).await?; + aux_ivf.add_partition(num_rows as u32); + } + } + + Ok((hnsw_metadata, aux_ivf)) + } + .await; + + // On error, abort and await the partition builds we never consumed so none + // keep running in the background; see `drain_partition_tasks`. + if build_result.is_err() { + for err in drain_partition_tasks(&mut tasks).await { + log::warn!( + "HNSW partition build task failed while draining after an earlier error: {err}" + ); + } + } + + build_result +} - aux_writer.write(&batches).await?; - aux_ivf.add_partition(num_rows as u32); +/// Abort and await every still-outstanding partition-build task, returning the +/// non-cancellation errors they surfaced. +/// +/// A dropped [`JoinHandle`] detaches its task, so every handle is aborted up front +/// before any is awaited: otherwise a task slow to observe its own cancellation +/// would keep running -- and keep writing into the scratch dir -- while an earlier +/// handle is still being awaited. Awaiting then resolves only once each task has +/// actually stopped. Cancellation errors are the expected result of the abort and +/// dropped; failures and panics from tasks that had already finished before the +/// abort are returned so the caller can surface them (a task still in flight is +/// cancelled, so this best-effort drain only reports errors already produced). +async fn drain_partition_tasks(tasks: &mut [Option>>]) -> Vec { + for task in tasks.iter() { + if let Some(handle) = task.as_ref() { + handle.abort(); } } - Ok((hnsw_metadata, aux_ivf)) + let mut errors = Vec::with_capacity(tasks.len()); + for task in tasks.iter_mut() { + let Some(handle) = task.take() else { + continue; + }; + match handle.await { + Ok(Ok(_)) => {} + Ok(Err(e)) => errors.push(e), + Err(join_err) if join_err.is_cancelled() => {} + Err(join_err) => errors.push(Error::io(format!( + "HNSW partition build task panicked: {join_err}" + ))), + } + } + errors } #[allow(clippy::too_many_arguments)] @@ -473,24 +544,29 @@ async fn build_hnsw_quantization_partition( let build_hnsw = build_and_write_hnsw(vectors.clone(), (*hnsw_params).clone(), metric_type, writer); + // Build PQ storage as a child future, joined below: it writes `aux_writer`'s + // file into the scratch dir, so it is cancelled together with this task when the + // error-path drain aborts it, and the join surfaces its errors. let build_store = match quantizer { Quantizer::Flat(_) => { return Err(Error::index( "Flat quantizer is not supported for IVF_HNSW".to_string(), )); } - Quantizer::Product(pq) => tokio::spawn(build_and_write_pq_storage( - metric_type, - row_ids, - code_array, - pq, - aux_writer.unwrap(), - )), - - _ => unreachable!("IVF_HNSW_SQ has been moved to v2 index builder"), + Quantizer::Product(pq) => { + let aux_writer = aux_writer.ok_or_else(|| { + Error::index("IVF_HNSW_PQ requires an auxiliary writer for PQ storage".to_string()) + })?; + build_and_write_pq_storage(metric_type, row_ids, code_array, pq, aux_writer) + } + _ => { + return Err(Error::index( + "IVF_HNSW_SQ is not supported in the legacy HNSW partition writer".to_string(), + )); + } }; - let index_rows = futures::join!(build_hnsw, build_store).0?; + let (index_rows, ()) = futures::try_join!(build_hnsw, build_store)?; assert!( index_rows >= num_rows, "index rows {} must be greater than or equal to num rows {}", @@ -509,7 +585,7 @@ async fn build_and_write_hnsw( let batch = params.build(vectors, distance_type).await?.to_batch()?; let metadata = batch.schema_ref().metadata().clone(); writer.write_record_batch(batch).await?; - writer.finish_with_metadata(&metadata).await + Ok(writer.finish_with_metadata(&metadata).await?.num_rows as usize) } async fn build_and_write_pq_storage( @@ -530,6 +606,9 @@ async fn build_and_write_pq_storage( mod tests { use super::*; + use std::path::PathBuf; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use crate::Dataset; use crate::index::vector::ivf::v2; use crate::index::{DatasetIndexExt, DatasetIndexInternalExt, vector::VectorIndexParams}; @@ -538,6 +617,8 @@ mod tests { use lance_core::utils::tempfile::TempStrDir; use lance_index::IndexType; use lance_index::metrics::NoOpMetricsCollector; + use lance_index::vector::ivf::IvfBuildParams; + use lance_index::vector::pq::PQBuildParams; use lance_testing::datagen::generate_random_array; #[tokio::test] @@ -589,4 +670,248 @@ mod tests { //let indices = /ds. } + + /// The scratch dir must outlive every partition task: dropping the caller-side + /// guard while tasks are still in flight must not remove it, because each task + /// still holds an `Arc` clone of the guard. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_scratch_dir_outlives_partition_tasks() { + let tmp_part_dir_guard = Arc::new(TempStdDir::default()); + let scratch_path = tmp_part_dir_guard.to_path_buf(); + + // Park each task until the caller-side guard is dropped, so the dir's + // survival is attributable solely to the clones the tasks still hold. + let running = Arc::new(AtomicUsize::new(0)); + let released = Arc::new(AtomicBool::new(false)); + + const NUM_TASKS: usize = 3; + let mut tasks = Vec::with_capacity(NUM_TASKS); + for _ in 0..NUM_TASKS { + let task_guard = tmp_part_dir_guard.clone(); + let running = running.clone(); + let released = released.clone(); + tasks.push(tokio::spawn(async move { + running.fetch_add(1, Ordering::SeqCst); + while !released.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + // Still holding `task_guard`, so the directory must be live. + task_guard.exists() + })); + } + + // Drop the caller-side guard only once every task is parked holding its clone. + while running.load(Ordering::SeqCst) < NUM_TASKS { + tokio::task::yield_now().await; + } + drop(tmp_part_dir_guard); + assert!( + scratch_path.exists(), + "scratch dir removed while partition tasks still held the guard" + ); + + released.store(true, Ordering::SeqCst); + for task in tasks { + assert!( + task.await.unwrap(), + "a partition task observed its scratch dir already removed" + ); + } + assert!( + !scratch_path.exists(), + "scratch dir was not removed after the last task's guard clone dropped" + ); + } + + /// The drain step must outlive every spawned task: it aborts and awaits each + /// one so that, once it returns, no task is still running against the scratch + /// directory. It must also surface late failures rather than swallow them. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_drain_partition_tasks_waits_and_reports_errors() { + // Scratch dir guard analogous to the one held by + // `write_hnsw_quantization_index_partitions`; tasks read from it, and it + // is dropped only after the drain completes. + let scratch_guard = TempStdDir::default(); + let scratch_path = scratch_guard.to_path_buf(); + + // Count of live task futures. `LiveGuard` decrements on both completion and + // cancellation, so a zero count after the drain proves every task terminated + // rather than being detached. + let live = Arc::new(AtomicUsize::new(0)); + let saw_missing_dir = Arc::new(AtomicBool::new(false)); + + struct LiveGuard(Arc); + impl Drop for LiveGuard { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::SeqCst); + } + } + + const NUM_SLOW: usize = 3; + let mut tasks: Vec>>> = Vec::with_capacity(NUM_SLOW + 1); + + // Tasks that never resolve on their own: the drain's abort is the only thing + // that can stop them, which is exactly what this test exercises. + for _ in 0..NUM_SLOW { + let live = live.clone(); + let saw_missing_dir = saw_missing_dir.clone(); + let scratch_path = scratch_path.clone(); + tasks.push(Some(tokio::spawn(async move { + live.fetch_add(1, Ordering::SeqCst); + let _guard = LiveGuard(live.clone()); + if !scratch_path.exists() { + saw_missing_dir.store(true, Ordering::SeqCst); + } + futures::future::pending::<()>().await; + Ok(0) + }))); + } + + // A task that fails; its error must be returned, not silently dropped. + // The drain aborts every handle up front, and an abort only preserves a + // task's output if the task has already finished -- an in-flight task is + // cancelled and its error lost. So wait until this task has actually + // completed before handing it to the drain; otherwise whether its error + // surfaces would depend on the scheduler and the test would be flaky. + let failing = + tokio::spawn(async move { Err(Error::io("late partition failure".to_string())) }); + while !failing.is_finished() { + tokio::task::yield_now().await; + } + tasks.push(Some(failing)); + + // Ensure all slow tasks are actually running before draining, so the + // drain has to await their cancellation rather than aborting them before + // they ever start. + while live.load(Ordering::SeqCst) < NUM_SLOW { + tokio::task::yield_now().await; + } + + let errors = drain_partition_tasks(&mut tasks).await; + + assert!(tasks.iter().all(Option::is_none), "handles left undrained"); + assert_eq!( + live.load(Ordering::SeqCst), + 0, + "a task was still running after the drain returned" + ); + assert!( + !saw_missing_dir.load(Ordering::SeqCst), + "scratch dir was removed while a task was still running" + ); + assert_eq!(errors.len(), 1, "expected exactly the one late failure"); + assert!( + errors[0].to_string().contains("late partition failure"), + "late failure was not surfaced: {}", + errors[0] + ); + + // The guard, not the drain, removes the scratch dir. + assert!(scratch_path.exists()); + drop(scratch_guard); + assert!(!scratch_path.exists()); + } + + /// `write_hnsw_quantization_index_partitions` stages each partition in a scratch + /// directory owned by a [`TempStdDir`] guard, which must remove it once the build + /// finishes so the OS temp dir does not grow without bound across legacy + /// IVF_HNSW_* builds. + /// + /// The OS temp dir is process-global, so an in-process check can't attribute a + /// leak to our own build. Instead we run the build in a child process with + /// `TMPDIR` pointed at an isolated dir we own, then assert nothing survives. + #[test] + fn test_hnsw_pq_scratch_dir_is_not_leaked() { + // Isolated temp root for the child. Owned here so it -- and anything the + // child leaks into it -- is removed when this guard drops at test end. + let isolated_root = TempStdDir::default(); + + let child_test = "index::vector::ivf::io::tests::build_legacy_hnsw_pq_in_child_process"; + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([child_test, "--exact", "--ignored", "--nocapture"]) + .env("TMPDIR", isolated_root.as_ref()) + .env("LANCE_HNSW_LEAK_TEST_ROOT", isolated_root.as_ref()) + .output() + .expect("failed to spawn child test process"); + assert!( + output.status.success(), + "child build process failed:\n--- stdout ---\n{}\n--- stderr ---\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr), + ); + + // Each build stages its partitions in one `.tmp*` dir directly under TMPDIR. + // Every guard removes its dir when it drops, so none should survive. + let leaked: Vec = std::fs::read_dir(&isolated_root) + .expect("read isolated temp root") + .flatten() + .map(|entry| entry.path()) + .filter(|path| { + path.is_dir() + && path + .file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.starts_with(".tmp")) + }) + .collect(); + + assert!( + leaked.is_empty(), + "legacy IVF_HNSW_PQ build leaked {} scratch director{} under the temp dir; \ + the TempStdDir guard should remove each one when it drops: {:?}", + leaked.len(), + if leaked.len() == 1 { "y" } else { "ies" }, + leaked, + ); + } + + /// Child half of [`test_hnsw_pq_scratch_dir_is_not_leaked`]. Ignored so it only + /// runs when the parent spawns it with `TMPDIR` and `LANCE_HNSW_LEAK_TEST_ROOT` + /// pointed at an isolated dir. Builds a few legacy IVF_HNSW_PQ indices; the + /// parent does the leak detection. + #[tokio::test] + #[ignore = "spawned as a child process by test_hnsw_pq_scratch_dir_is_not_leaked"] + async fn build_legacy_hnsw_pq_in_child_process() { + // Only do work when spawned by the parent; a bare `--ignored` run leaves + // the variable unset, so no-op rather than fail. + let Ok(root) = std::env::var("LANCE_HNSW_LEAK_TEST_ROOT") else { + return; + }; + + const DIM: usize = 32; + const ROWS: usize = 1024; + const NLIST: usize = 4; + const NUM_BUILDS: usize = 3; + + // Keep the dataset out of the temp dir's `.tmp*` namespace so the parent + // never confuses it with a leaked scratch directory. + let dataset_uri = format!("{root}/dataset"); + let values = generate_random_array(ROWS * DIM); + let fsl = Arc::new(FixedSizeListArray::try_new_from_values(values, DIM as i32).unwrap()); + let schema = Arc::new(Schema::new(vec![Field::new( + "vector", + fsl.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![fsl]).unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let ds = Dataset::write(batches, &dataset_uri, Default::default()) + .await + .unwrap(); + + for _ in 0..NUM_BUILDS { + crate::index::vector::ivf::build_ivf_hnsw_pq_index( + &ds, + "vector", + "idx", + uuid::Uuid::new_v4(), + MetricType::L2, + &IvfBuildParams::new(NLIST), + &HnswBuildParams::default(), + &PQBuildParams::new(4, 8), + ) + .await + .unwrap(); + } + } } diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 40227d2d020..1301871b560 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -3,12 +3,13 @@ //! IVF - Inverted File index. +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::marker::PhantomData; use std::{ any::Any, borrow::Cow, collections::{BinaryHeap, HashMap}, - sync::{Arc, Mutex}, + sync::{Arc, LazyLock, Mutex}, }; use crate::index::vector::{IndexFileVersion, builder::index_type_string}; @@ -121,6 +122,41 @@ pub(crate) struct IvfIndexState { pub(crate) rq_search_cache: RabitSearchCacheCell, } +/// Number of prepared partitions handed to a single `spawn_cpu` dispatch on the +/// streaming search path. +/// +/// The streaming path deliberately avoids per-partition CPU-task fan-out (a measured +/// 14-30% latency win, see #6475). Searching a batch of partitions per `spawn_cpu` +/// keeps most of that benefit — the per-dispatch overhead is paid once per +/// `STREAMING_SEARCH_BATCH_SIZE` partitions instead of once per partition — while +/// keeping the channel `recv`/`send` in async code so no CPU-pool thread ever parks on +/// a channel (which can deadlock the pool on small hosts, see #7642). `should_stop` is +/// still checked per partition, so early-stop granularity is unchanged. +/// +/// This is a tunable knob: larger batches amortize dispatch overhead further and keep +/// more work on a single CPU thread, at the cost of more prepared partitions held in +/// memory at once. The batch is an upper bound: the search loop greedily drains +/// whatever is already prepared rather than waiting for a full batch, so a slow +/// producer yields small batches (matching the old search-as-it-arrives latency) and +/// only a fast producer fills whole ones. Override with the +/// `LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE` environment variable. +pub(crate) const DEFAULT_STREAMING_SEARCH_BATCH_SIZE: usize = 16; + +pub(crate) static STREAMING_SEARCH_BATCH_SIZE: LazyLock = LazyLock::new(|| { + let batch_size = std::env::var("LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE") + .map(|value| { + value + .parse() + .expect("failed to parse LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE") + }) + .unwrap_or(DEFAULT_STREAMING_SEARCH_BATCH_SIZE); + assert!( + batch_size > 0, + "LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE must be greater than 0, got {batch_size}" + ); + batch_size +}); + struct PreparedPartitionSearch { query: Query, pre_filter: Arc, @@ -1621,8 +1657,11 @@ impl VectorIndex for IVFInd ))); } + // The prepared channel holds a full search batch so that partitions prepared + // while the previous batch is being searched are ready for the next greedy + // drain, instead of serializing producer and consumer through a single slot. let (prepared_tx, mut prepared_rx) = - mpsc::channel::>>(1); + mpsc::channel::>>(*STREAMING_SEARCH_BATCH_SIZE); let (batch_tx, batch_rx) = mpsc::channel::>(1); let prepare_index = self.clone(); @@ -1664,61 +1703,139 @@ impl VectorIndex for IVFInd let use_query_residual = self.use_query_residual; let use_residual_scratch = self.use_residual_scratch; let search_metrics = metrics.clone(); - let batch_tx_for_search = batch_tx.clone(); let search_control = control.clone(); let scratch_pool = self.scratch_pool.clone(); + // Search prepared partitions in batches. Each batch is searched in a single + // `spawn_cpu` dispatch (amortizing the per-dispatch overhead the single-worker + // design in #6475 avoided), but the channel `recv`/`send` stay in async code so + // no CPU-pool thread ever parks on a channel — parking one can deadlock the pool + // on small hosts (#7642). `should_stop` is checked per partition, so early-stop + // granularity is unchanged. + // + // Batches are formed greedily: wait for one prepared partition, then drain + // whatever else is already prepared, up to the batch size. Waiting for a full + // batch instead would delay the first search (and the early-stop feedback it + // produces) behind up to a whole batch of prepare I/O, which is significant + // when prepare parallelism is low. tokio::spawn(async move { - let search_result = spawn_cpu(move || -> DataFusionResult<()> { - scratch_pool.with_scratch(|scratch| { - while let Some(prepared) = prepared_rx.blocking_recv() { - let prepared = match prepared { - Ok(prepared) => prepared, - Err(err) => { - let _ = batch_tx_for_search - .blocking_send(Err(DataFusionError::from(err))); - return Ok(()); - } - }; + loop { + // Stop pulling as soon as the search is done — or the receiver of our + // results is gone — so the producer stops preparing partitions we + // would never search. + if search_control + .as_ref() + .is_some_and(|control| control.should_stop()) + || batch_tx.is_closed() + { + return; + } - if search_control - .as_ref() - .is_some_and(|control| control.should_stop()) - { - return Ok(()); + let mut prepared_batch = Vec::with_capacity(*STREAMING_SEARCH_BATCH_SIZE); + let mut prepare_error = None; + let mut producer_done = false; + match prepared_rx.recv().await { + Some(Ok(prepared)) => prepared_batch.push(prepared), + Some(Err(err)) => prepare_error = Some(DataFusionError::from(err)), + None => producer_done = true, + } + while prepare_error.is_none() + && !producer_done + && prepared_batch.len() < *STREAMING_SEARCH_BATCH_SIZE + { + match prepared_rx.try_recv() { + Ok(Ok(prepared)) => prepared_batch.push(prepared), + Ok(Err(err)) => { + prepare_error = Some(DataFusionError::from(err)); } - - let batch = { - Self::run_prepared_partition_search( - use_query_residual, - use_residual_scratch, - prepared, - search_metrics.as_ref(), - scratch, - ) + // Nothing else is prepared yet; search what we have rather + // than waiting for more. + Err(mpsc::error::TryRecvError::Empty) => break, + Err(mpsc::error::TryRecvError::Disconnected) => { + producer_done = true; } - .map_err(DataFusionError::from); - match batch { - Ok(batch) => { - if let Some(control) = search_control.as_ref() { - control.record_batch(&batch); + } + } + + if !prepared_batch.is_empty() { + let scratch_pool = scratch_pool.clone(); + let search_metrics = search_metrics.clone(); + let search_control = search_control.clone(); + // `is_closed` is synchronously callable, so a sender clone lets the + // CPU loop notice a dropped receiver between partitions instead of + // searching out the whole batch for a cancelled query. (A `select!` + // on `closed()` would not help here: `spawn_cpu` closures are not + // cancellable, so abandoning the await leaves the work running.) + let cancel_probe = batch_tx.clone(); + let search_output = spawn_cpu(move || { + let mut outputs: Vec> = + Vec::with_capacity(prepared_batch.len()); + // `stopped` means the whole search should end (an error, an + // early-stop signal, or cancellation), not just this batch. + let mut stopped = false; + scratch_pool.with_scratch(|scratch| { + for prepared in prepared_batch { + if search_control + .as_ref() + .is_some_and(|control| control.should_stop()) + || cancel_probe.is_closed() + { + stopped = true; + break; } - if batch_tx_for_search.blocking_send(Ok(batch)).is_err() { - return Ok(()); + match Self::run_prepared_partition_search( + use_query_residual, + use_residual_scratch, + prepared, + search_metrics.as_ref(), + scratch, + ) + .map_err(DataFusionError::from) + { + Ok(batch) => { + if let Some(control) = search_control.as_ref() { + control.record_batch(&batch); + } + outputs.push(Ok(batch)); + } + Err(err) => { + outputs.push(Err(err)); + stopped = true; + break; + } } } - Err(err) => { - let _ = batch_tx_for_search.blocking_send(Err(err)); - return Ok(()); - } + }); + Ok::<_, DataFusionError>((outputs, stopped)) + }) + .await; + + let (outputs, stopped) = match search_output { + Ok(output) => output, + // Defensive: the closure always returns Ok (search errors are + // captured per partition in `outputs`), so this arm should be + // unreachable. Forward and stop rather than drop silently. + Err(err) => { + let _ = batch_tx.send(Err(err)).await; + return; + } + }; + for output in outputs { + if batch_tx.send(output).await.is_err() { + return; } } - Ok(()) - }) - }) - .await; + if stopped { + return; + } + } - if let Err(err) = search_result { - let _ = batch_tx.send(Err(err)).await; + if let Some(err) = prepare_error { + let _ = batch_tx.send(Err(err)).await; + return; + } + if producer_done { + return; + } } }); @@ -1790,7 +1907,7 @@ impl VectorIndex for IVFInd todo!("this method is for only IVF_HNSW_* index"); } - async fn remap(&mut self, _mapping: &HashMap>) -> Result<()> { + async fn remap(&mut self, _mapping: &RowAddrRemap) -> Result<()> { Err(Error::index( "Remapping IVF in this way not supported".to_string(), )) @@ -4399,17 +4516,30 @@ mod tests { test_remap(params, nlist, recall_requirement).await; } - // RQ doesn't perform well for random data - // need to verify recall with real-world dataset (e.g. sift1m) + #[tokio::test] + async fn test_build_ivf_sq_dot_with_negative_values() { + let nlist = 4; + let ivf_params = IvfBuildParams::new(nlist); + let sq_params = SQBuildParams::default(); + let params = + VectorIndexParams::with_ivf_sq_params(DistanceType::Dot, ivf_params, sq_params); + + test_index_impl::(params, nlist, 0.75, -1.0..1.0, None).await; + } + + // These queries probe every partition, so recall here measures RaBitQ quantization + // error alone. At 1 bit per dimension it averages ~0.67 on this uniformly random, + // L2-normalized data, and each build draws a fresh random rotation, so no bar worth + // asserting sits clear of the spread. 5 bits lifts recall to ~0.97; its `ex_bits = 4` + // also covers a FastScan ex-code kernel that the multi-bit test below never reaches. #[rstest] - #[case(1, DistanceType::L2, 0.5)] - #[case(1, DistanceType::Cosine, 0.5)] - #[case(1, DistanceType::Dot, 0.5)] - #[case(4, DistanceType::L2, 0.5)] - #[case(4, DistanceType::Cosine, 0.5)] - #[case(4, DistanceType::Dot, 0.5)] + #[case(1, DistanceType::L2, 0.9)] + #[case(1, DistanceType::Cosine, 0.9)] + #[case(1, DistanceType::Dot, 0.9)] + #[case(4, DistanceType::L2, 0.9)] + #[case(4, DistanceType::Cosine, 0.9)] + #[case(4, DistanceType::Dot, 0.9)] #[tokio::test] - // #[ignore = "Temporarily skipping flaky 4-bit IVF_RQ tests"] async fn test_build_ivf_rq( #[case] nlist: usize, #[case] distance_type: DistanceType, @@ -4418,7 +4548,7 @@ mod tests { ) { let _ = env_logger::try_init(); let ivf_params = IvfBuildParams::new(nlist); - let rq_params = RQBuildParams::with_rotation_type(1, rotation_type); + let rq_params = RQBuildParams::with_rotation_type(5, rotation_type); let params = VectorIndexParams::with_ivf_rq_params(distance_type, ivf_params, rq_params); test_index(params.clone(), nlist, recall_requirement, None).await; if distance_type == DistanceType::Cosine { @@ -4571,6 +4701,22 @@ mod tests { test_remap(params, nlist, recall_requirement).await; } + #[tokio::test] + async fn test_create_ivf_hnsw_sq_dot_with_negative_values() { + let nlist = 4; + let ivf_params = IvfBuildParams::new(nlist); + let sq_params = SQBuildParams::default(); + let hnsw_params = HnswBuildParams::default(); + let params = VectorIndexParams::with_ivf_hnsw_sq_params( + DistanceType::Dot, + ivf_params, + hnsw_params, + sq_params, + ); + + test_index_impl::(params, nlist, 0.75, -1.0..1.0, None).await; + } + #[rstest] #[case(4, DistanceType::L2, 0.9)] #[case(4, DistanceType::Cosine, 0.9)] @@ -4689,10 +4835,7 @@ mod tests { .as_primitive::() .values() .to_vec(); - let results = dists - .into_iter() - .zip(row_ids.clone().into_iter()) - .collect::>(); + let results = dists.into_iter().zip(row_ids.clone()).collect::>(); let row_ids = row_ids.into_iter().collect::>(); let gt = multivec_ground_truth(&vectors, &query, k, params.metric_type); @@ -5085,10 +5228,7 @@ mod tests { .as_primitive::() .values() .to_vec(); - let results = dists - .into_iter() - .zip(row_ids.into_iter()) - .collect::>(); + let results = dists.into_iter().zip(row_ids).collect::>(); let row_ids = results.iter().map(|(_, id)| *id).collect::>(); assert!(row_ids.len() == k); diff --git a/rust/lance/src/index/vector/pq.rs b/rust/lance/src/index/vector/pq.rs index 6e335cddc80..217667afbe9 100644 --- a/rust/lance/src/index/vector/pq.rs +++ b/rust/lance/src/index/vector/pq.rs @@ -1,8 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_core::utils::row_addr_remap::RowAddrRemap; +use std::any::Any; use std::sync::Arc; -use std::{any::Any, collections::HashMap}; use arrow::compute::concat; use arrow_array::types::UInt64Type; @@ -438,14 +439,14 @@ impl VectorIndex for PQIndex { todo!("this method is for only IVF_HNSW_* index"); } - async fn remap(&mut self, mapping: &HashMap>) -> Result<()> { + async fn remap(&mut self, mapping: &RowAddrRemap) -> Result<()> { let num_vectors = self.row_ids.as_ref().unwrap().len(); let row_ids = self.row_ids.as_ref().unwrap().values().iter(); let transposed_codes = self.code.as_ref().unwrap(); let remapped = row_ids .enumerate() .filter_map(|(vec_idx, old_row_id)| { - let new_row_id = mapping.get(old_row_id).cloned(); + let new_row_id = mapping.get(*old_row_id); // If the row id is not in the mapping then this row is not remapped and we keep as is let new_row_id = new_row_id.unwrap_or(Some(*old_row_id)); new_row_id.map(|new_row_id| { @@ -644,6 +645,7 @@ pub(crate) fn build_pq_storage( mod tests { use super::*; + use std::collections::HashMap; use std::{ops::Range, sync::Mutex}; use arrow::datatypes::Float32Type; diff --git a/rust/lance/src/index/vector/utils.rs b/rust/lance/src/index/vector/utils.rs index 3046d0f3a83..79c8ccb6295 100644 --- a/rust/lance/src/index/vector/utils.rs +++ b/rust/lance/src/index/vector/utils.rs @@ -1,9 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::collections::HashSet; use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use arrow::array::ArrayData; use arrow::datatypes::DataType; @@ -18,6 +18,7 @@ use log::{info, warn}; use rand::rngs::SmallRng; use rand::seq::{IteratorRandom, SliceRandom}; use rand::{Rng, SeedableRng}; +use roaring::RoaringTreemap; use tokio::sync::Mutex; use crate::dataset::{Dataset, ProjectionRequest, TakeBuilder, row_offsets_to_row_addresses}; @@ -486,18 +487,39 @@ async fn sample_training_data( ); return vector_column_to_fsl(&batch, column); } + // Rows the consumer still needs. The fragment producer sizes each + // prefetch round to this outstanding demand, keeping reads bounded by + // the requested sample size. + let still_needed = Arc::new(AtomicUsize::new(sample_size_hint)); let scan = sample_training_data_scan_from_fragments( dataset, column, - sample_size_hint, num_rows, fragment_ids, + still_needed.clone(), )?; return match vector_field.data_type() { DataType::FixedSizeList(_, _) => { - sample_nullable_fsl(column, sample_size_hint, byte_width, vector_field, scan).await + sample_nullable_fsl( + column, + sample_size_hint, + byte_width, + vector_field, + scan, + Some(still_needed), + ) + .await + } + _ => { + sample_nullable_fallback( + column, + sample_size_hint, + is_nullable, + scan, + Some(still_needed), + ) + .await } - _ => sample_nullable_fallback(column, sample_size_hint, is_nullable, scan).await, }; } @@ -516,12 +538,20 @@ async fn sample_training_data( DataType::FixedSizeList(_, _) => { let scan = sample_training_data_scan(dataset, column, sample_size_hint, num_rows, byte_width)?; - sample_nullable_fsl(column, sample_size_hint, byte_width, vector_field, scan).await + sample_nullable_fsl( + column, + sample_size_hint, + byte_width, + vector_field, + scan, + None, + ) + .await } _ => { let scan = sample_training_data_scan(dataset, column, sample_size_hint, num_rows, byte_width)?; - sample_nullable_fallback(column, sample_size_hint, is_nullable, scan).await + sample_nullable_fallback(column, sample_size_hint, is_nullable, scan, None).await } } } @@ -550,12 +580,19 @@ fn sample_training_data_scan( /// sampling must first map random offsets within the selected fragments to row /// addresses and then `take` those rows. Both nullable FSL and multivector /// paths reuse this stream to avoid duplicating fragment sampling logic. +/// +/// Each round is sized to `still_needed` (the consumer's outstanding demand), +/// so a low-null column reads at most the requested sample. Visited offsets +/// are tracked in a [`RoaringTreemap`] because sparse or all-null columns +/// force the stream to visit most selected rows before it can terminate, and +/// a compressed bitmap keeps that persistent state near `num_rows / 8` bytes +/// even when fully populated. fn sample_training_data_scan_from_fragments( dataset: &Dataset, column: &str, - sample_size_hint: usize, num_rows: usize, fragment_ids: &[u32], + still_needed: Arc, ) -> Result> + Send>>> { if fragment_ids.is_empty() { return Err(Error::invalid_input( @@ -589,19 +626,36 @@ fn sample_training_data_scan_from_fragments( dataset, projection, selected_fragments, - HashSet::::with_capacity(sample_size_hint.min(num_rows)), + RoaringTreemap::new(), SmallRng::from_os_rng(), + still_needed, ), - move |(dataset, projection, selected_fragments, mut seen_offsets, mut rng)| async move { - if seen_offsets.len() >= num_rows { + move |( + dataset, + projection, + selected_fragments, + mut seen_offsets, + mut rng, + still_needed, + )| async move { + if seen_offsets.len() as usize >= num_rows { + return Ok(None); + } + let still = still_needed.load(Ordering::Relaxed); + if still == 0 { return Ok(None); } - let remaining = num_rows.saturating_sub(seen_offsets.len()); - let target = sample_size_hint.saturating_mul(2).min(remaining); + let remaining = num_rows.saturating_sub(seen_offsets.len() as usize); + // Sizing the round to the outstanding demand keeps a low-null + // column's reads bounded by the requested sample, matching the + // non-nullable path. + let target = still.min(remaining); let mut sampled_offsets = if remaining <= target.saturating_mul(4) { + // Few offsets remain unseen, so shuffling the unseen set is + // cheaper than repeatedly rejecting already-sampled offsets. let mut unseen_indices = (0..num_rows as u64) - .filter(|index| !seen_offsets.contains(index)) + .filter(|index| !seen_offsets.contains(*index)) .collect::>(); unseen_indices.shuffle(&mut rng); unseen_indices.truncate(target); @@ -635,7 +689,14 @@ fn sample_training_data_scan_from_fragments( .await?; Ok(Some(( batch, - (dataset, projection, selected_fragments, seen_offsets, rng), + ( + dataset, + projection, + selected_fragments, + seen_offsets, + rng, + still_needed, + ), ))) }, ); @@ -721,6 +782,7 @@ async fn sample_nullable_fsl( byte_width: usize, vector_field: &lance_core::datatypes::Field, mut scan: S, + still_needed: Option>, ) -> Result where S: Stream> + Unpin, @@ -731,6 +793,12 @@ where let mut rows_scanned: usize = 0; while num_non_null < sample_size_hint { + let remaining_rows = sample_size_hint - num_non_null; + // A fragment-limited producer sizes its next prefetch round to this + // outstanding demand. + if let Some(still_needed) = &still_needed { + still_needed.store(remaining_rows, Ordering::Relaxed); + } let Some(batch) = scan.next().await else { break; }; @@ -750,7 +818,16 @@ where continue; } let previous_num_non_null = num_non_null; - accumulate_fsl_values(&mut values_buf, &mut num_non_null, &array, byte_width, true)?; + // `remaining_rows` keeps `values_buf` within its pre-allocated + // `sample_size_hint * byte_width` capacity. + accumulate_fsl_values( + &mut values_buf, + &mut num_non_null, + &array, + byte_width, + true, + remaining_rows, + )?; info!( "Sample training data: batch {} read {} rows, accepted {} rows ({} scanned, {}/{} sampled after null filtering)", batch_count, @@ -762,6 +839,11 @@ where ); } + // Zero the demand so any further poll of the producer terminates instead + // of reading another round. + if let Some(still_needed) = &still_needed { + still_needed.store(0, Ordering::Relaxed); + } let num_rows_out = num_non_null.min(sample_size_hint); values_buf.truncate(num_rows_out * byte_width); @@ -797,7 +879,14 @@ async fn sample_fsl_uniform( for (chunk_idx, chunk) in indices.chunks(TAKE_CHUNK_SIZE).enumerate() { let batch = dataset.take(chunk, projection.clone()).await?; let array = get_column_from_batch(&batch, column)?; - accumulate_fsl_values(&mut values_buf, &mut total_rows, &array, byte_width, false)?; + accumulate_fsl_values( + &mut values_buf, + &mut total_rows, + &array, + byte_width, + false, + usize::MAX, + )?; info!( "Sample training data: batch {}/{} read {} rows ({}/{} sampled by uniform random sampling)", chunk_idx + 1, @@ -821,13 +910,20 @@ async fn sample_fsl_uniform( /// When `filter_nulls` is false and there are no nulls, copies raw bytes /// directly from the FSL values buffer (accounting for child array offset). /// When `filter_nulls` is true, uses Arrow's `filter` kernel to remove nulls. +/// At most `max_rows` rows are appended so callers can stop copying once their +/// sample is full; otherwise one oversized source batch can grow `values_buf` +/// far beyond the intended cap. fn accumulate_fsl_values( values_buf: &mut MutableBuffer, num_rows: &mut usize, array: &ArrayRef, byte_width: usize, filter_nulls: bool, + max_rows: usize, ) -> Result<()> { + if max_rows == 0 { + return Ok(()); + } let needs_filter = filter_nulls && array.null_count() > 0; if needs_filter { @@ -835,21 +931,29 @@ fn accumulate_fsl_values( let mask = arrow_array::BooleanArray::from(nulls.inner().clone()); let filtered = arrow::compute::filter(array, &mask)?; let fsl = filtered.as_fixed_size_list(); + let take = fsl.len().min(max_rows); + if take == 0 { + return Ok(()); + } let values_data = fsl.values().to_data(); - let value_bytes = &values_data.buffers()[0].as_slice()[..fsl.len() * byte_width]; + let value_bytes = &values_data.buffers()[0].as_slice()[..take * byte_width]; values_buf.extend_from_slice(value_bytes); - *num_rows += fsl.len(); + *num_rows += take; } else { // No nulls: copy raw bytes directly, accounting for child array offset. let fsl = array.as_fixed_size_list(); + let take = fsl.len().min(max_rows); + if take == 0 { + return Ok(()); + } let values = fsl.values(); let values_data = values.to_data(); let elem_size = byte_width / fsl.value_length() as usize; let offset_bytes = values_data.offset() * elem_size; - let total_bytes = fsl.len() * byte_width; + let total_bytes = take * byte_width; let buf = &values_data.buffers()[0].as_slice()[offset_bytes..offset_bytes + total_bytes]; values_buf.extend_from_slice(buf); - *num_rows += fsl.len(); + *num_rows += take; } Ok(()) } @@ -862,6 +966,7 @@ async fn sample_nullable_fallback( sample_size_hint: usize, is_nullable: bool, mut scan: S, + still_needed: Option>, ) -> Result where S: Stream> + Unpin, @@ -873,6 +978,12 @@ where let mut rows_scanned: usize = 0; while num_non_null < sample_size_hint { + let remaining_rows = sample_size_hint - num_non_null; + // A fragment-limited producer sizes its next prefetch round to this + // outstanding demand. + if let Some(still_needed) = &still_needed { + still_needed.store(remaining_rows, Ordering::Relaxed); + } let Some(batch) = scan.next().await else { break; }; @@ -898,7 +1009,14 @@ where } else { batch }; - let accepted_rows = batch.num_rows(); + // Slicing to the outstanding demand keeps the retained batches, and + // the post-loop `concat_batches`, bounded by the sample size. + let accepted_rows = batch.num_rows().min(remaining_rows); + let batch = if accepted_rows < batch.num_rows() { + batch.slice(0, accepted_rows) + } else { + batch + }; num_non_null += accepted_rows; info!( "Sample training data (fallback): batch {} read {} rows, accepted {} rows ({} scanned, {}/{} sampled)", @@ -912,6 +1030,12 @@ where filtered.push(batch); } + // Zero the demand so any further poll of the producer terminates instead + // of reading another round. + if let Some(still_needed) = &still_needed { + still_needed.store(0, Ordering::Relaxed); + } + let Some(schema) = schema else { return Err(Error::index("No non-null training data found".to_string())); }; @@ -1040,6 +1164,7 @@ mod tests { use crate::dataset::InsertBuilder; use arrow_array::{ArrayRef, Float32Array, types::Float32Type}; + use arrow_buffer::{BooleanBufferBuilder, NullBuffer}; use arrow_schema::{DataType, Field}; use lance_arrow::FixedSizeListArrayExt; use lance_datagen::{ArrayGeneratorExt, Dimension, RowCount, array, gen_batch}; @@ -1201,7 +1326,15 @@ mod tests { let mut buf = MutableBuffer::new(0); let mut num_rows = 0usize; let sliced_ref: ArrayRef = Arc::new(sliced); - accumulate_fsl_values(&mut buf, &mut num_rows, &sliced_ref, byte_width, false).unwrap(); + accumulate_fsl_values( + &mut buf, + &mut num_rows, + &sliced_ref, + byte_width, + false, + usize::MAX, + ) + .unwrap(); assert_eq!(num_rows, 4); let result: &[f32] = @@ -1327,4 +1460,140 @@ mod tests { let result = count_rows(&dataset, Some(&[ids[2], ids[0]])).await.unwrap(); assert_eq!(result, 250); } + + /// Nullable FSL with fragment-limited sampling must fill the requested sample + /// size when enough non-null rows exist, and terminate cleanly when all + /// selected rows are null. + #[tokio::test] + async fn test_maybe_sample_training_data_fsl_nullable_fragment_limited() { + let nrows: usize = 2000; + let dims: u32 = 8; + let sample_size: usize = 500; + + for (case, null_probability, expected_len) in + [("partial_nulls", 0.5, sample_size), ("all_nulls", 1.0, 0)] + { + let col_gen = array::rand_vec::(Dimension::from(dims)) + .with_random_nulls(null_probability); + let data = gen_batch() + .col("vec", col_gen) + .into_batch_rows(RowCount::from(nrows as u64)) + .unwrap(); + + let dataset = InsertBuilder::new("memory://") + .execute(vec![data]) + .await + .unwrap(); + + let fragment_ids: Vec = dataset + .get_fragments() + .iter() + .map(|f| f.id() as u32) + .collect(); + + let training_data = + maybe_sample_training_data(&dataset, "vec", sample_size, Some(&fragment_ids)) + .await + .unwrap(); + + assert_eq!(training_data.len(), expected_len, "{case}"); + assert_eq!(training_data.null_count(), 0, "{case}"); + assert_eq!(training_data.value_length(), dims as i32, "{case}"); + } + } + + /// Scan-side regression: each fragment-limited producer round must read at + /// most the consumer's outstanding demand. Driving the producer directly + /// with a fixed `still_needed` and inspecting the raw batch size catches + /// over-reads that a post-truncation output-length check cannot. + #[tokio::test] + async fn test_sample_fragment_scan_round_caps_at_still_needed() { + let nrows: usize = 4000; + let dims: u32 = 8; + let still: usize = 500; + + let col_gen = array::rand_vec::(Dimension::from(dims)).with_random_nulls(0.5); + let data = gen_batch() + .col("vec", col_gen) + .into_batch_rows(RowCount::from(nrows as u64)) + .unwrap(); + + let dataset = InsertBuilder::new("memory://fsl_scan_round_cap_test") + .execute(vec![data]) + .await + .unwrap(); + + let fragment_ids: Vec = dataset + .get_fragments() + .iter() + .map(|f| f.id() as u32) + .collect(); + let num_rows = count_rows(&dataset, Some(&fragment_ids)).await.unwrap(); + + // `still_needed` is left large enough that `num_rows` never bounds the + // round, so the batch size reflects the demand cap and nothing else. + let still_needed = Arc::new(AtomicUsize::new(still)); + let mut scan = sample_training_data_scan_from_fragments( + &dataset, + "vec", + num_rows, + &fragment_ids, + still_needed.clone(), + ) + .unwrap(); + + let batch = scan.next().await.unwrap().unwrap(); + assert!( + batch.num_rows() <= still, + "producer round read {} rows but only {} were outstanding", + batch.num_rows(), + still + ); + } + + #[test] + fn test_accumulate_fsl_values_respects_max_rows() { + let dim: usize = 4; + let total_rows: usize = 100; + let max_rows: usize = 16; + let byte_width = dim * std::mem::size_of::(); + + let values: Vec = (0..total_rows * dim).map(|i| i as f32).collect(); + let fsl = FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim as i32) + .unwrap(); + let arr: ArrayRef = Arc::new(fsl); + + let mut buf = MutableBuffer::new(0); + let mut num_rows = 0usize; + accumulate_fsl_values(&mut buf, &mut num_rows, &arr, byte_width, true, max_rows).unwrap(); + + assert_eq!(num_rows, max_rows); + assert_eq!(buf.len(), max_rows * byte_width); + + let values: Vec = (0..total_rows * dim).map(|i| i as f32).collect(); + let item_field = Arc::new(Field::new("item", DataType::Float32, true)); + + // Every other row is null, leaving 50 non-null rows. + let mut nulls_builder = BooleanBufferBuilder::new(total_rows); + for i in 0..total_rows { + nulls_builder.append(i % 2 == 0); + } + let nulls = NullBuffer::new(nulls_builder.finish()); + + let fsl = FixedSizeListArray::try_new( + item_field, + dim as i32, + Arc::new(Float32Array::from(values)), + Some(nulls), + ) + .unwrap(); + let arr: ArrayRef = Arc::new(fsl); + + let mut buf = MutableBuffer::new(0); + let mut num_rows = 0usize; + accumulate_fsl_values(&mut buf, &mut num_rows, &arr, byte_width, true, max_rows).unwrap(); + + assert_eq!(num_rows, max_rows); + assert_eq!(buf.len(), max_rows * byte_width); + } } diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index ce0d29d550b..35cd275b8ee 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -489,13 +489,12 @@ fn fix_schema(manifest: &mut Manifest) -> Result<()> { } // Now, we need to remap the field ids to be unique. - let mut field_id_seed = manifest.max_field_id() + 1; let mut old_field_id_mapping: HashMap = HashMap::new(); let mut fields_with_duplicate_ids = fields_with_duplicate_ids.into_iter().collect::>(); fields_with_duplicate_ids.sort_unstable(); - for field_id in fields_with_duplicate_ids { + for (field_id_seed, field_id) in (manifest.max_field_id() + 1..).zip(fields_with_duplicate_ids) + { old_field_id_mapping.insert(field_id, field_id_seed); - field_id_seed += 1; } let mut fragments = manifest.fragments.as_ref().clone(); @@ -592,17 +591,21 @@ pub(crate) async fn migrate_fragments( let mut data_files = fragment.files.clone(); - // For each of the data files in the fragment, we need to get the file size - let object_store = dataset.object_store.as_ref(); + // For each of the data files in the fragment, we need to get the file size. + // Resolve each file against its own storage base: multi-base datasets + // keep data files outside the dataset root (DataFile.base_id). let get_sizes = data_files .iter() .map(|file| { if let Some(size) = file.file_size_bytes.get() { Either::Left(futures::future::ready(Ok(size))) } else { - Either::Right(async { + let dataset = dataset.clone(); + Either::Right(async move { + let object_store = dataset.object_store_for_data_file(file).await?; + let data_dir = dataset.data_file_dir_for_base(file.base_id)?; object_store - .size(&dataset.base.clone().join("data").join(file.path.clone())) + .size(&data_dir.join(file.path.clone())) .map_ok(|size| { NonZero::new(size).ok_or_else(|| { Error::internal(format!("File {} has size 0", file.path)) @@ -1687,6 +1690,7 @@ mod tests { DataFile::new_legacy_from_fields("path1", vec![0, 1, 2], None), DataFile::new_legacy_from_fields("unused", vec![9], None), ], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1699,6 +1703,7 @@ mod tests { DataFile::new_legacy_from_fields("path2", vec![0, 1, 2], None), DataFile::new_legacy_from_fields("path3", vec![2], None), ], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1736,6 +1741,7 @@ mod tests { vec![0, 1, 10], None, )], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1748,6 +1754,7 @@ mod tests { DataFile::new_legacy_from_fields("path2", vec![0, 1, 2], None), DataFile::new_legacy_from_fields("path3", vec![10], None), ], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1838,6 +1845,7 @@ mod tests { let fragment = Fragment { id: 0, files: vec![data_file], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(100), diff --git a/rust/lance/src/io/commit/conflict_resolver.rs b/rust/lance/src/io/commit/conflict_resolver.rs index d95821dd130..7be9c37230e 100644 --- a/rust/lance/src/io/commit/conflict_resolver.rs +++ b/rust/lance/src/io/commit/conflict_resolver.rs @@ -7,7 +7,7 @@ use crate::index::mem_wal::{load_mem_wal_index_details, new_mem_wal_index_meta}; use crate::io::deletion::read_dataset_deletion_file; use crate::{ Dataset, - dataset::transaction::{Operation, Transaction, UpdateMode}, + dataset::transaction::{DataOverlayGroup, Operation, Transaction, UpdateMode}, }; use futures::{StreamExt, TryStreamExt}; use lance_core::{Error, Result, utils::deletion::DeletionVector}; @@ -15,7 +15,9 @@ use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; use lance_index::mem_wal::{MEM_WAL_INDEX_NAME, MergedGeneration}; use lance_select::{RowAddrTreeMap, RowSetOps}; use lance_table::format::IndexMetadata; +use lance_table::format::overlay::OverlayCoverage; use lance_table::{format::Fragment, io::deletion::write_deletion_file}; +use roaring::RoaringBitmap; use std::{ borrow::Cow, collections::{HashMap, HashSet}, @@ -137,6 +139,21 @@ impl<'a> TransactionRebase<'a> { conflicting_mem_wal_merged_gens: Vec::new(), }) } + Operation::DataOverlay { groups } => { + let modified_fragment_ids = + groups.iter().map(|g| g.fragment_id).collect::>(); + let initial_fragments = + initial_fragments_for_rebase(dataset, &transaction, &modified_fragment_ids) + .await; + Ok(Self { + transaction, + affected_rows, + initial_fragments, + modified_fragment_ids, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_merged_gens: Vec::new(), + }) + } Operation::Merge { fragments, .. } => { let modified_fragment_ids = fragments.iter().map(|f| f.id).collect::>(); let initial_fragments = @@ -219,6 +236,9 @@ impl<'a> TransactionRebase<'a> { Operation::DataReplacement { .. } => { self.check_data_replacement_txn(other_transaction, other_version) } + Operation::DataOverlay { .. } => { + self.check_data_overlay_txn(other_transaction, other_version) + } Operation::Merge { .. } => self.check_merge_txn(other_transaction, other_version), Operation::Restore { .. } => self.check_restore_txn(other_transaction, other_version), Operation::ReserveFragments { .. } => { @@ -251,6 +271,10 @@ impl<'a> TransactionRebase<'a> { | Operation::Project { .. } | Operation::Append { .. } | Operation::UpdateConfig { .. } + // A concurrent overlay is inert against the rows we delete + // (deletions take precedence over overlays) and otherwise + // preserves physical offsets, so it never conflicts. + | Operation::DataOverlay { .. } | Operation::UpdateBases { .. } => Ok(()), Operation::Rewrite { groups, .. } => { if groups @@ -346,6 +370,8 @@ impl<'a> TransactionRebase<'a> { if let Operation::Update { inserted_rows_filter: self_inserted_rows_filter, merged_generations: self_merged_generations, + new_fragments: self_new_fragments, + update_mode: self_update_mode, .. } = &self.transaction.operation { @@ -399,6 +425,49 @@ impl<'a> TransactionRebase<'a> { | Operation::Clone { .. } | Operation::UpdateConfig { .. } | Operation::UpdateBases { .. } => Ok(()), + Operation::DataOverlay { groups } => { + // Our update recomputed rows from the pre-overlay base, so if + // it commits over an overlay it would silently undo the + // overlay's values for any cell it recomputed. A row-moving + // update (RewriteRows) relocates the rows it touches out to + // new fragments; only the rows it actually moved lose their + // overlay, so we conflict only when the moved rows intersect + // the overlay's coverage. An in-place column rewrite + // (RewriteColumns) preserves offsets and just tombstones the + // overlaid fields at build time, so it never conflicts. + let moves_rows = !self_new_fragments.is_empty() + && matches!(self_update_mode, Some(UpdateMode::RewriteRows) | None); + if !moves_rows { + return Ok(()); + } + // `affected_rows` holds the physical offsets (per fragment) + // this update moved. The overlay's coverage is in the same + // physical-offset space, so we can intersect the two in + // memory. Without affected rows we cannot be precise, so we + // fall back to a fragment-granular conflict. + for group in groups { + if !self.modified_fragment_ids.contains(&group.fragment_id) { + continue; + } + let Some(affected_rows) = self.affected_rows else { + return Err( + self.retryable_conflict_err(other_transaction, other_version) + ); + }; + let Some(moved) = + affected_rows.get_fragment_bitmap(group.fragment_id as u32) + else { + continue; + }; + let coverage = overlay_group_coverage(group); + if !(moved & &coverage).is_empty() { + return Err( + self.retryable_conflict_err(other_transaction, other_version) + ); + } + } + Ok(()) + } Operation::Append { .. } => { // If current transaction has primary key conflict detection, // we can't safely commit against an Append because we don't @@ -514,6 +583,10 @@ impl<'a> TransactionRebase<'a> { match &other_transaction.operation { Operation::Append { .. } | Operation::Clone { .. } + // An overlay committed after this index's version is newer than + // the index; the query path excludes its covered cells via the + // version gate, so the build does not conflict. + | Operation::DataOverlay { .. } | Operation::UpdateBases { .. } => Ok(()), Operation::CreateIndex { new_indices: created_indices, @@ -695,6 +768,20 @@ impl<'a> TransactionRebase<'a> { Ok(()) } } + Operation::DataOverlay { groups } => { + // Rewriting a fragment changes its physical row addresses, so + // an overlay addressed by physical offset on that fragment is + // invalidated and must be re-applied against the new base. + if groups + .iter() + .map(|g| g.fragment_id) + .any(|id| self.modified_fragment_ids.contains(&id)) + { + Err(self.retryable_conflict_err(other_transaction, other_version)) + } else { + Ok(()) + } + } Operation::Rewrite { groups, frag_reuse_index: committed_fri, @@ -874,6 +961,7 @@ impl<'a> TransactionRebase<'a> { | Operation::CreateIndex { .. } | Operation::Rewrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::ReserveFragments { .. } @@ -907,7 +995,8 @@ impl<'a> TransactionRebase<'a> { | Operation::Merge { .. } | Operation::UpdateConfig { .. } | Operation::Clone { .. } - | Operation::DataReplacement { .. } => Ok(()), + | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } => Ok(()), } } @@ -923,6 +1012,9 @@ impl<'a> TransactionRebase<'a> { | Operation::UpdateConfig { .. } | Operation::ReserveFragments { .. } | Operation::Project { .. } + // Both a column replacement and an overlay preserve physical row + // addresses; the overlay is newer and wins its covered cells. + | Operation::DataOverlay { .. } | Operation::UpdateBases { .. } => Ok(()), Operation::Merge { .. } => { // Merge rewrites the whole fragment list; always conflict @@ -1055,6 +1147,119 @@ impl<'a> TransactionRebase<'a> { } } + /// Conflict checks for our DataOverlay transaction against a concurrent one. + /// + /// Overlays are intentionally permissive (see the Data Overlay Files spec): + /// they stack with other overlays and tolerate appends, index builds, data + /// replacement, deletes, and in-place column rewrites (Update with + /// `RewriteColumns`), because overlay coverage is addressed by physical offset + /// and the version gate keeps indexes correct. A concurrent operation + /// conflicts when it takes precedence over the overlay for cells the overlay + /// covers, dropping the overlay's values: retryably when it rewrites the + /// physical layout of one of our fragments (Rewrite, Merge) or re-creates the + /// covered rows from the pre-overlay base (a row-moving Update — checked + /// row-by-row in `finish_data_overlay`), or removes an overlaid fragment + /// outright (a Delete / Update that drops the fragment); and incompatibly for + /// whole-dataset replacements (Overwrite / Restore) and MemWAL state updates + /// (UpdateMemWalState), which do not rebase against data operations. + fn check_data_overlay_txn( + &mut self, + other_transaction: &Transaction, + other_version: u64, + ) -> Result<()> { + match &other_transaction.operation { + Operation::Append { .. } + | Operation::CreateIndex { .. } + | Operation::ReserveFragments { .. } + | Operation::Project { .. } + | Operation::UpdateConfig { .. } + | Operation::UpdateBases { .. } + | Operation::Clone { .. } + | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } => Ok(()), + // A concurrent Delete only tombstones rows via a deletion vector, + // which preserves physical offsets; the overlay value for a deleted + // offset is simply inert. Conflict only if the whole overlaid + // fragment was removed, orphaning the overlay. + Operation::Delete { + deleted_fragment_ids, + .. + } => { + if deleted_fragment_ids + .iter() + .any(|id| self.modified_fragment_ids.contains(id)) + { + Err(self.retryable_conflict_err(other_transaction, other_version)) + } else { + Ok(()) + } + } + // A concurrent Update that removed an overlaid fragment orphans the + // overlay outright — conflict. A row-moving update (RewriteRows) + // deletes the rows it touches and re-creates them in new fragments; + // the update took precedence and the re-created rows were computed + // from the pre-overlay base, so the overlay's values for those cells + // are lost. That is a per-row problem, not an offset one: only the + // moved rows are affected. Comparing the moved rows against the + // overlay's coverage needs the update's deletion vectors, so we mark + // the fragment here and verify row-by-row in `finish_data_overlay`. + // An in-place column rewrite (RewriteColumns) preserves rows and just + // tombstones the overlaid fields at build time, so it never conflicts. + Operation::Update { + removed_fragment_ids, + updated_fragments, + new_fragments, + update_mode, + .. + } => { + let removed_ours = removed_fragment_ids + .iter() + .any(|id| self.modified_fragment_ids.contains(id)); + if removed_ours { + return Err(self.retryable_conflict_err(other_transaction, other_version)); + } + let moves_rows = !new_fragments.is_empty() + && matches!(update_mode, Some(UpdateMode::RewriteRows) | None); + if moves_rows { + for updated in updated_fragments { + if let Some((_, needs_row_check)) = + self.initial_fragments.get_mut(&updated.id) + { + *needs_row_check = true; + } + } + } + Ok(()) + } + Operation::Rewrite { groups, .. } => { + // A rewrite (compaction / fold) of a fragment we are overlaying + // changes its physical row addresses, so our offsets would be + // invalid. Conflict only if it touches one of our fragments. + let touches_our_fragment = groups + .iter() + .flat_map(|g| g.old_fragments.iter()) + .any(|f| self.modified_fragment_ids.contains(&f.id)); + if touches_our_fragment { + Err(self.retryable_conflict_err(other_transaction, other_version)) + } else { + Ok(()) + } + } + Operation::Merge { .. } => { + // Merge rewrites the whole fragment list; always conflict. + Err(self.retryable_conflict_err(other_transaction, other_version)) + } + // Overwrite/Restore replace the dataset; UpdateMemWalState does not + // rebase against data operations (mirroring check_update_mem_wal_state_txn, + // which likewise treats a concurrent DataOverlay as incompatible). + Operation::Overwrite { .. } + | Operation::Restore { .. } + | Operation::UpdateMemWalState { .. } => { + Err(self.incompatible_conflict_err(other_transaction, other_version)) + } + } + } + fn check_merge_txn( &mut self, other_transaction: &Transaction, @@ -1072,7 +1277,8 @@ impl<'a> TransactionRebase<'a> { | Operation::Delete { .. } | Operation::Rewrite { .. } | Operation::Merge { .. } - | Operation::DataReplacement { .. } => { + | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } => { Err(self.retryable_conflict_err(other_transaction, other_version)) } Operation::Overwrite { .. } @@ -1096,6 +1302,7 @@ impl<'a> TransactionRebase<'a> { | Operation::CreateIndex { .. } | Operation::Rewrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::ReserveFragments { .. } @@ -1124,6 +1331,7 @@ impl<'a> TransactionRebase<'a> { | Operation::CreateIndex { .. } | Operation::Rewrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::ReserveFragments { .. } | Operation::Update { .. } @@ -1148,6 +1356,7 @@ impl<'a> TransactionRebase<'a> { | Operation::UpdateConfig { .. } | Operation::CreateIndex { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Rewrite { .. } | Operation::Clone { .. } | Operation::ReserveFragments { .. } @@ -1212,6 +1421,7 @@ impl<'a> TransactionRebase<'a> { | Operation::CreateIndex { .. } | Operation::Rewrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::ReserveFragments { .. } @@ -1284,6 +1494,7 @@ impl<'a> TransactionRebase<'a> { | Operation::Overwrite { .. } | Operation::Delete { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::Clone { .. } @@ -1374,6 +1585,7 @@ impl<'a> TransactionRebase<'a> { } Operation::CreateIndex { .. } => self.finish_create_index(dataset).await, Operation::Rewrite { .. } => self.finish_rewrite(dataset).await, + Operation::DataOverlay { .. } => self.finish_data_overlay(dataset).await, Operation::Append { .. } | Operation::Overwrite { .. } | Operation::DataReplacement { .. } @@ -1550,6 +1762,90 @@ impl<'a> TransactionRebase<'a> { } } + /// Verify no concurrent row-moving Update dropped the values of any cell + /// this overlay covers. `check_data_overlay_txn` flags (via the + /// `initial_fragments` needs-check bool) each overlaid fragment on which a + /// concurrent RewriteRows update relocated rows; here we read the deletion + /// vectors and conflict only when the moved rows intersect the overlay's + /// coverage. + /// + /// The moved rows are computed as the current deletion vector minus the + /// read-time one. In the rare case where both a concurrent Delete and a + /// concurrent Update touched the same flagged fragment, the Delete's rows are + /// also counted and may trigger an unnecessary retry — never data loss. Pure + /// concurrent deletes leave the fragment unflagged and are not examined here. + async fn finish_data_overlay(self, dataset: &Dataset) -> Result { + let fragments_to_check: HashSet = self + .initial_fragments + .iter() + .filter_map(|(id, (_, needs_check))| needs_check.then_some(*id)) + .collect(); + if fragments_to_check.is_empty() { + return Ok(Transaction { + read_version: dataset.manifest.version, + ..self.transaction + }); + } + + // Coverage (physical offsets, unioned across fields) per flagged fragment. + let Operation::DataOverlay { groups } = &self.transaction.operation else { + return Err(wrong_operation_err(&self.transaction.operation)); + }; + let mut coverage_by_fragment: HashMap = HashMap::new(); + for group in groups { + if !fragments_to_check.contains(&group.fragment_id) { + continue; + } + *coverage_by_fragment.entry(group.fragment_id).or_default() |= + overlay_group_coverage(group); + } + + for (fragment_id, coverage) in coverage_by_fragment { + let Some(current_fragment) = dataset + .fragments() + .as_slice() + .iter() + .find(|f| f.id == fragment_id) + else { + // The fragment is gone entirely; the overlay is orphaned. + return Err(crate::Error::retryable_commit_conflict_source( + dataset.manifest.version, + format!( + "This {} transaction was preempted: overlaid fragment {} was removed by a concurrent transaction. Please retry.", + self.transaction.uuid, fragment_id + ) + .into(), + )); + }; + let current_deletions = + read_fragment_deletion_bitmap(dataset, current_fragment).await?; + let initial_deletions = match self.initial_fragments.get(&fragment_id) { + Some((initial_fragment, _)) => { + read_fragment_deletion_bitmap(dataset, initial_fragment).await? + } + None => RoaringBitmap::new(), + }; + let moved_rows = ¤t_deletions - &initial_deletions; + let conflicting = &moved_rows & &coverage; + if !conflicting.is_empty() { + let sample: Vec = conflicting.iter().take(5).collect(); + return Err(crate::Error::retryable_commit_conflict_source( + dataset.manifest.version, + format!( + "This {} transaction was preempted by a concurrent update that moved overlaid rows on fragment {} (offsets {:?}). Please retry.", + self.transaction.uuid, fragment_id, sample.as_slice() + ) + .into(), + )); + } + } + + Ok(Transaction { + read_version: dataset.manifest.version, + ..self.transaction + }) + } + async fn finish_create_index(mut self, dataset: &Dataset) -> Result { if let Operation::CreateIndex { new_indices, @@ -1774,6 +2070,40 @@ async fn initial_fragments_for_rebase( .collect::>() } +/// Read a fragment's deletion vector as a bitmap of physical offsets, or an +/// empty bitmap when the fragment has no deletion file. +async fn read_fragment_deletion_bitmap( + dataset: &Dataset, + fragment: &Fragment, +) -> Result { + match &fragment.deletion_file { + Some(deletion_file) => { + let dv = read_dataset_deletion_file(dataset, fragment.id, deletion_file).await?; + Ok(RoaringBitmap::from(dv.as_ref())) + } + None => Ok(RoaringBitmap::new()), + } +} + +/// The physical offsets a group's overlays cover, unioned across every overlay +/// and every field. This is the set of cells whose values the overlay supplies, +/// used to test whether a concurrent row-moving Update actually invalidates the +/// overlay. +fn overlay_group_coverage(group: &DataOverlayGroup) -> RoaringBitmap { + let mut union = RoaringBitmap::new(); + for overlay in &group.overlays { + match &overlay.coverage { + OverlayCoverage::Shared(bitmap) => union |= bitmap.as_ref(), + OverlayCoverage::PerField(bitmaps) => { + for bitmap in bitmaps { + union |= bitmap.as_ref(); + } + } + } + } + union +} + fn wrong_operation_err(op: &Operation) -> Error { Error::internal(format!("function called against a wrong operation: {}", op)) } @@ -2781,6 +3111,442 @@ mod tests { } } + #[test] + fn test_data_overlay_conflicts() { + use crate::dataset::transaction::{DataOverlayGroup, UpdateMode}; + use ConflictResult::*; + use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; + use roaring::RoaringBitmap; + + // Our transaction overlays fragment 1. + let overlay_op = |fragment_id: u64| Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id, + overlays: vec![DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay.lance", vec![0], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 0, + }], + }], + }; + let update_removing = |removed_fragment_ids: Vec| Operation::Update { + removed_fragment_ids, + updated_fragments: vec![], + new_fragments: vec![], + fields_modified: vec![], + merged_generations: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode: None, + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + let delete = |updated: Vec, deleted: Vec| Operation::Delete { + updated_fragments: updated, + deleted_fragment_ids: deleted, + predicate: "x > 2".to_string(), + }; + // A row-moving update (RewriteRows) relocates the updated rows into + // new_fragments; an in-place column rewrite (RewriteColumns) leaves rows + // where they are. + let update_moving = |updated: Vec, new: Vec| Operation::Update { + removed_fragment_ids: vec![], + updated_fragments: updated, + new_fragments: new, + fields_modified: vec![], + merged_generations: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode: Some(UpdateMode::RewriteRows), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + let update_rewrite_columns = |updated: Vec| Operation::Update { + removed_fragment_ids: vec![], + updated_fragments: updated, + new_fragments: vec![], + fields_modified: vec![0], + merged_generations: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode: Some(UpdateMode::RewriteColumns), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + let rewrite_of = |old: &Fragment| Operation::Rewrite { + groups: vec![RewriteGroup { + old_fragments: vec![old.clone()], + new_fragments: vec![], + }], + rewritten_indices: vec![], + frag_reuse_index: None, + }; + + let fragment0 = Fragment::new(0); + let fragment1 = Fragment::new(1); + + // Each case is checked against our overlay on fragment 1. + let cases: Vec<(Operation, ConflictResult)> = vec![ + // Permissive: preserves physical offsets / leaves fragment 1 in place. + ( + Operation::Append { + fragments: vec![fragment0.clone()], + }, + Compatible, + ), + ( + Operation::CreateIndex { + new_indices: vec![], + removed_indices: vec![], + }, + Compatible, + ), + ( + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 1, + DataFile::new_legacy_from_fields("r.lance", vec![0], None), + )], + }, + Compatible, + ), + // Another overlay on the same fragment stacks rather than conflicts. + (overlay_op(1), Compatible), + // A Delete only tombstones rows (deletion vector) on fragment 1, and + // an in-place column rewrite preserves offsets, so both are compatible. + (delete(vec![fragment1.clone()], vec![]), Compatible), + (update_rewrite_columns(vec![fragment1.clone()]), Compatible), + (update_removing(vec![2]), Compatible), + // ...but removing our overlaid fragment 1 orphans the overlay -> conflict. + (delete(vec![], vec![1]), Retryable), + (update_removing(vec![1]), Retryable), + // A row-moving update re-creates the rows it touches from the + // pre-overlay base. Whether that actually drops any overlaid cell is + // a per-row question answered in `finish_data_overlay` (see + // test_data_overlay_finish_conflicts_with_row_moving_update), so the + // check itself defers rather than conflicting; a moving update on any + // fragment is compatible at this stage. + ( + update_moving(vec![fragment1.clone()], vec![fragment0.clone()]), + Compatible, + ), + ( + update_moving(vec![fragment0.clone()], vec![fragment0.clone()]), + Compatible, + ), + // Rewriting fragment 1 invalidates its physical offsets -> conflict; + // a rewrite of a different fragment does not. + (rewrite_of(&fragment1), Retryable), + (rewrite_of(&fragment0), Compatible), + // Merge rewrites the whole fragment list; Restore replaces the dataset. + ( + Operation::Merge { + fragments: vec![fragment1.clone()], + schema: lance_core::datatypes::Schema::default(), + }, + Retryable, + ), + (Operation::Restore { version: 1 }, NotCompatible), + // Overwrite/Restore replace the dataset, and UpdateMemWalState does + // not rebase against data operations — all hard conflicts. + ( + Operation::Overwrite { + fragments: vec![fragment0.clone()], + schema: lance_core::datatypes::Schema::default(), + config_upsert_values: None, + initial_bases: None, + }, + NotCompatible, + ), + ( + Operation::UpdateMemWalState { + merged_generations: vec![], + }, + NotCompatible, + ), + ]; + + for (other, expected) in cases { + let mut rebase = TransactionRebase { + transaction: Transaction::new(0, overlay_op(1), None), + initial_fragments: HashMap::new(), + modified_fragment_ids: modified_fragment_ids(&overlay_op(1)) + .collect::>(), + affected_rows: None, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_merged_gens: Vec::new(), + }; + let other_txn = Transaction::new(0, other.clone(), None); + let result = rebase.check_txn(&other_txn, 1); + match expected { + Compatible => assert!( + result.is_ok(), + "overlay should be compatible with {other:?}, got {result:?}" + ), + Retryable => assert!( + matches!(result, Err(Error::RetryableCommitConflict { .. })), + "overlay should retryably conflict with {other:?}, got {result:?}" + ), + NotCompatible => assert!( + matches!(result, Err(Error::IncompatibleTransaction { .. })), + "overlay should be incompatible with {other:?}, got {result:?}" + ), + } + } + } + + #[test] + fn test_rewrite_conflicts_with_data_overlay() { + // Reverse direction of test_data_overlay_conflicts: our transaction is a + // Rewrite and a concurrent DataOverlay has already committed. A rewrite + // changes the physical row addresses of the fragments it touches, so an + // overlay on one of those fragments is invalidated (retryable); an + // overlay on any other fragment is unaffected. + use crate::dataset::transaction::DataOverlayGroup; + use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; + use roaring::RoaringBitmap; + + let overlay_on = |fragment_id: u64| Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id, + overlays: vec![DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay.lance", vec![0], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 0, + }], + }], + }; + // Our transaction rewrites fragment 1. + let rewrite_op = Operation::Rewrite { + groups: vec![RewriteGroup { + old_fragments: vec![Fragment::new(1)], + new_fragments: vec![], + }], + rewritten_indices: vec![], + frag_reuse_index: None, + }; + + for (other, expect_conflict) in [(overlay_on(1), true), (overlay_on(0), false)] { + let mut rebase = TransactionRebase { + transaction: Transaction::new(0, rewrite_op.clone(), None), + initial_fragments: HashMap::new(), + modified_fragment_ids: modified_fragment_ids(&rewrite_op).collect::>(), + affected_rows: None, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_merged_gens: Vec::new(), + }; + let other_txn = Transaction::new(0, other.clone(), None); + let result = rebase.check_txn(&other_txn, 1); + if expect_conflict { + assert!( + matches!(result, Err(Error::RetryableCommitConflict { .. })), + "rewrite of fragment 1 should retryably conflict with {other:?}, got {result:?}" + ); + } else { + assert!( + result.is_ok(), + "rewrite of fragment 1 should not conflict with {other:?}, got {result:?}" + ); + } + } + } + + #[test] + fn test_update_conflicts_with_data_overlay() { + // Reverse direction of test_data_overlay_conflicts: our transaction is an + // Update and a concurrent DataOverlay has already committed. A row-moving + // update relocates the rows it touches, so an overlay on one of those + // fragments can no longer be applied (retryable); an overlay on any other + // fragment, or an in-place column rewrite, is compatible. + use crate::dataset::transaction::{DataOverlayGroup, UpdateMode}; + use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; + use roaring::RoaringBitmap; + + let overlay_on = |fragment_id: u64| Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id, + overlays: vec![DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay.lance", vec![0], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 0, + }], + }], + }; + // Our update always touches fragment 1. + let update = + |update_mode: Option, new_fragments: Vec| Operation::Update { + removed_fragment_ids: vec![], + updated_fragments: vec![Fragment::new(1)], + new_fragments, + fields_modified: vec![0], + merged_generations: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode, + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + + // The overlay covers physical offset 0 of its fragment. Row addresses + // pack the fragment id in the high 32 bits and the offset in the low 32. + let rows_on = |fragment_id: u64, offsets: &[u32]| { + let mut map = RowAddrTreeMap::new(); + map.insert_bitmap( + fragment_id as u32, + RoaringBitmap::from_iter(offsets.iter().copied()), + ); + map + }; + + // (update, committed overlay, moved rows the update carries, expect conflict) + let cases = [ + // Row-moving update whose moved rows include the overlaid cell -> the + // update would undo the overlay, so conflict. + ( + update(Some(UpdateMode::RewriteRows), vec![Fragment::new(2)]), + overlay_on(1), + Some(rows_on(1, &[0])), + true, + ), + // ...but if the moved rows miss the overlaid cell, the overlay survives. + ( + update(Some(UpdateMode::RewriteRows), vec![Fragment::new(2)]), + overlay_on(1), + Some(rows_on(1, &[5])), + false, + ), + // An overlay on a fragment the update did not touch is fine. + ( + update(Some(UpdateMode::RewriteRows), vec![Fragment::new(2)]), + overlay_on(0), + Some(rows_on(1, &[0])), + false, + ), + // An in-place column rewrite preserves rows -> compatible. + ( + update(Some(UpdateMode::RewriteColumns), vec![]), + overlay_on(1), + Some(rows_on(1, &[0])), + false, + ), + // Without affected rows we cannot be precise, so a row-moving update + // on the overlaid fragment falls back to a conservative conflict. + ( + update(Some(UpdateMode::RewriteRows), vec![Fragment::new(2)]), + overlay_on(1), + None, + true, + ), + ]; + + for (update_op, other, affected_rows, expect_conflict) in cases { + let mut rebase = TransactionRebase { + transaction: Transaction::new(0, update_op.clone(), None), + initial_fragments: HashMap::new(), + modified_fragment_ids: modified_fragment_ids(&update_op).collect::>(), + affected_rows: affected_rows.as_ref(), + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_merged_gens: Vec::new(), + }; + let other_txn = Transaction::new(0, other.clone(), None); + let result = rebase.check_txn(&other_txn, 1); + if expect_conflict { + assert!( + matches!(result, Err(Error::RetryableCommitConflict { .. })), + "update should retryably conflict with {other:?}, got {result:?}" + ); + } else { + assert!( + result.is_ok(), + "update should be compatible with {other:?}, got {result:?}" + ); + } + } + } + + #[tokio::test] + #[rstest::rstest] + #[case::coverage_overlaps_moved_row(vec![0u32], true)] + #[case::coverage_disjoint_from_moved_row(vec![3u32], false)] + async fn test_data_overlay_finish_conflicts_with_row_moving_update( + #[case] coverage_offsets: Vec, + #[case] expect_conflict: bool, + ) { + // 5 rows in one fragment. A concurrent RewriteRows update moves row 0 out + // to a new fragment (deleting it from fragment 0). Our overlay on fragment + // 0 conflicts only when its coverage includes the moved row; the decision + // is made in finish, which reads the deletion vectors. + use crate::dataset::transaction::{DataOverlayGroup, UpdateMode}; + use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; + use roaring::RoaringBitmap; + + let dataset = test_dataset(5, 1).await; + let mut fragment = dataset.fragments().as_slice()[0].clone(); + + let moved_fragment = Fragment::new(0) + .with_file( + "moved.lance", + vec![0], + vec![0], + &LanceFileVersion::Stable, + NonZero::new(10), + ) + .with_physical_rows(1); + let update_op = Operation::Update { + updated_fragments: vec![apply_deletion(&[0], &mut fragment, &dataset).await], + removed_fragment_ids: vec![], + new_fragments: vec![moved_fragment], + fields_modified: vec![], + merged_generations: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode: Some(UpdateMode::RewriteRows), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + let update_txn = Transaction::new_from_version(dataset.manifest.version, update_op); + + let overlay_op = Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id: 0, + overlays: vec![DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay.lance", vec![0], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter(coverage_offsets)), + committed_version: 0, + }], + }], + }; + let overlay_txn = Transaction::new_from_version(dataset.manifest.version, overlay_op); + + // Commit the update so the latest dataset reflects the moved (deleted) row. + let latest_dataset = CommitBuilder::new(Arc::new(dataset.clone())) + .execute(update_txn.clone()) + .await + .unwrap(); + + let mut rebase = TransactionRebase::try_new(&dataset, overlay_txn.clone(), None) + .await + .unwrap(); + // The check defers the row-level decision to finish, flagging fragment 0. + rebase.check_txn(&update_txn, 1).unwrap(); + assert_eq!( + rebase + .initial_fragments + .iter() + .map(|(id, (_, needs_check))| (*id, *needs_check)) + .collect::>(), + vec![(0, true)], + ); + + let res = rebase.finish(&latest_dataset).await; + if expect_conflict { + assert!( + matches!(res, Err(crate::Error::RetryableCommitConflict { .. })), + "overlay covering the moved row should conflict, got {res:?}" + ); + } else { + assert!( + res.is_ok(), + "overlay disjoint from the moved row should succeed, got {res:?}" + ); + } + } + #[test] fn test_create_index_conflicts_only_on_same_name() { let index0 = IndexMetadata { @@ -3255,6 +4021,7 @@ mod tests { Operation::DataReplacement { replacements } => { Box::new(replacements.iter().map(|r| r.0)) } + Operation::DataOverlay { groups } => Box::new(groups.iter().map(|g| g.fragment_id)), } } diff --git a/rust/lance/src/io/commit/s3_test.rs b/rust/lance/src/io/commit/s3_test.rs index b5b1a09c776..4be469ee368 100644 --- a/rust/lance/src/io/commit/s3_test.rs +++ b/rust/lance/src/io/commit/s3_test.rs @@ -341,7 +341,10 @@ async fn test_ddb_open_iops() { // Checkout original version dataset.checkout_version(1).await.unwrap(); let io_stats = dataset.object_store.as_ref().io_stats_incremental(); - // Checkout: 1 IOPS: manifest file - assert_io_eq!(io_stats, read_iops, 1); + // Checkout: 0 read IOPS. Version 1's manifest was already loaded and cached + // on this Session when the dataset was opened above, so the checkout serves + // the manifest body from the metadata cache. Version resolution is handled + // in DynamoDB and issues no S3 read. + assert_io_eq!(io_stats, read_iops, 0); assert_io_eq!(io_stats, write_iops, 0); } diff --git a/rust/lance/src/io/exec/count_pushdown.rs b/rust/lance/src/io/exec/count_pushdown.rs index d5d90b5881a..50804466cba 100644 --- a/rust/lance/src/io/exec/count_pushdown.rs +++ b/rust/lance/src/io/exec/count_pushdown.rs @@ -380,7 +380,8 @@ fn strip_row_preserving_wrappers(plan: &Arc) -> Option<&Filte inner.input() } else if let Some(inner) = current.as_any().downcast_ref::() { inner.input() - } else if let Some(proj) = current.as_any().downcast_ref::() { + } else { + let proj = current.as_any().downcast_ref::()?; // Only walk through projections that are row-preserving: every // output expression is a direct column reference back to the // input. (Empty projections trivially qualify — DataFusion uses @@ -398,8 +399,6 @@ fn strip_row_preserving_wrappers(plan: &Arc) -> Option<&Filte return None; } proj.input() - } else { - return None; }; current = next.as_ref(); } diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index b50cba1f9ce..f944a3a1f94 100644 --- a/rust/lance/src/io/exec/filtered_read.rs +++ b/rust/lance/src/io/exec/filtered_read.rs @@ -2,13 +2,12 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use std::any::Any; use std::collections::{BTreeMap, HashMap}; -use std::pin::Pin; use std::sync::Mutex; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::{ops::Range, sync::Arc}; use arrow_array::RecordBatch; -use arrow_schema::SchemaRef; +use arrow_schema::{Schema as ArrowSchema, SchemaRef}; use datafusion::common::runtime::SpawnedTask; use datafusion::common::stats::Precision; use datafusion::error::{DataFusionError, Result as DataFusionResult}; @@ -60,6 +59,13 @@ use crate::dataset::scanner::{ use super::utils::IoMetrics; +fn public_blob_v2_binary_projection_schema(projection: &Projection) -> SchemaRef { + let schema = projection.to_schema(); + let schema = crate::dataset::blob::public_blob_v2_binary_output_schema(&schema); + let schema: ArrowSchema = (&schema).into(); + Arc::new(schema) +} + #[derive(Debug)] pub struct EvaluatedIndex { index_result: IndexExprResult, @@ -385,7 +391,7 @@ impl FilteredReadStream { .try_collect::>() .await?; - let output_schema = Arc::new(options.projection.to_arrow_schema()); + let output_schema = public_blob_v2_binary_projection_schema(&options.projection); let obj_store = dataset.object_store.clone(); // Explicit options take precedence; otherwise fall back to the @@ -417,11 +423,14 @@ impl FilteredReadStream { let fragment_streams = futures::stream::iter(scoped_fragments) .map({ let scan_range_after_filter = scan_range_after_filter.clone(); + let dataset = dataset.clone(); move |scoped_fragment| { let metrics = global_metrics_clone.clone(); let limit = scan_range_after_filter.as_ref().map(|r| r.end); + let dataset = dataset.clone(); SpawnedTask::spawn( - Self::read_fragment(scoped_fragment, metrics, limit).in_current_span(), + Self::read_fragment(dataset, scoped_fragment, metrics, limit) + .in_current_span(), ) .map(|thread_result| thread_result.unwrap()) } @@ -1074,13 +1083,15 @@ impl FilteredReadStream { } // Reads a single fragment into a stream of batch tasks - #[instrument(name = "read_fragment", skip_all)] + #[instrument(name = "read_fragment", level = "debug", skip_all)] async fn read_fragment( + dataset: Arc, mut fragment_read_task: ScopedFragmentRead, global_metrics: Arc, fragment_soft_limit: Option, - ) -> Result>> { - let output_schema = Arc::new(fragment_read_task.projection.to_arrow_schema()); + ) -> Result>> { + let output_schema = + public_blob_v2_binary_projection_schema(fragment_read_task.projection.as_ref()); if let Some(filter) = &fragment_read_task.filter { let filter_cols = Planner::column_names_in_expr(filter); @@ -1095,10 +1106,22 @@ impl FilteredReadStream { } } - let read_schema = fragment_read_task.projection.to_bare_schema(); + let output_read_schema = Arc::new(fragment_read_task.projection.to_schema()); + let bare_read_schema = fragment_read_task.projection.to_bare_schema(); + let materialize_blob_v2_binary = + crate::dataset::blob::schema_has_blob_v2_binary_view(&bare_read_schema); + let read_schema = if materialize_blob_v2_binary { + crate::dataset::blob::blob_v2_descriptor_schema(&bare_read_schema) + } else { + bare_read_schema + }; + let mut frag_read_config = fragment_read_task.frag_read_config(); + if materialize_blob_v2_binary { + frag_read_config = frag_read_config.with_row_address(true); + } let mut fragment_reader = fragment_read_task .fragment - .open(&read_schema, fragment_read_task.frag_read_config()) + .open(&read_schema, frag_read_config) .await?; if fragment_read_task.with_deleted_rows { @@ -1112,8 +1135,9 @@ impl FilteredReadStream { let physical_filter = fragment_read_task .filter .map(|filter| { - let planner = - Planner::new(Arc::new(fragment_read_task.projection.to_arrow_schema())); + let planner = Planner::new(public_blob_v2_binary_projection_schema( + fragment_read_task.projection.as_ref(), + )); planner.create_physical_expr(&filter) }) .transpose()?; @@ -1136,7 +1160,7 @@ impl FilteredReadStream { let global_metrics = global_metrics.clone(); let fragment_counted = fragment_counted.clone(); let range_tracker = range_tracker.clone(); - batch_fut + let batch_fut = batch_fut .inspect_ok(move |batch| { let num_rows = batch.num_rows(); global_metrics.rows_scanned.add(num_rows); @@ -1151,7 +1175,23 @@ impl FilteredReadStream { global_metrics.ranges_scanned.add(additional_ranges); } }) - .boxed() + .boxed(); + if materialize_blob_v2_binary { + let dataset = dataset.clone(); + let output_read_schema = output_read_schema.clone(); + batch_fut + .and_then(move |batch| async move { + crate::dataset::blob::materialize_blob_v2_binary_batch( + &dataset, + output_read_schema.as_ref(), + batch, + ) + .await + }) + .boxed() + } else { + batch_fut + } }) .zip(futures::stream::repeat(( physical_filter.clone(), @@ -1159,12 +1199,11 @@ impl FilteredReadStream { ))) .map(|(batch_fut, args)| Self::wrap_with_filter(batch_fut, args.0, args.1)); - let result: Pin> + Send>> = - if let Some(limit) = fragment_soft_limit { - Box::pin(Self::apply_soft_limit(fragment_stream, limit)) - } else { - Box::pin(fragment_stream) - }; + let result = if let Some(limit) = fragment_soft_limit { + Self::apply_soft_limit(fragment_stream, limit).boxed() + } else { + fragment_stream.boxed() + }; Ok(result) } @@ -1482,6 +1521,19 @@ impl FilteredReadOptions { self.only_indexed_fragments = true; self } + + /// Specify the threading mode to use for the scan. + /// + /// This controls how decode work is parallelized. For the default single-partition + /// scan, the parameter of [`FilteredReadThreadingMode::OnePartitionMultipleThreads`] + /// bounds how many batch-decode tasks are buffered in flight (via `try_buffered`). + /// + /// The parallelism must be greater than 0. A value of 0 is rejected by + /// [`FilteredReadExec::try_new`]. + pub fn with_threading_mode(mut self, threading_mode: FilteredReadThreadingMode) -> Self { + self.threading_mode = threading_mode; + self + } } /// A plan node that reads a dataset, applying an optional filter and projection. @@ -1573,6 +1625,23 @@ impl FilteredReadExec { .into())); } + // A parallelism of 0 would cause `try_buffered(0)` to hang forever instead of erroring + match options.threading_mode { + FilteredReadThreadingMode::OnePartitionMultipleThreads(0) => { + return Err(Error::invalid_input_source( + "FilteredReadThreadingMode::OnePartitionMultipleThreads must be greater than 0, got 0" + .into(), + )); + } + FilteredReadThreadingMode::MultiplePartitions(0) => { + return Err(Error::invalid_input_source( + "FilteredReadThreadingMode::MultiplePartitions must be greater than 0, got 0" + .into(), + )); + } + _ => {} + } + if options.scan_range_after_filter.is_some() { // Validate that there's a filter when using scan_range_after_filter if options.full_filter.is_none() @@ -1595,7 +1664,7 @@ impl FilteredReadExec { )); } } - let output_schema = Arc::new(options.projection.to_arrow_schema()); + let output_schema = public_blob_v2_binary_projection_schema(&options.projection); let num_partitions = match options.threading_mode { FilteredReadThreadingMode::OnePartitionMultipleThreads(_) => 1, FilteredReadThreadingMode::MultiplePartitions(n) => n, @@ -1973,7 +2042,7 @@ impl ExecutionPlan for FilteredReadExec { .clone() .union_columns(filter_columns, OnMissing::Error)?; - let read_schema = Arc::new(read_projection.to_arrow_schema()); + let read_schema = public_blob_v2_binary_projection_schema(&read_projection); let planner = Arc::new(Planner::new(read_schema.clone())); let physical_filter = planner.create_physical_expr(filter)?; diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index add864c0ea9..ed2e3859f60 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -135,7 +135,7 @@ async fn search_segments( let mut searches = searches; while let Some((doc_ids, scores)) = searches.try_next().await? { - for (row_id, score) in doc_ids.into_iter().zip(scores.into_iter()) { + for (row_id, score) in doc_ids.into_iter().zip(scores) { if candidates.len() < limit { candidates.push(std::cmp::Reverse(ScoredDoc::new(row_id, score))); } else if candidates.peek().unwrap().0.score.0 < score { diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 01125ac1617..3b82ec85056 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +#[cfg(test)] +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::any::Any; use std::cmp::Ordering as CmpOrdering; use std::collections::{BinaryHeap, HashMap, HashSet}; @@ -2319,6 +2321,7 @@ mod tests { use crate::dataset::{WriteMode, WriteParams}; use crate::index::vector::VectorIndexParams; + use crate::index::vector::ivf::v2::STREAMING_SEARCH_BATCH_SIZE; use crate::io::exec::testing::TestingExec; fn base_query() -> Query { @@ -2492,7 +2495,7 @@ mod tests { Box::new(self.row_ids.iter()) } - async fn remap(&mut self, _mapping: &HashMap>) -> Result<()> { + async fn remap(&mut self, _mapping: &RowAddrRemap) -> Result<()> { Ok(()) } @@ -2632,7 +2635,6 @@ mod tests { _metrics: Arc, ) -> Result { let (batch_tx, batch_rx) = mpsc::channel(1); - let batch_tx_for_search = batch_tx.clone(); let prepared_partition_ids = (start_idx..end_idx) .map(|idx| partitions.value(idx) as usize) .collect::>(); @@ -2640,42 +2642,74 @@ mod tests { .lock() .unwrap() .extend(prepared_partition_ids.iter().copied()); + // Mirror the production streaming path (v2.rs): search prepared partitions + // in batches of STREAMING_SEARCH_BATCH_SIZE, one `spawn_cpu` per batch, with + // the channel send in async code so no CPU-pool thread parks (#7642). tokio::spawn(async move { - let search_result = spawn_cpu(move || -> DataFusionResult<()> { - for partition_id in prepared_partition_ids { - if control - .as_ref() - .is_some_and(|control| control.should_stop()) - { - return Ok(()); - } - let batch = self - .search_prepared_partition( - Box::new(partition_id), - &lance_index::metrics::NoOpMetricsCollector, - ) - .map_err(datafusion::error::DataFusionError::from); - match batch { - Ok(batch) => { - if let Some(control) = control.as_ref() { - control.record_batch(&batch); + for chunk in prepared_partition_ids.chunks(*STREAMING_SEARCH_BATCH_SIZE) { + if control + .as_ref() + .is_some_and(|control| control.should_stop()) + || batch_tx.is_closed() + { + return; + } + let chunk = chunk.to_vec(); + let index = self.clone(); + let control_for_search = control.clone(); + let cancel_probe = batch_tx.clone(); + let search_output = spawn_cpu(move || { + let mut outputs: Vec> = + Vec::with_capacity(chunk.len()); + let mut stopped = false; + for partition_id in chunk { + if control_for_search + .as_ref() + .is_some_and(|control| control.should_stop()) + || cancel_probe.is_closed() + { + stopped = true; + break; + } + match index + .search_prepared_partition( + Box::new(partition_id), + &lance_index::metrics::NoOpMetricsCollector, + ) + .map_err(datafusion::error::DataFusionError::from) + { + Ok(batch) => { + if let Some(control) = control_for_search.as_ref() { + control.record_batch(&batch); + } + outputs.push(Ok(batch)); } - if batch_tx_for_search.blocking_send(Ok(batch)).is_err() { - return Ok(()); + Err(err) => { + outputs.push(Err(err)); + stopped = true; + break; } } - Err(err) => { - let _ = batch_tx_for_search.blocking_send(Err(err)); - return Ok(()); - } } - } - Ok(()) - }) - .await; + Ok::<_, datafusion::error::DataFusionError>((outputs, stopped)) + }) + .await; - if let Err(err) = search_result { - let _ = batch_tx.send(Err(err)).await; + let (outputs, stopped) = match search_output { + Ok(output) => output, + Err(err) => { + let _ = batch_tx.send(Err(err)).await; + return; + } + }; + for output in outputs { + if batch_tx.send(output).await.is_err() { + return; + } + } + if stopped { + return; + } } }); @@ -2710,7 +2744,7 @@ mod tests { Box::new(self.row_ids.iter()) } - async fn remap(&mut self, _mapping: &HashMap>) -> Result<()> { + async fn remap(&mut self, _mapping: &RowAddrRemap) -> Result<()> { Ok(()) } @@ -2869,19 +2903,29 @@ mod tests { ); } + // All partitions fit in a single search batch, so they are searched in one + // `spawn_cpu` dispatch and therefore share one cpu thread. The partition count + // adapts to the configured batch size so the single-batch property holds under + // any valid `LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE`, including 1. #[tokio::test] async fn test_sequential_initial_search_prepares_all_then_searches_on_one_cpu_thread() { + let num_partitions = 3.min(*STREAMING_SEARCH_BATCH_SIZE); + let row_ids = (0..num_partitions).map(|i| 10 + i as u64).collect(); let (index, prepared_partitions, searched_partitions, search_threads) = - prepared_index(vec![10, 11, 12]); + prepared_index(row_ids); let mut query = base_query(); - query.minimum_nprobes = 3; + query.minimum_nprobes = num_partitions; let state = Arc::new(ANNIvfEarlySearchResults::new(1, query.k)); + let partition_idx = (0..num_partitions as u32).collect::>(); + let q_c_dists = (0..num_partitions) + .map(|i| i as f32 * 0.1) + .collect::>(); let batches = ANNIvfSubIndexExec::initial_search( index, query, - Arc::new(UInt32Array::from(vec![0, 1, 2])), - Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3])), + Arc::new(UInt32Array::from(partition_idx)), + Arc::new(Float32Array::from(q_c_dists)), empty_prefilter().await, prepared_metrics(), state, @@ -2891,11 +2935,12 @@ mod tests { .await .unwrap(); - assert_eq!(batches.len(), 3); - assert_eq!(*prepared_partitions.lock().unwrap(), vec![0, 1, 2]); - assert_eq!(*searched_partitions.lock().unwrap(), vec![0, 1, 2]); + let expected: Vec = (0..num_partitions).collect(); + assert_eq!(batches.len(), num_partitions); + assert_eq!(*prepared_partitions.lock().unwrap(), expected); + assert_eq!(*searched_partitions.lock().unwrap(), expected); let search_threads = search_threads.lock().unwrap().clone(); - assert_eq!(search_threads.len(), 3); + assert_eq!(search_threads.len(), num_partitions); assert!( search_threads.iter().all(|name| name.contains("lance-cpu")), "expected prepared searches to run on the cpu runtime, got threads {search_threads:?}", @@ -2906,6 +2951,52 @@ mod tests { ); } + // Regression guard for the batched streaming search (#7642): with more partitions + // than a single batch, the search spans multiple `spawn_cpu` dispatches. Verify that + // every partition is still prepared and searched in order across the batch boundary, + // and that all search work stays on the cpu runtime. + // + // Note: this does not reproduce the single-thread-pool deadlock the async recv/send + // fixes -- that requires a 1-thread CPU pool, which is a process-global singleton and + // impractical to force in a unit test (same limitation noted for the #7423 fix). + #[tokio::test] + async fn test_sequential_search_spans_multiple_cpu_batches() { + let num_partitions = *STREAMING_SEARCH_BATCH_SIZE + 3; + let row_ids = (0..num_partitions).map(|i| i as u64 * 10).collect(); + let (index, prepared_partitions, searched_partitions, search_threads) = + prepared_index(row_ids); + let mut query = base_query(); + query.minimum_nprobes = num_partitions; + let state = Arc::new(ANNIvfEarlySearchResults::new(1, query.k)); + + let partition_idx = (0..num_partitions as u32).collect::>(); + let q_c_dists = (0..num_partitions).map(|i| i as f32).collect::>(); + let batches = ANNIvfSubIndexExec::initial_search( + index, + query, + Arc::new(UInt32Array::from(partition_idx.clone())), + Arc::new(Float32Array::from(q_c_dists)), + empty_prefilter().await, + prepared_metrics(), + state, + usize::MAX, + ) + .try_collect::>() + .await + .unwrap(); + + let expected: Vec = (0..num_partitions).collect(); + assert_eq!(batches.len(), num_partitions); + assert_eq!(*prepared_partitions.lock().unwrap(), expected); + assert_eq!(*searched_partitions.lock().unwrap(), expected); + let search_threads = search_threads.lock().unwrap().clone(); + assert_eq!(search_threads.len(), num_partitions); + assert!( + search_threads.iter().all(|name| name.contains("lance-cpu")), + "expected prepared searches to run on the cpu runtime, got threads {search_threads:?}", + ); + } + #[tokio::test] async fn test_sequential_late_search_prepares_all_then_stops_search_early() { let (index, prepared_partitions, searched_partitions, _search_threads) = diff --git a/rust/lance/src/io/exec/projection.rs b/rust/lance/src/io/exec/projection.rs index 3106fcfac61..06bc0b8de67 100644 --- a/rust/lance/src/io/exec/projection.rs +++ b/rust/lance/src/io/exec/projection.rs @@ -44,7 +44,7 @@ pub fn project(input: Arc, projection: &ArrowSchema) -> Resul let field_names = projection.fields().iter().map(|f| f.name()).cloned(); - for (name, selection) in field_names.zip(selections.into_iter()) { + for (name, selection) in field_names.zip(selections) { let expr = selection_as_expr(&selection, input_schema.fields(), None); exprs.push((expr, name)); } diff --git a/rust/lance/src/io/exec/pushdown_scan.rs b/rust/lance/src/io/exec/pushdown_scan.rs index 0d49232c094..2ac243bff71 100644 --- a/rust/lance/src/io/exec/pushdown_scan.rs +++ b/rust/lance/src/io/exec/pushdown_scan.rs @@ -25,7 +25,8 @@ use datafusion::{ }; use datafusion_functions::core::expr_ext::FieldAccessor; use datafusion_physical_expr::EquivalenceProperties; -use futures::{FutureExt, Stream, StreamExt, TryStreamExt}; +use futures::stream::BoxStream; +use futures::{FutureExt, StreamExt, TryStreamExt}; use lance_arrow::{RecordBatchExt, SchemaExt}; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_core::{ROW_ADDR, ROW_ADDR_FIELD, ROW_ID_FIELD}; @@ -218,7 +219,7 @@ impl ExecutionPlan for LancePushdownScanExec { ) .await?; - frag_scanner.scan().await + frag_scanner.scan() }); let batch_stream = batch_stream @@ -325,7 +326,7 @@ impl FragmentScanner { }) } - pub async fn scan(self) -> Result> + 'static + Send> { + pub fn scan(self) -> Result>> { let batch_readahead = self.config.batch_readahead; let simplified_predicates = self.simplified_predicates()?; let ordered_output = self.config.ordered_output; diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index ee05ce7a86f..7a6726172d5 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -6,7 +6,7 @@ use std::sync::{Arc, LazyLock}; use super::utils::{IndexMetrics, InstrumentedRecordBatchStreamAdapter}; use crate::{ Dataset, - dataset::rowids::load_row_id_sequences, + dataset::rowids::{load_row_id_sequence, load_row_id_sequences}, index::{ prefilter::DatasetPreFilter, scalar_logical::{open_named_scalar_index, scalar_index_fragment_bitmap}, @@ -42,7 +42,8 @@ use lance_index::{ }, }; use lance_select::{ - IndexExprResult, RowAddrMask, RowAddrTreeMap, RowSetOps, result::IndexExprResultWireFormat, + IndexExprResult, NullableIndexExprResult, NullableRowAddrMask, NullableRowAddrSet, RowAddrMask, + RowAddrSelection, RowAddrTreeMap, RowSetOps, result::IndexExprResultWireFormat, }; use lance_table::format::Fragment; use roaring::RoaringBitmap; @@ -58,6 +59,112 @@ impl ScalarIndexLoader for Dataset { ) -> Result> { open_named_scalar_index(self, column, index_name, metrics).await } + + async fn row_addr_result_to_row_ids( + &self, + result: NullableIndexExprResult, + ) -> Result { + // Addresses and row ids only diverge under stable row ids; otherwise the + // address is the row id and there is nothing to translate. + if !self.manifest.uses_stable_row_ids() { + return Ok(result); + } + + let NullableIndexExprResult { lower, upper, .. } = result; + let lower = translate_addr_mask_to_row_ids(self, lower).await?; + let upper = translate_addr_mask_to_row_ids(self, upper).await?; + Ok(NullableIndexExprResult::new(lower, upper)) + } +} + +/// Translate an address-domain [`NullableRowAddrMask`] into the row-id domain +/// +/// Address-domain index results are always positive allow-lists (`AtMost`), so +/// a block-list here would mean a boolean op was applied before translation, +/// which is unsupported. +async fn translate_addr_mask_to_row_ids( + dataset: &Dataset, + mask: NullableRowAddrMask, +) -> Result { + match mask { + NullableRowAddrMask::AllowList(set) => Ok(NullableRowAddrMask::AllowList( + translate_addr_set_to_row_ids(dataset, set).await?, + )), + NullableRowAddrMask::BlockList(_) => Err(Error::internal( + "cannot translate a block-list address mask to the row-id domain", + )), + } +} + +async fn translate_addr_set_to_row_ids( + dataset: &Dataset, + set: NullableRowAddrSet, +) -> Result { + let selected = translate_addr_treemap_to_row_ids(dataset, set.selected_rows()).await?; + let nulls = translate_addr_treemap_to_row_ids(dataset, set.null_rows()).await?; + Ok(NullableRowAddrSet::new(selected, nulls)) +} + +/// Map a set of physical row addresses to their stable row ids +/// +/// For each fragment present in `addrs`, the live rows in physical order carry +/// the stable ids yielded by the fragment's [`RowIdSequence`] in the same +/// order. Zipping the two (skipping deleted physical offsets) gives the +/// `physical offset -> stable id` mapping. Addresses that point at deleted rows +/// have no live counterpart and are dropped, which is correct: those rows are +/// not part of the answer. +async fn translate_addr_treemap_to_row_ids( + dataset: &Dataset, + addrs: &RowAddrTreeMap, +) -> Result { + let mut row_ids = RowAddrTreeMap::new(); + for (fragment_id, selection) in addrs.iter() { + let file_fragment = dataset.get_fragment(*fragment_id as usize).ok_or_else(|| { + Error::internal(format!( + "fragment {fragment_id} referenced by an address-domain index result \ + was not found in the dataset" + )) + })?; + let sequence = load_row_id_sequence(dataset, file_fragment.metadata()).await?; + + match selection { + RowAddrSelection::Full => { + // The whole fragment is selected: every live row's id qualifies. + row_ids |= RowAddrTreeMap::from(sequence.as_ref()); + } + RowAddrSelection::Partial(offsets) => { + let Some(max_offset) = offsets.max() else { + continue; + }; + let (deletion_vector, num_physical_rows) = futures::try_join!( + file_fragment.get_deletion_vector(), + file_fragment.physical_rows() + )?; + let num_physical_rows = num_physical_rows as u32; + let mut ids = sequence.iter(); + for physical_offset in 0..num_physical_rows { + if physical_offset > max_offset { + break; + } + let deleted = deletion_vector + .as_ref() + .is_some_and(|dv| dv.contains(physical_offset)); + if deleted { + continue; + } + match ids.next() { + Some(id) => { + if offsets.contains(physical_offset) { + row_ids.insert(id); + } + } + None => break, + } + } + } + } + } + Ok(row_ids) } /// An execution node that performs a scalar index search @@ -103,7 +210,7 @@ impl ScalarIndexExec { )); Self { dataset, - expr, + expr: expr.optimize(), properties, metrics: ExecutionPlanMetricsSet::new(), result_format, @@ -849,13 +956,15 @@ mod tests { use crate::index::DatasetIndexExt; use arrow::datatypes::UInt64Type; + use arrow::record_batch::RecordBatchIterator; + use arrow_array::{ArrayRef, Int32Array, RecordBatch}; use arrow_schema::Schema; use datafusion::{ execution::TaskContext, physical_plan::ExecutionPlan, prelude::SessionConfig, scalar::ScalarValue, }; use futures::TryStreamExt; - use lance_core::utils::tempfile::TempStrDir; + use lance_core::utils::{address::RowAddress, tempfile::TempStrDir}; use lance_datagen::gen_batch; use lance_index::{ IndexType, @@ -864,10 +973,11 @@ mod tests { expression::{ScalarIndexExpr, ScalarIndexSearch}, }, }; - use lance_select::result::IndexExprResultWireFormat; + use lance_select::{RowAddrTreeMap, result::IndexExprResultWireFormat}; use crate::{ Dataset, + dataset::WriteParams, io::exec::scalar_index::MaterializeIndexExec, utils::test::{DatagenExt, FragmentCount, FragmentRowCount, NoContextTestFixture}, }; @@ -928,7 +1038,6 @@ mod tests { needs_recheck: false, fragment_bitmap: None, }); - let fragments = dataset.fragments().clone(); let plan = MaterializeIndexExec::new(dataset, query, fragments); @@ -949,6 +1058,51 @@ mod tests { assert_eq!(batches[0].num_rows(), 5); } + #[tokio::test] + async fn test_translate_addr_treemap_to_stable_row_ids() { + let test_dir = TempStrDir::default(); + let batch = RecordBatch::try_from_iter(vec![( + "id", + Arc::new(Int32Array::from((0..10).collect::>())) as ArrayRef, + )]) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let write_params = WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: 5, + ..Default::default() + }; + let dataset = Dataset::write(reader, test_dir.as_str(), Some(write_params)) + .await + .unwrap(); + let fragment_id = dataset.get_fragments()[1].id() as u32; + + let mut full_fragment = RowAddrTreeMap::new(); + full_fragment.insert_fragment(fragment_id); + let translated = super::translate_addr_treemap_to_row_ids(&dataset, &full_fragment) + .await + .unwrap(); + let row_ids = translated + .get_fragment_bitmap(0) + .unwrap() + .iter() + .collect::>(); + assert_eq!(row_ids, vec![5, 6, 7, 8, 9]); + + let mut partial_fragment = RowAddrTreeMap::new(); + partial_fragment.insert(RowAddress::new_from_parts(fragment_id, 1).into()); + partial_fragment.insert(RowAddress::new_from_parts(fragment_id, 3).into()); + let translated = super::translate_addr_treemap_to_row_ids(&dataset, &partial_fragment) + .await + .unwrap(); + let row_ids = translated + .get_fragment_bitmap(0) + .unwrap() + .iter() + .collect::>(); + assert_eq!(row_ids, vec![6, 8]); + } + /// `ScalarIndexExec::schema()` (and the stream it emits) must advertise /// the same schema the batch actually carries — otherwise downstream /// consumers that trust `ExecutionPlan::schema()` will see a different diff --git a/rust/lance/src/io/exec/scan.rs b/rust/lance/src/io/exec/scan.rs index 3ec63ce04cc..9065fb04b3d 100644 --- a/rust/lance/src/io/exec/scan.rs +++ b/rust/lance/src/io/exec/scan.rs @@ -26,7 +26,10 @@ use futures::{StreamExt, TryStreamExt}; use lance_arrow::SchemaExt; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_core::utils::tracing::StreamTracingExt; -use lance_core::{Error, ROW_ADDR_FIELD, ROW_ID_FIELD}; +use lance_core::{ + Error, ROW_ADDR_FIELD, ROW_CREATED_AT_VERSION_FIELD, ROW_ID_FIELD, + ROW_LAST_UPDATED_AT_VERSION_FIELD, +}; use lance_file::reader::FileReaderOptions; use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; use lance_table::format::Fragment; @@ -192,7 +195,36 @@ impl LanceStream { ) -> Result { let scan_metrics = ScanMetrics::new(metrics, partition); let timer = scan_metrics.baseline_metrics.elapsed_compute().timer(); - let project_schema = projection.clone(); + let materialize_blob_v2_binary = + crate::dataset::blob::schema_has_blob_v2_binary_view(projection.as_ref()); + let read_projection = if materialize_blob_v2_binary { + Arc::new(crate::dataset::blob::blob_v2_descriptor_schema( + projection.as_ref(), + )) + } else { + projection.clone() + }; + let project_schema = read_projection; + let output_projection = if materialize_blob_v2_binary { + let mut output_projection = projection.as_ref().clone(); + let mut system_fields = Vec::with_capacity(4); + if config.with_row_id { + system_fields.push(ROW_ID_FIELD.clone()); + } + if config.with_row_address { + system_fields.push(ROW_ADDR_FIELD.clone()); + } + if config.with_row_last_updated_at_version { + system_fields.push(ROW_LAST_UPDATED_AT_VERSION_FIELD.clone()); + } + if config.with_row_created_at_version { + system_fields.push(ROW_CREATED_AT_VERSION_FIELD.clone()); + } + output_projection.extend(&system_fields)?; + Arc::new(output_projection) + } else { + projection.clone() + }; let io_parallelism = dataset.object_store.io_parallelism(); // First, use the value specified by the user in the call // Second, use the default from the environment variable, if specified @@ -275,12 +307,14 @@ impl LanceStream { let scan_scheduler_clone = scan_scheduler.clone(); + let materialize_dataset = dataset; let config_for_stream = config.clone(); let batches = stream::iter(file_fragments.into_iter().enumerate()) .map(move |(priority, file_fragment)| { let project_schema = project_schema.clone(); let scan_scheduler = scan_scheduler.clone(); let config = config_for_stream.clone(); + let force_row_address = materialize_blob_v2_binary; #[allow(clippy::type_complexity)] let frag_task: BoxFuture< Result>>>>, @@ -288,7 +322,7 @@ impl LanceStream { (async move { let mut frag_config = FragReadConfig::default() .with_row_id(config.with_row_id) - .with_row_address(config.with_row_address) + .with_row_address(config.with_row_address || force_row_address) .with_row_last_updated_at_version( config.with_row_last_updated_at_version, ) @@ -349,6 +383,25 @@ impl LanceStream { ) .stream_in_current_span() .boxed(); + let inner_stream = if materialize_blob_v2_binary { + inner_stream + .and_then(move |batch| { + let dataset = materialize_dataset.clone(); + let output_projection = output_projection.clone(); + async move { + crate::dataset::blob::materialize_blob_v2_binary_batch( + &dataset, + output_projection.as_ref(), + batch, + ) + .await + .map_err(DataFusionError::from) + } + }) + .boxed() + } else { + inner_stream + }; timer.done(); Ok(Self { @@ -482,7 +535,9 @@ impl core::fmt::Debug for LanceStream { impl RecordBatchStream for LanceStream { fn schema(&self) -> SchemaRef { - let mut schema: ArrowSchema = self.projection.as_ref().into(); + let output_projection = + crate::dataset::blob::public_blob_v2_binary_output_schema(self.projection.as_ref()); + let mut schema: ArrowSchema = (&output_projection).into(); if self.config.with_row_id { schema = schema.try_with_column(ROW_ID_FIELD.clone()).unwrap(); } @@ -602,7 +657,9 @@ impl LanceScanExec { projection: Arc, config: LanceScanConfig, ) -> Self { - let mut output_schema: ArrowSchema = projection.as_ref().into(); + let output_projection = + crate::dataset::blob::public_blob_v2_binary_output_schema(projection.as_ref()); + let mut output_schema: ArrowSchema = (&output_projection).into(); if config.with_row_id { output_schema = output_schema.try_with_column(ROW_ID_FIELD.clone()).unwrap(); diff --git a/rust/lance/src/io/exec/take.rs b/rust/lance/src/io/exec/take.rs index c3642cdb043..6d3add80142 100644 --- a/rust/lance/src/io/exec/take.rs +++ b/rust/lance/src/io/exec/take.rs @@ -70,6 +70,10 @@ struct TakeStream { dataset: Arc, /// The fields to take from the input stream fields_to_take: Arc, + /// The descriptor-view schema used for storage reads when blob payloads + /// must be materialized after take. + read_fields: Arc, + materialize_blob_v2_binary: bool, /// The output schema, needed for us to merge the new columns /// into the input data in the correct order output_schema: SchemaRef, @@ -92,9 +96,20 @@ impl TakeStream { metrics: &ExecutionPlanMetricsSet, partition: usize, ) -> Self { + let materialize_blob_v2_binary = + crate::dataset::blob::schema_has_blob_v2_binary_view(fields_to_take.as_ref()); + let read_fields = if materialize_blob_v2_binary { + Arc::new(crate::dataset::blob::blob_v2_descriptor_schema( + fields_to_take.as_ref(), + )) + } else { + fields_to_take.clone() + }; Self { dataset, fields_to_take, + read_fields, + materialize_blob_v2_binary, output_schema, readers_cache: Arc::new(Mutex::new(HashMap::new())), scan_scheduler, @@ -131,14 +146,12 @@ impl TakeStream { )) })?; - let reader = Arc::new( - fragment - .open( - &self.fields_to_take, - FragReadConfig::default().with_scan_scheduler(self.scan_scheduler.clone()), - ) - .await?, - ); + let mut read_config = + FragReadConfig::default().with_scan_scheduler(self.scan_scheduler.clone()); + if self.materialize_blob_v2_binary { + read_config = read_config.with_row_address(true); + } + let reader = Arc::new(fragment.open(&self.read_fields, read_config).await?); let mut readers = self.readers_cache.lock().unwrap(); readers.insert(fragment_id, reader.clone()); @@ -355,6 +368,15 @@ impl TakeStream { (None, None) => {} } + if self.materialize_blob_v2_binary { + new_data = crate::dataset::blob::materialize_blob_v2_binary_batch( + &self.dataset, + self.fields_to_take.as_ref(), + new_data, + ) + .await?; + } + Ok(batch.merge_with_schema(&new_data, self.output_schema.as_ref())?) } @@ -487,10 +509,10 @@ impl TakeExec { projection ); - let output_schema = Arc::new(Self::calculate_output_schema( - dataset.schema(), - &input.schema(), - &projection, + let output_schema = + Self::calculate_output_schema(dataset.schema(), &input.schema(), &projection); + let output_schema = Arc::new(crate::dataset::blob::public_blob_v2_binary_output_schema( + &output_schema, )); let output_arrow = Arc::new(ArrowSchema::from(output_schema.as_ref())); let properties = Arc::new( diff --git a/rust/lance/src/lib.rs b/rust/lance/src/lib.rs index 729cf2ffbe7..66bc876bd88 100644 --- a/rust/lance/src/lib.rs +++ b/rust/lance/src/lib.rs @@ -81,6 +81,8 @@ pub mod datafusion; pub mod dataset; pub mod index; pub mod io; +#[cfg(feature = "metrics")] +pub mod metrics; pub mod session; pub mod table; pub mod utils; @@ -90,7 +92,11 @@ pub mod pb { include!(concat!(env!("OUT_DIR"), "/lance.pb.rs")); } -pub use blob::{BlobArrayBuilder, BlobFieldOptions, blob_field, blob_field_with_options}; +pub use blob::{ + BlobArrayBuilder, BlobDescriptor, BlobDescriptorArrayBuilder, BlobDescriptorColumn, + BlobFieldOptions, BlobRange, DedicatedBlobWriter, PackedBlobWriter, blob_field, + blob_field_with_options, +}; pub use dataset::Dataset; use lance_index::vector::DIST_COL; diff --git a/rust/lance/src/metrics.md b/rust/lance/src/metrics.md new file mode 100644 index 00000000000..05d9df0c0e3 --- /dev/null +++ b/rust/lance/src/metrics.md @@ -0,0 +1,40 @@ +Lance publishes metrics through the [`metrics`](https://docs.rs/metrics) crate +facade. Install any recorder (Prometheus, OpenTelemetry, etc.) in your +application and Lance will emit into it; when no recorder is installed, emission +is a cheap no-op. Metrics are only emitted when Lance is built with the +`metrics` feature. + +## Object store metrics + +These track I/O against the underlying object store. The `base` label +identifies the store; its cardinality is controlled by the +`LANCE_OBJECT_STORE_METRICS_LABEL` environment variable: + +- `scheme` (default) — the scheme only (`s3`, `gs`, `az`, `file`, `memory`); + low, bounded cardinality. +- `full` — the store's unique prefix (`s3$my-bucket`, `az$container@account` + where Azure's account also matters), so multiple buckets on the same cloud + can be told apart. Cardinality grows with the number of stores accessed. +- `off` — omit the `base` label entirely. + +`operation` is one of `get`, `put`, `put_part`, `head`, `list`, `delete`, +`copy`, `rename`, `complete_multipart`, or `abort_multipart`. + +Request counts are per logical operation: a `list` or `delete` that spans many +objects is one request, matching how backends batch them. + +| Metric | Type | Labels | Description | +|--------|------|--------|-------------| +| `lance_object_store_requests_total` | counter | `operation`, `base` | Object store requests issued. | +| `lance_object_store_request_bytes_total` | counter | `operation`, `base` | Bytes transferred by `get`/`put` requests. A `get` is counted once its response body has been fully read. | +| `lance_object_store_request_duration_seconds` | histogram | `operation`, `base` | Per-request latency, in seconds. For `get` this covers the full body transfer, not just time-to-first-byte. | +| `lance_object_store_errors_total` | counter | `operation`, `base` | Requests that returned an error. | +| `lance_object_store_in_flight_requests` | gauge | `operation`, `base` | Requests currently in flight. | +| `lance_object_store_throttle_total` | counter | `status`, `base` | Throttle responses (HTTP 429 / 503) seen at the HTTP layer, counted per attempt including retries. The `status` label is the numeric HTTP status. | +| `lance_object_store_retryable_responses_total` | counter | `status`, `base` | Retryable responses (HTTP 5xx / 429 / 408) seen at the HTTP layer, counted per attempt including retries. A superset of `throttle_total`; 409 (conflict) is excluded so commit conflicts are not counted. | + +`lance_object_store_throttle_total` and +`lance_object_store_retryable_responses_total` are recorded only for the native +cloud stores (S3, GCS, Azure); Opendal-backed stores bypass the HTTP client +where the counters are installed, so they report the other object store metrics +but not throttle/retryable counts. diff --git a/rust/lance/src/metrics.rs b/rust/lance/src/metrics.rs new file mode 100644 index 00000000000..1d48865e859 --- /dev/null +++ b/rust/lance/src/metrics.rs @@ -0,0 +1,10 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Metrics published by Lance. +#![doc = include_str!("metrics.md")] +//! +//! The metrics themselves are emitted from the relevant subsystems (for +//! example object store I/O is instrumented in +//! [`lance_io::object_store::metrics`]); this module exists to document the +//! full catalogue of metric names, types, and labels in one place. diff --git a/rust/lance/src/session/index_extension.rs b/rust/lance/src/session/index_extension.rs index 301213c6f06..0f4d2ecf310 100644 --- a/rust/lance/src/session/index_extension.rs +++ b/rust/lance/src/session/index_extension.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +#[cfg(test)] +use lance_core::utils::row_addr_remap::RowAddrRemap; use std::sync::Arc; use lance_core::Result; @@ -62,7 +64,6 @@ mod test { use std::{ any::Any, - collections::HashMap, sync::{Arc, atomic::AtomicBool}, }; @@ -182,7 +183,7 @@ mod test { unimplemented!() } - async fn remap(&mut self, _: &HashMap>) -> Result<()> { + async fn remap(&mut self, _: &RowAddrRemap) -> Result<()> { Ok(()) } diff --git a/rust/lance/src/utils/test.rs b/rust/lance/src/utils/test.rs index 3338eee07a8..f804a7cc38a 100644 --- a/rust/lance/src/utils/test.rs +++ b/rust/lance/src/utils/test.rs @@ -243,6 +243,7 @@ impl TestDatasetGenerator { Fragment { id: 0, files, + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(batch.num_rows()), diff --git a/test_data/fri_straddle_pre_6610/datagen/Cargo.lock b/test_data/fri_straddle_pre_6610/datagen/Cargo.lock index 6fcff711a06..5edd16c2fed 100644 --- a/test_data/fri_straddle_pre_6610/datagen/Cargo.lock +++ b/test_data/fri_straddle_pre_6610/datagen/Cargo.lock @@ -1124,9 +1124,9 @@ dependencies = [ [[package]] name = "cmov" -version = "0.5.3" +version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f88a43d011fc4a6876cb7344703e297c71dda42494fee094d5f7c76bf13f746" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" [[package]] name = "colorchoice"